mirror of
https://github.com/pydantic/pydantic-ai-harness.git
synced 2026-07-21 02:45:34 +00:00
David SFandGitHub
6c60f7cf0d
* feat(step_persistence): rescue last provider-valid resume point on error `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 * test(step_persistence): drop dead validator branch to restore 100% coverage The `gate` output-validator in `test_does_not_regress_to_stale_stash_after_model_request_save` had a `return value` branch that never executed: the run fails on request 3 before any output passes validation a second time. Simplify it to always raise `ModelRetry` (one retry is all the stash scenario needs), matching the sibling `test_rescues_newest_text_history_over_earlier_saved_snapshot`.
55 lines
2.2 KiB
Python
55 lines
2.2 KiB
Python
"""Shared async-context state for `StepPersistence` cross-capability coordination."""
|
|
|
|
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,
|
|
)
|
|
"""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.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.
|
|
"""
|
|
|
|
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`.
|
|
"""
|