diff --git a/.gitignore b/.gitignore index bc634fe..525330b 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,8 @@ .mcp.json .DS_Store .agents/settings.local.json +.agents/skills/branch-context/ +AGENTS.local.md CLAUDE.local.md LOCAL_WORKTREES.md diff --git a/pydantic_ai_harness/__init__.py b/pydantic_ai_harness/__init__.py index e51895a..0ccda88 100644 --- a/pydantic_ai_harness/__init__.py +++ b/pydantic_ai_harness/__init__.py @@ -1,4 +1,4 @@ -"""Pydantic AI capability library.""" +"""The batteries for your Pydantic AI agent -- the official capability library.""" from typing import TYPE_CHECKING @@ -8,7 +8,12 @@ if TYPE_CHECKING: from .logfire import ManagedPrompt from .shell import Shell -__all__ = ['CodeMode', 'FileSystem', 'ManagedPrompt', 'Shell'] +__all__ = [ + 'CodeMode', + 'FileSystem', + 'ManagedPrompt', + 'Shell', +] def __getattr__(name: str) -> object: diff --git a/pydantic_ai_harness/experimental/_warn.py b/pydantic_ai_harness/experimental/_warn.py index 3ddf1e7..9e42d1f 100644 --- a/pydantic_ai_harness/experimental/_warn.py +++ b/pydantic_ai_harness/experimental/_warn.py @@ -28,7 +28,7 @@ _SILENCE_HINT = ( def warn_experimental(feature: str) -> None: """Emit a `HarnessExperimentalWarning` for *feature*, including how to silence all of them. - One filter silences the whole category — every experimental capability — so users never + One filter silences the whole category (every experimental capability), so users never need a suppression line per capability. """ warnings.warn( diff --git a/pydantic_ai_harness/experimental/compaction/README.md b/pydantic_ai_harness/experimental/compaction/README.md index 67dd838..d2d0eea 100644 --- a/pydantic_ai_harness/experimental/compaction/README.md +++ b/pydantic_ai_harness/experimental/compaction/README.md @@ -3,7 +3,7 @@ > [!WARNING] > **Experimental.** These capabilities live under `pydantic_ai_harness.experimental` and may > change or be removed in any release, without a deprecation period. Import them from the -> experimental path — there is no top-level export: +> experimental path -- there is no top-level export: > > ```python > from pydantic_ai_harness.experimental.compaction import TieredCompaction @@ -24,7 +24,7 @@ window. Each is a Pydantic AI `Capability` that runs in the `before_model_reques **persist** into the run's message history, so a trim/clear/summary carries forward to later steps (it is not recomputed from the full history every turn). -All strategies preserve tool-call / tool-return **pairing** — core does not validate this, and a +All strategies preserve tool-call / tool-return **pairing** -- core does not validate this, and a provider rejects an orphaned pair. The zero-LLM strategies never call a model. ## The menu @@ -48,9 +48,9 @@ near-lossless). `TieredCompaction` triggers and stops on a single `target_tokens ## Cost: why summarization is the last resort Summarization turns input tokens into output tokens, which are billed at a premium and generated -serially — so it is genuinely expensive. The zero-LLM strategies touch only the cheaper input side. +serially -- so it is genuinely expensive. The zero-LLM strategies touch only the cheaper input side. The field consensus (Anthropic, OpenCode, Letta) is to clear/dedupe first and summarize only when -that is not enough — which is exactly what `TieredCompaction` encodes: +that is not enough -- which is exactly what `TieredCompaction` encodes: ```python from pydantic_ai import Agent @@ -77,14 +77,14 @@ agent = Agent( ``` A tier inside `TieredCompaction` is driven directly by the orchestrator, which re-measures after each -and stops once under `target_tokens` — so a tier's own `max_*` trigger is irrelevant there (set it to +and stops once under `target_tokens` -- so a tier's own `max_*` trigger is irrelevant there (set it to anything valid). Any object with `async def compact(messages, ctx) -> list[ModelMessage]` (`CompactionStrategy`) can be a tier, so you can plug in your own. ## Cache tradeoff (read before using `ClearToolResults`) Clearing or deduplicating rewrites message content, which invalidates the provider's prompt cache -from the edit point onward — the next request pays a cache-write. Use `ClearToolResults`' +from the edit point onward -- the next request pays a cache-write. Use `ClearToolResults`' `min_clear_tokens` to skip clearing that reclaims too little to be worth busting the cache. ## Model inheritance @@ -94,8 +94,8 @@ running agent's model. No token caps are imposed on the summary call. ## Usage accounting -The summary call is a real request to the model, so its full usage — tokens **and** the request -itself — is folded into the run's `ctx.usage`. This is deliberate: it keeps cost honest, keeps the +The summary call is a real request to the model, so its full usage -- tokens **and** the request +itself -- is folded into the run's `ctx.usage`. This is deliberate: it keeps cost honest, keeps the request count consistent (a model request that didn't count as one would be the surprise), and lets a `UsageLimits` request limit catch a runaway compaction. A run-request / iteration limiter will therefore see compaction calls among its requests. @@ -120,5 +120,5 @@ def my_file_key(call: ToolCallPart) -> str | None: ## Out of scope These strategies compress or drop context *inside* the window. Moving large tool outputs *out* of the -window — overflowing them to a file the agent (or a subagent) can query on demand — is a separate +window -- overflowing them to a file the agent (or a subagent) can query on demand -- is a separate capability, not lossy truncation. Prefer it over capping individual tool outputs. diff --git a/pydantic_ai_harness/experimental/compaction/_clear_tool_results.py b/pydantic_ai_harness/experimental/compaction/_clear_tool_results.py index db88778..94b1b7e 100644 --- a/pydantic_ai_harness/experimental/compaction/_clear_tool_results.py +++ b/pydantic_ai_harness/experimental/compaction/_clear_tool_results.py @@ -1,4 +1,4 @@ -"""`ClearToolResults` — zero-cost in-place clearing of old tool results.""" +"""`ClearToolResults` -- zero-cost in-place clearing of old tool results.""" from __future__ import annotations @@ -31,7 +31,7 @@ class ClearToolResults(AbstractCapability[AgentDepsT]): calls remain paired with their (now-blanked) results, so the history stays valid. No LLM calls are made. - This is the cheap first tier of compaction — tool results typically dominate + This is the cheap first tier of compaction -- tool results typically dominate context, and the agent can re-run a tool if it needs the data again. Cache tradeoff: clearing rewrites message content, which invalidates the provider's diff --git a/pydantic_ai_harness/experimental/compaction/_deduplicate_file_reads.py b/pydantic_ai_harness/experimental/compaction/_deduplicate_file_reads.py index 63d436d..8d8f6ff 100644 --- a/pydantic_ai_harness/experimental/compaction/_deduplicate_file_reads.py +++ b/pydantic_ai_harness/experimental/compaction/_deduplicate_file_reads.py @@ -1,4 +1,4 @@ -"""`DeduplicateFileReads` — zero-cost in-place clearing of superseded file reads.""" +"""`DeduplicateFileReads` -- zero-cost in-place clearing of superseded file reads.""" from __future__ import annotations @@ -25,7 +25,7 @@ class DeduplicateFileReads(AbstractCapability[AgentDepsT]): earlier reads are blanked with a placeholder. Tool-call pairing is preserved. No LLM calls are made. - File identity is supplied by the ``file_key`` seam — given a ``ToolCallPart`` it returns + File identity is supplied by the ``file_key`` seam -- given a ``ToolCallPart`` it returns a stable key for the file being read, or ``None`` if the call is not a file read. There is no default: file-read identification is agent-specific, and a wrong guess would drop live data. diff --git a/pydantic_ai_harness/experimental/compaction/_limit_warner.py b/pydantic_ai_harness/experimental/compaction/_limit_warner.py index bd2dc9e..f7913d4 100644 --- a/pydantic_ai_harness/experimental/compaction/_limit_warner.py +++ b/pydantic_ai_harness/experimental/compaction/_limit_warner.py @@ -1,4 +1,4 @@ -"""`LimitWarner` — injects warnings as the run approaches configured limits.""" +"""`LimitWarner` -- injects warnings as the run approaches configured limits.""" from __future__ import annotations diff --git a/pydantic_ai_harness/experimental/compaction/_shared.py b/pydantic_ai_harness/experimental/compaction/_shared.py index 2518f87..4116e5e 100644 --- a/pydantic_ai_harness/experimental/compaction/_shared.py +++ b/pydantic_ai_harness/experimental/compaction/_shared.py @@ -1,7 +1,7 @@ """Shared utilities for the compaction capabilities. Token estimation, the `CompactionStrategy` protocol, tool-pair-safe cutoff logic, first-user -preservation, and in-place tool-result clearing — anything used by more than one capability. +preservation, and in-place tool-result clearing -- anything used by more than one capability. """ from __future__ import annotations @@ -121,7 +121,7 @@ class CompactionStrategy(Protocol[AgentDepsT]): # --------------------------------------------------------------------------- -# Safe cutoff logic — preserves tool-call / tool-return pairs +# Safe cutoff logic -- preserves tool-call / tool-return pairs # --------------------------------------------------------------------------- _TOOL_PAIR_SEARCH_RANGE = 5 diff --git a/pydantic_ai_harness/experimental/compaction/_sliding_window.py b/pydantic_ai_harness/experimental/compaction/_sliding_window.py index be10078..b29b68b 100644 --- a/pydantic_ai_harness/experimental/compaction/_sliding_window.py +++ b/pydantic_ai_harness/experimental/compaction/_sliding_window.py @@ -1,4 +1,4 @@ -"""`SlidingWindow` — zero-cost trimming of the oldest messages.""" +"""`SlidingWindow` -- zero-cost trimming of the oldest messages.""" from __future__ import annotations diff --git a/pydantic_ai_harness/experimental/compaction/_summarizing_compaction.py b/pydantic_ai_harness/experimental/compaction/_summarizing_compaction.py index 7bc94e0..a6eb1a7 100644 --- a/pydantic_ai_harness/experimental/compaction/_summarizing_compaction.py +++ b/pydantic_ai_harness/experimental/compaction/_summarizing_compaction.py @@ -1,4 +1,4 @@ -"""`SummarizingCompaction` — LLM-powered summarization of older messages.""" +"""`SummarizingCompaction` -- LLM-powered summarization of older messages.""" from __future__ import annotations @@ -44,7 +44,7 @@ The user's overall goal and any standing constraints or preferences. Choices made and the reasoning, so they are not relitigated. ## Artifacts -Files, paths, identifiers, commands, and APIs touched — quote exact names. +Files, paths, identifiers, commands, and APIs touched -- quote exact names. ## Current state What is done and what is in progress right now. @@ -55,7 +55,7 @@ The immediate actions still required to finish the task. ## Open questions Unresolved questions or blockers. -Focus on results, not a replay of completed actions. Respond ONLY with the summary — no \ +Focus on results, not a replay of completed actions. Respond ONLY with the summary -- no \ preamble, no markdown fences. @@ -140,8 +140,8 @@ class SummarizingCompaction(AbstractCapability[AgentDepsT]): summarized using a dedicated model call and replaced with a compact, structured summary message, preserving recent context and tool-call integrity. - This is the expensive tier — summarization turns input tokens into (pricier) output - tokens — so it is best used behind cheaper passes (see `TieredCompaction`). + This is the expensive tier -- summarization turns input tokens into (pricier) output + tokens -- so it is best used behind cheaper passes (see `TieredCompaction`). The summary call's usage is folded into the parent run's usage (it counts as a real request), so cost accounting stays honest; note this also increments the run's request diff --git a/pydantic_ai_harness/experimental/compaction/_tiered_compaction.py b/pydantic_ai_harness/experimental/compaction/_tiered_compaction.py index c3e68f5..ec071df 100644 --- a/pydantic_ai_harness/experimental/compaction/_tiered_compaction.py +++ b/pydantic_ai_harness/experimental/compaction/_tiered_compaction.py @@ -1,4 +1,4 @@ -"""`TieredCompaction` — escalation orchestrator over a sequence of strategies.""" +"""`TieredCompaction` -- escalation orchestrator over a sequence of strategies.""" from __future__ import annotations @@ -26,7 +26,7 @@ class TieredCompaction(AbstractCapability[AgentDepsT]): tool results, deduplicate reads, then summarize) so the expensive summarization tier is only reached when the cheap passes cannot reclaim enough. - Each tier's own trigger is bypassed — `TieredCompaction` drives the tiers directly via + Each tier's own trigger is bypassed -- `TieredCompaction` drives the tiers directly via their ``compact`` method and decides when to stop. Example: diff --git a/pydantic_ai_harness/experimental/media/__init__.py b/pydantic_ai_harness/experimental/media/__init__.py new file mode 100644 index 0000000..f3aee4f --- /dev/null +++ b/pydantic_ai_harness/experimental/media/__init__.py @@ -0,0 +1,41 @@ +"""Content-addressed media stores for offloading large binary parts. + +Used by `pydantic_ai_harness.experimental.step_persistence` to keep snapshots small when +messages carry `BinaryContent` payloads. A forthcoming `MediaExternalizer` +capability will reuse these stores for in-flight wire-payload reduction +(rewriting `BinaryContent` to URL parts before the model sees them). +""" + +from pydantic_ai_harness.experimental._warn import warn_experimental +from pydantic_ai_harness.experimental.media._s3 import S3MediaStore +from pydantic_ai_harness.experimental.media._store import ( + DiskMediaStore, + KeyStrategy, + MediaContext, + MediaStore, + PublicUrlResolver, + SqliteMediaStore, + default_key_strategy, + make_static_public_url, + media_uri_for, + parse_media_uri, +) +from pydantic_ai_harness.experimental.media._walker import externalize_media, restore_media + +warn_experimental('media') + +__all__ = [ + 'DiskMediaStore', + 'KeyStrategy', + 'MediaContext', + 'MediaStore', + 'PublicUrlResolver', + 'S3MediaStore', + 'SqliteMediaStore', + 'default_key_strategy', + 'externalize_media', + 'make_static_public_url', + 'media_uri_for', + 'parse_media_uri', + 'restore_media', +] diff --git a/pydantic_ai_harness/experimental/media/_s3.py b/pydantic_ai_harness/experimental/media/_s3.py new file mode 100644 index 0000000..ed35e9b --- /dev/null +++ b/pydantic_ai_harness/experimental/media/_s3.py @@ -0,0 +1,309 @@ +"""S3-compatible media store with handrolled AWS SigV4 signing. + +Targets AWS S3, Cloudflare R2, MinIO, and other S3-compatible endpoints. +Implements path-style URLs (`//`) -- the lowest common +denominator across providers. Only PUT / GET / HEAD are implemented; +multipart, lifecycle, and listing operations are out of scope (any blob we +store is sub-5GB by construction so multipart is unnecessary). + +The SigV4 implementation follows the AWS reference algorithm but is local +to this module to keep `pydantic-ai-harness` free of `botocore` / `boto3`. +If a provider needs deeper integration, write a subclass that overrides +`_request`. +""" + +from __future__ import annotations + +import hashlib +import hmac +import re +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import TYPE_CHECKING +from urllib.parse import quote + +from pydantic_ai_harness.experimental.media._store import ( + _EMPTY_CONTEXT, # pyright: ignore[reportPrivateUsage] + KeyStrategy, + MediaContext, + PublicUrlResolver, + _resolve_public_url, # pyright: ignore[reportPrivateUsage] + default_key_strategy, + media_uri_for, + parse_media_uri, +) + +if TYPE_CHECKING: + import httpx + + +_ALGORITHM = 'AWS4-HMAC-SHA256' +_SERVICE = 's3' +# Conservative subset for x-amz-meta-* keys: ASCII letters/digits/dash; lowercase. +_META_KEY_RE = re.compile(r'^[a-zA-Z0-9-]+$') + + +def _utcnow() -> datetime: + return datetime.now(timezone.utc) + + +def _hex_sha256(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _hmac_sha256(key: bytes, msg: str) -> bytes: + return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest() + + +def _canonical_uri(path: str) -> str: + """Percent-encode each path segment but keep `/` separators. + + AWS canonicalization requires double-encoding for non-S3 services but + single-encoding for S3 -- we use single-encoding throughout. + """ + return '/' + '/'.join(quote(segment, safe='') for segment in path.lstrip('/').split('/')) + + +def _signing_key(secret_access_key: str, date_stamp: str, region: str) -> bytes: + k_date = _hmac_sha256(f'AWS4{secret_access_key}'.encode(), date_stamp) + k_region = _hmac_sha256(k_date, region) + k_service = _hmac_sha256(k_region, _SERVICE) + return _hmac_sha256(k_service, 'aws4_request') + + +def sign_request( + *, + method: str, + host: str, + path: str, + body: bytes, + region: str, + access_key_id: str, + secret_access_key: str, + content_type: str | None = None, + extra_signed_headers: Mapping[str, str] | None = None, + now: datetime | None = None, +) -> dict[str, str]: + """Return the headers (including `Authorization`) to send with an S3 request. + + `path` should start with `/` and already include any bucket prefix -- + e.g. `/my-bucket/object-key`. Query strings are not supported in v1 + (we never send them for PUT/GET/HEAD of a single object). + + `extra_signed_headers` (lowercase keys) are included in the canonical + request and signature -- e.g. `x-amz-meta-*` user metadata headers. + """ + timestamp = now or _utcnow() + amz_date = timestamp.strftime('%Y%m%dT%H%M%SZ') + date_stamp = timestamp.strftime('%Y%m%d') + + payload_hash = _hex_sha256(body) + headers_to_sign: dict[str, str] = { + 'host': host, + 'x-amz-content-sha256': payload_hash, + 'x-amz-date': amz_date, + } + if content_type is not None: + headers_to_sign['content-type'] = content_type + if extra_signed_headers: + for k, v in extra_signed_headers.items(): + headers_to_sign[k.lower()] = v + + sorted_keys = sorted(headers_to_sign) + canonical_headers = ''.join(f'{k}:{headers_to_sign[k]}\n' for k in sorted_keys) + signed_headers = ';'.join(sorted_keys) + + canonical_request = '\n'.join( + [ + method.upper(), + _canonical_uri(path), + '', + canonical_headers, + signed_headers, + payload_hash, + ] + ) + + credential_scope = f'{date_stamp}/{region}/{_SERVICE}/aws4_request' + string_to_sign = '\n'.join( + [ + _ALGORITHM, + amz_date, + credential_scope, + _hex_sha256(canonical_request.encode('utf-8')), + ] + ) + + signature = hmac.new( + _signing_key(secret_access_key, date_stamp, region), + string_to_sign.encode('utf-8'), + hashlib.sha256, + ).hexdigest() + + authorization = ( + f'{_ALGORITHM} ' + f'Credential={access_key_id}/{credential_scope}, ' + f'SignedHeaders={signed_headers}, ' + f'Signature={signature}' + ) + return {**headers_to_sign, 'authorization': authorization} + + +def _meta_headers_from_context(context: MediaContext) -> dict[str, str]: + """Map `context.metadata` to `x-amz-meta-*` headers (lowercase, ASCII-safe keys). + + Raises `ValueError` for keys containing characters disallowed by HTTP + header naming. Values are passed through verbatim -- callers must keep + them ASCII-clean (S3 returns 400 for non-ASCII without `*=` RFC 5987 + encoding, which we don't support in v1). + """ + headers: dict[str, str] = {} + for key, value in context.metadata.items(): + if not _META_KEY_RE.fullmatch(key): + raise ValueError( + f'metadata key {key!r} is not a valid HTTP header token; use ASCII letters / digits / dashes' + ) + headers[f'x-amz-meta-{key.lower()}'] = value + return headers + + +class S3MediaStore: + """S3-compatible content-addressed store using path-style URLs + SigV4. + + Works with AWS S3 (`endpoint='https://s3..amazonaws.com'`), + Cloudflare R2 (`endpoint='https://.r2.cloudflarestorage.com'`, + `region='auto'`), MinIO, and other compatible providers. + + Customisation: + + - `key_strategy=` -- `Callable[[str, MediaContext], str]`. Default + strategy yields `.bin`. Override when the bucket + layout demands a specific path / extension scheme. See the + `KeyStrategy` docstring for the read-time caveat. + - `public_url=` -- sync or async callable that takes + `(uri, MediaContext)` and returns a URL (or `None`). Use + `make_static_public_url(...)` for public bucket / CDN setups; + provide your own async callable for presigned URLs (TTL captured in + the closure; `MediaContext` available for content-type-specific + response headers etc.). Without a resolver `public_url(...)` returns + `None`. + - `client=` -- bring your own `httpx.AsyncClient` for connection pooling. + + **Metadata persistence**: `context.metadata` is sent as signed + `x-amz-meta-*` headers on PUT (subject to ASCII naming rules) and + read back via `get_metadata(uri)`, which strips the `x-amz-meta-` + prefix from the HEAD response headers. + """ + + def __init__( + self, + *, + bucket: str, + endpoint: str, + region: str, + access_key_id: str, + secret_access_key: str, + key_prefix: str = '', + client: httpx.AsyncClient | None = None, + key_strategy: KeyStrategy = default_key_strategy, + public_url: PublicUrlResolver | None = None, + ) -> None: + self._bucket = bucket + self._endpoint = endpoint.rstrip('/') + self._region = region + self._access_key_id = access_key_id + self._secret_access_key = secret_access_key + self._key_prefix = key_prefix + self._client = client + self._key_strategy = key_strategy + self._public_url_resolver = public_url + + def _object_path(self, uri: str, context: MediaContext) -> str: + key = self._key_strategy(uri, context) + return f'/{self._bucket}/{self._key_prefix}{key}' + + def _host(self) -> str: + host = self._endpoint.split('://', 1)[1] + return host.split('/', 1)[0] + + async def _request( + self, + method: str, + path: str, + *, + body: bytes = b'', + content_type: str | None = None, + meta_headers: Mapping[str, str] | None = None, + ) -> httpx.Response: + import httpx + + headers = sign_request( + method=method, + host=self._host(), + path=path, + body=body, + region=self._region, + access_key_id=self._access_key_id, + secret_access_key=self._secret_access_key, + content_type=content_type, + extra_signed_headers=meta_headers, + ) + # The signature covers `_canonical_uri(path)` (each segment percent-encoded). + # Send that exact byte sequence as the wire path via `raw_path`, otherwise + # httpx applies its own (looser) path encoding and a custom key_prefix / + # key_strategy emitting reserved chars (`@`, `(`, `=`, ...) would diverge + # from the signed path -> SignatureDoesNotMatch. + url = httpx.URL(self._endpoint).copy_with(raw_path=_canonical_uri(path).encode('ascii')) + if self._client is not None: + return await self._client.request(method, url, content=body, headers=headers) + async with httpx.AsyncClient() as client: + return await client.request(method, url, content=body, headers=headers) + + async def put(self, data: bytes, *, context: MediaContext = _EMPTY_CONTEXT) -> str: + uri = media_uri_for(data) + meta = _meta_headers_from_context(context) + response = await self._request( + 'PUT', + self._object_path(uri, context), + body=data, + content_type=context.media_type, + meta_headers=meta or None, + ) + if response.status_code // 100 != 2: + raise RuntimeError(f'S3 PUT failed for {uri}: {response.status_code} {response.text[:200]}') + return uri + + async def get(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> bytes: + digest = parse_media_uri(uri) + response = await self._request('GET', self._object_path(uri, context)) + if response.status_code == 404: + raise FileNotFoundError(f'media not found: {digest}') + if response.status_code // 100 != 2: + raise RuntimeError(f'S3 GET failed for {digest}: {response.status_code} {response.text[:200]}') + return response.content + + async def exists(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> bool: + digest = parse_media_uri(uri) + response = await self._request('HEAD', self._object_path(uri, context)) + if response.status_code == 404: + return False + if response.status_code // 100 != 2: + raise RuntimeError(f'S3 HEAD failed for {digest}: {response.status_code} {response.text[:200]}') + return True + + async def public_url(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> str | None: + return await _resolve_public_url(self._public_url_resolver, uri, context) + + async def get_metadata(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> Mapping[str, str]: + digest = parse_media_uri(uri) + response = await self._request('HEAD', self._object_path(uri, context)) + if response.status_code == 404: + raise FileNotFoundError(f'media not found: {digest}') + if response.status_code // 100 != 2: + raise RuntimeError(f'S3 HEAD failed for {digest}: {response.status_code} {response.text[:200]}') + out: dict[str, str] = {} + for key, value in response.headers.items(): + lower = key.lower() + if lower.startswith('x-amz-meta-'): + out[lower[len('x-amz-meta-') :]] = value + return out diff --git a/pydantic_ai_harness/experimental/media/_store.py b/pydantic_ai_harness/experimental/media/_store.py new file mode 100644 index 0000000..b7bc848 --- /dev/null +++ b/pydantic_ai_harness/experimental/media/_store.py @@ -0,0 +1,445 @@ +"""Local `MediaStore` protocol + `DiskMediaStore` / `SqliteMediaStore`. + +The shared URI scheme is `media+sha256://` -- content-addressed, +so the same blob written through any store resolves the same way and dedup +is automatic. +""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import re +import sqlite3 +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType +from typing import Protocol, TypeGuard, runtime_checkable + +import anyio.to_thread + +_URI_SCHEME = 'media+sha256://' +_HEX_RE = re.compile(r'^[0-9a-f]{64}$') + +# Sentinel: empty, shareable, immutable. Used as the default `context` on every +# `MediaStore` method. Pulled into a module constant so a default-bound empty +# context is always the same instance (cheap identity checks, no per-call +# allocation). +_EMPTY_METADATA: Mapping[str, str] = MappingProxyType({}) + + +@dataclass(frozen=True, kw_only=True) +class MediaContext: + """Per-operation context threaded through `MediaStore` methods + callables. + + Extensible bag -- new use cases (TTL hints, response-header overrides, + origin run ids for audit, etc.) add a field here without breaking any + method signature or user-supplied callable. Fields default to `None` / + empty so callers and resolvers can ignore what they don't care about. + + Conventions: + + - `media_type` -- IANA media type (`image/png`, `audio/wav`, ...). When + absent, callbacks may guess from `filename` (stdlib `mimetypes`) or + fall back to a generic default. Never required. + - `filename` -- original filename for the bytes, if known. Useful for + key strategies that want a recognisable extension, and for resolvers + that need to set `Content-Disposition` on presigned URLs. + - `metadata` -- free-form `dict[str, str]` of user-supplied tags. + Persisted by all three concrete stores (`SqliteMediaStore` → JSON + column, `S3MediaStore` → signed `x-amz-meta-*` headers, + `DiskMediaStore` → sidecar JSON file). Read back via + `MediaStore.get_metadata(uri)`. + """ + + media_type: str | None = None + filename: str | None = None + metadata: Mapping[str, str] = field(default_factory=lambda: _EMPTY_METADATA) + + +_EMPTY_CONTEXT = MediaContext() + + +PublicUrlResolver = Callable[[str, MediaContext], 'str | None | Awaitable[str | None]'] +"""User-supplied callable that turns a `media+sha256://` URI into a fetchable URL. + +Sync or async; the store auto-detects via `inspect.isawaitable` on the result. +Return `None` to signal "no URL for this URI" (the resolver may be content-aware +and choose to inline some payloads). + +Receives the full `MediaContext` so the resolver can vary its output by media +type, response headers, TTL, etc. without breaking the signature later. +""" + +KeyStrategy = Callable[[str, MediaContext], str] +"""User-supplied callable that turns a URI + context into a backend storage key. + +Used by `DiskMediaStore` and `S3MediaStore` to derive the on-disk path / S3 +object key from the canonical URI. `SqliteMediaStore` does not take one: its +primary key is the content digest, so a user-chosen key would break dedup. +Default strategy is `.bin` (see `default_key_strategy`); override when +an existing bucket layout or naming convention dictates otherwise. + +**Caveat**: if your strategy depends on `context.media_type` (e.g. to pick an +extension), `get(uri)` and `exists(uri)` won't find the blob unless callers +pass the same context at read time. For pure path-organisation strategies +(e.g. `'media/' + digest + '.bin'`) the context can be ignored entirely. +""" + + +def default_key_strategy(uri: str, context: MediaContext) -> str: + """Default storage key: `.bin`. Ignores context.""" + return f'{parse_media_uri(uri)}.bin' + + +async def _resolve_public_url( + resolver: PublicUrlResolver | None, + uri: str, + context: MediaContext, +) -> str | None: + if resolver is None: + return None + result = resolver(uri, context) + if inspect.isawaitable(result): + return await result + return result + + +def make_static_public_url( + base: str, + *, + key_prefix: str = '', + extension: str = '.bin', +) -> PublicUrlResolver: + """Build a static-URL resolver for the common "public bucket / CDN" case. + + The returned callable takes a `media+sha256://` URI and returns + `{base}/{key_prefix}{extension}`. Use this when the bytes live at a + known stable URL -- e.g. an R2 public bucket (`https://pub-xxx.r2.dev`) + or a CDN in front of S3. For presigned URLs or any logic that runs per + request, pass your own callable instead. + + Ignores `MediaContext` -- pass your own callable if URL needs to vary + with media type, TTL, etc. + """ + base = base.rstrip('/') + + def resolver(uri: str, context: MediaContext) -> str: + digest = parse_media_uri(uri) + return f'{base}/{key_prefix}{digest}{extension}' + + return resolver + + +def media_uri_for(data: bytes) -> str: + """Return the canonical `media+sha256://` URI for `data`. + + The URI is the same regardless of which store the bytes are written to -- + content-addressed so two stores holding the same bytes can be queried + interchangeably. + """ + return f'{_URI_SCHEME}{hashlib.sha256(data).hexdigest()}' + + +def parse_media_uri(uri: str) -> str: + """Return the lowercase hex digest from a `media+sha256://` URI. + + Raises `ValueError` if `uri` does not match the scheme or the digest is + not 64 lowercase hex characters. + """ + if not uri.startswith(_URI_SCHEME): + raise ValueError(f'not a media URI: {uri!r}') + digest = uri[len(_URI_SCHEME) :] + if not _HEX_RE.fullmatch(digest): + raise ValueError(f'invalid sha256 digest in {uri!r}') + return digest + + +@runtime_checkable +class MediaStore(Protocol): + """Async content-addressed bytes store. + + `put` returns the canonical URI (`media+sha256://`) for the bytes -- + callers do not pick the key. Implementations may dedup on the hash. + + Every method takes an optional `context: MediaContext` for forward + extensibility. New context fields don't break existing call sites or + user-supplied callables. + + `public_url(uri)` returns a URL the model can fetch directly, or `None` + if the store can't or hasn't been configured to produce one. The forthcoming + `MediaExternalizer` capability uses this to rewrite `BinaryContent` parts + to `ImageUrl`/`AudioUrl`/etc. before the model sees the message. + """ + + async def put(self, data: bytes, *, context: MediaContext = _EMPTY_CONTEXT) -> str: ... # pragma: no cover + + async def get(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> bytes: ... # pragma: no cover + + async def exists(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> bool: ... # pragma: no cover + + async def public_url( + self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT + ) -> str | None: ... # pragma: no cover + + async def get_metadata( + self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT + ) -> Mapping[str, str]: ... # pragma: no cover + + +class DiskMediaStore: + """Filesystem store. One file per blob; default layout is `/.bin`. + + Directory is created on first write. Dedup is automatic via the + content-hash filename. + + Customisation: + + - `key_strategy=` -- `Callable[[str, MediaContext], str]` overriding the + relative path inside `directory` (e.g. `'images/.png'`). + The returned path must be relative and free of `..` segments + (absolute paths and `..` are both rejected with `ValueError`). + Caveat: if the strategy uses `context.media_type` etc. to pick the + filename, the same context must be supplied on `get`/`exists` to + locate the blob. + - `public_url=` -- resolver exposing a URL the model can fetch. + Without it `public_url(...)` returns `None` (local filesystem paths + are not addressable from a model). + + **Metadata persistence**: `context.metadata` is written to a sidecar + JSON file (`.meta.json`) alongside the blob and read back + via `get_metadata(uri)`. Sidecars are written atomically (tmp + + rename) and are missing only when no metadata was supplied. + """ + + def __init__( + self, + directory: str | Path, + *, + key_strategy: KeyStrategy = default_key_strategy, + public_url: PublicUrlResolver | None = None, + ) -> None: + self._root = Path(directory) + self._key_strategy = key_strategy + self._public_url_resolver = public_url + + def _path_for(self, uri: str, context: MediaContext) -> Path: + relative = self._key_strategy(uri, context) + candidate = Path(relative) + # `Path('/root') / '/abs'` discards `/root` and returns `/abs` -- an absolute + # key would silently escape the store directory. Reject both shapes. + if candidate.is_absolute() or '..' in candidate.parts: + raise ValueError(f'key_strategy produced traversal-unsafe path: {relative!r}') + return self._root / candidate + + async def put(self, data: bytes, *, context: MediaContext = _EMPTY_CONTEXT) -> str: + return await anyio.to_thread.run_sync(self._sync_put, data, context) + + def _sync_put(self, data: bytes, context: MediaContext) -> str: + uri = media_uri_for(data) + path = self._path_for(uri, context) + path.parent.mkdir(parents=True, exist_ok=True) + if not path.exists(): + tmp = path.with_suffix(path.suffix + '.tmp') + tmp.write_bytes(data) + tmp.replace(path) + if context.metadata: + meta_path = path.with_name(path.name + '.meta.json') + meta_tmp = meta_path.with_suffix('.json.tmp') + meta_tmp.write_text(json.dumps(dict(context.metadata))) + meta_tmp.replace(meta_path) + return uri + + async def get(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> bytes: + return await anyio.to_thread.run_sync(self._sync_get, uri, context) + + def _sync_get(self, uri: str, context: MediaContext) -> bytes: + path = self._path_for(uri, context) + if not path.exists(): + raise FileNotFoundError(f'media not found: {uri}') + return path.read_bytes() + + async def exists(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> bool: + return await anyio.to_thread.run_sync(self._sync_exists, uri, context) + + def _sync_exists(self, uri: str, context: MediaContext) -> bool: + return self._path_for(uri, context).exists() + + async def public_url(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> str | None: + return await _resolve_public_url(self._public_url_resolver, uri, context) + + async def get_metadata(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> Mapping[str, str]: + return await anyio.to_thread.run_sync(self._sync_get_metadata, uri, context) + + def _sync_get_metadata(self, uri: str, context: MediaContext) -> Mapping[str, str]: + path = self._path_for(uri, context) + meta_path = path.with_name(path.name + '.meta.json') + if not meta_path.exists(): + return _EMPTY_METADATA + return _coerce_metadata_mapping(json.loads(meta_path.read_text())) + + +def _is_object_dict(raw: object) -> TypeGuard[dict[object, object]]: + return isinstance(raw, dict) + + +def _coerce_metadata_mapping(raw: object) -> Mapping[str, str]: + if not _is_object_dict(raw): + raise ValueError(f'metadata payload must be a JSON object, got {type(raw).__name__}') + out: dict[str, str] = {} + for key, value in raw.items(): + if not isinstance(key, str) or not isinstance(value, str): + raise ValueError(f'metadata entries must be str→str, got {key!r}: {value!r}') + out[key] = value + return out + + +class SqliteMediaStore: + """SQLite store. One row per blob in a `media` table keyed by sha256 hex. + + Pass either a path to a SQLite file (created on demand) or an existing + `sqlite3.Connection`. When a connection is passed the caller owns its + lifecycle; when a path is passed each call opens a short-lived connection + inside the worker thread (safe across event-loop threads). + + A caller-owned connection **must** be created with + `check_same_thread=False`. Store methods dispatch SQL onto worker threads + via `anyio.to_thread`, so the stdlib default (`check_same_thread=True`) + raises `sqlite3.ProgrammingError` on first use. The path form sets this + internally; a passed-in connection cannot, so it is the caller's + responsibility. + + The table layout is: + + ```sql + CREATE TABLE IF NOT EXISTS media ( + sha256 TEXT PRIMARY KEY, + media_type TEXT, + bytes BLOB NOT NULL, + size_bytes INTEGER NOT NULL, + metadata TEXT + ); + ``` + + `INSERT OR IGNORE` makes writes idempotent -- the second `put` with the + same content is a no-op, not an overwrite. `metadata` is stored as + canonical JSON of `context.metadata` (empty mapping → `'{}'`) and + read back via `get_metadata(uri)`. + + There is no `key_strategy=` parameter: SqliteMediaStore is + content-addressed and the digest IS the primary key -- any user-chosen + storage key would either break dedup or be a no-op. Use `table=` if + you need a non-default table name. + """ + + def __init__( + self, + *, + database: str | Path | None = None, + connection: sqlite3.Connection | None = None, + table: str = 'media', + public_url: PublicUrlResolver | None = None, + ) -> None: + if (database is None) == (connection is None): + raise ValueError('provide exactly one of `database=` or `connection=`') + self._database = Path(database) if database is not None else None + self._connection = connection + if not re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]*', table): + raise ValueError(f'invalid table name: {table!r}') + self._table = table + self._schema_ready = False + self._public_url_resolver = public_url + + def _ensure_schema(self, conn: sqlite3.Connection) -> None: + if self._schema_ready: + return + conn.execute( + f'CREATE TABLE IF NOT EXISTS {self._table} (' + 'sha256 TEXT PRIMARY KEY, ' + 'media_type TEXT, ' + 'bytes BLOB NOT NULL, ' + 'size_bytes INTEGER NOT NULL, ' + 'metadata TEXT)' + ) + conn.commit() + self._schema_ready = True + + def _open(self) -> sqlite3.Connection: + if self._connection is not None: + return self._connection + assert self._database is not None + self._database.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(self._database, isolation_level=None, check_same_thread=False) + conn.execute('PRAGMA journal_mode=WAL').close() + return conn + + def _maybe_close(self, conn: sqlite3.Connection) -> None: + if self._connection is None: + conn.close() + + async def put(self, data: bytes, *, context: MediaContext = _EMPTY_CONTEXT) -> str: + return await anyio.to_thread.run_sync(self._sync_put, data, context) + + def _sync_put(self, data: bytes, context: MediaContext) -> str: + uri = media_uri_for(data) + digest = parse_media_uri(uri) + metadata_json = json.dumps(dict(context.metadata)) + conn = self._open() + try: + self._ensure_schema(conn) + conn.execute( + f'INSERT OR IGNORE INTO {self._table} ' + '(sha256, media_type, bytes, size_bytes, metadata) VALUES (?, ?, ?, ?, ?)', + (digest, context.media_type, data, len(data), metadata_json), + ) + finally: + self._maybe_close(conn) + return uri + + async def get(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> bytes: + digest = parse_media_uri(uri) + return await anyio.to_thread.run_sync(self._sync_get, digest) + + def _sync_get(self, digest: str) -> bytes: + conn = self._open() + try: + self._ensure_schema(conn) + row = conn.execute(f'SELECT bytes FROM {self._table} WHERE sha256 = ?', (digest,)).fetchone() + finally: + self._maybe_close(conn) + if row is None: + raise FileNotFoundError(f'media not found: {digest}') + return bytes(row[0]) + + async def exists(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> bool: + digest = parse_media_uri(uri) + return await anyio.to_thread.run_sync(self._sync_exists, digest) + + def _sync_exists(self, digest: str) -> bool: + conn = self._open() + try: + self._ensure_schema(conn) + row = conn.execute(f'SELECT 1 FROM {self._table} WHERE sha256 = ?', (digest,)).fetchone() + finally: + self._maybe_close(conn) + return row is not None + + async def public_url(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> str | None: + return await _resolve_public_url(self._public_url_resolver, uri, context) + + async def get_metadata(self, uri: str, *, context: MediaContext = _EMPTY_CONTEXT) -> Mapping[str, str]: + digest = parse_media_uri(uri) + return await anyio.to_thread.run_sync(self._sync_get_metadata, digest) + + def _sync_get_metadata(self, digest: str) -> Mapping[str, str]: + conn = self._open() + try: + self._ensure_schema(conn) + row = conn.execute(f'SELECT metadata FROM {self._table} WHERE sha256 = ?', (digest,)).fetchone() + finally: + self._maybe_close(conn) + if row is None: + raise FileNotFoundError(f'media not found: {digest}') + return _coerce_metadata_mapping(json.loads(row[0])) diff --git a/pydantic_ai_harness/experimental/media/_walker.py b/pydantic_ai_harness/experimental/media/_walker.py new file mode 100644 index 0000000..ff56858 --- /dev/null +++ b/pydantic_ai_harness/experimental/media/_walker.py @@ -0,0 +1,128 @@ +"""JSON walkers that swap inline `BinaryContent` for externalized markers. + +Used by step-persistence backends to keep large media out of snapshot +payloads. The walkers operate on JSON-shaped trees (`list`/`dict`/scalars) +so they never need to know about `pydantic_ai.messages` types directly -- +they recognize the serialized shape (`{"kind": "binary", "data": ""}`) +emitted by `ModelMessagesTypeAdapter.dump_json`. + +Round-trip: `externalize_media` (pre-write) and `restore_media` (post-read) +are inverses for bytes ≥ threshold. Bytes below the threshold pass through +unchanged. +""" + +from __future__ import annotations + +import base64 +from typing import TypeGuard + +from pydantic_ai_harness.experimental.media._store import MediaContext, MediaStore + +_EXTERNAL_MARKER = '__harness_external_media__' + + +def _is_json_dict(node: object) -> TypeGuard[dict[str, object]]: + """Treat any `dict` returned by `json.loads` as `dict[str, object]`. + + JSON guarantees string keys, so the runtime `isinstance(..., dict)` check + is sufficient -- pyright just needs the explicit TypeGuard to propagate + the value-type narrowing through the walker. + """ + return isinstance(node, dict) + + +def _is_json_list(node: object) -> TypeGuard[list[object]]: + """Treat any `list` returned by `json.loads` as `list[object]`. + + Same rationale as `_is_json_dict` -- JSON lists contain JSON-compatible + scalars, and we walk every element regardless. + """ + return isinstance(node, list) + + +def _is_binary_part(node: dict[str, object]) -> bool: + return node.get('kind') == 'binary' and isinstance(node.get('data'), str) + + +def _is_external_marker(node: dict[str, object]) -> bool: + return node.get(_EXTERNAL_MARKER) is True + + +async def externalize_media(node: object, *, media_store: MediaStore, threshold_bytes: int) -> object: + """Walk `node` and replace inline `BinaryContent` parts ≥ threshold. + + Each qualifying part becomes a marker dict carrying the canonical + `media+sha256://` URI; the raw bytes are written to `media_store`. + Returns a new tree -- input is not mutated. + """ + if _is_json_list(node): + out_list: list[object] = [] + for item in node: + out_list.append(await externalize_media(item, media_store=media_store, threshold_bytes=threshold_bytes)) + return out_list + if _is_json_dict(node): + if _is_binary_part(node): + replaced = await _maybe_externalize_binary(node, media_store, threshold_bytes) + if replaced is not None: + return replaced + out_dict: dict[str, object] = {} + for key, value in node.items(): + out_dict[key] = await externalize_media(value, media_store=media_store, threshold_bytes=threshold_bytes) + return out_dict + return node + + +async def _maybe_externalize_binary( + node: dict[str, object], + media_store: MediaStore, + threshold_bytes: int, +) -> dict[str, object] | None: + # `_is_binary_part` already verified `data` is a string; the cast is safe. + data_value = node['data'] + assert isinstance(data_value, str) + raw = base64.b64decode(data_value) + if len(raw) < threshold_bytes: + return None + media_type_value = node.get('media_type') + media_type = media_type_value if isinstance(media_type_value, str) else None + uri = await media_store.put(raw, context=MediaContext(media_type=media_type)) + # Preserve every field except the externalized `data`, so the round-trip + # survives new `BinaryContent` fields added by pydantic_ai upstream. + marker = {key: value for key, value in node.items() if key != 'data'} + marker[_EXTERNAL_MARKER] = True + marker['uri'] = uri + return marker + + +async def restore_media(node: object, *, media_store: MediaStore) -> object: + """Inverse of `externalize_media`. Walks `node` and re-inlines external refs. + + Each marker dict's `uri` is resolved via `media_store.get` and the bytes + are re-encoded as a `kind=binary` part so the result round-trips + through `ModelMessagesTypeAdapter.validate_python`. + """ + if _is_json_list(node): + out_list: list[object] = [] + for item in node: + out_list.append(await restore_media(item, media_store=media_store)) + return out_list + if _is_json_dict(node): + if _is_external_marker(node): + return await _restore_external(node, media_store) + out_dict: dict[str, object] = {} + for key, value in node.items(): + out_dict[key] = await restore_media(value, media_store=media_store) + return out_dict + return node + + +async def _restore_external(node: dict[str, object], media_store: MediaStore) -> dict[str, object]: + uri_value = node.get('uri') + if not isinstance(uri_value, str): + raise ValueError(f'externalized media marker missing string uri: {node!r}') + raw = await media_store.get(uri_value) + # Inverse of `_maybe_externalize_binary`: drop the marker keys, restore `data`, + # keep every other field the original part carried. + restored = {key: value for key, value in node.items() if key not in (_EXTERNAL_MARKER, 'uri')} + restored['data'] = base64.b64encode(raw).decode('ascii') + return restored diff --git a/pydantic_ai_harness/experimental/planning/README.md b/pydantic_ai_harness/experimental/planning/README.md index 2f110f4..750c9d6 100644 --- a/pydantic_ai_harness/experimental/planning/README.md +++ b/pydantic_ai_harness/experimental/planning/README.md @@ -1,5 +1,24 @@ # Planning +> [!WARNING] +> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may +> change or be removed in any release, without a deprecation period. Import it from the +> experimental path -- there is no top-level export: +> +> ```python +> from pydantic_ai_harness.experimental.planning import Planning +> ``` +> +> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all** +> harness experimental warnings with a single filter (no per-capability lines needed): +> +> ```python +> import warnings +> from pydantic_ai_harness.experimental import HarnessExperimentalWarning +> +> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning) +> ``` + Give an agent a structured, self-updating task plan -- without ever invalidating the prompt cache. ## The problem diff --git a/pydantic_ai_harness/experimental/step_persistence/README.md b/pydantic_ai_harness/experimental/step_persistence/README.md new file mode 100644 index 0000000..bcdd447 --- /dev/null +++ b/pydantic_ai_harness/experimental/step_persistence/README.md @@ -0,0 +1,505 @@ +# StepPersistence + +> [!WARNING] +> **Experimental.** `StepPersistence` and the `media` stores live under +> `pydantic_ai_harness.experimental` and may change or be removed in any release, without a +> deprecation period. Import them from the experimental path -- there is no top-level export: +> +> ```python +> from pydantic_ai_harness.experimental.step_persistence import StepPersistence +> from pydantic_ai_harness.experimental.media import S3MediaStore +> ``` +> +> Importing either package emits a `HarnessExperimentalWarning`. Silence **all** harness +> experimental warnings with a single filter (no per-capability lines needed): +> +> ```python +> import warnings +> from pydantic_ai_harness.experimental import HarnessExperimentalWarning +> +> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning) +> ``` + +`StepPersistence` records what an agent did at each boundary, separate from +whether the run can be safely resumed. It is the persistence substrate for +orchestrators that delegate to sub-agents (e.g. an AICA orchestrator spawning +a `code_librarian` to investigate one symbol, then continuing that delegate's +investigation with a follow-up question). + +It is not a full graph-state checkpoint. Capability-state restore, workspace +snapshots, and graph-node resume are out of scope and tracked separately +(see `pydantic-ai-harness` issues #149 and #196). + +## What it gives you + +1. **Append-only step events** -- every interesting boundary (run start/end, + model request, tool call, failure) appends a `StepEvent`. A run killed + mid-tool-call still leaves a usable event trail. +2. **Continuable snapshots** -- a `ContinuableSnapshot` is saved only at + boundaries where the message history is **provider-valid**: every + `ToolCallPart` has a matching `ToolReturnPart` or `RetryPromptPart`, with + no orphan, duplicate, or out-of-order returns. Pass the snapshot's + `messages` back to `Agent.run(message_history=...)` to continue or fork. +3. **Tool-effect ledger** -- every tool call's lifecycle (`started`, + `completed`, `failed`) is recorded against `(run_id, tool_call_id)`. + After a crash, a tool with a `started` record and no terminal update + should be treated as `unknown_after_crash`: the side effect may or may + not have happened. +4. **Lineage metadata** -- `conversation_id` (sequence) and `parent_run_id` + (hierarchy) are independent axes. See [Three-level identity](#three-level-identity). + +## Quick start + +```python +from pydantic_ai import Agent +from pydantic_ai_harness.experimental.step_persistence import StepPersistence, InMemoryStepStore + +store = InMemoryStepStore() +librarian = Agent( + 'openai:gpt-5', + capabilities=[StepPersistence(store=store, agent_name='code_librarian')], +) + +await librarian.run('Find ThinkingPartDelta and confirm the callable allowance') +``` + +That is the whole setup. **`run_id` is always per-`Agent.run` call**, +matching pydantic_ai's `RunContext.run_id`. For multi-turn logical +grouping use `conversation_id=` -- that is the pydantic_ai-native +primitive for it (see [Three-level identity](#three-level-identity)). + +`run_id` resolution per call: + +- **Explicit `run_id='libr-1'`** → used as-is for this one call. + Single-shot use cases (deterministic id for testing, replay, debugging, + a one-off scripted run). Reusing one capability instance with the same + explicit `run_id` across multiple `.run()` calls **raises `ValueError` + in `before_run`** -- the tool-effect ledger is keyed by + `(run_id, tool_call_id)` and providers reuse deterministic tool-call + ids, so a silent collision would erase the `unknown_after_crash` + signal. Use `conversation_id=` for multi-turn grouping instead. +- **`agent_name` set, `run_id` unset** → `'{agent_name}-{8-char-hex}'`, + freshly materialised in `for_run` per `.run()` call. Reusing one + capability instance across runs yields distinct ids + (`code_librarian-a3b2`, `code_librarian-c9d1`, …). This is the + recommended default for delegate capabilities. +- **Neither set** → `ctx.run_id` (pydantic_ai's auto-generated id) per + `.run()` call, falling back to a UUID4. + +The orchestrator pattern -- one logical agent serving many turns -- uses +`conversation_id`, not a shared `run_id`: + +```python +orchestrator = Agent( + 'openai:gpt-5', + capabilities=[StepPersistence(store=store, agent_name='orchestrator')], +) + +for turn in turns: + await orchestrator.run(turn, conversation_id='orch-conv') + +# All turns of this orchestrator, chronological: +records = await store.list_runs(conversation_id='orch-conv') +``` + +## Three-level identity + +The capability mirrors pydantic_ai's identity stack: + +| Concept | Definition | Granularity | +| --- | --- | --- | +| `conversation_id` | The dialogue. Resolved by pydantic_ai from the `conversation_id=` argument to `Agent.run`, or the most recent `conversation_id` on `message_history`, or a fresh UUID7. | sequence of runs | +| `run_id` | One `Agent.run` invocation. | one step in the sequence | +| `step_index` | Graph-node count *within* a run (`ctx.run_step`). | one node within one run | + +`StepEvent.conversation_id` and `RunRecord.conversation_id` are populated +from `ctx.conversation_id`. So three `.run()` calls sharing one +`conversation_id` produce three distinct `run_id`s, all queryable as a +group: + +```python +runs = await store.list_runs(conversation_id='conv-abc') # 3 records, chronological +``` + +## Continuing a delegate's investigation + +pydantic_ai already has `message_history=` for "carry on with this prior +context". `StepPersistence` does not introduce a parallel mechanism -- it +exposes one helper that loads the most recent provider-valid snapshot: + +```python +from pydantic_ai_harness.experimental.step_persistence import continue_run + +# Earlier: tag the first turn with a conversation id so the follow-up can find it. +await librarian.run( + 'Find ThinkingPartDelta and confirm the callable allowance', + conversation_id='libr-conv', +) + +# Later (possibly a different process): +prior_run = (await store.list_runs(conversation_id='libr-conv'))[-1].run_id +history = await continue_run(store, run_id=prior_run) +await librarian.run( + 'Read _apply_provider_details_delta and check the path', + message_history=history, + conversation_id='libr-conv', # keep the conversation grouping +) +``` + +`fork_run(store, run_id=...)` returns the same shape but is intended when +the caller wants a branched logical run from that snapshot point (the new +run gets a fresh `run_id` and probably a fresh `conversation_id`). + +### What "safe to continue from" means + +`continue_run` only returns the messages of the **latest provider-valid +snapshot** for that `run_id`. Snapshots are written at two boundaries: + +- after every `CallToolsNode` completes (all tool calls returned), and +- at `after_run`. + +A run that crashed mid-tool-call has events (`tool_call_started`) but no +snapshot for that point. `continue_run` returns the snapshot from the +**previous** safe boundary, not the failed step. + +## Run lineage -- `parent_run_id` + +`parent_run_id` is a lineage label, not a functional dependency. It does +two things: + +- Every `StepEvent` and `RunRecord` carries it, so you can filter / group. +- `store.list_runs(parent_run_id='orch-1')` returns every delegate run + pointing at that orchestrator. + +It is **auto-inferred for in-process delegation**: when an orchestrator's +tool synchronously calls a delegate's `Agent.run(...)`, the delegate's +`StepPersistence` picks up the orchestrator's `run_id` via a `ContextVar` +that the orchestrator's `wrap_run` set. No threading required: + +```python +orchestrator = Agent( + 'openai:gpt-5', + capabilities=[StepPersistence(store=store, agent_name='orchestrator')], +) +librarian = Agent( + 'openai:gpt-5', + capabilities=[StepPersistence(store=store, agent_name='code_librarian')], +) + +@orchestrator.tool_plain +async def ask_librarian(question: str) -> str: + result = await librarian.run(question) # parent_run_id auto-filled + return result.output + +# Tag the orchestrator turn so the lookup below can find its run_id. +await orchestrator.run( + 'Where is ThinkingPartDelta defined?', + conversation_id='orch-conv', +) + +# All librarian runs now point at the orchestrator's run_id: +orch_run_id = (await store.list_runs(conversation_id='orch-conv'))[-1].run_id +delegates = await store.list_runs(parent_run_id=orch_run_id) +``` + +Set `parent_run_id=` explicitly to override (e.g. cross-process delegation +where `ContextVar`s do not propagate). + +`parent_run_id` is **distinct from `conversation_id`**. The orchestrator +and delegate usually live in *different* conversations (the orchestrator +talks to a user; the delegate talks to itself). But they share a +parent-child link. + +## Inspecting a run tree + +`list_runs` returns matches sorted by `started_at` ascending across both +backends -- pick the most recent with `[-1]`. + +```python +# Every delegate of one orchestrator run (chronological) +delegates = await store.list_runs(parent_run_id='orch-3f2a') + +# Every run in one dialogue (multi-turn conversation across many .run() calls) +turns = await store.list_runs(conversation_id='conv-abc') +latest_turn = turns[-1] + +# Filters combine (AND): +focused = await store.list_runs( + parent_run_id='orch-3f2a', + conversation_id='libr-conv', +) + +# Detail per run: +events = await store.list_events(run_id=delegates[0].run_id) +snapshot = await store.latest_snapshot(run_id=delegates[0].run_id) +unresolved = await store.list_unresolved_tool_effects(run_id=delegates[0].run_id) +``` + +## Failure recovery + +```python +# An earlier delegate run died mid-investigation. +events = await store.list_events(run_id='libr-3f2a') +unresolved = await store.list_unresolved_tool_effects(run_id='libr-3f2a') +for record in unresolved: + # status == 'started' with no terminal update -- unknown_after_crash. + print(f'tool {record.tool_name} ({record.tool_call_id}) may or may not have run') + print(f' idempotency_key={record.idempotency_key} ' + f'effect_summary={record.effect_summary}') + +# Decide whether to resume or branch: +history = await continue_run(store, run_id='libr-3f2a') +# If the unresolved tools were read-only and safe to redo: +await librarian.run('continue investigating', message_history=history, + conversation_id='libr-conv') +# If side effects might have happened and the orchestrator wants a fresh attempt: +history = await fork_run(store, run_id='libr-3f2a') +# ... pass to a new delegate run with a different agent_name / conversation_id. +``` + +Side-effect deduplication is the orchestrator's responsibility. Tools that +write external state should annotate their in-flight `ToolEffectRecord` +via `annotate_tool_effect`: + +```python +from pydantic_ai_harness.experimental.step_persistence import annotate_tool_effect + +@orchestrator.tool +async def set_label(ctx: RunContext[Deps], issue: int, label: str) -> str: + await annotate_tool_effect( + store, + ctx, + idempotency_key=f'issue-{issue}::label::{label}', + effect_summary=f'set label {label!r} on issue #{issue}', + ) + await github.set_label(issue, label) # the actual side effect + return 'ok' +``` + +The helper reads the active `run_id` from the `StepPersistence` +`ContextVar` and `tool_call_id` / `tool_name` from `ctx`, then merges the +metadata into the prior record. `after_tool_execute` preserves both +fields when it writes the terminal `completed` / `failed` entry. + +## Backends + +- `InMemoryStepStore` -- process-local; great for tests. +- `FileStepStore(directory)` -- directory layout under `//`: + - `run.json` -- `RunRecord` (lineage) + - `events.jsonl` -- append-only `StepEvent`s + - `tool_effects.jsonl` -- append-only `ToolEffectRecord`s, scoped to this run + - `snapshots/{seq}.json` -- `ContinuableSnapshot`s, named by a per-run + monotonic counter (NOT `step_index`, which would collide when the + same `run_id` is reused across `Agent.run` calls -- `ctx.run_step` + resets to 0 each call). +- `SqliteStepStore(database='runs.db')` -- single SQLite file with tables + `runs`, `events`, `snapshots`, `tool_effects`, and a sibling `media` + table for externalized blobs (see [Persisting media](#persisting-media) + below). WAL mode is enabled; `tool_effects` upserts per + `(run_id, tool_call_id)` so the latest state wins; snapshots use + `AUTOINCREMENT seq` to mirror `FileStepStore._next_snapshot_seq`. + Pass `connection=` instead of `database=` to share a `sqlite3.Connection` + with the rest of your application; the connection must be opened with + `check_same_thread=False` because hook calls are dispatched onto a + worker thread. + +All three implement the same async `StepStore` protocol, so capability +hooks never block the event loop on the file/sqlite backends (I/O is +dispatched via `anyio.to_thread`). + +`FileStepStore` validates `run_id` against `[A-Za-z0-9_.-]{1,200}` (and +rejects `..`) to prevent path traversal -- callers passing user-controlled +IDs should still sanitise first. + +## Persisting media + +`BinaryContent` payloads (images, audio, documents, video) inline as +base64 inside a snapshot would balloon every file/row containing the +message. Both `FileStepStore` and `SqliteStepStore` externalize any +`BinaryContent.data` ≥ **64 KiB** through a configured `MediaStore`, +leaving a URI reference in the snapshot. Round-trip is transparent -- +`latest_snapshot(...).messages[*]` returns `BinaryContent` with the +original bytes. + +| StepStore | Default `media_store` | Where blobs live | +| ------------------- | --------------------------------------- | ------------------------------------- | +| `InMemoryStepStore` | _(not applicable)_ | bytes stay in the in-memory snapshot | +| `FileStepStore` | `DiskMediaStore(/media/)` | `/media/.bin` | +| `SqliteStepStore` | `SqliteMediaStore(database=)` | sibling `media` table in the same DB | + +Override the destination by passing your own `MediaStore`: + +```python +from pydantic_ai_harness.experimental.step_persistence import FileStepStore +from pydantic_ai_harness.experimental.media import S3MediaStore + +store = FileStepStore( + 'runs', + media_store=S3MediaStore( + bucket='my-bucket', + endpoint='https://.r2.cloudflarestorage.com', + region='auto', + access_key_id=..., + secret_access_key=..., + ), + media_threshold_bytes=64 * 1024, # raise/lower if you want +) +``` + +Opt out entirely (keep bytes inline in the snapshot JSON/row): + +```python +FileStepStore('runs', media_store=None) +SqliteStepStore(database='runs.db', media_store=None) +``` + +URIs are `media+sha256://`, content-addressed. The same blob written +through any `MediaStore` resolves the same way -- dedup is automatic and +moving the underlying storage is a one-line swap. The shipped +implementations are: + +- `DiskMediaStore(directory)` -- one file per blob at + `/.bin`. +- `SqliteMediaStore(database=...)` or `SqliteMediaStore(connection=...)` -- + one row per blob (`INSERT OR IGNORE` for content-addressed dedup). +- `S3MediaStore(bucket=, endpoint=, region=, access_key_id=, secret_access_key=)` + -- path-style URLs + handrolled SigV4. Compatible with AWS S3, Cloudflare + R2 (`region='auto'`), MinIO, and other S3-compatible providers. PUT/GET/HEAD + only -- no multipart, lifecycle, or listing in v1. + +### Exposing externalized bytes as URLs + +Each store accepts a `public_url=` callable that turns the canonical +`media+sha256://` URI into a URL the model can fetch directly. The +forthcoming `MediaExternalizer` capability will use this to swap +`BinaryContent` parts for `ImageUrl` / `AudioUrl` / etc. before the +model sees the message -- letting providers fetch big media over the wire +without re-encoding bytes into the request body. + +Static base URL (public R2 bucket, CDN): + +```python +from pydantic_ai_harness.experimental.media import S3MediaStore, make_static_public_url + +store = S3MediaStore( + bucket='my-bucket', + endpoint='https://.r2.cloudflarestorage.com', + region='auto', + access_key_id=..., secret_access_key=..., + key_prefix='media/', + public_url=make_static_public_url('https://pub-abc.r2.dev', key_prefix='media/'), +) +``` + +Presigned / rotating-signature URL -- pass any async callable that takes +`(uri, MediaContext)`: + +```python +from pydantic_ai_harness.experimental.media import MediaContext, S3MediaStore + +async def presign(uri: str, ctx: MediaContext) -> str: + key = 'media/' + uri.removeprefix('media+sha256://') + '.bin' + return await my_signer.generate(key, ttl=3600, content_type=ctx.media_type) + +store = S3MediaStore(..., public_url=presign) +``` + +### `MediaContext`, extensible per-operation bag + +Every `MediaStore` method (`put`, `get`, `exists`, `public_url`, +`get_metadata`) and both user-supplied callables (`PublicUrlResolver`, +`KeyStrategy`) accept a `MediaContext`: + +```python +@dataclass(frozen=True, kw_only=True) +class MediaContext: + media_type: str | None = None # e.g. 'image/png' + filename: str | None = None # original filename, when known + metadata: Mapping[str, str] = {} # user-supplied tags +``` + +All fields default; new fields are added non-breakingly as use cases +emerge. Pass what you have, ignore the rest. + +**Persistence by store.** `get_metadata(uri)` round-trips the +user-supplied `metadata` mapping on all three stores. `media_type` is +also persisted but is not part of what `get_metadata` returns (it is +stored for the byte payload itself, e.g. as the `Content-Type`). + +- `SqliteMediaStore` writes `metadata` to a JSON column and `media_type` + to a dedicated column +- `S3MediaStore` sends `metadata` as signed `x-amz-meta-*` headers + (ASCII alphanumeric + dash key names) and `media_type` as + `Content-Type`; `get_metadata` reads the `x-amz-meta-*` values back + from the HEAD response +- `DiskMediaStore` writes a sidecar JSON file (`.meta.json`) + alongside each blob, atomic via tmp + rename. Sidecars are absent + only when the put carried no metadata + +### `key_strategy` -- controlling the backend storage path + +Default is `.bin`. `DiskMediaStore` and `S3MediaStore` accept +overrides to fit existing layouts; `SqliteMediaStore` does not (its +primary key is the digest, so a user-chosen key would either break +dedup or be a no-op): + +```python +from pydantic_ai_harness.experimental.media import DiskMediaStore, MediaContext + +def by_media_type(uri: str, ctx: MediaContext) -> str: + digest = uri.removeprefix('media+sha256://') + ext = {'image/png': '.png', 'image/jpeg': '.jpg'}.get(ctx.media_type or '', '.bin') + return f'images/{digest}{ext}' + +store = DiskMediaStore('runs', key_strategy=by_media_type) +``` + +**Caveat**: if your strategy depends on `context.media_type` (e.g. to +pick an extension), `get(uri)` and `exists(uri)` won't find the blob +unless the same context is supplied at read time. For pure +path-organisation strategies (no context dependency) the constraint +doesn't apply. + +`DiskMediaStore` rejects strategies that produce absolute paths or paths +containing `..` segments, to prevent escaping the store directory. + +Separately, all three stores accept a `public_url=` resolver, useful +when a CDN, local HTTP server, or signed-URL service fronts the bytes. +Without it `public_url(...)` returns `None` (the model never sees a URL +unless a resolver is configured and it returns a string). + +pyai providers transparently download bytes from a URL when the target +model doesn't natively accept that URL type, so emitting a URL is +always safe -- you only ever lose wire savings, never correctness. + +> **Note on the future `MediaExternalizer` capability.** When it lands, +> the composition will be +> `Agent(capabilities=[MediaExternalizer(store), StepPersistence(...)])` +> and `StepPersistence` will see already-URL-ified messages -- the +> externalize walk becomes a no-op. The existing API does not change. + +### Persisting unsupported backends + +DynamoDB, Postgres, Redis, GCS, and other backends are out of scope for +this release. Write your own `StepStore` (≈ ten methods on a Protocol) or +your own `MediaStore` (three methods) and pass it via `store=` / +`media_store=`. Please open an issue if you ship one -- we want to feed +the eventual shared adapter layer with N≥3 real implementations before +abstracting. + +## What this capability does **not** do + +- It does not restore capability per-run state, graph-node state, retry + counters, or in-flight streaming responses. +- It does not deduplicate replayed side effects automatically. Tools that + write artifacts, labels, PRs, or external state should call + `annotate_tool_effect(store, ctx, ...)` (see [Failure recovery](#failure-recovery)) + so the orchestrator can decide whether replay is safe. +- It does not clean up old snapshots/events. Retention is the caller's + responsibility. +- It does not emit OpenTelemetry spans. pydantic_ai's `Instrumentation` + capability already spans `agent run` / `chat` / `running tool` and + populates `gen_ai.agent.name`, `gen_ai.agent.call.id`, + `gen_ai.conversation.id` via baggage. A future change may add + step-persistence attributes to the active span; that is tracked as a + follow-up issue. diff --git a/pydantic_ai_harness/experimental/step_persistence/__init__.py b/pydantic_ai_harness/experimental/step_persistence/__init__.py new file mode 100644 index 0000000..3559889 --- /dev/null +++ b/pydantic_ai_harness/experimental/step_persistence/__init__.py @@ -0,0 +1,44 @@ +"""Step-event persistence: append-only event log, continuable snapshots, tool-effect ledger.""" + +from pydantic_ai_harness.experimental._warn import warn_experimental +from pydantic_ai_harness.experimental.step_persistence._capability import StepPersistence +from pydantic_ai_harness.experimental.step_persistence._helpers import ( + annotate_tool_effect, + continue_run, + fork_run, + is_provider_valid, +) +from pydantic_ai_harness.experimental.step_persistence._store import ( + FileStepStore, + InMemoryStepStore, + SqliteStepStore, + StepStore, +) +from pydantic_ai_harness.experimental.step_persistence._types import ( + ContinuableSnapshot, + EventKind, + RunRecord, + StepEvent, + ToolEffectRecord, + ToolEffectStatus, +) + +warn_experimental('step_persistence') + +__all__ = [ + 'ContinuableSnapshot', + 'EventKind', + 'FileStepStore', + 'InMemoryStepStore', + 'RunRecord', + 'SqliteStepStore', + 'StepEvent', + 'StepPersistence', + 'StepStore', + 'ToolEffectRecord', + 'ToolEffectStatus', + 'annotate_tool_effect', + 'continue_run', + 'fork_run', + 'is_provider_valid', +] diff --git a/pydantic_ai_harness/experimental/step_persistence/_capability.py b/pydantic_ai_harness/experimental/step_persistence/_capability.py new file mode 100644 index 0000000..d674613 --- /dev/null +++ b/pydantic_ai_harness/experimental/step_persistence/_capability.py @@ -0,0 +1,426 @@ +"""StepPersistence capability: append-only event log + continuable snapshots.""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from datetime import datetime, timezone +from typing import Any +from uuid import uuid4 + +from pydantic_ai import CallToolsNode +from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.capabilities.abstract import AgentNode, NodeResult, WrapRunHandler +from pydantic_ai.messages import ModelResponse, ToolCallPart +from pydantic_ai.models import ModelRequestContext +from pydantic_ai.run import AgentRunResult +from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition + +from pydantic_ai_harness.experimental.step_persistence._context import current_run_id, snapshot_saved +from pydantic_ai_harness.experimental.step_persistence._helpers import is_provider_valid +from pydantic_ai_harness.experimental.step_persistence._store import InMemoryStepStore, StepStore +from pydantic_ai_harness.experimental.step_persistence._types import ( + ContinuableSnapshot, + EventKind, + RunRecord, + StepEvent, + ToolEffectRecord, +) + + +def _empty_metadata() -> dict[str, str]: + return {} + + +@dataclass +class StepPersistence(AbstractCapability[AgentDepsT]): + """Append-only step log + continuable snapshots + tool-effect ledger. + + The capability emits a `StepEvent` at every interesting boundary + (run/model-request/tool-call start, completion, failure), records a + `ToolEffectRecord` per tool call so the orchestrator can decide whether + replay is safe, and saves a `ContinuableSnapshot` at every + provider-valid boundary -- the end of each `CallToolsNode` -- plus a + fallback save at `after_run` if the run reached no such boundary. + + A run that crashes between `before_tool_execute` and `after_tool_execute` + leaves a visible event trail and a `started` tool-effect record, but no + new continuable snapshot -- the latest snapshot reflects the last + provider-valid state. + + ```python + from pydantic_ai import Agent + from pydantic_ai_harness.experimental.step_persistence import StepPersistence, InMemoryStepStore + + store = InMemoryStepStore() + librarian = Agent( + 'openai:gpt-5', + capabilities=[StepPersistence(store=store, agent_name='code_librarian')], + ) + await librarian.run('Find ThinkingPartDelta and confirm the callable allowance') + ``` + + Use `continue_run(store, run_id=...)` / `fork_run(store, run_id=...)` + to load a prior snapshot, then pass the result to + `Agent.run(..., message_history=...)`. + """ + + store: StepStore = field(default_factory=InMemoryStepStore) + """Backend that records events, snapshots, and tool effects.""" + + agent_name: str | None = None + """Logical agent name (e.g. `code_librarian`, `reproducer`). + + Used as a stable prefix for the auto-derived `run_id` so store + inspection shows readable IDs like `code_librarian-a3b2`. + """ + + run_id: str | None = None + """Identifier for this one `Agent.run` call. + + `run_id` is per-call, matching `pydantic_ai.RunContext.run_id`. For + multi-turn logical grouping use `conversation_id` on `Agent.run(...)` -- + that is the pyai-native primitive for it. + + Resolution order (materialised in `for_run`): + + 1. **Explicit value** → used as-is. Single-shot use cases: + deterministic id for testing, replay, debugging. Reusing the + capability across multiple `.run()` calls with the same explicit + `run_id` raises `ValueError` in `before_run` -- the tool-effect + ledger keys on `(run_id, tool_call_id)` and providers reuse + deterministic tool-call ids, so a silent collision would erase + the `unknown_after_crash` signal. Use `conversation_id=` on + `Agent.run` for multi-turn grouping. + 2. **`agent_name` set, `run_id` unset** → `{agent_name}-{short-uuid}`, + freshly materialised per `.run()`. Reusing the capability instance + yields distinct ids. Recommended default for delegate capabilities. + 3. **Neither set** → `ctx.run_id` per `.run()`, falling back to UUID4. + """ + + parent_run_id: str | None = None + """Run that spawned this one. + + Auto-inferred from the enclosing `StepPersistence` `wrap_run` scope -- + when an orchestrator's tool synchronously calls a delegate's + `Agent.run(...)`, the delegate picks up the orchestrator's `run_id` + here without manual threading. Set explicitly to override (e.g. for + cross-process delegation where `ContextVar`s do not propagate). + """ + + metadata: dict[str, str] = field(default_factory=_empty_metadata) + """Free-form metadata stored on the `RunRecord` and on each event.""" + + @classmethod + def from_spec(cls, *args: Any, **kwargs: Any) -> StepPersistence[Any]: + """Construct from a serialised spec. + + Supports `backend='memory'` (default), `backend='file'` (with + `directory`), or `backend='sqlite'` (with `database`). Raises + `ValueError` for any other `backend` value -- silently falling + back to in-memory storage would turn a typo into accidental + non-durability. + """ + backend = kwargs.pop('backend', 'memory') + if backend == 'memory': + return cls(store=InMemoryStepStore(), **kwargs) + if backend == 'file': + from pydantic_ai_harness.experimental.step_persistence._store import FileStepStore + + directory = kwargs.pop('directory', '.step-persistence') + return cls(store=FileStepStore(directory), **kwargs) + if backend == 'sqlite': + from pydantic_ai_harness.experimental.step_persistence._store import SqliteStepStore + + database = kwargs.pop('database', '.step-persistence.db') + return cls(store=SqliteStepStore(database=database), **kwargs) + raise ValueError(f'unknown backend {backend!r}; expected `memory`, `file`, or `sqlite`') + + async def for_run(self, ctx: RunContext[AgentDepsT]) -> AbstractCapability[AgentDepsT]: + """Materialise `run_id` and `parent_run_id` for this `Agent.run` call. + + Reads the contextvar set by any enclosing `StepPersistence.wrap_run` + before the local run overwrites it, so a delegate's `parent_run_id` + ends up pointing at its orchestrator's `run_id`. + + A separate `ContextVar` is needed because pydantic_ai's own + cross-run signals (`RUN_ID_BAGGAGE_KEY` via OTel baggage, + `RunContext.run_id`, and `_CURRENT_RUN_CONTEXT`) are single-slot: + the inner `Instrumentation.wrap_run` overwrites them before any + nested capability sees the parent. The harness-local contextvar + lets us snapshot the parent here, *before* the local `wrap_run` + rebinds it. + """ + inferred_parent = self.parent_run_id if self.parent_run_id is not None else current_run_id.get() + resolved_run_id = self.run_id or self._derive_run_id(ctx) + if resolved_run_id == self.run_id and inferred_parent == self.parent_run_id: + return self + return replace(self, run_id=resolved_run_id, parent_run_id=inferred_parent) + + def _derive_run_id(self, ctx: RunContext[AgentDepsT]) -> str: + if self.agent_name is not None: + return f'{self.agent_name}-{uuid4().hex[:8]}' + return ctx.run_id or str(uuid4()) + + def _effective_run_id(self, ctx: RunContext[AgentDepsT]) -> str: + if self.run_id is not None: + return self.run_id + return ctx.run_id or str(uuid4()) + + def _make_event( + self, + ctx: RunContext[AgentDepsT], + *, + kind: EventKind, + tool_call_id: str | None = None, + tool_name: str | None = None, + error: str | None = None, + ) -> StepEvent: + return StepEvent( + run_id=self._effective_run_id(ctx), + kind=kind, + step_index=ctx.run_step, + conversation_id=ctx.conversation_id, + parent_run_id=self.parent_run_id, + agent_name=self.agent_name, + tool_call_id=tool_call_id, + tool_name=tool_name, + error=error, + metadata=dict(self.metadata), + ) + + async def wrap_run( + self, + ctx: RunContext[AgentDepsT], + *, + handler: WrapRunHandler, + ) -> AgentRunResult[Any]: + """Push this run's id onto the contextvar so nested delegates can read it.""" + token = current_run_id.set(self._effective_run_id(ctx)) + saved_token = snapshot_saved.set(False) + try: + return await handler() + finally: + snapshot_saved.reset(saved_token) + current_run_id.reset(token) + + async def before_run(self, ctx: RunContext[AgentDepsT]) -> None: + """Register run lineage and emit `run_started`. + + When the caller pinned an explicit `run_id`, reject reuse -- the + tool-effect ledger keys on `(run_id, tool_call_id)` and providers + reuse deterministic tool-call ids, so a second `Agent.run` with + the same explicit `run_id` would silently collide. The auto-derived + cases cannot trigger this check because each call materialises a + fresh id in `for_run`. + """ + run_id = self._effective_run_id(ctx) + if self.run_id is not None and await self.store.get_run(run_id=run_id) is not None: + raise ValueError( + f'StepPersistence: run_id {run_id!r} is already in the store. ' + 'Explicit `run_id` is single-shot; pass `conversation_id=` to ' + '`Agent.run` for multi-turn grouping instead.' + ) + await self.store.register_run( + RunRecord( + run_id=run_id, + conversation_id=ctx.conversation_id, + parent_run_id=self.parent_run_id, + agent_name=self.agent_name, + metadata=dict(self.metadata), + ) + ) + await self.store.append_event(self._make_event(ctx, kind='run_started')) + + async def after_run( + self, + ctx: RunContext[AgentDepsT], + *, + result: AgentRunResult[Any], + ) -> AgentRunResult[Any]: + """Emit `run_completed`, saving a final snapshot only as a fallback. + + The terminal `CallToolsNode` already saved the final provider-valid + snapshot via `after_node_run`, carrying the correct `step_index`. By + `after_run`, `ctx.run_step` is reset to 0, so re-saving here would both + duplicate the tail and stamp a misleading `step_index`. We only save + when the run produced no snapshot at all (no provider-valid node + boundary was reached), as a last-resort capture of the final state. + """ + if not snapshot_saved.get(): + messages = result.all_messages() + if is_provider_valid(messages): + await self.store.save_snapshot( + ContinuableSnapshot( + run_id=self._effective_run_id(ctx), + step_index=ctx.run_step, + messages=list(messages), + conversation_id=ctx.conversation_id, + parent_run_id=self.parent_run_id, + agent_name=self.agent_name, + ) + ) + await self.store.append_event(self._make_event(ctx, kind='run_completed')) + return result + + async def on_run_error( + self, + ctx: RunContext[AgentDepsT], + *, + error: BaseException, + ) -> AgentRunResult[Any]: + """Emit `run_failed` so a killed run leaves a visible event trail.""" + await self.store.append_event(self._make_event(ctx, kind='run_failed', error=repr(error))) + raise error + + async def before_model_request( + self, + ctx: RunContext[AgentDepsT], + request_context: ModelRequestContext, + ) -> ModelRequestContext: + await self.store.append_event(self._make_event(ctx, kind='model_request_started')) + return request_context + + async def after_model_request( + self, + ctx: RunContext[AgentDepsT], + *, + request_context: ModelRequestContext, + response: ModelResponse, + ) -> ModelResponse: + await self.store.append_event(self._make_event(ctx, kind='model_request_completed')) + return response + + async def on_model_request_error( + self, + ctx: RunContext[AgentDepsT], + *, + request_context: ModelRequestContext, + error: Exception, + ) -> ModelResponse: + await self.store.append_event(self._make_event(ctx, kind='model_request_failed', error=repr(error))) + raise error + + async def before_tool_execute( + self, + ctx: RunContext[AgentDepsT], + *, + call: ToolCallPart, + tool_def: ToolDefinition, + args: dict[str, Any], + ) -> dict[str, Any]: + run_id = self._effective_run_id(ctx) + await self.store.record_tool_effect( + ToolEffectRecord( + tool_call_id=call.tool_call_id, + tool_name=tool_def.name, + run_id=run_id, + status='started', + ) + ) + await self.store.append_event( + self._make_event( + ctx, + kind='tool_call_started', + tool_call_id=call.tool_call_id, + tool_name=tool_def.name, + ) + ) + return args + + async def after_tool_execute( + self, + ctx: RunContext[AgentDepsT], + *, + call: ToolCallPart, + tool_def: ToolDefinition, + args: dict[str, Any], + result: Any, + ) -> Any: + run_id = self._effective_run_id(ctx) + prior = await self.store.get_tool_effect(run_id=run_id, tool_call_id=call.tool_call_id) + await self.store.record_tool_effect( + ToolEffectRecord( + tool_call_id=call.tool_call_id, + tool_name=tool_def.name, + run_id=run_id, + status='completed', + started_at=prior.started_at if prior is not None else datetime.now(timezone.utc), + ended_at=datetime.now(timezone.utc), + idempotency_key=prior.idempotency_key if prior is not None else None, + effect_summary=prior.effect_summary if prior is not None else None, + ) + ) + await self.store.append_event( + self._make_event( + ctx, + kind='tool_call_completed', + tool_call_id=call.tool_call_id, + tool_name=tool_def.name, + ) + ) + return result + + async def on_tool_execute_error( + self, + ctx: RunContext[AgentDepsT], + *, + call: ToolCallPart, + tool_def: ToolDefinition, + args: dict[str, Any], + error: Exception, + ) -> Any: + run_id = self._effective_run_id(ctx) + prior = await self.store.get_tool_effect(run_id=run_id, tool_call_id=call.tool_call_id) + prior_summary = prior.effect_summary if prior is not None else None + await self.store.record_tool_effect( + ToolEffectRecord( + tool_call_id=call.tool_call_id, + tool_name=tool_def.name, + run_id=run_id, + status='failed', + started_at=prior.started_at if prior is not None else datetime.now(timezone.utc), + ended_at=datetime.now(timezone.utc), + idempotency_key=prior.idempotency_key if prior is not None else None, + effect_summary=prior_summary if prior_summary is not None else repr(error), + ) + ) + await self.store.append_event( + self._make_event( + ctx, + kind='tool_call_failed', + tool_call_id=call.tool_call_id, + tool_name=tool_def.name, + error=repr(error), + ) + ) + raise error + + async def after_node_run( + self, + ctx: RunContext[AgentDepsT], + *, + node: AgentNode[AgentDepsT], + result: NodeResult[AgentDepsT], + ) -> NodeResult[AgentDepsT]: + """Save a mid-run continuable snapshot after `CallToolsNode` succeeds. + + At that boundary every tool call from the preceding `ModelRequestNode` + has a matching tool return, so the history is provider-valid. + Snapshots are filtered through `is_provider_valid` defensively in case + a custom node reshapes history. + """ + if isinstance(node, CallToolsNode): + messages = list(ctx.messages) + if is_provider_valid(messages): + await self.store.save_snapshot( + ContinuableSnapshot( + run_id=self._effective_run_id(ctx), + step_index=ctx.run_step, + messages=messages, + conversation_id=ctx.conversation_id, + parent_run_id=self.parent_run_id, + agent_name=self.agent_name, + ) + ) + snapshot_saved.set(True) + return result diff --git a/pydantic_ai_harness/experimental/step_persistence/_context.py b/pydantic_ai_harness/experimental/step_persistence/_context.py new file mode 100644 index 0000000..7649820 --- /dev/null +++ b/pydantic_ai_harness/experimental/step_persistence/_context.py @@ -0,0 +1,34 @@ +"""Shared async-context state for `StepPersistence` cross-capability coordination.""" + +from __future__ import annotations + +from contextvars import ContextVar + +current_run_id: ContextVar[str | None] = ContextVar( + 'pydantic_ai_harness.experimental.step_persistence.current_run_id', + default=None, +) +"""Async-context-local pointer to the active `StepPersistence` `run_id`. + +Set by `StepPersistence.wrap_run` for the duration of a run; read by a +nested capability's `for_run` to auto-fill `parent_run_id`, and by +`annotate_tool_effect` to find the in-flight tool's run scope. + +Module-level rather than a class attribute so the helpers in `_helpers.py` +and the capability in `_capability.py` can share it without a circular +import. +""" + +snapshot_saved: ContextVar[bool] = ContextVar( + 'pydantic_ai_harness.experimental.step_persistence.snapshot_saved', + default=False, +) +"""Async-context-local flag: did `after_node_run` already save a snapshot this run? + +Set `False` in `wrap_run`, flipped `True` whenever `after_node_run` saves a +`CallToolsNode` snapshot. `after_run` reads it to skip a redundant terminal +snapshot -- the final `CallToolsNode` already captured the provider-valid tail +with the correct `step_index`, whereas `after_run` runs with `ctx.run_step` +reset to 0. Task-isolated like `current_run_id`, so concurrent runs don't +interfere. +""" diff --git a/pydantic_ai_harness/experimental/step_persistence/_helpers.py b/pydantic_ai_harness/experimental/step_persistence/_helpers.py new file mode 100644 index 0000000..6b02105 --- /dev/null +++ b/pydantic_ai_harness/experimental/step_persistence/_helpers.py @@ -0,0 +1,119 @@ +"""Helpers for continuation, forking, provider-validity checks, and tool-effect annotation.""" + +from __future__ import annotations + +from typing import Any + +from pydantic_ai.messages import ( + ModelMessage, + ModelResponse, + RetryPromptPart, + ToolCallPart, + ToolReturnPart, +) +from pydantic_ai.tools import RunContext + +from pydantic_ai_harness.experimental.step_persistence._context import current_run_id +from pydantic_ai_harness.experimental.step_persistence._store import StepStore +from pydantic_ai_harness.experimental.step_persistence._types import ToolEffectRecord + + +def is_provider_valid(messages: list[ModelMessage]) -> bool: + """Return True when `messages` can be safely passed to `Agent.run(message_history=...)`. + + A history is provider-valid when: + + 1. Every `ToolCallPart` has a matching `ToolReturnPart` or + tool-bound `RetryPromptPart` later in the conversation, and + 2. Every tool return / tool-bound retry resolves a currently-open + tool call (no orphans, duplicates, or out-of-order returns). + + A `RetryPromptPart` with `tool_name is None` is an output-validation + retry -- providers map it as a regular user message, not a tool result, + so it does not need to resolve an open call. + """ + open_calls: set[str] = set() + for msg in messages: + if isinstance(msg, ModelResponse): + for part in msg.parts: + if isinstance(part, ToolCallPart): + open_calls.add(part.tool_call_id) + else: + for part in msg.parts: + if isinstance(part, ToolReturnPart): + if part.tool_call_id not in open_calls: + return False + open_calls.discard(part.tool_call_id) + elif isinstance(part, RetryPromptPart) and part.tool_name is not None: + if part.tool_call_id not in open_calls: + return False + open_calls.discard(part.tool_call_id) + return not open_calls + + +async def continue_run(store: StepStore, *, run_id: str) -> list[ModelMessage]: + """Load the latest continuable snapshot for `run_id` as a message history. + + Pass the return value to `Agent.run(message_history=...)` to continue + a delegate's prior investigation instead of starting fresh. + + Raises `LookupError` if no continuable snapshot exists for `run_id` -- the + run may have crashed mid-tool-call, in which case there is event-log data + but no safe resume point. + """ + snapshot = await store.latest_snapshot(run_id=run_id) + if snapshot is None: + raise LookupError(f'no continuable snapshot for run_id {run_id!r}') + return list(snapshot.messages) + + +async def fork_run(store: StepStore, *, run_id: str) -> list[ModelMessage]: + """Return a copy of the latest snapshot's messages, intended for a new logical run. + + Semantically identical to `continue_run` at the data layer; the + distinction is in how the caller treats the returned history (new + `run_id`, new lineage entry, branching off prior context). + """ + return await continue_run(store, run_id=run_id) + + +async def annotate_tool_effect( + store: StepStore, + ctx: RunContext[Any], + *, + idempotency_key: str | None = None, + effect_summary: str | None = None, +) -> None: + """Attach `idempotency_key` and / or `effect_summary` to the in-flight tool's effect record. + + Call from inside a tool body when the tool writes external state + (artifacts, labels, PRs, network mutations) so an orchestrator + inspecting `list_unresolved_tool_effects` after a crash can tell + whether replay is safe. + + Resolves the active run from the `StepPersistence` `ContextVar` and + the tool identity from `ctx.tool_call_id` / `ctx.tool_name`. No-op + when called outside a step-persistence-wrapped tool call. The + capability's `after_tool_execute` preserves these fields when it + writes the terminal `completed` / `failed` record. + """ + run_id = current_run_id.get() + tool_call_id = ctx.tool_call_id + tool_name = ctx.tool_name + if run_id is None or tool_call_id is None or tool_name is None: + return + prior = await store.get_tool_effect(run_id=run_id, tool_call_id=tool_call_id) + if prior is None: + return + await store.record_tool_effect( + ToolEffectRecord( + tool_call_id=tool_call_id, + tool_name=tool_name, + run_id=run_id, + status=prior.status, + started_at=prior.started_at, + ended_at=prior.ended_at, + idempotency_key=idempotency_key if idempotency_key is not None else prior.idempotency_key, + effect_summary=effect_summary if effect_summary is not None else prior.effect_summary, + ) + ) diff --git a/pydantic_ai_harness/experimental/step_persistence/_store.py b/pydantic_ai_harness/experimental/step_persistence/_store.py new file mode 100644 index 0000000..6019225 --- /dev/null +++ b/pydantic_ai_harness/experimental/step_persistence/_store.py @@ -0,0 +1,1032 @@ +"""Storage backends for step events, snapshots, and tool-effect records.""" + +from __future__ import annotations + +import json +import re +import sqlite3 +from collections import defaultdict +from datetime import datetime +from pathlib import Path +from typing import Literal, Protocol, runtime_checkable + +import anyio.to_thread +from pydantic import TypeAdapter +from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter + +from pydantic_ai_harness.experimental.media import ( + DiskMediaStore, + MediaStore, + SqliteMediaStore, + externalize_media, + restore_media, +) +from pydantic_ai_harness.experimental.step_persistence._types import ( + ContinuableSnapshot, + EventKind, + RunRecord, + StepEvent, + ToolEffectRecord, + ToolEffectStatus, +) + +_DEFAULT_MEDIA_THRESHOLD_BYTES = 64 * 1024 +_AutoMedia = Literal['auto'] + +_VALID_ID_RE = re.compile(r'^[A-Za-z0-9_.-]{1,200}$') + +_EVENT_KIND_TABLE: dict[str, EventKind] = { + 'run_started': 'run_started', + 'run_completed': 'run_completed', + 'run_failed': 'run_failed', + 'model_request_started': 'model_request_started', + 'model_request_completed': 'model_request_completed', + 'model_request_failed': 'model_request_failed', + 'tool_call_started': 'tool_call_started', + 'tool_call_completed': 'tool_call_completed', + 'tool_call_failed': 'tool_call_failed', +} + +_TOOL_STATUS_TABLE: dict[str, ToolEffectStatus] = { + 'started': 'started', + 'completed': 'completed', + 'failed': 'failed', +} + + +def _validate_id(value: str, *, field: str) -> None: + r"""Reject identifiers that would let a caller escape the store directory. + + `FileStepStore` interpolates `run_id` into a path. Allowing `..`, `/`, + `\`, empty input, or oversized values would enable path traversal. + """ + if not _VALID_ID_RE.fullmatch(value) or '..' in value: + raise ValueError(f'invalid {field}: {value!r}') + + +@runtime_checkable +class StepStore(Protocol): + """Async protocol for step-persistence backends. + + All methods are async so the same protocol covers in-memory stores and + file/database stores that would otherwise block the event loop. + """ + + async def register_run(self, record: RunRecord) -> None: ... # pragma: no cover + + async def get_run(self, *, run_id: str) -> RunRecord | None: ... # pragma: no cover + + async def list_runs( + self, + *, + parent_run_id: str | None = None, + conversation_id: str | None = None, + ) -> list[RunRecord]: + """Return matching `RunRecord`s sorted by `started_at` ascending. + + Both filters are optional and AND-combine when both are set. The + chronological ordering is part of the protocol -- callers may pick + the most recent run with `[-1]`. + """ + ... # pragma: no cover + + async def append_event(self, event: StepEvent) -> None: ... # pragma: no cover + + async def list_events(self, *, run_id: str) -> list[StepEvent]: ... # pragma: no cover + + async def save_snapshot(self, snapshot: ContinuableSnapshot) -> None: ... # pragma: no cover + + async def latest_snapshot(self, *, run_id: str) -> ContinuableSnapshot | None: ... # pragma: no cover + + async def record_tool_effect(self, record: ToolEffectRecord) -> None: ... # pragma: no cover + + async def get_tool_effect( + self, *, run_id: str, tool_call_id: str + ) -> ToolEffectRecord | None: ... # pragma: no cover + + async def list_unresolved_tool_effects(self, *, run_id: str) -> list[ToolEffectRecord]: ... # pragma: no cover + + +class InMemoryStepStore: + """Process-local store, suitable for tests and single-process orchestrators.""" + + def __init__(self) -> None: + self._runs: dict[str, RunRecord] = {} + self._events: dict[str, list[StepEvent]] = defaultdict(list) + self._snapshots: dict[str, list[ContinuableSnapshot]] = defaultdict(list) + self._tool_effects: dict[tuple[str, str], ToolEffectRecord] = {} + + async def register_run(self, record: RunRecord) -> None: + self._runs[record.run_id] = record + + async def get_run(self, *, run_id: str) -> RunRecord | None: + return self._runs.get(run_id) + + async def list_runs( + self, + *, + parent_run_id: str | None = None, + conversation_id: str | None = None, + ) -> list[RunRecord]: + records = [ + r + for r in self._runs.values() + if (parent_run_id is None or r.parent_run_id == parent_run_id) + and (conversation_id is None or r.conversation_id == conversation_id) + ] + return sorted(records, key=lambda r: r.started_at) + + async def append_event(self, event: StepEvent) -> None: + self._events[event.run_id].append(event) + + async def list_events(self, *, run_id: str) -> list[StepEvent]: + return list(self._events.get(run_id, ())) + + async def save_snapshot(self, snapshot: ContinuableSnapshot) -> None: + self._snapshots[snapshot.run_id].append(snapshot) + + async def latest_snapshot(self, *, run_id: str) -> ContinuableSnapshot | None: + snaps = self._snapshots.get(run_id) + if not snaps: + return None + return snaps[-1] + + async def record_tool_effect(self, record: ToolEffectRecord) -> None: + self._tool_effects[(record.run_id, record.tool_call_id)] = record + + async def get_tool_effect(self, *, run_id: str, tool_call_id: str) -> ToolEffectRecord | None: + return self._tool_effects.get((run_id, tool_call_id)) + + async def list_unresolved_tool_effects(self, *, run_id: str) -> list[ToolEffectRecord]: + return [r for r in self._tool_effects.values() if r.run_id == run_id and r.status == 'started'] + + +def _opt_str(value: object) -> str | None: + if value is None: + return None + if not isinstance(value, str): + raise ValueError(f'expected str|None, got {type(value).__name__}') + return value + + +_STR_STR_DICT_ADAPTER: TypeAdapter[dict[str, str]] = TypeAdapter(dict[str, str]) +_OBJECT_DICT_ADAPTER: TypeAdapter[dict[str, object]] = TypeAdapter(dict[str, object]) + + +def _str_str_dict(value: object) -> dict[str, str]: + if value is None: + return {} + return _STR_STR_DICT_ADAPTER.validate_python(value) + + +def _load_json_object(text: str) -> dict[str, object]: + return _OBJECT_DICT_ADAPTER.validate_json(text) + + +def _event_to_dict(event: StepEvent) -> dict[str, object]: + return { + 'run_id': event.run_id, + 'kind': event.kind, + 'step_index': event.step_index, + 'timestamp': event.timestamp.isoformat(), + 'conversation_id': event.conversation_id, + 'parent_run_id': event.parent_run_id, + 'agent_name': event.agent_name, + 'tool_call_id': event.tool_call_id, + 'tool_name': event.tool_name, + 'error': event.error, + 'metadata': dict(event.metadata), + } + + +def _event_from_dict(data: dict[str, object]) -> StepEvent: + run_id = data['run_id'] + kind_raw = data['kind'] + step_raw = data['step_index'] + timestamp_raw = data['timestamp'] + if not ( + isinstance(run_id, str) + and isinstance(kind_raw, str) + and isinstance(step_raw, int) + and isinstance(timestamp_raw, str) + ): + raise ValueError('event payload has wrong types') + kind = _EVENT_KIND_TABLE.get(kind_raw) + if kind is None: + raise ValueError(f'unknown event kind: {kind_raw!r}') + return StepEvent( + run_id=run_id, + kind=kind, + step_index=step_raw, + timestamp=datetime.fromisoformat(timestamp_raw), + conversation_id=_opt_str(data.get('conversation_id')), + parent_run_id=_opt_str(data.get('parent_run_id')), + agent_name=_opt_str(data.get('agent_name')), + tool_call_id=_opt_str(data.get('tool_call_id')), + tool_name=_opt_str(data.get('tool_name')), + error=_opt_str(data.get('error')), + metadata=_str_str_dict(data.get('metadata')), + ) + + +def _run_to_dict(record: RunRecord) -> dict[str, object]: + return { + 'run_id': record.run_id, + 'conversation_id': record.conversation_id, + 'parent_run_id': record.parent_run_id, + 'agent_name': record.agent_name, + 'metadata': dict(record.metadata), + 'started_at': record.started_at.isoformat(), + } + + +def _run_from_dict(data: dict[str, object]) -> RunRecord: + run_id = data['run_id'] + started_at_raw = data['started_at'] + if not (isinstance(run_id, str) and isinstance(started_at_raw, str)): + raise ValueError('run record has wrong types') + return RunRecord( + run_id=run_id, + conversation_id=_opt_str(data.get('conversation_id')), + parent_run_id=_opt_str(data.get('parent_run_id')), + agent_name=_opt_str(data.get('agent_name')), + metadata=_str_str_dict(data.get('metadata')), + started_at=datetime.fromisoformat(started_at_raw), + ) + + +def _tool_effect_to_dict(record: ToolEffectRecord) -> dict[str, object]: + return { + 'tool_call_id': record.tool_call_id, + 'tool_name': record.tool_name, + 'run_id': record.run_id, + 'status': record.status, + 'started_at': record.started_at.isoformat(), + 'ended_at': record.ended_at.isoformat() if record.ended_at is not None else None, + 'idempotency_key': record.idempotency_key, + 'effect_summary': record.effect_summary, + } + + +def _tool_effect_from_dict(data: dict[str, object]) -> ToolEffectRecord: + tool_call_id = data['tool_call_id'] + tool_name = data['tool_name'] + run_id = data['run_id'] + status_raw = data['status'] + started_at_raw = data['started_at'] + if not ( + isinstance(tool_call_id, str) + and isinstance(tool_name, str) + and isinstance(run_id, str) + and isinstance(status_raw, str) + and isinstance(started_at_raw, str) + ): + raise ValueError('tool effect record has wrong types') + status = _TOOL_STATUS_TABLE.get(status_raw) + if status is None: + raise ValueError(f'unknown tool effect status: {status_raw!r}') + ended_at_raw = data.get('ended_at') + ended_at = datetime.fromisoformat(ended_at_raw) if isinstance(ended_at_raw, str) else None + return ToolEffectRecord( + tool_call_id=tool_call_id, + tool_name=tool_name, + run_id=run_id, + status=status, + started_at=datetime.fromisoformat(started_at_raw), + ended_at=ended_at, + idempotency_key=_opt_str(data.get('idempotency_key')), + effect_summary=_opt_str(data.get('effect_summary')), + ) + + +class FileStepStore: + """Filesystem-backed store using JSONL for events/tool effects and JSON for snapshots. + + Layout under the configured root directory: + + ```text + / + / + run.json + events.jsonl + tool_effects.jsonl + snapshots/ + .json + ``` + + Snapshot filenames are a per-run monotonic counter assigned at write + time (see `_next_snapshot_seq`), not `step_index` -- `ctx.run_step` + resets each `Agent.run` so a reused `run_id` would otherwise clash. + The actual `step_index` is preserved as a field inside the JSON. + + `run_id` is validated against `[A-Za-z0-9_.-]{1,200}` (with `..` rejected) + to prevent path traversal. Blocking I/O is dispatched to a worker thread + via `anyio.to_thread`, so capability hooks do not stall the event loop. + + `BinaryContent` payloads ≥ `media_threshold_bytes` (default 64 KiB) are + externalized through the configured `MediaStore` to keep snapshot JSON + small. The default backs onto `/media/.bin`. Pass + `media_store=None` to keep bytes inline, or pass a custom `MediaStore` + to redirect (e.g. `S3MediaStore(...)`). + """ + + def __init__( + self, + directory: str | Path, + *, + media_store: MediaStore | None | _AutoMedia = 'auto', + media_threshold_bytes: int = _DEFAULT_MEDIA_THRESHOLD_BYTES, + ) -> None: + self._root = Path(directory) + resolved: MediaStore | None + if media_store == 'auto': + resolved = DiskMediaStore(self._root / 'media') + else: + resolved = media_store + self._media_store: MediaStore | None = resolved + self._media_threshold_bytes = media_threshold_bytes + + def _run_dir(self, run_id: str) -> Path: + _validate_id(run_id, field='run_id') + return self._root / run_id + + async def register_run(self, record: RunRecord) -> None: + await anyio.to_thread.run_sync(self._sync_register_run, record) + + def _sync_register_run(self, record: RunRecord) -> None: + run_dir = self._run_dir(record.run_id) + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / 'snapshots').mkdir(exist_ok=True) + (run_dir / 'run.json').write_text(json.dumps(_run_to_dict(record)), encoding='utf-8') + + async def get_run(self, *, run_id: str) -> RunRecord | None: + return await anyio.to_thread.run_sync(self._sync_get_run, run_id) + + def _sync_get_run(self, run_id: str) -> RunRecord | None: + path = self._run_dir(run_id) / 'run.json' + if not path.exists(): + return None + return _run_from_dict(_load_json_object(path.read_text(encoding='utf-8'))) + + async def list_runs( + self, + *, + parent_run_id: str | None = None, + conversation_id: str | None = None, + ) -> list[RunRecord]: + return await anyio.to_thread.run_sync(self._sync_list_runs, parent_run_id, conversation_id) + + def _sync_list_runs( + self, + parent_run_id: str | None, + conversation_id: str | None, + ) -> list[RunRecord]: + if not self._root.exists(): + return [] + records: list[RunRecord] = [] + for sub in self._root.iterdir(): + run_file = sub / 'run.json' + if not run_file.exists(): + continue + record = _run_from_dict(_load_json_object(run_file.read_text(encoding='utf-8'))) + if parent_run_id is not None and record.parent_run_id != parent_run_id: + continue + if conversation_id is not None and record.conversation_id != conversation_id: + continue + records.append(record) + return sorted(records, key=lambda r: r.started_at) + + async def append_event(self, event: StepEvent) -> None: + await anyio.to_thread.run_sync(self._sync_append_event, event) + + def _sync_append_event(self, event: StepEvent) -> None: + run_dir = self._run_dir(event.run_id) + run_dir.mkdir(parents=True, exist_ok=True) + line = json.dumps(_event_to_dict(event)) + with (run_dir / 'events.jsonl').open('a', encoding='utf-8') as fp: + fp.write(line + '\n') + + async def list_events(self, *, run_id: str) -> list[StepEvent]: + return await anyio.to_thread.run_sync(self._sync_list_events, run_id) + + def _sync_list_events(self, run_id: str) -> list[StepEvent]: + path = self._run_dir(run_id) / 'events.jsonl' + if not path.exists(): + return [] + events: list[StepEvent] = [] + for raw in path.read_text(encoding='utf-8').splitlines(): + if raw.strip(): + events.append(_event_from_dict(_load_json_object(raw))) + return events + + async def save_snapshot(self, snapshot: ContinuableSnapshot) -> None: + messages_json: object = json.loads(ModelMessagesTypeAdapter.dump_json(snapshot.messages).decode('utf-8')) + if self._media_store is not None: + messages_json = await externalize_media( + messages_json, + media_store=self._media_store, + threshold_bytes=self._media_threshold_bytes, + ) + await anyio.to_thread.run_sync(self._sync_save_snapshot, snapshot, messages_json) + + def _sync_save_snapshot(self, snapshot: ContinuableSnapshot, messages_json: object) -> None: + run_dir = self._run_dir(snapshot.run_id) + snap_dir = run_dir / 'snapshots' + snap_dir.mkdir(parents=True, exist_ok=True) + payload: dict[str, object] = { + 'run_id': snapshot.run_id, + 'step_index': snapshot.step_index, + 'conversation_id': snapshot.conversation_id, + 'parent_run_id': snapshot.parent_run_id, + 'agent_name': snapshot.agent_name, + 'timestamp': snapshot.timestamp.isoformat(), + 'messages': messages_json, + } + seq = self._next_snapshot_seq(snap_dir) + (snap_dir / f'{seq}.json').write_text(json.dumps(payload), encoding='utf-8') + + @staticmethod + def _next_snapshot_seq(snap_dir: Path) -> int: + """Append-only monotonic counter per run directory. + + Using the snapshot's `step_index` as the filename would collide when + the same `run_id` is reused across `Agent.run` calls -- `ctx.run_step` + resets to 0 each call. The physical sequence is independent of + `step_index`, which lives inside the JSON payload. + """ + max_seq = -1 + for path in snap_dir.glob('*.json'): + try: + seq = int(path.stem) + except ValueError: + continue + if seq > max_seq: + max_seq = seq + return max_seq + 1 + + async def latest_snapshot(self, *, run_id: str) -> ContinuableSnapshot | None: + loaded = await anyio.to_thread.run_sync(self._sync_load_latest_snapshot, run_id) + if loaded is None: + return None + data, messages_json = loaded + if self._media_store is not None: + messages_json = await restore_media(messages_json, media_store=self._media_store) + messages: list[ModelMessage] = ModelMessagesTypeAdapter.validate_python(messages_json) + timestamp_raw = data['timestamp'] + step_raw = data['step_index'] + if not (isinstance(timestamp_raw, str) and isinstance(step_raw, int)): + raise ValueError('snapshot has wrong types') + return ContinuableSnapshot( + run_id=run_id, + step_index=step_raw, + messages=messages, + conversation_id=_opt_str(data.get('conversation_id')), + parent_run_id=_opt_str(data.get('parent_run_id')), + agent_name=_opt_str(data.get('agent_name')), + timestamp=datetime.fromisoformat(timestamp_raw), + ) + + def _sync_load_latest_snapshot(self, run_id: str) -> tuple[dict[str, object], object] | None: + snap_dir = self._run_dir(run_id) / 'snapshots' + if not snap_dir.exists(): + return None + candidates: list[tuple[int, Path]] = [] + for path in snap_dir.glob('*.json'): + try: + candidates.append((int(path.stem), path)) + except ValueError: + continue + if not candidates: + return None + _, latest_path = max(candidates, key=lambda c: c[0]) + data = _load_json_object(latest_path.read_text(encoding='utf-8')) + messages_json = data['messages'] + return data, messages_json + + async def record_tool_effect(self, record: ToolEffectRecord) -> None: + await anyio.to_thread.run_sync(self._sync_record_tool_effect, record) + + def _sync_record_tool_effect(self, record: ToolEffectRecord) -> None: + run_dir = self._run_dir(record.run_id) + run_dir.mkdir(parents=True, exist_ok=True) + line = json.dumps(_tool_effect_to_dict(record)) + with (run_dir / 'tool_effects.jsonl').open('a', encoding='utf-8') as fp: + fp.write(line + '\n') + + async def get_tool_effect(self, *, run_id: str, tool_call_id: str) -> ToolEffectRecord | None: + return await anyio.to_thread.run_sync(self._sync_get_tool_effect, run_id, tool_call_id) + + def _sync_get_tool_effect(self, run_id: str, tool_call_id: str) -> ToolEffectRecord | None: + path = self._run_dir(run_id) / 'tool_effects.jsonl' + if not path.exists(): + return None + latest: ToolEffectRecord | None = None + for raw in path.read_text(encoding='utf-8').splitlines(): + if not raw.strip(): + continue + record = _tool_effect_from_dict(_load_json_object(raw)) + if record.tool_call_id == tool_call_id: + latest = record + return latest + + async def list_unresolved_tool_effects(self, *, run_id: str) -> list[ToolEffectRecord]: + return await anyio.to_thread.run_sync(self._sync_list_unresolved_tool_effects, run_id) + + def _sync_list_unresolved_tool_effects(self, run_id: str) -> list[ToolEffectRecord]: + path = self._run_dir(run_id) / 'tool_effects.jsonl' + if not path.exists(): + return [] + latest_by_call: dict[str, ToolEffectRecord] = {} + for raw in path.read_text(encoding='utf-8').splitlines(): + if not raw.strip(): + continue + record = _tool_effect_from_dict(_load_json_object(raw)) + latest_by_call[record.tool_call_id] = record + return [r for r in latest_by_call.values() if r.status == 'started'] + + +_SQLITE_SCHEMA = """ +CREATE TABLE IF NOT EXISTS runs ( + run_id TEXT PRIMARY KEY, + conversation_id TEXT, + parent_run_id TEXT, + agent_name TEXT, + metadata TEXT NOT NULL, + started_at TEXT NOT NULL +); +CREATE INDEX IF NOT EXISTS idx_runs_conv ON runs(conversation_id); +CREATE INDEX IF NOT EXISTS idx_runs_parent ON runs(parent_run_id); +CREATE INDEX IF NOT EXISTS idx_runs_started ON runs(started_at); + +CREATE TABLE IF NOT EXISTS events ( + seq INTEGER PRIMARY KEY AUTOINCREMENT, + run_id TEXT NOT NULL, + kind TEXT NOT NULL, + step_index INTEGER NOT NULL, + timestamp TEXT NOT NULL, + conversation_id TEXT, + parent_run_id TEXT, + agent_name TEXT, + tool_call_id TEXT, + tool_name TEXT, + error TEXT, + metadata TEXT NOT NULL +); +CREATE INDEX IF NOT EXISTS idx_events_run ON events(run_id, seq); + +CREATE TABLE IF NOT EXISTS snapshots ( + seq INTEGER PRIMARY KEY AUTOINCREMENT, + run_id TEXT NOT NULL, + step_index INTEGER NOT NULL, + conversation_id TEXT, + parent_run_id TEXT, + agent_name TEXT, + timestamp TEXT NOT NULL, + messages TEXT NOT NULL +); +CREATE INDEX IF NOT EXISTS idx_snapshots_run ON snapshots(run_id, seq); + +CREATE TABLE IF NOT EXISTS tool_effects ( + run_id TEXT NOT NULL, + tool_call_id TEXT NOT NULL, + tool_name TEXT NOT NULL, + status TEXT NOT NULL, + started_at TEXT NOT NULL, + ended_at TEXT, + idempotency_key TEXT, + effect_summary TEXT, + PRIMARY KEY (run_id, tool_call_id) +); +""" + + +class SqliteStepStore: + """SQLite-backed store. Single file holds runs, events, snapshots, tool effects + media. + + Pass either `database=` (path; connections opened short-lived per call) or + `connection=` (caller-owned `sqlite3.Connection`). When `database=` is + used, the default `media_store` is a `SqliteMediaStore` against the same + file (sibling `media` table), so a single-file deployment is the default. + + A caller-owned `connection=` **must** be created with + `check_same_thread=False`. Store methods dispatch SQL onto worker threads + via `anyio.to_thread`, so the stdlib default (`check_same_thread=True`) + raises `sqlite3.ProgrammingError` on first use. The `database=` path sets + this internally; `connection=` cannot, so it is the caller's responsibility. + + Concurrency: WAL mode is enabled on connections opened from a `database=` + path (a caller-owned `connection=` keeps whatever journal mode the caller + set). Concurrent readers + a single writer are safe; this matches the + access pattern of one capability per `Agent.run`. + + Schema: + + - `runs(run_id PK, conversation_id, parent_run_id, agent_name, metadata, started_at)` + - `events(seq PK AUTOINCREMENT, run_id, kind, step_index, timestamp, ...)` + - `snapshots(seq PK AUTOINCREMENT, run_id, step_index, ..., messages)` -- + the `AUTOINCREMENT` `seq` mirrors `FileStepStore._next_snapshot_seq` + so `ctx.run_step` resets across `Agent.run` cannot collide. + - `tool_effects(run_id, tool_call_id PK)` -- upsert per + `(run_id, tool_call_id)`, matching `InMemoryStepStore` semantics. + - `media(sha256 PK, media_type, bytes, size_bytes)` -- `INSERT OR IGNORE` + for content-addressed dedup. + + The `runs.run_id` PK enforces the "explicit `run_id` is single-shot" + contract -- `register_run` raises `sqlite3.IntegrityError` on reuse, + which `StepPersistence.before_run` converts to a friendlier + `ValueError` via its own pre-check. + """ + + def __init__( + self, + *, + database: str | Path | None = None, + connection: sqlite3.Connection | None = None, + media_store: MediaStore | None | _AutoMedia = 'auto', + media_threshold_bytes: int = _DEFAULT_MEDIA_THRESHOLD_BYTES, + ) -> None: + if (database is None) == (connection is None): + raise ValueError('provide exactly one of `database=` or `connection=`') + self._database = Path(database) if database is not None else None + self._connection = connection + self._schema_ready = False + resolved: MediaStore | None + if media_store == 'auto': + if self._database is not None: + resolved = SqliteMediaStore(database=self._database) + else: + assert connection is not None + resolved = SqliteMediaStore(connection=connection) + else: + resolved = media_store + self._media_store: MediaStore | None = resolved + self._media_threshold_bytes = media_threshold_bytes + + def _open(self) -> sqlite3.Connection: + if self._connection is not None: + return self._connection + assert self._database is not None + self._database.parent.mkdir(parents=True, exist_ok=True) + conn = sqlite3.connect(self._database, isolation_level=None) + conn.execute('PRAGMA journal_mode=WAL') + return conn + + def _maybe_close(self, conn: sqlite3.Connection) -> None: + if self._connection is None: + conn.close() + + def _ensure_schema(self, conn: sqlite3.Connection) -> None: + if self._schema_ready: + return + conn.executescript(_SQLITE_SCHEMA) + conn.commit() + self._schema_ready = True + + async def register_run(self, record: RunRecord) -> None: + await anyio.to_thread.run_sync(self._sync_register_run, record) + + def _sync_register_run(self, record: RunRecord) -> None: + conn = self._open() + try: + self._ensure_schema(conn) + conn.execute( + 'INSERT INTO runs (run_id, conversation_id, parent_run_id, agent_name, metadata, started_at) ' + 'VALUES (?, ?, ?, ?, ?, ?)', + ( + record.run_id, + record.conversation_id, + record.parent_run_id, + record.agent_name, + json.dumps(dict(record.metadata)), + record.started_at.isoformat(), + ), + ) + finally: + self._maybe_close(conn) + + async def get_run(self, *, run_id: str) -> RunRecord | None: + return await anyio.to_thread.run_sync(self._sync_get_run, run_id) + + def _sync_get_run(self, run_id: str) -> RunRecord | None: + conn = self._open() + try: + self._ensure_schema(conn) + row = conn.execute( + 'SELECT run_id, conversation_id, parent_run_id, agent_name, metadata, started_at ' + 'FROM runs WHERE run_id = ?', + (run_id,), + ).fetchone() + finally: + self._maybe_close(conn) + if row is None: + return None + return _run_from_row(row) + + async def list_runs( + self, + *, + parent_run_id: str | None = None, + conversation_id: str | None = None, + ) -> list[RunRecord]: + return await anyio.to_thread.run_sync(self._sync_list_runs, parent_run_id, conversation_id) + + def _sync_list_runs( + self, + parent_run_id: str | None, + conversation_id: str | None, + ) -> list[RunRecord]: + conn = self._open() + try: + self._ensure_schema(conn) + clauses: list[str] = [] + params: list[object] = [] + if parent_run_id is not None: + clauses.append('parent_run_id = ?') + params.append(parent_run_id) + if conversation_id is not None: + clauses.append('conversation_id = ?') + params.append(conversation_id) + sql = 'SELECT run_id, conversation_id, parent_run_id, agent_name, metadata, started_at FROM runs' + if clauses: + sql += ' WHERE ' + ' AND '.join(clauses) + sql += ' ORDER BY started_at ASC' + rows = conn.execute(sql, params).fetchall() + finally: + self._maybe_close(conn) + return [_run_from_row(row) for row in rows] + + async def append_event(self, event: StepEvent) -> None: + await anyio.to_thread.run_sync(self._sync_append_event, event) + + def _sync_append_event(self, event: StepEvent) -> None: + conn = self._open() + try: + self._ensure_schema(conn) + conn.execute( + 'INSERT INTO events (' + 'run_id, kind, step_index, timestamp, conversation_id, parent_run_id, ' + 'agent_name, tool_call_id, tool_name, error, metadata' + ') VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)', + ( + event.run_id, + event.kind, + event.step_index, + event.timestamp.isoformat(), + event.conversation_id, + event.parent_run_id, + event.agent_name, + event.tool_call_id, + event.tool_name, + event.error, + json.dumps(dict(event.metadata)), + ), + ) + finally: + self._maybe_close(conn) + + async def list_events(self, *, run_id: str) -> list[StepEvent]: + return await anyio.to_thread.run_sync(self._sync_list_events, run_id) + + def _sync_list_events(self, run_id: str) -> list[StepEvent]: + conn = self._open() + try: + self._ensure_schema(conn) + rows = conn.execute( + 'SELECT run_id, kind, step_index, timestamp, conversation_id, parent_run_id, ' + 'agent_name, tool_call_id, tool_name, error, metadata ' + 'FROM events WHERE run_id = ? ORDER BY seq ASC', + (run_id,), + ).fetchall() + finally: + self._maybe_close(conn) + return [_event_from_row(row) for row in rows] + + async def save_snapshot(self, snapshot: ContinuableSnapshot) -> None: + messages_json: object = json.loads(ModelMessagesTypeAdapter.dump_json(snapshot.messages).decode('utf-8')) + if self._media_store is not None: + messages_json = await externalize_media( + messages_json, + media_store=self._media_store, + threshold_bytes=self._media_threshold_bytes, + ) + await anyio.to_thread.run_sync(self._sync_save_snapshot, snapshot, messages_json) + + def _sync_save_snapshot(self, snapshot: ContinuableSnapshot, messages_json: object) -> None: + conn = self._open() + try: + self._ensure_schema(conn) + conn.execute( + 'INSERT INTO snapshots (' + 'run_id, step_index, conversation_id, parent_run_id, agent_name, timestamp, messages' + ') VALUES (?, ?, ?, ?, ?, ?, ?)', + ( + snapshot.run_id, + snapshot.step_index, + snapshot.conversation_id, + snapshot.parent_run_id, + snapshot.agent_name, + snapshot.timestamp.isoformat(), + json.dumps(messages_json), + ), + ) + finally: + self._maybe_close(conn) + + async def latest_snapshot(self, *, run_id: str) -> ContinuableSnapshot | None: + row = await anyio.to_thread.run_sync(self._sync_load_latest_snapshot, run_id) + if row is None: + return None + step_index, conv_id, parent_id, agent_name, timestamp_iso, messages_json_text = row + messages_json: object = json.loads(messages_json_text) + if self._media_store is not None: + messages_json = await restore_media(messages_json, media_store=self._media_store) + messages: list[ModelMessage] = ModelMessagesTypeAdapter.validate_python(messages_json) + return ContinuableSnapshot( + run_id=run_id, + step_index=step_index, + messages=messages, + conversation_id=conv_id, + parent_run_id=parent_id, + agent_name=agent_name, + timestamp=datetime.fromisoformat(timestamp_iso), + ) + + def _sync_load_latest_snapshot( + self, run_id: str + ) -> tuple[int, str | None, str | None, str | None, str, str] | None: + conn = self._open() + try: + self._ensure_schema(conn) + row = conn.execute( + 'SELECT step_index, conversation_id, parent_run_id, agent_name, timestamp, messages ' + 'FROM snapshots WHERE run_id = ? ORDER BY seq DESC LIMIT 1', + (run_id,), + ).fetchone() + finally: + self._maybe_close(conn) + if row is None: + return None + step_index, conv_id, parent_id, agent_name, timestamp_iso, messages_json_text = row + if not (isinstance(step_index, int) and isinstance(timestamp_iso, str) and isinstance(messages_json_text, str)): + raise ValueError('snapshot row has wrong types') + return ( + step_index, + _opt_str(conv_id), + _opt_str(parent_id), + _opt_str(agent_name), + timestamp_iso, + messages_json_text, + ) + + async def record_tool_effect(self, record: ToolEffectRecord) -> None: + await anyio.to_thread.run_sync(self._sync_record_tool_effect, record) + + def _sync_record_tool_effect(self, record: ToolEffectRecord) -> None: + conn = self._open() + try: + self._ensure_schema(conn) + conn.execute( + 'INSERT INTO tool_effects (' + 'run_id, tool_call_id, tool_name, status, started_at, ended_at, idempotency_key, effect_summary' + ') VALUES (?, ?, ?, ?, ?, ?, ?, ?) ' + 'ON CONFLICT(run_id, tool_call_id) DO UPDATE SET ' + 'tool_name=excluded.tool_name, status=excluded.status, started_at=excluded.started_at, ' + 'ended_at=excluded.ended_at, idempotency_key=excluded.idempotency_key, ' + 'effect_summary=excluded.effect_summary', + ( + record.run_id, + record.tool_call_id, + record.tool_name, + record.status, + record.started_at.isoformat(), + record.ended_at.isoformat() if record.ended_at is not None else None, + record.idempotency_key, + record.effect_summary, + ), + ) + finally: + self._maybe_close(conn) + + async def get_tool_effect(self, *, run_id: str, tool_call_id: str) -> ToolEffectRecord | None: + return await anyio.to_thread.run_sync(self._sync_get_tool_effect, run_id, tool_call_id) + + def _sync_get_tool_effect(self, run_id: str, tool_call_id: str) -> ToolEffectRecord | None: + conn = self._open() + try: + self._ensure_schema(conn) + row = conn.execute( + 'SELECT run_id, tool_call_id, tool_name, status, started_at, ended_at, ' + 'idempotency_key, effect_summary FROM tool_effects ' + 'WHERE run_id = ? AND tool_call_id = ?', + (run_id, tool_call_id), + ).fetchone() + finally: + self._maybe_close(conn) + if row is None: + return None + return _tool_effect_from_row(row) + + async def list_unresolved_tool_effects(self, *, run_id: str) -> list[ToolEffectRecord]: + return await anyio.to_thread.run_sync(self._sync_list_unresolved_tool_effects, run_id) + + def _sync_list_unresolved_tool_effects(self, run_id: str) -> list[ToolEffectRecord]: + conn = self._open() + try: + self._ensure_schema(conn) + rows = conn.execute( + 'SELECT run_id, tool_call_id, tool_name, status, started_at, ended_at, ' + 'idempotency_key, effect_summary FROM tool_effects ' + "WHERE run_id = ? AND status = 'started'", + (run_id,), + ).fetchall() + finally: + self._maybe_close(conn) + return [_tool_effect_from_row(row) for row in rows] + + +def _run_from_row(row: tuple[object, ...]) -> RunRecord: + run_id, conv_id, parent_id, agent_name, metadata_text, started_at_iso = row + if not (isinstance(run_id, str) and isinstance(metadata_text, str) and isinstance(started_at_iso, str)): + raise ValueError('run row has wrong types') + return RunRecord( + run_id=run_id, + conversation_id=_opt_str(conv_id), + parent_run_id=_opt_str(parent_id), + agent_name=_opt_str(agent_name), + metadata=_str_str_dict(json.loads(metadata_text)), + started_at=datetime.fromisoformat(started_at_iso), + ) + + +def _event_from_row(row: tuple[object, ...]) -> StepEvent: + ( + run_id, + kind_raw, + step_raw, + timestamp_raw, + conv_id, + parent_id, + agent_name, + tool_call_id, + tool_name, + error, + metadata_text, + ) = row + if not ( + isinstance(run_id, str) + and isinstance(kind_raw, str) + and isinstance(step_raw, int) + and isinstance(timestamp_raw, str) + and isinstance(metadata_text, str) + ): + raise ValueError('event row has wrong types') + kind = _EVENT_KIND_TABLE.get(kind_raw) + if kind is None: + raise ValueError(f'unknown event kind: {kind_raw!r}') + return StepEvent( + run_id=run_id, + kind=kind, + step_index=step_raw, + timestamp=datetime.fromisoformat(timestamp_raw), + conversation_id=_opt_str(conv_id), + parent_run_id=_opt_str(parent_id), + agent_name=_opt_str(agent_name), + tool_call_id=_opt_str(tool_call_id), + tool_name=_opt_str(tool_name), + error=_opt_str(error), + metadata=_str_str_dict(json.loads(metadata_text)), + ) + + +def _tool_effect_from_row(row: tuple[object, ...]) -> ToolEffectRecord: + ( + run_id, + tool_call_id, + tool_name, + status_raw, + started_at_raw, + ended_at_raw, + idempotency_key, + effect_summary, + ) = row + if not ( + isinstance(run_id, str) + and isinstance(tool_call_id, str) + and isinstance(tool_name, str) + and isinstance(status_raw, str) + and isinstance(started_at_raw, str) + ): + raise ValueError('tool effect row has wrong types') + status = _TOOL_STATUS_TABLE.get(status_raw) + if status is None: + raise ValueError(f'unknown tool effect status: {status_raw!r}') + return ToolEffectRecord( + tool_call_id=tool_call_id, + tool_name=tool_name, + run_id=run_id, + status=status, + started_at=datetime.fromisoformat(started_at_raw), + ended_at=datetime.fromisoformat(ended_at_raw) if isinstance(ended_at_raw, str) else None, + idempotency_key=_opt_str(idempotency_key), + effect_summary=_opt_str(effect_summary), + ) diff --git a/pydantic_ai_harness/experimental/step_persistence/_types.py b/pydantic_ai_harness/experimental/step_persistence/_types.py new file mode 100644 index 0000000..b153220 --- /dev/null +++ b/pydantic_ai_harness/experimental/step_persistence/_types.py @@ -0,0 +1,129 @@ +"""Data types for step-event persistence.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Literal + +from pydantic_ai.messages import ModelMessage + +EventKind = Literal[ + 'run_started', + 'run_completed', + 'run_failed', + 'model_request_started', + 'model_request_completed', + 'model_request_failed', + 'tool_call_started', + 'tool_call_completed', + 'tool_call_failed', +] +"""Boundary that produced a `StepEvent`. + +Choose `kind` from this set so consumers can route on event type without +string typos. Append-only: never mutate an emitted event; record corrections +as a follow-up event. +""" + +ToolEffectStatus = Literal['started', 'completed', 'failed'] +"""Lifecycle status of a tool call recorded in the effect ledger. + +A `tool_call_id` whose latest record is `started` was in flight when the +process last wrote. Treat it as `unknown_after_crash` when replaying -- the +external side effect may or may not have happened. +""" + + +def _utcnow() -> datetime: + return datetime.now(timezone.utc) + + +def _empty_str_dict() -> dict[str, str]: + return {} + + +@dataclass(kw_only=True) +class StepEvent: + """A single append-only event recorded during agent execution. + + Events describe boundaries (run start/end, model request, tool call, + failure) but never carry recoverable state on their own. Pair with a + `ContinuableSnapshot` for resume; pair with `ToolEffectRecord` for + side-effect status. + + `conversation_id` mirrors pydantic_ai's three-level identity stack -- + conversation (the dialogue) → run (one `Agent.run` call) → step + (one graph node). Two `Agent.run` calls that share a + `conversation_id` produce two separate `run_id`s. + """ + + run_id: str + kind: EventKind + step_index: int + timestamp: datetime = field(default_factory=_utcnow) + conversation_id: str | None = None + parent_run_id: str | None = None + agent_name: str | None = None + tool_call_id: str | None = None + tool_name: str | None = None + error: str | None = None + metadata: dict[str, str] = field(default_factory=_empty_str_dict) + + +@dataclass(kw_only=True) +class ContinuableSnapshot: + """A provider-valid message-history snapshot safe to resume from. + + Only emitted at boundaries where every `ToolCallPart` in the history + has a matching `ToolReturnPart` or `RetryPromptPart`. Pass `messages` + to `Agent.run(..., message_history=...)` to continue or fork the run. + """ + + run_id: str + step_index: int + messages: list[ModelMessage] + conversation_id: str | None = None + parent_run_id: str | None = None + agent_name: str | None = None + timestamp: datetime = field(default_factory=_utcnow) + + +@dataclass(kw_only=True) +class ToolEffectRecord: + """Ledger entry for a tool call's side-effect status. + + Read-only tools and side-effectful tools share the same record shape; + the orchestrator decides whether replay is safe based on + `idempotency_key` and `effect_summary`. A record without a matching + `completed` or `failed` update after process restart should be treated + as `unknown_after_crash`. + """ + + tool_call_id: str + tool_name: str + run_id: str + status: ToolEffectStatus + started_at: datetime = field(default_factory=_utcnow) + ended_at: datetime | None = None + idempotency_key: str | None = None + effect_summary: str | None = None + + +@dataclass(kw_only=True) +class RunRecord: + """Lineage metadata for an agent run. + + `conversation_id` groups runs of the same dialogue (the user-visible + sequence). `parent_run_id` is the hierarchical link: which run spawned + this one. The two are independent axes -- a delegate may share a + conversation across attempts (different `run_id`s, same `conversation_id`) + while pointing at a different orchestrator run via `parent_run_id`. + """ + + run_id: str + conversation_id: str | None = None + parent_run_id: str | None = None + agent_name: str | None = None + metadata: dict[str, str] = field(default_factory=_empty_str_dict) + started_at: datetime = field(default_factory=_utcnow) diff --git a/pydantic_ai_harness/experimental/subagents/README.md b/pydantic_ai_harness/experimental/subagents/README.md index 12079cf..00ef71d 100644 --- a/pydantic_ai_harness/experimental/subagents/README.md +++ b/pydantic_ai_harness/experimental/subagents/README.md @@ -1,5 +1,24 @@ # SubAgents +> [!WARNING] +> **Experimental.** This capability lives under `pydantic_ai_harness.experimental` and may +> change or be removed in any release, without a deprecation period. Import it from the +> experimental path -- there is no top-level export: +> +> ```python +> from pydantic_ai_harness.experimental.subagents import SubAgents +> ``` +> +> Importing any experimental capability emits a `HarnessExperimentalWarning`. Silence **all** +> harness experimental warnings with a single filter (no per-capability lines needed): +> +> ```python +> import warnings +> from pydantic_ai_harness.experimental import HarnessExperimentalWarning +> +> warnings.filterwarnings('ignore', category=HarnessExperimentalWarning) +> ``` + Let an agent delegate self-contained tasks to named child agents. ## The problem diff --git a/pyproject.toml b/pyproject.toml index dcc4c38..f526d4a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ classifiers = [ 'Typing :: Typed', ] dependencies = [ + "httpx>=0.28.1", 'pydantic-ai-slim>=1.95.1', ] @@ -71,6 +72,7 @@ dev = [ 'inline-snapshot>=0.32.5', 'pydantic-ai-slim[spec]>=1.95.1', "pytest-examples>=0.0.18", + "pytest-recording>=0.13.4", ] lint = [ 'ruff>=0.14', diff --git a/tests/experimental/media/__init__.py b/tests/experimental/media/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_exists_present_after_put.yaml b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_exists_present_after_put.yaml new file mode 100644 index 0000000..a5877fa --- /dev/null +++ b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_exists_present_after_put.yaml @@ -0,0 +1,90 @@ +interactions: + - request: + body: pydantic-ai-harness step-persistence VCR cassette payload v1 + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '60' + User-Agent: + - python-httpx/0.28.1 + authorization: + - REDACTED + content-type: + - application/octet-stream + host: + - account.r2.cloudflarestorage.com + x-amz-content-sha256: + - 1743d339a10eaa072801e8bf3eca688d6a339cb26b9354194f15c131727274f1 + x-amz-date: + - REDACTED + method: PUT + uri: https://account.r2.cloudflarestorage.com/harness-test-bucket/harness-vcr/1743d339a10eaa072801e8bf3eca688d6a339cb26b9354194f15c131727274f1.bin + response: + body: + string: '' + headers: + Connection: + - keep-alive + Content-Length: + - '0' + Content-Type: + - text/plain;charset=UTF-8 + Date: + - Mon, 25 May 2026 04:00:29 GMT + ETag: + - '"076102a313b5c40b44923280cc2c04c3"' + Server: + - cloudflare + status: + code: 200 + message: OK + - request: + body: '' + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + User-Agent: + - python-httpx/0.28.1 + authorization: + - REDACTED + host: + - account.r2.cloudflarestorage.com + x-amz-content-sha256: + - e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 + x-amz-date: + - REDACTED + method: HEAD + uri: https://account.r2.cloudflarestorage.com/harness-test-bucket/harness-vcr/1743d339a10eaa072801e8bf3eca688d6a339cb26b9354194f15c131727274f1.bin + response: + body: + string: '' + headers: + Accept-Ranges: + - bytes + Connection: + - keep-alive + Content-Length: + - '60' + Content-Type: + - application/octet-stream + Date: + - Mon, 25 May 2026 04:00:29 GMT + ETag: + - '"076102a313b5c40b44923280cc2c04c3"' + Last-Modified: + - Mon, 25 May 2026 04:00:29 GMT + Server: + - cloudflare + status: + code: 200 + message: OK +version: 1 diff --git a/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_exists_returns_false_for_missing.yaml b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_exists_returns_false_for_missing.yaml new file mode 100644 index 0000000..43460f5 --- /dev/null +++ b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_exists_returns_false_for_missing.yaml @@ -0,0 +1,38 @@ +interactions: + - request: + body: '' + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + User-Agent: + - python-httpx/0.28.1 + authorization: + - REDACTED + host: + - account.r2.cloudflarestorage.com + x-amz-content-sha256: + - e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 + x-amz-date: + - REDACTED + method: HEAD + uri: https://account.r2.cloudflarestorage.com/harness-test-bucket/harness-vcr/0000000000000000000000000000000000000000000000000000000000000000.bin + response: + body: + string: '' + headers: + Connection: + - keep-alive + Content-Type: + - text/plain;charset=UTF-8 + Date: + - Mon, 25 May 2026 04:00:29 GMT + Server: + - cloudflare + status: + code: 404 + message: Not Found +version: 1 diff --git a/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_get_missing_raises.yaml b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_get_missing_raises.yaml new file mode 100644 index 0000000..36939df --- /dev/null +++ b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_get_missing_raises.yaml @@ -0,0 +1,38 @@ +interactions: + - request: + body: '' + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + User-Agent: + - python-httpx/0.28.1 + authorization: + - REDACTED + host: + - account.r2.cloudflarestorage.com + x-amz-content-sha256: + - e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 + x-amz-date: + - REDACTED + method: GET + uri: https://account.r2.cloudflarestorage.com/harness-test-bucket/harness-vcr/0000000000000000000000000000000000000000000000000000000000000000.bin + response: + body: + string: '' + headers: + Connection: + - keep-alive + Content-Type: + - application/xml + Date: + - Mon, 25 May 2026 04:00:31 GMT + Server: + - cloudflare + status: + code: 404 + message: Not Found +version: 1 diff --git a/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_get_round_trips_bytes.yaml b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_get_round_trips_bytes.yaml new file mode 100644 index 0000000..384de43 --- /dev/null +++ b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_get_round_trips_bytes.yaml @@ -0,0 +1,90 @@ +interactions: + - request: + body: pydantic-ai-harness step-persistence VCR cassette payload v1 + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '60' + User-Agent: + - python-httpx/0.28.1 + authorization: + - REDACTED + content-type: + - application/octet-stream + host: + - account.r2.cloudflarestorage.com + x-amz-content-sha256: + - 1743d339a10eaa072801e8bf3eca688d6a339cb26b9354194f15c131727274f1 + x-amz-date: + - REDACTED + method: PUT + uri: https://account.r2.cloudflarestorage.com/harness-test-bucket/harness-vcr/1743d339a10eaa072801e8bf3eca688d6a339cb26b9354194f15c131727274f1.bin + response: + body: + string: '' + headers: + Connection: + - keep-alive + Content-Length: + - '0' + Content-Type: + - text/plain;charset=UTF-8 + Date: + - Mon, 25 May 2026 04:00:30 GMT + ETag: + - '"076102a313b5c40b44923280cc2c04c3"' + Server: + - cloudflare + status: + code: 200 + message: OK + - request: + body: '' + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + User-Agent: + - python-httpx/0.28.1 + authorization: + - REDACTED + host: + - account.r2.cloudflarestorage.com + x-amz-content-sha256: + - e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855 + x-amz-date: + - REDACTED + method: GET + uri: https://account.r2.cloudflarestorage.com/harness-test-bucket/harness-vcr/1743d339a10eaa072801e8bf3eca688d6a339cb26b9354194f15c131727274f1.bin + response: + body: + string: pydantic-ai-harness step-persistence VCR cassette payload v1 + headers: + Accept-Ranges: + - bytes + Connection: + - keep-alive + Content-Length: + - '60' + Content-Type: + - application/octet-stream + Date: + - Mon, 25 May 2026 04:00:30 GMT + ETag: + - '"076102a313b5c40b44923280cc2c04c3"' + Last-Modified: + - Mon, 25 May 2026 04:00:30 GMT + Server: + - cloudflare + status: + code: 200 + message: OK +version: 1 diff --git a/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_put_succeeds.yaml b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_put_succeeds.yaml new file mode 100644 index 0000000..3d09d8c --- /dev/null +++ b/tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_put_succeeds.yaml @@ -0,0 +1,46 @@ +interactions: + - request: + body: pydantic-ai-harness step-persistence VCR cassette payload v1 + headers: + Accept: + - '*/*' + Accept-Encoding: + - gzip, deflate + Connection: + - keep-alive + Content-Length: + - '60' + User-Agent: + - python-httpx/0.28.1 + authorization: + - REDACTED + content-type: + - application/octet-stream + host: + - account.r2.cloudflarestorage.com + x-amz-content-sha256: + - 1743d339a10eaa072801e8bf3eca688d6a339cb26b9354194f15c131727274f1 + x-amz-date: + - REDACTED + method: PUT + uri: https://account.r2.cloudflarestorage.com/harness-test-bucket/harness-vcr/1743d339a10eaa072801e8bf3eca688d6a339cb26b9354194f15c131727274f1.bin + response: + body: + string: '' + headers: + Connection: + - keep-alive + Content-Length: + - '0' + Content-Type: + - text/plain;charset=UTF-8 + Date: + - Mon, 25 May 2026 04:00:28 GMT + ETag: + - '"076102a313b5c40b44923280cc2c04c3"' + Server: + - cloudflare + status: + code: 200 + message: OK +version: 1 diff --git a/tests/experimental/media/conftest.py b/tests/experimental/media/conftest.py new file mode 100644 index 0000000..bb53e5f --- /dev/null +++ b/tests/experimental/media/conftest.py @@ -0,0 +1,206 @@ +"""VCR configuration for `S3MediaStore` cassette tests. + +The cassettes are committed alongside the tests so CI can replay them +without R2 / AWS credentials. Recording is opt-in via +`pytest --record-mode=once` (or `=new_episodes`/`=all`) with real +`S3_*` env vars; replay is the default mode and what CI runs. + +Sanitisation policy: every recorded cassette is rewritten on disk to +swap the real R2 account-id subdomain and bucket name for fixed +placeholders, and drop the `Authorization` header. The test setup uses +the placeholder endpoint + bucket so request matching still succeeds on +replay. See `_rewrite_request` / `_rewrite_response` below. +""" + +# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportMissingTypeStubs=false +from __future__ import annotations + +import os +import re +from typing import TYPE_CHECKING, Any + +import pytest + +if TYPE_CHECKING: + from vcr.request import Request as VcrRequest # pyright: ignore[reportMissingTypeStubs] + + +# Public placeholders baked into the committed cassettes. Tests pass +# these *exact* values when constructing `S3MediaStore`, so the replay +# URI matches the recorded URI even though both were sanitised. +SANITIZED_HOST = 'account.r2.cloudflarestorage.com' +SANITIZED_BUCKET = 'harness-test-bucket' +SANITIZED_ENDPOINT = f'https://{SANITIZED_HOST}' +SANITIZED_REGION = 'auto' + + +def _real_account_host_pattern() -> re.Pattern[str] | None: # pragma: no cover + """Build a regex that matches the real R2 host so we can scrub it.""" + endpoint = os.environ.get('S3_ENDPOINT') + if not endpoint: + return None + match = re.match(r'https?://([^/]+)', endpoint) + if not match: + return None + return re.compile(re.escape(match.group(1))) + + +def _real_bucket_pattern() -> re.Pattern[str] | None: # pragma: no cover + bucket = os.environ.get('S3_BUCKET_NAME') + if not bucket: + return None + return re.compile(r'/' + re.escape(bucket) + r'/') + + +def _rewrite_request(request: VcrRequest) -> VcrRequest: # pragma: no cover + """Strip account-id, bucket name, and credentials from recorded request.""" + host_pat = _real_account_host_pattern() + if host_pat is not None: + request.uri = host_pat.sub(SANITIZED_HOST, request.uri) + bucket_pat = _real_bucket_pattern() + if bucket_pat is not None: + request.uri = bucket_pat.sub(f'/{SANITIZED_BUCKET}/', request.uri) + # VCR already drops Authorization via `filter_headers`, but Host is set + # by httpx independently — overwrite it so cassettes never carry the + # real account subdomain. + if 'host' in request.headers: + request.headers['host'] = SANITIZED_HOST + return request + + +_DROP_RESPONSE_HEADERS = frozenset( + { + 'cf-ray', + 'cf-cache-status', + 'x-amz-version-id', + 'x-amz-request-id', + 'x-amz-id-2', + 'x-amz-checksum-crc64nvme', + 'x-amz-checksum-crc32', + 'x-amz-checksum-crc32c', + 'x-amz-checksum-sha1', + 'x-amz-checksum-sha256', + } +) + + +def _rewrite_response(response: dict[str, Any]) -> dict[str, Any]: # pragma: no cover + """Sanitise the response: drop noisy / identifying headers and any error body. + + For non-2xx responses (typically the gzipped XML R2/AWS error envelope, + which can mention the bucket) we blank the body entirely and strip + `Content-Encoding`. Our `S3MediaStore` only inspects `status_code` for + 4xx and `response.text[:200]` for 5xx — no test in this module relies + on the error body shape. + """ + host_pat = _real_account_host_pattern() + bucket_pat = _real_bucket_pattern() + headers = response.get('headers', {}) + for header_name in list(headers.keys()): + if header_name.lower() in _DROP_RESPONSE_HEADERS: + del headers[header_name] + continue + values = headers[header_name] + if not isinstance(values, list): + continue + new_values: list[str] = [] + for v in values: + if not isinstance(v, str): + new_values.append(v) + continue + if host_pat is not None: + v = host_pat.sub(SANITIZED_HOST, v) + if bucket_pat is not None: + v = bucket_pat.sub(f'/{SANITIZED_BUCKET}/', v) + new_values.append(v) + headers[header_name] = new_values + + status = response.get('status', {}) + code = status.get('code') if isinstance(status, dict) else None + is_success = isinstance(code, int) and 200 <= code < 300 + body = response.get('body', {}) + if not is_success and isinstance(body, dict): + # Drop any provider error envelope — it can name the bucket inside + # the gzipped XML. Tests only read the status code on this path. + body['string'] = b'' + for header_name in list(headers.keys()): + if header_name.lower() in ('content-encoding', 'content-length', 'transfer-encoding'): + del headers[header_name] + return response + + +@pytest.fixture(scope='module') +def vcr_config() -> dict[str, Any]: + """Per-module VCR configuration. Cassettes live next to the tests. + + Matching: method + scheme + host + path + body. Headers (including + SigV4 `authorization` and `x-amz-date`) are NOT part of matching — + they regenerate per replay and would otherwise miss every time. + + Record mode is whatever `--record-mode` says (default `none`). + """ + return { + 'filter_headers': [ + ('authorization', 'REDACTED'), + ('x-amz-date', 'REDACTED'), + ], + 'before_record_request': _rewrite_request, + 'before_record_response': _rewrite_response, + 'match_on': ['method', 'scheme', 'host', 'path', 'body'], + } + + +@pytest.fixture +def anyio_backend() -> str: + """Restrict the S3 cassette tests to asyncio — we don't need trio cassettes.""" + return 'asyncio' + + +@pytest.fixture +def s3_credentials() -> dict[str, str]: + """Real R2 creds when env is set; sanitised placeholders otherwise. + + The placeholders match the values baked into the scrubbed cassettes + (see `_rewrite_request` / `_rewrite_response` above), so replay works + against `tests/media/cassettes/` with no env vars at all — exactly what + CI runs. + + **Why the placeholders double as a leakage canary:** if the scrubber + ever misses a value when re-recording, the cassette will contain the + real bucket / account id while the replay test still constructs URLs + from the placeholder constants — the URL matcher will fail and CI + will surface the leak. Reusing this pattern across the suite (always + pass placeholder values, scrub on write) catches accidental + credential / private-data exposure in committed cassettes. + + `region` is hardcoded to `'auto'` because R2 rejects every other name + and the SigV4 region is part of the credential scope (filtered from + the cassette `Authorization`, so it does not affect replay matching). + Override the fixture in another conftest if recording against AWS S3. + """ + return { + 'bucket': os.environ.get('S3_BUCKET_NAME', SANITIZED_BUCKET), + 'endpoint': os.environ.get('S3_ENDPOINT', SANITIZED_ENDPOINT), + 'region': 'auto', + 'access_key_id': os.environ.get('S3_ACCESS_KEY_ID', 'AKIAIOSFODNN7EXAMPLE'), + 'secret_access_key': os.environ.get('S3_SECRET_ACCESS_KEY', 'REDACTED-FOR-REPLAY'), + } + + +@pytest.fixture +def s3_store(s3_credentials: dict[str, str]) -> Any: + """`S3MediaStore` built from the credentials fixture, with a fixed key prefix. + + The key prefix is part of the URL path that lands in the cassette, so + keep it stable across re-records. + """ + from pydantic_ai_harness.experimental.media import S3MediaStore + + return S3MediaStore( + bucket=s3_credentials['bucket'], + endpoint=s3_credentials['endpoint'], + region=s3_credentials['region'], + access_key_id=s3_credentials['access_key_id'], + secret_access_key=s3_credentials['secret_access_key'], + key_prefix='harness-vcr/', + ) diff --git a/tests/experimental/media/test_media.py b/tests/experimental/media/test_media.py new file mode 100644 index 0000000..b02e678 --- /dev/null +++ b/tests/experimental/media/test_media.py @@ -0,0 +1,821 @@ +"""Tests for `pydantic_ai_harness.experimental.media`: stores + walker + SigV4 + S3 store.""" + +from __future__ import annotations + +import os +import sqlite3 +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from httpx import AsyncClient, MockTransport, Request, Response + +from pydantic_ai_harness.experimental.media import ( + DiskMediaStore, + MediaContext, + MediaStore, + S3MediaStore, + SqliteMediaStore, + externalize_media, + make_static_public_url, + media_uri_for, + parse_media_uri, + restore_media, +) +from pydantic_ai_harness.experimental.media._s3 import sign_request + +pytestmark = pytest.mark.anyio + + +class TestMediaUriHelpers: + def test_media_uri_for_returns_canonical_scheme(self) -> None: + uri = media_uri_for(b'hello world') + assert uri == 'media+sha256://b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9' + + def test_parse_media_uri_strips_scheme(self) -> None: + uri = media_uri_for(b'abc') + digest = parse_media_uri(uri) + assert len(digest) == 64 + assert all(c in '0123456789abcdef' for c in digest) + + def test_parse_media_uri_rejects_other_schemes(self) -> None: + with pytest.raises(ValueError, match='not a media URI'): + parse_media_uri('http://example.com/foo') + + def test_parse_media_uri_rejects_short_digest(self) -> None: + with pytest.raises(ValueError, match='invalid sha256 digest'): + parse_media_uri('media+sha256://deadbeef') + + def test_parse_media_uri_rejects_uppercase_hex(self) -> None: + with pytest.raises(ValueError, match='invalid sha256 digest'): + parse_media_uri('media+sha256://' + 'A' * 64) + + +class TestDiskMediaStore: + async def test_put_get_round_trip(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + uri = await store.put(b'hello bytes', context=MediaContext(media_type='application/octet-stream')) + assert uri.startswith('media+sha256://') + assert await store.get(uri) == b'hello bytes' + + async def test_put_with_empty_context_works(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + uri = await store.put(b'no context') + assert await store.get(uri) == b'no context' + + async def test_dedup_on_repeated_put(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + uri1 = await store.put(b'same content') + uri2 = await store.put(b'same content') + assert uri1 == uri2 + files = list(tmp_path.glob('*.bin')) + assert len(files) == 1 + + async def test_exists(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + uri = await store.put(b'present') + assert await store.exists(uri) is True + missing_uri = media_uri_for(b'never put') + assert await store.exists(missing_uri) is False + + async def test_get_missing_raises(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + with pytest.raises(FileNotFoundError): + await store.get(media_uri_for(b'never put')) + + async def test_custom_key_strategy_controls_path(self, tmp_path: Path) -> None: + def strategy(uri: str, ctx: MediaContext) -> str: + return f'images/{parse_media_uri(uri)}.png' + + store = DiskMediaStore(tmp_path, key_strategy=strategy) + uri = await store.put(b'pixels') + assert (tmp_path / 'images' / f'{parse_media_uri(uri)}.png').exists() + assert await store.get(uri) == b'pixels' + + async def test_key_strategy_blocks_traversal(self, tmp_path: Path) -> None: + def evil(uri: str, ctx: MediaContext) -> str: + return '../../../etc/passwd' + + store = DiskMediaStore(tmp_path, key_strategy=evil) + with pytest.raises(ValueError, match='traversal-unsafe'): + await store.put(b'attack') + + async def test_key_strategy_blocks_absolute_path(self, tmp_path: Path) -> None: + # `Path('/root') / '/abs'` silently returns `/abs` — without this check + # an absolute key escapes the store directory. + def evil(uri: str, ctx: MediaContext) -> str: + return '/etc/passwd' + + store = DiskMediaStore(tmp_path, key_strategy=evil) + with pytest.raises(ValueError, match='traversal-unsafe'): + await store.put(b'attack') + + async def test_metadata_round_trips_via_sidecar(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + uri = await store.put( + b'tagged', + context=MediaContext(metadata={'origin': 'user', 'tenant': 'acme'}), + ) + sidecars = list(tmp_path.glob('*.meta.json')) + assert len(sidecars) == 1 + assert await store.get_metadata(uri) == {'origin': 'user', 'tenant': 'acme'} + + async def test_metadata_absent_when_not_supplied(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + uri = await store.put(b'no tags') + assert list(tmp_path.glob('*.meta.json')) == [] + assert await store.get_metadata(uri) == {} + + async def test_get_metadata_rejects_non_object_sidecar(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + uri = await store.put(b'with meta', context=MediaContext(metadata={'k': 'v'})) + sidecar = next(iter(tmp_path.glob('*.meta.json'))) + sidecar.write_text('"not an object"') + with pytest.raises(ValueError, match='must be a JSON object'): + await store.get_metadata(uri) + + async def test_get_metadata_rejects_non_string_values(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + uri = await store.put(b'with meta', context=MediaContext(metadata={'k': 'v'})) + sidecar = next(iter(tmp_path.glob('*.meta.json'))) + sidecar.write_text('{"k": 1}') + with pytest.raises(ValueError, match='must be str→str'): + await store.get_metadata(uri) + + +class TestSqliteMediaStore: + async def test_put_get_round_trip(self, tmp_path: Path) -> None: + store = SqliteMediaStore(database=tmp_path / 'media.db') + uri = await store.put(b'hello bytes', context=MediaContext(media_type='image/png')) + assert await store.get(uri) == b'hello bytes' + + async def test_put_persists_metadata_column(self, tmp_path: Path) -> None: + store = SqliteMediaStore(database=tmp_path / 'media.db') + uri = await store.put( + b'tagged', + context=MediaContext(media_type='image/png', metadata={'origin': 'user', 'tenant': 'acme'}), + ) + conn = sqlite3.connect(tmp_path / 'media.db', check_same_thread=False) + try: + row = conn.execute('SELECT media_type, metadata FROM media').fetchone() + finally: + conn.close() + assert row[0] == 'image/png' + assert await store.get_metadata(uri) == {'origin': 'user', 'tenant': 'acme'} + + async def test_get_metadata_missing_raises(self, tmp_path: Path) -> None: + store = SqliteMediaStore(database=tmp_path / 'media.db') + with pytest.raises(FileNotFoundError): + await store.get_metadata(media_uri_for(b'never put')) + + async def test_with_shared_connection(self) -> None: + connection = sqlite3.connect(':memory:', check_same_thread=False) + try: + store = SqliteMediaStore(connection=connection) + uri = await store.put(b'shared conn data') + assert await store.get(uri) == b'shared conn data' + assert await store.exists(uri) is True + finally: + connection.close() + + async def test_dedup_via_insert_or_ignore(self, tmp_path: Path) -> None: + store = SqliteMediaStore(database=tmp_path / 'm.db') + uri = '' + for _ in range(3): + uri = await store.put(b'duplicated') + assert await store.get(uri) == b'duplicated' + + async def test_get_missing_raises(self, tmp_path: Path) -> None: + store = SqliteMediaStore(database=tmp_path / 'm.db') + with pytest.raises(FileNotFoundError): + await store.get(media_uri_for(b'missing')) + + def test_requires_exactly_one_of_database_or_connection(self) -> None: + with pytest.raises(ValueError, match='exactly one'): + SqliteMediaStore() + conn = sqlite3.connect(':memory:') + try: + with pytest.raises(ValueError, match='exactly one'): + SqliteMediaStore(database='x', connection=conn) + finally: + conn.close() + + def test_rejects_invalid_table_name(self) -> None: + with pytest.raises(ValueError, match='invalid table name'): + SqliteMediaStore(database='x.db', table='bad-name') + + +class TestExternalizeRestoreWalker: + async def test_round_trip_with_inline_binary(self, tmp_path: Path) -> None: + import base64 + import json as _json + + store = DiskMediaStore(tmp_path) + big_payload = b'\x00' * 70_000 + b64_payload = base64.b64encode(big_payload).decode('ascii') + node: object = { + 'parts': [ + { + 'kind': 'binary', + 'data': b64_payload, + 'media_type': 'image/png', + 'identifier': 'abc', + 'vendor_metadata': None, + } + ] + } + externalized = await externalize_media(node, media_store=store, threshold_bytes=64 * 1024) + externalized_text = _json.dumps(externalized) + assert '__harness_external_media__' in externalized_text + assert 'media+sha256://' in externalized_text + assert b64_payload not in externalized_text # bytes really went external + + restored = await restore_media(externalized, media_store=store) + restored_text = _json.dumps(restored) + assert '"kind": "binary"' in restored_text + assert b64_payload in restored_text # bytes restored exactly + + async def test_threshold_boundary_keeps_small_inline(self, tmp_path: Path) -> None: + import base64 + + store = DiskMediaStore(tmp_path) + small_payload = b'\x42' * 32 + node = { + 'kind': 'binary', + 'data': base64.b64encode(small_payload).decode('ascii'), + 'media_type': 'text/plain', + 'identifier': 's', + 'vendor_metadata': None, + } + externalized = await externalize_media(node, media_store=store, threshold_bytes=64 * 1024) + assert externalized == node + assert list(tmp_path.glob('*.bin')) == [] + + async def test_restore_raises_when_marker_missing_uri(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + bad_node = {'__harness_external_media__': True, 'media_type': 'image/png'} + with pytest.raises(ValueError, match='missing string uri'): + await restore_media(bad_node, media_store=store) + + async def test_round_trip_preserves_unknown_binary_fields(self, tmp_path: Path) -> None: + """A field the walker doesn't know about survives externalize -> restore.""" + import base64 + import json as _json + + store = DiskMediaStore(tmp_path) + big = base64.b64encode(b'\x01' * 70_000).decode('ascii') + node = { + 'kind': 'binary', + 'data': big, + 'media_type': 'image/png', + 'identifier': 'abc', + 'vendor_metadata': {'x': 1}, + 'future_field': 'keep-me', + } + externalized = await externalize_media(node, media_store=store, threshold_bytes=64 * 1024) + assert '"data"' not in _json.dumps(externalized) # bytes field went external + restored = await restore_media(externalized, media_store=store) + assert restored == node + + async def test_restore_raises_when_blob_is_missing(self, tmp_path: Path) -> None: + """A well-formed marker whose blob was pruned surfaces the store's error.""" + store = DiskMediaStore(tmp_path) + marker = { + '__harness_external_media__': True, + 'uri': 'media+sha256://' + '0' * 64, + 'kind': 'binary', + 'media_type': 'image/png', + } + with pytest.raises(FileNotFoundError): + await restore_media(marker, media_store=store) + + async def test_walker_passes_through_scalars(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + for value in [None, 1, 'foo', True]: + assert await externalize_media(value, media_store=store, threshold_bytes=1) == value + assert await restore_media(value, media_store=store) == value + + +class TestSigV4Signer: + def test_produces_required_headers(self) -> None: + headers = sign_request( + method='PUT', + host='examplebucket.s3.amazonaws.com', + path='/my-key', + body=b'payload', + region='us-east-1', + access_key_id='AKIAIOSFODNN7EXAMPLE', + secret_access_key='wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY', + content_type='image/png', + now=datetime(2013, 5, 24, tzinfo=timezone.utc), + ) + assert headers['host'] == 'examplebucket.s3.amazonaws.com' + assert headers['x-amz-date'] == '20130524T000000Z' + assert headers['x-amz-content-sha256'] == ('239f59ed55e737c77147cf55ad0c1b030b6d7ee748a7426952f9b852d5a935e5') + assert headers['content-type'] == 'image/png' + auth = headers['authorization'] + assert auth.startswith('AWS4-HMAC-SHA256 ') + assert 'Credential=AKIAIOSFODNN7EXAMPLE/20130524/us-east-1/s3/aws4_request' in auth + assert 'SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date' in auth + # Signature is deterministic for fixed inputs — locks down the algorithm. + assert 'Signature=' in auth + + def test_signature_is_deterministic(self) -> None: + common = dict( + method='GET', + host='bucket.s3.amazonaws.com', + path='/key', + body=b'', + region='us-east-1', + access_key_id='K', + secret_access_key='S', + content_type=None, + now=datetime(2024, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + ) + a = sign_request(**common) # type: ignore[arg-type] + b = sign_request(**common) # type: ignore[arg-type] + assert a == b + + def test_signature_changes_with_body(self) -> None: + kwargs = dict( + method='PUT', + host='bucket.s3.amazonaws.com', + path='/key', + region='us-east-1', + access_key_id='K', + secret_access_key='S', + content_type=None, + now=datetime(2024, 1, 1, tzinfo=timezone.utc), + ) + sig_a = sign_request(body=b'aaa', **kwargs)['authorization'] # type: ignore[arg-type] + sig_b = sign_request(body=b'bbb', **kwargs)['authorization'] # type: ignore[arg-type] + assert sig_a != sig_b + + +class TestS3MediaStoreWithMockTransport: + async def test_put_signs_request(self) -> None: + captured: list[Request] = [] + + async def handler(request: Request) -> Response: + captured.append(request) + return Response(200) + + transport = MockTransport(handler) + async with AsyncClient(transport=transport) as client: + store = S3MediaStore( + bucket='my-bucket', + endpoint='https://s3.us-east-1.amazonaws.com', + region='us-east-1', + access_key_id='AKIA-FAKE', + secret_access_key='secret-fake', + client=client, + ) + uri = await store.put(b'payload', context=MediaContext(media_type='image/png')) + + assert uri.startswith('media+sha256://') + assert len(captured) == 1 + request = captured[0] + assert request.method == 'PUT' + assert 'authorization' in request.headers + assert request.headers['authorization'].startswith('AWS4-HMAC-SHA256 ') + assert request.headers['x-amz-content-sha256'] + assert request.headers['content-type'] == 'image/png' + + async def test_wire_path_matches_signed_path_for_reserved_chars(self) -> None: + """A key_prefix with reserved chars must be percent-encoded identically on the wire. + + Otherwise the signed canonical path and the path S3 receives diverge -> + SignatureDoesNotMatch. The wire `raw_path` must equal `_canonical_uri(path)`. + """ + from pydantic_ai_harness.experimental.media._s3 import _canonical_uri # pyright: ignore[reportPrivateUsage] + + captured: list[Request] = [] + + async def handler(request: Request) -> Response: + captured.append(request) + return Response(200) + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='my-bucket', + endpoint='https://acct.r2.cloudflarestorage.com', + region='auto', + access_key_id='k', + secret_access_key='s', + key_prefix='runs(2024)/tenant@acme/', + client=client, + ) + uri = await store.put(b'payload') + + expected_path = _canonical_uri(store._object_path(uri, MediaContext())) # pyright: ignore[reportPrivateUsage] + assert '%28' in expected_path and '%40' in expected_path # ( and @ encoded + assert captured[0].url.raw_path.decode('ascii') == expected_path + + async def test_put_propagates_metadata_as_signed_x_amz_meta_headers(self) -> None: + captured: list[Request] = [] + + async def handler(request: Request) -> Response: + captured.append(request) + return Response(200) + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + await store.put( + b'meta-tagged', + context=MediaContext( + media_type='image/png', + metadata={'origin': 'pipeline-a', 'tenant': 'acme'}, + ), + ) + request = captured[0] + assert request.headers.get('x-amz-meta-origin') == 'pipeline-a' + assert request.headers.get('x-amz-meta-tenant') == 'acme' + # Metadata headers MUST be in SignedHeaders to be valid SigV4. + signed = request.headers['authorization'].split('SignedHeaders=')[1].split(',')[0] + assert 'x-amz-meta-origin' in signed + assert 'x-amz-meta-tenant' in signed + + async def test_put_rejects_invalid_metadata_key(self) -> None: + async def handler(request: Request) -> Response: # pragma: no cover + return Response(200) # never reached — validation raises pre-request + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + with pytest.raises(ValueError, match='not a valid HTTP header token'): + await store.put( + b'data', + context=MediaContext(metadata={'bad key with space': 'v'}), + ) + + async def test_get_resolves_uri(self) -> None: + async def handler(request: Request) -> Response: + return Response(200, content=b'fetched bytes') + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + data = await store.get(media_uri_for(b'fetched bytes')) + assert data == b'fetched bytes' + + async def test_get_404_raises(self) -> None: + async def handler(request: Request) -> Response: + return Response(404) + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + with pytest.raises(FileNotFoundError): + await store.get(media_uri_for(b'missing')) + + async def test_exists_uses_head(self) -> None: + captured_methods: list[str] = [] + + async def handler(request: Request) -> Response: + captured_methods.append(request.method) + return Response(200) + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + assert await store.exists(media_uri_for(b'anything')) is True + + assert captured_methods == ['HEAD'] + + async def test_exists_returns_false_on_404(self) -> None: + async def handler(request: Request) -> Response: + return Response(404) + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + assert await store.exists(media_uri_for(b'missing')) is False + + async def test_put_failure_raises(self) -> None: + async def handler(request: Request) -> Response: + return Response(500, content=b'internal error') + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + with pytest.raises(RuntimeError, match='S3 PUT failed'): + await store.put(b'data') + + async def test_get_failure_raises(self) -> None: + async def handler(request: Request) -> Response: + return Response(500, content=b'oops') + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + with pytest.raises(RuntimeError, match='S3 GET failed'): + await store.get(media_uri_for(b'x')) + + async def test_get_metadata_reads_x_amz_meta_headers(self) -> None: + async def handler(request: Request) -> Response: + return Response( + 200, + headers={ + 'x-amz-meta-origin': 'pipeline-a', + 'x-amz-meta-tenant': 'acme', + 'content-type': 'image/png', + 'content-length': '0', + }, + ) + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + meta = await store.get_metadata(media_uri_for(b'x')) + assert meta == {'origin': 'pipeline-a', 'tenant': 'acme'} + + async def test_get_metadata_404_raises(self) -> None: + async def handler(request: Request) -> Response: + return Response(404) + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + with pytest.raises(FileNotFoundError): + await store.get_metadata(media_uri_for(b'missing')) + + async def test_get_metadata_500_raises(self) -> None: + async def handler(request: Request) -> Response: + return Response(500, content=b'oops') + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + with pytest.raises(RuntimeError, match='S3 HEAD failed'): + await store.get_metadata(media_uri_for(b'x')) + + async def test_head_failure_raises(self) -> None: + async def handler(request: Request) -> Response: + return Response(500) + + async with AsyncClient(transport=MockTransport(handler)) as client: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + client=client, + ) + with pytest.raises(RuntimeError, match='S3 HEAD failed'): + await store.exists(media_uri_for(b'x')) + + async def test_no_client_branch_opens_one_per_call(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Without `client=`, the store opens a fresh `httpx.AsyncClient` per call.""" + import httpx + + captured: list[Request] = [] + + async def handler(request: Request) -> Response: + captured.append(request) + return Response(200) + + original_async_client = httpx.AsyncClient + + def patched_client(*args: object, **kwargs: object) -> httpx.AsyncClient: + return original_async_client(transport=MockTransport(handler)) + + monkeypatch.setattr(httpx, 'AsyncClient', patched_client) + + store = S3MediaStore( + bucket='my-bucket', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + key_prefix='runs/', + ) + await store.put(b'no-client data') + assert len(captured) == 1 + sample_uri = media_uri_for(b'sample') + digest = parse_media_uri(sample_uri) + path = store._object_path(sample_uri, MediaContext()) # type: ignore[reportPrivateUsage] + assert path == f'/my-bucket/runs/{digest}.bin' + + +@pytest.mark.skipif( # pragma: no cover + not all( + os.environ.get(k) + for k in ('S3_ENDPOINT', 'S3_BUCKET_NAME', 'S3_ACCESS_KEY_ID', 'S3_SECRET_ACCESS_KEY', 'S3_REGION') + ), + reason='live S3/R2 env vars not set', +) +class TestS3MediaStoreLive: # pragma: no cover + """Live integration against the configured S3 endpoint (e.g. R2). + + Activated only when all five S3_* env vars are set. Reads creds out of + the environment — the test harness inherits them when invoked with + `~/.claude/scripts/env-run .env -- make test`. + """ + + async def test_round_trip_against_live_endpoint(self) -> None: + bucket = os.environ['S3_BUCKET_NAME'] + store = S3MediaStore( + bucket=bucket, + endpoint=os.environ['S3_ENDPOINT'], + region=os.environ['S3_REGION'], + access_key_id=os.environ['S3_ACCESS_KEY_ID'], + secret_access_key=os.environ['S3_SECRET_ACCESS_KEY'], + key_prefix='harness-test/', + ) + payload = b'pydantic-ai-harness step-persistence live test ' + os.urandom(32) + ctx = MediaContext(media_type='application/octet-stream') + uri = await store.put(payload, context=ctx) + assert await store.exists(uri) is True + fetched = await store.get(uri) + assert fetched == payload + + +class TestPublicUrl: + """Verify `public_url` resolves through the user-supplied callable. + + Static prefix, async callable, and the absence of a resolver all go + through the same `MediaStore.public_url(uri)` shape — the future + `MediaExternalizer` capability uses this to swap `BinaryContent` for + URL message parts before the model sees them. + """ + + async def test_disk_store_without_resolver_returns_none(self, tmp_path: Path) -> None: + store = DiskMediaStore(tmp_path) + assert await store.public_url(media_uri_for(b'x')) is None + + async def test_sqlite_store_without_resolver_returns_none(self, tmp_path: Path) -> None: + store = SqliteMediaStore(database=tmp_path / 'm.db') + assert await store.public_url(media_uri_for(b'x')) is None + + async def test_s3_store_without_resolver_returns_none(self) -> None: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + ) + assert await store.public_url(media_uri_for(b'x')) is None + + async def test_sync_callable_resolver(self, tmp_path: Path) -> None: + store = DiskMediaStore( + tmp_path, + public_url=lambda uri, ctx: f'https://cdn.example.com/{parse_media_uri(uri)}.bin', + ) + uri = media_uri_for(b'payload') + result = await store.public_url(uri) + assert result == f'https://cdn.example.com/{parse_media_uri(uri)}.bin' + + async def test_async_callable_resolver(self, tmp_path: Path) -> None: + async def signer(uri: str, ctx: MediaContext) -> str: + return f'https://signed.example.com/{parse_media_uri(uri)}?sig=abc' + + store = DiskMediaStore(tmp_path, public_url=signer) + result = await store.public_url(media_uri_for(b'payload')) + assert result is not None + assert result.startswith('https://signed.example.com/') + assert '?sig=abc' in result + + async def test_resolver_receives_media_context(self, tmp_path: Path) -> None: + """The resolver sees the full `MediaContext` so it can vary by media type, metadata, etc.""" + seen: list[MediaContext] = [] + + def resolver(uri: str, ctx: MediaContext) -> str: + seen.append(ctx) + return 'https://example.com/x' + + store = DiskMediaStore(tmp_path, public_url=resolver) + await store.public_url( + media_uri_for(b'p'), + context=MediaContext(media_type='image/png', metadata={'tag': 'v'}), + ) + assert seen[0].media_type == 'image/png' + assert dict(seen[0].metadata) == {'tag': 'v'} + + async def test_resolver_can_return_none(self, tmp_path: Path) -> None: + """Resolvers may opt out per-URI (e.g. small payloads keep inline).""" + store = DiskMediaStore(tmp_path, public_url=lambda uri, ctx: None) + assert await store.public_url(media_uri_for(b'x')) is None + + async def test_make_static_public_url_helper(self, tmp_path: Path) -> None: + resolver = make_static_public_url('https://pub-abc.r2.dev', key_prefix='media/', extension='.bin') + store = DiskMediaStore(tmp_path, public_url=resolver) + uri = media_uri_for(b'payload') + digest = parse_media_uri(uri) + assert await store.public_url(uri) == f'https://pub-abc.r2.dev/media/{digest}.bin' + + async def test_make_static_public_url_strips_trailing_slash(self, tmp_path: Path) -> None: + resolver = make_static_public_url('https://cdn.example.com/') + store = DiskMediaStore(tmp_path, public_url=resolver) + digest = parse_media_uri(media_uri_for(b'p')) + assert await store.public_url(media_uri_for(b'p')) == f'https://cdn.example.com/{digest}.bin' + + async def test_s3_with_resolver_uses_it(self) -> None: + store = S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + key_prefix='media/', + public_url=make_static_public_url('https://pub-xyz.r2.dev', key_prefix='media/'), + ) + uri = media_uri_for(b'p') + digest = parse_media_uri(uri) + assert await store.public_url(uri) == f'https://pub-xyz.r2.dev/media/{digest}.bin' + + async def test_sqlite_with_resolver_uses_it(self, tmp_path: Path) -> None: + store = SqliteMediaStore( + database=tmp_path / 'm.db', + public_url=make_static_public_url('https://cdn.example.com'), + ) + uri = media_uri_for(b'p') + digest = parse_media_uri(uri) + assert await store.public_url(uri) == f'https://cdn.example.com/{digest}.bin' + + +def _assert_media_store_protocol(store: MediaStore) -> None: + """Mypy/pyright check: every concrete store satisfies the protocol.""" + assert isinstance(store, MediaStore) + + +def test_concrete_stores_satisfy_protocol(tmp_path: Path) -> None: + _assert_media_store_protocol(DiskMediaStore(tmp_path)) + _assert_media_store_protocol(SqliteMediaStore(database=tmp_path / 'm.db')) + _assert_media_store_protocol( + S3MediaStore( + bucket='b', + endpoint='https://example.com', + region='auto', + access_key_id='k', + secret_access_key='s', + ) + ) diff --git a/tests/experimental/media/test_s3_cassettes.py b/tests/experimental/media/test_s3_cassettes.py new file mode 100644 index 0000000..1414aa8 --- /dev/null +++ b/tests/experimental/media/test_s3_cassettes.py @@ -0,0 +1,63 @@ +"""VCR cassette tests for `S3MediaStore` against a real R2 endpoint. + +The cassettes under `tests/media/cassettes/` were recorded against a +Cloudflare R2 bucket and then sanitised (account id, bucket name, +`Authorization`, `x-amz-date`, identifying response headers, error +bodies) by hooks in `conftest.py::vcr_config`. CI replays them — no +creds needed. + +To re-record (or add new cases): + +``` +~/.claude/scripts/env-run .env -- uv run pytest \\ + tests/media/test_s3_cassettes.py \\ + --record-mode=once +``` + +Recording reads real R2 creds from the `s3_credentials` fixture; replay +falls back to the sanitised placeholders that match the committed +cassettes. Drift between real values and placeholders (e.g. scrubber +misses a field) shows up as a replay-time URL mismatch — the canary. +""" + +from __future__ import annotations + +import pytest + +from pydantic_ai_harness.experimental.media import MediaContext, S3MediaStore, media_uri_for + +pytestmark = pytest.mark.anyio + +# Deterministic payload so the URI / object key are stable across re-records. +_PAYLOAD = b'pydantic-ai-harness step-persistence VCR cassette payload v1' +_MISSING_URI = 'media+sha256://' + ('0' * 64) +_CONTEXT = MediaContext(media_type='application/octet-stream') + + +class TestS3MediaStoreCassettes: + """Replay-driven verification of `S3MediaStore` against a real R2 server.""" + + @pytest.mark.vcr + async def test_put_succeeds(self, s3_store: S3MediaStore) -> None: + uri = await s3_store.put(_PAYLOAD, context=_CONTEXT) + assert uri == media_uri_for(_PAYLOAD) + + @pytest.mark.vcr + async def test_exists_present_after_put(self, s3_store: S3MediaStore) -> None: + await s3_store.put(_PAYLOAD, context=_CONTEXT) + assert await s3_store.exists(media_uri_for(_PAYLOAD)) is True + + @pytest.mark.vcr + async def test_exists_returns_false_for_missing(self, s3_store: S3MediaStore) -> None: + assert await s3_store.exists(_MISSING_URI) is False + + @pytest.mark.vcr + async def test_get_round_trips_bytes(self, s3_store: S3MediaStore) -> None: + await s3_store.put(_PAYLOAD, context=_CONTEXT) + fetched = await s3_store.get(media_uri_for(_PAYLOAD)) + assert fetched == _PAYLOAD + + @pytest.mark.vcr + async def test_get_missing_raises(self, s3_store: S3MediaStore) -> None: + with pytest.raises(FileNotFoundError): + await s3_store.get(_MISSING_URI) diff --git a/tests/experimental/step_persistence/__init__.py b/tests/experimental/step_persistence/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/experimental/step_persistence/test_sqlite_and_media.py b/tests/experimental/step_persistence/test_sqlite_and_media.py new file mode 100644 index 0000000..6cf861b --- /dev/null +++ b/tests/experimental/step_persistence/test_sqlite_and_media.py @@ -0,0 +1,578 @@ +"""Tests for `SqliteStepStore` and the media externalization paths. + +Mirrors selected coverage from `test_step_persistence.py::TestFileStepStore` +for the SQLite backend, plus end-to-end media-externalization round-trips +through `FileStepStore` and `SqliteStepStore`. The full hook-level coverage +already lives in `test_step_persistence.py`; this module focuses on the +SQLite-specific behavior and the media plumbing. +""" + +from __future__ import annotations + +import base64 +import sqlite3 +from pathlib import Path + +import pytest +from pydantic_ai import Agent +from pydantic_ai.messages import ( + BinaryContent, + ModelMessage, + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, +) +from pydantic_ai.models.test import TestModel + +from pydantic_ai_harness.experimental.media import DiskMediaStore, SqliteMediaStore +from pydantic_ai_harness.experimental.step_persistence import ( + ContinuableSnapshot, + FileStepStore, + RunRecord, + SqliteStepStore, + StepEvent, + StepPersistence, + ToolEffectRecord, +) + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend() -> str: + return 'asyncio' + + +def _sample_messages_with_media(payload_size: int) -> list[ModelMessage]: + big = b'\xab' * payload_size + return [ + ModelRequest( + parts=[ + UserPromptPart( + content=[ + 'analyze this', + BinaryContent(data=big, media_type='image/png'), + ] + ) + ] + ), + ModelResponse(parts=[TextPart(content='done')]), + ] + + +# --------------------------------------------------------------------------- +# SqliteStepStore — direct protocol exercises +# --------------------------------------------------------------------------- + + +class TestSqliteStepStoreProtocol: + async def test_register_and_get_run(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + record = RunRecord(run_id='r1', conversation_id='c1', agent_name='agent', metadata={'k': 'v'}) + await store.register_run(record) + fetched = await store.get_run(run_id='r1') + assert fetched is not None + assert fetched.run_id == 'r1' + assert fetched.conversation_id == 'c1' + assert fetched.metadata == {'k': 'v'} + + async def test_register_duplicate_run_raises(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + await store.register_run(RunRecord(run_id='r1')) + with pytest.raises(sqlite3.IntegrityError): + await store.register_run(RunRecord(run_id='r1')) + + async def test_list_runs_chronological(self, tmp_path: Path) -> None: + from datetime import datetime, timedelta, timezone + + store = SqliteStepStore(database=tmp_path / 'runs.db') + base = datetime(2024, 1, 1, tzinfo=timezone.utc) + await store.register_run(RunRecord(run_id='r3', started_at=base + timedelta(seconds=3))) + await store.register_run(RunRecord(run_id='r1', started_at=base + timedelta(seconds=1))) + await store.register_run(RunRecord(run_id='r2', started_at=base + timedelta(seconds=2))) + + records = await store.list_runs() + assert [r.run_id for r in records] == ['r1', 'r2', 'r3'] + + async def test_list_runs_filters(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + await store.register_run(RunRecord(run_id='r1', conversation_id='a', parent_run_id='p')) + await store.register_run(RunRecord(run_id='r2', conversation_id='a', parent_run_id='q')) + await store.register_run(RunRecord(run_id='r3', conversation_id='b', parent_run_id='p')) + + by_conv = await store.list_runs(conversation_id='a') + assert {r.run_id for r in by_conv} == {'r1', 'r2'} + + by_parent = await store.list_runs(parent_run_id='p') + assert {r.run_id for r in by_parent} == {'r1', 'r3'} + + both = await store.list_runs(parent_run_id='p', conversation_id='a') + assert [r.run_id for r in both] == ['r1'] + + async def test_append_and_list_events(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + await store.register_run(RunRecord(run_id='r1')) + await store.append_event(StepEvent(run_id='r1', kind='run_started', step_index=0)) + await store.append_event( + StepEvent( + run_id='r1', + kind='tool_call_started', + step_index=1, + tool_call_id='t1', + tool_name='add', + metadata={'k': 'v'}, + ) + ) + events = await store.list_events(run_id='r1') + assert [e.kind for e in events] == ['run_started', 'tool_call_started'] + assert events[1].tool_call_id == 't1' + assert events[1].metadata == {'k': 'v'} + + async def test_save_and_load_snapshot(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db', media_store=None) + await store.register_run(RunRecord(run_id='r1')) + messages: list[ModelMessage] = [ + ModelRequest(parts=[UserPromptPart(content='hello')]), + ModelResponse(parts=[TextPart(content='hi back')]), + ] + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=2, messages=messages)) + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + assert snap.step_index == 2 + assert len(snap.messages) == 2 + + async def test_snapshot_seq_monotonic_across_reset_step(self, tmp_path: Path) -> None: + """A reused `run_id` with `step_index` reset to 0 must not clobber the prior snapshot. + + Mirrors `FileStepStore._next_snapshot_seq` — SQLite uses + `AUTOINCREMENT seq` for the same guarantee. + """ + store = SqliteStepStore(database=tmp_path / 'runs.db', media_store=None) + await store.register_run(RunRecord(run_id='r1')) + msgs: list[ModelMessage] = [ModelRequest(parts=[UserPromptPart(content='a')])] + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=5, messages=msgs)) + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=msgs)) + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + assert snap.step_index == 0 # the *last* write wins, not the highest step_index + + async def test_tool_effect_upsert_per_call_id(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='t1', tool_name='add', run_id='r1', status='started') + ) + await store.record_tool_effect( + ToolEffectRecord( + tool_call_id='t1', + tool_name='add', + run_id='r1', + status='completed', + effect_summary='ok', + ) + ) + effect = await store.get_tool_effect(run_id='r1', tool_call_id='t1') + assert effect is not None + assert effect.status == 'completed' + assert effect.effect_summary == 'ok' + + async def test_tool_effect_scoped_by_run(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='t1', tool_name='add', run_id='r1', status='started') + ) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='t1', tool_name='add', run_id='r2', status='completed') + ) + a = await store.get_tool_effect(run_id='r1', tool_call_id='t1') + b = await store.get_tool_effect(run_id='r2', tool_call_id='t1') + assert a is not None and a.status == 'started' + assert b is not None and b.status == 'completed' + + async def test_list_unresolved_tool_effects(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='t1', tool_name='add', run_id='r1', status='started') + ) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='t2', tool_name='mul', run_id='r1', status='completed') + ) + unresolved = await store.list_unresolved_tool_effects(run_id='r1') + assert [r.tool_call_id for r in unresolved] == ['t1'] + + async def test_missing_lookups_return_none_or_empty(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + assert await store.get_run(run_id='nope') is None + assert await store.latest_snapshot(run_id='nope') is None + assert await store.get_tool_effect(run_id='nope', tool_call_id='x') is None + assert await store.list_events(run_id='nope') == [] + assert await store.list_unresolved_tool_effects(run_id='nope') == [] + assert await store.list_runs() == [] + + async def test_shared_connection_mode(self, tmp_path: Path) -> None: + conn = sqlite3.connect(tmp_path / 'runs.db', check_same_thread=False) + try: + store = SqliteStepStore(connection=conn, media_store=None) + await store.register_run(RunRecord(run_id='r1')) + await store.append_event(StepEvent(run_id='r1', kind='run_started', step_index=0)) + assert (await store.get_run(run_id='r1')) is not None + assert len(await store.list_events(run_id='r1')) == 1 + finally: + conn.close() + + async def test_shared_connection_default_media_store_uses_same_connection(self, tmp_path: Path) -> None: + """`media_store='auto'` + `connection=` paths a SqliteMediaStore at the same conn.""" + conn = sqlite3.connect(tmp_path / 'runs.db', check_same_thread=False) + try: + store = SqliteStepStore(connection=conn) + await store.register_run(RunRecord(run_id='r1')) + await store.save_snapshot( + ContinuableSnapshot(run_id='r1', step_index=0, messages=_sample_messages_with_media(70_000)) + ) + (count,) = conn.execute('SELECT COUNT(*) FROM media').fetchone() + assert count == 1 + finally: + conn.close() + + def test_rejects_both_database_and_connection(self) -> None: + with pytest.raises(ValueError, match='exactly one'): + SqliteStepStore() + conn = sqlite3.connect(':memory:') + try: + with pytest.raises(ValueError, match='exactly one'): + SqliteStepStore(database='x', connection=conn) + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# Media integration: FileStepStore + DiskMediaStore (default) +# --------------------------------------------------------------------------- + + +class TestFileStepStoreMedia: + async def test_large_binary_externalized_to_media_dir(self, tmp_path: Path) -> None: + """Default FileStepStore wires a `DiskMediaStore(/media/)`. + + Large `BinaryContent` payloads end up as `/media/.bin`, + not inlined in the snapshot JSON. + """ + store = FileStepStore(tmp_path / 'runs', media_threshold_bytes=64 * 1024) + await store.register_run(RunRecord(run_id='r1')) + await store.save_snapshot( + ContinuableSnapshot(run_id='r1', step_index=0, messages=_sample_messages_with_media(70_000)) + ) + + media_dir = tmp_path / 'runs' / 'media' + assert media_dir.is_dir() + blobs = list(media_dir.glob('*.bin')) + assert len(blobs) == 1 + assert blobs[0].stat().st_size == 70_000 + + # Snapshot JSON references the URI, not the inline base64. + snap_files = list((tmp_path / 'runs' / 'r1' / 'snapshots').glob('*.json')) + assert len(snap_files) == 1 + snap_text = snap_files[0].read_text(encoding='utf-8') + assert '__harness_external_media__' in snap_text + assert base64.b64encode(b'\xab' * 70_000).decode('ascii') not in snap_text + + async def test_round_trip_restores_original_bytes(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path / 'runs') + await store.register_run(RunRecord(run_id='r1')) + payload = b'\xcd' * 80_000 + messages: list[ModelMessage] = [ + ModelRequest(parts=[UserPromptPart(content=[BinaryContent(data=payload, media_type='image/jpeg')])]), + ModelResponse(parts=[TextPart(content='ok')]), + ] + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=messages)) + + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + first = snap.messages[0] + assert isinstance(first, ModelRequest) + prompt = first.parts[0] + assert isinstance(prompt, UserPromptPart) + assert isinstance(prompt.content, list) + binary = prompt.content[0] + assert isinstance(binary, BinaryContent) + assert binary.data == payload + assert binary.media_type == 'image/jpeg' + + async def test_below_threshold_stays_inline(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path / 'runs', media_threshold_bytes=64 * 1024) + await store.register_run(RunRecord(run_id='r1')) + await store.save_snapshot( + ContinuableSnapshot(run_id='r1', step_index=0, messages=_sample_messages_with_media(1_024)) + ) + media_dir = tmp_path / 'runs' / 'media' + assert not media_dir.exists() or list(media_dir.glob('*.bin')) == [] + + async def test_opt_out_keeps_bytes_inline(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path / 'runs', media_store=None) + await store.register_run(RunRecord(run_id='r1')) + await store.save_snapshot( + ContinuableSnapshot(run_id='r1', step_index=0, messages=_sample_messages_with_media(70_000)) + ) + assert not (tmp_path / 'runs' / 'media').exists() + snap_text = next((tmp_path / 'runs' / 'r1' / 'snapshots').glob('*.json')).read_text(encoding='utf-8') + assert '__harness_external_media__' not in snap_text + + # Loading also works with media_store=None — no restore_media walk. + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + + async def test_explicit_media_store_override(self, tmp_path: Path) -> None: + custom_root = tmp_path / 'shared-media' + store = FileStepStore( + tmp_path / 'runs', + media_store=DiskMediaStore(custom_root), + ) + await store.register_run(RunRecord(run_id='r1')) + await store.save_snapshot( + ContinuableSnapshot(run_id='r1', step_index=0, messages=_sample_messages_with_media(70_000)) + ) + # External bytes live under the override, not under /media/ + assert list(custom_root.glob('*.bin')) + assert not (tmp_path / 'runs' / 'media').exists() + + +# --------------------------------------------------------------------------- +# Media integration: SqliteStepStore + same-DB SqliteMediaStore (default) +# --------------------------------------------------------------------------- + + +class TestSqliteStepStoreMedia: + async def test_large_binary_stored_in_same_db(self, tmp_path: Path) -> None: + db = tmp_path / 'runs.db' + store = SqliteStepStore(database=db) + await store.register_run(RunRecord(run_id='r1')) + await store.save_snapshot( + ContinuableSnapshot(run_id='r1', step_index=0, messages=_sample_messages_with_media(70_000)) + ) + + conn = sqlite3.connect(db, check_same_thread=False) + try: + row = conn.execute('SELECT COUNT(*), SUM(size_bytes) FROM media').fetchone() + assert row[0] == 1 + assert row[1] == 70_000 + # No sibling DB file or media directory. + assert not (tmp_path / 'media').exists() + finally: + conn.close() + + async def test_round_trip_restores_bytes(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'runs.db') + await store.register_run(RunRecord(run_id='r1')) + payload = b'\xef' * 70_000 + messages: list[ModelMessage] = [ + ModelRequest(parts=[UserPromptPart(content=[BinaryContent(data=payload, media_type='image/png')])]), + ModelResponse(parts=[TextPart(content='done')]), + ] + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=messages)) + + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + first = snap.messages[0] + assert isinstance(first, ModelRequest) + prompt = first.parts[0] + assert isinstance(prompt, UserPromptPart) + assert isinstance(prompt.content, list) + binary = prompt.content[0] + assert isinstance(binary, BinaryContent) + assert binary.data == payload + + async def test_dedup_across_snapshots(self, tmp_path: Path) -> None: + """The same payload appearing in two snapshots ends up as one row.""" + db = tmp_path / 'runs.db' + store = SqliteStepStore(database=db) + await store.register_run(RunRecord(run_id='r1')) + for step in range(3): + await store.save_snapshot( + ContinuableSnapshot(run_id='r1', step_index=step, messages=_sample_messages_with_media(70_000)) + ) + conn = sqlite3.connect(db, check_same_thread=False) + try: + (count,) = conn.execute('SELECT COUNT(*) FROM media').fetchone() + assert count == 1 + finally: + conn.close() + + async def test_external_media_store_override(self, tmp_path: Path) -> None: + """SqliteStepStore can be paired with a non-sqlite media store (e.g. disk).""" + disk = tmp_path / 'media' + store = SqliteStepStore(database=tmp_path / 'runs.db', media_store=DiskMediaStore(disk)) + await store.register_run(RunRecord(run_id='r1')) + await store.save_snapshot( + ContinuableSnapshot(run_id='r1', step_index=0, messages=_sample_messages_with_media(70_000)) + ) + assert len(list(disk.glob('*.bin'))) == 1 + # The `media` table is never created in the same DB because we routed + # blobs to a disk store; the StepStore schema doesn't include it. + conn = sqlite3.connect(tmp_path / 'runs.db', check_same_thread=False) + try: + row = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='media'").fetchone() + assert row is None + finally: + conn.close() + + +# --------------------------------------------------------------------------- +# End-to-end: Agent + StepPersistence + media externalization +# --------------------------------------------------------------------------- + + +class TestEndToEndMediaWithAgent: + async def test_agent_run_with_large_binary_input(self, tmp_path: Path) -> None: + """An Agent run carrying BinaryContent in its prompt round-trips through SP.""" + store = FileStepStore(tmp_path / 'runs') + agent: Agent[None, str] = Agent( + TestModel(), + capabilities=[StepPersistence(store=store, agent_name='vision')], + ) + + big = b'\xab' * 100_000 + result = await agent.run( + [ + 'classify this image', + BinaryContent(data=big, media_type='image/png'), + ] + ) + + # Sanity: agent produced something. + assert isinstance(result.output, str) + + runs = await store.list_runs() + assert len(runs) == 1 + run_id = runs[0].run_id + + # Snapshot exists, references external media. + snap = await store.latest_snapshot(run_id=run_id) + assert snap is not None + first = snap.messages[0] + assert isinstance(first, ModelRequest) + prompt = first.parts[0] + assert isinstance(prompt, UserPromptPart) + assert isinstance(prompt.content, list) + binary = next(p for p in prompt.content if isinstance(p, BinaryContent)) + assert binary.data == big + + # On disk: one media blob, no inlined base64 in any snapshot json. + blobs = list((tmp_path / 'runs' / 'media').glob('*.bin')) + assert len(blobs) == 1 + + +# --------------------------------------------------------------------------- +# Auto / opt-out / explicit media_store semantics on both stores +# --------------------------------------------------------------------------- + + +class TestSqliteRowDefenses: + """Exercise the defensive `isinstance`/`None` guards in the row decoders. + + These guards never trigger via the public API — they catch poke-by-hand + corruption of the DB. Hitting them explicitly is the cleanest way to + document the contract and the only way to satisfy 100% coverage. + """ + + def _open(self, db: Path) -> sqlite3.Connection: + return sqlite3.connect(db, check_same_thread=False, isolation_level=None) + + async def test_corrupted_run_row_raises(self, tmp_path: Path) -> None: + db = tmp_path / 'runs.db' + store = SqliteStepStore(database=db, media_store=None) + await store.register_run(RunRecord(run_id='r1')) + conn = self._open(db) + try: + # TEXT affinity coerces numbers to strings, but leaves BLOBs alone — + # binding `bytes` is the cleanest way to violate `isinstance(v, str)`. + conn.execute('UPDATE runs SET metadata = ? WHERE run_id = ?', (b'\x00bin', 'r1')) + finally: + conn.close() + with pytest.raises(ValueError, match='run row has wrong types'): + await store.get_run(run_id='r1') + + async def test_corrupted_event_row_wrong_type_raises(self, tmp_path: Path) -> None: + db = tmp_path / 'runs.db' + store = SqliteStepStore(database=db, media_store=None) + await store.register_run(RunRecord(run_id='r1')) + await store.append_event(StepEvent(run_id='r1', kind='run_started', step_index=0)) + conn = self._open(db) + try: + conn.execute('UPDATE events SET metadata = ? WHERE run_id = ?', (b'\x00bin', 'r1')) + finally: + conn.close() + with pytest.raises(ValueError, match='event row has wrong types'): + await store.list_events(run_id='r1') + + async def test_corrupted_event_kind_raises(self, tmp_path: Path) -> None: + db = tmp_path / 'runs.db' + store = SqliteStepStore(database=db, media_store=None) + await store.register_run(RunRecord(run_id='r1')) + await store.append_event(StepEvent(run_id='r1', kind='run_started', step_index=0)) + conn = self._open(db) + try: + conn.execute('UPDATE events SET kind = ? WHERE run_id = ?', ('fictional', 'r1')) + finally: + conn.close() + with pytest.raises(ValueError, match='unknown event kind'): + await store.list_events(run_id='r1') + + async def test_corrupted_tool_effect_wrong_type_raises(self, tmp_path: Path) -> None: + db = tmp_path / 'runs.db' + store = SqliteStepStore(database=db, media_store=None) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='t1', tool_name='add', run_id='r1', status='started') + ) + conn = self._open(db) + try: + conn.execute('UPDATE tool_effects SET tool_name = ? WHERE tool_call_id = ?', (b'\x00bin', 't1')) + finally: + conn.close() + with pytest.raises(ValueError, match='tool effect row has wrong types'): + await store.get_tool_effect(run_id='r1', tool_call_id='t1') + + async def test_corrupted_tool_effect_status_raises(self, tmp_path: Path) -> None: + db = tmp_path / 'runs.db' + store = SqliteStepStore(database=db, media_store=None) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='t1', tool_name='add', run_id='r1', status='started') + ) + conn = self._open(db) + try: + conn.execute('UPDATE tool_effects SET status = ? WHERE tool_call_id = ?', ('exploded', 't1')) + finally: + conn.close() + with pytest.raises(ValueError, match='unknown tool effect status'): + await store.get_tool_effect(run_id='r1', tool_call_id='t1') + + async def test_corrupted_snapshot_row_raises(self, tmp_path: Path) -> None: + db = tmp_path / 'runs.db' + store = SqliteStepStore(database=db, media_store=None) + await store.register_run(RunRecord(run_id='r1')) + msgs: list[ModelMessage] = [ModelRequest(parts=[UserPromptPart(content='x')])] + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=msgs)) + conn = self._open(db) + try: + conn.execute('UPDATE snapshots SET timestamp = ? WHERE run_id = ?', (b'\x00bin', 'r1')) + finally: + conn.close() + with pytest.raises(ValueError, match='snapshot row has wrong types'): + await store.latest_snapshot(run_id='r1') + + +class TestMediaStoreResolution: + def test_file_store_auto_creates_disk_media_store(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path / 'runs') + assert isinstance(store._media_store, DiskMediaStore) # type: ignore[reportPrivateUsage] + + def test_file_store_none_disables_media(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path / 'runs', media_store=None) + assert store._media_store is None # type: ignore[reportPrivateUsage] + + def test_sqlite_store_auto_uses_same_db(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'r.db') + assert isinstance(store._media_store, SqliteMediaStore) # type: ignore[reportPrivateUsage] + + def test_sqlite_store_none_disables_media(self, tmp_path: Path) -> None: + store = SqliteStepStore(database=tmp_path / 'r.db', media_store=None) + assert store._media_store is None # type: ignore[reportPrivateUsage] diff --git a/tests/experimental/step_persistence/test_step_persistence.py b/tests/experimental/step_persistence/test_step_persistence.py new file mode 100644 index 0000000..06ce4cc --- /dev/null +++ b/tests/experimental/step_persistence/test_step_persistence.py @@ -0,0 +1,1313 @@ +"""Tests for the `StepPersistence` capability. + +Exercises the public capability behavior through `Agent(...)`/`TestModel` and +covers the helper / store branches that are awkward to reach through a real +agent run (e.g. the path-traversal guard on `FileStepStore`). +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest +from pydantic_ai import Agent, RunContext +from pydantic_ai._agent_graph import GraphAgentState # pyright: ignore[reportPrivateUsage] +from pydantic_ai.messages import ( + ModelMessage, + ModelRequest, + ModelResponse, + RetryPromptPart, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) +from pydantic_ai.models import ModelRequestContext, ModelRequestParameters +from pydantic_ai.models.test import TestModel +from pydantic_ai.run import AgentRunResult +from pydantic_ai.tools import ToolDefinition +from pydantic_ai.usage import RunUsage + +from pydantic_ai_harness.experimental.step_persistence import ( + ContinuableSnapshot, + FileStepStore, + InMemoryStepStore, + RunRecord, + SqliteStepStore, + StepEvent, + StepPersistence, + StepStore, + ToolEffectRecord, + continue_run, + fork_run, + is_provider_valid, +) +from pydantic_ai_harness.experimental.step_persistence._store import _validate_id # pyright: ignore[reportPrivateUsage] + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend() -> str: + """Restrict async tests to asyncio (Agent.run uses `asyncio.create_task`).""" + return 'asyncio' + + +def build_run_context( + deps: object = None, + *, + run_id: str | None = None, + run_step: int = 0, + conversation_id: str | None = None, +) -> RunContext[Any]: + """Fabricate a minimal `RunContext` for direct hook invocation.""" + return RunContext[Any]( + deps=deps, + model=TestModel(), + usage=RunUsage(), + prompt=None, + messages=[], + run_step=run_step, + run_id=run_id, + conversation_id=conversation_id, + ) + + +def make_simple_agent(capabilities: list[Any]) -> Agent[None, str]: + agent: Agent[None, str] = Agent(TestModel(), capabilities=capabilities) + + @agent.tool_plain + def add(a: int, b: int) -> int: # pyright: ignore[reportUnusedFunction] + return a + b + + return agent + + +async def first_run_id(store: StepStore) -> str: + runs = await store.list_runs() + assert len(runs) >= 1 + return runs[0].run_id + + +# --------------------------------------------------------------------------- +# is_provider_valid +# --------------------------------------------------------------------------- + + +class TestIsProviderValid: + def test_empty_history_is_valid(self) -> None: + assert is_provider_valid([]) is True + + def test_matched_tool_call_is_valid(self) -> None: + messages: list[ModelMessage] = [ + ModelResponse(parts=[ToolCallPart(tool_name='add', args={}, tool_call_id='c1')]), + ModelRequest(parts=[ToolReturnPart(tool_name='add', content=3, tool_call_id='c1')]), + ] + assert is_provider_valid(messages) is True + + def test_unmatched_tool_call_is_invalid(self) -> None: + messages: list[ModelMessage] = [ + ModelResponse(parts=[ToolCallPart(tool_name='add', args={}, tool_call_id='c1')]) + ] + assert is_provider_valid(messages) is False + + def test_retry_prompt_resolves_a_tool_call(self) -> None: + messages: list[ModelMessage] = [ + ModelResponse(parts=[ToolCallPart(tool_name='add', args={}, tool_call_id='c1')]), + ModelRequest(parts=[RetryPromptPart(content='try again', tool_name='add', tool_call_id='c1')]), + ] + assert is_provider_valid(messages) is True + + def test_request_only_history_is_valid(self) -> None: + messages: list[ModelMessage] = [ModelRequest(parts=[UserPromptPart(content='hi')])] + assert is_provider_valid(messages) is True + + def test_text_only_response_is_valid(self) -> None: + messages: list[ModelMessage] = [ModelResponse(parts=[TextPart(content='hi')])] + assert is_provider_valid(messages) is True + + def test_orphan_tool_return_is_invalid(self) -> None: + """Return whose `tool_call_id` was never opened by any prior call -> reject.""" + messages: list[ModelMessage] = [ + ModelRequest(parts=[ToolReturnPart(tool_name='add', content=1, tool_call_id='ghost')]), + ] + assert is_provider_valid(messages) is False + + def test_duplicate_tool_return_is_invalid(self) -> None: + """Two returns for the same `tool_call_id` -> the second has no open call.""" + messages: list[ModelMessage] = [ + ModelResponse(parts=[ToolCallPart(tool_name='add', args={}, tool_call_id='c1')]), + ModelRequest(parts=[ToolReturnPart(tool_name='add', content=1, tool_call_id='c1')]), + ModelRequest(parts=[ToolReturnPart(tool_name='add', content=2, tool_call_id='c1')]), + ] + assert is_provider_valid(messages) is False + + def test_out_of_order_tool_return_is_invalid(self) -> None: + """Return appearing before its call in a later response -> reject.""" + messages: list[ModelMessage] = [ + ModelRequest(parts=[ToolReturnPart(tool_name='add', content=1, tool_call_id='c1')]), + ModelResponse(parts=[ToolCallPart(tool_name='add', args={}, tool_call_id='c1')]), + ] + assert is_provider_valid(messages) is False + + def test_orphan_retry_prompt_is_invalid(self) -> None: + """`RetryPromptPart` with no matching open call is also rejected.""" + messages: list[ModelMessage] = [ + ModelRequest(parts=[RetryPromptPart(content='retry', tool_name='add', tool_call_id='ghost')]), + ] + assert is_provider_valid(messages) is False + + +# --------------------------------------------------------------------------- +# continue_run / fork_run +# --------------------------------------------------------------------------- + + +class TestContinueAndForkRun: + async def test_continue_run_returns_snapshot_messages(self) -> None: + store = InMemoryStepStore() + msgs: list[ModelMessage] = [ModelRequest(parts=[UserPromptPart(content='hi')])] + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=msgs)) + + loaded = await continue_run(store, run_id='r1') + assert len(loaded) == 1 + assert loaded is not msgs # caller gets an independent list + + async def test_continue_run_raises_when_no_snapshot(self) -> None: + store = InMemoryStepStore() + with pytest.raises(LookupError, match="no continuable snapshot for run_id 'missing'"): + await continue_run(store, run_id='missing') + + async def test_fork_run_delegates_to_continue_run(self) -> None: + store = InMemoryStepStore() + msgs: list[ModelMessage] = [ModelRequest(parts=[UserPromptPart(content='hi')])] + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=msgs)) + + forked = await fork_run(store, run_id='r1') + assert len(forked) == 1 + + +# --------------------------------------------------------------------------- +# _validate_id +# --------------------------------------------------------------------------- + + +class TestValidateId: + @pytest.mark.parametrize('bad', ['../evil', 'a/b', '', 'a..b', '..', 'has space', 'x' * 201]) + def test_rejects_bad_ids(self, bad: str) -> None: + with pytest.raises(ValueError, match='invalid run_id'): + _validate_id(bad, field='run_id') + + @pytest.mark.parametrize('good', ['good', 'a.b-c_d', 'A1', 'x' * 200]) + def test_accepts_safe_ids(self, good: str) -> None: + _validate_id(good, field='run_id') + + +# --------------------------------------------------------------------------- +# InMemoryStepStore +# --------------------------------------------------------------------------- + + +class TestInMemoryStepStore: + async def test_register_and_get_run(self) -> None: + store = InMemoryStepStore() + await store.register_run(RunRecord(run_id='r1', agent_name='a')) + + record = await store.get_run(run_id='r1') + assert record is not None + assert record.agent_name == 'a' + assert await store.get_run(run_id='missing') is None + + async def test_list_runs_with_and_without_parent_filter(self) -> None: + store = InMemoryStepStore() + await store.register_run(RunRecord(run_id='r1', parent_run_id=None)) + await store.register_run(RunRecord(run_id='r2', parent_run_id='r1')) + await store.register_run(RunRecord(run_id='r3', parent_run_id='r1')) + await store.register_run(RunRecord(run_id='r4', parent_run_id='other')) + + assert {r.run_id for r in await store.list_runs()} == {'r1', 'r2', 'r3', 'r4'} + children = await store.list_runs(parent_run_id='r1') + assert {r.run_id for r in children} == {'r2', 'r3'} + + async def test_list_runs_filters_by_conversation_id(self) -> None: + store = InMemoryStepStore() + await store.register_run(RunRecord(run_id='r1', conversation_id='conv-A')) + await store.register_run(RunRecord(run_id='r2', conversation_id='conv-A')) + await store.register_run(RunRecord(run_id='r3', conversation_id='conv-B')) + + a_runs = await store.list_runs(conversation_id='conv-A') + assert {r.run_id for r in a_runs} == {'r1', 'r2'} + + async def test_list_runs_combines_parent_and_conversation_filters(self) -> None: + store = InMemoryStepStore() + await store.register_run(RunRecord(run_id='r1', parent_run_id='p', conversation_id='conv-A')) + await store.register_run(RunRecord(run_id='r2', parent_run_id='p', conversation_id='conv-B')) + await store.register_run(RunRecord(run_id='r3', parent_run_id='other', conversation_id='conv-A')) + + narrowed = await store.list_runs(parent_run_id='p', conversation_id='conv-A') + assert [r.run_id for r in narrowed] == ['r1'] + + async def test_append_and_list_events(self) -> None: + store = InMemoryStepStore() + await store.append_event(StepEvent(run_id='r1', kind='run_started', step_index=0)) + await store.append_event(StepEvent(run_id='r1', kind='run_completed', step_index=1)) + + events = await store.list_events(run_id='r1') + assert [e.kind for e in events] == ['run_started', 'run_completed'] + assert await store.list_events(run_id='missing') == [] + + async def test_latest_snapshot_returns_last_appended(self) -> None: + store = InMemoryStepStore() + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=[])) + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=2, messages=[])) + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=1, messages=[])) + + latest = await store.latest_snapshot(run_id='r1') + assert latest is not None + # InMemoryStepStore returns the last *appended* snapshot. + assert latest.step_index == 1 + assert await store.latest_snapshot(run_id='missing') is None + + async def test_tool_effects_started_then_completed(self) -> None: + store = InMemoryStepStore() + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c1', tool_name='add', run_id='r1', status='started') + ) + unresolved = await store.list_unresolved_tool_effects(run_id='r1') + assert [r.tool_call_id for r in unresolved] == ['c1'] + + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c1', tool_name='add', run_id='r1', status='completed') + ) + assert await store.list_unresolved_tool_effects(run_id='r1') == [] + + # mix completed and another started; only the started one is unresolved. + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c2', tool_name='add', run_id='r1', status='started') + ) + unresolved = await store.list_unresolved_tool_effects(run_id='r1') + assert [r.tool_call_id for r in unresolved] == ['c2'] + + async def test_get_tool_effect_returns_latest_or_none(self) -> None: + store = InMemoryStepStore() + assert await store.get_tool_effect(run_id='r1', tool_call_id='missing') is None + + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c1', tool_name='add', run_id='r1', status='started') + ) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c1', tool_name='add', run_id='r1', status='completed') + ) + + record = await store.get_tool_effect(run_id='r1', tool_call_id='c1') + assert record is not None + assert record.status == 'completed' + + +# --------------------------------------------------------------------------- +# FileStepStore +# --------------------------------------------------------------------------- + + +class TestFileStepStore: + async def test_runs_round_trip(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + await store.register_run(RunRecord(run_id='r1', parent_run_id='p1', agent_name='a', metadata={'k': 'v'})) + + record = await store.get_run(run_id='r1') + assert record is not None + assert record.parent_run_id == 'p1' + assert record.agent_name == 'a' + assert record.metadata == {'k': 'v'} + assert await store.get_run(run_id='missing') is None + + async def test_list_runs_returns_empty_when_root_missing(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path / 'does-not-exist') + assert await store.list_runs() == [] + + async def test_list_runs_skips_directories_without_run_json(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + (tmp_path / 'orphan').mkdir() + await store.register_run(RunRecord(run_id='real', agent_name='a')) + + runs = await store.list_runs() + assert [r.run_id for r in runs] == ['real'] + + async def test_list_runs_filters_by_parent(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + await store.register_run(RunRecord(run_id='r1', parent_run_id=None)) + await store.register_run(RunRecord(run_id='r2', parent_run_id='r1')) + + children = await store.list_runs(parent_run_id='r1') + assert [r.run_id for r in children] == ['r2'] + + async def test_list_runs_filters_by_conversation_id(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + await store.register_run(RunRecord(run_id='r1', conversation_id='conv-A')) + await store.register_run(RunRecord(run_id='r2', conversation_id='conv-B')) + + assert [r.run_id for r in await store.list_runs(conversation_id='conv-A')] == ['r1'] + + async def test_events_round_trip_skips_blank_lines(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + await store.append_event(StepEvent(run_id='r1', kind='run_started', step_index=0)) + await store.append_event( + StepEvent( + run_id='r1', + kind='tool_call_started', + step_index=1, + tool_call_id='c1', + tool_name='add', + metadata={'k': 'v'}, + ) + ) + + # Inject a blank line to exercise the strip() branch on read. + events_file = tmp_path / 'r1' / 'events.jsonl' + events_file.write_text(events_file.read_text(encoding='utf-8') + '\n', encoding='utf-8') + + events = await store.list_events(run_id='r1') + assert [e.kind for e in events] == ['run_started', 'tool_call_started'] + assert events[1].metadata == {'k': 'v'} + assert events[1].tool_call_id == 'c1' + + async def test_list_events_empty_when_file_missing(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + assert await store.list_events(run_id='nonexistent') == [] + + async def test_snapshot_round_trip(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + messages: list[ModelMessage] = [ + ModelRequest(parts=[UserPromptPart(content='hi')]), + ModelResponse(parts=[TextPart(content='ok')]), + ] + await store.save_snapshot( + ContinuableSnapshot( + run_id='r1', + step_index=0, + messages=messages, + conversation_id='c1', + parent_run_id='p1', + agent_name='a', + ) + ) + + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + assert snap.step_index == 0 + assert snap.conversation_id == 'c1' + assert snap.parent_run_id == 'p1' + assert snap.agent_name == 'a' + assert len(snap.messages) == 2 + + async def test_latest_snapshot_picks_most_recent_seq(self, tmp_path: Path) -> None: + """Latest is by physical write order (filename seq), not by `step_index`.""" + store = FileStepStore(tmp_path) + for step in (0, 2, 1): + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=step, messages=[])) + + # Drop a non-integer filename to exercise the ValueError branch. + (tmp_path / 'r1' / 'snapshots' / 'not-a-number.json').write_text('{}', encoding='utf-8') + + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + # The last save had step_index=1 and was written as the highest seq. + assert snap.step_index == 1 + + async def test_lower_step_index_save_supersedes_earlier(self, tmp_path: Path) -> None: + """A reused `run_id` whose later save has a LOWER `step_index` is still + treated as the latest -- the physical seq wins.""" + store = FileStepStore(tmp_path) + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=5, messages=[])) + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=2, messages=[])) + + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + assert snap.step_index == 2 + + async def test_snapshot_seq_counter_increments(self, tmp_path: Path) -> None: + """Three consecutive saves produce files `0.json`, `1.json`, `2.json`.""" + store = FileStepStore(tmp_path) + for _ in range(3): + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=[])) + + snap_dir = tmp_path / 'r1' / 'snapshots' + names = sorted(p.name for p in snap_dir.glob('*.json')) + assert names == ['0.json', '1.json', '2.json'] + + async def test_snapshot_seq_counter_skips_non_integer_filenames(self, tmp_path: Path) -> None: + """`_next_snapshot_seq` ignores `*.json` files whose stem is not an int.""" + store = FileStepStore(tmp_path) + snap_dir = tmp_path / 'r1' / 'snapshots' + snap_dir.mkdir(parents=True) + (snap_dir / 'not-a-number.json').write_text('{}', encoding='utf-8') + + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=[])) + + # Despite the foreign file, the new snapshot is written as `0.json`. + assert (snap_dir / '0.json').exists() + + async def test_snapshot_seq_counter_keeps_max_across_unordered_iter(self, tmp_path: Path) -> None: + """`_next_snapshot_seq` keeps the highest seq even when `glob` yields lower ones later.""" + store = FileStepStore(tmp_path) + snap_dir = tmp_path / 'r1' / 'snapshots' + snap_dir.mkdir(parents=True) + # Pre-populate multiple numeric files so `glob` iteration hits both + # the `seq > max_seq` true branch and its false branch (a lower seq + # seen after a higher one already set the max). Insertion order on + # APFS / ext4 is the directory iteration order: write the high stem + # first so the lower ones that follow hit the False branch. + # Lexicographic glob order puts `10.json` before `1.json`, so the + # iteration encounters a high seq first and then several lower seqs, + # forcing the `seq > max_seq` False branch. + for seq in (10, 1, 2, 9): + (snap_dir / f'{seq}.json').write_text('{}', encoding='utf-8') + + await store.save_snapshot(ContinuableSnapshot(run_id='r1', step_index=0, messages=[])) + + # Next seq must be 11 (one above the highest existing numeric stem). + assert (snap_dir / '11.json').exists() + + async def test_latest_snapshot_returns_none_when_dir_missing(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + assert await store.latest_snapshot(run_id='nope') is None + + async def test_latest_snapshot_returns_none_when_dir_empty(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + await store.register_run(RunRecord(run_id='r1')) # creates snapshots/ but no files + assert await store.latest_snapshot(run_id='r1') is None + + async def test_tool_effects_round_trip(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c1', tool_name='add', run_id='r1', status='started') + ) + await store.record_tool_effect( + ToolEffectRecord( + tool_call_id='c2', + tool_name='mul', + run_id='r1', + status='completed', + idempotency_key='k', + effect_summary='ok', + ) + ) + + # Blank line to exercise the strip branch on read. + path = tmp_path / 'r1' / 'tool_effects.jsonl' + path.write_text(path.read_text(encoding='utf-8') + '\n', encoding='utf-8') + + unresolved = await store.list_unresolved_tool_effects(run_id='r1') + assert [r.tool_call_id for r in unresolved] == ['c1'] + + async def test_list_unresolved_tool_effects_empty_when_file_missing(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + assert await store.list_unresolved_tool_effects(run_id='nonexistent') == [] + + async def test_get_tool_effect_returns_latest_for_run(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c1', tool_name='add', run_id='runA', status='started') + ) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c1', tool_name='add', run_id='runA', status='completed') + ) + await store.record_tool_effect( + ToolEffectRecord(tool_call_id='c2', tool_name='mul', run_id='runB', status='started') + ) + # Blank line to exercise the strip branch on read. + (tmp_path / 'runA' / 'tool_effects.jsonl').write_text( + (tmp_path / 'runA' / 'tool_effects.jsonl').read_text(encoding='utf-8') + '\n', + encoding='utf-8', + ) + + record = await store.get_tool_effect(run_id='runA', tool_call_id='c1') + assert record is not None + assert record.status == 'completed' + assert record.run_id == 'runA' + + other = await store.get_tool_effect(run_id='runB', tool_call_id='c2') + assert other is not None and other.status == 'started' + + assert await store.get_tool_effect(run_id='runA', tool_call_id='missing') is None + + async def test_get_tool_effect_returns_none_when_file_missing(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + assert await store.get_tool_effect(run_id='absent', tool_call_id='anything') is None + + async def test_register_run_rejects_bad_run_id(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + with pytest.raises(ValueError, match='invalid run_id'): + await store.register_run(RunRecord(run_id='../evil')) + + async def test_event_deserialization_rejects_unknown_kind(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + run_dir = tmp_path / 'r1' + run_dir.mkdir() + (run_dir / 'events.jsonl').write_text( + json.dumps( + { + 'run_id': 'r1', + 'kind': 'made_up', + 'step_index': 0, + 'timestamp': '2024-01-01T00:00:00+00:00', + } + ) + + '\n', + encoding='utf-8', + ) + with pytest.raises(ValueError, match='unknown event kind'): + await store.list_events(run_id='r1') + + async def test_event_deserialization_rejects_wrong_types(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + run_dir = tmp_path / 'r1' + run_dir.mkdir() + (run_dir / 'events.jsonl').write_text( + json.dumps( + { + 'run_id': 1, # wrong type + 'kind': 'run_started', + 'step_index': 0, + 'timestamp': '2024-01-01T00:00:00+00:00', + } + ) + + '\n', + encoding='utf-8', + ) + with pytest.raises(ValueError, match='event payload has wrong types'): + await store.list_events(run_id='r1') + + async def test_event_deserialization_rejects_non_string_optional(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + run_dir = tmp_path / 'r1' + run_dir.mkdir() + (run_dir / 'events.jsonl').write_text( + json.dumps( + { + 'run_id': 'r1', + 'kind': 'run_started', + 'step_index': 0, + 'timestamp': '2024-01-01T00:00:00+00:00', + 'agent_name': 5, # neither None nor str + } + ) + + '\n', + encoding='utf-8', + ) + with pytest.raises(ValueError, match='expected str'): + await store.list_events(run_id='r1') + + async def test_run_deserialization_rejects_wrong_types(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + run_dir = tmp_path / 'r1' + run_dir.mkdir() + (run_dir / 'run.json').write_text(json.dumps({'run_id': 1, 'started_at': 'x'}), encoding='utf-8') + with pytest.raises(ValueError, match='run record has wrong types'): + await store.get_run(run_id='r1') + + async def test_tool_effect_deserialization_rejects_wrong_types(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + run_dir = tmp_path / 'r1' + run_dir.mkdir() + (run_dir / 'tool_effects.jsonl').write_text( + json.dumps({'tool_call_id': 1, 'tool_name': 'add', 'run_id': 'r1', 'status': 'started', 'started_at': 'x'}) + + '\n', + encoding='utf-8', + ) + with pytest.raises(ValueError, match='tool effect record has wrong types'): + await store.list_unresolved_tool_effects(run_id='r1') + + async def test_tool_effect_deserialization_rejects_unknown_status(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + run_dir = tmp_path / 'r1' + run_dir.mkdir() + (run_dir / 'tool_effects.jsonl').write_text( + json.dumps( + { + 'tool_call_id': 'c1', + 'tool_name': 'add', + 'run_id': 'r1', + 'status': 'pending', + 'started_at': '2024-01-01T00:00:00+00:00', + } + ) + + '\n', + encoding='utf-8', + ) + with pytest.raises(ValueError, match='unknown tool effect status'): + await store.list_unresolved_tool_effects(run_id='r1') + + async def test_snapshot_deserialization_rejects_wrong_types(self, tmp_path: Path) -> None: + store = FileStepStore(tmp_path) + snap_dir = tmp_path / 'r1' / 'snapshots' + snap_dir.mkdir(parents=True) + (snap_dir / '0.json').write_text( + json.dumps({'step_index': 'wrong', 'timestamp': 'x', 'messages': []}), + encoding='utf-8', + ) + with pytest.raises(ValueError, match='snapshot has wrong types'): + await store.latest_snapshot(run_id='r1') + + +# --------------------------------------------------------------------------- +# Capability behavior via Agent + TestModel +# --------------------------------------------------------------------------- + + +class TestStepPersistenceCapability: + async def test_basic_run_records_lifecycle_and_snapshot(self) -> None: + store = InMemoryStepStore() + agent = make_simple_agent([StepPersistence(store=store, agent_name='librarian')]) + + result = await agent.run('add 1 and 2') + + rid = await first_run_id(store) + events = await store.list_events(run_id=rid) + kinds = [e.kind for e in events] + assert kinds[0] == 'run_started' + assert kinds[-1] == 'run_completed' + assert 'tool_call_started' in kinds + assert 'tool_call_completed' in kinds + + # RunRecord lineage was registered. + record = await store.get_run(run_id=rid) + assert record is not None + assert record.agent_name == 'librarian' + assert record.parent_run_id is None + + # Snapshot is provider-valid and round-trips through `continue_run`. + snap = await store.latest_snapshot(run_id=rid) + assert snap is not None + assert is_provider_valid(snap.messages) is True + assert len(snap.messages) == len(result.all_messages()) + + # All events tagged with the same run_id and agent_name. + assert {e.run_id for e in events} == {rid} + assert {e.agent_name for e in events} == {'librarian'} + + async def test_helper_based_continuation_replays_prior_messages(self) -> None: + """`continue_run(store, run_id=...) -> Agent.run(message_history=...)`.""" + store = InMemoryStepStore() + agent1 = make_simple_agent([StepPersistence(store=store, agent_name='a')]) + result1 = await agent1.run('add 1 and 2') + first_rid = await first_run_id(store) + + history = await continue_run(store, run_id=first_rid) + assert len(history) == len(result1.all_messages()) + + # Second run uses a fresh capability instance + the helper-provided history. + agent2 = make_simple_agent([StepPersistence(store=store, agent_name='b')]) + result2 = await agent2.run('add 3 and 4', message_history=history) + + msgs = result2.all_messages() + assert len(msgs) > len(history) + # The prior messages appear at the head of the second run's history. + for prior_msg, replayed in zip(history, msgs[: len(history)]): + assert type(prior_msg) is type(replayed) + + async def test_agent_name_derived_run_id_prefix(self) -> None: + """No explicit run_id + agent_name -> `{agent_name}-{8-hex}`.""" + store = InMemoryStepStore() + agent = make_simple_agent([StepPersistence(store=store, agent_name='librarian')]) + await agent.run('add 1 and 2') + + runs = await store.list_runs() + assert len(runs) == 1 + assert runs[0].run_id.startswith('librarian-') + # 'librarian-' + 8 hex chars. + assert len(runs[0].run_id) == len('librarian-') + 8 + + async def test_single_capability_instance_reused_gets_fresh_ids(self) -> None: + """One `StepPersistence(agent_name=...)` reused for two runs -> two distinct ids.""" + store = InMemoryStepStore() + cap: StepPersistence[None] = StepPersistence(store=store, agent_name='librarian') + + agent1: Agent[None, str] = Agent(TestModel(), capabilities=[cap]) + + @agent1.tool_plain + def add(a: int, b: int) -> int: # pyright: ignore[reportUnusedFunction] + return a + b + + await agent1.run('add 1 and 2') + await agent1.run('add 3 and 4') + + runs = await store.list_runs() + assert len(runs) == 2 + rids = {r.run_id for r in runs} + assert len(rids) == 2 + for rid in rids: + assert rid.startswith('librarian-') + + async def test_parent_run_id_inferred_via_contextvar(self) -> None: + """Orchestrator tool calls a delegate `Agent.run` -> delegate's `parent_run_id` + is auto-set to the orchestrator's `run_id` without manual threading.""" + store = InMemoryStepStore() + + delegate: Agent[None, str] = Agent( + TestModel(), + capabilities=[StepPersistence(store=store, agent_name='delegate')], + ) + + @delegate.tool_plain + def add(a: int, b: int) -> int: # pyright: ignore[reportUnusedFunction] + return a + b + + orchestrator: Agent[None, str] = Agent( + TestModel(), + capabilities=[StepPersistence(store=store, agent_name='orchestrator')], + ) + + @orchestrator.tool_plain + async def delegate_work() -> str: # pyright: ignore[reportUnusedFunction] + res = await delegate.run('add 1 and 2') + return res.output + + await orchestrator.run('coordinate') + + runs = await store.list_runs() + orch = next(r for r in runs if r.agent_name == 'orchestrator') + dele = next(r for r in runs if r.agent_name == 'delegate') + + assert dele.parent_run_id == orch.run_id + # And the delegate's events also carry that parent_run_id. + dele_events = await store.list_events(run_id=dele.run_id) + assert {e.parent_run_id for e in dele_events} == {orch.run_id} + + async def test_conversation_id_groups_two_runs(self) -> None: + """Passing the same `conversation_id` to two `Agent.run` calls -> store.list_runs + finds both.""" + store = InMemoryStepStore() + agent = make_simple_agent([StepPersistence(store=store, agent_name='c')]) + + await agent.run('add 1 and 2', conversation_id='conv-1') + await agent.run('add 3 and 4', conversation_id='conv-1') + + runs = await store.list_runs(conversation_id='conv-1') + assert len(runs) == 2 + assert {r.conversation_id for r in runs} == {'conv-1'} + + async def test_list_runs_parent_and_conversation_filters_combine(self) -> None: + store = InMemoryStepStore() + + delegate: Agent[None, str] = Agent( + TestModel(), + capabilities=[StepPersistence(store=store, agent_name='delegate')], + ) + + @delegate.tool_plain + def add(a: int, b: int) -> int: # pyright: ignore[reportUnusedFunction] + return a + b + + orchestrator: Agent[None, str] = Agent( + TestModel(), + capabilities=[StepPersistence(store=store, agent_name='orchestrator')], + ) + + @orchestrator.tool_plain + async def delegate_work() -> str: # pyright: ignore[reportUnusedFunction] + res = await delegate.run('add 1 and 2', conversation_id='target-conv') + return res.output + + await orchestrator.run('coordinate', conversation_id='target-conv') + orch_rid = next(r.run_id for r in await store.list_runs() if r.agent_name == 'orchestrator') + + # Sibling unrelated run under the same conversation but no parent. + unrelated = make_simple_agent([StepPersistence(store=store, agent_name='unrelated')]) + await unrelated.run('add 5 and 6', conversation_id='target-conv') + + narrowed = await store.list_runs(parent_run_id=orch_rid, conversation_id='target-conv') + assert [r.agent_name for r in narrowed] == ['delegate'] + + async def test_step_event_carries_conversation_id(self) -> None: + """`ctx.conversation_id` propagates onto each emitted `StepEvent`.""" + store = InMemoryStepStore() + agent = make_simple_agent([StepPersistence(store=store, agent_name='c')]) + await agent.run('add 1 and 2', conversation_id='conv-evt') + + rid = await first_run_id(store) + events = await store.list_events(run_id=rid) + assert events # sanity + assert {e.conversation_id for e in events} == {'conv-evt'} + + async def test_tool_effect_is_scoped_per_run_in_memory(self) -> None: + """Same `tool_call_id` across two runs must NOT share the effect record.""" + store = InMemoryStepStore() + agent = make_simple_agent([StepPersistence(store=store, agent_name='a')]) + + await agent.run('add 1 and 2') + await agent.run('add 3 and 4') + + runs = await store.list_runs() + assert len(runs) == 2 + run1_id, run2_id = runs[0].run_id, runs[1].run_id + + e1 = await store.get_tool_effect(run_id=run1_id, tool_call_id='pyd_ai_tool_call_id__add') + e2 = await store.get_tool_effect(run_id=run2_id, tool_call_id='pyd_ai_tool_call_id__add') + assert e1 is not None and e2 is not None + assert e1.run_id == run1_id + assert e2.run_id == run2_id + # The second run's record is independent: its started_at is not inherited from run 1. + assert e2.started_at >= e1.started_at + assert e2.started_at != e1.started_at or e2 is not e1 + + # Cross-lookups return only the owning run's record. + cross = await store.get_tool_effect(run_id=run1_id, tool_call_id='pyd_ai_tool_call_id__add') + assert cross is not None + assert cross.run_id == run1_id + + async def test_tool_effect_is_scoped_per_run_file_store(self, tmp_path: Path) -> None: + """Same correctness contract under `FileStepStore`.""" + store = FileStepStore(tmp_path) + agent = make_simple_agent([StepPersistence(store=store, agent_name='a')]) + + await agent.run('add 1 and 2') + await agent.run('add 3 and 4') + + runs = await store.list_runs() + assert len(runs) == 2 + run1_id, run2_id = runs[0].run_id, runs[1].run_id + + e1 = await store.get_tool_effect(run_id=run1_id, tool_call_id='pyd_ai_tool_call_id__add') + e2 = await store.get_tool_effect(run_id=run2_id, tool_call_id='pyd_ai_tool_call_id__add') + assert e1 is not None and e2 is not None + assert e1.run_id == run1_id + assert e2.run_id == run2_id + # The other run's directory does not leak into this run's lookup. + assert await store.get_tool_effect(run_id=run1_id, tool_call_id='missing') is None + + async def test_tool_failure_records_failed_status_and_event(self) -> None: + store = InMemoryStepStore() + agent: Agent[None, str] = Agent(TestModel(), capabilities=[StepPersistence(store=store)]) + + @agent.tool_plain + def boom() -> int: # pyright: ignore[reportUnusedFunction] + raise ValueError('kaboom') + + with pytest.raises(ValueError, match='kaboom'): + await agent.run('boom please') + + rid = await first_run_id(store) + events = await store.list_events(run_id=rid) + kinds = [e.kind for e in events] + assert 'tool_call_started' in kinds + assert 'tool_call_failed' in kinds + assert 'run_failed' in kinds + # The failure event records the exception repr. + failed_event = next(e for e in events if e.kind == 'tool_call_failed') + assert failed_event.error is not None and 'kaboom' in failed_event.error + + effect = await store.get_tool_effect(run_id=rid, tool_call_id='pyd_ai_tool_call_id__boom') + assert effect is not None + assert effect.status == 'failed' + assert effect.effect_summary is not None and 'kaboom' in effect.effect_summary + + async def test_explicit_run_id_wins_over_ctx_run_id(self) -> None: + store = InMemoryStepStore() + agent = make_simple_agent( + [StepPersistence(store=store, run_id='librarian-001', agent_name='librarian')], + ) + + await agent.run('add 1 and 2') + + # Caller-supplied run_id is the persisted identity, not the auto-generated ctx.run_id. + record = await store.get_run(run_id='librarian-001') + assert record is not None + assert record.agent_name == 'librarian' + events = await store.list_events(run_id='librarian-001') + assert {e.run_id for e in events} == {'librarian-001'} + + async def test_explicit_parent_run_id_overrides_contextvar(self) -> None: + """Manual `parent_run_id=` wins over the auto-inferred contextvar value.""" + store = InMemoryStepStore() + agent = make_simple_agent( + [StepPersistence(store=store, agent_name='delegate', parent_run_id='manual-parent')], + ) + + await agent.run('add 1 and 2') + + runs = await store.list_runs(parent_run_id='manual-parent') + assert len(runs) == 1 + assert runs[0].agent_name == 'delegate' + + async def test_from_spec_memory_backend(self) -> None: + cap = StepPersistence.from_spec() + assert isinstance(cap.store, InMemoryStepStore) + + async def test_from_spec_explicit_memory_backend_with_kwargs(self) -> None: + cap = StepPersistence.from_spec(backend='memory', agent_name='a') + assert isinstance(cap.store, InMemoryStepStore) + assert cap.agent_name == 'a' + + async def test_from_spec_file_backend(self, tmp_path: Path) -> None: + cap = StepPersistence.from_spec(backend='file', directory=tmp_path) + assert isinstance(cap.store, FileStepStore) + + async def test_from_spec_file_backend_default_directory(self) -> None: + cap = StepPersistence.from_spec(backend='file') + assert isinstance(cap.store, FileStepStore) + + +# --------------------------------------------------------------------------- +# Headline acceptance test +# --------------------------------------------------------------------------- + + +class TestCrashMidToolCallContract: + """The signature acceptance test from the PR comment. + + A run killed after a tool starts but before its return is persisted must + leave a visible event trail without exposing the killed point as a + valid `message_history` continuation. The latest snapshot must be older + than the in-flight call, and that snapshot must be provider-valid. + """ + + async def test_visible_trail_no_false_continuation_point(self) -> None: + store = InMemoryStepStore() + cap: StepPersistence[None] = StepPersistence(store=store, agent_name='delegate') + agent: Agent[None, str] = Agent(TestModel(), capabilities=[cap]) + + @agent.tool_plain + def add(a: int, b: int) -> int: # pyright: ignore[reportUnusedFunction] + return a + b + + # 1) Drive a full successful run so a provider-valid snapshot exists. + result = await agent.run('add 1 and 2') + rid = await first_run_id(store) + snap_before_crash = await store.latest_snapshot(run_id=rid) + assert snap_before_crash is not None + assert is_provider_valid(snap_before_crash.messages) is True + snap_step = snap_before_crash.step_index + + # 2) Simulate a crash mid-tool-call by calling `before_tool_execute` + # directly with a synthesised ToolCallPart and never firing + # `after_tool_execute` / `on_tool_execute_error`. Use the resolved + # `cap.run_id` if set; otherwise rely on the discovered rid via ctx. + crash_ctx = build_run_context(deps=None, run_id=rid, run_step=snap_step + 1) + crash_call = ToolCallPart(tool_name='add', args={'a': 9, 'b': 9}, tool_call_id='crash-call-1') + tool_def = ToolDefinition(name='add', description='Add two numbers.') + await cap.before_tool_execute(crash_ctx, call=crash_call, tool_def=tool_def, args={'a': 9, 'b': 9}) + + # 3) Assert the event log shows the started call with no terminal update. + events = await store.list_events(run_id=rid) + started = [e for e in events if e.kind == 'tool_call_started' and e.tool_call_id == 'crash-call-1'] + completed = [e for e in events if e.kind == 'tool_call_completed' and e.tool_call_id == 'crash-call-1'] + failed = [e for e in events if e.kind == 'tool_call_failed' and e.tool_call_id == 'crash-call-1'] + assert len(started) == 1 + assert completed == [] + assert failed == [] + + # 4) The unresolved-effect ledger surfaces the in-flight tool call. + unresolved = await store.list_unresolved_tool_effects(run_id=rid) + crash_records = [r for r in unresolved if r.tool_call_id == 'crash-call-1'] + assert len(crash_records) == 1 + assert crash_records[0].status == 'started' + + # 5) Resume point is the snapshot from step 1 — older than the crash — + # and is still provider-valid. + snap_after_crash = await store.latest_snapshot(run_id=rid) + assert snap_after_crash is not None + assert snap_after_crash.step_index == snap_step + assert is_provider_valid(snap_after_crash.messages) is True + # And the snapshot is consistent with what the prior successful run produced. + assert len(snap_after_crash.messages) == len(result.all_messages()) + + +# --------------------------------------------------------------------------- +# Hook-level branches awkward to reach through Agent +# --------------------------------------------------------------------------- + + +class TestCapabilityHookBranches: + async def test_effective_run_id_falls_back_to_capability_field(self) -> None: + """When `ctx.run_id` is missing, the capability uses its own `run_id`.""" + store = InMemoryStepStore() + cap: StepPersistence[None] = StepPersistence(store=store, run_id='configured', agent_name='a') + ctx_no_run_id = build_run_context(deps=None, run_id=None) + await cap.before_run(ctx_no_run_id) + + record = await store.get_run(run_id='configured') + assert record is not None + events = await store.list_events(run_id='configured') + assert [e.kind for e in events] == ['run_started'] + + async def test_after_run_skips_snapshot_when_history_not_provider_valid(self) -> None: + """`after_run` only persists a snapshot when the history is provider-valid.""" + store = InMemoryStepStore() + cap: StepPersistence[None] = StepPersistence(store=store) + ctx = build_run_context(deps=None, run_id='r1') + + unmatched: list[ModelMessage] = [ + ModelRequest(parts=[UserPromptPart(content='hi')]), + ModelResponse(parts=[ToolCallPart(tool_name='add', args={}, tool_call_id='orphan')]), + ] + result: AgentRunResult[str] = AgentRunResult( + output='out', + _state=GraphAgentState(message_history=unmatched, run_id='r1'), + ) + + await cap.after_run(ctx, result=result) + + assert await store.latest_snapshot(run_id='r1') is None + events = await store.list_events(run_id='r1') + assert [e.kind for e in events] == ['run_completed'] + + async def test_after_run_saves_fallback_snapshot_when_no_node_snapshot(self) -> None: + """With no `CallToolsNode` snapshot taken, `after_run` saves the final valid history.""" + store = InMemoryStepStore() + cap: StepPersistence[None] = StepPersistence(store=store) + ctx = build_run_context(deps=None, run_id='r1', run_step=3) + + valid: list[ModelMessage] = [ + ModelRequest(parts=[UserPromptPart(content='hi')]), + ModelResponse(parts=[TextPart(content='done')]), + ] + result: AgentRunResult[str] = AgentRunResult( + output='out', + _state=GraphAgentState(message_history=valid, run_id='r1'), + ) + + await cap.after_run(ctx, result=result) + + snap = await store.latest_snapshot(run_id='r1') + assert snap is not None + assert snap.step_index == 3 + assert len(snap.messages) == 2 + + async def test_on_model_request_error_records_event_and_reraises(self) -> None: + store = InMemoryStepStore() + cap: StepPersistence[None] = StepPersistence(store=store) + ctx = build_run_context(deps=None, run_id='r1') + request_context = ModelRequestContext( + model=ctx.model, + messages=[], + model_settings=None, + model_request_parameters=ModelRequestParameters(), + ) + boom = RuntimeError('nope') + + with pytest.raises(RuntimeError, match='nope'): + await cap.on_model_request_error(ctx, request_context=request_context, error=boom) + + events = await store.list_events(run_id='r1') + assert [e.kind for e in events] == ['model_request_failed'] + assert events[0].error is not None and 'nope' in events[0].error + + async def test_for_run_returns_self_when_resolution_is_no_op(self) -> None: + """When `run_id` is explicit and no contextvar is set, `for_run` returns `self`.""" + store = InMemoryStepStore() + cap: StepPersistence[None] = StepPersistence(store=store, run_id='fixed') + ctx = build_run_context(deps=None, run_id='ignored') + + result = await cap.for_run(ctx) + assert result is cap + + async def test_run_record_load_with_missing_metadata(self, tmp_path: Path) -> None: + """`_str_str_dict(None)` returns `{}` when metadata is absent in storage.""" + store = FileStepStore(tmp_path) + run_dir = tmp_path / 'r1' + run_dir.mkdir() + (run_dir / 'run.json').write_text( + json.dumps({'run_id': 'r1', 'started_at': '2024-01-01T00:00:00+00:00'}), + encoding='utf-8', + ) + + record = await store.get_run(run_id='r1') + assert record is not None + assert record.metadata == {} + + +# --------------------------------------------------------------------------- +# Round-2 review fixes (RetryPromptPart, list_runs ordering, effect metadata, +# from_spec validation, annotate_tool_effect) +# --------------------------------------------------------------------------- + + +class TestNonToolRetryPrompt: + def test_non_tool_retry_prompt_is_valid(self) -> None: + """`RetryPromptPart(tool_name=None)` is an output-validation retry, not a tool result.""" + messages: list[ModelMessage] = [ + ModelRequest(parts=[UserPromptPart(content='hi')]), + ModelResponse(parts=[TextPart(content='wrong shape')]), + ModelRequest(parts=[RetryPromptPart(content='try again', tool_name=None)]), + ] + assert is_provider_valid(messages) is True + + def test_tool_retry_prompt_still_needs_open_call(self) -> None: + """`RetryPromptPart(tool_name=...)` still requires a matching open `ToolCallPart`.""" + messages: list[ModelMessage] = [ + ModelRequest(parts=[RetryPromptPart(content='retry', tool_name='add', tool_call_id='ghost')]), + ] + assert is_provider_valid(messages) is False + + +class TestListRunsChronologicalOrdering: + async def test_in_memory_returns_started_at_order(self) -> None: + """`InMemoryStepStore.list_runs` sorts by `started_at`, not insertion order.""" + from datetime import datetime, timezone + + store = InMemoryStepStore() + await store.register_run(RunRecord(run_id='z-newer', started_at=datetime(2026, 5, 24, 12, tzinfo=timezone.utc))) + await store.register_run(RunRecord(run_id='a-older', started_at=datetime(2026, 5, 24, 10, tzinfo=timezone.utc))) + runs = await store.list_runs() + assert [r.run_id for r in runs] == ['a-older', 'z-newer'] + + async def test_file_store_returns_started_at_order(self, tmp_path: Path) -> None: + """`FileStepStore.list_runs` sorts by `started_at`, not by directory name.""" + from datetime import datetime, timezone + + store = FileStepStore(tmp_path) + # `a-new` is lexicographically first but chronologically last. + await store.register_run(RunRecord(run_id='z-old', started_at=datetime(2026, 5, 24, 10, tzinfo=timezone.utc))) + await store.register_run(RunRecord(run_id='a-new', started_at=datetime(2026, 5, 24, 12, tzinfo=timezone.utc))) + runs = await store.list_runs() + assert [r.run_id for r in runs] == ['z-old', 'a-new'] + + +class TestToolEffectMetadataPreservation: + async def test_completed_preserves_idempotency_key_and_effect_summary(self) -> None: + """Metadata written during the tool call survives the terminal `completed` record.""" + from pydantic_ai_harness.experimental.step_persistence import annotate_tool_effect + + store = InMemoryStepStore() + agent: Agent[None, str] = Agent(TestModel(), capabilities=[StepPersistence(store=store, run_id='r1')]) + + @agent.tool + async def write_label(ctx: RunContext[None], label: str) -> str: # pyright: ignore[reportUnusedFunction] + await annotate_tool_effect( + store, + ctx, + idempotency_key=f'label::{label}', + effect_summary=f'set label to {label}', + ) + return f'wrote {label}' + + await agent.run('apply a label please') + + effect = await store.get_tool_effect(run_id='r1', tool_call_id='pyd_ai_tool_call_id__write_label') + assert effect is not None + assert effect.status == 'completed' + assert effect.idempotency_key is not None and effect.idempotency_key.startswith('label::') + assert effect.effect_summary is not None and 'set label to' in effect.effect_summary + + async def test_failed_preserves_idempotency_key(self) -> None: + """Metadata written before a tool raises still appears on the `failed` record.""" + from pydantic_ai_harness.experimental.step_persistence import annotate_tool_effect + + store = InMemoryStepStore() + agent: Agent[None, str] = Agent(TestModel(), capabilities=[StepPersistence(store=store, run_id='r1')]) + + @agent.tool + async def boom(ctx: RunContext[None]) -> int: # pyright: ignore[reportUnusedFunction] + await annotate_tool_effect(store, ctx, idempotency_key='boom-key') + raise ValueError('kaboom') + + with pytest.raises(ValueError, match='kaboom'): + await agent.run('please boom') + + effect = await store.get_tool_effect(run_id='r1', tool_call_id='pyd_ai_tool_call_id__boom') + assert effect is not None + assert effect.status == 'failed' + assert effect.idempotency_key == 'boom-key' + # default summary still records the error when no summary was annotated: + assert effect.effect_summary is not None and 'kaboom' in effect.effect_summary + + async def test_annotate_tool_effect_outside_step_persistence_is_a_noop(self) -> None: + """No `current_run_id` → `annotate_tool_effect` returns without writing.""" + from pydantic_ai_harness.experimental.step_persistence import annotate_tool_effect + + store = InMemoryStepStore() + ctx = build_run_context(deps=None, run_id='r1') + # No StepPersistence active and ctx.tool_call_id is None. + await annotate_tool_effect(store, ctx, idempotency_key='ignored') + assert await store.get_tool_effect(run_id='r1', tool_call_id='whatever') is None + + async def test_annotate_tool_effect_noop_when_prior_record_missing(self) -> None: + """`current_run_id` set + ctx tool fields set, but no prior record → no-op.""" + from pydantic_ai_harness.experimental.step_persistence import annotate_tool_effect + from pydantic_ai_harness.experimental.step_persistence._context import current_run_id + + store = InMemoryStepStore() + ctx = RunContext[Any]( + deps=None, + model=TestModel(), + usage=RunUsage(), + prompt=None, + messages=[], + run_step=0, + run_id='r1', + tool_call_id='tc-1', + tool_name='write_label', + ) + token = current_run_id.set('r1') + try: + await annotate_tool_effect(store, ctx, idempotency_key='label::x') + finally: + current_run_id.reset(token) + # `before_tool_execute` hasn't fired, so the prior record doesn't exist. + # The helper returns without inventing one. + assert await store.get_tool_effect(run_id='r1', tool_call_id='tc-1') is None + + +class TestFromSpecBackendValidation: + def test_unknown_backend_raises(self) -> None: + """A typo like `backend='disk'` raises instead of silently using memory.""" + with pytest.raises(ValueError, match='unknown backend'): + StepPersistence.from_spec(backend='disk') + + def test_memory_backend_still_works(self) -> None: + cap: StepPersistence[Any] = StepPersistence.from_spec(backend='memory') + assert isinstance(cap.store, InMemoryStepStore) + + def test_sqlite_backend(self, tmp_path: Path) -> None: + cap: StepPersistence[Any] = StepPersistence.from_spec(backend='sqlite', database=str(tmp_path / 'runs.db')) + assert isinstance(cap.store, SqliteStepStore) + + +# --------------------------------------------------------------------------- +# Identity model contract: run_id is per-call; conversation_id groups turns +# --------------------------------------------------------------------------- + + +class TestRunIdIsPerCall: + """`run_id` is one `Agent.run` invocation; `conversation_id` groups turns.""" + + async def test_multi_turn_orchestrator_uses_conversation_id(self) -> None: + """Recommended multi-turn pattern: shared `conversation_id`, distinct `run_id` per call.""" + store = InMemoryStepStore() + agent = make_simple_agent([StepPersistence(store=store, agent_name='orchestrator')]) + + for prompt in ('first turn', 'second turn', 'third turn'): + await agent.run(prompt, conversation_id='orch-conv') + + records = await store.list_runs(conversation_id='orch-conv') + assert len(records) == 3 + # All distinct ids, all carrying the same conversation_id. + assert len({r.run_id for r in records}) == 3 + assert all(r.conversation_id == 'orch-conv' for r in records) + assert all(r.run_id.startswith('orchestrator-') for r in records) + + async def test_explicit_run_id_reuse_raises(self) -> None: + """Reusing an explicit `run_id` across `.run()` calls raises ValueError. + + The tool-effect ledger keys on `(run_id, tool_call_id)`, so a + silent second run would collide with the first. `before_run` + rejects the second run explicitly, pointing the caller to + `conversation_id` for multi-turn grouping. + """ + store = InMemoryStepStore() + agent = make_simple_agent([StepPersistence(store=store, run_id='shared')]) + + await agent.run('first') + + with pytest.raises(ValueError, match=r"run_id 'shared' is already in the store"): + await agent.run('second') + + # First run's records remain untouched. + record = await store.get_run(run_id='shared') + assert record is not None + effect = await store.get_tool_effect(run_id='shared', tool_call_id='pyd_ai_tool_call_id__add') + assert effect is not None + assert effect.status == 'completed' diff --git a/tests/experimental/test_warnings.py b/tests/experimental/test_warnings.py index 9b440d8..70ec94e 100644 --- a/tests/experimental/test_warnings.py +++ b/tests/experimental/test_warnings.py @@ -33,3 +33,9 @@ class TestExperimentalWarning: module = importlib.import_module('pydantic_ai_harness.experimental.compaction') with pytest.warns(HarnessExperimentalWarning): importlib.reload(module) + + @pytest.mark.parametrize('feature', ['step_persistence', 'media']) + def test_importing_step_persistence_or_media_warns(self, feature: str): + module = importlib.import_module(f'pydantic_ai_harness.experimental.{feature}') + with pytest.warns(HarnessExperimentalWarning, match=feature): + importlib.reload(module) diff --git a/uv.lock b/uv.lock index 9a847a2..5eabb4f 100644 --- a/uv.lock +++ b/uv.lock @@ -978,6 +978,7 @@ wheels = [ name = "pydantic-ai-harness" source = { editable = "." } dependencies = [ + { name = "httpx" }, { name = "pydantic-ai-slim" }, ] @@ -1011,6 +1012,7 @@ dev = [ { name = "pytest" }, { name = "pytest-anyio" }, { name = "pytest-examples" }, + { name = "pytest-recording" }, ] lint = [ { name = "pyright" }, @@ -1019,6 +1021,7 @@ lint = [ [package.metadata] requires-dist = [ + { name = "httpx", specifier = ">=0.28.1" }, { name = "logfire", marker = "extra == 'logfire'", specifier = ">=4.31.0" }, { name = "pydantic-ai-slim", specifier = ">=1.95.1" }, { name = "pydantic-ai-slim", extras = ["dbos"], marker = "extra == 'dbos'" }, @@ -1041,6 +1044,7 @@ dev = [ { name = "pytest", specifier = ">=9.0.0" }, { name = "pytest-anyio" }, { name = "pytest-examples", specifier = ">=0.0.18" }, + { name = "pytest-recording", specifier = ">=0.13.4" }, ] lint = [ { name = "pyright", specifier = ">=1.1.408" }, @@ -1359,6 +1363,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/09/52/7bbfb6e987d9a8a945f22941a8da63e3529465f1b106ef0e26f5df7c780d/pytest_examples-0.0.18-py3-none-any.whl", hash = "sha256:86c195b98c4e55049a0df3a0a990ca89123b7280473ab57608eecc6c47bcfe9c", size = 18169, upload-time = "2025-05-06T07:46:09.349Z" }, ] +[[package]] +name = "pytest-recording" +version = "0.13.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, + { name = "vcrpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/32/9c/f4027c5f1693847b06d11caf4b4f6bb09f22c1581ada4663877ec166b8c6/pytest_recording-0.13.4.tar.gz", hash = "sha256:568d64b2a85992eec4ae0a419c855d5fd96782c5fb016784d86f18053792768c", size = 26576, upload-time = "2025-05-08T10:41:11.231Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/42/c2/ce34735972cc42d912173e79f200fe66530225190c06655c5632a9d88f1e/pytest_recording-0.13.4-py3-none-any.whl", hash = "sha256:ad49a434b51b1c4f78e85b1e6b74fdcc2a0a581ca16e52c798c6ace971f7f439", size = 13723, upload-time = "2025-05-08T10:41:09.684Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -1790,6 +1807,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] +[[package]] +name = "vcrpy" +version = "8.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/07/bcfd5ebd7cb308026ab78a353e091bd699593358be49197d39d004e5ad83/vcrpy-8.1.1.tar.gz", hash = "sha256:58e3053e33b423f3594031cb758c3f4d1df931307f1e67928e30cf352df7709f", size = 85770, upload-time = "2026-01-04T19:22:03.886Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/d7/f79b05a5d728f8786876a7d75dfb0c5cae27e428081b2d60152fb52f155f/vcrpy-8.1.1-py3-none-any.whl", hash = "sha256:2d16f31ad56493efb6165182dd99767207031b0da3f68b18f975545ede8ac4b9", size = 42445, upload-time = "2026-01-04T19:22:02.532Z" }, +] + [[package]] name = "websockets" version = "16.0"