/pentestgpt_agent/tests/test_trial.py

import json
import sqlite3
from collections.abc import AsyncIterator
from dataclasses import replace
from pathlib import Path

import pytest
from unified_agent import (
    AgentEvent,
    CommandRun,
    RunOptions,
    SandboxPolicy,
    TurnCompleted,
    UnifiedAgent,
)

from pentestgpt_agent import trial
from pentestgpt_agent.agents import (
    EXECUTOR_INSTRUCTIONS,
    SUPERVISOR_INSTRUCTIONS,
    Executor,
    Supervisor,
)
from pentestgpt_agent.audit import _is_hidden_provider_memory_action, audit_run
from pentestgpt_agent.memory import MemoryKernel, RunSpec
from pentestgpt_agent.trace import EpisodeRunner, TraceStore
from pentestgpt_agent.trial import TrialConfig, run_trial


class TrialSupervisorBackend:
    name = "scripted"

    def __init__(self) -> None:
        self.calls = 0

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        if self.calls == 0:
            output = {
                "base_revision": 0,
                "new_tasks": [
                    {
                        "id": "capture-flag",
                        "kind": "verify",
                        "target": "http://target.test",
                        "objective": "Capture the benchmark flag.",
                        "done_when": "The exact flag is recorded.",
                        "basis_ids": [],
                        "depends_on": [],
                    }
                ],
                "next_task_id": "capture-flag",
                "finish": False,
                "finish_basis_ids": [],
                "summary": "Capture the flag directly.",
            }
        else:
            state = json.loads(prompt.split("\n\n", 1)[1])
            output = {
                "base_revision": 2,
                "new_tasks": [],
                "next_task_id": None,
                "finish": True,
                "finish_basis_ids": [state["observations"][-1]["id"]],
                "summary": "The exact benchmark flag was captured.",
            }
        self.calls += 1
        yield CommandRun(command="inspect-planning-workspace", exit_code=0, output="ok")
        yield TurnCompleted(success=True, structured_output=output)


class TrialExecutorBackend:
    name = "scripted"

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        yield CommandRun(command="read-flag", exit_code=0, output="FLAG{trial-oracle}")
        yield TurnCompleted(
            success=True,
            structured_output={
                "task_id": "capture-flag",
                "outcome": "done",
                "summary": "Captured the exact flag.",
                "evidence_excerpt": "FLAG{trial-oracle}",
            },
        )


class FailureSupervisorBackend:
    name = "scripted"

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        yield TurnCompleted(
            success=True,
            structured_output={
                "base_revision": 0,
                "new_tasks": [
                    {
                        "id": "test-target",
                        "kind": "test",
                        "target": "http://target.test",
                        "objective": "Test the target.",
                        "done_when": "The behavior is recorded.",
                        "basis_ids": [],
                        "depends_on": [],
                    }
                ],
                "next_task_id": "test-target",
                "finish": False,
                "finish_basis_ids": [],
                "summary": "Test the target.",
            },
        )


class FailureExecutorBackend:
    name = "scripted"

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        yield CommandRun(command="curl http://target.test", exit_code=0, output="partial")
        yield TurnCompleted(
            success=False,
            stop_reason="error_max_turns",
            error="Reached maximum number of turns (2)",
        )
        raise RuntimeError("Command failed with exit code 1")


class InvalidSupervisorBackend:
    name = "scripted"

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        yield TurnCompleted(success=True, structured_output={"invalid": "decision"})


class MustNotRunBackend:
    name = "scripted"

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        raise AssertionError("a completed resumed trial must not invoke a provider")
        yield  # pragma: no cover


def test_claude_provider_memory_is_disabled_for_fresh_episodes() -> None:
    assert trial._provider_environment("claude") == {
        "CLAUDE_CODE_DISABLE_AUTO_MEMORY": "1",
        "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
        "DISABLE_AUTOUPDATER": "1",
    }
    assert trial._provider_environment("codex") == {}


def test_audit_recognizes_provider_side_memory_writes() -> None:
    assert _is_hidden_provider_memory_action(
        {
            "type": "tool_call",
            "name": "Write",
            "input": {"file_path": "/home/pentester/.claude/projects/-trial/memory/MEMORY.md"},
        }
    )
    assert not _is_hidden_provider_memory_action(
        {
            "type": "tool_call",
            "name": "Write",
            "input": {"file_path": "/trial/workspaces/run/notes.md"},
        }
    )


def test_audit_emits_an_incomplete_record_for_an_empty_episode_directory(
    tmp_path: Path,
) -> None:
    run_dir = tmp_path / "runs" / "incomplete-trial"
    MemoryKernel(run_dir / "state.sqlite3").create_run(
        RunSpec(
            run_id="incomplete-trial",
            goal="Assess the target.",
            allowed_targets=("http://target.test",),
        )
    )
    (run_dir / "traces" / "executor-incomplete").mkdir(parents=True)

    audit = audit_run(run_dir, expected_flag="FLAG{not-present}")

    assert audit["passed"] is False
    assert audit["checks"]["all_episodes_complete"] is False
    assert audit["episodes"] == [
        {
            "episode_id": "executor-incomplete",
            "run_id": None,
            "input_episode_id": None,
            "role": None,
            "state_revision": None,
            "task_id": None,
            "attempt_id": None,
            "opened_at": None,
            "closed_at": None,
            "success": False,
            "duration_ms": None,
            "cost_usd": 0.0,
            "usage": {},
            "structured_output": None,
            "event_counts": {},
            "actions": [],
            "events": [],
            "truncated_tail": False,
            "complete": False,
            "integrity_errors": [
                "missing input.json",
                "missing output.json",
                "missing events.jsonl",
            ],
        }
    ]


def test_audit_accounts_for_an_initialization_crash_settled_by_run_transition(
    tmp_path: Path,
) -> None:
    run_dir = tmp_path / "runs" / "settled-supervisor-crash"
    memory = MemoryKernel(run_dir / "state.sqlite3")
    memory.create_run(
        RunSpec(
            run_id="settled-supervisor-crash",
            goal="Assess the target.",
            allowed_targets=("http://target.test",),
        )
    )
    (run_dir / "traces" / "supervisor-r0").mkdir(parents=True)
    memory.commit_run_failure(
        "settled-supervisor-crash",
        0,
        failure_kind="supervisor_contract",
        failure_message="initialization was interrupted twice",
    )

    audit = audit_run(run_dir, expected_flag="FLAG{not-present}")

    assert audit["checks"]["all_episodes_complete"] is True
    assert audit["checks"]["transition_timeline_complete"] is True
    assert audit["passed"] is False


def test_audit_emits_partial_events_and_integrity_errors_for_malformed_trace_files(
    tmp_path: Path,
) -> None:
    run_dir = tmp_path / "runs" / "malformed-trial"
    MemoryKernel(run_dir / "state.sqlite3").create_run(
        RunSpec(
            run_id="malformed-trial",
            goal="Assess the target.",
            allowed_targets=("http://target.test",),
        )
    )
    episode_dir = run_dir / "traces" / "executor-malformed"
    episode_dir.mkdir(parents=True)
    (episode_dir / "input.json").write_text("{", encoding="utf-8")
    (episode_dir / "events.jsonl").write_text(
        '{"sequence":1,"type":"assistant_text","text":"partial"}\n{',
        encoding="utf-8",
    )
    (episode_dir / "output.json").write_text("[", encoding="utf-8")

    audit = audit_run(run_dir, expected_flag="FLAG{not-present}")

    assert audit["passed"] is False
    episode = audit["episodes"][0]
    assert episode["complete"] is False
    assert episode["truncated_tail"] is True
    assert episode["events"] == [{"sequence": 1, "type": "assistant_text", "text": "partial"}]
    assert episode["integrity_errors"] == [
        "malformed input.json",
        "malformed output.json",
        "malformed events.jsonl at line 2",
    ]


@pytest.mark.asyncio
async def test_trial_writes_a_run_that_the_artifact_auditor_can_verify(
    tmp_path: Path,
) -> None:
    runs_root = tmp_path / "runs"
    traces = TraceStore(runs_root)
    supervisor = Supervisor(
        EpisodeRunner(
            UnifiedAgent(
                TrialSupervisorBackend(),
                workspace=tmp_path / "supervisor",
                sandbox=SandboxPolicy.FULL_ACCESS,
                instructions=SUPERVISOR_INSTRUCTIONS,
            ),
            traces,
        )
    )
    executor = Executor(
        EpisodeRunner(
            UnifiedAgent(
                TrialExecutorBackend(),
                workspace=tmp_path / "executor",
                sandbox=SandboxPolicy.FULL_ACCESS,
                instructions=EXECUTOR_INSTRUCTIONS,
            ),
            traces,
        )
    )
    config = TrialConfig(
        run_id="trial-run",
        goal="Capture the benchmark flag.",
        targets=("http://target.test",),
        backend="claude",
        model="claude-opus-4-8",
        runs_root=runs_root,
        workspace_root=tmp_path / "agents",
        effort="xhigh",
        max_decisions=4,
        supervisor_max_turns=2,
        executor_max_turns=2,
    )

    summary = await run_trial(config, supervisor=supervisor, executor=executor)
    audit = audit_run(runs_root / "trial-run", expected_flag="FLAG{trial-oracle}")

    assert summary["status"] == "completed"
    assert summary["model"] == "claude-opus-4-8"
    assert summary["effort"] == "xhigh"
    assert summary["runtime_policy_revision"] == 2
    assert [attempt["status"] for attempt in summary["attempts"]] == ["done"]
    assert [transition["kind"] for transition in summary["transitions"]] == [
        "run_created",
        "plan_committed",
        "attempt_committed",
        "plan_committed",
    ]
    assert (runs_root / "trial-run" / "summary.json").exists()
    assert audit["passed"] is True
    assert audit["checks"] == {
        "run_completed": True,
        "oracle_in_canonical_observation": True,
        "no_active_tasks": True,
        "no_active_attempts": True,
        "all_episodes_complete": True,
        "all_observations_grounded": True,
        "all_observations_are_direct_quotes": True,
        "all_observation_identities_match": True,
        "all_basis_ids_exist": True,
        "all_basis_producers_are_dependencies": True,
        "all_exploit_bases_current": True,
        "completion_basis_valid": True,
        "transition_timeline_complete": True,
        "all_failed_episodes_settled": True,
    }
    assert audit["schema_version"] == 2
    assert audit["totals"]["supervisor_actions"] == 2
    assert len(audit["episodes"]) == 3
    trial_identity = json.loads(
        (runs_root / "trial-run" / "trial-config.json").read_text(encoding="utf-8")
    )
    assert trial_identity["schema_version"] == 2
    assert trial_identity["supervisor_sandbox"] == "full_access"
    assert trial_identity["executor_sandbox"] == "full_access"

    executor_episode = next(
        episode
        for episode in (runs_root / "trial-run" / "traces").iterdir()
        if json.loads((episode / "input.json").read_text(encoding="utf-8"))["role"] == "executor"
    )
    events_path = executor_episode / "events.jsonl"
    events = [json.loads(line) for line in events_path.read_text(encoding="utf-8").splitlines()]
    command_receipt = next(event for event in events if event["type"] == "command_run")
    command_receipt["exit_code"] = 255
    events_path.write_text(
        "".join(json.dumps(event) + "\n" for event in events),
        encoding="utf-8",
    )
    nonzero_exit_audit = audit_run(
        runs_root / "trial-run",
        expected_flag="FLAG{trial-oracle}",
    )
    assert nonzero_exit_audit["checks"]["all_observations_grounded"] is True
    assert nonzero_exit_audit["checks"]["all_observations_are_direct_quotes"] is True

    database = runs_root / "trial-run" / "state.sqlite3"
    with sqlite3.connect(database) as connection:
        connection.execute(
            "UPDATE observations SET statement = ?",
            ("FLAG{trial-oracle} with uncaptured suffix",),
        )
    tampered = audit_run(runs_root / "trial-run", expected_flag="FLAG{trial-oracle}")
    assert tampered["checks"]["oracle_in_canonical_observation"] is True
    assert tampered["checks"]["all_observations_are_direct_quotes"] is False
    assert tampered["passed"] is False


@pytest.mark.asyncio
async def test_trial_interface_can_explicitly_resume_a_persisted_run(tmp_path: Path) -> None:
    runs_root = tmp_path / "runs"
    traces = TraceStore(runs_root)
    initial_config = TrialConfig(
        run_id="resumable-trial",
        goal="Capture the benchmark flag.",
        targets=("http://target.test",),
        backend="claude",
        model="claude-opus-4-8",
        runs_root=runs_root,
        workspace_root=tmp_path / "agents",
        max_decisions=4,
        supervisor_max_turns=2,
        executor_max_turns=2,
    )
    initial = await run_trial(
        initial_config,
        supervisor=Supervisor(
            EpisodeRunner(
                UnifiedAgent(
                    TrialSupervisorBackend(),
                    workspace=tmp_path / "supervisor",
                    sandbox=SandboxPolicy.FULL_ACCESS,
                    instructions=SUPERVISOR_INSTRUCTIONS,
                ),
                traces,
            )
        ),
        executor=Executor(
            EpisodeRunner(
                UnifiedAgent(
                    TrialExecutorBackend(),
                    workspace=tmp_path / "executor",
                    sandbox=SandboxPolicy.FULL_ACCESS,
                    instructions=EXECUTOR_INSTRUCTIONS,
                ),
                traces,
            )
        ),
    )
    assert initial["status"] == "completed"

    resume_config = replace(initial_config, resume=True)
    resumed = await run_trial(
        resume_config,
        supervisor=Supervisor(
            EpisodeRunner(
                UnifiedAgent(
                    MustNotRunBackend(),
                    workspace=tmp_path / "supervisor-resume",
                    sandbox=SandboxPolicy.FULL_ACCESS,
                    instructions=SUPERVISOR_INSTRUCTIONS,
                ),
                traces,
            )
        ),
        executor=Executor(
            EpisodeRunner(
                UnifiedAgent(
                    MustNotRunBackend(),
                    workspace=tmp_path / "executor-resume",
                    sandbox=SandboxPolicy.FULL_ACCESS,
                    instructions=EXECUTOR_INSTRUCTIONS,
                ),
                traces,
            )
        ),
    )

    assert resumed["status"] == "completed"
    assert resumed["revision"] == initial["revision"]
    assert resumed["episodes"] == initial["episodes"]

    with pytest.raises(ValueError, match="does not match requested config"):
        await run_trial(replace(resume_config, model="different-model"))

    with pytest.raises(ValueError, match="does not match requested config"):
        await run_trial(replace(resume_config, effort="high"))


def test_build_roles_forward_effort_and_grant_full_access(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    class CapturedAgent:
        def __init__(self, backend: object, **options: object) -> None:
            self.backend = backend
            self.effort = options["effort"]
            self.sandbox = options["sandbox"]
            self.instructions = options["instructions"]

    monkeypatch.setattr(trial, "UnifiedAgent", CapturedAgent)
    config = TrialConfig(
        run_id="effort-trial",
        goal="Assess the target.",
        targets=("http://target.test",),
        backend="claude",
        model="claude-opus-4-8",
        runs_root=tmp_path / "runs",
        workspace_root=tmp_path / "agents",
        effort="xhigh",
    )

    supervisor, executor = trial._build_roles(config, TraceStore(config.runs_root))

    assert supervisor.runner.agent.effort == "xhigh"
    assert executor.runner.agent.effort == "xhigh"
    assert supervisor.runner.agent.sandbox is SandboxPolicy.FULL_ACCESS
    assert executor.runner.agent.sandbox is SandboxPolicy.FULL_ACCESS


@pytest.mark.asyncio
async def test_trial_rejects_a_path_like_run_id_before_creating_artifacts(
    tmp_path: Path,
) -> None:
    config = TrialConfig(
        run_id="../escape",
        goal="Assess the target.",
        targets=("http://target.test",),
        backend="claude",
        model="claude-opus-4-8",
        runs_root=tmp_path / "runs",
        workspace_root=tmp_path / "agents",
    )

    with pytest.raises(ValueError, match="run_id must be"):
        await run_trial(config)

    assert not (tmp_path / "escape").exists()


@pytest.mark.asyncio
async def test_failed_trial_summary_has_typed_failure_and_no_active_lease(
    tmp_path: Path,
) -> None:
    runs_root = tmp_path / "runs"
    traces = TraceStore(runs_root)
    supervisor = Supervisor(
        EpisodeRunner(
            UnifiedAgent(
                FailureSupervisorBackend(),
                workspace=tmp_path / "supervisor",
                sandbox=SandboxPolicy.FULL_ACCESS,
                instructions=SUPERVISOR_INSTRUCTIONS,
            ),
            traces,
        )
    )
    executor = Executor(
        EpisodeRunner(
            UnifiedAgent(
                FailureExecutorBackend(),
                workspace=tmp_path / "executor",
                sandbox=SandboxPolicy.FULL_ACCESS,
                instructions=EXECUTOR_INSTRUCTIONS,
            ),
            traces,
        ),
        max_turns=2,
    )
    config = TrialConfig(
        run_id="failed-trial",
        goal="Assess the target.",
        targets=("http://target.test",),
        backend="claude",
        model="claude-opus-4-8",
        runs_root=runs_root,
        workspace_root=tmp_path / "agents",
        max_decisions=2,
        supervisor_max_turns=2,
        executor_max_turns=2,
    )

    summary = await run_trial(config, supervisor=supervisor, executor=executor)
    audit = audit_run(runs_root / "failed-trial", expected_flag="FLAG{not-present}")

    assert summary["status"] == "failed"
    assert summary["error"] == "max_turns: Reached maximum number of turns (2)"
    assert summary["attempts"][0]["status"] == "error"
    assert summary["attempts"][0]["failure_kind"] == "max_turns"
    assert summary["tasks"][0]["status"] == "failed"
    assert audit["checks"]["no_active_tasks"] is True
    assert audit["checks"]["no_active_attempts"] is True
    assert audit["checks"]["all_failed_episodes_settled"] is True
    assert audit["checks"]["transition_timeline_complete"] is True
    assert audit["passed"] is False


@pytest.mark.asyncio
async def test_run_level_supervisor_failure_is_present_in_the_trial_summary(
    tmp_path: Path,
) -> None:
    runs_root = tmp_path / "runs"
    traces = TraceStore(runs_root)
    supervisor = Supervisor(
        EpisodeRunner(
            UnifiedAgent(
                InvalidSupervisorBackend(),
                workspace=tmp_path / "supervisor",
                sandbox=SandboxPolicy.FULL_ACCESS,
                instructions=SUPERVISOR_INSTRUCTIONS,
            ),
            traces,
        )
    )
    executor = Executor(
        EpisodeRunner(
            UnifiedAgent(
                FailureExecutorBackend(),
                workspace=tmp_path / "executor",
                sandbox=SandboxPolicy.FULL_ACCESS,
                instructions=EXECUTOR_INSTRUCTIONS,
            ),
            traces,
        )
    )
    config = TrialConfig(
        run_id="supervisor-failure",
        goal="Assess the target.",
        targets=("http://target.test",),
        backend="claude",
        model="claude-opus-4-8",
        runs_root=runs_root,
        workspace_root=tmp_path / "agents",
    )

    summary = await run_trial(config, supervisor=supervisor, executor=executor)

    assert summary["status"] == "failed"
    assert summary["error"].startswith("supervisor_contract: Supervisor failed after 2 attempts")
    assert summary["transitions"][-1]["kind"] == "supervisor_failed"


def test_cli_returns_nonzero_for_a_canonical_failed_run(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    async def failed_run(config: TrialConfig) -> dict[str, object]:
        assert config.effort == "xhigh"
        return {"status": "failed", "run_id": config.run_id}

    monkeypatch.setattr(trial, "run_trial", failed_run)

    exit_code = trial.main(
        [
            "--goal",
            "Assess the target.",
            "--target",
            "http://target.test",
            "--run-id",
            "failed-cli",
            "--effort",
            "xhigh",
        ]
    )

    assert exit_code == 1