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

feat(step_persistence): rescue last provider-valid resume point on error

Browse Source 
`StepPersistence` previously left an errored run with no `ContinuableSnapshot`
for a provider-valid state that no node boundary had captured, so `continue_run`
could only resume from an earlier point or raise `LookupError`.

Add an error-path capture: when a run fails against a provider-valid history,
save one snapshot so the run still exposes its last safe resume point. Two
sites cover the two ways such a state is reached without a completed-node
snapshot:

- `on_model_request_error` saves the request payload directly -- e.g. a
  resolved tool cycle, which is only provider-valid as the next request is
  built, never at a `CallToolsNode` boundary.
- `on_run_error` reads a `latest_node_history` contextvar refreshed at every
  `after_node_run` boundary, rescuing a text response whose following
  `CallToolsNode` raises inside output validation. It cannot read `ctx.messages`
  directly: the `RunContext` passed to `on_run_error` holds the start-of-run
  history, which the graph rebinds away during the run.

Both saves are gated on `_is_resumable_history` (provider-valid plus at least
one `ModelResponse`), so a crash mid-tool-call leaves a dangling `ToolCallPart`
that is skipped and `latest_snapshot` never regresses to an unsendable point.
`on_run_error` compares by message count and skips when the store already holds
a newer snapshot, so a stale completed-node stash never supersedes a fresher
`on_model_request_error` save.

This is the gated (provider-validity) approach; the ungated simplification that
rides on pydantic-ai#6319 is deferred to a follow-up.

Closes #253
This commit is contained in:
David SF
2026-07-17 12:41:50 -05:00
parent e4839c654b
commit 59f4059877
5 changed files with 355 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/step-persistence.md
+5 -4
View File
@@ -144,12 +144,13 @@ asyncio.run(main())
### 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:
`continue_run` only returns the messages of the latest provider-valid snapshot for that `run_id`. Snapshots are written at these boundaries:
- after every `CallToolsNode` completes (all tool calls returned), and
- at `after_run`, as a fallback if the run reached no such boundary.
- after every `CallToolsNode` completes (all tool calls returned),
- at `after_run`, as a fallback if the run reached no such boundary, and
- when a run *fails* against a provider-valid history -- a model request that raises after a clean tool cycle, or output validation that raises after a text response -- so an errored run still exposes its last safe resume point.
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. If no continuable snapshot exists at all, `continue_run` raises `LookupError`.
A run that crashed mid-tool-call has events (`tool_call_started`) but no snapshot for that point: the dangling `ToolCallPart` is not provider-valid, so the error-path capture skips it. `continue_run` returns the snapshot from the previous safe boundary, not the failed step. If no continuable snapshot exists at all, `continue_run` raises `LookupError`.
## Run lineage: `parent_run_id`
pydantic_ai_harness/step_persistence/README.md
+8 -4
View File
@@ -146,13 +146,17 @@ 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:
snapshot for that `run_id`. Snapshots are written at these boundaries:
- after every `CallToolsNode` completes (all tool calls returned), and
- at `after_run`, as a fallback if the run reached no such boundary.
- after every `CallToolsNode` completes (all tool calls returned),
- at `after_run`, as a fallback if the run reached no such boundary, and
- when a run *fails* against a provider-valid history -- a model request that
raises after a clean tool cycle, or output validation that raises after a
text response -- so an errored run still exposes its last safe resume point.
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
snapshot for that point: the dangling `ToolCallPart` is not provider-valid, so
the error-path capture skips it. `continue_run` returns the snapshot from the
previous safe boundary, not the failed step. If no continuable snapshot
exists at all, `continue_run` raises `LookupError`.
pydantic_ai_harness/step_persistence/_capability.py
+100 -7
View File
@@ -10,12 +10,12 @@ 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.messages import ModelMessage, 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.step_persistence._context import current_run_id, snapshot_saved
from pydantic_ai_harness.step_persistence._context import current_run_id, latest_node_history, snapshot_saved
from pydantic_ai_harness.step_persistence._helpers import is_provider_valid
from pydantic_ai_harness.step_persistence._store import InMemoryStepStore, StepStore
from pydantic_ai_harness.step_persistence._types import (
@@ -31,6 +31,16 @@ def _empty_metadata() -> dict[str, str]:
return {}
def _is_resumable_history(messages: list[ModelMessage]) -> bool:
"""A history worth rescuing as a resume point on error.
Requires provider-validity (sendable to `Agent.run(message_history=...)`)
and at least one model response: a bare user prompt is equivalent to
restarting the run, so it is not worth persisting.
"""
return is_provider_valid(messages) and any(isinstance(message, ModelResponse) for message in messages)
@dataclass
class StepPersistence(AbstractCapability[AgentDepsT]):
"""Append-only step log + continuable snapshots + tool-effect ledger.
@@ -40,12 +50,16 @@ class StepPersistence(AbstractCapability[AgentDepsT]):
`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.
fallback save at `after_run` if the run reached no such boundary. A run
that *fails* against a provider-valid history also saves one, so an
errored run still exposes its last safe resume point (see
`on_model_request_error` and `on_run_error`).
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.
new continuable snapshot -- the dangling `ToolCallPart` is not
provider-valid, so the latest snapshot reflects the last provider-valid
state.
```python
from pydantic_ai import Agent
@@ -197,9 +211,11 @@ class StepPersistence(AbstractCapability[AgentDepsT]):
"""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)
history_token = latest_node_history.set(None)
try:
return await handler()
finally:
latest_node_history.reset(history_token)
snapshot_saved.reset(saved_token)
current_run_id.reset(token)
@@ -262,13 +278,73 @@ class StepPersistence(AbstractCapability[AgentDepsT]):
await self.store.append_event(self._make_event(ctx, kind='run_completed'))
return result
def _stash_provider_valid_history(self, ctx: RunContext[AgentDepsT], messages: list[ModelMessage]) -> None:
"""Record `messages` as the latest resume point for `on_run_error` when it is a meaningful one.
Called at node boundaries, where `ctx.messages` is the live history.
A contextvar carries it to `on_run_error`, which cannot read the live
history itself (its `RunContext` holds the start-of-run reference).
Only a completed node's write reaches `on_run_error`. `after_node_run`
never fires for a node that raises, so a failing node stashes nothing;
and a model request runs in an isolated context, so a contextvar write
inside it does not propagate to `on_run_error`. That is why the
model-request path saves directly instead (see `on_model_request_error`).
"""
if _is_resumable_history(messages):
latest_node_history.set((messages, ctx.run_step))
async def _save_continuable_snapshot(
self,
ctx: RunContext[AgentDepsT],
messages: list[ModelMessage],
step_index: int,
) -> None:
await self.store.save_snapshot(
ContinuableSnapshot(
run_id=self._effective_run_id(ctx),
step_index=step_index,
messages=messages,
conversation_id=ctx.conversation_id,
parent_run_id=self.parent_run_id,
agent_name=self.agent_name,
)
)
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."""
"""Rescue the last provider-valid resume point from a completed node, then emit `run_failed`.
Covers a text response whose following `CallToolsNode` raises inside
output validation: `after_node_run` stashed the provider-valid
`[prompt, text-response]` history, and this persists it.
Reads `latest_node_history` rather than `ctx.messages`: the
`RunContext` passed to `on_run_error` carries the start-of-run history,
not the live message list. The stash only ever holds a provider-valid,
past-the-prompt history, so a crash mid-tool-call leaves nothing to
rescue and `latest_snapshot` never regresses to an unsendable point.
The stash reflects the last *completed* node, so it can be older than a
snapshot `on_model_request_error` already saved for a failing request
later in the same run (that request runs in an isolated context and
saves to the store directly, so its newer history never reaches this
stash). Skip the save when the store already holds a newer resume point,
so a stale stash never supersedes it as `latest_snapshot`. Recency is
compared by message count, not `step_index`: absent a history-rewriting
processor a run's history only grows, so a longer capture is strictly
later, whereas `step_index` repeats across the boundaries within one
request cycle (a retried text response and its tool cycle share a step).
"""
stashed = latest_node_history.get()
if stashed is not None:
messages, step_index = stashed
existing = await self.store.latest_snapshot(run_id=self._effective_run_id(ctx))
if existing is None or len(messages) > len(existing.messages):
await self._save_continuable_snapshot(ctx, messages, step_index)
await self.store.append_event(self._make_event(ctx, kind='run_failed', error=repr(error)))
raise error
@@ -297,6 +373,18 @@ class StepPersistence(AbstractCapability[AgentDepsT]):
request_context: ModelRequestContext,
error: Exception,
) -> ModelResponse:
"""Rescue the request payload as a resume point when a model request fails.
The payload is the provider-valid history the run was about to send --
e.g. a resolved tool cycle after a clean `CallToolsNode`, which is
never a completed-node boundary (the tool return only enters the
history as this request is built). It is saved here, directly to the
store, because the contextvar path used by `on_run_error` cannot carry
a value out of a model request that raises.
"""
messages = list(request_context.messages)
if _is_resumable_history(messages):
await self._save_continuable_snapshot(ctx, messages, ctx.run_step)
await self.store.append_event(self._make_event(ctx, kind='model_request_failed', error=repr(error)))
raise error
@@ -408,9 +496,14 @@ class StepPersistence(AbstractCapability[AgentDepsT]):
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.
Every node boundary also refreshes `latest_node_history` so that
`on_run_error` can rescue the last provider-valid tail when a later
node raises before its own `after_node_run` fires.
"""
messages = list(ctx.messages)
self._stash_provider_valid_history(ctx, messages)
if isinstance(node, CallToolsNode):
messages = list(ctx.messages)
if is_provider_valid(messages):
await self.store.save_snapshot(
ContinuableSnapshot(
pydantic_ai_harness/step_persistence/_context.py
+20
View File
@@ -4,6 +4,8 @@ from __future__ import annotations
from contextvars import ContextVar
from pydantic_ai.messages import ModelMessage
current_run_id: ContextVar[str | None] = ContextVar(
'pydantic_ai_harness.step_persistence.current_run_id',
default=None,
@@ -32,3 +34,21 @@ 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.
"""
latest_node_history: ContextVar[tuple[list[ModelMessage], int] | None] = ContextVar(
'pydantic_ai_harness.step_persistence.latest_node_history',
default=None,
)
"""Async-context-local `(messages, step_index)` captured at the last completed node.
`after_node_run` refreshes this at every node boundary; `on_run_error` reads it
to rescue a provider-valid resume point that a node reached but that no
`after_node_run` persisted -- e.g. a text response whose subsequent
`CallToolsNode` raises inside output validation. The `RunContext` passed to
`on_run_error` carries the start-of-run history (a stale reference the graph
rebinds during the run), not the live message list, so the hook cannot read the
partial history from `ctx.messages` directly.
Reset to `None` in `wrap_run` so a run never inherits a prior run's tail.
Task-isolated like `current_run_id`.
"""
tests/step_persistence/test_step_persistence.py
+222 -1
View File
@@ -12,7 +12,7 @@ from pathlib import Path
from typing import Any
import pytest
from pydantic_ai import Agent, RunContext
from pydantic_ai import Agent, ModelRetry, RunContext
from pydantic_ai._agent_graph import GraphAgentState # pyright: ignore[reportPrivateUsage]
from pydantic_ai.messages import (
ModelMessage,
@@ -25,6 +25,7 @@ from pydantic_ai.messages import (
UserPromptPart,
)
from pydantic_ai.models import ModelRequestContext, ModelRequestParameters
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.run import AgentRunResult
from pydantic_ai.tools import ToolDefinition
@@ -1015,6 +1016,226 @@ class TestCrashMidToolCallContract:
assert len(snap_after_crash.messages) == len(result.all_messages())
class TestOnRunErrorSnapshot:
"""`on_run_error` rescues a provider-valid resume point no snapshot captured (#253).
A provider-valid history can exist at a boundary that no snapshot records:
the model request after a clean tool cycle (the resolved tool return only
enters the history as that request is built) or the `CallToolsNode` after a
text response (output validation raising there skips its `after_node_run`).
`on_run_error` persists the last such history, and skips when the crash left
a dangling tool call -- so `latest_snapshot` never regresses to an
unsendable point.
"""
async def test_rescues_provider_valid_history_when_next_model_request_raises(self) -> None:
"""Issue #253 acceptance test: request 1 tool call, request 2 raises.
The resolved tool cycle `[prompt, response, tool-return]` is
provider-valid with no open tool calls, but it is never a completed
node boundary, so without this behavior the run leaves no snapshot.
`on_run_error` rescues it from the `before_model_request` boundary.
"""
store = InMemoryStepStore()
def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
prior_responses = [m for m in messages if isinstance(m, ModelResponse)]
if not prior_responses:
return ModelResponse(parts=[ToolCallPart('lookup', {}, tool_call_id='lookup-1')])
raise RuntimeError('provider down on request 2')
agent: Agent[object, str] = Agent(
FunctionModel(model),
capabilities=[StepPersistence(store=store, agent_name='delegate')],
)
@agent.tool_plain
def lookup() -> str: # pyright: ignore[reportUnusedFunction]
return 'ok'
with pytest.raises(RuntimeError, match='provider down on request 2'):
await agent.run('go')
rid = await first_run_id(store)
assert 'run_failed' in [e.kind for e in await store.list_events(run_id=rid)]
snap = await store.latest_snapshot(run_id=rid)
assert snap is not None
assert is_provider_valid(snap.messages) is True
# The resolved tool cycle: the tool call has its return, no open calls.
assert any(
isinstance(part, ToolReturnPart) and part.tool_call_id == 'lookup-1'
for msg in snap.messages
if isinstance(msg, ModelRequest)
for part in msg.parts
)
async def test_rescues_text_response_when_output_validation_raises(self) -> None:
"""Text response then a hard output-validation error: the clean tail is rescued.
The terminal `CallToolsNode` raises inside output validation before its
`after_node_run` fires; `on_run_error` captures the provider-valid
`[prompt, text-response]` history from the `after_node_run` boundary.
"""
store = InMemoryStepStore()
agent: Agent[object, str] = Agent(
TestModel(custom_output_text='final answer'),
capabilities=[StepPersistence(store=store, agent_name='delegate')],
)
@agent.output_validator
def reject(value: str) -> str: # pyright: ignore[reportUnusedFunction]
raise RuntimeError('validator down')
with pytest.raises(RuntimeError, match='validator down'):
await agent.run('answer please')
rid = await first_run_id(store)
assert 'run_failed' in [e.kind for e in await store.list_events(run_id=rid)]
snap = await store.latest_snapshot(run_id=rid)
assert snap is not None
assert is_provider_valid(snap.messages) is True
assert isinstance(snap.messages[-1], ModelResponse)
assert any(isinstance(part, TextPart) for part in snap.messages[-1].parts)
async def test_skips_snapshot_when_crash_leaves_dangling_tool_call(self) -> None:
"""A crash mid-tool-call leaves a dangling call that is never a resume point."""
store = InMemoryStepStore()
def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
return ModelResponse(parts=[ToolCallPart('boom', {}, tool_call_id='boom-1')])
agent: Agent[object, str] = Agent(
FunctionModel(model),
capabilities=[StepPersistence(store=store, agent_name='delegate')],
)
@agent.tool_plain
def boom() -> str: # pyright: ignore[reportUnusedFunction]
raise ValueError('kaboom')
with pytest.raises(ValueError, match='kaboom'):
await agent.run('go')
rid = await first_run_id(store)
assert 'run_failed' in [e.kind for e in await store.list_events(run_id=rid)]
# The only history reached is the bare prompt (skipped) and the dangling
# `boom` call (provider-invalid, skipped): no false continuation point.
assert await store.latest_snapshot(run_id=rid) is None
async def test_does_not_regress_to_stale_stash_after_model_request_save(self) -> None:
"""A stale completed-node stash must not supersede a newer resolved-cycle save.
req1 text -> output-validation `ModelRetry` (run continues; the text
history is stashed into `latest_node_history`). req2 tool call resolves
cleanly. req3 raises: `on_model_request_error` saves the newer resolved
tool cycle, and `on_run_error` must not re-save the older text stash
over it (both save sites write to the same store, and `latest_snapshot`
returns the last write).
"""
store = InMemoryStepStore()
def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
n_resp = len([m for m in messages if isinstance(m, ModelResponse)])
if n_resp == 0:
return ModelResponse(parts=[TextPart('draft')])
if n_resp == 1:
return ModelResponse(parts=[ToolCallPart('lookup', {}, tool_call_id='lookup-1')])
raise RuntimeError('provider down on request 3')
agent: Agent[object, str] = Agent(
FunctionModel(model),
capabilities=[StepPersistence(store=store, agent_name='delegate')],
output_type=str,
)
retried = {'done': False}
@agent.output_validator
def gate(value: str) -> str: # pyright: ignore[reportUnusedFunction]
if not retried['done']:
retried['done'] = True
raise ModelRetry('call a tool first')
return value
@agent.tool_plain
def lookup() -> str: # pyright: ignore[reportUnusedFunction]
return 'ok'
with pytest.raises(RuntimeError, match='provider down on request 3'):
await agent.run('go')
rid = await first_run_id(store)
snap = await store.latest_snapshot(run_id=rid)
assert snap is not None
assert is_provider_valid(snap.messages) is True
# The newest safe resume point is the resolved tool cycle, not the older text history.
assert any(
isinstance(part, ToolReturnPart) and part.tool_call_id == 'lookup-1'
for msg in snap.messages
if isinstance(msg, ModelRequest)
for part in msg.parts
)
async def test_rescues_newest_text_history_over_earlier_saved_snapshot(self) -> None:
"""`on_run_error` still saves its stash when it is newer than the store's latest snapshot.
Two text responses each pass an output-validation `ModelRetry`, then the
second validator raises hard. The first cycle already saved a snapshot;
the `on_run_error` stash (the newer text history) supersedes it.
"""
store = InMemoryStepStore()
def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
n_resp = len([m for m in messages if isinstance(m, ModelResponse)])
return ModelResponse(parts=[TextPart('v1' if n_resp == 0 else 'v2')])
agent: Agent[object, str] = Agent(
FunctionModel(model),
capabilities=[StepPersistence(store=store, agent_name='delegate')],
output_type=str,
)
calls = {'n': 0}
@agent.output_validator
def gate(value: str) -> str: # pyright: ignore[reportUnusedFunction]
calls['n'] += 1
if calls['n'] == 1:
raise ModelRetry('retry once')
raise RuntimeError('validator down')
with pytest.raises(RuntimeError, match='validator down'):
await agent.run('go')
rid = await first_run_id(store)
snap = await store.latest_snapshot(run_id=rid)
assert snap is not None
assert is_provider_valid(snap.messages) is True
# The latest resume point is the newer history including `v2`, not the earlier `v1`-only snapshot.
assert any(
isinstance(part, TextPart) and part.content == 'v2'
for msg in snap.messages
if isinstance(msg, ModelResponse)
for part in msg.parts
)
async def test_on_run_error_without_stashed_history_saves_nothing(self) -> None:
"""Direct-hook: no provider-valid boundary reached -> no snapshot, `run_failed` still emitted.
Covers the `stashed is None` branch -- `on_run_error` invoked before any
boundary refreshed `latest_node_history`.
"""
store = InMemoryStepStore()
cap: StepPersistence[object] = StepPersistence(store=store)
ctx = build_run_context(run_id='r1', run_step=1)
with pytest.raises(RuntimeError, match='boom'):
await cap.on_run_error(ctx, error=RuntimeError('boom'))
assert await store.latest_snapshot(run_id='r1') is None
assert [e.kind for e in await store.list_events(run_id='r1')] == ['run_failed']
# ---------------------------------------------------------------------------
# Hook-level branches awkward to reach through Agent
# ---------------------------------------------------------------------------