diff --git a/doc/code/framework.md b/doc/code/framework.md index dd67790ed2..209ae5a77e 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -289,6 +289,13 @@ If you are contributing to PyRIT, that work will most likely land in one of the - Surface sources read what a location holds for a `SurfaceScorable`, such as files under one root directory. `FileWriteScorer` evaluates its `ContentWritten` condition against that evidence. The caller owns workspace isolation; a snapshot does not identify its writer. +- Program-execution sources run a program for a `ProgramExecutionScorable` through an injected + client from the experimental `pyrit._code_execution` module, and turn its results into scoring evidence. + `ProgramExecutionScorer` splits its `ProgramBehaves` condition so the source receives only the + runtime profile and invocations, never the expected results. A program that fails is observed + evidence; a sandbox that cannot run it is an acquisition error, so the verdict is undetermined + rather than false. + Each scoring call runs afresh; an observation identifies the run and supports explicit replay. - Observation lifecycle is optional. A scorer that implements `AttemptObserver`, whether it is a root scorer or wrapped inside one, is asked by the attack to start observing after setup, once its attempt ID is allocated and before anything is sent, and to stop when the attempt ends. A @@ -362,6 +369,20 @@ The below talks about responsibilities of most modules in the PyRIT library - Components that need credentials should go through these helpers rather than handling tokens themselves. +## Code Execution (Experimental) + +**Responsibility**: Connect to caller-configured execution backends and return typed invocation results. + +- The client API is private and experimental in `pyrit._code_execution`; it may change without a deprecation cycle. +- Targets, converters, dataset generators, and scorers can inject the same `SandboxClient` contract + from `pyrit._code_execution`. Execution inputs and results are private types in `pyrit.models.code_execution`. +- `runtime_profile_name` selects a client-owned configuration by name, not by supplying an image or command. + Clients advertise the available names through `get_runtime_profile_names()`. +- PyRIT does not own environments. The one-image Docker example is the only exception: + `DockerSandboxClient` manages disposable containers; callers supply the daemon and pinned image. +- **Does not own**: interpreting results, scoring criteria, message or seed creation, memory + persistence, or infrastructure setup. Those remain with the consumer or caller. + ## [Exceptions](../contributing/9_exception) **Responsibility**: Define PyRIT's exception hierarchy and the retry behavior built around it. diff --git a/doc/code/scoring/0_scoring.ipynb b/doc/code/scoring/0_scoring.ipynb index a58a705fd4..8c1f01439b 100644 --- a/doc/code/scoring/0_scoring.ipynb +++ b/doc/code/scoring/0_scoring.ipynb @@ -275,7 +275,8 @@ "the target need not be a language model. Its kind is `scorer_target_response`.\n", "Media observation capture remains deferred until its evidence can be snapshotted.\n", "Trace-backed tool observations are covered in [Tool-call scoring](5_tool_call_scorer.ipynb), and file-system\n", - "observations in [File-write scoring](6_file_write_scorer.ipynb).\n", + "observations in [File-write scoring](6_file_write_scorer.ipynb), and sandboxed program runs in\n", + "[Program-execution scoring](7_program_execution_scorer.ipynb).\n", "\n", "Replaying a judgment is different from evaluating a stored run against a new expectation.\n", "A retained target judgment answers the original expectation; changing that expectation\n", diff --git a/doc/code/scoring/0_scoring.py b/doc/code/scoring/0_scoring.py index 0dc5d1fa2b..006d7e7052 100644 --- a/doc/code/scoring/0_scoring.py +++ b/doc/code/scoring/0_scoring.py @@ -161,7 +161,8 @@ # the target need not be a language model. Its kind is `scorer_target_response`. # Media observation capture remains deferred until its evidence can be snapshotted. # Trace-backed tool observations are covered in [Tool-call scoring](5_tool_call_scorer.ipynb), and file-system -# observations in [File-write scoring](6_file_write_scorer.ipynb). +# observations in [File-write scoring](6_file_write_scorer.ipynb), and sandboxed program runs in +# [Program-execution scoring](7_program_execution_scorer.ipynb). # # Replaying a judgment is different from evaluating a stored run against a new expectation. # A retained target judgment answers the original expectation; changing that expectation diff --git a/doc/code/scoring/7_program_execution_scorer.ipynb b/doc/code/scoring/7_program_execution_scorer.ipynb new file mode 100644 index 0000000000..2c79fe18c4 --- /dev/null +++ b/doc/code/scoring/7_program_execution_scorer.ipynb @@ -0,0 +1,175 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Program execution and scoring\n", + "\n", + "`pyrit._code_execution` is experimental; its API may change without a deprecation cycle.\n", + "Its shared data types in `pyrit.models.code_execution` are also private and experimental.\n", + "The caller provides a Docker daemon and this exact pre-pulled image:\n", + "`docker pull python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1`.\n", + "PyRIT does not start Docker or pull images. This example needs no model credentials.\n", + "Docker must support cgroup memory, swap, CPU, and process limits. Runtime images must not\n", + "declare volumes. Unsupported backends and images are rejected before execution.\n", + "\n", + "## Execute without scoring or memory\n", + "\n", + "A runtime profile name selects a trusted client configuration, not an arbitrary image.\n", + "The client runs code verbatim and returns bounded output and exit information.\n", + "Targets, converters, and generators can use this same interface." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Direct execution: 9\n" + ] + } + ], + "source": [ + "from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient\n", + "from pyrit.models.code_execution import _ProgramInvocation\n", + "\n", + "client = DockerSandboxClient(\n", + " profiles={\n", + " \"python\": DockerRuntimeProfile(\n", + " image=\"python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1\",\n", + " command=(\"python3\", \"-I\"),\n", + " )\n", + " }\n", + ")\n", + "program = \"import sys\\nprint(int(sys.argv[1]) ** 2)\\n\"\n", + "invocation = _ProgramInvocation(arguments=(\"3\",))\n", + "assert \"python\" in client.get_runtime_profile_names()\n", + "run = await client.execute_async( # type: ignore\n", + " program=program, runtime_profile_name=\"python\", invocations=(invocation,)\n", + ")\n", + "assert run.events[0].exit_code == 0\n", + "assert run.events[0].stdout == \"9\\n\"\n", + "print(f\"Direct execution: {run.events[0].stdout.strip()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## Score and save the evidence\n", + "\n", + "The source receives only code and invocation inputs; expected results stay with the scorer.\n", + "Every `score_async` call starts a fresh execution, identified by its saved observation ID.\n", + "Wrong output, unexpected exits, and enforced resource/time limits fail a case. Infrastructure\n", + "failures or insufficient evidence leave the verdict undetermined." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Score: True - All 1 case(s) produced the expected result under python.\n" + ] + } + ], + "source": [ + "from pyrit.memory import CentralMemory\n", + "from pyrit.models import ContentScorable, Equals, ProgramBehaves, ProgramCase, ProgramResult, ScoringExpectation\n", + "from pyrit.score import ProgramExecutionScorer, ProgramExecutionSource\n", + "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", + "\n", + "await initialize_pyrit_async( # type: ignore\n", + " memory_db_type=IN_MEMORY, load_defaults=False, env_files=[], silent=True\n", + ")\n", + "case = ProgramCase(invocation=invocation, expected=ProgramResult(stdout=Equals(value=\"9\")))\n", + "expectation = ScoringExpectation(conditions=(ProgramBehaves(runtime_profile_name=\"python\", cases=(case,)),))\n", + "scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=client))\n", + "score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=expectation))[0] # type: ignore\n", + "assert score.get_value() is True\n", + "print(f\"Score: {score.get_value()} - {score.score_rationale}\")" + ] + }, + { + "cell_type": "markdown", + "id": "4", + "metadata": {}, + "source": [ + "## Replay without another execution\n", + "\n", + "Replay uses the saved outputs, not the Docker client. Expected results may change, but the\n", + "runtime profile name and invocation inputs must still match the saved run." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Replay: False - The program did not produce the expected result under python: case 0 output did not match.\n" + ] + } + ], + "source": [ + "memory = CentralMemory.get_memory_instance()\n", + "observation = (await memory.get_observations_async(observation_ids=score.observation_ids))[0]\n", + "different_case = ProgramCase(invocation=invocation, expected=ProgramResult(stdout=Equals(value=\"8\")))\n", + "different_expectation = ScoringExpectation(\n", + " conditions=(ProgramBehaves(runtime_profile_name=\"python\", cases=(different_case,)),)\n", + ")\n", + "replayed = (\n", + " await scorer.score_observation_async(observation=observation, expectation=different_expectation) # type: ignore\n", + ")[0]\n", + "assert replayed.get_value() is False\n", + "assert replayed.observation_ids == score.observation_ids\n", + "print(f\"Replay: {replayed.get_value()} - {replayed.score_rationale}\")" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "Real execution and fault coverage is in `tests/integration/score/test_program_execution_scorer_integration.py`.\n", + "Opt in with `RUN_ALL_TESTS=true`; the tests require the daemon and pinned image above." + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.2" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/doc/code/scoring/7_program_execution_scorer.py b/doc/code/scoring/7_program_execution_scorer.py new file mode 100644 index 0000000000..12775a87ec --- /dev/null +++ b/doc/code/scoring/7_program_execution_scorer.py @@ -0,0 +1,96 @@ +# --- +# jupyter: +# jupytext: +# cell_metadata_filter: -all +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# --- + +# %% [markdown] +# # Program execution and scoring +# +# `pyrit._code_execution` is experimental; its API may change without a deprecation cycle. +# Its shared data types in `pyrit.models.code_execution` are also private and experimental. +# The caller provides a Docker daemon and this exact pre-pulled image: +# `docker pull python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1`. +# PyRIT does not start Docker or pull images. This example needs no model credentials. +# Docker must support cgroup memory, swap, CPU, and process limits. Runtime images must not +# declare volumes. Unsupported backends and images are rejected before execution. +# +# ## Execute without scoring or memory +# +# A runtime profile name selects a trusted client configuration, not an arbitrary image. +# The client runs code verbatim and returns bounded output and exit information. +# Targets, converters, and generators can use this same interface. + +# %% +from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient +from pyrit.models.code_execution import _ProgramInvocation + +client = DockerSandboxClient( + profiles={ + "python": DockerRuntimeProfile( + image="python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1", + command=("python3", "-I"), + ) + } +) +program = "import sys\nprint(int(sys.argv[1]) ** 2)\n" +invocation = _ProgramInvocation(arguments=("3",)) +assert "python" in client.get_runtime_profile_names() +run = await client.execute_async( # type: ignore + program=program, runtime_profile_name="python", invocations=(invocation,) +) +assert run.events[0].exit_code == 0 +assert run.events[0].stdout == "9\n" +print(f"Direct execution: {run.events[0].stdout.strip()}") + +# %% [markdown] +# ## Score and save the evidence +# +# The source receives only code and invocation inputs; expected results stay with the scorer. +# Every `score_async` call starts a fresh execution, identified by its saved observation ID. +# Wrong output, unexpected exits, and enforced resource/time limits fail a case. Infrastructure +# failures or insufficient evidence leave the verdict undetermined. + +# %% +from pyrit.memory import CentralMemory +from pyrit.models import ContentScorable, Equals, ProgramBehaves, ProgramCase, ProgramResult, ScoringExpectation +from pyrit.score import ProgramExecutionScorer, ProgramExecutionSource +from pyrit.setup import IN_MEMORY, initialize_pyrit_async + +await initialize_pyrit_async( # type: ignore + memory_db_type=IN_MEMORY, load_defaults=False, env_files=[], silent=True +) +case = ProgramCase(invocation=invocation, expected=ProgramResult(stdout=Equals(value="9"))) +expectation = ScoringExpectation(conditions=(ProgramBehaves(runtime_profile_name="python", cases=(case,)),)) +scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=client)) +score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=expectation))[0] # type: ignore +assert score.get_value() is True +print(f"Score: {score.get_value()} - {score.score_rationale}") + +# %% [markdown] +# ## Replay without another execution +# +# Replay uses the saved outputs, not the Docker client. Expected results may change, but the +# runtime profile name and invocation inputs must still match the saved run. + +# %% +memory = CentralMemory.get_memory_instance() +observation = (await memory.get_observations_async(observation_ids=score.observation_ids))[0] +different_case = ProgramCase(invocation=invocation, expected=ProgramResult(stdout=Equals(value="8"))) +different_expectation = ScoringExpectation( + conditions=(ProgramBehaves(runtime_profile_name="python", cases=(different_case,)),) +) +replayed = ( + await scorer.score_observation_async(observation=observation, expectation=different_expectation) # type: ignore +)[0] +assert replayed.get_value() is False +assert replayed.observation_ids == score.observation_ids +print(f"Replay: {replayed.get_value()} - {replayed.score_rationale}") + +# %% [markdown] +# Real execution and fault coverage is in `tests/integration/score/test_program_execution_scorer_integration.py`. +# Opt in with `RUN_ALL_TESTS=true`; the tests require the daemon and pinned image above. diff --git a/doc/myst.yml b/doc/myst.yml index 6c1783749a..8f983dd9ac 100644 --- a/doc/myst.yml +++ b/doc/myst.yml @@ -162,6 +162,7 @@ project: - file: code/scoring/4_scorer_metrics.ipynb - file: code/scoring/5_tool_call_scorer.ipynb - file: code/scoring/6_file_write_scorer.ipynb + - file: code/scoring/7_program_execution_scorer.ipynb - file: code/memory/0_memory.md children: - file: code/memory/1_sqlite_memory.ipynb diff --git a/pyrit/_code_execution/__init__.py b/pyrit/_code_execution/__init__.py new file mode 100644 index 0000000000..84ed3c54a7 --- /dev/null +++ b/pyrit/_code_execution/__init__.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +# ruff: noqa: F401 + +""" +Experimental clients for caller-configured code-execution backends. + +This private package is an explicit opt-in. Its API may change without a deprecation cycle. +""" + +from typing import TYPE_CHECKING + +from pyrit.common.lazy_imports import get_lazy_dir, resolve_lazy_export + +if TYPE_CHECKING: + from pyrit._code_execution.client import SandboxClient, SandboxExecutionError + from pyrit._code_execution.docker_client import DockerRuntimeProfile, DockerSandboxClient + +_LAZY_EXPORTS: dict[str, str] = { + "SandboxClient": "pyrit._code_execution.client", + "SandboxExecutionError": "pyrit._code_execution.client", + "DockerRuntimeProfile": "pyrit._code_execution.docker_client", + "DockerSandboxClient": "pyrit._code_execution.docker_client", +} + +__all__ = list(_LAZY_EXPORTS) + + +def __getattr__(name: str) -> object: + return resolve_lazy_export( + name=name, + module_name=__name__, + module_globals=globals(), + exports=_LAZY_EXPORTS, + ) + + +def __dir__() -> list[str]: + return get_lazy_dir(module_globals=globals(), exports=_LAZY_EXPORTS) diff --git a/pyrit/_code_execution/client.py b/pyrit/_code_execution/client.py new file mode 100644 index 0000000000..ff1fdf51f8 --- /dev/null +++ b/pyrit/_code_execution/client.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""The experimental contract for connecting to an isolated code-execution backend.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from pyrit.models import ComponentIdentifier + from pyrit.models.code_execution import _ProgramInvocation, _RuntimeProfileName, _SandboxRun + + +class SandboxExecutionError(Exception): + """ + The sandbox could not run the program. + + This is an infrastructure failure, such as an unreachable service, a failed + authentication, or an environment that could not be provisioned. It is never raised for a + program that fails, times out, or exceeds a limit: those are observed behavior. + """ + + +@runtime_checkable +class SandboxClient(Protocol): + """ + Connect to a caller-configured backend that runs code in isolation. + + Each invocation runs in a fresh environment. The client requests execution and cleanup, + and waits for cleanup before a cancelled call returns. It does not set up the backend + service, build images, or retain environments between calls. + A runtime profile is a trusted name the client maps to its own configuration; the program + cannot choose an image or a command. The caller supplies the running backend and trusted + runtimes, and decides how to use the returned results. + """ + + def get_identifier(self) -> ComponentIdentifier: + """Return the nonsecret client configuration.""" + ... + + def get_runtime_profile_names(self) -> frozenset[_RuntimeProfileName]: + """Return the names that select this client's configured runtimes.""" + ... + + async def execute_async( + self, *, program: str, runtime_profile_name: _RuntimeProfileName, invocations: tuple[_ProgramInvocation, ...] + ) -> _SandboxRun: + """ + Run each invocation of the program once, in order, in a fresh isolated environment. + + Args: + program (str): The code to run, without extracting it from surrounding prose. + runtime_profile_name (_RuntimeProfileName): One of the names returned by + ``get_runtime_profile_names``. The client owns the configuration selected by this name. + invocations (tuple[_ProgramInvocation, ...]): A nonempty sequence of arguments and input. + + Returns: + _SandboxRun: The invocation results and any execution gaps or cleanup warnings. + + Raises: + ValueError: If the runtime profile is unsupported or there are no invocations. + SandboxExecutionError: If the sandbox could not run the program at all. + """ + ... diff --git a/pyrit/_code_execution/docker_client.py b/pyrit/_code_execution/docker_client.py new file mode 100644 index 0000000000..67ec96763b --- /dev/null +++ b/pyrit/_code_execution/docker_client.py @@ -0,0 +1,749 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Run programs in disposable, locked-down Docker containers.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +import re +import shutil +import stat +import tempfile +import time +import uuid +from dataclasses import dataclass, replace +from pathlib import Path +from typing import TYPE_CHECKING, TypeVar + +from pyrit._code_execution.client import SandboxExecutionError +from pyrit.common.task_utils import gather_with_cleanup_async +from pyrit.models import ComponentIdentifier +from pyrit.models.code_execution import _ProgramExecutionEvent, _SandboxRun + +if TYPE_CHECKING: + from collections.abc import Awaitable, Mapping + + from pyrit.models.code_execution import _ProgramInvocation, _RuntimeProfileName + +logger = logging.getLogger(__name__) + +_DIGEST_PATTERN = re.compile(r"@(sha256:[0-9a-f]{64})$") +_NEVER_STARTED = "0001-01-01T00:00:00Z" + +_T = TypeVar("_T") + + +async def _finish_work_async(work: Awaitable[_T]) -> _T: + """ + Drain owned file work or cleanup even if the caller is cancelled repeatedly. + + Returns: + _T: The completed work's result. + + Raises: + CancelledError: Once the work has ended, if the caller was cancelled. + """ + task = asyncio.ensure_future(work) + cancellation: asyncio.CancelledError | None = None + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError as error: + cancellation = cancellation or error + except Exception: + break + try: + result = task.result() + except Exception as error: + if cancellation is not None: + raise cancellation from error + raise + if cancellation is not None: + raise cancellation + return result + + +async def _close_stdin_async(stream: asyncio.StreamWriter) -> None: + """ + Close input without letting cancellation poison the writer's close notification. + + Raises: + CancelledError: After the buffered writes have been aborted and closure has finished. + """ + stream.close() + closing = asyncio.create_task(stream.wait_closed()) + try: + await asyncio.shield(closing) + except asyncio.CancelledError: + stream.transport.abort() + with contextlib.suppress(BrokenPipeError, ConnectionResetError): + await _finish_work_async(closing) + raise + except (BrokenPipeError, ConnectionResetError): + pass + + +@dataclass(frozen=True) +class DockerRuntimeProfile: + """ + Configure a trusted, digest-pinned image and its program interpreter. + + The command is followed by the program's path and the invocation's arguments. + """ + + image: str + command: tuple[str, ...] + program_filename: str = "main.py" + + def __post_init__(self) -> None: + """ + Validate the image pin and the program file name. + + Raises: + ValueError: If the image is not pinned by digest, the command is empty, or the + file name is not a plain name. + """ + if _DIGEST_PATTERN.search(self.image) is None: + raise ValueError(f"Runtime image {self.image!r} must be pinned by digest (name@sha256:<64 hex>).") + if not self.command or any(not part for part in self.command): + raise ValueError("A runtime profile command must be a nonempty argument list.") + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", self.program_filename): + raise ValueError(f"Program file name {self.program_filename!r} must be a plain file name.") + + @property + def image_digest(self) -> str: + """The ``sha256:`` digest the image is pinned to; ``__post_init__`` guarantees its presence.""" + return self.image.rsplit("@", 1)[1] + + +@dataclass(frozen=True) +class _ContainerState: + status: str + exit_code: int | None + oom_killed: bool + + +@dataclass(frozen=True) +class _ClientOutput: + stdout: bytes + stdout_truncated: bool + stderr: bytes + stderr_truncated: bool + + +@dataclass(frozen=True) +class _InvocationResult: + output: _ClientOutput + container: _ContainerState + timed_out: bool = False + + +class DockerSandboxClient: + """ + Run each invocation in a disposable container on a caller-provided Docker daemon. + + Profiles require pre-pulled, pinned images with no declared volumes. The daemon must support + cgroup memory, swap, CPU, and process limits; unsupported backends are rejected before execution. + Containers have no network or privileges, run as an unprivileged user, and have memory, CPU, + and process limits. The root and program are read-only; ``/tmp`` is a bounded tmpfs mount. + The profile replaces the image's entrypoint. + Daemon logging is disabled, and each stream retains at most ``max_output_bytes`` bytes. + The program mount sits inside a private host directory that is never mounted. + + One deadline covers startup, output capture, and the attached client's exit. The container's + state distinguishes program termination from client failure. Cleanup removes each container + before another invocation starts. + + This is not a hardened multi-tenant sandbox. Use a host dedicated to evaluation. + """ + + _READ_CHUNK_BYTES = 1 << 16 + _CONTROL_TIMEOUT_SECONDS = 30.0 + + def __init__( + self, + *, + profiles: Mapping[_RuntimeProfileName, DockerRuntimeProfile], + docker_executable: str = "docker", + timeout_seconds: float = 10.0, + memory: str = "256m", + cpus: float = 1.0, + pids_limit: int = 64, + max_output_bytes: int = 64_000, + ) -> None: + """ + Initialize with the allowlisted runtime profiles and per-container limits. + + Args: + profiles (Mapping[_RuntimeProfileName, DockerRuntimeProfile]): Names and their trusted + runtime configurations. + docker_executable (str): The Docker command line client to call. + timeout_seconds (float): Wall-clock limit for each invocation. + memory (str): Memory limit for each container, in Docker's notation, such as ``256m``. + cpus (float): CPU limit for each container. + pids_limit (int): Most processes a container may hold. + max_output_bytes (int): Most bytes of each stream kept per invocation. + + Raises: + ValueError: If no profile is given, a profile name is invalid, or a limit is not positive. + """ + if not profiles: + raise ValueError("DockerSandboxClient requires at least one runtime profile.") + invalid = [name for name in profiles if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}", name)] + if invalid: + raise ValueError(f"Invalid runtime profile names: {invalid}.") + if timeout_seconds <= 0 or cpus <= 0 or pids_limit < 1 or max_output_bytes < 1: + raise ValueError("timeout_seconds, cpus, pids_limit and max_output_bytes must be positive.") + if not re.fullmatch(r"[1-9][0-9]*[bkmg]?", memory): + raise ValueError(f"Memory limit {memory!r} must look like 256m or 1g.") + self._profiles = dict(profiles) + self._docker = docker_executable + self._timeout_seconds = timeout_seconds + self._memory = memory + self._cpus = cpus + self._pids_limit = pids_limit + self._max_output_bytes = max_output_bytes + self._temp_root = Path(tempfile.gettempdir()) + + def get_identifier(self) -> ComponentIdentifier: + """ + Return the profiles and limits. + + Returns: + ComponentIdentifier: The nonsecret client configuration. + """ + return ComponentIdentifier.of( + self, + params={ + "client_version": 3, + "profiles": { + name: { + "image": profile.image, + "command": list(profile.command), + "program_filename": profile.program_filename, + } + for name, profile in sorted(self._profiles.items()) + }, + "timeout_seconds": self._timeout_seconds, + "memory": self._memory, + "cpus": self._cpus, + "pids_limit": self._pids_limit, + "max_output_bytes": self._max_output_bytes, + }, + ) + + def get_runtime_profile_names(self) -> frozenset[_RuntimeProfileName]: + """ + Return the configured runtime profile names. + + Returns: + frozenset[_RuntimeProfileName]: Names that select the configured runtimes. + """ + return frozenset(self._profiles) + + async def execute_async( + self, *, program: str, runtime_profile_name: _RuntimeProfileName, invocations: tuple[_ProgramInvocation, ...] + ) -> _SandboxRun: + """ + Run each invocation in a fresh container, in order. + + Args: + program (str): The program text, written to a read-only file in the container. + runtime_profile_name (_RuntimeProfileName): A key from ``profiles``, also returned by + ``get_runtime_profile_names``. The key selects a trusted image and command. + invocations (tuple[_ProgramInvocation, ...]): The arguments and input of each run. + + Returns: + _SandboxRun: One event per invocation that ran. If the sandbox fails after some + invocations ran, the run holds those events and the reason it stopped. + + Raises: + ValueError: If the runtime profile is not configured or there are no invocations. + SandboxExecutionError: If the backend cannot enforce the limits, the image declares + volumes, or no invocation could be run. + """ + profile = self._profiles.get(runtime_profile_name) + if profile is None: + raise ValueError(f"Runtime profile name {runtime_profile_name!r} is not configured.") + if not invocations: + raise ValueError("At least one program invocation is required.") + await self._validate_backend_async(profile=profile) + workspace_dir = self._temp_root / f"pyrit-program-{uuid.uuid4().hex}" + program_dir = workspace_dir / "program" + program_path = program_dir / profile.program_filename + try: + await _finish_work_async( + asyncio.to_thread( + self._write_program, workspace_dir=workspace_dir, program_path=program_path, program=program + ) + ) + events: list[_ProgramExecutionEvent] = [] + for index, invocation in enumerate(invocations): + try: + event, cleanup_warning = await self._run_invocation_async( + invocation_index=index, + runtime_profile_name=runtime_profile_name, + profile=profile, + program_dir=program_dir, + invocation=invocation, + ) + except SandboxExecutionError as error: + if not events: + raise + return _SandboxRun(events=tuple(events), reasons=(f"sandbox error at invocation {index}: {error}",)) + events.append(event) + if cleanup_warning is not None: + # The environment may still be running, so no further invocation starts beside it. + remaining = len(invocations) - len(events) + reasons = (f"stopped after invocation {index}: {cleanup_warning}",) if remaining else () + return _SandboxRun(events=tuple(events), reasons=reasons, warnings=(cleanup_warning,)) + return _SandboxRun(events=tuple(events)) + finally: + await _finish_work_async( + asyncio.to_thread(self._remove_workspace, workspace_dir=workspace_dir, program_path=program_path) + ) + + async def _validate_backend_async(self, *, profile: DockerRuntimeProfile) -> None: + """ + Reject unsupported resource controls and image-declared writable volumes. + + Raises: + SandboxExecutionError: If the backend or image cannot meet the isolation requirements. + """ + status, output = await self._control_async( + "info", + "--format", + "{{.CgroupDriver}} {{.MemoryLimit}} {{.SwapLimit}} {{.CPUCfsPeriod}} {{.CPUCfsQuota}} {{.PidsLimit}}", + ) + if status != 0: + raise SandboxExecutionError(f"could not inspect Docker resource support: {output.strip()[:300]}") + fields = output.split() + if len(fields) != 6 or fields[0] not in ("cgroupfs", "systemd") or fields[1:] != ["true"] * 5: + raise SandboxExecutionError( + "Docker must support cgroup memory, swap, CPU, and process limits; " + f"check controller delegation on rootless hosts (reported: {output.strip()[:300]})." + ) + status, output = await self._control_async( + "image", "inspect", "--format", "{{json .Config.Volumes}}", profile.image + ) + if status != 0: + raise SandboxExecutionError(f"could not inspect the runtime image: {output.strip()[:300]}") + if output.strip() not in ("null", "{}"): + raise SandboxExecutionError( + "Runtime images must not declare volumes, which can bypass the bounded tmpfs storage policy." + ) + + @staticmethod + def _write_program(*, workspace_dir: Path, program_path: Path, program: str) -> None: + """Keep the host workspace private while making its inner program mount readable.""" + workspace_dir.mkdir(mode=0o700) + program_path.parent.mkdir(mode=0o700) + program_path.write_bytes(program.encode("utf-8")) + program_path.chmod(0o444) + program_path.parent.chmod(0o755) + + @staticmethod + def _remove_workspace(*, workspace_dir: Path, program_path: Path) -> None: + """Restore file write permission for Windows and remove the entire private workspace.""" + with contextlib.suppress(OSError): + program_path.chmod(stat.S_IREAD | stat.S_IWRITE) + shutil.rmtree(workspace_dir, ignore_errors=True) + if workspace_dir.exists(): + logger.warning("Could not remove the program workspace %s", workspace_dir) + + async def _run_invocation_async( + self, + *, + invocation_index: int, + runtime_profile_name: _RuntimeProfileName, + profile: DockerRuntimeProfile, + program_dir: Path, + invocation: _ProgramInvocation, + ) -> tuple[_ProgramExecutionEvent, str | None]: + """ + Run one invocation in a new container and remove the container afterwards. + + Returns: + tuple[_ProgramExecutionEvent, str | None]: What the invocation did, and a warning when + the container could not be removed. + + Raises: + SandboxExecutionError: If Docker cannot start or inspect the container, or the + container had not finished when the attached client exited. + """ + name = f"pyrit-sandbox-{uuid.uuid4().hex}" + resolved_command = (*profile.command, f"/program/{profile.program_filename}", *invocation.arguments) + argv = self._run_arguments( + name=name, profile=profile, program_dir=program_dir, resolved_command=resolved_command + ) + started = time.monotonic() + deadline = started + self._timeout_seconds + startup = asyncio.create_task(self._start_client_async(argv=argv)) + cleanup_warning: str | None = None + try: + try: + process = await asyncio.wait_for(asyncio.shield(startup), timeout=self._timeout_seconds) + except TimeoutError as error: + raise SandboxExecutionError("the Docker client did not finish starting before its deadline") from error + result = await self._wait_for_invocation_async( + process=process, name=name, stdin=invocation.stdin, deadline=deadline + ) + duration_ms = int((time.monotonic() - started) * 1000) + finally: + cleanup_warning = await _finish_work_async(self._cleanup_invocation_async(startup=startup, name=name)) + output, container = result.output, result.container + event = _ProgramExecutionEvent( + invocation_index=invocation_index, + runtime_profile_name=runtime_profile_name, + runtime_image_digest=profile.image_digest, + resolved_command=resolved_command, + exit_code=None if (result.timed_out or container.oom_killed) else container.exit_code, + stdout=output.stdout.decode("utf-8", errors="replace"), + stderr=output.stderr.decode("utf-8", errors="replace"), + stdout_truncated=output.stdout_truncated, + stderr_truncated=output.stderr_truncated, + timed_out=result.timed_out, + resource_limit_exceeded=container.oom_killed and not result.timed_out, + duration_ms=duration_ms, + ) + return event, cleanup_warning + + async def _start_client_async(self, *, argv: list[str]) -> asyncio.subprocess.Process: + """ + Start the attached Docker client for an invocation that already owns cleanup. + + Returns: + asyncio.subprocess.Process: The started client. + + Raises: + SandboxExecutionError: If the client cannot be started. + """ + try: + return await asyncio.create_subprocess_exec( + *argv, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + except OSError as error: + raise SandboxExecutionError(f"could not start {self._docker!r}: {error}") from error + + async def _cleanup_invocation_async( + self, *, startup: asyncio.Task[asyncio.subprocess.Process], name: str + ) -> str | None: + """ + Wait for startup before reaping the client and removing its container. + + Returns: + str | None: A warning if cleanup could not be confirmed. + """ + process: asyncio.subprocess.Process | None = None + try: + process = await startup + except SandboxExecutionError: + status, output = await self._control_async( + "container", "ls", "--all", "--filter", f"name=^/{name}$", "--format", "{{.Names}}" + ) + if status != 0: + warning = f"could not verify cleanup of container {name}: {output.strip()[:300]}" + logger.warning("Sandbox cleanup failed: %s", warning) + return warning + if not output.strip(): + return None + return await self._remove_container_async(process=process, name=name) + + async def _remove_container_async(self, *, process: asyncio.subprocess.Process | None, name: str) -> str | None: + """ + Reap the attached client and remove its container. + + Returns: + str | None: A warning if removal failed, otherwise None. + """ + warnings = [] + if process is not None: + try: + self._kill_client(process) + await process.communicate() + except OSError as error: + warnings.append(f"could not reap the Docker client for container {name}: {error}") + status, output = await self._control_async("rm", "--force", "--volumes", name) + if status != 0: + detail = output.strip()[:300] or "unknown Docker failure" + warnings.append(f"could not remove container {name}: {detail}") + if not warnings: + return None + warning = "; ".join(warnings) + logger.warning("Sandbox cleanup failed: %s", warning) + return warning + + def _run_arguments( + self, *, name: str, profile: DockerRuntimeProfile, program_dir: Path, resolved_command: tuple[str, ...] + ) -> list[str]: + """ + Build the ``docker run`` command line for one invocation. + + Returns: + list[str]: The arguments, starting with the Docker client. + """ + return [ + self._docker, + "run", + "--name", + name, + "--interactive", + "--pull", + "never", + "--log-driver", + "none", + "--network", + "none", + "--read-only", + "--cap-drop", + "ALL", + "--security-opt", + "no-new-privileges", + "--user", + "65534:65534", + "--memory", + self._memory, + "--memory-swap", + self._memory, + "--cpus", + str(self._cpus), + "--pids-limit", + str(self._pids_limit), + "--tmpfs", + "/tmp:rw,nosuid,nodev,size=16m", + "--workdir", + "/tmp", + "--mount", + f"type=bind,source={os.fspath(program_dir)},target=/program,readonly", + "--entrypoint", + resolved_command[0], + profile.image, + *resolved_command[1:], + ] + + async def _wait_for_invocation_async( + self, *, process: asyncio.subprocess.Process, name: str, stdin: str | None, deadline: float + ) -> _InvocationResult: + """ + Apply one deadline and distinguish a client failure from the program's termination. + + Returns: + _InvocationResult: The captured output and the container's termination facts. + + Raises: + SandboxExecutionError: If capture failed or the client and container ended inconsistently. + """ + capture = asyncio.create_task(self._collect_output_async(process=process, stdin=stdin)) + try: + try: + output = await asyncio.wait_for(asyncio.shield(capture), timeout=max(deadline - time.monotonic(), 0.0)) + except TimeoutError: + return await self._finish_timed_out_capture_async(process=process, name=name, capture=capture) + client_stopped = False + if process.returncode is None: + try: + await asyncio.wait_for(process.wait(), timeout=max(deadline - time.monotonic(), 0.0)) + except TimeoutError: + self._kill_client(process) + await process.wait() + client_stopped = True + output = replace(output, stdout_truncated=True, stderr_truncated=True) + logger.warning("Docker client for %s stalled after EOF; output completeness is unknown", name) + container = await self._inspect_async(name=name, cli_stderr=output.stderr) + if container.exit_code is None: + raise SandboxExecutionError( + f"the Docker client exited while the container was still {container.status}, " + "so the program did not finish" + ) + if not client_stopped and process.returncode != container.exit_code: + raise SandboxExecutionError( + f"the Docker client exited with {process.returncode}, but the container exited with " + f"{container.exit_code}; output capture may be incomplete" + ) + return _InvocationResult(output=output, container=container) + finally: + if not capture.done(): + capture.cancel() + await _finish_work_async(asyncio.gather(capture, return_exceptions=True)) + + async def _finish_timed_out_capture_async( + self, + *, + process: asyncio.subprocess.Process, + name: str, + capture: asyncio.Task[_ClientOutput], + ) -> _InvocationResult: + """ + Enforce the program deadline only when the container is still running. + + Returns: + _InvocationResult: A timed-out invocation with its retained output. + + Raises: + SandboxExecutionError: If the deadline expired because capture stalled after execution. + """ + container = await self._inspect_async(name=name, cli_stderr=b"") + if container.status != "running": + raise SandboxExecutionError( + f"output capture did not finish before the deadline, but the container was {container.status}; " + "a program timeout could not be established" + ) + status, detail = await self._control_async("kill", name) + if status != 0: + logger.warning("Could not kill timed-out container %s: %s", name, detail.strip()[:300]) + interrupted = not capture.done() + self._kill_client(process) + try: + output = await asyncio.wait_for(asyncio.shield(capture), timeout=self._CONTROL_TIMEOUT_SECONDS) + except TimeoutError as error: + raise SandboxExecutionError("output capture did not finish after stopping the Docker client") from error + if interrupted: + output = replace(output, stdout_truncated=True, stderr_truncated=True) + await process.wait() + container = await self._inspect_async(name=name, cli_stderr=output.stderr) + return _InvocationResult(output=output, container=container, timed_out=True) + + async def _collect_output_async(self, *, process: asyncio.subprocess.Process, stdin: str | None) -> _ClientOutput: + """ + Feed standard input and capture bounded output without managing the container. + + Returns: + _ClientOutput: The retained streams and their truncation markers. + + Raises: + SandboxExecutionError: If the client streams cannot be read. + """ + feeder = asyncio.create_task(self._feed_stdin_async(stream=process.stdin, data=stdin)) + try: + (stdout, stdout_truncated), (stderr, stderr_truncated) = await gather_with_cleanup_async( + (self._read_capped_async(process.stdout), self._read_capped_async(process.stderr)) + ) + await feeder + return _ClientOutput( + stdout=stdout, + stdout_truncated=stdout_truncated, + stderr=stderr, + stderr_truncated=stderr_truncated, + ) + except OSError as error: + raise SandboxExecutionError(f"could not capture Docker client output: {error}") from error + finally: + if not feeder.done(): + feeder.cancel() + await _finish_work_async(asyncio.gather(feeder, return_exceptions=True)) + + async def _inspect_async(self, *, name: str, cli_stderr: bytes) -> _ContainerState: + """ + Read the container's state without inferring anything from the attached client. + + Returns: + _ContainerState: The state, with an exit code only for a finished container. + + Raises: + SandboxExecutionError: If the container does not exist, never started, or has invalid state. + """ + returncode, output = await self._control_async( + "inspect", + "--format", + "{{.State.Status}} {{.State.ExitCode}} {{.State.OOMKilled}} {{.State.StartedAt}}", + name, + ) + fields = output.split() + if returncode != 0 or len(fields) != 4 or fields[3] == _NEVER_STARTED: + lines = [line.strip() for line in cli_stderr.decode("utf-8", errors="replace").splitlines() if line.strip()] + detail = lines[0][:500] if lines else output.strip()[:500] + raise SandboxExecutionError(f"the container did not start: {detail or 'unknown Docker failure'}") + try: + exit_code = int(fields[1]) + except ValueError as error: + raise SandboxExecutionError(f"the container reported an invalid exit code: {fields[1]!r}") from error + if fields[2] not in ("true", "false"): + raise SandboxExecutionError(f"the container reported an invalid memory-kill state: {fields[2]!r}") + return _ContainerState( + status=fields[0], exit_code=exit_code if fields[0] == "exited" else None, oom_killed=fields[2] == "true" + ) + + async def _control_async(self, *arguments: str) -> tuple[int, str]: + """ + Run a short Docker management command, bounded in time. + + Returns: + tuple[int, str]: The exit status and combined output; status -1 when it failed to run. + """ + try: + process = await asyncio.create_subprocess_exec( + self._docker, + *arguments, + stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT, + ) + except OSError as error: + logger.warning("Docker %s failed to start: %s", arguments[0], error) + return -1, str(error) + try: + output, _ = await asyncio.wait_for(process.communicate(), timeout=self._CONTROL_TIMEOUT_SECONDS) + except TimeoutError: + logger.warning("Docker %s timed out", arguments[0]) + return -1, "timed out" + finally: + # Runs on timeout and on cancellation alike, so no management command outlives its caller. + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() + await _finish_work_async(process.communicate()) + return process.returncode if process.returncode is not None else -1, output.decode("utf-8", errors="replace") + + async def _read_capped_async(self, stream: asyncio.StreamReader | None) -> tuple[bytes, bool]: + """ + Read a stream to its end, keeping at most the configured number of bytes. + + Returns: + tuple[bytes, bool]: The retained bytes and whether more were discarded. + """ + if stream is None: + return b"", False + kept = bytearray() + truncated = False + while chunk := await stream.read(self._READ_CHUNK_BYTES): + room = self._max_output_bytes - len(kept) + if len(chunk) > room: + truncated = True + kept.extend(chunk[: max(room, 0)]) + return bytes(kept), truncated + + @staticmethod + async def _feed_stdin_async(*, stream: asyncio.StreamWriter | None, data: str | None) -> None: + """ + Write the invocation's input, then close standard input so the program sees its end. + + Raises: + CancelledError: After input has been aborted and its writer closed. + """ + if stream is None: + return + try: + if data: + stream.write(data.encode("utf-8")) + await stream.drain() + except (BrokenPipeError, ConnectionResetError): + pass + except asyncio.CancelledError: + stream.transport.abort() + raise + finally: + await _close_stdin_async(stream) + + @staticmethod + def _kill_client(process: asyncio.subprocess.Process) -> None: + """Stop only this invocation's client, leaving container removal to its owner.""" + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + process.kill() diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index f3d6d98464..dbac7de10f 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -203,6 +203,11 @@ Observation, ObservationPayload, OutputMatches, + ProgramBehaves, + ProgramCase, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + ProgramResult, Regex, Scorable, ScorableUnion, @@ -279,6 +284,11 @@ "Regex": "pyrit.models.score", "TextMatcher": "pyrit.models.score", "OutputMatches": "pyrit.models.score", + "ProgramBehaves": "pyrit.models.score", + "ProgramCase": "pyrit.models.score", + "ProgramExecutionObservationPayload": "pyrit.models.score", + "ProgramExecutionScorable": "pyrit.models.score", + "ProgramResult": "pyrit.models.score", "ConversationScorable": "pyrit.models.score", "ConversationObservationPayload": "pyrit.models.score", "AllAvailableDatasetSize": "pyrit.models.scenario_dataset_size_estimate", diff --git a/pyrit/models/code_execution.py b/pyrit/models/code_execution.py new file mode 100644 index 0000000000..10a86ca0e2 --- /dev/null +++ b/pyrit/models/code_execution.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Private shared data for experimental code execution; these contracts may change.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Annotated + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +#: A trusted runtime profile name, such as ``python-3.12-restricted``. It names a sandbox +#: configuration the caller allowlisted; it is never an image name or a command. +_RuntimeProfileName = Annotated[str, Field(min_length=1, max_length=64, pattern=r"^[A-Za-z0-9][A-Za-z0-9._-]*$")] + + +class _ProgramInvocation(BaseModel): + """The arguments and standard input for one invocation of a program.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + arguments: tuple[Annotated[str, Field(pattern=r"^[^\x00]*$")], ...] = () + stdin: str | None = None + + +class _ProgramExecutionEvent(BaseModel): + """ + What one invocation did when the sandbox ran it. + + An event records an outcome, not a guarantee that the next execution will produce it. + A program can read the clock, use randomness, or depend on its runtime. + + Each output stream records whether its retained text is incomplete, including interrupted capture. + ``runtime_image_digest`` and ``resolved_command`` are optional provenance: a backend that cannot + expose them leaves them unset. + ``exit_code`` is ``None`` when the program did not exit on its own, for example because the + sandbox stopped it at its time or resource limit. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + invocation_index: int = Field(ge=0) + runtime_profile_name: _RuntimeProfileName + runtime_image_digest: str | None = Field(default=None, min_length=1) + resolved_command: tuple[str, ...] | None = Field(default=None, min_length=1) + exit_code: int | None + stdout: str = "" + stderr: str = "" + stdout_truncated: bool = False + stderr_truncated: bool = False + timed_out: bool = False + resource_limit_exceeded: bool = False + duration_ms: int = Field(ge=0) + + @model_validator(mode="after") + def _validate_exit_state(self) -> _ProgramExecutionEvent: + """ + Keep the exit state consistent with how the run ended. + + Returns: + _ProgramExecutionEvent: The validated event. + + Raises: + ValueError: If a run the sandbox stopped also reports an exit code of its own. + """ + if self.timed_out and self.exit_code is not None: + raise ValueError("A timed-out invocation did not exit on its own and has no exit code.") + return self + + @property + def output_truncated(self) -> bool: + """Whether either output stream could not be retained in full.""" + return self.stdout_truncated or self.stderr_truncated + + +@dataclass(frozen=True) +class _SandboxRun: + """ + Results from one request to execute a program's invocations. + + ``events`` holds one event per invocation that ran, in invocation order. ``reasons`` explains + why any invocations did not run. ``warnings`` reports cleanup problems that do not invalidate + the results already collected. + """ + + events: tuple[_ProgramExecutionEvent, ...] = () + reasons: tuple[str, ...] = () + warnings: tuple[str, ...] = () diff --git a/pyrit/models/score/__init__.py b/pyrit/models/score/__init__.py index 8cda49e477..44baf4d0e0 100644 --- a/pyrit/models/score/__init__.py +++ b/pyrit/models/score/__init__.py @@ -22,6 +22,9 @@ DivergesFromRepetition, MatchesObjective, OutputMatches, + ProgramBehaves, + ProgramCase, + ProgramResult, ToolCallRequirement, ToolsCalled, ) @@ -34,6 +37,7 @@ ConversationObservationPayload, Observation, ObservationPayload, + ProgramExecutionObservationPayload, ScorerTargetResponsePayload, SurfaceObservationPayload, ToolEventsObservationPayload, @@ -43,6 +47,7 @@ ContentScorable, ConversationScorable, MessageScorable, + ProgramExecutionScorable, Scorable, ScorableUnion, SurfaceScorable, @@ -69,6 +74,11 @@ ) _LAZY_EXPORTS: dict[str, str] = { + "ProgramBehaves": "pyrit.models.score.condition", + "ProgramCase": "pyrit.models.score.condition", + "ProgramExecutionObservationPayload": "pyrit.models.score.observation", + "ProgramExecutionScorable": "pyrit.models.score.scorable", + "ProgramResult": "pyrit.models.score.condition", "ConversationScorable": "pyrit.models.score.scorable", "ConversationObservationPayload": "pyrit.models.score.observation", "Contains": "pyrit.models.score.text_matcher", diff --git a/pyrit/models/score/condition.py b/pyrit/models/score/condition.py index 8491522d21..8ca274ed95 100644 --- a/pyrit/models/score/condition.py +++ b/pyrit/models/score/condition.py @@ -8,6 +8,10 @@ from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, SerializeAsAny, TypeAdapter, model_validator +from pyrit.models.code_execution import ( # noqa: TC001 (runtime-required by Pydantic) + _ProgramInvocation, + _RuntimeProfileName, +) from pyrit.models.score._trace_validation import ToolName # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.score.surface import SurfaceMatch # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.score.text_matcher import TextMatcher # noqa: TC001 (runtime-required by Pydantic) @@ -240,6 +244,65 @@ class ContentWritten(Condition): require_content_change: bool = False +class ProgramResult(BaseModel): + """ + The expected result of one invocation. + + ``exit_code`` defaults to 0; ``None`` accepts any exit code from a program that exited on + its own. A stream with no matcher is not checked. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + exit_code: int | None = 0 + stdout: TextMatcher | None = None + stderr: TextMatcher | None = None + + +class ProgramCase(BaseModel): + """One invocation and the result it is expected to produce.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + invocation: _ProgramInvocation = Field(default_factory=_ProgramInvocation) + expected: ProgramResult = Field(default_factory=ProgramResult) + + +class ProgramBehaves(Condition): + """ + The program produces the expected result for every case. + + The condition carries test stimuli and assertions together, because both vary by dataset + row. A scorer splits them before acquisition: the runtime profile and each invocation go to + the source, and each expected result stays with the scorer's matcher. A run that times out + or exceeds a resource limit fails its case. + """ + + condition_type: Literal["program_behaves"] = "program_behaves" + runtime_profile_name: _RuntimeProfileName + cases: tuple[ProgramCase, ...] + + @model_validator(mode="after") + def _validate_cases(self) -> ProgramBehaves: + """ + Require at least one case. + + Returns: + ProgramBehaves: The validated condition. + + Raises: + ValueError: If there are no cases. + """ + if not self.cases: + raise ValueError("ProgramBehaves requires at least one case.") + return self + + @property + def invocations(self) -> tuple[_ProgramInvocation, ...]: + """The stimuli of every case, in case order, without their expected results.""" + return tuple(case.invocation for case in self.cases) + + def _parse_conditions(value: Any) -> Any: """ Rebuild discriminator-tagged conditions without losing subclass fields. @@ -263,8 +326,11 @@ class _ConditionTupleSchema: def __get_pydantic_json_schema__(self, core_schema: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue: """Return the tuple schema with a discriminated union of concrete conditions.""" schema = handler(core_schema) + # A condition that uses one nested model in several fields is emitted as a reference to its + # definition; resolve it so every entry describes its fields inline. condition_schemas = [ - handler(TypeAdapter(condition_class).core_schema) for _, condition_class in sorted(_CONDITION_TYPES.items()) + handler.resolve_ref_schema(handler(TypeAdapter(condition_class).core_schema)) + for _, condition_class in sorted(_CONDITION_TYPES.items()) ] for condition_schema in condition_schemas: required_fields = condition_schema.setdefault("required", []) diff --git a/pyrit/models/score/observation.py b/pyrit/models/score/observation.py index 857261fc69..7a4034a58d 100644 --- a/pyrit/models/score/observation.py +++ b/pyrit/models/score/observation.py @@ -12,6 +12,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator, model_validator +from pyrit.models.code_execution import _ProgramExecutionEvent # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.identifiers.component_identifier import ( ComponentIdentifier, # noqa: TC001 (runtime-required by Pydantic field annotations) ) @@ -21,6 +22,7 @@ ContentScorable, ConversationScorable, MessageScorable, + ProgramExecutionScorable, ScorableUnion, # noqa: TC001 (runtime-required by Pydantic field annotations) SurfaceScorable, TraceScorable, @@ -419,11 +421,65 @@ def _validate_entries(self) -> SurfaceObservationPayload: return self +class ProgramExecutionObservationPayload(BaseModel): + """ + What one isolated execution attempt produced: one event per invocation that ran. + + ``program_sha256`` digests the exact code the sandbox received, so the evidence names the + program even if the message it came from is later edited. ``reasons`` explains why an + acquisition is not complete. A program that fails its invocations is still complete + evidence; only a sandbox that could not run it, or did not finish, leaves gaps. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + kind: Literal["program_execution"] = "program_execution" + schema_version: Literal[1] = 1 + program_sha256: str | None = Field(default=None, min_length=64, max_length=64, pattern=r"^[0-9a-f]{64}$") + events: tuple[_ProgramExecutionEvent, ...] = () + reasons: tuple[str, ...] = () + + @field_validator("schema_version", mode="before") + @classmethod + def _validate_schema_version(cls, value: Any) -> Any: + """ + Require the exact supported schema version, without numeric coercion. + + Returns: + Any: The supported schema version. + + Raises: + ValueError: If the version is not exactly the supported integer. + """ + if type(value) is not int or value != 1: + raise ValueError("Unsupported program execution payload schema_version; expected 1.") + return value + + @model_validator(mode="after") + def _validate_events(self) -> ProgramExecutionObservationPayload: + """ + Keep one event per invocation, in invocation order. + + Returns: + ProgramExecutionObservationPayload: The validated payload. + + Raises: + ValueError: If an invocation repeats, events are out of order, or events lack a program digest. + """ + indices = [event.invocation_index for event in self.events] + if indices != sorted(set(indices)): + raise ValueError("Program execution events must name each invocation once, in invocation order.") + if self.events and self.program_sha256 is None: + raise ValueError("Program execution events require the digest of the program that ran.") + return self + + ObservationPayload = Annotated[ ScorerTargetResponsePayload | ToolEventsObservationPayload | ConversationObservationPayload - | SurfaceObservationPayload, + | SurfaceObservationPayload + | ProgramExecutionObservationPayload, Field(discriminator="kind"), ] @@ -528,6 +584,36 @@ def _validate_surface_acquisition(self) -> Observation: raise ValueError("Unavailable or failed surface acquisition cannot contain entries.") return self + @model_validator(mode="after") + def _validate_program_acquisition(self) -> Observation: + """ + Require acquisition status to agree with the execution events. + + Returns: + Observation: The validated observation. + + Raises: + ValueError: If the anchor, events, or acquisition are inconsistent. + """ + if not isinstance(self.payload, ProgramExecutionObservationPayload): + return self + if not isinstance(self.scorable, ProgramExecutionScorable): + raise ValueError("Program execution observations require a ProgramExecutionScorable.") + events = self.payload.events + if any(event.invocation_index >= len(self.scorable.invocations) for event in events): + raise ValueError("Program execution events must belong to the scorable's invocations.") + if any(event.runtime_profile_name != self.scorable.runtime_profile_name for event in events): + raise ValueError("Program execution events must use the scorable's runtime profile.") + if self.acquisition is Acquisition.COMPLETE and ( + len(events) != len(self.scorable.invocations) or self.payload.reasons + ): + raise ValueError("Complete program acquisition requires one event per invocation and no gaps.") + if self.acquisition is Acquisition.PARTIAL and len(events) == len(self.scorable.invocations): + raise ValueError("Partial program acquisition cannot hold an event for every invocation.") + if self.acquisition in (Acquisition.UNAVAILABLE, Acquisition.ERROR) and events: + raise ValueError("Unavailable or failed program acquisition cannot contain events.") + return self + @property def response_message_piece_ids(self) -> tuple[uuid.UUID, ...]: """The ordered message references retained by this payload.""" @@ -566,7 +652,10 @@ def validate_evidence( Raises: ValueError: If scored or response evidence is missing, modified, or unsupported. """ - if isinstance(self.payload, (ToolEventsObservationPayload, SurfaceObservationPayload)): + if isinstance( + self.payload, + (ToolEventsObservationPayload, SurfaceObservationPayload, ProgramExecutionObservationPayload), + ): return if isinstance(self.payload, ConversationObservationPayload): if not isinstance(self.scorable, ConversationScorable): diff --git a/pyrit/models/score/scorable.py b/pyrit/models/score/scorable.py index e3de7f4474..d9b649e2b3 100644 --- a/pyrit/models/score/scorable.py +++ b/pyrit/models/score/scorable.py @@ -9,6 +9,10 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator +from pyrit.models.code_execution import ( # noqa: TC001 (runtime-required by Pydantic) + _ProgramInvocation, + _RuntimeProfileName, +) from pyrit.models.literals import PromptDataType # noqa: TC001 (runtime-required by Pydantic field annotations) from pyrit.models.score._trace_validation import TraceId # noqa: TC001 (runtime-required by Pydantic) from pyrit.models.score.surface import SurfaceMatch # noqa: TC001 (runtime-required by Pydantic) @@ -171,11 +175,53 @@ class SurfaceScorable(Scorable): match: SurfaceMatch = SurfaceMatch.EXACT +class ProgramExecutionScorable(Scorable): + """ + A program, its runtime, and the invocations to execute. + + This is an acquisition request, not a result: it holds no expected output, so the source + that runs it cannot see which output passes. ``program`` names the code to run: a stored + message with one text piece, loose text, or stored loose text. It is used verbatim; code is never + extracted from surrounding prose. + + Each acquisition runs afresh. The resulting observation's ID identifies that run; + replay uses the stored observation, not this input description. + """ + + scorable_type: Literal["program_execution"] = "program_execution" + program: Annotated[MessageScorable | ContentScorable | ContentEntryScorable, Field(discriminator="scorable_type")] + runtime_profile_name: _RuntimeProfileName + invocations: tuple[_ProgramInvocation, ...] + + @model_validator(mode="after") + def _validate_execution(self) -> ProgramExecutionScorable: + """ + Require text code and at least one invocation. + + Returns: + ProgramExecutionScorable: The validated scorable. + + Raises: + ValueError: If the program is not text, or there is nothing to run. + """ + if isinstance(self.program, (ContentScorable, ContentEntryScorable)) and self.program.data_type != "text": + raise ValueError("A program must be text content.") + if not self.invocations: + raise ValueError("A ProgramExecutionScorable requires at least one invocation.") + return self + + # Polymorphic union of scorables that can be stored on a Score. Every member declares a # ``scorable_type`` tag and Pydantic dispatches on it, so a new member is never mistaken for # an existing one and storage never depends on field shape. ScorableUnion = Annotated[ - MessageScorable | ContentScorable | ContentEntryScorable | ConversationScorable | TraceScorable | SurfaceScorable, + MessageScorable + | ContentScorable + | ContentEntryScorable + | ConversationScorable + | TraceScorable + | SurfaceScorable + | ProgramExecutionScorable, Field(discriminator="scorable_type"), ] diff --git a/pyrit/score/__init__.py b/pyrit/score/__init__.py index 2c6c0a47b4..48cea04304 100644 --- a/pyrit/score/__init__.py +++ b/pyrit/score/__init__.py @@ -48,6 +48,7 @@ from pyrit.score.observation.observation_source import ObservationSource from pyrit.score.observation.otel_span_exporter import InMemoryTraceExporter from pyrit.score.observation.otel_trace_source import OtelTraceSource + from pyrit.score.observation.program_execution_source import ProgramExecutionSource from pyrit.score.observation.trace_client import InMemoryTraceClient, TraceAcquisitionError, TraceClient from pyrit.score.response_handler import ( CallableResponseHandler, @@ -101,6 +102,7 @@ from pyrit.score.true_false.message_tool_call_scorer import MessageToolCallScorer from pyrit.score.true_false.otel_tool_call_scorer import OtelToolCallScorer from pyrit.score.true_false.output_matches_scorer import OutputMatchesScorer + from pyrit.score.true_false.program_execution_scorer import ProgramExecutionScorer from pyrit.score.true_false.prompt_shield_scorer import PromptShieldScorer from pyrit.score.true_false.question_answer_scorer import QuestionAnswerScorer from pyrit.score.true_false.regex.agent_threat_rules_scorer import AgentThreatRulesScorer @@ -161,6 +163,8 @@ from pyrit.score.true_false.wildguard_scorer import WildGuardScorer, render_wildguard_prompt _LAZY_EXPORTS: dict[str, str | tuple[str, str | None]] = { + "ProgramExecutionScorer": "pyrit.score.true_false.program_execution_scorer", + "ProgramExecutionSource": "pyrit.score.observation.program_execution_source", "OutputMatchesScorer": "pyrit.score.true_false.output_matches_scorer", "AnsiEscapeOutputScorer": "pyrit.score.true_false.regex.ansi_escape_output_scorer", "AnthraxKeywordScorer": "pyrit.score.true_false.regex.anthrax_keyword_scorer", diff --git a/pyrit/score/observation/__init__.py b/pyrit/score/observation/__init__.py index bbb4b94186..ddcf100056 100644 --- a/pyrit/score/observation/__init__.py +++ b/pyrit/score/observation/__init__.py @@ -16,9 +16,11 @@ from pyrit.score.observation.observation_source import ObservationSource from pyrit.score.observation.otel_span_exporter import InMemoryTraceExporter from pyrit.score.observation.otel_trace_source import OtelTraceSource + from pyrit.score.observation.program_execution_source import ProgramExecutionSource from pyrit.score.observation.trace_client import InMemoryTraceClient, TraceAcquisitionError, TraceClient _LAZY_EXPORTS: dict[str, str] = { + "ProgramExecutionSource": "pyrit.score.observation.program_execution_source", "AttemptObserver": "pyrit.score.observation.lifecycle", "ConversationSource": "pyrit.score.observation.conversation_source", "InMemoryTraceClient": "pyrit.score.observation.trace_client", diff --git a/pyrit/score/observation/execution.py b/pyrit/score/observation/execution.py index c7fd663ad3..b14b81ad97 100644 --- a/pyrit/score/observation/execution.py +++ b/pyrit/score/observation/execution.py @@ -15,6 +15,7 @@ MessagePiece, MessageScorable, Observation, + ProgramExecutionObservationPayload, Scorable, ScorableUnion, Score, @@ -38,7 +39,11 @@ class NonReplayableObservationError(ValueError): _ObservationEvidence: TypeAlias = ( - Message | ToolEventsObservationPayload | SurfaceObservationPayload | tuple[MessagePiece, ...] + Message + | ToolEventsObservationPayload + | SurfaceObservationPayload + | ProgramExecutionObservationPayload + | tuple[MessagePiece, ...] ) @@ -315,7 +320,9 @@ async def resolve_async(self, *, observation: Observation) -> _ObservationEviden NonReplayableObservationError: If referenced evidence is missing, modified, or unsupported. """ payload = observation.payload - if isinstance(payload, (ToolEventsObservationPayload, SurfaceObservationPayload)): + if isinstance( + payload, (ToolEventsObservationPayload, SurfaceObservationPayload, ProgramExecutionObservationPayload) + ): return payload pieces = await self._memory.get_message_pieces_async(prompt_ids=list(observation.evidence_message_piece_ids)) pieces_by_id = {piece.id: piece for piece in pieces} diff --git a/pyrit/score/observation/program_execution_source.py b/pyrit/score/observation/program_execution_source.py new file mode 100644 index 0000000000..878027a3af --- /dev/null +++ b/pyrit/score/observation/program_execution_source.py @@ -0,0 +1,200 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Acquire how a program behaves by asking an injected sandbox to run it.""" + +from __future__ import annotations + +import hashlib +import logging +from typing import TYPE_CHECKING + +from pyrit._code_execution import SandboxExecutionError +from pyrit.memory import CentralMemory +from pyrit.models import ( + Acquisition, + ComponentIdentifier, + ContentScorable, + MessageScorable, + Observation, + ProgramExecutionObservationPayload, +) +from pyrit.score.message_scorable_resolver import MessageScorableResolver + +if TYPE_CHECKING: + from pyrit._code_execution import SandboxClient + from pyrit.models import ProgramExecutionScorable + from pyrit.models.code_execution import _ProgramExecutionEvent + +logger = logging.getLogger(__name__) + + +class _NoProgramError(Exception): + """The named evidence holds nothing that can be submitted as a program.""" + + +class ProgramExecutionSource: + """ + Submit verbatim program text and invocation inputs to an injected ``SandboxClient``. + + Expected results never reach the source. Memory is needed only to resolve stored program + references. A message must contain exactly one text piece; blocked or failed responses + leave evidence unavailable. Code is not extracted from prose. + + Each acquisition runs afresh and returns a new observation. Program failures are execution + evidence; infrastructure failures are acquisition errors. A run that stops after some + invocations leaves partial evidence. Streams retain at most ``max_output_chars`` characters + with truncation flags, and cleanup warnings remain in ``sandbox_warnings`` metadata. + """ + + def __init__(self, *, sandbox: SandboxClient, max_output_chars: int = 64_000) -> None: + """ + Initialize with the sandbox that runs submitted programs. + + Args: + sandbox (SandboxClient): Runs programs under its allowlisted runtime profiles. + max_output_chars (int): The most characters of each stream retained per invocation. + + Raises: + ValueError: If the output limit is not positive. + """ + if max_output_chars < 1: + raise ValueError("max_output_chars must be positive.") + self._sandbox = sandbox + self._max_output_chars = max_output_chars + + def get_identifier(self) -> ComponentIdentifier: + """ + Return the acquisition version, output limit, and sandbox configuration. + + Returns: + ComponentIdentifier: The nonsecret source configuration. + """ + return ComponentIdentifier.of( + self, + params={"acquisition_version": 3, "max_output_chars": self._max_output_chars}, + children={"sandbox": self._sandbox.get_identifier()}, + ) + + async def acquire_async(self, *, scorable: ProgramExecutionScorable) -> Observation: + """ + Run every invocation afresh and return what the program did. + + Args: + scorable (ProgramExecutionScorable): The program, runtime profile, and invocations. + + Returns: + Observation: The execution evidence, with its acquisition state. + + Raises: + ValueError: If the runtime profile is not one the sandbox offers, the sandbox + returns events for invocations it was not given, or the program reference + names missing or several pieces. + """ + if scorable.runtime_profile_name not in self._sandbox.get_runtime_profile_names(): + raise ValueError( + f"Runtime profile name {scorable.runtime_profile_name!r} is not offered by the sandbox; " + f"available names: {sorted(self._sandbox.get_runtime_profile_names())}." + ) + try: + program = await self._resolve_program_async(scorable=scorable) + except _NoProgramError as error: + return self._observation(scorable=scorable, acquisition=Acquisition.UNAVAILABLE, reasons=(str(error),)) + try: + run = await self._sandbox.execute_async( + program=program, + runtime_profile_name=scorable.runtime_profile_name, + invocations=scorable.invocations, + ) + except SandboxExecutionError as error: + logger.warning("Sandbox could not run program under %s: %s", scorable.runtime_profile_name, error) + return self._observation( + scorable=scorable, acquisition=Acquisition.ERROR, reasons=(f"sandbox error: {error}",) + ) + events = tuple(self._bounded(event) for event in run.events) + if any(event.invocation_index >= len(scorable.invocations) for event in events): + raise ValueError("The sandbox returned an event for an invocation it was not given.") + program_sha256 = hashlib.sha256(program.encode("utf-8")).hexdigest() + if len(events) == len(scorable.invocations): + observation = self._observation( + scorable=scorable, acquisition=Acquisition.COMPLETE, events=events, program_sha256=program_sha256 + ) + else: + reasons = run.reasons or ("the sandbox stopped before running every invocation",) + acquisition = Acquisition.PARTIAL if events else Acquisition.ERROR + observation = self._observation( + scorable=scorable, + acquisition=acquisition, + events=events, + program_sha256=program_sha256 if events else None, + reasons=reasons, + ) + if run.warnings: + observation = observation.model_copy(update={"metadata": {"sandbox_warnings": "; ".join(run.warnings)}}) + return observation + + async def _resolve_program_async(self, *, scorable: ProgramExecutionScorable) -> str: + """ + Read the program text the scorable names. + + Returns: + str: The code to submit, verbatim. + + Raises: + _NoProgramError: If the named message holds a blocked or failed response. + ValueError: If the reference cannot be resolved to one text piece, or stored + content is missing. + """ + reference = scorable.program + if isinstance(reference, ContentScorable): + return reference.value + if isinstance(reference, MessageScorable) and len(reference.message_piece_ids) != 1: + raise ValueError("A program message reference must name exactly one message piece.") + message = await MessageScorableResolver().resolve_async( + scorable=reference, memory=CentralMemory.get_memory_instance() + ) + piece = message.get_piece() + if piece.has_error(): + raise _NoProgramError(f"the program message holds a {piece.response_error} response, not a program") + if piece.converted_value_data_type != "text": + raise ValueError(f"A program must be text, not {piece.converted_value_data_type}.") + return piece.converted_value + + def _bounded(self, event: _ProgramExecutionEvent) -> _ProgramExecutionEvent: + """ + Cut each stream to the retention limit, keeping any truncation the sandbox reported. + + Returns: + _ProgramExecutionEvent: The event as it is retained. + """ + limit = self._max_output_chars + if len(event.stdout) <= limit and len(event.stderr) <= limit: + return event + return event.model_copy( + update={ + "stdout": event.stdout[:limit], + "stderr": event.stderr[:limit], + "stdout_truncated": event.stdout_truncated or len(event.stdout) > limit, + "stderr_truncated": event.stderr_truncated or len(event.stderr) > limit, + } + ) + + def _observation( + self, + *, + scorable: ProgramExecutionScorable, + acquisition: Acquisition, + events: tuple[_ProgramExecutionEvent, ...] = (), + program_sha256: str | None = None, + reasons: tuple[str, ...] = (), + ) -> Observation: + return Observation( + source_identifier=self.get_identifier(), + acquisition=acquisition, + scorable=scorable, + payload=ProgramExecutionObservationPayload( + program_sha256=program_sha256, + events=events, + reasons=reasons, + ), + ) diff --git a/pyrit/score/true_false/program_execution_scorer.py b/pyrit/score/true_false/program_execution_scorer.py new file mode 100644 index 0000000000..90e3309e21 --- /dev/null +++ b/pyrit/score/true_false/program_execution_scorer.py @@ -0,0 +1,231 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Score whether a program behaves as expected, over evidence from an isolated execution.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from pyrit.models import ( + Acquisition, + Contains, + ContentEntryScorable, + ContentScorable, + MessageScorable, + ProgramBehaves, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + Score, + ScoreStatus, +) +from pyrit.score.observation.execution import NonReplayableObservationError, _collect_observation +from pyrit.score.text_matching import match_text +from pyrit.score.true_false.true_false_scorer import TrueFalseScorer + +if TYPE_CHECKING: + from pyrit.models import ( + ComponentIdentifier, + Observation, + ProgramResult, + Scorable, + ScoringExpectation, + TextMatcher, + ) + from pyrit.models.code_execution import _ProgramExecutionEvent + from pyrit.score.observation.execution import _ObservationEvidence + from pyrit.score.observation.observation_source import ObservationSource + + +@dataclass(frozen=True) +class _Evaluation: + value: bool | None + rationale: str + + +def _stream_matches(*, matcher: TextMatcher, text: str, truncated: bool) -> bool | None: + """ + Match one output stream, treating a cut stream as unknown unless a prefix proves a match. + + Returns: + bool | None: Whether the stream matches, or None when the retained text cannot tell. + """ + if truncated and not isinstance(matcher, Contains): + return None + if match_text(matcher=matcher, text=text): + return True + return None if truncated else False + + +def _evaluate_case(*, expected: ProgramResult, event: _ProgramExecutionEvent | None) -> _Evaluation: + """ + Evaluate one case once, keeping its explanation with the decision. + + Returns: + _Evaluation: The case's verdict and explanation. + """ + if event is None: + return _Evaluation(value=None, rationale="did not run") + if event.timed_out: + return _Evaluation(value=False, rationale="timed out") + if event.resource_limit_exceeded: + return _Evaluation(value=False, rationale="exceeded a resource limit") + if event.exit_code is None: + return _Evaluation(value=None, rationale="did not report an exit code") + if expected.exit_code is not None and event.exit_code != expected.exit_code: + return _Evaluation(value=False, rationale=f"exited with {event.exit_code}, expected {expected.exit_code}") + checks: list[bool | None] = [] + if expected.stdout is not None: + checks.append(_stream_matches(matcher=expected.stdout, text=event.stdout, truncated=event.stdout_truncated)) + if expected.stderr is not None: + checks.append(_stream_matches(matcher=expected.stderr, text=event.stderr, truncated=event.stderr_truncated)) + if False in checks: + return _Evaluation(value=False, rationale="output did not match") + if None in checks: + return _Evaluation(value=None, rationale="output was cut before the check could be settled") + return _Evaluation(value=True, rationale="") + + +def _evaluate_program( + *, condition: ProgramBehaves, payload: ProgramExecutionObservationPayload, acquisition: Acquisition +) -> _Evaluation: + """ + Index evidence once and reuse case evaluations for the verdict and rationale. + + Returns: + _Evaluation: The program's verdict and explanation. + """ + profile_name = condition.runtime_profile_name + if acquisition in (Acquisition.ERROR, Acquisition.UNAVAILABLE): + gaps = ", ".join(payload.reasons) or "no execution evidence" + return _Evaluation( + value=None, rationale=f"The program's behavior under {profile_name} could not be observed ({gaps})." + ) + events = {event.invocation_index: event for event in payload.events} + evaluations = [ + _evaluate_case(expected=case.expected, event=events.get(index)) for index, case in enumerate(condition.cases) + ] + if all(evaluation.value is True for evaluation in evaluations): + return _Evaluation( + value=True, rationale=f"All {len(evaluations)} case(s) produced the expected result under {profile_name}." + ) + summary = "; ".join( + f"case {index} {evaluation.rationale}" + for index, evaluation in enumerate(evaluations) + if evaluation.value is not True + ) + if any(evaluation.value is False for evaluation in evaluations): + return _Evaluation( + value=False, rationale=f"The program did not produce the expected result under {profile_name}: {summary}." + ) + gaps = ", ".join(payload.reasons) + suffix = f" ({gaps})" if gaps else "" + return _Evaluation( + value=None, + rationale=f"The program's behavior under {profile_name} could not be fully established: {summary}{suffix}.", + ) + + +class ProgramExecutionScorer(TrueFalseScorer): + """ + Compare execution evidence with the expected results in a ``ProgramBehaves`` condition. + + The source receives verbatim program text, a runtime profile name, and invocation inputs, + but never expected results. Each case checks its expected exit and output; a timeout or + resource limit fails the case. All cases passing yields true, any proven failure yields + false, and insufficient evidence leaves the verdict undetermined. + + Every scoring call runs afresh. ``score_observation_async`` replays saved evidence without + executing the program. + """ + + CONDITION_TYPE = ProgramBehaves + + def __init__(self, *, source: ObservationSource[ProgramExecutionScorable]) -> None: + """ + Initialize with a condition-independent, caller-configured execution source. + + Args: + source (ObservationSource[ProgramExecutionScorable]): Runs the programs scorables name. + """ + super().__init__() + self._source = source + + def _build_identifier(self) -> ComponentIdentifier: + return self._create_identifier( + params={"matching_version": 1}, + children={"source": self._source.get_identifier()}, + ) + + async def _score_scorable_async(self, *, scorable: Scorable, expectation: ScoringExpectation | None) -> list[Score]: + condition = self._get_required_condition(expectation=expectation, condition_type=ProgramBehaves) + if isinstance(scorable, ProgramExecutionScorable): + if (scorable.runtime_profile_name, scorable.invocations) != ( + condition.runtime_profile_name, + condition.invocations, + ): + raise ValueError( + "A ProgramExecutionScorable must run the runtime profile and invocations of the " + "ProgramBehaves condition." + ) + execution = scorable + elif isinstance(scorable, (MessageScorable, ContentScorable, ContentEntryScorable)): + execution = ProgramExecutionScorable( + program=scorable, + runtime_profile_name=condition.runtime_profile_name, + invocations=condition.invocations, + ) + else: + raise TypeError( + "ProgramExecutionScorer requires a MessageScorable, a ContentScorable, a ContentEntryScorable, " + "or a ProgramExecutionScorable." + ) + + observation = await self._source.acquire_async(scorable=execution) + if observation.scorable != execution: + raise ValueError("Program execution source changed the caller's evidence anchor.") + if not isinstance(observation.payload, ProgramExecutionObservationPayload): + raise ValueError("Program execution source returned incompatible evidence.") + _collect_observation(observation) + scores = self._score_observation(observation=observation, evidence=observation.payload, expectation=expectation) + for score in scores: + score.scorable = scorable + score.message_piece_id = self._piece_id_from_scorable(scorable) + return scores + + def _score_observation( + self, + *, + observation: Observation, + evidence: _ObservationEvidence, + expectation: ScoringExpectation | None, + ) -> list[Score]: + if not isinstance(evidence, ProgramExecutionObservationPayload) or not isinstance( + observation.scorable, ProgramExecutionScorable + ): + raise NonReplayableObservationError("Program behavior scoring requires a stored program execution.") + condition = self._get_required_condition(expectation=expectation, condition_type=ProgramBehaves) + execution = observation.scorable + if (execution.runtime_profile_name, execution.invocations) != ( + condition.runtime_profile_name, + condition.invocations, + ): + raise NonReplayableObservationError( + "The stored execution ran different invocations or a different runtime profile than the " + "ProgramBehaves condition." + ) + evaluation = _evaluate_program(condition=condition, payload=evidence, acquisition=observation.acquisition) + return [ + Score( + score_value=None if evaluation.value is None else str(evaluation.value).lower(), + status=ScoreStatus.UNDETERMINED if evaluation.value is None else ScoreStatus.COMPLETE, + score_type="true_false", + score_rationale=evaluation.rationale, + score_value_description="The program produced the expected result for every case.", + scorer_class_identifier=self.get_identifier(), + scorable=observation.scorable, + message_piece_id=self._piece_id_from_scorable(observation.scorable), + observation_ids=[observation.id], + ) + ] diff --git a/tests/integration/score/test_program_execution_scorer_integration.py b/tests/integration/score/test_program_execution_scorer_integration.py new file mode 100644 index 0000000000..0261fd21ea --- /dev/null +++ b/tests/integration/score/test_program_execution_scorer_integration.py @@ -0,0 +1,299 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Opt-in shared execution and scoring against a caller-provided Docker daemon and image.""" + +from __future__ import annotations + +import asyncio +import hashlib +from typing import TYPE_CHECKING +from unittest.mock import patch + +import pytest + +from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient +from pyrit.models import ( + Acquisition, + Contains, + ContentScorable, + Equals, + ProgramBehaves, + ProgramCase, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + ProgramResult, + ScoreStatus, + ScoringExpectation, +) +from pyrit.models.code_execution import _ProgramInvocation +from pyrit.score import ProgramExecutionScorer, ProgramExecutionSource + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator + from pathlib import Path + + from pyrit.memory import SQLiteMemory + +pytestmark = [pytest.mark.run_only_if_all_tests, pytest.mark.usefixtures("patch_central_database")] + +_IMAGE = "python@sha256:a6e34c598f2467ed0e9a8d349809fcd8b5c603269512df273a0bb1784edc11b1" +_PROFILE = "python-3.12" +_PROGRAM = "import sys\nn = int(sys.argv[1]) if len(sys.argv) > 1 else int(sys.stdin.read())\nprint(n * n)\n" + + +@pytest.fixture +async def docker_sandbox(tmp_path: Path) -> AsyncGenerator[DockerSandboxClient]: + client = DockerSandboxClient( + profiles={_PROFILE: DockerRuntimeProfile(image=_IMAGE, command=("python3", "-I"))}, + timeout_seconds=5, + max_output_bytes=128, + ) + client._temp_root = tmp_path + for command in (("version", "--format", "{{.Server.Version}}"), ("image", "inspect", _IMAGE)): + status, output = await client._control_async(*command) + if status != 0: + pytest.fail( + f"Live execution requires a running Docker daemon and a pre-pulled {_IMAGE}. " + f"Docker {' '.join(command)} failed: {output}" + ) + with patch.object(client, "_run_arguments", wraps=client._run_arguments) as arguments: + try: + yield client + finally: + problems = [] + for call in arguments.call_args_list: + name = call.kwargs["name"] + status, output = await client._control_async( + "container", "ls", "--all", "--filter", f"name=^/{name}$", "--format", "{{.Names}}" + ) + if status != 0: + problems.append(f"Could not verify container cleanup for {name}: {output}") + elif output.strip(): + problems.append(f"Container {name} remained after execution") + removed, detail = await client._control_async("rm", "--force", "--volumes", name) + if removed != 0: + problems.append(f"Could not remove test container {name}: {detail}") + remaining = await asyncio.to_thread(lambda: list(tmp_path.iterdir())) + for directory in remaining: + problems.append(f"Program directory remained after execution: {directory}") + if directory.name.startswith("pyrit-program-"): + await asyncio.to_thread( + client._remove_workspace, + workspace_dir=directory, + program_path=directory / "program" / "main.py", + ) + assert not problems, "\n".join(problems) + + +def _expectation(*, stdout: str = "ok") -> ScoringExpectation: + return ScoringExpectation( + conditions=( + ProgramBehaves( + runtime_profile_name=_PROFILE, + cases=(ProgramCase(expected=ProgramResult(stdout=Equals(value=stdout))),), + ), + ) + ) + + +async def test_real_client_returns_results_without_a_scorer_async(docker_sandbox: DockerSandboxClient) -> None: + run = await docker_sandbox.execute_async( + program="import sys\nprint(sys.argv[1] + ':' + sys.stdin.read())\n", + runtime_profile_name=_PROFILE, + invocations=( + _ProgramInvocation(arguments=("first",), stdin="one"), + _ProgramInvocation(arguments=("second",), stdin="two"), + ), + ) + assert [(event.invocation_index, event.exit_code, event.stdout) for event in run.events] == [ + (0, 0, "first:one\n"), + (1, 0, "second:two\n"), + ] + assert run.reasons == run.warnings == () + + +async def test_recorded_digest_matches_the_exact_executed_bytes_async( + docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory +) -> None: + program = ( + "import hashlib\nfrom pathlib import Path\nprint(hashlib.sha256(Path(__file__).read_bytes()).hexdigest())\n" + ) + digest = hashlib.sha256(program.encode("utf-8")).hexdigest() + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + score = ( + await scorer.score_async(scorable=ContentScorable(value=program), expectation=_expectation(stdout=digest)) + )[0] + assert score.get_value() is True + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.program_sha256 == digest + assert observation.payload.events[0].stdout.strip() == digest + + +async def test_real_execution_persistence_and_replay_async( + docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory +) -> None: + condition = ProgramBehaves( + runtime_profile_name=_PROFILE, + cases=( + ProgramCase( + invocation=_ProgramInvocation(arguments=("3",)), expected=ProgramResult(stdout=Equals(value="9")) + ), + ProgramCase( + invocation=_ProgramInvocation(stdin="12\n"), expected=ProgramResult(stdout=Equals(value="144")) + ), + ), + ) + expectation = ScoringExpectation(objective="Square each input", conditions=(condition,)) + execution = ProgramExecutionScorable( + program=ContentScorable(value=_PROGRAM), runtime_profile_name=_PROFILE, invocations=condition.invocations + ) + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + with patch.object(docker_sandbox, "execute_async", wraps=docker_sandbox.execute_async) as execute: + first = (await scorer.score_async(scorable=execution, expectation=expectation))[0] + second = (await scorer.score_async(scorable=execution, expectation=expectation))[0] + assert execute.await_count == 2 + assert first.observation_ids != second.observation_ids + stored = await sqlite_instance.get_scores_async(score_ids=[str(first.id), str(second.id)]) + assert len(stored) == 2 + for score in stored: + assert score.get_value() is True + assert score.status is ScoreStatus.COMPLETE + assert score.scored_expectation == expectation + assert score.scorable == execution + observation = (await sqlite_instance.get_observations_async(observation_ids=first.observation_ids))[0] + assert observation.acquisition is Acquisition.COMPLETE + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.program_sha256 == hashlib.sha256(_PROGRAM.encode("utf-8")).hexdigest() + assert [event.stdout for event in observation.payload.events] == ["9\n", "144\n"] + for event, invocation in zip(observation.payload.events, condition.invocations, strict=True): + assert event.runtime_profile_name == _PROFILE + assert event.runtime_image_digest == _IMAGE.rsplit("@", 1)[1] + assert event.resolved_command == ("python3", "-I", "/program/main.py", *invocation.arguments) + assert event.exit_code == 0 + assert not event.timed_out and not event.output_truncated + stricter = condition.model_copy( + update={ + "cases": ( + condition.cases[0], + ProgramCase( + invocation=condition.cases[1].invocation, expected=ProgramResult(stdout=Equals(value="145")) + ), + ) + } + ) + replayed = ( + await scorer.score_observation_async( + observation=observation, expectation=ScoringExpectation(conditions=(stricter,)) + ) + )[0] + assert replayed.get_value() is False + assert replayed.observation_ids == first.observation_ids + assert execute.await_count == 2 + saved_replay = (await sqlite_instance.get_scores_async(score_ids=[str(replayed.id)]))[0] + assert saved_replay.get_value() is False + assert saved_replay.observation_ids == first.observation_ids + + +@pytest.mark.parametrize(("program", "exit_code"), [("print('wrong')", 0), ("raise SystemExit(7)", 7), ("def (", 1)]) +async def test_real_program_failure_is_false_async( + docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory, program: str, exit_code: int +) -> None: + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + score = (await scorer.score_async(scorable=ContentScorable(value=program), expectation=_expectation()))[0] + assert score.status is ScoreStatus.COMPLETE + assert score.get_value() is False + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert observation.acquisition is Acquisition.COMPLETE + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events[0].exit_code == exit_code + + +async def test_real_timeout_is_false_async(docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory) -> None: + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + score = ( + await scorer.score_async(scorable=ContentScorable(value="while True:\n pass\n"), expectation=_expectation()) + )[0] + assert score.status is ScoreStatus.COMPLETE + assert score.get_value() is False + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events[0].timed_out + assert observation.payload.events[0].exit_code is None + + +async def test_real_output_is_bounded_async(docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory) -> None: + expectation = ScoringExpectation( + conditions=( + ProgramBehaves( + runtime_profile_name=_PROFILE, + cases=(ProgramCase(expected=ProgramResult(stdout=Contains(value="xxx"))),), + ), + ) + ) + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + score = ( + await scorer.score_async( + scorable=ContentScorable(value="import sys\nsys.stdout.write('x' * 8192)\nsys.stderr.write('e' * 8192)\n"), + expectation=expectation, + ) + )[0] + assert score.get_value() is True + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + event = observation.payload.events[0] + assert event.stdout == "x" * 128 and event.stdout_truncated + assert event.stderr == "e" * 128 and event.stderr_truncated + + +async def test_missing_image_is_undetermined_async( + docker_sandbox: DockerSandboxClient, sqlite_instance: SQLiteMemory +) -> None: + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + with patch.object( + docker_sandbox, + "_profiles", + {_PROFILE: DockerRuntimeProfile(image="python@sha256:" + "0" * 64, command=("python3", "-I"))}, + ): + score = (await scorer.score_async(scorable=ContentScorable(value="print('ok')"), expectation=_expectation()))[0] + assert score.status is ScoreStatus.UNDETERMINED + assert score.score_value is None + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert observation.acquisition is Acquisition.ERROR + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events == () + assert observation.payload.reasons + + +async def test_cancelling_real_execution_cleans_up_async(docker_sandbox: DockerSandboxClient) -> None: + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=docker_sandbox)) + with ( + patch.object(docker_sandbox, "_timeout_seconds", 30), + patch.object(docker_sandbox, "_run_arguments", wraps=docker_sandbox._run_arguments) as arguments, + ): + task = asyncio.create_task( + scorer.score_async(scorable=ContentScorable(value="while True:\n pass\n"), expectation=_expectation()) + ) + + async def wait_until_running_async() -> None: + while not task.done(): + if arguments.call_args is not None: + name = arguments.call_args.kwargs["name"] + status, output = await docker_sandbox._control_async( + "inspect", "--format", "{{.State.Running}}", name + ) + if status == 0 and output.strip() == "true": + return + await asyncio.sleep(0.01) + pytest.fail(f"Execution ended before cancellation: {task.result()}") + + try: + await asyncio.wait_for(wait_until_running_async(), timeout=15) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=15) + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/integration/score/test_scorer_notebooks.py b/tests/integration/score/test_scorer_notebooks.py index b8444482ac..1b0ff9db2c 100644 --- a/tests/integration/score/test_scorer_notebooks.py +++ b/tests/integration/score/test_scorer_notebooks.py @@ -17,9 +17,16 @@ @pytest.mark.parametrize( "file_name", - [file for file in os.listdir(nb_directory_path) if file.endswith(".ipynb") and file not in skipped_files], + [ + pytest.param( + file, + marks=pytest.mark.run_only_if_all_tests if file == "7_program_execution_scorer.ipynb" else (), + ) + for file in os.listdir(nb_directory_path) + if file.endswith(".ipynb") and file not in skipped_files + ], ) -def test_execute_notebooks(file_name): +def test_execute_notebooks(file_name: str) -> None: nb_path = pathlib.Path(nb_directory_path, file_name).resolve() with open(nb_path, encoding="utf-8") as f: nb = nbformat.read(f, as_version=4) diff --git a/tests/unit/_code_execution/test_docker_client.py b/tests/unit/_code_execution/test_docker_client.py new file mode 100644 index 0000000000..9f65d8219b --- /dev/null +++ b/tests/unit/_code_execution/test_docker_client.py @@ -0,0 +1,533 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Deterministic Docker fault injection; real execution is covered by integration tests.""" + +from __future__ import annotations + +import asyncio +import os +import stat +import threading +from contextlib import nullcontext +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient, SandboxExecutionError +from pyrit.models.code_execution import _ProgramInvocation, _SandboxRun + +if TYPE_CHECKING: + from pathlib import Path + +_IMAGE = "python@sha256:" + "c" * 64 +_STARTED = "2026-10-21T00:00:00Z" + + +@pytest.fixture +def client(tmp_path: Path) -> DockerSandboxClient: + instance = DockerSandboxClient( + profiles={"python": DockerRuntimeProfile(image=_IMAGE, command=("python3", "-I"))}, + docker_executable=str(tmp_path / "no-docker"), + ) + instance._temp_root = tmp_path + return instance + + +def _process(*, returncode: int | None = 0, output: bytes = b"ok", eof: bool = True) -> MagicMock: + process = MagicMock(spec=asyncio.subprocess.Process) + process.returncode, process.stdin = returncode, None + readers = [asyncio.StreamReader(), asyncio.StreamReader()] + for reader, data in zip(readers, (output, b""), strict=True): + reader.feed_data(data) + if eof: + reader.feed_eof() + process.stdout, process.stderr = readers + exited = asyncio.Event() + if returncode is not None: + exited.set() + + async def wait_async() -> int: + await exited.wait() + code = process.returncode + assert isinstance(code, int) + return code + + def kill() -> None: + process.returncode = -9 + exited.set() + for reader in readers: + reader.feed_eof() + + process.wait = AsyncMock(side_effect=wait_async) + process.kill.side_effect = kill + process.communicate = AsyncMock(return_value=(b"", b"")) + return process + + +def _control( + *, + state: str = "exited", + code: int = 0, + oom: bool = False, + removal_fails: bool = False, + capabilities: str = "cgroupfs true true true true true", + volumes: str = "null", +) -> AsyncMock: + async def command_async(*arguments: str) -> tuple[int, str]: + nonlocal state, code + if arguments[0] == "info": + return 0, capabilities + if arguments[:2] == ("image", "inspect"): + return 0, volumes + if arguments[0] == "inspect": + started = "0001-01-01T00:00:00Z" if state == "created" else _STARTED + return 0, f"{state} {code} {str(oom).lower()} {started}" + if arguments[0] == "kill": + state, code = "exited", 137 + return 0, "" + if arguments[:2] == ("container", "ls"): + return 0, "" + assert arguments[:3] == ("rm", "--force", "--volumes"), arguments + return (1, "daemon unavailable") if removal_fails else (0, "") + + return AsyncMock(side_effect=command_async) + + +async def _run_async(client: DockerSandboxClient) -> _SandboxRun: + return await client.execute_async( + program="print('ok')", runtime_profile_name="python", invocations=(_ProgramInvocation(),) + ) + + +@pytest.mark.parametrize( + "capabilities", + [ + "none true true true true true", + "cgroupfs false true true true true", + "cgroupfs true false true true true", + "systemd true true false true true", + "systemd true true true false true", + "systemd true true true true false", + "", + "unexpected response", + ], +) +async def test_unsupported_resource_controls_prevent_execution_async( + client: DockerSandboxClient, tmp_path: Path, capabilities: str +) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as launch, + patch.object(client, "_control_async", _control(capabilities=capabilities)), + pytest.raises(SandboxExecutionError, match="must support cgroup"), + ): + await _run_async(client) + launch.assert_not_awaited() + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +@pytest.mark.parametrize("volumes", ['{"/data":{}}', '{"/tmp":{}}', "invalid metadata"]) +async def test_image_volumes_prevent_execution_async(client: DockerSandboxClient, tmp_path: Path, volumes: str) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as launch, + patch.object(client, "_control_async", _control(volumes=volumes)), + pytest.raises(SandboxExecutionError, match="must not declare volumes"), + ): + await _run_async(client) + launch.assert_not_awaited() + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +async def test_systemd_controls_and_empty_volume_map_allow_execution_async(client: DockerSandboxClient) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=_process()), + patch.object(client, "_control_async", _control(capabilities="systemd true true true true true", volumes="{}")), + ): + assert (await _run_async(client)).events[0].exit_code == 0 + + +@pytest.mark.parametrize(("command", "message"), [("info", "resource support"), ("image", "runtime image")]) +async def test_preflight_failures_prevent_execution_async( + client: DockerSandboxClient, command: str, message: str +) -> None: + healthy = _control() + + async def control_async(*arguments: str) -> tuple[int, str]: + if arguments[0] == command: + return 1, "unavailable" + response: tuple[int, str] = await healthy(*arguments) + return response + + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as launch, + patch.object(client, "_control_async", side_effect=control_async), + pytest.raises(SandboxExecutionError, match=f"could not inspect.*{message}"), + ): + await _run_async(client) + launch.assert_not_awaited() + + +def test_run_arguments_enforce_runtime_and_isolation(client: DockerSandboxClient, tmp_path: Path) -> None: + command = ("python3", "-I", "/program/main.py", "argument") + argv = client._run_arguments( + name="test-container", profile=client._profiles["python"], program_dir=tmp_path, resolved_command=command + ) + for flag, value in { + "--pull": "never", + "--log-driver": "none", + "--network": "none", + "--cap-drop": "ALL", + "--security-opt": "no-new-privileges", + "--user": "65534:65534", + "--memory": "256m", + "--memory-swap": "256m", + "--cpus": "1.0", + "--pids-limit": "64", + }.items(): + assert argv[argv.index(flag) + 1] == value + assert "--read-only" in argv and "--privileged" not in argv + assert argv[argv.index("--entrypoint") + 1] == command[0] + assert argv.index("--entrypoint") < argv.index(_IMAGE) + assert argv[argv.index(_IMAGE) + 1 :] == list(command[1:]) + assert argv[argv.index("--mount") + 1].endswith(",target=/program,readonly") + + +@pytest.mark.parametrize( + ("state", "code", "oom", "client_code", "error"), + [ + ("exited", 0, False, 0, None), + ("exited", 125, False, 125, None), + ("exited", 137, True, 137, None), + ("running", 0, False, 1, "still running"), + ("created", 125, False, 125, "did not start"), + ("exited", 0, False, 1, "output capture may be incomplete"), + ], +) +async def test_container_and_client_outcomes_async( + client: DockerSandboxClient, + tmp_path: Path, + state: str, + code: int, + oom: bool, + client_code: int, + error: str | None, +) -> None: + process = _process(returncode=client_code) + control = _control(state=state, code=code, oom=oom) + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process), + patch.object(client, "_control_async", control), + patch.object(client, "_run_arguments", wraps=client._run_arguments) as arguments, + ): + if error is not None: + with pytest.raises(SandboxExecutionError, match=error): + await _run_async(client) + else: + event = (await _run_async(client)).events[0] + assert event.exit_code == (None if oom else code) + assert event.resource_limit_exceeded is oom + assert event.runtime_image_digest == "sha256:" + "c" * 64 + assert control.call_args_list[-1].args[:3] == ("rm", "--force", "--volumes") + mounted_dir = arguments.call_args.kwargs["program_dir"] + assert mounted_dir.name == "program" + assert mounted_dir.parent.parent == tmp_path + process.communicate.assert_awaited_once() + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +@pytest.mark.parametrize( + ("eof", "state", "error"), + [ + (False, "exited", "a program timeout could not be established"), + (True, "running", "still running"), + (True, "exited", None), + (False, "running", None), + ], +) +async def test_deadline_distinguishes_stalled_capture_client_and_program_async( + client: DockerSandboxClient, eof: bool, state: str, error: str | None +) -> None: + process, control = _process(returncode=None, eof=eof), _control(state=state) + with ( + patch.object(client, "_timeout_seconds", 0.05), + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process), + patch.object(client, "_control_async", control), + ): + if error is not None: + with pytest.raises(SandboxExecutionError, match=error): + await _run_async(client) + else: + event = (await _run_async(client)).events[0] + assert event.timed_out is (not eof) + assert event.exit_code == (0 if eof else None) + assert event.stdout_truncated and event.stderr_truncated + commands = [call.args[0] for call in control.call_args_list] + assert ("kill" in commands) is (state == "running" and not eof) + assert commands[-1] == "rm" + process.kill.assert_called_once() + + +@pytest.mark.parametrize("count", [1, 2]) +async def test_cleanup_failure_keeps_evidence_and_stops_further_invocations_async( + client: DockerSandboxClient, count: int +) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=_process()) as launch, + patch.object(client, "_control_async", _control(removal_fails=True)), + ): + run = await client.execute_async( + program="print('ok')", runtime_profile_name="python", invocations=(_ProgramInvocation(),) * count + ) + assert len(run.events) == 1 and run.events[0].exit_code == 0 + assert "could not remove container" in run.warnings[0] + assert bool(run.reasons) is (count > 1) + launch.assert_awaited_once() + + +async def test_failure_after_a_completed_invocation_returns_partial_results_async(client: DockerSandboxClient) -> None: + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, side_effect=[_process(), OSError("offline")]), + patch.object(client, "_control_async", _control()), + ): + run = await client.execute_async( + program="print('ok')", runtime_profile_name="python", invocations=(_ProgramInvocation(),) * 2 + ) + assert [event.invocation_index for event in run.events] == [0] + assert "invocation 1" in run.reasons[0] and "offline" in run.reasons[0] + + +@pytest.mark.parametrize("failed", [False, True]) +async def test_startup_is_owned_until_cleanup_async(client: DockerSandboxClient, tmp_path: Path, failed: bool) -> None: + process, entered, release = _process(returncode=None), asyncio.Event(), asyncio.Event() + created: list[str] = [] + healthy = _control() + + async def start_async(*args: str, **kwargs: object) -> MagicMock: + created.append(args[args.index("--name") + 1]) + entered.set() + await release.wait() + if failed: + raise OSError("stdio connection failed after container creation") + return process + + async def control_async(*arguments: str) -> tuple[int, str]: + if arguments[0] in ("info", "image"): + response: tuple[int, str] = await healthy(*arguments) + return response + if arguments[:2] == ("container", "ls"): + assert f"name=^/{created[0]}$" in arguments + return 0, created[0] + assert arguments == ("rm", "--force", "--volumes", created[0]) + return 0, "" + + with ( + patch("asyncio.create_subprocess_exec", side_effect=start_async), + patch.object(client, "_control_async", side_effect=control_async) as control, + ): + task = asyncio.create_task(_run_async(client)) + try: + await asyncio.wait_for(entered.wait(), timeout=5) + if not failed: + for _ in range(2): + task.cancel() + await asyncio.sleep(0.01) + assert not task.done() + assert all(call.args[0] in ("info", "image") for call in control.call_args_list) + release.set() + with pytest.raises(SandboxExecutionError if failed else asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5) + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + assert control.call_args.args == ("rm", "--force", "--volumes", created[0]) + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + if not failed: + process.kill.assert_called_once() + process.communicate.assert_awaited_once() + + +@pytest.mark.parametrize("newline", ["\n", "\r\n"]) +def test_program_bytes_and_readonly_file_cleanup(client: DockerSandboxClient, tmp_path: Path, newline: str) -> None: + workspace = tmp_path / "workspace" + path = workspace / "program" / "main.py" + program = "print('\u00e9')\nprint(42)\n".replace("\n", newline) + try: + client._write_program(workspace_dir=workspace, program_path=path, program=program) + assert path.read_bytes() == program.encode("utf-8") + if os.name == "posix": + assert stat.S_IMODE(workspace.stat().st_mode) == 0o700 + assert stat.S_IMODE(path.parent.stat().st_mode) == 0o755 + assert stat.S_IMODE(path.stat().st_mode) == 0o444 + finally: + client._remove_workspace(workspace_dir=workspace, program_path=path) + assert not workspace.exists() + + +async def test_partial_workspace_is_removed_after_write_failure_async( + client: DockerSandboxClient, tmp_path: Path +) -> None: + with ( + patch.object(client, "_control_async", _control()), + patch("pathlib.Path.write_bytes", side_effect=OSError("disk full")), + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as launch, + pytest.raises(OSError, match="disk full"), + ): + await _run_async(client) + launch.assert_not_awaited() + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +@pytest.mark.parametrize("stage", ["drain", "close"]) +async def test_input_cancellation_aborts_buffered_writes_async(stage: str) -> None: + stream = MagicMock(spec=asyncio.StreamWriter) + stream.transport = MagicMock(spec=asyncio.WriteTransport) + blocked, released = asyncio.Event(), asyncio.Event() + + async def drain_async() -> None: + if stage == "drain": + blocked.set() + await asyncio.Event().wait() + + async def close_async() -> None: + if stage == "close": + blocked.set() + await released.wait() + + stream.transport.abort.side_effect = released.set + stream.drain, stream.wait_closed = AsyncMock(side_effect=drain_async), AsyncMock(side_effect=close_async) + task = asyncio.create_task(DockerSandboxClient._feed_stdin_async(stream=stream, data="input")) + await asyncio.wait_for(blocked.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + stream.transport.abort.assert_called_once() + stream.wait_closed.assert_awaited_once() + + +@pytest.mark.parametrize("phase", ["write", "remove"]) +@pytest.mark.parametrize("cancel", [False, True]) +async def test_file_work_does_not_block_or_outlive_the_call_async( + client: DockerSandboxClient, tmp_path: Path, phase: str, cancel: bool +) -> None: + entered, release = threading.Event(), threading.Event() + operation = "_write_program" if phase == "write" else "_remove_workspace" + original = getattr(client, operation) + loop_thread = threading.get_ident() + + def gated(**kwargs: object) -> None: + assert threading.get_ident() != loop_thread + entered.set() + assert release.wait(5), "file work blocked the event loop" + original(**kwargs) + + with ( + patch.object(client, operation, side_effect=gated), + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=_process()), + patch.object(client, "_control_async", _control()), + ): + task = asyncio.create_task(_run_async(client)) + try: + for _ in range(200): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + if cancel: + for _ in range(2): + task.cancel() + await asyncio.sleep(0.01) + assert not task.done() + release.set() + if cancel: + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5) + else: + assert (await asyncio.wait_for(task, timeout=5)).events[0].exit_code == 0 + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + assert await asyncio.to_thread(lambda: list(tmp_path.iterdir())) == [] + + +@pytest.mark.parametrize("phase", ["container", "management"]) +async def test_repeated_cancellation_drains_process_cleanup_async(client: DockerSandboxClient, phase: str) -> None: + process = _process(returncode=None, eof=False) + running, cleaning, release = asyncio.Event(), asyncio.Event(), asyncio.Event() + calls = 0 + + async def communicate_async() -> tuple[bytes, bytes]: + nonlocal calls + calls += 1 + if phase == "management" and calls == 1: + running.set() + await asyncio.Event().wait() + cleaning.set() + await release.wait() + return b"", b"" + + async def collect_async(**kwargs: object) -> None: + running.set() + await asyncio.Event().wait() + + process.communicate = AsyncMock(side_effect=communicate_async) + control = _control() + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process), + patch.object(client, "_collect_output_async", side_effect=collect_async), + patch.object(client, "_control_async", control) if phase == "container" else nullcontext(), + ): + task = asyncio.create_task( + _run_async(client) if phase == "container" else client._control_async("inspect", "test") + ) + try: + await asyncio.wait_for(running.wait(), timeout=5) + task.cancel() + await asyncio.wait_for(cleaning.wait(), timeout=5) + task.cancel() + await asyncio.sleep(0.01) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=5) + finally: + release.set() + await asyncio.gather(task, return_exceptions=True) + process.kill.assert_called_once() + assert process.communicate.await_count == (2 if phase == "management" else 1) + if phase == "container": + removals = [call for call in control.call_args_list if call.args[0] == "rm"] + assert len(removals) == 1 + assert removals[0].args[:3] == ("rm", "--force", "--volumes") + + +@pytest.mark.parametrize("case", ["empty", "unknown", "missing-client"]) +async def test_invalid_requests_fail_before_execution_async(client: DockerSandboxClient, case: str) -> None: + with pytest.raises(SandboxExecutionError if case == "missing-client" else ValueError): + await client.execute_async( + program="print(1)", + runtime_profile_name="unknown" if case == "unknown" else "python", + invocations=() if case == "empty" else (_ProgramInvocation(),), + ) + + +@pytest.mark.parametrize( + "settings", [{"timeout_seconds": 0}, {"pids_limit": 0}, {"max_output_bytes": 0}, {"memory": "bad"}] +) +def test_invalid_limits(settings: dict[str, object]) -> None: + with pytest.raises(ValueError): + DockerSandboxClient(profiles={"python": DockerRuntimeProfile(image=_IMAGE, command=("python3",))}, **settings) # type: ignore[arg-type] + + +def test_profiles_are_explicit_trusted_configuration() -> None: + for settings in ( + {"image": "python:latest", "command": ("python3",)}, + {"image": _IMAGE, "command": ()}, + {"image": _IMAGE, "command": ("python3",), "program_filename": "../main.py"}, + ): + with pytest.raises(ValueError): + DockerRuntimeProfile(**settings) # type: ignore[arg-type] + with pytest.raises(ValueError): + DockerSandboxClient(profiles={}) + with pytest.raises(ValueError): + DockerSandboxClient(profiles={"bad name": DockerRuntimeProfile(image=_IMAGE, command=("python3",))}) diff --git a/tests/unit/common/test_lazy_package_imports.py b/tests/unit/common/test_lazy_package_imports.py index f8cb7a709a..655f73bed9 100644 --- a/tests/unit/common/test_lazy_package_imports.py +++ b/tests/unit/common/test_lazy_package_imports.py @@ -22,6 +22,18 @@ ) _LAZY_IMPORT_SPOT_CHECKS = [ + ( + "pyrit._code_execution", + "SandboxClient", + "pyrit._code_execution.client", + "pyrit._code_execution.docker_client", + ), + ( + "pyrit._code_execution", + "DockerSandboxClient", + "pyrit._code_execution.docker_client", + "pyrit.score", + ), ( "pyrit.analytics", "compute_technique_stats_async", @@ -456,6 +468,55 @@ def test_analytics_foundations_do_not_load_higher_layers() -> None: ) +def test_code_execution_is_independent_of_component_layers() -> None: + _assert_subprocess_succeeds( + """ + import importlib.util + import sys + import pyrit.models + import pyrit.models.code_execution as execution_models + from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient, SandboxClient + from pyrit.models.code_execution import ( + _ProgramExecutionEvent, _ProgramInvocation, _RuntimeProfileName, _SandboxRun, + ) + + for name in ("_ProgramExecutionEvent", "_ProgramInvocation", "_RuntimeProfileName", "_SandboxRun"): + assert name not in pyrit.models.__all__ + assert not hasattr(pyrit.models, name) + public_name = name.removeprefix("_") + assert public_name not in pyrit.models.__all__ + assert not hasattr(pyrit.models, public_name) + assert not hasattr(execution_models, public_name) + + client = DockerSandboxClient( + profiles={"python": DockerRuntimeProfile(image="python@sha256:" + "a" * 64, command=("python3",))} + ) + assert isinstance(client, SandboxClient) + assert client.get_runtime_profile_names() == frozenset({"python"}) + assert client.get_identifier().class_module == "pyrit._code_execution.docker_client" + assert importlib.util.find_spec("pyrit.code_execution") is None + event = _ProgramExecutionEvent( + invocation_index=0, runtime_profile_name="python", + exit_code=0, stdout="result", duration_ms=1, + ) + assert event.resolved_command is None + assert _SandboxRun(events=(event,)).events[0].stdout == "result" + assert _ProgramInvocation(stdin="input").stdin == "input" + + forbidden = ( + "pyrit.score", "pyrit.models.score", "pyrit.memory", "pyrit.executor", + "pyrit.prompt_target", "pyrit.converter", "pyrit.datasets", "pyrit.scenario", + "pyrit.backend", "sqlalchemy", + ) + loaded = [ + name for name in sys.modules + if any(name == prefix or name.startswith(prefix + ".") for prefix in forbidden) + ] + assert not loaded, loaded + """ + ) + + def test_observation_acquisition_and_scorer_do_not_load_otel_sdk() -> None: _assert_subprocess_succeeds( """ diff --git a/tests/unit/models/test_code_execution.py b/tests/unit/models/test_code_execution.py new file mode 100644 index 0000000000..6f3b55c64e --- /dev/null +++ b/tests/unit/models/test_code_execution.py @@ -0,0 +1,75 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Execution data is usable without scoring, messages, or memory.""" + +import pytest +from pydantic import TypeAdapter, ValidationError + +from pyrit.models.code_execution import _ProgramExecutionEvent, _ProgramInvocation, _RuntimeProfileName, _SandboxRun + + +def _event(**updates: object) -> _ProgramExecutionEvent: + return _ProgramExecutionEvent.model_validate( + { + "invocation_index": 0, + "runtime_profile_name": "python-3.12", + "resolved_command": ("python3", "-I", "/program/main.py"), + "exit_code": 0, + "duration_ms": 1, + **updates, + } + ) + + +def test_invocation_arguments_reject_nul() -> None: + with pytest.raises(ValidationError): + _ProgramInvocation(arguments=("a\x00b",)) + + +@pytest.mark.parametrize("profile", ["", "-leading-dash", "has space", "image:tag", "a/b", "x" * 65]) +def test_runtime_profile_name_rejects_invalid_names(profile: str) -> None: + with pytest.raises(ValidationError): + TypeAdapter(_RuntimeProfileName).validate_python(profile) + + +def test_timed_out_event_cannot_report_an_exit_code() -> None: + with pytest.raises(ValidationError, match="timed-out"): + _event(timed_out=True, exit_code=0) + assert _event(timed_out=True, exit_code=None).exit_code is None + + +def test_runtime_image_provenance_is_optional_but_not_blank() -> None: + assert _event().runtime_image_digest is None + assert _event(runtime_image_digest="sha256:" + "a" * 64).runtime_image_digest == "sha256:" + "a" * 64 + with pytest.raises(ValidationError): + _event(runtime_image_digest="") + + +def test_command_provenance_is_optional_but_not_empty() -> None: + assert _event(resolved_command=None).resolved_command is None + with pytest.raises(ValidationError): + _event(resolved_command=()) + + +@pytest.mark.parametrize(("stdout", "stderr"), [(False, False), (True, False), (False, True), (True, True)]) +def test_output_truncation_reports_either_stream(stdout: bool, stderr: bool) -> None: + assert _event(stdout_truncated=stdout, stderr_truncated=stderr).output_truncated is (stdout or stderr) + + +def test_invocation_index_is_nonnegative() -> None: + with pytest.raises(ValidationError): + _event(invocation_index=-1) + + +def test_sandbox_run_round_trips_shared_execution_data() -> None: + run = _SandboxRun( + events=(_event(stdout="result"),), + reasons=("the backend stopped before invocation 1",), + warnings=("cleanup was incomplete",), + ) + adapter = TypeAdapter(_SandboxRun) + serialized = adapter.dump_json(run) + assert b'"invocation_index":0' in serialized + assert b"case_index" not in serialized + assert adapter.validate_json(serialized) == run diff --git a/tests/unit/models/test_program.py b/tests/unit/models/test_program.py new file mode 100644 index 0000000000..ba06c355b1 --- /dev/null +++ b/tests/unit/models/test_program.py @@ -0,0 +1,258 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import uuid + +import pytest +from pydantic import ValidationError + +from pyrit.models import ( + Acquisition, + ComponentIdentifier, + Condition, + ContentEntryScorable, + ContentScorable, + Equals, + MessageScorable, + Observation, + ProgramBehaves, + ProgramCase, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + ProgramResult, + Score, + ScoringExpectation, + SurfaceScorable, + scorable_from_dict, +) +from pyrit.models.code_execution import _ProgramExecutionEvent, _ProgramInvocation + +_SOURCE = ComponentIdentifier(class_name="FakeSource", class_module="tests.fake") +_DIGEST = "a" * 64 +_IMAGE_DIGEST = "sha256:" + "b" * 64 + + +def _event(*, invocation_index: int = 0, profile: str = "python-3.12", **updates: object) -> _ProgramExecutionEvent: + values: dict[str, object] = { + "invocation_index": invocation_index, + "runtime_profile_name": profile, + "runtime_image_digest": _IMAGE_DIGEST, + "resolved_command": ("python3", "/program/main.py"), + "exit_code": 0, + "stdout": "ok", + "duration_ms": 5, + } + values.update(updates) + return _ProgramExecutionEvent.model_validate(values) + + +def _execution(*, invocations: int = 2, profile: str = "python-3.12") -> ProgramExecutionScorable: + return ProgramExecutionScorable( + program=ContentScorable(value="print('ok')"), + runtime_profile_name=profile, + invocations=tuple(_ProgramInvocation(arguments=(str(index),)) for index in range(invocations)), + ) + + +def _observation( + *, + scorable: ProgramExecutionScorable, + acquisition: Acquisition, + events: tuple[_ProgramExecutionEvent, ...] = (), + reasons: tuple[str, ...] = (), +) -> Observation: + return Observation( + source_identifier=_SOURCE, + acquisition=acquisition, + scorable=scorable, + payload=ProgramExecutionObservationPayload( + program_sha256=_DIGEST if events else None, events=events, reasons=reasons + ), + ) + + +def test_program_behaves_round_trips_through_condition_registry() -> None: + condition = ProgramBehaves( + runtime_profile_name="python-3.12", + cases=( + ProgramCase( + invocation=_ProgramInvocation(arguments=("3",), stdin="x"), + expected=ProgramResult(exit_code=None, stdout=Equals(value="9")), + ), + ), + ) + restored = Condition.model_validate(condition.model_dump(mode="json")) + assert restored == condition + assert ScoringExpectation(conditions=(condition,)).conditions == (condition,) + + +def test_program_behaves_invocations_drop_expected_results() -> None: + condition = ProgramBehaves( + runtime_profile_name="python-3.12", + cases=( + ProgramCase(invocation=_ProgramInvocation(arguments=("1",)), expected=ProgramResult(exit_code=3)), + ProgramCase(invocation=_ProgramInvocation(stdin="in")), + ), + ) + assert condition.invocations == (_ProgramInvocation(arguments=("1",)), _ProgramInvocation(stdin="in")) + + +def test_program_case_defaults_expect_a_clean_exit_with_no_arguments() -> None: + case = ProgramCase() + assert case.invocation == _ProgramInvocation() + assert case.expected == ProgramResult(exit_code=0) + + +def test_program_behaves_requires_a_case() -> None: + with pytest.raises(ValidationError, match="at least one case"): + ProgramBehaves(runtime_profile_name="python-3.12", cases=()) + + +def test_condition_validates_its_runtime_profile_name() -> None: + with pytest.raises(ValidationError): + ProgramBehaves(runtime_profile_name="bad name", cases=(ProgramCase(),)) + + +def test_execution_scorable_round_trips_as_a_stored_scorable() -> None: + message_program = ProgramExecutionScorable( + program=MessageScorable(message_piece_ids=(uuid.uuid4(),)), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), + ) + for scorable in (message_program, _execution()): + assert scorable_from_dict(scorable.model_dump(mode="json")) == scorable + + +def test_execution_scorable_accepts_stored_text_content() -> None: + stored = ProgramExecutionScorable( + program=ContentEntryScorable(content_id=uuid.uuid4()), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), + ) + assert scorable_from_dict(stored.model_dump(mode="json")) == stored + with pytest.raises(ValidationError, match="text"): + ProgramExecutionScorable( + program=ContentEntryScorable(content_id=uuid.uuid4(), data_type="image_path"), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), + ) + + +def test_execution_scorable_requires_text_and_an_invocation() -> None: + with pytest.raises(ValidationError, match="text"): + ProgramExecutionScorable( + program=ContentScorable(value="image.png", data_type="image_path"), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), + ) + with pytest.raises(ValidationError, match="at least one invocation"): + ProgramExecutionScorable( + program=ContentScorable(value="print(1)"), + runtime_profile_name="python-3.12", + invocations=(), + ) + + +def test_execution_scorable_rejects_a_program_that_is_not_text_or_a_message() -> None: + with pytest.raises(ValidationError): + ProgramExecutionScorable.model_validate( + { + "program": SurfaceScorable(uri="/x").model_dump(), + "runtime_profile_name": "python-3.12", + "invocations": [{}], + } + ) + + +def test_score_round_trips_an_execution_anchor() -> None: + scorable = _execution() + score = Score( + score_value="true", + score_type="true_false", + score_rationale="", + score_value_description="", + scorable=scorable, + scorer_class_identifier=_SOURCE, + ) + assert Score.model_validate(score.model_dump(mode="json")).scorable == scorable + + +@pytest.mark.parametrize( + "indices", + [(1, 0), (0, 0)], +) +def test_payload_requires_events_once_in_invocation_order(indices: tuple[int, ...]) -> None: + with pytest.raises(ValidationError, match="invocation order"): + ProgramExecutionObservationPayload( + program_sha256=_DIGEST, events=tuple(_event(invocation_index=index) for index in indices) + ) + + +def test_payload_events_require_the_program_digest() -> None: + with pytest.raises(ValidationError, match="digest"): + ProgramExecutionObservationPayload(events=(_event(),)) + + +def test_payload_rejects_a_coerced_schema_version() -> None: + with pytest.raises(ValidationError, match="schema_version"): + ProgramExecutionObservationPayload.model_validate({"schema_version": "1"}) + + +def test_complete_observation_holds_one_event_per_invocation() -> None: + scorable = _execution() + observation = _observation( + scorable=scorable, + acquisition=Acquisition.COMPLETE, + events=(_event(invocation_index=0), _event(invocation_index=1)), + ) + assert Observation.model_validate(observation.model_dump(mode="json")) == observation + with pytest.raises(ValidationError, match="Complete program acquisition"): + _observation(scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(invocation_index=0),)) + with pytest.raises(ValidationError, match="Complete program acquisition"): + _observation( + scorable=scorable, + acquisition=Acquisition.COMPLETE, + events=(_event(invocation_index=0), _event(invocation_index=1)), + reasons=("gap",), + ) + + +def test_partial_observation_is_missing_an_invocation() -> None: + scorable = _execution() + _observation( + scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(invocation_index=0),), reasons=("x",) + ) + with pytest.raises(ValidationError, match="Partial program acquisition"): + _observation( + scorable=scorable, + acquisition=Acquisition.PARTIAL, + events=(_event(invocation_index=0), _event(invocation_index=1)), + ) + + +@pytest.mark.parametrize("acquisition", [Acquisition.ERROR, Acquisition.UNAVAILABLE]) +def test_failed_observation_holds_no_events(acquisition: Acquisition) -> None: + scorable = _execution() + _observation(scorable=scorable, acquisition=acquisition, reasons=("sandbox down",)) + with pytest.raises(ValidationError, match="cannot contain events"): + _observation(scorable=scorable, acquisition=acquisition, events=(_event(invocation_index=0),)) + + +def test_observation_events_must_belong_to_the_scorable() -> None: + scorable = _execution(invocations=1) + with pytest.raises(ValidationError, match="belong to the scorable"): + _observation(scorable=scorable, acquisition=Acquisition.PARTIAL, events=(_event(invocation_index=1),)) + with pytest.raises(ValidationError, match="runtime profile"): + _observation(scorable=scorable, acquisition=Acquisition.COMPLETE, events=(_event(profile="node-20"),)) + + +def test_program_payload_requires_an_execution_anchor() -> None: + with pytest.raises(ValidationError, match="ProgramExecutionScorable"): + Observation( + source_identifier=_SOURCE, + acquisition=Acquisition.ERROR, + scorable=ContentScorable(value="print(1)"), + payload=ProgramExecutionObservationPayload(reasons=("x",)), + ) diff --git a/tests/unit/models/test_scorable.py b/tests/unit/models/test_scorable.py index 79f9cd5798..532b85d49b 100644 --- a/tests/unit/models/test_scorable.py +++ b/tests/unit/models/test_scorable.py @@ -13,12 +13,14 @@ Message, MessagePiece, MessageScorable, + ProgramExecutionScorable, Scorable, Score, SurfaceScorable, TraceScorable, scorable_from_dict, ) +from pyrit.models.code_execution import _ProgramInvocation from pyrit.models.score.scorable import SCORABLE_TYPES @@ -155,6 +157,11 @@ def test_every_union_member_round_trips_to_its_own_type(): TraceScorable(trace_ids=("1" * 32,)), SurfaceScorable(uri="/data/out.txt"), ConversationScorable(conversation_id="whole-conversation"), + ProgramExecutionScorable( + program=ContentScorable(value="print(1)"), + runtime_profile_name="python-3.12", + invocations=(_ProgramInvocation(),), + ), ] assert {type(case) for case in cases} == set(SCORABLE_TYPES) diff --git a/tests/unit/score/test_program_execution_scorer.py b/tests/unit/score/test_program_execution_scorer.py new file mode 100644 index 0000000000..e72dcd2a42 --- /dev/null +++ b/tests/unit/score/test_program_execution_scorer.py @@ -0,0 +1,660 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import asyncio +import hashlib +import uuid +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit._code_execution import DockerRuntimeProfile, DockerSandboxClient, SandboxExecutionError +from pyrit.models import ( + Acquisition, + ComponentIdentifier, + Contains, + ContentEntryScorable, + ContentScorable, + ContentWritten, + Equals, + MessagePiece, + MessageScorable, + ProgramBehaves, + ProgramCase, + ProgramExecutionObservationPayload, + ProgramExecutionScorable, + ProgramResult, + Regex, + ScoreStatus, + ScoringExpectation, + SurfaceScorable, +) +from pyrit.models.code_execution import _ProgramExecutionEvent, _ProgramInvocation, _SandboxRun +from pyrit.score import ( + NonReplayableObservationError, + ProgramExecutionScorer, + ProgramExecutionSource, + TrueFalseInverterScorer, +) +from pyrit.score.text_matching import match_text +from pyrit.score.true_false.program_execution_scorer import _evaluate_case, _evaluate_program + +if TYPE_CHECKING: + from pyrit.memory import SQLiteMemory + from pyrit.models import TextMatcher + +pytestmark = pytest.mark.usefixtures("patch_central_database") + +_PROFILE = "python-3.12" +_IMAGE_DIGEST = "sha256:" + "b" * 64 + + +def _event(*, invocation_index: int = 0, **updates: object) -> _ProgramExecutionEvent: + values: dict[str, object] = { + "invocation_index": invocation_index, + "runtime_profile_name": _PROFILE, + "runtime_image_digest": _IMAGE_DIGEST, + "resolved_command": ("python3", "/program/main.py"), + "exit_code": 0, + "stdout": "", + "duration_ms": 1, + } + values.update(updates) + return _ProgramExecutionEvent.model_validate(values) + + +class _FakeSandbox: + """Answers each invocation from a script keyed by its first argument, and records calls.""" + + def __init__(self, *, outputs: dict[str, dict[str, object]] | None = None, profiles: frozenset[str] | None = None): + self.outputs = outputs or {} + self.profiles = profiles or frozenset({_PROFILE}) + self.calls: list[tuple[str, str, tuple[_ProgramInvocation, ...]]] = [] + self.error: SandboxExecutionError | None = None + self.stop_after: int | None = None + + def get_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"fake": True}) + + def get_runtime_profile_names(self) -> frozenset[str]: + return self.profiles + + async def execute_async( + self, *, program: str, runtime_profile_name: str, invocations: tuple[_ProgramInvocation, ...] + ) -> _SandboxRun: + self.calls.append((program, runtime_profile_name, invocations)) + if self.error is not None: + raise self.error + ran = invocations if self.stop_after is None else invocations[: self.stop_after] + events = tuple( + _event( + invocation_index=index, **self.outputs.get(invocation.arguments[0] if invocation.arguments else "", {}) + ) + for index, invocation in enumerate(ran) + ) + reasons = ("the sandbox lost its worker",) if len(ran) < len(invocations) else () + return _SandboxRun(events=events, reasons=reasons) + + +def _case(argument: str, *, exit_code: int | None = 0, stdout: TextMatcher | None = None) -> ProgramCase: + return ProgramCase( + invocation=_ProgramInvocation(arguments=(argument,)), + expected=ProgramResult(exit_code=exit_code, stdout=stdout), + ) + + +def _expectation(*cases: ProgramCase, profile: str = _PROFILE) -> ScoringExpectation: + return ScoringExpectation(conditions=(ProgramBehaves(runtime_profile_name=profile, cases=cases),)) + + +_SQUARES = (_case("3", stdout=Equals(value="9")), _case("4", stdout=Equals(value="16"))) + + +def _scorer(sandbox: _FakeSandbox, *, max_output_chars: int = 64_000) -> ProgramExecutionScorer: + return ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=sandbox, max_output_chars=max_output_chars)) + + +@pytest.mark.parametrize( + ("expected", "event", "result"), + [ + (ProgramResult(), {}, True), + (ProgramResult(), {"exit_code": 1}, False), + (ProgramResult(exit_code=None), {"exit_code": 1}, True), + (ProgramResult(exit_code=None), {"exit_code": None}, None), + (ProgramResult(exit_code=None), {"exit_code": None, "timed_out": True}, False), + (ProgramResult(exit_code=None), {"exit_code": None, "resource_limit_exceeded": True}, False), + (ProgramResult(stdout=Equals(value="9")), {"stdout": "9\n"}, True), + (ProgramResult(stdout=Equals(value="9")), {"stdout": "10"}, False), + (ProgramResult(stdout=Equals(value="9")), {"stdout": "9", "stdout_truncated": True}, None), + (ProgramResult(stdout=Contains(value="key")), {"stdout": "a key", "stdout_truncated": True}, True), + (ProgramResult(stdout=Contains(value="key")), {"stdout": "a lock", "stdout_truncated": True}, None), + (ProgramResult(stdout=Contains(value="key")), {"stdout": "a lock"}, False), + (ProgramResult(stdout=Regex(value=r"^\d+$")), {"stdout": "42"}, True), + (ProgramResult(stderr=Contains(value="Error")), {"stderr": "ValueError"}, True), + (ProgramResult(stderr=Contains(value="Error")), {"stderr": ""}, False), + (ProgramResult(stdout=Equals(value="9")), {"stdout": "9", "stderr_truncated": True}, True), + (ProgramResult(exit_code=1, stdout=Equals(value="9")), {"stdout": "9"}, False), + ( + ProgramResult(stdout=Equals(value="9"), stderr=Equals(value="x")), + {"stdout": "8", "stderr_truncated": True}, + False, + ), + ], +) +def test_evaluate_case(*, expected: ProgramResult, event: dict[str, object], result: bool | None) -> None: + assert _evaluate_case(expected=expected, event=_event(**event)).value is result + + +def test_evaluate_program_needs_every_case() -> None: + condition = ProgramBehaves(runtime_profile_name=_PROFILE, cases=_SQUARES) + passing = (_event(invocation_index=0, stdout="9"), _event(invocation_index=1, stdout="16")) + payload = ProgramExecutionObservationPayload(program_sha256="a" * 64, events=passing) + assert _evaluate_program(condition=condition, payload=payload, acquisition=Acquisition.COMPLETE).value is True + + one_ran = ProgramExecutionObservationPayload(program_sha256="a" * 64, events=passing[:1]) + assert _evaluate_program(condition=condition, payload=one_ran, acquisition=Acquisition.PARTIAL).value is None + + failing_first = ProgramExecutionObservationPayload( + program_sha256="a" * 64, events=(_event(invocation_index=0, stdout="8"),) + ) + assert _evaluate_program(condition=condition, payload=failing_first, acquisition=Acquisition.PARTIAL).value is False + + empty = ProgramExecutionObservationPayload(reasons=("down",)) + for acquisition in (Acquisition.ERROR, Acquisition.UNAVAILABLE): + assert _evaluate_program(condition=condition, payload=empty, acquisition=acquisition).value is None + + +async def test_scores_true_when_every_case_passes_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "16"}}) + program = ContentScorable(value="import sys\nprint(int(sys.argv[1]) ** 2)\n") + + score = (await _scorer(sandbox).score_async(scorable=program, expectation=_expectation(*_SQUARES)))[0] + + assert score.score_value == "true" + assert score.status is ScoreStatus.COMPLETE + assert "All 2 case(s)" in score.score_rationale + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.scorable, ProgramExecutionScorable) + assert observation.acquisition is Acquisition.COMPLETE + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.program_sha256 == hashlib.sha256(program.value.encode()).hexdigest() + + +async def test_sandbox_receives_stimuli_but_not_expected_results_async() -> None: + sandbox = _FakeSandbox() + secret = "the-expected-answer" + expectation = _expectation( + ProgramCase( + invocation=_ProgramInvocation(arguments=("3",), stdin="in"), + expected=ProgramResult(stdout=Equals(value=secret)), + ) + ) + + await _scorer(sandbox).score_async(scorable=ContentScorable(value="print(1)"), expectation=expectation) + + program, profile, invocations = sandbox.calls[0] + assert (program, profile) == ("print(1)", _PROFILE) + assert invocations == (_ProgramInvocation(arguments=("3",), stdin="in"),) + assert secret not in repr(sandbox.calls) + + +async def test_scoring_and_replay_do_not_require_runtime_provenance_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9", "runtime_image_digest": None, "resolved_command": None}}) + scorer = _scorer(sandbox) + expectation = _expectation(_case("3", stdout=Equals(value="9"))) + score = (await scorer.score_async(scorable=ContentScorable(value="print(9)"), expectation=expectation))[0] + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + + assert score.get_value() is True + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events[0].runtime_image_digest is None + assert observation.payload.events[0].resolved_command is None + replayed = (await scorer.score_observation_async(observation=observation, expectation=expectation))[0] + assert replayed.get_value() is True + assert replayed.observation_ids == score.observation_ids + assert len(sandbox.calls) == 1 + saved = (await sqlite_instance.get_scores_async(score_ids=[str(replayed.id)]))[0] + assert saved.get_value() is True + assert saved.scored_expectation == expectation + + +async def test_cases_are_matched_once_for_verdict_and_rationale_async() -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "wrong"}}) + with patch("pyrit.score.true_false.program_execution_scorer.match_text", wraps=match_text) as matching: + score = ( + await _scorer(sandbox).score_async( + scorable=ContentScorable(value="print(1)"), expectation=_expectation(*_SQUARES) + ) + )[0] + assert score.get_value() is False + assert "case 1 output did not match" in score.score_rationale + assert matching.call_count == 2 + + +@pytest.mark.parametrize( + ("failure", "rationale"), + [ + ("lost-output", "output capture may be incomplete"), + ("resource-limits", "must support cgroup"), + ("image-volumes", "must not declare volumes"), + ], +) +async def test_unsafe_docker_execution_cannot_produce_a_passing_score_async( + sqlite_instance: SQLiteMemory, failure: str, rationale: str +) -> None: + client = DockerSandboxClient( + profiles={_PROFILE: DockerRuntimeProfile(image="python@" + _IMAGE_DIGEST, command=("python3", "-I"))} + ) + process = MagicMock(spec=asyncio.subprocess.Process) + process.returncode = 1 + process.stdin = None + for name, data in (("stdout", b"ok"), ("stderr", b"attachment failed")): + reader = asyncio.StreamReader() + reader.feed_data(data) + reader.feed_eof() + setattr(process, name, reader) + process.communicate = AsyncMock(return_value=(b"", b"")) + + async def control_async(*arguments: str) -> tuple[int, str]: + if arguments[0] == "info": + driver = "none" if failure == "resource-limits" else "cgroupfs" + return 0, f"{driver} true true true true true" + if arguments[:2] == ("image", "inspect"): + return 0, '{"/data":{}}' if failure == "image-volumes" else "null" + if arguments[0] == "inspect": + return 0, "exited 0 false 2026-10-21T00:00:00Z" + assert arguments[:3] == ("rm", "--force", "--volumes") + return 0, "" + + scorer = ProgramExecutionScorer(source=ProgramExecutionSource(sandbox=client)) + expectation = _expectation(_case("3", stdout=Equals(value="ok"))) + with ( + patch("asyncio.create_subprocess_exec", new_callable=AsyncMock, return_value=process) as launch, + patch.object(client, "_control_async", side_effect=control_async), + ): + score = (await scorer.score_async(scorable=ContentScorable(value="print('okextra')"), expectation=expectation))[ + 0 + ] + assert score.status is ScoreStatus.UNDETERMINED + assert score.score_value is None + assert rationale in score.score_rationale + assert launch.await_count == (1 if failure == "lost-output" else 0) + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert observation.acquisition is Acquisition.ERROR + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + assert observation.payload.events == () + + +async def test_scores_false_when_one_case_fails_async() -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "15"}}) + + score = ( + await _scorer(sandbox).score_async( + scorable=ContentScorable(value="print(1)"), expectation=_expectation(*_SQUARES) + ) + )[0] + + assert score.score_value == "false" + assert "case 1 output did not match" in score.score_rationale + + +@pytest.mark.parametrize( + ("outcome", "rationale"), + [ + ({"exit_code": 1}, "case 0 exited with 1, expected 0"), + ({"exit_code": None, "timed_out": True}, "case 0 timed out"), + ({"exit_code": None, "resource_limit_exceeded": True}, "case 0 exceeded a resource limit"), + ], +) +async def test_program_failure_is_a_false_verdict_not_undetermined_async( + outcome: dict[str, object], rationale: str +) -> None: + sandbox = _FakeSandbox(outputs={"1": outcome}) + + score = ( + await _scorer(sandbox).score_async(scorable=ContentScorable(value="x"), expectation=_expectation(_case("1"))) + )[0] + + assert (score.score_value, score.status) == ("false", ScoreStatus.COMPLETE) + assert rationale in score.score_rationale + + +async def test_sandbox_failure_is_undetermined_not_false_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox() + sandbox.error = SandboxExecutionError("service unavailable") + + score = ( + await _scorer(sandbox).score_async(scorable=ContentScorable(value="x"), expectation=_expectation(_case("1"))) + )[0] + + assert score.status is ScoreStatus.UNDETERMINED + assert score.score_value is None + assert "service unavailable" in score.score_rationale + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert observation.acquisition is Acquisition.ERROR + + +async def test_sandbox_stopping_early_leaves_passing_runs_undetermined_async() -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}}) + sandbox.stop_after = 1 + + score = ( + await _scorer(sandbox).score_async(scorable=ContentScorable(value="x"), expectation=_expectation(*_SQUARES)) + )[0] + + assert score.status is ScoreStatus.UNDETERMINED + assert "case 1 did not run" in score.score_rationale + assert "lost its worker" in score.score_rationale + + +async def test_sandbox_stopping_early_still_fails_a_failed_case_async() -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "8"}}) + sandbox.stop_after = 1 + + score = ( + await _scorer(sandbox).score_async(scorable=ContentScorable(value="x"), expectation=_expectation(*_SQUARES)) + )[0] + + assert score.score_value == "false" + + +async def test_output_beyond_the_limit_is_cut_and_marked_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"1": {"stdout": "x" * 50 + "key", "stderr": "e" * 60}}) + scorer = _scorer(sandbox, max_output_chars=10) + + equals = ( + await scorer.score_async( + scorable=ContentScorable(value="x"), + expectation=_expectation(_case("1", stdout=Equals(value="x" * 50 + "key"))), + ) + )[0] + contains = ( + await scorer.score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1", stdout=Contains(value="xxx"))) + ) + )[0] + + assert equals.status is ScoreStatus.UNDETERMINED + assert contains.score_value == "true" + observation = (await sqlite_instance.get_observations_async(observation_ids=equals.observation_ids))[0] + assert isinstance(observation.payload, ProgramExecutionObservationPayload) + event = observation.payload.events[0] + assert (event.stdout, event.stdout_truncated, len(event.stderr), event.stderr_truncated) == ( + "x" * 10, + True, + 10, + True, + ) + + +async def test_replay_judges_stored_output_without_running_again_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "16"}}) + scorer = _scorer(sandbox) + score = (await scorer.score_async(scorable=ContentScorable(value="x"), expectation=_expectation(*_SQUARES)))[0] + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + + same = (await scorer.score_observation_async(observation=observation, expectation=_expectation(*_SQUARES)))[0] + stricter = ( + await scorer.score_observation_async( + observation=observation, + expectation=_expectation(_case("3", stdout=Equals(value="9")), _case("4", stdout=Equals(value="17"))), + ) + )[0] + + assert same.score_value == "true" + assert stricter.score_value == "false" + assert same.observation_ids == [observation.id] + assert len(sandbox.calls) == 1 + + +async def test_replay_rejects_different_stimuli_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "16"}}) + scorer = _scorer(sandbox) + score = (await scorer.score_async(scorable=ContentScorable(value="x"), expectation=_expectation(*_SQUARES)))[0] + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + + with pytest.raises(NonReplayableObservationError, match="different invocations"): + await scorer.score_observation_async( + observation=observation, expectation=_expectation(_case("5", stdout=Equals(value="25"))) + ) + with pytest.raises(NonReplayableObservationError, match="different invocations"): + await scorer.score_observation_async( + observation=observation, expectation=_expectation(*_SQUARES, profile="python-3.13") + ) + + +@pytest.mark.parametrize("explicit", [False, True]) +@pytest.mark.parametrize("concurrent", [False, True]) +async def test_scoring_calls_store_independent_executions_async( + sqlite_instance: SQLiteMemory, explicit: bool, concurrent: bool +) -> None: + sandbox = _FakeSandbox() + scorer = _scorer(sandbox) + program = ContentScorable(value="x") + + execution = ProgramExecutionScorable( + program=program, + runtime_profile_name=_PROFILE, + invocations=(_ProgramInvocation(arguments=("1",)),), + ) + + scorable = execution if explicit else program + expectation = _expectation(_case("1")) + if concurrent: + first_scores, second_scores = await asyncio.gather( + scorer.score_async(scorable=scorable, expectation=expectation), + scorer.score_async(scorable=scorable, expectation=expectation), + ) + else: + first_scores = await scorer.score_async(scorable=scorable, expectation=expectation) + second_scores = await scorer.score_async(scorable=scorable, expectation=expectation) + first, second = first_scores[0], second_scores[0] + + assert len(sandbox.calls) == 2 + assert first.observation_ids != second.observation_ids + stored = await sqlite_instance.get_scores_async(score_ids=[str(first.id), str(second.id)]) + assert len(stored) == 2 + assert all(score.get_value() is True for score in stored) + + +async def test_a_run_with_no_events_is_an_acquisition_error_async() -> None: + sandbox = _FakeSandbox() + sandbox.stop_after = 0 + source = ProgramExecutionSource(sandbox=sandbox) + execution = ProgramExecutionScorable( + program=ContentScorable(value="x"), + runtime_profile_name=_PROFILE, + invocations=(_ProgramInvocation(),), + ) + + failed = await source.acquire_async(scorable=execution) + assert failed.acquisition is Acquisition.ERROR + assert isinstance(failed.payload, ProgramExecutionObservationPayload) + assert failed.payload.reasons == ("the sandbox lost its worker",) + + +async def test_explicit_execution_must_match_the_condition_async() -> None: + execution = ProgramExecutionScorable( + program=ContentScorable(value="x"), + runtime_profile_name=_PROFILE, + invocations=(_ProgramInvocation(arguments=("2",)),), + ) + + with pytest.raises(RuntimeError, match="runtime profile and invocations"): + await _scorer(_FakeSandbox()).score_async(scorable=execution, expectation=_expectation(_case("1"))) + + +async def test_runtime_profile_must_be_offered_by_the_sandbox_async() -> None: + sandbox = _FakeSandbox() + + with pytest.raises(RuntimeError, match="not offered by the sandbox"): + await _scorer(sandbox).score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1"), profile="node-20") + ) + assert sandbox.calls == [] + + +async def test_scores_a_stored_message_verbatim_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox() + piece = MessagePiece(role="assistant", original_value="print('hi')", conversation_id=str(uuid.uuid4())) + await sqlite_instance.add_message_to_memory_async(request=piece.to_message()) + scorable = MessageScorable(message_piece_ids=(piece.id,)) + + score = (await _scorer(sandbox).score_async(scorable=scorable, expectation=_expectation(_case("1"))))[0] + + assert score.score_value == "true" + assert sandbox.calls[0][0] == "print('hi')" + assert score.scorable == scorable + assert score.message_piece_id == piece.id + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert isinstance(observation.scorable, ProgramExecutionScorable) + assert observation.scorable.program == scorable + + +async def test_a_blocked_response_has_no_program_to_run_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox() + piece = MessagePiece( + role="assistant", + original_value="content filtered", + conversation_id=str(uuid.uuid4()), + response_error="blocked", + ) + await sqlite_instance.add_message_to_memory_async(request=piece.to_message()) + + score = ( + await _scorer(sandbox).score_async( + scorable=MessageScorable(message_piece_ids=(piece.id,)), expectation=_expectation(_case("1")) + ) + )[0] + + assert score.status is ScoreStatus.UNDETERMINED + assert "blocked response" in score.score_rationale + assert sandbox.calls == [] + + +async def test_a_message_reference_must_name_one_piece_async(sqlite_instance: SQLiteMemory) -> None: + conversation_id = str(uuid.uuid4()) + first = MessagePiece(role="assistant", original_value="a", conversation_id=conversation_id) + second = MessagePiece(role="assistant", original_value="b", conversation_id=conversation_id) + await sqlite_instance.add_message_to_memory_async(request=first.to_message()) + + with pytest.raises(RuntimeError, match="exactly one message piece"): + await _scorer(_FakeSandbox()).score_async( + scorable=MessageScorable(message_piece_ids=(first.id, second.id)), expectation=_expectation(_case("1")) + ) + + +async def test_rejects_other_scorables_async() -> None: + with pytest.raises(RuntimeError, match="requires a MessageScorable"): + await _scorer(_FakeSandbox()).score_async( + scorable=SurfaceScorable(uri="/out.txt"), expectation=_expectation(_case("1")) + ) + + +async def test_requires_the_program_behaves_condition_async() -> None: + with pytest.raises(ValueError, match="requires one ProgramBehaves condition"): + await _scorer(_FakeSandbox()).score_async( + scorable=ContentScorable(value="x"), + expectation=ScoringExpectation(conditions=(ContentWritten(uri="/out.txt"),)), + ) + + +async def test_inverter_keeps_undetermined_async() -> None: + sandbox = _FakeSandbox() + sandbox.error = SandboxExecutionError("down") + passing = _FakeSandbox() + + undetermined = ( + await TrueFalseInverterScorer(scorer=_scorer(sandbox)).score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1")) + ) + )[0] + inverted = ( + await TrueFalseInverterScorer(scorer=_scorer(passing)).score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1")) + ) + )[0] + + assert undetermined.status is ScoreStatus.UNDETERMINED + assert inverted.get_value() is False + + +def test_source_rejects_a_nonpositive_output_limit() -> None: + with pytest.raises(ValueError, match="positive"): + ProgramExecutionSource(sandbox=_FakeSandbox(), max_output_chars=0) + + +def test_identifier_names_the_source_and_sandbox() -> None: + identifier = _scorer(_FakeSandbox()).get_identifier() + source = identifier.children["source"] + assert isinstance(source, ComponentIdentifier) + assert source.class_name == "ProgramExecutionSource" + assert isinstance(source.children["sandbox"], ComponentIdentifier) + + +async def test_source_rejects_events_for_unknown_invocations_async() -> None: + class _Rogue(_FakeSandbox): + async def execute_async(self, *, program, runtime_profile_name, invocations): # type: ignore[no-untyped-def] + return _SandboxRun(events=(_event(invocation_index=5),)) + + execution = ProgramExecutionScorable( + program=ContentScorable(value="x"), + runtime_profile_name=_PROFILE, + invocations=(_ProgramInvocation(),), + ) + with pytest.raises(ValueError, match="not given"): + await ProgramExecutionSource(sandbox=_Rogue()).acquire_async(scorable=execution) + + +async def test_sandbox_warnings_are_kept_on_the_observation_async(sqlite_instance: SQLiteMemory) -> None: + class _LeakySandbox(_FakeSandbox): + async def execute_async(self, *, program, runtime_profile_name, invocations): # type: ignore[no-untyped-def] + return _SandboxRun( + events=(_event(),), warnings=("could not remove container pyrit-sandbox-1: daemon unavailable",) + ) + + score = ( + await _scorer(_LeakySandbox()).score_async( + scorable=ContentScorable(value="x"), expectation=_expectation(_case("1")) + ) + )[0] + + observation = (await sqlite_instance.get_observations_async(observation_ids=score.observation_ids))[0] + assert score.score_value == "true" + assert observation.acquisition is Acquisition.COMPLETE + assert "daemon unavailable" in observation.metadata["sandbox_warnings"] + + +async def test_scoring_a_saved_score_anchor_runs_the_stored_program_async(sqlite_instance: SQLiteMemory) -> None: + sandbox = _FakeSandbox(outputs={"3": {"stdout": "9"}, "4": {"stdout": "16"}}) + scorer = _scorer(sandbox) + program = ContentScorable(value="import sys\nprint(int(sys.argv[1]) ** 2)\n") + first = (await scorer.score_async(scorable=program, expectation=_expectation(*_SQUARES)))[0] + saved = (await sqlite_instance.get_scores_async(score_ids=[first.id]))[0] + assert isinstance(saved.scorable, ContentEntryScorable) + + again = (await scorer.score_async(scorable=saved.scorable, expectation=_expectation(*_SQUARES)))[0] + + assert again.score_value == "true" + assert [call[0] for call in sandbox.calls] == [program.value, program.value] + observation = (await sqlite_instance.get_observations_async(observation_ids=again.observation_ids))[0] + assert isinstance(observation.scorable, ProgramExecutionScorable) + assert observation.scorable.program == saved.scorable + + +async def test_each_acquisition_returns_new_evidence_without_memory_async() -> None: + sandbox = _FakeSandbox() + source = ProgramExecutionSource(sandbox=sandbox) + execution = ProgramExecutionScorable( + program=ContentScorable(value="x"), runtime_profile_name=_PROFILE, invocations=(_ProgramInvocation(),) + ) + with patch("pyrit.score.observation.program_execution_source.CentralMemory.get_memory_instance") as memory: + memory.side_effect = AssertionError("Loose program acquisition must not read memory") + first = await source.acquire_async(scorable=execution) + second = await source.acquire_async(scorable=execution) + assert first.id != second.id + assert first.scorable == second.scorable == execution + assert len(sandbox.calls) == 2