test(orchestration): add full mission loop e2e

agent/tests/test_mission_loop.py

@@ -0,0 +1,297 @@
+"""End-to-end coverage for the explicit Planner mission loop."""
+
+import json
+
+import pytest
+
+from agent.core.runner import AgentRunner, RunConfig
+from agent.orchestration.models import (
+    AgentRole,
+    CompletionPolicy,
+    TaskStatus,
+    ValidationVerdict,
+)
+from agent.orchestration.store import FileSystemArtifactStore, FileSystemTaskStore
+from agent.orchestration.wiring import wire_orchestration
+from agent.tools.builtin.knowledge import KnowledgeConfig
+from agent.trace.store import FileSystemTraceStore
+
+
+ROOT_OBJECTIVE = "complete the generic root"
+CHILD_OBJECTIVE_V1 = "produce the generic unit"
+CHILD_OBJECTIVE_V2 = "produce the revised generic unit"
+
+
+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 _no_tool():
+    return {"content": "", "tool_calls": None, "finish_reason": "stop"}
+
+
+def _latest_json_user_prompt(messages):
+    for message in reversed(messages):
+        if message.get("role") != "user":
+            continue
+        content = message.get("content")
+        if not isinstance(content, str):
+            continue
+        try:
+            value = json.loads(content)
+        except (TypeError, ValueError):
+            continue
+        if isinstance(value, dict) and "task_spec" in value:
+            return value
+    raise AssertionError("role sub-trace did not receive its structured task prompt")
+
+
+class MissionLoopLLM:
+    """Choose actions from role tools and durable Ledger state."""
+
+    def __init__(self, task_store):
+        self.task_store = task_store
+        self.root_trace_id = None
+        self.first_planner_exit_done = False
+        self.inspection_requested = False
+        self.call_number = 0
+        self.planner_histories = []
+        self.root_worker_prompt = None
+
+    def _call(self, name, arguments):
+        self.call_number += 1
+        return _tool_call(f"mission-call-{self.call_number}", name, arguments)
+
+    async def __call__(self, messages, tools, **_kwargs):
+        names = {tool["function"]["name"] for tool in tools or []}
+        if "submit_attempt" in names:
+            return self._worker(messages)
+        if "submit_validation" in names:
+            return self._validator(messages)
+        if "task_plan" in names:
+            return await self._planner(messages)
+        raise AssertionError(f"unexpected role tool set: {sorted(names)}")
+
+    def _worker(self, messages):
+        prompt = _latest_json_user_prompt(messages)
+        spec = prompt["task_spec"]
+        if spec["objective"] == ROOT_OBJECTIVE:
+            self.root_worker_prompt = prompt
+        return self._call("submit_attempt", {
+            "summary": f"submitted: {spec['objective']}",
+            "artifact_refs": [{
+                "uri": f"memory://result/{spec['version']}",
+                "version": str(spec["version"]),
+            }],
+            "evidence_refs": [],
+        })
+
+    def _validator(self, messages):
+        prompt = _latest_json_user_prompt(messages)
+        spec = prompt["task_spec"]
+        failed_first_child = (
+            spec["objective"] == CHILD_OBJECTIVE_V1 and spec["version"] == 1
+        )
+        verdict = "failed" if failed_first_child else "passed"
+        return self._call("submit_validation", {
+            "verdict": verdict,
+            "criterion_results": [{
+                "criterion_id": criterion["criterion_id"],
+                "verdict": verdict,
+                "reason": "checked against the immutable snapshot",
+            } for criterion in spec["acceptance_criteria"]],
+            "summary": f"validation {verdict}",
+            "evidence_refs": [],
+            "unverified_claims": [],
+            "risks": [] if verdict == "passed" else ["task spec needs revision"],
+            "recommendation": "accept" if verdict == "passed" else "replan",
+        })
+
+    async def _planner(self, messages):
+        self.planner_histories.append(list(messages))
+
+        # Exercise the Runner guard before the Planner performs any operation.
+        if not self.first_planner_exit_done:
+            self.first_planner_exit_done = True
+            return _no_tool()
+
+        ledger = await self.task_store.load(self.root_trace_id)
+        root = ledger.tasks[ledger.root_task_id]
+
+        if not self.inspection_requested:
+            self.inspection_requested = True
+            return self._call("task_plan", {"operation": "inspect"})
+
+        children = [
+            task for task in ledger.tasks.values()
+            if task.parent_task_id == root.task_id
+        ]
+        if not children:
+            return self._call("task_plan", {
+                "operation": "create",
+                "parent_task_id": root.task_id,
+                "tasks": [{
+                    "objective": CHILD_OBJECTIVE_V1,
+                    "acceptance_criteria": [{
+                        "criterion_id": "unit-ready",
+                        "description": "the generic unit is ready",
+                        "hard": True,
+                    }],
+                    "context_refs": ["memory://input/generic"],
+                }],
+            })
+
+        assert len(children) == 1
+        child = children[0]
+        if child.status in (TaskStatus.PENDING, TaskStatus.NEEDS_REPLAN):
+            return self._call("dispatch_tasks", {"task_ids": [child.task_id]})
+
+        if child.status == TaskStatus.AWAITING_DECISION:
+            validation = ledger.validations[child.validation_ids[-1]]
+            if validation.verdict == ValidationVerdict.FAILED:
+                return self._call("task_decide", {
+                    "task_id": child.task_id,
+                    "validation_id": validation.validation_id,
+                    "action": "revise",
+                    "reason": "clarify the immutable task specification",
+                    "payload": {
+                        "objective": CHILD_OBJECTIVE_V2,
+                        "acceptance_criteria": [{
+                            "criterion_id": "unit-ready",
+                            "description": "the revised generic unit is ready",
+                            "hard": True,
+                        }],
+                        "context_refs": ["memory://input/generic"],
+                    },
+                })
+            assert validation.verdict == ValidationVerdict.PASSED
+            return self._call("task_decide", {
+                "task_id": child.task_id,
+                "validation_id": validation.validation_id,
+                "action": "accept",
+                "reason": "the current child version passed validation",
+            })
+
+        assert child.status == TaskStatus.COMPLETED
+        if root.status == TaskStatus.NEEDS_REPLAN:
+            return self._call("dispatch_tasks", {"task_ids": [root.task_id]})
+        if root.status == TaskStatus.AWAITING_DECISION:
+            validation = ledger.validations[root.validation_ids[-1]]
+            assert validation.verdict == ValidationVerdict.PASSED
+            return self._call("task_decide", {
+                "task_id": root.task_id,
+                "validation_id": validation.validation_id,
+                "action": "accept",
+                "reason": "the root result passed independent validation",
+            })
+        assert root.status == TaskStatus.COMPLETED
+        return _no_tool()
+
+
+def _knowledge_off():
+    return KnowledgeConfig(
+        enable_extraction=False,
+        enable_completion_extraction=False,
+        enable_injection=False,
+    )
+
+
+@pytest.mark.asyncio
+async def test_real_runner_completes_replanned_root_mission(tmp_path):
+    root_trace_id = "mission-loop-root"
+    trace_store = FileSystemTraceStore(str(tmp_path))
+    task_store = FileSystemTaskStore(str(tmp_path))
+    fake_llm = MissionLoopLLM(task_store)
+    fake_llm.root_trace_id = root_trace_id
+    runner = AgentRunner(trace_store=trace_store, llm_call=fake_llm)
+    wire_orchestration(
+        runner,
+        task_store,
+        FileSystemArtifactStore(str(tmp_path)),
+    )
+
+    result = await runner.run_result(
+        [{"role": "user", "content": "run the generic mission"}],
+        RunConfig(
+            agent_type="planner",
+            completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
+            max_iterations=30,
+            new_trace_id=root_trace_id,
+            tools=["task_plan", "dispatch_tasks", "task_decide"],
+            tool_groups=[],
+            enable_memory=False,
+            enable_research_flow=False,
+            knowledge=_knowledge_off(),
+            root_task_spec={
+                "objective": ROOT_OBJECTIVE,
+                "acceptance_criteria": [{
+                    "criterion_id": "root-ready",
+                    "description": "the root result is ready",
+                    "hard": True,
+                }],
+                "context_refs": ["memory://input/root"],
+            },
+        ),
+    )
+    assert result["trace_id"] == root_trace_id
+    ledger = await task_store.load(result["trace_id"])
+    root = ledger.tasks[ledger.root_task_id]
+    children = [task for task in ledger.tasks.values() if task.parent_task_id == root.task_id]
+    assert result["status"] == "completed", {
+        "result": result,
+        "root_status": root.status,
+        "child_statuses": [task.status for task in children],
+        "planner_calls": len(fake_llm.planner_histories),
+    }
+    assert result["error"] is None
+
+    assert len([task for task in ledger.tasks.values() if task.parent_task_id is None]) == 1
+    assert len(children) == 1
+    child = children[0]
+    assert child.current_spec_version == 2
+    assert child.status == TaskStatus.COMPLETED
+    assert root.status == TaskStatus.COMPLETED
+    assert len(child.attempt_ids) == 2
+    assert len(root.attempt_ids) == 1
+
+    assert fake_llm.root_worker_prompt is not None
+    accepted = fake_llm.root_worker_prompt["accepted_child_results"]
+    assert [item["task_id"] for item in accepted] == [child.task_id]
+    assert accepted[0]["task_spec"]["version"] == 2
+    assert accepted[0]["submission"]["summary"] == (
+        f"submitted: {CHILD_OBJECTIVE_V2}"
+    )
+
+    planner_history = json.dumps(fake_llm.planner_histories, ensure_ascii=False, default=str)
+    assert len(ledger.validations) == 3
+    assert all(
+        validation.validation_id in planner_history
+        for validation in ledger.validations.values()
+    )
+    planner_messages = await trace_store.get_trace_messages(result["trace_id"])
+    assert any(
+        "[FRAMEWORK_MISSION_INCOMPLETE]" in str(message.content)
+        for message in planner_messages
+    )
+
+    traces = await trace_store.list_traces(limit=50)
+    planner_traces = [trace for trace in traces if trace.agent_role == AgentRole.PLANNER.value]
+    worker_traces = [trace for trace in traces if trace.agent_role == AgentRole.WORKER.value]
+    validator_traces = [trace for trace in traces if trace.agent_role == AgentRole.VALIDATOR.value]
+    assert [trace.trace_id for trace in planner_traces] == [result["trace_id"]]
+    assert len(worker_traces) == 3
+    assert len(validator_traces) == 3
+    assert len({trace.trace_id for trace in worker_traces + validator_traces}) == 6
+    assert all(trace.parent_trace_id == result["trace_id"] for trace in worker_traces + validator_traces)
+    assert all(trace.status == "completed" for trace in worker_traces + validator_traces)