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()