script-build-and-framework-tmp/script_build_host/tests/test_mission_workbench.py

from __future__ import annotations

import json
from dataclasses import asdict
from types import SimpleNamespace
from typing import Any

import pytest
from agent.orchestration import AgentRole, ArtifactRef, RoleContextRequest, TaskStatus

from script_build_host.application.mission_workbench import (
    MissionWorkbench,
    WorkbenchError,
    _fit_planner_state,
)
from script_build_host.domain.task_contracts import (
    ScriptCriterion,
    ScriptIntentClass,
    ScriptTaskBudget,
    ScriptTaskContractV1,
    ScriptTaskKind,
)
from script_build_host.domain.workbench import NotModified, PlannerState, semantic_handle


class _TaskStore:
    def __init__(self, ledger: Any) -> None:
        self.ledger = ledger

    async def load(self, root_trace_id: str) -> Any:
        assert root_trace_id == "root"
        return self.ledger


class _Receipts:
    def __init__(self) -> None:
        self.events: list[dict[str, Any]] = []

    async def append_event(self, _trace_id: str, event: str, payload: dict[str, Any]) -> None:
        self.events.append({"event": event, "event_id": len(self.events) + 1, **payload})

    async def get_events(self, _trace_id: str) -> list[dict[str, Any]]:
        return list(self.events)


class _Bindings:
    async def get_by_root(self, root_trace_id: str) -> Any:
        assert root_trace_id == "root"
        return SimpleNamespace(
            root_trace_id="root",
            script_build_id=900049,
            input_snapshot_id=3,
            active_direction_artifact_version_id=None,
        )


class _Snapshots:
    async def get(self, snapshot_id: str, *, script_build_id: int) -> Any:
        assert (snapshot_id, script_build_id) == ("3", 900049)
        return SimpleNamespace(
            snapshot_id="3",
            topic_id=49,
            topic={"topic": {"result": "one bounded topic"}},
            account={"account_name": "fixture"},
            persona_points=tuple({"name": f"point-{index}"} for index in range(80)),
            section_patterns=tuple({"name": f"pattern-{index}"} for index in range(80)),
            strategies=(),
        )


class _Contracts:
    def __init__(self, contract: ScriptTaskContractV1 | None = None) -> None:
        self.contract = contract

    async def read(self, root_trace_id: str, contract_ref: str) -> Any:
        assert root_trace_id == "root"
        assert contract_ref.startswith("script-build://task-contracts/sha256/")
        assert self.contract is not None
        return SimpleNamespace(contract=self.contract)


class _Artifact:
    def content_payload(self) -> dict[str, Any]:
        return {
            "paragraphs": [
                {
                    "paragraph_id": 77,
                    "name": "opening",
                    "full_description": "complete candidate prose",
                }
            ],
            "artifact_ref": "script-build://artifact-versions/77",
            "canonical_sha256": "sha256:" + "a" * 64,
        }


class _Artifacts:
    async def read_by_ref(self, ref: ArtifactRef, *, script_build_id: int) -> Any:
        assert script_build_id == 900049
        assert ref.uri == "script-build://artifact-versions/77"
        return SimpleNamespace(artifact=_Artifact())


class _Candidates:
    async def read_attempt_workspace(self, *, context: dict[str, Any]) -> dict[str, Any]:
        assert context["attempt_id"] == "attempt"
        return {
            "paragraphs": [
                {
                    "paragraph_id": 77,
                    "paragraph_index": 1,
                    "name": "opening",
                    "content_range": {"scope": "opening"},
                    "level": 1,
                    "theme": "theme",
                    "form": "form",
                    "function": "function",
                    "feeling": "feeling",
                    "description": "description",
                    "full_description": "complete prose",
                }
            ],
            "elements": [],
            "paragraph_element_links": [],
        }


def test_planner_state_compacts_optional_frontiers_instead_of_failing() -> None:
    state = PlannerState(
        state_revision="ledger:1",
        phase=2,
        root={"task_id": "root", "status": "waiting_children", "objective": "compose"},
        focused_task_id="root",
        active_subgraph=tuple(
            {"task_id": f"task-{index}", "objective": "x" * 500} for index in range(12)
        ),
        accepted_frontier=tuple(
            {"decision_id": f"decision-{index}", "summary": "y" * 1_200}
            for index in range(10)
        ),
        goal_progress={"total": 1, "covered": 0, "uncovered_goal_ids": ["goal"]},
        counts_by_status={"completed": 10},
        active_operations=(),
        latest_failures=(),
        retired_task_count=20,
        active_total=12,
        accepted_total=10,
        context_bundle={"cards": [{"summary": "z" * 1_000}]},
        phase_protocol={"phase": 2},
    )

    compacted = _fit_planner_state(state)

    assert compacted.accepted_has_more is True
    assert len(compacted.accepted_frontier) < len(state.accepted_frontier)
    assert len(json.dumps(compacted.to_payload(), ensure_ascii=False)) <= 24_000


def _root_task(*, children: list[str] | None = None) -> Any:
    return SimpleNamespace(
        task_id="root-task",
        parent_task_id=None,
        child_task_ids=list(children or []),
        display_path="1",
        status=TaskStatus.RUNNING,
        blocked_reason=None,
        current_spec=SimpleNamespace(objective="build script", context_refs=()),
        current_spec_version=1,
        attempt_ids=[],
        decision_ids=[],
    )


def _contract() -> ScriptTaskContractV1:
    return ScriptTaskContractV1(
        task_kind=ScriptTaskKind.PARAGRAPH,
        scope_ref="script-build://scopes/full/opening",
        intent_class=ScriptIntentClass.EXPLORE,
        objective="write the opening",
        input_decision_refs=(),
        base_artifact_ref=None,
        write_scope=("script-build://writes/paragraphs/opening",),
        gap_ref=None,
        output_schema="paragraph-artifact/v1",
        criteria=(ScriptCriterion("criterion", "complete prose"),),
        budget=ScriptTaskBudget(),
        goal_ids=("goal-1",),
        compiler_manifest_digest="sha256:" + "b" * 64,
    )


def _worker_task() -> Any:
    return SimpleNamespace(
        task_id="worker",
        parent_task_id="root-task",
        child_task_ids=[],
        display_path="1.1",
        status=TaskStatus.RUNNING,
        blocked_reason=None,
        current_spec=SimpleNamespace(
            objective="write the opening",
            context_refs=(
                "script-build://task-kinds/paragraph",
                "script-build://task-contracts/sha256/contract",
            ),
        ),
        current_spec_version=1,
        attempt_ids=["attempt"],
        decision_ids=[],
    )


@pytest.mark.asyncio
async def test_planner_projection_is_bounded_incremental_and_does_not_copy_retired_tree() -> None:
    tasks = {"root-task": _root_task(children=[f"done-{index}" for index in range(120)])}
    for index in range(120):
        tasks[f"done-{index}"] = SimpleNamespace(
            task_id=f"done-{index}",
            parent_task_id="root-task",
            child_task_ids=[],
            display_path=f"1.{index + 1}",
            status=TaskStatus.COMPLETED,
            blocked_reason=None,
            current_spec=SimpleNamespace(
                objective="retired detail " + "x" * 1000,
                context_refs=(),
            ),
            current_spec_version=1,
            attempt_ids=[],
            decision_ids=[],
        )
    ledger = SimpleNamespace(
        revision=49,
        root_task_id="root-task",
        focused_task_id="root-task",
        tasks=tasks,
        attempts={},
        decisions={},
        operations={},
    )
    receipts = _Receipts()
    workbench = MissionWorkbench(
        task_store=_TaskStore(ledger),
        bindings=_Bindings(),
        snapshots=_Snapshots(),
        artifacts=_Artifacts(),
        contracts=_Contracts(),
        receipt_store=receipts,
    )

    state = await workbench.planner_state("root")
    assert not isinstance(state, NotModified)
    assert state.retired_task_count == 120
    assert len(state.active_subgraph) == 1
    assert state.phase_protocol["phase_epoch"] == "script-build-phase-1"
    assert state.phase_protocol["current_task"]["allowed_child_kinds"] == ("direction",)
    encoded = await workbench.render("root", ledger)
    assert len(encoded) < 24_000
    assert "retired detail" not in encoded
    legacy_full_tree = json.dumps(
        [item.current_spec.objective for item in tasks.values()], ensure_ascii=False
    )
    assert len(encoded) * 2 < len(legacy_full_tree)
    for _ in range(15):
        assert isinstance(
            await workbench.planner_state("root", known_revision="ledger:49"),
            NotModified,
        )
    first_page = await workbench.search_mission_context(
        root_trace_id="root",
        role="planner",
        collection="tasks",
        page_size=20,
    )
    assert 0 < len(first_page["page"]["items"]) <= 20
    assert first_page["receipt"]["estimated_tokens"] <= 2_000
    assert first_page["page"]["has_more"] is True
    assert first_page["page"]["next_cursor"]
    assert len(json.dumps(first_page, ensure_ascii=False, separators=(",", ":"))) < 8_000
    await workbench.search_mission_context(
        root_trace_id="root",
        role="planner",
        collection="tasks",
        page_size=20,
    )
    progress = await workbench.progress_snapshot("root")
    assert len(progress["unique_observations"]) == 1
    await workbench.search_mission_context(
        root_trace_id="root",
        role="planner",
        collection="tasks",
        query="different semantic scope",
        page_size=20,
    )
    progress = await workbench.progress_snapshot("root")
    assert len(progress["unique_observations"]) == 2


@pytest.mark.asyncio
async def test_role_bootstrap_uses_semantic_targets_and_protects_artifact_identities() -> None:
    root = _root_task(children=["worker"])
    task = _worker_task()
    attempt = SimpleNamespace(attempt_id="attempt", failure=None)
    ledger = SimpleNamespace(
        revision=7,
        root_task_id="root-task",
        focused_task_id="worker",
        tasks={"root-task": root, "worker": task},
        attempts={"attempt": attempt},
        decisions={},
        operations={},
    )
    workbench = MissionWorkbench(
        task_store=_TaskStore(ledger),
        bindings=_Bindings(),
        snapshots=_Snapshots(),
        artifacts=_Artifacts(),
        contracts=_Contracts(_contract()),
        candidates=_Candidates(),
    )
    worker_request = RoleContextRequest(
        role=AgentRole.WORKER,
        root_trace_id="root",
        task_id="worker",
        spec_version=1,
        task_spec={},
        attempt_id="attempt",
        ledger_revision=7,
    )
    worker = await workbench.build(worker_request)
    encoded_worker = json.dumps(worker, ensure_ascii=False)
    assert worker["paragraph_targets"][0]["paragraph_target_key"].startswith("pt_")
    assert "paragraph_id" not in encoded_worker
    assert len(encoded_worker) < 32_000

    artifact_ref = ArtifactRef(
        uri="script-build://artifact-versions/77",
        kind="paragraph",
        version="77",
        digest="sha256:" + "a" * 64,
    )
    validator_request = RoleContextRequest(
        role=AgentRole.VALIDATOR,
        root_trace_id="root",
        task_id="worker",
        spec_version=1,
        task_spec={},
        attempt_id="attempt",
        ledger_revision=7,
        snapshot_id="snapshot",
        artifact_snapshot={
            "artifact_refs": [asdict(artifact_ref)],
            "evidence_refs": [],
        },
    )
    validator = await workbench.build(validator_request)
    encoded_validator = json.dumps(validator, ensure_ascii=False)
    assert "complete candidate prose" in encoded_validator
    assert "paragraph_id" not in encoded_validator
    assert "script-build://artifact-versions/77" not in encoded_validator
    assert "canonical_sha256" not in encoded_validator

    with pytest.raises(WorkbenchError) as physical:
        await workbench.read_mission_context(
            root_trace_id="root",
            role="planner",
            handle="script-build://artifact-versions/77",
        )
    assert physical.value.code == "CONTEXT_HANDLE_FORMAT_INVALID"

    unknown_handle = semantic_handle("unknown", "root", "not-present")
    with pytest.raises(WorkbenchError) as unknown:
        await workbench.read_mission_context(
            root_trace_id="root",
            role="planner",
            handle=unknown_handle,
        )
    assert unknown.value.code == "CONTEXT_HANDLE_UNKNOWN"

    with pytest.raises(WorkbenchError) as unauthorized:
        await workbench.read_mission_context(
            root_trace_id="root",
            role="worker",
            task_id="worker",
            attempt_id="attempt",
            handle=unknown_handle,
        )
    assert unauthorized.value.code == "CONTEXT_HANDLE_UNAUTHORIZED"