furyhawk/pydantic-ai-harness
1
0
Fork 0
mirror of https://github.com/pydantic/pydantic-ai-harness.git synced 2026-07-21 02:45:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add StepPersistence capability for step-event durability across delegates (#251)

Browse Source 
* Add StepPersistence capability for step-event durability across delegates

Supersedes PR #176 (SessionPersistence): orchestrators like pydanty need
visible event trails for delegate runs that may time out before a
"save full session after the run" hook can fire, and need to continue or
fork a delegate's prior investigation without rediscovering context. A
single after-run snapshot is too coarse for that use case.

The capability now records (a) append-only StepEvents at every boundary
(run/model-request/tool-call start, completion, failure), (b) a
ContinuableSnapshot only when message history is provider-valid (every
ToolCallPart has a matching ToolReturnPart / RetryPromptPart) — saved
mid-run after CallToolsNode and at after_run, and (c) a ToolEffectRecord
ledger so a run killed between before_tool_execute and
after_tool_execute leaves an `unknown_after_crash`-style record rather
than a falsely-continuable snapshot. Lineage metadata
(parent_run_id, agent_name) ties delegate runs back to their
orchestrator. `continue_run` / `fork_run` helpers load the latest
continuable snapshot for a run.

Backends: InMemoryStepStore (tests) and FileStepStore (JSONL events +
JSON snapshots, with run_id path-traversal validation and
anyio.to_thread for blocking I/O).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Address pydanty review + ergonomics: tighten correctness and identity model

Correctness fixes from pydanty's PR review:

- FileStepStore: snapshot filenames are now a per-run monotonic counter,
  not `ctx.run_step` — `run_step` resets each Agent.run, so re-using a
  `run_id` across calls would let an earlier run's higher step-index
  snapshot mask a later run's lower-step-index one.
- StepStore.get_tool_effect now takes both `run_id` and `tool_call_id`.
  TestModel and other providers can reuse deterministic tool-call ids
  across runs; the previous unscoped lookup let one run's effect leak
  into another's record (including `started_at`).
- is_provider_valid now rejects orphan, duplicate, and out-of-order
  tool returns — the old `set.discard` pattern silently accepted any
  return regardless of whether a matching call was open.

Identity model:

- `run_id` resolution: explicit > `{agent_name}-{8-char-hex}` > UUID.
  Materialised per Agent.run in `for_run`, so reusing one capability
  instance never silently merges runs.
- `parent_run_id` auto-inferred via a module-level ContextVar set in
  `wrap_run`, so an orchestrator's tool that synchronously calls
  `delegate.run(...)` produces a delegate `RunRecord.parent_run_id`
  pointing at the orchestrator's `run_id` with zero threading. Explicit
  `parent_run_id=` still wins.
- `conversation_id` propagated to `StepEvent` and `RunRecord`;
  `store.list_runs(conversation_id=..., parent_run_id=...)` supports
  filtering by either or both. Mirrors pydantic_ai's three-level
  identity (conversation -> run -> step) so "run 1, run 2, run 3" of
  one dialogue is queryable as a group via `conversation_id`.
- `continue_from=` field dropped from the capability. Continuation is
  now only via `continue_run(store, run_id=...)` -> standard
  `Agent.run(message_history=...)`. One way to pass history into
  pydantic_ai, no parallel capability flag.

README rewritten around the final API. New sections: three-level
identity, run lineage with auto-inferred parent, inspecting a run tree,
failure recovery.

Tests: 168 total (up from 64), 100% branch coverage on the package.
New coverage for the snapshot seq counter, cross-run tool-effect
isolation, orphan/duplicate/out-of-order return rejection, ContextVar
parent inference across nested agent.run, conversation_id propagation,
and the agent_name-derived run_id default.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Address pydanty round-2 review: retry validity, list ordering, effect metadata

Correctness:

- is_provider_valid no longer rejects non-tool RetryPromptParts. Pydantic
  AI emits `RetryPromptPart(tool_name=None)` for output-validation
  failures and providers map those as plain user messages, not tool
  results. The previous check required every RetryPromptPart to resolve
  an open tool call, so a run with one output retry produced no final
  continuable snapshot despite being fully valid.
- StepStore.list_runs now guarantees chronological (started_at ascending)
  ordering across both backends. FileStepStore was previously returning
  directory-name order (lexicographic), so the README's `[-1]` pattern
  for "latest run in conversation" could pick the older run when run ids
  did not sort by recency.
- after_tool_execute and on_tool_execute_error preserve idempotency_key
  and effect_summary from the prior `started` record. Previously the
  terminal record was written without those fields, so any annotation
  the tool body wrote was lost on completion.
- from_spec raises ValueError for unknown backends instead of silently
  falling back to in-memory storage. For a persistence capability,
  turning a typo into accidental non-durability is the wrong failure
  mode.

API:

- New annotate_tool_effect(store, ctx, *, idempotency_key=None,
  effect_summary=None) helper. Tool bodies that write external state
  call it to attach idempotency + effect metadata to the in-flight
  ToolEffectRecord without knowing the (run_id, tool_call_id) plumbing.
  Resolves run_id from a ContextVar set by wrap_run; reads tool_call_id
  / tool_name from RunContext.
- ContextVar moved from `_capability.py` into a new `_context.py` module
  so the helper and the capability can share it without circular imports
  and without crossing the private-name barrier.

Docs: README fixes a non-existent `list_runs(agent_name=None)` call,
documents the chronological-ordering guarantee, and replaces the
hand-wavy "populate fields on the ToolEffectRecord" line with a concrete
`annotate_tool_effect` example.

Tests: 178 total (was 168), 100% branch coverage on the package.
Added coverage for non-tool retry acceptance, chronological list_runs on
both backends, metadata preservation across completed/failed
transitions, annotate_tool_effect under realistic agent.tool, and
from_spec backend validation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* docs(step_persistence): align FileStepStore docstring with seq-counter layout

The class docstring still showed snapshots/<step_index>.json from the
pre-fix layout, but both the README and _next_snapshot_seq document the
monotonic counter. Bring the class docstring in line.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* docs(step_persistence): fix broken README examples + clarify run_id sharing

Pydanty round-3 review:

- README continuation and lineage examples queried `list_runs(conversation_id=...)`
  on conversations the earlier `.run(...)` calls never set, so the
  examples crashed with IndexError on `[-1]`. Pass the conversation_id to
  the earlier calls so the lookup actually works.
- The capability docstring claimed reusing a `StepPersistence` instance
  across `Agent.run` calls does NOT share the id. That is true only for
  the auto-derived (`agent_name`-prefixed or `ctx.run_id`) cases — an
  explicit `run_id=` is shared across every `.run()` by design, since
  that is the orchestrator pattern where the caller owns one logical
  identity across turns. Rewrite the resolution-order docs to spell out
  which cases share and which don't, and when to pick each.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* docs(step_persistence): align run_id semantics with pydantic_ai (per-call, not shared)

Pydanty round-4 review: the prior round documented explicit `run_id` as
shared across `.run()` calls on one capability instance — that framing
caused real correctness gaps. The `ToolEffectRecord` ledger is keyed by
`(run_id, tool_call_id)` and providers reuse deterministic tool-call ids
(e.g. `TestModel` emits `pyd_ai_tool_call_id__{name}`), so a second
`.run()` overwrites the first's effect record under the same key — the
`unknown_after_crash` signal from turn 1 disappears when turn 2 lands.

Realign:

- `run_id` is per-`Agent.run`, matching `pydantic_ai.RunContext.run_id`.
- For multi-turn logical grouping, use `conversation_id=` on
  `Agent.run(...)` — that is the pyai-native primitive. The orchestrator
  pattern is `conversation_id='orch'` with each turn auto-deriving its
  own `run_id`.
- Explicit `run_id=` remains supported but is documented as single-shot
  (testing, replay, debugging). Reusing it across calls is a caller
  contract violation, not an implementation feature.

Code is unchanged — the implementation was already correct under the
right contract. Only the docs were misleading.

Tests:
- `TestRunIdIsPerCall::test_multi_turn_orchestrator_uses_conversation_id`
  exercises the recommended pattern: three turns sharing a
  `conversation_id`, three distinct auto-derived `run_id`s, all
  queryable as a group.
- `TestRunIdIsPerCall::test_explicit_run_id_reuse_collides_ledger` locks
  down the misuse contract: reusing one explicit `run_id` across two
  `.run()` calls produces colliding effect records under the
  `(run_id, tool_call_id)` key. The behavior is documented; the test
  exists so a future refactor cannot silently change it.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* docs(step_persistence): explain why for_run uses a local ContextVar

Pyai-aligned review flagged this as a P3 explainer: pydantic_ai already
has three single-slot cross-run signals (RUN_ID_BAGGAGE_KEY, ctx.run_id,
_CURRENT_RUN_CONTEXT). All three get overwritten by the inner
Instrumentation.wrap_run before any nested capability can see the parent
identity. A separate harness-local ContextVar, snapshotted before our
own wrap_run rebinds it, is the only correct mechanism today. Spell
this out so the next reader doesn't try to 'simplify' it.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* feat(step_persistence): enforce explicit-run_id reuse with ValueError

Pydanty round-5 review accepted the docs-only contract but flagged that
"documented but not enforced" is a soft spot. Enforce it: `before_run`
calls `store.get_run(run_id=...)` when the user supplied an explicit
`run_id`, and raises `ValueError` if a record with that id already
exists. The auto-derived cases cannot trigger this check (each call
materialises a fresh id in `for_run`).

The check is one extra store read per Agent.run when an explicit run_id
is set, only. The error message points the caller at `conversation_id`
for multi-turn grouping.

Test renamed from `test_explicit_run_id_reuse_collides_ledger` to
`test_explicit_run_id_reuse_raises` — asserts the second `.run()` raises
and the first run's records survive untouched.

README + capability docstring updated: the misuse path is now "raises"
not "caller's contract."

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* chore: gitignore branch-context skill artifacts + AGENTS.local.md

Two patterns that match the existing CLAUDE.local.md ignore convention:

- AGENTS.local.md — canonical local-instructions file (CLAUDE.local.md
  is symlinked to it where the worktree follows the same
  AGENTS.md/CLAUDE.md symlink pattern).
- .agents/skills/branch-context/ — per-worktree decisions log
  (`pr-decisions.md`) and the skill's local SKILL.md. Pattern lifted
  from `~/pydantic/ai/base/.claude/skills/branch-context/` where pyai
  uses an identical setup.

Neither is intended to land in PRs — they record cross-iteration design
calls so future AICA sessions in this worktree don't silently undo them.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* feat(step_persistence): add SqliteStepStore + MediaStore externalization

Adds a new `pydantic_ai_harness.media` package (MediaStore protocol +
DiskMediaStore / SqliteMediaStore / S3MediaStore) and wires it into the
file/sqlite step-persistence backends so large BinaryContent payloads
get externalized out of snapshot JSON / table rows by default.

Defaults are zero-config: FileStepStore writes blobs under
`<root>/media/<sha256>.bin`; SqliteStepStore writes them to a sibling
`media` table in the same DB. Threshold is 64 KiB and URI scheme is
`media+sha256://<hex>` so blobs are content-addressed across stores.
Pass `media_store=None` to keep bytes inline, or a custom `MediaStore`
to redirect (e.g. `S3MediaStore` for R2 / AWS / MinIO).

S3MediaStore handrolls SigV4 over httpx to avoid a botocore/boto3
dependency. Verified working against Cloudflare R2.

`StepPersistence.from_spec(backend='sqlite', database=...)` now resolves.

180 → 261 tests, 100% branch coverage maintained.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* test(media): VCR cassettes for S3MediaStore against R2

Adds replay-driven tests under `tests/media/test_s3_cassettes.py` that
exercise `S3MediaStore.put/get/exists` against pre-recorded Cloudflare
R2 responses. CI runs them without any S3 creds via the committed
cassettes under `tests/media/cassettes/`.

Sanitisation policy:

- `before_record_request`/`before_record_response` swap the real R2
  account-id subdomain and bucket name for fixed placeholders
  (`account.r2.cloudflarestorage.com`, `harness-test-bucket`)
- `Authorization` and `x-amz-date` filtered to `REDACTED`
- CF-RAY, x-amz-version-id, x-amz-checksum-*, x-amz-request-id headers
  dropped (none load-bearing for tests; some carry identifying info)
- Non-2xx response bodies blanked (R2's gzipped XML error envelope
  leaks the bucket name; our code only checks status code)

The `s3_credentials` fixture uses `os.environ.get(NAME, PLACEHOLDER)`
per field, so real R2 creds are used when recording locally with `.env`
loaded, and the placeholder constants match the scrubbed cassettes
during replay. Because the placeholders are fixed, any scrubber miss
during a future re-record shows up as a replay URL mismatch — built-in
canary against credential / private-data leakage in committed cassettes.

Adds `pytest-recording` (pulls `vcrpy`) to the dev deps.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* feat(media): public_url resolver — protocol method + per-store callable

Adds `MediaStore.public_url(uri) -> str | None` plus a `public_url=`
constructor parameter on every concrete store. The parameter accepts a
sync or async callable; the store auto-detects and awaits if needed.

This is the bottom-layer primitive for the forthcoming `MediaExternalizer`
capability — that capability will call `store.public_url(...)` per
externalized blob and swap `BinaryContent` for `ImageUrl` / `AudioUrl`
parts before the model sees the message. The callable shape covers both
static URLs (public bucket / CDN — use `make_static_public_url`
helper) and dynamic URLs (presigned, per-request signing — pass any
async callable with TTL captured in its closure).

Why a callable rather than a static config: a public bucket's URL host
is not derivable from the bucket creds (R2 public buckets use
`pub-<hash>.r2.dev`, AWS public buckets use a different scheme than the
path-style endpoint we sign for). The URL is always user-supplied
information, so a callable is the right primitive — same shape for the
static and presigned cases, and `get` stays untouched (it serves the
harness's internal byte fetch, not the model's external HTTP fetch).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* refactor(media): introduce MediaContext + key_strategy for extensible operations

Adds `MediaContext` (frozen, kw-only dataclass with `media_type`, `filename`,
`metadata`) and threads it through every `MediaStore` method and both user
callables (`PublicUrlResolver`, `KeyStrategy`). New context fields can be
added non-breakingly; existing call sites and resolvers keep working.

Also adds:
- `KeyStrategy = Callable[[str, MediaContext], str]` for per-store layout
  control. Default `default_key_strategy` produces `<sha256>.bin`. Disk store
  validates the result against `..` traversal.
- `metadata` persistence on `SqliteMediaStore` (new JSON column) and
  `S3MediaStore` (signed `x-amz-meta-*` headers, ASCII key validation).
  Disk store explicitly does NOT persist metadata in v1 — sidecar / xattr
  options each have load-bearing drawbacks; we ship nothing rather than a
  half-true persistence promise.
- `make_static_public_url(...)` updated to the new `(uri, ctx)` signature.

The shift is motivated by the same principle as pydantic_ai's
`RunContext`: extension via fields on a context bag rather than via
breaking signature changes. Every new requirement (TTL hints for
presigned URLs, audit ids, response-header overrides, etc.) becomes a
field addition, not an API revision.

Cassettes from the previous commit replay unchanged — match-on does not
include the signed headers and the request URLs are stable.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* feat(media): close metadata round-trip; drop vestigial sqlite key_strategy

Adds `MediaStore.get_metadata(uri) -> Mapping[str, str]` to the protocol
and implements it on all three concrete stores:

- `DiskMediaStore`: writes a sidecar `<resolved>.meta.json` alongside
  the blob on put (atomic tmp + rename), reads it back on
  `get_metadata`. Returns `{}` when no metadata was supplied. v1 had
  documented this as a deliberate gap — sidecar JSON is straightforward
  and the xattr / ADS drawbacks don't apply.
- `SqliteMediaStore`: `SELECT metadata FROM <table> WHERE sha256=?` +
  `json.loads`. Raises `FileNotFoundError` for unknown URIs.
- `S3MediaStore`: HEAD + collects `x-amz-meta-*` response headers,
  strips the prefix. Reuses the existing 404 / non-2xx error shape.

Drops `key_strategy=` from `SqliteMediaStore`. The digest is the primary
key by content-addressing construction — a user-chosen key would either
break dedup or be a no-op. Kept on Disk + S3 where bucket / directory
layout is a real concern.

README + branch-context entries updated to reflect: all three stores
round-trip metadata; key_strategy is Disk + S3 only.

Coverage stays at 100% branch.

* fix(media): reject absolute paths from DiskMediaStore key_strategy

`Path(root) / absolute_path` returns `absolute_path` — the root is
silently discarded — so a custom `key_strategy` returning `/etc/passwd`
(or similar) escapes the store directory even though the previous check
only blocked `..`. Tighten the validator to reject both shapes.

Caught by pydanty during its #251 integration review.

* fix(step_persistence): dedupe terminal snapshot; document sqlite thread-affinity

The terminal CallToolsNode already saves the final provider-valid snapshot
with the correct step_index. after_run was re-saving the same tail stamped
with step_index=0 (ctx.run_step is reset by then), so latest_snapshot
reported a misleading step and every run wrote a duplicate. Track whether a
node snapshot was taken via a task-local ContextVar and make after_run a
fallback that only fires when the run reached no provider-valid boundary.

Also document that a caller-owned sqlite connection= must set
check_same_thread=False (store SQL runs on anyio worker threads), on both
SqliteStepStore and SqliteMediaStore, and correct the WAL-on-every-connection
claim.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(media): SigV4 wire path matches signed path; walker preserves unknown fields

S3MediaStore signed `_canonical_uri(path)` (each segment percent-encoded) but
sent the raw path, letting httpx apply looser encoding. A custom key_prefix /
key_strategy emitting reserved chars (`@`, `(`, `=`, ...) diverged from the
signed path -> SignatureDoesNotMatch. Send the canonical bytes via httpx
`raw_path` so signer and sender agree. Default `<hex>.bin` keys are unaffected.

The externalize/restore walker hardcoded the BinaryContent key set, silently
dropping any field pydantic_ai adds upstream. Copy the node and swap only
`data` <-> marker keys so unknown fields round-trip. Adds tests for reserved-char
path agreement, unknown-field preservation, and restore over a pruned blob.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* refactor: gate step_persistence + media behind experimental namespace

Move both packages under `pydantic_ai_harness.experimental`, matching the
convention introduced for compaction/planning/subagents (#191): the old
`pydantic_ai_harness.step_persistence` / `.media` paths and the top-level
re-exports are gone, so the only import path is now

    from pydantic_ai_harness.experimental.step_persistence import StepPersistence
    from pydantic_ai_harness.experimental.media import S3MediaStore

Both package __init__s call `warn_experimental(...)`, so importing either
emits a `HarnessExperimentalWarning` (silenced category-wide by one
filter). This keeps us from committing to a public surface before the
capability has real usage.

README gains the standard experimental banner; warning tests cover both
new packages. 100% branch coverage retained.

* docs(media): fix metadata/key_strategy drift, convert em-dashes

Subagent review of the experimental move surfaced doc drift:

- README "Persistence by store" implied `get_metadata` returns
  `media_type`; it returns only the user `metadata` mapping. Reworded.
- `KeyStrategy` docstring still listed `SqliteMediaStore` / "DB primary
  key" as a user; sqlite has no `key_strategy`. Dropped it.
- README understated `DiskMediaStore` traversal protection (it rejects
  absolute paths as well as `..`).
- README had a stale paragraph that read as `key_strategy` but described
  `public_url`, and omitted S3. Rewritten to name `public_url` and all
  three stores.
- README `MediaContext` method list now includes `get_metadata`.
- Converted em-dashes to `--` across the step_persistence + media trees
  per the writing-style rule (#270). Other experimental packages'
  pre-existing em-dashes are left for a separate sweep.

No code behavior change; 100% branch coverage retained.

* docs(experimental): add warning banner to planning/subagents, finish em-dash sweep

Uniformity pass across the experimental packages:

- Add the standard `[!WARNING]` experimental banner to planning/README.md
  and subagents/README.md (compaction and step_persistence already had
  it; these two lacked it).
- Convert remaining em-dashes to `--` in the compaction package and the
  shared `_warn.py`, per the writing-style rule (#270). The whole
  experimental source tree is now em-dash-free.

Docs only; 100% branch coverage retained.

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: David SF <david.sanchez@pydantic.dev>
This commit is contained in:
David SF
2026-06-10 23:55:06 -05:00
committed by GitHub
co-authored by Claude Opus 4.7 David SF
parent d347d08e3c
commit 6f2aa11438
39 changed files with 6605 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -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
pydantic_ai_harness/__init__.py
+7 -2
View File
@@ -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:
pydantic_ai_harness/experimental/_warn.py
+1 -1
View File
@@ -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(
pydantic_ai_harness/experimental/compaction/README.md
+9 -9
View File
@@ -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.
pydantic_ai_harness/experimental/compaction/_clear_tool_results.py
+2 -2
View File
@@ -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
pydantic_ai_harness/experimental/compaction/_deduplicate_file_reads.py
+2 -2
View File
@@ -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.
pydantic_ai_harness/experimental/compaction/_limit_warner.py
+1 -1
View File
@@ -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
pydantic_ai_harness/experimental/compaction/_shared.py
+2 -2
View File
@@ -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
pydantic_ai_harness/experimental/compaction/_sliding_window.py
+1 -1
View File
@@ -1,4 +1,4 @@
"""`SlidingWindow` — zero-cost trimming of the oldest messages."""
"""`SlidingWindow` -- zero-cost trimming of the oldest messages."""
from __future__ import annotations
pydantic_ai_harness/experimental/compaction/_summarizing_compaction.py
+5 -5
View File
@@ -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.
<messages>
@@ -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
pydantic_ai_harness/experimental/compaction/_tiered_compaction.py
+2 -2
View File
@@ -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:
pydantic_ai_harness/experimental/media/__init__.py
+41
View File
@@ -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',
]
pydantic_ai_harness/experimental/media/_s3.py
+309
View File
@@ -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 (`<endpoint>/<bucket>/<key>`) -- 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.<region>.amazonaws.com'`),
Cloudflare R2 (`endpoint='https://<account_id>.r2.cloudflarestorage.com'`,
`region='auto'`), MinIO, and other compatible providers.
Customisation:
- `key_strategy=` -- `Callable[[str, MediaContext], str]`. Default
strategy yields `<key_prefix><sha256>.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
pydantic_ai_harness/experimental/media/_store.py
+445
View File
@@ -0,0 +1,445 @@
"""Local `MediaStore` protocol + `DiskMediaStore` / `SqliteMediaStore`.
The shared URI scheme is `media+sha256://<lowercase-hex>` -- 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://<hex>` 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 `<sha256>.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: `<sha256>.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://<hex>` URI and returns
`{base}/{key_prefix}<hex>{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://<hex>` 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://<hex>` 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://<hex>`) 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 `<directory>/<sha256>.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/<digest>.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 (`<resolved>.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]))
pydantic_ai_harness/experimental/media/_walker.py
+128
View File
@@ -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": "<b64>"}`)
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
pydantic_ai_harness/experimental/planning/README.md
+19
View File
@@ -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
pydantic_ai_harness/experimental/step_persistence/README.md
+505
View File
@@ -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 `<directory>/<run_id>/`:
- `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(<root>/media/)` | `<root>/media/<sha256>.bin` |
| `SqliteStepStore` | `SqliteMediaStore(database=<same db>)` | 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://<account>.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://<hex>`, 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
`<directory>/<sha256>.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://<hex>` 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://<acc>.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 (`<resolved>.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 `<sha256>.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.
pydantic_ai_harness/experimental/step_persistence/__init__.py
+44
View File
@@ -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',
]
pydantic_ai_harness/experimental/step_persistence/_capability.py
+426
View File
@@ -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
pydantic_ai_harness/experimental/step_persistence/_context.py
+34
View File
@@ -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.
"""
pydantic_ai_harness/experimental/step_persistence/_helpers.py
+119
View File
@@ -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,
)
)
pydantic_ai_harness/experimental/step_persistence/_store.py
+1032
View File
File diff suppressed because it is too large Load Diff
pydantic_ai_harness/experimental/step_persistence/_types.py
+129
View File
@@ -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)
pydantic_ai_harness/experimental/subagents/README.md
+19
View File
@@ -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
pyproject.toml
+2
View File
@@ -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',
tests/experimental/media/__init__.py
View File
tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_exists_present_after_put.yaml
+90
View File
@@ -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
tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_exists_returns_false_for_missing.yaml
+38
View File
@@ -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
tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_get_missing_raises.yaml
+38
View File
@@ -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
tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_get_round_trips_bytes.yaml
+90
View File
@@ -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
tests/experimental/media/cassettes/test_s3_cassettes/TestS3MediaStoreCassettes.test_put_succeeds.yaml
+46
View File
@@ -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
tests/experimental/media/conftest.py
+206
View File
@@ -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/',
)
tests/experimental/media/test_media.py
+821
View File
@@ -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',
)
)
tests/experimental/media/test_s3_cassettes.py
+63
View File
@@ -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)
tests/experimental/step_persistence/__init__.py
View File
tests/experimental/step_persistence/test_sqlite_and_media.py
+578
View File
@@ -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(<root>/media/)`.
Large `BinaryContent` payloads end up as `<root>/media/<sha256>.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 <root>/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]
tests/experimental/step_persistence/test_step_persistence.py
+1313
View File
File diff suppressed because it is too large Load Diff
tests/experimental/test_warnings.py
+6
View File
@@ -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)
uv.lock
Generated
+30
View File
@@ -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"