/pentestgpt_agent/tests/test_executor.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_INSTRUCTIONS, Executor
from pentestgpt_agent.memory import (
    AttemptRecord,
    AttemptStatus,
    ObservationRecord,
    RunSnapshot,
    RunStatus,
    TaskLease,
)
from pentestgpt_agent.plan import TaskKind, TaskRecord, TaskStatus
from pentestgpt_agent.trace import EpisodeRunner, TraceStore


class CapturingExecutorBackend:
    name = "capturing"

    def __init__(self) -> None:
        self.prompt: str | None = None
        self.max_turns: int | None = None

    async def stream(self, prompt: str, opts: RunOptions) -> AsyncIterator[AgentEvent]:
        self.prompt = prompt
        self.max_turns = opts.max_turns
        yield TurnCompleted(
            success=True,
            structured_output={
                "task_id": "active-task",
                "outcome": "progress",
                "summary": "A bounded next action remains.",
                "evidence_excerpt": None,
            },
        )


def _active_snapshot(kind: TaskKind) -> tuple[RunSnapshot, TaskLease]:
    task = TaskRecord(
        id="active-task",
        kind=kind,
        target="http://target.test/input",
        objective="Assess only the named surface.",
        done_when="The named hypothesis is resolved.",
        basis_ids=(),
        depends_on=(),
        status=TaskStatus.ACTIVE,
        created_revision=1,
    )
    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,
        tasks=(task,),
    )
    return snapshot, TaskLease(
        run_id="run-1",
        task_id=task.id,
        attempt_id="attempt-1",
        revision=2,
    )


@pytest.mark.asyncio
async def test_executor_caps_test_episode_and_exposes_its_turn_budget(tmp_path: Path) -> None:
    backend = CapturingExecutorBackend()
    executor = Executor(
        EpisodeRunner(
            UnifiedAgent(
                backend,
                workspace=tmp_path / "workspace",
                sandbox=SandboxPolicy.WORKSPACE_WRITE,
                instructions=EXECUTOR_INSTRUCTIONS,
            ),
            TraceStore(tmp_path / "runs"),
        ),
        max_turns=12,
    )
    snapshot, lease = _active_snapshot(TaskKind.TEST)

    await executor.execute(snapshot, lease, episode_id="executor-1")

    assert backend.max_turns == 6
    assert backend.prompt is not None
    assert '"turn_budget": 5' in backend.prompt


@pytest.mark.asyncio
async def test_executor_retry_receives_bounded_non_evidentiary_diagnostics_and_own_evidence(
    tmp_path: Path,
) -> None:
    backend = CapturingExecutorBackend()
    executor = Executor(
        EpisodeRunner(
            UnifiedAgent(
                backend,
                workspace=tmp_path / "workspace",
                sandbox=SandboxPolicy.WORKSPACE_WRITE,
                instructions=EXECUTOR_INSTRUCTIONS,
            ),
            TraceStore(tmp_path / "runs"),
        ),
        max_turns=12,
    )
    base, lease = _active_snapshot(TaskKind.TEST)
    snapshot = RunSnapshot(
        **{
            **base.__dict__,
            "observations": (
                ObservationRecord(
                    id="obs-progress",
                    task_id="active-task",
                    attempt_id="attempt-previous",
                    statement="uid=0(root)",
                    trace_episode_id="executor-previous",
                    evidence_sequences=(7,),
                    created_revision=2,
                ),
            ),
            "attempts": (
                AttemptRecord(
                    id="attempt-previous",
                    task_id="active-task",
                    status=AttemptStatus.ERROR,
                    started_revision=1,
                    finished_revision=2,
                    trace_episode_id="executor-previous",
                    summary="Provider rejected the malformed result.",
                    failure_kind="validation",
                    failure_message="evidence was not an exact quote",
                ),
            ),
        }
    )

    await executor.execute(snapshot, lease, episode_id="executor-retry")

    assert backend.prompt is not None
    envelope = json.loads(backend.prompt.split("\n\n", 1)[1])
    assert envelope["prior_task_evidence"] == [{"id": "obs-progress", "evidence": "uid=0(root)"}]
    assert envelope["retry_diagnostic"] == {
        "status": "error",
        "failure_kind": "validation",
        "failure_message": "evidence was not an exact quote",
    }
    assert "Provider rejected the malformed result." not in backend.prompt


def test_executor_instructions_protect_the_result_turn_and_test_boundary() -> None:
    assert "Immediately return StructuredOutput when done_when is met" in EXECUTOR_INSTRUCTIONS
    assert "runtime reserves one additional transport turn" in EXECUTOR_INSTRUCTIONS
    assert "quote the complete contiguous result block" in EXECUTOR_INSTRUCTIONS
    assert "TEST never pursues or retrieves the run goal" in EXECUTOR_INSTRUCTIONS
    assert "only one command token survives and '$' is filtered" in EXECUTOR_INSTRUCTIONS
    assert "prefer shell input redirection" in EXECUTOR_INSTRUCTIONS
    assert "Tool-call timeout metadata is not an operating-system bound" in EXECUTOR_INSTRUCTIONS
    assert "Prefer userspace protocol clients over kernel filesystem mounts" in (
        EXECUTOR_INSTRUCTIONS
    )
    assert "Never issue an unprivileged kernel filesystem mount" in EXECUTOR_INSTRUCTIONS
    assert "sudo -n -l" in EXECUTOR_INSTRUCTIONS


@pytest.mark.parametrize(
    ("kind", "requested", "expected_task_turns", "expected_provider_turns"),
    (
        (TaskKind.DISCOVER, 12, 5, 6),
        (TaskKind.ENUMERATE, 12, 6, 7),
        (TaskKind.TEST, 12, 5, 6),
        (TaskKind.EXPLOIT, 12, 9, 10),
        (TaskKind.VERIFY, 12, 4, 5),
        (TaskKind.RECOVER, 12, 6, 7),
        (TaskKind.EXPLOIT, 3, 2, 3),
    ),
)
@pytest.mark.asyncio
async def test_executor_enforces_the_smaller_requested_or_per_kind_budget(
    tmp_path: Path,
    kind: TaskKind,
    requested: int,
    expected_task_turns: int,
    expected_provider_turns: int,
) -> None:
    backend = CapturingExecutorBackend()
    executor = Executor(
        EpisodeRunner(
            UnifiedAgent(
                backend,
                workspace=tmp_path / "workspace",
                sandbox=SandboxPolicy.WORKSPACE_WRITE,
                instructions=EXECUTOR_INSTRUCTIONS,
            ),
            TraceStore(tmp_path / "runs"),
        ),
        max_turns=requested,
    )
    snapshot, lease = _active_snapshot(kind)

    await executor.execute(snapshot, lease, episode_id=f"executor-{kind.value}")

    assert backend.max_turns == expected_provider_turns
    assert backend.prompt is not None
    assert f'"turn_budget": {expected_task_turns}' in backend.prompt