/pentestgpt_agent/tests/test_supervisor.py

import json
from collections.abc import AsyncIterator
from pathlib import Path

import pytest
from unified_agent import AgentEvent, RunOptions, SandboxPolicy, TurnCompleted, UnifiedAgent

from pentestgpt_agent.agents import (
    EXECUTOR_SCHEMA,
    SUPERVISOR_INSTRUCTIONS,
    SUPERVISOR_SCHEMA,
    AgentContractError,
    Supervisor,
    _supervisor_prompt,
    parse_supervisor_decision,
)
from pentestgpt_agent.memory import (
    AttemptRecord,
    AttemptStatus,
    MemoryKernel,
    ObservationRecord,
    RunSnapshot,
    RunSpec,
    RunStatus,
)
from pentestgpt_agent.plan import TaskKind, TaskRecord, TaskStatus
from pentestgpt_agent.trace import EpisodeRunner, TraceStore


class SupervisorBackend:
    name = "scripted"

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        assert opts.sandbox is SandboxPolicy.FULL_ACCESS
        assert opts.output_schema is not None
        task_schema = opts.output_schema["properties"]["new_tasks"]["items"]
        assert task_schema["properties"]["kind"]["enum"] == [
            "discover",
            "enumerate",
            "test",
            "exploit",
            "verify",
            "recover",
        ]
        assert "maxItems" not in opts.output_schema["properties"]["new_tasks"]
        assert "speculative backlog" in opts.output_schema["properties"]["new_tasks"]["description"]
        assert "uniqueItems" not in task_schema["properties"]["basis_ids"]
        assert "newest same-target TEST" in task_schema["properties"]["basis_ids"]["description"]
        assert "uniqueItems" not in task_schema["properties"]["depends_on"]
        assert "basis-producing task" in task_schema["properties"]["depends_on"]["description"]
        assert "canonical evidence" in opts.output_schema["properties"]["finish"]["description"]
        assert "finish_basis_ids" in opts.output_schema["required"]
        finish_basis_schema = opts.output_schema["properties"]["finish_basis_ids"]
        assert "uniqueItems" not in finish_basis_schema
        assert "supplied canonical observation IDs" in finish_basis_schema["description"]
        yield TurnCompleted(
            success=True,
            final_text="structured decision",
            structured_output={
                "base_revision": 0,
                "new_tasks": [
                    {
                        "id": "discover-http",
                        "kind": "discover",
                        "target": "http://127.0.0.1:8080",
                        "objective": "Inspect the HTTP service.",
                        "done_when": "The reachable HTTP surface is recorded.",
                        "basis_ids": [],
                        "depends_on": [],
                    }
                ],
                "next_task_id": "discover-http",
                "finish": False,
                "finish_basis_ids": [],
                "summary": "Begin with HTTP discovery.",
            },
        )


def test_provider_output_schemas_use_the_codex_supported_subset() -> None:
    unsupported = {
        "format",
        "maxItems",
        "maxLength",
        "maximum",
        "minItems",
        "minLength",
        "minimum",
        "multipleOf",
        "pattern",
        "uniqueItems",
    }

    def walk(value: object) -> None:
        if isinstance(value, dict):
            assert not unsupported.intersection(value)
            for child in value.values():
                walk(child)
        elif isinstance(value, list):
            for child in value:
                walk(child)

    walk(SUPERVISOR_SCHEMA)
    walk(EXECUTOR_SCHEMA)


@pytest.mark.asyncio
async def test_supervisor_turns_authoritative_state_into_a_typed_decision(
    tmp_path: Path,
) -> None:
    assert "An exact anomaly is a lead" in SUPERVISOR_INSTRUCTIONS
    assert "one syntax-preserving, goal-directed derivative" in SUPERVISOR_INSTRUCTIONS
    assert "Use only supplied observation IDs" in SUPERVISOR_INSTRUCTIONS
    assert "newest completed TEST observation" in SUPERVISOR_INSTRUCTIONS
    assert "Finish only when canonical evidence" in SUPERVISOR_INSTRUCTIONS
    assert "one bounded TEST task" in SUPERVISOR_INSTRUCTIONS
    assert "Do not create one task per nearby payload" in SUPERVISOR_INSTRUCTIONS
    assert "whitespace or argument-shape control" in SUPERVISOR_INSTRUCTIONS
    assert "single-token redirection or IFS-style payload" in SUPERVISOR_INSTRUCTIONS
    assert "Recent diagnostics are noncanonical" in SUPERVISOR_INSTRUCTIONS
    assert "finish_basis_ids is empty unless finish is true" in SUPERVISOR_INSTRUCTIONS
    assert "one or more supplied canonical observation IDs" in SUPERVISOR_INSTRUCTIONS
    assert "every tool exposed by the provider" in SUPERVISOR_INSTRUCTIONS
    assert "file read/write tools" in SUPERVISOR_INSTRUCTIONS
    assert "byte-for-byte copy one supplied allowed target" in SUPERVISOR_INSTRUCTIONS
    assert "Ports, schemes, vhosts, URLs, and paths belong only in the objective" in (
        SUPERVISOR_INSTRUCTIONS
    )
    snapshot = MemoryKernel(tmp_path / "state.sqlite3").create_run(
        RunSpec(
            run_id="run-1",
            goal="Assess the authorized target.",
            allowed_targets=("http://127.0.0.1:8080",),
        )
    )
    agent = UnifiedAgent(
        SupervisorBackend(),
        workspace=tmp_path / "workspace",
        sandbox=SandboxPolicy.FULL_ACCESS,
        instructions=SUPERVISOR_INSTRUCTIONS,
    )
    supervisor = Supervisor(EpisodeRunner(agent, TraceStore(tmp_path / "runs")))

    decision = await supervisor.decide(snapshot, episode_id="supervisor-1")

    assert decision.base_revision == snapshot.revision
    assert decision.next_task_id == "discover-http"
    assert decision.finish is False
    assert decision.finish_basis_ids == ()
    assert len(decision.new_tasks) == 1
    assert decision.new_tasks[0].kind is TaskKind.DISCOVER
    assert decision.new_tasks[0].target == "http://127.0.0.1:8080"


def test_supervisor_retrieval_keeps_a_bounded_working_set_and_compact_history() -> None:
    tasks = tuple(
        TaskRecord(
            id=f"task-{index}",
            kind=TaskKind.TEST,
            target="http://target.test/input",
            objective=f"Full objective {index}",
            done_when=f"Completion condition {index}",
            basis_ids=(),
            depends_on=(),
            status=TaskStatus.DONE,
            created_revision=index,
        )
        for index in range(1, 11)
    )
    snapshot = RunSnapshot(
        run_id="run-1",
        goal="Capture the flag.",
        allowed_targets=("http://target.test",),
        status=RunStatus.RUNNING,
        revision=12,
        max_attempts_per_task=2,
        tasks=tasks,
    )

    state = json.loads(_supervisor_prompt(snapshot).split("\n\n", 1)[1])

    assert [task["id"] for task in state["tasks"]] == [
        "task-7",
        "task-8",
        "task-9",
        "task-10",
    ]
    assert state["task_history"]["total_closed"] == 10
    assert state["task_history"]["counts_by_status"] == {"done": 10}
    assert [task["id"] for task in state["task_history"]["recent"]] == [
        "task-3",
        "task-4",
        "task-5",
        "task-6",
    ]
    assert len(state["task_history"]["recent"]) == 4
    assert all(task.get("objective") != "Full objective 1" for task in state["tasks"])


def test_supervisor_bounds_observations_diagnostics_and_required_context() -> None:
    closed_tasks = tuple(
        TaskRecord(
            id=task_id,
            kind=TaskKind.TEST,
            target="http://target.test/input",
            objective=f"Objective for {task_id}",
            done_when=f"Done condition for {task_id}",
            basis_ids=(),
            depends_on=(),
            status=TaskStatus.DONE,
            created_revision=index,
        )
        for index, task_id in enumerate(
            (
                "basis-task",
                "dependency-task",
                "closed-3",
                "closed-4",
                "closed-5",
                "closed-6",
                "closed-7",
                "closed-8",
            ),
            start=1,
        )
    )
    open_task = TaskRecord(
        id="open-task",
        kind=TaskKind.EXPLOIT,
        target="http://target.test/input",
        objective="Use the confirmed primitive.",
        done_when="The goal artifact is captured.",
        basis_ids=("obs-required",),
        depends_on=("basis-task", "dependency-task"),
        status=TaskStatus.READY,
        created_revision=9,
    )
    observations = tuple(
        ObservationRecord(
            id=observation_id,
            task_id="basis-task" if observation_id == "obs-required" else "closed-8",
            attempt_id=f"attempt-{index}",
            statement=f"Evidence {observation_id}",
            trace_episode_id=f"executor-{index}",
            evidence_sequences=(1,),
            created_revision=index,
        )
        for index, observation_id in enumerate(
            ("obs-required", *(f"obs-{number}" for number in range(2, 10))),
            start=1,
        )
    )
    attempts = tuple(
        AttemptRecord(
            id=f"attempt-{index}",
            task_id="closed-8",
            status=status,
            started_revision=index,
            finished_revision=index + 1,
            trace_episode_id=f"executor-{index}",
            summary=summary,
            failure_kind="provider" if status is AttemptStatus.ERROR else None,
            failure_message="provider failed" if status is AttemptStatus.ERROR else None,
        )
        for index, (status, summary) in enumerate(
            (
                (AttemptStatus.FAILED, "old failure"),
                (AttemptStatus.PROGRESS, "more work remains"),
                (AttemptStatus.BLOCKED, "missing prerequisite"),
                (AttemptStatus.DONE, "successful DONE summary must stay out"),
                (AttemptStatus.ERROR, "provider error"),
                (AttemptStatus.FAILED, "latest failure"),
            ),
            start=1,
        )
    )
    snapshot = RunSnapshot(
        run_id="run-1",
        goal="Capture the flag.",
        allowed_targets=("http://target.test",),
        status=RunStatus.RUNNING,
        revision=20,
        max_attempts_per_task=2,
        tasks=(*closed_tasks, open_task),
        observations=observations,
        attempts=attempts,
    )

    state = json.loads(_supervisor_prompt(snapshot).split("\n\n", 1)[1])

    assert [task["id"] for task in state["tasks"]] == [
        "closed-5",
        "closed-6",
        "closed-7",
        "closed-8",
        "open-task",
    ]
    assert [task["id"] for task in state["required_task_context"]] == [
        "basis-task",
        "dependency-task",
    ]
    assert [observation["id"] for observation in state["observations"]] == [
        "obs-required",
        "obs-4",
        "obs-5",
        "obs-6",
        "obs-7",
        "obs-8",
        "obs-9",
    ]
    assert [diagnostic["status"] for diagnostic in state["recent_diagnostics"]] == [
        "progress",
        "blocked",
        "error",
        "failed",
    ]
    assert "successful DONE summary must stay out" not in json.dumps(state)


def test_supervisor_prompt_bounds_optional_validation_feedback() -> None:
    snapshot = RunSnapshot(
        run_id="run-1",
        goal="Capture the flag.",
        allowed_targets=("http://target.test",),
        status=RunStatus.RUNNING,
        revision=2,
        max_attempts_per_task=2,
    )

    state_without_feedback = json.loads(_supervisor_prompt(snapshot).split("\n\n", 1)[1])
    state_with_feedback = json.loads(
        _supervisor_prompt(snapshot, feedback="x" * 2_000).split("\n\n", 1)[1]
    )

    assert "validation_feedback" not in state_without_feedback
    assert state_with_feedback["validation_feedback"] == "x" * 1_000


def test_live_supervisor_output_requires_finish_basis_ids() -> None:
    raw = {
        "base_revision": 2,
        "new_tasks": [],
        "next_task_id": "ready-task",
        "finish": False,
        "summary": "Continue the ready task.",
    }

    with pytest.raises(AgentContractError, match="unexpected or missing fields"):
        parse_supervisor_decision(raw)