script-build-and-framework-tmp/agent/tests/test_mission_loop.py

"""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"
ROOT_OBJECTIVE_V2 = "complete the revised 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"] in (ROOT_OBJECTIVE, ROOT_OBJECTIVE_V2):
            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 in (TaskStatus.PENDING, TaskStatus.NEEDS_REPLAN):
            if root.current_spec_version == 1:
                return self._call("task_decide", {
                    "task_id": root.task_id,
                    "action": "revise",
                    "reason": "clarify the root objective after child acceptance",
                    "payload": {
                        "objective": ROOT_OBJECTIVE_V2,
                        "acceptance_criteria": [{
                            "criterion_id": "root-ready",
                            "description": "the revised root result is ready",
                            "hard": True,
                        }],
                        "context_refs": ["memory://input/root"],
                    },
                })
            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 root.current_spec_version == 2
    assert root.current_spec.objective == ROOT_OBJECTIVE_V2
    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)