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

import pytest

from agent.core.presets import AgentPreset, register_preset
from agent.core.runner import AgentRunner, RunConfig
from agent.failures import FailureDetail, FailureDisposition, ToolExecutionError
from agent.orchestration.models import AgentRole, CompletionPolicy
from agent.trace.models import Trace
from agent.trace.store import FileSystemTraceStore
from agent.tools.models import ToolCapability
from agent.tools.registry import ToolRegistry


class StubCoordinator:
    def __init__(self, snapshots):
        self.snapshots = list(snapshots)
        self.ensure_calls = []
        self.completion_calls = 0

    async def ensure_ledger(self, trace_id, mission, *, allow_create=True):
        self.ensure_calls.append((trace_id, mission))

    async def mission_completion(self, root_trace_id):
        self.completion_calls += 1
        if len(self.snapshots) > 1:
            return self.snapshots.pop(0)
        return self.snapshots[0]


class RootFirstCoordinator(StubCoordinator):
    def __init__(self):
        super().__init__([{"status": "legacy"}])
        self.root_calls = 0

    async def root_completion(self, root_trace_id):
        self.root_calls += 1
        return {"status": "completed", "root_task_id": "root-task"}


def _knowledge_off():
    from agent.tools.builtin.knowledge import KnowledgeConfig

    return KnowledgeConfig(
        enable_extraction=False,
        enable_completion_extraction=False,
        enable_injection=False,
    )


@pytest.mark.asyncio
async def test_runner_prefers_root_completion_and_keeps_legacy_fallback(tmp_path):
    root_first = RootFirstCoordinator()
    runner = AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path / "root-first")),
        llm_call=None,
        task_coordinator=root_first,
    )
    assert (await runner._get_mission_completion("root"))["status"] == "completed"
    assert root_first.root_calls == 1
    assert root_first.completion_calls == 0

    legacy = StubCoordinator([{"status": "blocked"}])
    runner.task_coordinator = legacy
    assert (await runner._get_mission_completion("root"))["status"] == "blocked"
    assert legacy.completion_calls == 1


def _register_role_preset(name, role, max_iterations):
    register_preset(
        name,
        AgentPreset(
            role=role,
            allowed_tools=[],
            max_iterations=max_iterations,
            skills=[],
        ),
    )


def _explicit_config(name, max_iterations):
    root_task_spec = {
        "objective": "complete the mission",
        "acceptance_criteria": [{"description": "mission result is validated"}],
    }
    return RunConfig(
        agent_type=name,
        completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
        max_iterations=max_iterations,
        tools=[],
        tool_groups=[],
        root_task_spec=root_task_spec,
        knowledge=_knowledge_off(),
    )


@pytest.mark.asyncio
async def test_planner_no_tool_exit_is_persisted_and_loop_continues(tmp_path):
    preset = "test_completion_gate_planner"
    _register_role_preset(preset, AgentRole.PLANNER, 3)
    coordinator = StubCoordinator([
        {
            "status": "needs_replan",
            "root_task_id": "root-task",
            "pending_tasks": [{"task_id": "child", "status": "pending"}],
        },
        {
            "status": "completed",
            "root_task_id": "root-task",
            "pending_tasks": [],
            "result_summary": "accepted root result",
        },
    ])
    calls = []

    async def llm_call(**kwargs):
        calls.append(kwargs["messages"])
        return {"content": "", "tool_calls": None, "finish_reason": "stop"}

    store = FileSystemTraceStore(str(tmp_path))
    runner = AgentRunner(
        trace_store=store,
        llm_call=llm_call,
        task_coordinator=coordinator,
    )
    events = [
        event
        async for event in runner.run(
            [{"role": "user", "content": "run mission"}],
            _explicit_config(preset, 3),
        )
    ]

    final = events[-1]
    assert isinstance(final, Trace)
    assert final.status == "completed"
    assert final.result_summary == "accepted root result"
    assert coordinator.ensure_calls[0][1]["objective"] == "complete the mission"
    assert len(calls) == 2
    assert any(
        message.get("role") == "user"
        and "[FRAMEWORK_MISSION_INCOMPLETE]" in message.get("content", "")
        for message in calls[1]
    )
    assert calls[1][0]["role"] == "system"
    assert "Framework role contract: Planner" in calls[1][0]["content"]
    persisted = await store.get_trace_messages(final.trace_id)
    controls = [
        message for message in persisted
        if message.role == "user"
        and "[FRAMEWORK_MISSION_INCOMPLETE]" in str(message.content)
    ]
    assert len(controls) == 1


@pytest.mark.asyncio
async def test_incomplete_planner_resume_restores_explicit_identity(tmp_path):
    preset = "test_completion_gate_resume"
    _register_role_preset(preset, AgentRole.PLANNER, 1)
    coordinator = StubCoordinator([{
        "status": "needs_replan",
        "root_task_id": "root-task",
        "pending_tasks": [{"task_id": "root-task", "status": "needs_replan"}],
    }])

    async def llm_call(**_kwargs):
        return {"content": "", "tool_calls": None, "finish_reason": "stop"}

    store = FileSystemTraceStore(str(tmp_path))
    runner = AgentRunner(
        trace_store=store,
        llm_call=llm_call,
        task_coordinator=coordinator,
    )
    first = await runner.run_result(
        [{"role": "user", "content": "run mission"}],
        _explicit_config(preset, 1),
    )
    assert first["status"] == "incomplete"

    resume_config = RunConfig(
        trace_id=first["trace_id"],
        knowledge=_knowledge_off(),
    )
    resumed = await runner.run_result([], resume_config)

    assert resumed["status"] == "incomplete"
    assert resumed["error"] == "iteration_limit_reached"
    assert resume_config.agent_type == preset
    assert resume_config.completion_policy == CompletionPolicy.EXPLICIT_VALIDATION


@pytest.mark.asyncio
async def test_explicit_planner_rewind_and_wrong_role_resume_are_rejected_before_mutation(
    tmp_path,
):
    planner_preset = "test_completion_gate_rewind_planner"
    worker_preset = "test_completion_gate_wrong_role_worker"
    _register_role_preset(planner_preset, AgentRole.PLANNER, 2)
    _register_role_preset(worker_preset, AgentRole.WORKER, 2)
    store = FileSystemTraceStore(str(tmp_path))
    planner_trace = Trace(
        trace_id="planner-trace",
        mode="agent",
        task="mission",
        agent_type=planner_preset,
        agent_role=AgentRole.PLANNER.value,
        status="completed",
        head_sequence=0,
        last_sequence=10,
    )
    worker_trace = Trace(
        trace_id="worker-trace",
        mode="agent",
        task="worker",
        agent_type=worker_preset,
        agent_role=AgentRole.WORKER.value,
        status="completed",
    )
    await store.create_trace(planner_trace)
    await store.create_trace(worker_trace)
    coordinator = StubCoordinator([{"status": "completed"}])
    runner = AgentRunner(
        trace_store=store,
        llm_call=lambda **_kwargs: None,
        task_coordinator=coordinator,
    )

    with pytest.raises(ValueError, match="do not support rewind"):
        await runner.run_result(
            [],
            RunConfig(
                trace_id=planner_trace.trace_id,
                agent_type=planner_preset,
                completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
                after_sequence=5,
            ),
        )
    with pytest.raises(ValueError, match="Trace role mismatch"):
        await runner.run_result(
            [],
            RunConfig(
                trace_id=worker_trace.trace_id,
                agent_type=planner_preset,
                completion_policy=CompletionPolicy.EXPLICIT_VALIDATION,
            ),
        )

    assert (await store.get_trace(planner_trace.trace_id)).status == "completed"
    assert (await store.get_trace(worker_trace.trace_id)).status == "completed"
    assert coordinator.ensure_calls == []


@pytest.mark.asyncio
async def test_planner_iteration_exhaustion_is_incomplete(tmp_path):
    preset = "test_completion_gate_exhaustion"
    _register_role_preset(preset, AgentRole.PLANNER, 2)
    coordinator = StubCoordinator([{
        "status": "waiting_children",
        "root_task_id": "root-task",
        "pending_tasks": ["child"],
    }])

    async def llm_call(**_kwargs):
        return {"content": "", "tool_calls": None, "finish_reason": "stop"}

    runner = AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path)),
        llm_call=llm_call,
        task_coordinator=coordinator,
    )
    result = await runner.run_result(
        [{"role": "user", "content": "run mission"}],
        _explicit_config(preset, 2),
    )

    assert result["status"] == "incomplete"
    assert result["error"] == "iteration_limit_reached"
    assert result["summary"] is None


@pytest.mark.asyncio
async def test_blocked_root_stops_planner_as_incomplete(tmp_path):
    preset = "test_completion_gate_blocked"
    _register_role_preset(preset, AgentRole.PLANNER, 3)
    coordinator = StubCoordinator([{
        "status": "blocked",
        "root_task_id": "root-task",
        "blocked_reason": "external dependency",
        "pending_tasks": [],
    }])
    calls = 0

    async def llm_call(**_kwargs):
        nonlocal calls
        calls += 1
        return {"content": "blocked", "tool_calls": None, "finish_reason": "stop"}

    runner = AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path)),
        llm_call=llm_call,
        task_coordinator=coordinator,
    )
    result = await runner.run_result(
        [{"role": "user", "content": "run mission"}],
        _explicit_config(preset, 3),
    )

    assert calls == 1
    assert result["status"] == "incomplete"
    assert result["error"] == "mission_blocked"


@pytest.mark.asyncio
@pytest.mark.parametrize("role", [AgentRole.WORKER, AgentRole.VALIDATOR])
async def test_explicit_subtrace_without_terminal_submit_is_protocol_failure(
    tmp_path, role
):
    preset = "test_completion_gate_" + role.value
    _register_role_preset(preset, role, 2)

    async def llm_call(**_kwargs):
        return {"content": "done", "tool_calls": None, "finish_reason": "stop"}

    runner = AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path / role.value)),
        llm_call=llm_call,
        task_coordinator=StubCoordinator([{"status": "completed"}]),
    )
    result = await runner.run_result(
        [{"role": "user", "content": "sub task"}],
        _explicit_config(preset, 2),
    )

    assert result["status"] == "failed"
    assert result["error"] == "protocol_violation"
    assert result["failure"]["code"] == "PROTOCOL_VIOLATION"


@pytest.mark.asyncio
async def test_worker_replan_failure_terminates_without_protocol_overwrite(tmp_path):
    preset = "test_structured_failure_worker"
    register_preset(
        preset,
        AgentPreset(
            role=AgentRole.WORKER,
            allowed_tools=["reject_contract"],
            max_iterations=4,
            skills=[],
        ),
    )
    registry = ToolRegistry()

    async def reject_contract() -> str:
        raise ToolExecutionError(
            FailureDetail(
                code="INPUT_SCOPE_MISMATCH",
                message="Paragraph has no covering Structure",
                disposition=FailureDisposition.REPLAN_TASK,
            )
        )

    registry.register(
        reject_contract,
        capabilities=[ToolCapability.READ],
    )
    calls = 0

    async def llm_call(**_kwargs):
        nonlocal calls
        calls += 1
        return {
            "content": "",
            "tool_calls": [
                {
                    "id": "call-reject",
                    "type": "function",
                    "function": {"name": "reject_contract", "arguments": "{}"},
                }
            ],
            "finish_reason": "tool_calls",
        }

    store = FileSystemTraceStore(str(tmp_path))
    runner = AgentRunner(
        trace_store=store,
        tool_registry=registry,
        llm_call=llm_call,
        task_coordinator=StubCoordinator([{"status": "completed"}]),
    )
    config = _explicit_config(preset, 4)
    config.tools = ["reject_contract"]
    result = await runner.run_result([{"role": "user", "content": "work"}], config)

    assert calls == 1
    assert result["status"] == "failed"
    assert result["error"].startswith("INPUT_SCOPE_MISMATCH:")
    assert result["failure"]["code"] == "INPUT_SCOPE_MISMATCH"
    trace = await store.get_trace(result["trace_id"])
    assert trace.failure.code == "INPUT_SCOPE_MISMATCH"


@pytest.mark.asyncio
async def test_planner_observes_replan_failure_and_keeps_control(tmp_path):
    preset = "test_structured_failure_planner"
    register_preset(
        preset,
        AgentPreset(
            role=AgentRole.PLANNER,
            allowed_tools=["reject_contract"],
            max_iterations=3,
            skills=[],
        ),
    )
    registry = ToolRegistry()

    async def reject_contract() -> str:
        raise ToolExecutionError(
            FailureDetail(
                code="TASK_CONTRACT_NOT_EXECUTABLE",
                message="revise the frozen contract",
                disposition=FailureDisposition.REPLAN_TASK,
            )
        )

    registry.register(reject_contract, capabilities=[ToolCapability.TASK_CONTROL])
    calls = 0

    async def llm_call(**_kwargs):
        nonlocal calls
        calls += 1
        if calls == 1:
            return {
                "content": "",
                "tool_calls": [
                    {
                        "id": "call-reject",
                        "type": "function",
                        "function": {"name": "reject_contract", "arguments": "{}"},
                    }
                ],
                "finish_reason": "tool_calls",
            }
        return {"content": "replanned", "tool_calls": None, "finish_reason": "stop"}

    config = _explicit_config(preset, 3)
    config.tools = ["reject_contract"]
    result = await AgentRunner(
        trace_store=FileSystemTraceStore(str(tmp_path)),
        tool_registry=registry,
        llm_call=llm_call,
        task_coordinator=StubCoordinator([{"status": "completed"}]),
    ).run_result([{"role": "user", "content": "plan"}], config)

    assert calls == 2
    assert result["status"] == "completed"
    assert result["failure"] is None