From 6776ba318ea6d5938a181a20a7b237a9f5e1b35c Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Wed, 5 Aug 2026 17:00:03 -0700 Subject: [PATCH 01/16] Add MCP v2 server foundation Introduce the optional MCP server package, lazy CLI entry points, typed server discovery, capability modes, and in-memory plus stdio protocol coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- .github/workflows/regression.yml | 5 +- assert_ai/cli.py | 2 + assert_ai/mcp/__init__.py | 8 ++ assert_ai/mcp/__main__.py | 10 +++ assert_ai/mcp/_command.py | 83 ++++++++++++++++++ assert_ai/mcp/models.py | 61 ++++++++++++++ assert_ai/mcp/server.py | 123 +++++++++++++++++++++++++++ pyproject.toml | 20 +++++ tests/test_mcp_cli.py | 71 ++++++++++++++++ tests/test_mcp_server.py | 139 +++++++++++++++++++++++++++++++ 10 files changed, 520 insertions(+), 2 deletions(-) create mode 100644 assert_ai/mcp/__init__.py create mode 100644 assert_ai/mcp/__main__.py create mode 100644 assert_ai/mcp/_command.py create mode 100644 assert_ai/mcp/models.py create mode 100644 assert_ai/mcp/server.py create mode 100644 tests/test_mcp_cli.py create mode 100644 tests/test_mcp_server.py diff --git a/.github/workflows/regression.yml b/.github/workflows/regression.yml index 517250e15..8c641373d 100644 --- a/.github/workflows/regression.yml +++ b/.github/workflows/regression.yml @@ -45,9 +45,10 @@ jobs: # opentelemetry at module load time (target.trace.backend: otel). Demo # examples like that one ship as part of the unit-test surface, so the # CI install set has to include the demo's optional extras even though - # core ASSERT doesn't need them. + # core ASSERT doesn't need them. MCP v2 is installed explicitly because + # it cannot currently share the `examples` extra's MCP v1 dependency. run: | - python -m pip install -e ".[dev,otel]" + python -m pip install -e ".[dev,otel,mcp]" - name: Install viewer npm dependencies # tests/test_viewer_*.py shell out to `node` against viewer TypeScript diff --git a/assert_ai/cli.py b/assert_ai/cli.py index b70cfb51e..765689988 100644 --- a/assert_ai/cli.py +++ b/assert_ai/cli.py @@ -780,8 +780,10 @@ def cli(ctx: click.Context, verbose: bool, quiet: bool, log_file: Path | None, o # -- init (design an eval config with an LLM assistant) --------------------- from assert_ai.init._command import init # noqa: E402 +from assert_ai.mcp._command import mcp # noqa: E402 cli.add_command(init) +cli.add_command(mcp) @cli.command(short_help="Run a pipeline from a YAML config") diff --git a/assert_ai/mcp/__init__.py b/assert_ai/mcp/__init__.py new file mode 100644 index 000000000..fe0e90216 --- /dev/null +++ b/assert_ai/mcp/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Optional MCP adapter for ASSERT.""" + +ASSERT_MCP_API_VERSION = "1" + +__all__ = ["ASSERT_MCP_API_VERSION"] diff --git a/assert_ai/mcp/__main__.py b/assert_ai/mcp/__main__.py new file mode 100644 index 000000000..36108c253 --- /dev/null +++ b/assert_ai/mcp/__main__.py @@ -0,0 +1,10 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Run the ASSERT MCP server over stdio.""" + +from assert_ai.mcp._command import serve + + +if __name__ == "__main__": + serve() diff --git a/assert_ai/mcp/_command.py b/assert_ai/mcp/_command.py new file mode 100644 index 000000000..2cac131b6 --- /dev/null +++ b/assert_ai/mcp/_command.py @@ -0,0 +1,83 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Click commands for the optional ASSERT MCP server.""" + +from __future__ import annotations + +import importlib +from pathlib import Path +from types import ModuleType + +import click + +from assert_ai.mcp.models import CapabilityGroup, ServerMode + +_INSTALL_HINT = 'Install the MCP dependencies with: python -m pip install "assert-ai[mcp]"' +_EXPLICIT_GROUPS = [ + CapabilityGroup.DESIGN.value, + CapabilityGroup.PROBE.value, + CapabilityGroup.TRACE.value, + CapabilityGroup.ANALYSIS.value, + CapabilityGroup.ACS.value, + CapabilityGroup.EXPORT.value, +] + + +def _load_server_module() -> ModuleType: + """Import the MCP SDK-dependent server only when serving starts.""" + try: + return importlib.import_module("assert_ai.mcp.server") + except ModuleNotFoundError as exc: + if exc.name == "mcp" or (exc.name and exc.name.startswith("mcp.")): + raise click.ClickException(_INSTALL_HINT) from exc + raise + + +@click.group(short_help="Expose ASSERT workflows through an MCP server.") +def mcp() -> None: + """Manage the ASSERT Model Context Protocol server.""" + + +@click.command(short_help="Serve ASSERT over MCP stdio.") +@click.option( + "--workspace", + type=click.Path( + exists=True, + file_okay=False, + resolve_path=True, + path_type=Path, + ), + default=Path("."), + show_default=True, + help="Workspace containing eval configs and managed artifacts.", +) +@click.option( + "--mode", + type=click.Choice([mode.value for mode in ServerMode], case_sensitive=False), + default=ServerMode.INSPECT.value, + show_default=True, + help="Base capability set exposed by the server.", +) +@click.option( + "--enable-group", + "enabled_groups", + type=click.Choice(_EXPLICIT_GROUPS, case_sensitive=False), + multiple=True, + help="Enable an additional capability group. Repeat as needed.", +) +def serve(workspace: Path, mode: str, enabled_groups: tuple[str, ...]) -> None: + """Serve ASSERT over stdio; stdout is reserved for MCP protocol traffic.""" + server_module = _load_server_module() + try: + options = server_module.ServerOptions.create( + workspace_root=workspace, + mode=mode, + enabled_groups=enabled_groups, + ) + except (OSError, ValueError) as exc: + raise click.ClickException(str(exc)) from exc + server_module.run_stdio_server(options) + + +mcp.add_command(serve) diff --git a/assert_ai/mcp/models.py b/assert_ai/mcp/models.py new file mode 100644 index 000000000..11036862b --- /dev/null +++ b/assert_ai/mcp/models.py @@ -0,0 +1,61 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Typed public models for the ASSERT MCP adapter.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field + +from assert_ai.mcp import ASSERT_MCP_API_VERSION + + +class ServerMode(StrEnum): + """Predefined MCP capability bundles.""" + + INSPECT = "inspect" + AUTHOR = "author" + FULL = "full" + + +class CapabilityGroup(StrEnum): + """Stable names for independently gated MCP capabilities.""" + + INSPECT = "inspect" + AUTHOR = "author" + DESIGN = "design" + EXECUTE = "execute" + PROBE = "probe" + CURATE = "curate" + TRACE = "trace" + ANALYSIS = "analysis" + ACS = "acs" + EXPORT = "export" + + +class WorkspaceInfo(BaseModel): + """Workspace-relative roots managed by the MCP server.""" + + model_config = ConfigDict(frozen=True) + + root: Literal["."] = "." + configs_root: str = "evals" + artifacts_root: str = "artifacts" + results_root: str = "artifacts/results" + + +class ServerInfo(BaseModel): + """Discovery metadata returned by ``get_server_info``.""" + + model_config = ConfigDict(frozen=True) + + name: Literal["ASSERT"] = "ASSERT" + server_version: str + assert_mcp_api_version: Literal["1"] = ASSERT_MCP_API_VERSION + mode: ServerMode + enabled_capability_groups: list[CapabilityGroup] + workspace: WorkspaceInfo = Field(default_factory=WorkspaceInfo) + transports: list[Literal["stdio"]] = Field(default_factory=lambda: ["stdio"]) diff --git a/assert_ai/mcp/server.py b/assert_ai/mcp/server.py new file mode 100644 index 000000000..d3c1aaf3f --- /dev/null +++ b/assert_ai/mcp/server.py @@ -0,0 +1,123 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""MCP v2 server factory and stdio entry point for ASSERT.""" + +from __future__ import annotations + +from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Iterable + +from mcp.server import MCPServer +from mcp.types import ToolAnnotations + +from assert_ai.mcp.models import ( + CapabilityGroup, + ServerInfo, + ServerMode, + WorkspaceInfo, +) + +SERVER_NAME = "ASSERT" + +_MODE_GROUPS: dict[ServerMode, tuple[CapabilityGroup, ...]] = { + ServerMode.INSPECT: (CapabilityGroup.INSPECT,), + ServerMode.AUTHOR: ( + CapabilityGroup.INSPECT, + CapabilityGroup.AUTHOR, + ), + ServerMode.FULL: ( + CapabilityGroup.INSPECT, + CapabilityGroup.AUTHOR, + CapabilityGroup.DESIGN, + CapabilityGroup.EXECUTE, + CapabilityGroup.PROBE, + CapabilityGroup.CURATE, + ), +} +_GROUP_ORDER = {group: index for index, group in enumerate(CapabilityGroup)} +_AUTHOR_EXTENSION_GROUPS = { + CapabilityGroup.DESIGN, + CapabilityGroup.PROBE, +} + + +@dataclass(frozen=True) +class ServerOptions: + """Launch-time settings fixed for the lifetime of one MCP server.""" + + workspace_root: Path + mode: ServerMode = ServerMode.INSPECT + enabled_groups: tuple[CapabilityGroup, ...] = () + + @classmethod + def create( + cls, + *, + workspace_root: str | Path, + mode: str | ServerMode = ServerMode.INSPECT, + enabled_groups: Iterable[str | CapabilityGroup] = (), + ) -> "ServerOptions": + root = Path(workspace_root).expanduser().resolve(strict=True) + if not root.is_dir(): + raise ValueError(f"Workspace is not a directory: {root}") + parsed_mode = ServerMode(mode) + parsed_groups = tuple(CapabilityGroup(group) for group in enabled_groups) + invalid_groups = _AUTHOR_EXTENSION_GROUPS.intersection(parsed_groups) + if parsed_mode is ServerMode.INSPECT and invalid_groups: + names = ", ".join(sorted(group.value for group in invalid_groups)) + raise ValueError( + f"Capability group(s) {names} require --mode author or --mode full." + ) + return cls( + workspace_root=root, + mode=parsed_mode, + enabled_groups=parsed_groups, + ) + + @property + def capability_groups(self) -> tuple[CapabilityGroup, ...]: + groups = {*_MODE_GROUPS[self.mode], *self.enabled_groups} + return tuple(sorted(groups, key=_GROUP_ORDER.__getitem__)) + + +def _server_version() -> str: + try: + return version("assert-ai") + except PackageNotFoundError: + return "0.1.0" + + +def build_server(options: ServerOptions) -> MCPServer: + """Build an in-process MCP server for the configured workspace.""" + server = MCPServer( + SERVER_NAME, + description="Local, spec-driven evaluation workflows for AI agents.", + version=_server_version(), + ) + + @server.tool( + title="Get ASSERT server information", + annotations=ToolAnnotations( + read_only_hint=True, + open_world_hint=False, + ), + structured_output=True, + ) + def get_server_info() -> ServerInfo: + """Describe this ASSERT server's API, workspace, and capabilities.""" + return ServerInfo( + server_version=_server_version(), + mode=options.mode, + enabled_capability_groups=list(options.capability_groups), + workspace=WorkspaceInfo(), + ) + + return server + + +def run_stdio_server(options: ServerOptions) -> None: + """Run the configured server over stdio.""" + build_server(options).run("stdio") diff --git a/pyproject.toml b/pyproject.toml index 655a78b78..05da85b94 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,6 +91,13 @@ acs = [ "acs-generator>=0.3.1b0", "agent-control-specification>=0.3.1b0", ] +mcp = [ + # OpenAI Agents currently requires mcp<2, so this extra is intentionally + # separate from `examples` and the `all` meta-extra. psutil is required by + # the planned Windows worker process-tree cancellation path. + "mcp>=2,<3", + "psutil>=6,<8", +] examples = [ "autogen-agentchat>=0.7.5", "autogen-ext>=0.7.5", @@ -129,6 +136,7 @@ dev = [ [project.scripts] assert-ai = "assert_ai.cli:cli" +assert-ai-mcp = "assert_ai.mcp._command:serve" [project.urls] Homepage = "https://github.com/responsibleai/ASSERT" @@ -140,6 +148,18 @@ Documentation = "https://github.com/responsibleai/ASSERT#readme" requires = ["setuptools>=61.0"] build-backend = "setuptools.build_meta" +[tool.uv] +conflicts = [ + [ + { extra = "mcp" }, + { extra = "examples" }, + ], + [ + { extra = "mcp" }, + { extra = "all" }, + ], +] + [tool.setuptools] include-package-data = true diff --git a/tests/test_mcp_cli.py b/tests/test_mcp_cli.py new file mode 100644 index 000000000..d0dc52981 --- /dev/null +++ b/tests/test_mcp_cli.py @@ -0,0 +1,71 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from click.testing import CliRunner + +from assert_ai.cli import cli + + +def test_root_help_lists_mcp_without_importing_sdk() -> None: + runner = CliRunner() + with patch("assert_ai.mcp._command.importlib.import_module") as import_module: + result = runner.invoke(cli, ["--help"]) + + assert result.exit_code == 0, result.output + assert "mcp" in result.output + import_module.assert_not_called() + + +def test_mcp_serve_forwards_resolved_options() -> None: + runner = CliRunner() + run_stdio_server = Mock() + options = object() + server_options = SimpleNamespace(create=Mock(return_value=options)) + server_module = SimpleNamespace( + ServerOptions=server_options, + run_stdio_server=run_stdio_server, + ) + + with runner.isolated_filesystem(), patch( + "assert_ai.mcp._command._load_server_module", + return_value=server_module, + ): + result = runner.invoke( + cli, + [ + "mcp", + "serve", + "--workspace", + ".", + "--mode", + "author", + "--enable-group", + "design", + ], + ) + + assert result.exit_code == 0, result.output + create_kwargs = server_options.create.call_args.kwargs + assert create_kwargs["workspace_root"].is_absolute() + assert create_kwargs["mode"] == "author" + assert create_kwargs["enabled_groups"] == ("design",) + run_stdio_server.assert_called_once_with(options) + + +def test_mcp_serve_reports_missing_optional_dependency() -> None: + runner = CliRunner() + missing = ModuleNotFoundError("No module named 'mcp'", name="mcp") + with patch( + "assert_ai.mcp._command.importlib.import_module", + side_effect=missing, + ): + result = runner.invoke(cli, ["mcp", "serve"]) + + assert result.exit_code == 1 + assert 'python -m pip install "assert-ai[mcp]"' in result.output + assert "Traceback" not in result.output diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py new file mode 100644 index 000000000..3e521d5da --- /dev/null +++ b/tests/test_mcp_server.py @@ -0,0 +1,139 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import asyncio +import os +import sys +from contextlib import asynccontextmanager +from pathlib import Path +from typing import AsyncIterator + +import pytest + +pytest.importorskip("mcp") + +from mcp.client import Client +from mcp.client._transport import TransportStreams +from mcp.client.stdio import StdioServerParameters, stdio_client + +from assert_ai.mcp.models import CapabilityGroup, ServerMode +from assert_ai.mcp.server import ServerOptions, build_server + + +@asynccontextmanager +async def _stdio_transport( + workspace: Path, +) -> AsyncIterator[TransportStreams]: + parameters = StdioServerParameters( + command=sys.executable, + args=[ + "-m", + "assert_ai.mcp", + "--workspace", + str(workspace), + ], + env={ + "PYTHONPATH": os.pathsep.join( + filter( + None, + [ + str(Path(__file__).resolve().parents[1]), + os.environ.get("PYTHONPATH"), + ], + ) + ) + }, + ) + async with stdio_client(parameters) as streams: + yield streams + + +def test_server_options_resolve_workspace_and_capabilities(tmp_path: Path) -> None: + workspace = tmp_path / "workspace" + workspace.mkdir() + + options = ServerOptions.create( + workspace_root=workspace / ".." / "workspace", + mode="author", + enabled_groups=["design", "trace", "design"], + ) + + assert options.workspace_root == workspace.resolve() + assert options.mode is ServerMode.AUTHOR + assert options.capability_groups == ( + CapabilityGroup.INSPECT, + CapabilityGroup.AUTHOR, + CapabilityGroup.DESIGN, + CapabilityGroup.TRACE, + ) + + +def test_design_group_requires_author_or_full_mode(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="require --mode author or --mode full"): + ServerOptions.create( + workspace_root=tmp_path, + mode="inspect", + enabled_groups=["design"], + ) + + +def test_get_server_info_protocol_round_trip(tmp_path: Path) -> None: + async def run() -> tuple[set[str], object]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="full", + enabled_groups=["analysis"], + ) + async with Client(build_server(options), raise_exceptions=True) as client: + tools = await client.list_tools() + result = await client.call_tool("get_server_info", {}) + return {tool.name for tool in tools.tools}, result + + tool_names, result = asyncio.run(run()) + + assert tool_names == {"get_server_info"} + assert result.is_error is False + assert result.structured_content is not None + assert result.structured_content["assert_mcp_api_version"] == "1" + assert result.structured_content["mode"] == "full" + assert result.structured_content["workspace"]["root"] == "." + assert result.structured_content["enabled_capability_groups"] == [ + "inspect", + "author", + "design", + "execute", + "probe", + "curate", + "analysis", + ] + + +def test_get_server_info_publishes_structured_output_schema(tmp_path: Path) -> None: + async def run() -> object: + options = ServerOptions.create(workspace_root=tmp_path) + async with Client(build_server(options), raise_exceptions=True) as client: + tools = await client.list_tools() + return tools.tools[0] + + tool = asyncio.run(run()) + + assert tool.name == "get_server_info" + assert tool.output_schema is not None + assert "assert_mcp_api_version" in tool.output_schema["properties"] + assert tool.annotations is not None + assert tool.annotations.read_only_hint is True + assert tool.annotations.open_world_hint is False + + +def test_stdio_module_entry_point_keeps_protocol_wire_clean(tmp_path: Path) -> None: + async def run() -> object: + async with Client(_stdio_transport(tmp_path), raise_exceptions=True) as client: + return await client.call_tool("get_server_info", {}) + + result = asyncio.run(run()) + + assert result.is_error is False + assert result.structured_content is not None + assert result.structured_content["workspace"]["root"] == "." From 88224b3b48f32d1c3531c4cfdbe0f6844bd5204b Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Thu, 6 Aug 2026 14:07:32 -0700 Subject: [PATCH 02/16] Add strict MCP runtime containment Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/cli.py | 10 +- assert_ai/config.py | 176 ++++++- assert_ai/core/artifact_cache.py | 291 ++++++++++- assert_ai/core/azure_auth.py | 2 +- assert_ai/core/environment.py | 37 ++ assert_ai/core/model_client.py | 30 +- assert_ai/core/otel_session.py | 13 +- assert_ai/core/runtime_path_policy.py | 469 +++++++++++++++++ assert_ai/core/security.py | 16 +- assert_ai/core/session.py | 22 +- assert_ai/core/tool_backend.py | 505 +++++++++++++++++- assert_ai/core/workspace.py | 69 +++ assert_ai/init/_command.py | 16 +- assert_ai/mcp/_command.py | 28 +- assert_ai/mcp/server.py | 26 +- assert_ai/runner.py | 78 ++- assert_ai/stages/inference.py | 134 ++++- assert_ai/stages/judge.py | 92 +++- assert_ai/stages/stratification.py | 2 + assert_ai/stages/systematize.py | 2 + assert_ai/stages/test_set.py | 2 + scripts/benchmark.py | 4 + tests/test_artifact_cache.py | 12 +- tests/test_cli.py | 14 +- tests/test_environment_bootstrap.py | 82 +++ tests/test_init_command.py | 29 + tests/test_mcp_cli.py | 83 +++ tests/test_mcp_server.py | 13 + tests/test_runtime_path_policy.py | 727 ++++++++++++++++++++++++++ 29 files changed, 2836 insertions(+), 148 deletions(-) create mode 100644 assert_ai/core/environment.py create mode 100644 assert_ai/core/runtime_path_policy.py create mode 100644 assert_ai/core/workspace.py create mode 100644 tests/test_environment_bootstrap.py create mode 100644 tests/test_runtime_path_policy.py diff --git a/assert_ai/cli.py b/assert_ai/cli.py index 765689988..b70c037c8 100644 --- a/assert_ai/cli.py +++ b/assert_ai/cli.py @@ -857,12 +857,12 @@ def run( log_file=log_file, json_output=(output_format == "json"), ) + from assert_ai.core.environment import bootstrap_environment + + bootstrap_environment(discover_from_cwd=True) runner = _load_runner_module() - # Emit the resolved Azure auth mode AFTER runner.py has loaded ``.env`` - # and called ``refresh_azure_auth_mode(force=True)`` — so the log reflects - # the value that azure/* requests will actually use — and AFTER subcommand - # logging flags have been applied (above), so ``--quiet`` silences it and - # ``--output json`` formats it as JSON. + # Emit the resolved Azure auth mode after explicit environment bootstrap + # and subcommand logging configuration. from assert_ai.core.azure_auth import log_resolved_azure_auth_mode log_resolved_azure_auth_mode() rc = runner.run_pipeline( diff --git a/assert_ai/config.py b/assert_ai/config.py index aab274bf8..7b9501aba 100644 --- a/assert_ai/config.py +++ b/assert_ai/config.py @@ -33,6 +33,7 @@ ToolsConfig, TraceConfig, ) +from assert_ai.core.runtime_path_policy import RuntimePathPolicy ROOT = Path(__file__).resolve().parent.parent OUTPUT_PATH_KEYS = {"save_dir", "save_path"} @@ -119,12 +120,30 @@ def _resolve_path( artifacts_root: Path, cfg_dir: Path | None = None, use_artifacts_root: bool = False, + path_policy: RuntimePathPolicy | None = None, + field_name: str = "path", ) -> str: """Resolve one path against artifacts and config roots. Validates that relative paths do not escape their expected root directory via traversal sequences. """ + if path_policy is not None: + if use_artifacts_root: + return str( + path_policy.resolve_output( + path, + field_name=field_name, + ) + ) + return str( + path_policy.resolve_input( + path, + base_dir=cfg_dir or path_policy.config_root, + field_name=field_name, + ) + ) + artifacts_root = Path(artifacts_root).expanduser().resolve() cfg_dir = Path(cfg_dir).expanduser().resolve() if cfg_dir is not None else None candidate = Path(path).expanduser() @@ -170,8 +189,12 @@ def load_runtime_context( cfg_path: Path, *, stage_modules: dict[str, Any], + path_policy: RuntimePathPolicy | None = None, ) -> dict[str, Any]: """Build the shared runtime context used by every stage.""" + cfg_path = Path(cfg_path).expanduser().resolve() + if path_policy is not None: + cfg_path = path_policy.resolve_config_path(cfg_path) reject_unknown_keys( raw, field_name="config", @@ -198,14 +221,40 @@ def load_runtime_context( pipeline = parse_pipeline_config(raw) target = pipeline.target if pipeline else None - artifacts_root = Path(raw.get("artifacts_root") or "artifacts").expanduser() - if not artifacts_root.is_absolute(): - artifacts_root = (ROOT / artifacts_root).resolve() + if path_policy is not None: + artifacts_root = path_policy.artifacts_root + artifacts_root_raw = raw.get("artifacts_root") + if artifacts_root_raw: + configured_artifacts_root = path_policy.resolve_workspace_path( + artifacts_root_raw, + field_name="artifacts_root", + ) + path_policy.require_managed_root( + configured_artifacts_root, + artifacts_root, + field_name="artifacts_root", + ) else: - artifacts_root = artifacts_root.resolve() + artifacts_root = Path(raw.get("artifacts_root") or "artifacts").expanduser() + if not artifacts_root.is_absolute(): + artifacts_root = (ROOT / artifacts_root).resolve() + else: + artifacts_root = artifacts_root.resolve() results_dir_raw = raw.get("results_dir") - if results_dir_raw: + if path_policy is not None: + results_dir = path_policy.results_root + if results_dir_raw: + configured_results_dir = path_policy.resolve_output( + results_dir_raw, + field_name="results_dir", + ) + path_policy.require_managed_root( + configured_results_dir, + results_dir, + field_name="results_dir", + ) + elif results_dir_raw: results_dir = Path( _resolve_path( results_dir_raw, @@ -219,6 +268,18 @@ def load_runtime_context( suite_id = str(raw.get("suite") or datetime.now(timezone.utc).strftime("eval-%Y%m%dT%H%M%S")) _validate_identifier(suite_id, "suite") stages = _validate_pipeline_stages(pipeline_raw, stage_modules=stage_modules) + if path_policy is not None: + _validate_configured_path_fields( + stages, + cfg_path=cfg_path, + path_policy=path_policy, + ) + if target is not None and target.tools is not None and target.tools.toolset: + path_policy.resolve_input( + target.tools.toolset, + base_dir=cfg_path.parent, + field_name="pipeline.inference.target.tools.toolset", + ) if default_model_raw is not None: for stage_name, stage_cfg in stages: if stage_name in {"systematize", "test_set"} and "model" not in stage_cfg: @@ -281,11 +342,36 @@ def load_runtime_context( if context is not None and not isinstance(context, str): raise ValueError("context must be a string") - suite_root = (results_dir / suite_id).resolve() - _require_within(suite_root, results_dir, "suite_root") - run_root = (suite_root / run_id).resolve() if run_id else None - if run_root is not None: - _require_within(run_root, suite_root, "run_root") + if path_policy is not None: + suite_root = path_policy.resolve_managed_output( + results_dir / suite_id, + field_name="suite_root", + expected_root=results_dir, + reject_links=True, + ) + run_root = ( + path_policy.resolve_managed_output( + suite_root / run_id, + field_name="run_root", + expected_root=suite_root, + reject_links=True, + ) + if run_id + else None + ) + _validate_managed_stage_outputs( + stages, + stage_modules=stage_modules, + suite_root=suite_root, + run_root=run_root, + path_policy=path_policy, + ) + else: + suite_root = (results_dir / suite_id).resolve() + _require_within(suite_root, results_dir, "suite_root") + run_root = (suite_root / run_id).resolve() if run_id else None + if run_root is not None: + _require_within(run_root, suite_root, "run_root") return { "config_path": cfg_path, @@ -302,25 +388,95 @@ def load_runtime_context( "stages": stages, "target": target, "evaluation": pipeline.evaluation if pipeline else None, + "path_policy": path_policy, } +def _validate_configured_path_fields( + stages: list[tuple[str, dict[str, Any]]], + *, + cfg_path: Path, + path_policy: RuntimePathPolicy, +) -> None: + """Validate every explicitly configured stage path before execution.""" + for stage_name, stage_cfg in stages: + for key, value in stage_cfg.items(): + if not value or not key.endswith(("_path", "_dir")): + continue + field_name = f"pipeline.{stage_name}.{key}" + if key in OUTPUT_PATH_KEYS: + path_policy.resolve_output(value, field_name=field_name) + else: + path_policy.resolve_input( + value, + base_dir=cfg_path.parent, + field_name=field_name, + ) + + +def _validate_managed_stage_outputs( + stages: list[tuple[str, dict[str, Any]]], + *, + stage_modules: dict[str, Any], + suite_root: Path, + run_root: Path | None, + path_policy: RuntimePathPolicy, +) -> None: + """Confine explicit stage outputs to their suite or run.""" + for stage_name, stage_cfg in stages: + expected_root = ( + suite_root + if stage_modules[stage_name].SCOPE == "suite" + else run_root + ) + if expected_root is None: + continue + for key in OUTPUT_PATH_KEYS: + value = stage_cfg.get(key) + if not value: + continue + path_policy.resolve_managed_output( + value, + field_name=f"pipeline.{stage_name}.{key}", + expected_root=expected_root, + reject_links=True, + ) + + def resolve_stage_paths( cfg: dict[str, Any], *, cfg_path: Path, artifacts_root: Path, + path_policy: RuntimePathPolicy | None = None, + managed_output_root: Path | None = None, ) -> dict[str, Any]: """Resolve all *_path and *_dir values in one stage config mapping.""" resolved = dict(cfg) for key, value in list(resolved.items()): if not value or not key.endswith(("_path", "_dir")): continue + if ( + path_policy is not None + and managed_output_root is not None + and key in OUTPUT_PATH_KEYS + ): + resolved[key] = str( + path_policy.resolve_managed_output( + value, + field_name=key, + expected_root=managed_output_root, + reject_links=True, + ) + ) + continue resolved[key] = _resolve_path( value, artifacts_root=artifacts_root, cfg_dir=cfg_path.parent, use_artifacts_root=key in OUTPUT_PATH_KEYS, + path_policy=path_policy, + field_name=key, ) return resolved diff --git a/assert_ai/core/artifact_cache.py b/assert_ai/core/artifact_cache.py index 35f10e0e6..edbe28809 100644 --- a/assert_ai/core/artifact_cache.py +++ b/assert_ai/core/artifact_cache.py @@ -141,6 +141,56 @@ def supports_artifact_cache(ctx: dict[str, Any]) -> bool: return bool(ctx.get("suite_root") and ctx.get("config_path") and ctx.get("artifacts_root")) +def _managed_output_path( + ctx: dict[str, Any], + path: str | Path, + *, + field_name: str, + expected_root: str | Path | None = None, + reject_links: bool = False, +) -> Path: + policy = ctx.get("path_policy") + if policy is None: + return Path(path) + if expected_root is not None: + return policy.resolve_managed_output( + path, + field_name=field_name, + expected_root=expected_root, + reject_links=reject_links, + ) + return policy.resolve_output(path, field_name=field_name) + + +def _managed_suite_root(ctx: dict[str, Any]) -> Path: + policy = ctx.get("path_policy") + expected_root = policy.results_root if policy is not None else None + return _managed_output_path( + ctx, + ctx["suite_root"], + field_name="artifact cache suite root", + expected_root=expected_root, + reject_links=True, + ) + + +def _managed_input_path( + ctx: dict[str, Any], + path: str | Path, + *, + field_name: str, +) -> Path: + policy = ctx.get("path_policy") + if policy is None: + return Path(path) + config_path = Path(ctx["config_path"]) + return policy.resolve_input( + path, + base_dir=config_path.parent, + field_name=field_name, + ) + + def prepare_artifact_plan( *, ctx: dict[str, Any], @@ -152,15 +202,27 @@ def prepare_artifact_plan( if stage_name not in CACHEABLE_STAGES: raise ValueError(f"unsupported cacheable stage: {stage_name}") - suite_root = Path(ctx["suite_root"]) + suite_root = _managed_suite_root(ctx) fingerprint = build_artifact_fingerprint(ctx=ctx, stage_name=stage_name, raw_cfg=raw_cfg) - stage_root = suite_root / ARTIFACTS_DIR / stage_name + stage_root = _managed_output_path( + ctx, + suite_root / ARTIFACTS_DIR / stage_name, + field_name=f"{stage_name} artifact cache root", + expected_root=suite_root, + reject_links=True, + ) if not forced: match = _latest_matching_metadata(stage_name, stage_root, fingerprint.input_hash) if match is not None: version, metadata = match - artifact_dir = stage_root / version + artifact_dir = _managed_output_path( + ctx, + stage_root / version, + field_name=f"{stage_name} artifact cache version", + expected_root=stage_root, + reject_links=True, + ) return ArtifactPlan( stage_name=stage_name, version=version, @@ -258,8 +320,15 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: pipeline silently drift to stale legacy compatibility files. """ - suite_root = Path(ctx["suite_root"]) - latest = _load_json_object(suite_root / LATEST_FILE) + suite_root = _managed_suite_root(ctx) + latest_path = _managed_output_path( + ctx, + suite_root / LATEST_FILE, + field_name="artifact cache latest metadata", + expected_root=suite_root, + reject_links=True, + ) + latest = _load_json_object(latest_path, root=suite_root) artifacts = latest.get("artifacts") if isinstance(latest, dict) else None if not isinstance(artifacts, dict): return @@ -271,8 +340,20 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: version = ref.get("version") if not isinstance(version, str) or not version: continue - stage_root = suite_root / ARTIFACTS_DIR / stage_name - fallback_artifact_dir = stage_root / version + stage_root = _managed_output_path( + ctx, + suite_root / ARTIFACTS_DIR / stage_name, + field_name=f"{stage_name} artifact cache root", + expected_root=suite_root, + reject_links=True, + ) + fallback_artifact_dir = _managed_output_path( + ctx, + stage_root / version, + field_name=f"{stage_name} artifact cache version", + expected_root=stage_root, + reject_links=True, + ) resolved_artifact_dir = _resolve_ref_path(suite_root, ref.get("artifact_dir")) artifact_dir_fallback_used = ( resolved_artifact_dir is None or not resolved_artifact_dir.exists() @@ -280,6 +361,13 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: artifact_dir = ( fallback_artifact_dir if artifact_dir_fallback_used else resolved_artifact_dir ) + artifact_dir = _managed_output_path( + ctx, + artifact_dir, + field_name=f"{stage_name} artifact cache version", + expected_root=stage_root, + reject_links=True, + ) resolved_metadata_path = _resolve_ref_path( suite_root, ref.get("metadata_path") or ref.get("relative_metadata_path"), @@ -292,7 +380,14 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: if metadata_path_fallback_used else resolved_metadata_path ) - metadata = _load_json_object(metadata_path) + metadata_path = _managed_output_path( + ctx, + metadata_path, + field_name=f"{stage_name} artifact metadata", + expected_root=artifact_dir, + reject_links=True, + ) + metadata = _load_json_object(metadata_path, root=artifact_dir) if metadata and _metadata_outputs_exist(stage_name, artifact_dir, metadata): output_paths = _metadata_output_paths(stage_name, artifact_dir, metadata) # If the original ref's path entries pointed at locations that no @@ -364,8 +459,33 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: def finalize_artifact_plan(ctx: dict[str, Any], plan: ArtifactPlan) -> dict[str, Any]: """Write sidecar metadata and update latest/compatibility artifacts.""" - plan.artifact_dir.mkdir(parents=True, exist_ok=True) - file_hashes = _file_hashes(plan.output_paths) + suite_root = _managed_suite_root(ctx) + stage_root = _managed_output_path( + ctx, + suite_root / ARTIFACTS_DIR / plan.stage_name, + field_name=f"{plan.stage_name} artifact cache root", + expected_root=suite_root, + reject_links=True, + ) + artifact_dir = _managed_output_path( + ctx, + plan.artifact_dir, + field_name=f"{plan.stage_name} artifact cache version", + expected_root=stage_root, + reject_links=True, + ) + output_paths = { + key: _managed_output_path( + ctx, + path, + field_name=f"{plan.stage_name} artifact output '{key}'", + expected_root=artifact_dir, + reject_links=True, + ) + for key, path in plan.output_paths.items() + } + artifact_dir.mkdir(parents=True, exist_ok=True) + file_hashes = _file_hashes(output_paths) hashes: dict[str, Any] = { "config_hash": plan.fingerprint.config_hash, "input_hash": plan.fingerprint.input_hash, @@ -380,15 +500,22 @@ def finalize_artifact_plan(ctx: dict[str, Any], plan: ArtifactPlan) -> dict[str, "hashes": hashes, "inputs": plan.fingerprint.descriptor, "files": { - key: path.name for key, path in plan.output_paths.items() + key: path.name for key, path in output_paths.items() }, "file_hashes": file_hashes, } - write_json(plan.artifact_dir / ARTIFACT_METADATA_FILE, metadata) + metadata_path = _managed_output_path( + ctx, + artifact_dir / ARTIFACT_METADATA_FILE, + field_name=f"{plan.stage_name} artifact metadata", + expected_root=artifact_dir, + reject_links=True, + ) + write_json(metadata_path, metadata) ref = artifact_ref(ctx=ctx, plan=plan, metadata=metadata) ctx.setdefault("artifact_versions", {})[plan.stage_name] = ref update_latest(ctx, plan.stage_name, ref) - refresh_compatibility_files(ctx, plan.stage_name, plan.output_paths) + refresh_compatibility_files(ctx, plan.stage_name, output_paths) return ref @@ -420,7 +547,28 @@ def discard_artifact_plan(ctx: dict[str, Any], plan: ArtifactPlan) -> None: if plan.reused: return - artifact_dir = plan.artifact_dir + try: + suite_root = _managed_suite_root(ctx) + stage_root = _managed_output_path( + ctx, + suite_root / ARTIFACTS_DIR / plan.stage_name, + field_name=f"{plan.stage_name} artifact cache root", + expected_root=suite_root, + reject_links=True, + ) + artifact_dir = _managed_output_path( + ctx, + plan.artifact_dir, + field_name=f"{plan.stage_name} abandoned artifact cache version", + expected_root=stage_root, + reject_links=True, + ) + except ValueError as exc: + log.warning( + "[artifact-cache] refusing to clean up an unmanaged artifact path: %s", + exc, + ) + return if artifact_dir.exists() and artifact_dir.is_dir(): try: shutil.rmtree(artifact_dir) @@ -463,11 +611,31 @@ def refresh_compatibility_files( copy branch. """ - suite_root = Path(ctx["suite_root"]) + suite_root = _managed_suite_root(ctx) + stage_root = _managed_output_path( + ctx, + suite_root / ARTIFACTS_DIR / stage_name, + field_name=f"{stage_name} artifact cache root", + expected_root=suite_root, + reject_links=True, + ) for path in output_paths.values(): + path = _managed_output_path( + ctx, + path, + field_name=f"{stage_name} compatibility source", + expected_root=stage_root, + reject_links=True, + ) if not path.exists(): continue - dest = suite_root / path.name + dest = _managed_output_path( + ctx, + suite_root / path.name, + field_name=f"{stage_name} compatibility destination", + expected_root=suite_root, + reject_links=True, + ) if _is_local_edit(suite_root, stage_name, dest, path): log.warning( "[%s] Preserving local edits to %s: contents differ from the " @@ -531,7 +699,10 @@ def _was_cached_artifact( stage_root = suite_root / ARTIFACTS_DIR / stage_name for version_dir in _iter_version_dirs(stage_root): - metadata = _load_json_object(version_dir / ARTIFACT_METADATA_FILE) + metadata = _load_json_object( + version_dir / ARTIFACT_METADATA_FILE, + root=version_dir, + ) if not isinstance(metadata, dict): continue files_map = metadata.get("files") @@ -548,9 +719,18 @@ def _was_cached_artifact( def update_latest(ctx: dict[str, Any], stage_name: str, ref: dict[str, Any]) -> None: - suite_root = Path(ctx["suite_root"]) - latest_path = suite_root / LATEST_FILE - latest = _load_json_object(latest_path) or {"schema_version": 1, "artifacts": {}} + suite_root = _managed_suite_root(ctx) + latest_path = _managed_output_path( + ctx, + suite_root / LATEST_FILE, + field_name="artifact cache latest metadata", + expected_root=suite_root, + reject_links=True, + ) + latest = _load_json_object( + latest_path, + root=suite_root, + ) or {"schema_version": 1, "artifacts": {}} artifacts = latest.setdefault("artifacts", {}) if not isinstance(artifacts, dict): artifacts = {} @@ -567,7 +747,7 @@ def artifact_ref( ) -> dict[str, Any]: """Build the compact artifact reference stored in manifests/context.""" - suite_root = Path(ctx["suite_root"]) + suite_root = _managed_suite_root(ctx) primary_key = next(iter(_OUTPUT_FILES[plan.stage_name])) primary_path = plan.output_paths[primary_key] sidecar_path = plan.artifact_dir / ARTIFACT_METADATA_FILE @@ -603,7 +783,7 @@ def _ref_from_metadata( ) -> dict[str, Any]: """Build a ref payload from on-disk metadata (no plan/fingerprint needed).""" - suite_root = Path(ctx["suite_root"]) + suite_root = _managed_suite_root(ctx) sidecar_path = artifact_dir / ARTIFACT_METADATA_FILE hashes = metadata.get("hashes", {}) if isinstance(metadata, dict) else {} file_hashes = metadata.get("file_hashes", {}) if isinstance(metadata, dict) else {} @@ -758,7 +938,11 @@ def _artifact_or_file_dependency( if default_name: raw_path = str(Path(ctx["suite_root"]) / default_name) if isinstance(raw_path, str) and raw_path: - path = Path(raw_path) + path = _managed_input_path( + ctx, + raw_path, + field_name=f"{artifact_type} cache dependency", + ) if path.exists(): return { "path": str(path), @@ -803,7 +987,10 @@ def _latest_matching_metadata( ) -> tuple[str, dict[str, Any]] | None: matches: list[tuple[str, dict[str, Any]]] = [] for version_dir in _iter_version_dirs(stage_root): - metadata = _load_json_object(version_dir / ARTIFACT_METADATA_FILE) + metadata = _load_json_object( + version_dir / ARTIFACT_METADATA_FILE, + root=version_dir, + ) if not metadata: continue hashes = metadata.get("hashes") @@ -820,7 +1007,10 @@ def _recover_latest_valid_version( """Return the most recent intact version dir for a stage, if any.""" for version_dir in reversed(_iter_version_dirs(stage_root)): - metadata = _load_json_object(version_dir / ARTIFACT_METADATA_FILE) + metadata = _load_json_object( + version_dir / ARTIFACT_METADATA_FILE, + root=version_dir, + ) if metadata and _metadata_outputs_exist(stage_name, version_dir, metadata): return version_dir.name, version_dir, metadata return None @@ -855,7 +1045,11 @@ def _metadata_outputs_exist( return False for key in expected_keys: path = output_paths.get(key) - if path is None or not path.exists(): + if ( + path is None + or not _is_within(path.resolve(), version_dir.resolve()) + or not path.exists() + ): return False return True @@ -912,12 +1106,32 @@ def _allocate_version_dir(stage_root: Path) -> tuple[str, Path]: def _iter_version_dirs(stage_root: Path) -> list[Path]: if not stage_root.exists(): return [] + stage_root_resolved = stage_root.resolve() + version_dirs: list[Path] = [] + for path in stage_root.iterdir(): + if not re.fullmatch(r"v\d{4}", path.name) or not path.is_dir(): + continue + if not _is_within(path.resolve(), stage_root_resolved): + log.warning( + "Ignoring artifact version directory outside stage root: %s", + path, + ) + continue + version_dirs.append(path) return sorted( - [path for path in stage_root.iterdir() if path.is_dir() and re.fullmatch(r"v\d{4}", path.name)], + version_dirs, key=lambda path: path.name, ) +def _is_within(path: Path, root: Path) -> bool: + try: + path.relative_to(root) + return True + except ValueError: + return False + + def _resolve_ref_path(suite_root: Path, raw_path: Any) -> Path | None: if not isinstance(raw_path, str) or not raw_path: return None @@ -944,7 +1158,16 @@ def _resolve_ref_path(suite_root: Path, raw_path: Any) -> Path | None: raw_path, ) return None - return suite_root.joinpath(*parts) + resolved_path = suite_root.joinpath(*parts).resolve() + try: + resolved_path.relative_to(suite_root_resolved) + except ValueError: + log.warning( + "Refusing to resolve cache reference outside suite root: %r", + raw_path, + ) + return None + return resolved_path def _relative_to_suite(path: Path, suite_root: Path) -> str: @@ -1007,7 +1230,17 @@ def _file_hashes(output_paths: dict[str, Path]) -> dict[str, str]: return hashes -def _load_json_object(path: Path) -> dict[str, Any] | None: +def _load_json_object( + path: Path, + *, + root: Path | None = None, +) -> dict[str, Any] | None: + if root is not None: + resolved_path = path.resolve() + if not _is_within(resolved_path, root.resolve()): + log.warning("Refusing to read JSON outside expected root: %s", path) + return None + path = resolved_path try: text = path.read_text(encoding="utf-8") except FileNotFoundError: diff --git a/assert_ai/core/azure_auth.py b/assert_ai/core/azure_auth.py index 0fda84635..16ab6cdfd 100644 --- a/assert_ai/core/azure_auth.py +++ b/assert_ai/core/azure_auth.py @@ -241,7 +241,7 @@ def provider() -> str: # no extra function calls in the hot path) once warmed. # # Resolution is deliberately *lazy*: process entrypoints that load -# ``.env`` (the runner, ``assert-ai init``) call +# ``.env`` (``assert-ai run``, ``assert-ai init``, or MCP launch) call # ``refresh_azure_auth_mode(force=True)`` after ``load_dotenv`` so the # resolved mode reflects the dotenv-populated environment, not just # the shell vars present at module import. diff --git a/assert_ai/core/environment.py b/assert_ai/core/environment.py new file mode 100644 index 000000000..f889c6197 --- /dev/null +++ b/assert_ai/core/environment.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Explicit process-environment bootstrap for command-line entry points.""" + +from __future__ import annotations + +import sys +from pathlib import Path + + +def bootstrap_environment( + *, + env_file: Path | None = None, + discover_from_cwd: bool = False, +) -> None: + """Load one dotenv source, then refresh environment-sensitive caches.""" + if env_file is not None and discover_from_cwd: + raise ValueError("env_file and discover_from_cwd are mutually exclusive") + + from dotenv import find_dotenv, load_dotenv + + dotenv_path: str | Path | None = env_file + if discover_from_cwd: + dotenv_path = find_dotenv(usecwd=True) or None + if dotenv_path is not None: + load_dotenv(dotenv_path, override=False) + + from assert_ai.core.azure_auth import refresh_azure_auth_mode + + refresh_azure_auth_mode(force=True) + + model_client = sys.modules.get("assert_ai.core.model_client") + if model_client is not None: + refresh = getattr(model_client, "refresh_environment_settings", None) + if callable(refresh): + refresh() diff --git a/assert_ai/core/model_client.py b/assert_ai/core/model_client.py index 62476c241..8cfa88b4d 100644 --- a/assert_ai/core/model_client.py +++ b/assert_ai/core/model_client.py @@ -1784,7 +1784,7 @@ async def _call() -> Any: return result -# ── Import-time AZURE_API_BASE normalization ─────────────────── +# ── Environment-driven client settings ──────────────────────── # LiteLLM appends the OpenAI-API path itself (``/openai/deployments/…`` # or ``/openai/v1/responses``), so AZURE_API_BASE must be the bare # account endpoint. A trailing ``/openai/...`` or ``/openai/v1/...`` @@ -1816,18 +1816,20 @@ def _normalize_azure_api_base() -> None: os.environ["AZURE_API_BASE"] = normalized -_normalize_azure_api_base() +def refresh_environment_settings() -> None: + """Apply environment-driven client settings after explicit dotenv loading.""" + _normalize_azure_api_base() + if os.environ.get("ASSERT_PREFER_CHAT_COMPLETIONS", "").strip().lower() in ( + "1", + "true", + "yes", + ): + _activate_chat_completions_fallback( + "ASSERT_PREFER_CHAT_COMPLETIONS env var set", + proactive=True, + ) -# ── Import-time env-var seed ─────────────────────────────────── -# Users in Azure regions known to lack Responses API support -# (e.g. West Europe at time of writing) can pre-arm the fallback -# by exporting ``ASSERT_PREFER_CHAT_COMPLETIONS=1``. This avoids -# the one wasted Responses API round-trip + the user-visible WARN -# on every cold start, while keeping the reactive fallback as a -# safety net for regions that lose support later. -if os.environ.get("ASSERT_PREFER_CHAT_COMPLETIONS", "").strip().lower() in ("1", "true", "yes"): - _activate_chat_completions_fallback( - "ASSERT_PREFER_CHAT_COMPLETIONS env var set", - proactive=True, - ) +# Preserve direct-library-import behavior while allowing entry points that load +# dotenv later to refresh the same settings explicitly. +refresh_environment_settings() diff --git a/assert_ai/core/otel_session.py b/assert_ai/core/otel_session.py index ec407c008..9e17fc35c 100644 --- a/assert_ai/core/otel_session.py +++ b/assert_ai/core/otel_session.py @@ -24,7 +24,7 @@ import uuid from contextlib import nullcontext from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from assert_ai.core.async_utils import invoke_callable from assert_ai.core.collector import SpanCollector @@ -39,6 +39,9 @@ ) from assert_ai.core.session import TurnResult +if TYPE_CHECKING: + from assert_ai.core.runtime_path_policy import RuntimePathPolicy + class OTelTracedSession: """Session that invokes a callable target and captures OTel traces per turn. @@ -77,6 +80,7 @@ def __init__( max_events_per_turn: int = 50, live_otel: bool = False, config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, ) -> None: self._callable_ref = callable_ref self._collector = collector @@ -85,6 +89,7 @@ def __init__( self._message_timeout_s = message_timeout_s self._max_events_per_turn = max_events_per_turn self._config_path = config_path + self._path_policy = path_policy self._callable: Any = None self._supports_history = False self._session_id = "" @@ -130,7 +135,11 @@ async def open(self) -> None: _orig_stdout = sys.stdout sys.stdout = io.StringIO() try: - mod = import_callable_module(module_path, config_path=self._config_path) + mod = import_callable_module( + module_path, + config_path=self._config_path, + path_policy=self._path_policy, + ) finally: sys.stdout = _orig_stdout try: diff --git a/assert_ai/core/runtime_path_policy.py b/assert_ai/core/runtime_path_policy.py new file mode 100644 index 000000000..a80c3bca1 --- /dev/null +++ b/assert_ai/core/runtime_path_policy.py @@ -0,0 +1,469 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Workspace-aware runtime path resolution and containment policy.""" + +from __future__ import annotations + +import os +import stat +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path +from typing import Iterable + + +class RuntimePathErrorCode(StrEnum): + """Stable machine-readable categories for runtime path failures.""" + + INVALID_ROOT = "invalid_root" + OUTSIDE_CONFIG_ROOT = "outside_config_root" + OUTSIDE_INPUT_ROOT = "outside_input_root" + OUTSIDE_WORKSPACE = "outside_workspace" + OUTSIDE_ARTIFACTS_ROOT = "outside_artifacts_root" + OUTSIDE_EXPECTED_ROOT = "outside_expected_root" + MANAGED_ROOT_OVERRIDE = "managed_root_override" + MANAGED_PATH_LINK = "managed_path_link" + PATH_NOT_FOUND = "path_not_found" + NOT_A_FILE = "not_a_file" + + +class RuntimePathError(ValueError): + """Typed path-policy failure suitable for CLI and MCP error mapping.""" + + def __init__( + self, + code: RuntimePathErrorCode, + message: str, + *, + field_name: str, + path: Path | None = None, + expected_root: Path | None = None, + ) -> None: + super().__init__(message) + self.code = code + self.field_name = field_name + self.path = path + self.expected_root = expected_root + + +def _is_within(path: Path, root: Path) -> bool: + try: + path.relative_to(root) + return True + except ValueError: + return False + + +def _resolved(path: str | Path) -> Path: + return Path(path).expanduser().resolve() + + +def _deduplicate_paths(paths: Iterable[Path]) -> tuple[Path, ...]: + unique: list[Path] = [] + for path in paths: + if path not in unique: + unique.append(path) + return tuple(unique) + + +@dataclass(frozen=True, slots=True) +class RuntimePathPolicy: + """Resolve runtime paths against explicit workspace roots.""" + + workspace_root: Path + config_root: Path + artifacts_root: Path + results_root: Path + additional_read_roots: tuple[Path, ...] = () + allow_absolute_inputs: bool = False + force_managed_outputs: bool = True + + def __post_init__(self) -> None: + workspace_root = _resolved(self.workspace_root) + if not workspace_root.is_dir(): + raise RuntimePathError( + RuntimePathErrorCode.INVALID_ROOT, + f"Workspace root is not a directory: {workspace_root}", + field_name="workspace_root", + path=workspace_root, + ) + + config_root = _resolved(self.config_root) + artifacts_root = _resolved(self.artifacts_root) + results_root = _resolved(self.results_root) + additional_read_roots = _deduplicate_paths( + _resolved(root) for root in self.additional_read_roots + ) + + for field_name, root in ( + ("config_root", config_root), + ("artifacts_root", artifacts_root), + ("results_root", results_root), + ): + if not _is_within(root, workspace_root): + raise RuntimePathError( + RuntimePathErrorCode.INVALID_ROOT, + f"{field_name} must be inside workspace_root", + field_name=field_name, + path=root, + expected_root=workspace_root, + ) + + if self.force_managed_outputs and not _is_within(results_root, artifacts_root): + raise RuntimePathError( + RuntimePathErrorCode.INVALID_ROOT, + "results_root must be inside artifacts_root", + field_name="results_root", + path=results_root, + expected_root=artifacts_root, + ) + + object.__setattr__(self, "workspace_root", workspace_root) + object.__setattr__(self, "config_root", config_root) + object.__setattr__(self, "artifacts_root", artifacts_root) + object.__setattr__(self, "results_root", results_root) + object.__setattr__(self, "additional_read_roots", additional_read_roots) + + @property + def read_roots(self) -> tuple[Path, ...]: + return _deduplicate_paths( + ( + self.workspace_root, + self.config_root, + self.artifacts_root, + *self.additional_read_roots, + ) + ) + + def resolve_config_path( + self, + path: str | Path, + *, + must_exist: bool = False, + ) -> Path: + """Resolve a config path strictly under ``config_root``.""" + candidate = Path(path).expanduser() + if candidate.is_absolute(): + resolved = candidate.resolve() + else: + parts = candidate.parts + if parts and parts[0] == self.config_root.name: + candidate = Path(*parts[1:]) if len(parts) > 1 else Path() + resolved = (self.config_root / candidate).resolve() + self._require_within( + resolved, + self.config_root, + field_name="config", + code=RuntimePathErrorCode.OUTSIDE_CONFIG_ROOT, + ) + self._require_kind( + resolved, + field_name="config", + must_exist=must_exist, + file_only=must_exist, + ) + return resolved + + def resolve_input( + self, + path: str | Path, + *, + base_dir: Path, + field_name: str, + must_exist: bool = False, + file_only: bool = False, + ) -> Path: + """Resolve an input path without allowing relative root escapes.""" + candidate = Path(path).expanduser() + if candidate.is_absolute(): + resolved = candidate.resolve() + if not self.allow_absolute_inputs: + self._require_within_any_read_root(resolved, field_name=field_name) + else: + artifact_relative = self._artifact_relative(candidate) + root = self.artifacts_root if artifact_relative is not None else _resolved(base_dir) + self._require_within_any_read_root(root, field_name=f"{field_name} base directory") + suffix = artifact_relative if artifact_relative is not None else candidate + resolved = (root / suffix).resolve() + self._require_within( + resolved, + root, + field_name=field_name, + code=RuntimePathErrorCode.OUTSIDE_INPUT_ROOT, + ) + self._require_kind( + resolved, + field_name=field_name, + must_exist=must_exist, + file_only=file_only, + ) + return resolved + + def resolve_output( + self, + path: str | Path, + *, + field_name: str, + ) -> Path: + """Resolve an output path under the managed artifacts root.""" + resolved = self._output_candidate(path).resolve() + if self.force_managed_outputs: + self._require_within( + resolved, + self.artifacts_root, + field_name=field_name, + code=RuntimePathErrorCode.OUTSIDE_ARTIFACTS_ROOT, + ) + return resolved + + def resolve_managed_output( + self, + path: str | Path, + *, + field_name: str, + expected_root: str | Path, + reject_links: bool = False, + ) -> Path: + """Resolve an output within one operation-specific managed root.""" + expected_candidate = Path(expected_root).expanduser() + if not expected_candidate.is_absolute(): + expected_candidate = self._output_candidate(expected_candidate) + expected = expected_candidate.resolve() + raw_candidate = Path(path).expanduser() + if ( + raw_candidate.is_absolute() + or self._artifact_relative(raw_candidate) is not None + ): + candidate = self._output_candidate(raw_candidate) + else: + candidate = expected / raw_candidate + self._require_within( + expected, + self.artifacts_root, + field_name=f"{field_name} expected root", + code=RuntimePathErrorCode.OUTSIDE_ARTIFACTS_ROOT, + ) + resolved = candidate.resolve() + self._require_within( + resolved, + self.artifacts_root, + field_name=field_name, + code=RuntimePathErrorCode.OUTSIDE_ARTIFACTS_ROOT, + ) + self._require_within( + resolved, + expected, + field_name=field_name, + code=RuntimePathErrorCode.OUTSIDE_EXPECTED_ROOT, + ) + if reject_links: + self._require_no_links( + expected_candidate, + self.artifacts_root, + field_name=f"{field_name} expected root", + ) + self._require_no_links( + candidate, + expected, + field_name=field_name, + ) + return resolved + + def resolve_workspace_path( + self, + path: str | Path, + *, + field_name: str, + must_exist: bool = False, + file_only: bool = False, + ) -> Path: + """Resolve a path relative to the workspace and keep it contained.""" + candidate = Path(path).expanduser() + resolved = ( + candidate.resolve() + if candidate.is_absolute() + else (self.workspace_root / candidate).resolve() + ) + self.require_workspace_path(resolved, field_name=field_name) + self._require_kind( + resolved, + field_name=field_name, + must_exist=must_exist, + file_only=file_only, + ) + return resolved + + def require_managed_tree( + self, + path: str | Path, + *, + field_name: str, + expected_root: str | Path, + ) -> Path: + """Reject links or junctions anywhere in an existing managed tree.""" + root = self.resolve_managed_output( + path, + field_name=field_name, + expected_root=expected_root, + reject_links=True, + ) + if not root.is_dir(): + return root + for current_root, dir_names, file_names in os.walk( + root, + followlinks=False, + ): + current = Path(current_root) + for name in (*dir_names, *file_names): + self.resolve_managed_output( + current / name, + field_name=f"{field_name} entry", + expected_root=root, + reject_links=True, + ) + return root + + def require_workspace_path(self, path: str | Path, *, field_name: str) -> Path: + """Re-resolve and require a path to remain inside the workspace.""" + resolved = _resolved(path) + self._require_within( + resolved, + self.workspace_root, + field_name=field_name, + code=RuntimePathErrorCode.OUTSIDE_WORKSPACE, + ) + return resolved + + def module_search_roots(self, config_path: Path | None) -> tuple[tuple[str, Path], ...]: + """Return the only roots strict dynamic imports may add to ``sys.path``.""" + roots: list[tuple[str, Path]] = [] + if config_path is not None: + config_dir = self.require_workspace_path( + config_path.parent, + field_name="config module root", + ) + roots.append(("Relative to config", config_dir)) + if self.workspace_root not in {root for _, root in roots}: + roots.append(("Relative to workspace", self.workspace_root)) + return tuple(roots) + + def require_managed_root( + self, + configured: Path, + expected: Path, + *, + field_name: str, + ) -> None: + """Reject a config root override that differs from the managed root.""" + if configured != expected: + raise RuntimePathError( + RuntimePathErrorCode.MANAGED_ROOT_OVERRIDE, + f"{field_name} is managed by the runtime and cannot be overridden", + field_name=field_name, + path=configured, + expected_root=expected, + ) + + def _artifact_relative(self, path: Path) -> Path | None: + parts = path.parts + if not parts or parts[0] not in {"artifacts", self.artifacts_root.name}: + return None + return Path(*parts[1:]) if len(parts) > 1 else Path() + + def _output_candidate(self, path: str | Path) -> Path: + candidate = Path(path).expanduser() + if candidate.is_absolute(): + return candidate + artifact_relative = self._artifact_relative(candidate) + suffix = artifact_relative if artifact_relative is not None else candidate + return self.artifacts_root / suffix + + @staticmethod + def _require_no_links( + path: Path, + root: Path, + *, + field_name: str, + ) -> None: + normalized = Path(os.path.abspath(path)) + try: + relative = normalized.relative_to(root) + except ValueError: + return + current = root + for part in relative.parts: + current /= part + is_junction = getattr(current, "is_junction", None) + is_reparse_point = False + if os.name == "nt": + try: + attributes = os.lstat(current).st_file_attributes + except (AttributeError, FileNotFoundError, OSError): + attributes = 0 + is_reparse_point = bool( + attributes & stat.FILE_ATTRIBUTE_REPARSE_POINT + ) + if ( + current.is_symlink() + or (callable(is_junction) and is_junction()) + or is_reparse_point + ): + raise RuntimePathError( + RuntimePathErrorCode.MANAGED_PATH_LINK, + f"{field_name} cannot traverse a symbolic link or junction", + field_name=field_name, + path=current, + expected_root=root, + ) + + def _require_within_any_read_root(self, path: Path, *, field_name: str) -> None: + if any(_is_within(path, root) for root in self.read_roots): + return + raise RuntimePathError( + RuntimePathErrorCode.OUTSIDE_INPUT_ROOT, + f"{field_name} is outside the configured read roots", + field_name=field_name, + path=path, + ) + + @staticmethod + def _require_within( + path: Path, + root: Path, + *, + field_name: str, + code: RuntimePathErrorCode, + ) -> None: + if _is_within(path, root): + return + raise RuntimePathError( + code, + f"{field_name} escapes its expected root directory", + field_name=field_name, + path=path, + expected_root=root, + ) + + @staticmethod + def _require_kind( + path: Path, + *, + field_name: str, + must_exist: bool, + file_only: bool, + ) -> None: + if must_exist and not path.exists(): + raise RuntimePathError( + RuntimePathErrorCode.PATH_NOT_FOUND, + f"{field_name} does not exist: {path}", + field_name=field_name, + path=path, + ) + if file_only and path.exists() and not path.is_file(): + raise RuntimePathError( + RuntimePathErrorCode.NOT_A_FILE, + f"{field_name} is not a file: {path}", + field_name=field_name, + path=path, + ) diff --git a/assert_ai/core/security.py b/assert_ai/core/security.py index 0fccaf148..ea536c479 100644 --- a/assert_ai/core/security.py +++ b/assert_ai/core/security.py @@ -16,9 +16,12 @@ import socket import sys from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from urllib.parse import urlparse +if TYPE_CHECKING: + from assert_ai.core.runtime_path_policy import RuntimePathPolicy + log = logging.getLogger(__name__) @@ -67,12 +70,21 @@ def validate_module_ref(module_ref: str, *, config_path: Path | None = None) -> ) -def validate_sys_path_addition(path: Path, *, config_path: Path | None = None) -> None: +def validate_sys_path_addition( + path: Path, + *, + config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> None: """Validate that a sys.path addition is scoped to the workspace. Only allows paths that are within the config directory or current working directory. Raises ValueError for paths outside the expected workspace. """ + if path_policy is not None: + path_policy.require_workspace_path(path, field_name="module search path") + return + resolved = path.resolve() cwd = Path.cwd().resolve() diff --git a/assert_ai/core/session.py b/assert_ai/core/session.py index 5c21651fe..be9ca25dc 100644 --- a/assert_ai/core/session.py +++ b/assert_ai/core/session.py @@ -13,7 +13,7 @@ import re from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal from assert_ai.core.async_utils import invoke_callable from assert_ai.core.model_client import ( @@ -31,6 +31,9 @@ from assert_ai.core.tool_backend import load_tool_module from assert_ai.core.tools import build_target_tools +if TYPE_CHECKING: + from assert_ai.core.runtime_path_policy import RuntimePathPolicy + log = logging.getLogger(__name__) # Regex patterns for common credential formats in plain text @@ -480,11 +483,13 @@ def __init__( system_prompt: str | None = None, message_timeout_s: float | None = None, config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, ) -> None: self._callable_ref = callable_ref self._system_prompt = system_prompt self._message_timeout_s = message_timeout_s self._config_path = config_path + self._path_policy = path_policy self._callable = None self._supports_history = False @@ -498,7 +503,11 @@ async def open(self) -> None: validate_callable_ref(self._callable_ref) module_path, func_name = self._callable_ref.rsplit(":", 1) - mod = import_callable_module(module_path, config_path=self._config_path) + mod = import_callable_module( + module_path, + config_path=self._config_path, + path_policy=self._path_policy, + ) try: self._callable = getattr(mod, func_name) except AttributeError as exc: @@ -789,11 +798,18 @@ def __init__( startup_timeout_s: float | None = None, message_timeout_s: float | None = None, config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, ) -> None: from assert_ai.core.security import validate_module_ref validate_module_ref(connector_ref, config_path=config_path) - connector_cls = _discover_connector_class(load_tool_module(connector_ref, config_path=config_path)) + connector_cls = _discover_connector_class( + load_tool_module( + connector_ref, + config_path=config_path, + path_policy=path_policy, + ) + ) self._startup_timeout_s = startup_timeout_s self._message_timeout_s = message_timeout_s self._connector = connector_cls(scenario) diff --git a/assert_ai/core/tool_backend.py b/assert_ai/core/tool_backend.py index 3c5c97982..2e11d4b8d 100644 --- a/assert_ai/core/tool_backend.py +++ b/assert_ai/core/tool_backend.py @@ -5,22 +5,34 @@ from __future__ import annotations +import builtins import contextlib import importlib +import importlib.abc +import importlib.machinery import importlib.util import inspect import json import sys +import threading import types import uuid from hashlib import sha1 from pathlib import Path -from typing import Any, Union, get_args, get_origin, get_type_hints +from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hints from assert_ai.core.async_utils import invoke_callable +if TYPE_CHECKING: + from assert_ai.core.runtime_path_policy import RuntimePathPolicy -def _search_roots(config_path: Path | None) -> list[tuple[str, Path]]: + +def _search_roots( + config_path: Path | None, + path_policy: RuntimePathPolicy | None = None, +) -> list[tuple[str, Path]]: + if path_policy is not None: + return list(path_policy.module_search_roots(config_path)) roots: list[tuple[str, Path]] = [] if config_path is not None: roots.append(("Relative to config", config_path.parent.resolve())) @@ -28,34 +40,341 @@ def _search_roots(config_path: Path | None) -> list[tuple[str, Path]]: return roots -def _module_path_candidates(module_ref: str, *, config_path: Path | None) -> list[tuple[str, Path]]: +def _module_path_candidates( + module_ref: str, + *, + config_path: Path | None, + path_policy: RuntimePathPolicy | None = None, +) -> list[tuple[str, Path]]: dotted = Path(*module_ref.split(".")) file_name = dotted.with_suffix(".py") package_init = dotted / "__init__.py" candidates: list[tuple[str, Path]] = [] - for label, root in _search_roots(config_path): + for label, root in _search_roots(config_path, path_policy): for candidate in (root / file_name, root / package_init): if candidate.exists(): candidates.append((label, candidate)) return candidates -def _load_module_from_file(module_ref: str, path: Path) -> Any: +class _WorkspaceSourceLoader(importlib.machinery.SourceFileLoader): + def __init__( + self, + fullname: str, + path: str, + finder: "_WorkspaceModuleFinder", + ) -> None: + super().__init__(fullname, path) + self._finder = finder + + def exec_module(self, module: types.ModuleType) -> None: + workspace_builtins = dict(vars(builtins)) + workspace_builtins["__import__"] = self._finder.import_module + module.__dict__["__builtins__"] = workspace_builtins + super().exec_module(module) + + +class _WorkspaceModuleFinder(importlib.abc.MetaPathFinder): + def __init__( + self, + *, + namespace: str, + roots: tuple[Path, ...], + path_policy: RuntimePathPolicy, + ) -> None: + self.namespace = namespace + self.roots = roots + self.path_policy = path_policy + self._importlib_proxy = types.ModuleType("importlib") + self._importlib_proxy.__dict__.update(importlib.__dict__) + self._importlib_proxy.import_module = self.import_dynamic_module + + def find_spec( + self, + fullname: str, + path: Any = None, + target: Any = None, + ) -> importlib.machinery.ModuleSpec | None: + del target + prefix = f"{self.namespace}." + if not fullname.startswith(prefix): + return None + original_name = fullname[len(prefix):] + leaf_name = original_name.rsplit(".", 1)[-1] + search_dirs = ( + tuple(Path(value) for value in path) + if path is not None + else self.roots + ) + namespace_dirs: list[Path] = [] + for search_dir in search_dirs: + module_path = self._validated_path( + search_dir / f"{leaf_name}.py", + original_name, + ) + if module_path.is_file(): + loader = _WorkspaceSourceLoader( + fullname, + str(module_path), + self, + ) + return importlib.util.spec_from_file_location( + fullname, + module_path, + loader=loader, + ) + + package_dir = self._validated_path( + search_dir / leaf_name, + original_name, + ) + package_init = self._validated_path( + package_dir / "__init__.py", + original_name, + ) + if package_init.is_file(): + loader = _WorkspaceSourceLoader( + fullname, + str(package_init), + self, + ) + return importlib.util.spec_from_file_location( + fullname, + package_init, + loader=loader, + submodule_search_locations=[str(package_dir)], + ) + if package_dir.is_dir(): + namespace_dirs.append(package_dir) + + if namespace_dirs: + spec = importlib.machinery.ModuleSpec( + fullname, + loader=None, + is_package=True, + ) + spec.submodule_search_locations = [ + str(directory) for directory in namespace_dirs + ] + return spec + return None + + def import_module( + self, + name: str, + globals: dict[str, Any] | None = None, + locals: dict[str, Any] | None = None, + fromlist: tuple[str, ...] | list[str] | None = (), + level: int = 0, + ) -> Any: + from_items = fromlist or () + if level != 0 or not name: + return builtins.__import__( + name, + globals, + locals, + fromlist, + level, + ) + if name == "importlib": + for item in from_items: + if ( + isinstance(item, str) + and item != "*" + and not hasattr(self._importlib_proxy, item) + ): + try: + child = importlib.import_module(f"importlib.{item}") + except ModuleNotFoundError: + continue + setattr(self._importlib_proxy, item, child) + return self._importlib_proxy + if name.startswith("importlib."): + imported = builtins.__import__( + name, + globals, + locals, + fromlist, + level, + ) + if from_items: + return imported + child_name = name.split(".", 1)[1].split(".", 1)[0] + child = sys.modules.get(f"importlib.{child_name}") + if child is not None: + setattr(self._importlib_proxy, child_name, child) + return self._importlib_proxy + + top_level = name.split(".", 1)[0] + if not self._top_level_exists(top_level): + return builtins.__import__( + name, + globals, + locals, + fromlist, + level, + ) + + mapped_name = f"{self.namespace}.{name}" + module = importlib.import_module(mapped_name) + if from_items and getattr(module, "__path__", None) is not None: + for item in from_items: + if not isinstance(item, str) or item == "*": + continue + child_name = f"{mapped_name}.{item}" + if importlib.util.find_spec(child_name) is not None: + importlib.import_module(child_name) + if from_items: + return module + return importlib.import_module(f"{self.namespace}.{top_level}") + + def import_dynamic_module( + self, + name: str, + package: str | None = None, + ) -> types.ModuleType: + if name.startswith("."): + mapped_package = package + if ( + package + and not package.startswith(f"{self.namespace}.") + and self._top_level_exists(package.split(".", 1)[0]) + ): + mapped_package = f"{self.namespace}.{package}" + return importlib.import_module(name, mapped_package) + + top_level = name.split(".", 1)[0] + if self._top_level_exists(top_level): + return importlib.import_module(f"{self.namespace}.{name}") + return importlib.import_module(name, package) + + def _top_level_exists(self, name: str) -> bool: + for root in self.roots: + if ( + self._validated_path(root / f"{name}.py", name).is_file() + or self._validated_path(root / name / "__init__.py", name).is_file() + or self._validated_path(root / name, name).is_dir() + ): + return True + return False + + def _validated_path(self, path: Path, module_name: str) -> Path: + return self.path_policy.resolve_workspace_path( + path, + field_name=f"workspace module '{module_name}'", + ) + + +_WORKSPACE_FINDERS: dict[str, _WorkspaceModuleFinder] = {} +_WORKSPACE_FINDER_LOCK = threading.Lock() + + +def _workspace_module_namespace(root: Path) -> str: + root_digest = sha1(str(root.resolve()).encode("utf-8")).hexdigest() + return f"_assert_ai_workspace_{root_digest}" + + +def _ensure_workspace_finder( + *, + primary_root: Path, + path_policy: RuntimePathPolicy, +) -> _WorkspaceModuleFinder: + primary_root = path_policy.require_workspace_path( + primary_root, + field_name="module search root", + ) + roots = [primary_root] + if path_policy.workspace_root != primary_root: + roots.append(path_policy.workspace_root) + namespace = _workspace_module_namespace(primary_root) + + with _WORKSPACE_FINDER_LOCK: + finder = _WORKSPACE_FINDERS.get(namespace) + if finder is None: + finder = _WorkspaceModuleFinder( + namespace=namespace, + roots=tuple(roots), + path_policy=path_policy, + ) + _WORKSPACE_FINDERS[namespace] = finder + sys.meta_path.insert(0, finder) + spec = importlib.machinery.ModuleSpec( + namespace, + loader=None, + is_package=True, + ) + spec.submodule_search_locations = [str(root) for root in roots] + package = importlib.util.module_from_spec(spec) + sys.modules[namespace] = package + return finder + + +def _load_module_from_file( + module_ref: str, + path: Path, + *, + isolated_workspace_import: bool = False, + package_root: Path | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> Any: + if isolated_workspace_import: + if package_root is None or path_policy is None: + raise ValueError( + "An isolated workspace import requires a package root and path policy" + ) + finder = _ensure_workspace_finder( + primary_root=package_root, + path_policy=path_policy, + ) + module = importlib.import_module(f"{finder.namespace}.{module_ref}") + module_file = getattr(module, "__file__", None) + if module_file is None or Path(module_file).resolve() != path.resolve(): + raise ValueError( + f"Workspace module '{module_ref}' resolved to an unexpected source" + ) + return module + module_name = f"_assert_ai_module_{sha1(str(path).encode('utf-8')).hexdigest()}" - spec = importlib.util.spec_from_file_location(module_name, path) + existing = sys.modules.get(module_name) + if existing is not None: + existing_file = getattr(existing, "__file__", None) + if existing_file and Path(existing_file).resolve() == path.resolve(): + return existing + raise ValueError( + f"Module name '{module_name}' is already loaded from a different path" + ) + spec_kwargs = ( + {"submodule_search_locations": [str(path.parent)]} + if path.name == "__init__.py" + else {} + ) + spec = importlib.util.spec_from_file_location(module_name, path, **spec_kwargs) if spec is None or spec.loader is None: raise ValueError(f"Could not load module '{module_ref}' from {path}") module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module - spec.loader.exec_module(module) + try: + spec.loader.exec_module(module) + except BaseException: + sys.modules.pop(module_name, None) + raise return module @contextlib.contextmanager -def _temporary_sys_path(path: Path, *, config_path: Path | None = None): +def _temporary_sys_path( + path: Path, + *, + config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, +): from assert_ai.core.security import validate_sys_path_addition - validate_sys_path_addition(path, config_path=config_path) + validate_sys_path_addition( + path, + config_path=config_path, + path_policy=path_policy, + ) sys.path.insert(0, str(path)) try: yield @@ -73,6 +392,32 @@ def _is_direct_module_path(module_ref: str) -> bool: return module_ref.endswith((".py", "/__init__.py", "\\__init__.py")) +def _direct_workspace_module( + path: Path, + *, + config_path: Path | None, + path_policy: RuntimePathPolicy, +) -> tuple[str, Path]: + for _, root in path_policy.module_search_roots(config_path): + try: + relative = path.relative_to(root) + except ValueError: + continue + if relative.name == "__init__.py": + relative = relative.parent + else: + relative = relative.with_suffix("") + if not relative.parts: + continue + if any(not part.isidentifier() for part in relative.parts): + raise ValueError( + "Strict direct module paths must map to a dotted Python " + f"module name inside the workspace; got {path}" + ) + return ".".join(relative.parts), root + raise ValueError(f"Direct module path is outside the configured module roots: {path}") + + def _module_classes(module: Any) -> list[type[Any]]: return [ member @@ -81,59 +426,133 @@ def _module_classes(module: Any) -> list[type[Any]]: ] -def load_tool_module(module_ref: str, *, config_path: Path | None = None) -> Any: +def load_tool_module( + module_ref: str, + *, + config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> Any: from assert_ai.core.security import validate_module_ref validate_module_ref(module_ref, config_path=config_path) direct_path = Path(module_ref).expanduser() if _is_direct_module_path(module_ref): - if not direct_path.is_absolute() and config_path is not None: - direct_path = (config_path.parent / direct_path).resolve() + if not direct_path.is_absolute(): + if config_path is not None: + direct_path = (config_path.parent / direct_path).resolve() + elif path_policy is not None: + direct_path = (path_policy.workspace_root / direct_path).resolve() if not direct_path.exists(): raise ValueError(f"Tool module path does not exist: {direct_path}") # Validate direct path is within workspace - _validate_module_file_path(direct_path, config_path=config_path) + direct_path = _validate_module_file_path( + direct_path, + config_path=config_path, + path_policy=path_policy, + ) + if path_policy is not None: + isolated_ref, package_root = _direct_workspace_module( + direct_path, + config_path=config_path, + path_policy=path_policy, + ) + return _load_module_from_file( + isolated_ref, + direct_path, + isolated_workspace_import=True, + package_root=package_root, + path_policy=path_policy, + ) return _load_module_from_file(module_ref, direct_path) - return _smart_import(module_ref, config_path=config_path, kind="tool module") + return _smart_import( + module_ref, + config_path=config_path, + path_policy=path_policy, + kind="tool module", + ) def _smart_import( module_ref: str, *, config_path: Path | None, + path_policy: RuntimePathPolicy | None, kind: str, ) -> Any: - """Import ``module_ref`` with workspace-aware sys.path fallback. + """Import ``module_ref`` with workspace-aware fallback. - Resolution order: + Legacy resolution order: 1. Standard import via ``sys.path``. 2. Retry with the config directory temporarily on ``sys.path`` (if known). 3. Retry with the current working directory temporarily on ``sys.path``. 4. Direct file load via ``spec_from_file_location`` for ``/.py`` or ``//__init__.py`` under each search root. - On failure, raises ``ValueError`` listing every location that was searched. + With a strict path policy, only verified source files under the configured + workspace are loaded. On failure, raises ``ValueError`` listing every + location that was searched. ``kind`` is used only to make the error message specific (e.g. ``"tool module"``, ``"callable module"``). """ + if path_policy is not None: + module_parts = module_ref.split(".") + if not module_parts or any(not part.isidentifier() for part in module_parts): + raise ValueError( + f"Strict workspace imports require a dotted Python module name; " + f"got {module_ref!r}" + ) + attempted: list[str] = [] + dotted = Path(*module_parts) + for label, root in _search_roots(config_path, path_policy): + attempted.append(f"{len(attempted) + 1}. {label}: {root}") + for candidate in (root / dotted.with_suffix(".py"), root / dotted / "__init__.py"): + if not candidate.exists(): + continue + candidate = path_policy.resolve_workspace_path( + candidate, + field_name=kind, + must_exist=True, + file_only=True, + ) + return _load_module_from_file( + module_ref, + candidate, + isolated_workspace_import=True, + package_root=root, + path_policy=path_policy, + ) + searched = "\n ".join(attempted) + raise ValueError( + f"Could not import {kind} '{module_ref}' inside the configured workspace.\n" + f"Searched:\n {searched}", + ) + try: return importlib.import_module(module_ref) except ModuleNotFoundError as exc: if not _has_missing_target(exc, module_ref): raise attempted = ["1. Python path (sys.path)"] - for label, root in _search_roots(config_path): + for label, root in _search_roots(config_path, path_policy): attempted.append(f"{len(attempted) + 1}. {label}: {root}") - with _temporary_sys_path(root, config_path=config_path): + with _temporary_sys_path( + root, + config_path=config_path, + path_policy=path_policy, + ): try: return importlib.import_module(module_ref) except ModuleNotFoundError as retry_exc: if not _has_missing_target(retry_exc, module_ref): raise - for label, candidate in _module_path_candidates(module_ref, config_path=config_path): + for label, candidate in _module_path_candidates( + module_ref, + config_path=config_path, + path_policy=path_policy, + ): attempted.append(f"{len(attempted) + 1}. Direct file load: {candidate}") return _load_module_from_file(module_ref, candidate) searched = "\n ".join(attempted) @@ -144,7 +563,12 @@ def _smart_import( ) from exc -def import_callable_module(module_ref: str, *, config_path: Path | None = None) -> Any: +def import_callable_module( + module_ref: str, + *, + config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> Any: """Import the module portion of a ``module.path:function`` callable reference. Uses the same workspace-aware fallback as :func:`load_tool_module` so that a @@ -153,18 +577,36 @@ def import_callable_module(module_ref: str, *, config_path: Path | None = None) elsewhere. The caller is expected to have validated ``module_ref`` via :func:`assert_ai.core.security.validate_callable_ref` first. """ - return _smart_import(module_ref, config_path=config_path, kind="callable module") + return _smart_import( + module_ref, + config_path=config_path, + path_policy=path_policy, + kind="callable module", + ) -def _validate_module_file_path(path: Path, *, config_path: Path | None = None) -> None: +def _validate_module_file_path( + path: Path, + *, + config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> Path: """Validate that a direct module file path is within the workspace.""" + if path_policy is not None: + return path_policy.resolve_workspace_path( + path, + field_name="tool module path", + must_exist=True, + file_only=True, + ) + resolved = path.resolve() cwd = Path.cwd().resolve() # Allow paths within cwd try: resolved.relative_to(cwd) - return + return resolved except ValueError: pass @@ -173,7 +615,7 @@ def _validate_module_file_path(path: Path, *, config_path: Path | None = None) - config_dir = config_path.parent.resolve() try: resolved.relative_to(config_dir) - return + return resolved except ValueError: pass @@ -422,8 +864,17 @@ async def resolve( ) -def inspect_tool_module(module_ref: str, *, config_path: Path | None = None) -> tuple[type[Any], list[dict[str, Any]]]: - module = load_tool_module(module_ref, config_path=config_path) +def inspect_tool_module( + module_ref: str, + *, + config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> tuple[type[Any], list[dict[str, Any]]]: + module = load_tool_module( + module_ref, + config_path=config_path, + path_policy=path_policy, + ) tools_cls = _discover_tools_class(module) return tools_cls, _derive_tool_schemas(tools_cls) diff --git a/assert_ai/core/workspace.py b/assert_ai/core/workspace.py new file mode 100644 index 000000000..0f066c660 --- /dev/null +++ b/assert_ai/core/workspace.py @@ -0,0 +1,69 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Workspace layout and safe path references for application services.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + +from assert_ai.core.runtime_path_policy import RuntimePathPolicy + + +@dataclass(frozen=True, slots=True) +class WorkspaceService: + """Canonical workspace roots shared by MCP-facing services.""" + + root: Path + configs_root: Path + artifacts_root: Path + results_root: Path + path_policy: RuntimePathPolicy + + @classmethod + def create( + cls, + root: str | Path, + *, + additional_read_roots: Iterable[str | Path] = (), + ) -> "WorkspaceService": + workspace_root = Path(root).expanduser().resolve(strict=True) + configs_root = workspace_root / "evals" + artifacts_root = workspace_root / "artifacts" + results_root = artifacts_root / "results" + policy = RuntimePathPolicy( + workspace_root=workspace_root, + config_root=configs_root, + artifacts_root=artifacts_root, + results_root=results_root, + additional_read_roots=tuple(Path(path) for path in additional_read_roots), + allow_absolute_inputs=False, + force_managed_outputs=True, + ) + return cls( + root=workspace_root, + configs_root=policy.config_root, + artifacts_root=policy.artifacts_root, + results_root=policy.results_root, + path_policy=policy, + ) + + def resolve_file(self, path: str | Path, *, field_name: str) -> Path: + """Resolve an existing workspace-contained file.""" + return self.path_policy.resolve_workspace_path( + path, + field_name=field_name, + must_exist=True, + file_only=True, + ) + + def reference(self, path: str | Path) -> str: + """Return a workspace-relative, forward-slash reference.""" + resolved = self.path_policy.require_workspace_path( + path, + field_name="workspace reference", + ) + relative = resolved.relative_to(self.root) + return "." if not relative.parts else relative.as_posix() diff --git a/assert_ai/init/_command.py b/assert_ai/init/_command.py index 80df00f77..66658d350 100644 --- a/assert_ai/init/_command.py +++ b/assert_ai/init/_command.py @@ -125,24 +125,18 @@ def init( The assistant asks clarifying questions about your agent/system, eval goals, and constraints, then proposes a complete eval.yaml. """ - from dotenv import load_dotenv + from assert_ai.core.environment import bootstrap_environment + + bootstrap_environment(env_file=env_file if env_file.exists() else None) + from rich.console import Console from rich.syntax import Syntax from assert_ai.init._design_agent import run_design_loop from assert_ai.init._emit import emit_config - # Load env vars for LLM credentials - if env_file.exists(): - load_dotenv(env_file, override=False) - - # Refresh the Azure auth mode now that ``.env`` has populated the - # environment so the init LLM call respects the dotenv-supplied - # ``ASSERT_AZURE_USE_AAD`` / ``AZURE_API_KEY`` rather than whichever - # state was frozen at module import. - from assert_ai.core.azure_auth import log_resolved_azure_auth_mode, refresh_azure_auth_mode + from assert_ai.core.azure_auth import log_resolved_azure_auth_mode - refresh_azure_auth_mode(force=True) # Emit the resolved auth-mode line AFTER refresh so it reflects the # value the upcoming LLM call will use, and AFTER the parent ``cli`` # group has configured logging so ``--quiet``/``--output json`` apply. diff --git a/assert_ai/mcp/_command.py b/assert_ai/mcp/_command.py index 2cac131b6..19b8ed2e5 100644 --- a/assert_ai/mcp/_command.py +++ b/assert_ai/mcp/_command.py @@ -11,6 +11,8 @@ import click +from assert_ai.core.environment import bootstrap_environment +from assert_ai.core.workspace import WorkspaceService from assert_ai.mcp.models import CapabilityGroup, ServerMode _INSTALL_HINT = 'Install the MCP dependencies with: python -m pip install "assert-ai[mcp]"' @@ -66,12 +68,34 @@ def mcp() -> None: multiple=True, help="Enable an additional capability group. Repeat as needed.", ) -def serve(workspace: Path, mode: str, enabled_groups: tuple[str, ...]) -> None: +@click.option( + "--env-file", + type=click.Path(dir_okay=False, path_type=Path), + default=None, + help="Optional dotenv file contained within --workspace. No file is discovered by default.", +) +def serve( + workspace: Path, + mode: str, + enabled_groups: tuple[str, ...], + env_file: Path | None, +) -> None: """Serve ASSERT over stdio; stdout is reserved for MCP protocol traffic.""" + try: + workspace_service = WorkspaceService.create(workspace) + if env_file is not None: + resolved_env_file = workspace_service.resolve_file( + env_file, + field_name="--env-file", + ) + bootstrap_environment(env_file=resolved_env_file) + except (OSError, ValueError) as exc: + raise click.ClickException(str(exc)) from exc + server_module = _load_server_module() try: options = server_module.ServerOptions.create( - workspace_root=workspace, + workspace_root=workspace_service.root, mode=mode, enabled_groups=enabled_groups, ) diff --git a/assert_ai/mcp/server.py b/assert_ai/mcp/server.py index d3c1aaf3f..6bedd730e 100644 --- a/assert_ai/mcp/server.py +++ b/assert_ai/mcp/server.py @@ -5,7 +5,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Iterable @@ -13,6 +13,8 @@ from mcp.server import MCPServer from mcp.types import ToolAnnotations +from assert_ai.core.runtime_path_policy import RuntimePathPolicy +from assert_ai.core.workspace import WorkspaceService from assert_ai.mcp.models import ( CapabilityGroup, ServerInfo, @@ -51,6 +53,12 @@ class ServerOptions: workspace_root: Path mode: ServerMode = ServerMode.INSPECT enabled_groups: tuple[CapabilityGroup, ...] = () + workspace: WorkspaceService = field(init=False, repr=False) + + def __post_init__(self) -> None: + workspace = WorkspaceService.create(self.workspace_root) + object.__setattr__(self, "workspace_root", workspace.root) + object.__setattr__(self, "workspace", workspace) @classmethod def create( @@ -60,9 +68,6 @@ def create( mode: str | ServerMode = ServerMode.INSPECT, enabled_groups: Iterable[str | CapabilityGroup] = (), ) -> "ServerOptions": - root = Path(workspace_root).expanduser().resolve(strict=True) - if not root.is_dir(): - raise ValueError(f"Workspace is not a directory: {root}") parsed_mode = ServerMode(mode) parsed_groups = tuple(CapabilityGroup(group) for group in enabled_groups) invalid_groups = _AUTHOR_EXTENSION_GROUPS.intersection(parsed_groups) @@ -72,11 +77,15 @@ def create( f"Capability group(s) {names} require --mode author or --mode full." ) return cls( - workspace_root=root, + workspace_root=Path(workspace_root), mode=parsed_mode, enabled_groups=parsed_groups, ) + @property + def path_policy(self) -> RuntimePathPolicy: + return self.workspace.path_policy + @property def capability_groups(self) -> tuple[CapabilityGroup, ...]: groups = {*_MODE_GROUPS[self.mode], *self.enabled_groups} @@ -112,7 +121,12 @@ def get_server_info() -> ServerInfo: server_version=_server_version(), mode=options.mode, enabled_capability_groups=list(options.capability_groups), - workspace=WorkspaceInfo(), + workspace=WorkspaceInfo( + root=options.workspace.reference(options.workspace.root), + configs_root=options.workspace.reference(options.workspace.configs_root), + artifacts_root=options.workspace.reference(options.workspace.artifacts_root), + results_root=options.workspace.reference(options.workspace.results_root), + ), ) return server diff --git a/assert_ai/runner.py b/assert_ai/runner.py index 9e98264e9..858c99f47 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -15,10 +15,9 @@ import warnings from datetime import datetime, timezone from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import yaml -from dotenv import find_dotenv, load_dotenv from assert_ai.config import ( ConfigError, @@ -38,7 +37,6 @@ supports_artifact_cache, update_latest, ) -from assert_ai.core.azure_auth import refresh_azure_auth_mode from assert_ai.core.config_model import RunManifest, SuiteMetadata from assert_ai.core.io import write_json from assert_ai.core.model_client import ( @@ -57,18 +55,8 @@ from assert_ai.display import label_metric from assert_ai.stages import STAGES -# Walk up from cwd so the user's project `.env` is found when assert-ai is -# installed as a wheel. Bare `load_dotenv()` walks up from this file's -# directory, which lives inside the venv's site-packages and misses the -# project `.env`. -load_dotenv(find_dotenv(usecwd=True)) - -# Force-resolve the Azure auth mode now that ``.env`` has populated the -# environment. Without this, ``model_client``'s lazy resolution would -# fire on the first request (which is fine) — but doing it here lets -# entrypoints log the resolved mode up-front and surfaces missing -# ``azure-identity`` early. -refresh_azure_auth_mode(force=True) +if TYPE_CHECKING: + from assert_ai.core.runtime_path_policy import RuntimePathPolicy log = logging.getLogger(__name__) @@ -113,11 +101,21 @@ def _load_context( *, config: str, overrides: list[str] | None = None, + path_policy: RuntimePathPolicy | None = None, ) -> dict[str, Any]: """Load one config file into runtime context.""" - cfg_path = Path(config).resolve() + cfg_path = ( + path_policy.resolve_config_path(config, must_exist=True) + if path_policy is not None + else Path(config).resolve() + ) raw = _apply_config_overrides(load_config(cfg_path), overrides) - return load_runtime_context(raw, cfg_path, stage_modules=STAGES) + return load_runtime_context( + raw, + cfg_path, + stage_modules=STAGES, + path_policy=path_policy, + ) def _write_suite_metadata(ctx: dict[str, Any]) -> None: @@ -589,8 +587,12 @@ def run_pipeline( strict: bool = False, overrides: list[str] | None = None, concurrency: int | None = None, + path_policy: RuntimePathPolicy | None = None, ) -> int: - """Execute the configured stages sequentially and persist suite/run metadata.""" + """Execute configured stages. + + Programmatic callers are responsible for any desired dotenv bootstrap. + """ # Suppress litellm's internal async logging warnings — they fire because # litellm creates async coroutines for logging callbacks that never get # awaited in our synchronous runner context. Harmless but alarming. @@ -615,7 +617,11 @@ def run_pipeline( _install_async_cleanup_filters() try: - ctx = _load_context(config=config, overrides=overrides) + ctx = _load_context( + config=config, + overrides=overrides, + path_policy=path_policy, + ) ctx["strict"] = strict except (ConfigError, ValueError) as exc: log.error(f"[config error] {exc}") @@ -667,6 +673,15 @@ def run_pipeline( requested_force_stages = requested_force_stages.union(cascade) suite_root = Path(ctx["suite_root"]) + path_policy = ctx.get("path_policy") + if path_policy is not None: + suite_root = path_policy.resolve_managed_output( + suite_root, + field_name="suite root", + expected_root=path_policy.results_root, + reject_links=True, + ) + ctx["suite_root"] = suite_root suite_root.mkdir(parents=True, exist_ok=True) _write_suite_metadata(ctx) ctx.setdefault("artifact_versions", {}) @@ -722,6 +737,14 @@ def run_pipeline( stages_to_run.append((stage_name, module, raw_cfg)) run_root = Path(ctx["run_root"]) if ctx.get("run_root") else None + if run_root is not None and path_policy is not None: + run_root = path_policy.resolve_managed_output( + run_root, + field_name="run root", + expected_root=suite_root, + reject_links=True, + ) + ctx["run_root"] = run_root selected_run_stage = any(module.SCOPE == "run" for _, module, _ in stages_to_run) manifest = None if selected_run_stage and run_root is not None: @@ -729,7 +752,22 @@ def run_pipeline( manifest = _build_manifest(ctx) config_path = ctx.get("config_path") if config_path is not None and Path(config_path).is_file(): - shutil.copy2(config_path, run_root / "config.yaml") + config_path = ( + path_policy.resolve_config_path(config_path, must_exist=True) + if path_policy is not None + else Path(config_path) + ) + config_snapshot = ( + path_policy.resolve_managed_output( + run_root / "config.yaml", + field_name="run config snapshot", + expected_root=run_root, + reject_links=True, + ) + if path_policy is not None + else run_root / "config.yaml" + ) + shutil.copy2(config_path, config_snapshot) failed_stage: str | None = None pipeline_start = time.monotonic() stage_usage: dict[str, dict[str, Any]] = {} diff --git a/assert_ai/stages/inference.py b/assert_ai/stages/inference.py index 63dc71e66..6243f9b83 100644 --- a/assert_ai/stages/inference.py +++ b/assert_ai/stages/inference.py @@ -15,7 +15,7 @@ import traceback import uuid from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any log = logging.getLogger(__name__) @@ -67,6 +67,9 @@ from assert_ai.stages.test_set import TOOL_SOURCE_PER_TEST_CASE, TOOL_SOURCE_RUNTIME from assert_ai.viewer_read_model import build_run_viewer_artifacts +if TYPE_CHECKING: + from assert_ai.core.runtime_path_policy import RuntimePathPolicy + SCOPE = "run" SUITE_OUTPUT = None @@ -454,6 +457,7 @@ def _build_hosted_session( synthetic_prompt_template: str, tool_timeout_s: float | None = None, startup_timeout_s: float | None = None, + path_policy: RuntimePathPolicy | None = None, ) -> HostedSession: if not tools_config: return HostedSession( @@ -468,7 +472,11 @@ def _build_hosted_session( if module_ref is not None: if not isinstance(module_ref, str) or not module_ref.strip(): raise ValueError("tool-module tools require module") - tools_cls, schemas = inspect_tool_module(module_ref, config_path=config_path) + tools_cls, schemas = inspect_tool_module( + module_ref, + config_path=config_path, + path_policy=path_policy, + ) return HostedSession( model=model, generate_options=generate_options, @@ -492,8 +500,17 @@ def _build_hosted_session( if tools is None: if not isinstance(toolset_path, str) or not toolset_path.strip(): raise ValueError("simulated tools require target.tools.toolset or per-test-case tools") - resolved_path = Path(toolset_path).expanduser() - if not resolved_path.is_absolute(): + if path_policy is not None: + resolved_path = path_policy.resolve_input( + toolset_path, + base_dir=config_path.parent if config_path is not None else path_policy.config_root, + field_name="pipeline.inference.target.tools.toolset", + must_exist=True, + file_only=True, + ) + else: + resolved_path = Path(toolset_path).expanduser() + if path_policy is None and not resolved_path.is_absolute(): candidates = [] if config_path is not None: candidates.append((config_path.parent / resolved_path).resolve()) @@ -524,6 +541,7 @@ def _build_target_session( inference: InferenceConfig, max_tokens: int, config_path: Path | None, + path_policy: RuntimePathPolicy | None = None, call_label: str | None = None, ) -> HostedSession | ExternalSession | CallableSession | HTTPEndpointSession: """Create the runtime session for one test-case inference.""" @@ -549,12 +567,14 @@ def _build_target_session( group_by=target.trace.group_by, live_otel=True, config_path=config_path, + path_policy=path_policy, ) return CallableSession( callable_ref=target.callable, system_prompt=target.system_prompt, message_timeout_s=inference.tool_timeout_s, config_path=config_path, + path_policy=path_policy, ) if target.is_external: @@ -566,6 +586,7 @@ def _build_target_session( startup_timeout_s=inference.startup_timeout_s, message_timeout_s=inference.tool_timeout_s, config_path=config_path, + path_policy=path_policy, ) if not target.model: @@ -594,6 +615,7 @@ def _build_target_session( synthetic_prompt_template=TOOL_SIM_PROMPT, tool_timeout_s=inference.tool_timeout_s, startup_timeout_s=inference.startup_timeout_s, + path_policy=path_policy, ) @@ -604,6 +626,7 @@ async def _run_prompt_test_case( inference: InferenceConfig, max_tokens: int, config_path: Path | None, + path_policy: RuntimePathPolicy | None = None, ) -> Transcript: """Run one prompt test case against the target runtime.""" test_case_payload = test_case.get("seed") @@ -616,6 +639,7 @@ async def _run_prompt_test_case( inference=inference, max_tokens=max_tokens, config_path=config_path, + path_policy=path_policy, call_label=f"target:{test_case_id}", ) target_id = str(target.model.name) if target.model else (target.connector or target.callable or target.endpoint or "") @@ -914,6 +938,7 @@ async def _run_scenario_test_case( evaluation: EvaluationConfig, max_tokens: int, config_path: Path | None, + path_policy: RuntimePathPolicy | None = None, ) -> Transcript: """Run one scenario test case and capture its transcript.""" tester = evaluation.tester @@ -928,6 +953,7 @@ async def _run_scenario_test_case( inference=evaluation.inference, max_tokens=max_tokens, config_path=config_path, + path_policy=path_policy, call_label=f"target:{test_case_id}", ) transcript = Transcript( @@ -1003,6 +1029,9 @@ async def run_inference( target: TargetConfig, evaluation: EvaluationConfig | None = None, config_path: Path | None = None, + path_policy: RuntimePathPolicy | None = None, + managed_output_root: Path | None = None, + managed_test_set_root: Path | None = None, strict: bool = False, forced: bool = False, heartbeat: Any = None, @@ -1028,7 +1057,21 @@ async def run_inference( elif target.tools is not None and target.tools.simulator and not target.tools.toolset: raise ValueError("runtime tool_source requires target.tools.toolset when target.tools.simulator is set") fixed_system_prompt = str(target.system_prompt or "").strip() or None - resolved_test_set_path = resolve_path(test_set_path) + resolved_test_set_path = ( + path_policy.resolve_input( + test_set_path, + base_dir=( + config_path.parent + if config_path is not None + else path_policy.config_root + ), + field_name="inference test set", + must_exist=True, + file_only=True, + ) + if path_policy is not None + else resolve_path(test_set_path) + ) canonical_rows = normalize_test_case_rows(load_test_cases(resolved_test_set_path, strict=strict)) test_cases = _prepare_test_cases( canonical_rows, @@ -1042,15 +1085,54 @@ async def run_inference( # invalidate the recorded file_hashes in artifact.json. rewrite_test_set_path = False if rewrite_test_set_path: + if path_policy is not None: + if managed_test_set_root is None: + raise ValueError( + "managed_test_set_root is required to rewrite a test set " + "when a runtime path policy is active" + ) + resolved_test_set_path = path_policy.resolve_managed_output( + resolved_test_set_path, + field_name="canonicalized test set", + expected_root=managed_test_set_root, + reject_links=True, + ) write_jsonl(resolved_test_set_path, canonical_rows) resolved_run_id = str(run_id or uuid.uuid4().hex[:8]).lower() - out_dir = resolve_path(save_dir or (Path("artifacts/outputs") / resolved_run_id)) + if path_policy is not None: + if managed_output_root is None: + raise ValueError( + "managed_output_root is required when a runtime path policy is active" + ) + out_dir = path_policy.resolve_managed_output( + save_dir or (Path("artifacts/outputs") / resolved_run_id), + field_name="inference output directory", + expected_root=managed_output_root, + reject_links=True, + ) + else: + out_dir = resolve_path( + save_dir or (Path("artifacts/outputs") / resolved_run_id) + ) out_dir.mkdir(parents=True, exist_ok=True) + if path_policy is not None: + path_policy.require_managed_tree( + out_dir, + field_name="inference output directory", + expected_root=managed_output_root, + ) resolved_max_tokens = max_tokens if max_tokens is not None else DEFAULT_INFERENCE_MAX_TOKENS inference = evaluation.inference if evaluation is not None else InferenceConfig() indexed_test_cases = list(enumerate(test_cases)) inference_set_path = out_dir / INFERENCE_SET_FILE + if path_policy is not None: + inference_set_path = path_policy.resolve_managed_output( + inference_set_path, + field_name="inference output", + expected_root=out_dir, + reject_links=True, + ) # Resume: load already-completed test_case_ids and skip them. completed_test_case_ids: set[str] = set() @@ -1061,6 +1143,13 @@ async def run_inference( test_set_path=resolved_test_set_path, ) config_hash_path = out_dir / _INFERENCE_CONFIG_HASH_FILE + if path_policy is not None: + config_hash_path = path_policy.resolve_managed_output( + config_hash_path, + field_name="inference config hash", + expected_root=out_dir, + reject_links=True, + ) if inference_set_path.exists(): if forced: # User explicitly forced this stage (directly or via the runner's @@ -1115,6 +1204,7 @@ async def _worker(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: inference=inference, max_tokens=resolved_max_tokens, config_path=config_path, + path_policy=path_policy, ) elif kind == "scenario": if evaluation is None: @@ -1125,6 +1215,7 @@ async def _worker(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: evaluation=evaluation, max_tokens=resolved_max_tokens, config_path=config_path, + path_policy=path_policy, ) else: raise ValueError(f"unsupported test case type: {kind}") @@ -1196,6 +1287,13 @@ async def _guard(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: results.append(result) inference_row = result.get("inference_row") if inference_row is not None: + if path_policy is not None: + inference_set_path = path_policy.resolve_managed_output( + inference_set_path, + field_name="inference output", + expected_root=out_dir, + reject_links=True, + ) append_jsonl_row(inference_set_path, inference_row) error = result.get("error") # Scenario inferences catch target exceptions mid-conversation and @@ -1322,6 +1420,12 @@ async def _guard(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: "the target raised an exception mid-conversation", target_error_count, ) + if path_policy is not None: + path_policy.require_managed_tree( + out_dir, + field_name="inference output directory", + expected_root=managed_output_root, + ) build_run_viewer_artifacts(out_dir) return { @@ -1352,12 +1456,27 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: }, cfg_path=ctx["config_path"], artifacts_root=ctx["artifacts_root"], + path_policy=ctx.get("path_policy"), + managed_output_root=Path(ctx["run_root"]), ) test_set_artifact_ref = (ctx.get("artifact_versions") or {}).get("test_set") # Only rewrite the test_set file when there is no cached artifact to protect. # If the user supplied an explicit test_set_path AND we have no cache ref, we # still want the canonicalization pass to normalize their input file. rewrite_test_set_path = not isinstance(test_set_artifact_ref, dict) + path_policy = ctx.get("path_policy") + if path_policy is not None and rewrite_test_set_path: + try: + Path(cfg["test_set_path"]).relative_to(Path(ctx["suite_root"])) + except ValueError: + rewrite_test_set_path = False + else: + path_policy.resolve_managed_output( + cfg["test_set_path"], + field_name="canonicalized test set", + expected_root=Path(ctx["suite_root"]), + reject_links=True, + ) result = await run_inference( test_set_path=cfg["test_set_path"], save_dir=cfg["save_dir"], @@ -1366,6 +1485,9 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: target=ctx["target"], evaluation=ctx.get("evaluation"), config_path=ctx["config_path"], + path_policy=path_policy, + managed_output_root=Path(ctx["run_root"]), + managed_test_set_root=Path(ctx["suite_root"]), strict=cfg.get("strict", False), forced=bool(ctx.get("_stage_forced", False)), heartbeat=ctx.get("_heartbeat") if isinstance(ctx, dict) else None, diff --git a/assert_ai/stages/judge.py b/assert_ai/stages/judge.py index 90511328d..8341ff24e 100644 --- a/assert_ai/stages/judge.py +++ b/assert_ai/stages/judge.py @@ -11,7 +11,7 @@ import logging import traceback from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any log = logging.getLogger(__name__) @@ -27,6 +27,9 @@ from assert_ai.core.transcript import Transcript, TranscriptEvent, TranscriptMetadata from assert_ai.viewer_read_model import build_run_viewer_artifacts +if TYPE_CHECKING: + from assert_ai.core.runtime_path_policy import RuntimePathPolicy + SCOPE = "run" SUITE_OUTPUT = None @@ -95,6 +98,9 @@ async def run_judge( disabled_dimensions: list[str] | None = None, forced: bool = False, heartbeat: Any = None, + path_policy: RuntimePathPolicy | None = None, + config_path: Path | None = None, + managed_output_root: Path | None = None, ) -> dict[str, Any]: """Score inference rows and write score artifacts.""" judge_model = str(evaluation.judge.model.name) @@ -112,16 +118,62 @@ async def run_judge( if disabled_dimensions is not None else getattr(evaluation.judge, "disabled_dimensions", []) ) - resolved_inference_set_path = resolve_path(inference_set_path) + resolved_inference_set_path = ( + path_policy.resolve_input( + inference_set_path, + base_dir=( + config_path.parent + if config_path is not None + else path_policy.config_root + ), + field_name="judge inference set", + must_exist=True, + file_only=True, + ) + if path_policy is not None + else resolve_path(inference_set_path) + ) rows = load_jsonl(resolved_inference_set_path) if not rows: raise ValueError(f"No inference rows found in {inference_set_path}") - out_dir = resolve_path(save_dir or str(resolved_inference_set_path.parent)) + if path_policy is not None: + if managed_output_root is None: + raise ValueError( + "managed_output_root is required when a runtime path policy is active" + ) + out_dir = path_policy.resolve_managed_output( + save_dir or str(resolved_inference_set_path.parent), + field_name="judge output directory", + expected_root=managed_output_root, + reject_links=True, + ) + else: + out_dir = resolve_path(save_dir or str(resolved_inference_set_path.parent)) out_dir.mkdir(parents=True, exist_ok=True) + if path_policy is not None: + path_policy.require_managed_tree( + out_dir, + field_name="judge output directory", + expected_root=managed_output_root, + ) if not taxonomy_path: raise ValueError("judge stage requires taxonomy_path") - resolved_taxonomy_path = resolve_path(taxonomy_path) + resolved_taxonomy_path = ( + path_policy.resolve_input( + taxonomy_path, + base_dir=( + config_path.parent + if config_path is not None + else path_policy.config_root + ), + field_name="judge taxonomy", + must_exist=True, + file_only=True, + ) + if path_policy is not None + else resolve_path(taxonomy_path) + ) if not resolved_taxonomy_path.exists(): raise ValueError(f"Taxonomy file not found: {taxonomy_path}") try: @@ -339,6 +391,13 @@ async def worker(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: } scores_path = out_dir / SCORES_FILE + if path_policy is not None: + scores_path = path_policy.resolve_managed_output( + scores_path, + field_name="judge scores output", + expected_root=out_dir, + reject_links=True, + ) # Resume: load already-scored (kind, test_case_id) pairs and skip them, but only # if the judge configuration and inference-set file hasn't changed since the @@ -357,6 +416,13 @@ async def worker(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: inference_set_path=resolved_inference_set_path, ) config_hash_path = out_dir / _JUDGE_CONFIG_HASH_FILE + if path_policy is not None: + config_hash_path = path_policy.resolve_managed_output( + config_hash_path, + field_name="judge config hash", + expected_root=out_dir, + reject_links=True, + ) if scores_path.exists(): if forced: # User explicitly forced this stage (directly or via the runner's @@ -416,6 +482,13 @@ async def guard(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: result = await completed_task score = result.get("score_row") if score is not None: + if path_policy is not None: + scores_path = path_policy.resolve_managed_output( + scores_path, + field_name="judge scores output", + expected_root=out_dir, + reject_links=True, + ) append_jsonl_row(scores_path, score) written_rows += 1 error = result.get("error") @@ -434,6 +507,12 @@ async def guard(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: # Always rebuild viewer artifacts so the on-disk read model reflects the # current scores.jsonl, even when a row failed and we are about to raise. + if path_policy is not None: + path_policy.require_managed_tree( + out_dir, + field_name="judge output directory", + expected_root=managed_output_root, + ) build_run_viewer_artifacts(out_dir) # Per-row failures should not kill the stage as long as *some* rows # succeeded. The errors are surfaced via judge_failures in the @@ -507,6 +586,8 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, str]: }, cfg_path=ctx["config_path"], artifacts_root=ctx["artifacts_root"], + path_policy=ctx.get("path_policy"), + managed_output_root=Path(ctx["run_root"]), ) result = await run_judge( inference_set_path=cfg["inference_set_path"], @@ -517,6 +598,9 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, str]: disabled_dimensions=disabled_dimensions, forced=bool(ctx.get("_stage_forced", False)), heartbeat=ctx.get("_heartbeat") if isinstance(ctx, dict) else None, + path_policy=ctx.get("path_policy"), + config_path=Path(ctx["config_path"]), + managed_output_root=Path(ctx["run_root"]), ) return { "scores_path": result["scores_path"], diff --git a/assert_ai/stages/stratification.py b/assert_ai/stages/stratification.py index 91a0db5fc..eec87c109 100644 --- a/assert_ai/stages/stratification.py +++ b/assert_ai/stages/stratification.py @@ -396,6 +396,8 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: }, cfg_path=ctx["config_path"], artifacts_root=ctx["artifacts_root"], + path_policy=ctx.get("path_policy"), + managed_output_root=suite_root, ) result = await run_stratification( diff --git a/assert_ai/stages/systematize.py b/assert_ai/stages/systematize.py index bf5ed1c24..d3e54c980 100644 --- a/assert_ai/stages/systematize.py +++ b/assert_ai/stages/systematize.py @@ -167,6 +167,8 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: {"save_dir": save_dir}, cfg_path=ctx["config_path"], artifacts_root=ctx["artifacts_root"], + path_policy=ctx.get("path_policy"), + managed_output_root=Path(ctx["suite_root"]), ) behavior_name = ctx.get("behavior_name") or "behavior" diff --git a/assert_ai/stages/test_set.py b/assert_ai/stages/test_set.py index f774a29d4..b7f853939 100644 --- a/assert_ai/stages/test_set.py +++ b/assert_ai/stages/test_set.py @@ -1335,6 +1335,8 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: path_cfg, cfg_path=ctx["config_path"], artifacts_root=ctx["artifacts_root"], + path_policy=ctx.get("path_policy"), + managed_output_root=Path(ctx["suite_root"]), ) taxonomy_path = cfg["taxonomy_path"] stratification_dir = Path(cfg["save_path"]).parent diff --git a/scripts/benchmark.py b/scripts/benchmark.py index dbebbc132..63488b808 100644 --- a/scripts/benchmark.py +++ b/scripts/benchmark.py @@ -54,6 +54,10 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) +from assert_ai.core.environment import bootstrap_environment # noqa: E402 + +bootstrap_environment(discover_from_cwd=True) + from assert_ai.runner import run_pipeline # noqa: E402 from assert_ai.logging_config import configure_logging # noqa: E402 diff --git a/tests/test_artifact_cache.py b/tests/test_artifact_cache.py index 210082864..776c78539 100644 --- a/tests/test_artifact_cache.py +++ b/tests/test_artifact_cache.py @@ -344,7 +344,16 @@ def test_resolve_ref_path_rejects_parent_segments(self) -> None: self.assertIsNone(_resolve_ref_path(suite_root, "artifacts/../../escape")) inside = _resolve_ref_path(suite_root, "artifacts/systematize/v0001/taxonomy.json") assert inside is not None - self.assertEqual(inside, suite_root / "artifacts" / "systematize" / "v0001" / "taxonomy.json") + self.assertEqual( + inside, + ( + suite_root + / "artifacts" + / "systematize" + / "v0001" + / "taxonomy.json" + ).resolve(), + ) def test_override_cacheable_output_paths_redirects_user_save_dir(self) -> None: with TemporaryDirectory() as tmp_dir: @@ -964,4 +973,3 @@ def test_per_file_isolation(self) -> None: if __name__ == "__main__": unittest.main() - diff --git a/tests/test_cli.py b/tests/test_cli.py index 7d952d82d..9410f2fbd 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -156,6 +156,10 @@ def test_run_emits_auth_mode_log_after_runner_loads(self) -> None: call_order: list[str] = [] runner_mock = self._make_runner_mock() + def _bootstrap_environment(*, discover_from_cwd=False): + self.assertTrue(discover_from_cwd) + call_order.append("bootstrap_env") + def _load_runner_module(): call_order.append("load_runner") return runner_mock @@ -163,7 +167,10 @@ def _load_runner_module(): def _log_auth_mode(): call_order.append("log_auth") - with patch("assert_ai.cli._load_runner_module", side_effect=_load_runner_module), \ + with patch( + "assert_ai.core.environment.bootstrap_environment", + side_effect=_bootstrap_environment, + ), patch("assert_ai.cli._load_runner_module", side_effect=_load_runner_module), \ patch( "assert_ai.core.azure_auth.log_resolved_azure_auth_mode", side_effect=_log_auth_mode, @@ -173,9 +180,8 @@ def _log_auth_mode(): self.assertEqual(result.exit_code, 0, msg=result.output) self.assertEqual( call_order, - ["load_runner", "log_auth"], - msg="auth-mode log must fire AFTER runner.py loads (runner.py " - "triggers load_dotenv + refresh_azure_auth_mode).", + ["bootstrap_env", "load_runner", "log_auth"], + msg="environment bootstrap must precede runner import and auth logging.", ) def test_run_respects_subcommand_quiet_flag(self) -> None: diff --git a/tests/test_environment_bootstrap.py b/tests/test_environment_bootstrap.py new file mode 100644 index 000000000..89ef46ebd --- /dev/null +++ b/tests/test_environment_bootstrap.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import importlib +import os +import sys +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from assert_ai.core.environment import bootstrap_environment +from assert_ai.core.model_client import refresh_environment_settings + + +def test_cli_discovery_loads_once_and_refreshes_auth() -> None: + fake_model_client = SimpleNamespace(refresh_environment_settings=Mock()) + with ( + patch("dotenv.find_dotenv", return_value="C:\\workspace\\.env") as find_dotenv, + patch("dotenv.load_dotenv") as load_dotenv, + patch("assert_ai.core.azure_auth.refresh_azure_auth_mode") as refresh_auth, + patch.dict( + sys.modules, + {"assert_ai.core.model_client": fake_model_client}, + ), + ): + bootstrap_environment(discover_from_cwd=True) + + find_dotenv.assert_called_once_with(usecwd=True) + load_dotenv.assert_called_once_with("C:\\workspace\\.env", override=False) + refresh_auth.assert_called_once_with(force=True) + fake_model_client.refresh_environment_settings.assert_called_once_with() + + +def test_explicit_env_file_never_runs_discovery_with_path(tmp_path) -> None: + env_file = tmp_path / "workspace.env" + with ( + patch("dotenv.find_dotenv") as find_dotenv, + patch("dotenv.load_dotenv") as load_dotenv, + patch("assert_ai.core.azure_auth.refresh_azure_auth_mode"), + ): + bootstrap_environment(env_file=env_file) + + find_dotenv.assert_not_called() + load_dotenv.assert_called_once_with(env_file, override=False) + + +def test_env_file_and_discovery_are_mutually_exclusive(tmp_path) -> None: + with pytest.raises(ValueError, match="mutually exclusive"): + bootstrap_environment( + env_file=tmp_path / ".env", + discover_from_cwd=True, + ) + + +def test_importing_runner_does_not_discover_or_load_dotenv() -> None: + import assert_ai.runner as runner + + with ( + patch("dotenv.find_dotenv") as find_dotenv, + patch("dotenv.load_dotenv") as load_dotenv, + ): + importlib.reload(runner) + + find_dotenv.assert_not_called() + load_dotenv.assert_not_called() + + +def test_refreshes_model_client_settings_loaded_before_dotenv() -> None: + with patch.dict( + os.environ, + { + "AZURE_API_BASE": "https://example.openai.azure.com/openai/v1/", + "ASSERT_PREFER_CHAT_COMPLETIONS": "", + }, + clear=False, + ): + refresh_environment_settings() + + assert os.environ["AZURE_API_BASE"] == "https://example.openai.azure.com/" diff --git a/tests/test_init_command.py b/tests/test_init_command.py index b49fc5b4f..fbf1f4ebc 100644 --- a/tests/test_init_command.py +++ b/tests/test_init_command.py @@ -37,6 +37,35 @@ def _done_response(yaml_str: str = _MINIMAL_VALID_YAML) -> str: class InitCommandTest(unittest.TestCase): + @patch("assert_ai.core.environment.bootstrap_environment") + @patch("assert_ai.init._design_agent.chat_completion") + @patch("assert_ai.init._design_agent.build_system_message", return_value="sys") + def test_uses_shared_environment_bootstrap( + self, + _mock_sys, + mock_llm, + bootstrap_environment, + ) -> None: + mock_llm.return_value = _done_response() + runner = CliRunner() + with runner.isolated_filesystem(): + env_file = Path("credentials.env") + env_file.write_text("AZURE_API_KEY=placeholder\n", encoding="utf-8") + result = runner.invoke( + cli, + [ + "init", + "--describe", + "A chatbot", + "--non-interactive", + "--env-file", + str(env_file), + ], + ) + + self.assertEqual(result.exit_code, 0, result.output) + bootstrap_environment.assert_called_once_with(env_file=env_file) + @patch("assert_ai.init._design_agent.chat_completion") @patch("assert_ai.init._design_agent.build_system_message", return_value="sys") def test_non_interactive_generates_file(self, _mock_sys, mock_llm) -> None: diff --git a/tests/test_mcp_cli.py b/tests/test_mcp_cli.py index d0dc52981..78883de69 100644 --- a/tests/test_mcp_cli.py +++ b/tests/test_mcp_cli.py @@ -3,6 +3,7 @@ from __future__ import annotations +from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, patch @@ -69,3 +70,85 @@ def test_mcp_serve_reports_missing_optional_dependency() -> None: assert result.exit_code == 1 assert 'python -m pip install "assert-ai[mcp]"' in result.output assert "Traceback" not in result.output + + +def test_mcp_serve_loads_workspace_env_before_server() -> None: + runner = CliRunner() + calls: list[str] = [] + options = object() + server_module = SimpleNamespace( + ServerOptions=SimpleNamespace( + create=Mock(side_effect=lambda **_: calls.append("options") or options) + ), + run_stdio_server=Mock(side_effect=lambda _: calls.append("serve")), + ) + + with runner.isolated_filesystem(): + env_file = Path(".env") + env_file.write_text("AZURE_API_KEY=placeholder\n", encoding="utf-8") + with ( + patch( + "assert_ai.mcp._command.bootstrap_environment", + side_effect=lambda **_: calls.append("environment"), + ) as bootstrap, + patch( + "assert_ai.mcp._command._load_server_module", + return_value=server_module, + ), + ): + result = runner.invoke( + cli, + ["mcp", "serve", "--workspace", ".", "--env-file", ".env"], + ) + + assert result.exit_code == 0, result.output + assert calls == ["environment", "options", "serve"] + bootstrap.assert_called_once() + assert bootstrap.call_args.kwargs["env_file"].name == ".env" + assert bootstrap.call_args.kwargs["env_file"].is_absolute() + + +def test_mcp_serve_rejects_env_file_outside_workspace() -> None: + runner = CliRunner() + with runner.isolated_filesystem(): + workspace = Path("workspace") + workspace.mkdir() + Path("outside.env").write_text("AZURE_API_KEY=placeholder\n", encoding="utf-8") + with ( + patch("assert_ai.mcp._command.bootstrap_environment") as bootstrap, + patch("assert_ai.mcp._command._load_server_module") as load_server, + ): + result = runner.invoke( + cli, + [ + "mcp", + "serve", + "--workspace", + str(workspace), + "--env-file", + str(Path("..") / "outside.env"), + ], + ) + + assert result.exit_code == 1 + assert "escapes its expected root directory" in result.output + bootstrap.assert_not_called() + load_server.assert_not_called() + + +def test_mcp_serve_without_env_file_does_not_bootstrap() -> None: + runner = CliRunner() + server_module = SimpleNamespace( + ServerOptions=SimpleNamespace(create=Mock(return_value=object())), + run_stdio_server=Mock(), + ) + with runner.isolated_filesystem(), patch( + "assert_ai.mcp._command.bootstrap_environment", + ) as bootstrap, patch( + "assert_ai.mcp._command._load_server_module", + return_value=server_module, + ): + result = runner.invoke(cli, ["mcp", "serve"]) + + assert result.exit_code == 0, result.output + bootstrap.assert_not_called() diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 3e521d5da..4c61aad27 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -62,6 +62,8 @@ def test_server_options_resolve_workspace_and_capabilities(tmp_path: Path) -> No assert options.workspace_root == workspace.resolve() assert options.mode is ServerMode.AUTHOR + assert options.path_policy.workspace_root == workspace.resolve() + assert options.path_policy.force_managed_outputs is True assert options.capability_groups == ( CapabilityGroup.INSPECT, CapabilityGroup.AUTHOR, @@ -70,6 +72,16 @@ def test_server_options_resolve_workspace_and_capabilities(tmp_path: Path) -> No ) +def test_server_options_direct_constructor_preserves_workspace_root_api( + tmp_path: Path, +) -> None: + options = ServerOptions(workspace_root=tmp_path) + + assert options.workspace_root == tmp_path.resolve() + assert options.workspace.root == tmp_path.resolve() + assert options.path_policy.workspace_root == tmp_path.resolve() + + def test_design_group_requires_author_or_full_mode(tmp_path: Path) -> None: with pytest.raises(ValueError, match="require --mode author or --mode full"): ServerOptions.create( @@ -99,6 +111,7 @@ async def run() -> tuple[set[str], object]: assert result.structured_content["assert_mcp_api_version"] == "1" assert result.structured_content["mode"] == "full" assert result.structured_content["workspace"]["root"] == "." + assert "env_file" not in result.structured_content assert result.structured_content["enabled_capability_groups"] == [ "inspect", "author", diff --git a/tests/test_runtime_path_policy.py b/tests/test_runtime_path_policy.py new file mode 100644 index 000000000..aa610320b --- /dev/null +++ b/tests/test_runtime_path_policy.py @@ -0,0 +1,727 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import importlib +import os +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest + +from assert_ai.config import _resolve_path, load_runtime_context +from assert_ai.core.artifact_cache import ( + discard_artifact_plan, + prepare_artifact_plan, + refresh_compatibility_files, + update_latest, +) +from assert_ai.core.model_client import GenerateOptions +from assert_ai.core.runtime_path_policy import ( + RuntimePathError, + RuntimePathErrorCode, + RuntimePathPolicy, +) +from assert_ai.core.security import validate_sys_path_addition +from assert_ai.core.tool_backend import import_callable_module, load_tool_module +from assert_ai.core.workspace import WorkspaceService +from assert_ai.stages import STAGES +from assert_ai.stages.inference import _build_hosted_session + + +def _workspace(tmp_path: Path) -> tuple[WorkspaceService, Path]: + root = tmp_path / "workspace" + configs = root / "evals" + configs.mkdir(parents=True) + config_path = configs / "eval_config.yaml" + config_path.write_text("pipeline: {}\n", encoding="utf-8") + return WorkspaceService.create(root), config_path + + +def _minimal_config() -> dict: + return { + "suite": "suite-a", + "pipeline": { + "inference": { + "enabled": False, + } + }, + } + + +def _remove_workspace_imports(existing_modules: set[str]) -> None: + for name in set(sys.modules).difference(existing_modules): + if name.startswith("_assert_ai_workspace_"): + sys.modules.pop(name, None) + + +def test_workspace_service_exposes_only_relative_references(tmp_path: Path) -> None: + workspace, _ = _workspace(tmp_path) + + assert workspace.reference(workspace.root) == "." + assert workspace.reference(workspace.configs_root) == "evals" + assert workspace.reference(workspace.artifacts_root) == "artifacts" + assert workspace.reference(workspace.results_root) == "artifacts/results" + + +def test_config_path_is_contained_under_config_root(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + + assert workspace.path_policy.resolve_config_path("eval_config.yaml") == config_path + assert workspace.path_policy.resolve_config_path("evals/eval_config.yaml") == config_path + + with pytest.raises(RuntimePathError) as exc: + workspace.path_policy.resolve_config_path("../outside.yaml") + + assert exc.value.code is RuntimePathErrorCode.OUTSIDE_CONFIG_ROOT + + +def test_relative_input_cannot_escape_its_base_directory(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + shared = workspace.root / "shared.jsonl" + shared.write_text("{}\n", encoding="utf-8") + + with pytest.raises(RuntimePathError) as exc: + workspace.path_policy.resolve_input( + "../shared.jsonl", + base_dir=config_path.parent, + field_name="pipeline.inference.test_set_path", + ) + + assert exc.value.code is RuntimePathErrorCode.OUTSIDE_INPUT_ROOT + + +def test_absolute_input_requires_an_explicit_read_root(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + inside = workspace.root / "data.jsonl" + inside.write_text("{}\n", encoding="utf-8") + outside = tmp_path / "outside.jsonl" + outside.write_text("{}\n", encoding="utf-8") + + assert ( + workspace.path_policy.resolve_input( + inside, + base_dir=config_path.parent, + field_name="input", + ) + == inside + ) + with pytest.raises(RuntimePathError) as exc: + workspace.path_policy.resolve_input( + outside, + base_dir=config_path.parent, + field_name="input", + ) + + assert exc.value.code is RuntimePathErrorCode.OUTSIDE_INPUT_ROOT + + +def test_additional_read_root_allows_explicit_external_input(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + external_root = tmp_path / "approved-inputs" + external_root.mkdir() + external = external_root / "data.jsonl" + external.write_text("{}\n", encoding="utf-8") + policy = RuntimePathPolicy( + workspace_root=workspace.root, + config_root=workspace.configs_root, + artifacts_root=workspace.artifacts_root, + results_root=workspace.results_root, + additional_read_roots=(external_root,), + ) + + assert ( + policy.resolve_input( + external, + base_dir=config_path.parent, + field_name="input", + ) + == external + ) + + +def test_outputs_cannot_escape_managed_artifacts_root(tmp_path: Path) -> None: + workspace, _ = _workspace(tmp_path) + + with pytest.raises(RuntimePathError) as traversal: + workspace.path_policy.resolve_output( + "../outside", + field_name="pipeline.inference.save_dir", + ) + with pytest.raises(RuntimePathError) as absolute: + workspace.path_policy.resolve_output( + tmp_path / "outside", + field_name="pipeline.inference.save_dir", + ) + + assert traversal.value.code is RuntimePathErrorCode.OUTSIDE_ARTIFACTS_ROOT + assert absolute.value.code is RuntimePathErrorCode.OUTSIDE_ARTIFACTS_ROOT + + +def test_symlink_escape_is_rejected_after_resolution(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + outside = tmp_path / "outside" + outside.mkdir() + (outside / "data.jsonl").write_text("{}\n", encoding="utf-8") + link = config_path.parent / "linked" + try: + link.symlink_to(outside, target_is_directory=True) + except OSError as exc: + pytest.skip(f"symlink creation is unavailable: {exc}") + + with pytest.raises(RuntimePathError) as error: + workspace.path_policy.resolve_input( + "linked/data.jsonl", + base_dir=config_path.parent, + field_name="input", + ) + + assert error.value.code is RuntimePathErrorCode.OUTSIDE_INPUT_ROOT + + +def test_artifact_cache_cannot_allocate_through_symlink_escape( + tmp_path: Path, +) -> None: + workspace, config_path = _workspace(tmp_path) + suite_root = workspace.results_root / "suite-a" + suite_root.mkdir(parents=True) + outside = tmp_path / "outside-cache" + outside.mkdir() + cache_link = suite_root / "artifacts" + try: + cache_link.symlink_to(outside, target_is_directory=True) + except OSError as exc: + pytest.skip(f"symlink creation is unavailable: {exc}") + + with pytest.raises(RuntimePathError) as error: + prepare_artifact_plan( + ctx={ + "suite_root": suite_root, + "config_path": config_path, + "artifacts_root": workspace.artifacts_root, + "behavior_name": "behavior", + "behavior": "description", + "context": None, + "path_policy": workspace.path_policy, + }, + stage_name="systematize", + raw_cfg={}, + forced=False, + ) + + assert error.value.code is RuntimePathErrorCode.OUTSIDE_ARTIFACTS_ROOT + assert not any(outside.iterdir()) + + +def test_runtime_context_rejects_suite_link_to_another_suite( + tmp_path: Path, +) -> None: + workspace, config_path = _workspace(tmp_path) + workspace.results_root.mkdir(parents=True) + other_suite = workspace.results_root / "suite-b" + other_suite.mkdir() + suite_link = workspace.results_root / "suite-a" + try: + suite_link.symlink_to(other_suite, target_is_directory=True) + except OSError as exc: + pytest.skip(f"symlink creation is unavailable: {exc}") + + with pytest.raises(RuntimePathError) as error: + load_runtime_context( + _minimal_config(), + config_path, + stage_modules=STAGES, + path_policy=workspace.path_policy, + ) + + assert error.value.code is RuntimePathErrorCode.MANAGED_PATH_LINK + + +def test_artifact_cache_rejects_cross_suite_links_for_mutations( + tmp_path: Path, +) -> None: + workspace, config_path = _workspace(tmp_path) + suite_a = workspace.results_root / "suite-a" + suite_b = workspace.results_root / "suite-b" + suite_a.mkdir(parents=True) + suite_b.mkdir() + ctx = { + "suite_root": suite_a, + "config_path": config_path, + "artifacts_root": workspace.artifacts_root, + "behavior_name": "behavior", + "behavior": "description", + "context": None, + "path_policy": workspace.path_policy, + } + + protected_latest = suite_b / "protected-latest.json" + protected_latest.write_text('{"protected": true}\n', encoding="utf-8") + latest_link = suite_a / "latest.json" + try: + latest_link.symlink_to(protected_latest) + except OSError as exc: + pytest.skip(f"symlink creation is unavailable: {exc}") + + with pytest.raises(RuntimePathError) as latest_error: + update_latest(ctx, "systematize", {"version": "v0001"}) + assert latest_error.value.code is RuntimePathErrorCode.OUTSIDE_EXPECTED_ROOT + assert protected_latest.read_text(encoding="utf-8") == '{"protected": true}\n' + + latest_link.unlink() + plan = prepare_artifact_plan( + ctx=ctx, + stage_name="systematize", + raw_cfg={}, + forced=False, + ) + source = plan.output_paths["taxonomy"] + source.write_text('{"safe": true}\n', encoding="utf-8") + protected_copy = suite_b / "protected-taxonomy.json" + protected_copy.write_text('{"protected": true}\n', encoding="utf-8") + compatibility_link = suite_a / source.name + compatibility_link.symlink_to(protected_copy) + + with pytest.raises(RuntimePathError) as copy_error: + refresh_compatibility_files( + ctx, + "systematize", + plan.output_paths, + ) + assert copy_error.value.code is RuntimePathErrorCode.OUTSIDE_EXPECTED_ROOT + assert protected_copy.read_text(encoding="utf-8") == '{"protected": true}\n' + + compatibility_link.unlink() + protected_dir = suite_b / "protected-version" + protected_dir.mkdir() + protected_file = protected_dir / "keep.txt" + protected_file.write_text("keep\n", encoding="utf-8") + source.unlink() + plan.artifact_dir.rmdir() + plan.artifact_dir.symlink_to(protected_dir, target_is_directory=True) + + discard_artifact_plan(ctx, plan) + + assert protected_file.read_text(encoding="utf-8") == "keep\n" + + +def test_runtime_context_forces_managed_roots(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + raw = { + **_minimal_config(), + "artifacts_root": "custom-artifacts", + } + + with pytest.raises(RuntimePathError) as exc: + load_runtime_context( + raw, + config_path, + stage_modules=STAGES, + path_policy=workspace.path_policy, + ) + + assert exc.value.code is RuntimePathErrorCode.MANAGED_ROOT_OVERRIDE + + +def test_runtime_context_accepts_explicit_managed_roots(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + raw = { + **_minimal_config(), + "artifacts_root": "artifacts", + "results_dir": "results", + } + + context = load_runtime_context( + raw, + config_path, + stage_modules=STAGES, + path_policy=workspace.path_policy, + ) + + assert context["artifacts_root"] == workspace.artifacts_root + assert context["results_dir"] == workspace.results_root + assert context["path_policy"] is workspace.path_policy + + +def test_runtime_context_confines_run_outputs_to_current_run( + tmp_path: Path, +) -> None: + workspace, config_path = _workspace(tmp_path) + other_run = workspace.results_root / "suite-a" / "run-b" + raw = { + "suite": "suite-a", + "run": "run-a", + "pipeline": { + "inference": { + "target": {"callable": "agent:run"}, + "save_dir": str(other_run), + } + }, + } + + with pytest.raises(RuntimePathError) as error: + load_runtime_context( + raw, + config_path, + stage_modules=STAGES, + path_policy=workspace.path_policy, + ) + + assert error.value.code is RuntimePathErrorCode.OUTSIDE_EXPECTED_ROOT + + +def test_managed_tree_rejects_cross_run_file_link(tmp_path: Path) -> None: + workspace, _ = _workspace(tmp_path) + run_a = workspace.results_root / "suite-a" / "run-a" + run_b = workspace.results_root / "suite-a" / "run-b" + run_a.mkdir(parents=True) + run_b.mkdir() + protected = run_b / "inference_set.jsonl" + protected.write_text('{"protected": true}\n', encoding="utf-8") + output_link = run_a / "inference_set.jsonl" + try: + output_link.symlink_to(protected) + except OSError as exc: + pytest.skip(f"symlink creation is unavailable: {exc}") + + with pytest.raises(RuntimePathError) as error: + workspace.path_policy.require_managed_tree( + run_a, + field_name="run output", + expected_root=run_a, + ) + + assert error.value.code is RuntimePathErrorCode.OUTSIDE_EXPECTED_ROOT + assert protected.read_text(encoding="utf-8") == '{"protected": true}\n' + + +def test_every_explicit_stage_path_is_validated(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + raw = _minimal_config() + raw["pipeline"]["inference"]["file_path"] = "../outside.jsonl" + + with pytest.raises(RuntimePathError) as exc: + load_runtime_context( + raw, + config_path, + stage_modules=STAGES, + path_policy=workspace.path_policy, + ) + + assert exc.value.code is RuntimePathErrorCode.OUTSIDE_INPUT_ROOT + + +def test_toolset_is_revalidated_when_actually_loaded(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + outside = tmp_path / "tools.json" + outside.write_text("[]\n", encoding="utf-8") + + with pytest.raises(RuntimePathError) as exc: + _build_hosted_session( + model="mock/model", + tools_config={ + "_config_path": str(config_path), + "toolset": str(outside), + "simulator": "mock/simulator", + }, + scenario={}, + generate_options=GenerateOptions(), + max_tool_calls=1, + synthetic_prompt_template="{scenario}", + path_policy=workspace.path_policy, + ) + + assert exc.value.code is RuntimePathErrorCode.OUTSIDE_INPUT_ROOT + + +def test_strict_dynamic_import_searches_only_workspace(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + module_name = f"workspace_agent_{uuid.uuid4().hex[:8]}" + module_path = workspace.root / f"{module_name}.py" + module_path.write_text("def run(message):\n return message\n", encoding="utf-8") + + module = import_callable_module( + module_name, + config_path=config_path, + path_policy=workspace.path_policy, + ) + + assert module.run("ok") == "ok" + + +def test_strict_dynamic_import_does_not_fall_back_to_cwd(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + outside = tmp_path / "outside" + outside.mkdir() + module_name = f"outside_agent_{uuid.uuid4().hex[:8]}" + (outside / f"{module_name}.py").write_text( + "def run(message):\n return message\n", + encoding="utf-8", + ) + original_cwd = Path.cwd() + os.chdir(outside) + try: + with pytest.raises(ValueError, match="inside the configured workspace"): + import_callable_module( + module_name, + config_path=config_path, + path_policy=workspace.path_policy, + ) + finally: + os.chdir(original_cwd) + + +def test_strict_module_import_supports_package_relative_imports( + tmp_path: Path, +) -> None: + existing_modules = set(sys.modules) + workspace, config_path = _workspace(tmp_path) + package_name = "strict_runtime_package" + package = config_path.parent / package_name + package.mkdir() + (package / "__init__.py").write_text("", encoding="utf-8") + (package / "helpers.py").write_text("VALUE = 'workspace'\n", encoding="utf-8") + (package / "agent.py").write_text( + "from .helpers import VALUE\n", + encoding="utf-8", + ) + + try: + module = import_callable_module( + f"{package_name}.agent", + config_path=config_path, + path_policy=workspace.path_policy, + ) + assert module.VALUE == "workspace" + finally: + for name in tuple(sys.modules): + if name == package_name or name.startswith(f"{package_name}."): + sys.modules.pop(name, None) + _remove_workspace_imports(existing_modules) + + +def test_strict_module_import_isolates_identical_names_by_config_root( + tmp_path: Path, +) -> None: + existing_modules = set(sys.modules) + workspace, _ = _workspace(tmp_path) + first_root = workspace.configs_root / "first" + second_root = workspace.configs_root / "second" + first_root.mkdir() + second_root.mkdir() + first_config = first_root / "eval_config.yaml" + second_config = second_root / "eval_config.yaml" + first_config.write_text("pipeline: {}\n", encoding="utf-8") + second_config.write_text("pipeline: {}\n", encoding="utf-8") + (first_root / "helpers.py").write_text("VALUE = 'first'\n", encoding="utf-8") + (second_root / "helpers.py").write_text("VALUE = 'second'\n", encoding="utf-8") + (first_root / "agent.py").write_text( + "def get_value():\n import helpers\n return helpers.VALUE\n", + encoding="utf-8", + ) + (second_root / "agent.py").write_text( + "def get_value():\n import helpers\n return helpers.VALUE\n", + encoding="utf-8", + ) + + try: + with ThreadPoolExecutor(max_workers=2) as executor: + first_future = executor.submit( + import_callable_module, + "agent", + config_path=first_config, + path_policy=workspace.path_policy, + ) + second_future = executor.submit( + import_callable_module, + "agent", + config_path=second_config, + path_policy=workspace.path_policy, + ) + first = first_future.result() + second = second_future.result() + + assert first.get_value() == "first" + assert second.get_value() == "second" + assert first.__name__ != second.__name__ + finally: + _remove_workspace_imports(existing_modules) + + +def test_strict_module_import_isolates_absolute_package_imports( + tmp_path: Path, +) -> None: + existing_modules = set(sys.modules) + workspace, _ = _workspace(tmp_path) + package_name = "shared_runtime_package" + modules: list[object] = [] + + try: + for config_name, value in (("first", "first"), ("second", "second")): + config_root = workspace.configs_root / config_name + package = config_root / package_name + package.mkdir(parents=True) + config_path = config_root / "eval_config.yaml" + config_path.write_text("pipeline: {}\n", encoding="utf-8") + (package / "__init__.py").write_text("", encoding="utf-8") + (package / "helpers.py").write_text( + f"VALUE = {value!r}\n", + encoding="utf-8", + ) + (package / "agent.py").write_text( + f"from {package_name}.helpers import VALUE\n", + encoding="utf-8", + ) + modules.append( + import_callable_module( + f"{package_name}.agent", + config_path=config_path, + path_policy=workspace.path_policy, + ) + ) + + assert [module.VALUE for module in modules] == ["first", "second"] + finally: + _remove_workspace_imports(existing_modules) + + +def test_strict_direct_module_path_uses_isolated_lazy_imports( + tmp_path: Path, +) -> None: + existing_modules = set(sys.modules) + workspace, config_path = _workspace(tmp_path) + helper_name = "strict_direct_helper" + (config_path.parent / f"{helper_name}.py").write_text( + "VALUE = 'workspace'\n", + encoding="utf-8", + ) + module_path = config_path.parent / "direct_tools.py" + module_path.write_text( + "import importlib\n" + "from importlib import import_module\n\n" + f"def get_value():\n return importlib.import_module({helper_name!r}).VALUE\n\n" + f"def get_value_from_import():\n return import_module({helper_name!r}).VALUE\n", + encoding="utf-8", + ) + + external_root = tmp_path / "external-direct" + external_root.mkdir() + (external_root / f"{helper_name}.py").write_text( + "VALUE = 'external'\n", + encoding="utf-8", + ) + original_sys_path = list(sys.path) + sys.path.insert(0, str(external_root)) + try: + importlib.import_module(helper_name) + module = load_tool_module( + str(module_path), + config_path=config_path, + path_policy=workspace.path_policy, + ) + assert module.get_value() == "workspace" + assert module.get_value_from_import() == "workspace" + assert sys.path == [str(external_root), *original_sys_path] + finally: + sys.path[:] = original_sys_path + sys.modules.pop(helper_name, None) + _remove_workspace_imports(existing_modules) + + +def test_strict_direct_package_initializer_uses_workspace_module_name( + tmp_path: Path, +) -> None: + existing_modules = set(sys.modules) + workspace, config_path = _workspace(tmp_path) + package_init = config_path.parent / "__init__.py" + package_init.write_text("VALUE = 'package'\n", encoding="utf-8") + + try: + module = load_tool_module( + str(package_init), + config_path=config_path, + path_policy=workspace.path_policy, + ) + assert module.VALUE == "package" + finally: + _remove_workspace_imports(existing_modules) + + +def test_strict_module_import_ignores_preloaded_external_package( + tmp_path: Path, +) -> None: + existing_modules = set(sys.modules) + workspace, config_path = _workspace(tmp_path) + package_name = "strict_external_package" + workspace_package = config_path.parent / package_name + workspace_package.mkdir() + (workspace_package / "__init__.py").write_text("", encoding="utf-8") + (workspace_package / "helpers.py").write_text( + "VALUE = 'workspace'\n", + encoding="utf-8", + ) + (workspace_package / "agent.py").write_text( + f"from {package_name}.helpers import VALUE\n", + encoding="utf-8", + ) + + external_root = tmp_path / "external" + external_package = external_root / package_name + external_package.mkdir(parents=True) + (external_package / "__init__.py").write_text("", encoding="utf-8") + (external_package / "helpers.py").write_text( + "VALUE = 'external'\n", + encoding="utf-8", + ) + + sys.path.insert(0, str(external_root)) + try: + importlib.import_module(package_name) + module = import_callable_module( + f"{package_name}.agent", + config_path=config_path, + path_policy=workspace.path_policy, + ) + assert module.VALUE == "workspace" + finally: + sys.path.remove(str(external_root)) + for name in tuple(sys.modules): + if name == package_name or name.startswith(f"{package_name}."): + sys.modules.pop(name, None) + _remove_workspace_imports(existing_modules) + + +def test_direct_module_and_sys_path_must_stay_in_workspace(tmp_path: Path) -> None: + workspace, config_path = _workspace(tmp_path) + outside_module = tmp_path / "outside_tools.py" + outside_module.write_text("class Tools:\n pass\n", encoding="utf-8") + + with pytest.raises(RuntimePathError) as module_error: + load_tool_module( + str(outside_module), + config_path=config_path, + path_policy=workspace.path_policy, + ) + with pytest.raises(RuntimePathError) as sys_path_error: + validate_sys_path_addition( + tmp_path, + config_path=config_path, + path_policy=workspace.path_policy, + ) + + assert module_error.value.code is RuntimePathErrorCode.OUTSIDE_WORKSPACE + assert sys_path_error.value.code is RuntimePathErrorCode.OUTSIDE_WORKSPACE + + +def test_legacy_absolute_inputs_remain_supported(tmp_path: Path) -> None: + absolute = tmp_path / "outside.jsonl" + resolved = _resolve_path( + absolute, + artifacts_root=tmp_path / "artifacts", + cfg_dir=tmp_path / "configs", + ) + + assert Path(resolved) == absolute.resolve() From 11fd48d8e8673210be3cbfcf46948c002f2fa395 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Thu, 6 Aug 2026 20:44:26 -0700 Subject: [PATCH 03/16] Add canonical config services Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/config.py | 40 +- assert_ai/core/config_document.py | 426 +++++++++++++ assert_ai/core/io.py | 26 + assert_ai/core/runtime_path_policy.py | 18 +- assert_ai/init/_design_agent.py | 7 +- assert_ai/init/_validate.py | 190 +++--- assert_ai/services/__init__.py | 4 + assert_ai/services/configs.py | 852 ++++++++++++++++++++++++++ assert_ai/services/errors.py | 43 ++ assert_ai/stages/test_set.py | 2 + docs/config/schema.md | 6 + tests/test_config_document.py | 223 +++++++ tests/test_config_service.py | 262 ++++++++ 13 files changed, 1998 insertions(+), 101 deletions(-) create mode 100644 assert_ai/core/config_document.py create mode 100644 assert_ai/services/__init__.py create mode 100644 assert_ai/services/configs.py create mode 100644 assert_ai/services/errors.py create mode 100644 tests/test_config_document.py create mode 100644 tests/test_config_service.py diff --git a/assert_ai/config.py b/assert_ai/config.py index 7b9501aba..3c157c28b 100644 --- a/assert_ai/config.py +++ b/assert_ai/config.py @@ -33,20 +33,19 @@ ToolsConfig, TraceConfig, ) +from assert_ai.core.config_document import ( + EvalConfigDocumentError, + PIPELINE_STAGE_ORDER, + require_valid_eval_config_document, +) from assert_ai.core.runtime_path_policy import RuntimePathPolicy ROOT = Path(__file__).resolve().parent.parent OUTPUT_PATH_KEYS = {"save_dir", "save_path"} -PIPELINE_STAGE_ORDER = ( - "systematize", - "test_set", - "inference", - "judge", -) BEHAVIOR_REQUIRED_PIPELINE_STAGES = {"systematize"} -class ConfigError(Exception): +class ConfigError(ValueError): pass @@ -89,7 +88,7 @@ def require(condition: bool, message: str) -> None: def load_config(cfg_path: Path) -> dict[str, Any]: - """Load one YAML config file and require a mapping at the top level.""" + """Load and structurally validate one YAML config file.""" try: data = yaml.safe_load(cfg_path.read_text(encoding="utf-8")) except FileNotFoundError: @@ -99,6 +98,27 @@ def load_config(cfg_path: Path) -> dict[str, Any]: except yaml.YAMLError as exc: raise ConfigError(f"Invalid YAML in config file {cfg_path}: {exc}") from exc require(isinstance(data, dict), "Top-level YAML must be a mapping") + try: + reject_unknown_keys( + data, + field_name="config", + allowed={ + "suite", + "run", + "behavior", + "context", + "default_model", + "artifacts_root", + "results_dir", + "pipeline", + }, + ) + except ValueError as exc: + raise ConfigError(str(exc)) from exc + try: + require_valid_eval_config_document(data) + except EvalConfigDocumentError as exc: + raise ConfigError(f"Invalid config file {cfg_path}: {exc}") from exc return data @@ -209,6 +229,10 @@ def load_runtime_context( "pipeline", }, ) + try: + require_valid_eval_config_document(raw) + except EvalConfigDocumentError as exc: + raise ConfigError(f"Invalid config: {exc}") from exc default_model_raw = _get_default_model_mapping(raw) pipeline_raw = raw.get("pipeline") require(isinstance(pipeline_raw, dict), "'pipeline' must be a mapping") diff --git a/assert_ai/core/config_document.py b/assert_ai/core/config_document.py new file mode 100644 index 000000000..1f3966fce --- /dev/null +++ b/assert_ai/core/config_document.py @@ -0,0 +1,426 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Canonical machine-readable model for ASSERT evaluation YAML.""" + +from __future__ import annotations + +from copy import deepcopy +from enum import StrEnum +from typing import Any, Literal + +import yaml +from pydantic import ( + BaseModel, + ConfigDict, + Field, + ValidationError, + field_validator, + model_validator, +) + +EVAL_CONFIG_SCHEMA_VERSION = 1 +EVAL_CONFIG_SCHEMA_ID = "https://github.com/responsibleai/ASSERT/schemas/eval-config-v1.json" + +PIPELINE_STAGE_ORDER = ( + "systematize", + "test_set", + "inference", + "judge", +) + + +class _DocumentModel(BaseModel): + """Base class for strict YAML document nodes.""" + + model_config = ConfigDict(extra="forbid") + + +class ModelDocument(_DocumentModel): + """One LiteLLM model reference and generation options.""" + + name: str = Field(description="Provider-defined model identifier.") + temperature: float | None = Field(default=None) + max_tokens: int | None = Field(default=None, gt=0) + reasoning_effort: str | None = Field(default=None) + + +class BehaviorDocument(_DocumentModel): + """Behavior specification or reusable behavior preset reference.""" + + name: str | None = Field(default=None) + description: str | None = Field(default=None) + preset: str | None = Field(default=None) + + +class StageDocument(_DocumentModel): + """Fields shared by pipeline stage declarations.""" + + enabled: bool | None = Field(default=None) + file_path: str | None = Field( + default=None, + description="Compatibility field accepted by the shared pipeline loader.", + ) + + +class SystematizeDocument(StageDocument): + """Configuration for behavior systematization and taxonomy conversion.""" + + behavior_category_count: int | None = Field(default=None, gt=0) + web_search: bool | None = Field(default=None) + model: ModelDocument | None = Field(default=None) + save_dir: str | None = Field(default=None) + validators: Any = Field(default=None, deprecated=True) + validator_models: Any = Field(default=None, deprecated=True) + + +class SamplingDocument(_DocumentModel): + """Assignment sampling controls for one test-case kind.""" + + method: Literal["pairwise", "stratified", "full_factorial", "random"] = "pairwise" + stratify_by: list[str] | None = Field(default=None) + replication: Literal["balanced", "none"] | None = Field(default=None) + with_replacement: bool | None = Field(default=None) + + +class TestCaseGenerationDocument(_DocumentModel): + """Prompt or scenario generation settings.""" + + model: ModelDocument | None = Field(default=None) + sample_size: int | None = Field(default=None, ge=1, le=100_000) + timeout_s: float | None = Field(default=None, gt=0) + sampling: SamplingDocument | None = Field(default=None) + budget: Any = Field( + default=None, + deprecated=True, + description="Removed alias. Use sample_size.", + ) + + +class ScenarioGenerationDocument(TestCaseGenerationDocument): + """Scenario generation settings, including removed compatibility fields.""" + + modality: Any = Field( + default=None, + deprecated=True, + description="Removed field. Use test_set.tool_source.", + ) + + +class DimensionLevelDocument(_DocumentModel): + """One explicit level of a test-set variation dimension.""" + + name: str + definition: str + + +class DimensionDocument(_DocumentModel): + """One explicit or model-generated test-set variation dimension.""" + + name: str + description: str | None = Field(default=None) + levels: list[DimensionLevelDocument] | None = Field(default=None) + + +class StratifyDocument(_DocumentModel): + """Test-set dimension generation and level configuration.""" + + dimensions: list[DimensionDocument] | None = Field(default=None) + level_count: int | None = Field(default=None, gt=0) + model: ModelDocument | None = Field(default=None) + + +class TestSetDocument(StageDocument): + """Configuration for prompt and scenario test-set generation.""" + + taxonomy_path: str | None = Field(default=None) + save_path: str | None = Field(default=None) + stratify: StratifyDocument | None = Field(default=None) + tool_source: Literal["runtime", "per_test_case", "per_seed"] | None = Field( + default=None, + description="Tool source. per_seed is a deprecated alias for per_test_case.", + ) + model: ModelDocument | None = Field(default=None) + timeout_s: float | None = Field(default=None, gt=0) + prompt: TestCaseGenerationDocument | None = Field(default=None) + scenario: ScenarioGenerationDocument | None = Field(default=None) + validators: Any = Field(default=None, deprecated=True) + validator_model: Any = Field(default=None, deprecated=True) + + +class ToolsDocument(_DocumentModel): + """Prompt Agent tool backend or simulator configuration.""" + + module: str | None = Field(default=None) + toolset: str | None = Field(default=None) + simulator: str | None = Field(default=None) + + +class TraceDocument(_DocumentModel): + """OpenTelemetry trace capture configuration for a callable target.""" + + backend: str = "phoenix" + group_by: str = "session.id" + + +class TargetDocument(_DocumentModel): + """Hosted model, callable, endpoint, or connector target declaration.""" + + model: ModelDocument | None = Field(default=None) + system_prompt: str | None = Field(default=None) + tools: ToolsDocument | None = Field(default=None) + connector: str | None = Field(default=None) + callable: str | None = Field(default=None) + endpoint: str | None = Field(default=None) + trace: TraceDocument | None = Field(default=None) + + +class TesterDocument(_DocumentModel): + """Optional model-driven scenario tester.""" + + model: ModelDocument | None = Field(default=None) + max_turns: Any = Field( + default=None, + deprecated=True, + description="Removed field. Use pipeline.inference.max_turns.", + ) + + +class InferenceDocument(StageDocument): + """Configuration for executing test cases against a target.""" + + target: TargetDocument | None = Field(default=None) + tester: TesterDocument | None = Field(default=None) + max_turns: int | None = Field(default=None, gt=0) + max_tool_calls: int | None = Field(default=None, gt=0) + tool_timeout_s: float | None = Field(default=None, gt=0) + startup_timeout_s: float | None = Field(default=None, gt=0) + concurrency: int | None = Field(default=None, gt=0) + test_set_path: str | None = Field(default=None) + save_dir: str | None = Field(default=None) + strict: bool | None = Field(default=None) + + +class OrdinalScaleDocument(_DocumentModel): + """Ordered custom judge scale.""" + + type: str + values: dict[Any, str] + + +class JudgeDimensionDocument(_DocumentModel): + """One custom judge dimension and rubric.""" + + description: str + rubric: str + required_base: bool | None = Field(default=None) + allow_not_applicable: bool | None = Field(default=None) + scale: OrdinalScaleDocument | None = Field(default=None) + + +class JudgeDocument(StageDocument): + """Configuration for transcript judging.""" + + model: ModelDocument | None = Field(default=None) + n: int | None = Field(default=None, gt=0) + dimensions: dict[str, JudgeDimensionDocument] | None = Field(default=None) + disabled_dimensions: list[str] | None = Field(default=None) + inference_set_path: str | None = Field(default=None) + taxonomy_path: str | None = Field(default=None) + save_dir: str | None = Field(default=None) + preset: Any = Field(default=None) + + @field_validator( + "preset", + mode="before", + json_schema_input_type=str | list[str] | None, + ) + @classmethod + def _validate_preset(cls, value: Any) -> Any: + if value is None or isinstance(value, str): + return value + if isinstance(value, list): + for index, item in enumerate(value): + if not isinstance(item, str): + raise ValueError( + f"pipeline.judge.preset[{index}] must be a string" + ) + return value + raise ValueError( + "pipeline.judge.preset must be a string or a list of strings" + ) + + +class PipelineDocument(_DocumentModel): + """Canonical ordered ASSERT pipeline declaration.""" + + model_config = ConfigDict( + extra="forbid", + json_schema_extra={"minProperties": 1}, + ) + + # Non-optional annotations keep JSON Schema stage values object-only while + # defaults allow each stage key to be omitted. + systematize: SystematizeDocument = Field(default=None) # type: ignore[assignment] + test_set: TestSetDocument = Field(default=None) # type: ignore[assignment] + inference: InferenceDocument = Field(default=None) # type: ignore[assignment] + judge: JudgeDocument = Field(default=None) # type: ignore[assignment] + + @model_validator(mode="after") + def _require_stage(self) -> PipelineDocument: + if not any(getattr(self, stage_name) is not None for stage_name in PIPELINE_STAGE_ORDER): + raise ValueError("pipeline must define at least one stage") + return self + + +class EvalConfigDocument(_DocumentModel): + """Complete structural model for ``eval_config.yaml``.""" + + suite: str | None = Field(default=None) + run: str | None = Field(default=None) + behavior: BehaviorDocument | None = Field(default=None) + context: str | None = Field(default=None) + default_model: ModelDocument | None = Field(default=None) + artifacts_root: str | None = Field(default=None) + results_dir: str | None = Field(default=None) + pipeline: PipelineDocument + + +class ConfigValidationCode(StrEnum): + """Stable categories for machine-correctable config issues.""" + + INVALID_YAML = "INVALID_YAML" + REQUIRED_FIELD = "REQUIRED_FIELD" + UNKNOWN_FIELD = "UNKNOWN_FIELD" + INVALID_TYPE = "INVALID_TYPE" + INVALID_VALUE = "INVALID_VALUE" + SEMANTIC_ERROR = "SEMANTIC_ERROR" + WORKSPACE_VIOLATION = "WORKSPACE_VIOLATION" + DEPENDENCY_MISSING = "DEPENDENCY_MISSING" + DEPRECATED_FIELD = "DEPRECATED_FIELD" + + +class ConfigValidationIssue(BaseModel): + """One config issue located by an RFC 6901 JSON Pointer.""" + + model_config = ConfigDict(frozen=True) + + code: ConfigValidationCode + path: str + message: str + + +class ConfigValidationReport(BaseModel): + """Versioned structural validation result.""" + + model_config = ConfigDict(frozen=True) + + schema_version: int = EVAL_CONFIG_SCHEMA_VERSION + valid: bool + issues: tuple[ConfigValidationIssue, ...] = () + warnings: tuple[ConfigValidationIssue, ...] = () + + +class EvalConfigDocumentError(ValueError): + """Raised when a mapping does not conform to ``EvalConfigDocument``.""" + + def __init__(self, issues: tuple[ConfigValidationIssue, ...]) -> None: + self.issues = issues + super().__init__(format_config_validation_issues(issues)) + + +def _json_pointer(location: tuple[str | int, ...]) -> str: + if not location: + return "" + parts = [] + for part in location: + escaped = str(part).replace("~", "~0").replace("/", "~1") + parts.append(escaped) + return "/" + "/".join(parts) + + +def _issue_code(error_type: str) -> ConfigValidationCode: + if error_type == "missing": + return ConfigValidationCode.REQUIRED_FIELD + if error_type == "extra_forbidden": + return ConfigValidationCode.UNKNOWN_FIELD + if error_type.endswith(("_type", "_parsing")) or error_type in { + "bool_type", + "dict_type", + "float_type", + "int_type", + "list_type", + "model_type", + "string_type", + }: + return ConfigValidationCode.INVALID_TYPE + return ConfigValidationCode.INVALID_VALUE + + +def _validation_issues(exc: ValidationError) -> tuple[ConfigValidationIssue, ...]: + issues: list[ConfigValidationIssue] = [] + for error in exc.errors( + include_context=False, + include_input=False, + include_url=False, + ): + location = tuple(error.get("loc") or ()) + issues.append( + ConfigValidationIssue( + code=_issue_code(str(error.get("type") or "")), + path=_json_pointer(location), + message=str(error.get("msg") or "Invalid value"), + ) + ) + return tuple(issues) + + +def validate_eval_config_document(raw: Any) -> ConfigValidationReport: + """Validate a decoded YAML value without resolving paths or loading presets.""" + try: + EvalConfigDocument.model_validate(raw) + except ValidationError as exc: + issues = _validation_issues(exc) + return ConfigValidationReport(valid=False, issues=issues) + return ConfigValidationReport(valid=True) + + +def require_valid_eval_config_document(raw: Any) -> EvalConfigDocument: + """Return the parsed document or raise a stable issue-bearing error.""" + try: + return EvalConfigDocument.model_validate(raw) + except ValidationError as exc: + raise EvalConfigDocumentError(_validation_issues(exc)) from exc + + +def validate_eval_config_yaml(yaml_text: str) -> ConfigValidationReport: + """Validate YAML syntax and the decoded config document.""" + try: + raw = yaml.safe_load(yaml_text) + except yaml.YAMLError as exc: + issue = ConfigValidationIssue( + code=ConfigValidationCode.INVALID_YAML, + path="", + message=str(exc), + ) + return ConfigValidationReport(valid=False, issues=(issue,)) + return validate_eval_config_document(raw) + + +def format_config_validation_issues( + issues: tuple[ConfigValidationIssue, ...] | list[ConfigValidationIssue], +) -> str: + """Render stable compact issue lines for CLI and init feedback.""" + return "; ".join( + f"{issue.code.value} {issue.path or '/'}: {issue.message}" + for issue in issues + ) + + +def get_eval_config_json_schema() -> dict[str, Any]: + """Return the versioned JSON Schema used by MCP and other adapters.""" + schema = deepcopy(EvalConfigDocument.model_json_schema(mode="validation")) + schema["$schema"] = "https://json-schema.org/draft/2020-12/schema" + schema["$id"] = EVAL_CONFIG_SCHEMA_ID + schema["x-assert-schema-version"] = EVAL_CONFIG_SCHEMA_VERSION + return schema diff --git a/assert_ai/core/io.py b/assert_ai/core/io.py index 41fb775de..2a434a35c 100644 --- a/assert_ai/core/io.py +++ b/assert_ai/core/io.py @@ -50,6 +50,32 @@ def append_jsonl_row(path: Path, row: Dict[str, Any]) -> None: os.fsync(handle.fileno()) +def write_text_atomic(path: Path, text: str) -> None: + """Atomically replace a UTF-8 text file after flushing its contents.""" + path.parent.mkdir(parents=True, exist_ok=True) + tmp_name: str | None = None + try: + with tempfile.NamedTemporaryFile( + "wb", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as handle: + handle.write(text.encode("utf-8")) + handle.flush() + os.fsync(handle.fileno()) + tmp_name = handle.name + os.replace(tmp_name, path) + tmp_name = None + finally: + if tmp_name is not None: + try: + os.unlink(tmp_name) + except FileNotFoundError: + pass + + def _atomic_write_text(path: Path, text: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) tmp_name: str | None = None diff --git a/assert_ai/core/runtime_path_policy.py b/assert_ai/core/runtime_path_policy.py index a80c3bca1..ca7ec891a 100644 --- a/assert_ai/core/runtime_path_policy.py +++ b/assert_ai/core/runtime_path_policy.py @@ -141,16 +141,30 @@ def resolve_config_path( path: str | Path, *, must_exist: bool = False, + reject_links: bool = False, ) -> Path: """Resolve a config path strictly under ``config_root``.""" candidate = Path(path).expanduser() if candidate.is_absolute(): - resolved = candidate.resolve() + unresolved = candidate else: parts = candidate.parts if parts and parts[0] == self.config_root.name: candidate = Path(*parts[1:]) if len(parts) > 1 else Path() - resolved = (self.config_root / candidate).resolve() + unresolved = self.config_root / candidate + self._require_within( + Path(os.path.abspath(unresolved)), + self.config_root, + field_name="config", + code=RuntimePathErrorCode.OUTSIDE_CONFIG_ROOT, + ) + if reject_links: + self._require_no_links( + unresolved, + self.config_root, + field_name="config", + ) + resolved = unresolved.resolve() self._require_within( resolved, self.config_root, diff --git a/assert_ai/init/_design_agent.py b/assert_ai/init/_design_agent.py index 29aff1ec8..686599ca4 100644 --- a/assert_ai/init/_design_agent.py +++ b/assert_ai/init/_design_agent.py @@ -218,6 +218,7 @@ def run_design_loop( max_turns: int, console: Console, no_color: bool, + save_draft_on_failure: bool = True, ) -> str | None: """Run the design agent conversation loop. @@ -301,7 +302,7 @@ def run_design_loop( ) except (LLMAuthError, LLMInputError, LLMRateLimitError, LLMProviderError) as exc: log.error("LLM error: %s", exc) - if best_draft and best_errors: + if save_draft_on_failure and best_draft and best_errors: _save_draft(best_draft, best_errors, console) return None return best_draft @@ -419,13 +420,13 @@ def run_design_loop( # Exhausted turn budget. log.warning("Reached maximum turns (%d).", max_turns) - if best_draft: + if save_draft_on_failure and best_draft: _save_draft(best_draft, best_errors, console) return best_draft if best_draft and not best_errors else None except KeyboardInterrupt: console.print("") - if best_draft: + if save_draft_on_failure and best_draft: _save_draft(best_draft, best_errors, console) log.info("Interrupted. Draft saved.") else: diff --git a/assert_ai/init/_validate.py b/assert_ai/init/_validate.py index 6919d2c66..1dceb70dc 100644 --- a/assert_ai/init/_validate.py +++ b/assert_ai/init/_validate.py @@ -1,16 +1,21 @@ -"""Validation bridge — check a proposed YAML against config rules. - -Uses the same structural checks as ``assert_ai.config`` without requiring -stage modules or filesystem context so the design agent can validate -proposals mid-conversation. -""" +"""Headless validation bridge for config-design clients.""" from __future__ import annotations +import re from typing import Any import yaml +from assert_ai.core.config_document import ( + ConfigValidationCode, + ConfigValidationIssue, + format_config_validation_issues, + validate_eval_config_document, +) + +_SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*$") + def validate_proposed_yaml(yaml_str: str) -> tuple[bool, list[str]]: """Validate a proposed YAML string against the config schema. @@ -20,107 +25,103 @@ def validate_proposed_yaml(yaml_str: str) -> tuple[bool, list[str]]: try: data = yaml.safe_load(yaml_str) except yaml.YAMLError as exc: - return False, [f"Invalid YAML syntax: {exc}"] - + issue = ConfigValidationIssue( + code=ConfigValidationCode.INVALID_YAML, + path="", + message=str(exc), + ) + return False, [ + format_config_validation_issues([issue]) + ] if not isinstance(data, dict): - return False, ["Top-level YAML must be a mapping"] - + report = validate_eval_config_document(data) + issues = report.issues or ( + ConfigValidationIssue( + code=ConfigValidationCode.INVALID_TYPE, + path="", + message="Top-level YAML must be a mapping", + ), + ) + return False, [ + format_config_validation_issues([issue]) + for issue in issues + ] return validate_raw_config(data) def validate_raw_config(data: dict[str, Any]) -> tuple[bool, list[str]]: - """Structurally validate a raw config dict. - - Checks top-level keys, required fields, identifiers, behavior/context - shape, and pipeline stage names. Does **not** resolve file paths or - load stage modules — those are runtime concerns. - """ - from assert_ai.config import ( - ConfigError, - _SAFE_ID_RE, - PIPELINE_STAGE_ORDER, - reject_unknown_keys, - ) - - errors: list[str] = [] - - # -- top-level keys ------------------------------------------------------ - allowed_top = { - "suite", "run", "behavior", "context", "default_model", - "artifacts_root", "results_dir", "pipeline", - } - unknown = sorted(set(data) - allowed_top) - if unknown: - errors.append(f"Unknown top-level key(s): {', '.join(unknown)}") + """Validate document shape plus inexpensive identifier semantics.""" + report = validate_eval_config_document(data) + issues = list(report.issues) - # -- identifiers --------------------------------------------------------- for field in ("suite", "run"): val = data.get(field) if val is not None: val = str(val) if not _SAFE_ID_RE.match(val): - errors.append( - f"{field} must start with an alphanumeric character and " - f"contain only alphanumerics, dots, hyphens, or underscores; " - f"got: {val!r}" + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.INVALID_VALUE, + path=f"/{field}", + message=( + "must start with an alphanumeric character and contain only " + f"alphanumerics, dots, hyphens, or underscores; got: {val!r}" + ), + ) ) if ".." in val: - errors.append(f"{field} must not contain '..'") + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.INVALID_VALUE, + path=f"/{field}", + message="must not contain '..'", + ) + ) if len(val) > 255: - errors.append(f"{field} exceeds maximum length of 255 characters") + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.INVALID_VALUE, + path=f"/{field}", + message="exceeds maximum length of 255 characters", + ) + ) - # -- behavior ------------------------------------------------------------ behavior = data.get("behavior") - if behavior is None: - errors.append("'behavior' is required") - elif not isinstance(behavior, dict): - errors.append("behavior must be a mapping") - else: - beh_allowed = {"name", "description", "preset"} - beh_unknown = sorted(set(behavior) - beh_allowed) - if beh_unknown: - errors.append(f"behavior has unsupported field(s): {', '.join(beh_unknown)}") + pipeline = data.get("pipeline") + systematize = pipeline.get("systematize") if isinstance(pipeline, dict) else None + requires_behavior = ( + isinstance(systematize, dict) + and systematize.get("enabled", True) + ) + if behavior is None and requires_behavior: + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.REQUIRED_FIELD, + path="/behavior", + message="is required when systematize is enabled", + ) + ) + if isinstance(behavior, dict): name = behavior.get("name") if name is not None: name = str(name) if not _SAFE_ID_RE.match(name): - errors.append( - f"behavior.name must be a valid identifier; got: {name!r}" + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.INVALID_VALUE, + path="/behavior/name", + message=f"must be a valid identifier; got: {name!r}", + ) ) if not name and not behavior.get("preset"): - errors.append("behavior.name is required (or use behavior.preset)") - - # -- context ------------------------------------------------------------- - context = data.get("context") - if context is not None and not isinstance(context, str): - errors.append("context must be a string") - - # -- pipeline ------------------------------------------------------------ - pipeline = data.get("pipeline") - if pipeline is None: - errors.append("'pipeline' is required") - elif not isinstance(pipeline, dict): - errors.append("'pipeline' must be a mapping") - else: - valid_stages = set(PIPELINE_STAGE_ORDER) - unknown_stages = sorted(set(pipeline) - valid_stages) - if unknown_stages: - errors.append(f"Unknown pipeline stage(s): {', '.join(unknown_stages)}") - if not any(s in pipeline for s in valid_stages): - errors.append("'pipeline' must define at least one stage") - - # -- default_model ------------------------------------------------------- - dm = data.get("default_model") - if dm is not None: - if isinstance(dm, str): - pass # shorthand form — valid - elif isinstance(dm, dict): - if "name" not in dm: - errors.append("default_model.name is required") - else: - errors.append("default_model must be a string or mapping") + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.REQUIRED_FIELD, + path="/behavior/name", + message="is required (or use behavior.preset)", + ) + ) - # -- dimensions: reserved name check ------------------------------------ if isinstance(pipeline, dict): test_set = pipeline.get("test_set") if isinstance(test_set, dict): @@ -128,13 +129,26 @@ def validate_raw_config(data: dict[str, Any]) -> tuple[bool, list[str]]: if isinstance(stratify, dict): dims = stratify.get("dimensions") if isinstance(dims, list): - for dim in dims: + for index, dim in enumerate(dims): if isinstance(dim, dict): dim_name = dim.get("name", "") if dim_name == "behavior": - errors.append( - "Dimension name 'behavior' is reserved; " - "choose a different name" + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.INVALID_VALUE, + path=( + "/pipeline/test_set/stratify/" + f"dimensions/{index}/name" + ), + message=( + "Dimension name 'behavior' is reserved; " + "choose a different name" + ), + ) ) + errors = [ + format_config_validation_issues([issue]) + for issue in issues + ] return (not errors), errors diff --git a/assert_ai/services/__init__.py b/assert_ai/services/__init__.py new file mode 100644 index 000000000..770cc60a6 --- /dev/null +++ b/assert_ai/services/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Transport-neutral application services for ASSERT.""" diff --git a/assert_ai/services/configs.py b/assert_ai/services/configs.py new file mode 100644 index 000000000..eacfe2a7a --- /dev/null +++ b/assert_ai/services/configs.py @@ -0,0 +1,852 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Workspace-scoped evaluation config lifecycle.""" + +from __future__ import annotations + +import base64 +import hashlib +import importlib.util +import json +import os +import re +import time +from bisect import bisect +from contextlib import contextmanager +from copy import deepcopy +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Iterator, Mapping + +import yaml +from pydantic import BaseModel, ConfigDict, Field + +from assert_ai.config import ConfigError, load_runtime_context, parse_model_config +from assert_ai.core.config_document import ( + ConfigValidationCode, + ConfigValidationIssue, + ConfigValidationReport, + EVAL_CONFIG_SCHEMA_VERSION, + get_eval_config_json_schema, + validate_eval_config_document, +) +from assert_ai.core.io import write_text_atomic +from assert_ai.core.runtime_path_policy import RuntimePathError +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.stages import STAGES +from assert_ai.stages.test_set import validate_sampling_config + +_CONFIG_SUFFIXES = {".yaml", ".yml"} +_CURSOR_VERSION = 1 +_DEFAULT_MAX_CONFIG_BYTES = 1_048_576 +_DEFAULT_PAGE_SIZE = 50 +_DEFAULT_MAX_PAGE_SIZE = 200 +_LOCK_TIMEOUT_S = 10.0 + + +class _ServiceModel(BaseModel): + model_config = ConfigDict(frozen=True) + + +class ConfigCatalogEntry(_ServiceModel): + """Lightweight metadata for one managed config.""" + + config_ref: str + etag: str + size_bytes: int + modified_at: str + structurally_valid: bool + + +class ConfigPage(_ServiceModel): + """Bounded page of config catalog entries.""" + + items: tuple[ConfigCatalogEntry, ...] + next_cursor: str | None = None + + +class ConfigRecord(_ServiceModel): + """One normalized config and its validation state.""" + + config_ref: str + yaml: str + document: dict[str, Any] + etag: str + validation: ConfigValidationReport + + +class ConfigSaveResult(_ServiceModel): + """Identity and optimistic-concurrency token after a save.""" + + config_ref: str + etag: str + created: bool + validation: ConfigValidationReport + + +class ConfigDesignRequest(_ServiceModel): + """One headless design-agent request.""" + + description: str + model: str = "azure/gpt-5.4-mini" + seed_config_ref: str | None = None + seed_yaml: str | None = None + behavior_preset: str | None = None + judge_preset: str | None = None + dimension_hints: str | None = None + default_model_hint: str | None = None + max_turns: int = Field(default=5, ge=1, le=100) + + +class ConfigDraft(_ServiceModel): + """Model-generated draft that has not been persisted.""" + + yaml: str + document: dict[str, Any] + validation: ConfigValidationReport + + +@dataclass(slots=True) +class ConfigService: + """Read, validate, design, and atomically save managed eval configs.""" + + workspace: WorkspaceService + max_config_bytes: int = _DEFAULT_MAX_CONFIG_BYTES + default_page_size: int = _DEFAULT_PAGE_SIZE + max_page_size: int = _DEFAULT_MAX_PAGE_SIZE + + def get_schema(self) -> dict[str, Any]: + return get_eval_config_json_schema() + + def list_configs( + self, + *, + limit: int | None = None, + cursor: str | None = None, + ) -> ConfigPage: + page_size = self._page_size(limit) + after = _decode_cursor(cursor) if cursor else None + entries = self._catalog_entries() + refs = [entry.config_ref for entry in entries] + start = bisect(refs, after) if after is not None else 0 + page_items = entries[start : start + page_size] + next_cursor = None + if start + len(page_items) < len(entries) and page_items: + next_cursor = _encode_cursor(page_items[-1].config_ref) + return ConfigPage(items=tuple(page_items), next_cursor=next_cursor) + + def get_config(self, config_ref: str) -> ConfigRecord: + path = self._resolve_ref(config_ref, must_exist=True, reject_links=True) + raw_bytes = self._read_bytes(path) + yaml_text = self._decode(raw_bytes, config_ref=config_ref) + raw = _load_yaml_mapping(yaml_text) + normalized = _normalize_yaml(raw) + return ConfigRecord( + config_ref=self._config_ref(path), + yaml=normalized, + document=raw, + etag=_etag(raw_bytes), + validation=self.validate_document(raw, config_ref=self._config_ref(path)), + ) + + def validate_yaml( + self, + yaml_text: str, + *, + config_ref: str = "draft.yaml", + ) -> ConfigValidationReport: + self._check_payload_size(yaml_text.encode("utf-8")) + try: + raw = yaml.safe_load(yaml_text) + except yaml.YAMLError as exc: + return ConfigValidationReport( + valid=False, + issues=( + ConfigValidationIssue( + code=ConfigValidationCode.INVALID_YAML, + path="", + message=str(exc), + ), + ), + ) + return self.validate_document(raw, config_ref=config_ref) + + def validate_document( + self, + document: Any, + *, + config_ref: str = "draft.yaml", + ) -> ConfigValidationReport: + config_path = self._resolve_ref( + config_ref, + must_exist=False, + reject_links=True, + ) + structural = validate_eval_config_document(document) + warnings = ( + list(_compatibility_warnings(document)) + if isinstance(document, dict) + else [] + ) + if not structural.valid: + return ConfigValidationReport( + schema_version=structural.schema_version, + valid=False, + issues=structural.issues, + warnings=tuple(warnings), + ) + + assert isinstance(document, dict) + issues = _stage_semantic_issues(document) + issues.extend(_dependency_issues(document)) + runtime_document = deepcopy(document) + try: + load_runtime_context( + runtime_document, + config_path, + stage_modules=STAGES, + path_policy=self.workspace.path_policy, + ) + except RuntimePathError as exc: + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.WORKSPACE_VIOLATION, + path=_field_name_pointer(exc.field_name), + message=( + f"{exc.field_name} violates workspace path policy " + f"({exc.code.value})" + ), + ) + ) + except (ConfigError, ValueError, FileNotFoundError) as exc: + issues.append( + ConfigValidationIssue( + code=ConfigValidationCode.SEMANTIC_ERROR, + path=_semantic_error_pointer(str(exc)), + message=str(exc), + ) + ) + + return ConfigValidationReport( + schema_version=EVAL_CONFIG_SCHEMA_VERSION, + valid=not issues, + issues=tuple(_deduplicate_issues(issues)), + warnings=tuple(_deduplicate_issues(warnings)), + ) + + def save_config( + self, + config_ref: str, + *, + yaml_text: str | None = None, + document: Mapping[str, Any] | None = None, + expected_etag: str | None = None, + ) -> ConfigSaveResult: + if (yaml_text is None) == (document is None): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Provide exactly one of yaml_text or document", + ) + if yaml_text is not None: + self._check_payload_size(yaml_text.encode("utf-8")) + try: + raw = yaml.safe_load(yaml_text) + except yaml.YAMLError as exc: + report = ConfigValidationReport( + valid=False, + issues=( + ConfigValidationIssue( + code=ConfigValidationCode.INVALID_YAML, + path="", + message=str(exc), + ), + ), + ) + raise _invalid_config_error(report) from exc + else: + raw = deepcopy(dict(document or {})) + + report = self.validate_document(raw, config_ref=config_ref) + if not report.valid: + raise _invalid_config_error(report) + + assert isinstance(raw, dict) + try: + normalized = _normalize_yaml(raw) + except yaml.YAMLError as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Config document cannot be serialized as YAML", + ) from exc + encoded = normalized.encode("utf-8") + self._check_payload_size(encoded) + path = self._resolve_ref(config_ref, must_exist=False, reject_links=True) + self._ensure_parent(path) + + with self._config_lock(path): + path = self._resolve_ref(config_ref, must_exist=False, reject_links=True) + exists = path.is_file() + current_etag = _etag(self._read_bytes(path)) if exists else None + if exists and expected_etag is None: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "expected_etag is required when replacing an existing config", + details={"config_ref": self._config_ref(path)}, + ) + if expected_etag is not None and current_etag != expected_etag: + raise ServiceError( + ServiceErrorCode.STALE_ETAG, + "Config changed since it was read", + details={ + "config_ref": self._config_ref(path), + "current_etag": current_etag, + }, + ) + write_text_atomic(path, normalized) + saved_etag = _etag(encoded) + + return ConfigSaveResult( + config_ref=self._config_ref(path), + etag=saved_etag, + created=not exists, + validation=report, + ) + + def design_config(self, request: ConfigDesignRequest) -> ConfigDraft: + if request.seed_config_ref and request.seed_yaml is not None: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Provide seed_config_ref or seed_yaml, not both", + ) + seed_yaml = request.seed_yaml + validation_ref = "draft.yaml" + if request.seed_config_ref: + seed = self.get_config(request.seed_config_ref) + seed_yaml = seed.yaml + validation_ref = seed.config_ref + if seed_yaml is not None: + self._check_payload_size(seed_yaml.encode("utf-8")) + + from rich.console import Console + + from assert_ai.init._design_agent import run_design_loop + + yaml_result = run_design_loop( + model=request.model, + describe=request.description, + seed_yaml=seed_yaml, + seed_path=None, + behavior_preset=request.behavior_preset, + judge_preset=request.judge_preset, + dimension_hints=request.dimension_hints, + default_model_hint=request.default_model_hint, + non_interactive=True, + max_turns=request.max_turns, + console=Console(quiet=True, stderr=True), + no_color=True, + save_draft_on_failure=False, + ) + if yaml_result is None: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Design agent did not produce a valid config", + ) + raw = _load_yaml_mapping(yaml_result) + normalized = _normalize_yaml(raw) + return ConfigDraft( + yaml=normalized, + document=raw, + validation=self.validate_document(raw, config_ref=validation_ref), + ) + + def _catalog_entries(self) -> list[ConfigCatalogEntry]: + root = self.workspace.configs_root + if not root.exists(): + return [] + entries: list[ConfigCatalogEntry] = [] + for candidate in self._config_paths(): + relative = candidate.relative_to(root).as_posix() + path = self._resolve_ref(relative, must_exist=True, reject_links=True) + raw_bytes = self._read_bytes(path) + try: + yaml_text = raw_bytes.decode("utf-8") + except UnicodeDecodeError: + structurally_valid = False + else: + try: + decoded = yaml.safe_load(yaml_text) + except yaml.YAMLError: + structurally_valid = False + else: + structurally_valid = validate_eval_config_document(decoded).valid + stat_result = path.stat() + entries.append( + ConfigCatalogEntry( + config_ref=self._config_ref(path), + etag=_etag(raw_bytes), + size_bytes=len(raw_bytes), + modified_at=datetime.fromtimestamp( + stat_result.st_mtime, + tz=timezone.utc, + ).isoformat(), + structurally_valid=structurally_valid, + ) + ) + entries.sort(key=lambda entry: entry.config_ref) + return entries + + def _config_paths(self) -> list[Path]: + root = self.workspace.configs_root + pending = [root] + paths: list[Path] = [] + while pending: + directory = pending.pop() + try: + entries = list(os.scandir(directory)) + except FileNotFoundError: + continue + for entry in entries: + candidate = Path(entry.path) + relative = candidate.relative_to(root).as_posix() + try: + self.workspace.path_policy.resolve_config_path( + relative, + reject_links=True, + ) + except RuntimePathError as exc: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + "Managed config tree contains a symbolic link or junction", + details={"config_ref": relative}, + ) from exc + if entry.is_dir(follow_symlinks=False): + pending.append(candidate) + elif ( + entry.is_file(follow_symlinks=False) + and candidate.suffix.lower() in _CONFIG_SUFFIXES + ): + paths.append(candidate) + return sorted(paths) + + def _resolve_ref( + self, + config_ref: str, + *, + must_exist: bool, + reject_links: bool, + ) -> Path: + if not isinstance(config_ref, str) or not config_ref.strip(): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "config_ref must be a non-empty string", + ) + ref = config_ref.strip().replace("\\", "/") + if Path(ref).is_absolute(): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "config_ref must be relative to the managed config root", + ) + if Path(ref).suffix.lower() not in _CONFIG_SUFFIXES: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "config_ref must end in .yaml or .yml", + ) + try: + path = self.workspace.path_policy.resolve_config_path( + ref, + must_exist=must_exist, + reject_links=reject_links, + ) + except RuntimePathError as exc: + code = ( + ServiceErrorCode.NOT_FOUND + if exc.code.value in {"path_not_found", "not_a_file"} + else ServiceErrorCode.WORKSPACE_VIOLATION + ) + raise ServiceError(code, str(exc)) from exc + return path + + def _config_ref(self, path: Path) -> str: + return path.relative_to(self.workspace.configs_root).as_posix() + + def _ensure_parent(self, path: Path) -> None: + self.workspace.configs_root.mkdir(parents=True, exist_ok=True) + path.parent.mkdir(parents=True, exist_ok=True) + self.workspace.path_policy.resolve_config_path( + path, + reject_links=True, + ) + + def _read_bytes(self, path: Path) -> bytes: + try: + size = path.stat().st_size + except FileNotFoundError as exc: + raise ServiceError(ServiceErrorCode.NOT_FOUND, "Config not found") from exc + if size > self.max_config_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"Config exceeds the {self.max_config_bytes}-byte limit", + ) + data = path.read_bytes() + self._check_payload_size(data) + return data + + def _decode(self, data: bytes, *, config_ref: str) -> str: + try: + return data.decode("utf-8") + except UnicodeDecodeError as exc: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Config is not valid UTF-8: {config_ref}", + ) from exc + + def _check_payload_size(self, data: bytes) -> None: + if len(data) > self.max_config_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"Config exceeds the {self.max_config_bytes}-byte limit", + ) + + def _page_size(self, value: int | None) -> int: + size = self.default_page_size if value is None else value + if isinstance(size, bool) or not isinstance(size, int) or size < 1: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "limit must be a positive integer", + ) + if size > self.max_page_size: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"limit must be <= {self.max_page_size}", + ) + return size + + @contextmanager + def _config_lock(self, path: Path) -> Iterator[None]: + lock_name = hashlib.sha256(self._config_ref(path).encode("utf-8")).hexdigest() + lock_ref = f".locks/{lock_name}.lock" + lock_path = self.workspace.path_policy.resolve_config_path( + lock_ref, + reject_links=True, + ) + lock_path.parent.mkdir(parents=True, exist_ok=True) + lock_path = self.workspace.path_policy.resolve_config_path( + lock_ref, + reject_links=True, + ) + with _exclusive_file_lock(lock_path, timeout_s=_LOCK_TIMEOUT_S): + yield + + +def _load_yaml_mapping(yaml_text: str) -> dict[str, Any]: + try: + raw = yaml.safe_load(yaml_text) + except yaml.YAMLError as exc: + raise ServiceError(ServiceErrorCode.CONFIG_INVALID, "Invalid YAML") from exc + if not isinstance(raw, dict): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Top-level YAML must be a mapping", + ) + return raw + + +def _normalize_yaml(document: Mapping[str, Any]) -> str: + normalized = yaml.safe_dump( + dict(document), + default_flow_style=False, + sort_keys=False, + allow_unicode=True, + ) + return normalized if normalized.endswith("\n") else normalized + "\n" + + +def _etag(data: bytes) -> str: + return f"sha256:{hashlib.sha256(data).hexdigest()}" + + +def _encode_cursor(config_ref: str) -> str: + payload = json.dumps( + {"v": _CURSOR_VERSION, "after": config_ref}, + separators=(",", ":"), + ).encode("utf-8") + return base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=") + + +def _decode_cursor(cursor: str) -> str: + try: + padding = "=" * (-len(cursor) % 4) + payload = json.loads( + base64.urlsafe_b64decode(cursor + padding).decode("utf-8") + ) + except (ValueError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid config cursor", + ) from exc + if ( + not isinstance(payload, dict) + or payload.get("v") != _CURSOR_VERSION + or not isinstance(payload.get("after"), str) + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid config cursor", + ) + return payload["after"] + + +def _compatibility_warnings( + document: dict[str, Any], +) -> tuple[ConfigValidationIssue, ...]: + test_set = (document.get("pipeline") or {}).get("test_set") + if isinstance(test_set, dict) and test_set.get("tool_source") == "per_seed": + return ( + ConfigValidationIssue( + code=ConfigValidationCode.DEPRECATED_FIELD, + path="/pipeline/test_set/tool_source", + message="per_seed is deprecated; use per_test_case", + ), + ) + return () + + +def _stage_semantic_issues( + document: dict[str, Any], +) -> list[ConfigValidationIssue]: + issues: list[ConfigValidationIssue] = [] + pipeline = document["pipeline"] + default_model = document.get("default_model") + + systematize = pipeline.get("systematize") + if isinstance(systematize, dict) and systematize.get("enabled", True): + if "validators" in systematize or "validator_models" in systematize: + issues.append( + _semantic_issue( + "/pipeline/systematize", + "taxonomy validators are no longer supported", + ) + ) + _validate_model( + systematize.get("model") or default_model, + path="/pipeline/systematize/model", + required_message="systematize.model or default_model is required", + issues=issues, + ) + + test_set = pipeline.get("test_set") + if isinstance(test_set, dict) and test_set.get("enabled", True): + if "validators" in test_set or "validator_model" in test_set: + issues.append( + _semantic_issue( + "/pipeline/test_set", + "test_set validators are no longer supported", + ) + ) + prompt = test_set.get("prompt") + scenario = test_set.get("scenario") + if not prompt and not scenario: + issues.append( + _semantic_issue( + "/pipeline/test_set", + "test_set requires prompt and/or scenario configuration", + ) + ) + for kind, raw_kind in (("prompt", prompt), ("scenario", scenario)): + if not raw_kind: + continue + path = f"/pipeline/test_set/{kind}" + if "budget" in raw_kind: + issues.append( + _semantic_issue( + f"{path}/budget", + f"test_set.{kind}.budget was renamed to " + f"test_set.{kind}.sample_size", + ) + ) + if kind == "scenario" and "modality" in raw_kind: + issues.append( + _semantic_issue( + f"{path}/modality", + "test_set.scenario.modality is no longer supported; " + "use test_set.tool_source", + ) + ) + _validate_model( + raw_kind.get("model") + or test_set.get("model") + or default_model, + path=f"{path}/model", + required_message=f"test_set.{kind}.model is required", + issues=issues, + ) + try: + validate_sampling_config( + raw_kind.get("sampling"), + field_name=f"test_set.{kind}.sampling", + ) + except ValueError as exc: + issues.append(_semantic_issue(f"{path}/sampling", str(exc))) + + stratify = test_set.get("stratify") + if isinstance(stratify, dict) and stratify.get("model") is not None: + _validate_model( + stratify["model"], + path="/pipeline/test_set/stratify/model", + required_message="test_set.stratify.model is invalid", + issues=issues, + ) + + return issues + + +def _dependency_issues(document: dict[str, Any]) -> list[ConfigValidationIssue]: + inference = (document.get("pipeline") or {}).get("inference") + target = inference.get("target") if isinstance(inference, dict) else None + trace = target.get("trace") if isinstance(target, dict) else None + if ( + isinstance(trace, dict) + and trace.get("backend", "phoenix") == "phoenix" + and importlib.util.find_spec("phoenix") is None + ): + return [ + ConfigValidationIssue( + code=ConfigValidationCode.DEPENDENCY_MISSING, + path="/pipeline/inference/target/trace/backend", + message="Phoenix trace capture requires the otel optional dependency", + ) + ] + return [] + + +def _validate_model( + raw: Any, + *, + path: str, + required_message: str, + issues: list[ConfigValidationIssue], +) -> None: + if raw is None: + issues.append(_semantic_issue(path, required_message)) + return + try: + parse_model_config(raw, field_name=path.strip("/").replace("/", ".")) + except ValueError as exc: + issues.append(_semantic_issue(path, str(exc))) + + +def _semantic_issue(path: str, message: str) -> ConfigValidationIssue: + return ConfigValidationIssue( + code=ConfigValidationCode.SEMANTIC_ERROR, + path=path, + message=message, + ) + + +def _field_name_pointer(field_name: str) -> str: + if field_name.startswith("pipeline."): + return "/" + field_name.replace(".", "/") + if field_name in { + "artifacts_root", + "results_dir", + "suite", + "run", + "behavior", + "context", + }: + return f"/{field_name}" + return "" + + +def _semantic_error_pointer(message: str) -> str: + prefixes = ( + ("pipeline.", ""), + ("default_model", "/default_model"), + ("behavior.", "/behavior/"), + ("context", "/context"), + ("suite", "/suite"), + ("run", "/run"), + ("target.", "/pipeline/inference/target/"), + ("target ", "/pipeline/inference/target"), + ("trace.", "/pipeline/inference/target/trace/"), + ("test_set.", "/pipeline/test_set/"), + ("systematize.", "/pipeline/systematize/"), + ("inference.", "/pipeline/inference/"), + ("judge.", "/pipeline/judge/"), + ) + for prefix, replacement in prefixes: + if not message.startswith(prefix): + continue + token = re.split(r"\s", message, maxsplit=1)[0].rstrip(":") + if prefix == "pipeline.": + return "/" + token.replace(".", "/") + suffix = token[len(prefix):].replace(".", "/") + return replacement + suffix + return "" + + +def _deduplicate_issues( + issues: list[ConfigValidationIssue], +) -> list[ConfigValidationIssue]: + result: list[ConfigValidationIssue] = [] + seen: set[tuple[str, str, str]] = set() + for issue in issues: + key = (issue.code.value, issue.path, issue.message) + if key in seen: + continue + seen.add(key) + result.append(issue) + return result + + +def _invalid_config_error(report: ConfigValidationReport) -> ServiceError: + return ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Config validation failed", + details={"validation": report.model_dump(mode="json")}, + ) + + +@contextmanager +def _exclusive_file_lock(path: Path, *, timeout_s: float) -> Iterator[None]: + deadline = time.monotonic() + timeout_s + with path.open("a+b") as handle: + handle.seek(0, os.SEEK_END) + if handle.tell() == 0: + handle.write(b"\0") + handle.flush() + os.fsync(handle.fileno()) + handle.seek(0) + while True: + try: + if os.name == "nt": + import msvcrt + + msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) + else: + import fcntl + + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except OSError as exc: + if time.monotonic() >= deadline: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "Timed out waiting for the config write lock", + ) from exc + time.sleep(0.05) + try: + yield + finally: + handle.seek(0) + if os.name == "nt": + import msvcrt + + msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) + else: + import fcntl + + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) diff --git a/assert_ai/services/errors.py b/assert_ai/services/errors.py new file mode 100644 index 000000000..2bf1ced1c --- /dev/null +++ b/assert_ai/services/errors.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Stable application-service error taxonomy.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Any + + +class ServiceErrorCode(StrEnum): + INVALID_ARGUMENT = "INVALID_ARGUMENT" + NOT_FOUND = "NOT_FOUND" + CONFLICT = "CONFLICT" + STALE_ETAG = "STALE_ETAG" + STALE_CURSOR = "STALE_CURSOR" + CAPABILITY_DISABLED = "CAPABILITY_DISABLED" + DEPENDENCY_MISSING = "DEPENDENCY_MISSING" + WORKSPACE_VIOLATION = "WORKSPACE_VIOLATION" + CONFIG_INVALID = "CONFIG_INVALID" + PREFLIGHT_FAILED = "PREFLIGHT_FAILED" + TARGET_IMPORT_FAILED = "TARGET_IMPORT_FAILED" + JOB_NOT_CANCELLABLE = "JOB_NOT_CANCELLABLE" + JOB_INTERRUPTED = "JOB_INTERRUPTED" + RUN_FAILED = "RUN_FAILED" + ARTIFACT_TOO_LARGE = "ARTIFACT_TOO_LARGE" + INTERNAL = "INTERNAL" + + +class ServiceError(Exception): + """Expected application failure suitable for CLI or MCP adaptation.""" + + def __init__( + self, + code: ServiceErrorCode, + message: str, + *, + details: dict[str, Any] | None = None, + ) -> None: + super().__init__(message) + self.code = code + self.details = details or {} diff --git a/assert_ai/stages/test_set.py b/assert_ai/stages/test_set.py index b7f853939..18b72953c 100644 --- a/assert_ai/stages/test_set.py +++ b/assert_ai/stages/test_set.py @@ -1275,6 +1275,8 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: "scenario", "validators", "validator_model", + "enabled", + "file_path", }, ) if "validators" in raw_cfg or "validator_model" in raw_cfg: diff --git a/docs/config/schema.md b/docs/config/schema.md index 437170495..74991f405 100644 --- a/docs/config/schema.md +++ b/docs/config/schema.md @@ -2,6 +2,11 @@ This page documents the `eval_config.yaml` schema for the standard `behavior -> systematize -> test_set -> inference -> judge` pipeline. +ASSERT generates its machine-readable JSON Schema from +`assert_ai.core.config_document.EvalConfigDocument`. Runtime loading and +`assert-ai init` use that same structural model before applying stage semantics, +path policy, preset loading, and default-model injection. + ## Top-level keys ### `suite` @@ -270,6 +275,7 @@ Accepted keys: - `model` — model config. Required unless `default_model` is set. - `n` — positive integer. Default: `1`. - `preset` — optional string or list of strings. Loads judge dimension presets; inline `dimensions` override preset dimensions with the same name. +- `disabled_dimensions` — optional list of built-in dimension names to omit. Compatibility note: diff --git a/tests/test_config_document.py b/tests/test_config_document.py new file mode 100644 index 000000000..0050f66e8 --- /dev/null +++ b/tests/test_config_document.py @@ -0,0 +1,223 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from pathlib import Path + +import yaml + +from assert_ai.core.config_document import ( + ConfigValidationCode, + EvalConfigDocument, + get_eval_config_json_schema, + validate_eval_config_document, +) + +ROOT = Path(__file__).resolve().parent.parent + + +def test_complete_document_validates() -> None: + raw = { + "suite": "demo-suite", + "run": "run-1", + "behavior": { + "name": "safe_assistance", + "description": "The assistant follows the behavior.", + }, + "context": "A tool-using support agent.", + "default_model": { + "name": "azure/gpt-5.4", + "reasoning_effort": "medium", + }, + "artifacts_root": "artifacts", + "results_dir": "artifacts/results", + "pipeline": { + "systematize": { + "behavior_category_count": 10, + "web_search": False, + }, + "test_set": { + "tool_source": "runtime", + "prompt": { + "sample_size": 5, + "sampling": { + "method": "stratified", + "stratify_by": ["behavior"], + }, + }, + "scenario": { + "sample_size": 5, + "sampling": { + "method": "random", + "with_replacement": False, + }, + }, + "stratify": { + "dimensions": [ + { + "name": "user_type", + "levels": [ + {"name": "new", "definition": "A new user."}, + {"name": "returning", "definition": "A returning user."}, + ], + } + ] + }, + }, + "inference": { + "target": { + "callable": "agent:run", + "trace": {"backend": "phoenix", "group_by": "session.id"}, + }, + "tester": {"model": {"name": "azure/gpt-5.4-mini"}}, + "max_turns": 5, + "concurrency": 2, + }, + "judge": { + "preset": ["policy", "quality"], + "disabled_dimensions": ["overrefusal"], + "dimensions": { + "response_quality": { + "description": "How strong was the response?", + "rubric": "Score overall response quality.", + "allow_not_applicable": True, + "scale": { + "type": "ordinal", + "values": { + 1: "Poor", + 2: "Acceptable", + 3: "Strong", + }, + }, + } + }, + }, + }, + } + + report = validate_eval_config_document(raw) + + assert report.valid is True + assert report.issues == () + assert EvalConfigDocument.model_validate(raw).pipeline.inference is not None + + +def test_unknown_nested_field_has_stable_json_pointer() -> None: + report = validate_eval_config_document( + { + "pipeline": { + "inference": { + "target": { + "model": {"name": "azure/gpt-5.4"}, + "legacy/type": "model", + } + } + } + } + ) + + assert report.valid is False + assert [issue.model_dump() for issue in report.issues] == [ + { + "code": ConfigValidationCode.UNKNOWN_FIELD, + "path": "/pipeline/inference/target/legacy~1type", + "message": "Extra inputs are not permitted", + } + ] + + +def test_json_pointer_escapes_dynamic_judge_dimension_names() -> None: + report = validate_eval_config_document( + { + "pipeline": { + "judge": { + "dimensions": { + "quality/~strict": { + "description": "Quality", + "rubric": "Score quality", + "unexpected": True, + } + } + } + } + } + ) + + assert report.valid is False + assert report.issues[0].path == ( + "/pipeline/judge/dimensions/quality~1~0strict/unexpected" + ) + + +def test_generated_schema_is_versioned_and_strict() -> None: + schema = get_eval_config_json_schema() + + assert schema["$schema"] == "https://json-schema.org/draft/2020-12/schema" + assert schema["x-assert-schema-version"] == 1 + assert schema["additionalProperties"] is False + assert schema["required"] == ["pipeline"] + pipeline_schema = schema["$defs"]["PipelineDocument"] + assert pipeline_schema["additionalProperties"] is False + assert pipeline_schema["minProperties"] == 1 + assert set(pipeline_schema["properties"]) == { + "systematize", + "test_set", + "inference", + "judge", + } + + +def test_explicit_null_stage_is_rejected() -> None: + report = validate_eval_config_document({"pipeline": {"judge": None}}) + + assert report.valid is False + assert report.issues[0].code == ConfigValidationCode.INVALID_TYPE + assert report.issues[0].path == "/pipeline/judge" + + +def test_dimension_warning_threshold_is_not_a_schema_limit() -> None: + report = validate_eval_config_document( + { + "pipeline": { + "test_set": { + "stratify": { + "dimensions": [ + {"name": f"dimension_{index}", "description": "Generated"} + for index in range(11) + ] + } + } + } + } + ) + + assert report.valid is True + + +def test_all_customer_eval_configs_match_document_shape() -> None: + paths = sorted((ROOT / "examples").glob("**/eval_config*.yaml")) + paths.extend(sorted((ROOT / "examples").glob("**/behaviors/*.yaml"))) + assert paths + + failures: list[str] = [] + for path in paths: + raw = yaml.safe_load(path.read_text(encoding="utf-8")) + report = validate_eval_config_document(raw) + if not report.valid: + failures.append( + f"{path.relative_to(ROOT)}: " + + "; ".join(f"{issue.path}: {issue.message}" for issue in report.issues) + ) + + assert failures == [] + + +def test_document_schema_fields_are_documented() -> None: + docs = (ROOT / "docs" / "config" / "schema.md").read_text(encoding="utf-8") + schema = get_eval_config_json_schema() + + for field_name in schema["properties"]: + assert f"### `{field_name}`" in docs + for stage_name in schema["$defs"]["PipelineDocument"]["properties"]: + assert f"### `pipeline.{stage_name}`" in docs diff --git a/tests/test_config_service.py b/tests/test_config_service.py new file mode 100644 index 000000000..9ddccb51a --- /dev/null +++ b/tests/test_config_service.py @@ -0,0 +1,262 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +import pytest + +from assert_ai.core.config_document import ConfigValidationCode +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.configs import ( + ConfigDesignRequest, + ConfigService, +) +from assert_ai.services.errors import ServiceError, ServiceErrorCode + + +def _service(root: Path, *, max_config_bytes: int = 1_048_576) -> ConfigService: + return ConfigService( + workspace=WorkspaceService.create(root), + max_config_bytes=max_config_bytes, + ) + + +def _valid_document(*, suite: str = "demo") -> dict: + return { + "suite": suite, + "behavior": {"name": "safe_help"}, + "pipeline": { + "inference": { + "target": {"callable": "agent:run"}, + "test_set_path": "fixtures/test_set.jsonl", + } + }, + } + + +def test_save_get_list_and_replace_with_etag() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + service = _service(root) + + saved = service.save_config("nested/demo.yaml", document=_valid_document()) + record = service.get_config("nested/demo.yaml") + page = service.list_configs(limit=1) + + assert saved.created is True + assert saved.etag == record.etag + assert record.config_ref == "nested/demo.yaml" + assert record.document["suite"] == "demo" + assert record.validation.valid is True + assert [entry.config_ref for entry in page.items] == ["nested/demo.yaml"] + assert page.next_cursor is None + + with pytest.raises(ServiceError) as missing_etag: + service.save_config( + "nested/demo.yaml", + document=_valid_document(suite="changed"), + ) + assert missing_etag.value.code == ServiceErrorCode.CONFLICT + + with pytest.raises(ServiceError) as stale: + service.save_config( + "nested/demo.yaml", + document=_valid_document(suite="changed"), + expected_etag="sha256:stale", + ) + assert stale.value.code == ServiceErrorCode.STALE_ETAG + + replaced = service.save_config( + "nested/demo.yaml", + document=_valid_document(suite="changed"), + expected_etag=record.etag, + ) + assert replaced.created is False + assert service.get_config("nested/demo.yaml").document["suite"] == "changed" + + +def test_concurrent_replacements_allow_only_one_etag_winner() -> None: + with TemporaryDirectory() as tmp: + service = _service(Path(tmp)) + original = service.save_config("demo.yaml", document=_valid_document()) + + def replace(suite: str) -> str: + try: + service.save_config( + "demo.yaml", + document=_valid_document(suite=suite), + expected_etag=original.etag, + ) + except ServiceError as exc: + return exc.code.value + return "saved" + + with ThreadPoolExecutor(max_workers=2) as executor: + outcomes = list(executor.map(replace, ("first", "second"))) + + assert sorted(outcomes) == ["STALE_ETAG", "saved"] + + +def test_list_configs_is_paginated_and_cursor_is_opaque() -> None: + with TemporaryDirectory() as tmp: + service = _service(Path(tmp)) + for name in ("a.yaml", "b.yaml", "c.yaml"): + service.save_config(name, document=_valid_document(suite=name[0])) + + first = service.list_configs(limit=2) + second = service.list_configs(limit=2, cursor=first.next_cursor) + + assert [item.config_ref for item in first.items] == ["a.yaml", "b.yaml"] + assert first.next_cursor is not None + assert "b.yaml" not in first.next_cursor + assert [item.config_ref for item in second.items] == ["c.yaml"] + assert second.next_cursor is None + + with pytest.raises(ServiceError) as invalid: + service.list_configs(cursor="not-a-cursor") + assert invalid.value.code == ServiceErrorCode.INVALID_ARGUMENT + + +def test_save_rejects_invalid_config_without_writing() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + service = _service(root) + + with pytest.raises(ServiceError) as invalid: + service.save_config( + "bad.yaml", + yaml_text=( + "pipeline:\n" + " inference:\n" + " target:\n" + " model:\n" + " name: azure/gpt-5.4\n" + " legacy: true\n" + ), + ) + + assert invalid.value.code == ServiceErrorCode.CONFIG_INVALID + validation = invalid.value.details["validation"] + assert validation["issues"][0]["path"] == "/pipeline/inference/target/legacy" + assert not (root / "evals" / "bad.yaml").exists() + + +def test_validate_reports_semantics_deprecations_and_workspace_paths() -> None: + with TemporaryDirectory() as tmp: + service = _service(Path(tmp)) + + deprecated = _valid_document() + deprecated["default_model"] = {"name": "azure/gpt-5.4"} + deprecated["pipeline"]["test_set"] = { + "tool_source": "per_seed", + "prompt": {}, + } + report = service.validate_document(deprecated) + assert report.valid is False + assert report.warnings[0].code == ConfigValidationCode.DEPRECATED_FIELD + assert report.warnings[0].path == "/pipeline/test_set/tool_source" + assert any( + issue.path == "/pipeline/test_set" + and "prompt and/or scenario" in issue.message + for issue in report.issues + ) + + escaped = _valid_document() + escaped["artifacts_root"] = "../outside" + report = service.validate_document(escaped) + assert report.valid is False + assert report.issues[0].code == ConfigValidationCode.WORKSPACE_VIOLATION + assert report.issues[0].path == "/artifacts_root" + + +def test_validation_does_not_import_callable_target() -> None: + with TemporaryDirectory() as tmp: + service = _service(Path(tmp)) + raw = _valid_document() + raw["pipeline"]["inference"]["target"]["callable"] = ( + "does_not_exist.anywhere:run" + ) + + report = service.validate_document(raw) + + assert report.valid is True + + +def test_config_refs_are_contained_and_payloads_are_bounded() -> None: + with TemporaryDirectory() as tmp: + service = _service(Path(tmp), max_config_bytes=100) + + with pytest.raises(ServiceError) as escaped: + service.save_config("../escape.yaml", document=_valid_document()) + assert escaped.value.code == ServiceErrorCode.WORKSPACE_VIOLATION + + with pytest.raises(ServiceError) as absolute: + service.get_config(str((Path(tmp) / "evals" / "config.yaml").resolve())) + assert absolute.value.code == ServiceErrorCode.INVALID_ARGUMENT + + with pytest.raises(ServiceError) as too_large: + service.validate_yaml("x" * 101) + assert too_large.value.code == ServiceErrorCode.ARTIFACT_TOO_LARGE + + +def test_list_rejects_linked_config_entries() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + service = _service(root) + service.save_config("safe.yaml", document=_valid_document()) + outside = root / "outside.yaml" + outside.write_text("pipeline:\n inference: {}\n", encoding="utf-8") + link = root / "evals" / "linked.yaml" + try: + link.symlink_to(outside) + except OSError: + pytest.skip("symbolic links are unavailable") + + with pytest.raises(ServiceError) as linked: + service.list_configs() + + assert linked.value.code == ServiceErrorCode.WORKSPACE_VIOLATION + + +def test_design_config_is_headless_and_never_writes_a_draft() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + service = _service(root) + yaml_result = ( + "suite: designed\n" + "behavior:\n" + " name: safe_help\n" + "pipeline:\n" + " inference:\n" + " target:\n" + " callable: agent:run\n" + " test_set_path: fixtures/test_set.jsonl\n" + ) + + with patch( + "assert_ai.init._design_agent.run_design_loop", + return_value=yaml_result, + ) as design: + draft = service.design_config( + ConfigDesignRequest(description="Evaluate my agent") + ) + + assert draft.document["suite"] == "designed" + assert draft.validation.valid is True + assert not (root / "eval.draft.yaml").exists() + assert design.call_args.kwargs["non_interactive"] is True + assert design.call_args.kwargs["save_draft_on_failure"] is False + + +def test_schema_service_returns_versioned_document_schema() -> None: + with TemporaryDirectory() as tmp: + schema = _service(Path(tmp)).get_schema() + + assert schema["x-assert-schema-version"] == 1 + assert json.loads(json.dumps(schema))["required"] == ["pipeline"] From cd3e568a926720560e3563f479fd0d46c7ceea03 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Thu, 6 Aug 2026 20:59:39 -0700 Subject: [PATCH 04/16] Add typed pipeline outcomes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/core/run_result.py | 45 ++++++++++ assert_ai/runner.py | 143 ++++++++++++++++++++++++++++--- tests/test_run_result.py | 157 +++++++++++++++++++++++++++++++++++ 3 files changed, 333 insertions(+), 12 deletions(-) create mode 100644 assert_ai/core/run_result.py create mode 100644 tests/test_run_result.py diff --git a/assert_ai/core/run_result.py b/assert_ai/core/run_result.py new file mode 100644 index 000000000..9533c62fc --- /dev/null +++ b/assert_ai/core/run_result.py @@ -0,0 +1,45 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Typed terminal outcomes for ASSERT pipeline execution.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path +from typing import Any + + +class RunState(StrEnum): + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +@dataclass(frozen=True, slots=True) +class RunResult: + """Transport-neutral terminal result for one pipeline invocation.""" + + state: RunState + exit_code: int + suite_id: str | None = None + run_id: str | None = None + suite_root: Path | None = None + run_root: Path | None = None + failed_stage: str | None = None + error_code: str | None = None + error_message: str | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "state": self.state.value, + "exit_code": self.exit_code, + "suite_id": self.suite_id, + "run_id": self.run_id, + "suite_root": str(self.suite_root) if self.suite_root is not None else None, + "run_root": str(self.run_root) if self.run_root is not None else None, + "failed_stage": self.failed_stage, + "error_code": self.error_code, + "error_message": self.error_message, + } diff --git a/assert_ai/runner.py b/assert_ai/runner.py index 858c99f47..e364de0e4 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -52,6 +52,7 @@ PipelineWatchdog, run_stage_coro, ) +from assert_ai.core.run_result import RunResult, RunState from assert_ai.display import label_metric from assert_ai.stages import STAGES @@ -589,6 +590,55 @@ def run_pipeline( concurrency: int | None = None, path_policy: RuntimePathPolicy | None = None, ) -> int: + """Execute configured stages and return the legacy process exit code.""" + return run_pipeline_result( + config=config, + force_stages=force_stages, + strict=strict, + overrides=overrides, + concurrency=concurrency, + path_policy=path_policy, + ).exit_code + + +def run_pipeline_result( + *, + config: str, + force_stages: list[str] | None = None, + strict: bool = False, + overrides: list[str] | None = None, + concurrency: int | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> RunResult: + """Execute configured stages and return a structured terminal outcome.""" + try: + return _run_pipeline_result( + config=config, + force_stages=force_stages, + strict=strict, + overrides=overrides, + concurrency=concurrency, + path_policy=path_policy, + ) + except Exception: # noqa: BLE001 + log.error("[runner] Unexpected pipeline setup error", exc_info=True) + return RunResult( + state=RunState.FAILED, + exit_code=1, + error_code="INTERNAL", + error_message="Unexpected pipeline setup error", + ) + + +def _run_pipeline_result( + *, + config: str, + force_stages: list[str] | None = None, + strict: bool = False, + overrides: list[str] | None = None, + concurrency: int | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> RunResult: """Execute configured stages. Programmatic callers are responsible for any desired dotenv bootstrap. @@ -625,7 +675,12 @@ def run_pipeline( ctx["strict"] = strict except (ConfigError, ValueError) as exc: log.error(f"[config error] {exc}") - return 1 + return RunResult( + state=RunState.FAILED, + exit_code=1, + error_code="CONFIG_INVALID", + error_message=_result_error_message(str(exc), path_policy=path_policy), + ) # CLI --concurrency wins over the YAML-resolved value so a single run can be # widened or narrowed without editing the config. We mutate the live @@ -647,8 +702,15 @@ def run_pipeline( invalid_forced = sorted(requested_force_stages.difference(configured_stage_names)) if invalid_forced: joined = ", ".join(invalid_forced) - log.error(f"[config error] --force-stage stage(s) not present in config: {joined}") - return 1 + message = f"--force-stage stage(s) not present in config: {joined}" + log.error(f"[config error] {message}") + return _run_result_from_context( + ctx, + state=RunState.FAILED, + exit_code=1, + error_code="CONFIG_INVALID", + error_message=message, + ) # Cascade: forcing an upstream stage logically invalidates every stage # downstream of it. Without this, `--force-stage test_set` regenerates test_set @@ -768,7 +830,6 @@ def run_pipeline( else run_root / "config.yaml" ) shutil.copy2(config_path, config_snapshot) - failed_stage: str | None = None pipeline_start = time.monotonic() stage_usage: dict[str, dict[str, Any]] = {} @@ -833,10 +894,12 @@ def _run_stages_inner( stage_usage: dict[str, dict[str, Any]], heartbeat: ManifestHeartbeat | None, watchdog: PipelineWatchdog | None, -) -> int: +) -> RunResult: """Stage execution loop. Extracted so the outer function can manage heartbeat/watchdog lifecycle in a single try/finally.""" failed_stage: str | None = None + failed_error_code: str | None = None + failed_error_message: str | None = None for stage_name, module, raw_cfg in stages_to_run: if manifest is not None and module.SCOPE == "run": @@ -915,6 +978,11 @@ def _run_stages_inner( # Print just that message; suppress the multi-screen litellm/httpx # traceback unless the user opts into verbose output. ok = False + stage_error_code = "RUN_FAILED" + stage_error_message = _result_error_message( + str(exc), + path_policy=ctx.get("path_policy"), + ) log.error(f"[{stage_name}] {exc}") if os.environ.get("ASSERT_VERBOSE_ERRORS") == "1": log.debug("Full traceback:", exc_info=True) @@ -922,6 +990,8 @@ def _run_stages_inner( log.info("(set ASSERT_VERBOSE_ERRORS=1 to see the full traceback)") except Exception: # noqa: BLE001 ok = False + stage_error_code = "RUN_FAILED" + stage_error_message = f"Unexpected error while running {stage_name}" log.error(f"[{stage_name}] Unexpected error", exc_info=True) if not ok and stage_name in artifact_plans: @@ -970,6 +1040,8 @@ def _run_stages_inner( if not ok: failed_stage = stage_name + failed_error_code = stage_error_code + failed_error_message = stage_error_message break total_elapsed = time.monotonic() - pipeline_start @@ -1015,11 +1087,58 @@ def _run_stages_inner( else: log.error(f"Pipeline failed at {failed_stage} ({total_elapsed:.1f}s)") - if manifest is None: - return 0 if failed_stage is None else 1 + if manifest is not None: + manifest.ended_at = datetime.now(timezone.utc).isoformat() + manifest.status = "completed" if failed_stage is None else "failed" + _record_run_artifacts(manifest, ctx, run_root) + _write_manifest(manifest, run_root) + + if failed_stage is None: + return _run_result_from_context( + ctx, + state=RunState.COMPLETED, + exit_code=0, + ) + return _run_result_from_context( + ctx, + state=RunState.FAILED, + exit_code=1, + failed_stage=failed_stage, + error_code=failed_error_code or "RUN_FAILED", + error_message=failed_error_message or f"Pipeline failed at {failed_stage}", + ) + + +def _run_result_from_context( + ctx: dict[str, Any], + *, + state: RunState, + exit_code: int, + failed_stage: str | None = None, + error_code: str | None = None, + error_message: str | None = None, +) -> RunResult: + suite_root = ctx.get("suite_root") + run_root = ctx.get("run_root") + return RunResult( + state=state, + exit_code=exit_code, + suite_id=ctx.get("suite_id"), + run_id=ctx.get("run_id"), + suite_root=Path(suite_root) if suite_root is not None else None, + run_root=Path(run_root) if run_root is not None else None, + failed_stage=failed_stage, + error_code=error_code, + error_message=error_message, + ) - manifest.ended_at = datetime.now(timezone.utc).isoformat() - manifest.status = "completed" if failed_stage is None else "failed" - _record_run_artifacts(manifest, ctx, run_root) - _write_manifest(manifest, run_root) - return 0 if failed_stage is None else 1 + +def _result_error_message( + message: str, + *, + path_policy: RuntimePathPolicy | None, +) -> str: + if path_policy is None: + return message + workspace = str(path_policy.workspace_root) + return message.replace(workspace, ".") diff --git a/tests/test_run_result.py b/tests/test_run_result.py new file mode 100644 index 000000000..9de873344 --- /dev/null +++ b/tests/test_run_result.py @@ -0,0 +1,157 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +from assert_ai.core.model_client import LLMInputError +from assert_ai.core.run_result import RunState +from assert_ai.core.workspace import WorkspaceService +from assert_ai.runner import run_pipeline, run_pipeline_result + + +def test_invalid_config_returns_typed_failure_and_legacy_exit_code() -> None: + with TemporaryDirectory() as tmp: + config_path = Path(tmp) / "invalid.yaml" + config_path.write_text("pipeline: {}\n", encoding="utf-8") + + result = run_pipeline_result(config=str(config_path)) + + assert result.state == RunState.FAILED + assert result.exit_code == 1 + assert result.error_code == "CONFIG_INVALID" + assert result.failed_stage is None + assert result.suite_id is None + assert run_pipeline(config=str(config_path)) == 1 + + +def test_suite_only_success_returns_managed_identity_and_serializes() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + config_path = root / "config.yaml" + results = root / "results" + config_path.write_text( + "\n".join( + [ + "suite: suite-only", + f"results_dir: {results}", + "pipeline:", + " inference:", + " enabled: false", + ] + ) + + "\n", + encoding="utf-8", + ) + + result = run_pipeline_result(config=str(config_path)) + + assert result.state == RunState.COMPLETED + assert result.exit_code == 0 + assert result.suite_id == "suite-only" + assert result.run_id is None + assert result.suite_root == (results / "suite-only").resolve() + assert result.run_root is None + assert result.to_dict()["suite_root"] == str((results / "suite-only").resolve()) + + +def test_stage_failure_returns_failed_stage_without_raw_exception_text() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + config_path = root / "config.yaml" + config_path.write_text( + "\n".join( + [ + "suite: suite-a", + "run: run-a", + f"results_dir: {root / 'results'}", + "pipeline:", + " inference:", + " target:", + " callable: agent:run", + " test_set_path: fixture.jsonl", + ] + ) + + "\n", + encoding="utf-8", + ) + + async def fail_stage(*_: object, **__: object) -> dict: + raise RuntimeError("sensitive target detail") + + with patch("assert_ai.stages.inference.run", new=fail_stage): + result = run_pipeline_result(config=str(config_path)) + + assert result.state == RunState.FAILED + assert result.exit_code == 1 + assert result.failed_stage == "inference" + assert result.error_code == "RUN_FAILED" + assert result.error_message == "Unexpected error while running inference" + assert result.run_root == ( + root / "results" / "suite-a" / "run-a" + ).resolve() + manifest = json.loads( + (result.run_root / "manifest.json").read_text(encoding="utf-8") + ) + assert manifest["status"] == "failed" + assert manifest["stages"]["inference"] == "failed" + + +def test_classified_stage_error_hides_strict_workspace_root() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + workspace = WorkspaceService.create(root) + workspace.configs_root.mkdir() + config_path = workspace.configs_root / "config.yaml" + config_path.write_text( + "\n".join( + [ + "suite: suite-a", + "run: run-a", + "pipeline:", + " inference:", + " target:", + " callable: agent:run", + " test_set_path: fixture.jsonl", + ] + ) + + "\n", + encoding="utf-8", + ) + + async def fail_stage(*_: object, **__: object) -> dict: + raise LLMInputError(f"invalid request from {workspace.root}") + + with patch("assert_ai.stages.inference.run", new=fail_stage): + result = run_pipeline_result( + config="config.yaml", + path_policy=workspace.path_policy, + ) + + assert result.state == RunState.FAILED + assert result.error_message == "invalid request from ." + + +def test_unexpected_setup_failure_is_returned_not_raised() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + config_path = root / "config.yaml" + config_path.write_text( + "suite: suite-a\npipeline:\n inference:\n enabled: false\n", + encoding="utf-8", + ) + + with patch( + "assert_ai.runner._write_suite_metadata", + side_effect=OSError("disk failure"), + ): + result = run_pipeline_result(config=str(config_path)) + + assert result.state == RunState.FAILED + assert result.exit_code == 1 + assert result.error_code == "INTERNAL" + assert result.error_message == "Unexpected pipeline setup error" From 8ebdd5646628ea5c7dafd56a7240cec29a71b965 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Fri, 7 Aug 2026 10:04:21 -0700 Subject: [PATCH 05/16] Add indexed result services Persist lightweight run and suite summaries, build stale-aware JSONL indexes, and expose paginated result queries and comparisons through a shared repository used by the CLI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/cli.py | 580 ++------ assert_ai/core/jsonl_index.py | 396 +++++ assert_ai/runner.py | 99 +- assert_ai/services/result_metadata.py | 703 +++++++++ assert_ai/services/results.py | 1963 +++++++++++++++++++++++++ assert_ai/viewer_read_model.py | 32 +- tests/result_catalog_fixture.py | 218 +++ tests/test_jsonl_index.py | 132 ++ tests/test_result_metadata.py | 298 ++++ tests/test_result_service.py | 516 +++++++ tests/test_result_service_scale.py | 70 + tests/test_run_metadata.py | 31 + tests/test_runner_artifact_cache.py | 28 + 13 files changed, 4612 insertions(+), 454 deletions(-) create mode 100644 assert_ai/core/jsonl_index.py create mode 100644 assert_ai/services/result_metadata.py create mode 100644 assert_ai/services/results.py create mode 100644 tests/result_catalog_fixture.py create mode 100644 tests/test_jsonl_index.py create mode 100644 tests/test_result_metadata.py create mode 100644 tests/test_result_service.py create mode 100644 tests/test_result_service_scale.py diff --git a/assert_ai/cli.py b/assert_ai/cli.py index b70c037c8..d727e26a5 100644 --- a/assert_ai/cli.py +++ b/assert_ai/cli.py @@ -19,15 +19,10 @@ from rich.table import Table from assert_ai.core.config_model import DEFAULT_INFERENCE_CONCURRENCY -from assert_ai.core.io import load_json, load_jsonl, get_permissible_flag, row_behavior -from assert_ai.core.judge import get_verdict_dimension, infer_judge_status, is_valid_event_flag from assert_ai.display import label_metric, label_run_status, label_stage, label_stage_status, label_status from assert_ai.logging_config import configure_logging -from assert_ai.results import ( - compute_dimension_summary, - compute_policy_violation_by_permissibility, - detect_dimensions, -) +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.results import ResultRepository, RunReference from assert_ai.stages import STAGE_NAMES ROOT = Path(__file__).resolve().parent.parent @@ -405,28 +400,6 @@ def _complete_metric(_: click.Context, __: click.Parameter, incomplete: str) -> return [CompletionItem(name) for name in items if not incomplete or name.startswith(incomplete)] -def _current_stage_status(manifest: dict[str, Any] | None) -> tuple[str, str]: - if isinstance(manifest, dict): - manifest_status = manifest.get("status") - if isinstance(manifest_status, str) and manifest_status: - stages = manifest.get("stages") - if isinstance(stages, dict): - for stage_name, stage_status in stages.items(): - if stage_status == "running": - return manifest_status, str(stage_name) - return manifest_status, "-" - - return "unknown", "-" - - -def _detect_dimensions(rows: Iterable[dict[str, Any]]) -> list[str]: - return detect_dimensions(rows) - - -def _compute_dimension_summary(rows: Iterable[dict[str, Any]], metric: str) -> dict[str, Any]: - return compute_dimension_summary(rows, metric) - - def _dimension_rate(metrics: dict[str, Any], metric: str) -> float | None: dimensions = metrics.get("dimensions") if not isinstance(dimensions, dict): @@ -438,295 +411,103 @@ def _dimension_rate(metrics: dict[str, Any], metric: str) -> float | None: return float(rate) if isinstance(rate, (int, float)) else None -def _reject_ordinal_compare(run_summaries: Iterable[dict[str, Any]], metric: str) -> None: - for run_summary in run_summaries: - for key in ("prompt_metrics", "scenario_metrics"): - metrics = run_summary.get(key) - dimensions = metrics.get("dimensions") if isinstance(metrics, dict) else None - summary = dimensions.get(metric) if isinstance(dimensions, dict) else None - if isinstance(summary, dict) and summary.get("kind") == "ordinal": - _error( - f"Metric '{metric}' uses ordered grades and has no violation-rate delta. " - "Use 'assert-ai results status' or the viewer for its grade distribution." - ) - - -def _compute_prompt_metrics( - rows: list[dict[str, Any]], - behavior_categories: Iterable[dict[str, Any]] = (), -) -> dict[str, Any] | None: - if not rows: - return None - - scored_rows = [row for row in rows if infer_judge_status(row) == "ok"] - scored_total = len(scored_rows) - judge_failures = len(rows) - scored_total - - dimensions: dict[str, dict[str, Any]] = {} - for dim in _detect_dimensions(scored_rows): - dimensions[dim] = _compute_dimension_summary(scored_rows, dim) - - target = next( - ( - row.get("target") - for row in rows - if isinstance(row.get("target"), str) and row.get("target") - ), - ) or "-" - judge_model = next( - ( - row.get("judge_model") - for row in rows - if isinstance(row.get("judge_model"), str) and row.get("judge_model") - ), - ) or "-" - permissible_rows = [row for row in scored_rows if get_permissible_flag(row, default=False)] - permissibility_split = compute_policy_violation_by_permissibility( - scored_rows, - behavior_categories, - ) - - metrics: dict[str, Any] = { - "total": len(rows), - "scored_total": scored_total, - "judge_failures": judge_failures, - "judge_failure_rate": judge_failures / len(rows) if rows else 0.0, - "policy_violation_rate": _dimension_rate({"dimensions": dimensions}, "policy_violation"), - "overrefusal_rate": _dimension_rate({"dimensions": dimensions}, "overrefusal"), - "permissible_overrefusal_rate": _compute_dimension_summary(permissible_rows, "overrefusal")["rate"], - "dimensions": dimensions, - "target": target, - "judge_model": judge_model, - } - - if permissibility_split["permissible"] is not None: - permissible = permissibility_split["permissible"] - not_permissible = permissibility_split["not_permissible"] - assert not_permissible is not None - metrics.update( - { - "permissible_policy_violation_rate": permissible["rate"], - "not_permissible_policy_violation_rate": not_permissible["rate"], - "policy_violation_on_permissible": permissible, - "policy_violation_on_not_permissible": not_permissible, - } +def _load_run_summary(run_dir: Path) -> dict[str, Any] | None: + repository = ResultRepository(run_dir.parent.parent) + try: + detail = repository.load_run_detail( + run_dir.parent.name, + run_dir.name, ) - - return metrics - - -def _compute_scenario_metrics( - rows: list[dict[str, Any]], - behavior_categories: Iterable[dict[str, Any]] = (), -) -> dict[str, Any] | None: - if not rows: - return None - - scored_rows = [row for row in rows if infer_judge_status(row) == "ok"] - scored_total = len(scored_rows) - judge_failures = len(rows) - scored_total - - dimensions: dict[str, dict[str, Any]] = {} - for dim in _detect_dimensions(scored_rows): - dimensions[dim] = _compute_dimension_summary(scored_rows, dim) - - target = next( - ( - row.get("target") - for row in rows - if isinstance(row.get("target"), str) and row.get("target") - ), - None, - ) or "-" - tester_model = next( - ( - row.get("tester_model") - for row in rows - if isinstance(row.get("tester_model"), str) and row.get("tester_model") - ), - None, - ) or "-" - judge_model = next( - ( - row.get("judge_model") - for row in rows - if isinstance(row.get("judge_model"), str) and row.get("judge_model") - ), - ) or "-" - - permissibility_split = compute_policy_violation_by_permissibility( - scored_rows, - behavior_categories, - ) - - metrics: dict[str, Any] = { - "total": len(rows), - "scored_total": scored_total, - "judge_failures": judge_failures, - "judge_failure_rate": judge_failures / len(rows) if rows else 0.0, - "policy_violation_rate": _dimension_rate({"dimensions": dimensions}, "policy_violation"), - "overrefusal_rate": _dimension_rate({"dimensions": dimensions}, "overrefusal"), - "dimensions": dimensions, - "target": target, - "tester_model": tester_model, - "judge_model": judge_model, + except ServiceError as exc: + if exc.code == ServiceErrorCode.NOT_FOUND: + return None + _error(f"{exc.code}: {exc}") + quality = detail.get("quality") or {} + manifest = { + "status": detail.get("state"), + "started_at": detail.get("started_at"), + "ended_at": detail.get("ended_at"), + "stages": detail.get("stages") or {}, + "stage_timings": detail.get("stage_timings") or {}, } - - if permissibility_split["permissible"] is not None: - permissible = permissibility_split["permissible"] - not_permissible = permissibility_split["not_permissible"] - assert not_permissible is not None - metrics.update( - { - "permissible_policy_violation_rate": permissible["rate"], - "not_permissible_policy_violation_rate": not_permissible["rate"], - "policy_violation_on_permissible": permissible, - "policy_violation_on_not_permissible": not_permissible, - } - ) - - return metrics - - -def _load_run_summary(run_dir: Path) -> dict[str, Any] | None: - manifest = load_json(run_dir / "manifest.json") - score_rows = load_jsonl(run_dir / "scores.jsonl") - taxonomy = load_json(run_dir.parent / "taxonomy.json") - behavior_categories = (taxonomy or {}).get("behavior_categories") - if not isinstance(behavior_categories, list): - behavior_categories = [] - prompt_rows = [row for row in score_rows if not row.get("tester_model")] - scenario_rows = [row for row in score_rows if row.get("tester_model")] - - stages = (manifest or {}).get("stages", {}) - has_scores = isinstance(stages, dict) and stages.get("judge") is not None - has_data = bool(prompt_rows or scenario_rows) - if not has_data and not has_scores: - return None - if not has_data and (manifest or {}).get("status") == "failed": - return None - - status, current_stage = _current_stage_status(manifest) return { "run_id": run_dir.name, "path": str(run_dir), "manifest": manifest, - "status": status, - "current_stage": current_stage, - "started_at": (manifest or {}).get("started_at"), - "ended_at": (manifest or {}).get("ended_at"), - "prompt_metrics": _compute_prompt_metrics(prompt_rows, behavior_categories), - "scenario_metrics": _compute_scenario_metrics(scenario_rows, behavior_categories), - "prompt_rows": prompt_rows, - "scenario_rows": scenario_rows, + "status": detail.get("state") or "unknown", + "current_stage": detail.get("current_stage") or "-", + "started_at": detail.get("started_at"), + "ended_at": detail.get("ended_at"), + "prompt_metrics": quality.get("prompt"), + "scenario_metrics": quality.get("scenario"), } -def _count_test_case_types(path: Path) -> tuple[int, int]: - rows = load_jsonl(path) - prompt_count = 0 - scenario_count = 0 - for row in rows: - row_type = row.get("type") - if row_type == "prompt": - prompt_count += 1 - elif row_type == "scenario": - scenario_count += 1 - return prompt_count, scenario_count - - def _load_suite_summary(suite_dir: Path) -> dict[str, Any] | None: - suite_meta = load_json(suite_dir / "suite.json") - taxonomy = load_json(suite_dir / "taxonomy.json") - if suite_meta is None and taxonomy is None: - return None - - run_summaries = [] - for child in sorted(suite_dir.iterdir()) if suite_dir.exists() else []: - if not child.is_dir(): - continue - run_summary = _load_run_summary(child) - if run_summary is not None: - run_summaries.append(run_summary) - - has_results = any( - (run_summary.get("prompt_metrics") is not None) or (run_summary.get("scenario_metrics") is not None) - for run_summary in run_summaries - ) - prompt_test_case_count, scenario_test_case_count = _count_test_case_types(suite_dir / "test_set.jsonl") - - created_at = (suite_meta or {}).get("created_at") - - behavior_name = suite_dir.name - behavior_block = (taxonomy or {}).get("behavior") - if isinstance(behavior_block, dict) and isinstance(behavior_block.get("name"), str) and behavior_block.get("name"): - behavior_name = behavior_block["name"] - - if has_results: - status = "has_results" - elif prompt_test_case_count or scenario_test_case_count: - status = "test_set_ready" - else: - status = "systematized" + repository = ResultRepository(suite_dir.parent) + try: + detail = repository.get_suite(suite_dir.name) + run_summaries: list[dict[str, Any]] = [] + cursor: str | None = None + while True: + page = repository.list_run_catalog_entries( + suite_dir.name, + cursor=cursor, + ) + run_summaries.extend(page.items) + cursor = page.next_cursor + if cursor is None: + break + except ServiceError as exc: + if exc.code == ServiceErrorCode.NOT_FOUND: + return None + _error(f"{exc.code}: {exc}") + behavior = detail.get("behavior") or {} + counts = detail.get("test_case_counts") or {} return { "suite_id": suite_dir.name, "path": str(suite_dir), - "behavior_name": behavior_name, - "behavior_category_count": len((taxonomy or {}).get("behavior_categories") or []), - "prompt_test_case_count": prompt_test_case_count, - "scenario_test_case_count": scenario_test_case_count, + "behavior_name": behavior.get("name") or suite_dir.name, + "behavior_category_count": int( + detail.get("behavior_category_count") or 0 + ), + "prompt_test_case_count": int(counts.get("prompt") or 0), + "scenario_test_case_count": int(counts.get("scenario") or 0), "run_count": len(run_summaries), "runs": run_summaries, - "status": status, - "created_at": created_at, - "has_systematization": (suite_dir / "systematization.json").exists(), + "status": detail.get("status") or "unknown", + "created_at": detail.get("created_at"), + "has_systematization": "taxonomy" in (detail.get("sources") or {}), } def _load_all_suites(results_dir: Path) -> list[dict[str, Any]]: - if not results_dir.exists(): - return [] - suites = [] - for child in sorted(results_dir.iterdir()): - if not child.is_dir(): - continue - suite_summary = _load_suite_summary(child) - if suite_summary is not None: - suites.append(suite_summary) - suites.sort(key=lambda item: str(item.get("created_at") or ""), reverse=True) + repository = ResultRepository(results_dir) + suites: list[dict[str, Any]] = [] + cursor: str | None = None + while True: + page = repository.list_suite_catalog_entries(cursor=cursor) + suites.extend( + { + **entry, + "path": str(results_dir / str(entry["suite_id"])), + "runs": [], + } + for entry in page.items + ) + cursor = page.next_cursor + if cursor is None: + break return suites -def _behavior_category_metric_map(rows: Iterable[dict[str, Any]], metric: str) -> dict[str, dict[str, Any]]: - grouped: dict[str, dict[str, Any]] = {} - for row in rows: - if infer_judge_status(row) != "ok": - continue - value = get_verdict_dimension(row.get("verdict"), metric) - if not is_valid_event_flag(value): - continue - behavior_category = row_behavior(row) - bucket = grouped.setdefault( - behavior_category, - { - "true_count": 0, - "count": 0, - "permissible": get_permissible_flag(row), - }, - ) - bucket["true_count"] += int(value) - bucket["count"] += 1 - result = {} - for behavior_category, bucket in grouped.items(): - if bucket["count"] <= 0: - continue - result[behavior_category] = { - "rate": bucket["true_count"] / bucket["count"], - "count": bucket["count"], - "permissible": bucket["permissible"], - } - return result +def _comparison_metric_text(metrics: dict[str, Any], metric: str) -> str: + dimensions = metrics.get("dimensions") + summary = dimensions.get(metric) if isinstance(dimensions, dict) else None + if not isinstance(summary, dict): + return "-" + return _fmt_dimension_summary(summary)[0] @click.group( @@ -1231,76 +1012,43 @@ def _run_within_suite_compare( as_json: bool, no_color: bool, ) -> None: - """Original within-suite comparison logic.""" - results_root = _resolve_results_dir(results_dir) - suite_dir = results_root / suite - if not suite_dir.exists(): - _error(f"Suite not found: {suite}") - - run_summaries: list[dict[str, Any]] = [] - for run_id in runs: - run_summary = _load_run_summary(suite_dir / run_id) - if run_summary is None: - _error(f"Run not found or unreadable: {suite}/{run_id}") - run_summaries.append(run_summary) - - available_metrics: set[str] = set() - for run_summary in run_summaries: - available_metrics.update(_detect_dimensions(run_summary.get("prompt_rows") or [])) - if metric not in available_metrics: - _error(f"Metric '{metric}' was not found in the compared prompt judgments. Available: {sorted(available_metrics)}") - _reject_ordinal_compare(run_summaries, metric) - - behavior_category_deltas: list[dict[str, Any]] = [] - if all(run_summary.get("prompt_rows") for run_summary in run_summaries): - first_map = _behavior_category_metric_map(run_summaries[0]["prompt_rows"], metric) - last_map = _behavior_category_metric_map(run_summaries[-1]["prompt_rows"], metric) - for behavior_category in sorted(set(first_map) | set(last_map)): - first = first_map.get(behavior_category) - last = last_map.get(behavior_category) - if first is None or last is None: - continue - behavior_category_deltas.append( - { - "behavior_category": behavior_category, - "permissible": first.get("permissible"), - "first_rate": first["rate"], - "last_rate": last["rate"], - "delta": last["rate"] - first["rate"], - "first_count": first["count"], - "last_count": last["count"], - } - ) - behavior_category_deltas.sort(key=lambda row: abs(row["delta"]), reverse=True) - if limit >= 0: - behavior_category_deltas = behavior_category_deltas[:limit] - - run_rows = [] - for run_summary in run_summaries: - prompt_metrics = run_summary.get("prompt_metrics") or {} - scenario_metrics = run_summary.get("scenario_metrics") or {} - run_rows.append( - { - "run_id": run_summary["run_id"], - "status": run_summary["status"], - "started_at": run_summary.get("started_at"), - "prompt": prompt_metrics, - "scenario": scenario_metrics, - } + repository = ResultRepository(results_root) + try: + comparison = repository.compare_runs( + [RunReference(suite, run_id) for run_id in runs], + metric=metric, + behavior_limit=limit, ) + except ServiceError as exc: + _error(f"{exc.code}: {exc}") + + run_rows = [ + { + "run_id": row["run_id"], + "status": row["state"], + "started_at": row.get("started_at"), + "prompt": (row.get("quality") or {}).get("prompt") or {}, + "scenario": (row.get("quality") or {}).get("scenario") or {}, + "structural": row.get("structural") or {}, + } + for row in comparison["runs"] + ] payload = { "suite_id": suite, "metric": metric, "runs": run_rows, - "behavior_category_deltas": behavior_category_deltas, + "dimension_deltas": comparison["dimension_deltas"], + "behavior_category_deltas": comparison["behavior_category_deltas"], + "warnings": comparison["warnings"], } if as_json: _echo_json(payload) return console = _console(no_color=no_color) + metric_kind = comparison["dimension_deltas"][metric]["kind"] table = Table( title=f"Run Comparison ({suite}, {_metric_label(metric)})", box=None, @@ -1312,8 +1060,8 @@ def _run_within_suite_compare( table.add_column("Status", style="white", no_wrap=True) table.add_column("Started", style="dim", no_wrap=True) table.add_column("Target", style="white") - table.add_column(f"Prompt {_metric_label(metric).lower()} rate", style="white", no_wrap=True) - table.add_column(f"Scenario {_metric_label(metric).lower()} rate", style="white", no_wrap=True) + table.add_column(f"Prompt {_metric_label(metric).lower()}", style="white", no_wrap=True) + table.add_column(f"Scenario {_metric_label(metric).lower()}", style="white", no_wrap=True) table.add_column(label_metric("judge_failure_rate"), style="white", no_wrap=True) for row in payload["runs"]: prompt_metrics = row["prompt"] or {} @@ -1329,12 +1077,13 @@ def _run_within_suite_compare( label_run_status(row["status"] if isinstance(row.get("status"), str) else None), _format_timestamp(row.get("started_at")), str(target_model), - _fmt_percent(_dimension_rate(prompt_metrics, metric)), - _fmt_percent(_dimension_rate(scenario_metrics, metric)), + _comparison_metric_text(prompt_metrics, metric), + _comparison_metric_text(scenario_metrics, metric), _fmt_percent(fail_rate), ) console.print(table) + behavior_category_deltas = payload["behavior_category_deltas"] if behavior_category_deltas: delta_table = Table( title=f"Top behavior category deltas ({_metric_label(metric).lower()}: {runs[0]} -> {runs[-1]})", @@ -1351,12 +1100,22 @@ def _run_within_suite_compare( for row in behavior_category_deltas: delta_table.add_row( row["behavior_category"], - str(bool(row["permissible"])), + ( + str(bool(row["permissible"])) + if row.get("permissible") is not None + else "-" + ), _fmt_percent(row["first_rate"]), _fmt_percent(row["last_rate"]), _fmt_percent(row["delta"]), ) console.print(delta_table) + elif metric_kind == "ordinal": + console.print( + "Ordinal comparisons report grade distributions in --json output." + ) + for warning in comparison["warnings"]: + console.print(f"Warning: {warning}", style="yellow") @results.command("compare-suites", short_help="Compare runs across different suites (e.g., approach A vs B vs C)") @@ -1399,8 +1158,7 @@ def results_compare_suites( _error("Provide at least two SUITE/RUN arguments to compare.") results_root = _resolve_results_dir(results_dir) - run_summaries: list[dict[str, Any]] = [] - labels: list[str] = [] + refs: list[RunReference] = [] for suite_run in suite_runs: parts = suite_run.strip("/").split("/") @@ -1411,67 +1169,20 @@ def results_compare_suites( else: _error(f"Invalid format: '{suite_run}'. Use SUITE/RUN (e.g., my-suite/run-1).") return # unreachable, for type checker + refs.append(RunReference(suite_id, run_id)) - run_dir = results_root / suite_id / run_id - if not run_dir.exists(): - _error(f"Not found: {suite_id}/{run_id}") - run_summary = _load_run_summary(run_dir) - if run_summary is None: - _error(f"No scores in {suite_id}/{run_id}") - run_summary["suite_id"] = suite_id - run_summaries.append(run_summary) - labels.append(f"{suite_id}/{run_id}") - - _reject_ordinal_compare(run_summaries, metric) - - # Count structural visibility from inference rows - structural: list[dict[str, Any]] = [] - for i, suite_run in enumerate(suite_runs): - parts = suite_run.strip("/").split("/") - suite_id = parts[0] - run_id = parts[1] if len(parts) > 1 else "run-1" - inference_set_path = results_root / suite_id / run_id / "inference_set.jsonl" - inference_rows = load_jsonl(inference_set_path) - total_events = sum(len(r.get("events", [])) for r in inference_rows) - tool_events = sum( - 1 for r in inference_rows - for e in r.get("events", []) - if e.get("edit", {}).get("type") == "tool_call" - ) - msg_events = sum( - 1 for r in inference_rows - for e in r.get("events", []) - if e.get("edit", {}).get("type") == "add_message" - ) - with_tools = sum( - 1 for r in inference_rows - if any(e.get("edit", {}).get("type") == "tool_call" for e in r.get("events", [])) - ) - structural.append({ - "label": labels[i], - "inference_rows": len(inference_rows), - "total_events": total_events, - "msg_events": msg_events, - "tool_events": tool_events, - "with_tools": with_tools, - }) + repository = ResultRepository(results_root) + try: + comparison = repository.compare_runs(refs, metric=metric) + except ServiceError as exc: + _error(f"{exc.code}: {exc}") if as_json: - run_rows = [] - for i, run_summary in enumerate(run_summaries): - prompt_metrics = run_summary.get("prompt_metrics") or {} - run_rows.append({ - "label": labels[i], - "suite_id": run_summary["suite_id"], - "run_id": run_summary["run_id"], - "status": run_summary["status"], - "prompt": prompt_metrics, - "structural": structural[i], - }) - _echo_json({"metric": metric, "runs": run_rows}) + _echo_json(comparison) return console = _console(no_color=no_color) + metric_kind = comparison["dimension_deltas"][metric]["kind"] # Table 1: Judge quality table = Table( @@ -1485,23 +1196,26 @@ def results_compare_suites( table.add_column("Total", style="white", no_wrap=True) table.add_column("Scored", style="white", no_wrap=True) table.add_column(label_metric("judge_failure_rate"), style="white", no_wrap=True) - table.add_column(f"{_metric_label(metric)} rate", style="white", no_wrap=True) - table.add_column("Pass rate", style="white", no_wrap=True) - for i, run_summary in enumerate(run_summaries): - pm = run_summary.get("prompt_metrics") or {} + table.add_column(_metric_label(metric), style="white", no_wrap=True) + if metric_kind != "ordinal": + table.add_column("Pass rate", style="white", no_wrap=True) + for run_summary in comparison["runs"]: + pm = (run_summary.get("quality") or {}).get("prompt") or {} total = pm.get("total", 0) ok = pm.get("scored_total", 0) fail_rate = pm.get("judge_failure_rate") dim_rate = _dimension_rate(pm, metric) pass_rate = (1.0 - dim_rate) if dim_rate is not None else None - table.add_row( - labels[i], + row = [ + run_summary["label"], str(total), str(ok), _fmt_percent(fail_rate), - _fmt_percent(dim_rate), - _fmt_percent(pass_rate), - ) + _comparison_metric_text(pm, metric), + ] + if metric_kind != "ordinal": + row.append(_fmt_percent(pass_rate)) + table.add_row(*row) console.print(table) console.print() @@ -1519,16 +1233,22 @@ def results_compare_suites( struct_table.add_column("Messages", style="white", no_wrap=True) struct_table.add_column("Tool events", style="white", no_wrap=True) struct_table.add_column("With tools", style="white", no_wrap=True) - for s in structural: + for run_summary in comparison["runs"]: + structural = run_summary.get("structural") or {} struct_table.add_row( - s["label"], - str(s["inference_rows"]), - str(s["total_events"]), - str(s["msg_events"]), - str(s["tool_events"]), - f"{s['with_tools']}/{s['inference_rows']}", + run_summary["label"], + str(structural.get("inference_rows", 0)), + str(structural.get("total_events", 0)), + str(structural.get("message_events", 0)), + str(structural.get("tool_events", 0)), + ( + f"{structural.get('rows_with_tools', 0)}/" + f"{structural.get('inference_rows', 0)}" + ), ) console.print(struct_table) + for warning in comparison["warnings"]: + console.print(f"Warning: {warning}", style="yellow") @cli.group(cls=SuggestingGroup, short_help="Generate and validate ACS policies from ASSERT findings") diff --git a/assert_ai/core/jsonl_index.py b/assert_ai/core/jsonl_index.py new file mode 100644 index 000000000..813e161fc --- /dev/null +++ b/assert_ai/core/jsonl_index.py @@ -0,0 +1,396 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Versioned byte-offset indexes for canonical ASSERT JSONL artifacts.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path +from typing import Any + +from assert_ai.core.io import load_json, write_json + +JSONL_INDEX_SCHEMA_VERSION = 1 +DEFAULT_MAX_INDEXED_ROW_BYTES = 16 * 1024 * 1024 + + +class JsonlIndexErrorCode(StrEnum): + NOT_FOUND = "not_found" + INVALID_JSON = "invalid_json" + INVALID_ROW = "invalid_row" + INVALID_KEY = "invalid_key" + DUPLICATE_KEY = "duplicate_key" + INVALID_INDEX = "invalid_index" + STALE_INDEX = "stale_index" + SOURCE_CHANGED = "source_changed" + ROW_TOO_LARGE = "row_too_large" + + +class JsonlIndexError(ValueError): + """Typed JSONL scan, index, and lookup failure.""" + + def __init__( + self, + code: JsonlIndexErrorCode, + message: str, + *, + path: Path, + line_number: int | None = None, + key: str | None = None, + ) -> None: + super().__init__(message) + self.code = code + self.path = path + self.line_number = line_number + self.key = key + + +@dataclass(frozen=True, slots=True) +class JsonlRecord: + offset: int + length: int + line_number: int + row: dict[str, Any] + + +@dataclass(frozen=True, slots=True) +class JsonlScan: + path: Path + records: tuple[JsonlRecord, ...] + size_bytes: int + mtime_ns: int + sha256: str + + +def jsonl_index_path(source_path: Path) -> Path: + """Return the canonical sibling index path for one JSONL source.""" + return source_path.with_name(f"{source_path.stem}.index.json") + + +def jsonl_row_key(kind: str, test_case_id: str) -> str: + return f"{kind}:{test_case_id}" + + +def scan_jsonl( + path: Path, + *, + allow_trailing_partial: bool = False, + max_row_bytes: int | None = None, +) -> JsonlScan: + """Scan one JSONL file in binary mode and preserve exact byte ranges.""" + try: + before = path.stat() + except FileNotFoundError as exc: + raise JsonlIndexError( + JsonlIndexErrorCode.NOT_FOUND, + f"Missing JSONL artifact: {path}", + path=path, + ) from exc + + records: list[JsonlRecord] = [] + digest = hashlib.sha256() + offset = 0 + with path.open("rb") as handle: + for line_number, raw_line in enumerate(handle, 1): + length = len(raw_line) + digest.update(raw_line) + if max_row_bytes is not None and length > max_row_bytes: + raise JsonlIndexError( + JsonlIndexErrorCode.ROW_TOO_LARGE, + f"JSONL row exceeds {max_row_bytes} bytes in {path} on line {line_number}", + path=path, + line_number=line_number, + ) + stripped = raw_line.strip() + if not stripped: + offset += length + continue + try: + row = json.loads(stripped) + except json.JSONDecodeError as exc: + is_trailing_partial = ( + allow_trailing_partial + and not raw_line.endswith((b"\n", b"\r")) + and offset + length == before.st_size + ) + if is_trailing_partial: + offset += length + continue + raise JsonlIndexError( + JsonlIndexErrorCode.INVALID_JSON, + f"Invalid JSONL in {path} on line {line_number}: {exc}", + path=path, + line_number=line_number, + ) from exc + if not isinstance(row, dict): + raise JsonlIndexError( + JsonlIndexErrorCode.INVALID_ROW, + f"Expected JSON object in {path} on line {line_number}", + path=path, + line_number=line_number, + ) + records.append( + JsonlRecord( + offset=offset, + length=length, + line_number=line_number, + row=row, + ) + ) + offset += length + + after = path.stat() + if ( + before.st_size != after.st_size + or before.st_mtime_ns != after.st_mtime_ns + ): + raise JsonlIndexError( + JsonlIndexErrorCode.SOURCE_CHANGED, + f"JSONL source changed while it was being indexed: {path}", + path=path, + ) + return JsonlScan( + path=path, + records=tuple(records), + size_bytes=after.st_size, + mtime_ns=after.st_mtime_ns, + sha256=digest.hexdigest(), + ) + + +def build_jsonl_index( + source_path: Path, + *, + index_path: Path | None = None, + scan: JsonlScan | None = None, +) -> dict[str, Any]: + """Build and atomically persist a unique type/test-case lookup index.""" + source_path = source_path.resolve() + current_scan = scan or scan_jsonl(source_path) + if current_scan.path.resolve() != source_path: + raise ValueError("scan path does not match source_path") + + items: dict[str, dict[str, Any]] = {} + order: list[str] = [] + for record in current_scan.records: + kind, test_case_id = _row_identity(record.row, path=source_path) + key = jsonl_row_key(kind, test_case_id) + if key in items: + raise JsonlIndexError( + JsonlIndexErrorCode.DUPLICATE_KEY, + f"Duplicate {key} row in {source_path}", + path=source_path, + line_number=record.line_number, + key=key, + ) + items[key] = { + "type": kind, + "test_case_id": test_case_id, + "offset": record.offset, + "length": record.length, + "line_number": record.line_number, + } + order.append(key) + + payload = { + "schema_version": JSONL_INDEX_SCHEMA_VERSION, + "source": { + "name": source_path.name, + "size_bytes": current_scan.size_bytes, + "mtime_ns": current_scan.mtime_ns, + "sha256": current_scan.sha256, + }, + "key_fields": ["type", "test_case_id"], + "row_count": len(order), + "order": order, + "items": items, + } + write_json(index_path or jsonl_index_path(source_path), payload) + return payload + + +def load_jsonl_index( + source_path: Path, + *, + index_path: Path | None = None, + verify_hash: bool = False, +) -> dict[str, Any]: + """Load an index and reject incompatible or stale source metadata.""" + source_path = source_path.resolve() + resolved_index_path = index_path or jsonl_index_path(source_path) + try: + payload = load_json(resolved_index_path) + except (OSError, ValueError, json.JSONDecodeError) as exc: + raise JsonlIndexError( + JsonlIndexErrorCode.INVALID_INDEX, + f"Invalid or unreadable JSONL index: {resolved_index_path}", + path=resolved_index_path, + ) from exc + if not _valid_index_payload(payload): + raise JsonlIndexError( + JsonlIndexErrorCode.INVALID_INDEX, + f"Invalid or missing JSONL index: {resolved_index_path}", + path=resolved_index_path, + ) + assert payload is not None + source = payload["source"] + try: + stat_result = source_path.stat() + except FileNotFoundError as exc: + raise JsonlIndexError( + JsonlIndexErrorCode.NOT_FOUND, + f"Missing JSONL artifact: {source_path}", + path=source_path, + ) from exc + is_current = ( + source.get("name") == source_path.name + and source.get("size_bytes") == stat_result.st_size + and source.get("mtime_ns") == stat_result.st_mtime_ns + ) + if is_current and verify_hash: + is_current = source.get("sha256") == _file_sha256(source_path) + if not is_current: + raise JsonlIndexError( + JsonlIndexErrorCode.STALE_INDEX, + f"JSONL index is stale for {source_path}", + path=resolved_index_path, + ) + return payload + + +def read_indexed_jsonl_row( + source_path: Path, + *, + kind: str, + test_case_id: str, + index_path: Path | None = None, + max_row_bytes: int = DEFAULT_MAX_INDEXED_ROW_BYTES, +) -> dict[str, Any]: + """Seek directly to one indexed row and verify its stable identity.""" + source_path = source_path.resolve() + payload = load_jsonl_index(source_path, index_path=index_path) + key = jsonl_row_key(kind, test_case_id) + item = payload["items"].get(key) + if not isinstance(item, dict): + raise JsonlIndexError( + JsonlIndexErrorCode.NOT_FOUND, + f"JSONL row not found: {key}", + path=source_path, + key=key, + ) + offset = item.get("offset") + length = item.get("length") + if ( + not isinstance(offset, int) + or isinstance(offset, bool) + or offset < 0 + or not isinstance(length, int) + or isinstance(length, bool) + or length < 1 + ): + raise JsonlIndexError( + JsonlIndexErrorCode.INVALID_INDEX, + f"Invalid byte range for {key}", + path=index_path or jsonl_index_path(source_path), + key=key, + ) + if length > max_row_bytes: + raise JsonlIndexError( + JsonlIndexErrorCode.ROW_TOO_LARGE, + f"Indexed JSONL row exceeds {max_row_bytes} bytes: {key}", + path=source_path, + key=key, + ) + with source_path.open("rb") as handle: + handle.seek(offset) + raw = handle.read(length) + try: + row = json.loads(raw.strip()) + except json.JSONDecodeError as exc: + raise JsonlIndexError( + JsonlIndexErrorCode.STALE_INDEX, + f"Indexed JSONL row is no longer readable: {key}", + path=source_path, + key=key, + ) from exc + if not isinstance(row, dict) or _row_identity(row, path=source_path) != ( + kind, + test_case_id, + ): + raise JsonlIndexError( + JsonlIndexErrorCode.STALE_INDEX, + f"Indexed JSONL row identity changed: {key}", + path=source_path, + key=key, + ) + stat_result = source_path.stat() + source = payload["source"] + if ( + stat_result.st_size != source["size_bytes"] + or stat_result.st_mtime_ns != source["mtime_ns"] + ): + raise JsonlIndexError( + JsonlIndexErrorCode.STALE_INDEX, + f"JSONL source changed during indexed lookup: {source_path}", + path=source_path, + key=key, + ) + return row + + +def _row_identity(row: dict[str, Any], *, path: Path) -> tuple[str, str]: + kind = row.get("type") + test_case_id = row.get("test_case_id") + if not isinstance(kind, str) or not kind: + raise JsonlIndexError( + JsonlIndexErrorCode.INVALID_KEY, + f'{path.name}: expected field "type" (a non-empty string)', + path=path, + ) + if not isinstance(test_case_id, str) or not test_case_id: + raise JsonlIndexError( + JsonlIndexErrorCode.INVALID_KEY, + f'{path.name}: expected field "test_case_id" (a non-empty string)', + path=path, + ) + return kind, test_case_id + + +def _valid_index_payload(payload: dict[str, Any] | None) -> bool: + if not isinstance(payload, dict): + return False + if payload.get("schema_version") != JSONL_INDEX_SCHEMA_VERSION: + return False + source = payload.get("source") + items = payload.get("items") + order = payload.get("order") + row_count = payload.get("row_count") + return ( + isinstance(source, dict) + and isinstance(source.get("name"), str) + and isinstance(source.get("size_bytes"), int) + and isinstance(source.get("mtime_ns"), int) + and isinstance(source.get("sha256"), str) + and len(source["sha256"]) == 64 + and all(char in "0123456789abcdef" for char in source["sha256"]) + and isinstance(items, dict) + and isinstance(order, list) + and all(isinstance(key, str) for key in order) + and isinstance(row_count, int) + and row_count == len(order) + and row_count == len(items) + and len(set(order)) == len(order) + and set(order) == set(items) + ) + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/assert_ai/runner.py b/assert_ai/runner.py index e364de0e4..820ec8242 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -54,6 +54,11 @@ ) from assert_ai.core.run_result import RunResult, RunState from assert_ai.display import label_metric +from assert_ai.services.result_metadata import ( + refresh_stage_indexes, + write_run_summary, + write_suite_summary, +) from assert_ai.stages import STAGES if TYPE_CHECKING: @@ -175,6 +180,59 @@ def _record_run_artifacts(manifest: RunManifest, ctx: dict[str, Any], run_root: ) +def _refresh_stage_indexes( + ctx: dict[str, Any], + stage_name: str, + stage_result: dict[str, Any] | None = None, +) -> None: + try: + refresh_stage_indexes(ctx, stage_name, stage_result) + except Exception: # noqa: BLE001 + log.warning( + "[%s] Failed to refresh the derived JSONL index", + stage_name, + exc_info=True, + ) + + +def _refresh_run_summary( + ctx: dict[str, Any], + manifest: RunManifest, + *, + stage_usage: dict[str, dict[str, Any]] | None = None, + elapsed_s: float | None = None, + rebuild_indexes: bool, +) -> None: + try: + metrics = ( + _build_run_metrics(stage_usage or {}, elapsed_s) + if elapsed_s is not None + else None + ) + write_run_summary( + ctx, + manifest, + stage_summaries=ctx.get("_stage_summaries"), + metrics=metrics, + rebuild_indexes=rebuild_indexes, + ) + except Exception: # noqa: BLE001 + log.warning("Failed to refresh run_summary.json", exc_info=True) + + +def _refresh_suite_summary( + ctx: dict[str, Any], + *, + rebuild_indexes: bool, +) -> None: + if ctx.get("_suite_summary_blocked"): + return + try: + write_suite_summary(ctx, rebuild_indexes=rebuild_indexes) + except Exception: # noqa: BLE001 + log.warning("Failed to refresh suite_summary.json", exc_info=True) + + def _print_stage_start(stage_name: str, ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> None: """Print a human-readable stage header.""" tag = f"[{stage_name}]" @@ -747,10 +805,12 @@ def _run_pipeline_result( suite_root.mkdir(parents=True, exist_ok=True) _write_suite_metadata(ctx) ctx.setdefault("artifact_versions", {}) + ctx.setdefault("_stage_summaries", {}) artifact_plans: dict[str, Any] = {} cache_supported = supports_artifact_cache(ctx) if cache_supported: activate_latest_artifacts(ctx) + _refresh_suite_summary(ctx, rebuild_indexes=False) stages_to_run: list[tuple[str, Any, dict[str, Any]]] = [] for stage_name, raw_cfg in ctx["stages"]: @@ -773,6 +833,8 @@ def _run_pipeline_result( if plan.reused: refresh_compatibility_files(ctx, stage_name, plan.output_paths) update_latest(ctx, stage_name, ref) + _refresh_stage_indexes(ctx, stage_name) + _refresh_suite_summary(ctx, rebuild_indexes=True) log.info( f"[{stage_name}] Reused artifact {plan.version} " f"(input hashes match, use --force-stage {stage_name} to regenerate)" @@ -830,6 +892,13 @@ def _run_pipeline_result( else run_root / "config.yaml" ) shutil.copy2(config_path, config_snapshot) + _record_run_artifacts(manifest, ctx, run_root) + _write_manifest(manifest, run_root) + _refresh_run_summary( + ctx, + manifest, + rebuild_indexes=False, + ) pipeline_start = time.monotonic() stage_usage: dict[str, dict[str, Any]] = {} @@ -909,6 +978,13 @@ def _run_stages_inner( } _record_run_artifacts(manifest, ctx, run_root) _write_manifest(manifest, run_root) + _refresh_run_summary( + ctx, + manifest, + stage_usage=stage_usage, + elapsed_s=time.monotonic() - pipeline_start, + rebuild_indexes=False, + ) _print_stage_start(stage_name, ctx, raw_cfg) stage_start = time.monotonic() stage_result: dict[str, Any] = {} @@ -943,8 +1019,12 @@ def _run_stages_inner( module.run(ctx, raw_cfg), cleanup_timeout_s=300.0, ) or {} + stage_summary = (stage_result or {}).get("_summary") + if isinstance(stage_summary, dict): + ctx["_stage_summaries"][stage_name] = stage_summary + _refresh_stage_indexes(ctx, stage_name, stage_result) stage_errored_count = int( - ((stage_result or {}).get("_summary") or {}).get("errored_count", 0) or 0 + (stage_summary or {}).get("errored_count", 0) or 0 ) if ( cache_supported @@ -970,6 +1050,7 @@ def _run_stages_inner( "Re-run to fill the gap.", stage_name, stage_errored_count, ) + ctx["_suite_summary_blocked"] = True else: finalize_artifact_plan(ctx, artifact_plans[stage_name]) ok = True @@ -1034,9 +1115,17 @@ def _run_stages_inner( manifest.stage_timings[stage_name] = existing_timing _record_run_artifacts(manifest, ctx, run_root) _write_manifest(manifest, run_root) + _refresh_run_summary( + ctx, + manifest, + stage_usage=stage_usage, + elapsed_s=time.monotonic() - pipeline_start, + rebuild_indexes=ok, + ) if ok and module.SCOPE == "suite": _write_suite_metadata(ctx) + _refresh_suite_summary(ctx, rebuild_indexes=True) if not ok: failed_stage = stage_name @@ -1092,6 +1181,14 @@ def _run_stages_inner( manifest.status = "completed" if failed_stage is None else "failed" _record_run_artifacts(manifest, ctx, run_root) _write_manifest(manifest, run_root) + _refresh_run_summary( + ctx, + manifest, + stage_usage=stage_usage, + elapsed_s=total_elapsed, + rebuild_indexes=failed_stage is None, + ) + _refresh_suite_summary(ctx, rebuild_indexes=failed_stage is None) if failed_stage is None: return _run_result_from_context( diff --git a/assert_ai/services/result_metadata.py b/assert_ai/services/result_metadata.py new file mode 100644 index 000000000..dcb92726f --- /dev/null +++ b/assert_ai/services/result_metadata.py @@ -0,0 +1,703 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Persist lightweight suite/run metadata and derived JSONL indexes.""" + +from __future__ import annotations + +import hashlib +import json +from datetime import datetime, timezone +from pathlib import Path +from typing import Any +from urllib.parse import urlsplit, urlunsplit + +from assert_ai.core.io import load_json, load_jsonl, write_json +from assert_ai.core.jsonl_index import ( + JsonlIndexError, + JsonlIndexErrorCode, + build_jsonl_index, + jsonl_index_path, + load_jsonl_index, + scan_jsonl, +) +from assert_ai.results import ( + compute_prompt_metrics, + compute_scenario_metrics, +) + +SUITE_SUMMARY_SCHEMA_VERSION = 1 +RUN_SUMMARY_SCHEMA_VERSION = 1 + + +def refresh_stage_indexes( + ctx: dict[str, Any], + stage_name: str, + stage_result: dict[str, Any] | None = None, +) -> None: + """Create or refresh the canonical JSONL index produced by one stage.""" + stage_result = stage_result or {} + if stage_name == "test_set": + source = _optional_path( + stage_result.get("test_set_path") or ctx.get("test_set_path") + ) + elif stage_name == "inference": + source = _optional_path( + stage_result.get("inference_set_path") + or ctx.get("inference_set_path") + or _under_run_root(ctx, "inference_set.jsonl") + ) + if source is not None: + ctx["inference_set_path"] = str(source) + elif stage_name == "judge": + source = _optional_path( + stage_result.get("scores_path") + or ctx.get("scores_path") + or _under_run_root(ctx, "scores.jsonl") + ) + if source is not None: + ctx["scores_path"] = str(source) + else: + return + if source is not None and source.is_file(): + _ensure_jsonl_index(source) + + +def write_run_summary( + ctx: dict[str, Any], + manifest: Any, + *, + stage_summaries: dict[str, dict[str, Any]] | None = None, + metrics: dict[str, Any] | None = None, + rebuild_indexes: bool = True, +) -> dict[str, Any] | None: + """Write one metadata-only run summary without embedding source rows.""" + run_root = _optional_path(ctx.get("run_root")) + suite_root = _optional_path(ctx.get("suite_root")) + if run_root is None or suite_root is None: + return None + run_root.mkdir(parents=True, exist_ok=True) + + summary_path = run_root / "run_summary.json" + previous = _load_optional_json(summary_path) or {} + manifest_payload = ( + manifest.to_dict() + if hasattr(manifest, "to_dict") + else dict(manifest or {}) + ) + status = str(manifest_payload.get("status") or "unknown") + current_stage = _current_or_terminal_stage(manifest_payload) + + sources: dict[str, Any] = {} + indexes: dict[str, dict[str, Any]] = {} + counts: dict[str, Any] = {} + + source_paths = { + "taxonomy": _optional_path( + ctx.get("taxonomy_path") or suite_root / "taxonomy.json" + ), + "test_set": _optional_path( + ctx.get("test_set_path") or suite_root / "test_set.jsonl" + ), + "inference_set": _optional_path( + ctx.get("inference_set_path") or run_root / "inference_set.jsonl" + ), + "scores": _optional_path( + ctx.get("scores_path") or run_root / "scores.jsonl" + ), + } + + jsonl_payloads: dict[str, dict[str, Any]] = {} + for name, path in source_paths.items(): + if path is None or not path.is_file(): + continue + if path.suffix == ".jsonl": + index = _load_or_build_jsonl_index(path, rebuild=rebuild_indexes) + if index is not None: + jsonl_payloads[name] = index + indexes[name] = { + "schema_version": index["schema_version"], + **_path_reference( + jsonl_index_path(path), + suite_root=suite_root, + run_root=run_root, + ctx=ctx, + ), + } + sources[name] = { + **_path_reference( + path, + suite_root=suite_root, + run_root=run_root, + ctx=ctx, + ), + **index["source"], + "index_schema_version": index["schema_version"], + } + continue + sources[name] = _file_identity( + path, + suite_root=suite_root, + run_root=run_root, + ctx=ctx, + ) + + for name in ("test_set", "inference_set", "scores"): + index = jsonl_payloads.get(name) + if index is not None: + counts[name] = _index_counts(index) + + taxonomy = ( + _load_optional_json(source_paths["taxonomy"]) + if source_paths["taxonomy"] is not None + else None + ) + behavior_categories = ( + taxonomy.get("behavior_categories") + if isinstance(taxonomy, dict) + else None + ) + if not isinstance(behavior_categories, list): + behavior_categories = [] + + quality = previous.get("quality") + if rebuild_indexes and source_paths["scores"] is not None: + score_rows = load_jsonl(source_paths["scores"]) + prompt_rows = [row for row in score_rows if not row.get("tester_model")] + scenario_rows = [row for row in score_rows if row.get("tester_model")] + quality = { + "prompt": compute_prompt_metrics(prompt_rows, behavior_categories), + "scenario": compute_scenario_metrics( + scenario_rows, + behavior_categories, + ), + } + + payload = { + "schema_version": RUN_SUMMARY_SCHEMA_VERSION, + "suite_id": str(ctx.get("suite_id") or suite_root.name), + "run_id": str(ctx.get("run_id") or run_root.name), + "state": status, + "current_stage": current_stage, + "started_at": manifest_payload.get("started_at"), + "ended_at": manifest_payload.get("ended_at"), + "updated_at": _utc_now(), + "stages": manifest_payload.get("stages") or {}, + "stage_timings": manifest_payload.get("stage_timings") or {}, + "stage_summaries": _sanitize_managed_values( + stage_summaries or previous.get("stage_summaries") or {}, + ctx=ctx, + ), + "models": _model_references(ctx), + "counts": counts or previous.get("counts") or {}, + "quality": quality, + "metrics": metrics if metrics is not None else previous.get("metrics"), + "artifact_versions": ctx.get("artifact_versions") or {}, + "sources": sources or previous.get("sources") or {}, + "indexes": indexes or previous.get("indexes") or {}, + } + normalized_payload = _json_payload(payload) + write_json(summary_path, normalized_payload) + return normalized_payload + + +def write_suite_summary( + ctx: dict[str, Any], + *, + rebuild_indexes: bool = True, +) -> dict[str, Any] | None: + """Write one metadata-only suite catalog entry.""" + suite_root = _optional_path(ctx.get("suite_root")) + if suite_root is None: + return None + suite_root.mkdir(parents=True, exist_ok=True) + + summary_path = suite_root / "suite_summary.json" + previous = _load_optional_json(summary_path) or {} + suite_meta = _load_optional_json(suite_root / "suite.json") or {} + taxonomy_path = _optional_path( + ctx.get("taxonomy_path") or suite_root / "taxonomy.json" + ) + test_set_path = _optional_path( + ctx.get("test_set_path") or suite_root / "test_set.jsonl" + ) + + taxonomy = ( + _load_optional_json(taxonomy_path) + if taxonomy_path is not None and taxonomy_path.is_file() + else None + ) + categories = ( + taxonomy.get("behavior_categories") + if isinstance(taxonomy, dict) + else None + ) + if not isinstance(categories, list): + categories = [] + + sources: dict[str, Any] = {} + if taxonomy_path is not None and taxonomy_path.is_file(): + sources["taxonomy"] = _file_identity( + taxonomy_path, + suite_root=suite_root, + run_root=None, + ctx=ctx, + ) + + test_case_counts = previous.get("test_case_counts") or { + "total": 0, + "prompt": 0, + "scenario": 0, + "other": 0, + } + if test_set_path is not None and test_set_path.is_file(): + test_set_index = _load_or_build_jsonl_index( + test_set_path, + rebuild=rebuild_indexes, + ) + if test_set_index is not None: + test_case_counts = _index_counts(test_set_index) + sources["test_set"] = { + **_path_reference( + test_set_path, + suite_root=suite_root, + run_root=None, + ctx=ctx, + ), + **test_set_index["source"], + "index_schema_version": test_set_index["schema_version"], + "index": _path_reference( + jsonl_index_path(test_set_path), + suite_root=suite_root, + run_root=None, + ctx=ctx, + ), + } + + runs = _run_catalog_entries(suite_root) + latest_run = max( + runs, + key=lambda item: str( + item.get("ended_at") + or item.get("updated_at") + or item.get("started_at") + or "" + ), + default=None, + ) + has_results = any( + int(((entry.get("counts") or {}).get("scores") or {}).get("total", 0)) + > 0 + for entry in runs + ) + if has_results: + status = "has_results" + elif int(test_case_counts.get("total", 0)) > 0: + status = "test_set_ready" + elif taxonomy_path is not None and taxonomy_path.is_file(): + status = "systematized" + else: + status = "initialized" + + behavior_block = ( + taxonomy.get("behavior") + if isinstance(taxonomy, dict) + else None + ) + taxonomy_behavior_name = ( + behavior_block.get("name") + if isinstance(behavior_block, dict) + else None + ) + payload = { + "schema_version": SUITE_SUMMARY_SCHEMA_VERSION, + "suite_id": str(ctx.get("suite_id") or suite_root.name), + "status": status, + "behavior": { + "name": ( + ctx.get("behavior_name") + or taxonomy_behavior_name + or suite_root.name + ), + "description": ctx.get("behavior") or "", + }, + "behavior_category_count": len(categories), + "test_case_counts": test_case_counts, + "created_at": ( + suite_meta.get("created_at") + or previous.get("created_at") + or _utc_now() + ), + "updated_at": _utc_now(), + "run_count": len(runs), + "run_set_identity": suite_run_set_identity(suite_root), + "run_catalog_identity": suite_run_catalog_identity(suite_root), + "latest_run": ( + { + "run_id": latest_run.get("run_id"), + "state": latest_run.get("state"), + "started_at": latest_run.get("started_at"), + "ended_at": latest_run.get("ended_at"), + } + if latest_run is not None + else None + ), + "artifact_versions": ctx.get("artifact_versions") or {}, + "sources": sources or previous.get("sources") or {}, + } + normalized_payload = _json_payload(payload) + write_json(summary_path, normalized_payload) + return normalized_payload + + +def _ensure_jsonl_index(path: Path) -> dict[str, Any]: + return _load_or_build_jsonl_index(path, rebuild=True) or {} + + +def _load_or_build_jsonl_index( + path: Path, + *, + rebuild: bool, +) -> dict[str, Any] | None: + try: + return load_jsonl_index(path) + except JsonlIndexError as exc: + if ( + not rebuild + or exc.code + not in { + JsonlIndexErrorCode.INVALID_INDEX, + JsonlIndexErrorCode.STALE_INDEX, + } + ): + return None + try: + scan = scan_jsonl( + path, + allow_trailing_partial=path.name + in {"inference_set.jsonl", "scores.jsonl"}, + ) + return build_jsonl_index(path, scan=scan) + except JsonlIndexError: + return None + + +def _index_counts(index: dict[str, Any]) -> dict[str, int]: + counts = { + "total": 0, + "prompt": 0, + "scenario": 0, + "other": 0, + } + for item in index.get("items", {}).values(): + if not isinstance(item, dict): + continue + kind = item.get("type") + counts["total"] += 1 + if kind in {"prompt", "scenario"}: + counts[kind] += 1 + else: + counts["other"] += 1 + return counts + + +def _run_catalog_entries(suite_root: Path) -> list[dict[str, Any]]: + entries: list[dict[str, Any]] = [] + for child in sorted(suite_root.iterdir()) if suite_root.exists() else []: + if not child.is_dir() or child.name == "artifacts": + continue + summary = _load_optional_json(child / "run_summary.json") + if isinstance(summary, dict): + entries.append(summary) + continue + manifest = _load_optional_json(child / "manifest.json") + if not isinstance(manifest, dict): + continue + entries.append( + { + "run_id": child.name, + "state": manifest.get("status") or "unknown", + "started_at": manifest.get("started_at"), + "ended_at": manifest.get("ended_at"), + "counts": {}, + } + ) + return entries + + +def _model_references(ctx: dict[str, Any]) -> dict[str, Any]: + target = ctx.get("target") + evaluation = ctx.get("evaluation") + target_ref: dict[str, Any] | None = None + if target is not None: + model = getattr(target, "model", None) + if model is not None: + target_ref = { + "kind": "model", + "identifier": _sanitize_managed_string( + str(getattr(model, "name", model)), + ctx=ctx, + ), + } + else: + for kind in ("connector", "callable", "endpoint"): + value = getattr(target, kind, None) + if value: + target_ref = { + "kind": kind, + "identifier": _safe_target_identifier( + kind, + str(value), + ctx=ctx, + ), + } + break + tester = getattr(evaluation, "tester", None) if evaluation is not None else None + judge = getattr(evaluation, "judge", None) if evaluation is not None else None + return { + "target": target_ref, + "tester": ( + _optional_sanitized_model_name(tester, ctx=ctx) + if tester is not None + else None + ), + "judge": ( + _optional_sanitized_model_name(judge, ctx=ctx) + if judge is not None + else None + ), + } + + +def _current_or_terminal_stage(manifest: dict[str, Any]) -> str | None: + stages = manifest.get("stages") + if not isinstance(stages, dict): + return None + for stage_name, stage_status in stages.items(): + if stage_status == "running": + return str(stage_name) + for stage_name, stage_status in reversed(list(stages.items())): + if stage_status in {"completed", "failed", "cancelled"}: + return str(stage_name) + return None + + +def _file_identity( + path: Path, + *, + suite_root: Path, + run_root: Path | None, + ctx: dict[str, Any], +) -> dict[str, Any]: + stat_result = path.stat() + return { + **_path_reference( + path, + suite_root=suite_root, + run_root=run_root, + ctx=ctx, + ), + "name": path.name, + "size_bytes": stat_result.st_size, + "mtime_ns": stat_result.st_mtime_ns, + "sha256": _file_sha256(path), + } + + +def _path_reference( + path: Path, + *, + suite_root: Path, + run_root: Path | None, + ctx: dict[str, Any], +) -> dict[str, str]: + resolved = path.resolve() + roots: list[tuple[str, Path]] = [] + if run_root is not None: + roots.append(("run", run_root.resolve())) + roots.append(("suite", suite_root.resolve())) + path_policy = ctx.get("path_policy") + if path_policy is not None: + roots.append(("workspace", Path(path_policy.workspace_root).resolve())) + for scope, root in roots: + try: + relative = resolved.relative_to(root) + except ValueError: + continue + return {"scope": scope, "path": relative.as_posix()} + return {"scope": "external"} + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _optional_path(value: Any) -> Path | None: + if value is None or value == "": + return None + return Path(value) + + +def _under_run_root(ctx: dict[str, Any], filename: str) -> Path | None: + run_root = _optional_path(ctx.get("run_root")) + return run_root / filename if run_root is not None else None + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _load_optional_json(path: Path) -> dict[str, Any] | None: + try: + return load_json(path) + except (OSError, ValueError): + return None + + +def _json_payload(payload: dict[str, Any]) -> dict[str, Any]: + """Normalize mapping keys exactly as they will appear in persisted JSON.""" + normalized = json.loads(json.dumps(payload, ensure_ascii=False)) + assert isinstance(normalized, dict) + return normalized + + +def suite_run_catalog_identity(suite_root: Path) -> dict[str, Any]: + """Return a cheap identity for the suite's run set and run metadata.""" + entries: list[dict[str, Any]] = [] + if suite_root.exists(): + for child in sorted(suite_root.iterdir()): + if ( + not child.is_dir() + or child.name == "artifacts" + or child.name.startswith(".") + ): + continue + files: dict[str, dict[str, int]] = {} + for filename in ( + "run_summary.json", + "manifest.json", + "inference_set.jsonl", + "scores.jsonl", + ): + path = child / filename + if not path.is_file(): + continue + stat_result = path.stat() + files[filename] = { + "size_bytes": stat_result.st_size, + "mtime_ns": stat_result.st_mtime_ns, + } + if files: + entries.append({"run_id": child.name, "files": files}) + encoded = json.dumps( + entries, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return { + "run_count": len(entries), + "sha256": hashlib.sha256(encoded).hexdigest(), + } + + +def suite_run_set_identity(suite_root: Path) -> dict[str, Any]: + """Return a cheap identity that detects added or removed run directories.""" + run_ids = ( + sorted( + child.name + for child in suite_root.iterdir() + if child.is_dir() + and child.name != "artifacts" + and not child.name.startswith(".") + and any( + (child / filename).exists() + for filename in ( + "run_summary.json", + "manifest.json", + "inference_set.jsonl", + "scores.jsonl", + ) + ) + ) + if suite_root.exists() + else [] + ) + encoded = json.dumps( + run_ids, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + return { + "run_count": len(run_ids), + "sha256": hashlib.sha256(encoded).hexdigest(), + } + + +def _safe_target_identifier( + kind: str, + value: str, + *, + ctx: dict[str, Any], +) -> str: + if kind != "endpoint": + return _sanitize_managed_string(value, ctx=ctx) + try: + parsed = urlsplit(value) + if not parsed.scheme or not parsed.hostname: + return "" + host = parsed.hostname + if ":" in host and not host.startswith("["): + host = f"[{host}]" + if parsed.port is not None: + host = f"{host}:{parsed.port}" + return urlunsplit((parsed.scheme, host, "", "", "")) + except ValueError: + return "" + + +def _optional_sanitized_model_name( + owner: Any, + *, + ctx: dict[str, Any], +) -> str | None: + value = getattr(getattr(owner, "model", None), "name", None) + if value is None: + return None + return _sanitize_managed_string(str(value), ctx=ctx) + + +def _sanitize_managed_values(value: Any, *, ctx: dict[str, Any]) -> Any: + if isinstance(value, dict): + return { + str(key): _sanitize_managed_values(item, ctx=ctx) + for key, item in value.items() + } + if isinstance(value, list): + return [_sanitize_managed_values(item, ctx=ctx) for item in value] + if isinstance(value, tuple): + return [_sanitize_managed_values(item, ctx=ctx) for item in value] + if isinstance(value, str): + return _sanitize_managed_string(value, ctx=ctx) + return value + + +def _sanitize_managed_string(value: str, *, ctx: dict[str, Any]) -> str: + path_policy = ctx.get("path_policy") + if path_policy is None: + return value + roots = { + Path(path_policy.workspace_root), + Path(path_policy.config_root), + Path(path_policy.artifacts_root), + Path(path_policy.results_root), + } + sanitized = value + for root in sorted(roots, key=lambda path: len(str(path)), reverse=True): + for rendered in {str(root), root.as_posix()}: + sanitized = sanitized.replace(rendered, ".") + return sanitized diff --git a/assert_ai/services/results.py b/assert_ai/services/results.py new file mode 100644 index 000000000..60b2c0b7c --- /dev/null +++ b/assert_ai/services/results.py @@ -0,0 +1,1963 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Workspace-safe, paginated access to ASSERT result artifacts.""" + +from __future__ import annotations + +import base64 +import binascii +import hashlib +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence + +from assert_ai.core.io import ( + get_permissible_flag, + load_json, + row_behavior, + row_factors, +) +from assert_ai.core.jsonl_index import ( + DEFAULT_MAX_INDEXED_ROW_BYTES, + JsonlIndexError, + JsonlIndexErrorCode, + build_jsonl_index, + jsonl_index_path, + load_jsonl_index, + scan_jsonl, +) +from assert_ai.core.judge import get_verdict_dimension, infer_judge_status +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.result_metadata import ( + RUN_SUMMARY_SCHEMA_VERSION, + SUITE_SUMMARY_SCHEMA_VERSION, + suite_run_catalog_identity, + suite_run_set_identity, + write_run_summary, + write_suite_summary, +) + +if TYPE_CHECKING: + from assert_ai.core.runtime_path_policy import RuntimePathPolicy + +_CURSOR_VERSION = 1 +_IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,254}$") + + +@dataclass(frozen=True, slots=True) +class RunReference: + suite_id: str + run_id: str + + @property + def label(self) -> str: + return f"{self.suite_id}/{self.run_id}" + + +@dataclass(frozen=True, slots=True) +class ResultPage: + items: list[dict[str, Any]] + next_cursor: str | None + + def to_dict(self) -> dict[str, Any]: + return { + "items": self.items, + "next_cursor": self.next_cursor, + } + + +class ResultRepository: + """Read and repair result artifacts beneath one configured results root.""" + + def __init__( + self, + results_root: Path, + *, + path_policy: RuntimePathPolicy | None = None, + default_page_size: int = 50, + max_page_size: int = 200, + max_page_bytes: int = 1024 * 1024, + max_item_bytes: int = DEFAULT_MAX_INDEXED_ROW_BYTES, + ) -> None: + if default_page_size < 1: + raise ValueError("default_page_size must be positive") + if max_page_size < default_page_size: + raise ValueError("max_page_size must be >= default_page_size") + if max_page_bytes < 1 or max_item_bytes < 1: + raise ValueError("result size limits must be positive") + self.results_root = results_root.resolve() + self.path_policy = path_policy + if path_policy is not None: + try: + self.results_root = path_policy.resolve_managed_output( + self.results_root, + field_name="results root", + expected_root=path_policy.results_root, + reject_links=True, + ) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + str(exc), + ) from exc + self.default_page_size = default_page_size + self.max_page_size = max_page_size + self.max_page_bytes = max_page_bytes + self.max_item_bytes = max_item_bytes + + def list_suite_catalog_entries( + self, + *, + cursor: str | None = None, + page_size: int | None = None, + ) -> ResultPage: + entries: list[dict[str, Any]] = [] + if self.results_root.exists(): + for suite_dir in sorted(self.results_root.iterdir()): + if not suite_dir.is_dir() or suite_dir.name.startswith("."): + continue + managed_suite_dir = self._safe_child( + self.results_root, + suite_dir.name, + field_name="suite", + ) + summary = self._ensure_suite_summary( + managed_suite_dir, + verify_run_catalog=False, + ) + if summary is not None: + entries.append(self._suite_catalog_entry(summary)) + entries.sort( + key=lambda item: ( + str(item.get("created_at") or ""), + item["suite_id"], + ), + reverse=True, + ) + return self._catalog_page( + entries, + kind="suite_catalog", + cursor=cursor, + page_size=page_size, + query={}, + ) + + def get_suite(self, suite_id: str) -> dict[str, Any]: + suite_dir = self._suite_dir(suite_id, must_exist=True) + summary = self._ensure_suite_summary(suite_dir) + if summary is None: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Suite has no readable artifacts: {suite_id}", + ) + return self._public_summary(summary) + + def list_run_catalog_entries( + self, + suite_id: str, + *, + cursor: str | None = None, + page_size: int | None = None, + ) -> ResultPage: + suite_dir = self._suite_dir(suite_id, must_exist=True) + entries: list[dict[str, Any]] = [] + for run_dir in self._run_dirs(suite_dir): + summary = self._ensure_run_summary(suite_dir, run_dir) + if summary is not None: + entries.append(self._run_catalog_entry(summary)) + entries.sort( + key=lambda item: ( + str( + item.get("ended_at") + or item.get("updated_at") + or item.get("started_at") + or "" + ), + item["run_id"], + ), + reverse=True, + ) + return self._catalog_page( + entries, + kind="run_catalog", + cursor=cursor, + page_size=page_size, + query={"suite_id": suite_id}, + ) + + def load_run_detail( + self, + suite_id: str, + run_id: str, + *, + include_rows: bool = False, + ) -> dict[str, Any]: + suite_dir = self._suite_dir(suite_id, must_exist=True) + run_dir = self._run_dir(suite_dir, run_id, must_exist=True) + summary = self._ensure_run_summary(suite_dir, run_dir) + if summary is None: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Run has no readable artifacts: {suite_id}/{run_id}", + ) + detail = self._public_summary(summary) + if include_rows: + score_path = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="scores", + fallback=run_dir / "scores.jsonl", + ) + score_rows = self._all_rows(score_path) if score_path is not None else [] + detail["prompt_rows"] = [ + row for row in score_rows if not row.get("tester_model") + ] + detail["scenario_rows"] = [ + row for row in score_rows if row.get("tester_model") + ] + return detail + + def list_test_cases( + self, + suite_id: str, + *, + run_id: str | None = None, + cursor: str | None = None, + page_size: int | None = None, + kind: str | None = None, + behavior: str | None = None, + test_case_id: str | None = None, + factors: dict[str, Any] | None = None, + ) -> ResultPage: + suite_dir, run_dir, summary = self._source_context(suite_id, run_id) + source = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="test_set", + fallback=suite_dir / "test_set.jsonl", + ) + query = _compact_mapping( + { + "suite_id": suite_id, + "run_id": run_id, + "kind": kind, + "behavior": behavior, + "test_case_id": test_case_id, + "factors": factors or None, + } + ) + return self._query_jsonl( + source, + cursor=cursor, + page_size=page_size, + cursor_kind="test_cases", + query=query, + predicate=lambda row: self._matches_common( + row, + kind=kind, + behavior=behavior, + test_case_id=test_case_id, + factors=factors, + ), + ) + + def get_test_case( + self, + suite_id: str, + test_case_id: str, + *, + kind: str | None = None, + run_id: str | None = None, + ) -> dict[str, Any]: + suite_dir, run_dir, summary = self._source_context(suite_id, run_id) + source = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="test_set", + fallback=suite_dir / "test_set.jsonl", + ) + return self._lookup_row( + source, + test_case_id=test_case_id, + kind=kind, + ) + + def list_scores( + self, + suite_id: str, + run_id: str, + *, + cursor: str | None = None, + page_size: int | None = None, + kind: str | None = None, + behavior: str | None = None, + test_case_id: str | None = None, + dimension: str | None = None, + dimension_value: bool | int | str | None = None, + match_not_applicable: bool = False, + judge_status: str | None = None, + target: str | None = None, + stop_reason: str | None = None, + has_tool_use: bool | None = None, + factors: dict[str, Any] | None = None, + ) -> ResultPage: + suite_dir, run_dir, summary = self._source_context(suite_id, run_id) + assert run_dir is not None + source = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="scores", + fallback=run_dir / "scores.jsonl", + ) + inference_lookup = self._inference_lookup( + suite_dir, + run_dir, + summary, + ) if stop_reason is not None or has_tool_use is not None else None + query = _compact_mapping( + { + "suite_id": suite_id, + "run_id": run_id, + "kind": kind, + "behavior": behavior, + "test_case_id": test_case_id, + "dimension": dimension, + "dimension_value": dimension_value, + "match_not_applicable": match_not_applicable or None, + "judge_status": judge_status, + "target": target, + "stop_reason": stop_reason, + "has_tool_use": has_tool_use, + "factors": factors or None, + } + ) + + def matches(row: dict[str, Any]) -> bool: + if not self._matches_common( + row, + kind=kind, + behavior=behavior, + test_case_id=test_case_id, + factors=factors, + ): + return False + if judge_status is not None and infer_judge_status(row) != judge_status: + return False + if target is not None and row.get("target") != target: + return False + if dimension is not None: + value = get_verdict_dimension(row.get("verdict"), dimension) + if match_not_applicable: + applicability = ( + row.get("verdict", {}).get("dimension_applicability") + if isinstance(row.get("verdict"), dict) + else None + ) + if ( + not isinstance(applicability, dict) + or applicability.get(dimension) is not False + ): + return False + elif dimension_value is None: + if value is None: + return False + elif value != dimension_value: + return False + if inference_lookup is not None: + inference_row = inference_lookup(row) + if inference_row is None: + return False + if ( + stop_reason is not None + and inference_row.get("stop_reason") != stop_reason + ): + return False + if ( + has_tool_use is not None + and _row_has_tool_use(inference_row) is not has_tool_use + ): + return False + return True + + return self._query_jsonl( + source, + cursor=cursor, + page_size=page_size, + cursor_kind="scores", + query=query, + predicate=matches, + ) + + def list_failures( + self, + suite_id: str, + run_id: str, + *, + dimension: str = "policy_violation", + include_judge_failures: bool = True, + cursor: str | None = None, + page_size: int | None = None, + kind: str | None = None, + behavior: str | None = None, + ) -> ResultPage: + suite_dir, run_dir, summary = self._source_context(suite_id, run_id) + assert run_dir is not None + source = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="scores", + fallback=run_dir / "scores.jsonl", + ) + query = _compact_mapping( + { + "suite_id": suite_id, + "run_id": run_id, + "dimension": dimension, + "include_judge_failures": include_judge_failures, + "kind": kind, + "behavior": behavior, + } + ) + + def is_failure(row: dict[str, Any]) -> bool: + if not self._matches_common(row, kind=kind, behavior=behavior): + return False + status = infer_judge_status(row) + if include_judge_failures and status != "ok": + return True + return ( + status == "ok" + and get_verdict_dimension(row.get("verdict"), dimension) is True + ) + + return self._query_jsonl( + source, + cursor=cursor, + page_size=page_size, + cursor_kind="failures", + query=query, + predicate=is_failure, + ) + + def get_transcript( + self, + suite_id: str, + run_id: str, + test_case_id: str, + *, + kind: str | None = None, + ) -> dict[str, Any]: + suite_dir, run_dir, summary = self._source_context(suite_id, run_id) + assert run_dir is not None + inference_path = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="inference_set", + fallback=run_dir / "inference_set.jsonl", + ) + inference = self._lookup_row( + inference_path, + test_case_id=test_case_id, + kind=kind, + ) + resolved_kind = str(inference.get("type") or kind or "") + + test_case = self._optional_lookup( + self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="test_set", + fallback=suite_dir / "test_set.jsonl", + ), + test_case_id=test_case_id, + kind=resolved_kind or None, + ) + score = self._optional_lookup( + self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="scores", + fallback=run_dir / "scores.jsonl", + ), + test_case_id=test_case_id, + kind=resolved_kind or None, + ) + payload = { + "suite_id": suite_id, + "run_id": run_id, + "type": resolved_kind, + "test_case_id": test_case_id, + "test_case": test_case, + "inference": inference, + "score": score, + } + if len(_canonical_json(payload)) > self.max_item_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + "Transcript response exceeds the configured item limit", + details={"test_case_id": test_case_id}, + ) + return payload + + def compare_runs( + self, + run_refs: Sequence[RunReference | tuple[str, str]], + *, + metric: str = "policy_violation", + behavior_limit: int = 8, + ) -> dict[str, Any]: + refs = [ + ref if isinstance(ref, RunReference) else RunReference(*ref) + for ref in run_refs + ] + if len(refs) < 2: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "At least two runs are required for comparison", + ) + details = [ + self.load_run_detail(ref.suite_id, ref.run_id) + for ref in refs + ] + available_dimensions = _available_dimensions(details) + if metric not in available_dimensions: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"Metric '{metric}' was not found in the compared runs", + details={"available_metrics": sorted(available_dimensions)}, + ) + + rows_by_run: list[list[dict[str, Any]]] = [] + inference_by_run: list[list[dict[str, Any]]] = [] + run_payloads: list[dict[str, Any]] = [] + for ref, detail in zip(refs, details, strict=True): + suite_dir = self._suite_dir(ref.suite_id, must_exist=True) + run_dir = self._run_dir(suite_dir, ref.run_id, must_exist=True) + scores_path = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=detail, + source_name="scores", + fallback=run_dir / "scores.jsonl", + ) + inference_path = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=detail, + source_name="inference_set", + fallback=run_dir / "inference_set.jsonl", + ) + score_rows = self._all_rows(scores_path) if scores_path is not None else [] + inference_rows = ( + self._all_rows(inference_path) + if inference_path is not None + else [] + ) + rows_by_run.append(score_rows) + inference_by_run.append(inference_rows) + run_payloads.append( + { + "label": ref.label, + "suite_id": ref.suite_id, + "run_id": ref.run_id, + "state": detail.get("state"), + "started_at": detail.get("started_at"), + "ended_at": detail.get("ended_at"), + "quality": detail.get("quality") or {}, + "models": detail.get("models") or {}, + "usage": (detail.get("metrics") or {}).get("totals") or {}, + "elapsed_s": (detail.get("metrics") or {}).get("elapsed_s"), + "structural": _structural_summary(inference_rows), + } + ) + + dimension_deltas = _dimension_deltas( + refs, + details, + available_dimensions, + ) + behavior_deltas: list[dict[str, Any]] = [] + requested_summary = _first_dimension_summary(details, metric) + if not ( + isinstance(requested_summary, dict) + and requested_summary.get("kind") == "ordinal" + ): + first_map = _behavior_metric_map( + [row for row in rows_by_run[0] if not row.get("tester_model")], + metric, + ) + last_map = _behavior_metric_map( + [row for row in rows_by_run[-1] if not row.get("tester_model")], + metric, + ) + for behavior_name in sorted(set(first_map) | set(last_map)): + first = first_map.get(behavior_name) + last = last_map.get(behavior_name) + if first is None or last is None: + continue + behavior_deltas.append( + { + "behavior_category": behavior_name, + "permissible": first.get("permissible"), + "first_rate": first["rate"], + "last_rate": last["rate"], + "delta": last["rate"] - first["rate"], + "first_count": first["count"], + "last_count": last["count"], + } + ) + behavior_deltas.sort( + key=lambda item: abs(float(item["delta"])), + reverse=True, + ) + if behavior_limit >= 0: + behavior_deltas = behavior_deltas[:behavior_limit] + + warnings = _comparison_warnings(refs, details) + return { + "metric": metric, + "baseline": refs[0].label, + "runs": run_payloads, + "dimension_deltas": dimension_deltas, + "behavior_category_deltas": behavior_deltas, + "warnings": warnings, + } + + def _source_context( + self, + suite_id: str, + run_id: str | None, + ) -> tuple[Path, Path | None, dict[str, Any]]: + suite_dir = self._suite_dir(suite_id, must_exist=True) + if run_id is None: + summary = self._ensure_suite_summary(suite_dir) + if summary is None: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Suite has no readable artifacts: {suite_id}", + ) + return suite_dir, None, summary + run_dir = self._run_dir(suite_dir, run_id, must_exist=True) + summary = self._ensure_run_summary(suite_dir, run_dir) + if summary is None: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Run has no readable artifacts: {suite_id}/{run_id}", + ) + return suite_dir, run_dir, summary + + def _ensure_suite_summary( + self, + suite_dir: Path, + *, + verify_run_catalog: bool = True, + ) -> dict[str, Any] | None: + summary_path = suite_dir / "suite_summary.json" + summary = _load_optional_json(summary_path) + if ( + isinstance(summary, dict) + and summary.get("schema_version") == SUITE_SUMMARY_SCHEMA_VERSION + and self._summary_sources_current( + summary, + suite_dir=suite_dir, + run_dir=None, + ) + and ( + ( + summary.get("run_catalog_identity") + == suite_run_catalog_identity(suite_dir) + ) + if verify_run_catalog + else ( + summary.get("run_set_identity") + == suite_run_set_identity(suite_dir) + ) + ) + ): + return summary + if not self._suite_has_artifacts(suite_dir): + return None + for run_dir in self._run_dirs(suite_dir): + self._ensure_run_summary(suite_dir, run_dir) + ctx = self._legacy_context(suite_dir) + return write_suite_summary(ctx, rebuild_indexes=True) + + def _ensure_run_summary( + self, + suite_dir: Path, + run_dir: Path, + ) -> dict[str, Any] | None: + summary_path = run_dir / "run_summary.json" + summary = _load_optional_json(summary_path) + if ( + isinstance(summary, dict) + and summary.get("schema_version") == RUN_SUMMARY_SCHEMA_VERSION + and self._summary_sources_current( + summary, + suite_dir=suite_dir, + run_dir=run_dir, + ) + ): + return summary + if not self._run_has_artifacts(run_dir): + return None + ctx = self._legacy_context(suite_dir, run_dir=run_dir) + manifest = _load_optional_json(run_dir / "manifest.json") + if not isinstance(manifest, dict): + scores_exist = (run_dir / "scores.jsonl").is_file() + inference_exists = (run_dir / "inference_set.jsonl").is_file() + stages: dict[str, str] = {} + if inference_exists: + stages["inference"] = "completed" + if scores_exist: + stages["judge"] = "completed" + manifest = { + "status": "completed", + "started_at": None, + "ended_at": None, + "stages": stages, + } + return write_run_summary( + ctx, + manifest, + rebuild_indexes=True, + ) + + def _legacy_context( + self, + suite_dir: Path, + *, + run_dir: Path | None = None, + ) -> dict[str, Any]: + latest = _load_optional_json(suite_dir / "latest.json") or {} + refs = ( + dict(latest.get("artifacts") or {}) + if isinstance(latest, dict) + and isinstance(latest.get("artifacts") or {}, dict) + else {} + ) + if run_dir is not None: + run_artifacts = _load_optional_json(run_dir / "artifacts.json") or {} + run_refs = ( + run_artifacts.get("artifacts") + if isinstance(run_artifacts, dict) + else None + ) + if isinstance(run_refs, dict): + refs.update(run_refs) + refs = self._safe_artifact_refs(suite_dir, refs) + + taxonomy_path = self._artifact_ref_path( + suite_dir, + refs.get("systematize"), + fallback=suite_dir / "taxonomy.json", + ) + test_set_path = self._artifact_ref_path( + suite_dir, + refs.get("test_set"), + fallback=suite_dir / "test_set.jsonl", + ) + taxonomy = ( + _load_optional_json(taxonomy_path) + if taxonomy_path.is_file() + else None + ) + behavior = ( + taxonomy.get("behavior") + if isinstance(taxonomy, dict) + else None + ) + ctx: dict[str, Any] = { + "suite_id": suite_dir.name, + "suite_root": suite_dir, + "run_id": run_dir.name if run_dir is not None else None, + "run_root": run_dir, + "taxonomy_path": taxonomy_path, + "test_set_path": test_set_path, + "behavior_name": ( + behavior.get("name") + if isinstance(behavior, dict) + else suite_dir.name + ), + "behavior": ( + behavior.get("description", "") + if isinstance(behavior, dict) + else "" + ), + "artifact_versions": refs, + "path_policy": self.path_policy, + "target": None, + "evaluation": None, + } + if run_dir is not None: + ctx["inference_set_path"] = run_dir / "inference_set.jsonl" + ctx["scores_path"] = run_dir / "scores.jsonl" + return ctx + + def _source_path( + self, + *, + suite_dir: Path, + run_dir: Path | None, + summary: dict[str, Any], + source_name: str, + fallback: Path, + ) -> Path | None: + sources = summary.get("sources") + source = sources.get(source_name) if isinstance(sources, dict) else None + if isinstance(source, dict): + scope = source.get("scope") + raw_path = source.get("path") + root = ( + run_dir + if scope == "run" + else suite_dir + if scope == "suite" + else self.path_policy.workspace_root + if scope == "workspace" and self.path_policy is not None + else None + ) + if root is not None and isinstance(raw_path, str): + candidate = self._safe_relative_path( + Path(root), + raw_path, + field_name=f"{source_name} source", + ) + if candidate.is_file(): + return candidate + if fallback.is_file(): + return fallback + return None + + def _summary_sources_current( + self, + summary: dict[str, Any], + *, + suite_dir: Path, + run_dir: Path | None, + ) -> bool: + sources = summary.get("sources") + if not isinstance(sources, dict): + return True + for source in sources.values(): + if not isinstance(source, dict): + continue + scope = source.get("scope") + raw_path = source.get("path") + expected_size = source.get("size_bytes") + expected_mtime = source.get("mtime_ns") + if ( + not isinstance(raw_path, str) + or not isinstance(expected_size, int) + or not isinstance(expected_mtime, int) + ): + continue + root = ( + run_dir + if scope == "run" + else suite_dir + if scope == "suite" + else self.path_policy.workspace_root + if scope == "workspace" and self.path_policy is not None + else None + ) + if root is None: + continue + candidate = self._safe_relative_path( + Path(root), + raw_path, + field_name="summary source", + ) + try: + stat_result = candidate.stat() + except FileNotFoundError: + return False + if ( + stat_result.st_size != expected_size + or stat_result.st_mtime_ns != expected_mtime + ): + return False + return True + + def _artifact_ref_path( + self, + suite_dir: Path, + ref: Any, + *, + fallback: Path, + ) -> Path: + raw_path = ref.get("path") if isinstance(ref, dict) else None + if isinstance(raw_path, str): + candidate = self._safe_relative_path( + suite_dir, + raw_path, + field_name="artifact reference", + ) + if candidate.is_file(): + return candidate + return fallback + + def _query_jsonl( + self, + source: Path | None, + *, + cursor: str | None, + page_size: int | None, + cursor_kind: str, + query: dict[str, Any], + predicate: Callable[[dict[str, Any]], bool], + ) -> ResultPage: + if source is None or not source.is_file(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Result artifact not found for {cursor_kind}", + ) + index = self._ensure_index(source) + source_identity = str(index["source"]["sha256"]) + query_identity = hashlib.sha256(_canonical_json(query)).hexdigest() + start_offset = 0 + if cursor is not None: + cursor_payload = self._decode_cursor(cursor, expected_kind=cursor_kind) + if ( + cursor_payload.get("source_sha256") != source_identity + or cursor_payload.get("query_sha256") != query_identity + ): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + "The result source or query changed after this cursor was issued", + ) + start_offset = cursor_payload.get("offset") + if ( + not isinstance(start_offset, int) + or isinstance(start_offset, bool) + or start_offset < 0 + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid result cursor offset", + ) + + limit = self._page_size(page_size) + ordered_items = [ + index["items"][key] + for key in index["order"] + if isinstance(index["items"].get(key), dict) + ] + items: list[dict[str, Any]] = [] + response_bytes = 2 + next_offset: int | None = None + + with source.open("rb") as handle: + for item in ordered_items: + offset = item.get("offset") + if not isinstance(offset, int) or offset < start_offset: + continue + length = item.get("length") + if ( + isinstance(length, int) + and length > self.max_item_bytes + ): + resume_cursor = self._encode_cursor( + { + "kind": cursor_kind, + "source_sha256": source_identity, + "query_sha256": query_identity, + "offset": offset + length, + } + ) + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + "One result row exceeds the configured item limit", + details={ + "type": item.get("type"), + "test_case_id": item.get("test_case_id"), + "size_bytes": length, + "resume_cursor": resume_cursor, + }, + ) + row = self._read_index_item(handle, source, item) + if not predicate(row): + continue + row_size = len(_canonical_json(row)) + if items and ( + len(items) >= limit + or response_bytes + row_size > self.max_page_bytes + ): + next_offset = offset + break + if not items and response_bytes + row_size > self.max_page_bytes: + row = _oversized_row_stub( + row, + size_bytes=row_size, + ) + row_size = len(_canonical_json(row)) + items.append(row) + response_bytes += row_size + + current_stat = source.stat() + if ( + current_stat.st_size != index["source"]["size_bytes"] + or current_stat.st_mtime_ns != index["source"]["mtime_ns"] + ): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + "The result source changed during pagination", + ) + next_cursor = ( + self._encode_cursor( + { + "kind": cursor_kind, + "source_sha256": source_identity, + "query_sha256": query_identity, + "offset": next_offset, + } + ) + if next_offset is not None + else None + ) + return ResultPage(items=items, next_cursor=next_cursor) + + def _lookup_row( + self, + source: Path | None, + *, + test_case_id: str, + kind: str | None, + ) -> dict[str, Any]: + if source is None or not source.is_file(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + "Result artifact not found", + ) + index = self._ensure_index(source) + item = self._find_index_item( + index, + test_case_id=test_case_id, + kind=kind, + ) + with source.open("rb") as handle: + row = self._read_index_item(handle, source, item) + current_stat = source.stat() + if ( + current_stat.st_size != index["source"]["size_bytes"] + or current_stat.st_mtime_ns != index["source"]["mtime_ns"] + ): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + "The result source changed during indexed lookup", + ) + return row + + def _optional_lookup( + self, + source: Path | None, + *, + test_case_id: str, + kind: str | None, + ) -> dict[str, Any] | None: + try: + return self._lookup_row( + source, + test_case_id=test_case_id, + kind=kind, + ) + except ServiceError as exc: + if exc.code == ServiceErrorCode.NOT_FOUND: + return None + raise + + def _inference_lookup( + self, + suite_dir: Path, + run_dir: Path, + summary: dict[str, Any], + ) -> Callable[[dict[str, Any]], dict[str, Any] | None]: + inference_path = self._source_path( + suite_dir=suite_dir, + run_dir=run_dir, + summary=summary, + source_name="inference_set", + fallback=run_dir / "inference_set.jsonl", + ) + inference_index = ( + self._ensure_index(inference_path) + if inference_path is not None + else None + ) + cache: dict[str, dict[str, Any] | None] = {} + + def lookup(score_row: dict[str, Any]) -> dict[str, Any] | None: + kind = score_row.get("type") + test_case_id = score_row.get("test_case_id") + if not isinstance(kind, str) or not isinstance(test_case_id, str): + return None + key = f"{kind}:{test_case_id}" + if key not in cache: + item = ( + inference_index.get("items", {}).get(key) + if isinstance(inference_index, dict) + else None + ) + if ( + inference_path is None + or not isinstance(item, dict) + ): + cache[key] = None + else: + with inference_path.open("rb") as handle: + cache[key] = self._read_index_item( + handle, + inference_path, + item, + ) + return cache[key] + + return lookup + + def _ensure_index(self, source: Path) -> dict[str, Any]: + try: + return load_jsonl_index(source) + except JsonlIndexError as exc: + if exc.code not in { + JsonlIndexErrorCode.INVALID_INDEX, + JsonlIndexErrorCode.STALE_INDEX, + }: + raise self._jsonl_service_error(exc) from exc + try: + scan = scan_jsonl( + source, + allow_trailing_partial=_allows_trailing_partial(source), + ) + return build_jsonl_index(source, scan=scan) + except JsonlIndexError as exc: + raise self._jsonl_service_error(exc) from exc + + def _read_index_item( + self, + handle: Any, + source: Path, + item: dict[str, Any], + ) -> dict[str, Any]: + offset = item.get("offset") + length = item.get("length") + if ( + not isinstance(offset, int) + or isinstance(offset, bool) + or offset < 0 + or not isinstance(length, int) + or isinstance(length, bool) + or length < 1 + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + f"Invalid JSONL index byte range for {source.name}", + ) + if length > self.max_item_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + "Indexed result row exceeds the configured item limit", + details={ + "type": item.get("type"), + "test_case_id": item.get("test_case_id"), + }, + ) + handle.seek(offset) + raw = handle.read(length) + try: + row = json.loads(raw.strip()) + except json.JSONDecodeError as exc: + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + f"Indexed row is no longer readable in {source.name}", + ) from exc + if not isinstance(row, dict): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + f"Indexed row is no longer an object in {source.name}", + ) + if ( + row.get("type") != item.get("type") + or row.get("test_case_id") != item.get("test_case_id") + ): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + f"Indexed row identity changed in {source.name}", + ) + return row + + def _find_index_item( + self, + index: dict[str, Any], + *, + test_case_id: str, + kind: str | None, + ) -> dict[str, Any]: + if kind is not None: + exact = index.get("items", {}).get(f"{kind}:{test_case_id}") + if isinstance(exact, dict): + return exact + matches = [ + item + for item in index.get("items", {}).values() + if isinstance(item, dict) + and item.get("test_case_id") == test_case_id + and (kind is None or item.get("type") == kind) + ] + if not matches: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Test case not found: {test_case_id}", + ) + if len(matches) > 1: + raise ServiceError( + ServiceErrorCode.CONFLICT, + f"Test case ID is ambiguous; provide its type: {test_case_id}", + ) + return matches[0] + + def _all_rows(self, source: Path | None) -> list[dict[str, Any]]: + if source is None or not source.is_file(): + return [] + try: + return [ + record.row + for record in scan_jsonl( + source, + allow_trailing_partial=_allows_trailing_partial(source), + ).records + ] + except JsonlIndexError as exc: + raise self._jsonl_service_error(exc) from exc + + def _catalog_page( + self, + entries: list[dict[str, Any]], + *, + kind: str, + cursor: str | None, + page_size: int | None, + query: dict[str, Any], + ) -> ResultPage: + identity = hashlib.sha256(_canonical_json(entries)).hexdigest() + query_identity = hashlib.sha256(_canonical_json(query)).hexdigest() + offset = 0 + if cursor is not None: + payload = self._decode_cursor(cursor, expected_kind=kind) + if ( + payload.get("catalog_sha256") != identity + or payload.get("query_sha256") != query_identity + ): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + "The result catalog changed after this cursor was issued", + ) + offset = payload.get("offset") + if ( + not isinstance(offset, int) + or isinstance(offset, bool) + or offset < 0 + or offset > len(entries) + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid catalog cursor offset", + ) + limit = self._page_size(page_size) + page_items = entries[offset : offset + limit] + next_offset = offset + len(page_items) + next_cursor = ( + self._encode_cursor( + { + "kind": kind, + "catalog_sha256": identity, + "query_sha256": query_identity, + "offset": next_offset, + } + ) + if next_offset < len(entries) + else None + ) + return ResultPage(items=page_items, next_cursor=next_cursor) + + def _suite_catalog_entry( + self, + summary: dict[str, Any], + ) -> dict[str, Any]: + behavior = summary.get("behavior") + counts = summary.get("test_case_counts") + return { + "suite_id": summary.get("suite_id"), + "status": summary.get("status"), + "behavior_name": ( + behavior.get("name") + if isinstance(behavior, dict) + else None + ), + "behavior_category_count": int( + summary.get("behavior_category_count") or 0 + ), + "prompt_test_case_count": int( + (counts or {}).get("prompt") or 0 + ), + "scenario_test_case_count": int( + (counts or {}).get("scenario") or 0 + ), + "run_count": int(summary.get("run_count") or 0), + "created_at": summary.get("created_at"), + "updated_at": summary.get("updated_at"), + "latest_run": summary.get("latest_run"), + } + + def _run_catalog_entry( + self, + summary: dict[str, Any], + ) -> dict[str, Any]: + quality = summary.get("quality") or {} + return { + "suite_id": summary.get("suite_id"), + "run_id": summary.get("run_id"), + "status": summary.get("state"), + "current_stage": summary.get("current_stage"), + "started_at": summary.get("started_at"), + "ended_at": summary.get("ended_at"), + "updated_at": summary.get("updated_at"), + "prompt_metrics": quality.get("prompt"), + "scenario_metrics": quality.get("scenario"), + "models": summary.get("models") or {}, + "counts": summary.get("counts") or {}, + "metrics": summary.get("metrics"), + } + + def _matches_common( + self, + row: dict[str, Any], + *, + kind: str | None = None, + behavior: str | None = None, + test_case_id: str | None = None, + factors: dict[str, Any] | None = None, + ) -> bool: + if kind is not None and row.get("type") != kind: + return False + if test_case_id is not None and row.get("test_case_id") != test_case_id: + return False + if behavior is not None and row_behavior(row) != behavior: + return False + if factors: + row_dimensions = row_factors(row) + if any(row_dimensions.get(key) != value for key, value in factors.items()): + return False + return True + + def _suite_dir(self, suite_id: str, *, must_exist: bool) -> Path: + self._validate_identifier(suite_id, "suite_id") + candidate = self._safe_child( + self.results_root, + suite_id, + field_name="suite", + ) + if must_exist and not candidate.is_dir(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Suite not found: {suite_id}", + ) + return candidate + + def _run_dir( + self, + suite_dir: Path, + run_id: str, + *, + must_exist: bool, + ) -> Path: + self._validate_identifier(run_id, "run_id") + candidate = self._safe_child( + suite_dir, + run_id, + field_name="run", + ) + if must_exist and not candidate.is_dir(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Run not found: {suite_dir.name}/{run_id}", + ) + return candidate + + def _safe_child( + self, + root: Path, + name: str, + *, + field_name: str, + ) -> Path: + candidate = root / name + if self.path_policy is not None: + try: + return self.path_policy.resolve_managed_output( + candidate, + field_name=field_name, + expected_root=root, + reject_links=True, + ) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + str(exc), + ) from exc + resolved = candidate.resolve() + try: + resolved.relative_to(root.resolve()) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + f"{field_name} escapes the results root", + ) from exc + return resolved + + def _safe_relative_path( + self, + root: Path, + raw_path: str, + *, + field_name: str, + ) -> Path: + candidate = Path(raw_path) + if candidate.is_absolute(): + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + f"{field_name} must be relative", + ) + resolved = (root / candidate).resolve() + try: + resolved.relative_to(root.resolve()) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + f"{field_name} escapes its managed root", + ) from exc + if self.path_policy is not None: + try: + self.path_policy.require_managed_tree( + resolved, + field_name=field_name, + expected_root=root, + ) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + str(exc), + ) from exc + return resolved + + def _run_dirs(self, suite_dir: Path) -> list[Path]: + run_dirs: list[Path] = [] + for child in sorted(suite_dir.iterdir()): + if ( + not child.is_dir() + or child.name == "artifacts" + or child.name.startswith(".") + ): + continue + managed_child = self._safe_child( + suite_dir, + child.name, + field_name="run", + ) + if managed_child.is_dir() and self._run_has_artifacts(managed_child): + run_dirs.append(managed_child) + return run_dirs + + @staticmethod + def _suite_has_artifacts(suite_dir: Path) -> bool: + return any( + (suite_dir / filename).exists() + for filename in ( + "suite.json", + "suite_summary.json", + "taxonomy.json", + "test_set.jsonl", + "latest.json", + ) + ) + + @staticmethod + def _run_has_artifacts(run_dir: Path) -> bool: + return any( + (run_dir / filename).exists() + for filename in ( + "run_summary.json", + "manifest.json", + "inference_set.jsonl", + "scores.jsonl", + ) + ) + + @staticmethod + def _validate_identifier(value: str, field_name: str) -> None: + if not isinstance(value, str) or not _IDENTIFIER_RE.fullmatch(value): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{field_name} must be a safe result identifier", + ) + + def _page_size(self, page_size: int | None) -> int: + value = self.default_page_size if page_size is None else page_size + if ( + not isinstance(value, int) + or isinstance(value, bool) + or value < 1 + or value > self.max_page_size + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"page_size must be between 1 and {self.max_page_size}", + ) + return value + + @staticmethod + def _encode_cursor(payload: dict[str, Any]) -> str: + body = {"version": _CURSOR_VERSION, **payload} + return base64.urlsafe_b64encode(_canonical_json(body)).decode("ascii").rstrip("=") + + @staticmethod + def _decode_cursor( + cursor: str, + *, + expected_kind: str, + ) -> dict[str, Any]: + try: + padding = "=" * (-len(cursor) % 4) + decoded = base64.urlsafe_b64decode(cursor + padding) + payload = json.loads(decoded) + except ( + ValueError, + json.JSONDecodeError, + binascii.Error, + UnicodeError, + ) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid result cursor", + ) from exc + if ( + not isinstance(payload, dict) + or payload.get("version") != _CURSOR_VERSION + or payload.get("kind") != expected_kind + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Result cursor is not valid for this query", + ) + return payload + + def _safe_artifact_refs( + self, + suite_dir: Path, + refs: dict[str, Any], + ) -> dict[str, dict[str, Any]]: + safe_refs: dict[str, dict[str, Any]] = {} + for stage_name, raw_ref in refs.items(): + if not isinstance(stage_name, str) or not isinstance(raw_ref, dict): + continue + safe_ref: dict[str, Any] = {} + for key in ( + "artifact_type", + "version", + "input_hash", + "config_hash", + "behavior_hash", + ): + value = raw_ref.get(key) + if isinstance(value, (str, int, float, bool)) or value is None: + safe_ref[key] = value + for key in ("path", "artifact_dir", "metadata_path"): + value = raw_ref.get(key) + if not isinstance(value, str): + continue + try: + self._safe_relative_path( + suite_dir, + value, + field_name=f"{stage_name} artifact {key}", + ) + except ServiceError: + continue + safe_ref[key] = Path(value).as_posix() + file_hashes = raw_ref.get("file_hashes") + if isinstance(file_hashes, dict): + safe_ref["file_hashes"] = { + str(key): str(value) + for key, value in file_hashes.items() + if isinstance(key, str) and isinstance(value, str) + } + safe_refs[stage_name] = safe_ref + return safe_refs + + @staticmethod + def _public_summary(summary: dict[str, Any]) -> dict[str, Any]: + """Return a detached JSON-compatible payload.""" + payload = json.loads(json.dumps(summary, ensure_ascii=False)) + assert isinstance(payload, dict) + return payload + + @staticmethod + def _jsonl_service_error(exc: JsonlIndexError) -> ServiceError: + if exc.code == JsonlIndexErrorCode.NOT_FOUND: + code = ServiceErrorCode.NOT_FOUND + elif exc.code == JsonlIndexErrorCode.ROW_TOO_LARGE: + code = ServiceErrorCode.ARTIFACT_TOO_LARGE + else: + code = ServiceErrorCode.RUN_FAILED + return ServiceError( + code, + str(exc), + details={ + "jsonl_error_code": exc.code.value, + "line_number": exc.line_number, + }, + ) + + +def _compact_mapping(values: dict[str, Any]) -> dict[str, Any]: + return { + key: value + for key, value in values.items() + if value is not None + } + + +def _allows_trailing_partial(path: Path) -> bool: + return path.name in {"inference_set.jsonl", "scores.jsonl"} + + +def _load_optional_json(path: Path) -> dict[str, Any] | None: + try: + return load_json(path) + except (OSError, ValueError): + return None + + +def _canonical_json(value: Any) -> bytes: + return json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + + +def _row_has_tool_use(row: dict[str, Any]) -> bool: + events = row.get("events") + if not isinstance(events, list): + return False + return any( + isinstance(event, dict) + and isinstance(event.get("edit"), dict) + and event["edit"].get("type") in {"tool_call", "tool_result"} + for event in events + ) + + +def _oversized_row_stub( + row: dict[str, Any], + *, + size_bytes: int, +) -> dict[str, Any]: + return { + key: row[key] + for key in ( + "type", + "test_case_id", + "behavior", + "target", + "tester_model", + "judge_model", + "judge_status", + "stop_reason", + ) + if key in row + } | { + "content_omitted": True, + "size_bytes": size_bytes, + "retrieval_hint": ( + "Use get_test_case or get_transcript for this item." + ), + } + + +def _behavior_metric_map( + rows: Iterable[dict[str, Any]], + metric: str, +) -> dict[str, dict[str, Any]]: + grouped: dict[str, dict[str, Any]] = {} + for row in rows: + if infer_judge_status(row) != "ok": + continue + value = get_verdict_dimension(row.get("verdict"), metric) + if not isinstance(value, bool): + continue + behavior = row_behavior(row) + bucket = grouped.setdefault( + behavior, + { + "true_count": 0, + "count": 0, + "permissible": get_permissible_flag(row), + }, + ) + bucket["true_count"] += int(value) + bucket["count"] += 1 + return { + behavior: { + "rate": bucket["true_count"] / bucket["count"], + "count": bucket["count"], + "permissible": bucket["permissible"], + } + for behavior, bucket in grouped.items() + if bucket["count"] > 0 + } + + +def _structural_summary(rows: Iterable[dict[str, Any]]) -> dict[str, int]: + row_count = 0 + total_events = 0 + message_events = 0 + tool_events = 0 + rows_with_tools = 0 + rows_with_traces = 0 + for row in rows: + row_count += 1 + events = row.get("events") + if not isinstance(events, list): + events = [] + total_events += len(events) + row_has_tools = False + for event in events: + edit = event.get("edit") if isinstance(event, dict) else None + edit_type = edit.get("type") if isinstance(edit, dict) else None + if edit_type == "add_message": + message_events += 1 + if edit_type in {"tool_call", "tool_result"}: + tool_events += 1 + row_has_tools = True + rows_with_tools += int(row_has_tools) + rows_with_traces += int( + any( + key in row + for key in ( + "trace_id", + "span_id", + "trace", + "trace_refs", + "otel_trace", + ) + ) + ) + return { + "inference_rows": row_count, + "total_events": total_events, + "message_events": message_events, + "tool_events": tool_events, + "rows_with_tools": rows_with_tools, + "rows_with_traces": rows_with_traces, + } + + +def _available_dimensions(details: Iterable[dict[str, Any]]) -> set[str]: + names: set[str] = set() + for detail in details: + quality = detail.get("quality") + if not isinstance(quality, dict): + continue + for kind in ("prompt", "scenario"): + metrics = quality.get(kind) + dimensions = ( + metrics.get("dimensions") + if isinstance(metrics, dict) + else None + ) + if isinstance(dimensions, dict): + names.update(str(name) for name in dimensions) + return names + + +def _first_dimension_summary( + details: Iterable[dict[str, Any]], + dimension: str, +) -> dict[str, Any] | None: + for detail in details: + quality = detail.get("quality") + if not isinstance(quality, dict): + continue + for kind in ("prompt", "scenario"): + metrics = quality.get(kind) + dimensions = ( + metrics.get("dimensions") + if isinstance(metrics, dict) + else None + ) + summary = ( + dimensions.get(dimension) + if isinstance(dimensions, dict) + else None + ) + if isinstance(summary, dict): + return summary + return None + + +def _dimension_deltas( + refs: Sequence[RunReference], + details: Sequence[dict[str, Any]], + dimensions: set[str], +) -> dict[str, Any]: + baseline = details[0] + payload: dict[str, Any] = {} + for dimension in sorted(dimensions): + first_summary = _dimension_summary_by_kind(baseline, dimension) + kind = next( + ( + summary.get("kind") + for summary in first_summary.values() + if isinstance(summary, dict) and summary.get("kind") + ), + "binary", + ) + runs: list[dict[str, Any]] = [] + for ref, detail in zip(refs, details, strict=True): + summaries = _dimension_summary_by_kind(detail, dimension) + row: dict[str, Any] = { + "label": ref.label, + "prompt": summaries.get("prompt"), + "scenario": summaries.get("scenario"), + } + if kind != "ordinal": + row["prompt_rate_delta"] = _rate_delta( + summaries.get("prompt"), + first_summary.get("prompt"), + ) + row["scenario_rate_delta"] = _rate_delta( + summaries.get("scenario"), + first_summary.get("scenario"), + ) + else: + row["prompt_distribution_delta"] = _distribution_delta( + summaries.get("prompt"), + first_summary.get("prompt"), + ) + row["scenario_distribution_delta"] = _distribution_delta( + summaries.get("scenario"), + first_summary.get("scenario"), + ) + runs.append(row) + payload[dimension] = {"kind": kind, "runs": runs} + return payload + + +def _dimension_summary_by_kind( + detail: dict[str, Any], + dimension: str, +) -> dict[str, dict[str, Any] | None]: + quality = detail.get("quality") + result: dict[str, dict[str, Any] | None] = {} + for kind in ("prompt", "scenario"): + metrics = quality.get(kind) if isinstance(quality, dict) else None + dimensions = ( + metrics.get("dimensions") + if isinstance(metrics, dict) + else None + ) + summary = ( + dimensions.get(dimension) + if isinstance(dimensions, dict) + else None + ) + result[kind] = summary if isinstance(summary, dict) else None + return result + + +def _rate_delta( + current: dict[str, Any] | None, + baseline: dict[str, Any] | None, +) -> float | None: + current_rate = current.get("rate") if isinstance(current, dict) else None + baseline_rate = baseline.get("rate") if isinstance(baseline, dict) else None + if not isinstance(current_rate, (int, float)) or not isinstance( + baseline_rate, + (int, float), + ): + return None + return float(current_rate) - float(baseline_rate) + + +def _distribution_delta( + current: dict[str, Any] | None, + baseline: dict[str, Any] | None, +) -> dict[str, float]: + current_rates = current.get("rates") if isinstance(current, dict) else None + baseline_rates = baseline.get("rates") if isinstance(baseline, dict) else None + if not isinstance(current_rates, dict) or not isinstance(baseline_rates, dict): + return {} + return { + str(grade): float(current_rates.get(grade, 0.0)) + - float(baseline_rates.get(grade, 0.0)) + for grade in sorted(set(current_rates) | set(baseline_rates)) + } + + +def _comparison_warnings( + refs: Sequence[RunReference], + details: Sequence[dict[str, Any]], +) -> list[str]: + warnings: list[str] = [] + prompt_sizes = { + int( + (((detail.get("quality") or {}).get("prompt") or {}).get("total")) + or 0 + ) + for detail in details + } + scenario_sizes = { + int( + (((detail.get("quality") or {}).get("scenario") or {}).get("total")) + or 0 + ) + for detail in details + } + if len(prompt_sizes) > 1 or len(scenario_sizes) > 1: + warnings.append("Compared runs have different prompt or scenario sample sizes.") + + test_set_hashes = { + str( + (((detail.get("sources") or {}).get("test_set") or {}).get("sha256")) + or "" + ) + for detail in details + } + if len(test_set_hashes - {""}) > 1: + warnings.append("Compared runs reference different test-set source hashes.") + + targets = { + str( + (((detail.get("models") or {}).get("target") or {}).get("identifier")) + or "" + ) + for detail in details + } + if len(targets - {""}) > 1: + warnings.append("Compared runs use different targets.") + + if len({ref.suite_id for ref in refs}) > 1: + warnings.append("Cross-suite comparison may include different behavior taxonomies.") + return warnings diff --git a/assert_ai/viewer_read_model.py b/assert_ai/viewer_read_model.py index 46c55900b..e6a41a849 100644 --- a/assert_ai/viewer_read_model.py +++ b/assert_ai/viewer_read_model.py @@ -15,6 +15,7 @@ from assert_ai.core.io import write_json, row_behavior from assert_ai.core.judge import DIMENSION_APPLICABILITY_KEY +from assert_ai.core.jsonl_index import JsonlIndexError, scan_jsonl log = logging.getLogger(__name__) @@ -145,29 +146,14 @@ def _file_metadata(path: Path, *, relative_to: Path) -> dict[str, Any]: def _iter_jsonl_with_offsets(path: Path) -> list[tuple[int, int, dict[str, Any]]]: - if not path.exists(): - raise ViewerReadModelBuildError(f"Missing JSONL artifact: {path}") - - rows: list[tuple[int, int, dict[str, Any]]] = [] - offset = 0 - with path.open("rb") as handle: - for line_number, raw_line in enumerate(handle, 1): - length = len(raw_line) - stripped = raw_line.strip() - if not stripped: - offset += length - continue - try: - row = json.loads(stripped) - except json.JSONDecodeError as exc: - raise ViewerReadModelBuildError( - f"Invalid JSONL in {path} on line {line_number}: {exc}" - ) from exc - if not isinstance(row, dict): - raise ViewerReadModelBuildError(f"Expected JSON object in {path} on line {line_number}") - rows.append((offset, length, row)) - offset += length - return rows + try: + scan = scan_jsonl(path) + except JsonlIndexError as exc: + raise ViewerReadModelBuildError(str(exc)) from exc + return [ + (record.offset, record.length, record.row) + for record in scan.records + ] def _kind_and_test_case_id(row: dict[str, Any], *, path: Path) -> tuple[str, str]: diff --git a/tests/result_catalog_fixture.py b/tests/result_catalog_fixture.py new file mode 100644 index 000000000..9ee28c2fd --- /dev/null +++ b/tests/result_catalog_fixture.py @@ -0,0 +1,218 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +from assert_ai.core.io import write_json +from assert_ai.core.jsonl_index import build_jsonl_index +from assert_ai.services.result_metadata import ( + RUN_SUMMARY_SCHEMA_VERSION, + SUITE_SUMMARY_SCHEMA_VERSION, + suite_run_catalog_identity, + suite_run_set_identity, +) + + +@dataclass(frozen=True, slots=True) +class ResultCatalogFixture: + results_root: Path + large_suite_id: str + large_run_id: str + last_test_case_id: str + + +def create_result_catalog_fixture( + root: Path, + *, + suite_count: int = 100, + runs_per_suite: int = 10, + large_test_case_count: int = 10_000, +) -> ResultCatalogFixture: + results_root = root / "results" + large_suite_id = "suite-000" + large_run_id = "run-000" + last_test_case_id = f"case-{large_test_case_count - 1:05d}" + + for suite_index in range(suite_count): + suite_id = f"suite-{suite_index:03d}" + suite_root = results_root / suite_id + suite_root.mkdir(parents=True) + + for run_index in range(runs_per_suite): + run_id = f"run-{run_index:03d}" + run_root = suite_root / run_id + run_root.mkdir() + score_row = { + "type": "prompt", + "test_case_id": f"{suite_id}-{run_id}", + "target": "fixture-target", + "judge_model": "fixture-judge", + "judge_status": "ok", + "score_keys": ["policy_violation"], + "not_applicable_score_keys": [], + "verdict": { + "dimensions": {"policy_violation": False}, + "node_judgments": [], + }, + } + (run_root / "scores.jsonl").write_text( + json.dumps(score_row) + "\n", + encoding="utf-8", + ) + timestamp = ( + f"2026-08-{(suite_index % 28) + 1:02d}" + f"T00:{run_index:02d}:00+00:00" + ) + write_json( + run_root / "run_summary.json", + { + "schema_version": RUN_SUMMARY_SCHEMA_VERSION, + "suite_id": suite_id, + "run_id": run_id, + "state": "completed", + "current_stage": "judge", + "started_at": timestamp, + "ended_at": timestamp, + "updated_at": timestamp, + "stages": { + "inference": "completed", + "judge": "completed", + }, + "stage_timings": {}, + "stage_summaries": {}, + "models": { + "target": { + "kind": "model", + "identifier": "fixture-target", + }, + "tester": None, + "judge": "fixture-judge", + }, + "counts": { + "scores": { + "total": 1, + "prompt": 1, + "scenario": 0, + "other": 0, + } + }, + "quality": { + "prompt": { + "total": 1, + "scored_total": 1, + "judge_failures": 0, + "judge_failure_rate": 0.0, + "policy_violation_rate": 0.0, + "overrefusal_rate": None, + "dimensions": { + "policy_violation": { + "rate": 0.0, + "counts": {"0": 1, "1": 0}, + "count": 1, + "applicable_count": 1, + "not_applicable_count": 0, + "flagged_count": 0, + "clear_count": 1, + } + }, + "target": "fixture-target", + "judge_model": "fixture-judge", + }, + "scenario": None, + }, + "metrics": { + "schema_version": 1, + "elapsed_s": 1.0, + "totals": {"calls": 1}, + }, + "artifact_versions": {}, + "sources": {}, + "indexes": {}, + }, + ) + + sources: dict[str, object] = {} + test_case_counts = { + "total": 0, + "prompt": 0, + "scenario": 0, + "other": 0, + } + if suite_id == large_suite_id: + test_set_path = suite_root / "test_set.jsonl" + with test_set_path.open("w", encoding="utf-8") as handle: + for case_index in range(large_test_case_count): + handle.write( + json.dumps( + { + "type": "prompt", + "test_case_id": f"case-{case_index:05d}", + "dimensions": { + "behavior": "fixture-behavior", + "partition": case_index % 10, + }, + "seed": { + "prompt": f"Fixture prompt {case_index}", + }, + } + ) + + "\n" + ) + index = build_jsonl_index(test_set_path) + sources["test_set"] = { + "scope": "suite", + "path": "test_set.jsonl", + **index["source"], + "index_schema_version": index["schema_version"], + "index": { + "scope": "suite", + "path": "test_set.index.json", + }, + } + test_case_counts = { + "total": large_test_case_count, + "prompt": large_test_case_count, + "scenario": 0, + "other": 0, + } + + write_json( + suite_root / "suite_summary.json", + { + "schema_version": SUITE_SUMMARY_SCHEMA_VERSION, + "suite_id": suite_id, + "status": "has_results", + "behavior": { + "name": "fixture-behavior", + "description": "", + }, + "behavior_category_count": 1, + "test_case_counts": test_case_counts, + "created_at": ( + f"2026-08-{(suite_index % 28) + 1:02d}T00:00:00+00:00" + ), + "updated_at": ( + f"2026-08-{(suite_index % 28) + 1:02d}T01:00:00+00:00" + ), + "run_count": runs_per_suite, + "run_set_identity": suite_run_set_identity(suite_root), + "run_catalog_identity": suite_run_catalog_identity(suite_root), + "latest_run": { + "run_id": f"run-{runs_per_suite - 1:03d}", + "state": "completed", + }, + "artifact_versions": {}, + "sources": sources, + }, + ) + + return ResultCatalogFixture( + results_root=results_root, + large_suite_id=large_suite_id, + large_run_id=large_run_id, + last_test_case_id=last_test_case_id, + ) diff --git a/tests/test_jsonl_index.py b/tests/test_jsonl_index.py new file mode 100644 index 000000000..1630bc83e --- /dev/null +++ b/tests/test_jsonl_index.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from pathlib import Path +from tempfile import TemporaryDirectory + +import pytest + +from assert_ai.core.jsonl_index import ( + JsonlIndexError, + JsonlIndexErrorCode, + build_jsonl_index, + jsonl_index_path, + load_jsonl_index, + read_indexed_jsonl_row, + scan_jsonl, +) + + +def _write_rows(path: Path, rows: list[dict]) -> None: + path.write_text( + "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows), + encoding="utf-8", + ) + + +def test_build_and_seek_preserve_utf8_byte_offsets() -> None: + with TemporaryDirectory() as tmp: + source = Path(tmp) / "scores.jsonl" + rows = [ + {"type": "prompt", "test_case_id": "one", "text": "café"}, + {"type": "scenario", "test_case_id": "two", "text": "東京"}, + ] + _write_rows(source, rows) + + payload = build_jsonl_index(source) + + assert payload["row_count"] == 2 + assert payload["source"]["sha256"] + assert payload["order"] == ["prompt:one", "scenario:two"] + assert jsonl_index_path(source).exists() + assert read_indexed_jsonl_row( + source, + kind="scenario", + test_case_id="two", + ) == rows[1] + + +def test_scan_tracks_blank_lines_and_crlf_lengths() -> None: + with TemporaryDirectory() as tmp: + source = Path(tmp) / "test_set.jsonl" + first = b'{"type":"prompt","test_case_id":"one"}\r\n' + blank = b"\r\n" + second = b'{"type":"scenario","test_case_id":"two"}\r\n' + source.write_bytes(first + blank + second) + + scan = scan_jsonl(source) + + assert scan.records[0].offset == 0 + assert scan.records[0].length == len(first) + assert scan.records[1].offset == len(first) + len(blank) + assert scan.records[1].length == len(second) + + +def test_duplicate_keys_and_invalid_rows_fail_index_build() -> None: + with TemporaryDirectory() as tmp: + source = Path(tmp) / "inference_set.jsonl" + row = {"type": "prompt", "test_case_id": "duplicate"} + _write_rows(source, [row, row]) + + with pytest.raises(JsonlIndexError) as duplicate: + build_jsonl_index(source) + assert duplicate.value.code == JsonlIndexErrorCode.DUPLICATE_KEY + + _write_rows(source, [{"type": "prompt"}]) + with pytest.raises(JsonlIndexError) as missing_id: + build_jsonl_index(source) + assert missing_id.value.code == JsonlIndexErrorCode.INVALID_KEY + + +def test_invalid_json_and_trailing_partial_handling() -> None: + with TemporaryDirectory() as tmp: + source = Path(tmp) / "scores.jsonl" + source.write_bytes( + b'{"type":"prompt","test_case_id":"one"}\n{"type":"prompt"' + ) + + with pytest.raises(JsonlIndexError) as invalid: + scan_jsonl(source) + assert invalid.value.code == JsonlIndexErrorCode.INVALID_JSON + assert invalid.value.line_number == 2 + + scan = scan_jsonl(source, allow_trailing_partial=True) + assert len(scan.records) == 1 + + +def test_stale_index_is_rejected_after_source_change() -> None: + with TemporaryDirectory() as tmp: + source = Path(tmp) / "scores.jsonl" + _write_rows(source, [{"type": "prompt", "test_case_id": "one"}]) + build_jsonl_index(source) + + source.write_text( + '{"type":"prompt","test_case_id":"changed"}\n', + encoding="utf-8", + ) + + with pytest.raises(JsonlIndexError) as stale: + load_jsonl_index(source) + assert stale.value.code == JsonlIndexErrorCode.STALE_INDEX + + +def test_invalid_index_and_missing_row_are_typed() -> None: + with TemporaryDirectory() as tmp: + source = Path(tmp) / "scores.jsonl" + _write_rows(source, [{"type": "prompt", "test_case_id": "one"}]) + + with pytest.raises(JsonlIndexError) as missing_index: + load_jsonl_index(source) + assert missing_index.value.code == JsonlIndexErrorCode.INVALID_INDEX + + build_jsonl_index(source) + with pytest.raises(JsonlIndexError) as missing_row: + read_indexed_jsonl_row( + source, + kind="prompt", + test_case_id="missing", + ) + assert missing_row.value.code == JsonlIndexErrorCode.NOT_FOUND diff --git a/tests/test_result_metadata.py b/tests/test_result_metadata.py new file mode 100644 index 000000000..57083d69d --- /dev/null +++ b/tests/test_result_metadata.py @@ -0,0 +1,298 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from pathlib import Path +from tempfile import TemporaryDirectory +from types import SimpleNamespace +from unittest.mock import patch + +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.result_metadata import ( + write_run_summary, + write_suite_summary, +) + + +def _write_jsonl(path: Path, rows: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(row) + "\n" for row in rows), + encoding="utf-8", + ) + + +def _fixture(root: Path) -> tuple[dict, dict]: + suite_root = root / "results" / "suite-a" + run_root = suite_root / "run-a" + artifact_root = suite_root / "artifacts" / "test_set" / "v0002" + artifact_root.mkdir(parents=True) + run_root.mkdir() + (suite_root / "suite.json").write_text( + '{"created_at":"2026-08-12T00:00:00+00:00"}', + encoding="utf-8", + ) + taxonomy_path = suite_root / "artifacts" / "systematize" / "v0001" / "taxonomy.json" + taxonomy_path.parent.mkdir(parents=True) + taxonomy_path.write_text( + json.dumps( + { + "behavior": {"name": "safe_agent"}, + "behavior_categories": [ + {"name": "Allowed", "permissible": True}, + {"name": "Blocked", "permissible": False}, + ], + } + ), + encoding="utf-8", + ) + test_set_path = artifact_root / "test_set.jsonl" + _write_jsonl( + test_set_path, + [ + {"type": "prompt", "test_case_id": "p1"}, + {"type": "scenario", "test_case_id": "s1"}, + ], + ) + _write_jsonl( + suite_root / "test_set.jsonl", + [{"type": "prompt", "test_case_id": "stale"}], + ) + _write_jsonl( + run_root / "inference_set.jsonl", + [ + {"type": "prompt", "test_case_id": "p1"}, + {"type": "scenario", "test_case_id": "s1"}, + ], + ) + _write_jsonl( + run_root / "scores.jsonl", + [ + { + "type": "prompt", + "test_case_id": "p1", + "target": "target-model", + "judge_model": "judge-model", + "judge_status": "ok", + "score_keys": ["policy_violation", "severity"], + "not_applicable_score_keys": [], + "verdict": { + "dimensions": { + "policy_violation": False, + "severity": 1, + }, + "node_judgments": [], + }, + "dimension_scales": { + "severity": { + "type": "ordinal", + "values": [ + {"value": 1, "label": "low"}, + {"value": 2, "label": "high"}, + ], + } + }, + }, + { + "type": "scenario", + "test_case_id": "s1", + "target": "target-model", + "tester_model": "tester-model", + "judge_model": "judge-model", + "judge_status": "judge_failed", + "verdict": {}, + }, + ], + ) + target = SimpleNamespace( + model=SimpleNamespace(name="target-model"), + connector=None, + callable=None, + endpoint=None, + ) + evaluation = SimpleNamespace( + tester=SimpleNamespace(model=SimpleNamespace(name="tester-model")), + judge=SimpleNamespace(model=SimpleNamespace(name="judge-model")), + ) + ctx = { + "suite_id": "suite-a", + "run_id": "run-a", + "suite_root": suite_root, + "run_root": run_root, + "behavior_name": "safe_agent", + "behavior": "The agent follows policy.", + "taxonomy_path": taxonomy_path, + "test_set_path": test_set_path, + "artifact_versions": { + "test_set": { + "version": "v0002", + "path": "artifacts/test_set/v0002/test_set.jsonl", + } + }, + "target": target, + "evaluation": evaluation, + } + manifest = { + "status": "completed", + "started_at": "2026-08-12T00:01:00+00:00", + "ended_at": "2026-08-12T00:02:00+00:00", + "stages": {"inference": "completed", "judge": "completed"}, + "stage_timings": { + "judge": {"duration_secs": 1.5}, + }, + } + return ctx, manifest + + +def test_run_summary_persists_indexes_metrics_and_ordinal_dimensions() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + ctx, manifest = _fixture(root) + + payload = write_run_summary( + ctx, + manifest, + metrics={"schema_version": 1, "totals": {"calls": 2}}, + ) + + assert payload is not None + assert payload["schema_version"] == 1 + assert payload["state"] == "completed" + assert payload["current_stage"] == "judge" + assert payload["models"]["target"]["identifier"] == "target-model" + assert payload["counts"]["test_set"]["total"] == 2 + assert payload["counts"]["scores"]["scenario"] == 1 + assert ( + payload["quality"]["prompt"]["dimensions"]["severity"]["kind"] + == "ordinal" + ) + assert payload["quality"]["scenario"]["judge_failures"] == 1 + assert payload["sources"]["test_set"]["scope"] == "suite" + assert ( + payload["sources"]["test_set"]["path"] + == "artifacts/test_set/v0002/test_set.jsonl" + ) + assert str(root) not in json.dumps(payload) + assert "prompt_rows" not in payload + assert "scenario_rows" not in payload + assert ( + root + / "results" + / "suite-a" + / "run-a" + / "scores.index.json" + ).exists() + + +def test_suite_summary_uses_active_versioned_test_set_and_metadata_only_runs() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + ctx, manifest = _fixture(root) + write_run_summary(ctx, manifest) + + with patch( + "assert_ai.services.result_metadata.load_jsonl", + side_effect=AssertionError("suite catalog must not load score rows"), + ): + payload = write_suite_summary(ctx) + + assert payload is not None + assert payload["schema_version"] == 1 + assert payload["status"] == "has_results" + assert payload["test_case_counts"] == { + "total": 2, + "prompt": 1, + "scenario": 1, + "other": 0, + } + assert payload["run_count"] == 1 + assert payload["latest_run"]["run_id"] == "run-a" + assert ( + payload["sources"]["test_set"]["path"] + == "artifacts/test_set/v0002/test_set.jsonl" + ) + + +def test_running_boundary_preserves_last_valid_detail_without_rescanning_rows() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + ctx, manifest = _fixture(root) + completed = write_run_summary(ctx, manifest) + assert completed is not None + + running_manifest = { + **manifest, + "status": "running", + "ended_at": None, + "stages": {"inference": "completed", "judge": "running"}, + } + with patch( + "assert_ai.services.result_metadata.load_jsonl", + side_effect=AssertionError("running boundary must not rescan rows"), + ): + running = write_run_summary( + ctx, + running_manifest, + rebuild_indexes=False, + ) + + assert running is not None + assert running["state"] == "running" + assert running["current_stage"] == "judge" + assert running["quality"] == completed["quality"] + + +def test_endpoint_target_identifier_drops_credentials_path_and_query() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + ctx, manifest = _fixture(root) + ctx["target"] = SimpleNamespace( + model=None, + connector=None, + callable=None, + endpoint="https://user:secret@example.test/private?token=secret", + ) + + payload = write_run_summary(ctx, manifest) + + assert payload is not None + assert payload["models"]["target"] == { + "kind": "endpoint", + "identifier": "https://example.test", + } + serialized = json.dumps(payload) + assert "secret" not in serialized + assert "/private" not in serialized + + +def test_strict_summary_redacts_managed_paths_from_stage_and_model_metadata() -> None: + with TemporaryDirectory() as tmp: + workspace = WorkspaceService.create(tmp) + workspace.artifacts_root.mkdir() + ctx, manifest = _fixture(workspace.artifacts_root) + ctx["path_policy"] = workspace.path_policy + ctx["target"] = SimpleNamespace( + model=None, + connector=None, + callable=f"{workspace.root}\\agent.py:run", + endpoint=None, + ) + + payload = write_run_summary( + ctx, + manifest, + stage_summaries={ + "inference": { + "debug_path": str( + workspace.root / "private" / "trace.json" + ), + } + }, + ) + + assert payload is not None + serialized = json.dumps(payload) + assert str(workspace.root) not in serialized + assert payload["models"]["target"]["identifier"].startswith(".") diff --git a/tests/test_result_service.py b/tests/test_result_service.py new file mode 100644 index 000000000..2dda01dc5 --- /dev/null +++ b/tests/test_result_service.py @@ -0,0 +1,516 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +import shutil +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from assert_ai.cli import cli +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.results import ResultRepository, RunReference + + +def _write_json(path: Path, payload: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload), encoding="utf-8") + + +def _write_jsonl(path: Path, rows: list[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(row) + "\n" for row in rows), + encoding="utf-8", + ) + + +def _score( + test_case_id: str, + *, + kind: str = "prompt", + violation: bool = False, + severity: int = 1, + judge_status: str = "ok", +) -> dict: + verdict = ( + { + "dimensions": { + "policy_violation": violation, + "severity": severity, + }, + "node_judgments": [], + } + if judge_status == "ok" + else {} + ) + return { + "type": kind, + "test_case_id": test_case_id, + "behavior": "unsafe-action", + "target": "target-a", + "tester_model": "tester-a" if kind == "scenario" else "", + "judge_model": "judge-a", + "judge_status": judge_status, + "score_keys": ["policy_violation", "severity"], + "not_applicable_score_keys": [], + "dimension_scales": { + "severity": { + "type": "ordinal", + "values": [ + {"value": 1, "label": "low"}, + {"value": 2, "label": "high"}, + ], + } + }, + "verdict": verdict, + } + + +def _build_legacy_fixture(root: Path, *, second_run: bool = False) -> Path: + results_root = root / "results" + suite_root = results_root / "suite-a" + taxonomy = suite_root / "artifacts" / "systematize" / "v0001" / "taxonomy.json" + test_set = suite_root / "artifacts" / "test_set" / "v0002" / "test_set.jsonl" + _write_json( + suite_root / "suite.json", + {"created_at": "2026-08-12T00:00:00+00:00"}, + ) + _write_json( + taxonomy, + { + "behavior": {"name": "safe-agent"}, + "behavior_categories": [ + {"name": "unsafe-action", "permissible": False}, + ], + }, + ) + _write_jsonl( + test_set, + [ + { + "type": "prompt", + "test_case_id": "p1", + "dimensions": {"behavior": "unsafe-action", "region": "us"}, + "seed": {"prompt": "one"}, + }, + { + "type": "prompt", + "test_case_id": "p2", + "dimensions": {"behavior": "unsafe-action", "region": "eu"}, + "seed": {"prompt": "two"}, + }, + { + "type": "scenario", + "test_case_id": "s1", + "dimensions": {"behavior": "unsafe-action", "region": "us"}, + "seed": {"prompt": "three"}, + }, + ], + ) + _write_json( + suite_root / "latest.json", + { + "artifacts": { + "systematize": { + "path": "artifacts/systematize/v0001/taxonomy.json", + "version": "v0001", + }, + "test_set": { + "path": "artifacts/test_set/v0002/test_set.jsonl", + "version": "v0002", + }, + } + }, + ) + _write_jsonl( + suite_root / "test_set.jsonl", + [{"type": "prompt", "test_case_id": "stale"}], + ) + + run_ids = ["run-a", "run-b"] if second_run else ["run-a"] + for index, run_id in enumerate(run_ids): + run_root = suite_root / run_id + _write_json( + run_root / "manifest.json", + { + "status": "completed", + "started_at": f"2026-08-12T00:0{index + 1}:00+00:00", + "ended_at": f"2026-08-12T00:0{index + 2}:00+00:00", + "stages": { + "inference": "completed", + "judge": "completed", + }, + }, + ) + _write_json( + run_root / "artifacts.json", + { + "schema_version": 1, + "artifacts": { + "systematize": { + "path": "artifacts/systematize/v0001/taxonomy.json", + "version": "v0001", + }, + "test_set": { + "path": "artifacts/test_set/v0002/test_set.jsonl", + "version": "v0002", + }, + }, + }, + ) + _write_jsonl( + run_root / "inference_set.jsonl", + [ + { + "type": "prompt", + "test_case_id": "p1", + "stop_reason": "completed", + "events": [ + {"edit": {"type": "add_message"}}, + {"edit": {"type": "tool_call"}}, + ], + }, + { + "type": "prompt", + "test_case_id": "p2", + "stop_reason": "target_error", + "events": [], + }, + { + "type": "scenario", + "test_case_id": "s1", + "stop_reason": "completed", + "events": [{"edit": {"type": "add_message"}}], + }, + ], + ) + _write_jsonl( + run_root / "scores.jsonl", + [ + _score( + "p1", + violation=index == 0, + severity=1 if index == 0 else 2, + ), + _score("p2", judge_status="judge_failed"), + _score("s1", kind="scenario", severity=2), + ], + ) + return results_root + + +def test_catalogs_rebuild_legacy_once_and_then_remain_metadata_only() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + results_root = _build_legacy_fixture(root, second_run=True) + repository = ResultRepository(results_root, default_page_size=1) + + first = repository.list_suite_catalog_entries() + assert first.items[0]["suite_id"] == "suite-a" + assert first.items[0]["prompt_test_case_count"] == 2 + assert (results_root / "suite-a" / "suite_summary.json").exists() + assert ( + results_root / "suite-a" / "run-a" / "run_summary.json" + ).exists() + + with patch( + "assert_ai.services.results.scan_jsonl", + side_effect=AssertionError("catalog listing must not scan JSONL"), + ): + suites = repository.list_suite_catalog_entries() + runs = repository.list_run_catalog_entries("suite-a") + + assert len(suites.items) == 1 + assert len(runs.items) == 1 + assert runs.next_cursor is not None + second_page = repository.list_run_catalog_entries( + "suite-a", + cursor=runs.next_cursor, + ) + assert len(second_page.items) == 1 + + +def test_test_case_pagination_filters_and_stale_cursor() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + repository = ResultRepository(results_root, default_page_size=1) + repository.get_suite("suite-a") + + first = repository.list_test_cases( + "suite-a", + kind="prompt", + factors={"region": "us"}, + ) + assert [item["test_case_id"] for item in first.items] == ["p1"] + assert first.next_cursor is None + + page = repository.list_test_cases("suite-a") + assert page.next_cursor is not None + source = ( + results_root + / "suite-a" + / "artifacts" + / "test_set" + / "v0002" + / "test_set.jsonl" + ) + with source.open("a", encoding="utf-8") as handle: + handle.write( + json.dumps({"type": "prompt", "test_case_id": "p3"}) + "\n" + ) + + with pytest.raises(ServiceError) as stale: + repository.list_test_cases( + "suite-a", + cursor=page.next_cursor, + ) + assert stale.value.code == ServiceErrorCode.STALE_CURSOR + + +def test_score_queries_failures_and_transcript_use_indexes() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + repository = ResultRepository(results_root) + repository.get_suite("suite-a") + + filtered = repository.list_scores( + "suite-a", + "run-a", + stop_reason="completed", + has_tool_use=True, + ) + assert [row["test_case_id"] for row in filtered.items] == ["p1"] + + failures = repository.list_failures("suite-a", "run-a") + assert {row["test_case_id"] for row in failures.items} == {"p1", "p2"} + + with patch( + "assert_ai.services.results.scan_jsonl", + side_effect=AssertionError("single transcript lookup must use indexes"), + ): + transcript = repository.get_transcript( + "suite-a", + "run-a", + "p1", + kind="prompt", + ) + assert transcript["test_case"]["seed"]["prompt"] == "one" + assert transcript["inference"]["events"][1]["edit"]["type"] == "tool_call" + assert transcript["score"]["verdict"]["dimensions"]["policy_violation"] is True + + +def test_compare_runs_reports_binary_ordinal_structural_and_sample_warnings() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp), second_run=True) + repository = ResultRepository(results_root) + repository.get_suite("suite-a") + + comparison = repository.compare_runs( + [ + RunReference("suite-a", "run-a"), + RunReference("suite-a", "run-b"), + ] + ) + + policy = comparison["dimension_deltas"]["policy_violation"] + assert policy["kind"] == "binary" + assert policy["runs"][1]["prompt_rate_delta"] == -1.0 + severity = comparison["dimension_deltas"]["severity"] + assert severity["kind"] == "ordinal" + assert ( + severity["runs"][1]["prompt_distribution_delta"]["2"] + == 1.0 + ) + assert comparison["runs"][0]["structural"]["tool_events"] == 1 + assert comparison["behavior_category_deltas"][0]["delta"] == -1.0 + + +def test_repository_rejects_path_traversal_and_invalid_page_sizes() -> None: + with TemporaryDirectory() as tmp: + repository = ResultRepository(Path(tmp) / "results") + + with pytest.raises(ServiceError) as traversal: + repository.get_suite("../escape") + assert traversal.value.code == ServiceErrorCode.INVALID_ARGUMENT + + with pytest.raises(ServiceError) as page_size: + repository.list_suite_catalog_entries(page_size=1000) + assert page_size.value.code == ServiceErrorCode.INVALID_ARGUMENT + + +def test_changed_score_source_rebuilds_stale_run_summary() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + repository = ResultRepository(results_root) + original = repository.load_run_detail("suite-a", "run-a") + assert ( + original["quality"]["prompt"]["dimensions"]["policy_violation"]["rate"] + == 1.0 + ) + + scores_path = results_root / "suite-a" / "run-a" / "scores.jsonl" + _write_jsonl( + scores_path, + [ + _score("p1", violation=False), + _score("p2", violation=False), + _score("s1", kind="scenario", severity=2), + ], + ) + + rebuilt = repository.load_run_detail("suite-a", "run-a") + assert ( + rebuilt["quality"]["prompt"]["dimensions"]["policy_violation"]["rate"] + == 0.0 + ) + + +def test_partial_trailing_run_row_does_not_break_catalog_or_complete_row_queries() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + repository = ResultRepository(results_root) + repository.get_suite("suite-a") + scores_path = results_root / "suite-a" / "run-a" / "scores.jsonl" + with scores_path.open("ab") as handle: + handle.write(b'{"type":"prompt","test_case_id":"partial"') + + runs = repository.list_run_catalog_entries("suite-a") + scores = repository.list_scores("suite-a", "run-a") + detail = repository.load_run_detail( + "suite-a", + "run-a", + include_rows=True, + ) + + assert len(runs.items) == 1 + assert {row["test_case_id"] for row in scores.items} == { + "p1", + "p2", + "s1", + } + assert { + row["test_case_id"] + for row in detail["prompt_rows"] + detail["scenario_rows"] + } == {"p1", "p2", "s1"} + + +def test_suite_summary_detects_out_of_band_run_addition() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + repository = ResultRepository(results_root) + assert repository.get_suite("suite-a")["run_count"] == 1 + + suite_root = results_root / "suite-a" + shutil.copytree(suite_root / "run-a", suite_root / "run-b") + (suite_root / "run-b" / "run_summary.json").unlink() + + assert repository.get_suite("suite-a")["run_count"] == 2 + + +def test_oversized_page_row_returns_bounded_stub_and_remains_pageable() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + repository = ResultRepository( + results_root, + default_page_size=1, + max_page_bytes=250, + max_item_bytes=10_000, + ) + repository.get_suite("suite-a") + + first = repository.list_scores("suite-a", "run-a") + + assert first.items[0]["content_omitted"] is True + assert first.items[0]["test_case_id"] == "p1" + assert first.next_cursor is not None + second = repository.list_scores( + "suite-a", + "run-a", + cursor=first.next_cursor, + ) + assert second.items[0]["test_case_id"] == "p2" + + +def test_compare_negative_behavior_limit_preserves_unlimited_compatibility() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp), second_run=True) + repository = ResultRepository(results_root) + repository.get_suite("suite-a") + + comparison = repository.compare_runs( + [ + RunReference("suite-a", "run-a"), + RunReference("suite-a", "run-b"), + ], + behavior_limit=-1, + ) + + assert len(comparison["behavior_category_deltas"]) == 1 + + +def test_corrupt_derived_summaries_are_rebuilt_from_canonical_artifacts() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + suite_root = results_root / "suite-a" + run_root = suite_root / "run-a" + (suite_root / "suite_summary.json").write_text("{", encoding="utf-8") + (run_root / "run_summary.json").write_text("{", encoding="utf-8") + + repository = ResultRepository(results_root) + suite = repository.get_suite("suite-a") + run = repository.load_run_detail("suite-a", "run-a") + + assert suite["run_count"] == 1 + assert run["state"] == "completed" + + +def test_cli_results_list_uses_metadata_and_compare_supports_ordinal() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp), second_run=True) + repository = ResultRepository(results_root) + repository.get_suite("suite-a") + runner = CliRunner() + + with patch( + "assert_ai.services.results.scan_jsonl", + side_effect=AssertionError("CLI list must not scan JSONL"), + ): + listed = runner.invoke( + cli, + [ + "results", + "list", + "--results-dir", + str(results_root), + "--json", + ], + ) + assert listed.exit_code == 0, listed.output + list_payload = json.loads(listed.output) + assert list_payload["suites"][0]["suite_id"] == "suite-a" + assert list_payload["suites"][0]["runs"] == [] + + compared = runner.invoke( + cli, + [ + "results", + "compare", + "suite-a", + "run-a", + "run-b", + "--results-dir", + str(results_root), + "--metric", + "severity", + "--json", + ], + ) + assert compared.exit_code == 0, compared.output + comparison = json.loads(compared.output) + assert comparison["dimension_deltas"]["severity"]["kind"] == "ordinal" diff --git a/tests/test_result_service_scale.py b/tests/test_result_service_scale.py new file mode 100644 index 000000000..4839f2519 --- /dev/null +++ b/tests/test_result_service_scale.py @@ -0,0 +1,70 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import time +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +from assert_ai.services.results import ResultRepository +from tests.result_catalog_fixture import create_result_catalog_fixture + + +def test_large_catalog_is_metadata_only_and_single_case_lookup_is_indexed() -> None: + with TemporaryDirectory() as tmp: + fixture = create_result_catalog_fixture(Path(tmp)) + repository = ResultRepository( + fixture.results_root, + default_page_size=100, + max_page_size=200, + ) + original_open = Path.open + + def reject_score_reads( + path: Path, + mode: str = "r", + *args: object, + **kwargs: object, + ): + if path.name == "scores.jsonl" and "r" in mode: + raise AssertionError("catalog listing opened score rows") + return original_open(path, mode, *args, **kwargs) + + started = time.perf_counter() + with ( + patch.object(Path, "open", reject_score_reads), + patch( + "assert_ai.services.results.scan_jsonl", + side_effect=AssertionError("catalog listing scanned JSONL"), + ), + ): + suites = repository.list_suite_catalog_entries(page_size=100) + run_count = 0 + for suite in suites.items: + runs = repository.list_run_catalog_entries( + str(suite["suite_id"]), + page_size=20, + ) + run_count += len(runs.items) + assert runs.next_cursor is None + elapsed = time.perf_counter() - started + + assert len(suites.items) == 100 + assert suites.next_cursor is None + assert run_count == 1_000 + assert elapsed < 20.0 + + with patch( + "assert_ai.services.results.scan_jsonl", + side_effect=AssertionError("indexed lookup scanned test_set.jsonl"), + ): + test_case = repository.get_test_case( + fixture.large_suite_id, + fixture.last_test_case_id, + kind="prompt", + run_id=fixture.large_run_id, + ) + + assert test_case["test_case_id"] == fixture.last_test_case_id diff --git a/tests/test_run_metadata.py b/tests/test_run_metadata.py index 1da0755e9..e05c005f3 100644 --- a/tests/test_run_metadata.py +++ b/tests/test_run_metadata.py @@ -242,6 +242,37 @@ async def fake_run_judge(**_: object) -> dict[str, str]: saved_config = root / "results" / "suite-a" / "run-a" / "config.yaml" self.assertTrue(saved_config.exists()) self.assertEqual(saved_config.read_text(encoding="utf-8"), cfg_path.read_text(encoding="utf-8")) + run_summary = json.loads( + ( + root + / "results" + / "suite-a" + / "run-a" + / "run_summary.json" + ).read_text(encoding="utf-8") + ) + suite_summary = json.loads( + ( + root + / "results" + / "suite-a" + / "suite_summary.json" + ).read_text(encoding="utf-8") + ) + self.assertEqual(run_summary["state"], "completed") + self.assertEqual(run_summary["current_stage"], "judge") + self.assertEqual(run_summary["counts"]["scores"]["total"], 0) + self.assertTrue( + ( + root + / "results" + / "suite-a" + / "run-a" + / "scores.index.json" + ).exists() + ) + self.assertEqual(suite_summary["run_count"], 1) + self.assertEqual(suite_summary["latest_run"]["run_id"], "run-a") def test_run_pipeline_records_pid_host_and_heartbeat(self) -> None: """Manifest carries pid/host/heartbeat_at so the viewer can detect abandoned runs.""" diff --git a/tests/test_runner_artifact_cache.py b/tests/test_runner_artifact_cache.py index f33f6ae66..99c0b2d99 100644 --- a/tests/test_runner_artifact_cache.py +++ b/tests/test_runner_artifact_cache.py @@ -136,6 +136,28 @@ def test_identical_inputs_reuse_systematize_and_test_set(self) -> None: self.assertEqual(seen, ["systematize", "test_set"]) self.assertTrue((root / "results" / "suite-a" / "artifacts" / "systematize" / "v0001" / "taxonomy.json").exists()) self.assertFalse((root / "results" / "suite-a" / "artifacts" / "systematize" / "v0002").exists()) + test_set_index = ( + root + / "results" + / "suite-a" + / "artifacts" + / "test_set" + / "v0001" + / "test_set.index.json" + ) + self.assertTrue(test_set_index.exists()) + suite_summary = json.loads( + ( + root + / "results" + / "suite-a" + / "suite_summary.json" + ).read_text(encoding="utf-8") + ) + self.assertEqual( + suite_summary["sources"]["test_set"]["path"], + "artifacts/test_set/v0001/test_set.jsonl", + ) def test_behavior_change_regenerates_all_upstream_artifacts(self) -> None: with TemporaryDirectory() as tmp_dir: @@ -498,6 +520,12 @@ async def partial_test_set(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict if latest_path.exists(): latest = json.loads(latest_path.read_text(encoding="utf-8")) self.assertNotIn("test_set", latest.get("artifacts", {})) + suite_summary = json.loads( + (root / "results" / "suite-a" / "suite_summary.json").read_text( + encoding="utf-8" + ) + ) + self.assertNotIn("test_set", suite_summary.get("sources", {})) # Re-run with identical inputs but a fully-successful test_set # stage. The runner must NOT reuse the partial v0001 -- it From ef8c75ff69b207b002d30bcfa9bf4a6a1434a2c2 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Tue, 25 Aug 2026 08:16:16 -0700 Subject: [PATCH 06/16] Expose read-only MCP inspection surface Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/core/security.py | 57 +- assert_ai/core/session.py | 24 +- assert_ai/mcp/_command.py | 54 ++ assert_ai/mcp/errors.py | 161 +++++ assert_ai/mcp/models.py | 245 ++++++- assert_ai/mcp/resources.py | 379 +++++++++++ assert_ai/mcp/sanitize.py | 56 ++ assert_ai/mcp/server.py | 95 +++ assert_ai/mcp/tools/__init__.py | 8 + assert_ai/mcp/tools/inspect.py | 649 ++++++++++++++++++ assert_ai/mcp/uris.py | 75 +++ assert_ai/services/artifacts.py | 1104 +++++++++++++++++++++++++++++++ assert_ai/services/library.py | 239 +++++++ tests/test_artifact_service.py | 216 ++++++ tests/test_library_service.py | 60 ++ tests/test_mcp_cli.py | 18 + tests/test_mcp_server.py | 664 ++++++++++++++++++- tests/test_security.py | 10 + 18 files changed, 4064 insertions(+), 50 deletions(-) create mode 100644 assert_ai/mcp/errors.py create mode 100644 assert_ai/mcp/resources.py create mode 100644 assert_ai/mcp/sanitize.py create mode 100644 assert_ai/mcp/tools/__init__.py create mode 100644 assert_ai/mcp/tools/inspect.py create mode 100644 assert_ai/mcp/uris.py create mode 100644 assert_ai/services/artifacts.py create mode 100644 assert_ai/services/library.py create mode 100644 tests/test_artifact_service.py create mode 100644 tests/test_library_service.py diff --git a/assert_ai/core/security.py b/assert_ai/core/security.py index ea536c479..e49cc7b79 100644 --- a/assert_ai/core/security.py +++ b/assert_ai/core/security.py @@ -16,7 +16,7 @@ import socket import sys from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Iterable from urllib.parse import urlparse if TYPE_CHECKING: @@ -319,9 +319,60 @@ def validate_resolved_endpoint_ip(hostname: str, ip_str: str) -> None: re.IGNORECASE, ) +_CREDENTIAL_TEXT_PATTERNS = re.compile( + r"(" + r"Bearer\s+[A-Za-z0-9\-._~+/]+=*" + r"|Basic\s+[A-Za-z0-9+/]+=*" + r"|(?:sk|pk|api|key|token|secret)[-_][A-Za-z0-9\-._]{20,}" + r"|(?:api[_-]?key|auth[_-]?token|secret|password|access[_-]?token|refresh[_-]?token" + r"|client[_-]?secret|authorization)[\"':\s=]+[A-Za-z0-9\-._~+/]{16,}" + r")", + re.IGNORECASE, +) + +_SENSITIVE_TEXT_ASSIGNMENT = re.compile( + r"(?P" + r"[\"']?(?:api[_-]?key|auth[_-]?token|secret|password|credential|" + r"access[_-]?token|refresh[_-]?token|private[_-]?key|client[_-]?secret|" + r"authorization|azure[_-]?ad[_-]?token)[\"']?\s*[:=]\s*" + r")" + r"(?P\"[^\"\r\n]*\"|'[^'\r\n]*'|[^\s,}\]\r\n]+)", + re.IGNORECASE, +) + _REDACTED = "[REDACTED]" +def sanitize_text(text: str) -> str: + """Redact credential-like values embedded in plain text.""" + if not text: + return text + sanitized = _SENSITIVE_TEXT_ASSIGNMENT.sub( + lambda match: f'{match.group("prefix")}"{_REDACTED}"', + text, + ) + return _CREDENTIAL_TEXT_PATTERNS.sub(_REDACTED, sanitized) + + +def redact_path_prefixes(text: str, paths: Iterable[str | Path]) -> str: + """Replace configured path prefixes, including JSON-escaped variants.""" + if not text: + return text + variants: set[str] = set() + for path in paths: + raw = str(path) + variants.update((raw, Path(raw).as_posix(), raw.replace("\\", "\\\\"))) + flags = re.IGNORECASE if os.name == "nt" else 0 + redacted = text + for variant in sorted( + (value for value in variants if value), + key=len, + reverse=True, + ): + redacted = re.sub(re.escape(variant), ".", redacted, flags=flags) + return redacted + + def sanitize_payload(payload: Any, *, depth: int = 0, max_depth: int = 10) -> Any: """Recursively sanitize sensitive fields from a payload before writing to artifacts. @@ -346,7 +397,5 @@ def sanitize_payload(payload: Any, *, depth: int = 0, max_depth: int = 10) -> An return [sanitize_payload(item, depth=depth + 1, max_depth=max_depth) for item in payload] elif isinstance(payload, str): # Redact Bearer tokens in string values - if payload.startswith("Bearer ") or payload.startswith("Basic "): - return _REDACTED - return payload + return sanitize_text(payload) return payload diff --git a/assert_ai/core/session.py b/assert_ai/core/session.py index be9ca25dc..57a2d2bb7 100644 --- a/assert_ai/core/session.py +++ b/assert_ai/core/session.py @@ -10,7 +10,6 @@ import inspect import json import logging -import re from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any, Literal @@ -28,6 +27,7 @@ normalize_response, summarize_response, ) +from assert_ai.core.security import sanitize_text from assert_ai.core.tool_backend import load_tool_module from assert_ai.core.tools import build_target_tools @@ -36,29 +36,10 @@ log = logging.getLogger(__name__) -# Regex patterns for common credential formats in plain text -_CREDENTIAL_PATTERNS = re.compile( - r"(" - # Bearer/Basic tokens - r"Bearer\s+[A-Za-z0-9\-._~+/]+=*" - r"|Basic\s+[A-Za-z0-9+/]+=*" - # Common API key formats (sk-..., key-..., etc.) - r"|(?:sk|pk|api|key|token|secret)[-_][A-Za-z0-9\-._]{20,}" - # Generic long hex/base64 secrets following key-like prefixes - r"|(?:api[_-]?key|auth[_-]?token|secret|password|access[_-]?token|refresh[_-]?token" - r"|client[_-]?secret|authorization)[\"':\s=]+[A-Za-z0-9\-._~+/]{16,}" - r")", - re.IGNORECASE, -) - -_RESPONSE_REDACTED = "[REDACTED]" - def _sanitize_response_text(text: str) -> str: """Redact credential-like patterns from response text before persisting.""" - if not text: - return text - sanitized = _CREDENTIAL_PATTERNS.sub(_RESPONSE_REDACTED, text) + sanitized = sanitize_text(text) if sanitized != text: log.warning( "Credential-like patterns detected and redacted from HTTP endpoint response" @@ -68,6 +49,7 @@ def _sanitize_response_text(text: str) -> str: # ── Adapter types and helpers ────────────────────────────────── + @dataclass class AdapterEvent: role: Literal["assistant", "tool_call", "tool_result"] diff --git a/assert_ai/mcp/_command.py b/assert_ai/mcp/_command.py index 19b8ed2e5..6ec627329 100644 --- a/assert_ai/mcp/_command.py +++ b/assert_ai/mcp/_command.py @@ -74,11 +74,59 @@ def mcp() -> None: default=None, help="Optional dotenv file contained within --workspace. No file is discovered by default.", ) +@click.option( + "--default-page-size", + type=click.IntRange(min=1), + default=50, + show_default=True, + help="Default number of items returned by paginated inspect tools.", +) +@click.option( + "--max-page-size", + type=click.IntRange(min=1), + default=200, + show_default=True, + help="Maximum number of items accepted by paginated inspect tools.", +) +@click.option( + "--max-response-bytes", + type=click.IntRange(min=4096), + default=1024 * 1024, + show_default=True, + help="Maximum serialized bytes returned by one tool or resource.", +) +@click.option( + "--default-artifact-chunk-bytes", + type=click.IntRange(min=4), + default=64 * 1024, + show_default=True, + help="Default source-byte budget for artifact chunk reads.", +) +@click.option( + "--max-artifact-chunk-bytes", + type=click.IntRange(min=4), + default=256 * 1024, + show_default=True, + help="Maximum source-byte budget for one artifact chunk read.", +) +@click.option( + "--max-config-bytes", + type=click.IntRange(min=1), + default=256 * 1024, + show_default=True, + help="Maximum size of one managed config payload.", +) def serve( workspace: Path, mode: str, enabled_groups: tuple[str, ...], env_file: Path | None, + default_page_size: int, + max_page_size: int, + max_response_bytes: int, + default_artifact_chunk_bytes: int, + max_artifact_chunk_bytes: int, + max_config_bytes: int, ) -> None: """Serve ASSERT over stdio; stdout is reserved for MCP protocol traffic.""" try: @@ -98,6 +146,12 @@ def serve( workspace_root=workspace_service.root, mode=mode, enabled_groups=enabled_groups, + default_page_size=default_page_size, + max_page_size=max_page_size, + max_response_bytes=max_response_bytes, + default_artifact_chunk_bytes=default_artifact_chunk_bytes, + max_artifact_chunk_bytes=max_artifact_chunk_bytes, + max_config_bytes=max_config_bytes, ) except (OSError, ValueError) as exc: raise click.ClickException(str(exc)) from exc diff --git a/assert_ai/mcp/errors.py b/assert_ai/mcp/errors.py new file mode 100644 index 000000000..a8d9ce3db --- /dev/null +++ b/assert_ai/mcp/errors.py @@ -0,0 +1,161 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Stable application-error adaptation for MCP tools and resources.""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Callable +from functools import wraps +from typing import TypeVar +from uuid import uuid4 + +from mcp.server.mcpserver.exceptions import ResourceError, ResourceNotFoundError +from pydantic import BaseModel + +from assert_ai.core.workspace import WorkspaceService +from assert_ai.mcp.sanitize import sanitize_for_mcp +from assert_ai.services.errors import ServiceError, ServiceErrorCode + +log = logging.getLogger(__name__) + +_T = TypeVar("_T") + + +class _McpToolError(RuntimeError): + """Expected sanitized error that should survive nested adaptation.""" + + +def invoke_tool( + operation: Callable[[], _T], + *, + workspace: WorkspaceService, +) -> _T: + """Invoke one service operation with stable, sanitized tool errors.""" + try: + return operation() + except _McpToolError: + raise + except ServiceError as exc: + raise _McpToolError( + _service_error_payload(exc, workspace=workspace) + ) from exc + except Exception as exc: # noqa: BLE001 + correlation_id = uuid4().hex + log.exception("Unhandled MCP tool failure (%s)", correlation_id) + internal = ServiceError( + ServiceErrorCode.INTERNAL, + "An internal error occurred", + details={"correlation_id": correlation_id}, + ) + raise _McpToolError( + _service_error_payload(internal, workspace=workspace) + ) from exc + + +def adapt_tool_errors( + workspace: WorkspaceService, + *, + max_response_bytes: int | None = None, +) -> Callable[[Callable[..., _T]], Callable[..., _T]]: + """Decorate a complete tool body so adapter-side failures are sanitized.""" + + def decorator(operation: Callable[..., _T]) -> Callable[..., _T]: + @wraps(operation) + def wrapper(*args: object, **kwargs: object) -> _T: + return invoke_tool( + lambda: _bounded_result( + operation(*args, **kwargs), + max_response_bytes=max_response_bytes, + ), + workspace=workspace, + ) + + return wrapper + + return decorator + + +def _bounded_result( + result: _T, + *, + max_response_bytes: int | None, +) -> _T: + if max_response_bytes is None: + return result + payload = ( + result.model_dump(mode="json") + if isinstance(result, BaseModel) + else result + ) + payload_size_bytes = len( + json.dumps( + payload, + ensure_ascii=False, + separators=(",", ":"), + default=str, + ).encode("utf-8") + ) + estimated_wire_bytes = payload_size_bytes * 2 + 2048 + if estimated_wire_bytes > max_response_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + ( + "Tool response exceeds the configured response limit; " + "use a narrower query or a smaller page" + ), + details={ + "payload_size_bytes": payload_size_bytes, + "estimated_wire_bytes": estimated_wire_bytes, + "max_response_bytes": max_response_bytes, + }, + ) + return result + + +def invoke_resource( + operation: Callable[[], _T], + *, + workspace: WorkspaceService, +) -> _T: + """Invoke one service operation with resource-appropriate errors.""" + try: + return operation() + except ServiceError as exc: + error_type = ( + ResourceNotFoundError + if exc.code is ServiceErrorCode.NOT_FOUND + else ResourceError + ) + raise error_type(_service_error_payload(exc, workspace=workspace)) from exc + except Exception as exc: # noqa: BLE001 + correlation_id = uuid4().hex + log.exception("Unhandled MCP resource failure (%s)", correlation_id) + internal = ServiceError( + ServiceErrorCode.INTERNAL, + "An internal error occurred", + details={"correlation_id": correlation_id}, + ) + raise ResourceError( + _service_error_payload(internal, workspace=workspace) + ) from exc + + +def _service_error_payload( + error: ServiceError, + *, + workspace: WorkspaceService, +) -> str: + payload = { + "code": error.code.value, + "message": str(error), + "details": error.details, + } + return json.dumps( + sanitize_for_mcp(payload, workspace=workspace), + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ) diff --git a/assert_ai/mcp/models.py b/assert_ai/mcp/models.py index 11036862b..10c1f2530 100644 --- a/assert_ai/mcp/models.py +++ b/assert_ai/mcp/models.py @@ -6,11 +6,13 @@ from __future__ import annotations from enum import StrEnum -from typing import Literal +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field +from assert_ai.core.config_document import ConfigValidationReport from assert_ai.mcp import ASSERT_MCP_API_VERSION +from assert_ai.services.library import PresetKind class ServerMode(StrEnum): @@ -47,6 +49,19 @@ class WorkspaceInfo(BaseModel): results_root: str = "artifacts/results" +class ServerLimits(BaseModel): + """Response limits fixed when the MCP server starts.""" + + model_config = ConfigDict(frozen=True) + + default_page_size: int + max_page_size: int + max_response_bytes: int + default_artifact_chunk_bytes: int + max_artifact_chunk_bytes: int + max_config_bytes: int + + class ServerInfo(BaseModel): """Discovery metadata returned by ``get_server_info``.""" @@ -58,4 +73,232 @@ class ServerInfo(BaseModel): mode: ServerMode enabled_capability_groups: list[CapabilityGroup] workspace: WorkspaceInfo = Field(default_factory=WorkspaceInfo) + limits: ServerLimits + target_kinds: list[Literal["callable", "model"]] = Field( + default_factory=lambda: ["callable", "model"] + ) transports: list[Literal["stdio"]] = Field(default_factory=lambda: ["stdio"]) + protocol_notes: list[str] = Field( + default_factory=lambda: [ + "Long evaluations use ASSERT job polling rather than one blocking MCP request.", + "Resource and artifact identifiers are opaque and contain no host paths.", + ] + ) + + +class _McpModel(BaseModel): + model_config = ConfigDict(frozen=True) + + +class PresetCatalogItem(_McpModel): + """One built-in preset catalog entry.""" + + kind: PresetKind + name: str + version: str | None = None + tags: tuple[str, ...] = () + summary: str | None = None + description: str | None = None + resource_uri: str + + +class PresetCatalogPage(_McpModel): + """Bounded page of built-in presets.""" + + items: tuple[PresetCatalogItem, ...] + next_cursor: str | None = None + + +class PresetResult(_McpModel): + """Complete built-in preset definition.""" + + kind: PresetKind + name: str + version: str | None = None + tags: tuple[str, ...] = () + yaml: str + document: dict[str, Any] + resource_uri: str + + +class ConfigSchemaResult(_McpModel): + """Canonical machine-readable evaluation config schema.""" + + schema_version: int + json_schema: dict[str, Any] + resource_uri: Literal["assert://schema/eval-config"] = ( + "assert://schema/eval-config" + ) + + +class ConfigCatalogItem(_McpModel): + """One managed config catalog entry.""" + + config_ref: str + etag: str + size_bytes: int + modified_at: str + structurally_valid: bool + resource_uri: str + + +class ConfigCatalogPage(_McpModel): + """Bounded page of managed configs.""" + + items: tuple[ConfigCatalogItem, ...] + next_cursor: str | None = None + + +class ConfigResult(_McpModel): + """One sanitized managed evaluation config.""" + + config_ref: str + yaml: str + document: dict[str, Any] + etag: str + validation: ConfigValidationReport + resource_uri: str + + +class SuiteCatalogItem(_McpModel): + """Lightweight suite metadata.""" + + suite_id: str + status: str | None = None + behavior_name: str | None = None + behavior_category_count: int = 0 + prompt_test_case_count: int = 0 + scenario_test_case_count: int = 0 + run_count: int = 0 + created_at: str | None = None + updated_at: str | None = None + latest_run: dict[str, Any] | None = None + resources: dict[str, str] = Field(default_factory=dict) + + +class SuiteCatalogPage(_McpModel): + """Bounded page of suite metadata.""" + + items: tuple[SuiteCatalogItem, ...] + next_cursor: str | None = None + + +class SuiteResult(_McpModel): + """One suite's metadata-only detail.""" + + schema_version: int + suite_id: str + status: str + behavior: dict[str, Any] + behavior_category_count: int + test_case_counts: dict[str, Any] + created_at: str | None = None + updated_at: str | None = None + run_count: int + latest_run: dict[str, Any] | None = None + resources: dict[str, str] = Field(default_factory=dict) + + +class RunCatalogItem(_McpModel): + """Lightweight run metadata.""" + + suite_id: str + run_id: str + status: str | None = None + current_stage: str | None = None + started_at: str | None = None + ended_at: str | None = None + updated_at: str | None = None + prompt_metrics: dict[str, Any] | None = None + scenario_metrics: dict[str, Any] | None = None + models: dict[str, Any] = Field(default_factory=dict) + counts: dict[str, Any] = Field(default_factory=dict) + metrics: dict[str, Any] | None = None + resources: dict[str, str] = Field(default_factory=dict) + + +class RunCatalogPage(_McpModel): + """Bounded page of run metadata.""" + + items: tuple[RunCatalogItem, ...] + next_cursor: str | None = None + + +class RunResult(_McpModel): + """One run's metadata-only detail.""" + + schema_version: int + suite_id: str + run_id: str + state: str + current_stage: str | None = None + started_at: str | None = None + ended_at: str | None = None + updated_at: str | None = None + stages: dict[str, Any] = Field(default_factory=dict) + stage_timings: dict[str, Any] = Field(default_factory=dict) + stage_summaries: dict[str, Any] = Field(default_factory=dict) + models: dict[str, Any] = Field(default_factory=dict) + counts: dict[str, Any] = Field(default_factory=dict) + quality: dict[str, Any] | None = None + metrics: dict[str, Any] | None = None + resources: dict[str, str] = Field(default_factory=dict) + + +class RunReferenceInput(_McpModel): + """Suite/run pair accepted by comparison tools.""" + + suite_id: str + run_id: str + + +class RunComparisonResult(_McpModel): + """Structured comparison across two or more runs.""" + + metric: str + baseline: str + runs: list[dict[str, Any]] + dimension_deltas: dict[str, dict[str, Any]] + behavior_category_deltas: list[dict[str, Any]] + warnings: list[str] + + +class TestCasePage(_McpModel): + """Bounded page of test-case rows.""" + + items: list[dict[str, Any]] + next_cursor: str | None = None + + +class TestCaseResult(_McpModel): + """One complete test case.""" + + row: dict[str, Any] + resource_uri: str + + +class ScorePage(_McpModel): + """Bounded page of score rows.""" + + items: list[dict[str, Any]] + next_cursor: str | None = None + + +class FailurePage(_McpModel): + """Bounded page of failure rows.""" + + items: list[dict[str, Any]] + next_cursor: str | None = None + + +class TranscriptResult(_McpModel): + """One inference transcript joined with its test case and score.""" + + suite_id: str + run_id: str + type: str + test_case_id: str + test_case: dict[str, Any] | None = None + inference: dict[str, Any] + score: dict[str, Any] | None = None + resource_uri: str diff --git a/assert_ai/mcp/resources.py b/assert_ai/mcp/resources.py new file mode 100644 index 000000000..ff6e9cefc --- /dev/null +++ b/assert_ai/mcp/resources.py @@ -0,0 +1,379 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Lazy ASSERT MCP resources backed by inspect application services.""" + +from __future__ import annotations + +import json +from typing import Any + +import yaml +from mcp.server import MCPServer + +from assert_ai.core.config_document import EVAL_CONFIG_SCHEMA_VERSION +from assert_ai.mcp.errors import invoke_resource +from assert_ai.mcp.sanitize import sanitize_for_mcp +from assert_ai.mcp.tools.inspect import InspectServices +from assert_ai.services.errors import ServiceError, ServiceErrorCode + +_SCHEMA_URI = "assert://schema/eval-config" + + +def register_inspect_resources( + server: MCPServer, + services: InspectServices, + *, + inline_artifact_bytes: int, +) -> None: + """Register static and templated resources for the inspect group.""" + + workspace = services.workspace + + @server.resource( + _SCHEMA_URI, + name="eval-config-schema", + title="ASSERT eval config schema", + description="Canonical Draft 2020-12 JSON Schema for eval_config.yaml.", + mime_type="application/json", + ) + def eval_config_schema() -> str: + return invoke_resource( + lambda: _json_text( + { + "schema_version": EVAL_CONFIG_SCHEMA_VERSION, + "json_schema": services.configs.get_schema(), + }, + services=services, + ), + workspace=workspace, + ) + + @server.resource( + "assert://preset/{kind}/{name}", + name="preset", + title="ASSERT preset", + description="One built-in behavior or judge preset definition.", + mime_type="application/json", + ) + def preset(kind: str, name: str) -> str: + return invoke_resource( + lambda: _json_text( + services.library.get_preset(kind, name).document, + services=services, + ), + workspace=workspace, + ) + + @server.resource( + "assert://config/{config_ref}", + name="config", + title="ASSERT managed config", + description="One sanitized workspace-managed evaluation config.", + mime_type="application/yaml", + ) + def config(config_ref: str) -> str: + return invoke_resource( + lambda: _sanitized_config_yaml(config_ref, services=services), + workspace=workspace, + ) + + @server.resource( + "assert://suite/{suite_id}/taxonomy", + name="suite-taxonomy", + title="ASSERT suite taxonomy", + description="The active behavior taxonomy for one result suite.", + mime_type="application/json", + ) + def suite_taxonomy(suite_id: str) -> str: + return invoke_resource( + lambda: _named_artifact_resource( + suite_id, + "taxonomy", + services=services, + inline_artifact_bytes=inline_artifact_bytes, + ), + workspace=workspace, + ) + + @server.resource( + "assert://suite/{suite_id}/test-case/{test_case_id}{?kind,run_id}", + name="suite-test-case", + title="ASSERT test case", + description="One complete active suite test case.", + mime_type="application/json", + ) + def suite_test_case( + suite_id: str, + test_case_id: str, + kind: str | None = None, + run_id: str | None = None, + ) -> str: + return invoke_resource( + lambda: _json_text( + services.results.get_test_case( + suite_id, + test_case_id, + kind=kind, + run_id=run_id, + ), + services=services, + ), + workspace=workspace, + ) + + @server.resource( + "assert://run/{suite_id}/{run_id}/summary", + name="run-summary", + title="ASSERT run summary", + description="Metadata-only quality, timing, usage, and model summary.", + mime_type="application/json", + ) + def run_summary(suite_id: str, run_id: str) -> str: + return invoke_resource( + lambda: _json_text( + _public_run( + services.results.load_run_detail(suite_id, run_id) + ), + services=services, + ), + workspace=workspace, + ) + + @server.resource( + "assert://run/{suite_id}/{run_id}/manifest", + name="run-manifest", + title="ASSERT run manifest", + description="The persisted stage/status manifest for one run.", + mime_type="application/json", + ) + def run_manifest(suite_id: str, run_id: str) -> str: + return invoke_resource( + lambda: _named_artifact_resource( + suite_id, + "manifest", + run_id=run_id, + services=services, + inline_artifact_bytes=inline_artifact_bytes, + ), + workspace=workspace, + ) + + @server.resource( + "assert://run/{suite_id}/{run_id}/config", + name="run-config", + title="ASSERT run config", + description="The sanitized immutable config snapshot used by one run.", + mime_type="application/yaml", + ) + def run_config(suite_id: str, run_id: str) -> str: + return invoke_resource( + lambda: _named_artifact_resource( + suite_id, + "config", + run_id=run_id, + services=services, + inline_artifact_bytes=inline_artifact_bytes, + ), + workspace=workspace, + ) + + @server.resource( + "assert://run/{suite_id}/{run_id}/transcript/{test_case_id}{?kind}", + name="run-transcript", + title="ASSERT run transcript", + description="One inference transcript joined with its test case and score.", + mime_type="application/json", + ) + def run_transcript( + suite_id: str, + run_id: str, + test_case_id: str, + kind: str | None = None, + ) -> str: + return invoke_resource( + lambda: _json_text( + services.results.get_transcript( + suite_id, + run_id, + test_case_id, + kind=kind, + ), + services=services, + ), + workspace=workspace, + ) + + @server.resource( + "assert://artifact/{artifact_id}", + name="artifact", + title="ASSERT artifact", + description=( + "One small redacted text artifact, or metadata directing the caller " + "to read_artifact_chunk." + ), + mime_type="text/plain", + ) + def artifact(artifact_id: str) -> str: + return invoke_resource( + lambda: _artifact_resource( + artifact_id, + services=services, + inline_artifact_bytes=inline_artifact_bytes, + ), + workspace=workspace, + ) + + +def _sanitized_config_yaml( + config_ref: str, + *, + services: InspectServices, +) -> str: + record = services.configs.get_config(config_ref) + document = sanitize_for_mcp(record.document, workspace=services.workspace) + if not isinstance(document, dict): + raise TypeError("Expected a config mapping") + text = yaml.safe_dump( + document, + default_flow_style=False, + sort_keys=False, + allow_unicode=True, + ) + normalized = text if text.endswith("\n") else text + "\n" + return _bounded_text(normalized, services=services) + + +def _named_artifact_resource( + suite_id: str, + name: str, + *, + services: InspectServices, + inline_artifact_bytes: int, + run_id: str | None = None, +) -> str: + descriptor = services.artifacts.find_artifact( + suite_id, + name, + run_id=run_id, + ) + return _artifact_resource( + descriptor.artifact_id, + services=services, + inline_artifact_bytes=inline_artifact_bytes, + ) + + +def _artifact_resource( + artifact_id: str, + *, + services: InspectServices, + inline_artifact_bytes: int, +) -> str: + descriptor = services.artifacts.get_artifact(artifact_id) + if not descriptor.text: + return _bounded_text( + _artifact_redirect( + descriptor.model_dump(mode="json"), + readable=False, + reason=( + "Binary artifact reads are disabled because their contents " + "cannot be safely redacted." + ), + ), + services=services, + ) + if descriptor.size_bytes > services.artifacts.max_text_artifact_bytes: + return _bounded_text( + _artifact_redirect( + descriptor.model_dump(mode="json"), + readable=False, + reason=( + "This text artifact exceeds the generic read limit. " + "Use a dedicated paginated result tool when available." + ), + ), + services=services, + ) + if descriptor.size_bytes > inline_artifact_bytes: + return _bounded_text( + _artifact_redirect( + descriptor.model_dump(mode="json"), + readable=True, + ), + services=services, + ) + chunk = services.artifacts.read_artifact_chunk( + artifact_id, + chunk_size=max(4, descriptor.size_bytes), + ) + if not chunk.eof: + return _bounded_text( + _artifact_redirect( + descriptor.model_dump(mode="json"), + readable=True, + ), + services=services, + ) + return _bounded_text(chunk.data, services=services) + + +def _artifact_redirect( + descriptor: dict[str, Any], + *, + readable: bool, + reason: str | None = None, +) -> str: + payload: dict[str, Any] = { + "artifact": descriptor, + "inline": False, + "readable": readable, + } + if readable: + payload["next_step"] = ( + "Call read_artifact_chunk with artifact.artifact_id for bounded access." + ) + if reason is not None: + payload["reason"] = reason + return json.dumps( + payload, + ensure_ascii=False, + indent=2, + sort_keys=True, + ) + + +def _json_text(value: Any, *, services: InspectServices) -> str: + return _bounded_text( + json.dumps( + sanitize_for_mcp(value, workspace=services.workspace), + ensure_ascii=False, + indent=2, + sort_keys=True, + ), + services=services, + ) + + +def _public_run(summary: dict[str, Any]) -> dict[str, Any]: + payload = dict(summary) + for key in ("artifact_versions", "sources", "indexes"): + payload.pop(key, None) + return payload + + +def _bounded_text(text: str, *, services: InspectServices) -> str: + size_bytes = len(text.encode("utf-8")) + if size_bytes > services.max_response_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + ( + "Resource exceeds the configured response limit; " + "use its paginated or chunked tool" + ), + details={ + "size_bytes": size_bytes, + "max_response_bytes": services.max_response_bytes, + }, + ) + return text diff --git a/assert_ai/mcp/sanitize.py b/assert_ai/mcp/sanitize.py new file mode 100644 index 000000000..caf34b88a --- /dev/null +++ b/assert_ai/mcp/sanitize.py @@ -0,0 +1,56 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Customer-safe serialization helpers for MCP responses.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from pydantic import BaseModel + +from assert_ai.core.security import redact_path_prefixes, sanitize_payload +from assert_ai.core.workspace import WorkspaceService + + +def sanitize_for_mcp( + value: Any, + *, + workspace: WorkspaceService, +) -> Any: + """Redact credentials and host paths from one JSON-compatible value.""" + if isinstance(value, BaseModel): + value = value.model_dump(mode="json") + sanitized = sanitize_payload(value) + return _sanitize_paths(sanitized, workspace=workspace) + + +def _sanitize_paths(value: Any, *, workspace: WorkspaceService) -> Any: + if isinstance(value, dict): + return { + str(key): _sanitize_paths(item, workspace=workspace) + for key, item in value.items() + } + if isinstance(value, (list, tuple)): + return [ + _sanitize_paths(item, workspace=workspace) + for item in value + ] + if not isinstance(value, str): + return value + + replaced = redact_path_prefixes(value, (workspace.root,)) + if replaced != value: + return replaced + + try: + candidate = Path(value) + except (OSError, ValueError): + return value + if not candidate.is_absolute(): + return value + try: + return workspace.reference(candidate) + except ValueError: + return "[EXTERNAL_PATH]" diff --git a/assert_ai/mcp/server.py b/assert_ai/mcp/server.py index 6bedd730e..c89555d22 100644 --- a/assert_ai/mcp/server.py +++ b/assert_ai/mcp/server.py @@ -17,10 +17,17 @@ from assert_ai.core.workspace import WorkspaceService from assert_ai.mcp.models import ( CapabilityGroup, + ServerLimits, ServerInfo, ServerMode, WorkspaceInfo, ) +from assert_ai.mcp.resources import register_inspect_resources +from assert_ai.mcp.tools import InspectServices, register_inspect_tools +from assert_ai.services.artifacts import ArtifactRepository +from assert_ai.services.configs import ConfigService +from assert_ai.services.library import LibraryService +from assert_ai.services.results import ResultRepository SERVER_NAME = "ASSERT" @@ -53,9 +60,35 @@ class ServerOptions: workspace_root: Path mode: ServerMode = ServerMode.INSPECT enabled_groups: tuple[CapabilityGroup, ...] = () + default_page_size: int = 50 + max_page_size: int = 200 + max_response_bytes: int = 1024 * 1024 + default_artifact_chunk_bytes: int = 64 * 1024 + max_artifact_chunk_bytes: int = 256 * 1024 + max_config_bytes: int = 256 * 1024 workspace: WorkspaceService = field(init=False, repr=False) def __post_init__(self) -> None: + if self.default_page_size < 1: + raise ValueError("default_page_size must be positive") + if self.max_page_size < self.default_page_size: + raise ValueError("max_page_size must be >= default_page_size") + if self.max_response_bytes < 4096: + raise ValueError("max_response_bytes must be at least 4096") + if self.max_config_bytes < 1: + raise ValueError("max_config_bytes must be positive") + if self.max_config_bytes > self.max_response_bytes: + raise ValueError("max_config_bytes must not exceed max_response_bytes") + if self.default_artifact_chunk_bytes < 4: + raise ValueError("default_artifact_chunk_bytes must be at least 4") + if self.max_artifact_chunk_bytes < self.default_artifact_chunk_bytes: + raise ValueError( + "max_artifact_chunk_bytes must be >= default_artifact_chunk_bytes" + ) + if self.max_artifact_chunk_bytes * 2 > self.max_response_bytes: + raise ValueError( + "max_artifact_chunk_bytes must not exceed half max_response_bytes" + ) workspace = WorkspaceService.create(self.workspace_root) object.__setattr__(self, "workspace_root", workspace.root) object.__setattr__(self, "workspace", workspace) @@ -67,6 +100,12 @@ def create( workspace_root: str | Path, mode: str | ServerMode = ServerMode.INSPECT, enabled_groups: Iterable[str | CapabilityGroup] = (), + default_page_size: int = 50, + max_page_size: int = 200, + max_response_bytes: int = 1024 * 1024, + default_artifact_chunk_bytes: int = 64 * 1024, + max_artifact_chunk_bytes: int = 256 * 1024, + max_config_bytes: int = 256 * 1024, ) -> "ServerOptions": parsed_mode = ServerMode(mode) parsed_groups = tuple(CapabilityGroup(group) for group in enabled_groups) @@ -80,6 +119,12 @@ def create( workspace_root=Path(workspace_root), mode=parsed_mode, enabled_groups=parsed_groups, + default_page_size=default_page_size, + max_page_size=max_page_size, + max_response_bytes=max_response_bytes, + default_artifact_chunk_bytes=default_artifact_chunk_bytes, + max_artifact_chunk_bytes=max_artifact_chunk_bytes, + max_config_bytes=max_config_bytes, ) @property @@ -111,6 +156,8 @@ def build_server(options: ServerOptions) -> MCPServer: title="Get ASSERT server information", annotations=ToolAnnotations( read_only_hint=True, + destructive_hint=False, + idempotent_hint=True, open_world_hint=False, ), structured_output=True, @@ -127,6 +174,54 @@ def get_server_info() -> ServerInfo: artifacts_root=options.workspace.reference(options.workspace.artifacts_root), results_root=options.workspace.reference(options.workspace.results_root), ), + limits=ServerLimits( + default_page_size=options.default_page_size, + max_page_size=options.max_page_size, + max_response_bytes=options.max_response_bytes, + default_artifact_chunk_bytes=options.default_artifact_chunk_bytes, + max_artifact_chunk_bytes=options.max_artifact_chunk_bytes, + max_config_bytes=options.max_config_bytes, + ), + ) + + if CapabilityGroup.INSPECT in options.capability_groups: + results = ResultRepository( + options.workspace.results_root, + path_policy=options.path_policy, + default_page_size=options.default_page_size, + max_page_size=options.max_page_size, + max_page_bytes=options.max_response_bytes, + max_item_bytes=options.max_response_bytes, + ) + services = InspectServices( + workspace=options.workspace, + library=LibraryService( + default_page_size=options.default_page_size, + max_page_size=options.max_page_size, + ), + configs=ConfigService( + options.workspace, + max_config_bytes=options.max_config_bytes, + default_page_size=options.default_page_size, + max_page_size=options.max_page_size, + ), + results=results, + artifacts=ArtifactRepository( + options.workspace, + results, + default_page_size=options.default_page_size, + max_page_size=options.max_page_size, + default_chunk_bytes=options.default_artifact_chunk_bytes, + max_chunk_bytes=options.max_artifact_chunk_bytes, + max_text_artifact_bytes=options.max_response_bytes, + ), + max_response_bytes=options.max_response_bytes, + ) + register_inspect_tools(server, services) + register_inspect_resources( + server, + services, + inline_artifact_bytes=options.max_artifact_chunk_bytes, ) return server diff --git a/assert_ai/mcp/tools/__init__.py b/assert_ai/mcp/tools/__init__.py new file mode 100644 index 000000000..5cc7f45d1 --- /dev/null +++ b/assert_ai/mcp/tools/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Capability-group tool registrars for the ASSERT MCP server.""" + +from assert_ai.mcp.tools.inspect import InspectServices, register_inspect_tools + +__all__ = ["InspectServices", "register_inspect_tools"] diff --git a/assert_ai/mcp/tools/inspect.py b/assert_ai/mcp/tools/inspect.py new file mode 100644 index 000000000..323d7f69e --- /dev/null +++ b/assert_ai/mcp/tools/inspect.py @@ -0,0 +1,649 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Read-only ASSERT MCP tools.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import yaml +from mcp.server import MCPServer +from mcp.types import ToolAnnotations + +from assert_ai.core.config_document import ( + ConfigValidationReport, + EVAL_CONFIG_SCHEMA_VERSION, +) +from assert_ai.core.workspace import WorkspaceService +from assert_ai.mcp.errors import adapt_tool_errors, invoke_tool +from assert_ai.mcp.models import ( + ConfigCatalogItem, + ConfigCatalogPage, + ConfigResult, + ConfigSchemaResult, + FailurePage, + PresetCatalogItem, + PresetCatalogPage, + PresetResult, + RunCatalogItem, + RunCatalogPage, + RunComparisonResult, + RunReferenceInput, + RunResult, + ScorePage, + SuiteCatalogItem, + SuiteCatalogPage, + SuiteResult, + TestCasePage, + TestCaseResult, + TranscriptResult, +) +from assert_ai.mcp.sanitize import sanitize_for_mcp +from assert_ai.mcp.uris import ( + config_uri, + preset_uri, + run_config_uri, + run_manifest_uri, + run_summary_uri, + run_transcript_uri, + suite_taxonomy_uri, + suite_test_case_uri, +) +from assert_ai.services.artifacts import ( + ArtifactChunk, + ArtifactPage, + ArtifactRepository, +) +from assert_ai.services.configs import ConfigService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.library import LibraryService, PresetKind +from assert_ai.services.results import ResultRepository, RunReference + +_READ_ONLY_ANNOTATIONS = ToolAnnotations( + read_only_hint=True, + destructive_hint=False, + idempotent_hint=True, + open_world_hint=False, +) + + +@dataclass(frozen=True, slots=True) +class InspectServices: + """Application services shared by all inspect tools and resources.""" + + workspace: WorkspaceService + library: LibraryService + configs: ConfigService + results: ResultRepository + artifacts: ArtifactRepository + max_response_bytes: int + + +def register_inspect_tools( + server: MCPServer, + services: InspectServices, +) -> None: + """Register the complete read-only inspect capability group.""" + + workspace = services.workspace + + @server.tool( + title="List ASSERT presets", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def list_presets( + kind: PresetKind | None = None, + cursor: str | None = None, + page_size: int | None = None, + ) -> PresetCatalogPage: + """List built-in behavior and judge presets using bounded pagination.""" + page = invoke_tool( + lambda: services.library.list_presets( + kind=kind, + cursor=cursor, + page_size=page_size, + ), + workspace=workspace, + ) + return PresetCatalogPage( + items=tuple( + PresetCatalogItem( + **item.model_dump(mode="python"), + resource_uri=preset_uri(item.kind.value, item.name), + ) + for item in page.items + ), + next_cursor=page.next_cursor, + ) + + @server.tool( + title="Get an ASSERT preset", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def get_preset(kind: PresetKind, name: str) -> PresetResult: + """Get one complete built-in preset by kind and name.""" + record = invoke_tool( + lambda: services.library.get_preset(kind, name), + workspace=workspace, + ) + document = _safe_mapping(record.document, workspace=workspace) + return PresetResult( + kind=record.kind, + name=record.name, + version=record.version, + tags=record.tags, + yaml=_dump_yaml(document), + document=document, + resource_uri=preset_uri(record.kind.value, record.name), + ) + + @server.tool( + title="Get the ASSERT eval config schema", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def get_config_schema() -> ConfigSchemaResult: + """Get the canonical Draft 2020-12 schema for eval_config.yaml.""" + schema = invoke_tool( + services.configs.get_schema, + workspace=workspace, + ) + return ConfigSchemaResult( + schema_version=EVAL_CONFIG_SCHEMA_VERSION, + json_schema=_safe_mapping(schema, workspace=workspace), + ) + + @server.tool( + title="List managed ASSERT configs", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def list_configs( + cursor: str | None = None, + page_size: int | None = None, + ) -> ConfigCatalogPage: + """List workspace-managed eval configs using bounded pagination.""" + page = invoke_tool( + lambda: services.configs.list_configs( + cursor=cursor, + limit=page_size, + ), + workspace=workspace, + ) + return ConfigCatalogPage( + items=tuple( + ConfigCatalogItem( + **item.model_dump(mode="python"), + resource_uri=config_uri(item.config_ref), + ) + for item in page.items + ), + next_cursor=page.next_cursor, + ) + + @server.tool( + title="Get a managed ASSERT config", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def get_config(config_ref: str) -> ConfigResult: + """Get one normalized config, its ETag, and validation report.""" + record = invoke_tool( + lambda: services.configs.get_config(config_ref), + workspace=workspace, + ) + document = _safe_mapping(record.document, workspace=workspace) + validation = ConfigValidationReport.model_validate( + sanitize_for_mcp(record.validation, workspace=workspace) + ) + return ConfigResult( + config_ref=record.config_ref, + yaml=_dump_yaml(document), + document=document, + etag=record.etag, + validation=validation, + resource_uri=config_uri(record.config_ref), + ) + + @server.tool( + title="List ASSERT result suites", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def list_suites( + cursor: str | None = None, + page_size: int | None = None, + ) -> SuiteCatalogPage: + """List result suites without loading score or transcript rows.""" + page = invoke_tool( + lambda: services.results.list_suite_catalog_entries( + cursor=cursor, + page_size=page_size, + ), + workspace=workspace, + ) + items = [] + for raw_item in page.items: + item = _safe_mapping(raw_item, workspace=workspace) + suite_id = str(item["suite_id"]) + item["resources"] = { + "taxonomy": suite_taxonomy_uri(suite_id), + } + items.append(SuiteCatalogItem.model_validate(item)) + return SuiteCatalogPage( + items=tuple(items), + next_cursor=page.next_cursor, + ) + + @server.tool( + title="Get an ASSERT result suite", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def get_suite(suite_id: str) -> SuiteResult: + """Get metadata, behavior, counts, and resource links for one suite.""" + summary = invoke_tool( + lambda: services.results.get_suite(suite_id), + workspace=workspace, + ) + payload = _public_suite(summary, workspace=workspace) + payload["resources"] = { + "taxonomy": suite_taxonomy_uri(suite_id), + } + return SuiteResult.model_validate(payload) + + @server.tool( + title="List ASSERT runs", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def list_runs( + suite_id: str, + cursor: str | None = None, + page_size: int | None = None, + ) -> RunCatalogPage: + """List runs in one suite without loading score or transcript rows.""" + page = invoke_tool( + lambda: services.results.list_run_catalog_entries( + suite_id, + cursor=cursor, + page_size=page_size, + ), + workspace=workspace, + ) + items = [] + for raw_item in page.items: + item = _safe_mapping(raw_item, workspace=workspace) + run_id = str(item["run_id"]) + item["resources"] = _run_resources(suite_id, run_id) + items.append(RunCatalogItem.model_validate(item)) + return RunCatalogPage( + items=tuple(items), + next_cursor=page.next_cursor, + ) + + @server.tool( + title="Get an ASSERT run", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def get_run(suite_id: str, run_id: str) -> RunResult: + """Get metadata-only quality, timing, usage, and model details.""" + summary = invoke_tool( + lambda: services.results.load_run_detail(suite_id, run_id), + workspace=workspace, + ) + payload = _public_run(summary, workspace=workspace) + payload["resources"] = _run_resources(suite_id, run_id) + return RunResult.model_validate(payload) + + @server.tool( + title="Compare ASSERT runs", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def compare_runs( + run_refs: list[RunReferenceInput], + metric: str = "policy_violation", + behavior_limit: int = 8, + ) -> RunComparisonResult: + """Compare two or more within-suite or cross-suite runs.""" + if len(run_refs) > 20: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "compare_runs accepts at most 20 run references", + ) + if behavior_limit < 0 or behavior_limit > services.results.max_page_size: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + ( + "behavior_limit must be between 0 and " + f"{services.results.max_page_size}" + ), + ) + result = invoke_tool( + lambda: services.results.compare_runs( + [ + RunReference( + suite_id=reference.suite_id, + run_id=reference.run_id, + ) + for reference in run_refs + ], + metric=metric, + behavior_limit=behavior_limit, + ), + workspace=workspace, + ) + return RunComparisonResult.model_validate( + _safe_mapping(result, workspace=workspace) + ) + + @server.tool( + title="List ASSERT test cases", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def list_test_cases( + suite_id: str, + run_id: str | None = None, + cursor: str | None = None, + page_size: int | None = None, + kind: str | None = None, + behavior: str | None = None, + test_case_id: str | None = None, + factors: dict[str, Any] | None = None, + ) -> TestCasePage: + """Query a suite or run test set with stable, source-bound cursors.""" + page = invoke_tool( + lambda: services.results.list_test_cases( + suite_id, + run_id=run_id, + cursor=cursor, + page_size=page_size, + kind=kind, + behavior=behavior, + test_case_id=test_case_id, + factors=factors, + ), + workspace=workspace, + ) + return TestCasePage( + items=_safe_list(page.items, workspace=workspace), + next_cursor=page.next_cursor, + ) + + @server.tool( + title="Get an ASSERT test case", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def get_test_case( + suite_id: str, + test_case_id: str, + kind: str | None = None, + run_id: str | None = None, + ) -> TestCaseResult: + """Get one complete test case through its JSONL index.""" + row = invoke_tool( + lambda: services.results.get_test_case( + suite_id, + test_case_id, + kind=kind, + run_id=run_id, + ), + workspace=workspace, + ) + return TestCaseResult( + row=_safe_mapping(row, workspace=workspace), + resource_uri=suite_test_case_uri( + suite_id, + test_case_id, + kind=kind, + run_id=run_id, + ), + ) + + @server.tool( + title="List ASSERT scores", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def list_scores( + suite_id: str, + run_id: str, + cursor: str | None = None, + page_size: int | None = None, + kind: str | None = None, + behavior: str | None = None, + test_case_id: str | None = None, + dimension: str | None = None, + dimension_value: bool | int | str | None = None, + match_not_applicable: bool = False, + judge_status: str | None = None, + target: str | None = None, + stop_reason: str | None = None, + has_tool_use: bool | None = None, + factors: dict[str, Any] | None = None, + ) -> ScorePage: + """Query score rows by behavior, dimension, status, target, or tool use.""" + page = invoke_tool( + lambda: services.results.list_scores( + suite_id, + run_id, + cursor=cursor, + page_size=page_size, + kind=kind, + behavior=behavior, + test_case_id=test_case_id, + dimension=dimension, + dimension_value=dimension_value, + match_not_applicable=match_not_applicable, + judge_status=judge_status, + target=target, + stop_reason=stop_reason, + has_tool_use=has_tool_use, + factors=factors, + ), + workspace=workspace, + ) + return ScorePage( + items=_safe_list(page.items, workspace=workspace), + next_cursor=page.next_cursor, + ) + + @server.tool( + title="List ASSERT failures", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def list_failures( + suite_id: str, + run_id: str, + dimension: str = "policy_violation", + include_judge_failures: bool = True, + cursor: str | None = None, + page_size: int | None = None, + kind: str | None = None, + behavior: str | None = None, + ) -> FailurePage: + """List flagged score rows and optional judge failures.""" + page = invoke_tool( + lambda: services.results.list_failures( + suite_id, + run_id, + dimension=dimension, + include_judge_failures=include_judge_failures, + cursor=cursor, + page_size=page_size, + kind=kind, + behavior=behavior, + ), + workspace=workspace, + ) + return FailurePage( + items=_safe_list(page.items, workspace=workspace), + next_cursor=page.next_cursor, + ) + + @server.tool( + title="Get an ASSERT transcript", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def get_transcript( + suite_id: str, + run_id: str, + test_case_id: str, + kind: str | None = None, + ) -> TranscriptResult: + """Join one test case, inference transcript, and score verdict.""" + transcript = invoke_tool( + lambda: services.results.get_transcript( + suite_id, + run_id, + test_case_id, + kind=kind, + ), + workspace=workspace, + ) + payload = _safe_mapping(transcript, workspace=workspace) + payload["resource_uri"] = run_transcript_uri( + suite_id, + run_id, + test_case_id, + kind=kind, + ) + return TranscriptResult.model_validate(payload) + + @server.tool( + title="List ASSERT artifacts", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def list_artifacts( + suite_id: str, + run_id: str | None = None, + cursor: str | None = None, + page_size: int | None = None, + ) -> ArtifactPage: + """List manifest-backed artifacts without exposing filesystem paths.""" + return invoke_tool( + lambda: services.artifacts.list_artifacts( + suite_id, + run_id=run_id, + cursor=cursor, + page_size=page_size, + ), + workspace=workspace, + ) + + @server.tool( + title="Read an ASSERT artifact chunk", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors(workspace, max_response_bytes=services.max_response_bytes) + def read_artifact_chunk( + artifact_id: str, + offset: int = 0, + chunk_size: int | None = None, + ) -> ArtifactChunk: + """Read one bounded, redacted text or base64 binary artifact chunk.""" + return invoke_tool( + lambda: services.artifacts.read_artifact_chunk( + artifact_id, + offset=offset, + chunk_size=chunk_size, + ), + workspace=workspace, + ) + + +def _safe_mapping( + value: Any, + *, + workspace: WorkspaceService, +) -> dict[str, Any]: + sanitized = sanitize_for_mcp(value, workspace=workspace) + if not isinstance(sanitized, dict): + raise TypeError("Expected a mapping from the application service") + return sanitized + + +def _safe_list( + value: Any, + *, + workspace: WorkspaceService, +) -> list[dict[str, Any]]: + sanitized = sanitize_for_mcp(value, workspace=workspace) + if not isinstance(sanitized, list) or not all( + isinstance(item, dict) for item in sanitized + ): + raise TypeError("Expected a list of mappings from the application service") + return sanitized + + +def _dump_yaml(document: dict[str, Any]) -> str: + text = yaml.safe_dump( + document, + default_flow_style=False, + sort_keys=False, + allow_unicode=True, + ) + return text if text.endswith("\n") else text + "\n" + + +def _public_suite( + summary: dict[str, Any], + *, + workspace: WorkspaceService, +) -> dict[str, Any]: + payload = _safe_mapping(summary, workspace=workspace) + for key in ( + "artifact_versions", + "sources", + "run_set_identity", + "run_catalog_identity", + ): + payload.pop(key, None) + return payload + + +def _public_run( + summary: dict[str, Any], + *, + workspace: WorkspaceService, +) -> dict[str, Any]: + payload = _safe_mapping(summary, workspace=workspace) + for key in ("artifact_versions", "sources", "indexes"): + payload.pop(key, None) + return payload + + +def _run_resources(suite_id: str, run_id: str) -> dict[str, str]: + return { + "summary": run_summary_uri(suite_id, run_id), + "manifest": run_manifest_uri(suite_id, run_id), + "config": run_config_uri(suite_id, run_id), + } diff --git a/assert_ai/mcp/uris.py b/assert_ai/mcp/uris.py new file mode 100644 index 000000000..86b1844dc --- /dev/null +++ b/assert_ai/mcp/uris.py @@ -0,0 +1,75 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Canonical path-free ASSERT MCP resource URIs.""" + +from __future__ import annotations + +from urllib.parse import quote, urlencode + + +def preset_uri(kind: str, name: str) -> str: + return f"assert://preset/{quote(kind, safe='')}/{quote(name, safe='')}" + + +def config_uri(config_ref: str) -> str: + return f"assert://config/{quote(config_ref, safe='')}" + + +def suite_taxonomy_uri(suite_id: str) -> str: + return f"assert://suite/{quote(suite_id, safe='')}/taxonomy" + + +def suite_test_case_uri( + suite_id: str, + test_case_id: str, + *, + kind: str | None = None, + run_id: str | None = None, +) -> str: + uri = ( + f"assert://suite/{quote(suite_id, safe='')}/test-case/" + f"{quote(test_case_id, safe='')}" + ) + query = { + key: value + for key, value in (("kind", kind), ("run_id", run_id)) + if value is not None + } + return f"{uri}?{urlencode(query)}" if query else uri + + +def run_summary_uri(suite_id: str, run_id: str) -> str: + return ( + f"assert://run/{quote(suite_id, safe='')}/" + f"{quote(run_id, safe='')}/summary" + ) + + +def run_manifest_uri(suite_id: str, run_id: str) -> str: + return ( + f"assert://run/{quote(suite_id, safe='')}/" + f"{quote(run_id, safe='')}/manifest" + ) + + +def run_config_uri(suite_id: str, run_id: str) -> str: + return ( + f"assert://run/{quote(suite_id, safe='')}/" + f"{quote(run_id, safe='')}/config" + ) + + +def run_transcript_uri( + suite_id: str, + run_id: str, + test_case_id: str, + *, + kind: str | None = None, +) -> str: + uri = ( + f"assert://run/{quote(suite_id, safe='')}/" + f"{quote(run_id, safe='')}/transcript/" + f"{quote(test_case_id, safe='')}" + ) + return f"{uri}?{urlencode({'kind': kind})}" if kind is not None else uri diff --git a/assert_ai/services/artifacts.py b/assert_ai/services/artifacts.py new file mode 100644 index 000000000..ec35ffece --- /dev/null +++ b/assert_ai/services/artifacts.py @@ -0,0 +1,1104 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Manifest-backed, opaque access to managed ASSERT artifacts.""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import mimetypes +import re +from dataclasses import dataclass, field +from datetime import datetime, timezone +from enum import StrEnum +from pathlib import Path +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict + +from assert_ai.core.security import redact_path_prefixes, sanitize_text +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.results import ResultRepository + +_ARTIFACT_ID_PREFIX = "art1_" +_ARTIFACT_ID_VERSION = 1 +_CURSOR_PREFIX = "ac1_" +_CURSOR_VERSION = 1 +_IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,254}$") +_DEFAULT_PAGE_SIZE = 50 +_DEFAULT_MAX_PAGE_SIZE = 200 +_DEFAULT_CHUNK_BYTES = 64 * 1024 +_DEFAULT_MAX_CHUNK_BYTES = 256 * 1024 +_DEFAULT_MAX_TEXT_LINE_BYTES = 8 * 1024 * 1024 +_DEFAULT_MAX_METADATA_BYTES = 1024 * 1024 + + +class ArtifactScope(StrEnum): + """Result scope that owns an artifact reference.""" + + SUITE = "suite" + RUN = "run" + + +class _ArtifactModel(BaseModel): + model_config = ConfigDict(frozen=True) + + +class ArtifactDescriptor(_ArtifactModel): + """Path-free metadata for one manifest-backed artifact.""" + + artifact_id: str + name: str + aliases: tuple[str, ...] = () + scope: ArtifactScope + suite_id: str + run_id: str | None = None + media_type: str + size_bytes: int + modified_at: str + sha256: str | None = None + text: bool + redacted: bool + resource_uri: str + + +class ArtifactPage(_ArtifactModel): + """Bounded page of artifact descriptors.""" + + items: tuple[ArtifactDescriptor, ...] + next_cursor: str | None = None + + +class ArtifactChunk(_ArtifactModel): + """One bounded artifact chunk.""" + + artifact: ArtifactDescriptor + offset: int + next_offset: int | None + eof: bool + offset_basis: Literal["redacted_text"] + encoding: Literal["utf-8"] + data: str + bytes_returned: int + source_size_bytes: int + view_size_bytes: int | None = None + + +@dataclass(slots=True) +class _ArtifactCandidate: + name: str + path: Path + scope: ArtifactScope + suite_id: str + run_id: str | None + media_type: str + text: bool + size_bytes: int + mtime_ns: int + sha256: str | None = None + aliases: list[str] = field(default_factory=list) + + +class ArtifactRepository: + """Resolve only artifacts named by ASSERT summaries and manifests.""" + + def __init__( + self, + workspace: WorkspaceService, + results: ResultRepository, + *, + default_page_size: int = _DEFAULT_PAGE_SIZE, + max_page_size: int = _DEFAULT_MAX_PAGE_SIZE, + default_chunk_bytes: int = _DEFAULT_CHUNK_BYTES, + max_chunk_bytes: int = _DEFAULT_MAX_CHUNK_BYTES, + max_text_line_bytes: int = _DEFAULT_MAX_TEXT_LINE_BYTES, + max_metadata_bytes: int = _DEFAULT_MAX_METADATA_BYTES, + max_text_artifact_bytes: int = 1024 * 1024, + ) -> None: + if default_page_size < 1: + raise ValueError("default_page_size must be positive") + if max_page_size < default_page_size: + raise ValueError("max_page_size must be >= default_page_size") + if default_chunk_bytes < 1: + raise ValueError("default_chunk_bytes must be positive") + if max_chunk_bytes < default_chunk_bytes: + raise ValueError("max_chunk_bytes must be >= default_chunk_bytes") + if max_text_line_bytes < max_chunk_bytes: + raise ValueError("max_text_line_bytes must be >= max_chunk_bytes") + if max_metadata_bytes < 1: + raise ValueError("max_metadata_bytes must be positive") + if max_text_artifact_bytes < max_chunk_bytes: + raise ValueError( + "max_text_artifact_bytes must be >= max_chunk_bytes" + ) + self.workspace = workspace + self.results = results + self.default_page_size = default_page_size + self.max_page_size = max_page_size + self.default_chunk_bytes = default_chunk_bytes + self.max_chunk_bytes = max_chunk_bytes + self.max_text_line_bytes = max_text_line_bytes + self.max_metadata_bytes = max_metadata_bytes + self.max_text_artifact_bytes = max_text_artifact_bytes + + def list_artifacts( + self, + suite_id: str, + *, + run_id: str | None = None, + cursor: str | None = None, + page_size: int | None = None, + ) -> ArtifactPage: + candidates = self._catalog(suite_id, run_id=run_id) + descriptors = [self._descriptor(candidate) for candidate in candidates] + identity = _catalog_identity(candidates) + offset = 0 + if cursor is not None: + payload = _decode_cursor(cursor) + if ( + payload.get("suite_id") != suite_id + or payload.get("run_id") != run_id + or payload.get("catalog_sha256") != identity + ): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + "Artifact cursor no longer matches this catalog", + ) + offset = int(payload["offset"]) + if offset < 0 or offset > len(descriptors): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + "Artifact cursor is no longer valid", + ) + + limit = self._page_size(page_size) + items = tuple(descriptors[offset : offset + limit]) + next_offset = offset + len(items) + next_cursor = None + if next_offset < len(descriptors): + next_cursor = _encode_cursor( + suite_id=suite_id, + run_id=run_id, + catalog_sha256=identity, + offset=next_offset, + ) + return ArtifactPage(items=items, next_cursor=next_cursor) + + def get_artifact(self, artifact_id: str) -> ArtifactDescriptor: + return self._descriptor(self._resolve_artifact_id(artifact_id)) + + def find_artifact( + self, + suite_id: str, + name: str, + *, + run_id: str | None = None, + ) -> ArtifactDescriptor: + for candidate in self._catalog(suite_id, run_id=run_id): + if name == candidate.name or name in candidate.aliases: + return self._descriptor(candidate) + target = f"{suite_id}/{run_id}" if run_id is not None else suite_id + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Artifact not found: {target}/{name}", + ) + + def read_artifact_chunk( + self, + artifact_id: str, + *, + offset: int = 0, + chunk_size: int | None = None, + ) -> ArtifactChunk: + if offset < 0: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "offset must be non-negative", + ) + candidate = self._resolve_artifact_id(artifact_id) + descriptor = self._descriptor(candidate) + limit = self._chunk_size(chunk_size) + path = self._revalidate_candidate(candidate) + + if not candidate.text: + raise ServiceError( + ServiceErrorCode.CAPABILITY_DISABLED, + ( + "Binary artifact reads are disabled because their contents " + "cannot be safely redacted" + ), + ) + if candidate.size_bytes > self.max_text_artifact_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + ( + "Text artifact exceeds the generic read limit; " + "use the result-specific paginated tools" + ), + details={ + "size_bytes": candidate.size_bytes, + "max_text_artifact_bytes": self.max_text_artifact_bytes, + }, + ) + data, next_offset, eof, view_size = self._read_redacted_text( + path, + offset=offset, + chunk_size=limit, + ) + self._assert_unchanged(candidate, path) + return ArtifactChunk( + artifact=descriptor, + offset=offset, + next_offset=next_offset, + eof=eof, + offset_basis="redacted_text", + encoding="utf-8", + data=data.decode("utf-8"), + bytes_returned=len(data), + source_size_bytes=candidate.size_bytes, + view_size_bytes=view_size, + ) + + def _catalog( + self, + suite_id: str, + *, + run_id: str | None, + ) -> list[_ArtifactCandidate]: + self._validate_identifier(suite_id, "suite_id") + if run_id is not None: + self._validate_identifier(run_id, "run_id") + suite_dir = self._managed_dir( + self.results.results_root / suite_id, + expected_root=self.results.results_root, + field_name="artifact suite", + ) + if run_id is None: + summary = self.results.get_suite(suite_id) + candidates = self._suite_candidates(suite_dir, summary) + else: + run_dir = self._managed_dir( + suite_dir / run_id, + expected_root=suite_dir, + field_name="artifact run", + ) + summary = self.results.load_run_detail(suite_id, run_id) + candidates = self._run_candidates( + suite_dir, + run_dir, + summary, + ) + candidates.sort(key=lambda candidate: candidate.name) + return candidates + + def _suite_candidates( + self, + suite_dir: Path, + summary: dict[str, Any], + ) -> list[_ArtifactCandidate]: + candidates: list[_ArtifactCandidate] = [] + self._add_known_file( + candidates, + name="suite_summary", + path=suite_dir / "suite_summary.json", + expected_root=suite_dir, + suite_id=suite_dir.name, + ) + self._add_known_file( + candidates, + name="suite_metadata", + path=suite_dir / "suite.json", + expected_root=suite_dir, + suite_id=suite_dir.name, + ) + self._add_known_file( + candidates, + name="latest_artifacts", + path=suite_dir / "latest.json", + expected_root=suite_dir, + suite_id=suite_dir.name, + ) + self._add_summary_references( + candidates, + summary, + suite_dir=suite_dir, + run_dir=None, + suite_id=suite_dir.name, + run_id=None, + ) + self._add_active_artifact_versions( + candidates, + summary.get("artifact_versions"), + suite_dir=suite_dir, + suite_id=suite_dir.name, + ) + return candidates + + def _run_candidates( + self, + suite_dir: Path, + run_dir: Path, + summary: dict[str, Any], + ) -> list[_ArtifactCandidate]: + candidates: list[_ArtifactCandidate] = [] + for name, filename in ( + ("run_summary", "run_summary.json"), + ("manifest", "manifest.json"), + ("config", "config.yaml"), + ("metrics", "metrics.json"), + ("artifact_versions", "artifacts.json"), + ): + self._add_known_file( + candidates, + name=name, + path=run_dir / filename, + expected_root=run_dir, + suite_id=suite_dir.name, + run_id=run_dir.name, + ) + self._add_summary_references( + candidates, + summary, + suite_dir=suite_dir, + run_dir=run_dir, + suite_id=suite_dir.name, + run_id=run_dir.name, + ) + self._add_active_artifact_versions( + candidates, + summary.get("artifact_versions"), + suite_dir=suite_dir, + suite_id=suite_dir.name, + run_id=run_dir.name, + ) + self._add_viewer_artifacts( + candidates, + suite_dir=suite_dir, + run_dir=run_dir, + ) + return candidates + + def _add_summary_references( + self, + candidates: list[_ArtifactCandidate], + summary: dict[str, Any], + *, + suite_dir: Path, + run_dir: Path | None, + suite_id: str, + run_id: str | None, + ) -> None: + sources = summary.get("sources") + if isinstance(sources, dict): + for name, reference in sorted(sources.items()): + if not isinstance(name, str) or not isinstance(reference, dict): + continue + path, expected_root = self._reference_path( + reference, + suite_dir=suite_dir, + run_dir=run_dir, + ) + self._add_known_file( + candidates, + name=name, + path=path, + expected_root=expected_root, + suite_id=suite_id, + run_id=run_id, + sha256=_optional_sha256(reference.get("sha256")), + ) + index_reference = reference.get("index") + if isinstance(index_reference, dict): + index_path, index_root = self._reference_path( + index_reference, + suite_dir=suite_dir, + run_dir=run_dir, + ) + self._add_known_file( + candidates, + name=f"{name}_index", + path=index_path, + expected_root=index_root, + suite_id=suite_id, + run_id=run_id, + ) + + indexes = summary.get("indexes") + if isinstance(indexes, dict): + for name, reference in sorted(indexes.items()): + if not isinstance(name, str) or not isinstance(reference, dict): + continue + path, expected_root = self._reference_path( + reference, + suite_dir=suite_dir, + run_dir=run_dir, + ) + self._add_known_file( + candidates, + name=f"{name}_index", + path=path, + expected_root=expected_root, + suite_id=suite_id, + run_id=run_id, + ) + + def _add_active_artifact_versions( + self, + candidates: list[_ArtifactCandidate], + artifact_versions: Any, + *, + suite_dir: Path, + suite_id: str, + run_id: str | None = None, + ) -> None: + if not isinstance(artifact_versions, dict): + return + for stage_name, reference in sorted(artifact_versions.items()): + if not isinstance(stage_name, str) or not isinstance(reference, dict): + continue + artifact_dir_raw = reference.get("artifact_dir") + if not isinstance(artifact_dir_raw, str): + continue + artifact_dir = suite_dir / _relative_parts(artifact_dir_raw) + metadata_raw = reference.get("metadata_path") + metadata_path = ( + suite_dir / _relative_parts(metadata_raw) + if isinstance(metadata_raw, str) + else artifact_dir / "artifact.json" + ) + self._add_known_file( + candidates, + name=f"{stage_name}_artifact_metadata", + path=metadata_path, + expected_root=suite_dir, + suite_id=suite_id, + run_id=run_id, + ) + metadata = self._load_metadata( + self._managed_file( + metadata_path, + expected_root=suite_dir, + field_name=f"{stage_name} artifact metadata", + must_exist=False, + ), + label=f"{stage_name} artifact metadata", + ) + files = metadata.get("files") if isinstance(metadata, dict) else None + hashes = ( + metadata.get("file_hashes") + if isinstance(metadata, dict) + else None + ) + if not isinstance(files, dict): + primary_raw = reference.get("path") + if isinstance(primary_raw, str): + self._add_known_file( + candidates, + name=stage_name, + path=suite_dir / _relative_parts(primary_raw), + expected_root=suite_dir, + suite_id=suite_id, + run_id=run_id, + ) + continue + for output_name, filename in sorted(files.items()): + if not isinstance(output_name, str) or not isinstance(filename, str): + continue + self._add_known_file( + candidates, + name=f"{stage_name}_{output_name}", + aliases=(stage_name,) if output_name == stage_name else (), + path=artifact_dir / _relative_parts(filename), + expected_root=suite_dir, + suite_id=suite_id, + run_id=run_id, + sha256=( + _optional_sha256(hashes.get(output_name)) + if isinstance(hashes, dict) + else None + ), + ) + + def _add_viewer_artifacts( + self, + candidates: list[_ArtifactCandidate], + *, + suite_dir: Path, + run_dir: Path, + ) -> None: + manifest_path = run_dir / ".viewer" / "viewer_run_manifest.json" + self._add_known_file( + candidates, + name="viewer_manifest", + path=manifest_path, + expected_root=run_dir, + suite_id=suite_dir.name, + run_id=run_dir.name, + ) + safe_manifest_path = self._managed_file( + manifest_path, + expected_root=run_dir, + field_name="viewer manifest", + must_exist=False, + ) + manifest = self._load_metadata( + safe_manifest_path, + label="viewer manifest", + ) + if not isinstance(manifest, dict): + return + for section_name in ("source_files", "derived_files"): + section = manifest.get(section_name) + if not isinstance(section, dict): + continue + for artifact_name, reference in sorted(section.items()): + if not isinstance(artifact_name, str) or not isinstance(reference, dict): + continue + raw_path = reference.get("path") + if not isinstance(raw_path, str): + continue + self._add_known_file( + candidates, + name=f"viewer_{Path(artifact_name).stem}", + path=run_dir / _relative_parts(raw_path, allow_parent=True), + expected_root=suite_dir, + suite_id=suite_dir.name, + run_id=run_dir.name, + ) + + def _add_known_file( + self, + candidates: list[_ArtifactCandidate], + *, + name: str, + path: Path, + expected_root: Path, + suite_id: str, + run_id: str | None = None, + aliases: tuple[str, ...] = (), + sha256: str | None = None, + ) -> None: + managed = self._managed_file( + path, + expected_root=expected_root, + field_name=f"artifact {name}", + must_exist=False, + ) + if not managed.is_file(): + return + resolved = managed.resolve() + for candidate in candidates: + if candidate.path == resolved: + for alias in (name, *aliases): + if alias != candidate.name and alias not in candidate.aliases: + candidate.aliases.append(alias) + if candidate.sha256 is None: + candidate.sha256 = sha256 + return + stat_result = resolved.stat() + media_type, text = _media_type(resolved) + candidates.append( + _ArtifactCandidate( + name=name, + aliases=list(aliases), + path=resolved, + scope=ArtifactScope.RUN if run_id is not None else ArtifactScope.SUITE, + suite_id=suite_id, + run_id=run_id, + media_type=media_type, + text=text, + size_bytes=stat_result.st_size, + mtime_ns=stat_result.st_mtime_ns, + sha256=sha256, + ) + ) + + def _reference_path( + self, + reference: dict[str, Any], + *, + suite_dir: Path, + run_dir: Path | None, + ) -> tuple[Path, Path]: + scope = reference.get("scope") + raw_path = reference.get("path") + if not isinstance(raw_path, str) or not raw_path: + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Artifact metadata contains an invalid path reference", + ) + if scope == "run" and run_dir is not None: + root = run_dir + elif scope == "suite": + root = suite_dir + elif scope == "workspace": + root = self.workspace.root + else: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + "Artifact metadata references an unsupported scope", + ) + return root / _relative_parts(raw_path), root + + def _resolve_artifact_id(self, artifact_id: str) -> _ArtifactCandidate: + payload = _decode_artifact_id(artifact_id) + suite_id = str(payload["suite_id"]) + run_id = payload.get("run_id") + name = str(payload["name"]) + candidates = self._catalog( + suite_id, + run_id=str(run_id) if isinstance(run_id, str) else None, + ) + candidate = next( + (item for item in candidates if item.name == name), + None, + ) + if candidate is None: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + "Artifact ID no longer resolves to a managed artifact", + ) + if ( + candidate.size_bytes != payload.get("size_bytes") + or candidate.mtime_ns != payload.get("mtime_ns") + ): + raise ServiceError( + ServiceErrorCode.STALE_ETAG, + "Artifact changed since this ID was issued", + details={ + "suite_id": candidate.suite_id, + "run_id": candidate.run_id, + "name": candidate.name, + }, + ) + return candidate + + def _descriptor(self, candidate: _ArtifactCandidate) -> ArtifactDescriptor: + artifact_id = _encode_artifact_id(candidate) + return ArtifactDescriptor( + artifact_id=artifact_id, + name=candidate.name, + aliases=tuple(sorted(candidate.aliases)), + scope=candidate.scope, + suite_id=candidate.suite_id, + run_id=candidate.run_id, + media_type=candidate.media_type, + size_bytes=candidate.size_bytes, + modified_at=datetime.fromtimestamp( + candidate.mtime_ns / 1_000_000_000, + tz=timezone.utc, + ).isoformat(), + sha256=candidate.sha256, + text=candidate.text, + redacted=candidate.text, + resource_uri=f"assert://artifact/{artifact_id}", + ) + + def _read_redacted_text( + self, + path: Path, + *, + offset: int, + chunk_size: int, + ) -> tuple[bytes, int | None, bool, int | None]: + output = bytearray() + view_position = 0 + exhausted = False + has_more = False + + with path.open("rb") as handle: + while True: + line = handle.readline(self.max_text_line_bytes + 1) + if not line: + exhausted = True + break + if ( + len(line) > self.max_text_line_bytes + and not line.endswith((b"\n", b"\r")) + ): + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + "Text artifact contains a line larger than the configured limit", + ) + try: + sanitized = self._sanitize_text_line( + line.decode("utf-8") + ).encode("utf-8") + except UnicodeDecodeError as exc: + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Text artifact is not valid UTF-8", + ) from exc + line_end = view_position + len(sanitized) + if line_end <= offset: + view_position = line_end + continue + + start = max(0, offset - view_position) + if ( + start < len(sanitized) + and start > 0 + and sanitized[start] & 0xC0 == 0x80 + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "offset must align with a UTF-8 character boundary", + ) + remaining = chunk_size - len(output) + piece = _valid_utf8_prefix(sanitized[start : start + remaining]) + output.extend(piece) + consumed = start + len(piece) + view_position = line_end + if consumed < len(sanitized) or len(output) >= chunk_size: + has_more = consumed < len(sanitized) or bool(handle.read(1)) + break + + if offset > view_position and exhausted: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "offset exceeds the redacted artifact view", + ) + eof = not has_more + next_offset = None if eof else offset + len(output) + view_size = view_position if eof else None + return bytes(output), next_offset, eof, view_size + + def _revalidate_candidate(self, candidate: _ArtifactCandidate) -> Path: + suite_root = self.workspace.results_root / candidate.suite_id + run_root = ( + suite_root / candidate.run_id + if candidate.run_id is not None + else None + ) + if run_root is not None and candidate.path.is_relative_to(run_root.resolve()): + expected_root = run_root + elif candidate.path.is_relative_to(suite_root.resolve()): + expected_root = suite_root + else: + expected_root = self.workspace.root + path = self._managed_file( + candidate.path, + expected_root=expected_root, + field_name=f"artifact {candidate.name}", + must_exist=True, + ) + self._assert_unchanged(candidate, path) + return path + + def _sanitize_text_line(self, text: str) -> str: + return redact_path_prefixes( + sanitize_text(text), + (self.workspace.root,), + ) + + def _load_metadata( + self, + path: Path, + *, + label: str, + ) -> dict[str, Any] | None: + if not path.is_file(): + return None + try: + stat_result = path.stat() + if stat_result.st_size > self.max_metadata_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"{label.title()} exceeds the configured metadata limit", + ) + payload = json.loads(path.read_text(encoding="utf-8")) + except ServiceError: + raise + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + f"{label.title()} is not valid JSON", + ) from exc + if not isinstance(payload, dict): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + f"{label.title()} must contain a JSON object", + ) + return payload + + @staticmethod + def _assert_unchanged(candidate: _ArtifactCandidate, path: Path) -> None: + stat_result = path.stat() + if ( + stat_result.st_size != candidate.size_bytes + or stat_result.st_mtime_ns != candidate.mtime_ns + ): + raise ServiceError( + ServiceErrorCode.STALE_ETAG, + "Artifact changed while it was being read", + ) + + def _managed_dir( + self, + path: Path, + *, + expected_root: Path, + field_name: str, + ) -> Path: + try: + managed = self.workspace.path_policy.resolve_managed_output( + path, + field_name=field_name, + expected_root=expected_root, + reject_links=True, + ) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + str(exc), + ) from exc + if not managed.is_dir(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"{field_name.title()} not found", + ) + return managed + + def _managed_file( + self, + path: Path, + *, + expected_root: Path, + field_name: str, + must_exist: bool, + ) -> Path: + try: + managed = self.workspace.path_policy.resolve_managed_output( + path, + field_name=field_name, + expected_root=expected_root, + reject_links=True, + ) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + str(exc), + ) from exc + if must_exist and not managed.is_file(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Artifact not found: {field_name}", + ) + return managed + + def _page_size(self, requested: int | None) -> int: + if requested is None: + return self.default_page_size + if requested < 1 or requested > self.max_page_size: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"page_size must be between 1 and {self.max_page_size}", + ) + return requested + + def _chunk_size(self, requested: int | None) -> int: + if requested is None: + return self.default_chunk_bytes + if requested < 4 or requested > self.max_chunk_bytes: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"chunk_size must be between 4 and {self.max_chunk_bytes}", + ) + return requested + + @staticmethod + def _validate_identifier(value: str, field_name: str) -> None: + if not _IDENTIFIER_RE.fullmatch(value): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{field_name} contains unsupported characters", + ) + + +def _media_type(path: Path) -> tuple[str, bool]: + suffix = path.suffix.lower() + if suffix == ".jsonl": + return "application/x-ndjson", True + if suffix in {".yaml", ".yml"}: + return "application/yaml", True + guessed, _ = mimetypes.guess_type(path.name) + media_type = guessed or "application/octet-stream" + text = ( + media_type.startswith("text/") + or media_type in {"application/json", "application/xml"} + or suffix in {".md", ".log", ".txt"} + ) + return media_type, text + + +def _relative_parts(value: str, *, allow_parent: bool = False) -> Path: + if not value or "\x00" in value: + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + "Artifact metadata contains an invalid relative path", + ) + normalized = value.replace("\\", "/") + path = Path(normalized) + if path.is_absolute(): + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + "Artifact metadata contains an absolute path", + ) + if not allow_parent and any(part == ".." for part in path.parts): + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + "Artifact metadata contains parent traversal", + ) + return path + + +def _optional_sha256(value: Any) -> str | None: + if not isinstance(value, str) or not re.fullmatch(r"[0-9a-fA-F]{64}", value): + return None + return value.lower() + + +def _valid_utf8_prefix(data: bytes) -> bytes: + if not data: + return data + try: + data.decode("utf-8") + return data + except UnicodeDecodeError as exc: + if exc.reason == "unexpected end of data": + return data[: exc.start] + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Text artifact is not valid UTF-8", + ) from exc + + +def _encode_artifact_id(candidate: _ArtifactCandidate) -> str: + payload = json.dumps( + { + "v": _ARTIFACT_ID_VERSION, + "suite_id": candidate.suite_id, + "run_id": candidate.run_id, + "name": candidate.name, + "size_bytes": candidate.size_bytes, + "mtime_ns": candidate.mtime_ns, + }, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + token = base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=") + return f"{_ARTIFACT_ID_PREFIX}{token}" + + +def _decode_artifact_id(artifact_id: str) -> dict[str, Any]: + if ( + not artifact_id.startswith(_ARTIFACT_ID_PREFIX) + or len(artifact_id) > 4096 + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid artifact ID", + ) + token = artifact_id[len(_ARTIFACT_ID_PREFIX) :] + try: + padding = "=" * (-len(token) % 4) + payload = json.loads( + base64.urlsafe_b64decode(token + padding).decode("utf-8") + ) + except (ValueError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid artifact ID", + ) from exc + if ( + not isinstance(payload, dict) + or payload.get("v") != _ARTIFACT_ID_VERSION + or not isinstance(payload.get("suite_id"), str) + or payload.get("run_id") is not None + and not isinstance(payload.get("run_id"), str) + or not isinstance(payload.get("name"), str) + or not isinstance(payload.get("size_bytes"), int) + or not isinstance(payload.get("mtime_ns"), int) + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid artifact ID", + ) + return payload + + +def _catalog_identity(candidates: list[_ArtifactCandidate]) -> str: + payload = [ + { + "name": candidate.name, + "aliases": sorted(candidate.aliases), + "size_bytes": candidate.size_bytes, + "mtime_ns": candidate.mtime_ns, + } + for candidate in candidates + ] + return hashlib.sha256( + json.dumps( + payload, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + ).hexdigest() + + +def _encode_cursor( + *, + suite_id: str, + run_id: str | None, + catalog_sha256: str, + offset: int, +) -> str: + payload = json.dumps( + { + "v": _CURSOR_VERSION, + "suite_id": suite_id, + "run_id": run_id, + "catalog_sha256": catalog_sha256, + "offset": offset, + }, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + token = base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=") + return f"{_CURSOR_PREFIX}{token}" + + +def _decode_cursor(cursor: str) -> dict[str, Any]: + if not cursor.startswith(_CURSOR_PREFIX) or len(cursor) > 4096: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid artifact cursor", + ) + token = cursor[len(_CURSOR_PREFIX) :] + try: + padding = "=" * (-len(token) % 4) + payload = json.loads( + base64.urlsafe_b64decode(token + padding).decode("utf-8") + ) + except (ValueError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid artifact cursor", + ) from exc + if ( + not isinstance(payload, dict) + or payload.get("v") != _CURSOR_VERSION + or not isinstance(payload.get("suite_id"), str) + or payload.get("run_id") is not None + and not isinstance(payload.get("run_id"), str) + or not isinstance(payload.get("catalog_sha256"), str) + or not isinstance(payload.get("offset"), int) + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid artifact cursor", + ) + return payload diff --git a/assert_ai/services/library.py b/assert_ai/services/library.py new file mode 100644 index 000000000..e61a927a5 --- /dev/null +++ b/assert_ai/services/library.py @@ -0,0 +1,239 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Typed, bounded access to ASSERT's built-in preset library.""" + +from __future__ import annotations + +import base64 +import json +from enum import StrEnum +from typing import Any + +import yaml +from pydantic import BaseModel, ConfigDict + +from assert_ai.library.loader import discover, load_preset +from assert_ai.services.errors import ServiceError, ServiceErrorCode + +_CURSOR_VERSION = 1 +_DEFAULT_PAGE_SIZE = 50 +_DEFAULT_MAX_PAGE_SIZE = 200 + + +class PresetKind(StrEnum): + """Stable built-in preset categories.""" + + BEHAVIOR = "behavior" + JUDGE_PRESET = "judge_preset" + + +class _LibraryModel(BaseModel): + model_config = ConfigDict(frozen=True) + + +class PresetCatalogEntry(_LibraryModel): + """Lightweight metadata for one built-in preset.""" + + kind: PresetKind + name: str + version: str | None = None + tags: tuple[str, ...] = () + summary: str | None = None + description: str | None = None + + +class PresetPage(_LibraryModel): + """Bounded page of built-in preset metadata.""" + + items: tuple[PresetCatalogEntry, ...] + next_cursor: str | None = None + + +class PresetRecord(_LibraryModel): + """One complete built-in preset definition.""" + + kind: PresetKind + name: str + version: str | None = None + tags: tuple[str, ...] = () + yaml: str + document: dict[str, Any] + + +class LibraryService: + """Read the packaged preset library without exposing package paths.""" + + def __init__( + self, + *, + default_page_size: int = _DEFAULT_PAGE_SIZE, + max_page_size: int = _DEFAULT_MAX_PAGE_SIZE, + ) -> None: + if default_page_size < 1: + raise ValueError("default_page_size must be positive") + if max_page_size < default_page_size: + raise ValueError("max_page_size must be >= default_page_size") + self.default_page_size = default_page_size + self.max_page_size = max_page_size + + def list_presets( + self, + *, + kind: str | PresetKind | None = None, + cursor: str | None = None, + page_size: int | None = None, + ) -> PresetPage: + try: + parsed_kind = PresetKind(kind) if kind is not None else None + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"Unknown preset kind: {kind!r}", + ) from exc + entries = [ + self._catalog_entry(item) + for item in discover(parsed_kind.value if parsed_kind is not None else None) + ] + entries.sort(key=lambda item: (item.kind.value, item.name)) + + limit = self._page_size(page_size) + offset = 0 + if cursor is not None: + payload = _decode_cursor(cursor) + expected_kind = parsed_kind.value if parsed_kind is not None else None + if payload.get("kind") != expected_kind: + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + "Preset cursor does not match the requested filter", + ) + offset = int(payload["offset"]) + if offset < 0 or offset > len(entries): + raise ServiceError( + ServiceErrorCode.STALE_CURSOR, + "Preset cursor is no longer valid", + ) + + items = tuple(entries[offset : offset + limit]) + next_cursor = None + next_offset = offset + len(items) + if next_offset < len(entries): + next_cursor = _encode_cursor( + kind=parsed_kind.value if parsed_kind is not None else None, + offset=next_offset, + ) + return PresetPage(items=items, next_cursor=next_cursor) + + def get_preset( + self, + kind: str | PresetKind, + name: str, + ) -> PresetRecord: + try: + parsed_kind = PresetKind(kind) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"Unknown preset kind: {kind!r}", + ) from exc + if not name or "/" in name or "\\" in name or name in {".", ".."}: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Preset name must be a simple library identifier", + ) + try: + document = load_preset(parsed_kind.value, name) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Preset not found: {parsed_kind.value}/{name}", + ) from exc + + normalized = yaml.safe_dump( + document, + default_flow_style=False, + sort_keys=False, + allow_unicode=True, + ) + return PresetRecord( + kind=parsed_kind, + name=str(document.get("name") or name), + version=_optional_text(document.get("version")), + tags=_tags(document.get("tags")), + yaml=normalized if normalized.endswith("\n") else normalized + "\n", + document=json.loads(json.dumps(document, ensure_ascii=False)), + ) + + def _page_size(self, requested: int | None) -> int: + if requested is None: + return self.default_page_size + if requested < 1 or requested > self.max_page_size: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"page_size must be between 1 and {self.max_page_size}", + ) + return requested + + @staticmethod + def _catalog_entry(item: dict[str, Any]) -> PresetCatalogEntry: + return PresetCatalogEntry( + kind=PresetKind(str(item["kind"])), + name=str(item["name"]), + version=_optional_text(item.get("version")), + tags=_tags(item.get("tags")), + summary=_optional_text(item.get("summary")), + description=_bounded_text(item.get("description"), limit=512), + ) + + +def _tags(value: Any) -> tuple[str, ...]: + if not isinstance(value, list): + return () + return tuple(str(item) for item in value if isinstance(item, str)) + + +def _optional_text(value: Any) -> str | None: + return str(value) if isinstance(value, (str, int, float)) else None + + +def _bounded_text(value: Any, *, limit: int) -> str | None: + text = _optional_text(value) + if text is None or len(text) <= limit: + return text + return text[: limit - 3].rstrip() + "..." + + +def _encode_cursor(*, kind: str | None, offset: int) -> str: + payload = json.dumps( + { + "v": _CURSOR_VERSION, + "kind": kind, + "offset": offset, + }, + separators=(",", ":"), + ).encode("utf-8") + return base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=") + + +def _decode_cursor(cursor: str) -> dict[str, Any]: + try: + padding = "=" * (-len(cursor) % 4) + payload = json.loads( + base64.urlsafe_b64decode(cursor + padding).decode("utf-8") + ) + except (ValueError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid preset cursor", + ) from exc + if ( + not isinstance(payload, dict) + or payload.get("v") != _CURSOR_VERSION + or payload.get("kind") not in {None, *(kind.value for kind in PresetKind)} + or not isinstance(payload.get("offset"), int) + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid preset cursor", + ) + return payload diff --git a/tests/test_artifact_service.py b/tests/test_artifact_service.py new file mode 100644 index 000000000..8da46fa02 --- /dev/null +++ b/tests/test_artifact_service.py @@ -0,0 +1,216 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +from assert_ai.core.io import write_json +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.artifacts import ArtifactRepository +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.results import ResultRepository +from tests.result_catalog_fixture import create_result_catalog_fixture + + +def _repository(tmp_path: Path) -> tuple[ArtifactRepository, str, str]: + workspace = WorkspaceService.create(tmp_path) + fixture = create_result_catalog_fixture( + workspace.artifacts_root, + suite_count=1, + runs_per_suite=1, + large_test_case_count=5, + ) + run_root = ( + workspace.results_root + / fixture.large_suite_id + / fixture.large_run_id + ) + write_json( + run_root / "manifest.json", + { + "status": "completed", + "authorization": "not-a-real-secret", + "workspace": str(workspace.root), + }, + ) + (run_root / "config.yaml").write_text( + "pipeline: {}\napi_key: not-a-real-secret\n", + encoding="utf-8", + ) + (run_root / "binary.bin").write_bytes(b"\x00\x01\x02") + + results = ResultRepository( + workspace.results_root, + path_policy=workspace.path_policy, + default_page_size=2, + max_page_size=10, + ) + return ( + ArtifactRepository( + workspace, + results, + default_page_size=2, + max_page_size=10, + default_chunk_bytes=16, + max_chunk_bytes=2048, + ), + fixture.large_suite_id, + fixture.large_run_id, + ) + + +def test_artifact_catalog_uses_opaque_path_free_ids(tmp_path: Path) -> None: + repository, suite_id, run_id = _repository(tmp_path) + + page = repository.list_artifacts(suite_id, run_id=run_id) + + assert page.items + assert page.next_cursor is not None + assert all(item.artifact_id.startswith("art1_") for item in page.items) + serialized = json.dumps(page.model_dump(mode="json")) + assert str(tmp_path) not in serialized + assert "config.yaml" not in { + item.artifact_id for item in page.items + } + + +def test_artifact_text_chunks_are_redacted_and_resumable(tmp_path: Path) -> None: + repository, suite_id, run_id = _repository(tmp_path) + descriptor = repository.find_artifact( + suite_id, + "config", + run_id=run_id, + ) + + chunks: list[str] = [] + offset = 0 + while True: + chunk = repository.read_artifact_chunk( + descriptor.artifact_id, + offset=offset, + chunk_size=16, + ) + chunks.append(chunk.data) + if chunk.eof: + break + assert chunk.next_offset is not None + assert chunk.next_offset > offset + offset = chunk.next_offset + + content = "".join(chunks) + assert "not-a-real-secret" not in content + assert "[REDACTED]" in content + assert "pipeline" in content + assert yaml.safe_load(content)["api_key"] == "[REDACTED]" + + manifest = repository.find_artifact( + suite_id, + "manifest", + run_id=run_id, + ) + manifest_chunk = repository.read_artifact_chunk( + manifest.artifact_id, + chunk_size=2048, + ) + assert json.loads(manifest_chunk.data)["authorization"] == "[REDACTED]" + assert str(tmp_path) not in manifest_chunk.data + + +def test_artifact_binary_chunks_are_not_exposed(tmp_path: Path) -> None: + repository, suite_id, run_id = _repository(tmp_path) + run_root = repository.workspace.results_root / suite_id / run_id + summary = json.loads((run_root / "run_summary.json").read_text(encoding="utf-8")) + summary["sources"]["binary"] = { + "scope": "run", + "path": "binary.bin", + } + write_json(run_root / "run_summary.json", summary) + + descriptor = repository.find_artifact( + suite_id, + "binary", + run_id=run_id, + ) + with pytest.raises(ServiceError) as exc_info: + repository.read_artifact_chunk(descriptor.artifact_id) + + assert exc_info.value.code is ServiceErrorCode.CAPABILITY_DISABLED + + +def test_oversized_text_artifacts_are_not_quadratically_streamed( + tmp_path: Path, +) -> None: + repository, suite_id, run_id = _repository(tmp_path) + run_root = repository.workspace.results_root / suite_id / run_id + large_path = run_root / "large.log" + large_path.write_text("x" * 1025, encoding="utf-8") + summary = json.loads((run_root / "run_summary.json").read_text(encoding="utf-8")) + summary["sources"]["large"] = { + "scope": "run", + "path": "large.log", + } + write_json(run_root / "run_summary.json", summary) + repository.max_text_artifact_bytes = 1024 + + descriptor = repository.find_artifact( + suite_id, + "large", + run_id=run_id, + ) + with pytest.raises(ServiceError) as exc_info: + repository.read_artifact_chunk(descriptor.artifact_id) + + assert exc_info.value.code is ServiceErrorCode.ARTIFACT_TOO_LARGE + + +def test_artifact_ids_and_cursors_fail_stale_after_change( + tmp_path: Path, +) -> None: + repository, suite_id, run_id = _repository(tmp_path) + page = repository.list_artifacts( + suite_id, + run_id=run_id, + page_size=1, + ) + assert page.next_cursor is not None + descriptor = repository.find_artifact( + suite_id, + "config", + run_id=run_id, + ) + config_path = repository.workspace.results_root / suite_id / run_id / "config.yaml" + config_path.write_text("pipeline: {}\nchanged: true\n", encoding="utf-8") + + with pytest.raises(ServiceError) as exc_info: + repository.read_artifact_chunk(descriptor.artifact_id) + assert exc_info.value.code is ServiceErrorCode.STALE_ETAG + + with pytest.raises(ServiceError) as exc_info: + repository.list_artifacts( + suite_id, + run_id=run_id, + cursor=page.next_cursor, + page_size=1, + ) + assert exc_info.value.code is ServiceErrorCode.STALE_CURSOR + + +def test_artifact_manifest_cannot_escape_workspace(tmp_path: Path) -> None: + repository, suite_id, run_id = _repository(tmp_path) + run_root = repository.workspace.results_root / suite_id / run_id + summary = json.loads((run_root / "run_summary.json").read_text(encoding="utf-8")) + summary["sources"]["escape"] = { + "scope": "run", + "path": "../../../outside.txt", + } + write_json(run_root / "run_summary.json", summary) + + with pytest.raises(ServiceError) as exc_info: + repository.list_artifacts(suite_id, run_id=run_id) + + assert exc_info.value.code is ServiceErrorCode.WORKSPACE_VIOLATION diff --git a/tests/test_library_service.py b/tests/test_library_service.py new file mode 100644 index 000000000..7a7188029 --- /dev/null +++ b/tests/test_library_service.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import pytest + +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.library import LibraryService, PresetKind + + +def test_library_service_lists_bounded_path_free_presets() -> None: + service = LibraryService(default_page_size=2, max_page_size=3) + + first = service.list_presets(page_size=2) + second = service.list_presets(cursor=first.next_cursor, page_size=2) + + assert len(first.items) == 2 + assert first.next_cursor is not None + assert second.items + assert {item.kind for item in first.items} <= set(PresetKind) + assert all(not hasattr(item, "path") for item in first.items) + + +def test_library_service_cursor_is_bound_to_kind() -> None: + service = LibraryService(default_page_size=1, max_page_size=2) + first = service.list_presets(kind=PresetKind.BEHAVIOR, page_size=1) + assert first.next_cursor is not None + + with pytest.raises(ServiceError) as exc_info: + service.list_presets( + kind=PresetKind.JUDGE_PRESET, + cursor=first.next_cursor, + page_size=1, + ) + + assert exc_info.value.code is ServiceErrorCode.STALE_CURSOR + + +def test_library_service_gets_complete_preset() -> None: + service = LibraryService() + + preset = service.get_preset(PresetKind.BEHAVIOR, "prompt_injection") + + assert preset.kind is PresetKind.BEHAVIOR + assert preset.name == "prompt_injection" + assert preset.document["kind"] == "behavior" + assert "description:" in preset.yaml + + +def test_library_service_rejects_unknown_or_path_like_names() -> None: + service = LibraryService() + + with pytest.raises(ServiceError) as exc_info: + service.get_preset("behavior", "../prompt_injection") + assert exc_info.value.code is ServiceErrorCode.INVALID_ARGUMENT + + with pytest.raises(ServiceError) as exc_info: + service.get_preset("behavior", "missing-preset") + assert exc_info.value.code is ServiceErrorCode.NOT_FOUND diff --git a/tests/test_mcp_cli.py b/tests/test_mcp_cli.py index 78883de69..5ac393db5 100644 --- a/tests/test_mcp_cli.py +++ b/tests/test_mcp_cli.py @@ -47,6 +47,18 @@ def test_mcp_serve_forwards_resolved_options() -> None: "author", "--enable-group", "design", + "--default-page-size", + "10", + "--max-page-size", + "25", + "--max-response-bytes", + "8192", + "--default-artifact-chunk-bytes", + "1024", + "--max-artifact-chunk-bytes", + "2048", + "--max-config-bytes", + "4096", ], ) @@ -55,6 +67,12 @@ def test_mcp_serve_forwards_resolved_options() -> None: assert create_kwargs["workspace_root"].is_absolute() assert create_kwargs["mode"] == "author" assert create_kwargs["enabled_groups"] == ("design",) + assert create_kwargs["default_page_size"] == 10 + assert create_kwargs["max_page_size"] == 25 + assert create_kwargs["max_response_bytes"] == 8192 + assert create_kwargs["default_artifact_chunk_bytes"] == 1024 + assert create_kwargs["max_artifact_chunk_bytes"] == 2048 + assert create_kwargs["max_config_bytes"] == 4096 run_stdio_server.assert_called_once_with(options) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 4c61aad27..16de7aa62 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -4,11 +4,13 @@ from __future__ import annotations import asyncio +import hashlib +import json import os import sys from contextlib import asynccontextmanager from pathlib import Path -from typing import AsyncIterator +from typing import Any, AsyncIterator import pytest @@ -20,6 +22,40 @@ from assert_ai.mcp.models import CapabilityGroup, ServerMode from assert_ai.mcp.server import ServerOptions, build_server +from tests.result_catalog_fixture import create_result_catalog_fixture + +EXPECTED_INSPECT_TOOLS = { + "get_server_info", + "list_presets", + "get_preset", + "get_config_schema", + "list_configs", + "get_config", + "list_suites", + "get_suite", + "list_runs", + "get_run", + "compare_runs", + "list_test_cases", + "get_test_case", + "list_scores", + "list_failures", + "get_transcript", + "list_artifacts", + "read_artifact_chunk", +} + +EXPECTED_RESOURCE_TEMPLATES = { + "assert://preset/{kind}/{name}", + "assert://config/{config_ref}", + "assert://suite/{suite_id}/taxonomy", + "assert://suite/{suite_id}/test-case/{test_case_id}{?kind,run_id}", + "assert://run/{suite_id}/{run_id}/summary", + "assert://run/{suite_id}/{run_id}/manifest", + "assert://run/{suite_id}/{run_id}/config", + "assert://run/{suite_id}/{run_id}/transcript/{test_case_id}{?kind}", + "assert://artifact/{artifact_id}", +} @asynccontextmanager @@ -50,6 +86,181 @@ async def _stdio_transport( yield streams +def _write_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload), encoding="utf-8") + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(row) + "\n" for row in rows), + encoding="utf-8", + ) + + +def _score(test_case_id: str, *, violation: bool) -> dict[str, Any]: + return { + "type": "prompt", + "test_case_id": test_case_id, + "behavior": "unsafe-action", + "target": "fixture-target", + "judge_model": "fixture-judge", + "judge_status": "ok", + "score_keys": ["policy_violation"], + "not_applicable_score_keys": [], + "verdict": { + "dimensions": {"policy_violation": violation}, + "node_judgments": [], + }, + } + + +def _seed_workspace(root: Path) -> None: + config_path = root / "evals" / "nested" / "demo.yaml" + config_path.parent.mkdir(parents=True, exist_ok=True) + config_path.write_text( + "\n".join( + [ + "behavior:", + " name: demo_behavior", + " description: Demo behavior", + "context: 'authorization: not-a-real-secret'", + "pipeline:", + " inference:", + " target:", + " model: azure/demo", + "", + ] + ), + encoding="utf-8", + ) + + suite_root = root / "artifacts" / "results" / "suite-a" + _write_json( + suite_root / "suite.json", + {"created_at": "2026-08-12T00:00:00+00:00"}, + ) + _write_json( + suite_root / "taxonomy.json", + { + "behavior": { + "name": "demo_behavior", + "description": "Demo behavior", + }, + "behavior_categories": [ + {"name": "unsafe-action", "permissible": False}, + ], + }, + ) + _write_jsonl( + suite_root / "test_set.jsonl", + [ + { + "type": "prompt", + "test_case_id": "p1", + "dimensions": { + "behavior": "unsafe-action", + "region": "us", + }, + "seed": {"prompt": "First prompt"}, + }, + { + "type": "prompt", + "test_case_id": "p2", + "dimensions": { + "behavior": "unsafe-action", + "region": "eu", + }, + "seed": {"prompt": "Second prompt"}, + }, + ], + ) + + for index, run_id in enumerate(("run-a", "run-b")): + run_root = suite_root / run_id + _write_json( + run_root / "manifest.json", + { + "status": "completed", + "started_at": f"2026-08-12T00:0{index}:00+00:00", + "ended_at": f"2026-08-12T00:0{index + 1}:00+00:00", + "stages": { + "inference": "completed", + "judge": "completed", + }, + }, + ) + (run_root / "config.yaml").write_text( + "pipeline: {}\napi_key: not-a-real-secret\n", + encoding="utf-8", + ) + _write_jsonl( + run_root / "inference_set.jsonl", + [ + { + "type": "prompt", + "test_case_id": "p1", + "target": "fixture-target", + "stop_reason": "completed", + "events": [ + {"role": "user", "content": "First prompt"}, + { + "role": "assistant", + "content": "First response", + }, + { + "edit": { + "type": "tool_call", + "authorization": "not-a-real-secret", + } + }, + ], + }, + { + "type": "prompt", + "test_case_id": "p2", + "target": "fixture-target", + "stop_reason": "completed", + "events": [ + {"role": "user", "content": "Second prompt"}, + { + "role": "assistant", + "content": "Second response", + }, + ], + }, + ], + ) + _write_jsonl( + run_root / "scores.jsonl", + [ + _score("p1", violation=index == 0), + _score("p2", violation=False), + ], + ) + + +def _schema_digest(tool: Any) -> str: + payload = { + "input": tool.input_schema, + "output": tool.output_schema, + } + return hashlib.sha256( + json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest() + + +def _error_text(result: Any) -> str: + assert result.is_error is True + assert result.content + return str(result.content[0].text) + + def test_server_options_resolve_workspace_and_capabilities(tmp_path: Path) -> None: workspace = tmp_path / "workspace" workspace.mkdir() @@ -82,6 +293,16 @@ def test_server_options_direct_constructor_preserves_workspace_root_api( assert options.path_policy.workspace_root == tmp_path.resolve() +def test_server_options_validate_response_limits(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="must not exceed"): + ServerOptions( + workspace_root=tmp_path, + max_response_bytes=4096, + default_artifact_chunk_bytes=512, + max_artifact_chunk_bytes=3072, + ) + + def test_design_group_requires_author_or_full_mode(tmp_path: Path) -> None: with pytest.raises(ValueError, match="require --mode author or --mode full"): ServerOptions.create( @@ -91,27 +312,39 @@ def test_design_group_requires_author_or_full_mode(tmp_path: Path) -> None: ) +@pytest.mark.parametrize("mode", list(ServerMode)) +def test_inspect_tools_are_registered_in_every_base_mode( + tmp_path: Path, + mode: ServerMode, +) -> None: + async def run() -> set[str]: + options = ServerOptions.create(workspace_root=tmp_path, mode=mode) + async with Client(build_server(options), raise_exceptions=True) as client: + tools = await client.list_tools() + return {tool.name for tool in tools.tools} + + assert asyncio.run(run()) == EXPECTED_INSPECT_TOOLS + + def test_get_server_info_protocol_round_trip(tmp_path: Path) -> None: - async def run() -> tuple[set[str], object]: + async def run() -> object: options = ServerOptions.create( workspace_root=tmp_path, mode="full", enabled_groups=["analysis"], ) async with Client(build_server(options), raise_exceptions=True) as client: - tools = await client.list_tools() - result = await client.call_tool("get_server_info", {}) - return {tool.name for tool in tools.tools}, result + return await client.call_tool("get_server_info", {}) - tool_names, result = asyncio.run(run()) + result = asyncio.run(run()) - assert tool_names == {"get_server_info"} assert result.is_error is False assert result.structured_content is not None assert result.structured_content["assert_mcp_api_version"] == "1" assert result.structured_content["mode"] == "full" assert result.structured_content["workspace"]["root"] == "." assert "env_file" not in result.structured_content + assert result.structured_content["limits"]["max_page_size"] == 200 assert result.structured_content["enabled_capability_groups"] == [ "inspect", "author", @@ -123,30 +356,413 @@ async def run() -> tuple[set[str], object]: ] -def test_get_server_info_publishes_structured_output_schema(tmp_path: Path) -> None: - async def run() -> object: +def test_all_tools_publish_stable_schemas_and_read_only_annotations( + tmp_path: Path, +) -> None: + async def run() -> list[Any]: options = ServerOptions.create(workspace_root=tmp_path) async with Client(build_server(options), raise_exceptions=True) as client: - tools = await client.list_tools() - return tools.tools[0] + return (await client.list_tools()).tools - tool = asyncio.run(run()) + tools = asyncio.run(run()) - assert tool.name == "get_server_info" - assert tool.output_schema is not None - assert "assert_mcp_api_version" in tool.output_schema["properties"] - assert tool.annotations is not None - assert tool.annotations.read_only_hint is True - assert tool.annotations.open_world_hint is False + assert {tool.name for tool in tools} == EXPECTED_INSPECT_TOOLS + for tool in tools: + assert tool.input_schema is not None + assert tool.output_schema is not None + assert tool.annotations is not None + assert tool.annotations.read_only_hint is True + assert tool.annotations.destructive_hint is False + assert tool.annotations.idempotent_hint is True + assert tool.annotations.open_world_hint is False + digests = {tool.name: _schema_digest(tool) for tool in tools} + assert digests == { + # These hashes are protocol snapshots. Update them only for an + # intentional API-v1 schema change. + "compare_runs": "f7bfeca051f8f81bf3621936588ed906332076a3a34b550090f87c2944656ce5", + "get_config": "bf38188871cb818e0b0cf6e28183aa728ed8d041923a158f593832d2459bd13a", + "get_config_schema": "cca1d3a48240e20eff93a123b34d7ba92df3ed1df87f57f9eb217aa21515ec26", + "get_preset": "25352522a3ed4ff76217c5415453c6641ad229c2dd6e4e05df4621b89ce4819c", + "get_run": "e5216cd0085d049f8b49c54add913b6f83756c4ce59995317fe63e010ea44936", + "get_server_info": "d51f9ff9fe235b5c53f5db71a1bba11cfb35bbc27be1f8c3552f5bee8ecc5e8d", + "get_suite": "8f629c93e02b656052f637c3cbba9217834315a693c4f7935f6d961203b46fd0", + "get_test_case": "11380555caaa71d5992923815a499fc08b368c02f4d4836e4761630654589148", + "get_transcript": "aa09669e0cb99202e8dec0b858b4faa41742ecb616351c3956b0d0bd488717e8", + "list_artifacts": "3d3bede0b7209401b15d1f39d82671092c3097a05cd901122bd46c3c42edebfc", + "list_configs": "92f78db2533034e6bf80e1d95089460acdd40a18468d4eb06fdf055726dfef19", + "list_failures": "d3cc3f3bcc86c110754673297d28ac1e5ccbf698668c997de0bba2d0cbd425e2", + "list_presets": "2cb18ca86885b3dcfb230437413fef501624c0c989f26ea7f57fb56d3ff3d557", + "list_runs": "7280687daafcd7ff5d89756c9584ca06c432a44f5a98ce8ff3ae0e4427dcf40b", + "list_scores": "5c1951a3a3b91089b68b30e970a1b13f59bc2659a2c234db4451cbe2d5362a4d", + "list_suites": "80f696cdb1812636e4e98089ae7a5b9fceddfc8ae994ab8e9835307c8ea12108", + "list_test_cases": "5f46d6640668d017db8afb98d700113cfc671f93d43d57d425f82c773a6e0906", + "read_artifact_chunk": "895f33b2e66a44cca276f94348155563fec3eab4780fe49eff101835e0c15ab7", + } -def test_stdio_module_entry_point_keeps_protocol_wire_clean(tmp_path: Path) -> None: + +def test_complete_read_only_tool_workflow(tmp_path: Path) -> None: + _seed_workspace(tmp_path) + + async def run() -> dict[str, Any]: + async with Client( + build_server(ServerOptions.create(workspace_root=tmp_path)), + raise_exceptions=True, + ) as client: + preset_page = await client.call_tool( + "list_presets", + {"kind": "behavior", "page_size": 2}, + ) + preset_name = preset_page.structured_content["items"][0]["name"] + preset = await client.call_tool( + "get_preset", + {"kind": "behavior", "name": preset_name}, + ) + schema = await client.call_tool("get_config_schema", {}) + configs = await client.call_tool("list_configs", {}) + config = await client.call_tool( + "get_config", + {"config_ref": "nested/demo.yaml"}, + ) + suites = await client.call_tool("list_suites", {"page_size": 1}) + suite = await client.call_tool( + "get_suite", + {"suite_id": "suite-a"}, + ) + runs = await client.call_tool( + "list_runs", + {"suite_id": "suite-a", "page_size": 1}, + ) + run = await client.call_tool( + "get_run", + {"suite_id": "suite-a", "run_id": "run-a"}, + ) + comparison = await client.call_tool( + "compare_runs", + { + "run_refs": [ + {"suite_id": "suite-a", "run_id": "run-a"}, + {"suite_id": "suite-a", "run_id": "run-b"}, + ] + }, + ) + test_cases = await client.call_tool( + "list_test_cases", + { + "suite_id": "suite-a", + "page_size": 1, + "factors": {"region": "us"}, + }, + ) + test_case = await client.call_tool( + "get_test_case", + { + "suite_id": "suite-a", + "test_case_id": "p1", + "kind": "prompt", + }, + ) + scores = await client.call_tool( + "list_scores", + { + "suite_id": "suite-a", + "run_id": "run-a", + "dimension": "policy_violation", + "dimension_value": True, + }, + ) + failures = await client.call_tool( + "list_failures", + {"suite_id": "suite-a", "run_id": "run-a"}, + ) + transcript = await client.call_tool( + "get_transcript", + { + "suite_id": "suite-a", + "run_id": "run-a", + "test_case_id": "p1", + "kind": "prompt", + }, + ) + artifacts = await client.call_tool( + "list_artifacts", + {"suite_id": "suite-a", "run_id": "run-a"}, + ) + config_artifact = next( + item + for item in artifacts.structured_content["items"] + if item["name"] == "config" + ) + artifact_chunk = await client.call_tool( + "read_artifact_chunk", + { + "artifact_id": config_artifact["artifact_id"], + "chunk_size": 128, + }, + ) + return { + "preset_page": preset_page, + "preset": preset, + "schema": schema, + "configs": configs, + "config": config, + "suites": suites, + "suite": suite, + "runs": runs, + "run": run, + "comparison": comparison, + "test_cases": test_cases, + "test_case": test_case, + "scores": scores, + "failures": failures, + "transcript": transcript, + "artifacts": artifacts, + "artifact_chunk": artifact_chunk, + } + + results = asyncio.run(run()) + assert all(not result.is_error for result in results.values()) + assert results["preset"].structured_content["document"]["kind"] == "behavior" + assert results["schema"].structured_content["json_schema"]["$schema"].endswith( + "2020-12/schema" + ) + assert results["configs"].structured_content["items"][0]["config_ref"] == ( + "nested/demo.yaml" + ) + config_payload = results["config"].structured_content + assert "not-a-real-secret" not in json.dumps(config_payload) + assert "[REDACTED]" in config_payload["yaml"] + assert results["suites"].structured_content["items"][0]["suite_id"] == "suite-a" + assert results["suite"].structured_content["run_count"] == 2 + assert results["runs"].structured_content["next_cursor"] is not None + assert results["run"].structured_content["state"] == "completed" + assert results["comparison"].structured_content["baseline"] == "suite-a/run-a" + assert len(results["test_cases"].structured_content["items"]) == 1 + assert results["test_case"].structured_content["row"]["test_case_id"] == "p1" + assert len(results["scores"].structured_content["items"]) == 1 + assert len(results["failures"].structured_content["items"]) == 1 + transcript = results["transcript"].structured_content + assert transcript["inference"]["events"] + assert "not-a-real-secret" not in json.dumps(transcript) + artifacts = results["artifacts"].structured_content + assert str(tmp_path) not in json.dumps(artifacts) + chunk = results["artifact_chunk"].structured_content + assert chunk["encoding"] == "utf-8" + assert "not-a-real-secret" not in chunk["data"] + assert "[REDACTED]" in chunk["data"] + + +def test_resources_are_lazy_path_free_and_readable(tmp_path: Path) -> None: + _seed_workspace(tmp_path) + + async def run() -> tuple[set[str], set[str], dict[str, str]]: + async with Client( + build_server(ServerOptions.create(workspace_root=tmp_path)), + raise_exceptions=True, + ) as client: + resources = await client.list_resources() + templates = await client.list_resource_templates() + artifacts = await client.call_tool( + "list_artifacts", + {"suite_id": "suite-a", "run_id": "run-a"}, + ) + config_artifact = next( + item + for item in artifacts.structured_content["items"] + if item["name"] == "config" + ) + uris = { + "schema": "assert://schema/eval-config", + "preset": "assert://preset/behavior/prompt_injection", + "config": "assert://config/nested%2Fdemo.yaml", + "taxonomy": "assert://suite/suite-a/taxonomy", + "test_case": ( + "assert://suite/suite-a/test-case/p1?kind=prompt" + ), + "summary": "assert://run/suite-a/run-a/summary", + "manifest": "assert://run/suite-a/run-a/manifest", + "run_config": "assert://run/suite-a/run-a/config", + "transcript": ( + "assert://run/suite-a/run-a/transcript/p1?kind=prompt" + ), + "artifact": config_artifact["resource_uri"], + } + contents = {} + for name, uri in uris.items(): + result = await client.read_resource(uri) + contents[name] = result.contents[0].text + return ( + {str(resource.uri) for resource in resources.resources}, + { + template.uri_template + for template in templates.resource_templates + }, + contents, + ) + + static_resources, templates, contents = asyncio.run(run()) + + assert static_resources == {"assert://schema/eval-config"} + assert templates == EXPECTED_RESOURCE_TEMPLATES + assert "json_schema" in contents["schema"] + assert '"kind": "behavior"' in contents["preset"] + assert "demo_behavior" in contents["config"] + assert "not-a-real-secret" not in json.dumps(contents) + assert "unsafe-action" in contents["taxonomy"] + assert '"test_case_id": "p1"' in contents["test_case"] + assert '"run_id": "run-a"' in contents["summary"] + assert '"status": "completed"' in contents["manifest"] + assert "[REDACTED]" in contents["run_config"] + assert "First response" in contents["transcript"] + assert "[REDACTED]" in contents["artifact"] + assert str(tmp_path) not in json.dumps(contents) + + +@pytest.mark.parametrize( + ("mode", "raise_exceptions"), + [("auto", True), ("legacy", False)], +) +def test_service_errors_are_stable_tool_errors( + tmp_path: Path, + mode: str, + raise_exceptions: bool, +) -> None: async def run() -> object: - async with Client(_stdio_transport(tmp_path), raise_exceptions=True) as client: - return await client.call_tool("get_server_info", {}) + async with Client( + build_server(ServerOptions.create(workspace_root=tmp_path)), + mode=mode, + raise_exceptions=raise_exceptions, + ) as client: + return await client.call_tool( + "get_run", + {"suite_id": "missing", "run_id": "missing"}, + ) + + text = _error_text(asyncio.run(run())) + + assert '"code":"NOT_FOUND"' in text + assert str(tmp_path) not in text + + +def test_result_cursor_reports_stale_source_through_mcp(tmp_path: Path) -> None: + _seed_workspace(tmp_path) + + async def run() -> object: + async with Client( + build_server(ServerOptions.create(workspace_root=tmp_path)), + raise_exceptions=True, + ) as client: + first = await client.call_tool( + "list_test_cases", + {"suite_id": "suite-a", "page_size": 1}, + ) + cursor = first.structured_content["next_cursor"] + assert cursor is not None + test_set = ( + tmp_path + / "artifacts" + / "results" + / "suite-a" + / "test_set.jsonl" + ) + with test_set.open("a", encoding="utf-8") as handle: + handle.write( + json.dumps( + { + "type": "prompt", + "test_case_id": "p3", + "seed": {"prompt": "Third prompt"}, + } + ) + + "\n" + ) + return await client.call_tool( + "list_test_cases", + { + "suite_id": "suite-a", + "cursor": cursor, + "page_size": 1, + }, + ) result = asyncio.run(run()) - assert result.is_error is False - assert result.structured_content is not None - assert result.structured_content["workspace"]["root"] == "." + assert '"code":"STALE_CURSOR"' in _error_text(result) + + +def test_tool_response_limit_returns_bounded_error(tmp_path: Path) -> None: + async def run() -> object: + options = ServerOptions( + workspace_root=tmp_path, + max_response_bytes=4096, + default_artifact_chunk_bytes=512, + max_artifact_chunk_bytes=1024, + max_config_bytes=1024, + ) + async with Client(build_server(options), raise_exceptions=True) as client: + return await client.call_tool("get_config_schema", {}) + + text = _error_text(asyncio.run(run())) + + assert '"code":"ARTIFACT_TOO_LARGE"' in text + assert len(text.encode("utf-8")) < 1024 + + +@pytest.mark.timeout(60) +def test_scale_fixture_supports_read_only_mcp_workflow(tmp_path: Path) -> None: + fixture = create_result_catalog_fixture( + tmp_path / "artifacts", + suite_count=100, + runs_per_suite=10, + large_test_case_count=10_000, + ) + + async def run() -> tuple[object, object, object]: + async with Client( + build_server(ServerOptions.create(workspace_root=tmp_path)), + raise_exceptions=True, + ) as client: + suites = await client.call_tool("list_suites", {"page_size": 5}) + runs = await client.call_tool( + "list_runs", + {"suite_id": fixture.large_suite_id, "page_size": 5}, + ) + test_case = await client.call_tool( + "get_test_case", + { + "suite_id": fixture.large_suite_id, + "test_case_id": fixture.last_test_case_id, + "kind": "prompt", + }, + ) + return suites, runs, test_case + + suites, runs, test_case = asyncio.run(run()) + + assert len(suites.structured_content["items"]) == 5 + assert suites.structured_content["next_cursor"] is not None + assert len(runs.structured_content["items"]) == 5 + assert runs.structured_content["next_cursor"] is not None + assert ( + test_case.structured_content["row"]["test_case_id"] + == fixture.last_test_case_id + ) + + +def test_stdio_module_entry_point_keeps_protocol_wire_clean(tmp_path: Path) -> None: + async def run() -> tuple[object, object]: + async with Client(_stdio_transport(tmp_path), raise_exceptions=True) as client: + info = await client.call_tool("get_server_info", {}) + presets = await client.call_tool( + "list_presets", + {"page_size": 1}, + ) + return info, presets + + info, presets = asyncio.run(run()) + + assert info.is_error is False + assert info.structured_content is not None + assert info.structured_content["workspace"]["root"] == "." + assert presets.is_error is False + assert len(presets.structured_content["items"]) == 1 diff --git a/tests/test_security.py b/tests/test_security.py index 12cfd77f2..5dd7d9e47 100644 --- a/tests/test_security.py +++ b/tests/test_security.py @@ -14,6 +14,7 @@ from assert_ai.core.security import ( sanitize_payload, + sanitize_text, validate_callable_ref, validate_endpoint_url, validate_module_ref, @@ -346,6 +347,15 @@ def test_shallow_payload_unaffected_by_max_depth(self) -> None: class SanitizeResponseTextTest(unittest.TestCase): + def test_preserves_structured_json_while_redacting_assignment(self) -> None: + text = '{"authorization": "not-a-real-secret-value", "ok": true}' + result = sanitize_text(text) + + self.assertEqual( + json.loads(result), + {"authorization": "[REDACTED]", "ok": True}, + ) + def test_redacts_bearer_token(self) -> None: text = "Your token is Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.longtoken" result = _sanitize_response_text(text) From 86fbc6474cda167509c3fc150faa8889ec935828 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Tue, 25 Aug 2026 08:34:11 -0700 Subject: [PATCH 07/16] Add persisted run catalogs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/services/result_metadata.py | 80 ++++++++++++++- assert_ai/services/results.py | 119 +++++++++++++++++----- tests/result_catalog_fixture.py | 138 ++++++++++++++------------ tests/test_result_metadata.py | 13 +++ tests/test_result_service.py | 120 +++++++++++++++++++++- tests/test_result_service_scale.py | 13 ++- 6 files changed, 387 insertions(+), 96 deletions(-) diff --git a/assert_ai/services/result_metadata.py b/assert_ai/services/result_metadata.py index dcb92726f..ce2d8a050 100644 --- a/assert_ai/services/result_metadata.py +++ b/assert_ai/services/result_metadata.py @@ -28,6 +28,8 @@ SUITE_SUMMARY_SCHEMA_VERSION = 1 RUN_SUMMARY_SCHEMA_VERSION = 1 +RUN_CATALOG_SCHEMA_VERSION = 1 +RUN_CATALOG_FILENAME = "run_catalog.json" def refresh_stage_indexes( @@ -201,6 +203,74 @@ def write_run_summary( return normalized_payload +def run_catalog_entry( + summary: dict[str, Any], + *, + suite_id: str | None = None, +) -> dict[str, Any]: + """Project a run summary into the lightweight catalog contract.""" + quality = summary.get("quality") + if not isinstance(quality, dict): + quality = {} + return { + "suite_id": summary.get("suite_id") or suite_id, + "run_id": summary.get("run_id"), + "status": summary.get("state"), + "current_stage": summary.get("current_stage"), + "started_at": summary.get("started_at"), + "ended_at": summary.get("ended_at"), + "updated_at": summary.get("updated_at"), + "prompt_metrics": quality.get("prompt"), + "scenario_metrics": quality.get("scenario"), + "models": summary.get("models") or {}, + "counts": summary.get("counts") or {}, + "metrics": summary.get("metrics"), + } + + +def write_run_catalog( + suite_root: Path, + run_summaries: list[dict[str, Any]], + *, + catalog_identity: dict[str, Any] | None = None, +) -> dict[str, Any] | None: + """Atomically persist projected run metadata for one stable suite snapshot.""" + expected_identity = ( + suite_run_catalog_identity(suite_root) + if catalog_identity is None + else catalog_identity + ) + if suite_run_catalog_identity(suite_root) != expected_identity: + return None + + payload = _json_payload( + { + "schema_version": RUN_CATALOG_SCHEMA_VERSION, + "suite_id": suite_root.name, + "generated_at": _utc_now(), + "run_catalog_identity": expected_identity, + "items": [ + run_catalog_entry(summary, suite_id=suite_root.name) + for summary in run_summaries + ], + "summary_sources": { + str(summary.get("run_id")): ( + summary.get("sources") + if isinstance(summary.get("sources"), dict) + else {} + ) + for summary in run_summaries + if isinstance(summary.get("run_id"), str) + }, + } + ) + write_json(suite_root / RUN_CATALOG_FILENAME, payload) + + if suite_run_catalog_identity(suite_root) != expected_identity: + return None + return payload + + def write_suite_summary( ctx: dict[str, Any], *, @@ -274,7 +344,15 @@ def write_suite_summary( ), } + catalog_identity_before = suite_run_catalog_identity(suite_root) runs = _run_catalog_entries(suite_root) + catalog_identity = suite_run_catalog_identity(suite_root) + if catalog_identity_before == catalog_identity: + write_run_catalog( + suite_root, + runs, + catalog_identity=catalog_identity, + ) latest_run = max( runs, key=lambda item: str( @@ -331,7 +409,7 @@ def write_suite_summary( "updated_at": _utc_now(), "run_count": len(runs), "run_set_identity": suite_run_set_identity(suite_root), - "run_catalog_identity": suite_run_catalog_identity(suite_root), + "run_catalog_identity": catalog_identity, "latest_run": ( { "run_id": latest_run.get("run_id"), diff --git a/assert_ai/services/results.py b/assert_ai/services/results.py index 60b2c0b7c..517f56ca8 100644 --- a/assert_ai/services/results.py +++ b/assert_ai/services/results.py @@ -32,10 +32,14 @@ from assert_ai.core.judge import get_verdict_dimension, infer_judge_status from assert_ai.services.errors import ServiceError, ServiceErrorCode from assert_ai.services.result_metadata import ( + RUN_CATALOG_FILENAME, + RUN_CATALOG_SCHEMA_VERSION, RUN_SUMMARY_SCHEMA_VERSION, SUITE_SUMMARY_SCHEMA_VERSION, + run_catalog_entry, suite_run_catalog_identity, suite_run_set_identity, + write_run_catalog, write_run_summary, write_suite_summary, ) @@ -163,11 +167,7 @@ def list_run_catalog_entries( page_size: int | None = None, ) -> ResultPage: suite_dir = self._suite_dir(suite_id, must_exist=True) - entries: list[dict[str, Any]] = [] - for run_dir in self._run_dirs(suite_dir): - summary = self._ensure_run_summary(suite_dir, run_dir) - if summary is not None: - entries.append(self._run_catalog_entry(summary)) + entries = self._ensure_run_catalog(suite_dir) entries.sort( key=lambda item: ( str( @@ -693,6 +693,95 @@ def _ensure_suite_summary( ctx = self._legacy_context(suite_dir) return write_suite_summary(ctx, rebuild_indexes=True) + def _ensure_run_catalog( + self, + suite_dir: Path, + ) -> list[dict[str, Any]]: + catalog = _load_optional_json(suite_dir / RUN_CATALOG_FILENAME) + items = catalog.get("items") if isinstance(catalog, dict) else None + if ( + isinstance(catalog, dict) + and catalog.get("schema_version") == RUN_CATALOG_SCHEMA_VERSION + and catalog.get("suite_id") == suite_dir.name + and isinstance(items, list) + and all(isinstance(item, dict) for item in items) + and catalog.get("run_catalog_identity") + == suite_run_catalog_identity(suite_dir) + and self._run_catalog_sources_current(catalog, suite_dir) + ): + return [dict(item) for item in items] + + summaries: list[dict[str, Any]] = [] + for run_dir in self._run_dirs(suite_dir): + summary = self._ensure_run_summary(suite_dir, run_dir) + if summary is not None: + summaries.append(summary) + + catalog_identity = suite_run_catalog_identity(suite_dir) + try: + rebuilt = write_run_catalog( + suite_dir, + summaries, + catalog_identity=catalog_identity, + ) + except OSError: + rebuilt = None + rebuilt_items = ( + rebuilt.get("items") if isinstance(rebuilt, dict) else None + ) + if isinstance(rebuilt_items, list) and all( + isinstance(item, dict) for item in rebuilt_items + ): + return [dict(item) for item in rebuilt_items] + return [ + run_catalog_entry(summary, suite_id=suite_dir.name) + for summary in summaries + ] + + def _run_catalog_sources_current( + self, + catalog: dict[str, Any], + suite_dir: Path, + ) -> bool: + items = catalog.get("items") + summary_sources = catalog.get("summary_sources") + if not isinstance(items, list) or not isinstance( + summary_sources, + dict, + ): + return False + + run_ids: list[str] = [] + for item in items: + run_id = item.get("run_id") if isinstance(item, dict) else None + if not isinstance(run_id, str) or not _IDENTIFIER_RE.fullmatch( + run_id + ): + return False + run_ids.append(run_id) + if ( + len(run_ids) != len(set(run_ids)) + or set(run_ids) != set(summary_sources) + ): + return False + + for run_id in run_ids: + sources = summary_sources.get(run_id) + if not isinstance(sources, dict): + return False + run_dir = self._run_dir( + suite_dir, + run_id, + must_exist=False, + ) + if not run_dir.is_dir() or not self._summary_sources_current( + {"sources": sources}, + suite_dir=suite_dir, + run_dir=run_dir, + ): + return False + return True + def _ensure_run_summary( self, suite_dir: Path, @@ -1319,26 +1408,6 @@ def _suite_catalog_entry( "latest_run": summary.get("latest_run"), } - def _run_catalog_entry( - self, - summary: dict[str, Any], - ) -> dict[str, Any]: - quality = summary.get("quality") or {} - return { - "suite_id": summary.get("suite_id"), - "run_id": summary.get("run_id"), - "status": summary.get("state"), - "current_stage": summary.get("current_stage"), - "started_at": summary.get("started_at"), - "ended_at": summary.get("ended_at"), - "updated_at": summary.get("updated_at"), - "prompt_metrics": quality.get("prompt"), - "scenario_metrics": quality.get("scenario"), - "models": summary.get("models") or {}, - "counts": summary.get("counts") or {}, - "metrics": summary.get("metrics"), - } - def _matches_common( self, row: dict[str, Any], diff --git a/tests/result_catalog_fixture.py b/tests/result_catalog_fixture.py index 9ee28c2fd..e585c98f9 100644 --- a/tests/result_catalog_fixture.py +++ b/tests/result_catalog_fixture.py @@ -14,6 +14,7 @@ SUITE_SUMMARY_SCHEMA_VERSION, suite_run_catalog_identity, suite_run_set_identity, + write_run_catalog, ) @@ -41,6 +42,7 @@ def create_result_catalog_fixture( suite_id = f"suite-{suite_index:03d}" suite_root = results_root / suite_id suite_root.mkdir(parents=True) + run_summaries = [] for run_index in range(runs_per_suite): run_id = f"run-{run_index:03d}" @@ -67,73 +69,83 @@ def create_result_catalog_fixture( f"2026-08-{(suite_index % 28) + 1:02d}" f"T00:{run_index:02d}:00+00:00" ) - write_json( - run_root / "run_summary.json", - { - "schema_version": RUN_SUMMARY_SCHEMA_VERSION, - "suite_id": suite_id, - "run_id": run_id, - "state": "completed", - "current_stage": "judge", - "started_at": timestamp, - "ended_at": timestamp, - "updated_at": timestamp, - "stages": { - "inference": "completed", - "judge": "completed", - }, - "stage_timings": {}, - "stage_summaries": {}, - "models": { - "target": { - "kind": "model", - "identifier": "fixture-target", - }, - "tester": None, - "judge": "fixture-judge", - }, - "counts": { - "scores": { - "total": 1, - "prompt": 1, - "scenario": 0, - "other": 0, - } + run_summary = { + "schema_version": RUN_SUMMARY_SCHEMA_VERSION, + "suite_id": suite_id, + "run_id": run_id, + "state": "completed", + "current_stage": "judge", + "started_at": timestamp, + "ended_at": timestamp, + "updated_at": timestamp, + "stages": { + "inference": "completed", + "judge": "completed", + }, + "stage_timings": {}, + "stage_summaries": {}, + "models": { + "target": { + "kind": "model", + "identifier": "fixture-target", }, - "quality": { - "prompt": { - "total": 1, - "scored_total": 1, - "judge_failures": 0, - "judge_failure_rate": 0.0, - "policy_violation_rate": 0.0, - "overrefusal_rate": None, - "dimensions": { - "policy_violation": { - "rate": 0.0, - "counts": {"0": 1, "1": 0}, - "count": 1, - "applicable_count": 1, - "not_applicable_count": 0, - "flagged_count": 0, - "clear_count": 1, - } - }, - "target": "fixture-target", - "judge_model": "fixture-judge", + "tester": None, + "judge": "fixture-judge", + }, + "counts": { + "scores": { + "total": 1, + "prompt": 1, + "scenario": 0, + "other": 0, + } + }, + "quality": { + "prompt": { + "total": 1, + "scored_total": 1, + "judge_failures": 0, + "judge_failure_rate": 0.0, + "policy_violation_rate": 0.0, + "overrefusal_rate": None, + "dimensions": { + "policy_violation": { + "rate": 0.0, + "counts": {"0": 1, "1": 0}, + "count": 1, + "applicable_count": 1, + "not_applicable_count": 0, + "flagged_count": 0, + "clear_count": 1, + } }, - "scenario": None, + "target": "fixture-target", + "judge_model": "fixture-judge", }, - "metrics": { - "schema_version": 1, - "elapsed_s": 1.0, - "totals": {"calls": 1}, - }, - "artifact_versions": {}, - "sources": {}, - "indexes": {}, + "scenario": None, + }, + "metrics": { + "schema_version": 1, + "elapsed_s": 1.0, + "totals": {"calls": 1}, }, + "artifact_versions": {}, + "sources": {}, + "indexes": {}, + } + write_json(run_root / "run_summary.json", run_summary) + run_summaries.append(run_summary) + + run_catalog_identity = suite_run_catalog_identity(suite_root) + if ( + write_run_catalog( + suite_root, + run_summaries, + catalog_identity=run_catalog_identity, ) + is None + ): + raise RuntimeError("Result catalog fixture changed while being built") sources: dict[str, object] = {} test_case_counts = { @@ -200,7 +212,7 @@ def create_result_catalog_fixture( ), "run_count": runs_per_suite, "run_set_identity": suite_run_set_identity(suite_root), - "run_catalog_identity": suite_run_catalog_identity(suite_root), + "run_catalog_identity": run_catalog_identity, "latest_run": { "run_id": f"run-{runs_per_suite - 1:03d}", "state": "completed", diff --git a/tests/test_result_metadata.py b/tests/test_result_metadata.py index 57083d69d..150e583bc 100644 --- a/tests/test_result_metadata.py +++ b/tests/test_result_metadata.py @@ -213,6 +213,19 @@ def test_suite_summary_uses_active_versioned_test_set_and_metadata_only_runs() - payload["sources"]["test_set"]["path"] == "artifacts/test_set/v0002/test_set.jsonl" ) + run_catalog = json.loads( + ( + root + / "results" + / "suite-a" + / "run_catalog.json" + ).read_text(encoding="utf-8") + ) + assert run_catalog["run_catalog_identity"] == payload[ + "run_catalog_identity" + ] + assert run_catalog["items"][0]["run_id"] == "run-a" + assert run_catalog["items"][0]["status"] == "completed" def test_running_boundary_preserves_last_valid_detail_without_rescanning_rows() -> None: diff --git a/tests/test_result_service.py b/tests/test_result_service.py index 2dda01dc5..dce437d63 100644 --- a/tests/test_result_service.py +++ b/tests/test_result_service.py @@ -218,10 +218,32 @@ def test_catalogs_rebuild_legacy_once_and_then_remain_metadata_only() -> None: assert ( results_root / "suite-a" / "run-a" / "run_summary.json" ).exists() + assert ( + results_root / "suite-a" / "run_catalog.json" + ).exists() - with patch( - "assert_ai.services.results.scan_jsonl", - side_effect=AssertionError("catalog listing must not scan JSONL"), + original_open = Path.open + + def reject_run_summary_reads( + path: Path, + mode: str = "r", + *args: object, + **kwargs: object, + ): + if path.name == "run_summary.json" and "r" in mode: + raise AssertionError( + "warm catalog listing must use run_catalog.json" + ) + return original_open(path, mode, *args, **kwargs) + + with ( + patch.object(Path, "open", reject_run_summary_reads), + patch( + "assert_ai.services.results.scan_jsonl", + side_effect=AssertionError( + "catalog listing must not scan JSONL" + ), + ), ): suites = repository.list_suite_catalog_entries() runs = repository.list_run_catalog_entries("suite-a") @@ -411,6 +433,92 @@ def test_suite_summary_detects_out_of_band_run_addition() -> None: (suite_root / "run-b" / "run_summary.json").unlink() assert repository.get_suite("suite-a")["run_count"] == 2 + assert { + item["run_id"] + for item in repository.list_run_catalog_entries( + "suite-a" + ).items + } == {"run-a", "run-b"} + + +def test_run_catalog_detects_out_of_band_summary_change() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + repository = ResultRepository(results_root) + repository.get_suite("suite-a") + + summary_path = ( + results_root + / "suite-a" + / "run-a" + / "run_summary.json" + ) + summary = json.loads(summary_path.read_text(encoding="utf-8")) + summary["state"] = "failed" + summary["updated_at"] = "2026-08-25T00:00:00+00:00" + _write_json(summary_path, summary) + + runs = repository.list_run_catalog_entries("suite-a") + + assert runs.items[0]["status"] == "failed" + catalog = json.loads( + ( + results_root + / "suite-a" + / "run_catalog.json" + ).read_text(encoding="utf-8") + ) + assert catalog["items"][0]["status"] == "failed" + + +def test_run_catalog_detects_out_of_band_summary_source_change() -> None: + with TemporaryDirectory() as tmp: + results_root = _build_legacy_fixture(Path(tmp)) + repository = ResultRepository(results_root) + repository.get_suite("suite-a") + + taxonomy_path = ( + results_root + / "suite-a" + / "artifacts" + / "systematize" + / "v0001" + / "taxonomy.json" + ) + taxonomy = json.loads(taxonomy_path.read_text(encoding="utf-8")) + taxonomy["behavior_categories"].append( + {"name": "new-category", "permissible": True} + ) + _write_json(taxonomy_path, taxonomy) + + repository.list_run_catalog_entries("suite-a") + + expected_mtime = taxonomy_path.stat().st_mtime_ns + run_summary = json.loads( + ( + results_root + / "suite-a" + / "run-a" + / "run_summary.json" + ).read_text(encoding="utf-8") + ) + run_catalog = json.loads( + ( + results_root + / "suite-a" + / "run_catalog.json" + ).read_text(encoding="utf-8") + ) + assert ( + run_summary["sources"]["taxonomy"]["mtime_ns"] + == expected_mtime + ) + assert ( + run_catalog["summary_sources"]["run-a"]["taxonomy"][ + "mtime_ns" + ] + == expected_mtime + ) def test_oversized_page_row_returns_bounded_stub_and_remains_pageable() -> None: @@ -465,9 +573,15 @@ def test_corrupt_derived_summaries_are_rebuilt_from_canonical_artifacts() -> Non repository = ResultRepository(results_root) suite = repository.get_suite("suite-a") run = repository.load_run_detail("suite-a", "run-a") + (suite_root / "run_catalog.json").write_text( + "{", + encoding="utf-8", + ) + runs = repository.list_run_catalog_entries("suite-a") assert suite["run_count"] == 1 assert run["state"] == "completed" + assert runs.items[0]["status"] == "completed" def test_cli_results_list_uses_metadata_and_compare_supports_ordinal() -> None: diff --git a/tests/test_result_service_scale.py b/tests/test_result_service_scale.py index 4839f2519..7fde23309 100644 --- a/tests/test_result_service_scale.py +++ b/tests/test_result_service_scale.py @@ -22,19 +22,24 @@ def test_large_catalog_is_metadata_only_and_single_case_lookup_is_indexed() -> N ) original_open = Path.open - def reject_score_reads( + def reject_catalog_source_reads( path: Path, mode: str = "r", *args: object, **kwargs: object, ): - if path.name == "scores.jsonl" and "r" in mode: - raise AssertionError("catalog listing opened score rows") + if ( + path.name in {"scores.jsonl", "run_summary.json"} + and "r" in mode + ): + raise AssertionError( + f"catalog listing opened {path.name}" + ) return original_open(path, mode, *args, **kwargs) started = time.perf_counter() with ( - patch.object(Path, "open", reject_score_reads), + patch.object(Path, "open", reject_catalog_source_reads), patch( "assert_ai.services.results.scan_jsonl", side_effect=AssertionError("catalog listing scanned JSONL"), From 4ca50416db7ce6b5c81f2e009be29441541c6554 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Tue, 25 Aug 2026 09:48:56 -0700 Subject: [PATCH 08/16] Add MCP config authoring and preflight Expose typed config validation, ETag saves, pure run planning, model-backed design, and isolated target probing through explicit MCP capability groups. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/core/artifact_cache.py | 166 ++-- assert_ai/core/run_plan.py | 42 + assert_ai/mcp/_command.py | 52 ++ assert_ai/mcp/models.py | 47 +- assert_ai/mcp/server.py | 107 ++- assert_ai/mcp/tools/__init__.py | 17 +- assert_ai/mcp/tools/author.py | 349 ++++++++ assert_ai/runner.py | 37 +- assert_ai/services/_target_probe_worker.py | 261 ++++++ assert_ai/services/run_planning.py | 972 +++++++++++++++++++++ assert_ai/services/target_probe.py | 316 +++++++ pyproject.toml | 4 +- tests/test_artifact_cache.py | 41 + tests/test_mcp_cli.py | 37 + tests/test_mcp_server.py | 384 +++++++- tests/test_run_planning_service.py | 327 +++++++ tests/test_target_probe_service.py | 180 ++++ 17 files changed, 3244 insertions(+), 95 deletions(-) create mode 100644 assert_ai/core/run_plan.py create mode 100644 assert_ai/mcp/tools/author.py create mode 100644 assert_ai/services/_target_probe_worker.py create mode 100644 assert_ai/services/run_planning.py create mode 100644 assert_ai/services/target_probe.py create mode 100644 tests/test_run_planning_service.py create mode 100644 tests/test_target_probe_service.py diff --git a/assert_ai/core/artifact_cache.py b/assert_ai/core/artifact_cache.py index edbe28809..972b364f1 100644 --- a/assert_ai/core/artifact_cache.py +++ b/assert_ai/core/artifact_cache.py @@ -202,36 +202,23 @@ def prepare_artifact_plan( if stage_name not in CACHEABLE_STAGES: raise ValueError(f"unsupported cacheable stage: {stage_name}") + suite_root = _managed_suite_root(ctx) - fingerprint = build_artifact_fingerprint(ctx=ctx, stage_name=stage_name, raw_cfg=raw_cfg) - stage_root = _managed_output_path( - ctx, - suite_root / ARTIFACTS_DIR / stage_name, - field_name=f"{stage_name} artifact cache root", - expected_root=suite_root, - reject_links=True, + fingerprint = build_artifact_fingerprint( + ctx=ctx, + stage_name=stage_name, + raw_cfg=raw_cfg, ) - + stage_root = _artifact_stage_root(ctx, suite_root, stage_name) if not forced: - match = _latest_matching_metadata(stage_name, stage_root, fingerprint.input_hash) - if match is not None: - version, metadata = match - artifact_dir = _managed_output_path( - ctx, - stage_root / version, - field_name=f"{stage_name} artifact cache version", - expected_root=stage_root, - reject_links=True, - ) - return ArtifactPlan( - stage_name=stage_name, - version=version, - artifact_dir=artifact_dir, - output_paths=_output_paths(stage_name, artifact_dir), - fingerprint=fingerprint, - reused=True, - metadata=metadata, - ) + reusable = _find_reusable_artifact_plan( + ctx=ctx, + stage_name=stage_name, + fingerprint=fingerprint, + stage_root=stage_root, + ) + if reusable is not None: + return reusable version, artifact_dir = _allocate_version_dir(stage_root) return ArtifactPlan( @@ -245,6 +232,77 @@ def prepare_artifact_plan( ) +def find_reusable_artifact_plan( + *, + ctx: dict[str, Any], + stage_name: str, + raw_cfg: dict[str, Any], +) -> ArtifactPlan | None: + """Return a matching artifact plan without allocating or repairing files.""" + if stage_name not in CACHEABLE_STAGES: + raise ValueError(f"unsupported cacheable stage: {stage_name}") + suite_root = _managed_suite_root(ctx) + fingerprint = build_artifact_fingerprint( + ctx=ctx, + stage_name=stage_name, + raw_cfg=raw_cfg, + ) + stage_root = _artifact_stage_root(ctx, suite_root, stage_name) + return _find_reusable_artifact_plan( + ctx=ctx, + stage_name=stage_name, + fingerprint=fingerprint, + stage_root=stage_root, + ) + + +def _find_reusable_artifact_plan( + *, + ctx: dict[str, Any], + stage_name: str, + fingerprint: ArtifactFingerprint, + stage_root: Path, +) -> ArtifactPlan | None: + match = _latest_matching_metadata( + stage_name, + stage_root, + fingerprint.input_hash, + ) + if match is None: + return None + version, metadata = match + artifact_dir = _managed_output_path( + ctx, + stage_root / version, + field_name=f"{stage_name} artifact cache version", + expected_root=stage_root, + reject_links=True, + ) + return ArtifactPlan( + stage_name=stage_name, + version=version, + artifact_dir=artifact_dir, + output_paths=_output_paths(stage_name, artifact_dir), + fingerprint=fingerprint, + reused=True, + metadata=metadata, + ) + + +def _artifact_stage_root( + ctx: dict[str, Any], + suite_root: Path, + stage_name: str, +) -> Path: + return _managed_output_path( + ctx, + suite_root / ARTIFACTS_DIR / stage_name, + field_name=f"{stage_name} artifact cache root", + expected_root=suite_root, + reject_links=True, + ) + + def activate_artifact_plan(ctx: dict[str, Any], plan: ArtifactPlan) -> dict[str, Any]: """Put selected artifact paths/version metadata into runner context.""" @@ -310,7 +368,11 @@ def override_cacheable_output_paths( return cfg -def activate_latest_artifacts(ctx: dict[str, Any]) -> None: +def activate_latest_artifacts( + ctx: dict[str, Any], + *, + repair: bool = True, +) -> None: """Load latest artifact refs into context for run-only stage configs. When ``latest.json`` references an artifact directory that has been @@ -403,29 +465,32 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: metadata=metadata, primary_path=output_paths[next(iter(_OUTPUT_FILES[stage_name]))], ) - update_latest(ctx, stage_name, ref) - log.warning( - "latest.json %s entry referenced missing paths; rebuilt " - "ref pointing at the current on-disk location of version %s.", - stage_name, - version, - ) + if repair: + update_latest(ctx, stage_name, ref) + log.warning( + "latest.json %s entry referenced missing paths; rebuilt " + "ref pointing at the current on-disk location of version %s.", + stage_name, + version, + ) ctx.setdefault("artifact_versions", {})[stage_name] = ref ctx[_CONTEXT_DIR_KEYS[stage_name]] = str(artifact_dir) for output_key, context_key in _CONTEXT_PATH_KEYS[stage_name].items(): if output_key in output_paths: ctx[context_key] = str(output_paths[output_key]) - refresh_compatibility_files(ctx, stage_name, output_paths) + if repair: + refresh_compatibility_files(ctx, stage_name, output_paths) continue recovery = _recover_latest_valid_version(stage_name, stage_root) if recovery is None: - log.warning( - "latest.json references missing or incomplete %s artifact %s; " - "no valid prior version was found.", - stage_name, - version, - ) + if repair: + log.warning( + "latest.json references missing or incomplete %s artifact %s; " + "no valid prior version was found.", + stage_name, + version, + ) continue recovered_version, recovered_dir, recovered_metadata = recovery recovered_outputs = _metadata_output_paths( @@ -446,14 +511,15 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: for output_key, context_key in _CONTEXT_PATH_KEYS[stage_name].items(): if output_key in recovered_outputs: ctx[context_key] = str(recovered_outputs[output_key]) - refresh_compatibility_files(ctx, stage_name, recovered_outputs) - update_latest(ctx, stage_name, recovered_ref) - log.warning( - "latest.json %s entry was missing or incomplete; " - "recovered to version %s.", - stage_name, - recovered_version, - ) + if repair: + refresh_compatibility_files(ctx, stage_name, recovered_outputs) + update_latest(ctx, stage_name, recovered_ref) + log.warning( + "latest.json %s entry was missing or incomplete; " + "recovered to version %s.", + stage_name, + recovered_version, + ) def finalize_artifact_plan(ctx: dict[str, Any], plan: ArtifactPlan) -> dict[str, Any]: diff --git a/assert_ai/core/run_plan.py b/assert_ai/core/run_plan.py new file mode 100644 index 000000000..bbb7cddfc --- /dev/null +++ b/assert_ai/core/run_plan.py @@ -0,0 +1,42 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pure stage-selection helpers shared by preflight and execution.""" + +from __future__ import annotations + +from collections.abc import Iterable + +from assert_ai.core.config_document import PIPELINE_STAGE_ORDER + + +def resolve_forced_stages( + configured_stage_names: Iterable[str], + requested_force_stages: Iterable[str], +) -> tuple[str, ...]: + """Validate forced stages and apply the runner's downstream cascade.""" + configured = set(configured_stage_names) + requested = set(requested_force_stages) + invalid = sorted(requested.difference(configured)) + if invalid: + raise ValueError( + "Forced stage(s) not present in config: " + + ", ".join(invalid) + ) + + forced = set(requested) + forced_indices = [ + PIPELINE_STAGE_ORDER.index(name) + for name in requested + if name in PIPELINE_STAGE_ORDER + ] + if forced_indices: + first_forced = min(forced_indices) + forced.update( + name + for name in PIPELINE_STAGE_ORDER[first_forced:] + if name in configured + ) + return tuple( + name for name in PIPELINE_STAGE_ORDER if name in forced + ) diff --git a/assert_ai/mcp/_command.py b/assert_ai/mcp/_command.py index 6ec627329..8e3043af5 100644 --- a/assert_ai/mcp/_command.py +++ b/assert_ai/mcp/_command.py @@ -116,6 +116,46 @@ def mcp() -> None: show_default=True, help="Maximum size of one managed config payload.", ) +@click.option( + "--max-concurrency", + type=click.IntRange(min=1), + default=32, + show_default=True, + help="Maximum inference concurrency accepted by preflight.", +) +@click.option( + "--max-prompt-sample-size", + type=click.IntRange(min=1), + default=100_000, + show_default=True, + help="Maximum prompt sample size accepted by preflight.", +) +@click.option( + "--max-scenario-sample-size", + type=click.IntRange(min=1), + default=100_000, + show_default=True, + help="Maximum scenario sample size accepted by preflight.", +) +@click.option( + "--allowed-model", + "allowed_model_patterns", + multiple=True, + help="Optional allowed model glob. Repeat as needed.", +) +@click.option( + "--allowed-endpoint-host", + "allowed_endpoint_hosts", + multiple=True, + help="Optional allowed endpoint-host glob. Repeat as needed.", +) +@click.option( + "--target-probe-timeout-seconds", + type=click.FloatRange(min=0.1), + default=15.0, + show_default=True, + help="Operator timeout for isolated target imports.", +) def serve( workspace: Path, mode: str, @@ -127,6 +167,12 @@ def serve( default_artifact_chunk_bytes: int, max_artifact_chunk_bytes: int, max_config_bytes: int, + max_concurrency: int, + max_prompt_sample_size: int, + max_scenario_sample_size: int, + allowed_model_patterns: tuple[str, ...], + allowed_endpoint_hosts: tuple[str, ...], + target_probe_timeout_seconds: float, ) -> None: """Serve ASSERT over stdio; stdout is reserved for MCP protocol traffic.""" try: @@ -152,6 +198,12 @@ def serve( default_artifact_chunk_bytes=default_artifact_chunk_bytes, max_artifact_chunk_bytes=max_artifact_chunk_bytes, max_config_bytes=max_config_bytes, + max_concurrency=max_concurrency, + max_prompt_sample_size=max_prompt_sample_size, + max_scenario_sample_size=max_scenario_sample_size, + allowed_model_patterns=allowed_model_patterns, + allowed_endpoint_hosts=allowed_endpoint_hosts, + target_probe_timeout_s=target_probe_timeout_seconds, ) except (OSError, ValueError) as exc: raise click.ClickException(str(exc)) from exc diff --git a/assert_ai/mcp/models.py b/assert_ai/mcp/models.py index 10c1f2530..723b992d0 100644 --- a/assert_ai/mcp/models.py +++ b/assert_ai/mcp/models.py @@ -60,6 +60,14 @@ class ServerLimits(BaseModel): default_artifact_chunk_bytes: int max_artifact_chunk_bytes: int max_config_bytes: int + max_concurrency: int + max_prompt_sample_size: int + max_scenario_sample_size: int + model_allowlist_enabled: bool = False + endpoint_host_allowlist_enabled: bool = False + allowed_model_patterns: tuple[str, ...] = () + allowed_endpoint_hosts: tuple[str, ...] = () + target_probe_timeout_s: float class ServerInfo(BaseModel): @@ -74,8 +82,15 @@ class ServerInfo(BaseModel): enabled_capability_groups: list[CapabilityGroup] workspace: WorkspaceInfo = Field(default_factory=WorkspaceInfo) limits: ServerLimits - target_kinds: list[Literal["callable", "model"]] = Field( - default_factory=lambda: ["callable", "model"] + target_kinds: list[ + Literal["callable", "model", "connector", "endpoint"] + ] = Field( + default_factory=lambda: [ + "callable", + "model", + "connector", + "endpoint", + ] ) transports: list[Literal["stdio"]] = Field(default_factory=lambda: ["stdio"]) protocol_notes: list[str] = Field( @@ -160,6 +175,34 @@ class ConfigResult(_McpModel): resource_uri: str +class ConfigValidationResult(_McpModel): + """Layered validation report for a managed config or draft.""" + + source: Literal["config", "yaml", "document"] + config_ref: str + validation: ConfigValidationReport + + +class ConfigSaveToolResult(_McpModel): + """Identity and ETag after an atomic managed-config save.""" + + config_ref: str + etag: str + created: bool + validation: ConfigValidationReport + resource_uri: str + + +class ConfigDesignResult(_McpModel): + """Unpersisted model-generated config draft.""" + + yaml: str + document: dict[str, Any] + validation: ConfigValidationReport + model_cost_incurred: Literal[True] = True + persisted: Literal[False] = False + + class SuiteCatalogItem(_McpModel): """Lightweight suite metadata.""" diff --git a/assert_ai/mcp/server.py b/assert_ai/mcp/server.py index c89555d22..20d40f33a 100644 --- a/assert_ai/mcp/server.py +++ b/assert_ai/mcp/server.py @@ -23,11 +23,24 @@ WorkspaceInfo, ) from assert_ai.mcp.resources import register_inspect_resources -from assert_ai.mcp.tools import InspectServices, register_inspect_tools +from assert_ai.mcp.tools import ( + AuthorServices, + InspectServices, + ProbeServices, + register_author_tools, + register_design_tools, + register_inspect_tools, + register_probe_tools, +) from assert_ai.services.artifacts import ArtifactRepository from assert_ai.services.configs import ConfigService from assert_ai.services.library import LibraryService from assert_ai.services.results import ResultRepository +from assert_ai.services.run_planning import ( + PreflightPolicy, + RunPlanningService, +) +from assert_ai.services.target_probe import TargetProbeService SERVER_NAME = "ASSERT" @@ -66,6 +79,12 @@ class ServerOptions: default_artifact_chunk_bytes: int = 64 * 1024 max_artifact_chunk_bytes: int = 256 * 1024 max_config_bytes: int = 256 * 1024 + max_concurrency: int = 32 + max_prompt_sample_size: int = 100_000 + max_scenario_sample_size: int = 100_000 + allowed_model_patterns: tuple[str, ...] = () + allowed_endpoint_hosts: tuple[str, ...] = () + target_probe_timeout_s: float = 15.0 workspace: WorkspaceService = field(init=False, repr=False) def __post_init__(self) -> None: @@ -79,6 +98,21 @@ def __post_init__(self) -> None: raise ValueError("max_config_bytes must be positive") if self.max_config_bytes > self.max_response_bytes: raise ValueError("max_config_bytes must not exceed max_response_bytes") + if self.max_concurrency < 1: + raise ValueError("max_concurrency must be positive") + if self.max_prompt_sample_size < 1: + raise ValueError("max_prompt_sample_size must be positive") + if self.max_scenario_sample_size < 1: + raise ValueError("max_scenario_sample_size must be positive") + if self.target_probe_timeout_s <= 0: + raise ValueError("target_probe_timeout_s must be positive") + PreflightPolicy( + max_concurrency=self.max_concurrency, + max_prompt_sample_size=self.max_prompt_sample_size, + max_scenario_sample_size=self.max_scenario_sample_size, + allowed_model_patterns=self.allowed_model_patterns, + allowed_endpoint_hosts=self.allowed_endpoint_hosts, + ) if self.default_artifact_chunk_bytes < 4: raise ValueError("default_artifact_chunk_bytes must be at least 4") if self.max_artifact_chunk_bytes < self.default_artifact_chunk_bytes: @@ -106,6 +140,12 @@ def create( default_artifact_chunk_bytes: int = 64 * 1024, max_artifact_chunk_bytes: int = 256 * 1024, max_config_bytes: int = 256 * 1024, + max_concurrency: int = 32, + max_prompt_sample_size: int = 100_000, + max_scenario_sample_size: int = 100_000, + allowed_model_patterns: Iterable[str] = (), + allowed_endpoint_hosts: Iterable[str] = (), + target_probe_timeout_s: float = 15.0, ) -> "ServerOptions": parsed_mode = ServerMode(mode) parsed_groups = tuple(CapabilityGroup(group) for group in enabled_groups) @@ -125,6 +165,12 @@ def create( default_artifact_chunk_bytes=default_artifact_chunk_bytes, max_artifact_chunk_bytes=max_artifact_chunk_bytes, max_config_bytes=max_config_bytes, + max_concurrency=max_concurrency, + max_prompt_sample_size=max_prompt_sample_size, + max_scenario_sample_size=max_scenario_sample_size, + allowed_model_patterns=tuple(allowed_model_patterns), + allowed_endpoint_hosts=tuple(allowed_endpoint_hosts), + target_probe_timeout_s=target_probe_timeout_s, ) @property @@ -152,6 +198,13 @@ def build_server(options: ServerOptions) -> MCPServer: version=_server_version(), ) + configs = ConfigService( + options.workspace, + max_config_bytes=options.max_config_bytes, + default_page_size=options.default_page_size, + max_page_size=options.max_page_size, + ) + @server.tool( title="Get ASSERT server information", annotations=ToolAnnotations( @@ -181,6 +234,16 @@ def get_server_info() -> ServerInfo: default_artifact_chunk_bytes=options.default_artifact_chunk_bytes, max_artifact_chunk_bytes=options.max_artifact_chunk_bytes, max_config_bytes=options.max_config_bytes, + max_concurrency=options.max_concurrency, + max_prompt_sample_size=options.max_prompt_sample_size, + max_scenario_sample_size=options.max_scenario_sample_size, + model_allowlist_enabled=bool(options.allowed_model_patterns), + endpoint_host_allowlist_enabled=bool( + options.allowed_endpoint_hosts + ), + allowed_model_patterns=options.allowed_model_patterns, + allowed_endpoint_hosts=options.allowed_endpoint_hosts, + target_probe_timeout_s=options.target_probe_timeout_s, ), ) @@ -199,12 +262,7 @@ def get_server_info() -> ServerInfo: default_page_size=options.default_page_size, max_page_size=options.max_page_size, ), - configs=ConfigService( - options.workspace, - max_config_bytes=options.max_config_bytes, - default_page_size=options.default_page_size, - max_page_size=options.max_page_size, - ), + configs=configs, results=results, artifacts=ArtifactRepository( options.workspace, @@ -224,6 +282,41 @@ def get_server_info() -> ServerInfo: inline_artifact_bytes=options.max_artifact_chunk_bytes, ) + author_services = AuthorServices( + workspace=options.workspace, + configs=configs, + planning=RunPlanningService( + options.workspace, + configs, + policy=PreflightPolicy( + max_concurrency=options.max_concurrency, + max_prompt_sample_size=options.max_prompt_sample_size, + max_scenario_sample_size=options.max_scenario_sample_size, + allowed_model_patterns=options.allowed_model_patterns, + allowed_endpoint_hosts=options.allowed_endpoint_hosts, + ), + ), + max_response_bytes=options.max_response_bytes, + allowed_model_patterns=options.allowed_model_patterns, + ) + if CapabilityGroup.AUTHOR in options.capability_groups: + register_author_tools(server, author_services) + if CapabilityGroup.DESIGN in options.capability_groups: + register_design_tools(server, author_services) + if CapabilityGroup.PROBE in options.capability_groups: + register_probe_tools( + server, + ProbeServices( + workspace=options.workspace, + probe=TargetProbeService( + options.workspace, + configs, + timeout_s=options.target_probe_timeout_s, + ), + max_response_bytes=options.max_response_bytes, + ), + ) + return server diff --git a/assert_ai/mcp/tools/__init__.py b/assert_ai/mcp/tools/__init__.py index 5cc7f45d1..704c145cf 100644 --- a/assert_ai/mcp/tools/__init__.py +++ b/assert_ai/mcp/tools/__init__.py @@ -3,6 +3,21 @@ """Capability-group tool registrars for the ASSERT MCP server.""" +from assert_ai.mcp.tools.author import ( + AuthorServices, + ProbeServices, + register_author_tools, + register_design_tools, + register_probe_tools, +) from assert_ai.mcp.tools.inspect import InspectServices, register_inspect_tools -__all__ = ["InspectServices", "register_inspect_tools"] +__all__ = [ + "AuthorServices", + "InspectServices", + "ProbeServices", + "register_author_tools", + "register_design_tools", + "register_inspect_tools", + "register_probe_tools", +] diff --git a/assert_ai/mcp/tools/author.py b/assert_ai/mcp/tools/author.py new file mode 100644 index 000000000..0cd2e97b8 --- /dev/null +++ b/assert_ai/mcp/tools/author.py @@ -0,0 +1,349 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Config authoring, pure preflight, and isolated probe MCP tools.""" + +from __future__ import annotations + +from dataclasses import dataclass +from fnmatch import fnmatchcase +from typing import Annotated, Any + +from mcp.server import MCPServer +from mcp.types import ToolAnnotations +from pydantic import Field + +from assert_ai.core.workspace import WorkspaceService +from assert_ai.mcp.errors import adapt_tool_errors, invoke_tool +from assert_ai.mcp.models import ( + ConfigDesignResult, + ConfigSaveToolResult, + ConfigValidationResult, +) +from assert_ai.mcp.sanitize import sanitize_for_mcp +from assert_ai.mcp.uris import config_uri +from assert_ai.services.configs import ( + ConfigDesignRequest, + ConfigService, +) +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.run_planning import ( + EvaluationOverrides, + EvaluationPreflight, + RunPlanningService, +) +from assert_ai.services.target_probe import ( + TargetProbeResult, + TargetProbeService, +) + +_PURE_ANNOTATIONS = ToolAnnotations( + read_only_hint=True, + destructive_hint=False, + idempotent_hint=True, + open_world_hint=False, +) +_WRITE_ANNOTATIONS = ToolAnnotations( + read_only_hint=False, + destructive_hint=True, + idempotent_hint=False, + open_world_hint=False, +) +_DESIGN_ANNOTATIONS = ToolAnnotations( + read_only_hint=True, + destructive_hint=False, + idempotent_hint=False, + open_world_hint=True, +) +_PROBE_ANNOTATIONS = ToolAnnotations( + read_only_hint=True, + destructive_hint=False, + idempotent_hint=False, + open_world_hint=True, +) + + +@dataclass(frozen=True, slots=True) +class AuthorServices: + workspace: WorkspaceService + configs: ConfigService + planning: RunPlanningService + max_response_bytes: int + allowed_model_patterns: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class ProbeServices: + workspace: WorkspaceService + probe: TargetProbeService + max_response_bytes: int + + +def register_author_tools( + server: MCPServer, + services: AuthorServices, +) -> None: + """Register deterministic authoring and preflight tools.""" + + @server.tool( + title="Validate an ASSERT config", + annotations=_PURE_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def validate_config( + config_ref: str | None = None, + yaml_text: str | None = None, + document: dict[str, Any] | None = None, + validation_ref: str = "draft.yaml", + ) -> ConfigValidationResult: + """Validate exactly one managed config, YAML draft, or config document.""" + _require_exactly_one_config_source( + config_ref=config_ref, + yaml_text=yaml_text, + document=document, + ) + if config_ref is not None: + record = invoke_tool( + lambda: services.configs.get_config(config_ref), + workspace=services.workspace, + ) + return ConfigValidationResult.model_validate( + sanitize_for_mcp( + ConfigValidationResult( + source="config", + config_ref=record.config_ref, + validation=record.validation, + ), + workspace=services.workspace, + ) + ) + if yaml_text is not None: + report = invoke_tool( + lambda: services.configs.validate_yaml( + yaml_text, + config_ref=validation_ref, + ), + workspace=services.workspace, + ) + return ConfigValidationResult.model_validate( + sanitize_for_mcp( + ConfigValidationResult( + source="yaml", + config_ref=validation_ref, + validation=report, + ), + workspace=services.workspace, + ) + ) + assert document is not None + report = invoke_tool( + lambda: services.configs.validate_document( + document, + config_ref=validation_ref, + ), + workspace=services.workspace, + ) + return ConfigValidationResult.model_validate( + sanitize_for_mcp( + ConfigValidationResult( + source="document", + config_ref=validation_ref, + validation=report, + ), + workspace=services.workspace, + ) + ) + + @server.tool( + title="Save an ASSERT config", + annotations=_WRITE_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def save_config( + config_ref: str, + yaml_text: str | None = None, + document: dict[str, Any] | None = None, + expected_etag: str | None = None, + ) -> ConfigSaveToolResult: + """Validate and atomically create or replace one managed config.""" + if (yaml_text is None) == (document is None): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Provide exactly one of yaml_text or document", + ) + saved = invoke_tool( + lambda: services.configs.save_config( + config_ref, + yaml_text=yaml_text, + document=document, + expected_etag=expected_etag, + ), + workspace=services.workspace, + ) + return ConfigSaveToolResult.model_validate( + sanitize_for_mcp( + ConfigSaveToolResult( + config_ref=saved.config_ref, + etag=saved.etag, + created=saved.created, + validation=saved.validation, + resource_uri=config_uri(saved.config_ref), + ), + workspace=services.workspace, + ) + ) + + @server.tool( + title="Preflight an ASSERT evaluation", + annotations=_PURE_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def preflight_evaluation( + config_ref: str, + overrides: EvaluationOverrides | None = None, + ) -> EvaluationPreflight: + """Plan an exact effective run without importing targets or writing files.""" + plan = invoke_tool( + lambda: services.planning.preflight( + config_ref, + overrides=overrides, + ), + workspace=services.workspace, + ) + payload = plan.model_dump(mode="json") + credentials = payload.pop("credentials") + sanitized = sanitize_for_mcp( + payload, + workspace=services.workspace, + ) + sanitized["credentials"] = credentials + return EvaluationPreflight.model_validate(sanitized) + + +def register_design_tools( + server: MCPServer, + services: AuthorServices, +) -> None: + """Register the model-backed, non-persisting config designer.""" + + @server.tool( + title="Design an ASSERT config", + annotations=_DESIGN_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def design_config( + description: Annotated[str, Field(min_length=1)], + model: Annotated[str, Field(min_length=1)] = "azure/gpt-5.4-mini", + seed_config_ref: str | None = None, + seed_yaml: str | None = None, + behavior_preset: str | None = None, + judge_preset: str | None = None, + dimension_hints: str | None = None, + default_model_hint: str | None = None, + max_turns: Annotated[int, Field(ge=1, le=100)] = 5, + ) -> ConfigDesignResult: + """Call ASSERT's design model and return an unpersisted config draft.""" + description = description.strip() + model = model.strip() + if not description: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "description must not be blank", + ) + if not model: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "model must not be blank", + ) + if services.allowed_model_patterns and not any( + fnmatchcase(model, pattern) + for pattern in services.allowed_model_patterns + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"Design model {model!r} is not allowed by server policy", + ) + draft = invoke_tool( + lambda: services.configs.design_config( + ConfigDesignRequest( + description=description, + model=model, + seed_config_ref=seed_config_ref, + seed_yaml=seed_yaml, + behavior_preset=behavior_preset, + judge_preset=judge_preset, + dimension_hints=dimension_hints, + default_model_hint=default_model_hint, + max_turns=max_turns, + ) + ), + workspace=services.workspace, + ) + return ConfigDesignResult.model_validate( + sanitize_for_mcp( + ConfigDesignResult( + yaml=draft.yaml, + document=draft.document, + validation=draft.validation, + ), + workspace=services.workspace, + ) + ) + + +def register_probe_tools( + server: MCPServer, + services: ProbeServices, +) -> None: + """Register disposable-subprocess target probing.""" + + @server.tool( + title="Probe an ASSERT target", + annotations=_PROBE_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def probe_target(config_ref: str) -> TargetProbeResult: + """Import and inspect a managed config's target in an isolated process.""" + result = invoke_tool( + lambda: services.probe.probe(config_ref), + workspace=services.workspace, + ) + return TargetProbeResult.model_validate( + sanitize_for_mcp(result, workspace=services.workspace) + ) + + +def _require_exactly_one_config_source( + *, + config_ref: str | None, + yaml_text: str | None, + document: dict[str, Any] | None, +) -> None: + if sum( + value is not None + for value in (config_ref, yaml_text, document) + ) != 1: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Provide exactly one of config_ref, yaml_text, or document", + ) diff --git a/assert_ai/runner.py b/assert_ai/runner.py index 820ec8242..1c891f7e5 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -53,6 +53,7 @@ run_stage_coro, ) from assert_ai.core.run_result import RunResult, RunState +from assert_ai.core.run_plan import resolve_forced_stages from assert_ai.display import label_metric from assert_ai.services.result_metadata import ( refresh_stage_indexes, @@ -755,12 +756,16 @@ def _run_pipeline_result( "[runner] --concurrency ignored: this config has no inference stage to override." ) - requested_force_stages = set(force_stages or []) configured_stage_names = {stage_name for stage_name, _ in ctx["stages"]} - invalid_forced = sorted(requested_force_stages.difference(configured_stage_names)) - if invalid_forced: - joined = ", ".join(invalid_forced) - message = f"--force-stage stage(s) not present in config: {joined}" + try: + requested_force_stages = set( + resolve_forced_stages( + configured_stage_names, + force_stages or (), + ) + ) + except ValueError as exc: + message = str(exc).replace("Forced stage", "--force-stage stage", 1) log.error(f"[config error] {message}") return _run_result_from_context( ctx, @@ -770,28 +775,6 @@ def _run_pipeline_result( error_message=message, ) - # Cascade: forcing an upstream stage logically invalidates every stage - # downstream of it. Without this, `--force-stage test_set` regenerates test_set - # but inference silently keeps the old inference rows (its resume cache keys on - # test_case_id, and test case ids are deterministic so they collide with the prior - # run's content). Same hazard for judge against scores.jsonl. Computing - # the closure here keeps the workflow `--force-stage ` honest - # without forcing users to remember the full downstream chain. - if requested_force_stages: - forced_indices = [ - PIPELINE_STAGE_ORDER.index(name) - for name in requested_force_stages - if name in PIPELINE_STAGE_ORDER - ] - if forced_indices: - min_forced_index = min(forced_indices) - cascade = { - name - for name in PIPELINE_STAGE_ORDER[min_forced_index:] - if name in configured_stage_names - } - requested_force_stages = requested_force_stages.union(cascade) - suite_root = Path(ctx["suite_root"]) path_policy = ctx.get("path_policy") if path_policy is not None: diff --git a/assert_ai/services/_target_probe_worker.py b/assert_ai/services/_target_probe_worker.py new file mode 100644 index 000000000..ab4d30f17 --- /dev/null +++ b/assert_ai/services/_target_probe_worker.py @@ -0,0 +1,261 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Internal subprocess entry point for isolated target probing.""" + +from __future__ import annotations + +import inspect +import io +import json +import sys +from contextlib import redirect_stderr, redirect_stdout +from copy import deepcopy +from pathlib import Path +from typing import Any +from urllib.parse import urlsplit, urlunsplit + +from assert_ai.config import load_runtime_context +from assert_ai.core.security import ( + redact_path_prefixes, + sanitize_text, + validate_callable_ref, + validate_module_ref, +) +from assert_ai.core.session import _discover_connector_class +from assert_ai.core.tool_backend import ( + import_callable_module, + inspect_tool_module, + load_tool_module, +) +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.configs import ConfigService +from assert_ai.stages import STAGES + +_RESULT_MARKER = "ASSERT_TARGET_PROBE_RESULT=" +_MAX_REQUEST_BYTES = 1024 * 1024 + + +class _DiscardText(io.TextIOBase): + def write(self, value: str) -> int: + return len(value) + + def writable(self) -> bool: + return True + + +def main() -> int: + workspace: WorkspaceService | None = None + result_token = "" + payload: dict[str, Any] + try: + raw = sys.stdin.buffer.read(_MAX_REQUEST_BYTES + 1) + if len(raw) > _MAX_REQUEST_BYTES: + raise ValueError("Target probe request exceeds the worker limit") + request = json.loads(raw.decode("utf-8")) + if not isinstance(request, dict): + raise ValueError("Target probe request must be an object") + result_token = _required_string(request, "result_token") + workspace = WorkspaceService.create( + _required_string(request, "workspace_root") + ) + with ( + redirect_stdout(_DiscardText()), + redirect_stderr(_DiscardText()), + ): + payload = _probe(workspace, request) + except Exception as exc: # noqa: BLE001 - process boundary returns failure + message = sanitize_text(str(exc)) + if workspace is not None: + message = redact_path_prefixes( + message, + ( + workspace.root, + workspace.configs_root, + workspace.artifacts_root, + workspace.results_root, + ), + ) + payload = { + "ok": False, + "error_type": type(exc).__name__, + "message": message or "Target probe failed", + } + finally: + _cleanup_descendants() + + encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + sys.__stdout__.write( + f"{_RESULT_MARKER}{result_token}={encoded}\n" + ) + sys.__stdout__.flush() + return 0 if payload.get("ok") is True else 1 + + +def _probe( + workspace: WorkspaceService, + request: dict[str, Any], +) -> dict[str, Any]: + config_ref = _required_string(request, "config_ref") + max_config_bytes = request.get("max_config_bytes") + if ( + not isinstance(max_config_bytes, int) + or isinstance(max_config_bytes, bool) + or max_config_bytes < 1 + ): + raise ValueError("max_config_bytes must be a positive integer") + configs = ConfigService( + workspace, + max_config_bytes=max_config_bytes, + ) + record = configs.get_config(config_ref) + if not record.validation.valid: + raise ValueError("Target probe requires a valid config") + config_path = workspace.path_policy.resolve_config_path( + record.config_ref, + must_exist=True, + reject_links=True, + ) + ctx = load_runtime_context( + deepcopy(record.document), + config_path, + stage_modules=STAGES, + path_policy=workspace.path_policy, + ) + target = ctx.get("target") + inference_enabled = any( + stage_name == "inference" and raw_cfg.get("enabled", True) + for stage_name, raw_cfg in ctx["stages"] + ) + if not inference_enabled or target is None: + raise ValueError( + "Config has no enabled inference target to probe" + ) + + if target.model is not None: + details: dict[str, Any] = { + "model": str(target.model.name), + "trace_enabled": target.trace is not None, + } + tools = target.tools + if tools is not None and tools.module: + tools_class, schemas = inspect_tool_module( + tools.module, + config_path=config_path, + path_policy=workspace.path_policy, + ) + details["tools_module"] = tools.module + details["tools_class"] = tools_class.__name__ + details["tool_count"] = len(schemas) + if tools is not None and tools.toolset: + toolset_path = workspace.path_policy.resolve_input( + tools.toolset, + base_dir=config_path.parent, + field_name="pipeline.inference.target.tools.toolset", + must_exist=True, + file_only=True, + ) + details["toolset"] = workspace.reference(toolset_path) + return { + "ok": True, + "target_kind": "model", + "details": details, + } + + if target.callable: + validate_callable_ref(target.callable) + module_ref, function_name = target.callable.rsplit(":", 1) + module = import_callable_module( + module_ref, + config_path=config_path, + path_policy=workspace.path_policy, + ) + try: + function = getattr(module, function_name) + except AttributeError as exc: + raise ValueError( + f"Module {module_ref!r} has no attribute {function_name!r}" + ) from exc + if not callable(function): + raise ValueError( + f"Target attribute {target.callable!r} is not callable" + ) + signature = inspect.signature(function) + return { + "ok": True, + "target_kind": "callable", + "details": { + "reference": target.callable, + "is_async": inspect.iscoroutinefunction(function), + "accepts_history": "history" in signature.parameters, + "parameters": tuple(signature.parameters)[:20], + "trace_enabled": target.trace is not None, + }, + } + + if target.connector: + validate_module_ref(target.connector) + module = load_tool_module( + target.connector, + config_path=config_path, + path_policy=workspace.path_policy, + ) + connector_class = _discover_connector_class(module) + return { + "ok": True, + "target_kind": "connector", + "details": { + "reference": target.connector, + "adapter_class": connector_class.__name__, + }, + } + + endpoint = str(target.endpoint or "") + from assert_ai.core.session import HTTPEndpointSession + + HTTPEndpointSession(endpoint=endpoint) + return { + "ok": True, + "target_kind": "endpoint", + "details": {"origin": _endpoint_origin(endpoint)}, + } + + +def _endpoint_origin(value: str) -> str: + parsed = urlsplit(value) + if not parsed.scheme or not parsed.hostname: + raise ValueError("Endpoint target is invalid") + host = parsed.hostname + if ":" in host and not host.startswith("["): + host = f"[{host}]" + if parsed.port is not None: + host = f"{host}:{parsed.port}" + return urlunsplit((parsed.scheme, host, "", "", "")) + + +def _required_string(request: dict[str, Any], key: str) -> str: + value = request.get(key) + if not isinstance(value, str) or not value: + raise ValueError(f"{key} must be a non-empty string") + return value + + +def _cleanup_descendants() -> None: + try: + import psutil + except ImportError: + return + try: + children = psutil.Process().children(recursive=True) + for child in children: + child.terminate() + _, alive = psutil.wait_procs(children, timeout=1) + for child in alive: + child.kill() + psutil.wait_procs(alive, timeout=1) + except psutil.Error: + return + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/assert_ai/services/run_planning.py b/assert_ai/services/run_planning.py new file mode 100644 index 000000000..39ffced47 --- /dev/null +++ b/assert_ai/services/run_planning.py @@ -0,0 +1,972 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pure evaluation preflight and spend-policy checks.""" + +from __future__ import annotations + +import importlib.util +import os +from copy import deepcopy +from dataclasses import dataclass +from enum import StrEnum +from fnmatch import fnmatchcase +from pathlib import Path +from typing import Any, Literal +from urllib.parse import urlsplit, urlunsplit + +from pydantic import BaseModel, ConfigDict, Field + +from assert_ai.config import ConfigError, load_runtime_context +from assert_ai.core.artifact_cache import ( + activate_artifact_plan, + activate_latest_artifacts, + find_reusable_artifact_plan, + is_cacheable_stage, + supports_artifact_cache, +) +from assert_ai.core.config_document import ( + ConfigValidationIssue, + ConfigValidationReport, + PIPELINE_STAGE_ORDER, +) +from assert_ai.core.run_plan import resolve_forced_stages +from assert_ai.core.security import ( + validate_callable_ref, + validate_module_ref, +) +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.configs import ConfigService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.stages import STAGES + +StageName = Literal["systematize", "test_set", "inference", "judge"] + + +class _ServiceModel(BaseModel): + model_config = ConfigDict(frozen=True) + + +class ModelOverrides(_ServiceModel): + """Explicit model substitutions applied only to existing config roles.""" + + default_model: str | None = Field(default=None, min_length=1) + systematize_model: str | None = Field(default=None, min_length=1) + test_set_model: str | None = Field(default=None, min_length=1) + prompt_generator_model: str | None = Field(default=None, min_length=1) + scenario_generator_model: str | None = Field(default=None, min_length=1) + stratify_model: str | None = Field(default=None, min_length=1) + target_model: str | None = Field(default=None, min_length=1) + tester_model: str | None = Field(default=None, min_length=1) + judge_model: str | None = Field(default=None, min_length=1) + + +class EvaluationOverrides(_ServiceModel): + """Typed operational overrides accepted by pure preflight.""" + + suite: str | None = Field(default=None, min_length=1) + run: str | None = Field(default=None, min_length=1) + force_stages: tuple[StageName, ...] = () + strict: bool = False + concurrency: int | None = Field(default=None, ge=1) + prompt_sample_size: int | None = Field(default=None, ge=1) + scenario_sample_size: int | None = Field(default=None, ge=1) + models: ModelOverrides = Field(default_factory=ModelOverrides) + + +@dataclass(frozen=True, slots=True) +class PreflightPolicy: + """Operator-owned execution limits evaluated before a job can start.""" + + max_concurrency: int = 32 + max_prompt_sample_size: int = 100_000 + max_scenario_sample_size: int = 100_000 + allowed_model_patterns: tuple[str, ...] = () + allowed_endpoint_hosts: tuple[str, ...] = () + + def __post_init__(self) -> None: + if self.max_concurrency < 1: + raise ValueError("max_concurrency must be positive") + if self.max_prompt_sample_size < 1: + raise ValueError("max_prompt_sample_size must be positive") + if self.max_scenario_sample_size < 1: + raise ValueError("max_scenario_sample_size must be positive") + if any(not pattern.strip() for pattern in self.allowed_model_patterns): + raise ValueError("allowed_model_patterns cannot contain empty values") + if any(not host.strip() for host in self.allowed_endpoint_hosts): + raise ValueError("allowed_endpoint_hosts cannot contain empty values") + + +class StageAction(StrEnum): + DISABLED = "disabled" + REUSE = "reuse" + RUN = "run" + + +class PreflightIssue(_ServiceModel): + code: str + path: str = "" + message: str + + +class StagePreflight(_ServiceModel): + name: StageName + scope: Literal["suite", "run"] + action: StageAction + forced: bool = False + cacheable: bool = False + artifact_version: str | None = None + will_call_model: bool = False + reason: str + + +class ModelUse(_ServiceModel): + role: str + stage: StageName + model: str + provider: str + + +class CredentialRequirement(_ServiceModel): + provider: str + variables: dict[str, bool] + satisfied: bool | None + note: str + + +class TargetPreflight(_ServiceModel): + kind: Literal["model", "callable", "connector", "endpoint"] + identifier: str + trace_enabled: bool = False + static_validation: Literal["valid", "invalid"] + probe_required: bool = False + + +class ModelCallEstimate(_ServiceModel): + minimum: int | None = None + maximum: int | None = None + basis: str + + +class EvaluationPreflight(_ServiceModel): + config_ref: str + source_etag: str + effective_document: dict[str, Any] + validation: ConfigValidationReport + ready: bool + suite_id: str | None = None + run_id: str | None = None + strict: bool = False + concurrency: int | None = None + sample_sizes: dict[str, int | None] = Field(default_factory=dict) + target: TargetPreflight | None = None + stages: tuple[StagePreflight, ...] = () + models: tuple[ModelUse, ...] = () + credentials: tuple[CredentialRequirement, ...] = () + managed_outputs: dict[str, str] = Field(default_factory=dict) + estimated_model_calls: ModelCallEstimate = Field( + default_factory=lambda: ModelCallEstimate( + basis="No enabled stage uses a model.", + minimum=0, + maximum=0, + ) + ) + blocking_issues: tuple[PreflightIssue, ...] = () + warnings: tuple[PreflightIssue, ...] = () + + +@dataclass(slots=True) +class RunPlanningService: + """Build an exact, non-mutating execution plan for one managed config.""" + + workspace: WorkspaceService + configs: ConfigService + policy: PreflightPolicy = PreflightPolicy() + + def preflight( + self, + config_ref: str, + *, + overrides: EvaluationOverrides | None = None, + ) -> EvaluationPreflight: + record = self.configs.get_config(config_ref) + effective = deepcopy(record.document) + applied = overrides or EvaluationOverrides() + _apply_overrides(effective, applied) + validation = self.configs.validate_document( + effective, + config_ref=record.config_ref, + ) + blocking = [ + _validation_issue(issue) + for issue in validation.issues + ] + warnings = [ + _validation_issue(issue) + for issue in validation.warnings + ] + if not validation.valid: + return EvaluationPreflight( + config_ref=record.config_ref, + source_etag=record.etag, + effective_document=effective, + validation=validation, + ready=False, + strict=applied.strict, + blocking_issues=tuple(blocking), + warnings=tuple(warnings), + ) + + config_path = self.workspace.path_policy.resolve_config_path( + record.config_ref, + must_exist=True, + reject_links=True, + ) + try: + ctx = load_runtime_context( + deepcopy(effective), + config_path, + stage_modules=STAGES, + path_policy=self.workspace.path_policy, + ) + except (ConfigError, OSError, ValueError) as exc: + issue = PreflightIssue( + code=ServiceErrorCode.PREFLIGHT_FAILED.value, + message=str(exc), + ) + return EvaluationPreflight( + config_ref=record.config_ref, + source_etag=record.etag, + effective_document=effective, + validation=validation, + ready=False, + strict=applied.strict, + blocking_issues=(issue,), + warnings=tuple(warnings), + ) + + configured = [name for name, _ in ctx["stages"]] + try: + forced = set( + resolve_forced_stages( + configured, + applied.force_stages, + ) + ) + except ValueError as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + str(exc), + ) from exc + + ctx["strict"] = applied.strict + models = _collect_model_uses(effective) + target, target_issues = _target_preflight(ctx, self.policy) + blocking.extend(target_issues) + blocking.extend(_policy_issues(effective, models, ctx, self.policy)) + credentials, credential_issues, credential_warnings = ( + _credential_requirements(models) + ) + blocking.extend(credential_issues) + warnings.extend(credential_warnings) + stages = _stage_plan(ctx, forced, models) + managed_outputs = { + "artifacts_root": self.workspace.reference(ctx["artifacts_root"]), + "results_root": self.workspace.reference(ctx["results_dir"]), + "suite_root": self.workspace.reference(ctx["suite_root"]), + } + if ctx.get("run_root") is not None: + managed_outputs["run_root"] = self.workspace.reference( + ctx["run_root"] + ) + + concurrency = _effective_concurrency(ctx) + sample_sizes = _sample_sizes(effective) + estimate = _estimate_model_calls(stages) + return EvaluationPreflight( + config_ref=record.config_ref, + source_etag=record.etag, + effective_document=effective, + validation=validation, + ready=not blocking, + suite_id=str(ctx["suite_id"]), + run_id=( + str(ctx["run_id"]) + if ctx.get("run_id") is not None + else None + ), + strict=applied.strict, + concurrency=concurrency, + sample_sizes=sample_sizes, + target=target, + stages=tuple(stages), + models=tuple(models), + credentials=tuple(credentials), + managed_outputs=managed_outputs, + estimated_model_calls=estimate, + blocking_issues=tuple(_deduplicate_issues(blocking)), + warnings=tuple(_deduplicate_issues(warnings)), + ) + + +def _apply_overrides( + document: dict[str, Any], + overrides: EvaluationOverrides, +) -> None: + if overrides.suite is not None: + document["suite"] = overrides.suite.strip() + if overrides.run is not None: + document["run"] = overrides.run.strip() + pipeline = document.get("pipeline") + if not isinstance(pipeline, dict): + return + + if overrides.concurrency is not None: + inference = _require_mapping( + pipeline, + "inference", + "concurrency override requires pipeline.inference", + ) + inference["concurrency"] = overrides.concurrency + if overrides.prompt_sample_size is not None: + test_set = _require_mapping( + pipeline, + "test_set", + "prompt sample override requires pipeline.test_set", + ) + prompt = _require_mapping( + test_set, + "prompt", + "prompt sample override requires pipeline.test_set.prompt", + ) + prompt["sample_size"] = overrides.prompt_sample_size + if overrides.scenario_sample_size is not None: + test_set = _require_mapping( + pipeline, + "test_set", + "scenario sample override requires pipeline.test_set", + ) + scenario = _require_mapping( + test_set, + "scenario", + "scenario sample override requires pipeline.test_set.scenario", + ) + scenario["sample_size"] = overrides.scenario_sample_size + _apply_model_overrides(document, overrides.models) + + +def _apply_model_overrides( + document: dict[str, Any], + overrides: ModelOverrides, +) -> None: + values = overrides.model_dump(exclude_none=True) + if not values: + return + pipeline = document.get("pipeline") + if not isinstance(pipeline, dict): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Model overrides require a pipeline mapping", + ) + if overrides.default_model is not None: + _replace_model(document, "default_model", overrides.default_model) + if overrides.systematize_model is not None: + stage = _require_mapping( + pipeline, + "systematize", + "systematize_model requires pipeline.systematize", + ) + _replace_model(stage, "model", overrides.systematize_model) + if overrides.test_set_model is not None: + stage = _require_mapping( + pipeline, + "test_set", + "test_set_model requires pipeline.test_set", + ) + _replace_model(stage, "model", overrides.test_set_model) + for field_name, child_name in ( + ("prompt_generator_model", "prompt"), + ("scenario_generator_model", "scenario"), + ("stratify_model", "stratify"), + ): + value = getattr(overrides, field_name) + if value is None: + continue + stage = _require_mapping( + pipeline, + "test_set", + f"{field_name} requires pipeline.test_set", + ) + child = _require_mapping( + stage, + child_name, + f"{field_name} requires pipeline.test_set.{child_name}", + ) + _replace_model(child, "model", value) + if overrides.target_model is not None: + inference = _require_mapping( + pipeline, + "inference", + "target_model requires pipeline.inference", + ) + target = _require_mapping( + inference, + "target", + "target_model requires pipeline.inference.target", + ) + if any(target.get(kind) for kind in ("callable", "connector", "endpoint")): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "target_model cannot replace a callable, connector, or endpoint target", + ) + _replace_model(target, "model", overrides.target_model) + if overrides.tester_model is not None: + inference = _require_mapping( + pipeline, + "inference", + "tester_model requires pipeline.inference", + ) + tester = _require_mapping( + inference, + "tester", + "tester_model requires pipeline.inference.tester", + ) + _replace_model(tester, "model", overrides.tester_model) + if overrides.judge_model is not None: + judge = _require_mapping( + pipeline, + "judge", + "judge_model requires pipeline.judge", + ) + _replace_model(judge, "model", overrides.judge_model) + + +def _require_mapping( + owner: dict[str, Any], + key: str, + message: str, +) -> dict[str, Any]: + value = owner.get(key) + if not isinstance(value, dict): + raise ServiceError(ServiceErrorCode.INVALID_ARGUMENT, message) + return value + + +def _replace_model( + owner: dict[str, Any], + key: str, + model_name: str, +) -> None: + name = model_name.strip() + if not name: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{key} model name must not be blank", + ) + current = owner.get(key) + model = dict(current) if isinstance(current, dict) else {} + model["name"] = name + owner[key] = model + + +def _stage_plan( + ctx: dict[str, Any], + forced: set[str], + models: list[ModelUse], +) -> list[StagePreflight]: + plans: list[StagePreflight] = [] + if supports_artifact_cache(ctx): + ctx.setdefault("artifact_versions", {}) + activate_latest_artifacts(ctx, repair=False) + upstream_cache_miss = False + model_stages = {model.stage for model in models} + + for stage_name, raw_cfg in ctx["stages"]: + module = STAGES[stage_name] + enabled = raw_cfg.get("enabled", True) + cacheable = is_cacheable_stage(stage_name) + is_forced = stage_name in forced + action = StageAction.RUN + version = None + reason = "Enabled stage will execute." + if not enabled: + action = StageAction.DISABLED + reason = "Stage is disabled in the effective config." + elif cacheable and supports_artifact_cache(ctx): + if is_forced: + reason = "Forced regeneration invalidates cached output." + upstream_cache_miss = True + elif upstream_cache_miss: + reason = "An upstream cache miss invalidates this stage." + upstream_cache_miss = True + else: + reusable = find_reusable_artifact_plan( + ctx=ctx, + stage_name=stage_name, + raw_cfg=raw_cfg, + ) + if reusable is not None: + action = StageAction.REUSE + version = reusable.version + reason = "Input hashes match a complete artifact version." + activate_artifact_plan(ctx, reusable) + else: + reason = "No complete artifact version matches the inputs." + upstream_cache_miss = True + plans.append( + StagePreflight( + name=stage_name, + scope=module.SCOPE, + action=action, + forced=is_forced, + cacheable=cacheable, + artifact_version=version, + will_call_model=( + enabled + and action is not StageAction.REUSE + and stage_name in model_stages + ), + reason=reason, + ) + ) + return plans + + +def _collect_model_uses(document: dict[str, Any]) -> list[ModelUse]: + pipeline = document.get("pipeline") + if not isinstance(pipeline, dict): + return [] + default = _model_name(document.get("default_model")) + models: list[ModelUse] = [] + + systematize = pipeline.get("systematize") + if _enabled_stage(systematize): + _append_model( + models, + role="systematize", + stage="systematize", + model=_model_name(systematize.get("model")) or default, + ) + + test_set = pipeline.get("test_set") + if _enabled_stage(test_set): + stage_model = _model_name(test_set.get("model")) or default + for role, key in ( + ("test_set_prompt", "prompt"), + ("test_set_scenario", "scenario"), + ("test_set_stratify", "stratify"), + ): + child = test_set.get(key) + if isinstance(child, dict): + child_model = _model_name(child.get("model")) + if key == "stratify": + effective_model = child_model or default or stage_model + else: + effective_model = child_model or stage_model + _append_model( + models, + role=role, + stage="test_set", + model=effective_model, + ) + + inference = pipeline.get("inference") + if _enabled_stage(inference): + target = inference.get("target") + if isinstance(target, dict): + has_non_model_target = any( + target.get(kind) + for kind in ("callable", "connector", "endpoint") + ) + if not has_non_model_target: + _append_model( + models, + role="target", + stage="inference", + model=_model_name(target.get("model")) or default, + ) + tester = inference.get("tester") + if isinstance(tester, dict): + _append_model( + models, + role="tester", + stage="inference", + model=_model_name(tester.get("model")) or default, + ) + + judge = pipeline.get("judge") + if _enabled_stage(judge): + _append_model( + models, + role="judge", + stage="judge", + model=_model_name(judge.get("model")) or default, + ) + return models + + +def _append_model( + models: list[ModelUse], + *, + role: str, + stage: StageName, + model: str | None, +) -> None: + if model is None: + return + models.append( + ModelUse( + role=role, + stage=stage, + model=model, + provider=_model_provider(model), + ) + ) + + +def _enabled_stage(value: Any) -> bool: + return isinstance(value, dict) and value.get("enabled", True) + + +def _model_name(value: Any) -> str | None: + if not isinstance(value, dict): + return None + name = value.get("name") + return str(name) if isinstance(name, str) and name else None + + +def _model_provider(model: str) -> str: + if "/" in model: + return model.split("/", 1)[0].lower() + if model.lower().startswith(("gpt-", "o1", "o3", "o4")): + return "openai" + return "unknown" + + +def _target_preflight( + ctx: dict[str, Any], + policy: PreflightPolicy, +) -> tuple[TargetPreflight | None, list[PreflightIssue]]: + target = ctx.get("target") + if target is None: + return None, [] + issues: list[PreflightIssue] = [] + kind: Literal["model", "callable", "connector", "endpoint"] + identifier: str + probe_required = False + try: + if target.model is not None: + kind = "model" + identifier = str(target.model.name) + elif target.callable: + kind = "callable" + identifier = str(target.callable) + validate_callable_ref(identifier) + probe_required = True + elif target.connector: + kind = "connector" + identifier = str(target.connector) + validate_module_ref(identifier) + probe_required = True + else: + kind = "endpoint" + identifier = _endpoint_origin(str(target.endpoint or "")) + host = urlsplit(str(target.endpoint or "")).hostname + if host is None: + raise ValueError("Endpoint target must include a hostname") + if policy.allowed_endpoint_hosts and not any( + fnmatchcase(host.lower(), pattern.lower()) + for pattern in policy.allowed_endpoint_hosts + ): + issues.append( + PreflightIssue( + code="ENDPOINT_NOT_ALLOWED", + path="/pipeline/inference/target/endpoint", + message=( + f"Endpoint host {host!r} is not allowed by server policy" + ), + ) + ) + except ValueError as exc: + issues.append( + PreflightIssue( + code="TARGET_INVALID", + path="/pipeline/inference/target", + message=str(exc), + ) + ) + return None, issues + return ( + TargetPreflight( + kind=kind, + identifier=identifier, + trace_enabled=target.trace is not None, + static_validation="invalid" if issues else "valid", + probe_required=probe_required, + ), + issues, + ) + + +def _endpoint_origin(value: str) -> str: + parsed = urlsplit(value) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise ValueError("Endpoint target must use http or https and include a hostname") + host = parsed.hostname + if ":" in host and not host.startswith("["): + host = f"[{host}]" + if parsed.port is not None: + host = f"{host}:{parsed.port}" + return urlunsplit((parsed.scheme, host, "", "", "")) + + +def _policy_issues( + document: dict[str, Any], + models: list[ModelUse], + ctx: dict[str, Any], + policy: PreflightPolicy, +) -> list[PreflightIssue]: + issues: list[PreflightIssue] = [] + concurrency = _effective_concurrency(ctx) + if concurrency is not None and concurrency > policy.max_concurrency: + issues.append( + PreflightIssue( + code="CONCURRENCY_LIMIT_EXCEEDED", + path="/pipeline/inference/concurrency", + message=( + f"Concurrency {concurrency} exceeds the server limit " + f"of {policy.max_concurrency}" + ), + ) + ) + sample_sizes = _sample_sizes(document) + for kind, limit in ( + ("prompt", policy.max_prompt_sample_size), + ("scenario", policy.max_scenario_sample_size), + ): + value = sample_sizes[kind] + if value is not None and value > limit: + issues.append( + PreflightIssue( + code="SAMPLE_SIZE_LIMIT_EXCEEDED", + path=f"/pipeline/test_set/{kind}/sample_size", + message=( + f"{kind} sample size {value} exceeds the server limit " + f"of {limit}" + ), + ) + ) + if policy.allowed_model_patterns: + for model in models: + if any( + fnmatchcase(model.model, pattern) + for pattern in policy.allowed_model_patterns + ): + continue + issues.append( + PreflightIssue( + code="MODEL_NOT_ALLOWED", + path=_model_role_path(model.role), + message=( + f"Model {model.model!r} is not allowed by server policy" + ), + ) + ) + return issues + + +def _effective_concurrency(ctx: dict[str, Any]) -> int | None: + evaluation = ctx.get("evaluation") + inference = ( + getattr(evaluation, "inference", None) + if evaluation is not None + else None + ) + value = getattr(inference, "concurrency", None) + return int(value) if isinstance(value, int) else None + + +def _sample_sizes(document: dict[str, Any]) -> dict[str, int | None]: + pipeline = document.get("pipeline") + test_set = ( + pipeline.get("test_set") + if isinstance(pipeline, dict) + else None + ) + result: dict[str, int | None] = {"prompt": None, "scenario": None} + if not isinstance(test_set, dict): + return result + for kind in result: + value = test_set.get(kind) + sample_size = value.get("sample_size") if isinstance(value, dict) else None + result[kind] = sample_size if isinstance(sample_size, int) else None + return result + + +def _credential_requirements( + models: list[ModelUse], +) -> tuple[ + list[CredentialRequirement], + list[PreflightIssue], + list[PreflightIssue], +]: + requirements: list[CredentialRequirement] = [] + blockers: list[PreflightIssue] = [] + warnings: list[PreflightIssue] = [] + providers = sorted({model.provider for model in models}) + for provider in providers: + variables: dict[str, bool] + satisfied: bool | None + note: str + if provider == "azure": + variables = _configured_variables( + "AZURE_API_BASE", + "AZURE_API_KEY", + "ASSERT_AZURE_USE_AAD", + ) + aad_available = _module_available("azure.identity") + satisfied = variables["AZURE_API_BASE"] and ( + variables["AZURE_API_KEY"] or aad_available + ) + note = ( + "AZURE_API_BASE is required; authentication may use " + "AZURE_API_KEY or the Azure identity chain." + ) + elif provider == "azure_ai": + variables = _configured_variables( + "AZURE_AI_API_BASE", + "AZURE_AI_API_KEY", + "ASSERT_AZURE_USE_AAD", + ) + aad_available = _module_available("azure.identity") + satisfied = variables["AZURE_AI_API_BASE"] and ( + variables["AZURE_AI_API_KEY"] or aad_available + ) + note = ( + "AZURE_AI_API_BASE is required; authentication may use " + "AZURE_AI_API_KEY or the Azure identity chain." + ) + elif provider == "openai": + variables = _configured_variables("OPENAI_API_KEY") + satisfied = variables["OPENAI_API_KEY"] + note = "OPENAI_API_KEY is required for OpenAI-hosted models." + elif provider == "anthropic": + variables = _configured_variables("ANTHROPIC_API_KEY") + satisfied = variables["ANTHROPIC_API_KEY"] + note = "ANTHROPIC_API_KEY is required for Anthropic-hosted models." + elif provider in {"gemini", "google"}: + variables = _configured_variables( + "GEMINI_API_KEY", + "GOOGLE_API_KEY", + ) + satisfied = any(variables.values()) + note = "Configure GEMINI_API_KEY or GOOGLE_API_KEY." + elif provider in {"bedrock", "vertex_ai", "unknown"}: + variables = {} + satisfied = None + note = ( + "Credential readiness cannot be determined from environment " + "variable names alone for this provider." + ) + else: + variables = {} + satisfied = None + note = "Provider-specific credential requirements are not known." + + requirements.append( + CredentialRequirement( + provider=provider, + variables=variables, + satisfied=satisfied, + note=note, + ) + ) + if satisfied is False: + blockers.append( + PreflightIssue( + code="CREDENTIAL_CONFIGURATION_MISSING", + path="", + message=f"{provider} credential configuration is incomplete", + ) + ) + elif satisfied is None: + warnings.append( + PreflightIssue( + code="CREDENTIAL_CONFIGURATION_UNKNOWN", + path="", + message=( + f"Credential readiness for provider {provider!r} " + "must be confirmed by the operator" + ), + ) + ) + return requirements, blockers, warnings + + +def _configured_variables(*names: str) -> dict[str, bool]: + return {name: bool(os.environ.get(name)) for name in names} + + +def _module_available(name: str) -> bool: + try: + return importlib.util.find_spec(name) is not None + except ModuleNotFoundError: + return False + + +def _estimate_model_calls( + stages: list[StagePreflight], +) -> ModelCallEstimate: + active_model_stages = [ + stage + for stage in stages + if stage.will_call_model + ] + if not active_model_stages: + return ModelCallEstimate( + minimum=0, + maximum=0, + basis="No non-reused enabled stage uses a model.", + ) + return ModelCallEstimate( + minimum=0, + maximum=None, + basis=( + "Exact calls are not derivable before test-case generation and " + "depend on retries, tester turns, and target behavior." + ), + ) + + +def _model_role_path(role: str) -> str: + return { + "systematize": "/pipeline/systematize/model", + "test_set_prompt": "/pipeline/test_set/prompt/model", + "test_set_scenario": "/pipeline/test_set/scenario/model", + "test_set_stratify": "/pipeline/test_set/stratify/model", + "target": "/pipeline/inference/target/model", + "tester": "/pipeline/inference/tester/model", + "judge": "/pipeline/judge/model", + }.get(role, "") + + +def _validation_issue(issue: ConfigValidationIssue) -> PreflightIssue: + return PreflightIssue( + code=issue.code.value, + path=issue.path, + message=issue.message, + ) + + +def _deduplicate_issues( + issues: list[PreflightIssue], +) -> list[PreflightIssue]: + result: list[PreflightIssue] = [] + seen: set[tuple[str, str, str]] = set() + for issue in issues: + key = (issue.code, issue.path, issue.message) + if key in seen: + continue + seen.add(key) + result.append(issue) + return result diff --git a/assert_ai/services/target_probe.py b/assert_ai/services/target_probe.py new file mode 100644 index 000000000..9d437ac56 --- /dev/null +++ b/assert_ai/services/target_probe.py @@ -0,0 +1,316 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Timeout-bounded target probing in a disposable subprocess.""" + +from __future__ import annotations + +import json +import os +import secrets +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field + +from assert_ai.core.security import ( + redact_path_prefixes, + sanitize_payload, + sanitize_text, +) +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.configs import ConfigService +from assert_ai.services.errors import ServiceError, ServiceErrorCode + +_RESULT_MARKER = "ASSERT_TARGET_PROBE_RESULT=" +_MAX_WORKER_OUTPUT_BYTES = 1024 * 1024 + + +class TargetProbeResult(BaseModel): + """Sanitized result from an isolated import and shape check.""" + + model_config = ConfigDict(frozen=True) + + config_ref: str + target_kind: Literal["model", "callable", "connector", "endpoint"] + ready: Literal[True] = True + isolated: Literal[True] = True + duration_ms: int = Field(ge=0) + details: dict[str, Any] = Field(default_factory=dict) + + +@dataclass(slots=True) +class TargetProbeService: + """Probe managed target code without importing it into the server process.""" + + workspace: WorkspaceService + configs: ConfigService + timeout_s: float = 15.0 + + def __post_init__(self) -> None: + if self.timeout_s <= 0: + raise ValueError("target probe timeout must be positive") + + def probe(self, config_ref: str) -> TargetProbeResult: + record = self.configs.get_config(config_ref) + if not record.validation.valid: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Target probe requires a valid config", + details={ + "validation": record.validation.model_dump(mode="json"), + }, + ) + + request = { + "workspace_root": str(self.workspace.root), + "config_ref": record.config_ref, + "max_config_bytes": self.configs.max_config_bytes, + } + started = time.perf_counter() + payload = self._invoke_worker(request) + duration_ms = max( + 0, + round((time.perf_counter() - started) * 1000), + ) + if payload.get("ok") is not True: + message = self._sanitize_message( + str(payload.get("message") or "Target probe failed") + ) + details = self._sanitize_details( + { + "target_kind": payload.get("target_kind"), + "error_type": payload.get("error_type"), + } + ) + raise ServiceError( + ServiceErrorCode.TARGET_IMPORT_FAILED, + message, + details=details, + ) + + target_kind = payload.get("target_kind") + if target_kind not in {"model", "callable", "connector", "endpoint"}: + raise ServiceError( + ServiceErrorCode.TARGET_IMPORT_FAILED, + "Target probe worker returned an invalid target kind", + ) + details = payload.get("details") + if not isinstance(details, dict): + raise ServiceError( + ServiceErrorCode.TARGET_IMPORT_FAILED, + "Target probe worker returned invalid details", + ) + return TargetProbeResult( + config_ref=record.config_ref, + target_kind=target_kind, + duration_ms=duration_ms, + details=self._sanitize_details(details), + ) + + def _invoke_worker(self, request: dict[str, Any]) -> dict[str, Any]: + result_token = secrets.token_hex(16) + worker_request = { + **request, + "result_token": result_token, + } + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" + env["PYTHONDONTWRITEBYTECODE"] = "1" + creationflags = ( + subprocess.CREATE_NEW_PROCESS_GROUP + if os.name == "nt" + else 0 + ) + with ( + tempfile.TemporaryFile(mode="w+b") as stdout_file, + tempfile.TemporaryFile(mode="w+b") as stderr_file, + ): + process = subprocess.Popen( + [ + sys.executable, + "-m", + "assert_ai.services._target_probe_worker", + ], + cwd=self.workspace.root, + env=env, + stdin=subprocess.PIPE, + stdout=stdout_file, + stderr=stderr_file, + text=True, + encoding="utf-8", + creationflags=creationflags, + start_new_session=os.name != "nt", + ) + process_create_time = _process_create_time(process.pid) + try: + process.communicate( + json.dumps(worker_request, ensure_ascii=False), + timeout=self.timeout_s, + ) + except subprocess.TimeoutExpired as exc: + _terminate_process_tree( + process, + expected_create_time=process_create_time, + ) + raise ServiceError( + ServiceErrorCode.TARGET_IMPORT_FAILED, + ( + "Target probe exceeded the operator timeout " + f"of {self.timeout_s:g} seconds" + ), + details={"timed_out": True}, + ) from exc + + output = _read_tail( + stdout_file, + max_bytes=_MAX_WORKER_OUTPUT_BYTES, + ) + payload = _parse_worker_result( + output, + result_token=result_token, + ) + if payload is not None: + return payload + raise ServiceError( + ServiceErrorCode.TARGET_IMPORT_FAILED, + "Target probe worker did not return a valid result", + details={"exit_code": process.returncode}, + ) + + def _sanitize_message(self, message: str) -> str: + return redact_path_prefixes( + sanitize_text(message), + ( + self.workspace.root, + self.workspace.configs_root, + self.workspace.artifacts_root, + self.workspace.results_root, + ), + ) + + def _sanitize_details(self, details: dict[str, Any]) -> dict[str, Any]: + sanitized = sanitize_payload(details) + assert isinstance(sanitized, dict) + result = _redact_paths( + sanitized, + workspace=self.workspace, + ) + assert isinstance(result, dict) + return result + + +def _read_tail(handle: Any, *, max_bytes: int) -> str: + handle.flush() + handle.seek(0, os.SEEK_END) + size = handle.tell() + handle.seek(max(0, size - max_bytes)) + return handle.read(max_bytes).decode("utf-8", errors="replace") + + +def _parse_worker_result( + output: str, + *, + result_token: str, +) -> dict[str, Any] | None: + marker = f"{_RESULT_MARKER}{result_token}=" + marker_index = output.rfind(marker) + if marker_index < 0: + return None + encoded = output[marker_index + len(marker) :].splitlines()[0] + try: + payload = json.loads(encoded) + except (TypeError, ValueError, json.JSONDecodeError): + return None + return payload if isinstance(payload, dict) else None + + +def _process_create_time(pid: int) -> float | None: + try: + import psutil + except ImportError: + return None + try: + return float(psutil.Process(pid).create_time()) + except (OSError, psutil.Error): + return None + + +def _terminate_process_tree( + process: subprocess.Popen[str], + *, + expected_create_time: float | None, +) -> None: + try: + import psutil + except ImportError: + process.kill() + process.wait(timeout=5) + return + + try: + parent = psutil.Process(process.pid) + if ( + expected_create_time is not None + and abs(parent.create_time() - expected_create_time) > 0.001 + ): + return + descendants = parent.children(recursive=True) + for child in descendants: + child.terminate() + parent.terminate() + _, alive = psutil.wait_procs( + [*descendants, parent], + timeout=2, + ) + for item in alive: + item.kill() + psutil.wait_procs(alive, timeout=2) + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.TimeoutExpired): + pass + finally: + if process.poll() is None: + process.kill() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + +def _redact_paths( + value: Any, + *, + workspace: WorkspaceService, +) -> Any: + if isinstance(value, dict): + return { + str(key): _redact_paths(item, workspace=workspace) + for key, item in value.items() + } + if isinstance(value, list): + return [ + _redact_paths(item, workspace=workspace) + for item in value + ] + if isinstance(value, tuple): + return [ + _redact_paths(item, workspace=workspace) + for item in value + ] + if isinstance(value, str): + return redact_path_prefixes( + value, + ( + workspace.root, + workspace.configs_root, + workspace.artifacts_root, + workspace.results_root, + ), + ) + return value diff --git a/pyproject.toml b/pyproject.toml index 05da85b94..7903c760e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,8 +93,8 @@ acs = [ ] mcp = [ # OpenAI Agents currently requires mcp<2, so this extra is intentionally - # separate from `examples` and the `all` meta-extra. psutil is required by - # the planned Windows worker process-tree cancellation path. + # separate from `examples` and the `all` meta-extra. psutil is required for + # isolated target-probe cleanup and worker process-tree cancellation. "mcp>=2,<3", "psutil>=6,<8", ] diff --git a/tests/test_artifact_cache.py b/tests/test_artifact_cache.py index 776c78539..1bc115cf0 100644 --- a/tests/test_artifact_cache.py +++ b/tests/test_artifact_cache.py @@ -199,6 +199,47 @@ def test_activate_latest_recovers_when_referenced_artifact_dir_is_missing(self) self.assertIsNotNone(recovered) self.assertEqual(recovered["version"], v0001_plan.version) + def test_activate_latest_without_repair_does_not_write_recovered_state( + self, + ) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = self._ctx(root) + raw_cfg = { + "model": {"name": "azure/gpt-5.4"}, + "behavior_category_count": 2, + } + v0001_plan = self._finalize_policy(ctx, raw_cfg) + + changed_ctx = self._ctx(root) + changed_ctx["behavior"] = "Changed behavior text." + v0002_plan = self._finalize_policy(changed_ctx, raw_cfg) + shutil.rmtree(v0002_plan.artifact_dir) + + suite_root = Path(ctx["suite_root"]) + latest_path = suite_root / "latest.json" + compatibility_path = suite_root / "taxonomy.json" + compatibility_path.write_text( + '{"sentinel":"leave-unchanged"}', + encoding="utf-8", + ) + latest_before = latest_path.read_bytes() + compatibility_before = compatibility_path.read_bytes() + + recovery_ctx = self._ctx(root) + activate_latest_artifacts(recovery_ctx, repair=False) + + recovered = recovery_ctx.get("artifact_versions", {}).get( + "systematize" + ) + self.assertIsNotNone(recovered) + self.assertEqual(recovered["version"], v0001_plan.version) + self.assertEqual(latest_path.read_bytes(), latest_before) + self.assertEqual( + compatibility_path.read_bytes(), + compatibility_before, + ) + def test_activate_latest_skips_stage_when_no_valid_version_remains(self) -> None: with TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) diff --git a/tests/test_mcp_cli.py b/tests/test_mcp_cli.py index 5ac393db5..10f0d91c9 100644 --- a/tests/test_mcp_cli.py +++ b/tests/test_mcp_cli.py @@ -59,6 +59,18 @@ def test_mcp_serve_forwards_resolved_options() -> None: "2048", "--max-config-bytes", "4096", + "--max-concurrency", + "7", + "--max-prompt-sample-size", + "12", + "--max-scenario-sample-size", + "13", + "--allowed-model", + "azure/*", + "--allowed-endpoint-host", + "api.example.test", + "--target-probe-timeout-seconds", + "4.5", ], ) @@ -73,6 +85,12 @@ def test_mcp_serve_forwards_resolved_options() -> None: assert create_kwargs["default_artifact_chunk_bytes"] == 1024 assert create_kwargs["max_artifact_chunk_bytes"] == 2048 assert create_kwargs["max_config_bytes"] == 4096 + assert create_kwargs["max_concurrency"] == 7 + assert create_kwargs["max_prompt_sample_size"] == 12 + assert create_kwargs["max_scenario_sample_size"] == 13 + assert create_kwargs["allowed_model_patterns"] == ("azure/*",) + assert create_kwargs["allowed_endpoint_hosts"] == ("api.example.test",) + assert create_kwargs["target_probe_timeout_s"] == 4.5 run_stdio_server.assert_called_once_with(options) @@ -90,6 +108,25 @@ def test_mcp_serve_reports_missing_optional_dependency() -> None: assert "Traceback" not in result.output +def test_mcp_serve_rejects_nonpositive_preflight_limit() -> None: + runner = CliRunner() + + with patch("assert_ai.mcp._command._load_server_module") as load_server: + result = runner.invoke( + cli, + [ + "mcp", + "serve", + "--max-concurrency", + "0", + ], + ) + + assert result.exit_code == 2 + assert "not in the range" in result.output + load_server.assert_not_called() + + def test_mcp_serve_loads_workspace_env_before_server() -> None: runner = CliRunner() calls: list[str] = [] diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 16de7aa62..9fa322ae3 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -11,6 +11,7 @@ from contextlib import asynccontextmanager from pathlib import Path from typing import Any, AsyncIterator +from unittest.mock import Mock, patch import pytest @@ -20,8 +21,10 @@ from mcp.client._transport import TransportStreams from mcp.client.stdio import StdioServerParameters, stdio_client +from assert_ai.core.config_document import ConfigValidationReport from assert_ai.mcp.models import CapabilityGroup, ServerMode from assert_ai.mcp.server import ServerOptions, build_server +from assert_ai.services.configs import ConfigDraft from tests.result_catalog_fixture import create_result_catalog_fixture EXPECTED_INSPECT_TOOLS = { @@ -44,6 +47,15 @@ "list_artifacts", "read_artifact_chunk", } +EXPECTED_AUTHOR_TOOLS = EXPECTED_INSPECT_TOOLS | { + "validate_config", + "save_config", + "preflight_evaluation", +} +EXPECTED_FULL_TOOLS = EXPECTED_AUTHOR_TOOLS | { + "design_config", + "probe_target", +} EXPECTED_RESOURCE_TEMPLATES = { "assert://preset/{kind}/{name}", @@ -303,19 +315,70 @@ def test_server_options_validate_response_limits(tmp_path: Path) -> None: ) -def test_design_group_requires_author_or_full_mode(tmp_path: Path) -> None: +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("max_concurrency", 0, "max_concurrency must be positive"), + ( + "max_prompt_sample_size", + 0, + "max_prompt_sample_size must be positive", + ), + ( + "max_scenario_sample_size", + 0, + "max_scenario_sample_size must be positive", + ), + ( + "allowed_model_patterns", + (" ",), + "allowed_model_patterns cannot contain empty values", + ), + ( + "allowed_endpoint_hosts", + ("",), + "allowed_endpoint_hosts cannot contain empty values", + ), + ], +) +def test_server_options_validate_preflight_policy( + tmp_path: Path, + field: str, + value: object, + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + ServerOptions( + workspace_root=tmp_path, + **{field: value}, + ) + + +@pytest.mark.parametrize("group", ["design", "probe"]) +def test_author_extension_groups_require_author_or_full_mode( + tmp_path: Path, + group: str, +) -> None: with pytest.raises(ValueError, match="require --mode author or --mode full"): ServerOptions.create( workspace_root=tmp_path, mode="inspect", - enabled_groups=["design"], + enabled_groups=[group], ) -@pytest.mark.parametrize("mode", list(ServerMode)) -def test_inspect_tools_are_registered_in_every_base_mode( +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (ServerMode.INSPECT, EXPECTED_INSPECT_TOOLS), + (ServerMode.AUTHOR, EXPECTED_AUTHOR_TOOLS), + (ServerMode.FULL, EXPECTED_FULL_TOOLS), + ], +) +def test_tools_are_registered_for_each_base_mode( tmp_path: Path, mode: ServerMode, + expected: set[str], ) -> None: async def run() -> set[str]: options = ServerOptions.create(workspace_root=tmp_path, mode=mode) @@ -323,7 +386,32 @@ async def run() -> set[str]: tools = await client.list_tools() return {tool.name for tool in tools.tools} - assert asyncio.run(run()) == EXPECTED_INSPECT_TOOLS + assert asyncio.run(run()) == expected + + +@pytest.mark.parametrize( + ("group", "tool"), + [ + ("design", "design_config"), + ("probe", "probe_target"), + ], +) +def test_author_extension_groups_register_explicitly( + tmp_path: Path, + group: str, + tool: str, +) -> None: + async def run() -> set[str]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="author", + enabled_groups=[group], + ) + async with Client(build_server(options), raise_exceptions=True) as client: + tools = await client.list_tools() + return {item.name for item in tools.tools} + + assert asyncio.run(run()) == EXPECTED_AUTHOR_TOOLS | {tool} def test_get_server_info_protocol_round_trip(tmp_path: Path) -> None: @@ -332,6 +420,8 @@ async def run() -> object: workspace_root=tmp_path, mode="full", enabled_groups=["analysis"], + allowed_model_patterns=["azure/*"], + allowed_endpoint_hosts=["api.example.test"], ) async with Client(build_server(options), raise_exceptions=True) as client: return await client.call_tool("get_server_info", {}) @@ -345,6 +435,27 @@ async def run() -> object: assert result.structured_content["workspace"]["root"] == "." assert "env_file" not in result.structured_content assert result.structured_content["limits"]["max_page_size"] == 200 + assert result.structured_content["limits"]["max_concurrency"] == 32 + assert result.structured_content["limits"]["max_prompt_sample_size"] == 100_000 + assert result.structured_content["limits"]["max_scenario_sample_size"] == 100_000 + assert result.structured_content["limits"]["model_allowlist_enabled"] is True + assert ( + result.structured_content["limits"]["endpoint_host_allowlist_enabled"] + is True + ) + assert result.structured_content["limits"]["allowed_model_patterns"] == [ + "azure/*" + ] + assert result.structured_content["limits"]["allowed_endpoint_hosts"] == [ + "api.example.test" + ] + assert result.structured_content["limits"]["target_probe_timeout_s"] == 15.0 + assert result.structured_content["target_kinds"] == [ + "callable", + "model", + "connector", + "endpoint", + ] assert result.structured_content["enabled_capability_groups"] == [ "inspect", "author", @@ -385,7 +496,7 @@ async def run() -> list[Any]: "get_config_schema": "cca1d3a48240e20eff93a123b34d7ba92df3ed1df87f57f9eb217aa21515ec26", "get_preset": "25352522a3ed4ff76217c5415453c6641ad229c2dd6e4e05df4621b89ce4819c", "get_run": "e5216cd0085d049f8b49c54add913b6f83756c4ce59995317fe63e010ea44936", - "get_server_info": "d51f9ff9fe235b5c53f5db71a1bba11cfb35bbc27be1f8c3552f5bee8ecc5e8d", + "get_server_info": "19b16f3a1acbbe39914256977d48c594c6ac68a4ee6a406a00d70c09d3e72935", "get_suite": "8f629c93e02b656052f637c3cbba9217834315a693c4f7935f6d961203b46fd0", "get_test_case": "11380555caaa71d5992923815a499fc08b368c02f4d4836e4761630654589148", "get_transcript": "aa09669e0cb99202e8dec0b858b4faa41742ecb616351c3956b0d0bd488717e8", @@ -401,6 +512,239 @@ async def run() -> list[Any]: } +def test_author_tools_publish_stable_schemas_and_annotations( + tmp_path: Path, +) -> None: + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="full", + ) + async with Client(build_server(options), raise_exceptions=True) as client: + tools = (await client.list_tools()).tools + return {tool.name: tool for tool in tools} + + tools = asyncio.run(run()) + expected_annotations = { + "validate_config": (True, False, True, False), + "save_config": (False, True, False, False), + "preflight_evaluation": (True, False, True, False), + "design_config": (True, False, False, True), + "probe_target": (True, False, False, True), + } + expected_digests = { + "validate_config": ( + "b3068ce71b224596e9a8c25775ecaed0ab83feded27cc0fc9e26477153956203" + ), + "save_config": ( + "b09950417a44bf14c9bbf2702c1c00f23a18a0bfec03cab16a482733d8cf98c8" + ), + "preflight_evaluation": ( + "e86012ecf7fd9684714e25abc052341f22e6f9cac78d08d112b19c03111b7e29" + ), + "design_config": ( + "1cd55a1bba06468b0aaa785cf05a445e4bf16768ff6165254ceecf0181a8392e" + ), + "probe_target": ( + "0a46b06b94ef40d13b579fa3facb60024875c9a3348c9d80265c71b05837456b" + ), + } + + for name, annotations in expected_annotations.items(): + tool = tools[name] + actual = tool.annotations + assert actual is not None + assert ( + actual.read_only_hint, + actual.destructive_hint, + actual.idempotent_hint, + actual.open_world_hint, + ) == annotations + assert _schema_digest(tool) == expected_digests[name] + + +def test_complete_author_preflight_and_probe_workflow( + tmp_path: Path, +) -> None: + (tmp_path / "agent.py").write_text( + "def run(message, *, history=None):\n" + " return message\n", + encoding="utf-8", + ) + document = { + "suite": "author-suite", + "context": "authorization: not-a-real-secret", + "pipeline": { + "inference": { + "target": {"callable": "agent:run"}, + "test_set_path": "fixtures/test_set.jsonl", + } + }, + } + + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="full", + ) + async with Client(build_server(options), raise_exceptions=True) as client: + invalid = await client.call_tool( + "validate_config", + { + "yaml_text": ( + "pipeline: {}\n" + "api_key: 'not-a-real-secret\n" + ) + }, + ) + valid = await client.call_tool( + "validate_config", + {"document": document}, + ) + saved = await client.call_tool( + "save_config", + { + "config_ref": "nested/agent.yaml", + "document": document, + }, + ) + loaded = await client.call_tool( + "get_config", + {"config_ref": "nested/agent.yaml"}, + ) + preflight = await client.call_tool( + "preflight_evaluation", + { + "config_ref": "nested/agent.yaml", + "overrides": { + "run": "candidate-a", + "concurrency": 3, + }, + }, + ) + revised_document = { + **document, + "context": "Evaluate deterministic echo behavior.", + } + replaced = await client.call_tool( + "save_config", + { + "config_ref": "nested/agent.yaml", + "document": revised_document, + "expected_etag": loaded.structured_content["etag"], + }, + ) + probe = await client.call_tool( + "probe_target", + {"config_ref": "nested/agent.yaml"}, + ) + return { + "invalid": invalid.structured_content, + "valid": valid.structured_content, + "saved": saved.structured_content, + "loaded": loaded.structured_content, + "preflight": preflight.structured_content, + "replaced": replaced.structured_content, + "probe": probe.structured_content, + } + + results = asyncio.run(run()) + + assert results["invalid"]["validation"]["valid"] is False + assert "not-a-real-secret" not in json.dumps(results["invalid"]) + assert results["valid"]["validation"]["valid"] is True + assert results["saved"]["created"] is True + assert results["saved"]["resource_uri"] == ( + "assert://config/nested%2Fagent.yaml" + ) + assert results["preflight"]["ready"] is True + assert results["preflight"]["run_id"] == "candidate-a" + assert results["preflight"]["concurrency"] == 3 + assert results["preflight"]["target"]["kind"] == "callable" + assert "not-a-real-secret" not in json.dumps(results["preflight"]) + assert "[REDACTED]" in json.dumps(results["preflight"]) + assert results["loaded"]["etag"] == results["saved"]["etag"] + assert results["replaced"]["created"] is False + assert results["replaced"]["etag"] != results["saved"]["etag"] + assert results["probe"]["target_kind"] == "callable" + assert results["probe"]["details"]["reference"] == "agent:run" + assert not (tmp_path / "artifacts").exists() + + +def test_design_config_returns_an_unpersisted_draft(tmp_path: Path) -> None: + draft = ConfigDraft( + yaml=( + "pipeline: {}\n" + "api_key: not-a-real-secret\n" + ), + document={ + "pipeline": {}, + "api_key": "not-a-real-secret", + }, + validation=ConfigValidationReport(valid=True), + ) + + async def run() -> object: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="author", + enabled_groups=["design"], + ) + async with Client(build_server(options), raise_exceptions=True) as client: + return await client.call_tool( + "design_config", + { + "description": "Evaluate a local deterministic agent", + "max_turns": 3, + }, + ) + + with patch( + "assert_ai.services.configs.ConfigService.design_config", + new=Mock(return_value=draft), + ) as design_config: + result = asyncio.run(run()) + + assert "not-a-real-secret" not in json.dumps(result.structured_content) + assert "[REDACTED]" in result.structured_content["yaml"] + assert result.structured_content["persisted"] is False + assert result.structured_content["model_cost_incurred"] is True + assert design_config.call_count == 1 + request = design_config.call_args.args[0] + assert request.description == "Evaluate a local deterministic agent" + assert request.max_turns == 3 + assert not (tmp_path / "evals").exists() + + +def test_design_config_enforces_operator_model_allowlist( + tmp_path: Path, +) -> None: + async def run() -> object: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="author", + enabled_groups=["design"], + allowed_model_patterns=["openai/*"], + ) + async with Client(build_server(options), raise_exceptions=True) as client: + return await client.call_tool( + "design_config", + { + "description": "Draft an evaluation", + "model": "azure/gpt-5.4-mini", + }, + ) + + with patch( + "assert_ai.services.configs.ConfigService.design_config", + ) as design_config: + result = asyncio.run(run()) + + assert '"code":"INVALID_ARGUMENT"' in _error_text(result) + assert "not allowed by server policy" in _error_text(result) + design_config.assert_not_called() + + def test_complete_read_only_tool_workflow(tmp_path: Path) -> None: _seed_workspace(tmp_path) @@ -618,6 +962,34 @@ async def run() -> tuple[set[str], set[str], dict[str, str]]: assert str(tmp_path) not in json.dumps(contents) +@pytest.mark.parametrize( + ("mode", "raise_exceptions"), + [("auto", True), ("legacy", False)], +) +def test_author_errors_are_stable_across_protocol_modes( + tmp_path: Path, + mode: str, + raise_exceptions: bool, +) -> None: + async def run() -> object: + async with Client( + build_server( + ServerOptions.create( + workspace_root=tmp_path, + mode="author", + ) + ), + mode=mode, + raise_exceptions=raise_exceptions, + ) as client: + return await client.call_tool("validate_config", {}) + + text = _error_text(asyncio.run(run())) + + assert '"code":"INVALID_ARGUMENT"' in text + assert "Provide exactly one" in text + + @pytest.mark.parametrize( ("mode", "raise_exceptions"), [("auto", True), ("legacy", False)], diff --git a/tests/test_run_planning_service.py b/tests/test_run_planning_service.py new file mode 100644 index 000000000..9ec59b09e --- /dev/null +++ b/tests/test_run_planning_service.py @@ -0,0 +1,327 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import os +from copy import deepcopy +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +import pytest + +from assert_ai.config import load_runtime_context +from assert_ai.core.artifact_cache import ( + activate_artifact_plan, + finalize_artifact_plan, + prepare_artifact_plan, +) +from assert_ai.core.run_plan import resolve_forced_stages +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.configs import ConfigService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.run_planning import ( + EvaluationOverrides, + ModelOverrides, + PreflightPolicy, + RunPlanningService, + StageAction, +) +from assert_ai.stages import STAGES + + +def _document(*, model_target: bool = False) -> dict: + target = ( + {"model": {"name": "openai/gpt-test"}} + if model_target + else {"callable": "agent:run"} + ) + return { + "suite": "demo-suite", + "behavior": { + "name": "safe_help", + "description": "The agent should provide safe help.", + }, + "default_model": {"name": "openai/gpt-test"}, + "pipeline": { + "systematize": {}, + "test_set": { + "prompt": {"sample_size": 2}, + }, + "inference": { + "target": target, + "concurrency": 2, + }, + "judge": {}, + }, + } + + +def _services( + root: Path, + *, + policy: PreflightPolicy | None = None, +) -> tuple[ConfigService, RunPlanningService]: + workspace = WorkspaceService.create(root) + configs = ConfigService(workspace) + return ( + configs, + RunPlanningService( + workspace, + configs, + policy=policy or PreflightPolicy(), + ), + ) + + +def test_preflight_is_pure_and_matches_force_cascade() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + configs, planning = _services( + root, + policy=PreflightPolicy( + allowed_model_patterns=("openai/*",), + ), + ) + configs.save_config("demo.yaml", document=_document()) + before = { + path.relative_to(root).as_posix() + for path in root.rglob("*") + } + + with patch.dict( + os.environ, + {"OPENAI_API_KEY": "configured-for-test"}, + clear=False, + ): + result = planning.preflight( + "demo.yaml", + overrides=EvaluationOverrides( + run="candidate-a", + force_stages=("test_set",), + strict=True, + concurrency=3, + prompt_sample_size=4, + ), + ) + + after = { + path.relative_to(root).as_posix() + for path in root.rglob("*") + } + stages = {stage.name: stage for stage in result.stages} + assert result.ready is True + assert result.run_id == "candidate-a" + assert result.strict is True + assert result.concurrency == 3 + assert result.sample_sizes["prompt"] == 4 + assert result.target is not None + assert result.target.kind == "callable" + assert result.target.probe_required is True + assert stages["systematize"].forced is False + assert stages["test_set"].forced is True + assert stages["inference"].forced is True + assert stages["judge"].forced is True + assert all( + stage.action is StageAction.RUN + for stage in result.stages + ) + assert result.managed_outputs["suite_root"] == ( + "artifacts/results/demo-suite" + ) + assert result.managed_outputs["run_root"].endswith( + "/candidate-a" + ) + assert before == after + + +def test_preflight_reports_policy_and_credential_blockers() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + configs, planning = _services( + root, + policy=PreflightPolicy( + max_concurrency=2, + max_prompt_sample_size=5, + allowed_model_patterns=("azure/*",), + ), + ) + document = _document(model_target=True) + document["pipeline"]["inference"]["concurrency"] = 8 + document["pipeline"]["test_set"]["prompt"]["sample_size"] = 10 + configs.save_config("demo.yaml", document=document) + + with patch.dict( + os.environ, + {"OPENAI_API_KEY": ""}, + clear=False, + ): + result = planning.preflight("demo.yaml") + + codes = {issue.code for issue in result.blocking_issues} + assert result.ready is False + assert "CONCURRENCY_LIMIT_EXCEEDED" in codes + assert "SAMPLE_SIZE_LIMIT_EXCEEDED" in codes + assert "MODEL_NOT_ALLOWED" in codes + assert "CREDENTIAL_CONFIGURATION_MISSING" in codes + assert result.credentials[0].variables == { + "OPENAI_API_KEY": False, + } + + +def test_preflight_returns_structural_validation_without_side_effects() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + configs, planning = _services(root) + configs.workspace.configs_root.mkdir(parents=True) + bad_path = configs.workspace.configs_root / "bad.yaml" + bad_path.write_text( + "pipeline:\n inference:\n unknown: true\n", + encoding="utf-8", + ) + + result = planning.preflight("bad.yaml") + + assert result.ready is False + assert result.stages == () + assert result.blocking_issues[0].code == "UNKNOWN_FIELD" + assert not configs.workspace.artifacts_root.exists() + + +def test_model_override_cannot_replace_callable_target() -> None: + with TemporaryDirectory() as tmp: + configs, planning = _services(Path(tmp)) + configs.save_config("demo.yaml", document=_document()) + + with pytest.raises(ServiceError) as invalid: + planning.preflight( + "demo.yaml", + overrides=EvaluationOverrides( + models=ModelOverrides( + target_model="openai/replacement", + ) + ), + ) + + assert invalid.value.code == ServiceErrorCode.INVALID_ARGUMENT + + +def test_stratify_model_planning_matches_runtime_fallback_order() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + configs, planning = _services( + root, + policy=PreflightPolicy( + allowed_model_patterns=("azure/*",), + ), + ) + document = _document() + document["pipeline"]["systematize"]["enabled"] = False + document["pipeline"]["inference"]["enabled"] = False + document["pipeline"]["judge"]["enabled"] = False + document["pipeline"]["test_set"]["model"] = { + "name": "azure/test-set", + } + document["pipeline"]["test_set"]["stratify"] = {} + configs.save_config("demo.yaml", document=document) + + with patch.dict( + os.environ, + { + "OPENAI_API_KEY": "configured-for-test", + "AZURE_API_BASE": "https://example.openai.azure.com", + "AZURE_API_KEY": "configured-for-test", + }, + clear=False, + ): + result = planning.preflight("demo.yaml") + + models = {model.role: model.model for model in result.models} + assert models["test_set_prompt"] == "azure/test-set" + assert models["test_set_stratify"] == "openai/gpt-test" + model_issues = [ + issue + for issue in result.blocking_issues + if issue.code == "MODEL_NOT_ALLOWED" + ] + assert len(model_issues) == 1 + assert model_issues[0].path == "/pipeline/test_set/stratify/model" + + +def test_preflight_reuses_cache_without_writing_workspace() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + configs, planning = _services(root) + configs.save_config("demo.yaml", document=_document()) + record = configs.get_config("demo.yaml") + config_path = configs.workspace.path_policy.resolve_config_path( + record.config_ref, + must_exist=True, + reject_links=True, + ) + ctx = load_runtime_context( + deepcopy(record.document), + config_path, + stage_modules=STAGES, + path_policy=configs.workspace.path_policy, + ) + raw_cfg = dict( + next( + raw + for name, raw in ctx["stages"] + if name == "systematize" + ) + ) + plan = prepare_artifact_plan( + ctx=ctx, + stage_name="systematize", + raw_cfg=raw_cfg, + forced=False, + ) + activate_artifact_plan(ctx, plan) + plan.output_paths["taxonomy"].parent.mkdir( + parents=True, + exist_ok=True, + ) + plan.output_paths["taxonomy"].write_text( + '{"behavior_categories":[]}', + encoding="utf-8", + ) + plan.output_paths["systematization"].write_text( + "{}", + encoding="utf-8", + ) + finalize_artifact_plan(ctx, plan) + before = { + path.relative_to(root).as_posix(): path.read_bytes() + for path in root.rglob("*") + if path.is_file() + } + + with patch.dict( + os.environ, + {"OPENAI_API_KEY": "configured-for-test"}, + clear=False, + ): + result = planning.preflight("demo.yaml") + + after = { + path.relative_to(root).as_posix(): path.read_bytes() + for path in root.rglob("*") + if path.is_file() + } + stages = {stage.name: stage for stage in result.stages} + assert stages["systematize"].action is StageAction.REUSE + assert stages["systematize"].artifact_version == plan.version + assert before == after + + +def test_resolve_forced_stages_rejects_missing_and_cascades() -> None: + assert resolve_forced_stages( + ("systematize", "test_set", "inference", "judge"), + ("test_set",), + ) == ("test_set", "inference", "judge") + + with pytest.raises(ValueError, match="missing"): + resolve_forced_stages(("inference",), ("missing",)) diff --git a/tests/test_target_probe_service.py b/tests/test_target_probe_service.py new file mode 100644 index 000000000..260f0126f --- /dev/null +++ b/tests/test_target_probe_service.py @@ -0,0 +1,180 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from pathlib import Path +from tempfile import TemporaryDirectory + +import pytest + +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.configs import ConfigService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.target_probe import ( + TargetProbeService, + _parse_worker_result, +) + + +def _callable_document(reference: str) -> dict: + return { + "suite": "probe-suite", + "pipeline": { + "inference": { + "target": {"callable": reference}, + "test_set_path": "fixtures/test_set.jsonl", + } + }, + } + + +def _service( + root: Path, + *, + timeout_s: float = 15.0, +) -> tuple[ConfigService, TargetProbeService]: + workspace = WorkspaceService.create(root) + configs = ConfigService(workspace) + return configs, TargetProbeService( + workspace, + configs, + timeout_s=timeout_s, + ) + + +def test_callable_probe_imports_only_in_worker() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "agent.py").write_text( + "async def run(message, *, history=None):\n" + " return message\n", + encoding="utf-8", + ) + configs, probe = _service(root) + configs.save_config( + "demo.yaml", + document=_callable_document("agent:run"), + ) + + result = probe.probe("demo.yaml") + + assert result.target_kind == "callable" + assert result.isolated is True + assert result.details["reference"] == "agent:run" + assert result.details["is_async"] is True + assert result.details["accepts_history"] is True + assert result.details["parameters"] == [ + "message", + "history", + ] + assert not (root / "__pycache__").exists() + + +def test_probe_failure_redacts_credentials_and_workspace_paths() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "broken.py").write_text( + "raise RuntimeError(" + "f'AZURE_API_KEY=not-a-real-secret path={__file__}'" + ")\n", + encoding="utf-8", + ) + configs, probe = _service(root) + configs.save_config( + "demo.yaml", + document=_callable_document("broken:run"), + ) + + with pytest.raises(ServiceError) as failed: + probe.probe("demo.yaml") + + assert failed.value.code == ServiceErrorCode.TARGET_IMPORT_FAILED + assert "not-a-real-secret" not in str(failed.value) + assert str(root) not in str(failed.value) + assert "[REDACTED]" in str(failed.value) + + +def test_probe_timeout_terminates_worker() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "slow.py").write_text( + "from pathlib import Path\n" + "import subprocess\n" + "import sys\n" + "import time\n" + "child = subprocess.Popen(" + "[sys.executable, '-c', 'import time; time.sleep(60)']" + ")\n" + "Path('child.pid').write_text(str(child.pid), encoding='utf-8')\n" + "time.sleep(5)\n" + "def run(message):\n" + " return message\n", + encoding="utf-8", + ) + configs, probe = _service(root, timeout_s=1.5) + configs.save_config( + "demo.yaml", + document=_callable_document("slow:run"), + ) + + with pytest.raises(ServiceError) as timed_out: + probe.probe("demo.yaml") + + assert timed_out.value.code == ServiceErrorCode.TARGET_IMPORT_FAILED + assert timed_out.value.details == {"timed_out": True} + assert "1.5 seconds" in str(timed_out.value) + child_pid = int((root / "child.pid").read_text(encoding="utf-8")) + import psutil + + assert not psutil.pid_exists(child_pid) + + +def test_model_probe_is_static_and_reports_model() -> None: + with TemporaryDirectory() as tmp: + root = Path(tmp) + configs, probe = _service(root) + configs.save_config( + "demo.yaml", + document={ + "suite": "probe-suite", + "pipeline": { + "inference": { + "target": { + "model": {"name": "openai/gpt-test"}, + }, + "test_set_path": "fixtures/test_set.jsonl", + } + }, + }, + ) + + result = probe.probe("demo.yaml") + + assert result.target_kind == "model" + assert result.details == { + "model": "openai/gpt-test", + "trace_enabled": False, + } + + +def test_worker_result_parser_requires_the_invocation_token() -> None: + output = "\n".join( + ( + 'ASSERT_TARGET_PROBE_RESULT=expected={"ok":true}', + 'ASSERT_TARGET_PROBE_RESULT=forged={"ok":false}', + 'ASSERT_TARGET_PROBE_RESULT={"ok":false}', + ) + ) + + assert _parse_worker_result( + output, + result_token="expected", + ) == {"ok": True} + assert ( + _parse_worker_result( + output, + result_token="missing", + ) + is None + ) From 65c9012005c927d9d6c7eeb3d7d3bfa71a6cb26a Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Tue, 25 Aug 2026 14:56:06 -0700 Subject: [PATCH 09/16] Add persisted MCP evaluation jobs Add idempotent SQLite-backed jobs, isolated snapshot workers, bounded diagnostics, and start/list/get MCP tools with operator queue controls. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/core/runtime_path_policy.py | 20 +- assert_ai/core/runtime_safety.py | 9 +- assert_ai/mcp/_command.py | 27 + assert_ai/mcp/models.py | 3 + assert_ai/mcp/resources.py | 18 + assert_ai/mcp/server.py | 92 +- assert_ai/mcp/tools/__init__.py | 8 + assert_ai/mcp/tools/jobs.py | 132 +++ assert_ai/mcp/uris.py | 4 + assert_ai/runner.py | 101 +- assert_ai/services/_evaluation_worker.py | 352 ++++++ assert_ai/services/evaluations.py | 1330 ++++++++++++++++++++++ assert_ai/services/job_models.py | 154 +++ assert_ai/services/job_store.py | 836 ++++++++++++++ assert_ai/services/run_planning.py | 2 +- tests/test_evaluation_service.py | 357 ++++++ tests/test_job_store.py | 221 ++++ tests/test_mcp_cli.py | 9 + tests/test_mcp_server.py | 179 ++- tests/test_run_result.py | 42 +- tests/test_runtime_path_policy.py | 10 + tests/test_runtime_safety.py | 18 + 22 files changed, 3893 insertions(+), 31 deletions(-) create mode 100644 assert_ai/mcp/tools/jobs.py create mode 100644 assert_ai/services/_evaluation_worker.py create mode 100644 assert_ai/services/evaluations.py create mode 100644 assert_ai/services/job_models.py create mode 100644 assert_ai/services/job_store.py create mode 100644 tests/test_evaluation_service.py create mode 100644 tests/test_job_store.py diff --git a/assert_ai/core/runtime_path_policy.py b/assert_ai/core/runtime_path_policy.py index ca7ec891a..3f9cf190a 100644 --- a/assert_ai/core/runtime_path_policy.py +++ b/assert_ai/core/runtime_path_policy.py @@ -48,6 +48,8 @@ def __init__( def _is_within(path: Path, root: Path) -> bool: + path = _comparison_path(path) + root = _comparison_path(root) try: path.relative_to(root) return True @@ -55,6 +57,19 @@ def _is_within(path: Path, root: Path) -> bool: return False +def _comparison_path(path: Path) -> Path: + """Normalize equivalent Windows extended-length paths for comparison.""" + value = os.path.normpath(os.fspath(path)) + if os.name == "nt": + value = value.replace("/", "\\") + if value.startswith("\\\\?\\UNC\\"): + value = "\\\\" + value[8:] + elif value.startswith("\\\\?\\"): + value = value[4:] + value = os.path.normcase(value) + return Path(value) + + def _resolved(path: str | Path) -> Path: return Path(path).expanduser().resolve() @@ -400,9 +415,10 @@ def _require_no_links( *, field_name: str, ) -> None: - normalized = Path(os.path.abspath(path)) + normalized = _comparison_path(Path(os.path.abspath(path))) + comparison_root = _comparison_path(root) try: - relative = normalized.relative_to(root) + relative = normalized.relative_to(comparison_root) except ValueError: return current = root diff --git a/assert_ai/core/runtime_safety.py b/assert_ai/core/runtime_safety.py index fc9902b70..9b2a8ef41 100644 --- a/assert_ai/core/runtime_safety.py +++ b/assert_ai/core/runtime_safety.py @@ -153,9 +153,16 @@ def stop(self, *, write_final: bool = False) -> None: self._stop.set() self._thread.join(timeout=5.0) self._thread = None - if write_final: + if ( + write_final + and getattr(self._manifest, "status", None) == "running" + ): self._safe_write() + def write_now(self) -> bool: + """Serialize an immediate manifest write with heartbeat writes.""" + return self._safe_write() + def _loop(self) -> None: while not self._stop.wait(self._interval_s): wrote = self._safe_write() diff --git a/assert_ai/mcp/_command.py b/assert_ai/mcp/_command.py index 8e3043af5..53c2df2aa 100644 --- a/assert_ai/mcp/_command.py +++ b/assert_ai/mcp/_command.py @@ -123,6 +123,27 @@ def mcp() -> None: show_default=True, help="Maximum inference concurrency accepted by preflight.", ) +@click.option( + "--max-active-jobs", + type=click.IntRange(min=1), + default=1, + show_default=True, + help="Maximum number of evaluation workers active at once.", +) +@click.option( + "--max-queued-jobs", + type=click.IntRange(min=1), + default=100, + show_default=True, + help="Maximum number of evaluations waiting to start.", +) +@click.option( + "--max-job-log-bytes", + type=click.IntRange(min=4096, max=16 * 1024 * 1024), + default=1024 * 1024, + show_default=True, + help="Maximum retained bytes for each worker stdout/stderr log.", +) @click.option( "--max-prompt-sample-size", type=click.IntRange(min=1), @@ -168,6 +189,9 @@ def serve( max_artifact_chunk_bytes: int, max_config_bytes: int, max_concurrency: int, + max_active_jobs: int, + max_queued_jobs: int, + max_job_log_bytes: int, max_prompt_sample_size: int, max_scenario_sample_size: int, allowed_model_patterns: tuple[str, ...], @@ -199,6 +223,9 @@ def serve( max_artifact_chunk_bytes=max_artifact_chunk_bytes, max_config_bytes=max_config_bytes, max_concurrency=max_concurrency, + max_active_jobs=max_active_jobs, + max_queued_jobs=max_queued_jobs, + max_job_log_bytes=max_job_log_bytes, max_prompt_sample_size=max_prompt_sample_size, max_scenario_sample_size=max_scenario_sample_size, allowed_model_patterns=allowed_model_patterns, diff --git a/assert_ai/mcp/models.py b/assert_ai/mcp/models.py index b15024839..21c9f1b32 100644 --- a/assert_ai/mcp/models.py +++ b/assert_ai/mcp/models.py @@ -61,6 +61,9 @@ class ServerLimits(BaseModel): max_artifact_chunk_bytes: int max_config_bytes: int max_concurrency: int + max_active_jobs: int + max_queued_jobs: int + max_job_log_bytes: int max_prompt_sample_size: int max_scenario_sample_size: int model_allowlist_enabled: bool = False diff --git a/assert_ai/mcp/resources.py b/assert_ai/mcp/resources.py index ff6e9cefc..935615502 100644 --- a/assert_ai/mcp/resources.py +++ b/assert_ai/mcp/resources.py @@ -15,6 +15,7 @@ from assert_ai.mcp.errors import invoke_resource from assert_ai.mcp.sanitize import sanitize_for_mcp from assert_ai.mcp.tools.inspect import InspectServices +from assert_ai.mcp.tools.jobs import JobServices from assert_ai.services.errors import ServiceError, ServiceErrorCode _SCHEMA_URI = "assert://schema/eval-config" @@ -24,6 +25,7 @@ def register_inspect_resources( server: MCPServer, services: InspectServices, *, + job_services: JobServices, inline_artifact_bytes: int, ) -> None: """Register static and templated resources for the inspect group.""" @@ -78,6 +80,22 @@ def config(config_ref: str) -> str: workspace=workspace, ) + @server.resource( + "assert://job/{job_id}/log", + name="job-log", + title="ASSERT evaluation job log", + description="Bounded, filtered stdout and stderr tails for one worker.", + mime_type="text/plain", + ) + def job_log(job_id: str) -> str: + return invoke_resource( + lambda: job_services.evaluations.read_log( + job_id, + max_bytes=inline_artifact_bytes, + ), + workspace=workspace, + ) + @server.resource( "assert://suite/{suite_id}/taxonomy", name="suite-taxonomy", diff --git a/assert_ai/mcp/server.py b/assert_ai/mcp/server.py index 20d40f33a..51a4b9ca7 100644 --- a/assert_ai/mcp/server.py +++ b/assert_ai/mcp/server.py @@ -26,14 +26,22 @@ from assert_ai.mcp.tools import ( AuthorServices, InspectServices, + JobServices, ProbeServices, register_author_tools, register_design_tools, register_inspect_tools, + register_job_execute_tools, + register_job_inspect_tools, register_probe_tools, ) from assert_ai.services.artifacts import ArtifactRepository from assert_ai.services.configs import ConfigService +from assert_ai.services.evaluations import ( + EvaluationJobManager, + EvaluationService, +) +from assert_ai.services.job_store import JobStore from assert_ai.services.library import LibraryService from assert_ai.services.results import ResultRepository from assert_ai.services.run_planning import ( @@ -80,6 +88,9 @@ class ServerOptions: max_artifact_chunk_bytes: int = 256 * 1024 max_config_bytes: int = 256 * 1024 max_concurrency: int = 32 + max_active_jobs: int = 1 + max_queued_jobs: int = 100 + max_job_log_bytes: int = 1024 * 1024 max_prompt_sample_size: int = 100_000 max_scenario_sample_size: int = 100_000 allowed_model_patterns: tuple[str, ...] = () @@ -100,6 +111,14 @@ def __post_init__(self) -> None: raise ValueError("max_config_bytes must not exceed max_response_bytes") if self.max_concurrency < 1: raise ValueError("max_concurrency must be positive") + if self.max_active_jobs < 1: + raise ValueError("max_active_jobs must be positive") + if self.max_queued_jobs < 1: + raise ValueError("max_queued_jobs must be positive") + if not 4096 <= self.max_job_log_bytes <= 16 * 1024 * 1024: + raise ValueError( + "max_job_log_bytes must be between 4096 and 16777216" + ) if self.max_prompt_sample_size < 1: raise ValueError("max_prompt_sample_size must be positive") if self.max_scenario_sample_size < 1: @@ -141,6 +160,9 @@ def create( max_artifact_chunk_bytes: int = 256 * 1024, max_config_bytes: int = 256 * 1024, max_concurrency: int = 32, + max_active_jobs: int = 1, + max_queued_jobs: int = 100, + max_job_log_bytes: int = 1024 * 1024, max_prompt_sample_size: int = 100_000, max_scenario_sample_size: int = 100_000, allowed_model_patterns: Iterable[str] = (), @@ -166,6 +188,9 @@ def create( max_artifact_chunk_bytes=max_artifact_chunk_bytes, max_config_bytes=max_config_bytes, max_concurrency=max_concurrency, + max_active_jobs=max_active_jobs, + max_queued_jobs=max_queued_jobs, + max_job_log_bytes=max_job_log_bytes, max_prompt_sample_size=max_prompt_sample_size, max_scenario_sample_size=max_scenario_sample_size, allowed_model_patterns=tuple(allowed_model_patterns), @@ -204,6 +229,53 @@ def build_server(options: ServerOptions) -> MCPServer: default_page_size=options.default_page_size, max_page_size=options.max_page_size, ) + planning = RunPlanningService( + options.workspace, + configs, + policy=PreflightPolicy( + max_concurrency=options.max_concurrency, + max_prompt_sample_size=options.max_prompt_sample_size, + max_scenario_sample_size=options.max_scenario_sample_size, + allowed_model_patterns=options.allowed_model_patterns, + allowed_endpoint_hosts=options.allowed_endpoint_hosts, + ), + ) + jobs_db = options.path_policy.resolve_managed_output( + options.workspace.artifacts_root / "mcp" / "jobs.sqlite3", + field_name="evaluation job store", + expected_root=options.workspace.artifacts_root, + reject_links=True, + ) + job_store = JobStore( + jobs_db, + path_policy=options.path_policy, + expected_root=options.workspace.artifacts_root, + ) + execution_enabled = ( + CapabilityGroup.EXECUTE in options.capability_groups + ) + job_manager = EvaluationJobManager( + options.workspace, + job_store, + max_active_jobs=options.max_active_jobs, + max_log_bytes=options.max_job_log_bytes, + launch_enabled=execution_enabled, + ) + evaluations = EvaluationService( + options.workspace, + configs, + planning, + job_store, + job_manager, + default_page_size=options.default_page_size, + max_page_size=options.max_page_size, + max_queued_jobs=options.max_queued_jobs, + ) + job_services = JobServices( + workspace=options.workspace, + evaluations=evaluations, + max_response_bytes=options.max_response_bytes, + ) @server.tool( title="Get ASSERT server information", @@ -235,6 +307,9 @@ def get_server_info() -> ServerInfo: max_artifact_chunk_bytes=options.max_artifact_chunk_bytes, max_config_bytes=options.max_config_bytes, max_concurrency=options.max_concurrency, + max_active_jobs=options.max_active_jobs, + max_queued_jobs=options.max_queued_jobs, + max_job_log_bytes=options.max_job_log_bytes, max_prompt_sample_size=options.max_prompt_sample_size, max_scenario_sample_size=options.max_scenario_sample_size, model_allowlist_enabled=bool(options.allowed_model_patterns), @@ -276,26 +351,18 @@ def get_server_info() -> ServerInfo: max_response_bytes=options.max_response_bytes, ) register_inspect_tools(server, services) + register_job_inspect_tools(server, job_services) register_inspect_resources( server, services, + job_services=job_services, inline_artifact_bytes=options.max_artifact_chunk_bytes, ) author_services = AuthorServices( workspace=options.workspace, configs=configs, - planning=RunPlanningService( - options.workspace, - configs, - policy=PreflightPolicy( - max_concurrency=options.max_concurrency, - max_prompt_sample_size=options.max_prompt_sample_size, - max_scenario_sample_size=options.max_scenario_sample_size, - allowed_model_patterns=options.allowed_model_patterns, - allowed_endpoint_hosts=options.allowed_endpoint_hosts, - ), - ), + planning=planning, max_response_bytes=options.max_response_bytes, allowed_model_patterns=options.allowed_model_patterns, ) @@ -316,6 +383,9 @@ def get_server_info() -> ServerInfo: max_response_bytes=options.max_response_bytes, ), ) + if execution_enabled: + register_job_execute_tools(server, job_services) + job_manager.enqueue() return server diff --git a/assert_ai/mcp/tools/__init__.py b/assert_ai/mcp/tools/__init__.py index 704c145cf..22d2804c1 100644 --- a/assert_ai/mcp/tools/__init__.py +++ b/assert_ai/mcp/tools/__init__.py @@ -11,13 +11,21 @@ register_probe_tools, ) from assert_ai.mcp.tools.inspect import InspectServices, register_inspect_tools +from assert_ai.mcp.tools.jobs import ( + JobServices, + register_job_execute_tools, + register_job_inspect_tools, +) __all__ = [ "AuthorServices", "InspectServices", + "JobServices", "ProbeServices", "register_author_tools", "register_design_tools", "register_inspect_tools", + "register_job_execute_tools", + "register_job_inspect_tools", "register_probe_tools", ] diff --git a/assert_ai/mcp/tools/jobs.py b/assert_ai/mcp/tools/jobs.py new file mode 100644 index 000000000..4e3a83d62 --- /dev/null +++ b/assert_ai/mcp/tools/jobs.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Persisted evaluation job tools for the ASSERT MCP adapter.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from mcp.server import MCPServer +from mcp.types import ToolAnnotations + +from assert_ai.core.workspace import WorkspaceService +from assert_ai.mcp.errors import adapt_tool_errors, invoke_tool +from assert_ai.mcp.sanitize import sanitize_for_mcp +from assert_ai.services.evaluations import EvaluationService +from assert_ai.services.job_models import ( + JobDetail, + JobPage, + JobStartResult, + JobState, +) +from assert_ai.services.run_planning import EvaluationOverrides + +_READ_ONLY_ANNOTATIONS = ToolAnnotations( + read_only_hint=True, + destructive_hint=False, + idempotent_hint=True, + open_world_hint=False, +) +_START_ANNOTATIONS = ToolAnnotations( + read_only_hint=False, + destructive_hint=True, + idempotent_hint=True, + open_world_hint=True, +) + + +@dataclass(frozen=True, slots=True) +class JobServices: + """Application services and limits shared by evaluation job tools.""" + + workspace: WorkspaceService + evaluations: EvaluationService + max_response_bytes: int + + +def register_job_inspect_tools( + server: MCPServer, + services: JobServices, +) -> None: + """Register job discovery and polling for every inspect-capable mode.""" + + @server.tool( + title="List ASSERT evaluation jobs", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def list_jobs( + states: tuple[JobState, ...] = (), + cursor: str | None = None, + page_size: int | None = None, + ) -> JobPage: + """List persisted evaluation jobs with bounded pagination.""" + page = invoke_tool( + lambda: services.evaluations.list( + states=states, + cursor=cursor, + limit=page_size, + ), + workspace=services.workspace, + ) + return JobPage.model_validate( + sanitize_for_mcp(page, workspace=services.workspace) + ) + + @server.tool( + title="Get an ASSERT evaluation job", + annotations=_READ_ONLY_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def get_job(job_id: str) -> JobDetail: + """Get current state, progress, result, and resources for one job.""" + job = invoke_tool( + lambda: services.evaluations.get(job_id), + workspace=services.workspace, + ) + return JobDetail.model_validate( + sanitize_for_mcp(job, workspace=services.workspace) + ) + + +def register_job_execute_tools( + server: MCPServer, + services: JobServices, +) -> None: + """Register idempotent, non-blocking evaluation execution.""" + + @server.tool( + title="Start an ASSERT evaluation", + annotations=_START_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def start_evaluation( + config_ref: str, + request_id: str, + overrides: EvaluationOverrides | None = None, + ) -> JobStartResult: + """Validate, snapshot, and enqueue one evaluation without blocking.""" + started = invoke_tool( + lambda: services.evaluations.start( + config_ref, + request_id=request_id, + overrides=overrides, + ), + workspace=services.workspace, + ) + return JobStartResult.model_validate( + sanitize_for_mcp(started, workspace=services.workspace) + ) diff --git a/assert_ai/mcp/uris.py b/assert_ai/mcp/uris.py index 86b1844dc..09cf7d02c 100644 --- a/assert_ai/mcp/uris.py +++ b/assert_ai/mcp/uris.py @@ -16,6 +16,10 @@ def config_uri(config_ref: str) -> str: return f"assert://config/{quote(config_ref, safe='')}" +def job_log_uri(job_id: str) -> str: + return f"assert://job/{quote(job_id, safe='')}/log" + + def suite_taxonomy_uri(suite_id: str) -> str: return f"assert://suite/{quote(suite_id, safe='')}/taxonomy" diff --git a/assert_ai/runner.py b/assert_ai/runner.py index b64a32a1d..9eb46719e 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -13,6 +13,7 @@ import sys import time import warnings +from copy import deepcopy from datetime import datetime, timezone from pathlib import Path from typing import TYPE_CHECKING, Any @@ -38,7 +39,7 @@ update_latest, ) from assert_ai.core.config_model import RunManifest, SuiteMetadata -from assert_ai.core.io import write_json +from assert_ai.core.io import write_json, write_text_atomic from assert_ai.core.model_client import ( LLMAuthError, LLMInputError, @@ -54,6 +55,7 @@ ) from assert_ai.core.run_result import RunResult, RunState from assert_ai.core.run_plan import resolve_forced_stages +from assert_ai.core.yaml_io import dump_yaml from assert_ai.display import label_metric from assert_ai.services.result_metadata import ( refresh_stage_indexes, @@ -109,20 +111,34 @@ def _load_context( config: str, overrides: list[str] | None = None, path_policy: RuntimePathPolicy | None = None, + config_document: dict[str, Any] | None = None, ) -> dict[str, Any]: """Load one config file into runtime context.""" cfg_path = ( - path_policy.resolve_config_path(config, must_exist=True) + path_policy.resolve_config_path( + config, + must_exist=config_document is None, + ) if path_policy is not None else Path(config).resolve() ) - raw = _apply_config_overrides(load_config(cfg_path), overrides) - return load_runtime_context( + raw = _apply_config_overrides( + ( + deepcopy(config_document) + if config_document is not None + else load_config(cfg_path) + ), + overrides, + ) + ctx = load_runtime_context( raw, cfg_path, stage_modules=STAGES, path_policy=path_policy, ) + if config_document is not None: + ctx["_config_snapshot_document"] = raw + return ctx def _write_suite_metadata(ctx: dict[str, Any]) -> None: @@ -165,6 +181,18 @@ def _write_manifest(manifest: RunManifest, run_root: Path) -> None: write_json(manifest_path, manifest.to_dict()) +def _write_active_manifest( + manifest: RunManifest, + run_root: Path, + heartbeat: ManifestHeartbeat | None, +) -> None: + """Write without racing the heartbeat's atomic manifest replacement.""" + if heartbeat is not None: + heartbeat.write_now() + else: + _write_manifest(manifest, run_root) + + def _record_run_artifacts(manifest: RunManifest, ctx: dict[str, Any], run_root: Path) -> None: """Copy resolved artifact references into the run manifest and sidecar.""" @@ -712,6 +740,35 @@ def run_pipeline_result( ) +def run_pipeline_document_result( + *, + document: dict[str, Any], + config_path: str, + force_stages: list[str] | None = None, + strict: bool = False, + concurrency: int | None = None, + path_policy: RuntimePathPolicy | None = None, +) -> RunResult: + """Execute an immutable config document using its original path as a base.""" + try: + return _run_pipeline_result( + config=config_path, + force_stages=force_stages, + strict=strict, + concurrency=concurrency, + path_policy=path_policy, + config_document=document, + ) + except Exception: # noqa: BLE001 + log.error("[runner] Unexpected pipeline setup error", exc_info=True) + return RunResult( + state=RunState.FAILED, + exit_code=1, + error_code="INTERNAL", + error_message="Unexpected pipeline setup error", + ) + + def _run_pipeline_result( *, config: str, @@ -720,6 +777,7 @@ def _run_pipeline_result( overrides: list[str] | None = None, concurrency: int | None = None, path_policy: RuntimePathPolicy | None = None, + config_document: dict[str, Any] | None = None, ) -> RunResult: """Execute configured stages. @@ -753,6 +811,7 @@ def _run_pipeline_result( config=config, overrides=overrides, path_policy=path_policy, + config_document=config_document, ) ctx["strict"] = strict except (ConfigError, ValueError) as exc: @@ -881,12 +940,14 @@ def _run_pipeline_result( run_root.mkdir(parents=True, exist_ok=True) manifest = _build_manifest(ctx) config_path = ctx.get("config_path") - if config_path is not None and Path(config_path).is_file(): - config_path = ( - path_policy.resolve_config_path(config_path, must_exist=True) - if path_policy is not None - else Path(config_path) + snapshot_document = ctx.get("_config_snapshot_document") + if ( + isinstance(snapshot_document, dict) + or ( + config_path is not None + and Path(config_path).is_file() ) + ): config_snapshot = ( path_policy.resolve_managed_output( run_root / "config.yaml", @@ -897,7 +958,21 @@ def _run_pipeline_result( if path_policy is not None else run_root / "config.yaml" ) - shutil.copy2(config_path, config_snapshot) + if isinstance(snapshot_document, dict): + write_text_atomic( + config_snapshot, + dump_yaml(snapshot_document), + ) + else: + config_path = ( + path_policy.resolve_config_path( + config_path, + must_exist=True, + ) + if path_policy is not None + else Path(config_path) + ) + shutil.copy2(config_path, config_snapshot) _record_run_artifacts(manifest, ctx, run_root) _write_manifest(manifest, run_root) _refresh_run_summary( @@ -983,7 +1058,7 @@ def _run_stages_inner( "started_at": datetime.now(timezone.utc).isoformat(), } _record_run_artifacts(manifest, ctx, run_root) - _write_manifest(manifest, run_root) + _write_active_manifest(manifest, run_root, heartbeat) _refresh_run_summary( ctx, manifest, @@ -1120,7 +1195,7 @@ def _run_stages_inner( existing_timing["duration_secs"] = round(elapsed, 3) manifest.stage_timings[stage_name] = existing_timing _record_run_artifacts(manifest, ctx, run_root) - _write_manifest(manifest, run_root) + _write_active_manifest(manifest, run_root, heartbeat) _refresh_run_summary( ctx, manifest, @@ -1186,7 +1261,7 @@ def _run_stages_inner( manifest.ended_at = datetime.now(timezone.utc).isoformat() manifest.status = "completed" if failed_stage is None else "failed" _record_run_artifacts(manifest, ctx, run_root) - _write_manifest(manifest, run_root) + _write_active_manifest(manifest, run_root, heartbeat) _refresh_run_summary( ctx, manifest, diff --git a/assert_ai/services/_evaluation_worker.py b/assert_ai/services/_evaluation_worker.py new file mode 100644 index 000000000..d3c6fd911 --- /dev/null +++ b/assert_ai/services/_evaluation_worker.py @@ -0,0 +1,352 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Internal subprocess entry point for one persisted evaluation job.""" + +from __future__ import annotations + +import argparse +import contextlib +import hashlib +import io +import json +import logging +import re +import sys +import threading +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import yaml + +from assert_ai.core.io import write_json +from assert_ai.core.security import redact_path_prefixes, sanitize_text +from assert_ai.core.workspace import WorkspaceService + +_JOB_ID_RE = re.compile(r"^[0-9a-f]{32}$") +_SHA256_RE = re.compile(r"^sha256:[0-9a-f]{64}$") +_MAX_REQUEST_BYTES = 1024 * 1024 +_MAX_SNAPSHOT_BYTES = 16 * 1024 * 1024 +_MIN_LOG_BYTES = 4096 +_MAX_LOG_BYTES = 16 * 1024 * 1024 +_DEFAULT_LOG_BYTES = 1024 * 1024 +_TRUNCATION_MARKER = b"[earlier worker output truncated]\n" + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--workspace", required=True) + parser.add_argument("--job-id", required=True) + args = parser.parse_args(argv) + + workspace: WorkspaceService | None = None + result_path: Path | None = None + result_token: str | None = None + exit_code = 1 + try: + if not _JOB_ID_RE.fullmatch(args.job_id): + raise ValueError("Invalid evaluation job id") + workspace = WorkspaceService.create(args.workspace) + jobs_root = _jobs_root(workspace) + job_dir = workspace.path_policy.resolve_managed_output( + jobs_root / args.job_id, + field_name="evaluation job directory", + expected_root=jobs_root, + reject_links=True, + ) + request_path = workspace.path_policy.resolve_managed_output( + job_dir / "request.json", + field_name="evaluation job request", + expected_root=job_dir, + reject_links=True, + ) + snapshot_path = workspace.path_policy.resolve_managed_output( + job_dir / "config.yaml", + field_name="evaluation config snapshot", + expected_root=job_dir, + reject_links=True, + ) + result_path = workspace.path_policy.resolve_managed_output( + job_dir / "result.json", + field_name="evaluation job result", + expected_root=job_dir, + reject_links=True, + ) + request = _read_request(request_path) + if request.get("job_id") != args.job_id: + raise ValueError("Evaluation job request identity mismatch") + result_token = _required_string(request, "result_token") + config_ref = _required_string(request, "config_ref") + expected_snapshot_hash = _required_string( + request, + "config_sha256", + ) + if not _SHA256_RE.fullmatch(expected_snapshot_hash): + raise ValueError("config_sha256 must be a SHA-256 digest") + snapshot_bytes = _read_bytes( + snapshot_path, + max_bytes=_MAX_SNAPSHOT_BYTES, + label="Evaluation config snapshot", + ) + actual_snapshot_hash = ( + "sha256:" + hashlib.sha256(snapshot_bytes).hexdigest() + ) + if actual_snapshot_hash != expected_snapshot_hash: + raise ValueError("Evaluation config snapshot digest mismatch") + + document = yaml.safe_load(snapshot_bytes.decode("utf-8")) + if not isinstance(document, dict): + raise ValueError("Evaluation config snapshot must be a mapping") + config_path = workspace.path_policy.resolve_config_path( + config_ref, + reject_links=True, + ) + force_stages = request.get("force_stages") + if not isinstance(force_stages, list) or not all( + isinstance(item, str) for item in force_stages + ): + raise ValueError("force_stages must be a string array") + strict = request.get("strict") + if not isinstance(strict, bool): + raise ValueError("strict must be a boolean") + max_log_bytes = _log_limit(request.get("max_log_bytes")) + stdout_path = workspace.path_policy.resolve_managed_output( + job_dir / "stdout.log", + field_name="evaluation worker stdout", + expected_root=job_dir, + reject_links=True, + ) + stderr_path = workspace.path_policy.resolve_managed_output( + job_dir / "stderr.log", + field_name="evaluation worker stderr", + expected_root=job_dir, + reject_links=True, + ) + with ( + _BoundedTextLog(stdout_path, max_bytes=max_log_bytes) as stdout, + _BoundedTextLog(stderr_path, max_bytes=max_log_bytes) as stderr, + contextlib.redirect_stdout(stdout), + contextlib.redirect_stderr(stderr), + _capture_worker_logs(stderr), + ): + from assert_ai.runner import run_pipeline_document_result + + result = run_pipeline_document_result( + document=document, + config_path=str(config_path), + force_stages=force_stages, + strict=strict, + path_policy=workspace.path_policy, + ) + payload = { + "schema_version": 1, + "job_id": args.job_id, + "result_token": result_token, + "run_result": result.to_dict(), + } + exit_code = result.exit_code + except Exception as exc: # noqa: BLE001 - subprocess boundary + message = sanitize_text(str(exc)) or "Evaluation worker failed" + if workspace is not None: + message = redact_path_prefixes( + message, + ( + workspace.root, + workspace.configs_root, + workspace.artifacts_root, + workspace.results_root, + ), + ) + payload = { + "schema_version": 1, + "job_id": str(args.job_id), + "worker_error": { + "error_code": "INTERNAL", + "error_message": message, + }, + } + if result_token is not None: + payload["result_token"] = result_token + + if result_path is not None: + write_json(result_path, payload) + else: + sys.stderr.write("Evaluation worker could not resolve its result path\n") + return exit_code + + +def _jobs_root(workspace: WorkspaceService) -> Path: + root = workspace.artifacts_root / "mcp" / "jobs" + return workspace.path_policy.resolve_managed_output( + root, + field_name="evaluation jobs root", + expected_root=workspace.artifacts_root, + reject_links=True, + ) + + +def _read_request(path: Path) -> dict[str, Any]: + payload = json.loads( + _read_bytes( + path, + max_bytes=_MAX_REQUEST_BYTES, + label="Evaluation job request", + ).decode("utf-8") + ) + if not isinstance(payload, dict): + raise ValueError("Evaluation job request must be an object") + if payload.get("schema_version") != 1: + raise ValueError("Unsupported evaluation job request schema") + return payload + + +def _required_string(payload: dict[str, Any], key: str) -> str: + value = payload.get(key) + if not isinstance(value, str) or not value: + raise ValueError(f"{key} must be a non-empty string") + return value + + +def _read_bytes(path: Path, *, max_bytes: int, label: str) -> bytes: + with path.open("rb") as stream: + value = stream.read(max_bytes + 1) + if len(value) > max_bytes: + raise ValueError(f"{label} exceeds the worker limit") + return value + + +def _log_limit(value: Any) -> int: + if value is None: + return _DEFAULT_LOG_BYTES + if ( + not isinstance(value, int) + or isinstance(value, bool) + or not _MIN_LOG_BYTES <= value <= _MAX_LOG_BYTES + ): + raise ValueError( + f"max_log_bytes must be between {_MIN_LOG_BYTES} and " + f"{_MAX_LOG_BYTES}" + ) + return value + + +class _BoundedTextLog(io.TextIOBase): + """UTF-8 log sink that retains a bounded tail without using pipes.""" + + def __init__(self, path: Path, *, max_bytes: int) -> None: + self._path = path + self._max_bytes = max_bytes + self._lock = threading.Lock() + self._stream = path.open("w+b", buffering=0) + self._buffer = _BoundedBinaryLog(self) + + @property + def encoding(self) -> str: + return "utf-8" + + @property + def buffer(self) -> "_BoundedBinaryLog": + return self._buffer + + def writable(self) -> bool: + return True + + def isatty(self) -> bool: + return False + + def fileno(self) -> int: + return self._stream.fileno() + + def write(self, value: str) -> int: + if self.closed: + raise ValueError("I/O operation on closed worker log") + if not isinstance(value, str): + value = str(value) + value = sanitize_text(value) + encoded = value.encode("utf-8", errors="replace") + self._write_bytes(encoded) + return len(value) + + def _write_bytes(self, value: bytes) -> None: + with self._lock: + self._stream.seek(0, 2) + self._stream.write(value) + if self._stream.tell() > self._max_bytes: + self._truncate_to_tail() + self._stream.flush() + + def flush(self) -> None: + if self.closed: + return + with self._lock: + self._stream.flush() + + def close(self) -> None: + if self.closed: + return + with self._lock: + self._stream.seek(0, 2) + if self._stream.tell() > self._max_bytes: + self._truncate_to_tail() + super().close() + with self._lock: + self._stream.close() + + def _truncate_to_tail(self) -> None: + keep_bytes = max( + 1, + self._max_bytes // 2 - len(_TRUNCATION_MARKER), + ) + self._stream.seek(-keep_bytes, 2) + tail = self._stream.read() + tail = tail.decode("utf-8", errors="replace").encode("utf-8") + retained = (_TRUNCATION_MARKER + tail)[-self._max_bytes :] + self._stream.seek(0) + self._stream.write(retained) + self._stream.truncate() + + +class _BoundedBinaryLog: + """Binary facade used by code that writes to ``sys.stdout.buffer``.""" + + def __init__(self, text_log: _BoundedTextLog) -> None: + self._text_log = text_log + + def write(self, value: bytes | bytearray) -> int: + encoded = bytes(value) + sanitized = sanitize_text( + encoded.decode("utf-8", errors="replace") + ).encode("utf-8") + self._text_log._write_bytes(sanitized) + return len(encoded) + + def flush(self) -> None: + self._text_log.flush() + + def fileno(self) -> int: + return self._text_log.fileno() + + +@contextlib.contextmanager +def _capture_worker_logs( + stream: _BoundedTextLog, +) -> Iterator[None]: + root = logging.getLogger() + previous_level = root.level + handler = logging.StreamHandler(stream) + handler.setFormatter(logging.Formatter("%(message)s")) + handler.setLevel(logging.INFO) + root.addHandler(handler) + if previous_level > logging.INFO: + root.setLevel(logging.INFO) + try: + yield + finally: + root.removeHandler(handler) + root.setLevel(previous_level) + handler.close() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/assert_ai/services/evaluations.py b/assert_ai/services/evaluations.py new file mode 100644 index 000000000..ce9e43f1c --- /dev/null +++ b/assert_ai/services/evaluations.py @@ -0,0 +1,1330 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Persisted, idempotent evaluation execution services.""" + +from __future__ import annotations + +import base64 +import binascii +import hashlib +import json +import logging +import os +import secrets +import shutil +import subprocess +import sys +import threading +import time +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Sequence +from urllib.parse import quote + +from assert_ai.core.io import write_json, write_text_atomic +from assert_ai.core.security import ( + redact_path_prefixes, + sanitize_payload, + sanitize_text, +) +from assert_ai.core.workspace import WorkspaceService +from assert_ai.core.yaml_io import dump_yaml +from assert_ai.services.configs import ConfigService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.job_models import ( + JobCatalogEntry, + JobDetail, + JobPage, + JobRecord, + JobStartResult, + JobState, + JobTerminalResult, + NewJob, + TERMINAL_JOB_STATES, +) +from assert_ai.services.job_store import JobStore +from assert_ai.services.run_planning import ( + EvaluationOverrides, + RunPlanningService, + StageAction, +) + +_CURSOR_VERSION = 1 +_JOB_RESULT_MAX_BYTES = 1024 * 1024 +_JOB_ID_RETRIES = 5 +_LEASE_SECONDS = 60.0 +_LEASE_RENEW_SECONDS = 15.0 +_REQUEST_ID_MAX_LENGTH = 200 +_MIN_LOG_BYTES = 4096 +_MAX_LOG_BYTES = 16 * 1024 * 1024 + +log = logging.getLogger(__name__) + + +@dataclass(slots=True) +class EvaluationJobManager: + """Launch queued jobs and reconcile their terminal worker results.""" + + workspace: WorkspaceService + store: JobStore + max_active_jobs: int = 1 + max_log_bytes: int = 1024 * 1024 + launch_enabled: bool = True + lease_seconds: float = _LEASE_SECONDS + _owner: str = field( + default_factory=lambda: uuid.uuid4().hex, + init=False, + ) + _lock: threading.Lock = field( + default_factory=threading.Lock, + init=False, + repr=False, + ) + _scheduler: threading.Thread | None = field( + default=None, + init=False, + repr=False, + ) + _schedule_requested: bool = field( + default=False, + init=False, + repr=False, + ) + _processes: dict[str, subprocess.Popen[bytes]] = field( + default_factory=dict, + init=False, + repr=False, + ) + + def __post_init__(self) -> None: + if self.max_active_jobs < 1: + raise ValueError("max_active_jobs must be positive") + if not _MIN_LOG_BYTES <= self.max_log_bytes <= _MAX_LOG_BYTES: + raise ValueError( + f"max_log_bytes must be between {_MIN_LOG_BYTES} and " + f"{_MAX_LOG_BYTES}" + ) + if self.lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + + def enqueue(self) -> None: + """Wake a short-lived scheduler without holding an MCP request open.""" + if not self.launch_enabled: + return + with self._lock: + self._schedule_requested = True + if self._scheduler is not None and self._scheduler.is_alive(): + return + self._scheduler = threading.Thread( + target=self._schedule, + name="assert-mcp-job-scheduler", + daemon=True, + ) + self._scheduler.start() + + def reconcile(self, record: JobRecord) -> JobRecord: + """Adopt a worker result or mark a dead worker interrupted.""" + if record.state in TERMINAL_JOB_STATES: + return record + if record.state is JobState.QUEUED: + self.enqueue() + return record + result = self._read_result(record) + if result is not None: + try: + return self._adopt_result( + record, + result, + lease_owner=None, + ) + except Exception as exc: # noqa: BLE001 - persisted boundary + log.exception( + "Could not reconcile evaluation job %s", + record.job_id, + ) + return self._mark_internal_failure( + record, + exc, + lease_owner=None, + ) + if record.state is JobState.STARTING and record.pid is None: + if not _lease_expired(record.lease_expires_at): + return record + if ( + record.pid is not None + and record.process_create_time is not None + and _process_matches( + record.pid, + record.process_create_time, + ) + ): + return record + return self.store.mark_terminal( + record.job_id, + state=JobState.INTERRUPTED, + exit_code=record.exit_code, + failed_stage=record.failed_stage, + error_code=ServiceErrorCode.JOB_INTERRUPTED.value, + error_message=( + "Evaluation worker exited without a terminal result" + ), + result=None, + run_root=record.run_root, + ) + + def _schedule(self) -> None: + try: + while True: + with self._lock: + self._schedule_requested = False + try: + claimed = self.store.claim_next( + lease_owner=self._owner, + lease_seconds=self.lease_seconds, + max_active_jobs=self.max_active_jobs, + ) + except Exception: # noqa: BLE001 - daemon boundary + log.exception( + "Evaluation scheduler could not claim a queued job" + ) + return + if claimed is None: + with self._lock: + if self._schedule_requested: + continue + if self._scheduler is threading.current_thread(): + self._scheduler = None + return + try: + process = self._launch(claimed) + except Exception as exc: # noqa: BLE001 - process boundary + log.exception( + "Could not launch evaluation job %s", + claimed.job_id, + ) + self._mark_internal_failure( + claimed, + exc, + lease_owner=self._owner, + fallback="Evaluation worker could not be started", + ) + continue + with self._lock: + self._processes[claimed.job_id] = process + monitor = threading.Thread( + target=self._monitor, + args=(claimed.job_id, process), + name=f"assert-mcp-job-{claimed.job_id[:8]}", + daemon=True, + ) + monitor.start() + finally: + restart = False + with self._lock: + if self._scheduler is threading.current_thread(): + restart = ( + self.launch_enabled and self._schedule_requested + ) + self._scheduler = None + if restart: + self.enqueue() + + def _launch(self, record: JobRecord) -> subprocess.Popen[bytes]: + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" + env["PYTHONDONTWRITEBYTECODE"] = "1" + command = [ + sys.executable, + "-m", + "assert_ai.services._evaluation_worker", + "--workspace", + str(self.workspace.root), + "--job-id", + record.job_id, + ] + creationflags = ( + subprocess.CREATE_NEW_PROCESS_GROUP + if os.name == "nt" + else 0 + ) + process: subprocess.Popen[bytes] | None = None + try: + process = subprocess.Popen( + command, + cwd=str(self.workspace.root), + env=env, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + close_fds=True, + creationflags=creationflags, + start_new_session=os.name != "nt", + ) + create_time = _process_create_time(process.pid) + self.store.mark_running( + record.job_id, + lease_owner=self._owner, + pid=process.pid, + process_create_time=create_time, + lease_seconds=self.lease_seconds, + ) + except BaseException: + if process is not None: + _terminate_failed_launch(process) + raise + return process + + def _monitor( + self, + job_id: str, + process: subprocess.Popen[bytes], + ) -> None: + lost_lease = False + while True: + try: + process.wait(timeout=_LEASE_RENEW_SECONDS) + break + except subprocess.TimeoutExpired: + try: + renewed = self.store.renew_lease( + job_id, + lease_owner=self._owner, + lease_seconds=self.lease_seconds, + ) + except Exception: # noqa: BLE001 - daemon boundary + log.exception( + "Could not renew lease for evaluation job %s", + job_id, + ) + continue + if not renewed: + lost_lease = True + break + try: + if lost_lease: + log.warning( + "Stopped monitoring evaluation job %s after losing its lease", + job_id, + ) + return + record = self.store.get(job_id) + payload = self._read_result(record) + if payload is None: + self.store.mark_terminal( + job_id, + state=JobState.FAILED, + exit_code=process.returncode, + failed_stage=None, + error_code=ServiceErrorCode.RUN_FAILED.value, + error_message=( + "Evaluation worker exited without a valid result" + ), + result=None, + run_root=None, + lease_owner=self._owner, + ) + else: + self._adopt_result( + record, + payload, + lease_owner=self._owner, + ) + except Exception as exc: # noqa: BLE001 - daemon boundary + log.exception( + "Could not adopt terminal result for evaluation job %s", + job_id, + ) + try: + record = self.store.get(job_id) + self._mark_internal_failure( + record, + exc, + lease_owner=self._owner, + ) + except Exception: + log.exception( + "Could not persist failure for evaluation job %s", + job_id, + ) + finally: + with self._lock: + self._processes.pop(job_id, None) + self.enqueue() + + def _mark_internal_failure( + self, + record: JobRecord, + error: Exception, + *, + lease_owner: str | None, + fallback: str = "Evaluation result reconciliation failed", + ) -> JobRecord: + return self.store.mark_terminal( + record.job_id, + state=JobState.FAILED, + exit_code=1, + failed_stage=None, + error_code=ServiceErrorCode.INTERNAL.value, + error_message=_safe_error( + error, + workspace=self.workspace, + fallback=fallback, + ), + result=None, + run_root=record.run_root, + lease_owner=lease_owner, + ) + + def _read_result(self, record: JobRecord) -> dict[str, Any] | None: + result_path = self._job_file( + self._job_dir(record.job_id), + "result.json", + ) + if not result_path.is_file(): + return None + try: + if result_path.stat().st_size > _JOB_RESULT_MAX_BYTES: + return None + payload = json.loads(result_path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + if not isinstance(payload, dict): + return None + if payload.get("schema_version") != 1: + return None + if payload.get("job_id") != record.job_id: + return None + return payload + + def _adopt_result( + self, + record: JobRecord, + payload: dict[str, Any], + *, + lease_owner: str | None, + ) -> JobRecord: + job_dir = self._job_dir(record.job_id) + request = _read_json_file( + self._job_file(job_dir, "request.json"), + max_bytes=_JOB_RESULT_MAX_BYTES, + ) + if ( + not isinstance(request, dict) + or payload.get("result_token") != request.get("result_token") + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker result identity mismatch", + ) + raw_result = payload.get("run_result") + if not isinstance(raw_result, dict): + worker_error = payload.get("worker_error") + error_code = ServiceErrorCode.INTERNAL.value + message = None + if isinstance(worker_error, dict): + message = worker_error.get("error_message") + if worker_error.get("error_code") == "INTERNAL": + error_code = ServiceErrorCode.INTERNAL.value + return self.store.mark_terminal( + record.job_id, + state=JobState.FAILED, + exit_code=1, + failed_stage=None, + error_code=error_code, + error_message=_safe_text( + message, + workspace=self.workspace, + fallback="Evaluation worker failed", + ), + result=None, + run_root=None, + lease_owner=lease_owner, + ) + state_value = raw_result.get("state") + state_map = { + "completed": JobState.COMPLETED, + "failed": JobState.FAILED, + "cancelled": JobState.CANCELLED, + } + state = state_map.get(state_value) + if state is None: + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker returned an invalid state", + ) + result_suite_id = raw_result.get("suite_id") + result_run_id = raw_result.get("run_id") + failed_stage = _validated_optional_text( + raw_result.get("failed_stage"), + field_name="failed_stage", + ) + if ( + result_suite_id is not None + and not isinstance(result_suite_id, str) + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker returned an invalid suite id", + ) + if ( + result_run_id is not None + and not isinstance(result_run_id, str) + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker returned an invalid run id", + ) + if ( + result_suite_id is not None + and result_suite_id != record.suite_id + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker returned a mismatched suite id", + ) + if ( + result_run_id is not None + and result_run_id != record.run_id + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker returned a mismatched run id", + ) + may_omit_identity = ( + state is JobState.FAILED + and failed_stage is None + and result_suite_id is None + and result_run_id is None + ) + if not may_omit_identity and ( + result_suite_id != record.suite_id + or result_run_id != record.run_id + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker omitted its allocated run identity", + ) + exit_code = raw_result.get("exit_code") + if ( + not isinstance(exit_code, int) + or isinstance(exit_code, bool) + or exit_code < 0 + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker returned an invalid exit code", + ) + if ( + state is JobState.COMPLETED + and exit_code != 0 + ) or ( + state is not JobState.COMPLETED + and exit_code == 0 + ): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker returned an inconsistent exit code", + ) + run_root = ( + self._validated_run_root(record, raw_result) + if result_suite_id is not None + else None + ) + public_result = { + "state": state_value, + "exit_code": exit_code, + "failed_stage": failed_stage, + "error_code": _validated_optional_text( + raw_result.get("error_code"), + field_name="error_code", + ), + "error_message": _safe_text( + _validated_optional_text( + raw_result.get("error_message"), + field_name="error_message", + ), + workspace=self.workspace, + fallback=None, + ), + } + public_result = sanitize_payload(public_result) + return self.store.mark_terminal( + record.job_id, + state=state, + exit_code=public_result["exit_code"], + failed_stage=public_result["failed_stage"], + error_code=public_result["error_code"], + error_message=public_result["error_message"], + result=public_result, + run_root=str(run_root) if run_root is not None else None, + lease_owner=lease_owner, + ) + + def _validated_run_root( + self, + record: JobRecord, + result: dict[str, Any], + ) -> Path | None: + if record.run_id is None: + if result.get("run_root") is not None: + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Suite-only worker returned an unexpected run root", + ) + return None + expected_suite = self.workspace.path_policy.resolve_managed_output( + self.workspace.results_root / record.suite_id, + field_name="job suite root", + expected_root=self.workspace.results_root, + reject_links=True, + ) + expected_run = self.workspace.path_policy.resolve_managed_output( + expected_suite / record.run_id, + field_name="job run root", + expected_root=expected_suite, + reject_links=True, + ) + actual = result.get("run_root") + if actual is None or Path(str(actual)).resolve() != expected_run: + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + "Evaluation worker returned a mismatched run root", + ) + return expected_run + + def _job_dir(self, job_id: str) -> Path: + jobs_root = _jobs_root(self.workspace) + return self.workspace.path_policy.resolve_managed_output( + jobs_root / job_id, + field_name="evaluation job directory", + expected_root=jobs_root, + reject_links=True, + ) + + def _job_file(self, job_dir: Path, name: str) -> Path: + return self.workspace.path_policy.resolve_managed_output( + job_dir / name, + field_name=f"evaluation job {name}", + expected_root=job_dir, + reject_links=True, + ) + + +@dataclass(slots=True) +class EvaluationService: + """Author, persist, execute, and inspect evaluation jobs.""" + + workspace: WorkspaceService + configs: ConfigService + planning: RunPlanningService + store: JobStore + manager: EvaluationJobManager + default_page_size: int = 50 + max_page_size: int = 200 + max_queued_jobs: int = 100 + + def start( + self, + config_ref: str, + *, + request_id: str, + overrides: EvaluationOverrides | None = None, + ) -> JobStartResult: + if not self.manager.launch_enabled: + raise ServiceError( + ServiceErrorCode.CAPABILITY_DISABLED, + "Evaluation execution is disabled for this service", + ) + request_id = _validate_request_id(request_id) + applied = overrides or EvaluationOverrides() + config = self.configs.get_config(config_ref) + request_hash = _request_hash( + config_ref=config.config_ref, + config_etag=config.etag, + overrides=applied, + ) + existing = self.store.get_by_idempotency_key(request_id) + if existing is not None: + if existing.request_hash != request_hash: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "request_id is already bound to a different evaluation request", + details={"job_id": existing.job_id}, + ) + self.manager.enqueue() + return JobStartResult( + job=self.get(existing.job_id), + created=False, + ) + + initial = self.planning.preflight( + config.config_ref, + overrides=applied, + ) + _require_source_etag(initial.source_etag, config.etag) + _require_ready(initial) + suite_id = ( + applied.suite + or config.document.get("suite") + or _new_identity("mcp-suite") + ) + has_run_stage = any( + stage.scope == "run" + and stage.action is not StageAction.DISABLED + for stage in initial.stages + ) + run_id = ( + applied.run + or config.document.get("run") + or (_new_identity("run") if has_run_stage else None) + ) + effective_overrides = applied.model_copy( + update={"suite": suite_id, "run": run_id}, + ) + plan = self.planning.preflight( + config.config_ref, + overrides=effective_overrides, + ) + _require_source_etag(plan.source_etag, config.etag) + _require_ready(plan) + if plan.suite_id != suite_id or plan.run_id != run_id: + raise ServiceError( + ServiceErrorCode.INTERNAL, + "Preflight did not preserve allocated evaluation identity", + ) + if run_id is not None: + self._reject_existing_run(suite_id, run_id) + + yaml_text = dump_yaml(plan.effective_document) + config_sha256 = ( + "sha256:" + + hashlib.sha256(yaml_text.encode("utf-8")).hexdigest() + ) + new_job, job_dir = self._prepare_job( + config_ref=config.config_ref, + request_id=request_id, + request_hash=request_hash, + config_sha256=config_sha256, + suite_id=suite_id, + run_id=run_id, + plan=plan, + yaml_text=yaml_text, + ) + try: + created = self.store.create_or_get( + new_job, + max_queued_jobs=self.max_queued_jobs, + ) + except BaseException: + _remove_job_dir(job_dir) + raise + if not created.created: + _remove_job_dir(job_dir) + self.manager.enqueue() + return JobStartResult( + job=self.get(created.record.job_id), + created=created.created, + ) + + def get(self, job_id: str) -> JobDetail: + record = self.manager.reconcile(self.store.get(_validate_job_id(job_id))) + return self._detail(record) + + def list( + self, + *, + states: Sequence[JobState] = (), + cursor: str | None = None, + limit: int | None = None, + ) -> JobPage: + page_size = self.default_page_size if limit is None else limit + if ( + isinstance(page_size, bool) + or not isinstance(page_size, int) + or page_size < 1 + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "limit must be a positive integer", + ) + if page_size > self.max_page_size: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"limit must be <= {self.max_page_size}", + ) + before = _decode_cursor(cursor) if cursor else None + records = self.store.list_records( + limit=page_size + 1, + states=states, + before=before, + ) + visible = tuple( + ( + self.manager.reconcile(record) + if record.state not in TERMINAL_JOB_STATES + else record + ) + for record in records[:page_size] + ) + next_cursor = ( + _encode_cursor( + visible[-1].created_at, + visible[-1].job_id, + ) + if len(records) > page_size and visible + else None + ) + return JobPage( + items=tuple(_catalog_entry(record) for record in visible), + next_cursor=next_cursor, + ) + + def read_log(self, job_id: str, *, max_bytes: int) -> str: + """Return bounded, credential-filtered tails from one worker.""" + if ( + isinstance(max_bytes, bool) + or not isinstance(max_bytes, int) + or max_bytes < 4 + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "max_bytes must be an integer of at least 4", + ) + record = self.store.get(_validate_job_id(job_id)) + job_dir = self.manager._job_dir(record.job_id) + per_stream = max(128, (max_bytes - 128) // 2) + sections = [] + for label, name in ( + ("stdout", "stdout.log"), + ("stderr", "stderr.log"), + ): + path = self.manager._job_file(job_dir, name) + text = _read_text_tail(path, max_bytes=per_stream) + if text: + sections.append(f"--- {label} (filtered tail) ---\n{text}") + combined = ( + "\n".join(sections) + if sections + else "No worker output is available." + ) + combined = sanitize_text(combined) + combined = redact_path_prefixes( + combined, + ( + self.workspace.root, + self.workspace.configs_root, + self.workspace.artifacts_root, + self.workspace.results_root, + ), + ) + encoded = combined.encode("utf-8") + if len(encoded) > max_bytes: + combined = encoded[-max_bytes:].decode( + "utf-8", + errors="ignore", + ) + return combined + + def _prepare_job( + self, + *, + config_ref: str, + request_id: str, + request_hash: str, + config_sha256: str, + suite_id: str, + run_id: str | None, + plan: Any, + yaml_text: str, + ) -> tuple[NewJob, Path]: + jobs_root = _jobs_root(self.workspace) + jobs_root.mkdir(parents=True, exist_ok=True) + jobs_root = _jobs_root(self.workspace) + for _ in range(_JOB_ID_RETRIES): + job_id = uuid.uuid4().hex + job_dir = self.workspace.path_policy.resolve_managed_output( + jobs_root / job_id, + field_name="evaluation job directory", + expected_root=jobs_root, + reject_links=True, + ) + try: + job_dir.mkdir() + except FileExistsError: + continue + break + else: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "Could not allocate a unique evaluation job id", + ) + try: + snapshot = job_dir / "config.yaml" + request_path = job_dir / "request.json" + result_token = secrets.token_hex(32) + write_text_atomic(snapshot, yaml_text) + force_stages = [ + stage.name for stage in plan.stages if stage.forced + ] + write_json( + request_path, + { + "schema_version": 1, + "job_id": job_id, + "result_token": result_token, + "config_ref": config_ref, + "config_sha256": config_sha256, + "strict": bool(plan.strict), + "force_stages": force_stages, + "max_log_bytes": self.manager.max_log_bytes, + }, + ) + resource_keys = [] + # Reuse must not race a concurrent job that replaces active suite + # artifacts after this job's preflight selected them. + if any( + stage.scope == "suite" + and stage.action is not StageAction.DISABLED + for stage in plan.stages + ): + resource_keys.append(f"suite:{suite_id}") + if run_id is not None: + resource_keys.append(f"run:{suite_id}/{run_id}") + return ( + NewJob( + job_id=job_id, + idempotency_key=request_id, + request_hash=request_hash, + suite_id=suite_id, + run_id=run_id, + config_ref=config_ref, + config_sha256=config_sha256, + snapshot_path=str(snapshot), + request_path=str(request_path), + resource_keys=tuple(resource_keys), + ), + job_dir, + ) + except BaseException: + _remove_job_dir(job_dir) + raise + + def _reject_existing_run( + self, + suite_id: str, + run_id: str, + ) -> None: + suite_root = self.workspace.path_policy.resolve_managed_output( + self.workspace.results_root / suite_id, + field_name="job suite root", + expected_root=self.workspace.results_root, + reject_links=True, + ) + run_root = self.workspace.path_policy.resolve_managed_output( + suite_root / run_id, + field_name="job run root", + expected_root=suite_root, + reject_links=True, + ) + if run_root.exists(): + raise ServiceError( + ServiceErrorCode.CONFLICT, + "The requested suite/run output already exists", + ) + + def _detail(self, record: JobRecord) -> JobDetail: + manifest = self._manifest(record) + heartbeat_at = _optional_text(manifest.get("heartbeat_at")) + terminal_result = ( + JobTerminalResult.model_validate(record.result) + if record.result is not None + else None + ) + return JobDetail( + **_catalog_entry(record).model_dump(), + request_id=record.idempotency_key, + config_sha256=record.config_sha256, + heartbeat_at=heartbeat_at, + heartbeat_age_seconds=_heartbeat_age(heartbeat_at), + stages=( + dict(manifest.get("stages") or {}) + if isinstance(manifest.get("stages"), dict) + else {} + ), + stage_timings=( + dict(manifest.get("stage_timings") or {}) + if isinstance(manifest.get("stage_timings"), dict) + else {} + ), + progress=( + dict(manifest.get("progress") or {}) + if isinstance(manifest.get("progress"), dict) + else {} + ), + terminal_result=terminal_result, + error_code=record.error_code, + error_message=record.error_message, + resources=_job_resources(record), + ) + + def _manifest(self, record: JobRecord) -> dict[str, Any]: + if record.run_id is None: + return {} + suite_root = self.workspace.path_policy.resolve_managed_output( + self.workspace.results_root / record.suite_id, + field_name="job suite root", + expected_root=self.workspace.results_root, + reject_links=True, + ) + run_root = self.workspace.path_policy.resolve_managed_output( + suite_root / record.run_id, + field_name="job run root", + expected_root=suite_root, + reject_links=True, + ) + manifest_path = self.workspace.path_policy.resolve_managed_output( + run_root / "manifest.json", + field_name="job run manifest", + expected_root=self.workspace.results_root, + reject_links=True, + ) + payload = _read_json_file( + manifest_path, + max_bytes=_JOB_RESULT_MAX_BYTES, + ) + return payload if isinstance(payload, dict) else {} + + +def _catalog_entry(record: JobRecord) -> JobCatalogEntry: + return JobCatalogEntry( + job_id=record.job_id, + state=record.state, + revision=record.revision, + kind="evaluation", + config_ref=record.config_ref, + suite_id=record.suite_id, + run_id=record.run_id, + created_at=record.created_at, + started_at=record.started_at, + ended_at=record.ended_at, + ) + + +def _job_resources(record: JobRecord) -> dict[str, str]: + resources = { + "config": f"assert://config/{quote(record.config_ref, safe='')}", + "worker_log": f"assert://job/{quote(record.job_id, safe='')}/log", + } + if record.run_id is not None: + suite_id = quote(record.suite_id, safe="") + run_id = quote(record.run_id, safe="") + resources.update( + { + "run_summary": ( + f"assert://run/{suite_id}/{run_id}/summary" + ), + "run_manifest": ( + f"assert://run/{suite_id}/{run_id}/manifest" + ), + "run_config": ( + f"assert://run/{suite_id}/{run_id}/config" + ), + } + ) + return resources + + +def _jobs_root(workspace: WorkspaceService) -> Path: + return workspace.path_policy.resolve_managed_output( + workspace.artifacts_root / "mcp" / "jobs", + field_name="evaluation jobs root", + expected_root=workspace.artifacts_root, + reject_links=True, + ) + + +def _request_hash( + *, + config_ref: str, + config_etag: str, + overrides: EvaluationOverrides, +) -> str: + payload = json.dumps( + { + "config_ref": config_ref, + "config_etag": config_etag, + "overrides": overrides.model_dump(mode="json"), + }, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return "sha256:" + hashlib.sha256(payload).hexdigest() + + +def _require_ready(plan: Any) -> None: + if plan.ready: + return + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + "Evaluation preflight has blocking issues", + details={ + "validation": plan.validation.model_dump(mode="json"), + "blocking_issues": [ + issue.model_dump(mode="json") + for issue in plan.blocking_issues + ], + }, + ) + + +def _require_source_etag(actual: str, expected: str) -> None: + if actual == expected: + return + raise ServiceError( + ServiceErrorCode.STALE_ETAG, + "Config changed while the evaluation request was being prepared; retry", + details={ + "expected_etag": expected, + "current_etag": actual, + }, + ) + + +def _validate_request_id(value: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "request_id must be a non-empty string", + ) + normalized = value.strip() + if len(normalized) > _REQUEST_ID_MAX_LENGTH: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"request_id must be <= {_REQUEST_ID_MAX_LENGTH} characters", + ) + return normalized + + +def _validate_job_id(value: str) -> str: + try: + parsed = uuid.UUID(value) + except (AttributeError, TypeError, ValueError) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "job_id must be a valid ASSERT job id", + ) from exc + normalized = parsed.hex + if normalized != value: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "job_id must use the canonical ASSERT job-id format", + ) + return normalized + + +def _new_identity(prefix: str) -> str: + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S") + return f"{prefix}-{timestamp}-{secrets.token_hex(8)}" + + +def _encode_cursor(created_at: str, job_id: str) -> str: + payload = json.dumps( + { + "v": _CURSOR_VERSION, + "created_at": created_at, + "job_id": job_id, + }, + separators=(",", ":"), + ).encode("utf-8") + return base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=") + + +def _decode_cursor(value: str) -> tuple[str, str]: + if not isinstance(value, str) or len(value) > 4096: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid job cursor", + ) + try: + padding = "=" * (-len(value) % 4) + payload = json.loads( + base64.urlsafe_b64decode(value + padding).decode("utf-8") + ) + if ( + not isinstance(payload, dict) + or payload.get("v") != _CURSOR_VERSION + or not isinstance(payload.get("created_at"), str) + or not isinstance(payload.get("job_id"), str) + ): + raise ValueError + parsed_at = datetime.fromisoformat(payload["created_at"]) + if parsed_at.tzinfo is None: + raise ValueError + _validate_job_id(payload["job_id"]) + except ( + binascii.Error, + ValueError, + TypeError, + json.JSONDecodeError, + ) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Invalid job cursor", + ) from exc + return payload["created_at"], payload["job_id"] + + +def _read_json_file(path: Path, *, max_bytes: int) -> Any: + try: + if not path.is_file() or path.stat().st_size > max_bytes: + return None + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + + +def _read_text_tail(path: Path, *, max_bytes: int) -> str: + if not path.is_file(): + return "" + try: + with path.open("rb") as stream: + size = stream.seek(0, 2) + stream.seek(max(0, size - max_bytes)) + value = stream.read(max_bytes) + except OSError as exc: + raise ServiceError( + ServiceErrorCode.INTERNAL, + "Could not read the evaluation worker log", + ) from exc + return value.decode("utf-8", errors="replace") + + +def _remove_job_dir(path: Path) -> None: + try: + if path.is_dir(): + shutil.rmtree(path) + except OSError: + log.warning( + "Could not remove unused evaluation job directory %s", + path, + exc_info=True, + ) + + +def _process_create_time(pid: int) -> float: + import psutil + + return float(psutil.Process(pid).create_time()) + + +def _terminate_failed_launch(process: subprocess.Popen[bytes]) -> None: + try: + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + except (OSError, subprocess.SubprocessError): + log.exception( + "Could not terminate partially launched evaluation worker %s", + process.pid, + ) + + +def _process_matches(pid: int, create_time: float) -> bool: + try: + import psutil + + process = psutil.Process(pid) + return ( + process.is_running() + and abs(process.create_time() - create_time) < 0.01 + ) + except psutil.Error: + return False + + +def _lease_expired(value: str | None) -> bool: + if value is None: + return True + try: + parsed = datetime.fromisoformat(value) + if parsed.tzinfo is None: + return True + return parsed <= datetime.now(timezone.utc) + except (TypeError, ValueError): + return True + + +def _heartbeat_age(value: str | None) -> float | None: + if value is None: + return None + try: + heartbeat = datetime.fromisoformat(value) + if heartbeat.tzinfo is None: + return None + except (TypeError, ValueError): + return None + return max( + 0.0, + (datetime.now(timezone.utc) - heartbeat).total_seconds(), + ) + + +def _safe_error( + error: Exception, + *, + workspace: WorkspaceService, + fallback: str, +) -> str: + return _safe_text( + str(error), + workspace=workspace, + fallback=fallback, + ) or fallback + + +def _safe_text( + value: Any, + *, + workspace: WorkspaceService, + fallback: str | None, +) -> str | None: + if value is None: + return fallback + sanitized = sanitize_text(str(value)) + sanitized = redact_path_prefixes( + sanitized, + ( + workspace.root, + workspace.configs_root, + workspace.artifacts_root, + workspace.results_root, + ), + ) + return sanitized or fallback + + +def _optional_text(value: Any) -> str | None: + return str(value) if value is not None else None + + +def _validated_optional_text( + value: Any, + *, + field_name: str, +) -> str | None: + if value is None: + return None + if not isinstance(value, str): + raise ServiceError( + ServiceErrorCode.RUN_FAILED, + f"Evaluation worker returned an invalid {field_name}", + ) + return value diff --git a/assert_ai/services/job_models.py b/assert_ai/services/job_models.py new file mode 100644 index 000000000..130806d22 --- /dev/null +++ b/assert_ai/services/job_models.py @@ -0,0 +1,154 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Transport-neutral models for persisted evaluation jobs.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field + + +JOB_SCHEMA_VERSION = 1 + + +class JobState(StrEnum): + """Persisted orchestration states for evaluation workers.""" + + QUEUED = "queued" + STARTING = "starting" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLING = "cancelling" + CANCELLED = "cancelled" + INTERRUPTED = "interrupted" + + +TERMINAL_JOB_STATES = frozenset( + { + JobState.COMPLETED, + JobState.FAILED, + JobState.CANCELLED, + JobState.INTERRUPTED, + } +) + + +@dataclass(frozen=True, slots=True) +class JobRecord: + """Complete internal representation of one persisted job row.""" + + job_id: str + idempotency_key: str + request_hash: str + kind: str + state: JobState + created_at: str + started_at: str | None + ended_at: str | None + suite_id: str + run_id: str | None + config_ref: str + config_sha256: str + snapshot_path: str + request_path: str + run_root: str | None + pid: int | None + process_create_time: float | None + exit_code: int | None + failed_stage: str | None + error_code: str | None + error_message: str | None + cancel_requested_at: str | None + result: dict[str, Any] | None + resource_keys: tuple[str, ...] + revision: int + lease_owner: str | None + lease_expires_at: str | None + + +@dataclass(frozen=True, slots=True) +class NewJob: + """Values fixed before a queued job becomes visible.""" + + job_id: str + idempotency_key: str + request_hash: str + suite_id: str + run_id: str | None + config_ref: str + config_sha256: str + snapshot_path: str + request_path: str + resource_keys: tuple[str, ...] + kind: str = "evaluation" + + +@dataclass(frozen=True, slots=True) +class CreateJobResult: + record: JobRecord + created: bool + + +class _ServiceModel(BaseModel): + model_config = ConfigDict(frozen=True) + + +class JobTerminalResult(_ServiceModel): + """Sanitized terminal outcome retained with a job.""" + + state: Literal["completed", "failed", "cancelled"] + exit_code: int + failed_stage: str | None = None + error_code: str | None = None + error_message: str | None = None + + +class JobCatalogEntry(_ServiceModel): + """Lightweight persisted-job metadata.""" + + schema_version: Literal[1] = JOB_SCHEMA_VERSION + job_id: str + state: JobState + revision: int = Field(ge=0) + kind: Literal["evaluation"] = "evaluation" + config_ref: str + suite_id: str + run_id: str | None = None + created_at: str + started_at: str | None = None + ended_at: str | None = None + + +class JobPage(_ServiceModel): + """Bounded page of persisted jobs.""" + + items: tuple[JobCatalogEntry, ...] + next_cursor: str | None = None + + +class JobDetail(JobCatalogEntry): + """Detailed status for one evaluation job.""" + + request_id: str + config_sha256: str + heartbeat_at: str | None = None + heartbeat_age_seconds: float | None = Field(default=None, ge=0) + stages: dict[str, Any] = Field(default_factory=dict) + stage_timings: dict[str, Any] = Field(default_factory=dict) + progress: dict[str, Any] = Field(default_factory=dict) + terminal_result: JobTerminalResult | None = None + error_code: str | None = None + error_message: str | None = None + resources: dict[str, str] = Field(default_factory=dict) + + +class JobStartResult(_ServiceModel): + """Idempotent response returned after an evaluation is accepted.""" + + job: JobDetail + created: bool diff --git a/assert_ai/services/job_store.py b/assert_ai/services/job_store.py new file mode 100644 index 000000000..d92516b3c --- /dev/null +++ b/assert_ai/services/job_store.py @@ -0,0 +1,836 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Transactional SQLite persistence for evaluation jobs.""" + +from __future__ import annotations + +import json +import sqlite3 +import threading +import time +from contextlib import contextmanager +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any, Iterator, Sequence + +from assert_ai.core.runtime_path_policy import RuntimePathPolicy +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.job_models import ( + CreateJobResult, + JobRecord, + JobState, + NewJob, + TERMINAL_JOB_STATES, +) + +_BUSY_TIMEOUT_MS = 5_000 +_JOB_STORE_SCHEMA_VERSION = 1 +_ACTIVE_STATES = (JobState.STARTING.value, JobState.RUNNING.value) +_MAX_EVENTS_PER_JOB = 1000 + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS jobs( + job_id TEXT PRIMARY KEY, + idempotency_key TEXT NOT NULL UNIQUE, + request_hash TEXT NOT NULL, + kind TEXT NOT NULL, + state TEXT NOT NULL, + created_at TEXT NOT NULL, + started_at TEXT, + ended_at TEXT, + suite_id TEXT NOT NULL, + run_id TEXT, + config_ref TEXT NOT NULL, + config_sha256 TEXT NOT NULL, + snapshot_path TEXT NOT NULL, + request_path TEXT NOT NULL, + run_root TEXT, + pid INTEGER, + process_create_time REAL, + exit_code INTEGER, + failed_stage TEXT, + error_code TEXT, + error_message TEXT, + cancel_requested_at TEXT, + result_json TEXT, + resource_keys_json TEXT NOT NULL, + revision INTEGER NOT NULL DEFAULT 0, + lease_owner TEXT, + lease_expires_at TEXT +); +CREATE UNIQUE INDEX IF NOT EXISTS jobs_suite_run_unique + ON jobs(suite_id, run_id) + WHERE run_id IS NOT NULL; +CREATE INDEX IF NOT EXISTS jobs_state_created + ON jobs(state, created_at, job_id); +CREATE TABLE IF NOT EXISTS job_events( + job_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + timestamp TEXT NOT NULL, + event_type TEXT NOT NULL, + payload_json TEXT NOT NULL, + PRIMARY KEY(job_id, sequence), + FOREIGN KEY(job_id) REFERENCES jobs(job_id) ON DELETE CASCADE +); +CREATE TABLE IF NOT EXISTS resource_locks( + resource_key TEXT PRIMARY KEY, + job_id TEXT NOT NULL, + acquired_at TEXT NOT NULL, + lease_expires_at TEXT NOT NULL, + FOREIGN KEY(job_id) REFERENCES jobs(job_id) ON DELETE CASCADE +); +""" + + +class JobStore: + """Concurrency-safe persisted job registry with expiring leases.""" + + def __init__( + self, + path: str | Path, + *, + path_policy: RuntimePathPolicy | None = None, + expected_root: str | Path | None = None, + ) -> None: + if (path_policy is None) != (expected_root is None): + raise ValueError( + "path_policy and expected_root must be provided together" + ) + self.path = Path(path) + self._path_policy = path_policy + self._expected_root = ( + Path(expected_root) if expected_root is not None else None + ) + self._initialize_lock = threading.Lock() + self._initialized = False + + @property + def exists(self) -> bool: + return self._database_path().is_file() + + def create_or_get( + self, + new_job: NewJob, + *, + max_queued_jobs: int, + ) -> CreateJobResult: + if max_queued_jobs < 1: + raise ValueError("max_queued_jobs must be positive") + self.initialize() + with self._transaction() as connection: + existing = connection.execute( + "SELECT * FROM jobs WHERE idempotency_key = ?", + (new_job.idempotency_key,), + ).fetchone() + if existing is not None: + record = _record(existing) + if record.request_hash != new_job.request_hash: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "request_id is already bound to a different evaluation request", + details={"job_id": record.job_id}, + ) + return CreateJobResult(record=record, created=False) + + queued_count = int( + connection.execute( + "SELECT COUNT(*) FROM jobs WHERE state = ?", + (JobState.QUEUED.value,), + ).fetchone()[0] + ) + if queued_count >= max_queued_jobs: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "The evaluation queue has reached its operator limit", + details={"max_queued_jobs": max_queued_jobs}, + ) + + created_at = _now() + try: + connection.execute( + """ + INSERT INTO jobs( + job_id, idempotency_key, request_hash, kind, state, + created_at, suite_id, run_id, config_ref, + config_sha256, snapshot_path, request_path, + resource_keys_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + new_job.job_id, + new_job.idempotency_key, + new_job.request_hash, + new_job.kind, + JobState.QUEUED.value, + created_at, + new_job.suite_id, + new_job.run_id, + new_job.config_ref, + new_job.config_sha256, + new_job.snapshot_path, + new_job.request_path, + _json(new_job.resource_keys), + ), + ) + except sqlite3.IntegrityError as exc: + collision = connection.execute( + "SELECT job_id FROM jobs WHERE suite_id = ? AND run_id = ?", + (new_job.suite_id, new_job.run_id), + ).fetchone() + if collision is not None: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "The requested suite/run output is already assigned", + details={"job_id": str(collision["job_id"])}, + ) from exc + raise + self._append_event( + connection, + new_job.job_id, + "queued", + {"state": JobState.QUEUED.value}, + timestamp=created_at, + ) + row = connection.execute( + "SELECT * FROM jobs WHERE job_id = ?", + (new_job.job_id,), + ).fetchone() + assert row is not None + return CreateJobResult(record=_record(row), created=True) + + def get(self, job_id: str) -> JobRecord: + if not self.exists: + raise ServiceError(ServiceErrorCode.NOT_FOUND, "Job not found") + with self._connection() as connection: + row = connection.execute( + "SELECT * FROM jobs WHERE job_id = ?", + (job_id,), + ).fetchone() + if row is None: + raise ServiceError(ServiceErrorCode.NOT_FOUND, "Job not found") + return _record(row) + + def get_by_idempotency_key( + self, + idempotency_key: str, + ) -> JobRecord | None: + if not self.exists: + return None + with self._connection() as connection: + row = connection.execute( + "SELECT * FROM jobs WHERE idempotency_key = ?", + (idempotency_key,), + ).fetchone() + return _record(row) if row is not None else None + + def list_records( + self, + *, + limit: int, + states: Sequence[JobState] = (), + before: tuple[str, str] | None = None, + ) -> tuple[JobRecord, ...]: + if limit < 1: + raise ValueError("limit must be positive") + if not self.exists: + return () + conditions: list[str] = [] + values: list[Any] = [] + if states: + placeholders = ", ".join("?" for _ in states) + conditions.append(f"state IN ({placeholders})") + values.extend(state.value for state in states) + if before is not None: + conditions.append( + "(created_at < ? OR (created_at = ? AND job_id < ?))" + ) + values.extend((before[0], before[0], before[1])) + where = f"WHERE {' AND '.join(conditions)}" if conditions else "" + values.append(limit) + with self._connection() as connection: + rows = connection.execute( + f""" + SELECT * FROM jobs + {where} + ORDER BY created_at DESC, job_id DESC + LIMIT ? + """, + tuple(values), + ).fetchall() + return tuple(_record(row) for row in rows) + + def claim_next( + self, + *, + lease_owner: str, + lease_seconds: float, + max_active_jobs: int, + ) -> JobRecord | None: + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + if max_active_jobs < 1: + raise ValueError("max_active_jobs must be positive") + if not self.exists: + return None + now = _now() + expires_at = _after(lease_seconds) + with self._transaction() as connection: + active = int( + connection.execute( + "SELECT COUNT(*) FROM jobs WHERE state IN (?, ?)", + _ACTIVE_STATES, + ).fetchone()[0] + ) + if active >= max_active_jobs: + return None + candidates = connection.execute( + """ + SELECT * FROM jobs + WHERE state = ? + ORDER BY created_at, job_id + """, + (JobState.QUEUED.value,), + ).fetchall() + for row in candidates: + record = _record(row) + if not self._resources_available( + connection, + record.resource_keys, + ): + continue + changed = connection.execute( + """ + UPDATE jobs + SET state = ?, lease_owner = ?, lease_expires_at = ?, + revision = revision + 1 + WHERE job_id = ? AND state = ? + """, + ( + JobState.STARTING.value, + lease_owner, + expires_at, + record.job_id, + JobState.QUEUED.value, + ), + ).rowcount + if changed != 1: + continue + for resource_key in record.resource_keys: + connection.execute( + """ + INSERT INTO resource_locks( + resource_key, job_id, acquired_at, + lease_expires_at + ) VALUES (?, ?, ?, ?) + """, + ( + resource_key, + record.job_id, + now, + expires_at, + ), + ) + self._append_event( + connection, + record.job_id, + "starting", + {"state": JobState.STARTING.value}, + timestamp=now, + ) + claimed = connection.execute( + "SELECT * FROM jobs WHERE job_id = ?", + (record.job_id,), + ).fetchone() + assert claimed is not None + return _record(claimed) + return None + + def mark_running( + self, + job_id: str, + *, + lease_owner: str, + pid: int, + process_create_time: float, + lease_seconds: float, + ) -> JobRecord: + now = _now() + expires_at = _after(lease_seconds) + with self._transaction() as connection: + changed = connection.execute( + """ + UPDATE jobs + SET state = ?, started_at = COALESCE(started_at, ?), + pid = ?, process_create_time = ?, + lease_expires_at = ?, revision = revision + 1 + WHERE job_id = ? AND state = ? AND lease_owner = ? + """, + ( + JobState.RUNNING.value, + now, + pid, + process_create_time, + expires_at, + job_id, + JobState.STARTING.value, + lease_owner, + ), + ).rowcount + if changed != 1: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "Job can no longer transition to running", + ) + connection.execute( + """ + UPDATE resource_locks + SET lease_expires_at = ? + WHERE job_id = ? + """, + (expires_at, job_id), + ) + self._append_event( + connection, + job_id, + "running", + {"state": JobState.RUNNING.value, "pid": pid}, + timestamp=now, + ) + return self._get_in_transaction(connection, job_id) + + def renew_lease( + self, + job_id: str, + *, + lease_owner: str, + lease_seconds: float, + ) -> bool: + expires_at = _after(lease_seconds) + with self._transaction() as connection: + changed = connection.execute( + """ + UPDATE jobs + SET lease_expires_at = ? + WHERE job_id = ? AND lease_owner = ? + AND state IN (?, ?) + """, + ( + expires_at, + job_id, + lease_owner, + JobState.STARTING.value, + JobState.RUNNING.value, + ), + ).rowcount + if changed != 1: + return False + connection.execute( + """ + UPDATE resource_locks + SET lease_expires_at = ? + WHERE job_id = ? + """, + (expires_at, job_id), + ) + return True + + def mark_terminal( + self, + job_id: str, + *, + state: JobState, + exit_code: int | None, + failed_stage: str | None, + error_code: str | None, + error_message: str | None, + result: dict[str, Any] | None, + run_root: str | None, + lease_owner: str | None = None, + ) -> JobRecord: + if state not in TERMINAL_JOB_STATES: + raise ValueError("terminal state required") + now = _now() + with self._transaction() as connection: + current = self._get_in_transaction(connection, job_id) + if current.state in TERMINAL_JOB_STATES: + return current + if current.state not in { + JobState.STARTING, + JobState.RUNNING, + JobState.CANCELLING, + }: + raise ServiceError( + ServiceErrorCode.CONFLICT, + f"Job cannot transition from {current.state.value} to {state.value}", + ) + if lease_owner is not None and current.lease_owner != lease_owner: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "Job lease is owned by another manager", + ) + connection.execute( + """ + UPDATE jobs + SET state = ?, ended_at = ?, exit_code = ?, + failed_stage = ?, error_code = ?, error_message = ?, + result_json = ?, run_root = ?, lease_owner = NULL, + lease_expires_at = NULL, revision = revision + 1 + WHERE job_id = ? + """, + ( + state.value, + now, + exit_code, + failed_stage, + error_code, + error_message, + _json(result) if result is not None else None, + run_root, + job_id, + ), + ) + connection.execute( + "DELETE FROM resource_locks WHERE job_id = ?", + (job_id,), + ) + self._append_event( + connection, + job_id, + state.value, + { + "state": state.value, + "exit_code": exit_code, + "failed_stage": failed_stage, + "error_code": error_code, + }, + timestamp=now, + ) + return self._get_in_transaction(connection, job_id) + + def append_event( + self, + job_id: str, + event_type: str, + payload: dict[str, Any], + ) -> int: + with self._transaction() as connection: + self._get_in_transaction(connection, job_id) + return self._append_event( + connection, + job_id, + event_type, + payload, + timestamp=_now(), + ) + + def list_events( + self, + job_id: str, + *, + after_sequence: int = 0, + limit: int = 200, + ) -> tuple[dict[str, Any], ...]: + if limit < 1: + raise ValueError("limit must be positive") + if limit > _MAX_EVENTS_PER_JOB: + raise ValueError( + f"limit must be <= {_MAX_EVENTS_PER_JOB}" + ) + self.get(job_id) + with self._connection() as connection: + rows = connection.execute( + """ + SELECT sequence, timestamp, event_type, payload_json + FROM job_events + WHERE job_id = ? AND sequence > ? + ORDER BY sequence + LIMIT ? + """, + (job_id, after_sequence, limit), + ).fetchall() + return tuple( + { + "sequence": int(row["sequence"]), + "timestamp": str(row["timestamp"]), + "event_type": str(row["event_type"]), + "payload": _json_object( + row["payload_json"], + field_name="job event payload", + default={}, + ), + } + for row in rows + ) + + def initialize(self) -> None: + if self._initialized: + return + with self._initialize_lock: + if self._initialized: + return + path = self._database_path() + path.parent.mkdir(parents=True, exist_ok=True) + self._database_path() + deadline = time.monotonic() + (_BUSY_TIMEOUT_MS / 1000) + while True: + try: + with self._connection() as connection: + connection.execute("PRAGMA journal_mode = WAL") + connection.executescript(_SCHEMA) + version = int( + connection.execute( + "PRAGMA user_version" + ).fetchone()[0] + ) + if version not in { + 0, + _JOB_STORE_SCHEMA_VERSION, + }: + raise ServiceError( + ServiceErrorCode.INTERNAL, + "Unsupported job store schema version", + ) + connection.execute( + "PRAGMA user_version = " + f"{_JOB_STORE_SCHEMA_VERSION}" + ) + self._initialized = True + return + except sqlite3.OperationalError as exc: + if ( + "locked" not in str(exc).lower() + or time.monotonic() >= deadline + ): + raise + time.sleep(0.05) + + @contextmanager + def _connection(self) -> Iterator[sqlite3.Connection]: + connection = sqlite3.connect( + self._database_path(), + timeout=_BUSY_TIMEOUT_MS / 1000, + isolation_level=None, + ) + connection.row_factory = sqlite3.Row + connection.execute(f"PRAGMA busy_timeout = {_BUSY_TIMEOUT_MS}") + connection.execute("PRAGMA foreign_keys = ON") + try: + yield connection + finally: + connection.close() + + def _database_path(self) -> Path: + if self._path_policy is None: + return self.path + assert self._expected_root is not None + return self._path_policy.resolve_managed_output( + self.path, + field_name="evaluation job store", + expected_root=self._expected_root, + reject_links=True, + ) + + @contextmanager + def _transaction(self) -> Iterator[sqlite3.Connection]: + with self._connection() as connection: + connection.execute("BEGIN IMMEDIATE") + try: + yield connection + except BaseException: + connection.rollback() + raise + else: + connection.commit() + + @staticmethod + def _resources_available( + connection: sqlite3.Connection, + resource_keys: tuple[str, ...], + ) -> bool: + if not resource_keys: + return True + placeholders = ", ".join("?" for _ in resource_keys) + row = connection.execute( + f""" + SELECT 1 FROM resource_locks + WHERE resource_key IN ({placeholders}) + LIMIT 1 + """, + resource_keys, + ).fetchone() + return row is None + + @staticmethod + def _get_in_transaction( + connection: sqlite3.Connection, + job_id: str, + ) -> JobRecord: + row = connection.execute( + "SELECT * FROM jobs WHERE job_id = ?", + (job_id,), + ).fetchone() + if row is None: + raise ServiceError(ServiceErrorCode.NOT_FOUND, "Job not found") + return _record(row) + + @staticmethod + def _append_event( + connection: sqlite3.Connection, + job_id: str, + event_type: str, + payload: dict[str, Any], + *, + timestamp: str, + ) -> int: + sequence = int( + connection.execute( + """ + SELECT COALESCE(MAX(sequence), 0) + 1 + FROM job_events + WHERE job_id = ? + """, + (job_id,), + ).fetchone()[0] + ) + connection.execute( + """ + INSERT INTO job_events( + job_id, sequence, timestamp, event_type, payload_json + ) VALUES (?, ?, ?, ?, ?) + """, + ( + job_id, + sequence, + timestamp, + event_type, + _json(payload), + ), + ) + cutoff = sequence - _MAX_EVENTS_PER_JOB + if cutoff > 0: + connection.execute( + """ + DELETE FROM job_events + WHERE job_id = ? AND sequence <= ? + """, + (job_id, cutoff), + ) + return sequence + + +def _record(row: sqlite3.Row) -> JobRecord: + return JobRecord( + job_id=str(row["job_id"]), + idempotency_key=str(row["idempotency_key"]), + request_hash=str(row["request_hash"]), + kind=str(row["kind"]), + state=JobState(str(row["state"])), + created_at=str(row["created_at"]), + started_at=_optional_str(row["started_at"]), + ended_at=_optional_str(row["ended_at"]), + suite_id=str(row["suite_id"]), + run_id=_optional_str(row["run_id"]), + config_ref=str(row["config_ref"]), + config_sha256=str(row["config_sha256"]), + snapshot_path=str(row["snapshot_path"]), + request_path=str(row["request_path"]), + run_root=_optional_str(row["run_root"]), + pid=int(row["pid"]) if row["pid"] is not None else None, + process_create_time=( + float(row["process_create_time"]) + if row["process_create_time"] is not None + else None + ), + exit_code=( + int(row["exit_code"]) + if row["exit_code"] is not None + else None + ), + failed_stage=_optional_str(row["failed_stage"]), + error_code=_optional_str(row["error_code"]), + error_message=_optional_str(row["error_message"]), + cancel_requested_at=_optional_str(row["cancel_requested_at"]), + result=_json_object( + row["result_json"], + field_name="job result", + default=None, + ), + resource_keys=_json_string_tuple( + row["resource_keys_json"], + field_name="job resource keys", + ), + revision=int(row["revision"]), + lease_owner=_optional_str(row["lease_owner"]), + lease_expires_at=_optional_str(row["lease_expires_at"]), + ) + + +def _optional_str(value: Any) -> str | None: + return str(value) if value is not None else None + + +def _json(value: Any) -> str: + return json.dumps( + value, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + + +def _load_json(value: Any, default: Any, *, field_name: str) -> Any: + if value is None: + return default + try: + return json.loads(str(value)) + except (TypeError, ValueError) as exc: + raise ServiceError( + ServiceErrorCode.INTERNAL, + f"Job store contains invalid {field_name}", + ) from exc + + +def _json_object( + value: Any, + *, + field_name: str, + default: dict[str, Any] | None, +) -> dict[str, Any] | None: + payload = _load_json( + value, + default, + field_name=field_name, + ) + if payload is not None and not isinstance(payload, dict): + raise ServiceError( + ServiceErrorCode.INTERNAL, + f"Job store contains invalid {field_name}", + ) + return payload + + +def _json_string_tuple( + value: Any, + *, + field_name: str, +) -> tuple[str, ...]: + payload = _load_json(value, [], field_name=field_name) + if not isinstance(payload, list) or not all( + isinstance(item, str) for item in payload + ): + raise ServiceError( + ServiceErrorCode.INTERNAL, + f"Job store contains invalid {field_name}", + ) + return tuple(payload) + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _after(seconds: float) -> str: + return ( + datetime.now(timezone.utc) + timedelta(seconds=seconds) + ).isoformat() diff --git a/assert_ai/services/run_planning.py b/assert_ai/services/run_planning.py index 995fd47df..57a0aebe7 100644 --- a/assert_ai/services/run_planning.py +++ b/assert_ai/services/run_planning.py @@ -44,7 +44,7 @@ class _ServiceModel(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = ConfigDict(frozen=True, extra="forbid") class ModelOverrides(_ServiceModel): diff --git a/tests/test_evaluation_service.py b/tests/test_evaluation_service.py new file mode 100644 index 000000000..c9425f2e9 --- /dev/null +++ b/tests/test_evaluation_service.py @@ -0,0 +1,357 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +import time +from pathlib import Path +from unittest.mock import patch + +import pytest + +from assert_ai.core.io import write_json +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services._evaluation_worker import ( + _BoundedTextLog, + main as worker_main, +) +from assert_ai.services.configs import ConfigService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.evaluations import ( + EvaluationJobManager, + EvaluationService, +) +from assert_ai.services.job_models import JobState +from assert_ai.services.job_store import JobStore +from assert_ai.services.run_planning import ( + EvaluationOverrides, + RunPlanningService, +) + + +def _service(root: Path) -> tuple[ConfigService, EvaluationService]: + workspace = WorkspaceService.create(root) + configs = ConfigService(workspace) + planning = RunPlanningService(workspace, configs) + store = JobStore(workspace.artifacts_root / "mcp" / "jobs.sqlite3") + manager = EvaluationJobManager(workspace, store, max_active_jobs=1) + return configs, EvaluationService( + workspace, + configs, + planning, + store, + manager, + default_page_size=10, + max_page_size=20, + max_queued_jobs=10, + ) + + +def _write_inference_fixture(root: Path) -> dict: + fixture = root / "evals" / "fixture.jsonl" + fixture.parent.mkdir(parents=True, exist_ok=True) + fixture.write_text( + json.dumps( + { + "type": "prompt", + "test_case_id": "case-1", + "behavior": "local behavior", + "seed": {"description": "hello"}, + } + ) + + "\n", + encoding="utf-8", + ) + (root / "agent.py").write_text( + "def run(message, *, history=None):\n" + " del history\n" + " return f'local: {message}'\n", + encoding="utf-8", + ) + return { + "suite": "mcp-job-suite", + "pipeline": { + "inference": { + "target": { + "callable": "agent:run", + }, + "test_set_path": "fixture.jsonl", + "concurrency": 1, + } + }, + } + + +def _wait_terminal( + service: EvaluationService, + job_id: str, + *, + timeout_s: float = 30, +): + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + detail = service.get(job_id) + if detail.state in { + JobState.COMPLETED, + JobState.FAILED, + JobState.INTERRUPTED, + }: + return detail + time.sleep(0.05) + raise AssertionError("evaluation job did not finish") + + +def test_inference_only_job_completes_and_is_idempotent( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + + started = service.start( + "demo.yaml", + request_id="request-one", + ) + terminal = _wait_terminal(service, started.job.job_id) + repeated = service.start( + "demo.yaml", + request_id="request-one", + ) + _, restarted_service = _service(tmp_path) + repeated_after_restart = restarted_service.start( + "demo.yaml", + request_id="request-one", + ) + + assert started.created is True + assert terminal.state is JobState.COMPLETED + assert terminal.terminal_result is not None + assert terminal.terminal_result.exit_code == 0 + assert terminal.run_id is not None + assert repeated.created is False + assert repeated.job.job_id == started.job.job_id + assert repeated_after_restart.created is False + assert repeated_after_restart.job.job_id == started.job.job_id + run_root = ( + tmp_path + / "artifacts" + / "results" + / terminal.suite_id + / terminal.run_id + ) + inference_rows = [ + json.loads(line) + for line in (run_root / "inference_set.jsonl") + .read_text(encoding="utf-8") + .splitlines() + ] + assert inference_rows[0]["events"][-1]["edit"]["message"]["content"] == ( + "local: hello" + ) + snapshot = ( + run_root / "config.yaml" + ).read_text(encoding="utf-8") + assert f"run: {terminal.run_id}" in snapshot + assert service.list().items[0].job_id == started.job.job_id + + +def test_request_id_conflict_does_not_launch_duplicate( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + first = service.start("demo.yaml", request_id="same") + + with pytest.raises(ServiceError) as conflict: + service.start( + "demo.yaml", + request_id="same", + overrides=EvaluationOverrides(run="different"), + ) + + assert conflict.value.code is ServiceErrorCode.CONFLICT + assert conflict.value.details == {"job_id": first.job.job_id} + assert _wait_terminal(service, first.job.job_id).state is JobState.COMPLETED + + +def test_config_change_during_start_requires_a_retry( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + document = _write_inference_fixture(tmp_path) + saved = configs.save_config("demo.yaml", document=document) + original_preflight = service.planning.preflight + calls = 0 + + def racing_preflight( + _planning: RunPlanningService, + config_ref: str, + *, + overrides: EvaluationOverrides | None = None, + ): + nonlocal calls + calls += 1 + plan = original_preflight(config_ref, overrides=overrides) + if calls == 1: + configs.save_config( + "demo.yaml", + document={ + **document, + "context": "changed during start", + }, + expected_etag=saved.etag, + ) + return plan + + with ( + patch.object( + RunPlanningService, + "preflight", + new=racing_preflight, + ), + pytest.raises(ServiceError) as stale, + ): + service.start("demo.yaml", request_id="request") + + assert stale.value.code is ServiceErrorCode.STALE_ETAG + assert not service.store.exists + + +def test_failed_snapshot_write_removes_unregistered_job_directory( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + + with ( + patch( + "assert_ai.services.evaluations.write_json", + side_effect=OSError("disk full"), + ), + pytest.raises(OSError, match="disk full"), + ): + service.start("demo.yaml", request_id="request") + + jobs_root = tmp_path / "artifacts" / "mcp" / "jobs" + assert list(jobs_root.iterdir()) == [] + assert not service.store.exists + + +def test_missing_runtime_input_is_persisted_as_run_failure( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + document = _write_inference_fixture(tmp_path) + document["pipeline"]["inference"]["test_set_path"] = ( + "missing.jsonl" + ) + configs.save_config("demo.yaml", document=document) + started = service.start("demo.yaml", request_id="request") + terminal = _wait_terminal(service, started.job.job_id) + + assert terminal.state is JobState.FAILED + assert terminal.error_code == "RUN_FAILED" + + +def test_worker_rejects_a_tampered_config_snapshot( + tmp_path: Path, +) -> None: + job_id = "a" * 32 + job_dir = tmp_path / "artifacts" / "mcp" / "jobs" / job_id + job_dir.mkdir(parents=True) + (job_dir / "config.yaml").write_text( + "pipeline: {}\n", + encoding="utf-8", + ) + write_json( + job_dir / "request.json", + { + "schema_version": 1, + "job_id": job_id, + "result_token": "token", + "config_ref": "demo.yaml", + "config_sha256": "sha256:" + ("0" * 64), + "strict": False, + "force_stages": [], + "max_log_bytes": 4096, + }, + ) + + exit_code = worker_main( + ["--workspace", str(tmp_path), "--job-id", job_id] + ) + + assert exit_code == 1 + result = json.loads( + (job_dir / "result.json").read_text(encoding="utf-8") + ) + assert result["result_token"] == "token" + assert result["worker_error"]["error_code"] == "INTERNAL" + assert "digest mismatch" in result["worker_error"]["error_message"] + + +def test_malformed_worker_result_becomes_a_persisted_failure( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + with patch.object(EvaluationJobManager, "enqueue"): + started = service.start("demo.yaml", request_id="request") + record = service.store.claim_next( + lease_owner="test-manager", + lease_seconds=30, + max_active_jobs=1, + ) + assert record is not None + record = service.store.mark_running( + record.job_id, + lease_owner="test-manager", + pid=123, + process_create_time=456, + lease_seconds=30, + ) + request = json.loads(Path(record.request_path).read_text(encoding="utf-8")) + write_json( + Path(record.request_path).parent / "result.json", + { + "schema_version": 1, + "job_id": record.job_id, + "result_token": request["result_token"], + "run_result": { + "state": "unknown", + "exit_code": 0, + }, + }, + ) + + detail = service.get(record.job_id) + listed = service.list().items[0] + + assert detail.state is JobState.FAILED + assert detail.error_code == "INTERNAL" + assert "invalid state" in detail.error_message + assert listed.state is JobState.FAILED + + +def test_worker_log_retains_a_bounded_tail(tmp_path: Path) -> None: + log_path = tmp_path / "worker.log" + + with _BoundedTextLog(log_path, max_bytes=4096) as worker_log: + worker_log.write("first-line\n") + worker_log.write("x" * 10_000) + + contents = log_path.read_text(encoding="utf-8") + assert log_path.stat().st_size <= 4096 + assert contents.startswith("[earlier worker output truncated]") diff --git a/tests/test_job_store.py b/tests/test_job_store.py new file mode 100644 index 000000000..81dcb3926 --- /dev/null +++ b/tests/test_job_store.py @@ -0,0 +1,221 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest + +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.job_models import JobState, NewJob +from assert_ai.services.job_store import JobStore + + +def _new_job( + suffix: str, + *, + request_id: str | None = None, + request_hash: str | None = None, + suite_id: str | None = None, + run_id: str | None = None, + resource_keys: tuple[str, ...] = (), +) -> NewJob: + return NewJob( + job_id=f"job-{suffix}", + idempotency_key=request_id or f"request-{suffix}", + request_hash=request_hash or f"hash-{suffix}", + suite_id=suite_id or f"suite-{suffix}", + run_id=run_id if run_id is not None else f"run-{suffix}", + config_ref="demo.yaml", + config_sha256=f"sha256:{suffix}", + snapshot_path=f"artifacts/mcp/jobs/job-{suffix}/config.yaml", + request_path=f"artifacts/mcp/jobs/job-{suffix}/request.json", + resource_keys=resource_keys, + ) + + +def test_create_is_idempotent_and_detects_request_conflicts( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + first = store.create_or_get( + _new_job("one"), + max_queued_jobs=10, + ) + repeated = store.create_or_get( + _new_job( + "other-id", + request_id="request-one", + request_hash="hash-one", + ), + max_queued_jobs=10, + ) + + assert first.created is True + assert repeated.created is False + assert repeated.record.job_id == "job-one" + assert store.list_events("job-one")[0]["event_type"] == "queued" + + with pytest.raises(ServiceError) as conflict: + store.create_or_get( + _new_job( + "conflict", + request_id="request-one", + request_hash="different", + ), + max_queued_jobs=10, + ) + assert conflict.value.code == ServiceErrorCode.CONFLICT + + +def test_suite_run_assignment_is_unique(tmp_path: Path) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get(_new_job("one"), max_queued_jobs=10) + + with pytest.raises(ServiceError) as conflict: + store.create_or_get( + _new_job( + "two", + suite_id="suite-one", + run_id="run-one", + ), + max_queued_jobs=10, + ) + + assert conflict.value.code == ServiceErrorCode.CONFLICT + assert conflict.value.details == {"job_id": "job-one"} + + +def test_create_respects_queued_job_limit(tmp_path: Path) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get(_new_job("one"), max_queued_jobs=1) + + with pytest.raises(ServiceError) as conflict: + store.create_or_get(_new_job("two"), max_queued_jobs=1) + + assert conflict.value.code == ServiceErrorCode.CONFLICT + assert conflict.value.details == {"max_queued_jobs": 1} + + +def test_claim_transition_and_terminal_result_persist( + tmp_path: Path, +) -> None: + path = tmp_path / "jobs.sqlite3" + store = JobStore(path) + store.create_or_get( + _new_job( + "one", + resource_keys=("suite:suite-one", "run:suite-one/run-one"), + ), + max_queued_jobs=10, + ) + + claimed = store.claim_next( + lease_owner="manager-a", + lease_seconds=30, + max_active_jobs=1, + ) + assert claimed is not None + assert claimed.state is JobState.STARTING + running = store.mark_running( + claimed.job_id, + lease_owner="manager-a", + pid=123, + process_create_time=456.0, + lease_seconds=30, + ) + assert running.state is JobState.RUNNING + assert running.revision == 2 + completed = store.mark_terminal( + claimed.job_id, + state=JobState.COMPLETED, + exit_code=0, + failed_stage=None, + error_code=None, + error_message=None, + result={"state": "completed", "exit_code": 0}, + run_root="artifacts/results/suite-one/run-one", + lease_owner="manager-a", + ) + + assert completed.state is JobState.COMPLETED + assert completed.result == {"state": "completed", "exit_code": 0} + reopened = JobStore(path).get(claimed.job_id) + assert reopened == completed + assert [event["event_type"] for event in store.list_events(claimed.job_id)] == [ + "queued", + "starting", + "running", + "completed", + ] + + +def test_claim_respects_global_active_limit_and_resource_locks( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get( + _new_job("one", resource_keys=("suite:shared",)), + max_queued_jobs=10, + ) + store.create_or_get( + _new_job("two", resource_keys=("suite:shared",)), + max_queued_jobs=10, + ) + first = store.claim_next( + lease_owner="manager-a", + lease_seconds=30, + max_active_jobs=1, + ) + assert first is not None + assert ( + store.claim_next( + lease_owner="manager-b", + lease_seconds=30, + max_active_jobs=1, + ) + is None + ) + store.mark_terminal( + first.job_id, + state=JobState.FAILED, + exit_code=1, + failed_stage=None, + error_code="RUN_FAILED", + error_message="failed", + result=None, + run_root=None, + lease_owner="manager-a", + ) + second = store.claim_next( + lease_owner="manager-b", + lease_seconds=30, + max_active_jobs=1, + ) + assert second is not None + assert second.job_id == "job-two" + + +def test_concurrent_idempotent_create_has_one_winner( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + + def create(index: int) -> tuple[str, bool]: + result = store.create_or_get( + _new_job( + str(index), + request_id="same-request", + request_hash="same-hash", + ), + max_queued_jobs=10, + ) + return result.record.job_id, result.created + + with ThreadPoolExecutor(max_workers=4) as executor: + outcomes = list(executor.map(create, range(4))) + + assert sum(created for _, created in outcomes) == 1 + assert len({job_id for job_id, _ in outcomes}) == 1 diff --git a/tests/test_mcp_cli.py b/tests/test_mcp_cli.py index 10f0d91c9..42f7f7a10 100644 --- a/tests/test_mcp_cli.py +++ b/tests/test_mcp_cli.py @@ -61,6 +61,12 @@ def test_mcp_serve_forwards_resolved_options() -> None: "4096", "--max-concurrency", "7", + "--max-active-jobs", + "2", + "--max-queued-jobs", + "9", + "--max-job-log-bytes", + "5000", "--max-prompt-sample-size", "12", "--max-scenario-sample-size", @@ -86,6 +92,9 @@ def test_mcp_serve_forwards_resolved_options() -> None: assert create_kwargs["max_artifact_chunk_bytes"] == 2048 assert create_kwargs["max_config_bytes"] == 4096 assert create_kwargs["max_concurrency"] == 7 + assert create_kwargs["max_active_jobs"] == 2 + assert create_kwargs["max_queued_jobs"] == 9 + assert create_kwargs["max_job_log_bytes"] == 5000 assert create_kwargs["max_prompt_sample_size"] == 12 assert create_kwargs["max_scenario_sample_size"] == 13 assert create_kwargs["allowed_model_patterns"] == ("azure/*",) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 76b7dc720..92110d317 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -29,6 +29,8 @@ EXPECTED_INSPECT_TOOLS = { "get_server_info", + "list_jobs", + "get_job", "list_presets", "get_preset", "get_config_schema", @@ -55,11 +57,13 @@ EXPECTED_FULL_TOOLS = EXPECTED_AUTHOR_TOOLS | { "design_config", "probe_target", + "start_evaluation", } EXPECTED_RESOURCE_TEMPLATES = { "assert://preset/{kind}/{name}", "assert://config/{config_ref}", + "assert://job/{job_id}/log", "assert://suite/{suite_id}/taxonomy", "assert://suite/{suite_id}/test-case/{test_case_id}{?kind,run_id}", "assert://run/{suite_id}/{run_id}/summary", @@ -253,6 +257,45 @@ def _seed_workspace(root: Path) -> None: ) +def _seed_evaluation_workspace(root: Path) -> None: + evals_root = root / "evals" + evals_root.mkdir(parents=True, exist_ok=True) + _write_jsonl( + evals_root / "fixture.jsonl", + [ + { + "type": "prompt", + "test_case_id": "case-1", + "behavior": "local behavior", + "seed": {"description": "hello"}, + } + ], + ) + (root / "agent.py").write_text( + "def run(message, *, history=None):\n" + " del history\n" + " print('api_key=not-a-real-secret')\n" + " return f'local: {message}'\n", + encoding="utf-8", + ) + (evals_root / "job.yaml").write_text( + json.dumps( + { + "suite": "mcp-job-suite", + "pipeline": { + "inference": { + "target": {"callable": "agent:run"}, + "test_set_path": "fixture.jsonl", + "concurrency": 1, + } + }, + }, + indent=2, + ), + encoding="utf-8", + ) + + def _schema_digest(tool: Any) -> str: payload = { "input": tool.input_schema, @@ -319,6 +362,13 @@ def test_server_options_validate_response_limits(tmp_path: Path) -> None: ("field", "value", "message"), [ ("max_concurrency", 0, "max_concurrency must be positive"), + ("max_active_jobs", 0, "max_active_jobs must be positive"), + ("max_queued_jobs", 0, "max_queued_jobs must be positive"), + ( + "max_job_log_bytes", + 1024, + "max_job_log_bytes must be between", + ), ( "max_prompt_sample_size", 0, @@ -436,6 +486,9 @@ async def run() -> object: assert "env_file" not in result.structured_content assert result.structured_content["limits"]["max_page_size"] == 200 assert result.structured_content["limits"]["max_concurrency"] == 32 + assert result.structured_content["limits"]["max_active_jobs"] == 1 + assert result.structured_content["limits"]["max_queued_jobs"] == 100 + assert result.structured_content["limits"]["max_job_log_bytes"] == 1024 * 1024 assert result.structured_content["limits"]["max_prompt_sample_size"] == 100_000 assert result.structured_content["limits"]["max_scenario_sample_size"] == 100_000 assert result.structured_content["limits"]["model_allowlist_enabled"] is True @@ -495,15 +548,17 @@ async def run() -> list[Any]: "compare_runs": "f7bfeca051f8f81bf3621936588ed906332076a3a34b550090f87c2944656ce5", "get_config": "bf38188871cb818e0b0cf6e28183aa728ed8d041923a158f593832d2459bd13a", "get_config_schema": "cca1d3a48240e20eff93a123b34d7ba92df3ed1df87f57f9eb217aa21515ec26", + "get_job": "76794436d4665712dfbd226a4c44738f1b4e8ab6ff3ed3f7eb184311f1f60cb0", "get_preset": "81db6723ad5065ce8a0a402d29dc2f9df7657d302e3ebe8b377f54c9d62353d0", "get_run": "e5216cd0085d049f8b49c54add913b6f83756c4ce59995317fe63e010ea44936", - "get_server_info": "59f160f8840051916a5e0623fe0b46ea4bb6bba5b0ecc78202325a5b1ba4bc0d", + "get_server_info": "51e1d08335b4c8c6cb5eb70b6857563ab2dead550347f83dac64d89d8d417069", "get_suite": "8f629c93e02b656052f637c3cbba9217834315a693c4f7935f6d961203b46fd0", "get_test_case": "11380555caaa71d5992923815a499fc08b368c02f4d4836e4761630654589148", "get_transcript": "aa09669e0cb99202e8dec0b858b4faa41742ecb616351c3956b0d0bd488717e8", "list_artifacts": "3d3bede0b7209401b15d1f39d82671092c3097a05cd901122bd46c3c42edebfc", "list_configs": "92f78db2533034e6bf80e1d95089460acdd40a18468d4eb06fdf055726dfef19", "list_failures": "d3cc3f3bcc86c110754673297d28ac1e5ccbf698668c997de0bba2d0cbd425e2", + "list_jobs": "4570a8790f6f5c42fa015c49056111f4b7f00744a2300e3a79af9a68f08d2530", "list_presets": "55faa31adbf7f689eb5efbf1211fa73b2474d1a0e4549ec0a836ea69919c46b1", "list_runs": "7280687daafcd7ff5d89756c9584ca06c432a44f5a98ce8ff3ae0e4427dcf40b", "list_scores": "5c1951a3a3b91089b68b30e970a1b13f59bc2659a2c234db4451cbe2d5362a4d", @@ -532,6 +587,7 @@ async def run() -> dict[str, Any]: "preflight_evaluation": (True, False, True, False), "design_config": (True, False, False, True), "probe_target": (True, False, False, True), + "start_evaluation": (False, True, True, True), } expected_digests = { "validate_config": ( @@ -541,7 +597,7 @@ async def run() -> dict[str, Any]: "b09950417a44bf14c9bbf2702c1c00f23a18a0bfec03cab16a482733d8cf98c8" ), "preflight_evaluation": ( - "f53c4526b5df97e62f33b61a7c3a2eee375fc09eb44adabc7f5a9703d62eec64" + "c9a686f7879c7e06a8f32c210cc02ed8e30c4fb8c6473cf77999971d114c9805" ), "design_config": ( "1cd55a1bba06468b0aaa785cf05a445e4bf16768ff6165254ceecf0181a8392e" @@ -549,6 +605,9 @@ async def run() -> dict[str, Any]: "probe_target": ( "406d97ba84821a4e2661779dcc1211d208d2485013f8dea761c04f1fcdf59e63" ), + "start_evaluation": ( + "b594f1ef3510f31503960291e7ed5b57967876b5c5d4d628b0d5ee86cc000c0c" + ), } for name, annotations in expected_annotations.items(): @@ -672,6 +731,122 @@ async def run() -> dict[str, Any]: assert not (tmp_path / "artifacts").exists() +def test_complete_persisted_evaluation_workflow_through_mcp( + tmp_path: Path, +) -> None: + _seed_evaluation_workspace(tmp_path) + + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="full", + max_active_jobs=1, + max_queued_jobs=2, + ) + async with Client(build_server(options), raise_exceptions=True) as client: + empty = await client.call_tool("list_jobs", {}) + started = await client.call_tool( + "start_evaluation", + { + "config_ref": "job.yaml", + "request_id": "mcp-integration-request", + }, + ) + repeated = await client.call_tool( + "start_evaluation", + { + "config_ref": "job.yaml", + "request_id": "mcp-integration-request", + }, + ) + conflict = await client.call_tool( + "start_evaluation", + { + "config_ref": "job.yaml", + "request_id": "mcp-integration-request", + "overrides": {"run": "different-run"}, + }, + ) + invalid_override = await client.call_tool( + "start_evaluation", + { + "config_ref": "job.yaml", + "request_id": "invalid-override", + "overrides": {"unsupported": True}, + }, + ) + job_id = started.structured_content["job"]["job_id"] + deadline = asyncio.get_running_loop().time() + 30 + while True: + detail = await client.call_tool( + "get_job", + {"job_id": job_id}, + ) + if detail.structured_content["state"] in { + "completed", + "failed", + "interrupted", + }: + break + if asyncio.get_running_loop().time() >= deadline: + raise AssertionError("MCP evaluation job did not finish") + await asyncio.sleep(0.05) + jobs = await client.call_tool( + "list_jobs", + {"states": ["completed"], "page_size": 1}, + ) + run = await client.call_tool( + "get_run", + { + "suite_id": detail.structured_content["suite_id"], + "run_id": detail.structured_content["run_id"], + }, + ) + job_log = await client.read_resource( + detail.structured_content["resources"]["worker_log"] + ) + return { + "empty": empty, + "started": started, + "repeated": repeated, + "conflict": conflict, + "invalid_override": invalid_override, + "detail": detail, + "jobs": jobs, + "run": run, + "job_log": job_log.contents[0].text, + } + + results = asyncio.run(run()) + + assert results["empty"].structured_content == { + "items": [], + "next_cursor": None, + } + started = results["started"].structured_content + repeated = results["repeated"].structured_content + assert started["created"] is True + assert repeated["created"] is False + assert repeated["job"]["job_id"] == started["job"]["job_id"] + assert '"code":"CONFLICT"' in _error_text(results["conflict"]) + assert results["invalid_override"].is_error is True + detail = results["detail"].structured_content + assert detail["state"] == "completed" + assert detail["terminal_result"]["exit_code"] == 0 + assert detail["resources"]["config"] == "assert://config/job.yaml" + assert detail["resources"]["run_summary"].endswith("/summary") + assert "pid" not in detail + assert str(tmp_path) not in json.dumps(detail) + assert results["jobs"].structured_content["items"][0]["job_id"] == ( + detail["job_id"] + ) + assert results["run"].structured_content["state"] == "completed" + assert "filtered tail" in results["job_log"] + assert str(tmp_path) not in results["job_log"] + assert "not-a-real-secret" not in results["job_log"] + assert "[REDACTED]" in results["job_log"] + + def test_design_config_returns_an_unpersisted_draft(tmp_path: Path) -> None: draft = ConfigDraft( yaml=( diff --git a/tests/test_run_result.py b/tests/test_run_result.py index 9de873344..e027c6008 100644 --- a/tests/test_run_result.py +++ b/tests/test_run_result.py @@ -8,10 +8,16 @@ from tempfile import TemporaryDirectory from unittest.mock import patch +import yaml + from assert_ai.core.model_client import LLMInputError from assert_ai.core.run_result import RunState from assert_ai.core.workspace import WorkspaceService -from assert_ai.runner import run_pipeline, run_pipeline_result +from assert_ai.runner import ( + run_pipeline, + run_pipeline_document_result, + run_pipeline_result, +) def test_invalid_config_returns_typed_failure_and_legacy_exit_code() -> None: @@ -136,6 +142,40 @@ async def fail_stage(*_: object, **__: object) -> dict: assert result.error_message == "invalid request from ." +def test_document_run_preserves_logical_config_base_and_snapshot( + tmp_path: Path, +) -> None: + workspace = WorkspaceService.create(tmp_path) + logical_path = workspace.configs_root / "nested" / "config.yaml" + document = { + "suite": "suite-a", + "run": "run-a", + "pipeline": { + "inference": { + "target": {"callable": "agent:run"}, + "test_set_path": "fixture.jsonl", + } + }, + } + + async def complete_stage(*_: object, **__: object) -> dict: + return {} + + with patch("assert_ai.stages.inference.run", new=complete_stage): + result = run_pipeline_document_result( + document=document, + config_path=str(logical_path), + path_policy=workspace.path_policy, + ) + + assert result.state is RunState.COMPLETED + assert result.run_root is not None + snapshot = yaml.safe_load( + (result.run_root / "config.yaml").read_text(encoding="utf-8") + ) + assert snapshot == document + + def test_unexpected_setup_failure_is_returned_not_raised() -> None: with TemporaryDirectory() as tmp: root = Path(tmp) diff --git a/tests/test_runtime_path_policy.py b/tests/test_runtime_path_policy.py index aa610320b..133977f96 100644 --- a/tests/test_runtime_path_policy.py +++ b/tests/test_runtime_path_policy.py @@ -24,6 +24,7 @@ RuntimePathError, RuntimePathErrorCode, RuntimePathPolicy, + _is_within, ) from assert_ai.core.security import validate_sys_path_addition from assert_ai.core.tool_backend import import_callable_module, load_tool_module @@ -67,6 +68,15 @@ def test_workspace_service_exposes_only_relative_references(tmp_path: Path) -> N assert workspace.reference(workspace.results_root) == "artifacts/results" +@pytest.mark.skipif(os.name != "nt", reason="Windows path representation") +def test_extended_length_path_is_compared_as_the_same_windows_path() -> None: + root = Path("C:/workspace") + extended_child = Path("//?/C:/workspace/artifacts/results") + + assert _is_within(extended_child, root) + assert not _is_within(Path("//?/C:/outside"), root) + + def test_config_path_is_contained_under_config_root(tmp_path: Path) -> None: workspace, config_path = _workspace(tmp_path) diff --git a/tests/test_runtime_safety.py b/tests/test_runtime_safety.py index f9f8e882e..b1752a3d8 100644 --- a/tests/test_runtime_safety.py +++ b/tests/test_runtime_safety.py @@ -380,6 +380,24 @@ def _flaky_write(m: _StubManifest, root: Path) -> None: assert call_count[0] >= 2, "heartbeat should have retried despite errors" +def test_heartbeat_does_not_rewrite_a_terminal_manifest( + tmp_path: Path, +) -> None: + manifest = _StubManifest(status="failed") + writes: list[str] = [] + heartbeat = ManifestHeartbeat( + manifest, + tmp_path, + lambda current, _: writes.append(current.status), + interval_s=10, + ) + + heartbeat.start() + heartbeat.stop(write_final=True) + + assert writes == [] + + # --------------------------------------------------------------------------- # PipelineWatchdog # --------------------------------------------------------------------------- From f1a12786a206eae84e6f0f53457096707741f7e0 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Wed, 26 Aug 2026 10:35:11 -0700 Subject: [PATCH 10/16] Add MCP job cancellation and recovery Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/core/run_control.py | 131 +++ assert_ai/mcp/_command.py | 9 + assert_ai/mcp/models.py | 2 + assert_ai/mcp/server.py | 13 +- assert_ai/mcp/tools/jobs.py | 50 + assert_ai/runner.py | 313 +++++- assert_ai/services/_evaluation_worker.py | 303 +++++- assert_ai/services/evaluations.py | 1127 ++++++++++++++++++++-- assert_ai/services/job_models.py | 6 + assert_ai/services/job_store.py | 263 ++++- assert_ai/services/run_planning.py | 36 +- assert_ai/stages/inference.py | 13 + assert_ai/stages/judge.py | 18 +- assert_ai/stages/systematize.py | 7 + assert_ai/stages/test_set.py | 18 + tests/test_evaluation_service.py | 592 +++++++++++- tests/test_job_store.py | 203 ++++ tests/test_mcp_cli.py | 3 + tests/test_mcp_server.py | 168 +++- tests/test_run_result.py | 90 ++ 20 files changed, 3238 insertions(+), 127 deletions(-) create mode 100644 assert_ai/core/run_control.py diff --git a/assert_ai/core/run_control.py b/assert_ai/core/run_control.py new file mode 100644 index 000000000..b4fe01ac6 --- /dev/null +++ b/assert_ai/core/run_control.py @@ -0,0 +1,131 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Cooperative pipeline cancellation and transport-neutral run events.""" + +from __future__ import annotations + +import logging +import threading +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Protocol + +log = logging.getLogger(__name__) + + +class RunCancelled(RuntimeError): + """Raised at a safe checkpoint after cancellation is requested.""" + + def __init__(self, *, stage: str | None = None) -> None: + super().__init__("Evaluation cancellation requested") + self.stage = stage + + +@dataclass(slots=True) +class RunControl: + """Cancellation token checked between safe units of pipeline work.""" + + cancel_requested: Callable[[], bool] + cancel_acknowledged: Callable[[str | None], None] | None = None + _acknowledged: bool = field(default=False, init=False, repr=False) + _acknowledge_lock: threading.Lock = field( + default_factory=threading.Lock, + init=False, + repr=False, + ) + + @classmethod + def from_marker( + cls, + marker: str | Path, + *, + cancel_acknowledged: Callable[[str | None], None] | None = None, + ) -> "RunControl": + path = Path(marker) + return cls( + cancel_requested=path.is_file, + cancel_acknowledged=cancel_acknowledged, + ) + + def raise_if_cancelled(self, *, stage: str | None = None) -> None: + if self.cancel_requested(): + self._acknowledge(stage) + raise RunCancelled(stage=stage) + + def _acknowledge(self, stage: str | None) -> None: + callback = self.cancel_acknowledged + if callback is None: + return + with self._acknowledge_lock: + if self._acknowledged: + return + self._acknowledged = True + try: + callback(stage) + except Exception: # noqa: BLE001 - acknowledgement is diagnostic + log.warning( + "Could not acknowledge evaluation cancellation", + exc_info=True, + ) + + +@dataclass(frozen=True, slots=True) +class PipelineStarted: + suite_id: str | None + run_id: str | None + stages: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class StagePlanned: + name: str + scope: str + action: str + + +@dataclass(frozen=True, slots=True) +class StageStarted: + name: str + scope: str + + +@dataclass(frozen=True, slots=True) +class StageProgress: + name: str + values: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True) +class StageFinished: + name: str + scope: str + state: str + duration_seconds: float | None = None + summary: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True) +class PipelineFinished: + state: str + exit_code: int + failed_stage: str | None = None + error_code: str | None = None + error_message: str | None = None + + +class RunObserver(Protocol): + """Receives lifecycle events without depending on a transport.""" + + def pipeline_started(self, event: PipelineStarted) -> None: ... + + def stage_planned(self, event: StagePlanned) -> None: ... + + def stage_started(self, event: StageStarted) -> None: ... + + def stage_progress(self, event: StageProgress) -> None: ... + + def stage_finished(self, event: StageFinished) -> None: ... + + def pipeline_finished(self, event: PipelineFinished) -> None: ... diff --git a/assert_ai/mcp/_command.py b/assert_ai/mcp/_command.py index 53c2df2aa..8bf6d67fa 100644 --- a/assert_ai/mcp/_command.py +++ b/assert_ai/mcp/_command.py @@ -144,6 +144,13 @@ def mcp() -> None: show_default=True, help="Maximum retained bytes for each worker stdout/stderr log.", ) +@click.option( + "--cancellation-grace-seconds", + type=click.FloatRange(min=0.1), + default=10.0, + show_default=True, + help="Grace period before an unresponsive worker tree is terminated.", +) @click.option( "--max-prompt-sample-size", type=click.IntRange(min=1), @@ -192,6 +199,7 @@ def serve( max_active_jobs: int, max_queued_jobs: int, max_job_log_bytes: int, + cancellation_grace_seconds: float, max_prompt_sample_size: int, max_scenario_sample_size: int, allowed_model_patterns: tuple[str, ...], @@ -226,6 +234,7 @@ def serve( max_active_jobs=max_active_jobs, max_queued_jobs=max_queued_jobs, max_job_log_bytes=max_job_log_bytes, + cancellation_grace_seconds=cancellation_grace_seconds, max_prompt_sample_size=max_prompt_sample_size, max_scenario_sample_size=max_scenario_sample_size, allowed_model_patterns=allowed_model_patterns, diff --git a/assert_ai/mcp/models.py b/assert_ai/mcp/models.py index 21c9f1b32..321dc8895 100644 --- a/assert_ai/mcp/models.py +++ b/assert_ai/mcp/models.py @@ -64,6 +64,7 @@ class ServerLimits(BaseModel): max_active_jobs: int max_queued_jobs: int max_job_log_bytes: int + cancellation_grace_seconds: float max_prompt_sample_size: int max_scenario_sample_size: int model_allowlist_enabled: bool = False @@ -100,6 +101,7 @@ class ServerInfo(BaseModel): protocol_notes: list[str] = Field( default_factory=lambda: [ "Long evaluations use ASSERT job polling rather than one blocking MCP request.", + "Cancellation is cooperative before identity-checked process-tree termination.", "Resource and artifact identifiers are opaque and contain no host paths.", ] ) diff --git a/assert_ai/mcp/server.py b/assert_ai/mcp/server.py index 51a4b9ca7..34ae26582 100644 --- a/assert_ai/mcp/server.py +++ b/assert_ai/mcp/server.py @@ -91,6 +91,7 @@ class ServerOptions: max_active_jobs: int = 1 max_queued_jobs: int = 100 max_job_log_bytes: int = 1024 * 1024 + cancellation_grace_seconds: float = 10.0 max_prompt_sample_size: int = 100_000 max_scenario_sample_size: int = 100_000 allowed_model_patterns: tuple[str, ...] = () @@ -119,6 +120,10 @@ def __post_init__(self) -> None: raise ValueError( "max_job_log_bytes must be between 4096 and 16777216" ) + if self.cancellation_grace_seconds <= 0: + raise ValueError( + "cancellation_grace_seconds must be positive" + ) if self.max_prompt_sample_size < 1: raise ValueError("max_prompt_sample_size must be positive") if self.max_scenario_sample_size < 1: @@ -163,6 +168,7 @@ def create( max_active_jobs: int = 1, max_queued_jobs: int = 100, max_job_log_bytes: int = 1024 * 1024, + cancellation_grace_seconds: float = 10.0, max_prompt_sample_size: int = 100_000, max_scenario_sample_size: int = 100_000, allowed_model_patterns: Iterable[str] = (), @@ -191,6 +197,7 @@ def create( max_active_jobs=max_active_jobs, max_queued_jobs=max_queued_jobs, max_job_log_bytes=max_job_log_bytes, + cancellation_grace_seconds=cancellation_grace_seconds, max_prompt_sample_size=max_prompt_sample_size, max_scenario_sample_size=max_scenario_sample_size, allowed_model_patterns=tuple(allowed_model_patterns), @@ -260,6 +267,7 @@ def build_server(options: ServerOptions) -> MCPServer: max_active_jobs=options.max_active_jobs, max_log_bytes=options.max_job_log_bytes, launch_enabled=execution_enabled, + cancellation_grace_seconds=options.cancellation_grace_seconds, ) evaluations = EvaluationService( options.workspace, @@ -310,6 +318,9 @@ def get_server_info() -> ServerInfo: max_active_jobs=options.max_active_jobs, max_queued_jobs=options.max_queued_jobs, max_job_log_bytes=options.max_job_log_bytes, + cancellation_grace_seconds=( + options.cancellation_grace_seconds + ), max_prompt_sample_size=options.max_prompt_sample_size, max_scenario_sample_size=options.max_scenario_sample_size, model_allowlist_enabled=bool(options.allowed_model_patterns), @@ -385,7 +396,7 @@ def get_server_info() -> ServerInfo: ) if execution_enabled: register_job_execute_tools(server, job_services) - job_manager.enqueue() + job_manager.start() return server diff --git a/assert_ai/mcp/tools/jobs.py b/assert_ai/mcp/tools/jobs.py index 4e3a83d62..eb504d9f6 100644 --- a/assert_ai/mcp/tools/jobs.py +++ b/assert_ai/mcp/tools/jobs.py @@ -34,6 +34,12 @@ idempotent_hint=True, open_world_hint=True, ) +_CANCEL_ANNOTATIONS = ToolAnnotations( + read_only_hint=False, + destructive_hint=True, + idempotent_hint=True, + open_world_hint=False, +) @dataclass(frozen=True, slots=True) @@ -130,3 +136,47 @@ def start_evaluation( return JobStartResult.model_validate( sanitize_for_mcp(started, workspace=services.workspace) ) + + @server.tool( + title="Cancel an ASSERT evaluation", + annotations=_CANCEL_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def cancel_job(job_id: str) -> JobDetail: + """Request cooperative cancellation with process-tree escalation.""" + job = invoke_tool( + lambda: services.evaluations.cancel(job_id), + workspace=services.workspace, + ) + return JobDetail.model_validate( + sanitize_for_mcp(job, workspace=services.workspace) + ) + + @server.tool( + title="Retry an ASSERT evaluation", + annotations=_START_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def retry_job( + job_id: str, + request_id: str, + ) -> JobStartResult: + """Retry a terminal job from its earliest unsafe stage.""" + started = invoke_tool( + lambda: services.evaluations.retry( + job_id, + request_id=request_id, + ), + workspace=services.workspace, + ) + return JobStartResult.model_validate( + sanitize_for_mcp(started, workspace=services.workspace) + ) diff --git a/assert_ai/runner.py b/assert_ai/runner.py index 9eb46719e..5edfe6819 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -48,6 +48,17 @@ UsageAccumulator, track_usage, ) +from assert_ai.core.run_control import ( + PipelineFinished, + PipelineStarted, + RunCancelled, + RunControl, + RunObserver, + StageFinished, + StagePlanned, + StageProgress, + StageStarted, +) from assert_ai.core.runtime_safety import ( ManifestHeartbeat, PipelineWatchdog, @@ -68,6 +79,7 @@ from assert_ai.core.runtime_path_policy import RuntimePathPolicy log = logging.getLogger(__name__) +_OBSERVER_PROGRESS_INTERVAL_SECONDS = 0.5 def _set_nested(raw: dict[str, Any], path: list[str], value: Any) -> None: @@ -699,6 +711,8 @@ def run_pipeline( overrides: list[str] | None = None, concurrency: int | None = None, path_policy: RuntimePathPolicy | None = None, + control: RunControl | None = None, + observer: RunObserver | None = None, ) -> int: """Execute configured stages and return the legacy process exit code.""" return run_pipeline_result( @@ -708,6 +722,8 @@ def run_pipeline( overrides=overrides, concurrency=concurrency, path_policy=path_policy, + control=control, + observer=observer, ).exit_code @@ -719,6 +735,8 @@ def run_pipeline_result( overrides: list[str] | None = None, concurrency: int | None = None, path_policy: RuntimePathPolicy | None = None, + control: RunControl | None = None, + observer: RunObserver | None = None, ) -> RunResult: """Execute configured stages and return a structured terminal outcome.""" try: @@ -729,7 +747,25 @@ def run_pipeline_result( overrides=overrides, concurrency=concurrency, path_policy=path_policy, + control=control, + observer=observer, + ) + except RunCancelled as exc: + result = RunResult( + state=RunState.CANCELLED, + exit_code=130, + failed_stage=exc.stage, ) + _notify_observer( + observer, + "pipeline_finished", + PipelineFinished( + state=result.state.value, + exit_code=result.exit_code, + failed_stage=result.failed_stage, + ), + ) + return result except Exception: # noqa: BLE001 log.error("[runner] Unexpected pipeline setup error", exc_info=True) return RunResult( @@ -748,6 +784,8 @@ def run_pipeline_document_result( strict: bool = False, concurrency: int | None = None, path_policy: RuntimePathPolicy | None = None, + control: RunControl | None = None, + observer: RunObserver | None = None, ) -> RunResult: """Execute an immutable config document using its original path as a base.""" try: @@ -758,7 +796,25 @@ def run_pipeline_document_result( concurrency=concurrency, path_policy=path_policy, config_document=document, + control=control, + observer=observer, + ) + except RunCancelled as exc: + result = RunResult( + state=RunState.CANCELLED, + exit_code=130, + failed_stage=exc.stage, ) + _notify_observer( + observer, + "pipeline_finished", + PipelineFinished( + state=result.state.value, + exit_code=result.exit_code, + failed_stage=result.failed_stage, + ), + ) + return result except Exception: # noqa: BLE001 log.error("[runner] Unexpected pipeline setup error", exc_info=True) return RunResult( @@ -778,6 +834,8 @@ def _run_pipeline_result( concurrency: int | None = None, path_policy: RuntimePathPolicy | None = None, config_document: dict[str, Any] | None = None, + control: RunControl | None = None, + observer: RunObserver | None = None, ) -> RunResult: """Execute configured stages. @@ -823,6 +881,17 @@ def _run_pipeline_result( error_message=_result_error_message(str(exc), path_policy=path_policy), ) + ctx["_run_control"] = control + _notify_observer( + observer, + "pipeline_started", + PipelineStarted( + suite_id=ctx.get("suite_id"), + run_id=ctx.get("run_id"), + stages=tuple(name for name, _ in ctx["stages"]), + ), + ) + # CLI --concurrency wins over the YAML-resolved value so a single run can be # widened or narrowed without editing the config. We mutate the live # InferenceConfig instance (it's a regular dataclass, not frozen) because @@ -879,11 +948,19 @@ def _run_pipeline_result( stages_to_run: list[tuple[str, Any, dict[str, Any]]] = [] for stage_name, raw_cfg in ctx["stages"]: + module = STAGES[stage_name] if not raw_cfg.get("enabled", True): + _notify_observer( + observer, + "stage_planned", + StagePlanned( + name=stage_name, + scope=module.SCOPE, + action="disabled", + ), + ) continue - module = STAGES[stage_name] - if module.SCOPE == "suite": if cache_supported and is_cacheable_stage(stage_name): forced = stage_name in requested_force_stages @@ -904,6 +981,15 @@ def _run_pipeline_result( f"[{stage_name}] Reused artifact {plan.version} " f"(input hashes match, use --force-stage {stage_name} to regenerate)" ) + _notify_observer( + observer, + "stage_planned", + StagePlanned( + name=stage_name, + scope=module.SCOPE, + action="reused", + ), + ) continue # Force cacheable stages to write into their versioned artifact # directory regardless of any save_dir/save_path the user set @@ -921,8 +1007,26 @@ def _run_pipeline_result( log.info( f"[{stage_name}] Skipped (output exists, use --force-stage {stage_name} to regenerate)" ) + _notify_observer( + observer, + "stage_planned", + StagePlanned( + name=stage_name, + scope=module.SCOPE, + action="skipped", + ), + ) continue + _notify_observer( + observer, + "stage_planned", + StagePlanned( + name=stage_name, + scope=module.SCOPE, + action="pending", + ), + ) stages_to_run.append((stage_name, module, raw_cfg)) run_root = Path(ctx["run_root"]) if ctx.get("run_root") else None @@ -1005,10 +1109,12 @@ def _run_pipeline_result( check_interval_s=60.0, ) heartbeat.attach_watchdog(watchdog) - ctx["_heartbeat"] = heartbeat ctx["_watchdog"] = watchdog heartbeat.start() watchdog.start() + progress = _RunProgress(heartbeat=heartbeat, observer=observer) + if heartbeat is not None or observer is not None: + ctx["_heartbeat"] = progress try: return _run_stages_inner( @@ -1023,6 +1129,9 @@ def _run_pipeline_result( stage_usage=stage_usage, heartbeat=heartbeat, watchdog=watchdog, + control=control, + observer=observer, + progress=progress, ) finally: if heartbeat is not None: @@ -1044,12 +1153,16 @@ def _run_stages_inner( stage_usage: dict[str, dict[str, Any]], heartbeat: ManifestHeartbeat | None, watchdog: PipelineWatchdog | None, + control: RunControl | None, + observer: RunObserver | None, + progress: "_RunProgress", ) -> RunResult: """Stage execution loop. Extracted so the outer function can manage heartbeat/watchdog lifecycle in a single try/finally.""" failed_stage: str | None = None failed_error_code: str | None = None failed_error_message: str | None = None + cancelled = False for stage_name, module, raw_cfg in stages_to_run: if manifest is not None and module.SCOPE == "run": @@ -1069,6 +1182,11 @@ def _run_stages_inner( _print_stage_start(stage_name, ctx, raw_cfg) stage_start = time.monotonic() stage_result: dict[str, Any] = {} + _notify_observer( + observer, + "stage_started", + StageStarted(name=stage_name, scope=module.SCOPE), + ) # Tick the watchdog so it doesn't fire mid-stage on the previous # stage's idle clock, and reset the heartbeat's progress payload # so a stage that doesn't report progress (e.g. systematize, test_set) @@ -1076,9 +1194,8 @@ def _run_stages_inner( # in the manifest. if watchdog is not None: watchdog.tick() - if heartbeat is not None: - heartbeat.clear_progress() - heartbeat.set_progress(stage=stage_name) + progress.clear_progress() + progress.set_progress(stage=stage_name) # Pass the per-stage "was this forced" flag through ctx so stages # like inference/judge can distinguish a real cache-mismatch warning # from a redundant one (the user already opted into discarding via @@ -1087,6 +1204,8 @@ def _run_stages_inner( ctx["_stage_forced"] = stage_name in requested_force_stages usage_acc: UsageAccumulator | None = None try: + if control is not None: + control.raise_if_cancelled(stage=stage_name) with track_usage() as usage_acc: # run_stage_coro replaces asyncio.run with bounded teardown: # if the stage's event loop can't shut down its default @@ -1100,9 +1219,13 @@ def _run_stages_inner( module.run(ctx, raw_cfg), cleanup_timeout_s=300.0, ) or {} + if control is not None: + control.raise_if_cancelled(stage=stage_name) stage_summary = (stage_result or {}).get("_summary") if isinstance(stage_summary, dict): ctx["_stage_summaries"][stage_name] = stage_summary + if control is not None: + control.raise_if_cancelled(stage=stage_name) _refresh_stage_indexes(ctx, stage_name, stage_result) stage_errored_count = int( (stage_summary or {}).get("errored_count", 0) or 0 @@ -1133,8 +1256,15 @@ def _run_stages_inner( ) ctx["_suite_summary_blocked"] = True else: + if control is not None: + control.raise_if_cancelled(stage=stage_name) finalize_artifact_plan(ctx, artifact_plans[stage_name]) ok = True + except RunCancelled: + ok = False + cancelled = True + stage_error_code = None + stage_error_message = None except (LLMAuthError, LLMInputError, LLMRateLimitError, LLMProviderError) as exc: # Classified LLM errors already carry a clean, actionable message. # Print just that message; suppress the multi-screen litellm/httpx @@ -1184,12 +1314,19 @@ def _run_stages_inner( # that owned it. if watchdog is not None: watchdog.tick() - if heartbeat is not None: - heartbeat.clear_progress() + progress.clear_progress() if manifest is not None and module.SCOPE == "run": - manifest.stages[stage_name] = "completed" if ok else "failed" - manifest.status = "running" if ok else "failed" + manifest.stages[stage_name] = ( + "completed" + if ok + else ("cancelled" if cancelled else "failed") + ) + manifest.status = ( + "running" + if ok + else ("cancelled" if cancelled else "failed") + ) existing_timing = manifest.stage_timings.get(stage_name) or {} existing_timing["ended_at"] = datetime.now(timezone.utc).isoformat() existing_timing["duration_secs"] = round(elapsed, 3) @@ -1208,6 +1345,26 @@ def _run_stages_inner( _write_suite_metadata(ctx) _refresh_suite_summary(ctx, rebuild_indexes=True) + _notify_observer( + observer, + "stage_finished", + StageFinished( + name=stage_name, + scope=module.SCOPE, + state=( + "completed" + if ok + else ("cancelled" if cancelled else "failed") + ), + duration_seconds=round(elapsed, 3), + summary=( + dict(stage_result.get("_summary") or {}) + if isinstance(stage_result, dict) + else {} + ), + ), + ) + if not ok: failed_stage = stage_name failed_error_code = stage_error_code @@ -1235,7 +1392,9 @@ def _run_stages_inner( except Exception: # noqa: BLE001 log.debug("Failed to write metrics.json", exc_info=True) - if failed_stage is None: + if cancelled: + log.info(f"Pipeline cancelled ({total_elapsed:.1f}s)") + elif failed_stage is None: log.info(f"Pipeline completed ({total_elapsed:.1f}s)") if run_root is not None: _log_run_headline(run_root) @@ -1259,7 +1418,11 @@ def _run_stages_inner( if manifest is not None: manifest.ended_at = datetime.now(timezone.utc).isoformat() - manifest.status = "completed" if failed_stage is None else "failed" + manifest.status = ( + "cancelled" + if cancelled + else ("completed" if failed_stage is None else "failed") + ) _record_run_artifacts(manifest, ctx, run_root) _write_active_manifest(manifest, run_root, heartbeat) _refresh_run_summary( @@ -1267,24 +1430,128 @@ def _run_stages_inner( manifest, stage_usage=stage_usage, elapsed_s=total_elapsed, - rebuild_indexes=failed_stage is None, + rebuild_indexes=failed_stage is None and not cancelled, ) - _refresh_suite_summary(ctx, rebuild_indexes=failed_stage is None) + _refresh_suite_summary( + ctx, + rebuild_indexes=failed_stage is None and not cancelled, + ) - if failed_stage is None: - return _run_result_from_context( + if cancelled: + result = _run_result_from_context( + ctx, + state=RunState.CANCELLED, + exit_code=130, + failed_stage=failed_stage, + ) + elif failed_stage is None: + result = _run_result_from_context( ctx, state=RunState.COMPLETED, exit_code=0, ) - return _run_result_from_context( - ctx, - state=RunState.FAILED, - exit_code=1, - failed_stage=failed_stage, - error_code=failed_error_code or "RUN_FAILED", - error_message=failed_error_message or f"Pipeline failed at {failed_stage}", + else: + result = _run_result_from_context( + ctx, + state=RunState.FAILED, + exit_code=1, + failed_stage=failed_stage, + error_code=failed_error_code or "RUN_FAILED", + error_message=failed_error_message or f"Pipeline failed at {failed_stage}", + ) + _notify_observer( + observer, + "pipeline_finished", + PipelineFinished( + state=result.state.value, + exit_code=result.exit_code, + failed_stage=result.failed_stage, + error_code=result.error_code, + error_message=result.error_message, + ), ) + return result + + +class _RunProgress: + """Fan progress updates out to manifests and run observers.""" + + def __init__( + self, + *, + heartbeat: ManifestHeartbeat | None, + observer: RunObserver | None, + ) -> None: + self._heartbeat = heartbeat + self._observer = observer + self._last_observer_update = 0.0 + self._pending_observer_event: StageProgress | None = None + + def set_progress(self, **fields: Any) -> None: + if self._heartbeat is not None: + self._heartbeat.set_progress(**fields) + stage = fields.get("stage") + if isinstance(stage, str) and stage: + event = StageProgress(name=stage, values=dict(fields)) + now = time.monotonic() + completed = fields.get("completed") + total = fields.get("total") + terminal_update = ( + isinstance(completed, int) + and not isinstance(completed, bool) + and isinstance(total, int) + and not isinstance(total, bool) + and completed >= total + ) + if ( + self._last_observer_update == 0.0 + or terminal_update + or now - self._last_observer_update + >= _OBSERVER_PROGRESS_INTERVAL_SECONDS + ): + self._emit_observer_progress(event, now=now) + else: + self._pending_observer_event = event + + def clear_progress(self) -> None: + if self._pending_observer_event is not None: + self._emit_observer_progress( + self._pending_observer_event, + now=time.monotonic(), + ) + if self._heartbeat is not None: + self._heartbeat.clear_progress() + + def _emit_observer_progress( + self, + event: StageProgress, + *, + now: float, + ) -> None: + _notify_observer( + self._observer, + "stage_progress", + event, + ) + self._last_observer_update = now + self._pending_observer_event = None + + +def _notify_observer( + observer: RunObserver | None, + method_name: str, + event: Any, +) -> None: + if observer is None: + return + try: + getattr(observer, method_name)(event) + except Exception: # noqa: BLE001 - diagnostics must not fail a run + log.warning( + "Run observer failed while handling %s", + method_name, + exc_info=True, + ) def _run_result_from_context( diff --git a/assert_ai/services/_evaluation_worker.py b/assert_ai/services/_evaluation_worker.py index d3c6fd911..ce007a41c 100644 --- a/assert_ai/services/_evaluation_worker.py +++ b/assert_ai/services/_evaluation_worker.py @@ -15,14 +15,33 @@ import sys import threading from collections.abc import Iterator +from dataclasses import asdict +from datetime import datetime, timezone from pathlib import Path from typing import Any import yaml +from assert_ai.core.config_document import PIPELINE_STAGE_ORDER from assert_ai.core.io import write_json -from assert_ai.core.security import redact_path_prefixes, sanitize_text +from assert_ai.core.run_control import ( + PipelineFinished, + PipelineStarted, + RunCancelled, + RunControl, + StageFinished, + StagePlanned, + StageProgress, + StageStarted, +) +from assert_ai.core.run_result import RunResult, RunState +from assert_ai.core.security import ( + redact_path_prefixes, + sanitize_payload, + sanitize_text, +) from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.job_store import JobStore _JOB_ID_RE = re.compile(r"^[0-9a-f]{32}$") _SHA256_RE = re.compile(r"^sha256:[0-9a-f]{64}$") @@ -123,22 +142,72 @@ def main(argv: list[str] | None = None) -> int: expected_root=job_dir, reject_links=True, ) + cancel_path = workspace.path_policy.resolve_managed_output( + job_dir / "cancel.requested", + field_name="evaluation cancellation marker", + expected_root=job_dir, + reject_links=True, + ) + cancel_acknowledged_path = ( + workspace.path_policy.resolve_managed_output( + job_dir / "cancel.acknowledged", + field_name="evaluation cancellation acknowledgement", + expected_root=job_dir, + reject_links=True, + ) + ) + store = JobStore( + jobs_root.parent / "jobs.sqlite3", + path_policy=workspace.path_policy, + expected_root=workspace.artifacts_root, + ) + observer = _JobRunObserver( + store=store, + job_id=args.job_id, + workspace=workspace, + ) + control = RunControl.from_marker( + cancel_path, + cancel_acknowledged=lambda stage: ( + _acknowledge_cancellation( + cancel_acknowledged_path, + store=store, + job_id=args.job_id, + stage=stage, + ) + ), + ) with ( _BoundedTextLog(stdout_path, max_bytes=max_log_bytes) as stdout, _BoundedTextLog(stderr_path, max_bytes=max_log_bytes) as stderr, contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr), _capture_worker_logs(stderr), + observer, ): - from assert_ai.runner import run_pipeline_document_result - - result = run_pipeline_document_result( - document=document, - config_path=str(config_path), - force_stages=force_stages, - strict=strict, - path_policy=workspace.path_policy, - ) + try: + control.raise_if_cancelled( + stage=_first_enabled_stage(document) + ) + except RunCancelled as cancelled: + result = _cancelled_before_runner( + document, + workspace=workspace, + observer=observer, + failed_stage=cancelled.stage, + ) + else: + from assert_ai.runner import run_pipeline_document_result + + result = run_pipeline_document_result( + document=document, + config_path=str(config_path), + force_stages=force_stages, + strict=strict, + path_policy=workspace.path_policy, + control=control, + observer=observer, + ) payload = { "schema_version": 1, "job_id": args.job_id, @@ -208,6 +277,128 @@ def _required_string(payload: dict[str, Any], key: str) -> str: return value +def _acknowledge_cancellation( + path: Path, + *, + store: JobStore, + job_id: str, + stage: str | None, +) -> None: + acknowledged_at = datetime.now(timezone.utc).isoformat() + write_json( + path, + { + "schema_version": 1, + "job_id": job_id, + "acknowledged_at": acknowledged_at, + "stage": stage, + }, + ) + store.append_event( + job_id, + "cancel_observed", + { + "state": "cancelling", + "stage": stage, + "acknowledged_at": acknowledged_at, + }, + ) + + +def _first_enabled_stage(document: dict[str, Any]) -> str | None: + pipeline = document.get("pipeline") + if not isinstance(pipeline, dict): + return None + for stage_name in PIPELINE_STAGE_ORDER: + stage = pipeline.get(stage_name) + if isinstance(stage, dict) and stage.get("enabled", True): + return stage_name + return None + + +def _cancelled_before_runner( + document: dict[str, Any], + *, + workspace: WorkspaceService, + observer: "_JobRunObserver", + failed_stage: str | None, +) -> RunResult: + suite_id = document.get("suite") + run_id = document.get("run") + suite_id = suite_id if isinstance(suite_id, str) else None + run_id = run_id if isinstance(run_id, str) else None + suite_root = None + if suite_id is not None: + suite_root = workspace.path_policy.resolve_managed_output( + workspace.results_root / suite_id, + field_name="cancelled evaluation suite root", + expected_root=workspace.results_root, + reject_links=True, + ) + run_root = ( + workspace.path_policy.resolve_managed_output( + suite_root / run_id, + field_name="cancelled evaluation run root", + expected_root=suite_root, + reject_links=True, + ) + if suite_root is not None and run_id is not None + else None + ) + stages = tuple( + stage_name + for stage_name in PIPELINE_STAGE_ORDER + if isinstance(document.get("pipeline"), dict) + and isinstance(document["pipeline"].get(stage_name), dict) + and document["pipeline"][stage_name].get("enabled", True) + ) + observer.pipeline_started( + PipelineStarted( + suite_id=suite_id, + run_id=run_id, + stages=stages, + ) + ) + if failed_stage is not None: + scope = ( + "suite" + if failed_stage in {"systematize", "test_set"} + else "run" + ) + observer.stage_planned( + StagePlanned( + name=failed_stage, + scope=scope, + action="pending", + ) + ) + observer.stage_finished( + StageFinished( + name=failed_stage, + scope=scope, + state="cancelled", + duration_seconds=0.0, + ) + ) + result = RunResult( + state=RunState.CANCELLED, + exit_code=130, + suite_id=suite_id, + run_id=run_id, + suite_root=suite_root, + run_root=run_root, + failed_stage=failed_stage, + ) + observer.pipeline_finished( + PipelineFinished( + state=result.state.value, + exit_code=result.exit_code, + failed_stage=result.failed_stage, + ) + ) + return result + + def _read_bytes(path: Path, *, max_bytes: int, label: str) -> bytes: with path.open("rb") as stream: value = stream.read(max_bytes + 1) @@ -328,6 +519,98 @@ def fileno(self) -> int: return self._text_log.fileno() +class _JobRunObserver: + """Persist runner lifecycle events and a suite-only heartbeat.""" + + def __init__( + self, + *, + store: JobStore, + job_id: str, + workspace: WorkspaceService, + heartbeat_seconds: float = 15.0, + ) -> None: + self._store = store + self._job_id = job_id + self._workspace = workspace + self._heartbeat_seconds = heartbeat_seconds + self._stop = threading.Event() + self._thread: threading.Thread | None = None + + def __enter__(self) -> "_JobRunObserver": + self._thread = threading.Thread( + target=self._heartbeat_loop, + name="assert-mcp-job-heartbeat", + daemon=True, + ) + self._thread.start() + return self + + def __exit__(self, *_: object) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=5) + self._thread = None + + def pipeline_started(self, event: PipelineStarted) -> None: + self._append("pipeline_started", event) + + def stage_planned(self, event: StagePlanned) -> None: + self._append("stage_planned", event) + + def stage_started(self, event: StageStarted) -> None: + self._append("stage_started", event) + + def stage_progress(self, event: StageProgress) -> None: + self._append("stage_progress", event) + + def stage_finished(self, event: StageFinished) -> None: + self._append("stage_finished", event) + + def pipeline_finished(self, event: PipelineFinished) -> None: + self._append("pipeline_finished", event) + + def _heartbeat_loop(self) -> None: + while not self._stop.wait(self._heartbeat_seconds): + self._append("heartbeat", {"state": "running"}) + + def _append(self, event_type: str, event: object) -> None: + try: + raw = asdict(event) if hasattr(event, "__dataclass_fields__") else event + payload = sanitize_payload(raw) + if not isinstance(payload, dict): + payload = {} + payload = _redact_payload_paths(payload, self._workspace) + self._store.append_event( + self._job_id, + event_type, + payload, + ) + except Exception: # noqa: BLE001 - diagnostic boundary + logging.getLogger(__name__).exception( + "Could not persist evaluation job event %s", + event_type, + ) + + +def _redact_payload_paths( + payload: dict[str, Any], + workspace: WorkspaceService, +) -> dict[str, Any]: + serialized = json.dumps(payload, ensure_ascii=False) + redacted = redact_path_prefixes( + serialized, + ( + workspace.root, + workspace.configs_root, + workspace.artifacts_root, + workspace.results_root, + ), + ) + parsed = json.loads(redacted) + return parsed if isinstance(parsed, dict) else {} + + @contextlib.contextmanager def _capture_worker_logs( stream: _BoundedTextLog, diff --git a/assert_ai/services/evaluations.py b/assert_ai/services/evaluations.py index ce9e43f1c..7620bfcdc 100644 --- a/assert_ai/services/evaluations.py +++ b/assert_ai/services/evaluations.py @@ -12,6 +12,7 @@ import logging import os import secrets +import signal import shutil import subprocess import sys @@ -24,7 +25,11 @@ from typing import Any, Sequence from urllib.parse import quote +import yaml + from assert_ai.core.io import write_json, write_text_atomic +from assert_ai.core.jsonl_index import JsonlIndexError, scan_jsonl +from assert_ai.core.config_document import PIPELINE_STAGE_ORDER from assert_ai.core.security import ( redact_path_prefixes, sanitize_payload, @@ -54,9 +59,15 @@ _CURSOR_VERSION = 1 _JOB_RESULT_MAX_BYTES = 1024 * 1024 +_JOB_SNAPSHOT_MAX_BYTES = 16 * 1024 * 1024 _JOB_ID_RETRIES = 5 _LEASE_SECONDS = 60.0 _LEASE_RENEW_SECONDS = 15.0 +_DEFAULT_CANCELLATION_GRACE_SECONDS = 10.0 +_PROCESS_EXIT_TIMEOUT_SECONDS = 5.0 +_RECOVERY_POLL_SECONDS = 0.25 +_MAX_RECOVERY_SLEEP_SECONDS = 30.0 +_CANCELLATION_POLL_SECONDS = 0.1 _REQUEST_ID_MAX_LENGTH = 200 _MIN_LOG_BYTES = 4096 _MAX_LOG_BYTES = 16 * 1024 * 1024 @@ -74,6 +85,9 @@ class EvaluationJobManager: max_log_bytes: int = 1024 * 1024 launch_enabled: bool = True lease_seconds: float = _LEASE_SECONDS + cancellation_grace_seconds: float = ( + _DEFAULT_CANCELLATION_GRACE_SECONDS + ) _owner: str = field( default_factory=lambda: uuid.uuid4().hex, init=False, @@ -98,6 +112,21 @@ class EvaluationJobManager: init=False, repr=False, ) + _monitored_jobs: set[str] = field( + default_factory=set, + init=False, + repr=False, + ) + _cancellation_jobs: set[str] = field( + default_factory=set, + init=False, + repr=False, + ) + _recovery: threading.Thread | None = field( + default=None, + init=False, + repr=False, + ) def __post_init__(self) -> None: if self.max_active_jobs < 1: @@ -109,6 +138,37 @@ def __post_init__(self) -> None: ) if self.lease_seconds <= 0: raise ValueError("lease_seconds must be positive") + if self.cancellation_grace_seconds <= 0: + raise ValueError( + "cancellation_grace_seconds must be positive" + ) + + def start(self) -> None: + """Recover persisted work and then schedule queued jobs.""" + if not self.launch_enabled: + return + with self._lock: + if self._recovery is not None and self._recovery.is_alive(): + return + self._recovery = threading.Thread( + target=self._recover_startup, + name="assert-mcp-job-recovery", + daemon=True, + ) + self._recovery.start() + + def cancel(self, job_id: str) -> JobRecord: + """Persist cancellation and enforce it outside the request thread.""" + if not self.launch_enabled: + raise ServiceError( + ServiceErrorCode.CAPABILITY_DISABLED, + "Evaluation execution is disabled for this service", + ) + record = self.store.request_cancel(job_id) + if record.state is JobState.CANCELLING: + self._write_cancel_marker(record) + self._ensure_cancellation(record) + return record def enqueue(self) -> None: """Wake a short-lived scheduler without holding an MCP request open.""" @@ -129,16 +189,55 @@ def reconcile(self, record: JobRecord) -> JobRecord: """Adopt a worker result or mark a dead worker interrupted.""" if record.state in TERMINAL_JOB_STATES: return record + if not self.launch_enabled: + return record if record.state is JobState.QUEUED: self.enqueue() return record + if ( + record.state is JobState.STARTING + and record.pid is None + and record.lease_owner == self._owner + and not _lease_expired(record.lease_expires_at) + ): + return record + adopted_worker = False + if record.lease_owner != self._owner: + if not _lease_expired(record.lease_expires_at): + return record + adopted = self.store.adopt_lease( + record.job_id, + lease_owner=self._owner, + lease_seconds=self.lease_seconds, + ) + if adopted is None: + return self.store.get(record.job_id) + record = adopted + adopted_worker = True + elif _lease_expired(record.lease_expires_at): + adopted = self.store.adopt_lease( + record.job_id, + lease_owner=self._owner, + lease_seconds=self.lease_seconds, + ) + if adopted is not None: + record = adopted + adopted_worker = True + process_alive = ( + record.pid is not None + and record.process_create_time is not None + and _process_matches( + record.pid, + record.process_create_time, + ) + ) result = self._read_result(record) - if result is not None: + if result is not None and not process_alive: try: return self._adopt_result( record, result, - lease_owner=None, + lease_owner=self._owner, ) except Exception as exc: # noqa: BLE001 - persisted boundary log.exception( @@ -148,20 +247,26 @@ def reconcile(self, record: JobRecord) -> JobRecord: return self._mark_internal_failure( record, exc, - lease_owner=None, + lease_owner=self._owner, ) - if record.state is JobState.STARTING and record.pid is None: - if not _lease_expired(record.lease_expires_at): - return record - if ( - record.pid is not None - and record.process_create_time is not None - and _process_matches( - record.pid, - record.process_create_time, - ) - ): + if process_alive: + if adopted_worker: + self._ensure_recovered_monitor(record) + if record.state is JobState.CANCELLING: + try: + self._write_cancel_marker(record) + self._ensure_cancellation(record) + except Exception: + log.exception( + "Could not enforce cancellation for evaluation job %s", + record.job_id, + ) return record + if record.state is JobState.CANCELLING: + return self._mark_cancelled_without_result( + record, + lease_owner=self._owner, + ) return self.store.mark_terminal( record.job_id, state=JobState.INTERRUPTED, @@ -173,6 +278,7 @@ def reconcile(self, record: JobRecord) -> JobRecord: ), result=None, run_root=record.run_root, + lease_owner=self._owner, ) def _schedule(self) -> None: @@ -180,6 +286,13 @@ def _schedule(self) -> None: while True: with self._lock: self._schedule_requested = False + try: + self._sweep_cancelling_jobs() + except Exception: # noqa: BLE001 - daemon boundary + log.exception( + "Evaluation scheduler could not sweep cancelling jobs" + ) + return try: claimed = self.store.claim_next( lease_owner=self._owner, @@ -214,6 +327,7 @@ def _schedule(self) -> None: continue with self._lock: self._processes[claimed.job_id] = process + self._monitored_jobs.add(claimed.job_id) monitor = threading.Thread( target=self._monitor, args=(claimed.job_id, process), @@ -232,6 +346,18 @@ def _schedule(self) -> None: if restart: self.enqueue() + def _sweep_cancelling_jobs(self) -> None: + for record in self.store.list_nonterminal_records(): + if record.state is not JobState.CANCELLING: + continue + try: + self.reconcile(record) + except Exception: # noqa: BLE001 - scheduler boundary + log.exception( + "Could not reconcile cancelling evaluation job %s", + record.job_id, + ) + def _launch(self, record: JobRecord) -> subprocess.Popen[bytes]: env = os.environ.copy() env["PYTHONUNBUFFERED"] = "1" @@ -311,21 +437,29 @@ def _monitor( ) return record = self.store.get(job_id) + if record.state in TERMINAL_JOB_STATES: + return payload = self._read_result(record) if payload is None: - self.store.mark_terminal( - job_id, - state=JobState.FAILED, - exit_code=process.returncode, - failed_stage=None, - error_code=ServiceErrorCode.RUN_FAILED.value, - error_message=( - "Evaluation worker exited without a valid result" - ), - result=None, - run_root=None, - lease_owner=self._owner, - ) + if record.state is JobState.CANCELLING: + self._mark_cancelled_without_result( + record, + lease_owner=self._owner, + ) + else: + self.store.mark_terminal( + job_id, + state=JobState.FAILED, + exit_code=process.returncode, + failed_stage=None, + error_code=ServiceErrorCode.RUN_FAILED.value, + error_message=( + "Evaluation worker exited without a valid result" + ), + result=None, + run_root=None, + lease_owner=self._owner, + ) else: self._adopt_result( record, @@ -352,8 +486,403 @@ def _monitor( finally: with self._lock: self._processes.pop(job_id, None) + self._monitored_jobs.discard(job_id) + self.enqueue() + + def _recover_startup(self) -> None: + try: + while self.launch_enabled: + next_lease_check: float | None = None + has_queued_job = False + try: + records = self.store.list_nonterminal_records() + except Exception: # noqa: BLE001 - daemon boundary + log.exception("Could not scan evaluation jobs for recovery") + return + for record in records: + if record.state is JobState.QUEUED: + has_queued_job = True + continue + try: + current = self.reconcile(record) + except Exception: # noqa: BLE001 - persisted boundary + log.exception( + "Could not recover evaluation job %s", + record.job_id, + ) + continue + if ( + current.state not in TERMINAL_JOB_STATES + and current.lease_owner != self._owner + ): + lease_wait = _lease_seconds_remaining( + current.lease_expires_at + ) + next_lease_check = ( + lease_wait + if next_lease_check is None + else min(next_lease_check, lease_wait) + ) + if has_queued_job: + self.enqueue() + if next_lease_check is None: + return + time.sleep( + min( + _MAX_RECOVERY_SLEEP_SECONDS, + max(_RECOVERY_POLL_SECONDS, next_lease_check), + ) + ) + finally: + with self._lock: + if self._recovery is threading.current_thread(): + self._recovery = None + + def _ensure_recovered_monitor(self, record: JobRecord) -> None: + if record.pid is None or record.process_create_time is None: + return + with self._lock: + if record.job_id in self._monitored_jobs: + return + self._monitored_jobs.add(record.job_id) + monitor = threading.Thread( + target=self._monitor_recovered, + args=( + record.job_id, + record.pid, + record.process_create_time, + ), + name=f"assert-mcp-recovered-{record.job_id[:8]}", + daemon=True, + ) + monitor.start() + + def _monitor_recovered( + self, + job_id: str, + pid: int, + process_create_time: float, + ) -> None: + poll_seconds = min( + _LEASE_RENEW_SECONDS, + max(0.05, self.lease_seconds / 3), + ) + try: + while _process_matches(pid, process_create_time): + time.sleep(poll_seconds) + if not self.store.renew_lease( + job_id, + lease_owner=self._owner, + lease_seconds=self.lease_seconds, + ): + return + record = self.store.get(job_id) + if record.state in TERMINAL_JOB_STATES: + return + payload = self._read_result(record) + if payload is not None: + self._adopt_result( + record, + payload, + lease_owner=self._owner, + ) + elif record.state is JobState.CANCELLING: + self._mark_cancelled_without_result( + record, + lease_owner=self._owner, + ) + else: + self.store.mark_terminal( + job_id, + state=JobState.INTERRUPTED, + exit_code=record.exit_code, + failed_stage=record.failed_stage, + error_code=ServiceErrorCode.JOB_INTERRUPTED.value, + error_message=( + "Evaluation worker exited without a terminal result" + ), + result=None, + run_root=record.run_root, + lease_owner=self._owner, + ) + except Exception: # noqa: BLE001 - daemon boundary + log.exception( + "Could not monitor recovered evaluation job %s", + job_id, + ) + finally: + with self._lock: + self._monitored_jobs.discard(job_id) + self.enqueue() + + def _ensure_cancellation(self, record: JobRecord) -> None: + with self._lock: + if record.job_id in self._cancellation_jobs: + return + self._cancellation_jobs.add(record.job_id) + thread = threading.Thread( + target=self._enforce_cancellation, + args=(record.job_id,), + name=f"assert-mcp-cancel-{record.job_id[:8]}", + daemon=True, + ) + thread.start() + + def _enforce_cancellation(self, job_id: str) -> None: + observation_deadline = ( + time.monotonic() + self.cancellation_grace_seconds + ) + teardown_deadline: float | None = None + try: + while True: + record = self.store.get(job_id) + if record.state in TERMINAL_JOB_STATES: + return + process_alive = ( + record.pid is not None + and record.process_create_time is not None + and _process_matches( + record.pid, + record.process_create_time, + ) + ) + if not process_alive: + current = self.reconcile(record) + if current.state in TERMINAL_JOB_STATES: + return + time.sleep( + ( + _CANCELLATION_POLL_SECONDS + if current.lease_owner == self._owner + else _cancellation_wait_seconds(current) + ) + ) + continue + + now = time.monotonic() + if ( + teardown_deadline is None + and self._cancellation_acknowledged(record) + ): + teardown_deadline = ( + now + _PROCESS_EXIT_TIMEOUT_SECONDS + ) + deadline = ( + teardown_deadline + if teardown_deadline is not None + else observation_deadline + ) + if now < deadline: + time.sleep( + min(_CANCELLATION_POLL_SECONDS, deadline - now) + ) + continue + + if record.lease_owner != self._owner: + if not _lease_expired(record.lease_expires_at): + time.sleep(_cancellation_wait_seconds(record)) + continue + adopted = self.store.adopt_lease( + record.job_id, + lease_owner=self._owner, + lease_seconds=self.lease_seconds, + ) + if adopted is None: + time.sleep(_CANCELLATION_POLL_SECONDS) + continue + record = adopted + if ( + record.pid is None + or record.process_create_time is None + or not _process_matches( + record.pid, + record.process_create_time, + ) + ): + continue + self.store.append_event( + job_id, + "termination_escalated", + {"state": JobState.CANCELLING.value}, + ) + _terminate_process_tree( + record.pid, + record.process_create_time, + timeout_seconds=_PROCESS_EXIT_TIMEOUT_SECONDS, + ) + settle_deadline = ( + time.monotonic() + _PROCESS_EXIT_TIMEOUT_SECONDS + ) + while time.monotonic() < settle_deadline: + current = self.store.get(job_id) + if current.state in TERMINAL_JOB_STATES: + return + if ( + current.pid is None + or current.process_create_time is None + or not _process_matches( + current.pid, + current.process_create_time, + ) + ): + payload = self._read_result(current) + if payload is not None: + self._adopt_result( + current, + payload, + lease_owner=self._owner, + ) + else: + self._mark_cancelled_without_result( + current, + lease_owner=self._owner, + ) + return + time.sleep(0.05) + current = self.store.get(job_id) + if ( + current.state not in TERMINAL_JOB_STATES + and current.pid is not None + and current.process_create_time is not None + and _process_matches( + current.pid, + current.process_create_time, + ) + ): + raise RuntimeError( + "Evaluation worker remained alive after termination" + ) + self.reconcile(current) + return + except Exception: # noqa: BLE001 - daemon boundary + log.exception( + "Could not enforce cancellation for evaluation job %s", + job_id, + ) + finally: + with self._lock: + self._cancellation_jobs.discard(job_id) self.enqueue() + def _cancellation_acknowledged(self, record: JobRecord) -> bool: + marker = self._job_file( + self._job_dir(record.job_id), + "cancel.acknowledged", + ) + return marker.is_file() + + def _write_cancel_marker(self, record: JobRecord) -> None: + job_dir = self._job_dir(record.job_id) + marker = self._job_file(job_dir, "cancel.requested") + write_text_atomic( + marker, + json.dumps( + { + "schema_version": 1, + "job_id": record.job_id, + "requested_at": record.cancel_requested_at, + }, + separators=(",", ":"), + sort_keys=True, + ) + + "\n", + ) + + def _mark_cancelled_without_result( + self, + record: JobRecord, + *, + lease_owner: str | None = None, + ) -> JobRecord: + failed_stage = record.failed_stage or _active_stage_from_events( + self.store.list_events(record.job_id, limit=1000) + ) + if failed_stage is None: + acknowledgement = _read_json_file( + self._job_file( + self._job_dir(record.job_id), + "cancel.acknowledged", + ), + max_bytes=_JOB_RESULT_MAX_BYTES, + ) + if isinstance(acknowledgement, dict): + acknowledged_stage = acknowledgement.get("stage") + if isinstance(acknowledged_stage, str): + failed_stage = acknowledged_stage + try: + run_root = self._write_cancelled_manifest( + record, + failed_stage=failed_stage, + ) + except Exception: # noqa: BLE001 - terminal persistence wins + log.exception( + "Could not write cancelled manifest for evaluation job %s", + record.job_id, + ) + run_root = record.run_root + result = { + "state": JobState.CANCELLED.value, + "exit_code": 130, + "failed_stage": failed_stage, + "error_code": None, + "error_message": ( + "Evaluation stopped after cancellation was requested" + ), + } + return self.store.mark_terminal( + record.job_id, + state=JobState.CANCELLED, + exit_code=130, + failed_stage=failed_stage, + error_code=None, + error_message=result["error_message"], + result=result, + run_root=run_root, + lease_owner=lease_owner, + ) + + def _write_cancelled_manifest( + self, + record: JobRecord, + *, + failed_stage: str | None, + ) -> str | None: + if record.run_id is None: + return record.run_root + suite_root = self.workspace.path_policy.resolve_managed_output( + self.workspace.results_root / record.suite_id, + field_name="cancelled job suite root", + expected_root=self.workspace.results_root, + reject_links=True, + ) + run_root = self.workspace.path_policy.resolve_managed_output( + suite_root / record.run_id, + field_name="cancelled job run root", + expected_root=suite_root, + reject_links=True, + ) + manifest_path = self.workspace.path_policy.resolve_managed_output( + run_root / "manifest.json", + field_name="cancelled job manifest", + expected_root=run_root, + reject_links=True, + ) + manifest = _read_json_file( + manifest_path, + max_bytes=_JOB_RESULT_MAX_BYTES, + ) + if not isinstance(manifest, dict): + return str(run_root) if run_root.is_dir() else record.run_root + manifest["status"] = "cancelled" + manifest["ended_at"] = datetime.now(timezone.utc).isoformat() + manifest["heartbeat_at"] = manifest["ended_at"] + stages = manifest.get("stages") + if isinstance(stages, dict) and failed_stage is not None: + if stages.get(failed_stage) == "running": + stages[failed_stage] = "cancelled" + write_json(manifest_path, manifest) + return str(run_root) + def _mark_internal_failure( self, record: JobRecord, @@ -493,11 +1022,12 @@ def _adopt_result( ServiceErrorCode.RUN_FAILED, "Evaluation worker returned a mismatched run id", ) - may_omit_identity = ( - state is JobState.FAILED - and failed_stage is None - and result_suite_id is None - and result_run_id is None + identity_omitted = ( + result_suite_id is None and result_run_id is None + ) + may_omit_identity = identity_omitted and ( + state is JobState.CANCELLED + or (state is JobState.FAILED and failed_stage is None) ) if not may_omit_identity and ( result_suite_id != record.suite_id @@ -528,11 +1058,20 @@ def _adopt_result( ServiceErrorCode.RUN_FAILED, "Evaluation worker returned an inconsistent exit code", ) - run_root = ( - self._validated_run_root(record, raw_result) - if result_suite_id is not None - else None - ) + if result_suite_id is not None: + run_root = self._validated_run_root(record, raw_result) + elif state is JobState.CANCELLED: + run_root_value = self._write_cancelled_manifest( + record, + failed_stage=failed_stage, + ) + run_root = ( + Path(run_root_value) + if run_root_value is not None + else None + ) + else: + run_root = None public_result = { "state": state_value, "exit_code": exit_code, @@ -729,6 +1268,111 @@ def start( created=created.created, ) + def cancel(self, job_id: str) -> JobDetail: + """Request idempotent cooperative cancellation for one job.""" + record = self.manager.cancel(_validate_job_id(job_id)) + return self._detail(record) + + def retry( + self, + job_id: str, + *, + request_id: str, + ) -> JobStartResult: + """Create an idempotent retry from an immutable terminal snapshot.""" + if not self.manager.launch_enabled: + raise ServiceError( + ServiceErrorCode.CAPABILITY_DISABLED, + "Evaluation execution is disabled for this service", + ) + job_id = _validate_job_id(job_id) + request_id = _validate_request_id(request_id) + original = self.manager.reconcile(self.store.get(job_id)) + request_hash = _retry_request_hash( + retry_of=original.job_id, + config_sha256=original.config_sha256, + ) + existing = self.store.get_by_idempotency_key(request_id) + if existing is not None: + if existing.request_hash != request_hash: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "request_id is already bound to a different evaluation request", + details={"job_id": existing.job_id}, + ) + self.manager.enqueue() + return JobStartResult( + job=self.get(existing.job_id), + created=False, + ) + if original.state not in { + JobState.FAILED, + JobState.CANCELLED, + JobState.INTERRUPTED, + }: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Only failed, cancelled, or interrupted jobs can be retried", + ) + document, request = self._retry_snapshot(original) + retry_stage = self._retry_stage(original, document) + + run_id = _new_identity("run") if original.run_id is not None else None + overrides = EvaluationOverrides.model_validate( + { + "suite": original.suite_id, + "run": run_id, + "force_stages": [retry_stage], + "strict": bool(request.get("strict", False)), + } + ) + plan = self.planning.preflight_document( + original.config_ref, + document, + source_etag=original.config_sha256, + overrides=overrides, + ) + _require_ready(plan) + if plan.suite_id != original.suite_id or plan.run_id != run_id: + raise ServiceError( + ServiceErrorCode.INTERNAL, + "Retry preflight did not preserve allocated evaluation identity", + ) + if run_id is not None: + self._reject_existing_run(original.suite_id, run_id) + + yaml_text = dump_yaml(plan.effective_document) + config_sha256 = ( + "sha256:" + + hashlib.sha256(yaml_text.encode("utf-8")).hexdigest() + ) + new_job, job_dir = self._prepare_job( + config_ref=original.config_ref, + request_id=request_id, + request_hash=request_hash, + config_sha256=config_sha256, + suite_id=original.suite_id, + run_id=run_id, + plan=plan, + yaml_text=yaml_text, + retry_of=original.job_id, + ) + try: + created = self.store.create_or_get( + new_job, + max_queued_jobs=self.max_queued_jobs, + ) + except BaseException: + _remove_job_dir(job_dir) + raise + if not created.created: + _remove_job_dir(job_dir) + self.manager.enqueue() + return JobStartResult( + job=self.get(created.record.job_id), + created=created.created, + ) + def get(self, job_id: str) -> JobDetail: record = self.manager.reconcile(self.store.get(_validate_job_id(job_id))) return self._detail(record) @@ -828,6 +1472,121 @@ def read_log(self, job_id: str, *, max_bytes: int) -> str: ) return combined + def _retry_snapshot( + self, + record: JobRecord, + ) -> tuple[dict[str, Any], dict[str, Any]]: + job_dir = self.manager._job_dir(record.job_id) + snapshot_path = self.manager._job_file(job_dir, "config.yaml") + request_path = self.manager._job_file(job_dir, "request.json") + try: + if snapshot_path.stat().st_size > _JOB_SNAPSHOT_MAX_BYTES: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable evaluation snapshot exceeds its size limit", + ) + snapshot_bytes = snapshot_path.read_bytes() + except OSError as exc: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable evaluation snapshot is unavailable", + ) from exc + actual_hash = ( + "sha256:" + hashlib.sha256(snapshot_bytes).hexdigest() + ) + if actual_hash != record.config_sha256: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable evaluation snapshot failed its integrity check", + ) + try: + document = yaml.safe_load(snapshot_bytes.decode("utf-8")) + except (UnicodeDecodeError, yaml.YAMLError) as exc: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable evaluation snapshot is invalid", + ) from exc + request = _read_json_file( + request_path, + max_bytes=_JOB_RESULT_MAX_BYTES, + ) + if not isinstance(document, dict) or not isinstance(request, dict): + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable evaluation job inputs are invalid", + ) + if ( + request.get("job_id") != record.job_id + or request.get("config_sha256") != record.config_sha256 + or not isinstance(request.get("strict"), bool) + ): + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable evaluation job inputs failed their integrity check", + ) + return document, request + + def _retry_stage( + self, + record: JobRecord, + document: dict[str, Any], + ) -> str: + pipeline = document.get("pipeline") + configured = [ + stage + for stage in PIPELINE_STAGE_ORDER + if isinstance(pipeline, dict) + and isinstance(pipeline.get(stage), dict) + and pipeline[stage].get("enabled", True) + ] + if not configured: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The evaluation snapshot has no enabled stage to retry", + ) + candidate = ( + record.failed_stage + if record.failed_stage in configured + else _active_stage_from_events( + self.store.list_events(record.job_id, limit=1000) + ) + ) + if candidate not in configured: + candidate = configured[0] + + candidates = [candidate] + suite_root = self.workspace.path_policy.resolve_managed_output( + self.workspace.results_root / record.suite_id, + field_name="retry suite root", + expected_root=self.workspace.results_root, + reject_links=True, + ) + retry_index = PIPELINE_STAGE_ORDER.index(candidate) + if ( + "test_set" in configured + and PIPELINE_STAGE_ORDER.index("test_set") < retry_index + ): + test_set = suite_root / "test_set.jsonl" + if not _valid_jsonl(test_set): + candidates.append("test_set") + if ( + record.run_id is not None + and "inference" in configured + and PIPELINE_STAGE_ORDER.index("inference") < retry_index + ): + run_root = self.workspace.path_policy.resolve_managed_output( + suite_root / record.run_id, + field_name="retry source run", + expected_root=suite_root, + reject_links=True, + ) + if not _valid_jsonl(run_root / "inference_set.jsonl"): + candidates.append("inference") + return min( + candidates, + key=PIPELINE_STAGE_ORDER.index, + ) + def _prepare_job( self, *, @@ -839,6 +1598,7 @@ def _prepare_job( run_id: str | None, plan: Any, yaml_text: str, + retry_of: str | None = None, ) -> tuple[NewJob, Path]: jobs_root = _jobs_root(self.workspace) jobs_root.mkdir(parents=True, exist_ok=True) @@ -880,6 +1640,7 @@ def _prepare_job( "strict": bool(plan.strict), "force_stages": force_stages, "max_log_bytes": self.manager.max_log_bytes, + "retry_of": retry_of, }, ) resource_keys = [] @@ -905,6 +1666,7 @@ def _prepare_job( snapshot_path=str(snapshot), request_path=str(request_path), resource_keys=tuple(resource_keys), + retry_of=retry_of, ), job_dir, ) @@ -937,33 +1699,60 @@ def _reject_existing_run( def _detail(self, record: JobRecord) -> JobDetail: manifest = self._manifest(record) - heartbeat_at = _optional_text(manifest.get("heartbeat_at")) + event_projection = _event_projection( + self.store.list_events(record.job_id, limit=1000) + ) + heartbeat_at = ( + event_projection["heartbeat_at"] + or _optional_text(manifest.get("heartbeat_at")) + ) + stages = ( + dict(manifest.get("stages") or {}) + if isinstance(manifest.get("stages"), dict) + else {} + ) + stages.update(event_projection["stages"]) + stage_timings = ( + dict(manifest.get("stage_timings") or {}) + if isinstance(manifest.get("stage_timings"), dict) + else {} + ) + for stage_name, timing in event_projection["stage_timings"].items(): + existing = dict(stage_timings.get(stage_name) or {}) + existing.update(timing) + stage_timings[stage_name] = existing + manifest_progress = ( + dict(manifest.get("progress") or {}) + if isinstance(manifest.get("progress"), dict) + else {} + ) terminal_result = ( JobTerminalResult.model_validate(record.result) if record.result is not None else None ) + if ( + terminal_result is not None + and terminal_result.failed_stage is not None + ): + stages.setdefault( + terminal_result.failed_stage, + ( + "cancelled" + if terminal_result.state == "cancelled" + else "failed" + ), + ) return JobDetail( **_catalog_entry(record).model_dump(), request_id=record.idempotency_key, config_sha256=record.config_sha256, + cancel_requested_at=record.cancel_requested_at, heartbeat_at=heartbeat_at, heartbeat_age_seconds=_heartbeat_age(heartbeat_at), - stages=( - dict(manifest.get("stages") or {}) - if isinstance(manifest.get("stages"), dict) - else {} - ), - stage_timings=( - dict(manifest.get("stage_timings") or {}) - if isinstance(manifest.get("stage_timings"), dict) - else {} - ), - progress=( - dict(manifest.get("progress") or {}) - if isinstance(manifest.get("progress"), dict) - else {} - ), + stages=stages, + stage_timings=stage_timings, + progress=event_projection["progress"] or manifest_progress, terminal_result=terminal_result, error_code=record.error_code, error_message=record.error_message, @@ -1004,6 +1793,7 @@ def _catalog_entry(record: JobRecord) -> JobCatalogEntry: state=record.state, revision=record.revision, kind="evaluation", + retry_of=record.retry_of, config_ref=record.config_ref, suite_id=record.suite_id, run_id=record.run_id, @@ -1065,6 +1855,123 @@ def _request_hash( return "sha256:" + hashlib.sha256(payload).hexdigest() +def _retry_request_hash( + *, + retry_of: str, + config_sha256: str, +) -> str: + payload = json.dumps( + { + "operation": "retry_evaluation", + "retry_of": retry_of, + "config_sha256": config_sha256, + }, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return "sha256:" + hashlib.sha256(payload).hexdigest() + + +def _active_stage_from_events( + events: Sequence[dict[str, Any]], +) -> str | None: + active: str | None = None + terminal_stage: str | None = None + for event in events: + payload = event.get("payload") + if not isinstance(payload, dict): + continue + event_type = event.get("event_type") + if event_type == "pipeline_finished": + failed_stage = payload.get("failed_stage") + if isinstance(failed_stage, str) and failed_stage: + terminal_stage = failed_stage + continue + name = payload.get("name") + if not isinstance(name, str): + continue + if event_type == "stage_started": + active = name + elif event_type == "stage_finished": + if payload.get("state") in {"cancelled", "failed"}: + terminal_stage = name + if active == name: + active = None + return terminal_stage or active + + +def _valid_jsonl(path: Path) -> bool: + try: + scan = scan_jsonl(path, allow_trailing_partial=False) + except (JsonlIndexError, OSError): + return False + if not scan.records: + return False + identities: set[tuple[str, str]] = set() + for record in scan.records: + kind = record.row.get("type") + test_case_id = record.row.get("test_case_id") + if ( + not isinstance(kind, str) + or not kind + or not isinstance(test_case_id, str) + or not test_case_id + ): + return False + identity = (kind, test_case_id) + if identity in identities: + return False + identities.add(identity) + return True + + +def _event_projection( + events: Sequence[dict[str, Any]], +) -> dict[str, Any]: + stages: dict[str, str] = {} + stage_timings: dict[str, dict[str, Any]] = {} + progress: dict[str, Any] = {} + heartbeat_at: str | None = None + for event in events: + event_type = event.get("event_type") + timestamp = event.get("timestamp") + payload = event.get("payload") + if not isinstance(payload, dict): + continue + if isinstance(timestamp, str): + heartbeat_at = timestamp + name = payload.get("name") + if event_type == "stage_planned" and isinstance(name, str): + action = payload.get("action") + if isinstance(action, str): + stages[name] = action + elif event_type == "stage_started" and isinstance(name, str): + stages[name] = "running" + stage_timings.setdefault(name, {})["started_at"] = timestamp + elif event_type == "stage_progress" and isinstance(name, str): + values = payload.get("values") + if isinstance(values, dict): + progress = dict(values) + elif event_type == "stage_finished" and isinstance(name, str): + state = payload.get("state") + if isinstance(state, str): + stages[name] = state + timing = stage_timings.setdefault(name, {}) + timing["ended_at"] = timestamp + duration = payload.get("duration_seconds") + if isinstance(duration, (int, float)) and not isinstance( + duration, bool + ): + timing["duration_secs"] = duration + progress = {} + return { + "stages": stages, + "stage_timings": stage_timings, + "progress": progress, + "heartbeat_at": heartbeat_at, + } + + def _require_ready(plan: Any) -> None: if plan.ready: return @@ -1224,19 +2131,94 @@ def _process_create_time(pid: int) -> float: def _terminate_failed_launch(process: subprocess.Popen[bytes]) -> None: try: if process.poll() is None: - process.terminate() - try: - process.wait(timeout=5) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=5) - except (OSError, subprocess.SubprocessError): + _terminate_process_tree( + process.pid, + _process_create_time(process.pid), + timeout_seconds=_PROCESS_EXIT_TIMEOUT_SECONDS, + ) + except (OSError, subprocess.SubprocessError, ValueError): log.exception( "Could not terminate partially launched evaluation worker %s", process.pid, ) +def _terminate_process_tree( + pid: int, + create_time: float, + *, + timeout_seconds: float, +) -> None: + """Terminate one identity-verified worker and all descendants.""" + import psutil + + if timeout_seconds <= 0: + raise ValueError("timeout_seconds must be positive") + try: + process = psutil.Process(pid) + if ( + not process.is_running() + or abs(process.create_time() - create_time) >= 0.01 + ): + return + descendants = process.children(recursive=True) + except psutil.Error: + return + + if os.name != "nt": + try: + os.killpg(pid, signal.SIGTERM) + except ProcessLookupError: + return + except OSError: + log.warning( + "Could not terminate evaluation process group %s", + pid, + exc_info=True, + ) + for child in descendants: + try: + child.terminate() + except psutil.Error: + pass + try: + process.terminate() + except psutil.Error: + pass + else: + for child in descendants: + try: + child.terminate() + except psutil.Error: + pass + try: + process.terminate() + except psutil.Error: + pass + + targets = [*descendants, process] + _, alive = psutil.wait_procs(targets, timeout=timeout_seconds) + if not alive: + return + + if os.name != "nt": + if process in alive and _process_matches(pid, create_time): + try: + os.killpg(pid, signal.SIGKILL) + except (ProcessLookupError, OSError): + pass + for remaining in alive: + try: + remaining.kill() + except psutil.Error: + pass + _, still_alive = psutil.wait_procs(alive, timeout=timeout_seconds) + if still_alive: + raise RuntimeError( + "Evaluation process tree did not exit after forced termination" + ) + + def _process_matches(pid: int, create_time: float) -> bool: try: import psutil @@ -1262,6 +2244,31 @@ def _lease_expired(value: str | None) -> bool: return True +def _lease_seconds_remaining(value: str | None) -> float: + if value is None: + return 0.0 + try: + parsed = datetime.fromisoformat(value) + if parsed.tzinfo is None: + return 0.0 + except (TypeError, ValueError): + return 0.0 + return max( + 0.0, + (parsed - datetime.now(timezone.utc)).total_seconds(), + ) + + +def _cancellation_wait_seconds(record: JobRecord) -> float: + lease_wait = _lease_seconds_remaining(record.lease_expires_at) + if lease_wait <= 0: + return _CANCELLATION_POLL_SECONDS + return min( + _MAX_RECOVERY_SLEEP_SECONDS, + max(_CANCELLATION_POLL_SECONDS, lease_wait), + ) + + def _heartbeat_age(value: str | None) -> float | None: if value is None: return None diff --git a/assert_ai/services/job_models.py b/assert_ai/services/job_models.py index 130806d22..a9e2fb63d 100644 --- a/assert_ai/services/job_models.py +++ b/assert_ai/services/job_models.py @@ -12,6 +12,8 @@ from pydantic import BaseModel, ConfigDict, Field +# Public job responses remain API-v1 compatible. Persisted SQLite migrations +# use their own independent schema version in job_store.py. JOB_SCHEMA_VERSION = 1 @@ -46,6 +48,7 @@ class JobRecord: idempotency_key: str request_hash: str kind: str + retry_of: str | None state: JobState created_at: str started_at: str | None @@ -85,6 +88,7 @@ class NewJob: snapshot_path: str request_path: str resource_keys: tuple[str, ...] + retry_of: str | None = None kind: str = "evaluation" @@ -116,6 +120,7 @@ class JobCatalogEntry(_ServiceModel): state: JobState revision: int = Field(ge=0) kind: Literal["evaluation"] = "evaluation" + retry_of: str | None = None config_ref: str suite_id: str run_id: str | None = None @@ -136,6 +141,7 @@ class JobDetail(JobCatalogEntry): request_id: str config_sha256: str + cancel_requested_at: str | None = None heartbeat_at: str | None = None heartbeat_age_seconds: float | None = Field(default=None, ge=0) stages: dict[str, Any] = Field(default_factory=dict) diff --git a/assert_ai/services/job_store.py b/assert_ai/services/job_store.py index d92516b3c..91154a06c 100644 --- a/assert_ai/services/job_store.py +++ b/assert_ai/services/job_store.py @@ -25,8 +25,12 @@ ) _BUSY_TIMEOUT_MS = 5_000 -_JOB_STORE_SCHEMA_VERSION = 1 -_ACTIVE_STATES = (JobState.STARTING.value, JobState.RUNNING.value) +_JOB_STORE_SCHEMA_VERSION = 2 +_ACTIVE_STATES = ( + JobState.STARTING.value, + JobState.RUNNING.value, + JobState.CANCELLING.value, +) _MAX_EVENTS_PER_JOB = 1000 _SCHEMA = """ @@ -35,6 +39,7 @@ idempotency_key TEXT NOT NULL UNIQUE, request_hash TEXT NOT NULL, kind TEXT NOT NULL, + retry_of TEXT, state TEXT NOT NULL, created_at TEXT NOT NULL, started_at TEXT, @@ -151,17 +156,19 @@ def create_or_get( connection.execute( """ INSERT INTO jobs( - job_id, idempotency_key, request_hash, kind, state, + job_id, idempotency_key, request_hash, kind, retry_of, + state, created_at, suite_id, run_id, config_ref, config_sha256, snapshot_path, request_path, resource_keys_json - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( new_job.job_id, new_job.idempotency_key, new_job.request_hash, new_job.kind, + new_job.retry_of, JobState.QUEUED.value, created_at, new_job.suite_id, @@ -202,6 +209,7 @@ def create_or_get( def get(self, job_id: str) -> JobRecord: if not self.exists: raise ServiceError(ServiceErrorCode.NOT_FOUND, "Job not found") + self.initialize() with self._connection() as connection: row = connection.execute( "SELECT * FROM jobs WHERE job_id = ?", @@ -217,6 +225,7 @@ def get_by_idempotency_key( ) -> JobRecord | None: if not self.exists: return None + self.initialize() with self._connection() as connection: row = connection.execute( "SELECT * FROM jobs WHERE idempotency_key = ?", @@ -235,6 +244,7 @@ def list_records( raise ValueError("limit must be positive") if not self.exists: return () + self.initialize() conditions: list[str] = [] values: list[Any] = [] if states: @@ -273,12 +283,13 @@ def claim_next( raise ValueError("max_active_jobs must be positive") if not self.exists: return None + self.initialize() now = _now() expires_at = _after(lease_seconds) with self._transaction() as connection: active = int( connection.execute( - "SELECT COUNT(*) FROM jobs WHERE state IN (?, ?)", + "SELECT COUNT(*) FROM jobs WHERE state IN (?, ?, ?)", _ACTIVE_STATES, ).fetchone()[0] ) @@ -355,9 +366,25 @@ def mark_running( process_create_time: float, lease_seconds: float, ) -> JobRecord: + self.initialize() now = _now() expires_at = _after(lease_seconds) with self._transaction() as connection: + current = self._get_in_transaction(connection, job_id) + if ( + current.lease_owner != lease_owner + or current.state + not in {JobState.STARTING, JobState.CANCELLING} + ): + raise ServiceError( + ServiceErrorCode.CONFLICT, + "Job can no longer attach its evaluation worker", + ) + next_state = ( + JobState.CANCELLING + if current.state is JobState.CANCELLING + else JobState.RUNNING + ) changed = connection.execute( """ UPDATE jobs @@ -367,13 +394,13 @@ def mark_running( WHERE job_id = ? AND state = ? AND lease_owner = ? """, ( - JobState.RUNNING.value, + next_state.value, now, pid, process_create_time, expires_at, job_id, - JobState.STARTING.value, + current.state.value, lease_owner, ), ).rowcount @@ -393,8 +420,12 @@ def mark_running( self._append_event( connection, job_id, - "running", - {"state": JobState.RUNNING.value, "pid": pid}, + ( + "running" + if next_state is JobState.RUNNING + else "cancelling_worker_started" + ), + {"state": next_state.value, "pid": pid}, timestamp=now, ) return self._get_in_transaction(connection, job_id) @@ -406,6 +437,7 @@ def renew_lease( lease_owner: str, lease_seconds: float, ) -> bool: + self.initialize() expires_at = _after(lease_seconds) with self._transaction() as connection: changed = connection.execute( @@ -413,7 +445,7 @@ def renew_lease( UPDATE jobs SET lease_expires_at = ? WHERE job_id = ? AND lease_owner = ? - AND state IN (?, ?) + AND state IN (?, ?, ?) """, ( expires_at, @@ -421,6 +453,7 @@ def renew_lease( lease_owner, JobState.STARTING.value, JobState.RUNNING.value, + JobState.CANCELLING.value, ), ).rowcount if changed != 1: @@ -435,6 +468,157 @@ def renew_lease( ) return True + def request_cancel(self, job_id: str) -> JobRecord: + """Persist an idempotent cancellation request.""" + self.initialize() + now = _now() + with self._transaction() as connection: + current = self._get_in_transaction(connection, job_id) + if current.state is JobState.CANCELLED: + return current + if current.state in TERMINAL_JOB_STATES: + raise ServiceError( + ServiceErrorCode.JOB_NOT_CANCELLABLE, + f"Job is already {current.state.value}", + ) + if current.state is JobState.CANCELLING: + return current + if current.state is JobState.QUEUED: + result = { + "state": JobState.CANCELLED.value, + "exit_code": 130, + "failed_stage": None, + "error_code": None, + "error_message": None, + } + connection.execute( + """ + UPDATE jobs + SET state = ?, cancel_requested_at = ?, ended_at = ?, + exit_code = ?, result_json = ?, + lease_owner = NULL, lease_expires_at = NULL, + revision = revision + 1 + WHERE job_id = ? AND state = ? + """, + ( + JobState.CANCELLED.value, + now, + now, + 130, + _json(result), + job_id, + JobState.QUEUED.value, + ), + ) + self._append_event( + connection, + job_id, + "cancelled", + { + "state": JobState.CANCELLED.value, + "exit_code": 130, + }, + timestamp=now, + ) + else: + connection.execute( + """ + UPDATE jobs + SET state = ?, cancel_requested_at = ?, + revision = revision + 1 + WHERE job_id = ? AND state IN (?, ?) + """, + ( + JobState.CANCELLING.value, + now, + job_id, + JobState.STARTING.value, + JobState.RUNNING.value, + ), + ) + self._append_event( + connection, + job_id, + "cancel_requested", + {"state": JobState.CANCELLING.value}, + timestamp=now, + ) + return self._get_in_transaction(connection, job_id) + + def adopt_lease( + self, + job_id: str, + *, + lease_owner: str, + lease_seconds: float, + ) -> JobRecord | None: + """Claim an expired active-job lease during startup recovery.""" + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + self.initialize() + now = _now() + expires_at = _after(lease_seconds) + with self._transaction() as connection: + changed = connection.execute( + """ + UPDATE jobs + SET lease_owner = ?, lease_expires_at = ?, + revision = revision + 1 + WHERE job_id = ? + AND state IN (?, ?, ?) + AND ( + lease_expires_at IS NULL + OR lease_expires_at <= ? + OR lease_owner = ? + ) + """, + ( + lease_owner, + expires_at, + job_id, + JobState.STARTING.value, + JobState.RUNNING.value, + JobState.CANCELLING.value, + now, + lease_owner, + ), + ).rowcount + if changed != 1: + return None + connection.execute( + """ + UPDATE resource_locks + SET lease_expires_at = ? + WHERE job_id = ? + """, + (expires_at, job_id), + ) + self._append_event( + connection, + job_id, + "manager_recovered", + {"state": self._get_in_transaction(connection, job_id).state.value}, + timestamp=now, + ) + return self._get_in_transaction(connection, job_id) + + def list_nonterminal_records(self) -> tuple[JobRecord, ...]: + """Return every queued or active job for deterministic recovery.""" + if not self.exists: + return () + self.initialize() + placeholders = ", ".join("?" for _ in TERMINAL_JOB_STATES) + with self._connection() as connection: + rows = connection.execute( + f""" + SELECT * FROM jobs + WHERE state NOT IN ({placeholders}) + ORDER BY created_at, job_id + """, + tuple(state.value for state in TERMINAL_JOB_STATES), + ).fetchall() + return tuple(_record(row) for row in rows) + def mark_terminal( self, job_id: str, @@ -450,6 +634,7 @@ def mark_terminal( ) -> JobRecord: if state not in TERMINAL_JOB_STATES: raise ValueError("terminal state required") + self.initialize() now = _now() with self._transaction() as connection: current = self._get_in_transaction(connection, job_id) @@ -514,6 +699,7 @@ def append_event( event_type: str, payload: dict[str, Any], ) -> int: + self.initialize() with self._transaction() as connection: self._get_in_transaction(connection, job_id) return self._append_event( @@ -585,12 +771,28 @@ def initialize(self) -> None: ) if version not in { 0, + 1, _JOB_STORE_SCHEMA_VERSION, }: raise ServiceError( ServiceErrorCode.INTERNAL, "Unsupported job store schema version", ) + columns = { + str(row["name"]) + for row in connection.execute( + "PRAGMA table_info(jobs)" + ).fetchall() + } + if "retry_of" not in columns: + try: + connection.execute( + "ALTER TABLE jobs " + "ADD COLUMN retry_of TEXT" + ) + except sqlite3.OperationalError as exc: + if "duplicate column" not in str(exc).lower(): + raise connection.execute( "PRAGMA user_version = " f"{_JOB_STORE_SCHEMA_VERSION}" @@ -707,15 +909,47 @@ def _append_event( _json(payload), ), ) - cutoff = sequence - _MAX_EVENTS_PER_JOB - if cutoff > 0: + count = int( + connection.execute( + "SELECT COUNT(*) FROM job_events WHERE job_id = ?", + (job_id,), + ).fetchone()[0] + ) + excess = count - _MAX_EVENTS_PER_JOB + if excess > 0: connection.execute( """ DELETE FROM job_events - WHERE job_id = ? AND sequence <= ? + WHERE job_id = ? AND sequence IN ( + SELECT sequence FROM job_events + WHERE job_id = ? + AND event_type IN ('heartbeat', 'stage_progress') + ORDER BY sequence + LIMIT ? + ) """, - (job_id, cutoff), + (job_id, job_id, excess), ) + remaining = int( + connection.execute( + "SELECT COUNT(*) FROM job_events WHERE job_id = ?", + (job_id,), + ).fetchone()[0] + ) + overflow = remaining - _MAX_EVENTS_PER_JOB + if overflow > 0: + connection.execute( + """ + DELETE FROM job_events + WHERE job_id = ? AND sequence IN ( + SELECT sequence FROM job_events + WHERE job_id = ? + ORDER BY sequence + LIMIT ? + ) + """, + (job_id, job_id, overflow), + ) return sequence @@ -725,6 +959,7 @@ def _record(row: sqlite3.Row) -> JobRecord: idempotency_key=str(row["idempotency_key"]), request_hash=str(row["request_hash"]), kind=str(row["kind"]), + retry_of=_optional_str(row["retry_of"]), state=JobState(str(row["state"])), created_at=str(row["created_at"]), started_at=_optional_str(row["started_at"]), diff --git a/assert_ai/services/run_planning.py b/assert_ai/services/run_planning.py index 57a0aebe7..37d87e03a 100644 --- a/assert_ai/services/run_planning.py +++ b/assert_ai/services/run_planning.py @@ -190,12 +190,28 @@ def preflight( overrides: EvaluationOverrides | None = None, ) -> EvaluationPreflight: record = self.configs.get_config(config_ref) - effective = deepcopy(record.document) + return self.preflight_document( + record.config_ref, + record.document, + source_etag=record.etag, + overrides=overrides, + ) + + def preflight_document( + self, + config_ref: str, + document: dict[str, Any], + *, + source_etag: str, + overrides: EvaluationOverrides | None = None, + ) -> EvaluationPreflight: + """Preflight an immutable document using a managed logical base path.""" + effective = deepcopy(document) applied = overrides or EvaluationOverrides() _apply_overrides(effective, applied) validation = self.configs.validate_document( effective, - config_ref=record.config_ref, + config_ref=config_ref, ) blocking = [ _validation_issue(issue) @@ -207,8 +223,8 @@ def preflight( ] if not validation.valid: return EvaluationPreflight( - config_ref=record.config_ref, - source_etag=record.etag, + config_ref=config_ref, + source_etag=source_etag, effective_document=effective, validation=validation, ready=False, @@ -218,8 +234,8 @@ def preflight( ) config_path = self.workspace.path_policy.resolve_config_path( - record.config_ref, - must_exist=True, + config_ref, + must_exist=False, reject_links=True, ) try: @@ -235,8 +251,8 @@ def preflight( message=str(exc), ) return EvaluationPreflight( - config_ref=record.config_ref, - source_etag=record.etag, + config_ref=config_ref, + source_etag=source_etag, effective_document=effective, validation=validation, ready=False, @@ -288,8 +304,8 @@ def preflight( sample_sizes = _sample_sizes(effective) estimate = _estimate_model_calls(stages) return EvaluationPreflight( - config_ref=record.config_ref, - source_etag=record.etag, + config_ref=config_ref, + source_etag=source_etag, effective_document=effective, validation=validation, ready=not blocking, diff --git a/assert_ai/stages/inference.py b/assert_ai/stages/inference.py index 216c23214..6c58db910 100644 --- a/assert_ai/stages/inference.py +++ b/assert_ai/stages/inference.py @@ -44,6 +44,7 @@ ) from assert_ai.core.model_client import GenerateOptions, Message, ModelResponse, build_llm_call_trace, generate, to_jsonable from assert_ai.core.model_client import LLMAuthError, LLMContentFilterError, LLMInputError, LLMRateLimitError, LLMProviderError +from assert_ai.core.run_control import RunCancelled, RunControl from assert_ai.core.session import ( CallableSession, ExternalSession, @@ -1124,6 +1125,7 @@ async def run_inference( strict: bool = False, forced: bool = False, heartbeat: Any = None, + run_control: RunControl | None = None, rewrite_test_set_path: bool = True, ) -> dict[str, Any]: """Run all test-case inferences and write the transcript artifact.""" @@ -1289,6 +1291,8 @@ async def _worker(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: """ output_index, test_case_row = test_case try: + if run_control is not None: + run_control.raise_if_cancelled(stage="inference") kind = test_case_row["type"] if kind == "prompt": transcript = await _run_prompt_test_case( @@ -1312,7 +1316,11 @@ async def _worker(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: ) else: raise ValueError(f"unsupported test case type: {kind}") + if run_control is not None: + run_control.raise_if_cancelled(stage="inference") return {"output_index": output_index, "inference_row": transcript.to_dict()} + except RunCancelled: + raise except LLMContentFilterError as exc: # Adversarial-eval test cases (XPIA, PII, security attacks) can # legitimately trip the tester or target model's content @@ -1417,6 +1425,8 @@ async def _guard(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: ) except Exception: # noqa: BLE001 heartbeat = None + if run_control is not None: + run_control.raise_if_cancelled(stage="inference") idx = result["output_index"] test_case_row = test_cases[idx] kind = test_case_row.get("type", "") @@ -1519,6 +1529,8 @@ async def _guard(test_case: tuple[int, dict[str, Any]]) -> dict[str, Any]: field_name="inference output directory", expected_root=managed_output_root, ) + if run_control is not None: + run_control.raise_if_cancelled(stage="inference") build_run_viewer_artifacts(out_dir) return { @@ -1584,6 +1596,7 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: strict=cfg.get("strict", False), forced=bool(ctx.get("_stage_forced", False)), heartbeat=ctx.get("_heartbeat") if isinstance(ctx, dict) else None, + run_control=ctx.get("_run_control") if isinstance(ctx, dict) else None, rewrite_test_set_path=rewrite_test_set_path, ) target_obj = ctx["target"] diff --git a/assert_ai/stages/judge.py b/assert_ai/stages/judge.py index 8341ff24e..9e88a27cf 100644 --- a/assert_ai/stages/judge.py +++ b/assert_ai/stages/judge.py @@ -24,6 +24,7 @@ run_transcript_judge as run_llm_judge, ) from assert_ai.core.model_client import LLMAuthError, LLMContentFilterError, LLMInputError, LLMRateLimitError, LLMProviderError +from assert_ai.core.run_control import RunCancelled, RunControl from assert_ai.core.transcript import Transcript, TranscriptEvent, TranscriptMetadata from assert_ai.viewer_read_model import build_run_viewer_artifacts @@ -98,6 +99,7 @@ async def run_judge( disabled_dimensions: list[str] | None = None, forced: bool = False, heartbeat: Any = None, + run_control: RunControl | None = None, path_policy: RuntimePathPolicy | None = None, config_path: Path | None = None, managed_output_root: Path | None = None, @@ -297,10 +299,17 @@ async def worker(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: """ output_index, row = item try: + if run_control is not None: + run_control.raise_if_cancelled(stage="judge") + score = await score_row(row) + if run_control is not None: + run_control.raise_if_cancelled(stage="judge") return { "output_index": output_index, - "score_row": await score_row(row), + "score_row": score, } + except RunCancelled: + raise except LLMContentFilterError as exc: # Adversarial-eval workloads routinely send transcripts the # judge's content filter will reject (XPIA payloads, PII @@ -476,6 +485,8 @@ async def guard(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: ) except Exception: # noqa: BLE001 heartbeat = None + if run_control is not None: + run_control.raise_if_cancelled(stage="judge") errors: list[Exception] = [] written_rows = 0 for completed_task in asyncio.as_completed(tasks): @@ -504,6 +515,8 @@ async def guard(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: ) except Exception: # noqa: BLE001 heartbeat = None + if run_control is not None: + run_control.raise_if_cancelled(stage="judge") # Always rebuild viewer artifacts so the on-disk read model reflects the # current scores.jsonl, even when a row failed and we are about to raise. @@ -513,6 +526,8 @@ async def guard(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: field_name="judge output directory", expected_root=managed_output_root, ) + if run_control is not None: + run_control.raise_if_cancelled(stage="judge") build_run_viewer_artifacts(out_dir) # Per-row failures should not kill the stage as long as *some* rows # succeeded. The errors are surfaced via judge_failures in the @@ -598,6 +613,7 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, str]: disabled_dimensions=disabled_dimensions, forced=bool(ctx.get("_stage_forced", False)), heartbeat=ctx.get("_heartbeat") if isinstance(ctx, dict) else None, + run_control=ctx.get("_run_control") if isinstance(ctx, dict) else None, path_policy=ctx.get("path_policy"), config_path=Path(ctx["config_path"]), managed_output_root=Path(ctx["run_root"]), diff --git a/assert_ai/stages/systematize.py b/assert_ai/stages/systematize.py index d3e54c980..61a7a93d8 100644 --- a/assert_ai/stages/systematize.py +++ b/assert_ai/stages/systematize.py @@ -187,6 +187,9 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: reasoning_effort=model_cfg.reasoning_effort, ) sys_path = str(Path(cfg["save_dir"]) / "systematization.json") + run_control = ctx.get("_run_control") + if run_control is not None: + run_control.raise_if_cancelled(stage="systematize") log.debug(f"systematize: model={model_cfg.name}, behavior_category_count={behavior_category_count}, web_search={web_search}") log.info("[systematize] [1/2] Researching behavior taxonomy...") async with log_heartbeat("[systematize] [1/2] Researching behavior taxonomy"): @@ -198,6 +201,8 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: web_search=web_search, context=context, ) + if run_control is not None: + run_control.raise_if_cancelled(stage="systematize") log.info("[systematize] [1/2] Behavior taxonomy complete") taxonomy_path_str = str(Path(cfg["save_dir"]) / "taxonomy.json") @@ -215,6 +220,8 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: model_cfg=convert_model_cfg, behavior_category_count_hint=behavior_category_count, ) + if run_control is not None: + run_control.raise_if_cancelled(stage="systematize") return { "systematize_dir": cfg["save_dir"], diff --git a/assert_ai/stages/test_set.py b/assert_ai/stages/test_set.py index 18b72953c..2a9efeb6e 100644 --- a/assert_ai/stages/test_set.py +++ b/assert_ai/stages/test_set.py @@ -19,6 +19,7 @@ from assert_ai.config import parse_model_config, reject_unknown_keys, resolve_stage_paths from assert_ai.core.async_utils import gather_limited +from assert_ai.core.run_control import RunCancelled, RunControl from assert_ai.core.config_model import ( DEFAULT_GENERATION_MAX_TOKENS, DEFAULT_GENERATION_TEMPERATURE, @@ -957,6 +958,7 @@ async def _generate_records( sampling: dict[str, Any] | None = None, seed: int = 0, concurrency: int = 8, + run_control: RunControl | None = None, ) -> dict[str, Any]: """Call the LLM once per covering-array tuple and return records. @@ -1003,6 +1005,8 @@ async def _process(job: TestCaseJob) -> dict[str, Any]: slug = slugify(str(job.behavior.get("name") or "")) behavior_name = str(job.behavior.get("name") or "") try: + if run_control is not None: + run_control.raise_if_cancelled(stage="test_set") prompt = build_generation_prompt( kind=kind, taxonomy=taxonomy, @@ -1028,6 +1032,8 @@ async def _process(job: TestCaseJob) -> dict[str, Any]: call_label=f"test_set:{kind}:{slug}", ), ) + if run_control is not None: + run_control.raise_if_cancelled(stage="test_set") payload = response.parsed if not isinstance(payload, dict) or not isinstance(payload.get("test_set"), list): raise ValueError(f"{kind} test-case generation returned invalid test_set payload") @@ -1063,6 +1069,8 @@ async def _process(job: TestCaseJob) -> dict[str, Any]: ) ) return {"order": job.order, "records": records} + except RunCancelled: + raise except LLMAuthError: raise except (LLMInputError, LLMRateLimitError, LLMProviderError) as exc: @@ -1126,6 +1134,7 @@ async def run_test_set( stratification: dict[str, Any] | None = None, seed: int = 0, concurrency: int = 8, + run_control: RunControl | None = None, ) -> dict[str, Any]: """Generate prompt and/or scenario test_set.""" if prompt is None and scenario is None: @@ -1168,6 +1177,7 @@ async def run_test_set( sampling=cfg.get("sampling"), seed=seed, concurrency=concurrency, + run_control=run_control, ) for kind, cfg in kinds_cfgs ]) @@ -1175,6 +1185,8 @@ async def run_test_set( all_records = [rec for r in results for rec in r["records"]] errored_count = sum(int(r.get("errored_count", 0)) for r in results) all_records = normalize_test_case_rows(all_records) + if run_control is not None: + run_control.raise_if_cancelled(stage="test_set") write_jsonl(out_path, all_records) # --- Coverage check ------------------------------------------------------- @@ -1341,6 +1353,9 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: managed_output_root=Path(ctx["suite_root"]), ) taxonomy_path = cfg["taxonomy_path"] + run_control = ctx.get("_run_control") + if run_control is not None: + run_control.raise_if_cancelled(stage="test_set") stratification_dir = Path(cfg["save_path"]).parent stratification_result = await run_stratification( taxonomy_path=taxonomy_path, @@ -1352,6 +1367,8 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: reasoning_effort=stratify_model_cfg.reasoning_effort if stratify_model_cfg is not None else None, temperature=stratify_model_cfg.temperature if stratify_model_cfg is not None else None, ) + if run_control is not None: + run_control.raise_if_cancelled(stage="test_set") stratification_path = Path(stratification_result["stratification_path"]) raw_stratification = json.loads(stratification_path.read_text(encoding="utf-8")) @@ -1371,6 +1388,7 @@ async def run(ctx: dict[str, Any], raw_cfg: dict[str, Any]) -> dict[str, Any]: target=ctx.get("target"), tool_source=tool_source, stratification=raw_stratification, + run_control=run_control, ) return { "test_set_path": result["test_set_path"], diff --git a/tests/test_evaluation_service.py b/tests/test_evaluation_service.py index c9425f2e9..30ee4380c 100644 --- a/tests/test_evaluation_service.py +++ b/tests/test_evaluation_service.py @@ -4,7 +4,11 @@ from __future__ import annotations import json +import os +import subprocess +import sys import time +from dataclasses import replace from pathlib import Path from unittest.mock import patch @@ -21,6 +25,7 @@ from assert_ai.services.evaluations import ( EvaluationJobManager, EvaluationService, + _active_stage_from_events, ) from assert_ai.services.job_models import JobState from assert_ai.services.job_store import JobStore @@ -30,12 +35,23 @@ ) -def _service(root: Path) -> tuple[ConfigService, EvaluationService]: +def _service( + root: Path, + *, + cancellation_grace_seconds: float = 10.0, + lease_seconds: float = 60.0, +) -> tuple[ConfigService, EvaluationService]: workspace = WorkspaceService.create(root) configs = ConfigService(workspace) planning = RunPlanningService(workspace, configs) store = JobStore(workspace.artifacts_root / "mcp" / "jobs.sqlite3") - manager = EvaluationJobManager(workspace, store, max_active_jobs=1) + manager = EvaluationJobManager( + workspace, + store, + max_active_jobs=1, + cancellation_grace_seconds=cancellation_grace_seconds, + lease_seconds=lease_seconds, + ) return configs, EvaluationService( workspace, configs, @@ -95,6 +111,7 @@ def _wait_terminal( if detail.state in { JobState.COMPLETED, JobState.FAILED, + JobState.CANCELLED, JobState.INTERRUPTED, }: return detail @@ -102,6 +119,41 @@ def _wait_terminal( raise AssertionError("evaluation job did not finish") +def _wait_state( + service: EvaluationService, + job_id: str, + state: JobState, + *, + timeout_s: float = 10, +): + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + detail = service.get(job_id) + if detail.state is state: + return detail + time.sleep(0.05) + raise AssertionError(f"evaluation job did not reach {state.value}") + + +def _wait_stage_state( + service: EvaluationService, + job_id: str, + stage: str, + state: str, + *, + timeout_s: float = 15, +): + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + detail = service.get(job_id) + if detail.stages.get(stage) == state: + return detail + time.sleep(0.05) + raise AssertionError( + f"evaluation stage {stage} did not reach {state}" + ) + + def test_inference_only_job_completes_and_is_idempotent( tmp_path: Path, ) -> None: @@ -158,6 +210,34 @@ def test_inference_only_job_completes_and_is_idempotent( assert service.list().items[0].job_id == started.job.job_id +def test_suite_only_job_reports_observer_state_without_a_run_manifest( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "suite-only.yaml", + document={ + "suite": "suite-only", + "pipeline": {"inference": {"enabled": False}}, + }, + ) + + started = service.start( + "suite-only.yaml", + request_id="suite-only", + ) + terminal = _wait_terminal(service, started.job.job_id) + + assert terminal.state is JobState.COMPLETED + assert terminal.run_id is None + assert terminal.stages == {"inference": "disabled"} + assert terminal.heartbeat_at is not None + assert terminal.resources == { + "config": "assert://config/suite-only.yaml", + "worker_log": f"assert://job/{terminal.job_id}/log", + } + + def test_request_id_conflict_does_not_launch_duplicate( tmp_path: Path, ) -> None: @@ -310,14 +390,14 @@ def test_malformed_worker_result_becomes_a_persisted_failure( with patch.object(EvaluationJobManager, "enqueue"): started = service.start("demo.yaml", request_id="request") record = service.store.claim_next( - lease_owner="test-manager", + lease_owner=service.manager._owner, lease_seconds=30, max_active_jobs=1, ) assert record is not None record = service.store.mark_running( record.job_id, - lease_owner="test-manager", + lease_owner=service.manager._owner, pid=123, process_create_time=456, lease_seconds=30, @@ -355,3 +435,507 @@ def test_worker_log_retains_a_bounded_tail(tmp_path: Path) -> None: contents = log_path.read_text(encoding="utf-8") assert log_path.stat().st_size <= 4096 assert contents.startswith("[earlier worker output truncated]") + + +def test_running_job_cancels_cooperatively_and_persists_progress( + tmp_path: Path, +) -> None: + configs, service = _service( + tmp_path, + cancellation_grace_seconds=3, + ) + document = _write_inference_fixture(tmp_path) + fixture = tmp_path / "evals" / "fixture.jsonl" + fixture.write_text( + "".join( + json.dumps( + { + "type": "prompt", + "test_case_id": f"case-{index}", + "behavior": "local behavior", + "seed": {"description": f"hello {index}"}, + } + ) + + "\n" + for index in range(8) + ), + encoding="utf-8", + ) + (tmp_path / "agent.py").write_text( + "import time\n" + "def run(message, *, history=None):\n" + " del history\n" + " time.sleep(0.15)\n" + " return message\n", + encoding="utf-8", + ) + configs.save_config("demo.yaml", document=document) + started = service.start("demo.yaml", request_id="cancel-cooperative") + _wait_stage_state( + service, + started.job.job_id, + "inference", + "running", + ) + + cancelling = service.cancel(started.job.job_id) + terminal = _wait_terminal(service, started.job.job_id) + + assert cancelling.state is JobState.CANCELLING + assert cancelling.cancel_requested_at is not None + assert terminal.state is JobState.CANCELLED + assert terminal.terminal_result is not None + assert terminal.terminal_result.exit_code == 130 + assert terminal.terminal_result.failed_stage == "inference" + assert terminal.stages["inference"] == "cancelled" + assert terminal.heartbeat_at is not None + events = service.store.list_events(started.job.job_id, limit=1000) + event_types = {event["event_type"] for event in events} + assert { + "cancel_observed", + "pipeline_started", + "stage_planned", + "stage_started", + "stage_progress", + "stage_finished", + "pipeline_finished", + "cancel_requested", + "cancelled", + }.issubset(event_types) + assert "termination_escalated" not in event_types + assert ( + tmp_path + / "artifacts" + / "mcp" + / "jobs" + / started.job.job_id + / "result.json" + ).is_file() + + +def test_cancelled_stage_projection_prefers_terminal_events() -> None: + assert _active_stage_from_events( + ( + { + "event_type": "stage_started", + "payload": {"name": "inference"}, + }, + { + "event_type": "stage_finished", + "payload": { + "name": "inference", + "state": "cancelled", + }, + }, + { + "event_type": "pipeline_finished", + "payload": { + "state": "cancelled", + "failed_stage": "inference", + }, + }, + ) + ) == "inference" + + +def test_inspect_manager_does_not_reconcile_or_cancel_active_job( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + with patch.object(EvaluationJobManager, "enqueue"): + started = service.start("demo.yaml", request_id="inspect-only") + claimed = service.store.claim_next( + lease_owner="execute-manager", + lease_seconds=30, + max_active_jobs=1, + ) + assert claimed is not None + running = service.store.mark_running( + claimed.job_id, + lease_owner="execute-manager", + pid=2_147_483_647, + process_create_time=1, + lease_seconds=30, + ) + cancelling = service.store.request_cancel(running.job_id) + inspect_manager = EvaluationJobManager( + service.workspace, + service.store, + launch_enabled=False, + ) + + observed = inspect_manager.reconcile(cancelling) + + assert observed == cancelling + assert service.store.get(started.job.job_id).state is JobState.CANCELLING + assert not ( + tmp_path + / "artifacts" + / "mcp" + / "jobs" + / started.job.job_id + / "cancel.requested" + ).exists() + + +def test_cancellation_terminates_worker_descendants( + tmp_path: Path, +) -> None: + psutil = pytest.importorskip("psutil") + configs, service = _service( + tmp_path, + cancellation_grace_seconds=0.2, + ) + document = _write_inference_fixture(tmp_path) + (tmp_path / "agent.py").write_text( + "import pathlib, subprocess, sys, time\n" + "def run(message, *, history=None):\n" + " del message, history\n" + " child = subprocess.Popen([sys.executable, '-c', " + "'import time; time.sleep(60)'])\n" + " pathlib.Path('child.pid').write_text(str(child.pid), " + "encoding='utf-8')\n" + " time.sleep(60)\n" + " return 'late'\n", + encoding="utf-8", + ) + configs.save_config("demo.yaml", document=document) + started = service.start("demo.yaml", request_id="cancel-tree") + child_pid_path = tmp_path / "child.pid" + deadline = time.monotonic() + 15 + while not child_pid_path.is_file() and time.monotonic() < deadline: + time.sleep(0.05) + assert child_pid_path.is_file() + child_pid = int(child_pid_path.read_text(encoding="utf-8")) + + try: + service.cancel(started.job.job_id) + terminal = _wait_terminal( + service, + started.job.job_id, + timeout_s=15, + ) + assert terminal.state is JobState.CANCELLED + deadline = time.monotonic() + 5 + while psutil.pid_exists(child_pid) and time.monotonic() < deadline: + time.sleep(0.05) + assert not psutil.pid_exists(child_pid) + finally: + if psutil.pid_exists(child_pid): + process = psutil.Process(child_pid) + process.kill() + process.wait(timeout=5) + + +def test_startup_recovery_marks_dead_worker_interrupted( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path, lease_seconds=0.05) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + with patch.object(EvaluationJobManager, "enqueue"): + started = service.start("demo.yaml", request_id="recover-dead") + claimed = service.store.claim_next( + lease_owner="old-manager", + lease_seconds=0.05, + max_active_jobs=1, + ) + assert claimed is not None + service.store.mark_running( + claimed.job_id, + lease_owner="old-manager", + pid=2_147_483_647, + process_create_time=1, + lease_seconds=0.05, + ) + time.sleep(0.06) + recovered = EvaluationJobManager( + service.workspace, + service.store, + lease_seconds=0.1, + ) + + recovered.start() + service.manager = recovered + terminal = _wait_terminal(service, started.job.job_id) + + assert terminal.state is JobState.INTERRUPTED + assert terminal.error_code == ServiceErrorCode.JOB_INTERRUPTED.value + + +def test_startup_recovery_adopts_and_monitors_a_live_worker( + tmp_path: Path, +) -> None: + psutil = pytest.importorskip("psutil") + configs, service = _service(tmp_path, lease_seconds=0.05) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + with patch.object(EvaluationJobManager, "enqueue"): + started = service.start("demo.yaml", request_id="recover-live") + claimed = service.store.claim_next( + lease_owner="old-manager", + lease_seconds=0.05, + max_active_jobs=1, + ) + assert claimed is not None + process = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(60)"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + creationflags=( + subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0 + ), + start_new_session=os.name != "nt", + ) + try: + running = service.store.mark_running( + claimed.job_id, + lease_owner="old-manager", + pid=process.pid, + process_create_time=psutil.Process(process.pid).create_time(), + lease_seconds=0.05, + ) + time.sleep(0.06) + recovered = EvaluationJobManager( + service.workspace, + service.store, + lease_seconds=0.15, + ) + recovered.start() + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + current = service.store.get(running.job_id) + if current.lease_owner == recovered._owner: + break + time.sleep(0.05) + else: + raise AssertionError("live worker lease was not recovered") + + process.terminate() + process.wait(timeout=5) + service.manager = recovered + terminal = _wait_terminal(service, started.job.job_id) + + assert terminal.state is JobState.INTERRUPTED + assert "manager_recovered" in { + event["event_type"] + for event in service.store.list_events( + started.job.job_id, + limit=1000, + ) + } + finally: + if process.poll() is None: + process.kill() + process.wait(timeout=5) + + +def test_startup_recovery_waits_for_the_observed_peer_lease( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + with patch.object(EvaluationJobManager, "enqueue"): + started = service.start("demo.yaml", request_id="peer-lease") + claimed = service.store.claim_next( + lease_owner="peer-manager", + lease_seconds=10, + max_active_jobs=1, + ) + assert claimed is not None + recovering = EvaluationJobManager( + service.workspace, + service.store, + lease_seconds=1, + ) + + with ( + patch.object( + service.store, + "list_nonterminal_records", + side_effect=[(claimed,), ()], + ), + patch.object(EvaluationJobManager, "enqueue"), + patch( + "assert_ai.services.evaluations.time.sleep" + ) as sleep, + ): + recovering._recover_startup() + + sleep.assert_called_once() + assert sleep.call_args.args[0] > 1 + + +def test_scheduler_sweep_releases_a_dead_cancelling_job( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + with patch.object(EvaluationJobManager, "enqueue"): + first = service.start("demo.yaml", request_id="first") + second = service.start("demo.yaml", request_id="second") + claimed = service.store.claim_next( + lease_owner=service.manager._owner, + lease_seconds=30, + max_active_jobs=1, + ) + assert claimed is not None + running = service.store.mark_running( + claimed.job_id, + lease_owner=service.manager._owner, + pid=2_147_483_647, + process_create_time=1, + lease_seconds=30, + ) + service.store.request_cancel(running.job_id) + + service.manager._sweep_cancelling_jobs() + + assert service.store.get(first.job.job_id).state is JobState.CANCELLED + next_job = service.store.claim_next( + lease_owner="next-manager", + lease_seconds=30, + max_active_jobs=1, + ) + assert next_job is not None + assert next_job.job_id == second.job.job_id + + +def test_failed_job_retry_is_idempotent_and_uses_immutable_snapshot( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + document = _write_inference_fixture(tmp_path) + document["pipeline"]["inference"]["test_set_path"] = "missing.jsonl" + saved = configs.save_config("demo.yaml", document=document) + started = service.start("demo.yaml", request_id="original") + failed = _wait_terminal(service, started.job.job_id) + assert failed.state is JobState.FAILED + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + expected_etag=saved.etag, + ) + + retried = service.retry( + failed.job_id, + request_id="retry-request", + ) + Path(service.store.get(failed.job_id).snapshot_path).unlink() + repeated = service.retry( + failed.job_id, + request_id="retry-request", + ) + retry_terminal = _wait_terminal(service, retried.job.job_id) + + assert retried.created is True + assert repeated.created is False + assert repeated.job.job_id == retried.job.job_id + assert retried.job.retry_of == failed.job_id + assert retried.job.run_id != failed.run_id + assert retry_terminal.state is JobState.FAILED + retry_record = service.store.get(retried.job.job_id) + request = json.loads( + Path(retry_record.request_path).read_text(encoding="utf-8") + ) + snapshot = Path(retry_record.snapshot_path).read_text(encoding="utf-8") + assert request["retry_of"] == failed.job_id + assert request["force_stages"] == ["inference"] + assert "missing.jsonl" in snapshot + + +def test_retry_moves_back_to_corrupt_upstream_jsonl( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + started = service.start("demo.yaml", request_id="corrupt-source") + completed = _wait_terminal(service, started.job.job_id) + record = replace( + service.store.get(completed.job_id), + state=JobState.FAILED, + failed_stage="judge", + ) + assert record.run_id is not None + run_root = ( + tmp_path + / "artifacts" + / "results" + / record.suite_id + / record.run_id + ) + (run_root / "inference_set.jsonl").write_text( + '{"type":"prompt","test_case_id":"case-1"}\n{"type":', + encoding="utf-8", + ) + document = { + "pipeline": { + "inference": {"enabled": True}, + "judge": {"enabled": True}, + } + } + + assert service._retry_stage(record, document) == "inference" + + +def test_cancelled_worker_result_may_omit_identity_before_setup( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + with patch.object(EvaluationJobManager, "enqueue"): + started = service.start("demo.yaml", request_id="cancel-before-setup") + claimed = service.store.claim_next( + lease_owner=service.manager._owner, + lease_seconds=30, + max_active_jobs=1, + ) + assert claimed is not None + running = service.store.mark_running( + claimed.job_id, + lease_owner=service.manager._owner, + pid=2_147_483_647, + process_create_time=1, + lease_seconds=30, + ) + request = json.loads(Path(running.request_path).read_text(encoding="utf-8")) + write_json( + Path(running.request_path).parent / "result.json", + { + "schema_version": 1, + "job_id": running.job_id, + "result_token": request["result_token"], + "run_result": { + "state": "cancelled", + "exit_code": 130, + "failed_stage": "inference", + }, + }, + ) + + terminal = service.get(running.job_id) + + assert terminal.state is JobState.CANCELLED + assert terminal.terminal_result is not None + assert terminal.terminal_result.failed_stage == "inference" diff --git a/tests/test_job_store.py b/tests/test_job_store.py index 81dcb3926..8ff6ee465 100644 --- a/tests/test_job_store.py +++ b/tests/test_job_store.py @@ -3,6 +3,7 @@ from __future__ import annotations +import sqlite3 from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -21,6 +22,7 @@ def _new_job( suite_id: str | None = None, run_id: str | None = None, resource_keys: tuple[str, ...] = (), + retry_of: str | None = None, ) -> NewJob: return NewJob( job_id=f"job-{suffix}", @@ -33,6 +35,7 @@ def _new_job( snapshot_path=f"artifacts/mcp/jobs/job-{suffix}/config.yaml", request_path=f"artifacts/mcp/jobs/job-{suffix}/request.json", resource_keys=resource_keys, + retry_of=retry_of, ) @@ -219,3 +222,203 @@ def create(index: int) -> tuple[str, bool]: assert sum(created for _, created in outcomes) == 1 assert len({job_id for job_id, _ in outcomes}) == 1 + + +def test_cancel_queued_job_is_terminal_and_idempotent( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get(_new_job("one"), max_queued_jobs=10) + + cancelled = store.request_cancel("job-one") + repeated = store.request_cancel("job-one") + + assert cancelled.state is JobState.CANCELLED + assert cancelled.cancel_requested_at is not None + assert cancelled.ended_at is not None + assert cancelled.result == { + "state": "cancelled", + "exit_code": 130, + "failed_stage": None, + "error_code": None, + "error_message": None, + } + assert repeated == cancelled + + +def test_cancel_active_job_preserves_lock_until_terminal( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get( + _new_job("one", resource_keys=("suite:shared",)), + max_queued_jobs=10, + ) + claimed = store.claim_next( + lease_owner="manager-a", + lease_seconds=30, + max_active_jobs=1, + ) + assert claimed is not None + + cancelling = store.request_cancel(claimed.job_id) + + assert cancelling.state is JobState.CANCELLING + assert ( + store.claim_next( + lease_owner="manager-b", + lease_seconds=30, + max_active_jobs=1, + ) + is None + ) + terminal = store.mark_terminal( + claimed.job_id, + state=JobState.CANCELLED, + exit_code=130, + failed_stage=None, + error_code=None, + error_message=None, + result={"state": "cancelled", "exit_code": 130}, + run_root=None, + ) + assert terminal.state is JobState.CANCELLED + + +def test_expired_active_lease_can_be_adopted(tmp_path: Path) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get(_new_job("one"), max_queued_jobs=10) + claimed = store.claim_next( + lease_owner="manager-a", + lease_seconds=0.01, + max_active_jobs=1, + ) + assert claimed is not None + running = store.mark_running( + claimed.job_id, + lease_owner="manager-a", + pid=123, + process_create_time=456, + lease_seconds=0.01, + ) + assert running.state is JobState.RUNNING + + import time + + time.sleep(0.02) + adopted = store.adopt_lease( + running.job_id, + lease_owner="manager-b", + lease_seconds=30, + ) + + assert adopted is not None + assert adopted.lease_owner == "manager-b" + assert adopted.revision == running.revision + 1 + + +def test_retry_provenance_is_persisted(tmp_path: Path) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + created = store.create_or_get( + _new_job("retry", retry_of="job-original"), + max_queued_jobs=10, + ) + + assert created.record.retry_of == "job-original" + assert JobStore(tmp_path / "jobs.sqlite3").get( + created.record.job_id + ).retry_of == "job-original" + + +def test_v1_store_is_migrated_before_a_mutating_operation( + tmp_path: Path, +) -> None: + path = tmp_path / "jobs.sqlite3" + with sqlite3.connect(path) as connection: + connection.executescript( + """ + CREATE TABLE jobs( + job_id TEXT PRIMARY KEY, + idempotency_key TEXT NOT NULL UNIQUE, + request_hash TEXT NOT NULL, + kind TEXT NOT NULL, + state TEXT NOT NULL, + created_at TEXT NOT NULL, + started_at TEXT, + ended_at TEXT, + suite_id TEXT NOT NULL, + run_id TEXT, + config_ref TEXT NOT NULL, + config_sha256 TEXT NOT NULL, + snapshot_path TEXT NOT NULL, + request_path TEXT NOT NULL, + run_root TEXT, + pid INTEGER, + process_create_time REAL, + exit_code INTEGER, + failed_stage TEXT, + error_code TEXT, + error_message TEXT, + cancel_requested_at TEXT, + result_json TEXT, + resource_keys_json TEXT NOT NULL, + revision INTEGER NOT NULL DEFAULT 0, + lease_owner TEXT, + lease_expires_at TEXT + ); + CREATE TABLE job_events( + job_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + timestamp TEXT NOT NULL, + event_type TEXT NOT NULL, + payload_json TEXT NOT NULL, + PRIMARY KEY(job_id, sequence) + ); + CREATE TABLE resource_locks( + resource_key TEXT PRIMARY KEY, + job_id TEXT NOT NULL, + acquired_at TEXT NOT NULL, + lease_expires_at TEXT NOT NULL + ); + INSERT INTO jobs( + job_id, idempotency_key, request_hash, kind, state, + created_at, suite_id, run_id, config_ref, config_sha256, + snapshot_path, request_path, resource_keys_json + ) VALUES ( + 'job-v1', 'request-v1', 'hash-v1', 'evaluation', 'queued', + '2026-01-01T00:00:00+00:00', 'suite-v1', 'run-v1', + 'demo.yaml', 'sha256:v1', 'config.yaml', 'request.json', + '[]' + ); + PRAGMA user_version = 1; + """ + ) + + cancelled = JobStore(path).request_cancel("job-v1") + + assert cancelled.state is JobState.CANCELLED + assert cancelled.retry_of is None + with sqlite3.connect(path) as connection: + assert connection.execute("PRAGMA user_version").fetchone()[0] == 2 + columns = { + row[1] + for row in connection.execute("PRAGMA table_info(jobs)") + } + assert "retry_of" in columns + + +def test_event_retention_prefers_lifecycle_events(tmp_path: Path) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get(_new_job("one"), max_queued_jobs=10) + + for completed in range(1_005): + store.append_event( + "job-one", + "stage_progress", + {"name": "inference", "completed": completed}, + ) + + events = store.list_events("job-one", limit=1000) + assert len(events) == 1000 + assert events[0]["event_type"] == "queued" + assert events[-1]["payload"]["completed"] == 1004 diff --git a/tests/test_mcp_cli.py b/tests/test_mcp_cli.py index 42f7f7a10..e50739868 100644 --- a/tests/test_mcp_cli.py +++ b/tests/test_mcp_cli.py @@ -67,6 +67,8 @@ def test_mcp_serve_forwards_resolved_options() -> None: "9", "--max-job-log-bytes", "5000", + "--cancellation-grace-seconds", + "2.5", "--max-prompt-sample-size", "12", "--max-scenario-sample-size", @@ -95,6 +97,7 @@ def test_mcp_serve_forwards_resolved_options() -> None: assert create_kwargs["max_active_jobs"] == 2 assert create_kwargs["max_queued_jobs"] == 9 assert create_kwargs["max_job_log_bytes"] == 5000 + assert create_kwargs["cancellation_grace_seconds"] == 2.5 assert create_kwargs["max_prompt_sample_size"] == 12 assert create_kwargs["max_scenario_sample_size"] == 13 assert create_kwargs["allowed_model_patterns"] == ("azure/*",) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 92110d317..d6d2aad64 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -58,6 +58,8 @@ "design_config", "probe_target", "start_evaluation", + "cancel_job", + "retry_job", } EXPECTED_RESOURCE_TEMPLATES = { @@ -489,6 +491,7 @@ async def run() -> object: assert result.structured_content["limits"]["max_active_jobs"] == 1 assert result.structured_content["limits"]["max_queued_jobs"] == 100 assert result.structured_content["limits"]["max_job_log_bytes"] == 1024 * 1024 + assert result.structured_content["limits"]["cancellation_grace_seconds"] == 10.0 assert result.structured_content["limits"]["max_prompt_sample_size"] == 100_000 assert result.structured_content["limits"]["max_scenario_sample_size"] == 100_000 assert result.structured_content["limits"]["model_allowlist_enabled"] is True @@ -548,17 +551,17 @@ async def run() -> list[Any]: "compare_runs": "f7bfeca051f8f81bf3621936588ed906332076a3a34b550090f87c2944656ce5", "get_config": "bf38188871cb818e0b0cf6e28183aa728ed8d041923a158f593832d2459bd13a", "get_config_schema": "cca1d3a48240e20eff93a123b34d7ba92df3ed1df87f57f9eb217aa21515ec26", - "get_job": "76794436d4665712dfbd226a4c44738f1b4e8ab6ff3ed3f7eb184311f1f60cb0", + "get_job": "d7a6f643b5fceebcf28b995c8ef1a5aa72bb326551b5546af53ca5b43721a27f", "get_preset": "81db6723ad5065ce8a0a402d29dc2f9df7657d302e3ebe8b377f54c9d62353d0", "get_run": "e5216cd0085d049f8b49c54add913b6f83756c4ce59995317fe63e010ea44936", - "get_server_info": "51e1d08335b4c8c6cb5eb70b6857563ab2dead550347f83dac64d89d8d417069", + "get_server_info": "bae685a1691062b93cc8377d829b779807ad729c576d99cba356f642292a9ed1", "get_suite": "8f629c93e02b656052f637c3cbba9217834315a693c4f7935f6d961203b46fd0", "get_test_case": "11380555caaa71d5992923815a499fc08b368c02f4d4836e4761630654589148", "get_transcript": "aa09669e0cb99202e8dec0b858b4faa41742ecb616351c3956b0d0bd488717e8", "list_artifacts": "3d3bede0b7209401b15d1f39d82671092c3097a05cd901122bd46c3c42edebfc", "list_configs": "92f78db2533034e6bf80e1d95089460acdd40a18468d4eb06fdf055726dfef19", "list_failures": "d3cc3f3bcc86c110754673297d28ac1e5ccbf698668c997de0bba2d0cbd425e2", - "list_jobs": "4570a8790f6f5c42fa015c49056111f4b7f00744a2300e3a79af9a68f08d2530", + "list_jobs": "dd2a59740c543efda3d7141af96678b7a321b0ea8dbe3524779e79b297073f07", "list_presets": "55faa31adbf7f689eb5efbf1211fa73b2474d1a0e4549ec0a836ea69919c46b1", "list_runs": "7280687daafcd7ff5d89756c9584ca06c432a44f5a98ce8ff3ae0e4427dcf40b", "list_scores": "5c1951a3a3b91089b68b30e970a1b13f59bc2659a2c234db4451cbe2d5362a4d", @@ -588,6 +591,8 @@ async def run() -> dict[str, Any]: "design_config": (True, False, False, True), "probe_target": (True, False, False, True), "start_evaluation": (False, True, True, True), + "cancel_job": (False, True, True, False), + "retry_job": (False, True, True, True), } expected_digests = { "validate_config": ( @@ -606,7 +611,13 @@ async def run() -> dict[str, Any]: "406d97ba84821a4e2661779dcc1211d208d2485013f8dea761c04f1fcdf59e63" ), "start_evaluation": ( - "b594f1ef3510f31503960291e7ed5b57967876b5c5d4d628b0d5ee86cc000c0c" + "143c75380354a425936844f627569d66e721f604d4722af0cd40dc2959d4084a" + ), + "cancel_job": ( + "714906227e19526bca0b5ba965574812c7ff427557530644138092e230229e0f" + ), + "retry_job": ( + "db52c7589a88bcbdef299b210c03f134050beea74117ab95e9a7f1d8e182a455" ), } @@ -847,6 +858,155 @@ async def run() -> dict[str, Any]: assert "[REDACTED]" in results["job_log"] +def test_mcp_can_cancel_a_running_evaluation(tmp_path: Path) -> None: + _seed_evaluation_workspace(tmp_path) + (tmp_path / "agent.py").write_text( + "import time\n" + "def run(message, *, history=None):\n" + " del history\n" + " time.sleep(1)\n" + " return message\n", + encoding="utf-8", + ) + + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="full", + cancellation_grace_seconds=3, + ) + async with Client(build_server(options), raise_exceptions=True) as client: + started = await client.call_tool( + "start_evaluation", + { + "config_ref": "job.yaml", + "request_id": "cancel-through-mcp", + }, + ) + job_id = started.structured_content["job"]["job_id"] + deadline = asyncio.get_running_loop().time() + 15 + while True: + detail = await client.call_tool( + "get_job", + {"job_id": job_id}, + ) + if detail.structured_content["state"] == "running": + break + if asyncio.get_running_loop().time() >= deadline: + raise AssertionError("MCP evaluation did not start") + await asyncio.sleep(0.05) + cancelling = await client.call_tool( + "cancel_job", + {"job_id": job_id}, + ) + while True: + detail = await client.call_tool( + "get_job", + {"job_id": job_id}, + ) + if detail.structured_content["state"] == "cancelled": + break + if asyncio.get_running_loop().time() >= deadline: + raise AssertionError("MCP evaluation did not cancel") + await asyncio.sleep(0.05) + return { + "cancelling": cancelling.structured_content, + "terminal": detail.structured_content, + } + + results = asyncio.run(run()) + + assert results["cancelling"]["state"] == "cancelling" + assert results["cancelling"]["cancel_requested_at"] is not None + assert results["terminal"]["state"] == "cancelled" + assert results["terminal"]["terminal_result"]["exit_code"] == 130 + assert results["terminal"]["stages"]["inference"] == "cancelled" + + +def test_mcp_retry_is_idempotent_and_records_provenance( + tmp_path: Path, +) -> None: + _seed_evaluation_workspace(tmp_path) + config_path = tmp_path / "evals" / "job.yaml" + document = json.loads(config_path.read_text(encoding="utf-8")) + document["pipeline"]["inference"]["test_set_path"] = "missing.jsonl" + config_path.write_text(json.dumps(document), encoding="utf-8") + + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="full", + ) + async with Client(build_server(options), raise_exceptions=True) as client: + started = await client.call_tool( + "start_evaluation", + { + "config_ref": "job.yaml", + "request_id": "retry-original", + }, + ) + original_id = started.structured_content["job"]["job_id"] + deadline = asyncio.get_running_loop().time() + 20 + while True: + original = await client.call_tool( + "get_job", + {"job_id": original_id}, + ) + if original.structured_content["state"] == "failed": + break + if asyncio.get_running_loop().time() >= deadline: + raise AssertionError("Original MCP evaluation did not fail") + await asyncio.sleep(0.05) + retried = await client.call_tool( + "retry_job", + { + "job_id": original_id, + "request_id": "retry-attempt", + }, + ) + repeated = await client.call_tool( + "retry_job", + { + "job_id": original_id, + "request_id": "retry-attempt", + }, + ) + retry_id = retried.structured_content["job"]["job_id"] + deadline = asyncio.get_running_loop().time() + 20 + while True: + retry_detail = await client.call_tool( + "get_job", + {"job_id": retry_id}, + ) + if retry_detail.structured_content["state"] == "failed": + break + if asyncio.get_running_loop().time() >= deadline: + raise AssertionError("Retried MCP evaluation did not fail") + await asyncio.sleep(0.05) + not_cancellable = await client.call_tool( + "cancel_job", + {"job_id": original_id}, + ) + return { + "original_id": original_id, + "retried": retried, + "repeated": repeated, + "not_cancellable": not_cancellable, + } + + results = asyncio.run(run()) + + retried = results["retried"].structured_content + repeated = results["repeated"].structured_content + assert retried["created"] is True + assert retried["job"]["retry_of"] == results["original_id"] + assert repeated["created"] is False + assert repeated["job"]["job_id"] == retried["job"]["job_id"] + assert '"code":"JOB_NOT_CANCELLABLE"' in _error_text( + results["not_cancellable"] + ) + + def test_design_config_returns_an_unpersisted_draft(tmp_path: Path) -> None: draft = ConfigDraft( yaml=( diff --git a/tests/test_run_result.py b/tests/test_run_result.py index e027c6008..38f3f70c4 100644 --- a/tests/test_run_result.py +++ b/tests/test_run_result.py @@ -8,9 +8,11 @@ from tempfile import TemporaryDirectory from unittest.mock import patch +import pytest import yaml from assert_ai.core.model_client import LLMInputError +from assert_ai.core.run_control import RunCancelled, RunControl from assert_ai.core.run_result import RunState from assert_ai.core.workspace import WorkspaceService from assert_ai.runner import ( @@ -195,3 +197,91 @@ def test_unexpected_setup_failure_is_returned_not_raised() -> None: assert result.exit_code == 1 assert result.error_code == "INTERNAL" assert result.error_message == "Unexpected pipeline setup error" + + +def test_cooperative_cancellation_writes_terminal_manifest_and_events( + tmp_path: Path, +) -> None: + config_path = tmp_path / "config.yaml" + results = tmp_path / "results" + config_path.write_text( + "\n".join( + [ + "suite: suite-a", + "run: run-a", + f"results_dir: {results}", + "pipeline:", + " inference:", + " target:", + " callable: agent:run", + " test_set_path: fixture.jsonl", + ] + ) + + "\n", + encoding="utf-8", + ) + + class Observer: + def __init__(self) -> None: + self.events: list[tuple[str, object]] = [] + + def pipeline_started(self, event: object) -> None: + self.events.append(("pipeline_started", event)) + + def stage_planned(self, event: object) -> None: + self.events.append(("stage_planned", event)) + + def stage_started(self, event: object) -> None: + self.events.append(("stage_started", event)) + + def stage_progress(self, event: object) -> None: + self.events.append(("stage_progress", event)) + + def stage_finished(self, event: object) -> None: + self.events.append(("stage_finished", event)) + + def pipeline_finished(self, event: object) -> None: + self.events.append(("pipeline_finished", event)) + + observer = Observer() + with patch("assert_ai.stages.inference.run") as stage: + result = run_pipeline_result( + config=str(config_path), + control=RunControl(cancel_requested=lambda: True), + observer=observer, + ) + + assert result.state is RunState.CANCELLED + assert result.exit_code == 130 + assert result.failed_stage == "inference" + stage.assert_not_called() + manifest = json.loads( + (results / "suite-a" / "run-a" / "manifest.json").read_text( + encoding="utf-8" + ) + ) + assert manifest["status"] == "cancelled" + assert manifest["stages"]["inference"] == "cancelled" + assert [name for name, _ in observer.events] == [ + "pipeline_started", + "stage_planned", + "stage_started", + "stage_progress", + "stage_finished", + "pipeline_finished", + ] + + +def test_run_control_acknowledges_cancellation_once() -> None: + acknowledged: list[str | None] = [] + control = RunControl( + cancel_requested=lambda: True, + cancel_acknowledged=acknowledged.append, + ) + + with pytest.raises(RunCancelled): + control.raise_if_cancelled(stage="inference") + with pytest.raises(RunCancelled): + control.raise_if_cancelled(stage="judge") + + assert acknowledged == ["inference"] From 30b29f267b9ea2cfbba43deba5939287fdecffd3 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Wed, 26 Aug 2026 12:52:31 -0700 Subject: [PATCH 11/16] Add MCP curation and trace judging Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/core/artifact_cache.py | 56 +- assert_ai/core/io.py | 7 +- assert_ai/core/otel.py | 187 +++- assert_ai/mcp/_command.py | 9 + assert_ai/mcp/models.py | 33 + assert_ai/mcp/server.py | 47 +- assert_ai/mcp/tools/__init__.py | 10 + assert_ai/mcp/tools/curation.py | 149 +++ assert_ai/mcp/tools/inspect.py | 13 + assert_ai/mcp/tools/jobs.py | 21 +- assert_ai/mcp/tools/traces.py | 98 ++ assert_ai/services/_evaluation_worker.py | 371 ++++++- assert_ai/services/configs.py | 50 +- assert_ai/services/curation.py | 1187 ++++++++++++++++++++++ assert_ai/services/evaluations.py | 854 +++++++++++++++- assert_ai/services/job_models.py | 22 +- assert_ai/services/job_store.py | 253 ++++- assert_ai/services/locking.py | 65 ++ assert_ai/stages/judge.py | 7 + tests/test_curation_service.py | 555 ++++++++++ tests/test_evaluation_service.py | 533 +++++++++- tests/test_framework_agnostic.py | 5 + tests/test_job_store.py | 146 ++- tests/test_mcp_cli.py | 3 + tests/test_mcp_server.py | 473 ++++++++- 25 files changed, 5005 insertions(+), 149 deletions(-) create mode 100644 assert_ai/mcp/tools/curation.py create mode 100644 assert_ai/mcp/tools/traces.py create mode 100644 assert_ai/services/curation.py create mode 100644 assert_ai/services/locking.py create mode 100644 tests/test_curation_service.py diff --git a/assert_ai/core/artifact_cache.py b/assert_ai/core/artifact_cache.py index 972b364f1..43597b9f9 100644 --- a/assert_ai/core/artifact_cache.py +++ b/assert_ai/core/artifact_cache.py @@ -220,6 +220,24 @@ def prepare_artifact_plan( if reusable is not None: return reusable + return allocate_artifact_plan( + ctx=ctx, + stage_name=stage_name, + fingerprint=fingerprint, + ) + + +def allocate_artifact_plan( + *, + ctx: dict[str, Any], + stage_name: str, + fingerprint: ArtifactFingerprint, +) -> ArtifactPlan: + """Reserve a fresh immutable artifact version.""" + if stage_name not in CACHEABLE_STAGES: + raise ValueError(f"unsupported cacheable stage: {stage_name}") + suite_root = _managed_suite_root(ctx) + stage_root = _artifact_stage_root(ctx, suite_root, stage_name) version, artifact_dir = _allocate_version_dir(stage_root) return ArtifactPlan( stage_name=stage_name, @@ -522,7 +540,14 @@ def activate_latest_artifacts( ) -def finalize_artifact_plan(ctx: dict[str, Any], plan: ArtifactPlan) -> dict[str, Any]: +def finalize_artifact_plan( + ctx: dict[str, Any], + plan: ArtifactPlan, + *, + provenance: dict[str, Any] | None = None, + activate: bool = True, + preserve_local_edits: bool = True, +) -> dict[str, Any]: """Write sidecar metadata and update latest/compatibility artifacts.""" suite_root = _managed_suite_root(ctx) @@ -570,6 +595,8 @@ def finalize_artifact_plan(ctx: dict[str, Any], plan: ArtifactPlan) -> dict[str, }, "file_hashes": file_hashes, } + if provenance is not None: + metadata["provenance"] = _normalize_value(provenance) metadata_path = _managed_output_path( ctx, artifact_dir / ARTIFACT_METADATA_FILE, @@ -580,8 +607,14 @@ def finalize_artifact_plan(ctx: dict[str, Any], plan: ArtifactPlan) -> dict[str, write_json(metadata_path, metadata) ref = artifact_ref(ctx=ctx, plan=plan, metadata=metadata) ctx.setdefault("artifact_versions", {})[plan.stage_name] = ref - update_latest(ctx, plan.stage_name, ref) - refresh_compatibility_files(ctx, plan.stage_name, output_paths) + if activate: + update_latest(ctx, plan.stage_name, ref) + refresh_compatibility_files( + ctx, + plan.stage_name, + output_paths, + preserve_local_edits=preserve_local_edits, + ) return ref @@ -657,6 +690,8 @@ def refresh_compatibility_files( ctx: dict[str, Any], stage_name: str, output_paths: dict[str, Path], + *, + preserve_local_edits: bool = True, ) -> None: """Copy selected version outputs back to legacy suite-root filenames. @@ -702,7 +737,10 @@ def refresh_compatibility_files( expected_root=suite_root, reject_links=True, ) - if _is_local_edit(suite_root, stage_name, dest, path): + if ( + preserve_local_edits + and _is_local_edit(suite_root, stage_name, dest, path) + ): log.warning( "[%s] Preserving local edits to %s: contents differ from the " "cached artifact at %s and do not match any previously cached " @@ -785,6 +823,14 @@ def _was_cached_artifact( def update_latest(ctx: dict[str, Any], stage_name: str, ref: dict[str, Any]) -> None: + update_latest_artifacts(ctx, {stage_name: ref}) + + +def update_latest_artifacts( + ctx: dict[str, Any], + refs: dict[str, dict[str, Any]], +) -> None: + """Atomically activate one or more artifact references.""" suite_root = _managed_suite_root(ctx) latest_path = _managed_output_path( ctx, @@ -801,7 +847,7 @@ def update_latest(ctx: dict[str, Any], stage_name: str, ref: dict[str, Any]) -> if not isinstance(artifacts, dict): artifacts = {} latest["artifacts"] = artifacts - artifacts[stage_name] = ref + artifacts.update(refs) write_json(latest_path, latest) diff --git a/assert_ai/core/io.py b/assert_ai/core/io.py index 2a434a35c..f853ca81a 100644 --- a/assert_ai/core/io.py +++ b/assert_ai/core/io.py @@ -52,6 +52,11 @@ def append_jsonl_row(path: Path, row: Dict[str, Any]) -> None: def write_text_atomic(path: Path, text: str) -> None: """Atomically replace a UTF-8 text file after flushing its contents.""" + write_bytes_atomic(path, text.encode("utf-8")) + + +def write_bytes_atomic(path: Path, data: bytes) -> None: + """Atomically replace a file after flushing its exact bytes.""" path.parent.mkdir(parents=True, exist_ok=True) tmp_name: str | None = None try: @@ -62,7 +67,7 @@ def write_text_atomic(path: Path, text: str) -> None: suffix=".tmp", delete=False, ) as handle: - handle.write(text.encode("utf-8")) + handle.write(data) handle.flush() os.fsync(handle.fileno()) tmp_name = handle.name diff --git a/assert_ai/core/otel.py b/assert_ai/core/otel.py index 3585ff5f3..5400d49a5 100644 --- a/assert_ai/core/otel.py +++ b/assert_ai/core/otel.py @@ -137,21 +137,56 @@ def parse_otel_traces( } """ spans = _parse_otlp_json(Path(path)) + return _inference_rows_from_spans(spans, group_by=group_by) + + +def parse_otel_trace_document( + document: dict[str, Any], + *, + group_by: str = "session.id", +) -> list[dict[str, Any]]: + """Parse an already-snapshotted OTLP JSON document into inference rows.""" + spans = _parse_otlp_document(document) + return _inference_rows_from_spans(spans, group_by=group_by) + + +def _inference_rows_from_spans( + spans: list[OTelSpan], + *, + group_by: str, +) -> list[dict[str, Any]]: grouped = _group_spans(spans, group_by) rows = [] for session_id, session_spans in grouped.items(): session_spans.sort(key=lambda s: s.start_time_ns) events, aggregate = _spans_to_events(session_spans) + trace_spans: dict[str, list[str]] = {} + for span in session_spans: + if span.trace_id: + trace_spans.setdefault(span.trace_id, []) + if span.span_id: + trace_spans[span.trace_id].append(span.span_id) + trace_refs = [ + { + "trace_id": trace_id, + "span_ids": list(dict.fromkeys(span_ids)), + } + for trace_id, span_ids in sorted(trace_spans.items()) + ] rows.append({ "metadata": { "type": "otel_import", "session_id": session_id, "runtime_mode": "otel_traced", + "trace_refs": trace_refs, }, "events": events, - "raw": aggregate, + "raw": { + **aggregate, + "trace_refs": trace_refs, + }, }) return rows @@ -166,26 +201,124 @@ def _parse_otlp_json(path: Path) -> list[OTelSpan]: except json.JSONDecodeError as exc: raise ValueError(f"Malformed JSON in OTLP trace file {path}: {exc}") from exc + if not isinstance(data, dict): + raise ValueError(f"OTLP trace file must contain a JSON object: {path}") + return _parse_otlp_document(data) + + +def _parse_otlp_document(data: dict[str, Any]) -> list[OTelSpan]: + """Parse one decoded OTLP JSON document.""" spans: list[OTelSpan] = [] - for resource_span in data.get("resourceSpans", []): - for scope_span in resource_span.get("scopeSpans", []): - for raw in scope_span.get("spans", []): - attrs = _flatten_attributes(raw.get("attributes", [])) - spans.append(OTelSpan( - trace_id=raw.get("traceId", ""), - span_id=raw.get("spanId", ""), - parent_span_id=raw.get("parentSpanId"), - name=raw.get("name", ""), - kind=_classify_span_kind(attrs), - start_time_ns=int(raw.get("startTimeUnixNano", 0)), - end_time_ns=int(raw.get("endTimeUnixNano", 0)), - attributes=attrs, - status=raw.get("status", {}).get("code", "OK"), - events=raw.get("events", []) or [], - )) + resource_spans = _otel_array(data.get("resourceSpans", []), "resourceSpans") + for resource_index, resource_span in enumerate(resource_spans): + resource = _otel_object( + resource_span, + f"resourceSpans[{resource_index}]", + ) + scope_spans = _otel_array( + resource.get("scopeSpans", []), + f"resourceSpans[{resource_index}].scopeSpans", + ) + for scope_index, scope_span in enumerate(scope_spans): + scope = _otel_object( + scope_span, + f"resourceSpans[{resource_index}].scopeSpans[{scope_index}]", + ) + raw_spans = _otel_array( + scope.get("spans", []), + ( + f"resourceSpans[{resource_index}]." + f"scopeSpans[{scope_index}].spans" + ), + ) + for span_index, raw_span in enumerate(raw_spans): + location = ( + f"resourceSpans[{resource_index}]." + f"scopeSpans[{scope_index}].spans[{span_index}]" + ) + raw = _otel_object(raw_span, location) + attributes = _otel_array( + raw.get("attributes") or [], + f"{location}.attributes", + ) + status = _otel_object( + raw.get("status") or {}, + f"{location}.status", + ) + events = _otel_array( + raw.get("events") or [], + f"{location}.events", + ) + attrs = _flatten_attributes(attributes) + spans.append( + OTelSpan( + trace_id=_otel_string( + raw.get("traceId", ""), + f"{location}.traceId", + ), + span_id=_otel_string( + raw.get("spanId", ""), + f"{location}.spanId", + ), + parent_span_id=_otel_optional_string( + raw.get("parentSpanId"), + f"{location}.parentSpanId", + ), + name=_otel_string( + raw.get("name", ""), + f"{location}.name", + ), + kind=_classify_span_kind(attrs), + start_time_ns=_otel_integer( + raw.get("startTimeUnixNano", 0), + f"{location}.startTimeUnixNano", + ), + end_time_ns=_otel_integer( + raw.get("endTimeUnixNano", 0), + f"{location}.endTimeUnixNano", + ), + attributes=attrs, + status=status.get("code", "OK"), + events=events, + ) + ) return spans +def _otel_array(value: Any, field_name: str) -> list[Any]: + if not isinstance(value, list): + raise ValueError(f"OTLP field {field_name} must be an array") + return value + + +def _otel_object(value: Any, field_name: str) -> dict[str, Any]: + if not isinstance(value, dict): + raise ValueError(f"OTLP field {field_name} must be an object") + return value + + +def _otel_string(value: Any, field_name: str) -> str: + if not isinstance(value, str): + raise ValueError(f"OTLP field {field_name} must be a string") + return value + + +def _otel_optional_string(value: Any, field_name: str) -> str | None: + if value is None: + return None + return _otel_string(value, field_name) + + +def _otel_integer(value: Any, field_name: str) -> int: + if isinstance(value, int) and not isinstance(value, bool) and value >= 0: + return value + if isinstance(value, str) and value.isdecimal(): + return int(value) + raise ValueError( + f"OTLP field {field_name} must be a non-negative integer" + ) + + def _classify_span_kind(attrs: dict[str, Any]) -> str: """Determine an ASSERT span kind from OpenInference or GenAI attributes. @@ -211,9 +344,13 @@ def _classify_span_kind(attrs: dict[str, Any]) -> str: def _flatten_attributes(attrs: list[dict]) -> dict[str, Any]: """Convert OTLP attribute array [{key, value}] to flat dict.""" result: dict[str, Any] = {} - for attr in attrs: + for index, attr in enumerate(attrs): + if not isinstance(attr, dict): + raise ValueError(f"OTLP attribute {index} must be an object") key = attr.get("key", "") value = attr.get("value", {}) + if not isinstance(key, str) or not isinstance(value, dict): + raise ValueError(f"OTLP attribute {index} has an invalid key or value") if "stringValue" in value: result[key] = value["stringValue"] elif "intValue" in value: @@ -223,14 +360,26 @@ def _flatten_attributes(attrs: list[dict]) -> dict[str, Any]: elif "boolValue" in value: result[key] = value["boolValue"] elif "arrayValue" in value: + array_value = value["arrayValue"] + if not isinstance(array_value, dict): + raise ValueError( + f"OTLP attribute {index} has an invalid arrayValue" + ) + values = array_value.get("values", []) + if not isinstance(values, list): + raise ValueError( + f"OTLP attribute {index} arrayValue.values must be an array" + ) result[key] = [ - _extract_value(v) for v in value["arrayValue"].get("values", []) + _extract_value(v) for v in values ] return result def _extract_value(value_obj: dict) -> Any: """Extract a scalar value from an OTLP Value object.""" + if not isinstance(value_obj, dict): + raise ValueError("OTLP array attribute values must be objects") for key in ("stringValue", "intValue", "doubleValue", "boolValue"): if key in value_obj: return value_obj[key] diff --git a/assert_ai/mcp/_command.py b/assert_ai/mcp/_command.py index 8bf6d67fa..f7c41def7 100644 --- a/assert_ai/mcp/_command.py +++ b/assert_ai/mcp/_command.py @@ -144,6 +144,13 @@ def mcp() -> None: show_default=True, help="Maximum retained bytes for each worker stdout/stderr log.", ) +@click.option( + "--max-trace-input-bytes", + type=click.IntRange(min=1024, max=64 * 1024 * 1024), + default=64 * 1024 * 1024, + show_default=True, + help="Maximum size of one imported OTLP JSON trace file.", +) @click.option( "--cancellation-grace-seconds", type=click.FloatRange(min=0.1), @@ -199,6 +206,7 @@ def serve( max_active_jobs: int, max_queued_jobs: int, max_job_log_bytes: int, + max_trace_input_bytes: int, cancellation_grace_seconds: float, max_prompt_sample_size: int, max_scenario_sample_size: int, @@ -234,6 +242,7 @@ def serve( max_active_jobs=max_active_jobs, max_queued_jobs=max_queued_jobs, max_job_log_bytes=max_job_log_bytes, + max_trace_input_bytes=max_trace_input_bytes, cancellation_grace_seconds=cancellation_grace_seconds, max_prompt_sample_size=max_prompt_sample_size, max_scenario_sample_size=max_scenario_sample_size, diff --git a/assert_ai/mcp/models.py b/assert_ai/mcp/models.py index 321dc8895..cb649a869 100644 --- a/assert_ai/mcp/models.py +++ b/assert_ai/mcp/models.py @@ -64,6 +64,7 @@ class ServerLimits(BaseModel): max_active_jobs: int max_queued_jobs: int max_job_log_bytes: int + max_trace_input_bytes: int cancellation_grace_seconds: float max_prompt_sample_size: int max_scenario_sample_size: int @@ -209,6 +210,37 @@ class ConfigDesignResult(_McpModel): persisted: Literal[False] = False +class TestCaseRevisionInput(_McpModel): + """Partial revision for one stable test-case identity.""" + + test_case_id: str = Field(min_length=1, max_length=255) + updates: dict[str, Any] + + +class CuratedArtifactResult(_McpModel): + """One immutable artifact version created by curation.""" + + artifact_type: str + version: str + etag: str + source_etag: str + source_version: str | None = None + artifact_ref: str + metadata_ref: str + + +class CurationToolResult(_McpModel): + """Result of an atomic suite curation operation.""" + + schema_version: Literal[1] = 1 + suite_id: str + change_summary: str + artifacts: tuple[CuratedArtifactResult, ...] + invalidated_stages: tuple[str, ...] + affected_test_case_ids: tuple[str, ...] = () + warnings: tuple[str, ...] = () + + class SuiteCatalogItem(_McpModel): """Lightweight suite metadata.""" @@ -245,6 +277,7 @@ class SuiteResult(_McpModel): updated_at: str | None = None run_count: int latest_run: dict[str, Any] | None = None + active_artifact_etags: dict[str, str] = Field(default_factory=dict) resources: dict[str, str] = Field(default_factory=dict) diff --git a/assert_ai/mcp/server.py b/assert_ai/mcp/server.py index 34ae26582..85daf0a22 100644 --- a/assert_ai/mcp/server.py +++ b/assert_ai/mcp/server.py @@ -25,18 +25,23 @@ from assert_ai.mcp.resources import register_inspect_resources from assert_ai.mcp.tools import ( AuthorServices, + CurationServices, InspectServices, JobServices, ProbeServices, register_author_tools, + register_curation_tools, register_design_tools, register_inspect_tools, + register_job_control_tools, register_job_execute_tools, register_job_inspect_tools, register_probe_tools, + register_trace_tools, ) from assert_ai.services.artifacts import ArtifactRepository from assert_ai.services.configs import ConfigService +from assert_ai.services.curation import CurationService from assert_ai.services.evaluations import ( EvaluationJobManager, EvaluationService, @@ -91,6 +96,7 @@ class ServerOptions: max_active_jobs: int = 1 max_queued_jobs: int = 100 max_job_log_bytes: int = 1024 * 1024 + max_trace_input_bytes: int = 64 * 1024 * 1024 cancellation_grace_seconds: float = 10.0 max_prompt_sample_size: int = 100_000 max_scenario_sample_size: int = 100_000 @@ -120,6 +126,10 @@ def __post_init__(self) -> None: raise ValueError( "max_job_log_bytes must be between 4096 and 16777216" ) + if not 1024 <= self.max_trace_input_bytes <= 64 * 1024 * 1024: + raise ValueError( + "max_trace_input_bytes must be between 1024 and 67108864" + ) if self.cancellation_grace_seconds <= 0: raise ValueError( "cancellation_grace_seconds must be positive" @@ -168,6 +178,7 @@ def create( max_active_jobs: int = 1, max_queued_jobs: int = 100, max_job_log_bytes: int = 1024 * 1024, + max_trace_input_bytes: int = 64 * 1024 * 1024, cancellation_grace_seconds: float = 10.0, max_prompt_sample_size: int = 100_000, max_scenario_sample_size: int = 100_000, @@ -197,6 +208,7 @@ def create( max_active_jobs=max_active_jobs, max_queued_jobs=max_queued_jobs, max_job_log_bytes=max_job_log_bytes, + max_trace_input_bytes=max_trace_input_bytes, cancellation_grace_seconds=cancellation_grace_seconds, max_prompt_sample_size=max_prompt_sample_size, max_scenario_sample_size=max_scenario_sample_size, @@ -258,15 +270,23 @@ def build_server(options: ServerOptions) -> MCPServer: path_policy=options.path_policy, expected_root=options.workspace.artifacts_root, ) - execution_enabled = ( - CapabilityGroup.EXECUTE in options.capability_groups - ) + execution_enabled = CapabilityGroup.EXECUTE in options.capability_groups + trace_enabled = CapabilityGroup.TRACE in options.capability_groups + jobs_enabled = execution_enabled or trace_enabled job_manager = EvaluationJobManager( options.workspace, job_store, max_active_jobs=options.max_active_jobs, max_log_bytes=options.max_job_log_bytes, - launch_enabled=execution_enabled, + launch_enabled=jobs_enabled, + job_kinds=tuple( + kind + for kind, enabled in ( + ("evaluation", execution_enabled), + ("trace_judging", trace_enabled), + ) + if enabled + ), cancellation_grace_seconds=options.cancellation_grace_seconds, ) evaluations = EvaluationService( @@ -278,6 +298,7 @@ def build_server(options: ServerOptions) -> MCPServer: default_page_size=options.default_page_size, max_page_size=options.max_page_size, max_queued_jobs=options.max_queued_jobs, + max_trace_input_bytes=options.max_trace_input_bytes, ) job_services = JobServices( workspace=options.workspace, @@ -318,6 +339,7 @@ def get_server_info() -> ServerInfo: max_active_jobs=options.max_active_jobs, max_queued_jobs=options.max_queued_jobs, max_job_log_bytes=options.max_job_log_bytes, + max_trace_input_bytes=options.max_trace_input_bytes, cancellation_grace_seconds=( options.cancellation_grace_seconds ), @@ -394,8 +416,25 @@ def get_server_info() -> ServerInfo: max_response_bytes=options.max_response_bytes, ), ) + if CapabilityGroup.CURATE in options.capability_groups: + register_curation_tools( + server, + CurationServices( + workspace=options.workspace, + curation=CurationService( + options.workspace, + job_store=job_store, + ), + max_response_bytes=options.max_response_bytes, + ), + ) if execution_enabled: register_job_execute_tools(server, job_services) + if jobs_enabled: + register_job_control_tools(server, job_services) + if trace_enabled: + register_trace_tools(server, job_services) + if jobs_enabled: job_manager.start() return server diff --git a/assert_ai/mcp/tools/__init__.py b/assert_ai/mcp/tools/__init__.py index 22d2804c1..df39d4d2d 100644 --- a/assert_ai/mcp/tools/__init__.py +++ b/assert_ai/mcp/tools/__init__.py @@ -10,22 +10,32 @@ register_design_tools, register_probe_tools, ) +from assert_ai.mcp.tools.curation import ( + CurationServices, + register_curation_tools, +) from assert_ai.mcp.tools.inspect import InspectServices, register_inspect_tools from assert_ai.mcp.tools.jobs import ( JobServices, + register_job_control_tools, register_job_execute_tools, register_job_inspect_tools, ) +from assert_ai.mcp.tools.traces import register_trace_tools __all__ = [ "AuthorServices", + "CurationServices", "InspectServices", "JobServices", "ProbeServices", "register_author_tools", + "register_curation_tools", "register_design_tools", "register_inspect_tools", + "register_job_control_tools", "register_job_execute_tools", "register_job_inspect_tools", "register_probe_tools", + "register_trace_tools", ] diff --git a/assert_ai/mcp/tools/curation.py b/assert_ai/mcp/tools/curation.py new file mode 100644 index 000000000..06ef63348 --- /dev/null +++ b/assert_ai/mcp/tools/curation.py @@ -0,0 +1,149 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Versioned generated-artifact curation MCP tools.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Annotated, Any + +from mcp.server import MCPServer +from mcp.types import ToolAnnotations +from pydantic import Field + +from assert_ai.core.workspace import WorkspaceService +from assert_ai.mcp.errors import adapt_tool_errors, invoke_tool +from assert_ai.mcp.models import CurationToolResult, TestCaseRevisionInput +from assert_ai.mcp.sanitize import sanitize_for_mcp +from assert_ai.services.curation import ( + CurationResult, + CurationService, + TestCaseRevision, +) + +_CURATION_ANNOTATIONS = ToolAnnotations( + read_only_hint=False, + destructive_hint=True, + idempotent_hint=False, + open_world_hint=False, +) + + +@dataclass(frozen=True, slots=True) +class CurationServices: + """Dependencies shared by curation tool handlers.""" + + workspace: WorkspaceService + curation: CurationService + max_response_bytes: int + + +def register_curation_tools( + server: MCPServer, + services: CurationServices, +) -> None: + """Register immutable taxonomy and test-case revision tools.""" + + @server.tool( + title="Revise an ASSERT taxonomy", + annotations=_CURATION_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def revise_taxonomy( + suite_id: str, + taxonomy: dict[str, Any], + expected_etag: Annotated[str, Field(min_length=1)], + change_summary: Annotated[str, Field(min_length=1, max_length=500)], + ) -> CurationToolResult: + """Create and activate an immutable taxonomy revision.""" + result = invoke_tool( + lambda: services.curation.revise_taxonomy( + suite_id, + taxonomy, + expected_etag=expected_etag, + change_summary=change_summary, + ), + workspace=services.workspace, + ) + return _result(result, services=services) + + @server.tool( + title="Revise one ASSERT test case", + annotations=_CURATION_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def revise_test_case( + suite_id: str, + test_case_id: Annotated[str, Field(min_length=1, max_length=255)], + updates: dict[str, Any], + expected_etag: Annotated[str, Field(min_length=1)], + change_summary: Annotated[str, Field(min_length=1, max_length=500)], + ) -> CurationToolResult: + """Create and activate one immutable test-case revision.""" + result = invoke_tool( + lambda: services.curation.revise_test_case( + suite_id, + test_case_id, + updates, + expected_etag=expected_etag, + change_summary=change_summary, + ), + workspace=services.workspace, + ) + return _result(result, services=services) + + @server.tool( + title="Revise multiple ASSERT test cases", + annotations=_CURATION_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def bulk_revise_test_cases( + suite_id: str, + revisions: Annotated[ + tuple[TestCaseRevisionInput, ...], + Field(min_length=1), + ], + expected_etag: Annotated[str, Field(min_length=1)], + change_summary: Annotated[str, Field(min_length=1, max_length=500)], + ) -> CurationToolResult: + """Create and activate one immutable multi-test-case revision.""" + result = invoke_tool( + lambda: services.curation.bulk_revise_test_cases( + suite_id, + tuple( + TestCaseRevision( + test_case_id=revision.test_case_id, + updates=revision.updates, + ) + for revision in revisions + ), + expected_etag=expected_etag, + change_summary=change_summary, + ), + workspace=services.workspace, + ) + return _result(result, services=services) + + +def _result( + result: CurationResult, + *, + services: CurationServices, +) -> CurationToolResult: + payload = result.model_dump(mode="json") + return CurationToolResult.model_validate( + sanitize_for_mcp(payload, workspace=services.workspace) + ) diff --git a/assert_ai/mcp/tools/inspect.py b/assert_ai/mcp/tools/inspect.py index 323d7f69e..f5ead84e9 100644 --- a/assert_ai/mcp/tools/inspect.py +++ b/assert_ai/mcp/tools/inspect.py @@ -620,6 +620,19 @@ def _public_suite( workspace: WorkspaceService, ) -> dict[str, Any]: payload = _safe_mapping(summary, workspace=workspace) + sources = payload.get("sources") + artifact_etags: dict[str, str] = {} + if isinstance(sources, dict): + for name in ("taxonomy", "test_set"): + source = sources.get(name) + sha256 = source.get("sha256") if isinstance(source, dict) else None + if ( + isinstance(sha256, str) + and len(sha256) == 64 + and all(character in "0123456789abcdef" for character in sha256) + ): + artifact_etags[name] = f"sha256:{sha256}" + payload["active_artifact_etags"] = artifact_etags for key in ( "artifact_versions", "sources", diff --git a/assert_ai/mcp/tools/jobs.py b/assert_ai/mcp/tools/jobs.py index eb504d9f6..3b8f22645 100644 --- a/assert_ai/mcp/tools/jobs.py +++ b/assert_ai/mcp/tools/jobs.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Persisted evaluation job tools for the ASSERT MCP adapter.""" +"""Persisted job tools for the ASSERT MCP adapter.""" from __future__ import annotations @@ -44,7 +44,7 @@ @dataclass(frozen=True, slots=True) class JobServices: - """Application services and limits shared by evaluation job tools.""" + """Application services and limits shared by persisted job tools.""" workspace: WorkspaceService evaluations: EvaluationService @@ -58,7 +58,7 @@ def register_job_inspect_tools( """Register job discovery and polling for every inspect-capable mode.""" @server.tool( - title="List ASSERT evaluation jobs", + title="List ASSERT jobs", annotations=_READ_ONLY_ANNOTATIONS, structured_output=True, ) @@ -71,7 +71,7 @@ def list_jobs( cursor: str | None = None, page_size: int | None = None, ) -> JobPage: - """List persisted evaluation jobs with bounded pagination.""" + """List persisted jobs with bounded pagination.""" page = invoke_tool( lambda: services.evaluations.list( states=states, @@ -85,7 +85,7 @@ def list_jobs( ) @server.tool( - title="Get an ASSERT evaluation job", + title="Get an ASSERT job", annotations=_READ_ONLY_ANNOTATIONS, structured_output=True, ) @@ -137,8 +137,15 @@ def start_evaluation( sanitize_for_mcp(started, workspace=services.workspace) ) + +def register_job_control_tools( + server: MCPServer, + services: JobServices, +) -> None: + """Register cancellation and retry for enabled persisted-job kinds.""" + @server.tool( - title="Cancel an ASSERT evaluation", + title="Cancel an ASSERT job", annotations=_CANCEL_ANNOTATIONS, structured_output=True, ) @@ -157,7 +164,7 @@ def cancel_job(job_id: str) -> JobDetail: ) @server.tool( - title="Retry an ASSERT evaluation", + title="Retry an ASSERT job", annotations=_START_ANNOTATIONS, structured_output=True, ) diff --git a/assert_ai/mcp/tools/traces.py b/assert_ai/mcp/tools/traces.py new file mode 100644 index 000000000..880f17104 --- /dev/null +++ b/assert_ai/mcp/tools/traces.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Imported OpenTelemetry trace-judging MCP tools.""" + +from __future__ import annotations + +from mcp.server import MCPServer +from mcp.types import ToolAnnotations + +from assert_ai.mcp.errors import adapt_tool_errors, invoke_tool +from assert_ai.mcp.sanitize import sanitize_for_mcp +from assert_ai.mcp.tools.jobs import JobServices +from assert_ai.services.job_models import JobStartResult, TraceJudgingPreflight + +_PREFLIGHT_ANNOTATIONS = ToolAnnotations( + read_only_hint=True, + destructive_hint=False, + idempotent_hint=True, + open_world_hint=False, +) +_START_ANNOTATIONS = ToolAnnotations( + read_only_hint=False, + destructive_hint=True, + idempotent_hint=True, + open_world_hint=True, +) + + +def register_trace_tools( + server: MCPServer, + services: JobServices, +) -> None: + """Register pure trace preflight and persisted trace-job launch.""" + + @server.tool( + title="Preflight imported trace judging", + annotations=_PREFLIGHT_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def preflight_trace_judging( + config_ref: str, + trace_ref: str, + group_by: str = "session.id", + suite_id: str | None = None, + run_id: str | None = None, + ) -> TraceJudgingPreflight: + """Validate OTLP JSON and estimate the judge calls without writing.""" + plan = invoke_tool( + lambda: services.evaluations.preflight_trace_judging( + config_ref, + trace_ref, + group_by=group_by, + suite_id=suite_id, + run_id=run_id, + ), + workspace=services.workspace, + ) + return TraceJudgingPreflight.model_validate( + sanitize_for_mcp(plan, workspace=services.workspace) + ) + + @server.tool( + title="Start imported trace judging", + annotations=_START_ANNOTATIONS, + structured_output=True, + ) + @adapt_tool_errors( + services.workspace, + max_response_bytes=services.max_response_bytes, + ) + def start_trace_judging( + config_ref: str, + trace_ref: str, + request_id: str, + group_by: str = "session.id", + suite_id: str | None = None, + run_id: str | None = None, + ) -> JobStartResult: + """Snapshot OTLP JSON and enqueue judging without blocking.""" + started = invoke_tool( + lambda: services.evaluations.start_trace_judging( + config_ref, + trace_ref, + request_id=request_id, + group_by=group_by, + suite_id=suite_id, + run_id=run_id, + ), + workspace=services.workspace, + ) + return JobStartResult.model_validate( + sanitize_for_mcp(started, workspace=services.workspace) + ) diff --git a/assert_ai/services/_evaluation_worker.py b/assert_ai/services/_evaluation_worker.py index ce007a41c..9f64310ce 100644 --- a/assert_ai/services/_evaluation_worker.py +++ b/assert_ai/services/_evaluation_worker.py @@ -14,6 +14,7 @@ import re import sys import threading +import time from collections.abc import Iterator from dataclasses import asdict from datetime import datetime, timezone @@ -23,7 +24,8 @@ import yaml from assert_ai.core.config_document import PIPELINE_STAGE_ORDER -from assert_ai.core.io import write_json +from assert_ai.core.io import write_bytes_atomic, write_json, write_jsonl +from assert_ai.core.otel import parse_otel_trace_document from assert_ai.core.run_control import ( PipelineFinished, PipelineStarted, @@ -47,6 +49,7 @@ _SHA256_RE = re.compile(r"^sha256:[0-9a-f]{64}$") _MAX_REQUEST_BYTES = 1024 * 1024 _MAX_SNAPSHOT_BYTES = 16 * 1024 * 1024 +_MAX_TRACE_BYTES = 64 * 1024 * 1024 _MIN_LOG_BYTES = 4096 _MAX_LOG_BYTES = 16 * 1024 * 1024 _DEFAULT_LOG_BYTES = 1024 * 1024 @@ -95,6 +98,9 @@ def main(argv: list[str] | None = None) -> int: request = _read_request(request_path) if request.get("job_id") != args.job_id: raise ValueError("Evaluation job request identity mismatch") + kind = request.get("kind", "evaluation") + if kind not in {"evaluation", "trace_judging"}: + raise ValueError("Unsupported evaluation job kind") result_token = _required_string(request, "result_token") config_ref = _required_string(request, "config_ref") expected_snapshot_hash = _required_string( @@ -187,7 +193,11 @@ def main(argv: list[str] | None = None) -> int: ): try: control.raise_if_cancelled( - stage=_first_enabled_stage(document) + stage=( + "trace_import" + if kind == "trace_judging" + else _first_enabled_stage(document) + ) ) except RunCancelled as cancelled: result = _cancelled_before_runner( @@ -197,17 +207,28 @@ def main(argv: list[str] | None = None) -> int: failed_stage=cancelled.stage, ) else: - from assert_ai.runner import run_pipeline_document_result - - result = run_pipeline_document_result( - document=document, - config_path=str(config_path), - force_stages=force_stages, - strict=strict, - path_policy=workspace.path_policy, - control=control, - observer=observer, - ) + if kind == "trace_judging": + result = _run_trace_judging( + document, + request=request, + job_dir=job_dir, + config_path=config_path, + workspace=workspace, + control=control, + observer=observer, + ) + else: + from assert_ai.runner import run_pipeline_document_result + + result = run_pipeline_document_result( + document=document, + config_path=str(config_path), + force_stages=force_stages, + strict=strict, + path_policy=workspace.path_policy, + control=control, + observer=observer, + ) payload = { "schema_version": 1, "job_id": args.job_id, @@ -245,6 +266,330 @@ def main(argv: list[str] | None = None) -> int: return exit_code +def _run_trace_judging( + document: dict[str, Any], + *, + request: dict[str, Any], + job_dir: Path, + config_path: Path, + workspace: WorkspaceService, + control: RunControl, + observer: "_JobRunObserver", +) -> RunResult: + suite_id = _required_string(document, "suite") + run_id = _required_string(document, "run") + group_by = _required_string(request, "group_by") + trace_path = workspace.path_policy.resolve_managed_output( + job_dir / "trace.json", + field_name="immutable OTLP trace input", + expected_root=job_dir, + reject_links=True, + ) + taxonomy_snapshot = workspace.path_policy.resolve_managed_output( + job_dir / "taxonomy.json", + field_name="immutable trace taxonomy", + expected_root=job_dir, + reject_links=True, + ) + trace_bytes = _verified_snapshot( + trace_path, + expected_sha256=_required_string(request, "trace_sha256"), + max_bytes=_MAX_TRACE_BYTES, + label="OTLP trace input", + ) + taxonomy_bytes = _verified_snapshot( + taxonomy_snapshot, + expected_sha256=_required_string(request, "taxonomy_sha256"), + max_bytes=_MAX_SNAPSHOT_BYTES, + label="Trace taxonomy", + ) + suite_root = workspace.path_policy.resolve_managed_output( + workspace.results_root / suite_id, + field_name="trace judge suite root", + expected_root=workspace.results_root, + reject_links=True, + ) + run_root = workspace.path_policy.resolve_managed_output( + suite_root / run_id, + field_name="trace judge run root", + expected_root=suite_root, + reject_links=True, + ) + inference_path = workspace.path_policy.resolve_managed_output( + run_root / "inference_set.jsonl", + field_name="trace judge inference set", + expected_root=run_root, + reject_links=True, + ) + run_taxonomy_path = workspace.path_policy.resolve_managed_output( + run_root / "taxonomy.json", + field_name="trace judge taxonomy", + expected_root=run_root, + reject_links=True, + ) + + observer.pipeline_started( + PipelineStarted( + suite_id=suite_id, + run_id=run_id, + stages=("trace_import", "judge"), + ) + ) + observer.stage_planned( + StagePlanned( + name="trace_import", + scope="run", + action="run", + ) + ) + observer.stage_started(StageStarted(name="trace_import", scope="run")) + started = time.monotonic() + try: + control.raise_if_cancelled(stage="trace_import") + raw_document = json.loads(trace_bytes.decode("utf-8")) + if not isinstance(raw_document, dict): + raise ValueError("OTLP trace input must contain a JSON object") + rows = _normalize_trace_rows( + parse_otel_trace_document(raw_document, group_by=group_by) + ) + if not rows: + raise ValueError("OTLP trace input contains no trace sessions") + control.raise_if_cancelled(stage="trace_import") + run_root.mkdir(parents=True, exist_ok=True) + write_jsonl(inference_path, rows) + write_bytes_atomic(run_taxonomy_path, taxonomy_bytes) + control.raise_if_cancelled(stage="trace_import") + except RunCancelled: + duration = max(0.0, time.monotonic() - started) + observer.stage_finished( + StageFinished( + name="trace_import", + scope="run", + state="cancelled", + duration_seconds=duration, + ) + ) + result = RunResult( + state=RunState.CANCELLED, + exit_code=130, + suite_id=suite_id, + run_id=run_id, + suite_root=suite_root, + run_root=run_root, + failed_stage="trace_import", + error_message="Trace import was cancelled", + ) + observer.pipeline_finished( + PipelineFinished( + state=result.state.value, + exit_code=result.exit_code, + failed_stage=result.failed_stage, + error_message=result.error_message, + ) + ) + return result + except (OSError, UnicodeDecodeError, json.JSONDecodeError, TypeError, ValueError) as exc: + duration = max(0.0, time.monotonic() - started) + message = sanitize_text(str(exc)) or "OTLP trace import failed" + message = redact_path_prefixes( + message, + ( + workspace.root, + workspace.configs_root, + workspace.artifacts_root, + workspace.results_root, + ), + ) + observer.stage_finished( + StageFinished( + name="trace_import", + scope="run", + state="failed", + duration_seconds=duration, + summary={"error": message}, + ) + ) + result = RunResult( + state=RunState.FAILED, + exit_code=1, + suite_id=suite_id, + run_id=run_id, + suite_root=suite_root, + run_root=run_root, + failed_stage="trace_import", + error_code="RUN_FAILED", + error_message=message, + ) + observer.pipeline_finished( + PipelineFinished( + state=result.state.value, + exit_code=result.exit_code, + failed_stage=result.failed_stage, + error_code=result.error_code, + error_message=result.error_message, + ) + ) + return result + + duration = max(0.0, time.monotonic() - started) + observer.stage_progress( + StageProgress( + name="trace_import", + values={ + "completed": len(rows), + "total": len(rows), + "unit": "sessions", + }, + ) + ) + observer.stage_finished( + StageFinished( + name="trace_import", + scope="run", + state="completed", + duration_seconds=duration, + summary={"session_count": len(rows)}, + ) + ) + + from assert_ai.runner import run_pipeline_document_result + + return run_pipeline_document_result( + document=document, + config_path=str(config_path), + force_stages=["judge"], + strict=False, + path_policy=workspace.path_policy, + control=control, + observer=_TraceContinuationObserver(observer), + ) + + +def _verified_snapshot( + path: Path, + *, + expected_sha256: str, + max_bytes: int, + label: str, +) -> bytes: + if not _SHA256_RE.fullmatch(expected_sha256): + raise ValueError(f"{label} digest is invalid") + value = _read_bytes(path, max_bytes=max_bytes, label=label) + actual = "sha256:" + hashlib.sha256(value).hexdigest() + if actual != expected_sha256: + raise ValueError(f"{label} digest mismatch") + return value + + +def _normalize_trace_rows( + rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + normalized: list[dict[str, Any]] = [] + seen: set[str] = set() + for index, row in enumerate(rows, 1): + metadata = row.get("metadata") + metadata = dict(metadata) if isinstance(metadata, dict) else {} + session_id = str(metadata.get("session_id") or f"session-{index}") + raw_refs = metadata.get("trace_refs") + trace_refs = _normalize_trace_refs(raw_refs) + identity_material = json.dumps( + { + "session_id": session_id, + "trace_refs": trace_refs, + }, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + digest = hashlib.sha256(identity_material).hexdigest()[:12] + test_case_id = f"trace_{index:06d}_{digest}" + if test_case_id in seen: + raise ValueError("Imported trace identities are not unique") + seen.add(test_case_id) + metadata.update( + { + "type": "otel_import", + "session_id": session_id, + "runtime_mode": "otel_traced", + "trace_refs": trace_refs, + } + ) + events = ( + row.get("events") + if isinstance(row.get("events"), list) + else [] + ) + normalized.append( + { + "type": "prompt", + "test_case_id": test_case_id, + "behavior": "", + "target": "otel_import", + "tester_model": "", + "metadata": metadata, + "trace_refs": trace_refs, + "events": events, + "stop_reason": "trace_empty" if not events else None, + "raw": ( + row.get("raw") + if isinstance(row.get("raw"), dict) + else {} + ), + } + ) + return normalized + + +def _normalize_trace_refs(value: Any) -> list[dict[str, Any]]: + if not isinstance(value, list): + return [] + refs: list[dict[str, Any]] = [] + for item in value: + if not isinstance(item, dict): + continue + trace_id = item.get("trace_id") + span_ids = item.get("span_ids") + if not isinstance(trace_id, str) or not trace_id: + continue + refs.append( + { + "trace_id": trace_id, + "span_ids": [ + span_id + for span_id in ( + span_ids if isinstance(span_ids, list) else [] + ) + if isinstance(span_id, str) and span_id + ], + } + ) + return refs + + +class _TraceContinuationObserver: + """Forward runner events after the trace-import pipeline start.""" + + def __init__(self, delegate: "_JobRunObserver") -> None: + self._delegate = delegate + + def pipeline_started(self, event: PipelineStarted) -> None: + del event + + def stage_planned(self, event: StagePlanned) -> None: + self._delegate.stage_planned(event) + + def stage_started(self, event: StageStarted) -> None: + self._delegate.stage_started(event) + + def stage_progress(self, event: StageProgress) -> None: + self._delegate.stage_progress(event) + + def stage_finished(self, event: StageFinished) -> None: + self._delegate.stage_finished(event) + + def pipeline_finished(self, event: PipelineFinished) -> None: + self._delegate.pipeline_finished(event) + + def _jobs_root(workspace: WorkspaceService) -> Path: root = workspace.artifacts_root / "mcp" / "jobs" return workspace.path_policy.resolve_managed_output( diff --git a/assert_ai/services/configs.py b/assert_ai/services/configs.py index eacfe2a7a..d0d7b4f5f 100644 --- a/assert_ai/services/configs.py +++ b/assert_ai/services/configs.py @@ -11,7 +11,6 @@ import json import os import re -import time from bisect import bisect from contextlib import contextmanager from copy import deepcopy @@ -36,6 +35,7 @@ from assert_ai.core.runtime_path_policy import RuntimePathError from assert_ai.core.workspace import WorkspaceService from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.locking import exclusive_file_lock from assert_ai.stages import STAGES from assert_ai.stages.test_set import validate_sampling_config @@ -537,7 +537,11 @@ def _config_lock(self, path: Path) -> Iterator[None]: lock_ref, reject_links=True, ) - with _exclusive_file_lock(lock_path, timeout_s=_LOCK_TIMEOUT_S): + with exclusive_file_lock( + lock_path, + timeout_s=_LOCK_TIMEOUT_S, + conflict_message="Timed out waiting for the config write lock", + ): yield @@ -808,45 +812,3 @@ def _invalid_config_error(report: ConfigValidationReport) -> ServiceError: "Config validation failed", details={"validation": report.model_dump(mode="json")}, ) - - -@contextmanager -def _exclusive_file_lock(path: Path, *, timeout_s: float) -> Iterator[None]: - deadline = time.monotonic() + timeout_s - with path.open("a+b") as handle: - handle.seek(0, os.SEEK_END) - if handle.tell() == 0: - handle.write(b"\0") - handle.flush() - os.fsync(handle.fileno()) - handle.seek(0) - while True: - try: - if os.name == "nt": - import msvcrt - - msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) - else: - import fcntl - - fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) - break - except OSError as exc: - if time.monotonic() >= deadline: - raise ServiceError( - ServiceErrorCode.CONFLICT, - "Timed out waiting for the config write lock", - ) from exc - time.sleep(0.05) - try: - yield - finally: - handle.seek(0) - if os.name == "nt": - import msvcrt - - msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) - else: - import fcntl - - fcntl.flock(handle.fileno(), fcntl.LOCK_UN) diff --git a/assert_ai/services/curation.py b/assert_ai/services/curation.py new file mode 100644 index 000000000..cf290cb6b --- /dev/null +++ b/assert_ai/services/curation.py @@ -0,0 +1,1187 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Immutable, workspace-scoped curation of generated suite artifacts.""" + +from __future__ import annotations + +import hashlib +import json +import logging +import re +import threading +import uuid +from contextlib import contextmanager +from copy import deepcopy +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable, Iterator, Mapping, Sequence + +from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator + +from assert_ai.core.artifact_cache import ( + ARTIFACTS_DIR, + LATEST_FILE, + ArtifactFingerprint, + ArtifactPlan, + allocate_artifact_plan, + discard_artifact_plan, + finalize_artifact_plan, + refresh_compatibility_files, + hash_payload, + update_latest_artifacts, +) +from assert_ai.core.io import ( + row_behavior, + write_json, + write_jsonl, + write_text_atomic, +) +from assert_ai.core.jsonl_index import JsonlIndexError, scan_jsonl +from assert_ai.core.runtime_path_policy import RuntimePathError +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.job_store import JobStore +from assert_ai.services.locking import exclusive_file_lock +from assert_ai.services.result_metadata import write_suite_summary + +log = logging.getLogger(__name__) + +_IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_VERSION_RE = re.compile(r"^v[0-9]{4,}$") +_LOCK_TIMEOUT_S = 10.0 +_OPERATION_LEASE_S = 60.0 +_MAX_TAXONOMY_BYTES = 1_048_576 +_MAX_TEST_SET_BYTES = 16_777_216 +_STAGE_FILES: dict[str, tuple[str, ...]] = { + "systematize": ("taxonomy.json", "systematization.json"), + "test_set": ("test_set.jsonl", "stratification.json"), +} + + +class _ServiceModel(BaseModel): + model_config = ConfigDict(frozen=True) + + +class BehaviorDescription(_ServiceModel): + """Behavior block persisted in a taxonomy artifact.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str = Field(min_length=1) + definition: str = Field(min_length=1) + + +class DefinitionOfTerm(_ServiceModel): + """One term definition persisted in a taxonomy artifact.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + term: str = Field(min_length=1) + definition: str = Field(min_length=1) + examples: tuple[str, ...] + + +class BehaviorCategory(_ServiceModel): + """One editable behavior category.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + name: str = Field(min_length=1) + definition: str = Field(min_length=1) + examples: tuple[str, ...] + permissible: bool + + +class TaxonomyDocument(_ServiceModel): + """Canonical taxonomy document accepted by curation.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + behavior: BehaviorDescription + definition_of_terms: tuple[DefinitionOfTerm, ...] + behavior_categories: tuple[BehaviorCategory, ...] + meta: dict[str, Any] | None = None + + @model_validator(mode="after") + def _unique_names(self) -> "TaxonomyDocument": + names = [category.name for category in self.behavior_categories] + if len(names) != len(set(names)): + raise ValueError("behavior category names must be unique") + return self + + +class TestCaseRevision(_ServiceModel): + """Partial update for one stable test-case identity.""" + + test_case_id: str = Field(min_length=1, max_length=255) + updates: dict[str, Any] + + +class CuratedArtifactVersion(_ServiceModel): + """One immutable artifact version created by a curation operation.""" + + artifact_type: str + version: str + etag: str + source_etag: str + source_version: str | None = None + artifact_ref: str + metadata_ref: str + + +class CurationResult(_ServiceModel): + """Result of one atomic suite curation operation.""" + + schema_version: int = 1 + suite_id: str + change_summary: str + artifacts: tuple[CuratedArtifactVersion, ...] + invalidated_stages: tuple[str, ...] = ("inference", "judge") + affected_test_case_ids: tuple[str, ...] = () + warnings: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class _ArtifactSource: + stage_name: str + primary_path: Path + artifact_dir: Path + version: str | None + metadata: dict[str, Any] | None + etag: str + + +class CurationService: + """Create immutable taxonomy and test-set revisions.""" + + def __init__( + self, + workspace: WorkspaceService, + *, + job_store: JobStore | None = None, + ) -> None: + self.workspace = workspace + self.job_store = job_store + + def revise_taxonomy( + self, + suite_id: str, + taxonomy: Mapping[str, Any], + *, + expected_etag: str, + change_summary: str, + ) -> CurationResult: + """Create and atomically activate a taxonomy revision. + + The operation also rebases the active test set into a new immutable + version. This records the taxonomy dependency without rewriting a + completed run or forcing avoidable test-set regeneration. + """ + suite_root = self._suite_root(suite_id) + summary = _change_summary(change_summary) + try: + document = TaxonomyDocument.model_validate(dict(taxonomy)) + except ValidationError as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Taxonomy validation failed", + details={"issues": _validation_issues(exc)}, + ) from exc + taxonomy_payload = document.model_dump(mode="json", exclude_none=True) + _require_json_size( + taxonomy_payload, + max_bytes=_MAX_TAXONOMY_BYTES, + label="Taxonomy revision", + indent=2, + ) + + with self._suite_mutation(suite_id, suite_root) as ensure_lock: + source = self._active_source(suite_root, "systematize") + _require_etag(source.etag, expected_etag) + current = self._load_taxonomy(source.primary_path) + current_names = tuple( + category.name for category in current.behavior_categories + ) + revised_names = tuple( + category.name for category in document.behavior_categories + ) + if revised_names != current_names: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Category additions, removals, renames, and reordering " + "require test-set regeneration and are not supported by " + "revise_taxonomy", + details={ + "current_categories": list(current_names), + "revised_categories": list(revised_names), + }, + ) + if document == current: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "The taxonomy revision does not change the active artifact", + ) + + test_set_source = self._optional_active_source( + suite_root, + "test_set", + ) + plans: list[ArtifactPlan] = [] + refs: dict[str, dict[str, Any]] = {} + sources: dict[str, _ArtifactSource] = {"systematize": source} + created_at = _utc_now() + try: + taxonomy_plan = allocate_artifact_plan( + ctx=self._context(suite_id, suite_root), + stage_name="systematize", + fingerprint=_fingerprint(source), + ) + plans.append(taxonomy_plan) + write_json( + taxonomy_plan.output_paths["taxonomy"], + taxonomy_payload, + ) + self._copy_secondary( + source, + "systematization.json", + taxonomy_plan.output_paths["systematization"], + ) + taxonomy_ref = finalize_artifact_plan( + self._context(suite_id, suite_root), + taxonomy_plan, + provenance=_provenance( + source, + summary, + created_at=created_at, + ), + activate=False, + ) + refs["systematize"] = taxonomy_ref + + if test_set_source is not None: + test_set_plan = allocate_artifact_plan( + ctx=self._context(suite_id, suite_root), + stage_name="test_set", + fingerprint=_rebased_test_set_fingerprint( + test_set_source, + taxonomy_ref, + ), + ) + plans.append(test_set_plan) + self._copy_text( + test_set_source.primary_path, + test_set_plan.output_paths["test_set"], + max_bytes=_MAX_TEST_SET_BYTES, + ) + self._copy_secondary( + test_set_source, + "stratification.json", + test_set_plan.output_paths["stratification"], + ) + test_set_ref = finalize_artifact_plan( + self._context(suite_id, suite_root), + test_set_plan, + provenance={ + **_provenance( + test_set_source, + summary, + created_at=created_at, + ), + "operation": "taxonomy_rebase", + "taxonomy_version": taxonomy_plan.version, + "taxonomy_etag": _file_etag( + taxonomy_plan.output_paths["taxonomy"] + ), + }, + activate=False, + ) + refs["test_set"] = test_set_ref + sources["test_set"] = test_set_source + + ensure_lock() + return self._activate( + suite_id=suite_id, + suite_root=suite_root, + refs=refs, + plans=plans, + sources=sources, + change_summary=summary, + ) + except BaseException: + for plan in plans: + discard_artifact_plan( + self._context(suite_id, suite_root), + plan, + ) + raise + + def revise_test_case( + self, + suite_id: str, + test_case_id: str, + updates: Mapping[str, Any], + *, + expected_etag: str, + change_summary: str, + ) -> CurationResult: + """Create and activate one immutable test-case revision.""" + revision = TestCaseRevision( + test_case_id=test_case_id, + updates=dict(updates), + ) + return self.bulk_revise_test_cases( + suite_id, + (revision,), + expected_etag=expected_etag, + change_summary=change_summary, + ) + + def bulk_revise_test_cases( + self, + suite_id: str, + revisions: Sequence[TestCaseRevision], + *, + expected_etag: str, + change_summary: str, + ) -> CurationResult: + """Create and atomically activate one test-set revision.""" + if not revisions: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "At least one test-case revision is required", + ) + ids = tuple(revision.test_case_id for revision in revisions) + if len(ids) != len(set(ids)): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Each test_case_id may be revised only once per operation", + ) + suite_root = self._suite_root(suite_id) + summary = _change_summary(change_summary) + with self._suite_mutation(suite_id, suite_root) as ensure_lock: + source = self._active_source(suite_root, "test_set") + _require_etag(source.etag, expected_etag) + rows = self._load_test_cases(source.primary_path) + index = { + str(row["test_case_id"]): offset + for offset, row in enumerate(rows) + } + missing = [test_case_id for test_case_id in ids if test_case_id not in index] + if missing: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + "One or more test cases were not found", + details={"test_case_ids": missing}, + ) + + taxonomy_names = self._taxonomy_names(suite_root) + changed = False + for revision in revisions: + updates = dict(revision.updates) + if not updates: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"updates must not be empty for {revision.test_case_id}", + ) + replacement_id = updates.get("test_case_id") + if ( + replacement_id is not None + and replacement_id != revision.test_case_id + ): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "Test-case identities are immutable", + details={"test_case_id": revision.test_case_id}, + ) + offset = index[revision.test_case_id] + revised = {**rows[offset], **updates} + revised["test_case_id"] = revision.test_case_id + _validate_test_case(revised, taxonomy_names=taxonomy_names) + changed = changed or revised != rows[offset] + rows[offset] = revised + + if not changed: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "The test-case revisions do not change the active artifact", + ) + _validate_test_case_set(rows) + _require_jsonl_size( + rows, + max_bytes=_MAX_TEST_SET_BYTES, + label="Test-set revision", + ) + plan = allocate_artifact_plan( + ctx=self._context(suite_id, suite_root), + stage_name="test_set", + fingerprint=_fingerprint(source), + ) + try: + write_jsonl(plan.output_paths["test_set"], rows) + self._copy_secondary( + source, + "stratification.json", + plan.output_paths["stratification"], + ) + ref = finalize_artifact_plan( + self._context(suite_id, suite_root), + plan, + provenance={ + **_provenance( + source, + summary, + created_at=_utc_now(), + ), + "operation": "test_case_revision", + "test_case_ids": list(ids), + }, + activate=False, + ) + ensure_lock() + return self._activate( + suite_id=suite_id, + suite_root=suite_root, + refs={"test_set": ref}, + plans=[plan], + sources={"test_set": source}, + change_summary=summary, + affected_test_case_ids=ids, + ) + except BaseException: + discard_artifact_plan( + self._context(suite_id, suite_root), + plan, + ) + raise + + def _activate( + self, + *, + suite_id: str, + suite_root: Path, + refs: dict[str, dict[str, Any]], + plans: Sequence[ArtifactPlan], + sources: Mapping[str, _ArtifactSource], + change_summary: str, + affected_test_case_ids: tuple[str, ...] = (), + ) -> CurationResult: + ctx = self._context(suite_id, suite_root) + artifacts = tuple( + CuratedArtifactVersion( + artifact_type=plan.stage_name, + version=plan.version, + etag=_file_etag( + plan.output_paths[ + "taxonomy" + if plan.stage_name == "systematize" + else "test_set" + ] + ), + source_etag=sources[plan.stage_name].etag, + source_version=sources[plan.stage_name].version, + artifact_ref=str(refs[plan.stage_name]["path"]), + metadata_ref=str(refs[plan.stage_name]["metadata_path"]), + ) + for plan in plans + ) + update_latest_artifacts(ctx, refs) + warnings: list[str] = [] + try: + latest = self._load_json_object( + suite_root / LATEST_FILE, + required=True, + max_bytes=_MAX_TAXONOMY_BYTES, + ) + except (ServiceError, RuntimePathError) as exc: + warning = ( + "Artifacts were activated, but active artifact metadata " + "could not be reloaded" + ) + log.warning("%s: %s", warning, exc) + warnings.append(warning) + latest = None + latest_artifacts = latest.get("artifacts") if latest is not None else None + ctx["artifact_versions"] = ( + dict(latest_artifacts) + if isinstance(latest_artifacts, dict) + else dict(refs) + ) + for plan in plans: + try: + refresh_compatibility_files( + ctx, + plan.stage_name, + plan.output_paths, + preserve_local_edits=False, + ) + except ( + OSError, + RuntimePathError, + ServiceError, + TypeError, + ValueError, + ) as exc: + warning = ( + f"Activated {plan.stage_name} {plan.version}, but could not " + "refresh one or more legacy compatibility files" + ) + log.warning("%s: %s", warning, exc) + warnings.append(warning) + try: + if "systematize" in refs: + ctx["taxonomy_path"] = str( + next( + plan.output_paths["taxonomy"] + for plan in plans + if plan.stage_name == "systematize" + ) + ) + if "test_set" in refs: + ctx["test_set_path"] = str( + next( + plan.output_paths["test_set"] + for plan in plans + if plan.stage_name == "test_set" + ) + ) + write_suite_summary(ctx, rebuild_indexes=True) + except ( + OSError, + UnicodeDecodeError, + json.JSONDecodeError, + JsonlIndexError, + RuntimePathError, + ServiceError, + KeyError, + TypeError, + ValueError, + ) as exc: + warning = "Artifacts were activated, but suite summary refresh failed" + log.warning("%s: %s", warning, exc) + warnings.append(warning) + + return CurationResult( + suite_id=suite_id, + change_summary=change_summary, + artifacts=artifacts, + affected_test_case_ids=affected_test_case_ids, + warnings=tuple(warnings), + ) + + @contextmanager + def _suite_mutation( + self, + suite_id: str, + suite_root: Path, + ) -> Iterator[Callable[[], None]]: + lock_path = self.workspace.path_policy.resolve_managed_output( + suite_root / ".curation.lock", + field_name="suite curation lock", + expected_root=suite_root, + reject_links=True, + ) + owner = f"curation:{uuid.uuid4().hex}" + resource_keys = (f"suite:{suite_id}",) + stop_renewal = threading.Event() + lease_lost = threading.Event() + renewal_thread: threading.Thread | None = None + operation_acquired = False + + def ensure_lock() -> None: + renewed = ( + self.job_store.renew_operation_locks( + resource_keys, + owner=owner, + lease_seconds=_OPERATION_LEASE_S, + ) + if self.job_store is not None and not lease_lost.is_set() + else False + ) + if self.job_store is not None and not renewed: + lease_lost.set() + raise ServiceError( + ServiceErrorCode.CONFLICT, + "The suite curation lease was lost before activation", + ) + + try: + if self.job_store is not None: + operation_acquired = self.job_store.acquire_operation_locks( + resource_keys, + owner=owner, + lease_seconds=_OPERATION_LEASE_S, + ) + if not operation_acquired: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "The suite is currently being changed by an evaluation " + "or another curation operation", + ) + renewal_thread = threading.Thread( + target=self._renew_operation_lock, + args=(resource_keys, owner, stop_renewal, lease_lost), + name=f"assert-curation-lock-{suite_id}", + daemon=True, + ) + renewal_thread.start() + with exclusive_file_lock( + lock_path, + timeout_s=_LOCK_TIMEOUT_S, + conflict_message="Timed out waiting for the suite curation lock", + ): + yield ensure_lock + finally: + stop_renewal.set() + if renewal_thread is not None: + renewal_thread.join( + timeout=min(10.0, max(1.0, _OPERATION_LEASE_S)) + ) + if renewal_thread.is_alive(): + log.error( + "Suite curation lease-renewal thread did not stop" + ) + if self.job_store is not None and operation_acquired: + try: + self.job_store.release_operation_locks( + owner=owner, + resource_keys=resource_keys, + ) + except Exception: + log.exception( + "Failed to release the suite curation lease; " + "it will expire automatically" + ) + + def _renew_operation_lock( + self, + resource_keys: tuple[str, ...], + owner: str, + stop: threading.Event, + lease_lost: threading.Event, + ) -> None: + assert self.job_store is not None + interval = max(0.05, _OPERATION_LEASE_S / 3) + while not stop.wait(interval): + try: + renewed = self.job_store.renew_operation_locks( + resource_keys, + owner=owner, + lease_seconds=_OPERATION_LEASE_S, + ) + except Exception: + log.exception("Failed to renew the suite curation lease") + lease_lost.set() + return + if not renewed: + log.error("Lost the suite curation lease before activation") + lease_lost.set() + return + + def _suite_root(self, suite_id: str) -> Path: + if not isinstance(suite_id, str) or not _IDENTIFIER_RE.fullmatch(suite_id): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "suite_id must contain only letters, numbers, '.', '_', or '-'", + ) + suite_root = self.workspace.path_policy.resolve_managed_output( + self.workspace.results_root / suite_id, + field_name="curation suite", + expected_root=self.workspace.results_root, + reject_links=True, + ) + if not suite_root.is_dir(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Suite not found: {suite_id}", + ) + return suite_root + + def _context(self, suite_id: str, suite_root: Path) -> dict[str, Any]: + return { + "suite_id": suite_id, + "suite_root": str(suite_root), + "results_root": str(self.workspace.results_root), + "artifacts_root": str(self.workspace.artifacts_root), + "path_policy": self.workspace.path_policy, + } + + def _active_source( + self, + suite_root: Path, + stage_name: str, + ) -> _ArtifactSource: + source = self._optional_active_source(suite_root, stage_name) + if source is None: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"No active {stage_name} artifact was found", + ) + return source + + def _optional_active_source( + self, + suite_root: Path, + stage_name: str, + ) -> _ArtifactSource | None: + primary_name = _STAGE_FILES[stage_name][0] + latest = self._load_json_object( + suite_root / LATEST_FILE, + required=False, + max_bytes=_MAX_TAXONOMY_BYTES, + ) + artifacts = latest.get("artifacts") if latest is not None else None + ref = artifacts.get(stage_name) if isinstance(artifacts, dict) else None + if isinstance(ref, dict): + version = ref.get("version") + if not isinstance(version, str) or not _VERSION_RE.fullmatch(version): + raise ServiceError( + ServiceErrorCode.INTERNAL, + f"Active {stage_name} artifact has an invalid version", + ) + artifact_dir = self.workspace.path_policy.resolve_managed_output( + suite_root / ARTIFACTS_DIR / stage_name / version, + field_name=f"active {stage_name} artifact directory", + expected_root=suite_root, + reject_links=True, + ) + primary_path = self.workspace.path_policy.resolve_managed_output( + artifact_dir / primary_name, + field_name=f"active {stage_name} artifact", + expected_root=artifact_dir, + reject_links=True, + ) + metadata = self._load_json_object( + artifact_dir / "artifact.json", + required=True, + max_bytes=_MAX_TAXONOMY_BYTES, + ) + if not primary_path.is_file(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Active {stage_name} artifact is missing", + ) + if ( + metadata.get("artifact_type") != stage_name + or metadata.get("version") != version + ): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active {stage_name} artifact metadata is inconsistent", + ) + etag = _file_etag(primary_path) + file_hashes = metadata.get("file_hashes") + primary_key = ( + "taxonomy" if stage_name == "systematize" else "test_set" + ) + expected_hash = ( + file_hashes.get(primary_key) + if isinstance(file_hashes, dict) + else None + ) + if ( + isinstance(expected_hash, str) + and re.fullmatch(r"[0-9a-f]{64}", expected_hash) + and etag != f"sha256:{expected_hash}" + ): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active {stage_name} artifact failed its integrity check", + ) + return _ArtifactSource( + stage_name=stage_name, + primary_path=primary_path, + artifact_dir=artifact_dir, + version=version, + metadata=metadata, + etag=etag, + ) + + primary_path = self.workspace.path_policy.resolve_managed_output( + suite_root / primary_name, + field_name=f"legacy {stage_name} artifact", + expected_root=suite_root, + reject_links=True, + ) + if not primary_path.is_file(): + return None + return _ArtifactSource( + stage_name=stage_name, + primary_path=primary_path, + artifact_dir=suite_root, + version=None, + metadata=None, + etag=_file_etag(primary_path), + ) + + def _load_taxonomy(self, path: Path) -> TaxonomyDocument: + raw = self._load_json_object( + path, + required=True, + max_bytes=_MAX_TAXONOMY_BYTES, + ) + try: + return TaxonomyDocument.model_validate(raw) + except ValidationError as exc: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "The active taxonomy is invalid", + details={"issues": _validation_issues(exc)}, + ) from exc + + def _load_test_cases(self, path: Path) -> list[dict[str, Any]]: + if path.stat().st_size > _MAX_TEST_SET_BYTES: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"Test set exceeds the {_MAX_TEST_SET_BYTES}-byte curation limit", + ) + try: + scan = scan_jsonl(path, allow_trailing_partial=False) + except (OSError, JsonlIndexError) as exc: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "The active test set is not valid JSONL", + ) from exc + rows = [dict(record.row) for record in scan.records] + _validate_test_case_set(rows) + return rows + + def _taxonomy_names(self, suite_root: Path) -> set[str]: + source = self._optional_active_source(suite_root, "systematize") + if source is None: + return set() + taxonomy = self._load_taxonomy(source.primary_path) + return { + category.name + for category in taxonomy.behavior_categories + } + + def _load_json_object( + self, + path: Path, + *, + required: bool, + max_bytes: int, + ) -> dict[str, Any] | None: + path = self.workspace.path_policy.resolve_managed_output( + path, + field_name="curation JSON artifact", + expected_root=self.workspace.results_root, + reject_links=True, + ) + if not path.is_file(): + if required: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Required artifact is missing: {path.name}", + ) + return None + if path.stat().st_size > max_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"Artifact exceeds the {max_bytes}-byte curation limit", + ) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Artifact is not valid JSON: {path.name}", + ) from exc + if not isinstance(value, dict): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Artifact must contain a JSON object: {path.name}", + ) + return value + + def _copy_secondary( + self, + source: _ArtifactSource, + name: str, + destination: Path, + ) -> None: + path = self.workspace.path_policy.resolve_managed_output( + source.artifact_dir / name, + field_name=f"{source.stage_name} companion artifact", + expected_root=source.artifact_dir, + reject_links=True, + ) + if not path.is_file(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Required companion artifact is missing: {name}", + ) + self._copy_text(path, destination, max_bytes=_MAX_TEST_SET_BYTES) + + @staticmethod + def _copy_text( + source: Path, + destination: Path, + *, + max_bytes: int, + ) -> None: + if source.stat().st_size > max_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"Artifact exceeds the {max_bytes}-byte curation limit", + ) + try: + text = source.read_bytes().decode("utf-8") + except (OSError, UnicodeDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Artifact is not valid UTF-8: {source.name}", + ) from exc + write_text_atomic(destination, text) + + +def _fingerprint(source: _ArtifactSource) -> ArtifactFingerprint: + hashes = source.metadata.get("hashes") if source.metadata else None + config_hash = ( + hashes.get("config_hash") + if isinstance(hashes, dict) + and isinstance(hashes.get("config_hash"), str) + else source.etag.removeprefix("sha256:") + ) + input_hash = ( + hashes.get("input_hash") + if isinstance(hashes, dict) + and isinstance(hashes.get("input_hash"), str) + else source.etag.removeprefix("sha256:") + ) + behavior_hash = ( + hashes.get("behavior_hash") + if isinstance(hashes, dict) + and isinstance(hashes.get("behavior_hash"), str) + else None + ) + inputs = source.metadata.get("inputs") if source.metadata else None + descriptor = ( + dict(inputs) + if isinstance(inputs, dict) + else { + "curated_legacy_source": True, + "source_etag": source.etag, + } + ) + return ArtifactFingerprint( + stage_name=source.stage_name, + behavior_hash=behavior_hash, + config_hash=config_hash, + input_hash=input_hash, + descriptor=descriptor, + ) + + +def _rebased_test_set_fingerprint( + source: _ArtifactSource, + taxonomy_ref: Mapping[str, Any], +) -> ArtifactFingerprint: + fingerprint = _fingerprint(source) + descriptor = deepcopy(fingerprint.descriptor) + dependencies = descriptor.setdefault("dependencies", {}) + if not isinstance(dependencies, dict): + dependencies = {} + descriptor["dependencies"] = dependencies + dependencies["taxonomy"] = { + "artifact_type": taxonomy_ref.get("artifact_type", "systematize"), + "version": taxonomy_ref.get("version"), + "input_hash": taxonomy_ref.get("input_hash"), + "path": taxonomy_ref.get("path"), + } + input_hash = hash_payload( + { + "stage_name": "test_set", + "behavior_hash": fingerprint.behavior_hash, + "config_hash": fingerprint.config_hash, + "dependencies": dependencies, + "prompts": descriptor.get("prompts", {}), + } + ) + return ArtifactFingerprint( + stage_name="test_set", + behavior_hash=fingerprint.behavior_hash, + config_hash=fingerprint.config_hash, + input_hash=input_hash, + descriptor=descriptor, + ) + + +def _provenance( + source: _ArtifactSource, + change_summary: str, + *, + created_at: str, +) -> dict[str, Any]: + return { + "operation": "curation", + "edited_from": { + "artifact_type": source.stage_name, + "version": source.version, + "etag": source.etag, + }, + "edited_at": created_at, + "change_summary": change_summary, + } + + +def _validate_test_case_set(rows: Sequence[dict[str, Any]]) -> None: + seen: set[str] = set() + for row in rows: + _validate_test_case(row, taxonomy_names=set()) + test_case_id = str(row["test_case_id"]) + if test_case_id in seen: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Duplicate test_case_id: {test_case_id}", + ) + seen.add(test_case_id) + + +def _validate_test_case( + row: Mapping[str, Any], + *, + taxonomy_names: set[str], +) -> None: + test_case_id = row.get("test_case_id") + if ( + not isinstance(test_case_id, str) + or not test_case_id + or len(test_case_id) > 255 + ): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Every test case requires a stable string test_case_id", + ) + row_type = row.get("type") + if row_type not in {"prompt", "scenario"}: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Test case {test_case_id} type must be prompt or scenario", + ) + behavior = row_behavior(dict(row)) + if taxonomy_names and behavior and behavior not in taxonomy_names: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"Test case {test_case_id} references an unknown behavior category", + details={"behavior": behavior}, + ) + try: + json.dumps(row, ensure_ascii=False, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"Test case {test_case_id} is not valid JSON", + ) from exc + + +def _require_etag(current_etag: str, expected_etag: str) -> None: + normalized = expected_etag.strip() + if normalized and not normalized.startswith("sha256:"): + normalized = f"sha256:{normalized}" + if not normalized: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "expected_etag is required", + ) + if normalized != current_etag: + raise ServiceError( + ServiceErrorCode.STALE_ETAG, + "The artifact changed after it was read", + details={"current_etag": current_etag}, + ) + + +def _require_json_size( + value: Any, + *, + max_bytes: int, + label: str, + indent: int | None = None, +) -> None: + try: + size_bytes = len( + json.dumps( + value, + ensure_ascii=False, + allow_nan=False, + indent=indent, + ).encode("utf-8") + ) + except (TypeError, ValueError) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{label} is not valid JSON", + ) from exc + if size_bytes > max_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"{label} exceeds the {max_bytes}-byte curation limit", + ) + + +def _require_jsonl_size( + rows: Sequence[Mapping[str, Any]], + *, + max_bytes: int, + label: str, +) -> None: + size_bytes = 0 + for row in rows: + try: + size_bytes += len( + json.dumps( + row, + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + ) + 1 + except (TypeError, ValueError) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{label} is not valid JSON", + ) from exc + if size_bytes > max_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"{label} exceeds the {max_bytes}-byte curation limit", + ) + + +def _change_summary(value: str) -> str: + summary = value.strip() + if not summary: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "change_summary must not be empty", + ) + if len(summary) > 500: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "change_summary must be <= 500 characters", + ) + return summary + + +def _file_etag(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return f"sha256:{digest.hexdigest()}" + + +def _validation_issues(exc: ValidationError) -> list[dict[str, str]]: + return [ + { + "path": "/" + "/".join(str(part) for part in error["loc"]), + "message": str(error["msg"]), + } + for error in exc.errors(include_url=False, include_input=False) + ] + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() diff --git a/assert_ai/services/evaluations.py b/assert_ai/services/evaluations.py index 7620bfcdc..96a3f3bd7 100644 --- a/assert_ai/services/evaluations.py +++ b/assert_ai/services/evaluations.py @@ -11,6 +11,7 @@ import json import logging import os +import re import secrets import signal import shutil @@ -19,17 +20,21 @@ import threading import time import uuid +from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime, timezone +from fnmatch import fnmatchcase from pathlib import Path from typing import Any, Sequence from urllib.parse import quote import yaml -from assert_ai.core.io import write_json, write_text_atomic +from assert_ai.config import parse_model_config +from assert_ai.core.io import write_bytes_atomic, write_json, write_text_atomic from assert_ai.core.jsonl_index import JsonlIndexError, scan_jsonl from assert_ai.core.config_document import PIPELINE_STAGE_ORDER +from assert_ai.core.otel import parse_otel_trace_document from assert_ai.core.security import ( redact_path_prefixes, sanitize_payload, @@ -37,7 +42,7 @@ ) from assert_ai.core.workspace import WorkspaceService from assert_ai.core.yaml_io import dump_yaml -from assert_ai.services.configs import ConfigService +from assert_ai.services.configs import ConfigRecord, ConfigService from assert_ai.services.errors import ServiceError, ServiceErrorCode from assert_ai.services.job_models import ( JobCatalogEntry, @@ -49,6 +54,7 @@ JobTerminalResult, NewJob, TERMINAL_JOB_STATES, + TraceJudgingPreflight, ) from assert_ai.services.job_store import JobStore from assert_ai.services.run_planning import ( @@ -71,10 +77,39 @@ _REQUEST_ID_MAX_LENGTH = 200 _MIN_LOG_BYTES = 4096 _MAX_LOG_BYTES = 16 * 1024 * 1024 +_DEFAULT_MAX_TRACE_INPUT_BYTES = 64 * 1024 * 1024 +_GROUP_BY_RE = re.compile(r"^[A-Za-z0-9_.-]{1,128}$") +_OUTPUT_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_SUPPORTED_JOB_KINDS = frozenset({"evaluation", "trace_judging"}) log = logging.getLogger(__name__) +@dataclass(frozen=True, slots=True) +class _TraceInputs: + config: ConfigRecord + trace_path: Path + trace_ref: str + trace_bytes: bytes + trace_etag: str + + +@dataclass(frozen=True, slots=True) +class _TracePlan: + inputs: _TraceInputs + suite_id: str | None + run_id: str | None + group_by: str + session_count: int + estimated_judge_calls: int + judge_model: str + taxonomy_path: Path + taxonomy_ref: str + taxonomy_bytes: bytes + taxonomy_etag: str + warnings: tuple[str, ...] + + @dataclass(slots=True) class EvaluationJobManager: """Launch queued jobs and reconcile their terminal worker results.""" @@ -84,6 +119,7 @@ class EvaluationJobManager: max_active_jobs: int = 1 max_log_bytes: int = 1024 * 1024 launch_enabled: bool = True + job_kinds: tuple[str, ...] = ("evaluation", "trace_judging") lease_seconds: float = _LEASE_SECONDS cancellation_grace_seconds: float = ( _DEFAULT_CANCELLATION_GRACE_SECONDS @@ -142,6 +178,12 @@ def __post_init__(self) -> None: raise ValueError( "cancellation_grace_seconds must be positive" ) + self.job_kinds = tuple(dict.fromkeys(self.job_kinds)) + if ( + (self.launch_enabled and not self.job_kinds) + or any(kind not in _SUPPORTED_JOB_KINDS for kind in self.job_kinds) + ): + raise ValueError("job_kinds must contain supported job kinds") def start(self) -> None: """Recover persisted work and then schedule queued jobs.""" @@ -164,6 +206,12 @@ def cancel(self, job_id: str) -> JobRecord: ServiceErrorCode.CAPABILITY_DISABLED, "Evaluation execution is disabled for this service", ) + current = self.store.get(job_id) + if current.kind not in self.job_kinds: + raise ServiceError( + ServiceErrorCode.CAPABILITY_DISABLED, + f"The {current.kind} job kind is not controllable by this server", + ) record = self.store.request_cancel(job_id) if record.state is JobState.CANCELLING: self._write_cancel_marker(record) @@ -189,7 +237,7 @@ def reconcile(self, record: JobRecord) -> JobRecord: """Adopt a worker result or mark a dead worker interrupted.""" if record.state in TERMINAL_JOB_STATES: return record - if not self.launch_enabled: + if not self.launch_enabled or record.kind not in self.job_kinds: return record if record.state is JobState.QUEUED: self.enqueue() @@ -298,6 +346,7 @@ def _schedule(self) -> None: lease_owner=self._owner, lease_seconds=self.lease_seconds, max_active_jobs=self.max_active_jobs, + job_kinds=self.job_kinds, ) except Exception: # noqa: BLE001 - daemon boundary log.exception( @@ -347,7 +396,9 @@ def _schedule(self) -> None: self.enqueue() def _sweep_cancelling_jobs(self) -> None: - for record in self.store.list_nonterminal_records(): + for record in self.store.list_nonterminal_records( + job_kinds=self.job_kinds, + ): if record.state is not JobState.CANCELLING: continue try: @@ -495,7 +546,9 @@ def _recover_startup(self) -> None: next_lease_check: float | None = None has_queued_job = False try: - records = self.store.list_nonterminal_records() + records = self.store.list_nonterminal_records( + job_kinds=self.job_kinds, + ) except Exception: # noqa: BLE001 - daemon boundary log.exception("Could not scan evaluation jobs for recovery") return @@ -1164,6 +1217,7 @@ class EvaluationService: default_page_size: int = 50 max_page_size: int = 200 max_queued_jobs: int = 100 + max_trace_input_bytes: int = _DEFAULT_MAX_TRACE_INPUT_BYTES def start( self, @@ -1172,7 +1226,10 @@ def start( request_id: str, overrides: EvaluationOverrides | None = None, ) -> JobStartResult: - if not self.manager.launch_enabled: + if ( + not self.manager.launch_enabled + or "evaluation" not in self.manager.job_kinds + ): raise ServiceError( ServiceErrorCode.CAPABILITY_DISABLED, "Evaluation execution is disabled for this service", @@ -1268,6 +1325,143 @@ def start( created=created.created, ) + def preflight_trace_judging( + self, + config_ref: str, + trace_ref: str, + *, + group_by: str = "session.id", + suite_id: str | None = None, + run_id: str | None = None, + ) -> TraceJudgingPreflight: + """Validate a trace import and report its exact judge workload.""" + inputs = self._trace_inputs(config_ref, trace_ref) + resolved_suite = _optional_output_id( + suite_id or inputs.config.document.get("suite"), + field_name="suite_id", + ) + resolved_run = _optional_output_id( + run_id or inputs.config.document.get("run"), + field_name="run_id", + ) + plan = self._trace_plan( + inputs, + group_by=group_by, + suite_id=resolved_suite, + run_id=resolved_run, + ) + return TraceJudgingPreflight( + ready=True, + config_ref=inputs.config.config_ref, + config_etag=inputs.config.etag, + trace_ref=inputs.trace_ref, + trace_etag=inputs.trace_etag, + trace_size_bytes=len(inputs.trace_bytes), + group_by=plan.group_by, + session_count=plan.session_count, + estimated_judge_calls=plan.estimated_judge_calls, + suite_id=plan.suite_id, + run_id=plan.run_id, + judge_model=plan.judge_model, + taxonomy_ref=plan.taxonomy_ref, + warnings=plan.warnings, + ) + + def start_trace_judging( + self, + config_ref: str, + trace_ref: str, + *, + request_id: str, + group_by: str = "session.id", + suite_id: str | None = None, + run_id: str | None = None, + ) -> JobStartResult: + """Snapshot and enqueue one persisted OTLP trace-judging job.""" + if ( + not self.manager.launch_enabled + or "trace_judging" not in self.manager.job_kinds + ): + raise ServiceError( + ServiceErrorCode.CAPABILITY_DISABLED, + "Trace judging is disabled for this service", + ) + request_id = _validate_request_id(request_id) + inputs = self._trace_inputs(config_ref, trace_ref) + requested_suite = _optional_output_id( + suite_id, + field_name="suite_id", + ) + requested_run = _optional_output_id( + run_id, + field_name="run_id", + ) + group_by = _validate_group_by(group_by) + request_hash = _trace_request_hash( + config_ref=inputs.config.config_ref, + config_etag=inputs.config.etag, + trace_ref=inputs.trace_ref, + trace_etag=inputs.trace_etag, + group_by=group_by, + suite_id=requested_suite, + run_id=requested_run, + ) + existing = self.store.get_by_idempotency_key(request_id) + if existing is not None: + if existing.request_hash != request_hash: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "request_id is already bound to a different trace-judging request", + details={"job_id": existing.job_id}, + ) + self.manager.enqueue() + return JobStartResult( + job=self.get(existing.job_id), + created=False, + ) + + allocated_suite = _optional_output_id( + requested_suite + or inputs.config.document.get("suite") + or _new_identity("trace-suite"), + field_name="suite_id", + ) + allocated_run = _optional_output_id( + requested_run + or inputs.config.document.get("run") + or _new_identity("trace-run"), + field_name="run_id", + ) + assert allocated_suite is not None + assert allocated_run is not None + plan = self._trace_plan( + inputs, + group_by=group_by, + suite_id=allocated_suite, + run_id=allocated_run, + ) + self._reject_existing_run(allocated_suite, allocated_run) + new_job, job_dir = self._prepare_trace_job( + plan, + request_id=request_id, + request_hash=request_hash, + ) + try: + created = self.store.create_or_get( + new_job, + max_queued_jobs=self.max_queued_jobs, + ) + except BaseException: + _remove_job_dir(job_dir) + raise + if not created.created: + _remove_job_dir(job_dir) + self.manager.enqueue() + return JobStartResult( + job=self.get(created.record.job_id), + created=created.created, + ) + def cancel(self, job_id: str) -> JobDetail: """Request idempotent cooperative cancellation for one job.""" record = self.manager.cancel(_validate_job_id(job_id)) @@ -1287,10 +1481,17 @@ def retry( ) job_id = _validate_job_id(job_id) request_id = _validate_request_id(request_id) - original = self.manager.reconcile(self.store.get(job_id)) + original = self.store.get(job_id) + if original.kind not in self.manager.job_kinds: + raise ServiceError( + ServiceErrorCode.CAPABILITY_DISABLED, + f"The {original.kind} job kind is not controllable by this server", + ) + original = self.manager.reconcile(original) request_hash = _retry_request_hash( retry_of=original.job_id, config_sha256=original.config_sha256, + kind=original.kind, ) existing = self.store.get_by_idempotency_key(request_id) if existing is not None: @@ -1314,6 +1515,17 @@ def retry( ServiceErrorCode.INVALID_ARGUMENT, "Only failed, cancelled, or interrupted jobs can be retried", ) + if original.kind == "trace_judging": + return self._retry_trace_judging( + original, + request_id=request_id, + request_hash=request_hash, + ) + if original.kind != "evaluation": + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The persisted job kind is not supported", + ) document, request = self._retry_snapshot(original) retry_stage = self._retry_stage(original, document) @@ -1472,6 +1684,445 @@ def read_log(self, job_id: str, *, max_bytes: int) -> str: ) return combined + def _trace_inputs( + self, + config_ref: str, + trace_ref: str, + ) -> _TraceInputs: + config = self.configs.get_config(config_ref) + if not config.validation.valid: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Trace judge config validation failed", + details={ + "validation": config.validation.model_dump(mode="json"), + }, + ) + trace_path = self.workspace.resolve_file( + trace_ref, + field_name="OTLP trace input", + ) + _reject_environment_file(trace_path) + if trace_path.suffix.lower() != ".json": + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "OTLP trace input must be a JSON file", + ) + trace_bytes = _read_stable_bytes( + trace_path, + max_bytes=self.max_trace_input_bytes, + label="OTLP trace input", + ) + return _TraceInputs( + config=config, + trace_path=trace_path, + trace_ref=self.workspace.reference(trace_path), + trace_bytes=trace_bytes, + trace_etag=_sha256_etag(trace_bytes), + ) + + def _trace_plan( + self, + inputs: _TraceInputs, + *, + group_by: str, + suite_id: str | None, + run_id: str | None, + ) -> _TracePlan: + group_by = _validate_group_by(group_by) + document = inputs.config.document + pipeline = document.get("pipeline") + judge = pipeline.get("judge") if isinstance(pipeline, dict) else None + if not isinstance(judge, dict) or judge.get("enabled", True) is False: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Trace judging requires an enabled pipeline.judge stage", + ) + raw_model = judge.get("model") or document.get("default_model") + try: + judge_model = parse_model_config( + raw_model, + field_name="pipeline.judge.model", + ).name + except (TypeError, ValueError) as exc: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "pipeline.judge.model or default_model is required", + ) from exc + allowed_patterns = self.planning.policy.allowed_model_patterns + if allowed_patterns and not any( + fnmatchcase(judge_model, pattern) + for pattern in allowed_patterns + ): + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + "The judge model is not allowed by server policy", + details={"model": judge_model}, + ) + + taxonomy_path = self._trace_taxonomy_path( + inputs.config, + suite_id=suite_id, + ) + _reject_environment_file(taxonomy_path) + taxonomy_bytes = _read_stable_bytes( + taxonomy_path, + max_bytes=_JOB_SNAPSHOT_MAX_BYTES, + label="Trace judge taxonomy", + ) + _validate_taxonomy_bytes(taxonomy_bytes) + try: + trace_document = json.loads(inputs.trace_bytes.decode("utf-8")) + if not isinstance(trace_document, dict): + raise ValueError("OTLP payload must be an object") + parsed_rows = parse_otel_trace_document( + trace_document, + group_by=group_by, + ) + except ( + UnicodeDecodeError, + json.JSONDecodeError, + TypeError, + ValueError, + ) as exc: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "OTLP trace input could not be parsed", + ) from exc + if not parsed_rows: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "OTLP trace input contains no trace sessions", + ) + max_sessions = self.planning.policy.max_prompt_sample_size + if len(parsed_rows) > max_sessions: + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + f"Trace input contains {len(parsed_rows)} sessions, exceeding " + f"the server limit of {max_sessions}", + ) + empty_sessions = sum( + 1 + for row in parsed_rows + if not isinstance(row.get("events"), list) or not row["events"] + ) + warnings = ( + ( + f"{empty_sessions} imported session(s) contain no " + "judgeable transcript events", + ) + if empty_sessions + else () + ) + judge_n = judge.get("n", 1) + if ( + isinstance(judge_n, bool) + or not isinstance(judge_n, int) + or judge_n < 1 + ): + judge_n = 1 + return _TracePlan( + inputs=inputs, + suite_id=suite_id, + run_id=run_id, + group_by=group_by, + session_count=len(parsed_rows), + estimated_judge_calls=len(parsed_rows) * judge_n, + judge_model=judge_model, + taxonomy_path=taxonomy_path, + taxonomy_ref=self.workspace.reference(taxonomy_path), + taxonomy_bytes=taxonomy_bytes, + taxonomy_etag=_sha256_etag(taxonomy_bytes), + warnings=warnings, + ) + + def _trace_taxonomy_path( + self, + config: ConfigRecord, + *, + suite_id: str | None, + ) -> Path: + pipeline = config.document.get("pipeline") + judge = pipeline.get("judge") if isinstance(pipeline, dict) else None + raw_path = judge.get("taxonomy_path") if isinstance(judge, dict) else None + config_path = self.workspace.path_policy.resolve_config_path( + config.config_ref, + reject_links=True, + ) + if isinstance(raw_path, str) and raw_path.strip(): + return self.workspace.path_policy.resolve_input( + raw_path, + base_dir=config_path.parent, + field_name="trace judge taxonomy", + must_exist=True, + file_only=True, + ) + if suite_id is None: + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + "Provide suite_id or configure pipeline.judge.taxonomy_path", + ) + suite_root = self.workspace.path_policy.resolve_managed_output( + self.workspace.results_root / suite_id, + field_name="trace judge suite", + expected_root=self.workspace.results_root, + reject_links=True, + ) + latest_path = self.workspace.path_policy.resolve_managed_output( + suite_root / "latest.json", + field_name="trace judge active artifacts", + expected_root=suite_root, + reject_links=True, + ) + latest = _read_json_file( + latest_path, + max_bytes=_JOB_RESULT_MAX_BYTES, + ) + artifacts = latest.get("artifacts") if isinstance(latest, dict) else None + systematize = ( + artifacts.get("systematize") + if isinstance(artifacts, dict) + else None + ) + version = ( + systematize.get("version") + if isinstance(systematize, dict) + else None + ) + if isinstance(version, str) and re.fullmatch(r"v[0-9]{4,}", version): + taxonomy_path = suite_root / "artifacts" / "systematize" / version / "taxonomy.json" + else: + taxonomy_path = suite_root / "taxonomy.json" + resolved = self.workspace.path_policy.resolve_managed_output( + taxonomy_path, + field_name="trace judge taxonomy", + expected_root=suite_root, + reject_links=True, + ) + if not resolved.is_file(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + "Trace judge taxonomy was not found", + ) + return resolved + + def _prepare_trace_job( + self, + plan: _TracePlan, + *, + request_id: str, + request_hash: str, + retry_of: str | None = None, + base_document: dict[str, Any] | None = None, + ) -> tuple[NewJob, Path]: + assert plan.suite_id is not None + assert plan.run_id is not None + return self._prepare_trace_job_values( + config_ref=plan.inputs.config.config_ref, + base_document=base_document or plan.inputs.config.document, + trace_ref=plan.inputs.trace_ref, + trace_bytes=plan.inputs.trace_bytes, + trace_etag=plan.inputs.trace_etag, + taxonomy_ref=plan.taxonomy_ref, + taxonomy_bytes=plan.taxonomy_bytes, + taxonomy_etag=plan.taxonomy_etag, + group_by=plan.group_by, + session_count=plan.session_count, + suite_id=plan.suite_id, + run_id=plan.run_id, + request_id=request_id, + request_hash=request_hash, + retry_of=retry_of, + ) + + def _prepare_trace_job_values( + self, + *, + config_ref: str, + base_document: dict[str, Any], + trace_ref: str, + trace_bytes: bytes, + trace_etag: str, + taxonomy_ref: str, + taxonomy_bytes: bytes, + taxonomy_etag: str, + group_by: str, + session_count: int, + suite_id: str, + run_id: str, + request_id: str, + request_hash: str, + retry_of: str | None, + ) -> tuple[NewJob, Path]: + job_id, job_dir = _allocate_job_dir(self.workspace) + try: + run_root = self.workspace.path_policy.resolve_managed_output( + self.workspace.results_root / suite_id / run_id, + field_name="trace judge run", + expected_root=self.workspace.results_root, + reject_links=True, + ) + run_taxonomy = self.workspace.path_policy.resolve_managed_output( + run_root / "taxonomy.json", + field_name="trace judge taxonomy snapshot", + expected_root=run_root, + reject_links=True, + ) + taxonomy_relative = self.workspace.reference(run_taxonomy) + effective = deepcopy(base_document) + pipeline = effective.get("pipeline") + judge = ( + deepcopy(pipeline.get("judge")) + if isinstance(pipeline, dict) + and isinstance(pipeline.get("judge"), dict) + else None + ) + if judge is None: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Trace judge snapshot has no judge stage", + ) + judge["enabled"] = True + judge["taxonomy_path"] = taxonomy_relative + judge.pop("inference_set_path", None) + judge.pop("save_dir", None) + effective["pipeline"] = {"judge": judge} + effective["suite"] = suite_id + effective["run"] = run_id + effective.pop("artifacts_root", None) + effective.pop("results_dir", None) + yaml_text = dump_yaml(effective) + config_sha256 = _sha256_etag(yaml_text.encode("utf-8")) + + snapshot = job_dir / "config.yaml" + request_path = job_dir / "request.json" + trace_snapshot = job_dir / "trace.json" + taxonomy_snapshot = job_dir / "taxonomy.json" + write_text_atomic(snapshot, yaml_text) + write_bytes_atomic(trace_snapshot, trace_bytes) + write_bytes_atomic(taxonomy_snapshot, taxonomy_bytes) + write_json( + request_path, + { + "schema_version": 1, + "kind": "trace_judging", + "job_id": job_id, + "result_token": secrets.token_hex(32), + "config_ref": config_ref, + "config_sha256": config_sha256, + "strict": False, + "force_stages": ["judge"], + "max_log_bytes": self.manager.max_log_bytes, + "trace_ref": trace_ref, + "trace_sha256": trace_etag, + "trace_size_bytes": len(trace_bytes), + "taxonomy_ref": taxonomy_ref, + "taxonomy_sha256": taxonomy_etag, + "group_by": group_by, + "session_count": session_count, + "retry_of": retry_of, + }, + ) + return ( + NewJob( + job_id=job_id, + idempotency_key=request_id, + request_hash=request_hash, + suite_id=suite_id, + run_id=run_id, + config_ref=config_ref, + config_sha256=config_sha256, + snapshot_path=str(snapshot), + request_path=str(request_path), + resource_keys=(f"run:{suite_id}/{run_id}",), + retry_of=retry_of, + kind="trace_judging", + ), + job_dir, + ) + except BaseException: + _remove_job_dir(job_dir) + raise + + def _retry_trace_judging( + self, + original: JobRecord, + *, + request_id: str, + request_hash: str, + ) -> JobStartResult: + document, request = self._retry_snapshot(original) + if request.get("kind") != "trace_judging": + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable trace job request has an invalid kind", + ) + job_dir = self.manager._job_dir(original.job_id) + trace_bytes = _read_integrity_snapshot( + self.manager._job_file(job_dir, "trace.json"), + expected_etag=request.get("trace_sha256"), + max_bytes=self.max_trace_input_bytes, + label="immutable OTLP trace input", + ) + taxonomy_bytes = _read_integrity_snapshot( + self.manager._job_file(job_dir, "taxonomy.json"), + expected_etag=request.get("taxonomy_sha256"), + max_bytes=_JOB_SNAPSHOT_MAX_BYTES, + label="immutable trace taxonomy", + ) + group_by = _validate_group_by(request.get("group_by")) + session_count = request.get("session_count") + if ( + isinstance(session_count, bool) + or not isinstance(session_count, int) + or session_count < 1 + ): + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable trace job session count is invalid", + ) + trace_ref = request.get("trace_ref") + taxonomy_ref = request.get("taxonomy_ref") + if not isinstance(trace_ref, str) or not isinstance(taxonomy_ref, str): + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable trace job references are invalid", + ) + run_id = _new_identity("trace-run") + self._reject_existing_run(original.suite_id, run_id) + new_job, new_job_dir = self._prepare_trace_job_values( + config_ref=original.config_ref, + base_document=document, + trace_ref=trace_ref, + trace_bytes=trace_bytes, + trace_etag=str(request["trace_sha256"]), + taxonomy_ref=taxonomy_ref, + taxonomy_bytes=taxonomy_bytes, + taxonomy_etag=str(request["taxonomy_sha256"]), + group_by=group_by, + session_count=session_count, + suite_id=original.suite_id, + run_id=run_id, + request_id=request_id, + request_hash=request_hash, + retry_of=original.job_id, + ) + try: + created = self.store.create_or_get( + new_job, + max_queued_jobs=self.max_queued_jobs, + ) + except BaseException: + _remove_job_dir(new_job_dir) + raise + if not created.created: + _remove_job_dir(new_job_dir) + self.manager.enqueue() + return JobStartResult( + job=self.get(created.record.job_id), + created=created.created, + ) + def _retry_snapshot( self, record: JobRecord, @@ -1600,27 +2251,7 @@ def _prepare_job( yaml_text: str, retry_of: str | None = None, ) -> tuple[NewJob, Path]: - jobs_root = _jobs_root(self.workspace) - jobs_root.mkdir(parents=True, exist_ok=True) - jobs_root = _jobs_root(self.workspace) - for _ in range(_JOB_ID_RETRIES): - job_id = uuid.uuid4().hex - job_dir = self.workspace.path_policy.resolve_managed_output( - jobs_root / job_id, - field_name="evaluation job directory", - expected_root=jobs_root, - reject_links=True, - ) - try: - job_dir.mkdir() - except FileExistsError: - continue - break - else: - raise ServiceError( - ServiceErrorCode.CONFLICT, - "Could not allocate a unique evaluation job id", - ) + job_id, job_dir = _allocate_job_dir(self.workspace) try: snapshot = job_dir / "config.yaml" request_path = job_dir / "request.json" @@ -1792,7 +2423,7 @@ def _catalog_entry(record: JobRecord) -> JobCatalogEntry: job_id=record.job_id, state=record.state, revision=record.revision, - kind="evaluation", + kind=record.kind, retry_of=record.retry_of, config_ref=record.config_ref, suite_id=record.suite_id, @@ -1859,10 +2490,11 @@ def _retry_request_hash( *, retry_of: str, config_sha256: str, + kind: str, ) -> str: payload = json.dumps( { - "operation": "retry_evaluation", + "operation": f"retry_{kind}", "retry_of": retry_of, "config_sha256": config_sha256, }, @@ -2038,6 +2670,168 @@ def _new_identity(prefix: str) -> str: return f"{prefix}-{timestamp}-{secrets.token_hex(8)}" +def _optional_output_id(value: Any, *, field_name: str) -> str | None: + if value is None: + return None + if not isinstance(value, str) or not _OUTPUT_ID_RE.fullmatch(value): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{field_name} must contain only letters, numbers, '.', '_', or '-'", + ) + return value + + +def _validate_group_by(value: Any) -> str: + if not isinstance(value, str) or not _GROUP_BY_RE.fullmatch(value): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + "group_by must be a 1-128 character OpenTelemetry attribute name", + ) + return value + + +def _trace_request_hash( + *, + config_ref: str, + config_etag: str, + trace_ref: str, + trace_etag: str, + group_by: str, + suite_id: str | None, + run_id: str | None, +) -> str: + payload = json.dumps( + { + "operation": "trace_judging", + "config_ref": config_ref, + "config_etag": config_etag, + "trace_ref": trace_ref, + "trace_etag": trace_etag, + "group_by": group_by, + "suite_id": suite_id, + "run_id": run_id, + }, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return _sha256_etag(payload) + + +def _sha256_etag(value: bytes) -> str: + return "sha256:" + hashlib.sha256(value).hexdigest() + + +def _read_stable_bytes( + path: Path, + *, + max_bytes: int, + label: str, +) -> bytes: + try: + before = path.stat() + with path.open("rb") as handle: + value = handle.read(max_bytes + 1) + after = path.stat() + except OSError as exc: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"{label} is unavailable", + ) from exc + if len(value) > max_bytes: + raise ServiceError( + ServiceErrorCode.ARTIFACT_TOO_LARGE, + f"{label} exceeds the {max_bytes}-byte limit", + ) + if ( + before.st_size != after.st_size + or before.st_mtime_ns != after.st_mtime_ns + ): + raise ServiceError( + ServiceErrorCode.CONFLICT, + f"{label} changed while it was being read", + ) + return value + + +def _read_integrity_snapshot( + path: Path, + *, + expected_etag: Any, + max_bytes: int, + label: str, +) -> bytes: + if ( + not isinstance(expected_etag, str) + or not re.fullmatch(r"sha256:[0-9a-f]{64}", expected_etag) + ): + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + f"The {label} digest is invalid", + ) + value = _read_stable_bytes(path, max_bytes=max_bytes, label=label) + if _sha256_etag(value) != expected_etag: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + f"The {label} failed its integrity check", + ) + return value + + +def _validate_taxonomy_bytes(value: bytes) -> None: + try: + taxonomy = json.loads(value.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Trace judge taxonomy is not valid JSON", + ) from exc + if ( + not isinstance(taxonomy, dict) + or not isinstance(taxonomy.get("behavior"), dict) + or not isinstance(taxonomy.get("behavior_categories"), list) + ): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "Trace judge taxonomy has an invalid structure", + ) + + +def _reject_environment_file(path: Path) -> None: + for part in path.parts: + name = part.lower() + if name == ".env" or name.startswith(".env.") or name.endswith(".env"): + raise ServiceError( + ServiceErrorCode.WORKSPACE_VIOLATION, + "Environment files cannot be used as trace-judging inputs", + ) + + +def _allocate_job_dir( + workspace: WorkspaceService, +) -> tuple[str, Path]: + jobs_root = _jobs_root(workspace) + jobs_root.mkdir(parents=True, exist_ok=True) + jobs_root = _jobs_root(workspace) + for _ in range(_JOB_ID_RETRIES): + job_id = uuid.uuid4().hex + job_dir = workspace.path_policy.resolve_managed_output( + jobs_root / job_id, + field_name="evaluation job directory", + expected_root=jobs_root, + reject_links=True, + ) + try: + job_dir.mkdir() + except FileExistsError: + continue + return job_id, job_dir + raise ServiceError( + ServiceErrorCode.CONFLICT, + "Could not allocate a unique evaluation job id", + ) + + def _encode_cursor(created_at: str, job_id: str) -> str: payload = json.dumps( { diff --git a/assert_ai/services/job_models.py b/assert_ai/services/job_models.py index a9e2fb63d..7aef294f6 100644 --- a/assert_ai/services/job_models.py +++ b/assert_ai/services/job_models.py @@ -119,7 +119,7 @@ class JobCatalogEntry(_ServiceModel): job_id: str state: JobState revision: int = Field(ge=0) - kind: Literal["evaluation"] = "evaluation" + kind: Literal["evaluation", "trace_judging"] = "evaluation" retry_of: str | None = None config_ref: str suite_id: str @@ -158,3 +158,23 @@ class JobStartResult(_ServiceModel): job: JobDetail created: bool + + +class TraceJudgingPreflight(_ServiceModel): + """Pure validation and spend estimate for an imported-trace judge job.""" + + schema_version: Literal[1] = 1 + ready: bool + config_ref: str + config_etag: str + trace_ref: str + trace_etag: str + trace_size_bytes: int + group_by: str + session_count: int + estimated_judge_calls: int + suite_id: str | None = None + run_id: str | None = None + judge_model: str + taxonomy_ref: str + warnings: tuple[str, ...] = () diff --git a/assert_ai/services/job_store.py b/assert_ai/services/job_store.py index 91154a06c..501d89336 100644 --- a/assert_ai/services/job_store.py +++ b/assert_ai/services/job_store.py @@ -25,7 +25,7 @@ ) _BUSY_TIMEOUT_MS = 5_000 -_JOB_STORE_SCHEMA_VERSION = 2 +_JOB_STORE_SCHEMA_VERSION = 3 _ACTIVE_STATES = ( JobState.STARTING.value, JobState.RUNNING.value, @@ -85,6 +85,12 @@ lease_expires_at TEXT NOT NULL, FOREIGN KEY(job_id) REFERENCES jobs(job_id) ON DELETE CASCADE ); +CREATE TABLE IF NOT EXISTS operation_locks( + resource_key TEXT PRIMARY KEY, + owner TEXT NOT NULL, + acquired_at TEXT NOT NULL, + lease_expires_at TEXT NOT NULL +); """ @@ -276,6 +282,7 @@ def claim_next( lease_owner: str, lease_seconds: float, max_active_jobs: int, + job_kinds: Sequence[str] = (), ) -> JobRecord | None: if lease_seconds <= 0: raise ValueError("lease_seconds must be positive") @@ -287,6 +294,10 @@ def claim_next( now = _now() expires_at = _after(lease_seconds) with self._transaction() as connection: + connection.execute( + "DELETE FROM operation_locks WHERE lease_expires_at <= ?", + (now,), + ) active = int( connection.execute( "SELECT COUNT(*) FROM jobs WHERE state IN (?, ?, ?)", @@ -295,14 +306,26 @@ def claim_next( ) if active >= max_active_jobs: return None - candidates = connection.execute( - """ - SELECT * FROM jobs - WHERE state = ? - ORDER BY created_at, job_id - """, - (JobState.QUEUED.value,), - ).fetchall() + kinds = tuple(dict.fromkeys(job_kinds)) + if kinds: + kind_placeholders = ", ".join("?" for _ in kinds) + candidates = connection.execute( + f""" + SELECT * FROM jobs + WHERE state = ? AND kind IN ({kind_placeholders}) + ORDER BY created_at, job_id + """, + (JobState.QUEUED.value, *kinds), + ).fetchall() + else: + candidates = connection.execute( + """ + SELECT * FROM jobs + WHERE state = ? + ORDER BY created_at, job_id + """, + (JobState.QUEUED.value,), + ).fetchall() for row in candidates: record = _record(row) if not self._resources_available( @@ -357,6 +380,131 @@ def claim_next( return _record(claimed) return None + def acquire_operation_locks( + self, + resource_keys: Sequence[str], + *, + owner: str, + lease_seconds: float, + ) -> bool: + """Reserve resources against job claims for one short operation.""" + keys = tuple(dict.fromkeys(resource_keys)) + if not keys: + raise ValueError("at least one resource key is required") + if not owner: + raise ValueError("owner is required") + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + self.initialize() + now = _now() + expires_at = _after(lease_seconds) + placeholders = ", ".join("?" for _ in keys) + with self._transaction() as connection: + connection.execute( + "DELETE FROM operation_locks WHERE lease_expires_at <= ?", + (now,), + ) + active_job = any( + self._operation_conflicts_with_active_job( + connection, + resource_key, + ) + for resource_key in keys + ) + active_operation = connection.execute( + f""" + SELECT 1 FROM operation_locks + WHERE resource_key IN ({placeholders}) + LIMIT 1 + """, + keys, + ).fetchone() + if active_job or active_operation is not None: + return False + connection.executemany( + """ + INSERT INTO operation_locks( + resource_key, owner, acquired_at, lease_expires_at + ) VALUES (?, ?, ?, ?) + """, + ( + (resource_key, owner, now, expires_at) + for resource_key in keys + ), + ) + return True + + def release_operation_locks( + self, + *, + owner: str, + resource_keys: Sequence[str] = (), + ) -> None: + """Release operation locks owned by one caller.""" + if not owner: + raise ValueError("owner is required") + self.initialize() + keys = tuple(dict.fromkeys(resource_keys)) + with self._transaction() as connection: + if keys: + placeholders = ", ".join("?" for _ in keys) + connection.execute( + f""" + DELETE FROM operation_locks + WHERE owner = ? AND resource_key IN ({placeholders}) + """, + (owner, *keys), + ) + else: + connection.execute( + "DELETE FROM operation_locks WHERE owner = ?", + (owner,), + ) + + def renew_operation_locks( + self, + resource_keys: Sequence[str], + *, + owner: str, + lease_seconds: float, + ) -> bool: + """Extend unexpired operation locks when every key is still owned.""" + keys = tuple(dict.fromkeys(resource_keys)) + if not keys: + raise ValueError("at least one resource key is required") + if not owner: + raise ValueError("owner is required") + if lease_seconds <= 0: + raise ValueError("lease_seconds must be positive") + self.initialize() + now = _now() + expires_at = _after(lease_seconds) + placeholders = ", ".join("?" for _ in keys) + with self._transaction() as connection: + owned = connection.execute( + f""" + SELECT COUNT(*) AS count + FROM operation_locks + WHERE owner = ? + AND lease_expires_at > ? + AND resource_key IN ({placeholders}) + """, + (owner, now, *keys), + ).fetchone() + if owned is None or int(owned["count"]) != len(keys): + return False + changed = connection.execute( + f""" + UPDATE operation_locks + SET lease_expires_at = ? + WHERE owner = ? + AND lease_expires_at > ? + AND resource_key IN ({placeholders}) + """, + (expires_at, owner, now, *keys), + ).rowcount + return changed == len(keys) + def mark_running( self, job_id: str, @@ -602,20 +750,33 @@ def adopt_lease( ) return self._get_in_transaction(connection, job_id) - def list_nonterminal_records(self) -> tuple[JobRecord, ...]: + def list_nonterminal_records( + self, + *, + job_kinds: Sequence[str] = (), + ) -> tuple[JobRecord, ...]: """Return every queued or active job for deterministic recovery.""" if not self.exists: return () self.initialize() placeholders = ", ".join("?" for _ in TERMINAL_JOB_STATES) + values: list[str] = [ + state.value for state in TERMINAL_JOB_STATES + ] + kinds = tuple(dict.fromkeys(job_kinds)) + kind_clause = "" + if kinds: + kind_placeholders = ", ".join("?" for _ in kinds) + kind_clause = f" AND kind IN ({kind_placeholders})" + values.extend(kinds) with self._connection() as connection: rows = connection.execute( f""" SELECT * FROM jobs - WHERE state NOT IN ({placeholders}) + WHERE state NOT IN ({placeholders}){kind_clause} ORDER BY created_at, job_id """, - tuple(state.value for state in TERMINAL_JOB_STATES), + tuple(values), ).fetchall() return tuple(_record(row) for row in rows) @@ -772,6 +933,7 @@ def initialize(self) -> None: if version not in { 0, 1, + 2, _JOB_STORE_SCHEMA_VERSION, }: raise ServiceError( @@ -861,7 +1023,62 @@ def _resources_available( """, resource_keys, ).fetchone() - return row is None + if row is not None: + return False + operation_keys = tuple( + dict.fromkeys( + ( + *resource_keys, + *( + f"suite:{suite_id}" + for suite_id in ( + _suite_id_from_resource_key(key) + for key in resource_keys + ) + if suite_id is not None + ), + ) + ) + ) + operation_placeholders = ", ".join("?" for _ in operation_keys) + operation = connection.execute( + f""" + SELECT 1 FROM operation_locks + WHERE resource_key IN ({operation_placeholders}) + LIMIT 1 + """, + operation_keys, + ).fetchone() + return operation is None + + @staticmethod + def _operation_conflicts_with_active_job( + connection: sqlite3.Connection, + resource_key: str, + ) -> bool: + if resource_key.startswith("suite:"): + suite_id = resource_key.removeprefix("suite:") + row = connection.execute( + """ + SELECT 1 FROM resource_locks + WHERE resource_key = ? OR resource_key LIKE ? + LIMIT 1 + """, + ( + resource_key, + f"run:{suite_id}/%", + ), + ).fetchone() + else: + row = connection.execute( + """ + SELECT 1 FROM resource_locks + WHERE resource_key = ? + LIMIT 1 + """, + (resource_key,), + ).fetchone() + return row is not None @staticmethod def _get_in_transaction( @@ -953,6 +1170,16 @@ def _append_event( return sequence +def _suite_id_from_resource_key(resource_key: str) -> str | None: + if resource_key.startswith("suite:"): + return resource_key.removeprefix("suite:") + if resource_key.startswith("run:"): + value = resource_key.removeprefix("run:") + suite_id, separator, _ = value.partition("/") + return suite_id if separator and suite_id else None + return None + + def _record(row: sqlite3.Row) -> JobRecord: return JobRecord( job_id=str(row["job_id"]), diff --git a/assert_ai/services/locking.py b/assert_ai/services/locking.py new file mode 100644 index 000000000..6cdb9e84f --- /dev/null +++ b/assert_ai/services/locking.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Cross-process locks for short managed-file updates.""" + +from __future__ import annotations + +import os +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator + +from assert_ai.services.errors import ServiceError, ServiceErrorCode + + +@contextmanager +def exclusive_file_lock( + path: Path, + *, + timeout_s: float, + conflict_message: str, +) -> Iterator[None]: + """Hold an advisory lock on one workspace-managed lock file.""" + deadline = time.monotonic() + timeout_s + with path.open("a+b") as handle: + handle.seek(0, os.SEEK_END) + if handle.tell() == 0: + handle.write(b"\0") + handle.flush() + os.fsync(handle.fileno()) + handle.seek(0) + while True: + try: + if os.name == "nt": + import msvcrt + + msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) + else: + import fcntl + + fcntl.flock( + handle.fileno(), + fcntl.LOCK_EX | fcntl.LOCK_NB, + ) + break + except OSError as exc: + if time.monotonic() >= deadline: + raise ServiceError( + ServiceErrorCode.CONFLICT, + conflict_message, + ) from exc + time.sleep(0.05) + try: + yield + finally: + handle.seek(0) + if os.name == "nt": + import msvcrt + + msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) + else: + import fcntl + + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) diff --git a/assert_ai/stages/judge.py b/assert_ai/stages/judge.py index 9e88a27cf..7277e5715 100644 --- a/assert_ai/stages/judge.py +++ b/assert_ai/stages/judge.py @@ -49,6 +49,7 @@ "target_input_refused", "target_error", "tester_error", + "trace_empty", }) @@ -219,6 +220,9 @@ async def score_row(row: dict[str, Any]) -> dict[str, Any]: dimensions = row_factors(row) if dimensions: skipped["dimensions"] = dimensions + trace_refs = row.get("trace_refs") + if isinstance(trace_refs, list): + skipped["trace_refs"] = trace_refs return skipped transcript_metadata = TranscriptMetadata( kind=str(row.get("type") or ""), @@ -282,6 +286,9 @@ async def score_row(row: dict[str, Any]) -> dict[str, Any]: dimensions = row_factors(row) if dimensions: score_row["dimensions"] = dimensions + trace_refs = row.get("trace_refs") + if isinstance(trace_refs, list): + score_row["trace_refs"] = trace_refs if judge_result.get("multi_judge") is not None: score_row["multi_judge"] = judge_result["multi_judge"] return score_row diff --git a/tests/test_curation_service.py b/tests/test_curation_service.py new file mode 100644 index 000000000..6acfde421 --- /dev/null +++ b/tests/test_curation_service.py @@ -0,0 +1,555 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import hashlib +import json +import threading +from copy import deepcopy +from pathlib import Path +from unittest.mock import patch + +import pytest + +import assert_ai.services.curation as curation_module +from assert_ai.core.artifact_cache import ( + activate_artifact_plan, + finalize_artifact_plan, + prepare_artifact_plan, +) +from assert_ai.core.io import write_json, write_jsonl +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.curation import ( + CurationService, + TestCaseRevision as CaseRevision, +) +from assert_ai.services.errors import ServiceError, ServiceErrorCode +from assert_ai.services.job_models import NewJob +from assert_ai.services.job_store import JobStore + + +def _taxonomy() -> dict: + return { + "behavior": { + "name": "safe_travel", + "definition": "The planner follows travel safety constraints.", + }, + "definition_of_terms": [], + "behavior_categories": [ + { + "name": "safe_booking", + "definition": "Books only policy-compliant travel.", + "examples": ["Reject an unsafe itinerary."], + "permissible": True, + }, + { + "name": "unsafe_booking", + "definition": "Books travel that violates policy.", + "examples": ["Ignore a safety restriction."], + "permissible": False, + }, + ], + "meta": { + "source": "systematization", + "slug": "safe_travel", + }, + } + + +def _rows() -> list[dict]: + return [ + { + "type": "prompt", + "test_case_id": "test_case_000001", + "prompt": "Book a safe flight.", + "dimensions": {"behavior": "safe_booking"}, + }, + { + "type": "scenario", + "test_case_id": "test_case_000002", + "prompt": "Ignore the restriction.", + "dimensions": {"behavior": "unsafe_booking"}, + "tools": [], + }, + ] + + +def _seed_suite( + tmp_path: Path, +) -> tuple[WorkspaceService, JobStore, Path]: + workspace = WorkspaceService.create(tmp_path) + workspace.configs_root.mkdir(parents=True) + workspace.results_root.mkdir(parents=True) + config_path = workspace.configs_root / "demo.yaml" + config_path.write_text("suite: suite-a\n", encoding="utf-8") + suite_root = workspace.results_root / "suite-a" + suite_root.mkdir() + ctx = { + "suite_id": "suite-a", + "suite_root": str(suite_root), + "results_root": str(workspace.results_root), + "artifacts_root": str(workspace.artifacts_root), + "config_path": str(config_path), + "path_policy": workspace.path_policy, + "behavior_name": "safe_travel", + "behavior": "The planner follows travel safety constraints.", + "context": "A travel planning agent.", + "dimensions": {}, + "artifact_versions": {}, + } + + taxonomy_plan = prepare_artifact_plan( + ctx=ctx, + stage_name="systematize", + raw_cfg={"model": {"name": "test/model"}}, + forced=True, + ) + activate_artifact_plan(ctx, taxonomy_plan) + write_json(taxonomy_plan.output_paths["taxonomy"], _taxonomy()) + write_json( + taxonomy_plan.output_paths["systematization"], + { + "behavior": "safe_travel", + "systematization": "Original systematization", + "summary_items": [], + }, + ) + finalize_artifact_plan(ctx, taxonomy_plan) + + test_set_plan = prepare_artifact_plan( + ctx=ctx, + stage_name="test_set", + raw_cfg={ + "model": {"name": "test/model"}, + "prompt": {"sample_size": 1}, + "scenario": {"sample_size": 1}, + }, + forced=True, + ) + activate_artifact_plan(ctx, test_set_plan) + write_jsonl(test_set_plan.output_paths["test_set"], _rows()) + write_json(test_set_plan.output_paths["stratification"], {"counts": {}}) + finalize_artifact_plan(ctx, test_set_plan) + + store = JobStore(workspace.artifacts_root / "mcp" / "jobs.sqlite3") + return workspace, store, suite_root + + +def _etag(path: Path) -> str: + return f"sha256:{hashlib.sha256(path.read_bytes()).hexdigest()}" + + +def test_revise_taxonomy_creates_immutable_versions_and_rebases_test_set( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + source_taxonomy = ( + suite_root / "artifacts" / "systematize" / "v0001" / "taxonomy.json" + ) + original_taxonomy = source_taxonomy.read_bytes() + original_test_set = ( + suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + ).read_bytes() + revised = deepcopy(_taxonomy()) + revised["behavior_categories"][0]["definition"] = "Revised definition." + + result = CurationService(workspace, job_store=store).revise_taxonomy( + "suite-a", + revised, + expected_etag=_etag(source_taxonomy), + change_summary="Clarify the safe-booking rubric.", + ) + + assert [(item.artifact_type, item.version) for item in result.artifacts] == [ + ("systematize", "v0002"), + ("test_set", "v0002"), + ] + latest = json.loads((suite_root / "latest.json").read_text(encoding="utf-8")) + assert latest["artifacts"]["systematize"]["version"] == "v0002" + assert latest["artifacts"]["test_set"]["version"] == "v0002" + assert source_taxonomy.read_bytes() == original_taxonomy + assert ( + suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + ).read_bytes() == original_test_set + assert ( + suite_root / "artifacts" / "test_set" / "v0002" / "test_set.jsonl" + ).read_bytes() == original_test_set + assert json.loads( + (suite_root / "taxonomy.json").read_text(encoding="utf-8") + )["behavior_categories"][0]["definition"] == "Revised definition." + assert json.loads( + (suite_root / "taxonomy.json").read_text(encoding="utf-8") + )["meta"]["source"] == "systematization" + + metadata = json.loads( + ( + suite_root + / "artifacts" + / "systematize" + / "v0002" + / "artifact.json" + ).read_text(encoding="utf-8") + ) + assert metadata["provenance"]["edited_from"]["version"] == "v0001" + assert metadata["provenance"]["change_summary"] == ( + "Clarify the safe-booking rubric." + ) + old_test_metadata = json.loads( + ( + suite_root / "artifacts" / "test_set" / "v0001" / "artifact.json" + ).read_text(encoding="utf-8") + ) + new_test_metadata = json.loads( + ( + suite_root / "artifacts" / "test_set" / "v0002" / "artifact.json" + ).read_text(encoding="utf-8") + ) + assert ( + new_test_metadata["hashes"]["config_hash"] + == old_test_metadata["hashes"]["config_hash"] + ) + assert ( + new_test_metadata["hashes"]["input_hash"] + != old_test_metadata["hashes"]["input_hash"] + ) + summary = json.loads( + (suite_root / "suite_summary.json").read_text(encoding="utf-8") + ) + assert summary["artifact_versions"]["systematize"]["version"] == "v0002" + assert summary["artifact_versions"]["test_set"]["version"] == "v0002" + + +def test_revise_taxonomy_rejects_stale_etag_and_category_shape_changes( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + service = CurationService(workspace, job_store=store) + + with pytest.raises(ServiceError) as stale: + service.revise_taxonomy( + "suite-a", + _taxonomy(), + expected_etag="sha256:" + ("0" * 64), + change_summary="Attempt a stale edit.", + ) + assert stale.value.code == ServiceErrorCode.STALE_ETAG + + reordered = deepcopy(_taxonomy()) + reordered["behavior_categories"].reverse() + source = suite_root / "artifacts" / "systematize" / "v0001" / "taxonomy.json" + with pytest.raises(ServiceError) as invalid: + service.revise_taxonomy( + "suite-a", + reordered, + expected_etag=_etag(source), + change_summary="Reorder categories.", + ) + assert invalid.value.code == ServiceErrorCode.INVALID_ARGUMENT + assert not ( + suite_root / "artifacts" / "systematize" / "v0002" + ).exists() + + with pytest.raises(ServiceError) as unchanged: + service.revise_taxonomy( + "suite-a", + _taxonomy(), + expected_etag=_etag(source), + change_summary="Attempt a no-op revision.", + ) + assert unchanged.value.code == ServiceErrorCode.INVALID_ARGUMENT + + +def test_bulk_revise_test_cases_preserves_ids_order_and_old_version( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + source = suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + original = source.read_bytes() + + result = CurationService( + workspace, + job_store=store, + ).bulk_revise_test_cases( + "suite-a", + ( + CaseRevision( + test_case_id="test_case_000002", + updates={"prompt": "Revised unsafe request."}, + ), + CaseRevision( + test_case_id="test_case_000001", + updates={"prompt": "Revised safe request."}, + ), + ), + expected_etag=_etag(source).removeprefix("sha256:"), + change_summary="Make both prompts more explicit.", + ) + + assert result.affected_test_case_ids == ( + "test_case_000002", + "test_case_000001", + ) + assert source.read_bytes() == original + revised_rows = [ + json.loads(line) + for line in ( + suite_root / "artifacts" / "test_set" / "v0002" / "test_set.jsonl" + ).read_text(encoding="utf-8").splitlines() + ] + assert [row["test_case_id"] for row in revised_rows] == [ + "test_case_000001", + "test_case_000002", + ] + assert revised_rows[0]["prompt"] == "Revised safe request." + assert revised_rows[1]["prompt"] == "Revised unsafe request." + + +def test_post_activation_summary_failure_keeps_new_version( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + source = ( + suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + ) + + with patch( + "assert_ai.services.curation.write_suite_summary", + side_effect=OSError("fixture summary failure"), + ): + result = CurationService( + workspace, + job_store=store, + ).revise_test_case( + "suite-a", + "test_case_000001", + {"prompt": "Revised prompt."}, + expected_etag=_etag(source), + change_summary="Exercise post-activation failure handling.", + ) + + latest = json.loads((suite_root / "latest.json").read_text(encoding="utf-8")) + assert latest["artifacts"]["test_set"]["version"] == "v0002" + assert ( + suite_root / "artifacts" / "test_set" / "v0002" / "test_set.jsonl" + ).is_file() + assert result.warnings == ( + "Artifacts were activated, but suite summary refresh failed", + ) + + +def test_post_activation_lock_release_failure_keeps_new_version( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + source = ( + suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + ) + + with patch.object( + store, + "release_operation_locks", + side_effect=OSError("fixture lock release failure"), + ): + result = CurationService( + workspace, + job_store=store, + ).revise_test_case( + "suite-a", + "test_case_000001", + {"prompt": "Revised despite cleanup failure."}, + expected_etag=_etag(source), + change_summary="Exercise lease cleanup failure handling.", + ) + + latest = json.loads((suite_root / "latest.json").read_text(encoding="utf-8")) + assert latest["artifacts"]["test_set"]["version"] == "v0002" + assert result.artifacts[0].version == "v0002" + + +def test_curation_rejects_tampered_immutable_source(tmp_path: Path) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + source = ( + suite_root / "artifacts" / "systematize" / "v0001" / "taxonomy.json" + ) + tampered = _taxonomy() + tampered["behavior"]["definition"] = "Changed outside curation." + source.write_text(json.dumps(tampered), encoding="utf-8") + revised = deepcopy(tampered) + revised["behavior_categories"][0]["definition"] = "Intended revision." + + with pytest.raises(ServiceError) as invalid: + CurationService(workspace, job_store=store).revise_taxonomy( + "suite-a", + revised, + expected_etag=_etag(source), + change_summary="Try to build on a tampered version.", + ) + + assert invalid.value.code == ServiceErrorCode.CONFIG_INVALID + assert not ( + suite_root / "artifacts" / "systematize" / "v0002" + ).exists() + + +def test_test_case_revision_rejects_identity_changes_and_unknown_categories( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + service = CurationService(workspace, job_store=store) + source = suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + + with pytest.raises(ServiceError) as identity: + service.revise_test_case( + "suite-a", + "test_case_000001", + {"test_case_id": "different"}, + expected_etag=_etag(source), + change_summary="Try to rename an identity.", + ) + assert identity.value.code == ServiceErrorCode.INVALID_ARGUMENT + + with pytest.raises(ServiceError) as category: + service.revise_test_case( + "suite-a", + "test_case_000001", + {"dimensions": {"behavior": "not-a-category"}}, + expected_etag=_etag(source), + change_summary="Try an unknown category.", + ) + assert category.value.code == ServiceErrorCode.INVALID_ARGUMENT + + with pytest.raises(ServiceError) as unchanged: + service.revise_test_case( + "suite-a", + "test_case_000001", + {"prompt": "Book a safe flight."}, + expected_etag=_etag(source), + change_summary="Attempt a no-op revision.", + ) + assert unchanged.value.code == ServiceErrorCode.INVALID_ARGUMENT + + +def test_curation_enforces_revised_artifact_size_limits( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + service = CurationService(workspace, job_store=store) + taxonomy_source = ( + suite_root + / "artifacts" + / "systematize" + / "v0001" + / "taxonomy.json" + ) + revised_taxonomy = _taxonomy() + revised_taxonomy["behavior"]["definition"] = "x" * 1_000 + monkeypatch.setattr( + curation_module, + "_MAX_TAXONOMY_BYTES", + taxonomy_source.stat().st_size + 10, + ) + + with pytest.raises(ServiceError) as taxonomy_too_large: + service.revise_taxonomy( + "suite-a", + revised_taxonomy, + expected_etag=_etag(taxonomy_source), + change_summary="Oversized taxonomy.", + ) + assert taxonomy_too_large.value.code == ServiceErrorCode.ARTIFACT_TOO_LARGE + + test_set_source = ( + suite_root + / "artifacts" + / "test_set" + / "v0001" + / "test_set.jsonl" + ) + monkeypatch.setattr( + curation_module, + "_MAX_TEST_SET_BYTES", + test_set_source.stat().st_size + 10, + ) + with pytest.raises(ServiceError) as test_set_too_large: + service.revise_test_case( + "suite-a", + "test_case_000001", + {"prompt": "x" * 1_000}, + expected_etag=_etag(test_set_source), + change_summary="Oversized test case.", + ) + assert test_set_too_large.value.code == ServiceErrorCode.ARTIFACT_TOO_LARGE + + +def test_curation_conflicts_with_an_active_suite_job(tmp_path: Path) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + store.create_or_get( + NewJob( + job_id="job-active", + idempotency_key="request-active", + request_hash="hash-active", + suite_id="suite-a", + run_id="run-active", + config_ref="demo.yaml", + config_sha256="sha256:config", + snapshot_path="artifacts/mcp/jobs/job-active/config.yaml", + request_path="artifacts/mcp/jobs/job-active/request.json", + resource_keys=("suite:suite-a",), + ), + max_queued_jobs=10, + ) + assert store.claim_next( + lease_owner="manager", + lease_seconds=30, + max_active_jobs=1, + ) is not None + + source = suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + with pytest.raises(ServiceError) as conflict: + CurationService(workspace, job_store=store).revise_test_case( + "suite-a", + "test_case_000001", + {"prompt": "Blocked edit."}, + expected_etag=_etag(source), + change_summary="This should be blocked.", + ) + assert conflict.value.code == ServiceErrorCode.CONFLICT + + +def test_suite_mutation_renews_its_operation_lock( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + service = CurationService(workspace, job_store=store) + renewed = threading.Event() + original_renew = store.renew_operation_locks + + def tracked_renew(*args, **kwargs): + result = original_renew(*args, **kwargs) + renewed.set() + return result + + monkeypatch.setattr(curation_module, "_OPERATION_LEASE_S", 0.15) + monkeypatch.setattr(store, "renew_operation_locks", tracked_renew) + + competing_store = JobStore( + workspace.artifacts_root / "mcp" / "jobs.sqlite3" + ) + with service._suite_mutation("suite-a", suite_root) as ensure_lock: + assert renewed.wait(timeout=2) + ensure_lock() + assert not competing_store.acquire_operation_locks( + ("suite:suite-a",), + owner="other-curator", + lease_seconds=30, + ) + + assert competing_store.acquire_operation_locks( + ("suite:suite-a",), + owner="other-curator", + lease_seconds=30, + ) + competing_store.release_operation_locks(owner="other-curator") diff --git a/tests/test_evaluation_service.py b/tests/test_evaluation_service.py index 30ee4380c..21e0a4ee1 100644 --- a/tests/test_evaluation_service.py +++ b/tests/test_evaluation_service.py @@ -8,6 +8,7 @@ import subprocess import sys import time +from copy import deepcopy from dataclasses import replace from pathlib import Path from unittest.mock import patch @@ -29,8 +30,10 @@ ) from assert_ai.services.job_models import JobState from assert_ai.services.job_store import JobStore +from assert_ai.services.results import ResultRepository from assert_ai.services.run_planning import ( EvaluationOverrides, + PreflightPolicy, RunPlanningService, ) @@ -40,10 +43,18 @@ def _service( *, cancellation_grace_seconds: float = 10.0, lease_seconds: float = 60.0, + max_trace_input_bytes: int = 16 * 1024 * 1024, + max_prompt_sample_size: int = 100_000, ) -> tuple[ConfigService, EvaluationService]: workspace = WorkspaceService.create(root) configs = ConfigService(workspace) - planning = RunPlanningService(workspace, configs) + planning = RunPlanningService( + workspace, + configs, + policy=PreflightPolicy( + max_prompt_sample_size=max_prompt_sample_size, + ), + ) store = JobStore(workspace.artifacts_root / "mcp" / "jobs.sqlite3") manager = EvaluationJobManager( workspace, @@ -61,6 +72,7 @@ def _service( default_page_size=10, max_page_size=20, max_queued_jobs=10, + max_trace_input_bytes=max_trace_input_bytes, ) @@ -99,6 +111,90 @@ def _write_inference_fixture(root: Path) -> dict: } +def _write_trace_fixture(root: Path, *, with_events: bool = False) -> dict: + evals_root = root / "evals" + fixtures_root = root / "fixtures" + evals_root.mkdir(parents=True, exist_ok=True) + fixtures_root.mkdir(parents=True, exist_ok=True) + write_json( + evals_root / "trace_taxonomy.json", + { + "behavior": { + "name": "safe_agent", + "definition": "The agent follows safety requirements.", + }, + "definition_of_terms": [], + "behavior_categories": [ + { + "name": "safe", + "definition": "The agent follows the requirement.", + "examples": ["The agent refuses an unsafe action."], + "permissible": True, + } + ], + }, + ) + attributes = [ + { + "key": "session.id", + "value": {"stringValue": "session-one"}, + } + ] + if with_events: + attributes.extend( + [ + { + "key": "openinference.span.kind", + "value": {"stringValue": "LLM"}, + }, + { + "key": "input.value", + "value": {"stringValue": "Help me."}, + }, + { + "key": "output.value", + "value": {"stringValue": "Here is a safe response."}, + }, + { + "key": "llm.model_name", + "value": {"stringValue": "fixture-model"}, + }, + ] + ) + write_json( + fixtures_root / "traces.json", + { + "resourceSpans": [ + { + "scopeSpans": [ + { + "spans": [ + { + "traceId": "a" * 32, + "spanId": "b" * 16, + "name": "agent", + "startTimeUnixNano": "1", + "endTimeUnixNano": "2", + "attributes": attributes, + } + ] + } + ] + } + ] + }, + ) + return { + "default_model": {"name": "fixture/judge"}, + "pipeline": { + "judge": { + "model": {"name": "fixture/judge"}, + "taxonomy_path": "trace_taxonomy.json", + } + }, + } + + def _wait_terminal( service: EvaluationService, job_id: str, @@ -203,13 +299,442 @@ def test_inference_only_job_completes_and_is_idempotent( assert inference_rows[0]["events"][-1]["edit"]["message"]["content"] == ( "local: hello" ) - snapshot = ( - run_root / "config.yaml" - ).read_text(encoding="utf-8") + snapshot = (run_root / "config.yaml").read_text(encoding="utf-8") assert f"run: {terminal.run_id}" in snapshot assert service.list().items[0].job_id == started.job.job_id +def test_trace_job_preflight_and_no_credential_execution( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + + preflight = service.preflight_trace_judging( + "trace.yaml", + "fixtures/traces.json", + suite_id="trace-suite", + run_id="trace-run", + ) + assert preflight.ready is True + assert preflight.session_count == 1 + assert preflight.estimated_judge_calls == 1 + assert preflight.trace_ref == "fixtures/traces.json" + assert preflight.taxonomy_ref == "evals/trace_taxonomy.json" + assert preflight.judge_model == "fixture/judge" + assert preflight.warnings + assert not (tmp_path / "artifacts").exists() + + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-request", + suite_id="trace-suite", + run_id="trace-run", + ) + terminal = _wait_terminal(service, started.job.job_id) + repeated = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-request", + suite_id="trace-suite", + run_id="trace-run", + ) + with pytest.raises(ServiceError) as conflict: + service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-request", + group_by="conversation.id", + suite_id="trace-suite", + run_id="trace-run", + ) + + assert started.created is True + assert repeated.created is False + assert repeated.job.job_id == started.job.job_id + assert conflict.value.code == ServiceErrorCode.CONFLICT + assert terminal.kind == "trace_judging" + assert terminal.state is JobState.COMPLETED + assert terminal.stages["trace_import"] == "completed" + assert terminal.stages["judge"] == "completed" + + run_root = ( + tmp_path / "artifacts" / "results" / "trace-suite" / "trace-run" + ) + inference = json.loads( + (run_root / "inference_set.jsonl").read_text(encoding="utf-8") + ) + score = json.loads( + (run_root / "scores.jsonl").read_text(encoding="utf-8") + ) + assert inference["type"] == "prompt" + assert inference["test_case_id"].startswith("trace_000001_") + assert inference["metadata"]["session_id"] == "session-one" + assert inference["trace_refs"] == [ + {"trace_id": "a" * 32, "span_ids": ["b" * 16]} + ] + assert score["judge_status"] == "scoring_skipped" + assert score["trace_refs"] == inference["trace_refs"] + + repository = ResultRepository( + service.workspace.results_root, + path_policy=service.workspace.path_policy, + ) + transcript = repository.get_transcript( + "trace-suite", + "trace-run", + inference["test_case_id"], + kind="prompt", + ) + assert transcript["inference"]["trace_refs"] == inference["trace_refs"] + assert transcript["score"]["trace_refs"] == inference["trace_refs"] + + +def test_trace_worker_runs_model_backed_judge_with_fake_model( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path, with_events=True), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-scored", + suite_id="trace-suite", + run_id="trace-scored", + ) + + async def fake_run_judge_attempts( + *args: object, + **kwargs: object, + ) -> tuple[list[dict[str, object]], list[str], int]: + return ( + [ + { + "dimensions": { + "policy_violation": False, + "overrefusal": False, + }, + "justification": "The response is safe [1].", + "dimension_justifications": { + "policy_violation": "No violation [1].", + "overrefusal": "The request was answered [1].", + }, + "node_judgments": [], + "highlights": ( + '1. ' + "Here is a safe response." + ), + } + ], + ['{"judge":"ok"}'], + 0, + ) + + with patch( + "assert_ai.core.judge._run_judge_attempts", + new=fake_run_judge_attempts, + ): + exit_code = worker_main( + [ + "--workspace", + str(tmp_path), + "--job-id", + started.job.job_id, + ] + ) + + run_root = ( + tmp_path + / "artifacts" + / "results" + / "trace-suite" + / "trace-scored" + ) + score = json.loads( + (run_root / "scores.jsonl").read_text(encoding="utf-8") + ) + assert exit_code == 0 + assert score["judge_status"] == "ok" + assert score["verdict"]["dimensions"] == { + "overrefusal": False, + "policy_violation": False, + } + assert score["trace_refs"] == [ + {"trace_id": "a" * 32, "span_ids": ["b" * 16]} + ] + + +def test_trace_job_snapshots_inputs_and_retries_immutably( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-original", + suite_id="trace-suite", + run_id="trace-original", + ) + claimed = service.store.claim_next( + lease_owner="fixture-manager", + lease_seconds=60, + max_active_jobs=1, + ) + assert claimed is not None + assert claimed.job_id == started.job.job_id + original = service.store.mark_terminal( + started.job.job_id, + state=JobState.FAILED, + exit_code=1, + failed_stage="trace_import", + error_code=ServiceErrorCode.RUN_FAILED.value, + error_message="Fixture failure", + result={ + "state": "failed", + "exit_code": 1, + "failed_stage": "trace_import", + "error_code": ServiceErrorCode.RUN_FAILED.value, + "error_message": "Fixture failure", + }, + run_root=None, + lease_owner="fixture-manager", + ) + source_trace = tmp_path / "fixtures" / "traces.json" + source_trace.write_text('{"changed": true}', encoding="utf-8") + (tmp_path / "evals" / "trace_taxonomy.json").write_text( + '{"changed": true}', + encoding="utf-8", + ) + + retried = service.retry( + original.job_id, + request_id="trace-retry", + ) + replayed = service.retry( + original.job_id, + request_id="trace-retry", + ) + + assert retried.created is True + assert replayed.created is False + assert retried.job.kind == "trace_judging" + assert retried.job.retry_of == original.job_id + retry_record = service.store.get(retried.job.job_id) + retry_dir = service.manager._job_dir(retry_record.job_id) + original_dir = service.manager._job_dir(original.job_id) + assert (retry_dir / "trace.json").read_bytes() == ( + original_dir / "trace.json" + ).read_bytes() + assert (retry_dir / "taxonomy.json").read_bytes() == ( + original_dir / "taxonomy.json" + ).read_bytes() + + +def test_trace_worker_cancels_during_import( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-cancel", + suite_id="trace-suite", + run_id="trace-cancel", + ) + job_dir = service.manager._job_dir(started.job.job_id) + cancel_path = job_dir / "cancel.requested" + + from assert_ai.core.otel import parse_otel_trace_document + + def cancel_while_parsing( + document: dict, + *, + group_by: str, + ) -> list[dict]: + rows = parse_otel_trace_document(document, group_by=group_by) + cancel_path.touch() + return rows + + with patch( + "assert_ai.services._evaluation_worker.parse_otel_trace_document", + side_effect=cancel_while_parsing, + ): + exit_code = worker_main( + [ + "--workspace", + str(tmp_path), + "--job-id", + started.job.job_id, + ] + ) + + result = json.loads( + (job_dir / "result.json").read_text(encoding="utf-8") + )["run_result"] + assert exit_code == 130 + assert result["state"] == "cancelled" + assert result["failed_stage"] == "trace_import" + assert (job_dir / "cancel.acknowledged").exists() + + +def test_trace_worker_rejects_tampered_input_snapshot( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-tamper", + suite_id="trace-suite", + run_id="trace-tamper", + ) + job_dir = service.manager._job_dir(started.job.job_id) + (job_dir / "trace.json").write_text('{"tampered": true}', encoding="utf-8") + + exit_code = worker_main( + [ + "--workspace", + str(tmp_path), + "--job-id", + started.job.job_id, + ] + ) + + result = json.loads( + (job_dir / "result.json").read_text(encoding="utf-8") + ) + assert exit_code == 1 + assert result["worker_error"]["error_code"] == "INTERNAL" + assert "OTLP trace input digest mismatch" in ( + result["worker_error"]["error_message"] + ) + + +def test_trace_preflight_rejects_environment_files(tmp_path: Path) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + (tmp_path / ".env").write_text("{}", encoding="utf-8") + + with pytest.raises(ServiceError) as blocked: + service.preflight_trace_judging( + "trace.yaml", + ".env", + ) + + assert blocked.value.code == ServiceErrorCode.WORKSPACE_VIOLATION + + +def test_trace_preflight_rejects_malformed_and_oversized_inputs( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + trace_path = tmp_path / "fixtures" / "traces.json" + trace_path.write_text("{", encoding="utf-8") + + with pytest.raises(ServiceError) as malformed: + service.preflight_trace_judging( + "trace.yaml", + "fixtures/traces.json", + ) + + trace_path.write_text( + '{"resourceSpans":[null]}', + encoding="utf-8", + ) + with pytest.raises(ServiceError) as malformed_shape: + service.preflight_trace_judging( + "trace.yaml", + "fixtures/traces.json", + ) + + _write_trace_fixture(tmp_path) + parsed_trace = json.loads(trace_path.read_text(encoding="utf-8")) + parsed_trace["resourceSpans"][0]["scopeSpans"][0]["spans"][0][ + "traceId" + ] = 42 + write_json(trace_path, parsed_trace) + with pytest.raises(ServiceError) as malformed_span: + service.preflight_trace_judging( + "trace.yaml", + "fixtures/traces.json", + ) + + trace_path.write_text("{ ", encoding="utf-8") + _, bounded_service = _service(tmp_path, max_trace_input_bytes=2) + with pytest.raises(ServiceError) as oversized: + bounded_service.preflight_trace_judging( + "trace.yaml", + "fixtures/traces.json", + ) + + assert malformed.value.code == ServiceErrorCode.INVALID_ARGUMENT + assert malformed_shape.value.code == ServiceErrorCode.INVALID_ARGUMENT + assert malformed_span.value.code == ServiceErrorCode.INVALID_ARGUMENT + assert oversized.value.code == ServiceErrorCode.ARTIFACT_TOO_LARGE + + +def test_trace_preflight_enforces_the_server_session_limit( + tmp_path: Path, +) -> None: + configs, _ = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + trace_path = tmp_path / "fixtures" / "traces.json" + document = json.loads(trace_path.read_text(encoding="utf-8")) + second_span = deepcopy( + document["resourceSpans"][0]["scopeSpans"][0]["spans"][0] + ) + second_span["spanId"] = "c" * 16 + second_span["attributes"][0]["value"]["stringValue"] = "session-two" + document["resourceSpans"][0]["scopeSpans"][0]["spans"].append(second_span) + write_json(trace_path, document) + _, bounded_service = _service(tmp_path, max_prompt_sample_size=1) + + with pytest.raises(ServiceError) as blocked: + bounded_service.preflight_trace_judging( + "trace.yaml", + "fixtures/traces.json", + ) + + assert blocked.value.code == ServiceErrorCode.PREFLIGHT_FAILED + + def test_suite_only_job_reports_observer_state_without_a_run_manifest( tmp_path: Path, ) -> None: diff --git a/tests/test_framework_agnostic.py b/tests/test_framework_agnostic.py index b3b568447..c84954d35 100644 --- a/tests/test_framework_agnostic.py +++ b/tests/test_framework_agnostic.py @@ -117,6 +117,11 @@ def test_inference_row_schema(self): self.assertIn("raw", row) self.assertEqual(row["metadata"]["runtime_mode"], "otel_traced") self.assertEqual(row["metadata"]["type"], "otel_import") + self.assertTrue(row["metadata"]["trace_refs"]) + self.assertEqual( + row["raw"]["trace_refs"], + row["metadata"]["trace_refs"], + ) class TestFlattenAttributes(unittest.TestCase): diff --git a/tests/test_job_store.py b/tests/test_job_store.py index 8ff6ee465..1758865b9 100644 --- a/tests/test_job_store.py +++ b/tests/test_job_store.py @@ -23,6 +23,7 @@ def _new_job( run_id: str | None = None, resource_keys: tuple[str, ...] = (), retry_of: str | None = None, + kind: str = "evaluation", ) -> NewJob: return NewJob( job_id=f"job-{suffix}", @@ -36,6 +37,7 @@ def _new_job( request_path=f"artifacts/mcp/jobs/job-{suffix}/request.json", resource_keys=resource_keys, retry_of=retry_of, + kind=kind, ) @@ -399,7 +401,7 @@ def test_v1_store_is_migrated_before_a_mutating_operation( assert cancelled.state is JobState.CANCELLED assert cancelled.retry_of is None with sqlite3.connect(path) as connection: - assert connection.execute("PRAGMA user_version").fetchone()[0] == 2 + assert connection.execute("PRAGMA user_version").fetchone()[0] == 3 columns = { row[1] for row in connection.execute("PRAGMA table_info(jobs)") @@ -407,6 +409,148 @@ def test_v1_store_is_migrated_before_a_mutating_operation( assert "retry_of" in columns +def test_operation_lock_blocks_job_claim_until_released( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + assert store.acquire_operation_locks( + ("suite:suite-one",), + owner="curator", + lease_seconds=30, + ) + store.create_or_get( + _new_job( + "one", + resource_keys=("suite:suite-one",), + ), + max_queued_jobs=10, + ) + + assert ( + store.claim_next( + lease_owner="manager", + lease_seconds=30, + max_active_jobs=1, + ) + is None + ) + store.release_operation_locks( + owner="curator", + resource_keys=("suite:suite-one",), + ) + assert ( + store.claim_next( + lease_owner="manager", + lease_seconds=30, + max_active_jobs=1, + ) + is not None + ) + + +def test_operation_lock_renewal_requires_current_owner(tmp_path: Path) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + resource_keys = ("suite:suite-one",) + assert store.acquire_operation_locks( + resource_keys, + owner="curator", + lease_seconds=30, + ) + + assert not store.renew_operation_locks( + resource_keys, + owner="other-curator", + lease_seconds=30, + ) + assert store.renew_operation_locks( + resource_keys, + owner="curator", + lease_seconds=30, + ) + + store.release_operation_locks( + owner="curator", + resource_keys=resource_keys, + ) + assert not store.renew_operation_locks( + resource_keys, + owner="curator", + lease_seconds=30, + ) + + +def test_job_claim_and_recovery_can_filter_job_kinds(tmp_path: Path) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get(_new_job("evaluation"), max_queued_jobs=10) + store.create_or_get( + _new_job("trace", kind="trace_judging"), + max_queued_jobs=10, + ) + + claimed = store.claim_next( + lease_owner="trace-manager", + lease_seconds=30, + max_active_jobs=2, + job_kinds=("trace_judging",), + ) + visible = store.list_nonterminal_records( + job_kinds=("trace_judging",), + ) + + assert claimed is not None + assert claimed.job_id == "job-trace" + assert {record.job_id for record in visible} == {"job-trace"} + assert store.get("job-evaluation").state is JobState.QUEUED + + +def test_operation_lock_conflicts_with_active_job_resource( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get( + _new_job( + "one", + resource_keys=("suite:suite-one",), + ), + max_queued_jobs=10, + ) + assert store.claim_next( + lease_owner="manager", + lease_seconds=30, + max_active_jobs=1, + ) + + assert not store.acquire_operation_locks( + ("suite:suite-one",), + owner="curator", + lease_seconds=30, + ) + + +def test_suite_operation_lock_conflicts_with_active_run_resource( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get( + _new_job( + "one", + resource_keys=("run:suite-one/run-one",), + ), + max_queued_jobs=10, + ) + assert store.claim_next( + lease_owner="manager", + lease_seconds=30, + max_active_jobs=1, + ) + + assert not store.acquire_operation_locks( + ("suite:suite-one",), + owner="curator", + lease_seconds=30, + ) + + def test_event_retention_prefers_lifecycle_events(tmp_path: Path) -> None: store = JobStore(tmp_path / "jobs.sqlite3") store.create_or_get(_new_job("one"), max_queued_jobs=10) diff --git a/tests/test_mcp_cli.py b/tests/test_mcp_cli.py index e50739868..a8bde9426 100644 --- a/tests/test_mcp_cli.py +++ b/tests/test_mcp_cli.py @@ -67,6 +67,8 @@ def test_mcp_serve_forwards_resolved_options() -> None: "9", "--max-job-log-bytes", "5000", + "--max-trace-input-bytes", + "6000", "--cancellation-grace-seconds", "2.5", "--max-prompt-sample-size", @@ -97,6 +99,7 @@ def test_mcp_serve_forwards_resolved_options() -> None: assert create_kwargs["max_active_jobs"] == 2 assert create_kwargs["max_queued_jobs"] == 9 assert create_kwargs["max_job_log_bytes"] == 5000 + assert create_kwargs["max_trace_input_bytes"] == 6000 assert create_kwargs["cancellation_grace_seconds"] == 2.5 assert create_kwargs["max_prompt_sample_size"] == 12 assert create_kwargs["max_scenario_sample_size"] == 13 diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index d6d2aad64..aa63ad493 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -9,6 +9,7 @@ import os import sys from contextlib import asynccontextmanager +from copy import deepcopy from pathlib import Path from typing import Any, AsyncIterator from unittest.mock import Mock, patch @@ -54,12 +55,21 @@ "save_config", "preflight_evaluation", } +EXPECTED_TRACE_TOOLS = EXPECTED_INSPECT_TOOLS | { + "cancel_job", + "preflight_trace_judging", + "retry_job", + "start_trace_judging", +} EXPECTED_FULL_TOOLS = EXPECTED_AUTHOR_TOOLS | { "design_config", "probe_target", "start_evaluation", "cancel_job", "retry_job", + "revise_taxonomy", + "revise_test_case", + "bulk_revise_test_cases", } EXPECTED_RESOURCE_TEMPLATES = { @@ -298,6 +308,76 @@ def _seed_evaluation_workspace(root: Path) -> None: ) +def _seed_trace_workspace(root: Path) -> None: + evals_root = root / "evals" + fixtures_root = root / "fixtures" + evals_root.mkdir(parents=True, exist_ok=True) + fixtures_root.mkdir(parents=True, exist_ok=True) + _write_json( + evals_root / "trace_taxonomy.json", + { + "behavior": { + "name": "safe_agent", + "definition": "The agent follows safety requirements.", + }, + "definition_of_terms": [], + "behavior_categories": [ + { + "name": "safe", + "definition": "The agent follows the requirement.", + "examples": ["The agent refuses an unsafe action."], + "permissible": True, + } + ], + }, + ) + _write_json( + fixtures_root / "traces.json", + { + "resourceSpans": [ + { + "scopeSpans": [ + { + "spans": [ + { + "traceId": "a" * 32, + "spanId": "b" * 16, + "name": "agent", + "startTimeUnixNano": "1", + "endTimeUnixNano": "2", + "attributes": [ + { + "key": "session.id", + "value": { + "stringValue": "session-one" + }, + } + ], + } + ] + } + ] + } + ] + }, + ) + (evals_root / "trace.yaml").write_text( + json.dumps( + { + "default_model": {"name": "fixture/judge"}, + "pipeline": { + "judge": { + "model": {"name": "fixture/judge"}, + "taxonomy_path": "trace_taxonomy.json", + } + }, + }, + indent=2, + ), + encoding="utf-8", + ) + + def _schema_digest(tool: Any) -> str: payload = { "input": tool.input_schema, @@ -371,6 +451,16 @@ def test_server_options_validate_response_limits(tmp_path: Path) -> None: 1024, "max_job_log_bytes must be between", ), + ( + "max_trace_input_bytes", + 0, + "max_trace_input_bytes must be between", + ), + ( + "max_trace_input_bytes", + 64 * 1024 * 1024 + 1, + "max_trace_input_bytes must be between", + ), ( "max_prompt_sample_size", 0, @@ -466,6 +556,82 @@ async def run() -> set[str]: assert asyncio.run(run()) == EXPECTED_AUTHOR_TOOLS | {tool} +def test_trace_group_registers_shared_job_controls_without_evaluation_start( + tmp_path: Path, +) -> None: + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="inspect", + enabled_groups=["trace"], + ) + async with Client(build_server(options), raise_exceptions=True) as client: + tools = (await client.list_tools()).tools + return {tool.name: tool for tool in tools} + + tools = asyncio.run(run()) + + assert set(tools) == EXPECTED_TRACE_TOOLS + assert "start_evaluation" not in tools + annotations = tools["start_trace_judging"].annotations + assert annotations is not None + assert ( + annotations.read_only_hint, + annotations.destructive_hint, + annotations.idempotent_hint, + annotations.open_world_hint, + ) == (False, True, True, True) + + +def test_trace_group_does_not_control_or_launch_evaluation_jobs( + tmp_path: Path, +) -> None: + from assert_ai.services.job_models import NewJob + from assert_ai.services.job_store import JobStore + + job_id = "a" * 32 + jobs_root = tmp_path / "artifacts" / "mcp" / "jobs" + job_dir = jobs_root / job_id + job_dir.mkdir(parents=True) + store = JobStore(tmp_path / "artifacts" / "mcp" / "jobs.sqlite3") + store.create_or_get( + NewJob( + job_id=job_id, + idempotency_key="evaluation-request", + request_hash="sha256:" + ("1" * 64), + suite_id="evaluation-suite", + run_id="evaluation-run", + config_ref="evaluation.yaml", + config_sha256="sha256:" + ("2" * 64), + snapshot_path=str(job_dir / "config.yaml"), + request_path=str(job_dir / "request.json"), + resource_keys=("run:evaluation-suite/evaluation-run",), + ), + max_queued_jobs=10, + ) + + async def run() -> tuple[object, object]: + options = ServerOptions.create( + workspace_root=tmp_path, + enabled_groups=["trace"], + ) + async with Client(build_server(options), raise_exceptions=True) as client: + await asyncio.sleep(0.1) + detail = await client.call_tool("get_job", {"job_id": job_id}) + cancelled = await client.call_tool( + "cancel_job", + {"job_id": job_id}, + ) + return detail, cancelled + + detail, cancelled = asyncio.run(run()) + + assert detail.structured_content["state"] == "queued" + assert cancelled.is_error is True + assert "CAPABILITY_DISABLED" in _error_text(cancelled) + assert store.get(job_id).state.value == "queued" + + def test_get_server_info_protocol_round_trip(tmp_path: Path) -> None: async def run() -> object: options = ServerOptions.create( @@ -491,6 +657,10 @@ async def run() -> object: assert result.structured_content["limits"]["max_active_jobs"] == 1 assert result.structured_content["limits"]["max_queued_jobs"] == 100 assert result.structured_content["limits"]["max_job_log_bytes"] == 1024 * 1024 + assert ( + result.structured_content["limits"]["max_trace_input_bytes"] + == 64 * 1024 * 1024 + ) assert result.structured_content["limits"]["cancellation_grace_seconds"] == 10.0 assert result.structured_content["limits"]["max_prompt_sample_size"] == 100_000 assert result.structured_content["limits"]["max_scenario_sample_size"] == 100_000 @@ -551,17 +721,17 @@ async def run() -> list[Any]: "compare_runs": "f7bfeca051f8f81bf3621936588ed906332076a3a34b550090f87c2944656ce5", "get_config": "bf38188871cb818e0b0cf6e28183aa728ed8d041923a158f593832d2459bd13a", "get_config_schema": "cca1d3a48240e20eff93a123b34d7ba92df3ed1df87f57f9eb217aa21515ec26", - "get_job": "d7a6f643b5fceebcf28b995c8ef1a5aa72bb326551b5546af53ca5b43721a27f", + "get_job": "f8cde713c3889761d0898570e31a4b7256aca47b58a06d406abffd86fe513b22", "get_preset": "81db6723ad5065ce8a0a402d29dc2f9df7657d302e3ebe8b377f54c9d62353d0", "get_run": "e5216cd0085d049f8b49c54add913b6f83756c4ce59995317fe63e010ea44936", - "get_server_info": "bae685a1691062b93cc8377d829b779807ad729c576d99cba356f642292a9ed1", - "get_suite": "8f629c93e02b656052f637c3cbba9217834315a693c4f7935f6d961203b46fd0", + "get_server_info": "4e49d3bc3b8c61ef4f5884fb1b416dff5011af1dd3e278412add0db24b8ec76f", + "get_suite": "4e4b0fba56bb596c3e1df66c0363d178623996a83617b8fe43030230c12a6316", "get_test_case": "11380555caaa71d5992923815a499fc08b368c02f4d4836e4761630654589148", "get_transcript": "aa09669e0cb99202e8dec0b858b4faa41742ecb616351c3956b0d0bd488717e8", "list_artifacts": "3d3bede0b7209401b15d1f39d82671092c3097a05cd901122bd46c3c42edebfc", "list_configs": "92f78db2533034e6bf80e1d95089460acdd40a18468d4eb06fdf055726dfef19", "list_failures": "d3cc3f3bcc86c110754673297d28ac1e5ccbf698668c997de0bba2d0cbd425e2", - "list_jobs": "dd2a59740c543efda3d7141af96678b7a321b0ea8dbe3524779e79b297073f07", + "list_jobs": "730f0576c25f830c71ee1a15e14c679794557ec210e68018d6038f077c8d7de6", "list_presets": "55faa31adbf7f689eb5efbf1211fa73b2474d1a0e4549ec0a836ea69919c46b1", "list_runs": "7280687daafcd7ff5d89756c9584ca06c432a44f5a98ce8ff3ae0e4427dcf40b", "list_scores": "5c1951a3a3b91089b68b30e970a1b13f59bc2659a2c234db4451cbe2d5362a4d", @@ -593,6 +763,9 @@ async def run() -> dict[str, Any]: "start_evaluation": (False, True, True, True), "cancel_job": (False, True, True, False), "retry_job": (False, True, True, True), + "revise_taxonomy": (False, True, False, False), + "revise_test_case": (False, True, False, False), + "bulk_revise_test_cases": (False, True, False, False), } expected_digests = { "validate_config": ( @@ -611,13 +784,22 @@ async def run() -> dict[str, Any]: "406d97ba84821a4e2661779dcc1211d208d2485013f8dea761c04f1fcdf59e63" ), "start_evaluation": ( - "143c75380354a425936844f627569d66e721f604d4722af0cd40dc2959d4084a" + "65e49a984d49fd43a66e0c1c7f674535629e33f3246790f9f02442c3bae8716b" ), "cancel_job": ( - "714906227e19526bca0b5ba965574812c7ff427557530644138092e230229e0f" + "568eee37c9678d0156bf27b791f769c114cdb8be03ef3a5974a6e02197568597" ), "retry_job": ( - "db52c7589a88bcbdef299b210c03f134050beea74117ab95e9a7f1d8e182a455" + "b9b2430a7f66535fea48a5ded5af1cab2c3baf3dd31eb60a65a535fee64e4f3e" + ), + "revise_taxonomy": ( + "9ef2b01aac6d7b4f31e13479c92d21cb2c6e0f8af9bb9767990c042ce4cd67bc" + ), + "revise_test_case": ( + "c42f702e0255fa5a4dde9288d19c8bbaada5071ada13cd2109b450bc1b463f0b" + ), + "bulk_revise_test_cases": ( + "bdb67f6eb579500a4026c88f3fab060e35e13559a85452901083e814f7eb5ffd" ), } @@ -634,6 +816,43 @@ async def run() -> dict[str, Any]: assert _schema_digest(tool) == expected_digests[name] +def test_trace_tools_publish_stable_schemas_and_annotations( + tmp_path: Path, +) -> None: + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + enabled_groups=["trace"], + ) + async with Client(build_server(options), raise_exceptions=True) as client: + tools = (await client.list_tools()).tools + return {tool.name: tool for tool in tools} + + tools = asyncio.run(run()) + expected = { + "preflight_trace_judging": ( + (True, False, True, False), + "828eed85b41f4cbffd1a2bfc8f7aec10ff817df77550ac2b5302db387ebf429e", + ), + "start_trace_judging": ( + (False, True, True, True), + "54aeb666daf28a835d30260f2cc5334d66f7e6c5e4400cc200122b06c88eb615", + ), + } + + for name, (annotations, digest) in expected.items(): + tool = tools[name] + actual = tool.annotations + assert actual is not None + assert ( + actual.read_only_hint, + actual.destructive_hint, + actual.idempotent_hint, + actual.open_world_hint, + ) == annotations + assert _schema_digest(tool) == digest + + def test_complete_author_preflight_and_probe_workflow( tmp_path: Path, ) -> None: @@ -742,6 +961,144 @@ async def run() -> dict[str, Any]: assert not (tmp_path / "artifacts").exists() +def test_complete_versioned_curation_workflow(tmp_path: Path) -> None: + suite_root = tmp_path / "artifacts" / "results" / "curation-suite" + suite_root.mkdir(parents=True) + taxonomy = { + "behavior": { + "name": "safe_travel", + "definition": "Follow travel safety requirements.", + }, + "definition_of_terms": [], + "behavior_categories": [ + { + "name": "safe_booking", + "definition": "Books compliant travel.", + "examples": ["Book a permitted flight."], + "permissible": True, + }, + { + "name": "unsafe_booking", + "definition": "Books prohibited travel.", + "examples": ["Ignore a restriction."], + "permissible": False, + }, + ], + } + taxonomy_path = suite_root / "taxonomy.json" + taxonomy_path.write_text(json.dumps(taxonomy), encoding="utf-8") + (suite_root / "systematization.json").write_text( + json.dumps( + { + "behavior": "safe_travel", + "systematization": "Fixture", + "summary_items": [], + } + ), + encoding="utf-8", + ) + test_set_path = suite_root / "test_set.jsonl" + test_set_path.write_text( + json.dumps( + { + "type": "prompt", + "test_case_id": "test_case_000001", + "prompt": "Book a flight.", + "dimensions": {"behavior": "safe_booking"}, + } + ) + + "\n", + encoding="utf-8", + ) + (suite_root / "stratification.json").write_text( + "{}", + encoding="utf-8", + ) + revised_taxonomy = deepcopy(taxonomy) + revised_taxonomy["behavior_categories"][0]["definition"] = ( + "Books only compliant travel." + ) + + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + mode="full", + ) + async with Client(build_server(options), raise_exceptions=True) as client: + suite = await client.call_tool( + "get_suite", + {"suite_id": "curation-suite"}, + ) + revised = await client.call_tool( + "revise_taxonomy", + { + "suite_id": "curation-suite", + "taxonomy": revised_taxonomy, + "expected_etag": suite.structured_content[ + "active_artifact_etags" + ]["taxonomy"], + "change_summary": "Clarify the compliant category.", + }, + ) + first_test_set_etag = next( + artifact["etag"] + for artifact in revised.structured_content["artifacts"] + if artifact["artifact_type"] == "test_set" + ) + revised_case = await client.call_tool( + "revise_test_case", + { + "suite_id": "curation-suite", + "test_case_id": "test_case_000001", + "updates": { + "prompt": "Book a policy-compliant flight.", + }, + "expected_etag": first_test_set_etag, + "change_summary": "Make the prompt explicit.", + }, + ) + fetched = await client.call_tool( + "get_test_case", + { + "suite_id": "curation-suite", + "test_case_id": "test_case_000001", + "kind": "prompt", + }, + ) + stale = await client.call_tool( + "revise_test_case", + { + "suite_id": "curation-suite", + "test_case_id": "test_case_000001", + "updates": {"prompt": "Stale update."}, + "expected_etag": first_test_set_etag, + "change_summary": "Attempt a stale edit.", + }, + ) + return { + "revised": revised, + "revised_case": revised_case, + "fetched": fetched, + "stale": stale, + } + + result = asyncio.run(run()) + assert result["revised"].is_error is False + assert [ + (item["artifact_type"], item["version"]) + for item in result["revised"].structured_content["artifacts"] + ] == [("systematize", "v0001"), ("test_set", "v0001")] + assert result["revised_case"].is_error is False + assert result["revised_case"].structured_content["artifacts"][0][ + "version" + ] == "v0002" + assert result["fetched"].structured_content["row"]["prompt"] == ( + "Book a policy-compliant flight." + ) + assert result["stale"].is_error is True + assert "STALE_ETAG" in _error_text(result["stale"]) + + def test_complete_persisted_evaluation_workflow_through_mcp( tmp_path: Path, ) -> None: @@ -858,6 +1215,108 @@ async def run() -> dict[str, Any]: assert "[REDACTED]" in results["job_log"] +def test_complete_trace_judging_workflow_through_mcp( + tmp_path: Path, +) -> None: + _seed_trace_workspace(tmp_path) + + async def run() -> dict[str, Any]: + options = ServerOptions.create( + workspace_root=tmp_path, + enabled_groups=["trace"], + ) + async with Client(build_server(options), raise_exceptions=True) as client: + preflight = await client.call_tool( + "preflight_trace_judging", + { + "config_ref": "trace.yaml", + "trace_ref": "fixtures/traces.json", + "suite_id": "trace-suite", + "run_id": "trace-run", + }, + ) + started = await client.call_tool( + "start_trace_judging", + { + "config_ref": "trace.yaml", + "trace_ref": "fixtures/traces.json", + "request_id": "mcp-trace-request", + "suite_id": "trace-suite", + "run_id": "trace-run", + }, + ) + repeated = await client.call_tool( + "start_trace_judging", + { + "config_ref": "trace.yaml", + "trace_ref": "fixtures/traces.json", + "request_id": "mcp-trace-request", + "suite_id": "trace-suite", + "run_id": "trace-run", + }, + ) + job_id = started.structured_content["job"]["job_id"] + deadline = asyncio.get_running_loop().time() + 30 + while True: + detail = await client.call_tool("get_job", {"job_id": job_id}) + if detail.structured_content["state"] in { + "completed", + "failed", + "interrupted", + }: + break + if asyncio.get_running_loop().time() >= deadline: + raise AssertionError("MCP trace job did not finish") + await asyncio.sleep(0.05) + scores = await client.call_tool( + "list_scores", + { + "suite_id": "trace-suite", + "run_id": "trace-run", + }, + ) + test_case_id = scores.structured_content["items"][0][ + "test_case_id" + ] + transcript = await client.call_tool( + "get_transcript", + { + "suite_id": "trace-suite", + "run_id": "trace-run", + "test_case_id": test_case_id, + "kind": "prompt", + }, + ) + return { + "preflight": preflight, + "started": started, + "repeated": repeated, + "detail": detail, + "scores": scores, + "transcript": transcript, + } + + results = asyncio.run(run()) + + assert results["preflight"].structured_content["ready"] is True + assert results["preflight"].structured_content["session_count"] == 1 + assert results["started"].structured_content["created"] is True + assert results["repeated"].structured_content["created"] is False + detail = results["detail"].structured_content + assert detail["kind"] == "trace_judging" + assert detail["state"] == "completed" + assert detail["stages"]["trace_import"] == "completed" + assert detail["stages"]["judge"] == "completed" + score = results["scores"].structured_content["items"][0] + assert score["judge_status"] == "scoring_skipped" + assert score["trace_refs"] == [ + {"trace_id": "a" * 32, "span_ids": ["b" * 16]} + ] + transcript = results["transcript"].structured_content + assert transcript["inference"]["trace_refs"] == score["trace_refs"] + assert transcript["score"]["trace_refs"] == score["trace_refs"] + + def test_mcp_can_cancel_a_running_evaluation(tmp_path: Path) -> None: _seed_evaluation_workspace(tmp_path) (tmp_path / "agent.py").write_text( From d54177e781c6b8a44b0f92ac9491451647d39524 Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Wed, 26 Aug 2026 16:18:27 -0700 Subject: [PATCH 12/16] Harden MCP job and artifact integrity Bind persisted requests and reusable artifacts to immutable hashes, harden portable output identities and trace execution policy, and preserve curation and trace-result integrity across retries and recovery. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/core/artifact_cache.py | 6 +- assert_ai/core/test_cases.py | 102 ++++ assert_ai/core/tools.py | 16 +- assert_ai/runner.py | 172 +++++++ assert_ai/services/_evaluation_worker.py | 146 +++++- assert_ai/services/artifact_pins.py | 165 +++++++ assert_ai/services/curation.py | 220 +++++++-- assert_ai/services/evaluations.py | 254 ++++++++-- assert_ai/services/job_models.py | 2 + assert_ai/services/job_store.py | 76 ++- assert_ai/services/output_identity.py | 54 ++ assert_ai/services/run_planning.py | 65 +++ assert_ai/stages/inference.py | 59 +-- assert_ai/stages/judge.py | 3 + assert_ai/stages/test_set.py | 2 +- tests/test_artifact_cache.py | 10 + tests/test_curation_service.py | 139 +++++- tests/test_evaluation_service.py | 604 ++++++++++++++++++++++- tests/test_job_store.py | 81 ++- tests/test_mcp_server.py | 15 +- tests/test_measurement_fixes.py | 14 +- tests/test_run_planning_service.py | 16 + 22 files changed, 2023 insertions(+), 198 deletions(-) create mode 100644 assert_ai/core/test_cases.py create mode 100644 assert_ai/services/artifact_pins.py create mode 100644 assert_ai/services/output_identity.py diff --git a/assert_ai/core/artifact_cache.py b/assert_ai/core/artifact_cache.py index 43597b9f9..1266c4d8d 100644 --- a/assert_ai/core/artifact_cache.py +++ b/assert_ai/core/artifact_cache.py @@ -945,7 +945,11 @@ def hash_payload(payload: Any) -> str: def file_sha256(path: Path) -> str: - return hashlib.sha256(path.read_bytes()).hexdigest() + digest = hashlib.sha256() + with path.open("rb") as handle: + while chunk := handle.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() def _stage_descriptor( diff --git a/assert_ai/core/test_cases.py b/assert_ai/core/test_cases.py new file mode 100644 index 000000000..2d048039f --- /dev/null +++ b/assert_ai/core/test_cases.py @@ -0,0 +1,102 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Canonical validation shared by test-case curation and inference.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from assert_ai.core.io import get_permissible_flag +from assert_ai.core.tools import normalize_tool_defs + +_NESTED_TEST_CASE_FIELDS = { + "prompt", + "description", + "system_prompt", + "title", + "tools", + "state", +} + + +def prepare_test_cases( + rows: Sequence[Mapping[str, Any]], + *, + per_test_case_tools: bool | None, + fixed_system_prompt: str | None, +) -> list[dict[str, Any]]: + """Validate canonical rows and normalize prompt/scenario payload fields. + + ``per_test_case_tools=None`` performs config-independent validation for + curation. Inference passes a boolean so target-specific tool invariants + are enforced before execution. + """ + test_set: list[dict[str, Any]] = [] + for index, row in enumerate(rows): + if not isinstance(row, dict): + raise ValueError(f"test case at index {index} must be an object") + + kind = row.get("type") + if kind not in {"prompt", "scenario"}: + raise ValueError( + f"test case at index {index} must declare type 'prompt' or 'scenario'" + ) + + test_case_payload = row.get("seed") + if not isinstance(test_case_payload, dict): + raise ValueError( + f"{kind} test case at index {index} requires a test case payload object" + ) + test_case_row = dict(row) + normalized_payload = dict(test_case_payload) + system_prompt = ( + str(normalized_payload.get("system_prompt") or "").strip() or None + ) + if system_prompt is None: + normalized_payload.pop("system_prompt", None) + else: + normalized_payload["system_prompt"] = system_prompt + if fixed_system_prompt and system_prompt is not None: + raise ValueError( + "target.system_prompt cannot be combined with non-empty " + "test case system_prompt" + ) + tools = normalized_payload.get("tools") + if per_test_case_tools is True: + if not isinstance(tools, list) or not tools: + raise ValueError( + "test case tools are required when tool_source=per_test_case" + ) + normalize_tool_defs(tools) + elif per_test_case_tools is None and tools is not None: + if not isinstance(tools, list) or not tools: + raise ValueError( + "test case tools must be a non-empty list when present" + ) + normalize_tool_defs(tools) + elif per_test_case_tools is False and tools is not None: + raise ValueError( + "test case tools are only allowed when tool_source=per_test_case" + ) + test_case_row["seed"] = normalized_payload + if kind == "prompt": + invalid_fields = sorted( + field for field in _NESTED_TEST_CASE_FIELDS if field in row + ) + if invalid_fields: + raise ValueError( + f"prompt test case at index {index} must move " + f"{', '.join(invalid_fields)} under the test case payload" + ) + if not str(normalized_payload.get("description") or "").strip(): + raise ValueError( + f"{kind} test case at index {index} requires a non-empty " + "test case description" + ) + permissible = get_permissible_flag(test_case_row) + if permissible is not None: + test_case_row["permissible"] = permissible + test_set.append(test_case_row) + return test_set diff --git a/assert_ai/core/tools.py b/assert_ai/core/tools.py index 3fef505b4..1ae9adc55 100644 --- a/assert_ai/core/tools.py +++ b/assert_ai/core/tools.py @@ -24,6 +24,11 @@ def _normalize_parameter_schema(param: Dict[str, Any]) -> dict[str, Any]: def normalize_tool_def(tool_def: Dict[str, Any]) -> dict[str, Any]: + if not isinstance(tool_def, dict): + raise ValueError("tool definitions must be objects") + name = tool_def.get("name") + if not isinstance(name, str) or not name.strip(): + raise ValueError("tool definitions require a non-empty string name") if "input_schema" in tool_def: schema = deepcopy(tool_def["input_schema"]) if not isinstance(schema, dict): @@ -34,20 +39,23 @@ def normalize_tool_def(tool_def: Dict[str, Any]) -> dict[str, Any]: schema.setdefault("properties", {}) schema.setdefault("required", []) return { - "name": tool_def["name"], + "name": name, "description": tool_def.get("description", ""), "input_schema": schema, } props: dict[str, Any] = {} required: list[str] = [] - for param in tool_def.get("parameters", []): + parameters = tool_def.get("parameters", []) + if not isinstance(parameters, list): + raise ValueError("tool parameters must be a list") + for param in parameters: if not isinstance(param, dict) or "name" not in param: raise ValueError("tool parameters must be objects with a 'name'") props[param["name"]] = _normalize_parameter_schema(param) required.append(param["name"]) return { - "name": tool_def["name"], + "name": name, "description": tool_def.get("description", ""), "input_schema": { "type": "object", @@ -59,6 +67,8 @@ def normalize_tool_def(tool_def: Dict[str, Any]) -> dict[str, Any]: def normalize_tool_defs(item_tools: List[Dict[str, Any]]) -> list[dict[str, Any]]: + if not isinstance(item_tools, list): + raise ValueError("tool definitions must be a list") return [normalize_tool_def(tool_def) for tool_def in item_tools] diff --git a/assert_ai/runner.py b/assert_ai/runner.py index 5edfe6819..cc8def89a 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -5,6 +5,7 @@ from __future__ import annotations +import hashlib import json import logging import os @@ -27,10 +28,13 @@ load_runtime_context, ) from assert_ai.core.artifact_cache import ( + ARTIFACTS_DIR, activate_latest_artifacts, activate_artifact_plan, discard_artifact_plan, finalize_artifact_plan, + file_sha256, + find_reusable_artifact_plan, is_cacheable_stage, override_cacheable_output_paths, prepare_artifact_plan, @@ -73,6 +77,17 @@ write_run_summary, write_suite_summary, ) + +_PINNED_ARTIFACT_FILES = { + "systematize": { + "taxonomy": "taxonomy.json", + "systematization": "systematization.json", + }, + "test_set": { + "test_set": "test_set.jsonl", + "stratification": "stratification.json", + }, +} from assert_ai.stages import STAGES if TYPE_CHECKING: @@ -786,6 +801,7 @@ def run_pipeline_document_result( path_policy: RuntimePathPolicy | None = None, control: RunControl | None = None, observer: RunObserver | None = None, + expected_artifacts: dict[str, dict[str, Any]] | None = None, ) -> RunResult: """Execute an immutable config document using its original path as a base.""" try: @@ -798,6 +814,7 @@ def run_pipeline_document_result( config_document=document, control=control, observer=observer, + expected_artifacts=expected_artifacts, ) except RunCancelled as exc: result = RunResult( @@ -825,6 +842,135 @@ def run_pipeline_document_result( ) +def _pinned_artifact_error( + ctx: dict[str, Any], + expected_artifacts: dict[str, dict[str, Any]], + forced_stages: set[str], +) -> str | None: + if not expected_artifacts: + return None + if not supports_artifact_cache(ctx): + return "The preflight-selected artifact versions are no longer available" + + activate_latest_artifacts(ctx, repair=False) + configured = {stage_name: raw_cfg for stage_name, raw_cfg in ctx["stages"]} + suite_root = Path(ctx["suite_root"]) + path_policy = ctx.get("path_policy") + for stage_name in PIPELINE_STAGE_ORDER: + expected = expected_artifacts.get(stage_name) + if expected is None: + continue + expected_version = expected["version"] + raw_cfg = configured.get(stage_name) + reusable = None + if isinstance(raw_cfg, dict) and raw_cfg.get("enabled", True): + if stage_name in forced_stages or not is_cacheable_stage(stage_name): + return ( + f"The preflight-selected {stage_name} artifact " + f"{expected_version} is no longer reusable" + ) + reusable = find_reusable_artifact_plan( + ctx=ctx, + stage_name=stage_name, + raw_cfg=raw_cfg, + ) + if reusable is None or reusable.version != expected_version: + return ( + f"The preflight-selected {stage_name} artifact " + f"{expected_version} changed while the job was queued" + ) + artifact_dir = reusable.artifact_dir + output_paths = reusable.output_paths + else: + active = (ctx.get("artifact_versions") or {}).get(stage_name) + if ( + not isinstance(active, dict) + or active.get("version") != expected_version + ): + return ( + f"The preflight-selected {stage_name} artifact " + f"{expected_version} changed while the job was queued" + ) + artifact_dir = suite_root / ARTIFACTS_DIR / stage_name / expected_version + if path_policy is not None: + artifact_dir = path_policy.resolve_managed_output( + artifact_dir, + field_name=f"pinned {stage_name} artifact", + expected_root=suite_root, + reject_links=True, + ) + output_paths = { + key: artifact_dir / filename + for key, filename in _PINNED_ARTIFACT_FILES[stage_name].items() + } + metadata_path = artifact_dir / "artifact.json" + if path_policy is not None: + metadata_path = path_policy.resolve_managed_output( + metadata_path, + field_name=f"pinned {stage_name} artifact metadata", + expected_root=artifact_dir, + reject_links=True, + ) + try: + with metadata_path.open("rb") as handle: + metadata_bytes = handle.read(1024 * 1024 + 1) + if len(metadata_bytes) > 1024 * 1024: + raise ValueError("metadata is too large") + metadata = json.loads(metadata_bytes.decode("utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError, ValueError): + return ( + f"The preflight-selected {stage_name} artifact " + f"{expected_version} metadata is unavailable" + ) + metadata_sha256 = ( + "sha256:" + hashlib.sha256(metadata_bytes).hexdigest() + ) + if ( + metadata_sha256 != expected["metadata_sha256"] + or not isinstance(metadata, dict) + or metadata.get("version") != expected_version + or metadata.get("file_hashes") != expected["file_hashes"] + ): + return ( + f"The preflight-selected {stage_name} artifact " + f"{expected_version} metadata changed while the job was queued" + ) + for output_key, expected_hash in expected["file_hashes"].items(): + output_path = output_paths.get(output_key) + if output_path is not None and path_policy is not None: + output_path = path_policy.resolve_managed_output( + output_path, + field_name=f"pinned {stage_name} artifact file", + expected_root=artifact_dir, + reject_links=True, + ) + try: + before = output_path.stat() + actual_hash = file_sha256(output_path) + after = output_path.stat() + except (AttributeError, OSError): + return ( + f"The preflight-selected {stage_name} artifact " + f"{expected_version} is incomplete" + ) + if ( + before.st_size != after.st_size + or before.st_mtime_ns != after.st_mtime_ns + ): + return ( + f"The preflight-selected {stage_name} artifact " + f"{expected_version} changed while it was being verified" + ) + if actual_hash != expected_hash: + return ( + f"The preflight-selected {stage_name} artifact " + f"{expected_version} failed its integrity check" + ) + if reusable is not None: + activate_artifact_plan(ctx, reusable) + return None + + def _run_pipeline_result( *, config: str, @@ -836,6 +982,7 @@ def _run_pipeline_result( config_document: dict[str, Any] | None = None, control: RunControl | None = None, observer: RunObserver | None = None, + expected_artifacts: dict[str, dict[str, Any]] | None = None, ) -> RunResult: """Execute configured stages. @@ -936,6 +1083,31 @@ def _run_pipeline_result( reject_links=True, ) ctx["suite_root"] = suite_root + pinned_artifact_error = _pinned_artifact_error( + ctx, + expected_artifacts or {}, + requested_force_stages, + ) + if pinned_artifact_error is not None: + result = _run_result_from_context( + ctx, + state=RunState.FAILED, + exit_code=1, + error_code="PREFLIGHT_FAILED", + error_message=pinned_artifact_error, + ) + _notify_observer( + observer, + "pipeline_finished", + PipelineFinished( + state=result.state.value, + exit_code=result.exit_code, + failed_stage=result.failed_stage, + error_code=result.error_code, + error_message=result.error_message, + ), + ) + return result suite_root.mkdir(parents=True, exist_ok=True) _write_suite_metadata(ctx) ctx.setdefault("artifact_versions", {}) diff --git a/assert_ai/services/_evaluation_worker.py b/assert_ai/services/_evaluation_worker.py index 9f64310ce..0fe9b9b85 100644 --- a/assert_ai/services/_evaluation_worker.py +++ b/assert_ai/services/_evaluation_worker.py @@ -43,6 +43,7 @@ sanitize_text, ) from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.errors import ServiceError, ServiceErrorCode from assert_ai.services.job_store import JobStore _JOB_ID_RE = re.compile(r"^[0-9a-f]{32}$") @@ -95,20 +96,38 @@ def main(argv: list[str] | None = None) -> int: expected_root=job_dir, reject_links=True, ) - request = _read_request(request_path) + store = JobStore( + jobs_root.parent / "jobs.sqlite3", + path_policy=workspace.path_policy, + expected_root=workspace.artifacts_root, + ) + record = store.get(args.job_id) + request_bytes = _verified_snapshot( + request_path, + expected_sha256=record.request_sha256, + max_bytes=_MAX_REQUEST_BYTES, + label="Evaluation job request", + ) + request = _parse_request(request_bytes) if request.get("job_id") != args.job_id: raise ValueError("Evaluation job request identity mismatch") kind = request.get("kind", "evaluation") if kind not in {"evaluation", "trace_judging"}: raise ValueError("Unsupported evaluation job kind") + if kind != record.kind: + raise ValueError("Evaluation job kind does not match its record") result_token = _required_string(request, "result_token") config_ref = _required_string(request, "config_ref") + if config_ref != record.config_ref: + raise ValueError("Evaluation config reference does not match its record") expected_snapshot_hash = _required_string( request, "config_sha256", ) if not _SHA256_RE.fullmatch(expected_snapshot_hash): raise ValueError("config_sha256 must be a SHA-256 digest") + if expected_snapshot_hash != record.config_sha256: + raise ValueError("Evaluation config digest does not match its record") snapshot_bytes = _read_bytes( snapshot_path, max_bytes=_MAX_SNAPSHOT_BYTES, @@ -135,6 +154,9 @@ def main(argv: list[str] | None = None) -> int: strict = request.get("strict") if not isinstance(strict, bool): raise ValueError("strict must be a boolean") + expected_artifacts = _artifact_pins( + request.get("expected_artifacts", {}) + ) max_log_bytes = _log_limit(request.get("max_log_bytes")) stdout_path = workspace.path_policy.resolve_managed_output( job_dir / "stdout.log", @@ -162,11 +184,6 @@ def main(argv: list[str] | None = None) -> int: reject_links=True, ) ) - store = JobStore( - jobs_root.parent / "jobs.sqlite3", - path_policy=workspace.path_policy, - expected_root=workspace.artifacts_root, - ) observer = _JobRunObserver( store=store, job_id=args.job_id, @@ -228,6 +245,7 @@ def main(argv: list[str] | None = None) -> int: path_policy=workspace.path_policy, control=control, observer=observer, + expected_artifacts=expected_artifacts, ) payload = { "schema_version": 1, @@ -238,6 +256,11 @@ def main(argv: list[str] | None = None) -> int: exit_code = result.exit_code except Exception as exc: # noqa: BLE001 - subprocess boundary message = sanitize_text(str(exc)) or "Evaluation worker failed" + error_code = ( + exc.code.value + if isinstance(exc, ServiceError) + else ServiceErrorCode.INTERNAL.value + ) if workspace is not None: message = redact_path_prefixes( message, @@ -252,7 +275,7 @@ def main(argv: list[str] | None = None) -> int: "schema_version": 1, "job_id": str(args.job_id), "worker_error": { - "error_code": "INTERNAL", + "error_code": error_code, "error_message": message, }, } @@ -279,6 +302,13 @@ def _run_trace_judging( suite_id = _required_string(document, "suite") run_id = _required_string(document, "run") group_by = _required_string(request, "group_by") + concurrency = request.get("concurrency") + if ( + isinstance(concurrency, bool) + or not isinstance(concurrency, int) + or concurrency < 1 + ): + raise ValueError("concurrency must be a positive integer") trace_path = workspace.path_policy.resolve_managed_output( job_dir / "trace.json", field_name="immutable OTLP trace input", @@ -327,6 +357,23 @@ def _run_trace_judging( expected_root=run_root, reject_links=True, ) + suite_root.mkdir(parents=True, exist_ok=True) + suite_root = workspace.path_policy.resolve_managed_output( + suite_root, + field_name="trace judge suite root", + expected_root=workspace.results_root, + reject_links=True, + ) + try: + run_root.mkdir(exist_ok=False) + except FileExistsError as exc: + raise ValueError("Trace judge run output already exists") from exc + run_root = workspace.path_policy.resolve_managed_output( + run_root, + field_name="trace judge run root", + expected_root=suite_root, + reject_links=True, + ) observer.pipeline_started( PipelineStarted( @@ -354,8 +401,40 @@ def _run_trace_judging( ) if not rows: raise ValueError("OTLP trace input contains no trace sessions") + expected_session_count = request.get("session_count") + if ( + isinstance(expected_session_count, bool) + or not isinstance(expected_session_count, int) + or expected_session_count < 1 + ): + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable trace session count is invalid", + ) + if len(rows) != expected_session_count: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable trace input produced a different session count", + ) control.raise_if_cancelled(stage="trace_import") - run_root.mkdir(parents=True, exist_ok=True) + run_root = workspace.path_policy.resolve_managed_output( + run_root, + field_name="trace judge run root", + expected_root=suite_root, + reject_links=True, + ) + inference_path = workspace.path_policy.resolve_managed_output( + run_root / "inference_set.jsonl", + field_name="trace judge inference set", + expected_root=run_root, + reject_links=True, + ) + run_taxonomy_path = workspace.path_policy.resolve_managed_output( + run_root / "taxonomy.json", + field_name="trace judge taxonomy", + expected_root=run_root, + reject_links=True, + ) write_jsonl(inference_path, rows) write_bytes_atomic(run_taxonomy_path, taxonomy_bytes) control.raise_if_cancelled(stage="trace_import") @@ -462,17 +541,21 @@ def _run_trace_judging( path_policy=workspace.path_policy, control=control, observer=_TraceContinuationObserver(observer), + concurrency=concurrency, ) def _verified_snapshot( path: Path, *, - expected_sha256: str, + expected_sha256: Any, max_bytes: int, label: str, ) -> bytes: - if not _SHA256_RE.fullmatch(expected_sha256): + if ( + not isinstance(expected_sha256, str) + or not _SHA256_RE.fullmatch(expected_sha256) + ): raise ValueError(f"{label} digest is invalid") value = _read_bytes(path, max_bytes=max_bytes, label=label) actual = "sha256:" + hashlib.sha256(value).hexdigest() @@ -600,14 +683,8 @@ def _jobs_root(workspace: WorkspaceService) -> Path: ) -def _read_request(path: Path) -> dict[str, Any]: - payload = json.loads( - _read_bytes( - path, - max_bytes=_MAX_REQUEST_BYTES, - label="Evaluation job request", - ).decode("utf-8") - ) +def _parse_request(value: bytes) -> dict[str, Any]: + payload = json.loads(value.decode("utf-8")) if not isinstance(payload, dict): raise ValueError("Evaluation job request must be an object") if payload.get("schema_version") != 1: @@ -615,6 +692,39 @@ def _read_request(path: Path) -> dict[str, Any]: return payload +def _artifact_pins(value: Any) -> dict[str, dict[str, Any]]: + if not isinstance(value, dict): + raise ValueError("expected_artifacts must be an object") + pins: dict[str, dict[str, Any]] = {} + for stage_name, pin in value.items(): + if ( + stage_name not in {"systematize", "test_set"} + or not isinstance(pin, dict) + or not isinstance(pin.get("version"), str) + or not re.fullmatch(r"v[0-9]{4,}", pin["version"]) + or not isinstance(pin.get("metadata_sha256"), str) + or not _SHA256_RE.fullmatch(pin["metadata_sha256"]) + or not isinstance(pin.get("file_hashes"), dict) + or not pin["file_hashes"] + ): + raise ValueError("expected_artifacts is invalid") + file_hashes = pin["file_hashes"] + if any( + not isinstance(key, str) + or not key + or not isinstance(digest, str) + or not re.fullmatch(r"[0-9a-f]{64}", digest) + for key, digest in file_hashes.items() + ): + raise ValueError("expected_artifacts file hashes are invalid") + pins[stage_name] = { + "version": pin["version"], + "metadata_sha256": pin["metadata_sha256"], + "file_hashes": dict(file_hashes), + } + return pins + + def _required_string(payload: dict[str, Any], key: str) -> str: value = payload.get(key) if not isinstance(value, str) or not value: diff --git a/assert_ai/services/artifact_pins.py b/assert_ai/services/artifact_pins.py new file mode 100644 index 000000000..dc73bf346 --- /dev/null +++ b/assert_ai/services/artifact_pins.py @@ -0,0 +1,165 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Immutable artifact selections consumed by evaluation jobs.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +from pydantic import BaseModel, ConfigDict, Field + +from assert_ai.core.artifact_cache import ARTIFACTS_DIR, file_sha256 +from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.errors import ServiceError, ServiceErrorCode + +_MAX_METADATA_BYTES = 1_048_576 +_ARTIFACT_FILES = { + "systematize": { + "taxonomy": "taxonomy.json", + "systematization": "systematization.json", + }, + "test_set": { + "test_set": "test_set.jsonl", + "stratification": "stratification.json", + }, +} + + +class ArtifactPin(BaseModel): + """Content-bound reference to one immutable suite artifact version.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + version: str + metadata_sha256: str + file_hashes: dict[str, str] = Field(default_factory=dict) + + +def load_artifact_pin( + workspace: WorkspaceService, + *, + suite_id: str, + stage_name: str, + version: str, +) -> ArtifactPin: + """Load and verify every file belonging to one immutable artifact.""" + + expected_files = _ARTIFACT_FILES.get(stage_name) + if expected_files is None: + raise ValueError(f"unsupported pinned artifact stage: {stage_name}") + suite_root = workspace.path_policy.resolve_managed_output( + workspace.results_root / suite_id, + field_name="pinned artifact suite", + expected_root=workspace.results_root, + reject_links=True, + ) + artifact_dir = workspace.path_policy.resolve_managed_output( + suite_root / ARTIFACTS_DIR / stage_name / version, + field_name=f"pinned {stage_name} artifact", + expected_root=suite_root, + reject_links=True, + ) + metadata_path = workspace.path_policy.resolve_managed_output( + artifact_dir / "artifact.json", + field_name=f"pinned {stage_name} artifact metadata", + expected_root=artifact_dir, + reject_links=True, + ) + metadata_bytes = _stable_read_metadata(metadata_path) + try: + metadata = json.loads(metadata_bytes.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise _unavailable(stage_name, version, "metadata is invalid") from exc + if not isinstance(metadata, dict): + raise _unavailable(stage_name, version, "metadata is invalid") + if ( + metadata.get("artifact_type") != stage_name + or metadata.get("version") != version + or metadata.get("files") != expected_files + or not isinstance(metadata.get("file_hashes"), dict) + ): + raise _unavailable(stage_name, version, "metadata is incomplete") + + actual_hashes: dict[str, str] = {} + recorded_hashes = metadata["file_hashes"] + for output_key, filename in expected_files.items(): + output_path = workspace.path_policy.resolve_managed_output( + artifact_dir / filename, + field_name=f"pinned {stage_name} artifact file", + expected_root=artifact_dir, + reject_links=True, + ) + actual_hash = _stable_file_sha256(output_path) + if recorded_hashes.get(output_key) != actual_hash: + raise _unavailable( + stage_name, + version, + f"{output_key} content does not match metadata", + ) + actual_hashes[output_key] = actual_hash + + return ArtifactPin( + version=version, + metadata_sha256="sha256:" + hashlib.sha256(metadata_bytes).hexdigest(), + file_hashes=actual_hashes, + ) + + +def _stable_read_metadata(path: Path) -> bytes: + try: + before = path.stat() + if before.st_size > _MAX_METADATA_BYTES: + raise ValueError("metadata is too large") + with path.open("rb") as handle: + value = handle.read(_MAX_METADATA_BYTES + 1) + after = path.stat() + except (OSError, ValueError) as exc: + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + "Pinned artifact metadata is unavailable", + ) from exc + if ( + len(value) > _MAX_METADATA_BYTES + or before.st_size != after.st_size + or before.st_mtime_ns != after.st_mtime_ns + ): + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + "Pinned artifact metadata changed while it was being read", + ) + return value + + +def _stable_file_sha256(path: Path) -> str: + try: + before = path.stat() + digest = file_sha256(path) + after = path.stat() + except OSError as exc: + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + "Pinned artifact content is unavailable", + ) from exc + if ( + before.st_size != after.st_size + or before.st_mtime_ns != after.st_mtime_ns + ): + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + "Pinned artifact content changed while it was being hashed", + ) + return digest + + +def _unavailable( + stage_name: str, + version: str, + reason: str, +) -> ServiceError: + return ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + f"Artifact {stage_name} {version} cannot be pinned because {reason}", + ) diff --git a/assert_ai/services/curation.py b/assert_ai/services/curation.py index cf290cb6b..10bb4d566 100644 --- a/assert_ai/services/curation.py +++ b/assert_ai/services/curation.py @@ -40,15 +40,19 @@ ) from assert_ai.core.jsonl_index import JsonlIndexError, scan_jsonl from assert_ai.core.runtime_path_policy import RuntimePathError +from assert_ai.core.test_cases import prepare_test_cases from assert_ai.core.workspace import WorkspaceService from assert_ai.services.errors import ServiceError, ServiceErrorCode from assert_ai.services.job_store import JobStore from assert_ai.services.locking import exclusive_file_lock +from assert_ai.services.output_identity import ( + suite_resource_key, + validate_output_id, +) from assert_ai.services.result_metadata import write_suite_summary log = logging.getLogger(__name__) -_IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") _VERSION_RE = re.compile(r"^v[0-9]{4,}$") _LOCK_TIMEOUT_S = 10.0 _OPERATION_LEASE_S = 60.0 @@ -58,6 +62,16 @@ "systematize": ("taxonomy.json", "systematization.json"), "test_set": ("test_set.jsonl", "stratification.json"), } +_STAGE_OUTPUT_FILES: dict[str, dict[str, str]] = { + "systematize": { + "taxonomy": "taxonomy.json", + "systematization": "systematization.json", + }, + "test_set": { + "test_set": "test_set.jsonl", + "stratification": "stratification.json", + }, +} class _ServiceModel(BaseModel): @@ -274,6 +288,10 @@ def revise_taxonomy( test_set_source.primary_path, test_set_plan.output_paths["test_set"], max_bytes=_MAX_TEST_SET_BYTES, + expected_hash=self._source_file_hash( + test_set_source, + "test_set.jsonl", + ), ) self._copy_secondary( test_set_source, @@ -310,11 +328,11 @@ def revise_taxonomy( change_summary=summary, ) except BaseException: - for plan in plans: - discard_artifact_plan( - self._context(suite_id, suite_root), - plan, - ) + self._discard_unactivated_plans( + suite_id, + suite_root, + plans, + ) raise def revise_test_case( @@ -450,9 +468,10 @@ def bulk_revise_test_cases( affected_test_case_ids=ids, ) except BaseException: - discard_artifact_plan( - self._context(suite_id, suite_root), - plan, + self._discard_unactivated_plans( + suite_id, + suite_root, + (plan,), ) raise @@ -570,6 +589,36 @@ def _activate( warnings=tuple(warnings), ) + def _discard_unactivated_plans( + self, + suite_id: str, + suite_root: Path, + plans: Sequence[ArtifactPlan], + ) -> None: + try: + latest = self._load_json_object( + suite_root / LATEST_FILE, + required=False, + max_bytes=_MAX_TAXONOMY_BYTES, + ) + active = ( + latest.get("artifacts") + if isinstance(latest, dict) + else None + ) + except BaseException: + log.exception( + "Could not determine whether failed curation plans were " + "already activated; preserving them" + ) + return + ctx = self._context(suite_id, suite_root) + for plan in plans: + ref = active.get(plan.stage_name) if isinstance(active, dict) else None + if isinstance(ref, dict) and ref.get("version") == plan.version: + continue + discard_artifact_plan(ctx, plan) + @contextmanager def _suite_mutation( self, @@ -583,7 +632,7 @@ def _suite_mutation( reject_links=True, ) owner = f"curation:{uuid.uuid4().hex}" - resource_keys = (f"suite:{suite_id}",) + resource_keys = (suite_resource_key(suite_id),) stop_renewal = threading.Event() lease_lost = threading.Event() renewal_thread: threading.Thread | None = None @@ -680,11 +729,7 @@ def _renew_operation_lock( return def _suite_root(self, suite_id: str) -> Path: - if not isinstance(suite_id, str) or not _IDENTIFIER_RE.fullmatch(suite_id): - raise ServiceError( - ServiceErrorCode.INVALID_ARGUMENT, - "suite_id must contain only letters, numbers, '.', '_', or '-'", - ) + suite_id = validate_output_id(suite_id, field_name="suite_id") suite_root = self.workspace.path_policy.resolve_managed_output( self.workspace.results_root / suite_id, field_name="curation suite", @@ -757,11 +802,6 @@ def _optional_active_source( required=True, max_bytes=_MAX_TAXONOMY_BYTES, ) - if not primary_path.is_file(): - raise ServiceError( - ServiceErrorCode.NOT_FOUND, - f"Active {stage_name} artifact is missing", - ) if ( metadata.get("artifact_type") != stage_name or metadata.get("version") != version @@ -770,25 +810,12 @@ def _optional_active_source( ServiceErrorCode.CONFIG_INVALID, f"Active {stage_name} artifact metadata is inconsistent", ) - etag = _file_etag(primary_path) - file_hashes = metadata.get("file_hashes") - primary_key = ( - "taxonomy" if stage_name == "systematize" else "test_set" - ) - expected_hash = ( - file_hashes.get(primary_key) - if isinstance(file_hashes, dict) - else None + self._verify_versioned_source( + artifact_dir, + stage_name=stage_name, + metadata=metadata, ) - if ( - isinstance(expected_hash, str) - and re.fullmatch(r"[0-9a-f]{64}", expected_hash) - and etag != f"sha256:{expected_hash}" - ): - raise ServiceError( - ServiceErrorCode.CONFIG_INVALID, - f"Active {stage_name} artifact failed its integrity check", - ) + etag = _file_etag(primary_path) return _ArtifactSource( stage_name=stage_name, primary_path=primary_path, @@ -815,6 +842,49 @@ def _optional_active_source( etag=_file_etag(primary_path), ) + def _verify_versioned_source( + self, + artifact_dir: Path, + *, + stage_name: str, + metadata: Mapping[str, Any], + ) -> None: + files = metadata.get("files") + file_hashes = metadata.get("file_hashes") + expected_files = _STAGE_OUTPUT_FILES[stage_name] + if not isinstance(files, dict) or not isinstance(file_hashes, dict): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active {stage_name} artifact has incomplete integrity metadata", + ) + for output_key, filename in expected_files.items(): + expected_hash = file_hashes.get(output_key) + if ( + files.get(output_key) != filename + or not isinstance(expected_hash, str) + or not re.fullmatch(r"[0-9a-f]{64}", expected_hash) + ): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active {stage_name} artifact has invalid integrity metadata", + ) + path = self.workspace.path_policy.resolve_managed_output( + artifact_dir / filename, + field_name=f"active {stage_name} artifact file", + expected_root=artifact_dir, + reject_links=True, + ) + if not path.is_file(): + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Active {stage_name} artifact is missing: {filename}", + ) + if _file_etag(path) != f"sha256:{expected_hash}": + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active {stage_name} artifact failed its integrity check", + ) + def _load_taxonomy(self, path: Path) -> TaxonomyDocument: raw = self._load_json_object( path, @@ -913,7 +983,45 @@ def _copy_secondary( ServiceErrorCode.NOT_FOUND, f"Required companion artifact is missing: {name}", ) - self._copy_text(path, destination, max_bytes=_MAX_TEST_SET_BYTES) + self._copy_text( + path, + destination, + max_bytes=_MAX_TEST_SET_BYTES, + expected_hash=self._source_file_hash(source, name), + ) + + @staticmethod + def _source_file_hash( + source: _ArtifactSource, + filename: str, + ) -> str | None: + if source.metadata is None: + return None + files = source.metadata.get("files") + file_hashes = source.metadata.get("file_hashes") + if not isinstance(files, dict) or not isinstance(file_hashes, dict): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active {source.stage_name} artifact has incomplete integrity metadata", + ) + matching_keys = [ + key for key, value in files.items() if value == filename + ] + if len(matching_keys) != 1: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active {source.stage_name} artifact has invalid file metadata", + ) + expected_hash = file_hashes.get(matching_keys[0]) + if ( + not isinstance(expected_hash, str) + or not re.fullmatch(r"[0-9a-f]{64}", expected_hash) + ): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active {source.stage_name} artifact has invalid integrity metadata", + ) + return expected_hash @staticmethod def _copy_text( @@ -921,15 +1029,32 @@ def _copy_text( destination: Path, *, max_bytes: int, + expected_hash: str | None = None, ) -> None: - if source.stat().st_size > max_bytes: + try: + with source.open("rb") as handle: + value = handle.read(max_bytes + 1) + except OSError as exc: + raise ServiceError( + ServiceErrorCode.NOT_FOUND, + f"Artifact is unavailable: {source.name}", + ) from exc + if len(value) > max_bytes: raise ServiceError( ServiceErrorCode.ARTIFACT_TOO_LARGE, f"Artifact exceeds the {max_bytes}-byte curation limit", ) + if ( + expected_hash is not None + and hashlib.sha256(value).hexdigest() != expected_hash + ): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Active artifact failed its integrity check: {source.name}", + ) try: - text = source.read_bytes().decode("utf-8") - except (OSError, UnicodeDecodeError) as exc: + text = value.decode("utf-8") + except UnicodeDecodeError as exc: raise ServiceError( ServiceErrorCode.CONFIG_INVALID, f"Artifact is not valid UTF-8: {source.name}", @@ -1061,6 +1186,17 @@ def _validate_test_case( ServiceErrorCode.CONFIG_INVALID, f"Test case {test_case_id} type must be prompt or scenario", ) + try: + prepare_test_cases( + (row,), + per_test_case_tools=None, + fixed_system_prompt=None, + ) + except (TypeError, ValueError) as exc: + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + f"Test case {test_case_id} is invalid: {exc}", + ) from exc behavior = row_behavior(dict(row)) if taxonomy_names and behavior and behavior not in taxonomy_names: raise ServiceError( diff --git a/assert_ai/services/evaluations.py b/assert_ai/services/evaluations.py index 96a3f3bd7..ddb4a8d8c 100644 --- a/assert_ai/services/evaluations.py +++ b/assert_ai/services/evaluations.py @@ -31,6 +31,7 @@ import yaml from assert_ai.config import parse_model_config +from assert_ai.core.config_model import DEFAULT_INFERENCE_CONCURRENCY from assert_ai.core.io import write_bytes_atomic, write_json, write_text_atomic from assert_ai.core.jsonl_index import JsonlIndexError, scan_jsonl from assert_ai.core.config_document import PIPELINE_STAGE_ORDER @@ -42,6 +43,7 @@ ) from assert_ai.core.workspace import WorkspaceService from assert_ai.core.yaml_io import dump_yaml +from assert_ai.services.artifact_pins import load_artifact_pin from assert_ai.services.configs import ConfigRecord, ConfigService from assert_ai.services.errors import ServiceError, ServiceErrorCode from assert_ai.services.job_models import ( @@ -57,6 +59,11 @@ TraceJudgingPreflight, ) from assert_ai.services.job_store import JobStore +from assert_ai.services.output_identity import ( + run_resource_key, + suite_resource_key, + validate_output_id, +) from assert_ai.services.run_planning import ( EvaluationOverrides, RunPlanningService, @@ -79,7 +86,6 @@ _MAX_LOG_BYTES = 16 * 1024 * 1024 _DEFAULT_MAX_TRACE_INPUT_BYTES = 64 * 1024 * 1024 _GROUP_BY_RE = re.compile(r"^[A-Za-z0-9_.-]{1,128}$") -_OUTPUT_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") _SUPPORTED_JOB_KINDS = frozenset({"evaluation", "trace_judging"}) log = logging.getLogger(__name__) @@ -103,6 +109,7 @@ class _TracePlan: session_count: int estimated_judge_calls: int judge_model: str + concurrency: int taxonomy_path: Path taxonomy_ref: str taxonomy_bytes: bytes @@ -237,10 +244,11 @@ def reconcile(self, record: JobRecord) -> JobRecord: """Adopt a worker result or mark a dead worker interrupted.""" if record.state in TERMINAL_JOB_STATES: return record - if not self.launch_enabled or record.kind not in self.job_kinds: + if not self.launch_enabled or record.kind not in _SUPPORTED_JOB_KINDS: return record if record.state is JobState.QUEUED: - self.enqueue() + if record.kind in self.job_kinds: + self.enqueue() return record if ( record.state is JobState.STARTING @@ -396,9 +404,7 @@ def _schedule(self) -> None: self.enqueue() def _sweep_cancelling_jobs(self) -> None: - for record in self.store.list_nonterminal_records( - job_kinds=self.job_kinds, - ): + for record in self.store.list_nonterminal_records(): if record.state is not JobState.CANCELLING: continue try: @@ -546,15 +552,15 @@ def _recover_startup(self) -> None: next_lease_check: float | None = None has_queued_job = False try: - records = self.store.list_nonterminal_records( - job_kinds=self.job_kinds, - ) + records = self.store.list_nonterminal_records() except Exception: # noqa: BLE001 - daemon boundary log.exception("Could not scan evaluation jobs for recovery") return for record in records: if record.state is JobState.QUEUED: - has_queued_job = True + has_queued_job = ( + has_queued_job or record.kind in self.job_kinds + ) continue try: current = self.reconcile(record) @@ -989,13 +995,12 @@ def _adopt_result( lease_owner: str | None, ) -> JobRecord: job_dir = self._job_dir(record.job_id) - request = _read_json_file( + request = _read_bound_request( + record, self._job_file(job_dir, "request.json"), - max_bytes=_JOB_RESULT_MAX_BYTES, ) if ( - not isinstance(request, dict) - or payload.get("result_token") != request.get("result_token") + payload.get("result_token") != request.get("result_token") ): raise ServiceError( ServiceErrorCode.RUN_FAILED, @@ -1262,20 +1267,22 @@ def start( ) _require_source_etag(initial.source_etag, config.etag) _require_ready(initial) - suite_id = ( + suite_id = validate_output_id( applied.suite or config.document.get("suite") - or _new_identity("mcp-suite") + or _new_identity("mcp-suite"), + field_name="suite_id", ) has_run_stage = any( stage.scope == "run" and stage.action is not StageAction.DISABLED for stage in initial.stages ) - run_id = ( + run_id = _optional_output_id( applied.run or config.document.get("run") - or (_new_identity("run") if has_run_stage else None) + or (_new_identity("run") if has_run_stage else None), + field_name="run_id", ) effective_overrides = applied.model_copy( update={"suite": suite_id, "run": run_id}, @@ -1738,17 +1745,7 @@ def _trace_plan( ServiceErrorCode.CONFIG_INVALID, "Trace judging requires an enabled pipeline.judge stage", ) - raw_model = judge.get("model") or document.get("default_model") - try: - judge_model = parse_model_config( - raw_model, - field_name="pipeline.judge.model", - ).name - except (TypeError, ValueError) as exc: - raise ServiceError( - ServiceErrorCode.CONFIG_INVALID, - "pipeline.judge.model or default_model is required", - ) from exc + judge_model = _trace_judge_model(document) allowed_patterns = self.planning.policy.allowed_model_patterns if allowed_patterns and not any( fnmatchcase(judge_model, pattern) @@ -1759,6 +1756,10 @@ def _trace_plan( "The judge model is not allowed by server policy", details={"model": judge_model}, ) + concurrency = _trace_concurrency( + document, + maximum=self.planning.policy.max_concurrency, + ) taxonomy_path = self._trace_taxonomy_path( inputs.config, @@ -1829,6 +1830,7 @@ def _trace_plan( session_count=len(parsed_rows), estimated_judge_calls=len(parsed_rows) * judge_n, judge_model=judge_model, + concurrency=concurrency, taxonomy_path=taxonomy_path, taxonomy_ref=self.workspace.reference(taxonomy_path), taxonomy_bytes=taxonomy_bytes, @@ -1928,6 +1930,7 @@ def _prepare_trace_job( taxonomy_etag=plan.taxonomy_etag, group_by=plan.group_by, session_count=plan.session_count, + concurrency=plan.concurrency, suite_id=plan.suite_id, run_id=plan.run_id, request_id=request_id, @@ -1948,6 +1951,7 @@ def _prepare_trace_job_values( taxonomy_etag: str, group_by: str, session_count: int, + concurrency: int, suite_id: str, run_id: str, request_id: str, @@ -2001,7 +2005,7 @@ def _prepare_trace_job_values( write_text_atomic(snapshot, yaml_text) write_bytes_atomic(trace_snapshot, trace_bytes) write_bytes_atomic(taxonomy_snapshot, taxonomy_bytes) - write_json( + request_sha256 = _write_request_snapshot( request_path, { "schema_version": 1, @@ -2020,6 +2024,7 @@ def _prepare_trace_job_values( "taxonomy_sha256": taxonomy_etag, "group_by": group_by, "session_count": session_count, + "concurrency": concurrency, "retry_of": retry_of, }, ) @@ -2028,13 +2033,14 @@ def _prepare_trace_job_values( job_id=job_id, idempotency_key=request_id, request_hash=request_hash, + request_sha256=request_sha256, suite_id=suite_id, run_id=run_id, config_ref=config_ref, config_sha256=config_sha256, snapshot_path=str(snapshot), request_path=str(request_path), - resource_keys=(f"run:{suite_id}/{run_id}",), + resource_keys=(run_resource_key(suite_id, run_id),), retry_of=retry_of, kind="trace_judging", ), @@ -2070,17 +2076,69 @@ def _retry_trace_judging( max_bytes=_JOB_SNAPSHOT_MAX_BYTES, label="immutable trace taxonomy", ) + _validate_taxonomy_bytes(taxonomy_bytes) group_by = _validate_group_by(request.get("group_by")) + try: + trace_document = json.loads(trace_bytes.decode("utf-8")) + if not isinstance(trace_document, dict): + raise ValueError("OTLP payload must be an object") + parsed_rows = parse_otel_trace_document( + trace_document, + group_by=group_by, + ) + except ( + UnicodeDecodeError, + json.JSONDecodeError, + TypeError, + ValueError, + ) as exc: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable OTLP trace input is invalid", + ) from exc session_count = request.get("session_count") if ( isinstance(session_count, bool) or not isinstance(session_count, int) or session_count < 1 + or len(parsed_rows) != session_count ): raise ServiceError( ServiceErrorCode.JOB_INTERRUPTED, "The immutable trace job session count is invalid", ) + maximum_sessions = self.planning.policy.max_prompt_sample_size + if session_count > maximum_sessions: + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + f"Trace input contains {session_count} sessions, exceeding " + f"the current server limit of {maximum_sessions}", + ) + judge_model = _trace_judge_model(document) + allowed_patterns = self.planning.policy.allowed_model_patterns + if allowed_patterns and not any( + fnmatchcase(judge_model, pattern) + for pattern in allowed_patterns + ): + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + "The judge model is not allowed by current server policy", + details={"model": judge_model}, + ) + stored_concurrency = request.get("concurrency") + if ( + isinstance(stored_concurrency, bool) + or not isinstance(stored_concurrency, int) + or stored_concurrency < 1 + ): + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable trace job concurrency is invalid", + ) + concurrency = min( + stored_concurrency, + self.planning.policy.max_concurrency, + ) trace_ref = request.get("trace_ref") taxonomy_ref = request.get("taxonomy_ref") if not isinstance(trace_ref, str) or not isinstance(taxonomy_ref, str): @@ -2101,6 +2159,7 @@ def _retry_trace_judging( taxonomy_etag=str(request["taxonomy_sha256"]), group_by=group_by, session_count=session_count, + concurrency=concurrency, suite_id=original.suite_id, run_id=run_id, request_id=request_id, @@ -2157,10 +2216,7 @@ def _retry_snapshot( ServiceErrorCode.JOB_INTERRUPTED, "The immutable evaluation snapshot is invalid", ) from exc - request = _read_json_file( - request_path, - max_bytes=_JOB_RESULT_MAX_BYTES, - ) + request = _read_bound_request(record, request_path) if not isinstance(document, dict) or not isinstance(request, dict): raise ServiceError( ServiceErrorCode.JOB_INTERRUPTED, @@ -2168,6 +2224,8 @@ def _retry_snapshot( ) if ( request.get("job_id") != record.job_id + or request.get("kind", "evaluation") != record.kind + or request.get("config_ref") != record.config_ref or request.get("config_sha256") != record.config_sha256 or not isinstance(request.get("strict"), bool) ): @@ -2260,7 +2318,12 @@ def _prepare_job( force_stages = [ stage.name for stage in plan.stages if stage.forced ] - write_json( + expected_artifacts = _artifact_pins( + self.workspace, + suite_id=suite_id, + plan=plan, + ) + request_sha256 = _write_request_snapshot( request_path, { "schema_version": 1, @@ -2270,6 +2333,7 @@ def _prepare_job( "config_sha256": config_sha256, "strict": bool(plan.strict), "force_stages": force_stages, + "expected_artifacts": expected_artifacts, "max_log_bytes": self.manager.max_log_bytes, "retry_of": retry_of, }, @@ -2282,14 +2346,17 @@ def _prepare_job( and stage.action is not StageAction.DISABLED for stage in plan.stages ): - resource_keys.append(f"suite:{suite_id}") + resource_keys.append(suite_resource_key(suite_id)) + if expected_artifacts and suite_resource_key(suite_id) not in resource_keys: + resource_keys.append(suite_resource_key(suite_id)) if run_id is not None: - resource_keys.append(f"run:{suite_id}/{run_id}") + resource_keys.append(run_resource_key(suite_id, run_id)) return ( NewJob( job_id=job_id, idempotency_key=request_id, request_hash=request_hash, + request_sha256=request_sha256, suite_id=suite_id, run_id=run_id, config_ref=config_ref, @@ -2673,12 +2740,53 @@ def _new_identity(prefix: str) -> str: def _optional_output_id(value: Any, *, field_name: str) -> str | None: if value is None: return None - if not isinstance(value, str) or not _OUTPUT_ID_RE.fullmatch(value): + return validate_output_id(value, field_name=field_name) + + +def _trace_judge_model(document: dict[str, Any]) -> str: + pipeline = document.get("pipeline") + judge = pipeline.get("judge") if isinstance(pipeline, dict) else None + raw_model = ( + judge.get("model") if isinstance(judge, dict) else None + ) or document.get("default_model") + try: + return parse_model_config( + raw_model, + field_name="pipeline.judge.model", + ).name + except (TypeError, ValueError) as exc: raise ServiceError( - ServiceErrorCode.INVALID_ARGUMENT, - f"{field_name} must contain only letters, numbers, '.', '_', or '-'", + ServiceErrorCode.CONFIG_INVALID, + "pipeline.judge.model or default_model is required", + ) from exc + + +def _trace_concurrency( + document: dict[str, Any], + *, + maximum: int, +) -> int: + pipeline = document.get("pipeline") + inference = ( + pipeline.get("inference") + if isinstance(pipeline, dict) + else None + ) + configured = ( + inference.get("concurrency", DEFAULT_INFERENCE_CONCURRENCY) + if isinstance(inference, dict) + else DEFAULT_INFERENCE_CONCURRENCY + ) + if ( + isinstance(configured, bool) + or not isinstance(configured, int) + or configured < 1 + ): + raise ServiceError( + ServiceErrorCode.CONFIG_INVALID, + "pipeline.inference.concurrency must be a positive integer", ) - return value + return min(configured, maximum) def _validate_group_by(value: Any) -> str: @@ -2778,6 +2886,68 @@ def _read_integrity_snapshot( return value +def _read_bound_request( + record: JobRecord, + path: Path, +) -> dict[str, Any]: + value = _read_integrity_snapshot( + path, + expected_etag=record.request_sha256, + max_bytes=_JOB_RESULT_MAX_BYTES, + label="immutable evaluation job request", + ) + try: + request = json.loads(value.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable evaluation job request is invalid", + ) from exc + if not isinstance(request, dict): + raise ServiceError( + ServiceErrorCode.JOB_INTERRUPTED, + "The immutable evaluation job request must contain an object", + ) + return request + + +def _artifact_pins( + workspace: WorkspaceService, + *, + suite_id: str, + plan: Any, +) -> dict[str, dict[str, Any]]: + pins: dict[str, dict[str, Any]] = {} + for stage_name, expected in plan.consumed_artifacts.items(): + current = load_artifact_pin( + workspace, + suite_id=suite_id, + stage_name=stage_name, + version=expected.version, + ) + if current != expected: + raise ServiceError( + ServiceErrorCode.PREFLIGHT_FAILED, + f"The preflight-selected {stage_name} artifact changed " + "before the job was registered", + ) + pins[stage_name] = expected.model_dump(mode="json") + return pins + + +def _write_request_snapshot( + path: Path, + payload: dict[str, Any], +) -> str: + value = json.dumps( + payload, + ensure_ascii=False, + indent=2, + ).encode("utf-8") + write_bytes_atomic(path, value) + return _sha256_etag(value) + + def _validate_taxonomy_bytes(value: bytes) -> None: try: taxonomy = json.loads(value.decode("utf-8")) diff --git a/assert_ai/services/job_models.py b/assert_ai/services/job_models.py index 7aef294f6..71261d8c8 100644 --- a/assert_ai/services/job_models.py +++ b/assert_ai/services/job_models.py @@ -47,6 +47,7 @@ class JobRecord: job_id: str idempotency_key: str request_hash: str + request_sha256: str | None kind: str retry_of: str | None state: JobState @@ -81,6 +82,7 @@ class NewJob: job_id: str idempotency_key: str request_hash: str + request_sha256: str suite_id: str run_id: str | None config_ref: str diff --git a/assert_ai/services/job_store.py b/assert_ai/services/job_store.py index 501d89336..75f2421ff 100644 --- a/assert_ai/services/job_store.py +++ b/assert_ai/services/job_store.py @@ -25,7 +25,7 @@ ) _BUSY_TIMEOUT_MS = 5_000 -_JOB_STORE_SCHEMA_VERSION = 3 +_JOB_STORE_SCHEMA_VERSION = 4 _ACTIVE_STATES = ( JobState.STARTING.value, JobState.RUNNING.value, @@ -38,6 +38,7 @@ job_id TEXT PRIMARY KEY, idempotency_key TEXT NOT NULL UNIQUE, request_hash TEXT NOT NULL, + request_sha256 TEXT NOT NULL, kind TEXT NOT NULL, retry_of TEXT, state TEXT NOT NULL, @@ -158,21 +159,39 @@ def create_or_get( ) created_at = _now() + if new_job.run_id is not None: + collision = connection.execute( + """ + SELECT job_id FROM jobs + WHERE suite_id = ? COLLATE NOCASE + AND run_id = ? COLLATE NOCASE + LIMIT 1 + """, + (new_job.suite_id, new_job.run_id), + ).fetchone() + if collision is not None: + raise ServiceError( + ServiceErrorCode.CONFLICT, + "The requested suite/run output is already assigned", + details={"job_id": str(collision["job_id"])}, + ) try: connection.execute( """ INSERT INTO jobs( - job_id, idempotency_key, request_hash, kind, retry_of, + job_id, idempotency_key, request_hash, request_sha256, + kind, retry_of, state, created_at, suite_id, run_id, config_ref, config_sha256, snapshot_path, request_path, resource_keys_json - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( new_job.job_id, new_job.idempotency_key, new_job.request_hash, + new_job.request_sha256, new_job.kind, new_job.retry_of, JobState.QUEUED.value, @@ -183,12 +202,16 @@ def create_or_get( new_job.config_sha256, new_job.snapshot_path, new_job.request_path, - _json(new_job.resource_keys), + _json(_canonical_resource_keys(new_job.resource_keys)), ), ) except sqlite3.IntegrityError as exc: collision = connection.execute( - "SELECT job_id FROM jobs WHERE suite_id = ? AND run_id = ?", + """ + SELECT job_id FROM jobs + WHERE suite_id = ? COLLATE NOCASE + AND run_id = ? COLLATE NOCASE + """, (new_job.suite_id, new_job.run_id), ).fetchone() if collision is not None: @@ -388,7 +411,7 @@ def acquire_operation_locks( lease_seconds: float, ) -> bool: """Reserve resources against job claims for one short operation.""" - keys = tuple(dict.fromkeys(resource_keys)) + keys = _canonical_resource_keys(resource_keys) if not keys: raise ValueError("at least one resource key is required") if not owner: @@ -414,7 +437,7 @@ def acquire_operation_locks( active_operation = connection.execute( f""" SELECT 1 FROM operation_locks - WHERE resource_key IN ({placeholders}) + WHERE resource_key COLLATE NOCASE IN ({placeholders}) LIMIT 1 """, keys, @@ -444,14 +467,15 @@ def release_operation_locks( if not owner: raise ValueError("owner is required") self.initialize() - keys = tuple(dict.fromkeys(resource_keys)) + keys = _canonical_resource_keys(resource_keys) with self._transaction() as connection: if keys: placeholders = ", ".join("?" for _ in keys) connection.execute( f""" DELETE FROM operation_locks - WHERE owner = ? AND resource_key IN ({placeholders}) + WHERE owner = ? + AND resource_key COLLATE NOCASE IN ({placeholders}) """, (owner, *keys), ) @@ -469,7 +493,7 @@ def renew_operation_locks( lease_seconds: float, ) -> bool: """Extend unexpired operation locks when every key is still owned.""" - keys = tuple(dict.fromkeys(resource_keys)) + keys = _canonical_resource_keys(resource_keys) if not keys: raise ValueError("at least one resource key is required") if not owner: @@ -487,7 +511,7 @@ def renew_operation_locks( FROM operation_locks WHERE owner = ? AND lease_expires_at > ? - AND resource_key IN ({placeholders}) + AND resource_key COLLATE NOCASE IN ({placeholders}) """, (owner, now, *keys), ).fetchone() @@ -499,7 +523,7 @@ def renew_operation_locks( SET lease_expires_at = ? WHERE owner = ? AND lease_expires_at > ? - AND resource_key IN ({placeholders}) + AND resource_key COLLATE NOCASE IN ({placeholders}) """, (expires_at, owner, now, *keys), ).rowcount @@ -934,6 +958,7 @@ def initialize(self) -> None: 0, 1, 2, + 3, _JOB_STORE_SCHEMA_VERSION, }: raise ServiceError( @@ -955,6 +980,15 @@ def initialize(self) -> None: except sqlite3.OperationalError as exc: if "duplicate column" not in str(exc).lower(): raise + if "request_sha256" not in columns: + try: + connection.execute( + "ALTER TABLE jobs " + "ADD COLUMN request_sha256 TEXT" + ) + except sqlite3.OperationalError as exc: + if "duplicate column" not in str(exc).lower(): + raise connection.execute( "PRAGMA user_version = " f"{_JOB_STORE_SCHEMA_VERSION}" @@ -1018,7 +1052,7 @@ def _resources_available( row = connection.execute( f""" SELECT 1 FROM resource_locks - WHERE resource_key IN ({placeholders}) + WHERE resource_key COLLATE NOCASE IN ({placeholders}) LIMIT 1 """, resource_keys, @@ -1044,7 +1078,7 @@ def _resources_available( operation = connection.execute( f""" SELECT 1 FROM operation_locks - WHERE resource_key IN ({operation_placeholders}) + WHERE resource_key COLLATE NOCASE IN ({operation_placeholders}) LIMIT 1 """, operation_keys, @@ -1058,22 +1092,25 @@ def _operation_conflicts_with_active_job( ) -> bool: if resource_key.startswith("suite:"): suite_id = resource_key.removeprefix("suite:") + run_prefix = f"run:{suite_id}/" row = connection.execute( """ SELECT 1 FROM resource_locks - WHERE resource_key = ? OR resource_key LIKE ? + WHERE resource_key = ? COLLATE NOCASE + OR substr(resource_key, 1, length(?)) = ? COLLATE NOCASE LIMIT 1 """, ( resource_key, - f"run:{suite_id}/%", + run_prefix, + run_prefix, ), ).fetchone() else: row = connection.execute( """ SELECT 1 FROM resource_locks - WHERE resource_key = ? + WHERE resource_key = ? COLLATE NOCASE LIMIT 1 """, (resource_key,), @@ -1180,11 +1217,16 @@ def _suite_id_from_resource_key(resource_key: str) -> str | None: return None +def _canonical_resource_keys(resource_keys: Sequence[str]) -> tuple[str, ...]: + return tuple(dict.fromkeys(key.casefold() for key in resource_keys)) + + def _record(row: sqlite3.Row) -> JobRecord: return JobRecord( job_id=str(row["job_id"]), idempotency_key=str(row["idempotency_key"]), request_hash=str(row["request_hash"]), + request_sha256=_optional_str(row["request_sha256"]), kind=str(row["kind"]), retry_of=_optional_str(row["retry_of"]), state=JobState(str(row["state"])), diff --git a/assert_ai/services/output_identity.py b/assert_ai/services/output_identity.py new file mode 100644 index 000000000..30280d76f --- /dev/null +++ b/assert_ai/services/output_identity.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Portable validation and locking identities for managed suite/run paths.""" + +from __future__ import annotations + +import re +from typing import Any + +from assert_ai.services.errors import ServiceError, ServiceErrorCode + +_OUTPUT_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_WINDOWS_RESERVED_NAMES = frozenset( + { + "CON", + "PRN", + "AUX", + "NUL", + *(f"COM{index}" for index in range(1, 10)), + *(f"LPT{index}" for index in range(1, 10)), + } +) + + +def validate_output_id(value: Any, *, field_name: str) -> str: + """Return a portable path component or raise a stable service error.""" + if not isinstance(value, str) or not _OUTPUT_ID_RE.fullmatch(value): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{field_name} must contain only letters, numbers, '.', '_', or '-'", + ) + if value.endswith("."): + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{field_name} must not end with a period", + ) + device_stem = value.split(".", 1)[0].upper() + if device_stem in _WINDOWS_RESERVED_NAMES: + raise ServiceError( + ServiceErrorCode.INVALID_ARGUMENT, + f"{field_name} uses a reserved Windows device name", + ) + return value + + +def suite_resource_key(suite_id: str) -> str: + """Return a case-insensitive, cross-platform suite lock identity.""" + return f"suite:{suite_id.casefold()}" + + +def run_resource_key(suite_id: str, run_id: str) -> str: + """Return a case-insensitive, cross-platform run lock identity.""" + return f"run:{suite_id.casefold()}/{run_id.casefold()}" diff --git a/assert_ai/services/run_planning.py b/assert_ai/services/run_planning.py index 37d87e03a..02b47155e 100644 --- a/assert_ai/services/run_planning.py +++ b/assert_ai/services/run_planning.py @@ -36,6 +36,7 @@ validate_module_ref, ) from assert_ai.core.workspace import WorkspaceService +from assert_ai.services.artifact_pins import ArtifactPin, load_artifact_pin from assert_ai.services.configs import ConfigService from assert_ai.services.errors import ServiceError, ServiceErrorCode from assert_ai.stages import STAGES @@ -161,6 +162,7 @@ class EvaluationPreflight(_ServiceModel): sample_sizes: dict[str, int | None] = Field(default_factory=dict) target: TargetPreflight | None = None stages: tuple[StagePreflight, ...] = () + consumed_artifacts: dict[str, ArtifactPin] = Field(default_factory=dict) models: tuple[ModelUse, ...] = () credentials: tuple[CredentialRequirement, ...] = () managed_outputs: dict[str, str] = Field(default_factory=dict) @@ -290,6 +292,21 @@ def preflight_document( blocking.extend(credential_issues) warnings.extend(credential_warnings) stages = _stage_plan(ctx, forced, models) + try: + consumed_artifacts = _consumed_artifact_pins( + self.workspace, + suite_id=str(ctx["suite_id"]), + stages=stages, + artifact_versions=ctx.get("artifact_versions"), + ) + except ServiceError as exc: + blocking.append( + PreflightIssue( + code=exc.code.value, + message=str(exc), + ) + ) + consumed_artifacts = {} managed_outputs = { "artifacts_root": self.workspace.reference(ctx["artifacts_root"]), "results_root": self.workspace.reference(ctx["results_dir"]), @@ -320,6 +337,7 @@ def preflight_document( sample_sizes=sample_sizes, target=target, stages=tuple(stages), + consumed_artifacts=consumed_artifacts, models=tuple(models), credentials=tuple(credentials), managed_outputs=managed_outputs, @@ -329,6 +347,53 @@ def preflight_document( ) +def _consumed_artifact_pins( + workspace: WorkspaceService, + *, + suite_id: str, + stages: list[StagePreflight], + artifact_versions: Any, +) -> dict[str, ArtifactPin]: + actions = {stage.name: stage.action for stage in stages} + enabled = { + stage.name + for stage in stages + if stage.action is not StageAction.DISABLED + } + required = { + stage.name + for stage in stages + if stage.action is StageAction.REUSE + } + if {"test_set", "judge"} & enabled and actions.get("systematize") in { + None, + StageAction.DISABLED, + }: + required.add("systematize") + if "inference" in enabled and actions.get("test_set") in { + None, + StageAction.DISABLED, + }: + required.add("test_set") + + refs = artifact_versions if isinstance(artifact_versions, dict) else {} + pins: dict[str, ArtifactPin] = {} + for stage_name in ("systematize", "test_set"): + if stage_name not in required: + continue + ref = refs.get(stage_name) + version = ref.get("version") if isinstance(ref, dict) else None + if not isinstance(version, str) or not version: + continue + pins[stage_name] = load_artifact_pin( + workspace, + suite_id=suite_id, + stage_name=stage_name, + version=version, + ) + return pins + + def _apply_overrides( document: dict[str, Any], overrides: EvaluationOverrides, diff --git a/assert_ai/stages/inference.py b/assert_ai/stages/inference.py index 6c58db910..e6c5eea93 100644 --- a/assert_ai/stages/inference.py +++ b/assert_ai/stages/inference.py @@ -33,7 +33,6 @@ from assert_ai.core.io import ( INFERENCE_SET_FILE, append_jsonl_row, - get_permissible_flag, load_jsonl, load_prompt_text, load_test_cases, @@ -42,6 +41,7 @@ write_jsonl, row_factors, ) +from assert_ai.core.test_cases import prepare_test_cases from assert_ai.core.model_client import GenerateOptions, Message, ModelResponse, build_llm_call_trace, generate, to_jsonable from assert_ai.core.model_client import LLMAuthError, LLMContentFilterError, LLMInputError, LLMRateLimitError, LLMProviderError from assert_ai.core.run_control import RunCancelled, RunControl @@ -55,7 +55,7 @@ serialize_response, ) from assert_ai.core.tool_backend import ToolBackendResolver, inspect_tool_module -from assert_ai.core.tools import load_toolset_file, normalize_tool_defs +from assert_ai.core.tools import load_toolset_file from assert_ai.core.transcript import ( AddMessageEdit, Message as TranscriptMessage, @@ -439,56 +439,11 @@ def _prepare_test_cases( tool_source: str, fixed_system_prompt: str | None, ) -> list[dict[str, Any]]: - """Validate canonical test-case rows and normalize prompt/scenario-specific fields.""" - test_set: list[dict[str, Any]] = [] - nested_test_case_fields = {"prompt", "description", "system_prompt", "title", "tools", "state"} - for index, row in enumerate(rows): - if not isinstance(row, dict): - raise ValueError(f"test case at index {index} must be an object") - - kind = row.get("type") - if kind not in {"prompt", "scenario"}: - raise ValueError(f"test case at index {index} must declare type 'prompt' or 'scenario'") - - test_case_payload = row.get("seed") - if not isinstance(test_case_payload, dict): - raise ValueError(f"{kind} test case at index {index} requires a test case payload object") - test_case_row = dict(row) - normalized_payload = dict(test_case_payload) - system_prompt = str(normalized_payload.get("system_prompt") or "").strip() or None - if system_prompt is None: - normalized_payload.pop("system_prompt", None) - else: - normalized_payload["system_prompt"] = system_prompt - if fixed_system_prompt and system_prompt is not None: - raise ValueError("target.system_prompt cannot be combined with non-empty test case system_prompt") - tools = normalized_payload.get("tools") - if tool_source == TOOL_SOURCE_PER_TEST_CASE: - if not isinstance(tools, list) or not tools: - raise ValueError("test case tools are required when tool_source=per_test_case") - normalize_tool_defs(tools) - elif tools is not None: - raise ValueError("test case tools are only allowed when tool_source=per_test_case") - test_case_row["seed"] = normalized_payload - if kind == "prompt": - invalid_fields = sorted(field for field in nested_test_case_fields if field in row) - if invalid_fields: - raise ValueError( - f"prompt test case at index {index} must move {', '.join(invalid_fields)} under the test case payload" - ) - if not str(normalized_payload.get("description") or "").strip(): - raise ValueError( - f"prompt test case at index {index} requires a non-empty test case description" - ) - elif not str(normalized_payload.get("description") or "").strip(): - raise ValueError( - f"scenario test case at index {index} requires a non-empty test case description" - ) - permissible = get_permissible_flag(test_case_row) - if permissible is not None: - test_case_row["permissible"] = permissible - test_set.append(test_case_row) - return test_set + return prepare_test_cases( + rows, + per_test_case_tools=(tool_source == TOOL_SOURCE_PER_TEST_CASE), + fixed_system_prompt=fixed_system_prompt, + ) def _build_hosted_session( diff --git a/assert_ai/stages/judge.py b/assert_ai/stages/judge.py index 7277e5715..d598e44f4 100644 --- a/assert_ai/stages/judge.py +++ b/assert_ai/stages/judge.py @@ -371,6 +371,9 @@ async def worker(item: tuple[int, dict[str, Any]]) -> dict[str, Any]: score_row_filter_skipped["dimension_scales"] = judge_contract["dimension_scales"] if dimensions: score_row_filter_skipped["dimensions"] = dimensions + trace_refs = row.get("trace_refs") + if isinstance(trace_refs, list): + score_row_filter_skipped["trace_refs"] = trace_refs return { "output_index": output_index, "score_row": score_row_filter_skipped, diff --git a/assert_ai/stages/test_set.py b/assert_ai/stages/test_set.py index 2a9efeb6e..256341ab1 100644 --- a/assert_ai/stages/test_set.py +++ b/assert_ai/stages/test_set.py @@ -256,7 +256,7 @@ def normalize_generated_test_case( raise ValueError("generated test case requires non-empty tools when test_set.tool_source=per_test_case") try: normalize_tool_defs(raw_tools) - except (KeyError, TypeError) as exc: + except (KeyError, TypeError, ValueError) as exc: raise ValueError("generated test case contains invalid tool definitions") from exc payload["tools"] = raw_tools elif raw_tools: diff --git a/tests/test_artifact_cache.py b/tests/test_artifact_cache.py index 1bc115cf0..846e97d62 100644 --- a/tests/test_artifact_cache.py +++ b/tests/test_artifact_cache.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import hashlib import logging import shutil import unittest @@ -13,6 +14,7 @@ activate_latest_artifacts, artifact_ref, discard_artifact_plan, + file_sha256, finalize_artifact_plan, hash_payload, override_cacheable_output_paths, @@ -63,6 +65,14 @@ def test_hash_payload_is_stable_across_dict_key_order(self) -> None: hash_payload({"a": 1, "b": [2, {"c": 3, "d": 4}]}), ) + def test_file_sha256_hashes_across_streaming_chunks(self) -> None: + with TemporaryDirectory() as tmp_dir: + path = Path(tmp_dir) / "large.jsonl" + payload = (b"0123456789abcdef" * 65537) + b"tail" + path.write_bytes(payload) + + self.assertEqual(file_sha256(path), hashlib.sha256(payload).hexdigest()) + def test_prepare_reuses_latest_matching_artifact(self) -> None: with TemporaryDirectory() as tmp_dir: root = Path(tmp_dir) diff --git a/tests/test_curation_service.py b/tests/test_curation_service.py index 6acfde421..c5be5a7c2 100644 --- a/tests/test_curation_service.py +++ b/tests/test_curation_service.py @@ -62,15 +62,14 @@ def _rows() -> list[dict]: { "type": "prompt", "test_case_id": "test_case_000001", - "prompt": "Book a safe flight.", + "seed": {"description": "Book a safe flight."}, "dimensions": {"behavior": "safe_booking"}, }, { "type": "scenario", "test_case_id": "test_case_000002", - "prompt": "Ignore the restriction.", + "seed": {"description": "Ignore the restriction."}, "dimensions": {"behavior": "unsafe_booking"}, - "tools": [], }, ] @@ -275,11 +274,11 @@ def test_bulk_revise_test_cases_preserves_ids_order_and_old_version( ( CaseRevision( test_case_id="test_case_000002", - updates={"prompt": "Revised unsafe request."}, + updates={"seed": {"description": "Revised unsafe request."}}, ), CaseRevision( test_case_id="test_case_000001", - updates={"prompt": "Revised safe request."}, + updates={"seed": {"description": "Revised safe request."}}, ), ), expected_etag=_etag(source).removeprefix("sha256:"), @@ -301,8 +300,8 @@ def test_bulk_revise_test_cases_preserves_ids_order_and_old_version( "test_case_000001", "test_case_000002", ] - assert revised_rows[0]["prompt"] == "Revised safe request." - assert revised_rows[1]["prompt"] == "Revised unsafe request." + assert revised_rows[0]["seed"]["description"] == "Revised safe request." + assert revised_rows[1]["seed"]["description"] == "Revised unsafe request." def test_post_activation_summary_failure_keeps_new_version( @@ -323,7 +322,7 @@ def test_post_activation_summary_failure_keeps_new_version( ).revise_test_case( "suite-a", "test_case_000001", - {"prompt": "Revised prompt."}, + {"seed": {"description": "Revised prompt."}}, expected_etag=_etag(source), change_summary="Exercise post-activation failure handling.", ) @@ -357,7 +356,7 @@ def test_post_activation_lock_release_failure_keeps_new_version( ).revise_test_case( "suite-a", "test_case_000001", - {"prompt": "Revised despite cleanup failure."}, + {"seed": {"description": "Revised despite cleanup failure."}}, expected_etag=_etag(source), change_summary="Exercise lease cleanup failure handling.", ) @@ -367,6 +366,39 @@ def test_post_activation_lock_release_failure_keeps_new_version( assert result.artifacts[0].version == "v0002" +def test_post_activation_base_exception_keeps_new_version( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + source = ( + suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + ) + + with ( + patch( + "assert_ai.services.curation.refresh_compatibility_files", + side_effect=KeyboardInterrupt, + ), + pytest.raises(KeyboardInterrupt), + ): + CurationService( + workspace, + job_store=store, + ).revise_test_case( + "suite-a", + "test_case_000001", + {"seed": {"description": "Committed before interruption."}}, + expected_etag=_etag(source), + change_summary="Interrupt post-activation cleanup.", + ) + + latest = json.loads((suite_root / "latest.json").read_text(encoding="utf-8")) + assert latest["artifacts"]["test_set"]["version"] == "v0002" + assert ( + suite_root / "artifacts" / "test_set" / "v0002" / "test_set.jsonl" + ).is_file() + + def test_curation_rejects_tampered_immutable_source(tmp_path: Path) -> None: workspace, store, suite_root = _seed_suite(tmp_path) source = ( @@ -392,6 +424,67 @@ def test_curation_rejects_tampered_immutable_source(tmp_path: Path) -> None: ).exists() +def test_curation_rejects_tampered_companion_artifact( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + source = ( + suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + ) + companion = ( + suite_root + / "artifacts" + / "test_set" + / "v0001" + / "stratification.json" + ) + companion.write_text('{"tampered":true}', encoding="utf-8") + + with pytest.raises(ServiceError) as invalid: + CurationService(workspace, job_store=store).revise_test_case( + "suite-a", + "test_case_000001", + {"seed": {"description": "Revision over corrupt source."}}, + expected_etag=_etag(source), + change_summary="Reject a corrupt companion.", + ) + + assert invalid.value.code == ServiceErrorCode.CONFIG_INVALID + with pytest.raises(ServiceError) as invalid_tools: + CurationService(workspace, job_store=store).revise_test_case( + "suite-a", + "test_case_000001", + { + "seed": { + "description": "Malformed tools.", + "tools": "not-a-list", + } + }, + expected_etag=_etag(source), + change_summary="Attempt malformed per-test-case tools.", + ) + + assert invalid_tools.value.code == ServiceErrorCode.CONFIG_INVALID + with pytest.raises(ServiceError) as missing_tool_name: + CurationService(workspace, job_store=store).revise_test_case( + "suite-a", + "test_case_000001", + { + "seed": { + "description": "Malformed tool entry.", + "tools": [{}], + } + }, + expected_etag=_etag(source), + change_summary="Attempt a tool without a name.", + ) + + assert missing_tool_name.value.code == ServiceErrorCode.CONFIG_INVALID + assert not ( + suite_root / "artifacts" / "test_set" / "v0002" + ).exists() + + def test_test_case_revision_rejects_identity_changes_and_unknown_categories( tmp_path: Path, ) -> None: @@ -423,13 +516,34 @@ def test_test_case_revision_rejects_identity_changes_and_unknown_categories( service.revise_test_case( "suite-a", "test_case_000001", - {"prompt": "Book a safe flight."}, + {"seed": {"description": "Book a safe flight."}}, expected_etag=_etag(source), change_summary="Attempt a no-op revision.", ) assert unchanged.value.code == ServiceErrorCode.INVALID_ARGUMENT +def test_test_case_revision_rejects_rows_inference_cannot_execute( + tmp_path: Path, +) -> None: + workspace, store, suite_root = _seed_suite(tmp_path) + source = suite_root / "artifacts" / "test_set" / "v0001" / "test_set.jsonl" + + with pytest.raises(ServiceError) as invalid: + CurationService(workspace, job_store=store).revise_test_case( + "suite-a", + "test_case_000001", + {"seed": None}, + expected_etag=_etag(source), + change_summary="Attempt an invalid seed payload.", + ) + + assert invalid.value.code == ServiceErrorCode.CONFIG_INVALID + assert not ( + suite_root / "artifacts" / "test_set" / "v0002" + ).exists() + + def test_curation_enforces_revised_artifact_size_limits( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -476,7 +590,7 @@ def test_curation_enforces_revised_artifact_size_limits( service.revise_test_case( "suite-a", "test_case_000001", - {"prompt": "x" * 1_000}, + {"seed": {"description": "x" * 1_000}}, expected_etag=_etag(test_set_source), change_summary="Oversized test case.", ) @@ -490,6 +604,7 @@ def test_curation_conflicts_with_an_active_suite_job(tmp_path: Path) -> None: job_id="job-active", idempotency_key="request-active", request_hash="hash-active", + request_sha256="sha256:" + ("0" * 64), suite_id="suite-a", run_id="run-active", config_ref="demo.yaml", @@ -511,7 +626,7 @@ def test_curation_conflicts_with_an_active_suite_job(tmp_path: Path) -> None: CurationService(workspace, job_store=store).revise_test_case( "suite-a", "test_case_000001", - {"prompt": "Blocked edit."}, + {"seed": {"description": "Blocked edit."}}, expected_etag=_etag(source), change_summary="This should be blocked.", ) diff --git a/tests/test_evaluation_service.py b/tests/test_evaluation_service.py index 21e0a4ee1..631e68d0d 100644 --- a/tests/test_evaluation_service.py +++ b/tests/test_evaluation_service.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import json import os import subprocess @@ -15,20 +16,29 @@ import pytest +from assert_ai.config import load_runtime_context +from assert_ai.core.artifact_cache import ( + activate_artifact_plan, + finalize_artifact_plan, + prepare_artifact_plan, +) from assert_ai.core.io import write_json +from assert_ai.core.run_result import RunResult, RunState from assert_ai.core.workspace import WorkspaceService from assert_ai.services._evaluation_worker import ( _BoundedTextLog, main as worker_main, ) +from assert_ai.services.artifact_pins import load_artifact_pin from assert_ai.services.configs import ConfigService +from assert_ai.services.curation import CurationService from assert_ai.services.errors import ServiceError, ServiceErrorCode from assert_ai.services.evaluations import ( EvaluationJobManager, EvaluationService, _active_stage_from_events, ) -from assert_ai.services.job_models import JobState +from assert_ai.services.job_models import JobState, NewJob from assert_ai.services.job_store import JobStore from assert_ai.services.results import ResultRepository from assert_ai.services.run_planning import ( @@ -36,6 +46,7 @@ PreflightPolicy, RunPlanningService, ) +from assert_ai.stages import STAGES def _service( @@ -45,6 +56,8 @@ def _service( lease_seconds: float = 60.0, max_trace_input_bytes: int = 16 * 1024 * 1024, max_prompt_sample_size: int = 100_000, + max_concurrency: int = 32, + allowed_model_patterns: tuple[str, ...] = (), ) -> tuple[ConfigService, EvaluationService]: workspace = WorkspaceService.create(root) configs = ConfigService(workspace) @@ -53,6 +66,8 @@ def _service( configs, policy=PreflightPolicy( max_prompt_sample_size=max_prompt_sample_size, + max_concurrency=max_concurrency, + allowed_model_patterns=allowed_model_patterns, ), ) store = JobStore(workspace.artifacts_root / "mcp" / "jobs.sqlite3") @@ -304,6 +319,209 @@ def test_inference_only_job_completes_and_is_idempotent( assert service.list().items[0].job_id == started.job.job_id +def test_queued_job_rejects_a_different_curated_artifact_version( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + seed_document = { + "suite": "pinned-suite", + "behavior": { + "name": "safe_help", + "description": "The agent should provide safe help.", + }, + "default_model": {"name": "openai/gpt-test"}, + "pipeline": { + "systematize": { + "model": {"name": "openai/gpt-test"}, + }, + "test_set": { + "model": {"name": "openai/gpt-test"}, + "prompt": {"sample_size": 1}, + }, + }, + } + configs.save_config("seed.yaml", document=seed_document) + record = configs.get_config("seed.yaml") + config_path = service.workspace.path_policy.resolve_config_path( + record.config_ref, + must_exist=True, + reject_links=True, + ) + context = load_runtime_context( + deepcopy(record.document), + config_path, + stage_modules=STAGES, + path_policy=service.workspace.path_policy, + ) + raw_systematize = dict( + next(raw for name, raw in context["stages"] if name == "systematize") + ) + taxonomy_artifact = prepare_artifact_plan( + ctx=context, + stage_name="systematize", + raw_cfg=raw_systematize, + forced=False, + ) + activate_artifact_plan(context, taxonomy_artifact) + taxonomy = { + "behavior": { + "name": "safe_help", + "definition": "The agent should provide safe help.", + }, + "definition_of_terms": [], + "behavior_categories": [ + { + "name": "safe", + "definition": "The response follows the requirement.", + "examples": ["Provide safe help."], + "permissible": True, + } + ], + } + write_json(taxonomy_artifact.output_paths["taxonomy"], taxonomy) + write_json( + taxonomy_artifact.output_paths["systematization"], + { + "behavior": "safe_help", + "systematization": "Fixture", + "summary_items": [], + }, + ) + finalize_artifact_plan(context, taxonomy_artifact) + raw_test_set = dict( + next(raw for name, raw in context["stages"] if name == "test_set") + ) + test_set_artifact = prepare_artifact_plan( + ctx=context, + stage_name="test_set", + raw_cfg=raw_test_set, + forced=False, + ) + activate_artifact_plan(context, test_set_artifact) + test_set_artifact.output_paths["test_set"].write_text( + json.dumps( + { + "type": "prompt", + "test_case_id": "case-1", + "seed": {"description": "Provide safe help."}, + "dimensions": {"behavior": "safe"}, + } + ) + + "\n", + encoding="utf-8", + ) + write_json(test_set_artifact.output_paths["stratification"], {"counts": {}}) + finalize_artifact_plan(context, test_set_artifact) + (tmp_path / "agent.py").write_text( + "def run(message, *, history=None):\n" + " del history\n" + " return message\n", + encoding="utf-8", + ) + configs.save_config( + "pinned.yaml", + document={ + "suite": "pinned-suite", + "pipeline": { + "inference": { + "target": {"callable": "agent:run"}, + "concurrency": 1, + }, + }, + }, + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start( + "pinned.yaml", + request_id="pinned-artifact", + ) + job_record = service.store.get(started.job.job_id) + request = json.loads( + Path(job_record.request_path).read_text(encoding="utf-8") + ) + assert request["expected_artifacts"]["test_set"]["version"] == ( + test_set_artifact.version + ) + test_set_etag = "sha256:" + hashlib.sha256( + test_set_artifact.output_paths["test_set"].read_bytes() + ).hexdigest() + CurationService( + service.workspace, + job_store=service.store, + ).revise_test_case( + "pinned-suite", + "case-1", + {"seed": {"description": "Curated after preflight."}}, + expected_etag=test_set_etag, + change_summary="Change the queued job's active test set.", + ) + + exit_code = worker_main( + [ + "--workspace", + str(tmp_path), + "--job-id", + started.job.job_id, + ] + ) + + result = json.loads( + ( + Path(job_record.request_path).parent / "result.json" + ).read_text(encoding="utf-8") + )["run_result"] + assert exit_code == 1 + assert result["state"] == "failed" + assert result["error_code"] == "PREFLIGHT_FAILED" + assert "changed while the job was queued" in result["error_message"] + + +def test_large_reusable_artifact_can_be_pinned(tmp_path: Path) -> None: + workspace = WorkspaceService.create(tmp_path) + artifact_dir = ( + workspace.results_root + / "large-suite" + / "artifacts" + / "test_set" + / "v0001" + ) + artifact_dir.mkdir(parents=True) + test_set = artifact_dir / "test_set.jsonl" + stratification = artifact_dir / "stratification.json" + test_set.write_bytes(b"x" * ((16 * 1024 * 1024) + 1)) + stratification.write_text("{}", encoding="utf-8") + test_set_hash = hashlib.sha256(test_set.read_bytes()).hexdigest() + stratification_hash = hashlib.sha256( + stratification.read_bytes() + ).hexdigest() + write_json( + artifact_dir / "artifact.json", + { + "schema_version": 1, + "artifact_type": "test_set", + "version": "v0001", + "files": { + "test_set": "test_set.jsonl", + "stratification": "stratification.json", + }, + "file_hashes": { + "test_set": test_set_hash, + "stratification": stratification_hash, + }, + }, + ) + + pin = load_artifact_pin( + workspace, + suite_id="large-suite", + stage_name="test_set", + version="v0001", + ) + + assert pin.file_hashes["test_set"] == test_set_hash + + def test_trace_job_preflight_and_no_credential_execution( tmp_path: Path, ) -> None: @@ -545,6 +763,116 @@ def test_trace_job_snapshots_inputs_and_retries_immutably( ).read_bytes() +def test_trace_retry_reapplies_current_model_policy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-policy-original", + suite_id="trace-suite", + run_id="trace-policy-original", + ) + claimed = service.store.claim_next( + lease_owner="fixture-manager", + lease_seconds=60, + max_active_jobs=1, + ) + assert claimed is not None + service.store.mark_terminal( + claimed.job_id, + state=JobState.FAILED, + exit_code=1, + failed_stage="judge", + error_code=ServiceErrorCode.RUN_FAILED.value, + error_message="Fixture failure", + result={"state": "failed", "exit_code": 1}, + run_root=None, + lease_owner="fixture-manager", + ) + service.planning = RunPlanningService( + service.workspace, + configs, + policy=PreflightPolicy( + allowed_model_patterns=("approved/*",), + ), + ) + + with pytest.raises(ServiceError) as blocked: + service.retry( + started.job.job_id, + request_id="trace-policy-retry", + ) + + assert blocked.value.code == ServiceErrorCode.PREFLIGHT_FAILED + assert "current server policy" in str(blocked.value) + + +def test_trace_retry_reapplies_current_session_limit( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + trace_path = tmp_path / "fixtures" / "traces.json" + trace_document = json.loads(trace_path.read_text(encoding="utf-8")) + spans = trace_document["resourceSpans"][0]["scopeSpans"][0]["spans"] + second = deepcopy(spans[0]) + second["spanId"] = "c" * 16 + second["attributes"][0]["value"]["stringValue"] = "session-two" + spans.append(second) + write_json(trace_path, trace_document) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-size-original", + suite_id="trace-suite", + run_id="trace-size-original", + ) + claimed = service.store.claim_next( + lease_owner="fixture-manager", + lease_seconds=60, + max_active_jobs=1, + ) + assert claimed is not None + service.store.mark_terminal( + claimed.job_id, + state=JobState.FAILED, + exit_code=1, + failed_stage="judge", + error_code=ServiceErrorCode.RUN_FAILED.value, + error_message="Fixture failure", + result={"state": "failed", "exit_code": 1}, + run_root=None, + lease_owner="fixture-manager", + ) + service.planning = RunPlanningService( + service.workspace, + configs, + policy=PreflightPolicy(max_prompt_sample_size=1), + ) + + with pytest.raises(ServiceError) as blocked: + service.retry( + started.job.job_id, + request_id="trace-size-retry", + ) + + assert blocked.value.code == ServiceErrorCode.PREFLIGHT_FAILED + assert "current server limit" in str(blocked.value) + + def test_trace_worker_cancels_during_import( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -637,6 +965,155 @@ def test_trace_worker_rejects_tampered_input_snapshot( ) +def test_trace_worker_rejects_changed_parsed_session_count( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-parser-change", + suite_id="trace-suite", + run_id="trace-parser-change", + ) + job_dir = service.manager._job_dir(started.job.job_id) + + from assert_ai.core.otel import parse_otel_trace_document + + def duplicate_sessions( + document: dict, + *, + group_by: str, + ) -> list[dict]: + rows = parse_otel_trace_document(document, group_by=group_by) + return [*rows, dict(rows[0])] + + with patch( + "assert_ai.services._evaluation_worker.parse_otel_trace_document", + side_effect=duplicate_sessions, + ): + exit_code = worker_main( + [ + "--workspace", + str(tmp_path), + "--job-id", + started.job.job_id, + ] + ) + + result = json.loads( + (job_dir / "result.json").read_text(encoding="utf-8") + ) + assert exit_code == 1 + assert result["worker_error"]["error_code"] == "JOB_INTERRUPTED" + assert "different session count" in result["worker_error"]["error_message"] + + +def test_trace_worker_rejects_preexisting_run_output( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-output-race", + suite_id="trace-suite", + run_id="trace-output-race", + ) + run_root = ( + tmp_path + / "artifacts" + / "results" + / "trace-suite" + / "trace-output-race" + ) + run_root.mkdir(parents=True) + sentinel = run_root / "sentinel.txt" + sentinel.write_text("preserve", encoding="utf-8") + + exit_code = worker_main( + [ + "--workspace", + str(tmp_path), + "--job-id", + started.job.job_id, + ] + ) + + assert exit_code == 1 + assert sentinel.read_text(encoding="utf-8") == "preserve" + assert not (run_root / "inference_set.jsonl").exists() + + +def test_trace_worker_uses_policy_capped_concurrency( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path, max_concurrency=1) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="trace-concurrency", + suite_id="trace-suite", + run_id="trace-concurrency", + ) + record = service.store.get(started.job.job_id) + request = json.loads(Path(record.request_path).read_text(encoding="utf-8")) + captured: dict[str, object] = {} + + def fake_runner(**kwargs: object) -> RunResult: + captured.update(kwargs) + run_root = ( + tmp_path + / "artifacts" + / "results" + / "trace-suite" + / "trace-concurrency" + ) + return RunResult( + state=RunState.COMPLETED, + exit_code=0, + suite_id="trace-suite", + run_id="trace-concurrency", + suite_root=run_root.parent, + run_root=run_root, + ) + + with patch( + "assert_ai.runner.run_pipeline_document_result", + side_effect=fake_runner, + ): + exit_code = worker_main( + [ + "--workspace", + str(tmp_path), + "--job-id", + started.job.job_id, + ] + ) + + assert request["concurrency"] == 1 + assert exit_code == 0 + assert captured["concurrency"] == 1 + + def test_trace_preflight_rejects_environment_files(tmp_path: Path) -> None: configs, service = _service(tmp_path) configs.save_config( @@ -839,7 +1316,7 @@ def test_failed_snapshot_write_removes_unregistered_job_directory( with ( patch( - "assert_ai.services.evaluations.write_json", + "assert_ai.services.evaluations._write_request_snapshot", side_effect=OSError("disk full"), ), pytest.raises(OSError, match="disk full"), @@ -890,6 +1367,27 @@ def test_worker_rejects_a_tampered_config_snapshot( "max_log_bytes": 4096, }, ) + request_sha256 = "sha256:" + hashlib.sha256( + (job_dir / "request.json").read_bytes() + ).hexdigest() + JobStore( + tmp_path / "artifacts" / "mcp" / "jobs.sqlite3" + ).create_or_get( + NewJob( + job_id=job_id, + idempotency_key="tampered-config", + request_hash="sha256:" + ("1" * 64), + request_sha256=request_sha256, + suite_id="suite", + run_id="run", + config_ref="demo.yaml", + config_sha256="sha256:" + ("0" * 64), + snapshot_path=str(job_dir / "config.yaml"), + request_path=str(job_dir / "request.json"), + resource_keys=("run:suite/run",), + ), + max_queued_jobs=1, + ) exit_code = worker_main( ["--workspace", str(tmp_path), "--job-id", job_id] @@ -904,6 +1402,48 @@ def test_worker_rejects_a_tampered_config_snapshot( assert "digest mismatch" in result["worker_error"]["error_message"] +def test_worker_rejects_coordinated_request_and_config_tampering( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "demo.yaml", + document=_write_inference_fixture(tmp_path), + ) + monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) + started = service.start("demo.yaml", request_id="bound-request") + record = service.store.get(started.job.job_id) + snapshot_path = Path(record.snapshot_path) + request_path = Path(record.request_path) + tampered_snapshot = b"pipeline: {}\n" + snapshot_path.write_bytes(tampered_snapshot) + request = json.loads(request_path.read_text(encoding="utf-8")) + request["config_sha256"] = "sha256:" + hashlib.sha256( + tampered_snapshot + ).hexdigest() + write_json(request_path, request) + + exit_code = worker_main( + [ + "--workspace", + str(tmp_path), + "--job-id", + started.job.job_id, + ] + ) + + result = json.loads( + ( + request_path.parent / "result.json" + ).read_text(encoding="utf-8") + ) + assert exit_code == 1 + assert "job request digest mismatch" in ( + result["worker_error"]["error_message"] + ) + + def test_malformed_worker_result_becomes_a_persisted_failure( tmp_path: Path, ) -> None: @@ -1194,6 +1734,66 @@ def test_startup_recovery_marks_dead_worker_interrupted( assert terminal.error_code == ServiceErrorCode.JOB_INTERRUPTED.value +def test_startup_recovery_cleans_up_disabled_job_kinds( + tmp_path: Path, +) -> None: + configs, service = _service(tmp_path) + configs.save_config( + "trace.yaml", + document=_write_trace_fixture(tmp_path), + ) + configs.save_config( + "evaluation.yaml", + document=_write_inference_fixture(tmp_path), + ) + with patch.object(EvaluationJobManager, "enqueue"): + trace = service.start_trace_judging( + "trace.yaml", + "fixtures/traces.json", + request_id="disabled-trace", + suite_id="trace-suite", + run_id="disabled-trace", + ) + evaluation = service.start( + "evaluation.yaml", + request_id="enabled-evaluation", + ) + claimed = service.store.claim_next( + lease_owner="old-trace-manager", + lease_seconds=0.01, + max_active_jobs=1, + job_kinds=("trace_judging",), + ) + assert claimed is not None + service.store.mark_running( + claimed.job_id, + lease_owner="old-trace-manager", + pid=2_147_483_647, + process_create_time=1, + lease_seconds=0.01, + ) + time.sleep(0.02) + execute_only = EvaluationJobManager( + service.workspace, + service.store, + job_kinds=("evaluation",), + lease_seconds=0.1, + ) + + with patch.object(EvaluationJobManager, "enqueue"): + execute_only._recover_startup() + + assert service.store.get(trace.job.job_id).state is JobState.INTERRUPTED + next_job = service.store.claim_next( + lease_owner="execute-manager", + lease_seconds=30, + max_active_jobs=1, + job_kinds=("evaluation",), + ) + assert next_job is not None + assert next_job.job_id == evaluation.job.job_id + + def test_startup_recovery_adopts_and_monitors_a_live_worker( tmp_path: Path, ) -> None: diff --git a/tests/test_job_store.py b/tests/test_job_store.py index 1758865b9..5a4a671bf 100644 --- a/tests/test_job_store.py +++ b/tests/test_job_store.py @@ -12,6 +12,7 @@ from assert_ai.services.errors import ServiceError, ServiceErrorCode from assert_ai.services.job_models import JobState, NewJob from assert_ai.services.job_store import JobStore +from assert_ai.services.output_identity import validate_output_id def _new_job( @@ -29,6 +30,7 @@ def _new_job( job_id=f"job-{suffix}", idempotency_key=request_id or f"request-{suffix}", request_hash=request_hash or f"hash-{suffix}", + request_sha256="sha256:" + ("0" * 64), suite_id=suite_id or f"suite-{suffix}", run_id=run_id if run_id is not None else f"run-{suffix}", config_ref="demo.yaml", @@ -93,6 +95,33 @@ def test_suite_run_assignment_is_unique(tmp_path: Path) -> None: assert conflict.value.details == {"job_id": "job-one"} +def test_suite_run_assignment_is_portably_case_insensitive( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get( + _new_job("one", suite_id="SuiteA", run_id="RunA"), + max_queued_jobs=10, + ) + + with pytest.raises(ServiceError) as conflict: + store.create_or_get( + _new_job("two", suite_id="suitea", run_id="runa"), + max_queued_jobs=10, + ) + + assert conflict.value.code == ServiceErrorCode.CONFLICT + assert conflict.value.details == {"job_id": "job-one"} + + +@pytest.mark.parametrize("value", ["run.", "CON", "nul.txt", "LPT9"]) +def test_output_ids_reject_windows_path_aliases(value: str) -> None: + with pytest.raises(ServiceError) as invalid: + validate_output_id(value, field_name="run_id") + + assert invalid.value.code == ServiceErrorCode.INVALID_ARGUMENT + + def test_create_respects_queued_job_limit(tmp_path: Path) -> None: store = JobStore(tmp_path / "jobs.sqlite3") store.create_or_get(_new_job("one"), max_queued_jobs=1) @@ -392,21 +421,44 @@ def test_v1_store_is_migrated_before_a_mutating_operation( 'demo.yaml', 'sha256:v1', 'config.yaml', 'request.json', '[]' ); + INSERT INTO jobs( + job_id, idempotency_key, request_hash, kind, state, + created_at, suite_id, run_id, config_ref, config_sha256, + snapshot_path, request_path, resource_keys_json + ) VALUES ( + 'job-v1-case-alias', 'request-v1-case-alias', + 'hash-v1-case-alias', 'evaluation', 'queued', + '2026-01-01T00:00:01+00:00', 'Suite-V1', 'Run-V1', + 'demo.yaml', 'sha256:v1-case-alias', 'config.yaml', + 'request.json', '[]' + ); PRAGMA user_version = 1; """ ) - cancelled = JobStore(path).request_cancel("job-v1") + migrated = JobStore(path) + cancelled = migrated.request_cancel("job-v1") assert cancelled.state is JobState.CANCELLED assert cancelled.retry_of is None with sqlite3.connect(path) as connection: - assert connection.execute("PRAGMA user_version").fetchone()[0] == 3 + assert connection.execute("PRAGMA user_version").fetchone()[0] == 4 columns = { row[1] for row in connection.execute("PRAGMA table_info(jobs)") } assert "retry_of" in columns + assert "request_sha256" in columns + with pytest.raises(ServiceError) as collision: + migrated.create_or_get( + _new_job( + "new-case-alias", + suite_id="SUITE-V1", + run_id="RUN-V1", + ), + max_queued_jobs=10, + ) + assert collision.value.code == ServiceErrorCode.CONFLICT def test_operation_lock_blocks_job_claim_until_released( @@ -551,6 +603,31 @@ def test_suite_operation_lock_conflicts_with_active_run_resource( ) +def test_suite_operation_lock_treats_underscore_as_literal( + tmp_path: Path, +) -> None: + store = JobStore(tmp_path / "jobs.sqlite3") + store.create_or_get( + _new_job( + "one", + suite_id="suitex", + resource_keys=("run:suitex/run-one",), + ), + max_queued_jobs=10, + ) + assert store.claim_next( + lease_owner="manager", + lease_seconds=30, + max_active_jobs=1, + ) + + assert store.acquire_operation_locks( + ("suite:suite_",), + owner="curator", + lease_seconds=30, + ) + + def test_event_retention_prefers_lifecycle_events(tmp_path: Path) -> None: store = JobStore(tmp_path / "jobs.sqlite3") store.create_or_get(_new_job("one"), max_queued_jobs=10) diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index aa63ad493..81d48386b 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -599,6 +599,7 @@ def test_trace_group_does_not_control_or_launch_evaluation_jobs( job_id=job_id, idempotency_key="evaluation-request", request_hash="sha256:" + ("1" * 64), + request_sha256="sha256:" + ("3" * 64), suite_id="evaluation-suite", run_id="evaluation-run", config_ref="evaluation.yaml", @@ -775,7 +776,7 @@ async def run() -> dict[str, Any]: "b09950417a44bf14c9bbf2702c1c00f23a18a0bfec03cab16a482733d8cf98c8" ), "preflight_evaluation": ( - "c9a686f7879c7e06a8f32c210cc02ed8e30c4fb8c6473cf77999971d114c9805" + "a26252422f3aaebd06087da2c5cf4c7fc47c73e6a78ad80a29a6eaf543ffa014" ), "design_config": ( "1cd55a1bba06468b0aaa785cf05a445e4bf16768ff6165254ceecf0181a8392e" @@ -1003,7 +1004,7 @@ def test_complete_versioned_curation_workflow(tmp_path: Path) -> None: { "type": "prompt", "test_case_id": "test_case_000001", - "prompt": "Book a flight.", + "seed": {"description": "Book a flight."}, "dimensions": {"behavior": "safe_booking"}, } ) @@ -1051,7 +1052,9 @@ async def run() -> dict[str, Any]: "suite_id": "curation-suite", "test_case_id": "test_case_000001", "updates": { - "prompt": "Book a policy-compliant flight.", + "seed": { + "description": "Book a policy-compliant flight.", + }, }, "expected_etag": first_test_set_etag, "change_summary": "Make the prompt explicit.", @@ -1070,7 +1073,9 @@ async def run() -> dict[str, Any]: { "suite_id": "curation-suite", "test_case_id": "test_case_000001", - "updates": {"prompt": "Stale update."}, + "updates": { + "seed": {"description": "Stale update."}, + }, "expected_etag": first_test_set_etag, "change_summary": "Attempt a stale edit.", }, @@ -1092,7 +1097,7 @@ async def run() -> dict[str, Any]: assert result["revised_case"].structured_content["artifacts"][0][ "version" ] == "v0002" - assert result["fetched"].structured_content["row"]["prompt"] == ( + assert result["fetched"].structured_content["row"]["seed"]["description"] == ( "Book a policy-compliant flight." ) assert result["stale"].is_error is True diff --git a/tests/test_measurement_fixes.py b/tests/test_measurement_fixes.py index d4a0ad8d5..dfce87a4a 100644 --- a/tests/test_measurement_fixes.py +++ b/tests/test_measurement_fixes.py @@ -1150,7 +1150,15 @@ async def fake_run_judge_attempts( edit=AddMessageEdit(message=Message(role="assistant", content=body)), ) ) - handle.write(json.dumps(transcript.to_dict(), ensure_ascii=False) + "\n") + row = transcript.to_dict() + if sid == "seed-refused": + row["trace_refs"] = [ + { + "trace_id": "a" * 32, + "span_ids": ["b" * 16], + } + ] + handle.write(json.dumps(row, ensure_ascii=False) + "\n") with patch("assert_ai.core.judge._run_judge_attempts", new=fake_run_judge_attempts): result = asyncio.run( @@ -1180,6 +1188,10 @@ async def fake_run_judge_attempts( self.assertEqual(refused["judge_status"], "filter_skipped") self.assertIn("judge_input_refused", refused["judge_error"]) self.assertEqual(refused["verdict"], {}) + self.assertEqual( + refused["trace_refs"], + [{"trace_id": "a" * 32, "span_ids": ["b" * 16]}], + ) for ok_seed in ("seed-ok", "seed-ok-2"): self.assertEqual(by_test_case[ok_seed]["judge_status"], "ok") diff --git a/tests/test_run_planning_service.py b/tests/test_run_planning_service.py index 7d16afe45..fd52c0732 100644 --- a/tests/test_run_planning_service.py +++ b/tests/test_run_planning_service.py @@ -366,8 +366,24 @@ def test_preflight_reuses_cache_without_writing_workspace() -> None: stages = {stage.name: stage for stage in result.stages} assert stages["systematize"].action is StageAction.REUSE assert stages["systematize"].artifact_version == plan.version + assert result.consumed_artifacts["systematize"].version == plan.version assert before == after + with patch.dict( + os.environ, + {"OPENAI_API_KEY": "configured-for-test"}, + clear=False, + ): + forced = planning.preflight( + "demo.yaml", + overrides=EvaluationOverrides( + force_stages=("systematize",), + ), + ) + + assert forced.stages[0].action is StageAction.RUN + assert "systematize" not in forced.consumed_artifacts + def test_resolve_forced_stages_rejects_missing_and_cascades() -> None: assert resolve_forced_stages( From f3cace20c81e82218e6568684f0f053e0807d4dc Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Wed, 26 Aug 2026 16:28:07 -0700 Subject: [PATCH 13/16] Give regression CI enough time to run Increase the Tier 1 timeout so fresh dependency resolution does not consume the entire job budget before the repository test suite starts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- .github/workflows/regression.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/regression.yml b/.github/workflows/regression.yml index 66d159948..bc1a7c6fb 100644 --- a/.github/workflows/regression.yml +++ b/.github/workflows/regression.yml @@ -45,7 +45,7 @@ jobs: tier1-unit: name: "Tier 1: Unit Tests" runs-on: ubuntu-latest - timeout-minutes: 5 + timeout-minutes: 20 steps: - uses: actions/checkout@v4 From 82698b5aeaf07889155566ab5f054722dfd8fb0d Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Wed, 26 Aug 2026 16:44:00 -0700 Subject: [PATCH 14/16] Separate MCP v2 from Phoenix dependencies Run Tier 1 against the supported MCP v2 environment instead of combining it with Phoenix dependencies that still require MCP v1. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- .github/workflows/regression.yml | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/.github/workflows/regression.yml b/.github/workflows/regression.yml index bc1a7c6fb..0f0889180 100644 --- a/.github/workflows/regression.yml +++ b/.github/workflows/regression.yml @@ -60,15 +60,12 @@ jobs: cache-dependency-path: viewer/package-lock.json - name: Install Python dependencies - # The `otel` extra is required because tests/test_incident_triage_smoke.py - # imports examples/incident_triage_agent/agent.py, which imports - # opentelemetry at module load time (target.trace.backend: otel). Demo - # examples like that one ship as part of the unit-test surface, so the - # CI install set has to include the demo's optional extras even though - # core ASSERT doesn't need them. MCP v2 is installed explicitly because - # it cannot currently share the `examples` extra's MCP v1 dependency. + # The base package already includes the OpenTelemetry SDK imported by + # the incident-triage smoke test. Do not combine the full `otel` extra + # with `mcp`: Phoenix currently brings an MCP v1 dependency, while the + # ASSERT MCP server intentionally requires v2. run: | - python -m pip install -e ".[dev,otel,mcp]" + python -m pip install -e ".[dev,mcp]" - name: Check the behavior library # Guards two things the reference library cannot enforce by itself: From 8d39623f851f1633204e5bb3976e820d8ba70a7e Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Thu, 27 Aug 2026 09:03:55 -0700 Subject: [PATCH 15/16] Cover downstream artifact selection in preflight Assert that inference-only preflight exposes the exact immutable test-set version later bound into the persisted job request. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- tests/test_evaluation_service.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_evaluation_service.py b/tests/test_evaluation_service.py index 631e68d0d..6e7c34942 100644 --- a/tests/test_evaluation_service.py +++ b/tests/test_evaluation_service.py @@ -431,6 +431,10 @@ def test_queued_job_rejects_a_different_curated_artifact_version( }, }, ) + preflight = service.planning.preflight("pinned.yaml") + assert preflight.consumed_artifacts["test_set"].version == ( + test_set_artifact.version + ) monkeypatch.setattr(EvaluationJobManager, "enqueue", lambda self: None) started = service.start( "pinned.yaml", From c372baf06e41a330b0da23884c86485ead14680b Mon Sep 17 00:00:00 2001 From: AaronAspinwall123 Date: Thu, 27 Aug 2026 09:55:23 -0700 Subject: [PATCH 16/16] Preserve expected errors on MCP 2.1 Mark sanitized application failures as SDK ToolError instances so MCP 2.1 returns their stable model-correctable payload while continuing to redact unexpected crashes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb4bac77-8743-43a4-8f2f-d1f7184cb965 --- assert_ai/mcp/errors.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/assert_ai/mcp/errors.py b/assert_ai/mcp/errors.py index a8d9ce3db..ab231ef93 100644 --- a/assert_ai/mcp/errors.py +++ b/assert_ai/mcp/errors.py @@ -12,7 +12,11 @@ from typing import TypeVar from uuid import uuid4 -from mcp.server.mcpserver.exceptions import ResourceError, ResourceNotFoundError +from mcp.server.mcpserver.exceptions import ( + ResourceError, + ResourceNotFoundError, + ToolError, +) from pydantic import BaseModel from assert_ai.core.workspace import WorkspaceService @@ -24,7 +28,7 @@ _T = TypeVar("_T") -class _McpToolError(RuntimeError): +class _McpToolError(ToolError): """Expected sanitized error that should survive nested adaptation."""