Agent/tests/test_application_runtime_integration.py

import asyncio
import hashlib
import json
import tempfile
import unittest
from unittest.mock import patch

from fastapi import HTTPException
from pydantic import ValidationError

from cyber_agent.application import (
    AgentApplication,
    AgentRole,
    ApplicationRegistry,
    ApplicationRuntime,
    ApplicationServices,
    ContextMaterial,
    ContextResourceDescriptor,
    ProviderRef,
    RoleRunLimits,
    RunLimits,
    ToolSet,
    ToolSpec,
)
from cyber_agent.application.models import canonical_json
from cyber_agent.application.bootstrap import (
    APPLICATION_FACTORIES_ENV,
    build_runtime_from_environment,
)
from cyber_agent.core.context_policy import (
    ContextPolicyError,
    build_child_context_access,
    normalize_root_task_anchor,
)
from cyber_agent.core.run_snapshot import (
    RunConfigSnapshotV2,
    load_run_config_snapshot,
    persist_run_config_snapshot,
)
from cyber_agent.core.resource_budget import ResourceBudgetController
from cyber_agent.core.task_protocol import ContextRef, TaskBrief, new_task_protocol
from cyber_agent.tools import get_tool_registry
from cyber_agent.tools.builtin.knowledge import KnowledgeConfig
from cyber_agent.tools.registry import ToolRegistry
from cyber_agent.trace.models import Trace
from cyber_agent.trace.run_api import (
    CreateRequest,
    _restore_runner_and_config,
    create_and_run,
    set_application_runtime,
)
from cyber_agent.trace.store import FileSystemTraceStore


ROOT_ANCHOR = {
    "objective": "Produce an evidence-grounded article",
    "completion_criteria": ["Every claim is traceable"],
    "constraints": ["Do not invent evidence"],
}


class CountingContextProvider:
    def __init__(self, *, count=1, bad_hash=False):
        self.count = count
        self.bad_hash = bad_hash
        self.list_calls = 0
        self.resolve_calls = 0

    async def list_context(self, request):
        self.list_calls += 1
        descriptors = []
        for index in range(self.count):
            content = {"source": index, "fact": f"verified-{index}"}
            digest = hashlib.sha256(
                canonical_json(content).encode("utf-8")
            ).hexdigest()
            if self.bad_hash and index == 0:
                digest = "0" * 64
            descriptors.append(ContextResourceDescriptor(
                application_ref=request.application_ref,
                root_trace_id=request.root_trace_id,
                trace_id=request.trace_id,
                uid=request.uid,
                role_id=request.role_id,
                source_id=f"source-{index:02d}",
                source_version="1",
                content_hash=digest,
                kind="evidence",
                summary=f"Evidence {index}",
                inheritance="task_only" if index == 0 else "inheritable",
                priority=100 if index == 0 else index,
                estimated_chars=len(canonical_json(content)),
            ))
        return descriptors

    async def resolve_context(self, request, descriptor):
        del request
        self.resolve_calls += 1
        index = int(descriptor.source_id.rsplit("-", 1)[1])
        return ContextMaterial(
            descriptor=descriptor,
            content={"source": index, "fact": f"verified-{index}"},
        )


async def unused_llm(**kwargs):
    raise AssertionError(f"LLM must not be called in this test: {kwargs}")


def declaration(*, prompt="Bound role prompt", max_iterations=12):
    return AgentApplication(
        application_id="creative.reference",
        application_version="0.1.0",
        root_role="writer",
        roles=(AgentRole(
            role_id="writer",
            model="bound-model",
            temperature=0.2,
            system_prompt=prompt,
            allowed_child_roles=("writer",),
            run_limits=RoleRunLimits(max_iterations=max_iterations),
        ),),
        context_provider_ref=ProviderRef(
            provider_id="reference-context",
            provider_version="1",
        ),
        run_limits=RunLimits(max_iterations=20),
    )


class ApplicationRuntimeIntegrationTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp.name)
        self.provider = CountingContextProvider()
        self.registry = ApplicationRegistry()
        self.binding = self.registry.register(
            declaration(),
            ApplicationServices(
                tool_registry=ToolRegistry(),
                context_provider=self.provider,
            ),
        )
        self.runtime = ApplicationRuntime(
            registry=self.registry,
            trace_store=self.store,
            llm_call=unused_llm,
        )
        set_application_runtime(self.runtime)

    async def asyncTearDown(self):
        set_application_runtime(None)
        self.temp.cleanup()

    async def _create_trace(self, *, max_iterations=None):
        runner, config = self.runtime.new_run(
            "creative.reference",
            "0.1.0",
            uid="user-1",
            root_task_anchor=ROOT_ANCHOR,
            max_iterations=max_iterations,
        )
        with patch.dict("os.environ", {"AGENT_MODE": "recursive"}, clear=False):
            trace, _goal, _sequence = await runner._prepare_new_trace(
                [{"role": "user", "content": "Start"}],
                config,
            )
        return trace, runner, config

    async def test_new_run_persists_v2_binding_and_resume_does_not_reload_context(self):
        trace, _runner, config = await self._create_trace(max_iterations=7)
        snapshot = load_run_config_snapshot(trace.context)

        self.assertIsInstance(snapshot, RunConfigSnapshotV2)
        self.assertEqual(self.binding.application_ref.model_dump(), snapshot.application_ref)
        self.assertEqual("writer", snapshot.role_id)
        self.assertEqual(7, snapshot.max_iterations)
        self.assertEqual(7, snapshot.effective_run_limits["max_iterations"])
        self.assertEqual(1, self.provider.list_calls)
        self.assertEqual(1, self.provider.resolve_calls)
        snapshots = trace.context["context_access"]["snapshots"]
        self.assertEqual(1, len(snapshots))
        self.assertEqual(
            "application_resource",
            next(iter(snapshots.values()))["kind"],
        )
        self.assertEqual("Bound role prompt", config.system_prompt)

        restored_runner, restored_config = await self.runtime.restore(trace.trace_id)
        self.assertIsNot(restored_runner, _runner)
        self.assertEqual("Bound role prompt", restored_config.system_prompt)
        self.assertEqual(1, self.provider.list_calls)
        self.assertEqual(1, self.provider.resolve_calls)

    async def test_application_http_handler_uses_the_bound_runtime(self):
        with patch.dict("os.environ", {"AGENT_MODE": "recursive"}, clear=False):
            response = await create_and_run(CreateRequest(
                messages=[{"role": "user", "content": "Start through HTTP"}],
                application_id="creative.reference",
                application_version="0.1.0",
                uid="http-user",
                root_task_anchor=ROOT_ANCHOR,
            ))
        trace = await self.store.get_trace(response.trace_id)
        self.assertIsNotNone(trace)
        snapshot = load_run_config_snapshot(trace.context)
        self.assertIsInstance(snapshot, RunConfigSnapshotV2)
        self.assertEqual(
            self.binding.application_ref.model_dump(mode="json"),
            snapshot.application_ref,
        )
        await asyncio.sleep(0)

    async def test_restore_rejects_registry_hash_and_trace_binding_tampering(self):
        trace, base_runner, _config = await self._create_trace()
        changed_registry = ApplicationRegistry()
        changed_registry.register(
            declaration(prompt="Changed role prompt"),
            ApplicationServices(
                tool_registry=ToolRegistry(),
                context_provider=CountingContextProvider(),
            ),
        )
        changed_runtime = ApplicationRuntime(
            registry=changed_registry,
            trace_store=self.store,
            llm_call=unused_llm,
        )
        with self.assertRaisesRegex(ValueError, "config hash"):
            await changed_runtime.restore(trace.trace_id)

        persisted = await self.store.get_trace(trace.trace_id)
        persisted.context["application_role_hash"] = "f" * 64
        await self.store.update_trace(trace.trace_id, context=persisted.context)
        with self.assertRaisesRegex(ValueError, "Trace context"):
            await self.runtime.restore(trace.trace_id)

        with self.assertRaises(HTTPException) as raised:
            await _restore_runner_and_config(
                trace=await self.store.get_trace(trace.trace_id),
                base_runner=base_runner,
            )
        self.assertEqual(409, raised.exception.status_code)

    async def test_runtime_dependency_failure_happens_before_trace_creation(self):
        calls = 0

        def unavailable():
            nonlocal calls
            calls += 1
            raise RuntimeError("credential unavailable")

        registry = ApplicationRegistry()
        registry.register(
            declaration(),
            ApplicationServices(
                tool_registry=ToolRegistry(),
                context_provider=CountingContextProvider(),
                runtime_validator=unavailable,
            ),
        )
        runtime = ApplicationRuntime(
            registry=registry,
            trace_store=self.store,
            llm_call=unused_llm,
        )
        with self.assertRaisesRegex(RuntimeError, "credential unavailable"):
            runtime.new_run("creative.reference", "0.1.0")
        self.assertEqual(1, calls)
        self.assertEqual([], await self.store.list_traces())

    async def test_context_is_stably_capped_and_task_only_is_not_delegated(self):
        provider = CountingContextProvider(count=10)
        registry = ApplicationRegistry()
        binding = registry.register(
            declaration(),
            ApplicationServices(
                tool_registry=ToolRegistry(),
                context_provider=provider,
            ),
        )
        runtime = ApplicationRuntime(
            registry=registry,
            trace_store=self.store,
            llm_call=unused_llm,
        )
        self.runtime = runtime
        self.registry = registry
        self.binding = binding
        self.provider = provider
        set_application_runtime(runtime)
        trace, _runner, _config = await self._create_trace()

        snapshots = list(trace.context["context_access"]["snapshots"].values())
        self.assertEqual(8, len(snapshots))
        selected_sources = [
            item["content"]["application_resource"]["source_id"]
            for item in snapshots
        ]
        self.assertEqual(
            ["source-00", *[f"source-{index:02d}" for index in range(9, 2, -1)]],
            selected_sources,
        )

        refs = [
            ContextRef(ref_id=item["ref_id"], version=item["version"])
            for item in snapshots
        ]
        child_brief = TaskBrief(
            objective="Write the checked section",
            reason="The root needs this section",
            completion_criteria=["Use the supplied evidence"],
            expected_outputs=["One section"],
            context_refs=refs,
        )
        parent_state = new_task_protocol()
        with self.assertRaisesRegex(ContextPolicyError, "task_only"):
            build_child_context_access(
                parent_context=trace.context,
                parent_trace_id=trace.trace_id,
                root_trace_id=trace.trace_id,
                uid="user-1",
                parent_task_state=parent_state,
                child_task_brief=child_brief,
                root_task_anchor=normalize_root_task_anchor(ROOT_ANCHOR),
            )
        child_brief = child_brief.model_copy(update={
            "context_refs": refs[1:],
        })
        child_access = build_child_context_access(
            parent_context=trace.context,
            parent_trace_id=trace.trace_id,
            root_trace_id=trace.trace_id,
            uid="user-1",
            parent_task_state=parent_state,
            child_task_brief=child_brief,
            root_task_anchor=normalize_root_task_anchor(ROOT_ANCHOR),
        )
        delegated = list(child_access["snapshots"].values())
        self.assertEqual(7, len(delegated))
        self.assertTrue(all(
            item["content"]["application_resource"]["inheritance"]
            == "inheritable"
            for item in delegated
            if item["kind"] == "application_resource"
        ))

    async def test_context_hash_mismatch_leaves_no_trace(self):
        provider = CountingContextProvider(bad_hash=True)
        registry = ApplicationRegistry()
        registry.register(
            declaration(),
            ApplicationServices(
                tool_registry=ToolRegistry(),
                context_provider=provider,
            ),
        )
        runtime = ApplicationRuntime(
            registry=registry,
            trace_store=self.store,
            llm_call=unused_llm,
        )
        runner, config = runtime.new_run(
            "creative.reference",
            "0.1.0",
            root_task_anchor=ROOT_ANCHOR,
        )
        with patch.dict("os.environ", {"AGENT_MODE": "recursive"}, clear=False):
            with self.assertRaisesRegex(ValueError, "hash mismatch"):
                await runner._prepare_new_trace([], config)
        self.assertEqual([], await self.store.list_traces())


class ApplicationCreateRequestTest(unittest.TestCase):
    def test_application_pair_and_override_rules(self):
        with self.assertRaisesRegex(ValidationError, "provided together"):
            CreateRequest(
                messages=[{"role": "user", "content": "start"}],
                application_id="creative.reference",
            )
        with self.assertRaisesRegex(ValidationError, "mutually exclusive"):
            CreateRequest(
                messages=[{"role": "user", "content": "start"}],
                application_id="creative.reference",
                application_version="0.1.0",
                project_name="old_example",
            )
        for field, value in (
            ("model", "override"),
            ("temperature", 1.0),
            ("tools", []),
        ):
            with self.subTest(field=field):
                with self.assertRaisesRegex(ValidationError, "cannot override"):
                    CreateRequest(
                        messages=[{"role": "user", "content": "start"}],
                        application_id="creative.reference",
                        application_version="0.1.0",
                        **{field: value},
                    )
        with self.assertRaisesRegex(ValidationError, "system messages"):
            CreateRequest(
                messages=[{"role": "system", "content": "override"}],
                application_id="creative.reference",
                application_version="0.1.0",
            )

    def test_deployment_factory_builds_the_http_runtime_registry(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            store = FileSystemTraceStore(temp_dir)
            runtime = build_runtime_from_environment(
                trace_store=store,
                llm_call=unused_llm,
                environ={
                    APPLICATION_FACTORIES_ENV: (
                        "examples.application_reference.application:"
                        "build_reference_components"
                    ),
                },
            )
            self.assertIsNotNone(runtime)
            binding = runtime.registry.resolve("application_reference", "1")
            self.assertEqual("editor", binding.application.root_role)
            self.assertIsNone(build_runtime_from_environment(
                trace_store=store,
                llm_call=unused_llm,
                environ={},
            ))


def _tool_names(schemas):
    return {item["function"]["name"] for item in schemas or []}


def _has_tool_result(messages, tool_name):
    return any(
        item.get("role") == "tool" and item.get("name") == tool_name
        for item in messages
    )


def _tool_call(call_id, name, arguments):
    return {
        "content": "",
        "tool_calls": [{
            "id": call_id,
            "type": "function",
            "function": {
                "name": name,
                "arguments": json.dumps(arguments),
            },
        }],
        "finish_reason": "tool_calls",
    }


def _ready_progress(call_id):
    return _tool_call(call_id, "update_task_progress", {
        "expected_revision": 1,
        "progress": {
            "phase": "ready_to_submit",
            "questions": [],
            "blockers": [],
            "findings": [],
            "hypotheses": [],
            "work_items": [],
            "decision_rationale": "The application role completed its contract.",
        },
    })


class ApplicationChildPropagationTest(unittest.IsolatedAsyncioTestCase):
    async def test_child_inherits_application_but_uses_target_role_binding(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            store = FileSystemTraceStore(temp_dir)
            declarations = tuple(
                ToolSpec(tool_name=name, implementation_version="1")
                for name in (
                    "agent",
                    "update_task_progress",
                    "submit_task_report",
                    "review_task_result",
                )
            )
            application = AgentApplication(
                application_id="creative.roles",
                application_version="0.1.0",
                root_role="director",
                roles=(
                    AgentRole(
                        role_id="director",
                        model="root-model",
                        system_prompt="ROOT_APPLICATION_PROMPT",
                        tool_sets=("root-tools",),
                        allowed_child_roles=("checker",),
                    ),
                    AgentRole(
                        role_id="checker",
                        model="child-model",
                        system_prompt="CHILD_APPLICATION_PROMPT",
                        tool_sets=("child-tools",),
                    ),
                ),
                tool_sets=(
                    ToolSet(
                        tool_set_id="root-tools",
                        tools=tuple(
                            item
                            for item in declarations
                            if item.tool_name in {
                                "agent",
                                "update_task_progress",
                                "review_task_result",
                            }
                        ),
                    ),
                    ToolSet(
                        tool_set_id="child-tools",
                        tools=tuple(
                            item
                            for item in declarations
                            if item.tool_name in {
                                "update_task_progress",
                                "submit_task_report",
                            }
                        ),
                    ),
                ),
                run_limits=RunLimits(max_iterations=20),
            )
            registry = ApplicationRegistry()
            binding = registry.register(
                application,
                ApplicationServices(tool_registry=get_tool_registry()),
            )
            prompts_by_model = {}
            delegated = False
            reviewed = False

            async def llm_call(**kwargs):
                nonlocal delegated, reviewed
                messages = kwargs["messages"]
                model = kwargs["model"]
                prompts_by_model.setdefault(model, set()).update(
                    str(item.get("content", ""))
                    for item in messages
                    if item.get("role") == "system"
                )
                if any(
                    "recursive_validation_protocol" in str(item.get("content", ""))
                    for item in messages
                    if item.get("role") == "system"
                ):
                    packet = next(
                        json.loads(item["content"])
                        for item in messages
                        if item.get("role") == "user"
                        and "recursive_validation_protocol" in str(item.get("content", ""))
                    )
                    scope = packet["validation_scope"]
                    return {
                        "content": json.dumps({
                            "outcome": "passed",
                            "scope": scope,
                            "checks": [
                                {
                                    "check_id": check["check_id"],
                                    "status": "passed",
                                    "evidence_refs": [],
                                    "issue": None,
                                }
                                for check in packet["validation_plan"]["checks"]
                                if check["scope"] == scope
                            ],
                            "reason": "bound application result is valid",
                            "retry_from": None,
                        }),
                        "tool_calls": [],
                        "finish_reason": "stop",
                    }
                names = _tool_names(kwargs.get("tools"))
                if model == "child-model":
                    if not _has_tool_result(messages, "update_task_progress"):
                        return _ready_progress("child-progress")
                    if not _has_tool_result(messages, "submit_task_report"):
                        return _tool_call("child-report", "submit_task_report", {
                            "task_report": {
                                "summary": "The fact is checked.",
                                "outcome": "satisfied",
                                "validation": {
                                    "hard_passed": True,
                                    "open_issues": [],
                                },
                                "next_step_suggestion": {
                                    "direction": "NONE",
                                    "reason": "The local task is complete.",
                                },
                                "outputs": [{"kind": "fact", "value": "checked"}],
                                "evidence": [],
                                "remaining_issues": [],
                            },
                        })
                    return {"content": "child submitted", "tool_calls": []}
                if names == {"review_task_result"}:
                    root = next(
                        item
                        for item in await store.list_traces(limit=20)
                        if item.parent_trace_id is None
                    )
                    child_id = next(iter(
                        root.context["task_protocol"]["pending_reviews"]
                    ))
                    reviewed = True
                    return _tool_call("review-child", "review_task_result", {
                        "child_trace_id": child_id,
                        "decision": "ASCEND",
                        "reason": "The checked result satisfies the delegated task.",
                    })
                if not delegated:
                    delegated = True
                    return _tool_call("delegate", "agent", {
                        "agent_type": "checker",
                        "task_brief": {
                            "objective": "Check one fact",
                            "reason": "The article needs verified evidence",
                            "completion_criteria": ["Return one checked fact"],
                            "expected_outputs": ["Checked fact"],
                        },
                    })
                if reviewed and not _has_tool_result(messages, "update_task_progress"):
                    return _ready_progress("root-progress")
                return {"content": "root final", "tool_calls": []}

            runtime = ApplicationRuntime(
                registry=registry,
                trace_store=store,
                llm_call=llm_call,
            )
            runner, config = runtime.new_run(
                "creative.roles",
                "0.1.0",
                uid="user-1",
                root_task_anchor=ROOT_ANCHOR,
            )
            config.knowledge = KnowledgeConfig(
                enable_extraction=False,
                enable_completion_extraction=False,
                enable_injection=False,
            )
            with patch.dict("os.environ", {"AGENT_MODE": "recursive"}, clear=False):
                result = await runner.run_result(
                    [{"role": "user", "content": "Create the checked article"}],
                    config,
                )

            self.assertEqual("completed", result["status"])
            self.assertTrue(reviewed)
            root = await store.get_trace(result["trace_id"])
            children = await store.list_traces(
                parent_trace_id=root.trace_id,
                created_by_tool="agent",
                limit=10,
            )
            self.assertEqual(1, len(children))
            child = children[0]
            root_snapshot = load_run_config_snapshot(root.context)
            child_snapshot = load_run_config_snapshot(child.context)
            self.assertIsInstance(child_snapshot, RunConfigSnapshotV2)
            self.assertEqual(root_snapshot.application_ref, child_snapshot.application_ref)
            self.assertEqual("director", root_snapshot.role_id)
            self.assertEqual("checker", child_snapshot.role_id)
            self.assertEqual(
                binding.role("checker").role_hash,
                child_snapshot.role_hash,
            )
            self.assertEqual("child-model", child.model)
            self.assertEqual(
                {"update_task_progress", "submit_task_report"},
                set(child_snapshot.tools),
            )
            self.assertTrue(any(
                "ROOT_APPLICATION_PROMPT" in prompt
                for prompt in prompts_by_model["root-model"]
            ))
            self.assertTrue(any(
                "CHILD_APPLICATION_PROMPT" in prompt
                for prompt in prompts_by_model["child-model"]
            ))
            usage = await ResourceBudgetController(store).get_usage(root.trace_id)
            self.assertEqual(2, usage.total_agents)


if __name__ == "__main__":
    unittest.main()