mirror of
https://github.com/pydantic/pydantic-ai.git
synced 2026-10-11 21:59:51 +00:00
Add AbsurdDurability to the harness as a replacement for pydantic-ai-absurd (#8946)
This commit is contained in:
Aditya Vardhan
@@ -1133,6 +1133,7 @@ dependencies = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "absurd-sdk", marker = "extra == 'absurd'", specifier = ">=0.5.0,<1" },
|
||||
{ name = "agent-client-protocol", marker = "extra == 'acp'", specifier = ">=0.11,<0.12" },
|
||||
{ name = "aws-durable-execution-sdk-python", marker = "python_full_version >= '3.11' and extra == 'aws-lambda'", specifier = ">=1.7.0,<2" },
|
||||
{ name = "browser-use", marker = "python_full_version >= '3.11' and extra == 'browser-use'", specifier = ">=0.13.6,<0.14" },
|
||||
@@ -1179,7 +1180,7 @@ requires-dist = [
|
||||
{ name = "stackone-defender", extras = ["onnx"], marker = "python_full_version >= '3.11' and extra == 'prompt-injection-defender-ml'", specifier = ">=0.7.3,<0.8" },
|
||||
{ name = "youdotcom", marker = "extra == 'youdotcom'", specifier = ">=3.1.2,<3.2" },
|
||||
]
|
||||
provides-extras = ["acp", "anthropic", "aws-lambda", "browser-use", "cli", "code-mode", "codemode", "coder", "day-ai", "dbos", "dynamic-workflow", "e2b", "exa", "github", "google-workspace", "grain", "linear", "logfire", "logfire-mcp", "modal", "mongodb", "notion", "ordinal", "playwright", "posthog", "prompt-injection-defender", "prompt-injection-defender-ml", "pylon", "researcher", "skills", "slack", "sprites", "stackone", "temporal", "youdotcom"]
|
||||
provides-extras = ["absurd", "acp", "anthropic", "aws-lambda", "browser-use", "cli", "code-mode", "codemode", "coder", "day-ai", "dbos", "dynamic-workflow", "e2b", "exa", "github", "google-workspace", "grain", "linear", "logfire", "logfire-mcp", "modal", "mongodb", "notion", "ordinal", "playwright", "posthog", "prompt-injection-defender", "prompt-injection-defender-ml", "pylon", "researcher", "skills", "slack", "sprites", "stackone", "temporal", "youdotcom"]
|
||||
|
||||
[[package]]
|
||||
name = "pydantic-ai-slim"
|
||||
|
||||
@@ -571,6 +571,21 @@ jobs:
|
||||
BLOCKBUSTER_ENABLED: false
|
||||
# `tests/test_examples.py` type-checks the docs examples with pyright; once per run is enough.
|
||||
TYPECHECK_EXAMPLES: ${{ matrix.python-version == '3.14' }}
|
||||
# `tests/harness/absurd` runs on Absurd's schema; an unreachable server fails rather than skips.
|
||||
ABSURD_REQUIRE_LIVE: "1"
|
||||
ABSURD_TEST_DATABASE_URL: postgresql://postgres:[email protected]:5432/postgres
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16
|
||||
ports:
|
||||
- 5432:5432
|
||||
env:
|
||||
POSTGRES_PASSWORD: postgres
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U postgres"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
steps:
|
||||
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
|
||||
with:
|
||||
|
||||
@@ -150,7 +150,7 @@ class ResearchWorkflow(PydanticAIWorkflow):
|
||||
return result.output
|
||||
```
|
||||
|
||||
[DBOS](https://pydantic.dev/docs/ai/capabilities/durable_execution/dbos/) and [Prefect](https://pydantic.dev/docs/ai/capabilities/durable_execution/prefect/) attach the same way, first-party and co-maintained, with [Restate, AWS Lambda, Kitaru, and Airflow](https://pydantic.dev/docs/ai/capabilities/durable_execution/overview/) integrations besides.
|
||||
[DBOS](https://pydantic.dev/docs/ai/capabilities/durable_execution/dbos/) and [Prefect](https://pydantic.dev/docs/ai/capabilities/durable_execution/prefect/) attach the same way, first-party and co-maintained, with [Restate, AWS Lambda, Kitaru, Airflow, and Absurd](https://pydantic.dev/docs/ai/capabilities/durable_execution/overview/) integrations besides.
|
||||
|
||||
**Build this →** [Durable Execution](https://pydantic.dev/docs/ai/capabilities/durable_execution/overview/)
|
||||
|
||||
@@ -234,7 +234,7 @@ That [standalone image API](https://pydantic.dev/docs/ai/guides/image-generation
|
||||
|
||||
- **[Every interface](https://pydantic.dev/docs/ai/overview/interfaces/).** One agent definition runs as a [CLI](https://pydantic.dev/docs/ai/integrations/cli/), a [built-in web chat](https://pydantic.dev/docs/ai/guides/web/), or [realtime speech](https://pydantic.dev/docs/ai/realtime/overview/) (OpenAI Realtime, Gemini Live, Azure, xAI Grok Voice); [UI event streams](https://pydantic.dev/docs/ai/integrations/ui/overview/) (AG-UI, Vercel AI) connect it to your own frontend or anything else; [ACP](https://pydantic.dev/docs/ai/harness/acp/) serves it as an editor agent; and [GitHub Agentic Workflows](https://pydantic.dev/docs/ai/harness/gh-aw/) runs it headless on issues, pull requests or a schedule.
|
||||
|
||||
- **Durable execution.** [Durable execution](https://pydantic.dev/docs/ai/capabilities/durable_execution/overview/) on seven engines: Temporal, DBOS, Prefect, Restate, AWS Lambda, Kitaru, and Airflow, the first five co-maintained with the vendor teams. Agents survive restarts and run for days on the engine you already operate, with [human-in-the-loop approval](https://pydantic.dev/docs/ai/tools-toolsets/deferred-tools/#human-in-the-loop-tool-approval) built in.
|
||||
- **Durable execution.** [Durable execution](https://pydantic.dev/docs/ai/capabilities/durable_execution/overview/) on eight engines: Temporal, DBOS, Prefect, Restate, AWS Lambda, Kitaru, Airflow, and Absurd, the first five co-maintained with the vendor teams. Agents survive restarts and run for days on the engine you already operate, with [human-in-the-loop approval](https://pydantic.dev/docs/ai/tools-toolsets/deferred-tools/#human-in-the-loop-tool-approval) built in.
|
||||
|
||||
- **Coming from another framework?** The [comparisons](https://pydantic.dev/docs/ai/comparisons/overview/) show where Pydantic AI differs from LangChain, Google ADK, the Claude Agent SDK and seven more, and the [migration skills](https://pydantic.dev/docs/ai/comparisons/migrate-from-other-frameworks/) let your coding agent port an existing application over.
|
||||
|
||||
|
||||
@@ -153,6 +153,7 @@ Outside the loop: how runs persist, survive failures, and get observed and confi
|
||||
|---|---|---|
|
||||
| [Durable execution](../durable_execution/overview.md) | Core | Runs that survive restarts and failures on [Temporal](../durable_execution/temporal.md), [DBOS](../durable_execution/dbos.md), or [Prefect](../durable_execution/prefect.md), with [Restate](../durable_execution/restate.md), [Kitaru](../durable_execution/kitaru.md), and [Airflow](../durable_execution/airflow.md) integrations |
|
||||
| [AWS Lambda durability](https://pydantic.dev/docs/ai/harness/aws-lambda/) | Harness | Checkpoint model requests and tool calls into AWS Lambda durable function steps |
|
||||
| [Absurd durability](https://pydantic.dev/docs/ai/harness/absurd/) | Harness | Checkpoint model requests, MCP calls, and tool calls into Absurd steps on PostgreSQL |
|
||||
| [Step Persistence](https://pydantic.dev/docs/ai/harness/step-persistence/) | Harness | Save, restore, resume (`continue_run`), and fork (`fork_run`) runs; file/SQLite/Mongo backends |
|
||||
| [Instrumentation](instrumentation.md) | Core | OpenTelemetry GenAI spans for every model and tool call; the raw material for [Logfire](https://pydantic.dev/logfire) traces |
|
||||
| [Managed Prompt](https://pydantic.dev/docs/ai/harness/managed-prompt/) | Harness | Back instructions with a [Logfire](https://pydantic.dev/logfire)-managed prompt; version and roll out without redeploying |
|
||||
|
||||
@@ -14,7 +14,7 @@ The complete implementations for [Temporal][pydantic_ai.durable_exec.temporal.Te
|
||||
[Prefect][pydantic_ai.durable_exec.prefect.PrefectDurability] are useful references. The external
|
||||
[Restate](https://github.com/restatedev/sdk-python/tree/main/python/restate/ext/pydantic),
|
||||
[AWS Lambda](https://github.com/pydantic/pydantic-ai/tree/main/src/pydantic_ai_harness/pydantic_ai_harness/aws_lambda),
|
||||
and Absurd
|
||||
and [Absurd](https://github.com/pydantic/pydantic-ai/tree/main/src/pydantic_ai_harness/pydantic_ai_harness/absurd)
|
||||
integrations show the same public builder with JSON journals.
|
||||
|
||||
## Choose a backend tier
|
||||
|
||||
@@ -14,7 +14,7 @@ Pydantic AI allows you to build durable agents that can preserve their progress
|
||||
!!! note "Durability is not storage"
|
||||
A durable engine keeps one run alive across crashes and restarts. It does not store your chat threads: saving a conversation and picking it up later is a different problem with a much lighter answer, laid out in [Persistence](../persistence.md).
|
||||
|
||||
Pydantic AI supports seven durable execution solutions, plus a [builder](./backends.md) for any other engine. Five are co-maintained by the Pydantic and vendor teams:
|
||||
Pydantic AI supports eight durable execution solutions, plus a [builder](./backends.md) for any other engine. Five are co-maintained by the Pydantic and vendor teams:
|
||||
|
||||
- [Temporal](./temporal.md)
|
||||
- [DBOS](./dbos.md)
|
||||
@@ -26,3 +26,4 @@ Additional external SDK integrations:
|
||||
|
||||
- [Kitaru](./kitaru.md)
|
||||
- [Apache Airflow](./airflow.md)
|
||||
- [Absurd](https://pydantic.dev/docs/ai/harness/absurd/)
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
---
|
||||
title: Absurd Durability
|
||||
description: Make a Pydantic AI agent run durable on Absurd, a Postgres-based durable-execution engine -- model requests, MCP calls, and function tool calls are checkpointed into steps so a crashed worker resumes mid-run.
|
||||
---
|
||||
|
||||
# Absurd Durability
|
||||
|
||||
`AbsurdDurability` makes an agent run durable on [Absurd](https://github.com/earendil-works/absurd),
|
||||
a Postgres-based durable-execution engine. Run the agent inside an Absurd task handler and every
|
||||
model request, MCP call, and function tool call is checkpointed as a step, so a worker that crashes
|
||||
mid-run resumes from the last completed step instead of re-spending tokens on finished work. Outside
|
||||
a task the capability is transparent.
|
||||
|
||||
[Source](https://github.com/pydantic/pydantic-ai/tree/main/src/pydantic_ai_harness/pydantic_ai_harness/absurd/)
|
||||
|
||||
> While Pydantic AI Harness is on 0.x releases, the API may change between minor releases; when it does, deprecation warnings and release-note migration guidance tell you (or your agent) exactly how to upgrade. See the [version policy](index.md#version-policy).
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip/uv-add "pydantic-ai-harness[absurd]" "pydantic-ai-slim[openai]"
|
||||
```
|
||||
|
||||
Absurd stores its state in Postgres. Install the Absurd schema once per database (see the
|
||||
[Absurd repository](https://github.com/earendil-works/absurd)), then create a queue:
|
||||
|
||||
```python {test="skip"}
|
||||
from absurd_sdk import AsyncAbsurd
|
||||
|
||||
absurd = AsyncAbsurd('postgresql://localhost/absurd', queue_name='agents')
|
||||
await absurd.create_queue()
|
||||
```
|
||||
|
||||
## Quick start
|
||||
|
||||
Give the agent a `name`, register a task handler that runs it, then spawn tasks from a producer and
|
||||
execute them in a worker:
|
||||
|
||||
```python {test="skip"}
|
||||
from absurd_sdk import AsyncAbsurd, AsyncTaskContext, JsonValue
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.absurd import AbsurdDurability
|
||||
|
||||
absurd = AsyncAbsurd('postgresql://localhost/absurd', queue_name='agents')
|
||||
agent = Agent('openai:gpt-5', name='analyst', capabilities=[AbsurdDurability()])
|
||||
|
||||
|
||||
@absurd.register_task(name='analyse')
|
||||
async def analyse(params: JsonValue, ctx: AsyncTaskContext) -> JsonValue:
|
||||
assert isinstance(params, dict)
|
||||
result = await agent.run(params['prompt'])
|
||||
return {'output': result.output}
|
||||
|
||||
|
||||
# Producer: enqueue a task.
|
||||
await absurd.spawn('analyse', {'prompt': 'Summarize the Q3 report.'})
|
||||
|
||||
# Worker: claim and run tasks (in its own process).
|
||||
await absurd.start_worker()
|
||||
```
|
||||
|
||||
The handler must be async: a synchronous `TaskContext` raises a `UserError`.
|
||||
|
||||
## What gets checkpointed
|
||||
|
||||
After a crash, Absurd re-runs the task handler from the top. Plain Python in the handler runs again,
|
||||
but each completed step returns its stored result instead of re-issuing the model request or tool
|
||||
call. Step names are built from the agent's `name` and each toolset's `id`:
|
||||
|
||||
| Step name | Operation |
|
||||
|---|---|
|
||||
| `{name}__model.request` | one model request segment |
|
||||
| `{name}__model.request_stream` | one streamed model request segment |
|
||||
| `{name}__model.compact_messages` | one model message-compaction operation |
|
||||
| `{name}__model.cancel_suspended_response` | tearing down a suspended response |
|
||||
| `{name}__capability__{capability_id}.{operation}` | an operation contributed by another capability |
|
||||
| `{name}__function_toolset__{id}.validate_args` | running a function tool's `args_validator` |
|
||||
| `{name}__function_toolset__{id}.call_tool:{tool}` | a function tool call |
|
||||
| `{name}__mcp_server__{id}.get_tools` | listing an MCP server's tools |
|
||||
| `{name}__mcp_server__{id}.get_instructions` | an MCP server's instructions |
|
||||
| `{name}__mcp_server__{id}.call_tool` | an MCP tool call |
|
||||
| `{name}__event_stream_handler` | one event delivered to an `event_stream_handler` |
|
||||
|
||||
A model other than the agent's default adds its id to the step name
|
||||
(`{name}__model.request.{model_id}`). A step name that recurs, such as a second `agent.run()` in
|
||||
the same task or the same tool called twice, gets an encounter-order suffix (`#2`, `#3`, ...). Await
|
||||
one run before starting the next in the same task.
|
||||
|
||||
These are not checkpointed and run again on replay:
|
||||
|
||||
- a tool call that raises `ModelRetry`, `ToolFailed`, `CallDeferred`, or `ApprovalRequired`;
|
||||
- tools from a `DynamicToolset`.
|
||||
|
||||
## Constraints
|
||||
|
||||
- The agent needs a `name` (or pass `name=` to `AbsurdDurability`), and every function, MCP, or
|
||||
dynamic toolset needs a unique `id`; `Agent(...)` raises a `UserError` otherwise. Don't rename
|
||||
either once deployed: in-flight tasks re-run the steps whose names changed.
|
||||
- A checkpointed tool's return value is stored as JSON, so it must be JSON-serializable.
|
||||
- A step is checkpointed after it runs, so a crash between a tool's side effect and its checkpoint
|
||||
re-runs the tool. Keep side effects in tools and in an `event_stream_handler` idempotent.
|
||||
- Function, MCP, or dynamic toolsets passed per run with `run(toolsets=...)` inside a task raise a
|
||||
`UserError`; register them on the agent instead. `ExternalToolset` is allowed.
|
||||
- Streaming inside a task is not live: the model stream is consumed inside the step, then its
|
||||
captured events are replayed to the caller.
|
||||
|
||||
## Parallel execution
|
||||
|
||||
Tool calls run sequentially by default. Pass `parallel_execution_mode='parallel_ordered_events'` to
|
||||
run them concurrently; plain `'parallel'` is not supported. Stay sequential if a capability hook
|
||||
such as `before_tool_execute` or `wrap_tool_execute` awaits before the tool runs, as that can
|
||||
reorder concurrent calls of the same tool on replay.
|
||||
|
||||
## Code Mode
|
||||
|
||||
`AbsurdDurability` composes with [Code Mode](code-mode.md): each tool call made from `run_code` is
|
||||
its own step, served from its checkpoint on replay, while the `run_code` body itself re-runs.
|
||||
|
||||
## Migrating from `pydantic-ai-absurd`
|
||||
|
||||
Import `AbsurdDurability` and `AbsurdParallelExecutionMode` from `pydantic_ai_harness.absurd`
|
||||
instead of `pydantic_ai_absurd`. Step names and checkpoint payloads match
|
||||
[`pydantic-ai-absurd`](https://github.com/Kludex/pydantic-ai-absurd) 0.8, so tasks in flight during
|
||||
the switch resume without repeating work. A toolset without an `id` now needs one, and adding it
|
||||
re-runs that toolset's steps in tasks in flight. `AbsurdAgent`, `AbsurdModel`,
|
||||
`AbsurdFunctionToolset`, and `AbsurdMCPToolset` are not ported; move to the capability first.
|
||||
|
||||
## Relation to Step Persistence
|
||||
|
||||
`AbsurdDurability` resumes a single run after a crash. [Step Persistence](step-persistence.md)
|
||||
records runs so they can be resumed, forked, or replayed later as separate invocations. The two
|
||||
compose.
|
||||
|
||||
## Further reading
|
||||
|
||||
- [Pydantic AI capabilities](https://pydantic.dev/docs/ai/capabilities/overview/)
|
||||
- [Absurd](https://github.com/earendil-works/absurd)
|
||||
|
||||
## API reference
|
||||
|
||||
::: pydantic_ai_harness.absurd.AbsurdDurability
|
||||
@@ -244,6 +244,7 @@ Outside the loop: how runs persist, survive failures, and get observed and confi
|
||||
|---|---|---|
|
||||
| [Durable execution](../durable_execution/overview.md) | Core | Runs that survive restarts and failures on [Temporal](../durable_execution/temporal.md), [DBOS](../durable_execution/dbos.md), or [Prefect](../durable_execution/prefect.md), with [Restate](../durable_execution/restate.md), [Kitaru](../durable_execution/kitaru.md), and [Airflow](../durable_execution/airflow.md) integrations. See [what works on each engine](durable-execution.md) for Coder, Shell, and FileSystem |
|
||||
| [AWS Lambda durability](aws-lambda.md) | Harness | Checkpoint model requests and tool calls into AWS Lambda durable function steps |
|
||||
| [Absurd durability](absurd.md) | Harness | Checkpoint model requests, MCP calls, and tool calls into Absurd steps on PostgreSQL |
|
||||
| [Step Persistence](step-persistence.md) | Harness | Save, restore, resume (`continue_run`), and fork (`fork_run`) runs; file/SQLite/Mongo backends |
|
||||
| [Instrumentation](../capabilities/instrumentation.md) | Core | OpenTelemetry GenAI spans for every model and tool call; the raw material for [Logfire](https://pydantic.dev/logfire) traces |
|
||||
| [Logfire MCP](logfire-mcp.md) | Harness | Query Logfire telemetry and manage observability resources. |
|
||||
|
||||
+2
-2
@@ -161,7 +161,7 @@ From simple typed data extraction to complex, long-running multi-agent collabora
|
||||
return result.output
|
||||
```
|
||||
|
||||
[DBOS](durable_execution/dbos.md) and [Prefect](durable_execution/prefect.md) attach the same way, first-party and co-maintained, with [Restate, AWS Lambda, Kitaru, and Airflow](durable_execution/overview.md) integrations besides.
|
||||
[DBOS](durable_execution/dbos.md) and [Prefect](durable_execution/prefect.md) attach the same way, first-party and co-maintained, with [Restate, AWS Lambda, Kitaru, Airflow, and Absurd](durable_execution/overview.md) integrations besides.
|
||||
|
||||
**Build this →** [Durable Execution](durable_execution/overview.md)
|
||||
|
||||
@@ -247,7 +247,7 @@ From simple typed data extraction to complex, long-running multi-agent collabora
|
||||
|
||||
- **[Every interface](interfaces.md).** One agent definition runs as a [CLI](cli.md), a [built-in web chat](web.md), or [realtime speech](realtime/overview.md); [UI event streams](ui/overview.md) (AG-UI, Vercel AI) connect it to your own frontend or anything else; [ACP](https://pydantic.dev/docs/ai/harness/acp/) serves it as an editor agent; and [GitHub Agentic Workflows](https://pydantic.dev/docs/ai/harness/gh-aw/) runs it headless on issues, pull requests or a schedule.
|
||||
|
||||
- **Durable execution.** [Durable execution](durable_execution/overview.md) on seven engines: Temporal, DBOS, Prefect, Restate, AWS Lambda, Kitaru, and Airflow, the first five co-maintained with the vendor teams. Agents survive restarts and run for days on the engine you already operate, with [human-in-the-loop approval](deferred-tools.md#human-in-the-loop-tool-approval) built in.
|
||||
- **Durable execution.** [Durable execution](durable_execution/overview.md) on eight engines: Temporal, DBOS, Prefect, Restate, AWS Lambda, Kitaru, Airflow, and Absurd, the first five co-maintained with the vendor teams. Agents survive restarts and run for days on the engine you already operate, with [human-in-the-loop approval](deferred-tools.md#human-in-the-loop-tool-approval) built in.
|
||||
|
||||
- **Coming from another framework?** The [comparisons](comparisons/index.md) show where Pydantic AI differs from LangChain, Google ADK, the Claude Agent SDK and seven more, and the [migration skills](framework-migration.md) let your coding agent port an existing application over.
|
||||
|
||||
|
||||
@@ -620,6 +620,9 @@ navigation:
|
||||
- page: "AWS Lambda"
|
||||
path: "harness/aws-lambda.md"
|
||||
slug: "harness/aws-lambda"
|
||||
- page: "Absurd"
|
||||
path: "harness/absurd.md"
|
||||
slug: "harness/absurd"
|
||||
- page: "Kitaru"
|
||||
path: "durable_execution/kitaru.md"
|
||||
slug: "capabilities/durable_execution/kitaru"
|
||||
|
||||
+1
-1
@@ -603,7 +603,7 @@ snap = ["create"]
|
||||
|
||||
[tool.codespell]
|
||||
# Ref: https://github.com/codespell-project/codespell#using-a-config-file
|
||||
skip = '.git*,*.svg,*.lock,*.css,*.yaml'
|
||||
skip = '.git*,*.svg,*.lock,*.css,*.yaml,tests/harness/absurd/fixtures/absurd.sql'
|
||||
check-hidden = true
|
||||
# Ignore "formatting" like **L**anguage
|
||||
ignore-regex = '\*\*[A-Z]\*\*[a-z]+\b'
|
||||
|
||||
@@ -259,6 +259,7 @@ Outside the loop: how runs persist, survive failures, and get observed and confi
|
||||
|---|---|---|
|
||||
| [Durable execution](https://ai.pydantic.dev/capabilities/durable_execution/overview/) | Core | Runs that survive restarts and failures on [Temporal](https://ai.pydantic.dev/capabilities/durable_execution/temporal/), [DBOS](https://ai.pydantic.dev/capabilities/durable_execution/dbos/), or [Prefect](https://ai.pydantic.dev/capabilities/durable_execution/prefect/), with [Restate](https://ai.pydantic.dev/capabilities/durable_execution/restate/), [Kitaru](https://ai.pydantic.dev/capabilities/durable_execution/kitaru/), and [Airflow](https://ai.pydantic.dev/capabilities/durable_execution/airflow/) integrations. See [what works on each engine](https://pydantic.dev/docs/ai/harness/durable-execution/) for Coder, Shell, and FileSystem |
|
||||
| [AWS Lambda durability](pydantic_ai_harness/aws_lambda/) | Harness | Checkpoint model requests and tool calls into AWS Lambda durable function steps |
|
||||
| [Absurd durability](pydantic_ai_harness/absurd/) | Harness | Checkpoint model requests, MCP calls, and tool calls into Absurd steps on PostgreSQL |
|
||||
| [Step Persistence](pydantic_ai_harness/step_persistence/) | Harness | Save, restore, resume (`continue_run`), and fork (`fork_run`) runs; file/SQLite/Mongo backends |
|
||||
| [Instrumentation](https://ai.pydantic.dev/capabilities/instrumentation/) | Core | OpenTelemetry GenAI spans for every model and tool call; the raw material for [Logfire](https://pydantic.dev/logfire) traces |
|
||||
| [Logfire MCP](pydantic_ai_harness/logfire_mcp/) | Harness | Query Logfire telemetry and manage observability resources. |
|
||||
|
||||
+1
-1
@@ -12,7 +12,7 @@ Import `SubAgent` and `SubAgents` from `pydantic_ai_harness.subagents`. Harness
|
||||
| `RubricMiddleware` grading loop | an explicit evaluator/revision loop or focused capability | Preserve the invocation rubric, grader model and tools, revision feedback, iteration limit, terminal verdict, callbacks, and progress events. `TrajectoryJudge` provides cadence-based mid-run steering, so use it only when that different lifecycle is accepted. |
|
||||
| approval or `interrupt_on` | Pydantic AI `DeferredToolRequests`, `DeferredToolResults`, `ToolApproved`, and `ToolDenied` | Include `DeferredToolRequests` in `output_type`. When returned, persist `result.all_messages()`, resolve the request with `build_results(...)`, then call `run(message_history=..., deferred_tool_results=...)`. Use `HandleDeferredToolCalls` only for an in-process handler whose lifecycle is sufficient. The LangGraph `Command(resume=...)` protocol is not reproduced. |
|
||||
| provider-valid run snapshots | `StepPersistence` | It records step events, continuable snapshots, lineage, and tool effects; it is not arbitrary LangGraph graph-state checkpointing. |
|
||||
| durable model and tool execution | `TemporalDurability`, `DBOSDurability`, `PrefectDurability`, `AWSLambdaDurability`, or another matching integration | Select by the deployment contract and install its extra. Wire the integration's workflow or handler boundary, not only its capability; verify crash and redeploy recovery, replay, timers, signals, and side-effect idempotency in that runtime. |
|
||||
| durable model and tool execution | `TemporalDurability`, `DBOSDurability`, `PrefectDurability`, `AWSLambdaDurability`, `AbsurdDurability`, or another matching integration | Select by the deployment contract and install its extra. Wire the integration's workflow or handler boundary, not only its capability; verify crash and redeploy recovery, replay, timers, signals, and side-effect idempotency in that runtime. |
|
||||
| retries, fallbacks, and call limits | Pydantic AI tool/output retries, provider or model transport retries, `pydantic_ai.models.fallback.FallbackModel`, `UsageLimits`, and application or durable-workflow policy | Classify expected versus unexpected errors and retry scope. Match backoff, fallback order, tool-error visibility, counters, terminal propagation, and side-effect idempotency. |
|
||||
| LangGraph and Deep Agents event streaming | Pydantic AI streaming APIs plus Harness event handlers | Translate the application's selected modes and projections. Namespaced graphs and Deep Agents subagent handles with child identity, status, messages, tool calls, nested work, and output are not preserved automatically. |
|
||||
|
||||
|
||||
+2
-2
@@ -208,7 +208,7 @@ Load the references for the capabilities the task uses; each is self-contained.
|
||||
| Add guardrails, prompt-injection screening, spend limits, questions to the user, reminders, or a trajectory judge; repair malformed tool arguments | [Control and Safety](./references/CONTROL-AND-SAFETY.md) |
|
||||
| Research the web with Exa or You.com, use the `Researcher` stack, or drive a browser | [Research and Browsing](./references/RESEARCH-AND-BROWSING.md) |
|
||||
| Connect GitHub, Linear, Notion, Slack, Google Workspace, PostHog, Logfire, or another hosted service | [Hosted Integrations](./references/HOSTED-INTEGRATIONS.md) |
|
||||
| Save, resume, or fork runs; run under AWS Lambda; use managed prompts, runtime-created capabilities, ACP, GitHub Agentic Workflows, or agent specs | [Runtime and Extension](./references/RUNTIME-AND-EXTENSION.md) |
|
||||
| Save, resume, or fork runs; run under AWS Lambda or Absurd; use managed prompts, runtime-created capabilities, ACP, GitHub Agentic Workflows, or agent specs | [Runtime and Extension](./references/RUNTIME-AND-EXTENSION.md) |
|
||||
| Test an agent that uses harness capabilities, or debug a failing one | [Testing and Debugging](./references/TESTING-AND-DEBUGGING.md) |
|
||||
|
||||
## Install
|
||||
@@ -273,7 +273,7 @@ any:
|
||||
| [Control and Safety](./references/CONTROL-AND-SAFETY.md) | `RepairToolArguments` (`.repair_tool_arguments`); `InputGuardrail`, `OutputGuardrail`, `ToolGuardrail` (`.guardrails`); `PromptInjectionDefender` (`.prompt_injection_defender`, `[prompt-injection-defender]`); `SpendLimits` (`.spend`); `AskUser` (`.ask_user`); `SystemReminders` (`.system_reminders`); `TrajectoryJudge` (`.trajectory_judge`) |
|
||||
| [Research and Browsing](./references/RESEARCH-AND-BROWSING.md) | `Researcher` (`.researcher`, `[researcher]`); `ExaSearch`, `ExaAgent` (`.exa`, `[exa]`); `YouSearch`, `YouResearch` (`.youdotcom`, `[youdotcom]`); `BrowserUse` (`.browser_use`, `[browser-use]`); `PlaywrightBrowser` (`.playwright`, `[playwright]`) |
|
||||
| [Hosted Integrations](./references/HOSTED-INTEGRATIONS.md) | `GitHub` (`.github`, `[github]`); `Linear` (`.linear`, `[linear]`); `Notion` (`.notion`, `[notion]`); `GoogleWorkspace` (`.google_workspace`, `[google-workspace]`); `Slack` (`.slack`, `[slack]`); `StackOne` (`.stackone`, `[stackone]`); `Ordinal` (`.ordinal`, `[ordinal]`); `Grain` (`.grain`, `[grain]`); `DayAI` (`.day_ai`, `[day-ai]`); `PostHog` (`.posthog`, `[posthog]`); `Pylon` (`.pylon`, `[pylon]`); `LogfireMCP` (`.logfire_mcp`, `[logfire-mcp]`) |
|
||||
| [Runtime and Extension](./references/RUNTIME-AND-EXTENSION.md) | `StepPersistence` (`.step_persistence`, `[mongodb]` for MongoDB); `AWSLambdaDurability` (`.aws_lambda`, `[aws-lambda]`); `ManagedPrompt` (`.logfire`, `[logfire]`); `CapabilityCreation` (`.capability_creation`); experimental ACP server `run_acp_stdio` (`.experimental.acp`, `[acp]`) |
|
||||
| [Runtime and Extension](./references/RUNTIME-AND-EXTENSION.md) | `StepPersistence` (`.step_persistence`, `[mongodb]` for MongoDB); `AWSLambdaDurability` (`.aws_lambda`, `[aws-lambda]`); `AbsurdDurability` (`.absurd`, `[absurd]`); `ManagedPrompt` (`.logfire`, `[logfire]`); `CapabilityCreation` (`.capability_creation`); experimental ACP server `run_acp_stdio` (`.experimental.acp`, `[acp]`) |
|
||||
|
||||
For offline tests and debugging of any of these, load [Testing and Debugging](./references/TESTING-AND-DEBUGGING.md).
|
||||
|
||||
|
||||
+56
-10
@@ -1,10 +1,11 @@
|
||||
# Runtime and Extension
|
||||
|
||||
How a harness agent is persisted, made durable, configured, extended, and served. This covers saving and
|
||||
resuming runs (`StepPersistence`), AWS Lambda durable functions (`AWSLambdaDurability`), harness
|
||||
capabilities under core durable execution, Logfire-managed instructions (`ManagedPrompt`), agent-written
|
||||
capabilities (`CapabilityCreation`), loading harness capabilities from YAML/JSON specs, serving an agent
|
||||
to editors over ACP, and running one as a GitHub Agentic Workflow.
|
||||
resuming runs (`StepPersistence`), AWS Lambda durable functions (`AWSLambdaDurability`), Absurd durable
|
||||
tasks (`AbsurdDurability`), harness capabilities under core durable execution, Logfire-managed
|
||||
instructions (`ManagedPrompt`), agent-written capabilities (`CapabilityCreation`), loading harness
|
||||
capabilities from YAML/JSON specs, serving an agent to editors over ACP, and running one as a GitHub
|
||||
Agentic Workflow.
|
||||
|
||||
## Choose
|
||||
|
||||
@@ -13,6 +14,7 @@ to editors over ACP, and running one as a GitHub Agentic Workflow.
|
||||
| Resume, continue, or fork a run from saved history; audit tool side effects after a crash | `StepPersistence` |
|
||||
| Survive worker crashes with automatic replay (Temporal, DBOS, Prefect) | core durability capability; most harness capabilities work inside it |
|
||||
| Checkpoint every model/tool step on AWS Lambda durable functions | `AWSLambdaDurability` |
|
||||
| Checkpoint every model/tool step in Postgres with Absurd | `AbsurdDurability` |
|
||||
| Edit, version, and roll out the system prompt from Logfire without redeploying | `ManagedPrompt` |
|
||||
| Let the agent write new capabilities that load on the next run | `CapabilityCreation` |
|
||||
| Define the agent in YAML/JSON with harness capabilities | `Agent.from_file(..., custom_capability_types=[...])` |
|
||||
@@ -228,6 +230,49 @@ Gotchas:
|
||||
- `ctx.enqueue()` is unavailable inside a durable step. Do not detach work with `asyncio.create_task()`.
|
||||
- Budget: 3,000 operations and 100 MB of checkpointed state. Return references, not blobs.
|
||||
|
||||
## AbsurdDurability
|
||||
|
||||
Checkpoints every model request, function tool call, and MCP call as a step of an
|
||||
[Absurd](https://github.com/earendil-works/absurd) task, stored in Postgres. When a worker dies, the
|
||||
task re-runs from the top and completed steps return their stored results.
|
||||
|
||||
```bash
|
||||
uv add "pydantic-ai-harness[absurd]" "pydantic-ai-slim[openai]"
|
||||
```
|
||||
|
||||
```python {test="skip"}
|
||||
from absurd_sdk import AsyncAbsurd, AsyncTaskContext, JsonValue
|
||||
from pydantic_ai import Agent
|
||||
|
||||
from pydantic_ai_harness.absurd import AbsurdDurability
|
||||
|
||||
absurd = AsyncAbsurd('postgresql://localhost/absurd', queue_name='agents')
|
||||
agent = Agent('openai:gpt-5', name='analyst', capabilities=[AbsurdDurability()])
|
||||
|
||||
|
||||
@absurd.register_task(name='analyse')
|
||||
async def analyse(params: JsonValue, ctx: AsyncTaskContext) -> JsonValue:
|
||||
assert isinstance(params, dict)
|
||||
result = await agent.run(params['prompt'])
|
||||
return {'output': result.output}
|
||||
```
|
||||
|
||||
Parameters: `AbsurdDurability(*, models=None, event_stream_handler=None, name=None,
|
||||
parallel_execution_mode='sequential')`. `models` maps ids to extra models for `agent.run(model='<id>')`.
|
||||
`parallel_execution_mode` also accepts `'parallel_ordered_events'`, but not `'parallel'`.
|
||||
|
||||
Gotchas:
|
||||
|
||||
- A run is durable only inside an async Absurd task handler; a synchronous `TaskContext` raises
|
||||
`UserError`.
|
||||
- The agent needs a `name` (or `name=`), and every function, MCP, or dynamic toolset needs a unique
|
||||
`id`. Both are part of every step name, so renaming either makes in-flight tasks re-run those steps.
|
||||
- Tool return values are stored as JSON. `ModelRetry`, `ToolFailed`, `CallDeferred`, and
|
||||
`ApprovalRequired` are not checkpointed, and `DynamicToolset`s are not wrapped, so that work re-runs.
|
||||
- Function, MCP, or dynamic toolsets passed with `run(toolsets=...)` inside a task raise `UserError`;
|
||||
`ExternalToolset` is fine.
|
||||
- Keep tool side effects idempotent: a crash after a tool runs but before its step is saved re-runs it.
|
||||
|
||||
## ManagedPrompt
|
||||
|
||||
Resolves a Logfire-managed prompt once per run and uses it as the agent's instructions, with the label
|
||||
@@ -331,12 +376,12 @@ agent = Agent.from_spec(spec, custom_capability_types=[Planning, StepPersistence
|
||||
`Agent.from_file('agent.yaml', custom_capability_types=[...])` works the same way (install
|
||||
`pydantic-ai-slim[spec]` for YAML). Capability entries take the forms `Name`, `{Name: positional_arg}`,
|
||||
or `{Name: {kwargs}}`, and each class's `from_spec` builds the instance. Opted out (they hold callables
|
||||
or live agents, and raise `ValueError` if listed): `AWSLambdaDurability`, `CapabilityCreation`,
|
||||
`DynamicWorkflow`, `InputGuardrail`, `OutputGuardrail`, `SubAgents`, `SystemReminders`,
|
||||
`ToolGuardrail`, `TrajectoryJudge`. `AskUser` has a spec name but needs an `answerer` callable, so in
|
||||
practice it is passed in code too. Pass those instances with the `capabilities=` keyword of
|
||||
`Agent.from_spec`/`Agent.from_file`, which adds them to the spec's list. Keep secrets out of spec files
|
||||
and let capabilities read their env vars.
|
||||
or live agents, and raise `ValueError` if listed): `AWSLambdaDurability`, `AbsurdDurability`,
|
||||
`CapabilityCreation`, `DynamicWorkflow`, `InputGuardrail`, `OutputGuardrail`, `SubAgents`,
|
||||
`SystemReminders`, `ToolGuardrail`, `TrajectoryJudge`. `AskUser` has a spec name but needs an
|
||||
`answerer` callable, so in practice it is passed in code too. Pass those instances with the
|
||||
`capabilities=` keyword of `Agent.from_spec`/`Agent.from_file`, which adds them to the spec's list.
|
||||
Keep secrets out of spec files and let capabilities read their env vars.
|
||||
|
||||
A spec whose `model` is a provider string (`'openai:gpt-5'`) needs that provider's API key when the
|
||||
agent is built. In tests, pass `defer_model_check=True` and run under `agent.override(model='test')`.
|
||||
@@ -412,6 +457,7 @@ capabilities, because the CLI passes no `custom_capability_types`. Use a Python
|
||||
- https://pydantic.dev/docs/ai/harness/step-persistence/
|
||||
- https://pydantic.dev/docs/ai/harness/durable-execution/
|
||||
- https://pydantic.dev/docs/ai/harness/aws-lambda/
|
||||
- https://pydantic.dev/docs/ai/harness/absurd/
|
||||
- https://pydantic.dev/docs/ai/harness/managed-prompt/
|
||||
- https://pydantic.dev/docs/ai/harness/capability-creation/
|
||||
- https://pydantic.dev/docs/ai/harness/acp/
|
||||
|
||||
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
|
||||
from ._warn import HarnessDeprecationWarning
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .absurd import AbsurdDurability
|
||||
from .advisor import Advisor
|
||||
from .ask_user import AskUser
|
||||
from .background_tools import BackgroundTools
|
||||
@@ -70,6 +71,7 @@ if TYPE_CHECKING:
|
||||
from .youdotcom import YouResearch, YouSearch
|
||||
|
||||
__all__ = [
|
||||
'AbsurdDurability',
|
||||
'Advisor',
|
||||
'AskUser',
|
||||
'BackgroundTools',
|
||||
@@ -144,6 +146,7 @@ __all__ = [
|
||||
]
|
||||
|
||||
_CAPABILITY_EXPORTS = {
|
||||
'AbsurdDurability': 'absurd',
|
||||
'Advisor': 'advisor',
|
||||
'AskUser': 'ask_user',
|
||||
'BackgroundTools': 'background_tools',
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
# Absurd Durability
|
||||
|
||||
`AbsurdDurability` makes an agent run durable on [Absurd](https://github.com/earendil-works/absurd),
|
||||
a Postgres-based durable-execution engine. Run the agent inside an Absurd task handler and every
|
||||
model request, MCP call, and function tool call is checkpointed as a step, so a worker that crashes
|
||||
mid-run resumes from the last completed step instead of re-spending tokens on finished work. Outside
|
||||
a task the capability is transparent.
|
||||
|
||||
[Source](https://github.com/pydantic/pydantic-ai/tree/main/src/pydantic_ai_harness/pydantic_ai_harness/absurd/)
|
||||
|
||||
> While Pydantic AI Harness is on 0.x releases, the API may change between minor releases; when it does, deprecation warnings and release-note migration guidance tell you (or your agent) exactly how to upgrade. See the [version policy](https://github.com/pydantic/pydantic-ai-harness#version-policy).
|
||||
|
||||
## Installation
|
||||
|
||||
uv:
|
||||
|
||||
```bash
|
||||
uv add "pydantic-ai-harness[absurd]" "pydantic-ai-slim[openai]"
|
||||
```
|
||||
|
||||
pip:
|
||||
|
||||
```bash
|
||||
pip install "pydantic-ai-harness[absurd]" "pydantic-ai-slim[openai]"
|
||||
```
|
||||
|
||||
Absurd stores its state in Postgres. Install the Absurd schema once per database (see the
|
||||
[Absurd repository](https://github.com/earendil-works/absurd)), then create a queue:
|
||||
|
||||
```python {test="skip"}
|
||||
from absurd_sdk import AsyncAbsurd
|
||||
|
||||
absurd = AsyncAbsurd('postgresql://localhost/absurd', queue_name='agents')
|
||||
await absurd.create_queue()
|
||||
```
|
||||
|
||||
## Quick start
|
||||
|
||||
Give the agent a `name`, register a task handler that runs it, then spawn tasks from a producer and
|
||||
execute them in a worker:
|
||||
|
||||
```python {test="skip"}
|
||||
from absurd_sdk import AsyncAbsurd, AsyncTaskContext, JsonValue
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.absurd import AbsurdDurability
|
||||
|
||||
absurd = AsyncAbsurd('postgresql://localhost/absurd', queue_name='agents')
|
||||
agent = Agent('openai:gpt-5', name='analyst', capabilities=[AbsurdDurability()])
|
||||
|
||||
|
||||
@absurd.register_task(name='analyse')
|
||||
async def analyse(params: JsonValue, ctx: AsyncTaskContext) -> JsonValue:
|
||||
assert isinstance(params, dict)
|
||||
result = await agent.run(params['prompt'])
|
||||
return {'output': result.output}
|
||||
|
||||
|
||||
# Producer: enqueue a task.
|
||||
await absurd.spawn('analyse', {'prompt': 'Summarize the Q3 report.'})
|
||||
|
||||
# Worker: claim and run tasks (in its own process).
|
||||
await absurd.start_worker()
|
||||
```
|
||||
|
||||
The handler must be async: a synchronous `TaskContext` raises a `UserError`.
|
||||
|
||||
## What gets checkpointed
|
||||
|
||||
After a crash, Absurd re-runs the task handler from the top. Plain Python in the handler runs again,
|
||||
but each completed step returns its stored result instead of re-issuing the model request or tool
|
||||
call. Step names are built from the agent's `name` and each toolset's `id`:
|
||||
|
||||
| Step name | Operation |
|
||||
|---|---|
|
||||
| `{name}__model.request` | one model request segment |
|
||||
| `{name}__model.request_stream` | one streamed model request segment |
|
||||
| `{name}__model.compact_messages` | one model message-compaction operation |
|
||||
| `{name}__model.cancel_suspended_response` | tearing down a suspended response |
|
||||
| `{name}__capability__{capability_id}.{operation}` | an operation contributed by another capability |
|
||||
| `{name}__function_toolset__{id}.validate_args` | running a function tool's `args_validator` |
|
||||
| `{name}__function_toolset__{id}.call_tool:{tool}` | a function tool call |
|
||||
| `{name}__mcp_server__{id}.get_tools` | listing an MCP server's tools |
|
||||
| `{name}__mcp_server__{id}.get_instructions` | an MCP server's instructions |
|
||||
| `{name}__mcp_server__{id}.call_tool` | an MCP tool call |
|
||||
| `{name}__event_stream_handler` | one event delivered to an `event_stream_handler` |
|
||||
|
||||
A model other than the agent's default adds its id to the step name
|
||||
(`{name}__model.request.{model_id}`). A step name that recurs, such as a second `agent.run()` in
|
||||
the same task or the same tool called twice, gets an encounter-order suffix (`#2`, `#3`, ...). Await
|
||||
one run before starting the next in the same task.
|
||||
|
||||
These are not checkpointed and run again on replay:
|
||||
|
||||
- a tool call that raises `ModelRetry`, `ToolFailed`, `CallDeferred`, or `ApprovalRequired`;
|
||||
- tools from a `DynamicToolset`.
|
||||
|
||||
## Constraints
|
||||
|
||||
- The agent needs a `name` (or pass `name=` to `AbsurdDurability`), and every function, MCP, or
|
||||
dynamic toolset needs a unique `id`; `Agent(...)` raises a `UserError` otherwise. Don't rename
|
||||
either once deployed: in-flight tasks re-run the steps whose names changed.
|
||||
- A checkpointed tool's return value is stored as JSON, so it must be JSON-serializable.
|
||||
- A step is checkpointed after it runs, so a crash between a tool's side effect and its checkpoint
|
||||
re-runs the tool. Keep side effects in tools and in an `event_stream_handler` idempotent.
|
||||
- Function, MCP, or dynamic toolsets passed per run with `run(toolsets=...)` inside a task raise a
|
||||
`UserError`; register them on the agent instead. `ExternalToolset` is allowed.
|
||||
- Streaming inside a task is not live: the model stream is consumed inside the step, then its
|
||||
captured events are replayed to the caller.
|
||||
|
||||
## Parallel execution
|
||||
|
||||
Tool calls run sequentially by default. Pass `parallel_execution_mode='parallel_ordered_events'` to
|
||||
run them concurrently; plain `'parallel'` is not supported. Stay sequential if a capability hook
|
||||
such as `before_tool_execute` or `wrap_tool_execute` awaits before the tool runs, as that can
|
||||
reorder concurrent calls of the same tool on replay.
|
||||
|
||||
## Code Mode
|
||||
|
||||
`AbsurdDurability` composes with the harness `CodeMode` capability: each tool call made from
|
||||
`run_code` is its own step, served from its checkpoint on replay, while the `run_code` body itself
|
||||
re-runs.
|
||||
|
||||
## Migrating from `pydantic-ai-absurd`
|
||||
|
||||
Import `AbsurdDurability` and `AbsurdParallelExecutionMode` from `pydantic_ai_harness.absurd`
|
||||
instead of `pydantic_ai_absurd`. Step names and checkpoint payloads match
|
||||
[`pydantic-ai-absurd`](https://github.com/Kludex/pydantic-ai-absurd) 0.8, so tasks in flight during
|
||||
the switch resume without repeating work. A toolset without an `id` now needs one, and adding it
|
||||
re-runs that toolset's steps in tasks in flight. `AbsurdAgent`, `AbsurdModel`,
|
||||
`AbsurdFunctionToolset`, and `AbsurdMCPToolset` are not ported; move to the capability first.
|
||||
|
||||
## Relation to Step Persistence
|
||||
|
||||
`AbsurdDurability` resumes a single run after a crash. [Step Persistence](../step_persistence/)
|
||||
records runs so they can be resumed, forked, or replayed later as separate invocations. The two
|
||||
compose.
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Absurd durability capability: checkpoint an agent's I/O into Absurd steps for crash-resume."""
|
||||
|
||||
from pydantic_ai_harness.absurd._capability import AbsurdDurability, AbsurdParallelExecutionMode
|
||||
|
||||
__all__ = ['AbsurdDurability', 'AbsurdParallelExecutionMode']
|
||||
@@ -0,0 +1,107 @@
|
||||
"""Durable execution for Pydantic AI agents on the Absurd engine."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
try:
|
||||
import absurd_sdk # noqa: F401 # pyright: ignore[reportUnusedImport]
|
||||
except ImportError as _import_error: # pragma: no cover
|
||||
raise ImportError(
|
||||
'Please install the `absurd-sdk` package to use the Absurd durability capability, '
|
||||
'you can use the `absurd` optional group -- `pip install "pydantic-ai-harness[absurd]"`'
|
||||
) from _import_error
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, ClassVar, Literal
|
||||
|
||||
from pydantic_ai.agent import EventStreamHandler, ParallelExecutionMode
|
||||
from pydantic_ai.capabilities import WrapRunHandler
|
||||
from pydantic_ai.durable_exec import JSON_CODEC, BaseDurabilityCapability, DurabilityEngineSpec
|
||||
from pydantic_ai.models import Model
|
||||
from pydantic_ai.run import AgentRunResult
|
||||
from pydantic_ai.tools import AgentDepsT, RunContext
|
||||
|
||||
from ._context import ENGINE_NAME, current_async_task_context
|
||||
from ._operation_backend import AbsurdOperationBackend
|
||||
|
||||
AbsurdParallelExecutionMode = Literal['sequential', 'parallel_ordered_events']
|
||||
"""Tool-call execution modes usable with Absurd: a subset of `ParallelExecutionMode`. `'parallel'` is
|
||||
excluded because it emits tool events in completion order, which can differ on replay."""
|
||||
|
||||
|
||||
@dataclass(init=False)
|
||||
class AbsurdDurability(BaseDurabilityCapability[AgentDepsT]):
|
||||
"""Capability that makes an agent durable by checkpointing its I/O into Absurd steps.
|
||||
|
||||
Run the agent inside an async Absurd task handler: every model request, MCP call, and function
|
||||
tool call becomes a step, so a retried task resumes from the last completed step. Outside a task
|
||||
the capability is transparent.
|
||||
|
||||
Example:
|
||||
```python {test="skip"}
|
||||
from absurd_sdk import AsyncAbsurd, AsyncTaskContext, JsonValue
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai_harness.absurd import AbsurdDurability
|
||||
|
||||
absurd = AsyncAbsurd('postgresql://localhost/absurd', queue_name='agents')
|
||||
agent = Agent('openai:gpt-5', name='analyst', capabilities=[AbsurdDurability()])
|
||||
|
||||
|
||||
@absurd.register_task(name='analyse')
|
||||
async def analyse(params: JsonValue, ctx: AsyncTaskContext) -> JsonValue:
|
||||
assert isinstance(params, dict)
|
||||
result = await agent.run(params['prompt'])
|
||||
return {'output': result.output}
|
||||
```
|
||||
"""
|
||||
|
||||
engine_spec: ClassVar = DurabilityEngineSpec(
|
||||
engine_name=ENGINE_NAME,
|
||||
durable_unit_noun='step',
|
||||
durable_container_noun='task',
|
||||
codec=JSON_CODEC,
|
||||
wrapped_toolset_kinds=frozenset({'function', 'mcp'}),
|
||||
toolset_lifecycles={'function': 'enter-never', 'mcp': 'enter-never'},
|
||||
unsupported_runtime_toolset_kinds=frozenset({'function', 'mcp', 'dynamic'}),
|
||||
)
|
||||
# Absurd has no non-retryable error type, so a serialization failure retries under the task's `RetryStrategy`.
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
models: Mapping[str, Model] | None = None,
|
||||
event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
|
||||
name: str | None = None,
|
||||
parallel_execution_mode: AbsurdParallelExecutionMode = 'sequential',
|
||||
) -> None:
|
||||
"""Create an `AbsurdDurability` capability.
|
||||
|
||||
Args:
|
||||
models: Optional additional models keyed by ID for runtime model switching via
|
||||
`agent.run(model='<id>')`. The agent's primary model is always registered as
|
||||
`'default'`; the ID is folded into the checkpoint step name so a replay resolves
|
||||
to the same model.
|
||||
event_stream_handler: Optional event stream handler. Model events are handled live
|
||||
inside the model-request step; each tool event is handled in its own checkpointed
|
||||
step.
|
||||
name: Unique agent name used as the prefix for every checkpoint step. Defaults to the
|
||||
agent's `name` when the capability is bound.
|
||||
parallel_execution_mode: Tool-call execution mode applied to every run. `'parallel'` is
|
||||
excluded; see `AbsurdParallelExecutionMode`.
|
||||
"""
|
||||
super().__init__(models=models, event_stream_handler=event_stream_handler, name=name)
|
||||
self._parallel_execution_mode: ParallelExecutionMode = parallel_execution_mode
|
||||
|
||||
@property
|
||||
def in_durable_context(self) -> bool:
|
||||
return current_async_task_context() is not None
|
||||
|
||||
def get_durable_operation_backend(self) -> AbsurdOperationBackend:
|
||||
return AbsurdOperationBackend(agent_name=self.name, default_model_id=self.default_model_id)
|
||||
|
||||
async def wrap_run(self, ctx: RunContext[AgentDepsT], *, handler: WrapRunHandler) -> AgentRunResult[Any]:
|
||||
"""Apply the configured parallel-execution mode to every run, inside a task or not."""
|
||||
agent = self.agent
|
||||
assert agent is not None
|
||||
with agent.parallel_tool_call_execution_mode(self._parallel_execution_mode):
|
||||
return await handler()
|
||||
@@ -0,0 +1,19 @@
|
||||
from absurd_sdk import AsyncTaskContext, get_current_context
|
||||
|
||||
from pydantic_ai.exceptions import UserError
|
||||
|
||||
ENGINE_NAME = 'Absurd'
|
||||
|
||||
|
||||
def current_async_task_context() -> AsyncTaskContext | None:
|
||||
"""Return the active Absurd async task context, or `None` outside a task."""
|
||||
ctx = get_current_context()
|
||||
if ctx is None:
|
||||
return None
|
||||
if isinstance(ctx, AsyncTaskContext):
|
||||
return ctx
|
||||
raise UserError(
|
||||
f'{ENGINE_NAME} durability requires an async Absurd task context, but a synchronous '
|
||||
'`TaskContext` is active. Agent runs are async, so run the agent from an async task '
|
||||
'handler (`AsyncTaskContext`), not a synchronous one.'
|
||||
)
|
||||
@@ -0,0 +1,87 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TypeGuard
|
||||
|
||||
from pydantic_ai._utils import is_str_dict
|
||||
from pydantic_ai.durable_exec import (
|
||||
DurableOperationId,
|
||||
JournalCallableOperationBackend,
|
||||
RoleBasedOperationConfig,
|
||||
ToolsetCallToolId,
|
||||
)
|
||||
|
||||
from ._context import current_async_task_context
|
||||
|
||||
# A tool-call checkpoint holds the raw return value; a `ToolReturn` has no raw form, so it goes under this key.
|
||||
_ENVELOPE_KEY = '__pydantic_ai_harness_absurd_tool_result__'
|
||||
_RAW_RESULT_KINDS = frozenset({'tool_return', 'tool_content_result'})
|
||||
# Control flow (`ModelRetry`, `CallDeferred`, ...) is not checkpointed, so the call runs again on replay.
|
||||
_CONTROL_FLOW_KINDS = frozenset(
|
||||
{'approval_required', 'call_deferred', 'model_retry', 'validation_error', 'tool_failed'}
|
||||
)
|
||||
|
||||
# Absurd steps take no per-operation options.
|
||||
_NO_CONFIG = RoleBasedOperationConfig[None](model=None, event=None, capability=None, tool=None)
|
||||
|
||||
|
||||
def _is_tool_return_object(value: object) -> bool:
|
||||
return is_str_dict(value) and value.get('kind') == 'tool-return'
|
||||
|
||||
|
||||
def _is_envelope(stored: object) -> TypeGuard[dict[str, object]]:
|
||||
if not (is_str_dict(stored) and stored.keys() == {_ENVELOPE_KEY}):
|
||||
return False
|
||||
payload = stored[_ENVELOPE_KEY]
|
||||
return is_str_dict(payload) and payload.get('kind') in _RAW_RESULT_KINDS and 'result' in payload
|
||||
|
||||
|
||||
def _to_checkpoint(payload: dict[str, object]) -> object:
|
||||
"""Reduce an encoded `CallToolResult` to the stored checkpoint: the raw return value."""
|
||||
result = payload['result']
|
||||
# A raw value shaped like the envelope is enveloped too, so it reads back as itself.
|
||||
if _is_envelope(result) or (payload['kind'] == 'tool_return' and _is_tool_return_object(result)):
|
||||
return {_ENVELOPE_KEY: payload}
|
||||
return result
|
||||
|
||||
|
||||
def _from_checkpoint(stored: object) -> object:
|
||||
"""Rebuild the encoded `CallToolResult` for a stored checkpoint, raw or enveloped."""
|
||||
if _is_envelope(stored):
|
||||
return stored[_ENVELOPE_KEY]
|
||||
if _is_tool_return_object(stored):
|
||||
# A raw dict whose `kind` would otherwise be decoded as a `ToolReturn`.
|
||||
return {'kind': 'tool_content_result', 'result': stored}
|
||||
return {'kind': 'tool_return', 'result': stored}
|
||||
|
||||
|
||||
class AbsurdOperationBackend(JournalCallableOperationBackend[None]):
|
||||
def __init__(self, *, agent_name: str, default_model_id: str | None) -> None:
|
||||
super().__init__(agent_name=agent_name, default_model_id=default_model_id, config=_NO_CONFIG)
|
||||
|
||||
async def execute(
|
||||
self,
|
||||
*,
|
||||
operation_id: DurableOperationId,
|
||||
name: str,
|
||||
body: Callable[[], Awaitable[object]],
|
||||
cache_key: tuple[object, ...],
|
||||
config: None,
|
||||
) -> object:
|
||||
del cache_key, config
|
||||
task_ctx = current_async_task_context()
|
||||
assert task_ctx is not None
|
||||
if not isinstance(operation_id, ToolsetCallToolId):
|
||||
return await task_ctx.step(name, body)
|
||||
|
||||
handle = await task_ctx.begin_step(name)
|
||||
if handle.done:
|
||||
return _from_checkpoint(handle.state)
|
||||
payload = await body()
|
||||
assert is_str_dict(payload)
|
||||
kind = payload.get('kind')
|
||||
if kind in _CONTROL_FLOW_KINDS:
|
||||
return payload
|
||||
# Fail on a result kind this backend doesn't know, rather than silently not checkpointing it.
|
||||
assert kind in _RAW_RESULT_KINDS, kind
|
||||
return _from_checkpoint(await task_ctx.complete_step(handle, _to_checkpoint(payload)))
|
||||
@@ -65,6 +65,7 @@ acp = ["agent-client-protocol>=0.11,<0.12"]
|
||||
exa = ["exa-py>=2.16.0"]
|
||||
# 2.0 changed tool-result serialization in a way the durable bridge does not handle yet.
|
||||
aws-lambda = ["aws-durable-execution-sdk-python>=1.7.0,<2 ; python_full_version >= '3.11'"]
|
||||
absurd = ["absurd-sdk>=0.5.0,<1"]
|
||||
# 3.2 types `page_age` as `datetime | str`, and the toolset formats it as a `datetime`.
|
||||
youdotcom = ["youdotcom>=3.1.2,<3.2"]
|
||||
mongodb = ["pymongo>=4.17.0"]
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Enter real Absurd task contexts, the way a worker does.
|
||||
|
||||
`running_task_context` spawns and claims a task and enters its `AsyncTaskContext`, so the agent
|
||||
runs as it would inside a handler. `reenter_running_task` fails that run and re-claims the task,
|
||||
as Absurd does on a retry: the new context is hydrated from the checkpoints in Postgres, so a
|
||||
replay is Absurd's own. `checkpoints` reads the checkpoint table back.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip('absurd_sdk')
|
||||
|
||||
from absurd_sdk import (
|
||||
AsyncAbsurd,
|
||||
AsyncTaskContext,
|
||||
ClaimedTask,
|
||||
JsonValue,
|
||||
_create_async_task_context, # pyright: ignore[reportPrivateUsage]
|
||||
_current_task_context, # pyright: ignore[reportPrivateUsage]
|
||||
)
|
||||
from psycopg import AsyncConnection, sql
|
||||
from psycopg.rows import TupleRow
|
||||
|
||||
|
||||
def _conn(absurd: AsyncAbsurd) -> AsyncConnection[TupleRow]:
|
||||
conn: AsyncConnection[TupleRow] | None = absurd._conn # pyright: ignore[reportPrivateUsage]
|
||||
assert conn is not None
|
||||
return conn
|
||||
|
||||
|
||||
def _queue(absurd: AsyncAbsurd) -> str:
|
||||
return absurd._queue_name # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
async def _noop(params: JsonValue, ctx: AsyncTaskContext) -> JsonValue: # pragma: no cover
|
||||
return None
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _entered(absurd: AsyncAbsurd, task: ClaimedTask) -> AsyncGenerator[AsyncTaskContext]:
|
||||
ctx = await _create_async_task_context(task['task_id'], _conn(absurd), _queue(absurd), task, 120)
|
||||
token = _current_task_context.set(ctx)
|
||||
try:
|
||||
yield ctx
|
||||
finally:
|
||||
_current_task_context.reset(token)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def running_task_context(absurd: AsyncAbsurd, task_name: str = 'noop') -> AsyncGenerator[AsyncTaskContext]:
|
||||
"""Spawn, claim and enter a task context for the duration of the block.
|
||||
|
||||
The task stays `running` when the block exits, so `reenter_running_task` can retry it.
|
||||
"""
|
||||
absurd.register_task(name=task_name)(_noop)
|
||||
spawned = await absurd.spawn(task_name, None, max_attempts=2)
|
||||
[task] = await absurd.claim_tasks(batch_size=1)
|
||||
assert task['task_id'] == spawned['task_id']
|
||||
async with _entered(absurd, task) as ctx:
|
||||
yield ctx
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def reenter_running_task(absurd: AsyncAbsurd, task_id: str) -> AsyncGenerator[AsyncTaskContext]:
|
||||
"""Fail the task's current run and enter the retry, whose checkpoints come from Postgres."""
|
||||
conn = _conn(absurd)
|
||||
queue = _queue(absurd)
|
||||
cursor = await conn.execute(
|
||||
sql.SQL('SELECT run_id FROM absurd.{} WHERE task_id = %s AND state = %s').format(sql.Identifier(f'r_{queue}')),
|
||||
(task_id, 'running'),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
assert row is not None
|
||||
await conn.execute('SELECT absurd.fail_run(%s, %s, %s, %s)', (queue, row[0], '{"type": "test.Replay"}', None))
|
||||
[task] = await absurd.claim_tasks(batch_size=1)
|
||||
assert task['task_id'] == task_id
|
||||
async with _entered(absurd, task) as ctx:
|
||||
yield ctx
|
||||
|
||||
|
||||
async def checkpoints(absurd: AsyncAbsurd, task_id: str) -> dict[str, JsonValue]:
|
||||
"""The task's stored checkpoints, by step name, in the order they were written."""
|
||||
cursor = await _conn(absurd).execute(
|
||||
sql.SQL('SELECT checkpoint_name, state FROM absurd.{} WHERE task_id = %s ORDER BY updated_at').format(
|
||||
sql.Identifier(f'c_{_queue(absurd)}')
|
||||
),
|
||||
(task_id,),
|
||||
)
|
||||
return {name: state for name, state in await cursor.fetchall()}
|
||||
@@ -0,0 +1,85 @@
|
||||
"""A real Absurd schema on PostgreSQL, and task contexts entered the way an Absurd worker enters them.
|
||||
|
||||
Every test runs against Absurd's own SQL, so step naming, encounter-order disambiguation and
|
||||
checkpoint storage are Absurd's, not a stand-in's. The database is named by `ABSURD_TEST_DATABASE_URL`:
|
||||
|
||||
docker run -d -p 5432:5432 -e POSTGRES_PASSWORD=postgres postgres:16
|
||||
export ABSURD_TEST_DATABASE_URL=postgresql://postgres:[email protected]:5432/postgres
|
||||
|
||||
Without it the tests skip. With it but no reachable server they also skip, unless
|
||||
`ABSURD_REQUIRE_LIVE` is set (CI does), where they fail instead.
|
||||
|
||||
`fixtures/absurd.sql` is Absurd's `sql/absurd.sql` (https://github.com/earendil-works/absurd,
|
||||
Apache-2.0). Refresh it from the Absurd release matching the `absurd-sdk` floor in
|
||||
`pydantic-ai-harness[absurd]`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import AsyncGenerator
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from ...conftest import try_import
|
||||
|
||||
with try_import() as imports_successful:
|
||||
from absurd_sdk import AsyncAbsurd
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.rows import TupleRow
|
||||
|
||||
if TYPE_CHECKING:
|
||||
AsyncConn = AsyncConnection[TupleRow]
|
||||
|
||||
ABSURD_SQL = Path(__file__).parent / 'fixtures' / 'absurd.sql'
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def db_dsn() -> str:
|
||||
"""Install Absurd's schema, once, in the database named by `ABSURD_TEST_DATABASE_URL`."""
|
||||
import psycopg
|
||||
|
||||
dsn = os.environ.get('ABSURD_TEST_DATABASE_URL')
|
||||
if not dsn:
|
||||
pytest.skip('ABSURD_TEST_DATABASE_URL is not set')
|
||||
try:
|
||||
with psycopg.connect(dsn, autocommit=True, connect_timeout=5) as conn:
|
||||
# xdist workers share the database, and the schema script is not re-runnable.
|
||||
conn.execute('SELECT pg_advisory_lock(8126)')
|
||||
if conn.execute("SELECT to_regnamespace('absurd') IS NULL").fetchone() == (True,):
|
||||
conn.execute(ABSURD_SQL.read_bytes())
|
||||
conn.execute('SELECT pg_advisory_unlock(8126)')
|
||||
# Only when the configured server is unreachable, which CI never is.
|
||||
except psycopg.OperationalError as exc: # pragma: no cover
|
||||
message = f'PostgreSQL is unreachable at ABSURD_TEST_DATABASE_URL: {exc}'
|
||||
if os.environ.get('ABSURD_REQUIRE_LIVE', '').lower() in {'1', 'true', 'yes'}:
|
||||
pytest.fail(message)
|
||||
pytest.skip(message)
|
||||
return dsn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def async_conn(db_dsn: str) -> AsyncGenerator[AsyncConn]:
|
||||
"""An autocommit connection to the test database."""
|
||||
async with await AsyncConnection.connect(db_dsn, autocommit=True) as conn:
|
||||
yield conn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def queue_name() -> str:
|
||||
"""A queue name unique to the test, so runs on a kept database cannot collide."""
|
||||
return f'test_{uuid4().hex[:8]}'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def absurd(async_conn: AsyncConn, queue_name: str) -> AsyncGenerator[AsyncAbsurd]:
|
||||
"""An Absurd client on a queue of its own, dropped after the test."""
|
||||
client = AsyncAbsurd(async_conn, queue_name=queue_name)
|
||||
await client.create_queue()
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
await client.drop_queue()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,283 @@
|
||||
{
|
||||
"plain": {
|
||||
"analyst__mcp_server__calc.get_tools": {
|
||||
"add": {
|
||||
"kind": "function",
|
||||
"name": "add",
|
||||
"strict": null,
|
||||
"timeout": null,
|
||||
"metadata": {
|
||||
"meta": {
|
||||
"fastmcp": {
|
||||
"tags": []
|
||||
}
|
||||
},
|
||||
"task": false,
|
||||
"annotations": null
|
||||
},
|
||||
"tool_kind": null,
|
||||
"sequential": false,
|
||||
"toolset_id": null,
|
||||
"description": null,
|
||||
"with_native": null,
|
||||
"capability_id": null,
|
||||
"defer_loading": false,
|
||||
"return_schema": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"result"
|
||||
],
|
||||
"properties": {
|
||||
"result": {
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"x-fastmcp-wrap-result": true
|
||||
},
|
||||
"unless_native": null,
|
||||
"outer_typed_dict_key": null,
|
||||
"include_return_schema": null,
|
||||
"parameters_json_schema": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"a",
|
||||
"b"
|
||||
],
|
||||
"properties": {
|
||||
"a": {
|
||||
"type": "integer"
|
||||
},
|
||||
"b": {
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}
|
||||
}
|
||||
},
|
||||
"analyst__model.request": {
|
||||
"kind": "response",
|
||||
"model_name": "fn",
|
||||
"parts": [
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "report_finding",
|
||||
"args": {
|
||||
"title": "p99 up",
|
||||
"severity": "high"
|
||||
},
|
||||
"tool_call_id": "c1"
|
||||
},
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "query_metrics",
|
||||
"args": {
|
||||
"query": "errors"
|
||||
},
|
||||
"tool_call_id": "c2"
|
||||
}
|
||||
]
|
||||
},
|
||||
"analyst__function_toolset__findings.call_tool:report_finding": "recorded p99 up (high)",
|
||||
"analyst__function_toolset__metrics.call_tool:query_metrics": {
|
||||
"rows": 7
|
||||
},
|
||||
"analyst__model.request#2": {
|
||||
"kind": "response",
|
||||
"model_name": "fn",
|
||||
"parts": [
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "add",
|
||||
"args": {
|
||||
"a": 2,
|
||||
"b": 3
|
||||
},
|
||||
"tool_call_id": "c3"
|
||||
},
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "search_knowledge",
|
||||
"args": {
|
||||
"term": "deploys"
|
||||
},
|
||||
"tool_call_id": "c4"
|
||||
}
|
||||
]
|
||||
},
|
||||
"analyst__mcp_server__calc.call_tool": 5,
|
||||
"analyst__function_toolset__knowledge.call_tool:search_knowledge": {
|
||||
"hits": [
|
||||
"deploys-1",
|
||||
"deploys-2"
|
||||
]
|
||||
},
|
||||
"analyst__model.request#3": {
|
||||
"kind": "response",
|
||||
"model_name": "fn",
|
||||
"parts": [
|
||||
{
|
||||
"part_kind": "text",
|
||||
"content": "wrapping up"
|
||||
},
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "final_result",
|
||||
"args": {
|
||||
"report": "saw ['add', 'query_metrics', 'report_finding', 'search_knowledge']",
|
||||
"outcome": "completed"
|
||||
},
|
||||
"tool_call_id": "c5"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"with_retry": {
|
||||
"analyst__mcp_server__calc.get_tools": {
|
||||
"add": {
|
||||
"kind": "function",
|
||||
"name": "add",
|
||||
"strict": null,
|
||||
"timeout": null,
|
||||
"metadata": {
|
||||
"meta": {
|
||||
"fastmcp": {
|
||||
"tags": []
|
||||
}
|
||||
},
|
||||
"task": false,
|
||||
"annotations": null
|
||||
},
|
||||
"tool_kind": null,
|
||||
"sequential": false,
|
||||
"toolset_id": null,
|
||||
"description": null,
|
||||
"with_native": null,
|
||||
"capability_id": null,
|
||||
"defer_loading": false,
|
||||
"return_schema": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"result"
|
||||
],
|
||||
"properties": {
|
||||
"result": {
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"x-fastmcp-wrap-result": true
|
||||
},
|
||||
"unless_native": null,
|
||||
"outer_typed_dict_key": null,
|
||||
"include_return_schema": null,
|
||||
"parameters_json_schema": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"a",
|
||||
"b"
|
||||
],
|
||||
"properties": {
|
||||
"a": {
|
||||
"type": "integer"
|
||||
},
|
||||
"b": {
|
||||
"type": "integer"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}
|
||||
}
|
||||
},
|
||||
"analyst__model.request": {
|
||||
"kind": "response",
|
||||
"model_name": "fn",
|
||||
"parts": [
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "report_finding",
|
||||
"args": {
|
||||
"title": "p99 up",
|
||||
"severity": "bogus-high"
|
||||
},
|
||||
"tool_call_id": "c0"
|
||||
}
|
||||
]
|
||||
},
|
||||
"analyst__model.request#2": {
|
||||
"kind": "response",
|
||||
"model_name": "fn",
|
||||
"parts": [
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "report_finding",
|
||||
"args": {
|
||||
"title": "p99 up",
|
||||
"severity": "high"
|
||||
},
|
||||
"tool_call_id": "c1"
|
||||
},
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "query_metrics",
|
||||
"args": {
|
||||
"query": "errors"
|
||||
},
|
||||
"tool_call_id": "c2"
|
||||
}
|
||||
]
|
||||
},
|
||||
"analyst__function_toolset__findings.call_tool:report_finding#2": "recorded p99 up (high)",
|
||||
"analyst__function_toolset__metrics.call_tool:query_metrics": {
|
||||
"rows": 7
|
||||
},
|
||||
"analyst__model.request#3": {
|
||||
"kind": "response",
|
||||
"model_name": "fn",
|
||||
"parts": [
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "add",
|
||||
"args": {
|
||||
"a": 2,
|
||||
"b": 3
|
||||
},
|
||||
"tool_call_id": "c3"
|
||||
},
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "search_knowledge",
|
||||
"args": {
|
||||
"term": "deploys"
|
||||
},
|
||||
"tool_call_id": "c4"
|
||||
}
|
||||
]
|
||||
},
|
||||
"analyst__mcp_server__calc.call_tool": 5,
|
||||
"analyst__function_toolset__knowledge.call_tool:search_knowledge": {
|
||||
"hits": [
|
||||
"deploys-1",
|
||||
"deploys-2"
|
||||
]
|
||||
},
|
||||
"analyst__model.request#4": {
|
||||
"kind": "response",
|
||||
"model_name": "fn",
|
||||
"parts": [
|
||||
{
|
||||
"part_kind": "text",
|
||||
"content": "wrapping up"
|
||||
},
|
||||
{
|
||||
"part_kind": "tool-call",
|
||||
"tool_name": "final_result",
|
||||
"args": {
|
||||
"report": "saw ['add', 'query_metrics', 'report_finding', 'search_knowledge']",
|
||||
"outcome": "completed"
|
||||
},
|
||||
"tool_call_id": "c5"
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,818 @@
|
||||
"""`AbsurdDurability` against a real Absurd schema on PostgreSQL.
|
||||
|
||||
Each test enters a real task context (`_task.running_task_context`), and a replay fails the run
|
||||
and re-claims the task (`_task.reenter_running_task`), so step naming, encounter-order
|
||||
disambiguation and checkpoint storage are Absurd's own.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterable, AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
pytest.importorskip('absurd_sdk')
|
||||
pytest.importorskip('fastmcp')
|
||||
|
||||
from absurd_sdk import (
|
||||
AsyncAbsurd,
|
||||
JsonValue,
|
||||
TaskContext,
|
||||
_current_task_context, # pyright: ignore[reportPrivateUsage]
|
||||
)
|
||||
from fastmcp import FastMCP
|
||||
from inline_snapshot import snapshot
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from pydantic_ai import Agent, ToolReturn
|
||||
from pydantic_ai.capabilities import AbstractCapability, durable_operation
|
||||
from pydantic_ai.durable_exec._toolset import CallToolResult
|
||||
from pydantic_ai.exceptions import ModelRetry, UserError
|
||||
from pydantic_ai.mcp import MCPToolset
|
||||
from pydantic_ai.messages import (
|
||||
AgentStreamEvent,
|
||||
FunctionToolCallEvent,
|
||||
ModelMessage,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
PartDeltaEvent,
|
||||
PartStartEvent,
|
||||
RetryPromptPart,
|
||||
TextPart,
|
||||
ToolCallPart,
|
||||
ToolReturnPart,
|
||||
)
|
||||
from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
|
||||
from pydantic_ai.tools import RunContext
|
||||
from pydantic_ai.toolsets import AbstractToolset, DynamicToolset, ExternalToolset, FunctionToolset
|
||||
from pydantic_ai_harness.absurd import AbsurdDurability
|
||||
from pydantic_ai_harness.absurd._operation_backend import (
|
||||
_CONTROL_FLOW_KINDS, # pyright: ignore[reportPrivateUsage]
|
||||
_RAW_RESULT_KINDS, # pyright: ignore[reportPrivateUsage]
|
||||
)
|
||||
|
||||
from ._task import checkpoints, reenter_running_task, running_task_context
|
||||
|
||||
|
||||
def _make_model(counter: dict[str, int] | None = None, text: str = 'ok') -> FunctionModel:
|
||||
tally = counter if counter is not None else {'calls': 0}
|
||||
|
||||
def fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
tally['calls'] += 1
|
||||
return ModelResponse(parts=[TextPart(content=text)])
|
||||
|
||||
async def stream_fn(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str]:
|
||||
tally['calls'] += 1
|
||||
yield text
|
||||
|
||||
return FunctionModel(fn, stream_function=stream_fn, model_name='fn')
|
||||
|
||||
|
||||
def _tool_calling_model(*calls: ToolCallPart, counter: dict[str, int] | None = None) -> FunctionModel:
|
||||
"""Makes `calls` in one response, then answers `'done'` once they have returned or asked to retry."""
|
||||
tally = counter if counter is not None else {'calls': 0}
|
||||
|
||||
def fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
tally['calls'] += 1
|
||||
if any(isinstance(p, ToolReturnPart | RetryPromptPart) for m in messages for p in m.parts):
|
||||
return ModelResponse(parts=[TextPart(content='done')])
|
||||
return ModelResponse(parts=list(calls))
|
||||
|
||||
return FunctionModel(fn, model_name='fn')
|
||||
|
||||
|
||||
def _late_toolset(calls: dict[str, int]) -> FunctionToolset[object]:
|
||||
toolset = FunctionToolset[object](id='late')
|
||||
|
||||
@toolset.tool_plain
|
||||
def late() -> str:
|
||||
calls['calls'] += 1
|
||||
return 'late result'
|
||||
|
||||
return toolset
|
||||
|
||||
|
||||
def _dynamic_toolset(tool_calls: dict[str, int], *, id: str | None) -> DynamicToolset[object]:
|
||||
def build(ctx: RunContext[object]) -> FunctionToolset[object]:
|
||||
inner: FunctionToolset[object] = FunctionToolset(id='inner')
|
||||
|
||||
@inner.tool_plain
|
||||
def greet(name: str) -> str:
|
||||
tool_calls['calls'] += 1
|
||||
return f'hi {name}'
|
||||
|
||||
return inner
|
||||
|
||||
return DynamicToolset(build, id=id)
|
||||
|
||||
|
||||
_RUNTIME_TOOLSET_ERROR = 'cannot be added at runtime with Absurd'
|
||||
|
||||
|
||||
class TestDurability:
|
||||
async def test_requires_name(self) -> None:
|
||||
with pytest.raises(UserError, match='unique `name`'):
|
||||
Agent(_make_model(), capabilities=[AbsurdDurability()])
|
||||
|
||||
async def test_requires_model(self) -> None:
|
||||
with pytest.raises(UserError, match='needs to have a `model`'):
|
||||
Agent(name='a', capabilities=[AbsurdDurability()])
|
||||
|
||||
async def test_reserved_default_model_id_raises(self) -> None:
|
||||
with pytest.raises(UserError, match="'default' is reserved"):
|
||||
Agent(_make_model(), name='a', capabilities=[AbsurdDurability(models={'default': _make_model()})])
|
||||
|
||||
@pytest.mark.parametrize('kind', ['function', 'mcp', 'dynamic'])
|
||||
async def test_leaf_toolset_without_id_raises(self, kind: str) -> None:
|
||||
toolsets: dict[str, AbstractToolset[object]] = {
|
||||
'function': FunctionToolset[object](),
|
||||
'mcp': MCPToolset[object](FastMCP(name='calc')),
|
||||
'dynamic': _dynamic_toolset({'calls': 0}, id=None),
|
||||
}
|
||||
toolset = toolsets[kind]
|
||||
with pytest.raises(UserError, match='need to have a unique `id`'):
|
||||
Agent(_make_model(), name='a', toolsets=[toolset], capabilities=[AbsurdDurability()])
|
||||
|
||||
async def test_same_toolset_instance_in_two_places_is_wrapped_once(self, absurd: AsyncAbsurd) -> None:
|
||||
# Both mounts checkpoint through the one wrapper, so the two calls take the same step name in
|
||||
# encounter order.
|
||||
toolset = FunctionToolset[object](id='shared')
|
||||
|
||||
@toolset.tool_plain
|
||||
def echo(value: str) -> str:
|
||||
return value
|
||||
|
||||
agent = Agent(
|
||||
_tool_calling_model(ToolCallPart('a_echo', {'value': 'a'}), ToolCallPart('b_echo', {'value': 'b'})),
|
||||
name='a',
|
||||
toolsets=[toolset.prefixed('a'), toolset.prefixed('b')],
|
||||
capabilities=[AbsurdDurability()],
|
||||
)
|
||||
async with running_task_context(absurd) as ctx:
|
||||
await agent.run('hi')
|
||||
|
||||
stored = await checkpoints(absurd, ctx.task_id)
|
||||
step = 'a__function_toolset__shared.call_tool:echo'
|
||||
assert {k: v for k, v in stored.items() if 'call_tool' in k} == {step: 'a', f'{step}#2': 'b'}
|
||||
|
||||
async def test_duplicate_toolset_id_raises(self) -> None:
|
||||
toolsets = [FunctionToolset[object](id='tools'), FunctionToolset[object](id='tools')]
|
||||
with pytest.raises(UserError, match='same `id`'):
|
||||
Agent(_make_model(), name='a', toolsets=toolsets, capabilities=[AbsurdDurability()])
|
||||
|
||||
async def test_name_from_capability(self) -> None:
|
||||
agent = Agent(_make_model(), capabilities=[AbsurdDurability(name='custom')])
|
||||
bound = AbsurdDurability.from_agent(agent)
|
||||
assert bound is not None
|
||||
assert bound.name == 'custom'
|
||||
|
||||
async def test_from_agent_without_capability_returns_none(self) -> None:
|
||||
assert AbsurdDurability.from_agent(Agent(_make_model(), name='a')) is None
|
||||
|
||||
async def test_from_agent_multiple_raises(self) -> None:
|
||||
agent = Agent(_make_model(), name='a', capabilities=[AbsurdDurability(), AbsurdDurability()])
|
||||
with pytest.raises(UserError, match='at most one'):
|
||||
AbsurdDurability.from_agent(agent)
|
||||
|
||||
async def test_run_outside_task_is_transparent(self) -> None:
|
||||
counter = {'calls': 0}
|
||||
agent = Agent(_make_model(counter), name='a', capabilities=[AbsurdDurability()])
|
||||
result = await agent.run('hi')
|
||||
assert result.output == 'ok'
|
||||
assert counter['calls'] == 1
|
||||
|
||||
async def test_replay_serves_cached_model_response(self, absurd: AsyncAbsurd) -> None:
|
||||
counter = {'calls': 0}
|
||||
agent = Agent(_make_model(counter), name='crash', capabilities=[AbsurdDurability()])
|
||||
|
||||
async with running_task_context(absurd, 'crash') as ctx:
|
||||
first = await agent.run('hi')
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
replayed = await agent.run('hi')
|
||||
|
||||
assert counter['calls'] == 1
|
||||
assert replayed.output == first.output == 'ok'
|
||||
|
||||
async def test_replay_does_not_rerun_function_tool(self, absurd: AsyncAbsurd) -> None:
|
||||
tool_calls = {'calls': 0}
|
||||
toolset = FunctionToolset[object](id='tools')
|
||||
|
||||
@toolset.tool_plain
|
||||
def charge_card(amount: int) -> str:
|
||||
tool_calls['calls'] += 1
|
||||
return f'charged {amount}'
|
||||
|
||||
agent = Agent(
|
||||
_tool_calling_model(ToolCallPart('charge_card', {'amount': 42})),
|
||||
name='billing',
|
||||
toolsets=[toolset],
|
||||
capabilities=[AbsurdDurability()],
|
||||
)
|
||||
|
||||
async with running_task_context(absurd, 'billing') as ctx:
|
||||
first = await agent.run('charge it')
|
||||
# The raw tool return value is what is stored.
|
||||
stored = await checkpoints(absurd, ctx.task_id)
|
||||
assert stored['billing__function_toolset__tools.call_tool:charge_card'] == 'charged 42'
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
replayed = await agent.run('charge it')
|
||||
|
||||
assert tool_calls['calls'] == 1
|
||||
assert replayed.output == first.output == 'done'
|
||||
|
||||
async def test_registered_model_selected_per_run(self, absurd: AsyncAbsurd) -> None:
|
||||
# The selected model checkpoints under its own id-scoped step name, and a replay serves it.
|
||||
primary = {'calls': 0}
|
||||
cheap = {'calls': 0}
|
||||
agent = Agent(
|
||||
_make_model(primary, text='primary'),
|
||||
name='a',
|
||||
capabilities=[AbsurdDurability(models={'cheap': _make_model(cheap, text='cheap')})],
|
||||
)
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
default_result = await agent.run('hi')
|
||||
cheap_result = await agent.run('hi', model='cheap')
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
await agent.run('hi')
|
||||
replayed = await agent.run('hi', model='cheap')
|
||||
|
||||
assert default_result.output == 'primary'
|
||||
assert cheap_result.output == replayed.output == 'cheap'
|
||||
assert primary['calls'] == cheap['calls'] == 1
|
||||
assert list(await checkpoints(absurd, ctx.task_id)) == ['a__model.request', 'a__model.request.cheap']
|
||||
|
||||
@pytest.mark.parametrize('kind', ['function', 'dynamic'])
|
||||
async def test_runtime_toolset_rejected(self, absurd: AsyncAbsurd, kind: str) -> None:
|
||||
calls = {'calls': 0}
|
||||
late = _late_toolset(calls) if kind == 'function' else _dynamic_toolset(calls, id='late')
|
||||
agent: Agent[object, str] = Agent(_make_model(), name='a', capabilities=[AbsurdDurability()])
|
||||
async with running_task_context(absurd):
|
||||
with pytest.raises(UserError, match=_RUNTIME_TOOLSET_ERROR):
|
||||
await agent.run('hi', toolsets=[late])
|
||||
assert calls == {'calls': 0}
|
||||
|
||||
async def test_override_toolsets_rejected_inside_task(self, absurd: AsyncAbsurd) -> None:
|
||||
calls = {'calls': 0}
|
||||
agent: Agent[object, str] = Agent(
|
||||
_tool_calling_model(ToolCallPart('late')), name='a', capabilities=[AbsurdDurability()]
|
||||
)
|
||||
async with running_task_context(absurd):
|
||||
with agent.override(toolsets=[_late_toolset(calls)]):
|
||||
with pytest.raises(UserError, match=_RUNTIME_TOOLSET_ERROR):
|
||||
await agent.run('hi')
|
||||
assert calls['calls'] == 0
|
||||
|
||||
async def test_override_toolsets_respected_outside_task(self) -> None:
|
||||
calls = {'calls': 0}
|
||||
agent: Agent[object, str] = Agent(
|
||||
_tool_calling_model(ToolCallPart('late')), name='a', capabilities=[AbsurdDurability()]
|
||||
)
|
||||
with agent.override(toolsets=[_late_toolset(calls)]):
|
||||
result = await agent.run('hi')
|
||||
assert result.output == 'done'
|
||||
assert calls['calls'] == 1
|
||||
|
||||
async def test_override_tools_rejected_inside_task(self, absurd: AsyncAbsurd) -> None:
|
||||
calls = {'calls': 0}
|
||||
|
||||
def late() -> str: # pragma: no cover - rejected before it can run
|
||||
calls['calls'] += 1
|
||||
return 'late result'
|
||||
|
||||
agent: Agent[object, str] = Agent(
|
||||
_tool_calling_model(ToolCallPart('late')), name='a', capabilities=[AbsurdDurability()]
|
||||
)
|
||||
async with running_task_context(absurd):
|
||||
with agent.override(tools=[late]):
|
||||
with pytest.raises(UserError, match=_RUNTIME_TOOLSET_ERROR):
|
||||
await agent.run('hi')
|
||||
assert calls['calls'] == 0
|
||||
|
||||
async def test_override_tools_respected_outside_task(self) -> None:
|
||||
# The overriding toolset shares the agent's own toolset id, and must not be swapped for its wrapper.
|
||||
calls = {'calls': 0}
|
||||
|
||||
def late() -> str:
|
||||
calls['calls'] += 1
|
||||
return 'late result'
|
||||
|
||||
agent: Agent[object, str] = Agent(
|
||||
_tool_calling_model(ToolCallPart('late')), name='a', capabilities=[AbsurdDurability()]
|
||||
)
|
||||
with agent.override(tools=[late]):
|
||||
result = await agent.run('hi')
|
||||
assert result.output == 'done'
|
||||
assert calls['calls'] == 1
|
||||
|
||||
async def test_runtime_toolset_still_rejected_alongside_capability_toolset(self, absurd: AsyncAbsurd) -> None:
|
||||
# Skipping capability-owned wrappers must not let a genuine runtime toolset through.
|
||||
owned = FunctionToolset[object](id='owned')
|
||||
|
||||
class DemoCapability(AbstractCapability[object]):
|
||||
def get_toolset(self) -> FunctionToolset[object]:
|
||||
return owned
|
||||
|
||||
agent: Agent[object, str] = Agent(_make_model(), name='a', capabilities=[DemoCapability(), AbsurdDurability()])
|
||||
async with running_task_context(absurd):
|
||||
with pytest.raises(UserError, match=_RUNTIME_TOOLSET_ERROR):
|
||||
await agent.run('hi', toolsets=[_late_toolset({'calls': 0})])
|
||||
|
||||
async def test_capability_owned_toolset_is_durable(self, absurd: AsyncAbsurd) -> None:
|
||||
tool_calls = {'calls': 0}
|
||||
toolset = FunctionToolset[object](id='owned')
|
||||
|
||||
@toolset.tool_plain
|
||||
def charge_card(amount: int) -> str:
|
||||
tool_calls['calls'] += 1
|
||||
return f'charged {amount}'
|
||||
|
||||
class DemoCapability(AbstractCapability[object]):
|
||||
def get_toolset(self) -> FunctionToolset[object]:
|
||||
return toolset
|
||||
|
||||
agent: Agent[object, str] = Agent(
|
||||
_tool_calling_model(ToolCallPart('charge_card', {'amount': 5})),
|
||||
name='owner',
|
||||
capabilities=[DemoCapability(), AbsurdDurability()],
|
||||
)
|
||||
|
||||
async with running_task_context(absurd, 'owner') as ctx:
|
||||
first = await agent.run('charge it')
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
replayed = await agent.run('charge it')
|
||||
|
||||
assert tool_calls['calls'] == 1
|
||||
assert replayed.output == first.output == 'done'
|
||||
|
||||
async def test_runtime_external_toolset_allowed(self, absurd: AsyncAbsurd) -> None:
|
||||
agent: Agent[object, str] = Agent(_make_model(), name='a', capabilities=[AbsurdDurability()])
|
||||
async with running_task_context(absurd):
|
||||
result = await agent.run('hi', toolsets=[ExternalToolset[object](tool_defs=[])])
|
||||
assert result.output == 'ok'
|
||||
|
||||
async def test_construction_external_toolset_passes_through_unwrapped(self) -> None:
|
||||
external = ExternalToolset[object](tool_defs=[])
|
||||
agent = Agent(_make_model(), name='a', toolsets=[external], capabilities=[AbsurdDurability()])
|
||||
assert any(t is external for t in agent.toolsets)
|
||||
|
||||
async def test_event_stream_handler_receives_events(self, absurd: AsyncAbsurd) -> None:
|
||||
events: list[AgentStreamEvent] = []
|
||||
|
||||
async def handler(run_ctx: RunContext[object], stream: AsyncIterable[AgentStreamEvent]) -> None:
|
||||
async for event in stream:
|
||||
events.append(event)
|
||||
|
||||
async def stream_fn(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | DeltaToolCalls]:
|
||||
if len(messages) == 1:
|
||||
yield {0: DeltaToolCall(name='greet', json_args='{}')}
|
||||
else:
|
||||
yield 'done'
|
||||
|
||||
toolset = FunctionToolset[object](id='tools')
|
||||
|
||||
@toolset.tool_plain
|
||||
def greet() -> str:
|
||||
return 'hello'
|
||||
|
||||
agent = Agent(
|
||||
FunctionModel(stream_function=stream_fn, model_name='fn'),
|
||||
name='a',
|
||||
toolsets=[toolset],
|
||||
capabilities=[AbsurdDurability(event_stream_handler=handler)],
|
||||
)
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
result = await agent.run('hi')
|
||||
|
||||
assert result.output == 'done'
|
||||
assert any(isinstance(e, PartStartEvent | PartDeltaEvent) for e in events)
|
||||
assert any(isinstance(e, FunctionToolCallEvent) for e in events)
|
||||
assert 'a__event_stream_handler' in await checkpoints(absurd, ctx.task_id)
|
||||
|
||||
async def test_run_stream_events_inside_task(self, absurd: AsyncAbsurd) -> None:
|
||||
counter = {'calls': 0}
|
||||
agent = Agent(_make_model(counter), name='a', capabilities=[AbsurdDurability()])
|
||||
async with running_task_context(absurd) as ctx:
|
||||
async with agent.run_stream_events('hi') as stream:
|
||||
events = [event async for event in stream]
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
async with agent.run_stream_events('hi') as stream:
|
||||
replayed = [event async for event in stream]
|
||||
assert any(isinstance(e, PartStartEvent) for e in events)
|
||||
assert replayed == events
|
||||
assert counter['calls'] == 1
|
||||
|
||||
async def test_run_stream_inside_task_replays_buffered_stream(self, absurd: AsyncAbsurd) -> None:
|
||||
counter = {'calls': 0}
|
||||
agent = Agent(_make_model(counter), name='a', capabilities=[AbsurdDurability()])
|
||||
async with running_task_context(absurd) as ctx:
|
||||
async with agent.run_stream('hi') as result:
|
||||
assert await result.get_output() == 'ok'
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
async with agent.run_stream('hi') as result:
|
||||
assert await result.get_output() == 'ok'
|
||||
assert counter['calls'] == 1
|
||||
|
||||
async def test_iter_inside_task(self, absurd: AsyncAbsurd) -> None:
|
||||
agent = Agent(_make_model(), name='a', capabilities=[AbsurdDurability()])
|
||||
async with running_task_context(absurd) as ctx:
|
||||
async with agent.iter('hi') as run:
|
||||
async for _ in run:
|
||||
pass
|
||||
assert run.result is not None
|
||||
assert run.result.output == 'ok'
|
||||
assert list(await checkpoints(absurd, ctx.task_id)) == ['a__model.request']
|
||||
|
||||
async def test_cancel_suspended_response_is_checkpointed(self, absurd: AsyncAbsurd) -> None:
|
||||
# The model returns a `'suspended'` response, the continuation fails, and the graph tears the
|
||||
# suspended job down via `cancel_suspended_response`.
|
||||
cancelled: list[ModelResponse] = []
|
||||
|
||||
class CancellableModel(FunctionModel):
|
||||
async def cancel_suspended_response(self, response: ModelResponse) -> None:
|
||||
cancelled.append(response)
|
||||
|
||||
def fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
if not any(isinstance(m, ModelResponse) and m.state == 'suspended' for m in messages):
|
||||
return ModelResponse(parts=[TextPart(content='partial')], state='suspended')
|
||||
raise RuntimeError('continuation failed')
|
||||
|
||||
agent = Agent(CancellableModel(fn, model_name='fn'), name='a', capabilities=[AbsurdDurability()])
|
||||
async with running_task_context(absurd) as ctx:
|
||||
with pytest.raises(RuntimeError, match='continuation failed'):
|
||||
await agent.run('hi')
|
||||
|
||||
assert [response.state for response in cancelled] == ['suspended']
|
||||
assert 'a__model.cancel_suspended_response' in await checkpoints(absurd, ctx.task_id)
|
||||
|
||||
async def test_sync_context_raises(self) -> None:
|
||||
agent = Agent(_make_model(), name='a', capabilities=[AbsurdDurability()])
|
||||
sync_ctx: Any = object.__new__(TaskContext)
|
||||
token = _current_task_context.set(sync_ctx)
|
||||
try:
|
||||
with pytest.raises(UserError, match='requires an async Absurd task context'):
|
||||
await agent.run('hi')
|
||||
finally:
|
||||
_current_task_context.reset(token)
|
||||
|
||||
|
||||
class TestCapabilityOperation:
|
||||
async def test_operation_is_checkpointed_and_replayed(self, absurd: AsyncAbsurd) -> None:
|
||||
calls: list[str] = []
|
||||
|
||||
class Recorder(AbstractCapability[object]):
|
||||
id = 'recorder'
|
||||
|
||||
async def before_run(self, ctx: RunContext[object]) -> None:
|
||||
await self.record(ctx, 'started')
|
||||
|
||||
@durable_operation('record')
|
||||
async def record(self, ctx: RunContext[object], value: str) -> None:
|
||||
del ctx
|
||||
calls.append(value)
|
||||
|
||||
agent = Agent(_make_model(), name='cap', capabilities=[Recorder(), AbsurdDurability()])
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
await agent.run('hi')
|
||||
assert 'cap__capability__recorder.record' in await checkpoints(absurd, ctx.task_id)
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
await agent.run('hi')
|
||||
|
||||
assert calls == ['started']
|
||||
|
||||
|
||||
class TestToolResults:
|
||||
def test_every_core_result_kind_is_classified(self) -> None:
|
||||
# A result kind core adds must be sorted into checkpointed or not before it can reach a task.
|
||||
kinds = set(TypeAdapter(CallToolResult).json_schema()['discriminator']['mapping'])
|
||||
assert kinds == _RAW_RESULT_KINDS | _CONTROL_FLOW_KINDS
|
||||
|
||||
async def test_model_retry_is_not_checkpointed_and_the_tool_reruns_on_replay(self, absurd: AsyncAbsurd) -> None:
|
||||
model_calls = {'calls': 0}
|
||||
tool_calls = {'calls': 0}
|
||||
toolset = FunctionToolset(id='tools')
|
||||
|
||||
@toolset.tool_plain
|
||||
def flaky() -> str:
|
||||
tool_calls['calls'] += 1
|
||||
raise ModelRetry('nope, try again')
|
||||
|
||||
model = _tool_calling_model(ToolCallPart('flaky'), counter=model_calls)
|
||||
agent = Agent(model, name='retry', toolsets=[toolset], capabilities=[AbsurdDurability()])
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
first = await agent.run('go')
|
||||
# The `ModelRetry` propagates out of the step, so nothing is stored.
|
||||
assert list(await checkpoints(absurd, ctx.task_id)) == ['retry__model.request', 'retry__model.request#2']
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
second = await agent.run('go')
|
||||
|
||||
# Both model responses come from their checkpoints; only the tool runs again.
|
||||
assert first.output == second.output == 'done'
|
||||
assert model_calls['calls'] == tool_calls['calls'] == 2
|
||||
|
||||
async def test_tool_return_object_round_trips_through_replay(self, absurd: AsyncAbsurd) -> None:
|
||||
toolset = FunctionToolset(id='tools')
|
||||
|
||||
@toolset.tool_plain
|
||||
def lookup() -> ToolReturn:
|
||||
return ToolReturn(return_value='value', content='extra context', metadata={'source': 'db'})
|
||||
|
||||
agent = Agent(
|
||||
_tool_calling_model(ToolCallPart('lookup')),
|
||||
name='tr',
|
||||
toolsets=[toolset],
|
||||
capabilities=[AbsurdDurability()],
|
||||
)
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
first = await agent.run('go')
|
||||
# A `ToolReturn` has no raw form, so it is stored under a reserved key.
|
||||
stored = await checkpoints(absurd, ctx.task_id)
|
||||
assert stored['tr__function_toolset__tools.call_tool:lookup'] == snapshot(
|
||||
{
|
||||
'__pydantic_ai_harness_absurd_tool_result__': {
|
||||
'result': {
|
||||
'return_value': 'value',
|
||||
'content': 'extra context',
|
||||
'metadata': {'source': 'db'},
|
||||
'tools': None,
|
||||
'kind': 'tool-return',
|
||||
},
|
||||
'kind': 'tool_return',
|
||||
}
|
||||
}
|
||||
)
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
second = await agent.run('go')
|
||||
|
||||
def returned(messages: list[ModelMessage]) -> list[tuple[object, object]]:
|
||||
return [(p.content, p.metadata) for m in messages for p in m.parts if isinstance(p, ToolReturnPart)]
|
||||
|
||||
assert returned(second.all_messages()) == returned(first.all_messages()) == [('value', {'source': 'db'})]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'value',
|
||||
[
|
||||
{'kind': 'tool-return', 'rows': 3},
|
||||
{'__pydantic_ai_harness_absurd_tool_result__': 'hello'},
|
||||
{
|
||||
'__pydantic_ai_harness_absurd_tool_result__': {
|
||||
'kind': 'tool_return',
|
||||
'result': {'kind': 'tool-return', 'return_value': 'x', 'content': 'y'},
|
||||
}
|
||||
},
|
||||
],
|
||||
ids=['tool-return-kind', 'reserved-key', 'envelope-shaped'],
|
||||
)
|
||||
async def test_raw_dict_that_looks_encoded_round_trips(self, absurd: AsyncAbsurd, value: dict[str, object]) -> None:
|
||||
calls = {'calls': 0}
|
||||
toolset = FunctionToolset(id='tools')
|
||||
|
||||
@toolset.tool_plain
|
||||
def query() -> dict[str, object]:
|
||||
calls['calls'] += 1
|
||||
return value
|
||||
|
||||
agent = Agent(
|
||||
_tool_calling_model(ToolCallPart('query')),
|
||||
name='raw',
|
||||
toolsets=[toolset],
|
||||
capabilities=[AbsurdDurability()],
|
||||
)
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
first = await agent.run('go')
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
second = await agent.run('go')
|
||||
|
||||
def returned(messages: list[ModelMessage]) -> list[object]:
|
||||
return [p.content for m in messages for p in m.parts if isinstance(p, ToolReturnPart)]
|
||||
|
||||
assert returned(second.all_messages()) == returned(first.all_messages()) == [value]
|
||||
assert calls['calls'] == 1
|
||||
|
||||
|
||||
class TestCrashMidRun:
|
||||
async def test_model_step_served_from_checkpoint_while_failed_tool_reruns(self, absurd: AsyncAbsurd) -> None:
|
||||
# The model step completes and is checkpointed, then a tool raises a real (non-`ModelRetry`)
|
||||
# error that fails the attempt. On retry the model step is served from its checkpoint while
|
||||
# the tool re-runs.
|
||||
model_calls = {'calls': 0}
|
||||
tool_attempts = {'calls': 0}
|
||||
toolset = FunctionToolset(id='tools')
|
||||
|
||||
@toolset.tool_plain
|
||||
def flaky() -> str:
|
||||
tool_attempts['calls'] += 1
|
||||
if tool_attempts['calls'] == 1:
|
||||
raise RuntimeError('worker died mid-tool')
|
||||
return 'recovered'
|
||||
|
||||
model = _tool_calling_model(ToolCallPart('flaky'), counter=model_calls)
|
||||
agent = Agent(model, name='crash', toolsets=[toolset], capabilities=[AbsurdDurability()])
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
with pytest.raises(RuntimeError, match='worker died mid-tool'):
|
||||
await agent.run('go')
|
||||
assert list(await checkpoints(absurd, ctx.task_id)) == ['crash__model.request']
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
result = await agent.run('go')
|
||||
|
||||
assert result.output == 'done'
|
||||
assert model_calls['calls'] == tool_attempts['calls'] == 2
|
||||
assert list(await checkpoints(absurd, ctx.task_id)) == [
|
||||
'crash__model.request',
|
||||
'crash__function_toolset__tools.call_tool:flaky',
|
||||
'crash__model.request#2',
|
||||
]
|
||||
|
||||
|
||||
class TestStepNames:
|
||||
async def test_string_default_model_gets_unsuffixed_step_name(self, absurd: AsyncAbsurd) -> None:
|
||||
agent = Agent('test', name='strdef', capabilities=[AbsurdDurability()])
|
||||
async with running_task_context(absurd) as ctx:
|
||||
await agent.run('hi')
|
||||
assert list(await checkpoints(absurd, ctx.task_id)) == ['strdef__model.request']
|
||||
|
||||
async def test_two_runs_in_one_task_disambiguate_by_encounter_order(self, absurd: AsyncAbsurd) -> None:
|
||||
counter = {'calls': 0}
|
||||
agent = Agent(_make_model(counter), name='a', capabilities=[AbsurdDurability()])
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
await agent.run('hi')
|
||||
await agent.run('hi again')
|
||||
assert list(await checkpoints(absurd, ctx.task_id)) == ['a__model.request', 'a__model.request#2']
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
await agent.run('hi')
|
||||
await agent.run('hi again')
|
||||
|
||||
assert counter['calls'] == 2
|
||||
|
||||
|
||||
class TestParallelExecutionMode:
|
||||
async def test_step_slots_follow_scheduling_order_not_completion(self, absurd: AsyncAbsurd) -> None:
|
||||
# Two concurrent calls of the same tool, where the first-scheduled call completes last.
|
||||
# Absurd assigns the `#1`/`#2` slot when the step begins (`begin_step`), before the tool body
|
||||
# runs, so slots follow the model's tool-call order and a replay serves each call its own
|
||||
# result.
|
||||
toolset = FunctionToolset(id='tools')
|
||||
second_done = anyio.Event()
|
||||
|
||||
@toolset.tool_plain
|
||||
async def record(marker: str) -> str:
|
||||
if marker == 'first':
|
||||
await second_done.wait()
|
||||
else:
|
||||
second_done.set()
|
||||
return marker
|
||||
|
||||
agent = Agent(
|
||||
_tool_calling_model(
|
||||
ToolCallPart('record', {'marker': 'first'}, 'r1'), ToolCallPart('record', {'marker': 'second'}, 'r2')
|
||||
),
|
||||
name='par',
|
||||
toolsets=[toolset],
|
||||
capabilities=[AbsurdDurability(parallel_execution_mode='parallel_ordered_events')],
|
||||
)
|
||||
step = 'par__function_toolset__tools.call_tool:record'
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
first = await agent.run('go')
|
||||
stored = await checkpoints(absurd, ctx.task_id)
|
||||
assert (stored[step], stored[f'{step}#2']) == ('first', 'second')
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
second = await agent.run('go')
|
||||
|
||||
def returned(messages: list[ModelMessage]) -> list[tuple[str, object]]:
|
||||
return [(p.tool_call_id, p.content) for m in messages for p in m.parts if isinstance(p, ToolReturnPart)]
|
||||
|
||||
assert returned(second.all_messages()) == returned(first.all_messages()) == [('r1', 'first'), ('r2', 'second')]
|
||||
|
||||
|
||||
class TestMcpSessions:
|
||||
async def test_replay_opens_no_mcp_session(self, absurd: AsyncAbsurd) -> None:
|
||||
# The wrapper does not enter the server itself, so a replay served entirely from checkpoints
|
||||
# never connects to it.
|
||||
from tests.durable_exec.counting_mcp import counting_mcp_server
|
||||
|
||||
server, counts = counting_mcp_server(instructions='Echo things.')
|
||||
agent = Agent(
|
||||
_tool_calling_model(ToolCallPart('echo', {'text': 'hi'})),
|
||||
name='echo',
|
||||
toolsets=[MCPToolset[object](server, id='echo', include_instructions=True)],
|
||||
capabilities=[AbsurdDurability()],
|
||||
)
|
||||
async with running_task_context(absurd) as ctx:
|
||||
await agent.run('go')
|
||||
after_first_run = dict(counts)
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
result = await agent.run('go')
|
||||
|
||||
assert result.output == 'done'
|
||||
assert counts == after_first_run
|
||||
# The replayed run gets the server's instructions from their checkpoint.
|
||||
requests = [m for m in result.all_messages() if isinstance(m, ModelRequest)]
|
||||
assert requests[0].instructions is not None and 'Echo things.' in requests[0].instructions
|
||||
|
||||
|
||||
class TestCheckpointFormat:
|
||||
async def test_stream_checkpoint_replays(self, absurd: AsyncAbsurd) -> None:
|
||||
counter = {'calls': 0}
|
||||
agent = Agent(_make_model(counter), name='gold', capabilities=[AbsurdDurability()])
|
||||
# Pins the `{response, events}` payload shape of a stream checkpoint.
|
||||
payload: JsonValue = {
|
||||
'response': {
|
||||
'parts': [{'content': 'from-checkpoint', 'part_kind': 'text'}],
|
||||
'model_name': 'fn',
|
||||
'kind': 'response',
|
||||
},
|
||||
'events': [
|
||||
{'index': 0, 'part': {'content': 'from-checkpoint', 'part_kind': 'text'}, 'event_kind': 'part_start'}
|
||||
],
|
||||
}
|
||||
|
||||
async def write() -> JsonValue:
|
||||
return payload
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
await ctx.step('gold__model.request_stream', write)
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
async with agent.run_stream('hi') as result:
|
||||
assert await result.get_output() == 'from-checkpoint'
|
||||
|
||||
assert counter['calls'] == 0
|
||||
|
||||
|
||||
class TestDynamicToolset:
|
||||
"""Function and MCP toolsets are checkpointed; a construction-time `DynamicToolset` runs as-is."""
|
||||
|
||||
async def test_runs_uncheckpointed_inside_a_task(self, absurd: AsyncAbsurd) -> None:
|
||||
tool_calls = {'calls': 0}
|
||||
agent = Agent(
|
||||
_tool_calling_model(ToolCallPart('greet', {'name': 'ada'})),
|
||||
name='d',
|
||||
toolsets=[_dynamic_toolset(tool_calls, id='dyn')],
|
||||
capabilities=[AbsurdDurability()],
|
||||
)
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
first = await agent.run('greet ada')
|
||||
assert list(await checkpoints(absurd, ctx.task_id)) == ['d__model.request', 'd__model.request#2']
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
second = await agent.run('greet ada')
|
||||
|
||||
# The model responses replay from their checkpoints; the dynamic tool runs again.
|
||||
assert first.output == second.output == 'done'
|
||||
assert tool_calls['calls'] == 2
|
||||
|
||||
async def test_toolset_decorated_after_construction_rejected_inside_task(self, absurd: AsyncAbsurd) -> None:
|
||||
agent = Agent(_make_model(), name='decorated', capabilities=[AbsurdDurability()])
|
||||
|
||||
@agent.toolset(id='decorated-tools')
|
||||
def build(ctx: RunContext[object]) -> FunctionToolset[object]: # pragma: no cover
|
||||
return FunctionToolset[object](id='inner')
|
||||
|
||||
async with running_task_context(absurd):
|
||||
with pytest.raises(UserError, match=_RUNTIME_TOOLSET_ERROR):
|
||||
await agent.run('hi')
|
||||
|
||||
|
||||
class TestCodeMode:
|
||||
async def test_tool_call_inside_run_code_is_checkpointed(self, absurd: AsyncAbsurd) -> None:
|
||||
pytest.importorskip('pydantic_monty')
|
||||
from pydantic_ai_harness import CodeMode
|
||||
|
||||
calls = {'calls': 0}
|
||||
toolset = FunctionToolset(id='tools')
|
||||
|
||||
@toolset.tool_plain
|
||||
def search(query: str) -> str:
|
||||
calls['calls'] += 1
|
||||
return f'results for {query}'
|
||||
|
||||
def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
for part in (p for m in messages for p in m.parts):
|
||||
if isinstance(part, ToolReturnPart) and part.tool_name == 'run_code':
|
||||
return ModelResponse(parts=[TextPart(content=f'done: {part.content}')])
|
||||
code = "result = await search(query='x')\nresult"
|
||||
return ModelResponse(parts=[ToolCallPart('run_code', {'code': code}, 'tc1')])
|
||||
|
||||
agent = Agent(
|
||||
FunctionModel(model_fn), name='composed', toolsets=[toolset], capabilities=[CodeMode(), AbsurdDurability()]
|
||||
)
|
||||
|
||||
async with running_task_context(absurd) as ctx:
|
||||
first = await agent.run('go')
|
||||
stored = await checkpoints(absurd, ctx.task_id)
|
||||
assert stored['composed__function_toolset__tools.call_tool:search'] == 'results for x'
|
||||
async with reenter_running_task(absurd, ctx.task_id):
|
||||
second = await agent.run('go')
|
||||
|
||||
# The `run_code` body re-runs on replay, but its `search` call is served from the checkpoint.
|
||||
assert first.output == second.output == 'done: results for x'
|
||||
assert calls['calls'] == 1
|
||||
@@ -0,0 +1,205 @@
|
||||
"""Pins the stored checkpoint format against recorded runs, on a real Absurd PostgreSQL schema.
|
||||
|
||||
The fixture holds checkpoints recorded from runs of the agent built by `_agent` below (a logical
|
||||
string model ID resolved by a capability, function toolsets, a capability-contributed toolset, an MCP
|
||||
server, and structured output). Model payloads are trimmed to the fields a replay reads; tool payloads
|
||||
are verbatim.
|
||||
|
||||
A recorded run must resume without re-running any checkpointed step, and a fresh run must write the
|
||||
same step names and payloads, so tasks in flight across a deploy resume without repeating work.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip('absurd_sdk')
|
||||
pytest.importorskip('fastmcp')
|
||||
|
||||
from absurd_sdk import AsyncAbsurd, AsyncTaskContext, JsonValue
|
||||
from fastmcp import FastMCP
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.rows import TupleRow
|
||||
from pydantic import BaseModel
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.capabilities import AbstractCapability, ResolveModelId
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.mcp import MCPToolset
|
||||
from pydantic_ai.messages import ModelMessage, ModelResponse, RetryPromptPart, TextPart, ToolCallPart, ToolReturnPart
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
from pydantic_ai.toolsets import FunctionToolset
|
||||
from pydantic_ai_harness.absurd import AbsurdDurability
|
||||
|
||||
from ._task import checkpoints
|
||||
|
||||
GOLDEN: dict[str, dict[str, JsonValue]] = json.loads(
|
||||
(Path(__file__).parent / 'fixtures' / 'recorded_checkpoints.json').read_text()
|
||||
)
|
||||
# Fixture key -> whether the first `report_finding` call raises `ModelRetry`.
|
||||
CASES = {'plain': False, 'with_retry': True}
|
||||
|
||||
|
||||
class Report(BaseModel):
|
||||
outcome: Literal['completed', 'partial', 'failed']
|
||||
report: str
|
||||
|
||||
|
||||
def _agent(retry_first: bool, executions: list[str]) -> Agent[object, Report]:
|
||||
def script(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
executions.append('model')
|
||||
done = [p.tool_name for m in messages for p in m.parts if isinstance(p, ToolReturnPart)]
|
||||
retried = any(isinstance(p, RetryPromptPart) for m in messages for p in m.parts)
|
||||
if not done and retry_first and not retried:
|
||||
return ModelResponse(
|
||||
parts=[ToolCallPart('report_finding', {'title': 'p99 up', 'severity': 'bogus-high'}, 'c0')]
|
||||
)
|
||||
if not done:
|
||||
return ModelResponse(
|
||||
parts=[
|
||||
ToolCallPart('report_finding', {'title': 'p99 up', 'severity': 'high'}, 'c1'),
|
||||
ToolCallPart('query_metrics', {'query': 'errors'}, 'c2'),
|
||||
]
|
||||
)
|
||||
if 'add' not in done:
|
||||
return ModelResponse(
|
||||
parts=[
|
||||
ToolCallPart('add', {'a': 2, 'b': 3}, 'c3'),
|
||||
ToolCallPart('search_knowledge', {'term': 'deploys'}, 'c4'),
|
||||
]
|
||||
)
|
||||
return ModelResponse(
|
||||
parts=[
|
||||
TextPart('wrapping up'),
|
||||
ToolCallPart('final_result', {'outcome': 'completed', 'report': f'saw {sorted(done)}'}, 'c5'),
|
||||
]
|
||||
)
|
||||
|
||||
model = FunctionModel(script, model_name='fn')
|
||||
|
||||
findings = FunctionToolset[object](id='findings')
|
||||
|
||||
@findings.tool_plain
|
||||
def report_finding(title: str, severity: str) -> str:
|
||||
if severity.startswith('bogus'):
|
||||
executions.append('report_finding:retry')
|
||||
raise ModelRetry('severity must be low or high')
|
||||
executions.append('report_finding')
|
||||
return f'recorded {title} ({severity})'
|
||||
|
||||
metrics = FunctionToolset[object](id='metrics')
|
||||
|
||||
@metrics.tool_plain
|
||||
def query_metrics(query: str) -> dict[str, int]:
|
||||
executions.append('query_metrics')
|
||||
return {'rows': 7}
|
||||
|
||||
knowledge = FunctionToolset[object](id='knowledge')
|
||||
|
||||
@knowledge.tool_plain
|
||||
def search_knowledge(term: str) -> dict[str, list[str]]:
|
||||
executions.append('search_knowledge')
|
||||
return {'hits': [f'{term}-1', f'{term}-2']}
|
||||
|
||||
class ProjectKnowledge(AbstractCapability[object]):
|
||||
def get_toolset(self) -> FunctionToolset[object]:
|
||||
return knowledge
|
||||
|
||||
server: FastMCP[object] = FastMCP(name='calc')
|
||||
|
||||
@server.tool
|
||||
def add(a: int, b: int) -> int:
|
||||
executions.append('add')
|
||||
return a + b
|
||||
|
||||
return Agent(
|
||||
'custom:analyst',
|
||||
name='analyst',
|
||||
output_type=Report,
|
||||
toolsets=[findings, metrics, MCPToolset[object](server, id='calc')],
|
||||
capabilities=[
|
||||
ResolveModelId(lambda ctx, model_id: model if model_id == 'custom:analyst' else None),
|
||||
ProjectKnowledge(),
|
||||
AbsurdDurability(),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
EXPECTED_OUTPUT = {
|
||||
'outcome': 'completed',
|
||||
'report': "saw ['add', 'query_metrics', 'report_finding', 'search_knowledge']",
|
||||
}
|
||||
|
||||
|
||||
def _register(absurd: AsyncAbsurd, agent: Agent[object, Report]) -> None:
|
||||
@absurd.register_task(name='analyse')
|
||||
async def analyse(params: JsonValue, ctx: AsyncTaskContext) -> JsonValue:
|
||||
result = await agent.run('Investigate the latency spike.')
|
||||
return result.output.model_dump(mode='json')
|
||||
|
||||
|
||||
@pytest.mark.parametrize('case', CASES)
|
||||
class TestRecordedCheckpoints:
|
||||
async def test_resumes_a_recorded_run(
|
||||
self, case: str, absurd: AsyncAbsurd, async_conn: AsyncConnection[TupleRow], queue_name: str
|
||||
) -> None:
|
||||
executions: list[str] = []
|
||||
_register(absurd, _agent(CASES[case], executions))
|
||||
spawned = await absurd.spawn(
|
||||
'analyse', None, max_attempts=2, retry_strategy={'kind': 'fixed', 'base_seconds': 0}
|
||||
)
|
||||
[claimed] = await absurd.claim_tasks(batch_size=1)
|
||||
for name, state in GOLDEN[case].items():
|
||||
await async_conn.execute(
|
||||
'SELECT absurd.set_task_checkpoint_state(%s, %s, %s, %s, %s)',
|
||||
(queue_name, spawned['task_id'], name, json.dumps(state), claimed['run_id']),
|
||||
)
|
||||
# The recording worker died after its last checkpoint; the next attempt resumes the task.
|
||||
await async_conn.execute(
|
||||
'SELECT absurd.fail_run(%s, %s, %s)', (queue_name, claimed['run_id'], '{"type": "crash"}')
|
||||
)
|
||||
await absurd.work_batch(batch_size=1)
|
||||
|
||||
result = await absurd.fetch_task_result(spawned['task_id'])
|
||||
assert result is not None and result.state == 'completed'
|
||||
assert result.result == EXPECTED_OUTPUT
|
||||
# Only the `ModelRetry` call re-runs: it is never checkpointed.
|
||||
assert executions == (['report_finding:retry'] if CASES[case] else [])
|
||||
|
||||
async def test_fresh_run_writes_the_same_checkpoints(self, case: str, absurd: AsyncAbsurd) -> None:
|
||||
_register(absurd, _agent(CASES[case], []))
|
||||
spawned = await absurd.spawn('analyse', None)
|
||||
await absurd.work_batch(batch_size=1)
|
||||
|
||||
stored = await checkpoints(absurd, spawned['task_id'])
|
||||
golden = GOLDEN[case]
|
||||
assert sorted(stored) == sorted(golden)
|
||||
# Tool results are compared verbatim. Model responses and MCP tool listings also carry fields
|
||||
# (timestamps, usage, fields added in later releases) that differ run to run or release to
|
||||
# release, so those are compared on the fields a replay reads.
|
||||
for name, expected in golden.items():
|
||||
if '.get_tools' in name:
|
||||
assert _tool_schemas(stored[name]) == _tool_schemas(expected), name
|
||||
else:
|
||||
assert _project(stored[name], expected) == expected, name
|
||||
|
||||
|
||||
def _project(value: JsonValue, like: JsonValue) -> JsonValue:
|
||||
"""Keep only the parts of `value` that `like` has."""
|
||||
if isinstance(like, dict):
|
||||
assert isinstance(value, dict)
|
||||
return {key: _project(value[key], sub) for key, sub in like.items()}
|
||||
if isinstance(like, list):
|
||||
assert isinstance(value, list) and len(value) == len(like)
|
||||
return [_project(item, sub) for item, sub in zip(value, like)]
|
||||
return value
|
||||
|
||||
|
||||
def _tool_schemas(listing: JsonValue) -> JsonValue:
|
||||
assert isinstance(listing, dict)
|
||||
keys = ('name', 'description', 'kind', 'parameters_json_schema')
|
||||
return {name: _project(tool, {key: None for key in keys}) for name, tool in listing.items()}
|
||||
@@ -0,0 +1,62 @@
|
||||
"""`AbsurdDurability` on a real Absurd worker: a retried attempt is served from checkpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip('absurd_sdk')
|
||||
|
||||
from absurd_sdk import AsyncAbsurd, AsyncTaskContext, JsonValue
|
||||
|
||||
from pydantic_ai import Agent, ModelMessage, ModelResponse
|
||||
from pydantic_ai.messages import TextPart, ToolCallPart, ToolReturnPart
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
from pydantic_ai.toolsets import FunctionToolset
|
||||
from pydantic_ai_harness.absurd import AbsurdDurability
|
||||
|
||||
|
||||
async def test_worker_retry_replays_postgres_checkpoints(absurd: AsyncAbsurd) -> None:
|
||||
"""A retried attempt serves the model and tool calls from the first attempt's checkpoints."""
|
||||
calls = {'attempts': 0, 'model': 0, 'tool': 0}
|
||||
toolset = FunctionToolset[object](id='billing')
|
||||
|
||||
@toolset.tool_plain
|
||||
def charge_card(amount: int) -> str:
|
||||
calls['tool'] += 1
|
||||
return f'charged {amount}'
|
||||
|
||||
def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
|
||||
calls['model'] += 1
|
||||
if any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts):
|
||||
return ModelResponse(parts=[TextPart(content='done')])
|
||||
return ModelResponse(parts=[ToolCallPart(tool_name='charge_card', args={'amount': 42})])
|
||||
|
||||
agent = Agent[object, str](
|
||||
FunctionModel(model, model_name='fn'),
|
||||
name='checkout',
|
||||
toolsets=[toolset],
|
||||
capabilities=[AbsurdDurability()],
|
||||
)
|
||||
|
||||
@absurd.register_task(name='checkout')
|
||||
async def checkout(params: JsonValue, ctx: AsyncTaskContext) -> JsonValue:
|
||||
calls['attempts'] += 1
|
||||
result = await agent.run('charge it')
|
||||
if calls['attempts'] == 1:
|
||||
raise RuntimeError('simulated worker crash after checkpoints were committed')
|
||||
return {'output': result.output}
|
||||
|
||||
spawned = await absurd.spawn(
|
||||
'checkout',
|
||||
None,
|
||||
max_attempts=2,
|
||||
retry_strategy={'kind': 'fixed', 'base_seconds': 0},
|
||||
)
|
||||
await absurd.work_batch(batch_size=1)
|
||||
await absurd.work_batch(batch_size=1)
|
||||
result = await absurd.fetch_task_result(spawned['task_id'])
|
||||
|
||||
assert result is not None
|
||||
assert result.state == 'completed'
|
||||
assert result.result == {'output': 'done'}
|
||||
assert calls == {'attempts': 2, 'model': 2, 'tool': 1}
|
||||
@@ -318,6 +318,10 @@ COMBINE_POLICY: dict[str, Policy] = {
|
||||
'a durability engine is one per agent; `from_agent` rejects a second when the engine looks '
|
||||
'itself up, before any id is consulted'
|
||||
),
|
||||
'AbsurdDurability': Rejected(
|
||||
'a durability engine is one per agent; `from_agent` rejects a second when the engine looks '
|
||||
'itself up, before any id is consulted'
|
||||
),
|
||||
# -- No default `id`, but two never coexist anyway: their tool names collide. --
|
||||
'FileSystem': Collides(
|
||||
'its toolset registers `read_file` and friends under fixed names',
|
||||
|
||||
@@ -133,6 +133,7 @@ _CAPABILITY_PAGE_META = {
|
||||
'advisor.md': ('advisor', 'Advisor'),
|
||||
'ask-user.md': ('ask_user', 'Ask User'),
|
||||
'aws-lambda.md': ('aws_lambda', 'AWS Lambda Durability'),
|
||||
'absurd.md': ('absurd', 'Absurd Durability'),
|
||||
'background-tools.md': ('background_tools', 'Background Tools'),
|
||||
'code-mode.md': ('code_mode', 'Code Mode'),
|
||||
'coder.md': ('coder', 'Coder'),
|
||||
|
||||
@@ -89,6 +89,18 @@ constraints = [
|
||||
excludes = ["browser-use"]
|
||||
build-constraints = [{ name = "hatchling", specifier = "<1.32" }]
|
||||
|
||||
[[package]]
|
||||
name = "absurd-sdk"
|
||||
version = "0.5.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "psycopg", extra = ["binary"] },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/18/6a/2d4199ddc3d846ba13c51aa9c8614b9ee42a5af5144ef456ef11f4f6aa20/absurd_sdk-0.5.0.tar.gz", hash = "sha256:b22fbe8d10f649d0ed427ba16d885a4a452e4240f412591e7f087fe8584a6f00", size = 14973, upload-time = "2026-08-04T13:06:40.682Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/3b/12/0d1c4e23dc00069a4666fc25c165514657fc62eb723a7bee3ce22ef416ae/absurd_sdk-0.5.0-py3-none-any.whl", hash = "sha256:a21465deaf159ba29ab4d15d87f429ae5dcf6f0e3098bce11ae9bc3e267fd03a", size = 15407, upload-time = "2026-08-04T13:06:39.406Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ag-ui-protocol"
|
||||
version = "0.1.19"
|
||||
@@ -6417,6 +6429,9 @@ dependencies = [
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
absurd = [
|
||||
{ name = "absurd-sdk" },
|
||||
]
|
||||
acp = [
|
||||
{ name = "agent-client-protocol" },
|
||||
]
|
||||
@@ -6528,6 +6543,7 @@ youdotcom = [
|
||||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "absurd-sdk", marker = "extra == 'absurd'", specifier = ">=0.5.0,<1" },
|
||||
{ name = "agent-client-protocol", marker = "extra == 'acp'", specifier = ">=0.11,<0.12" },
|
||||
{ name = "aws-durable-execution-sdk-python", marker = "python_full_version >= '3.11' and extra == 'aws-lambda'", specifier = ">=1.7.0,<2" },
|
||||
{ name = "browser-use", marker = "python_full_version >= '3.11' and extra == 'browser-use'", specifier = ">=0.13.6,<0.14" },
|
||||
@@ -6574,7 +6590,7 @@ requires-dist = [
|
||||
{ name = "stackone-defender", extras = ["onnx"], marker = "python_full_version >= '3.11' and extra == 'prompt-injection-defender-ml'", specifier = ">=0.7.3,<0.8" },
|
||||
{ name = "youdotcom", marker = "extra == 'youdotcom'", specifier = ">=3.1.2,<3.2" },
|
||||
]
|
||||
provides-extras = ["acp", "anthropic", "aws-lambda", "browser-use", "cli", "code-mode", "codemode", "coder", "day-ai", "dbos", "dynamic-workflow", "e2b", "exa", "github", "google-workspace", "grain", "linear", "logfire", "logfire-mcp", "modal", "mongodb", "notion", "ordinal", "playwright", "posthog", "prompt-injection-defender", "prompt-injection-defender-ml", "pylon", "researcher", "skills", "slack", "sprites", "stackone", "temporal", "youdotcom"]
|
||||
provides-extras = ["absurd", "acp", "anthropic", "aws-lambda", "browser-use", "cli", "code-mode", "codemode", "coder", "day-ai", "dbos", "dynamic-workflow", "e2b", "exa", "github", "google-workspace", "grain", "linear", "logfire", "logfire-mcp", "modal", "mongodb", "notion", "ordinal", "playwright", "posthog", "prompt-injection-defender", "prompt-injection-defender-ml", "pylon", "researcher", "skills", "slack", "sprites", "stackone", "temporal", "youdotcom"]
|
||||
|
||||
[[package]]
|
||||
name = "pydantic-ai-slim"
|
||||
|
||||
Reference in New Issue
Block a user