Skip to content

Add Pydantic AI Durable Capability Example#552

Draft
nvasiu wants to merge 1 commit into
mainfrom
pydantic-ai-durable-capability
Draft

Add Pydantic AI Durable Capability Example#552
nvasiu wants to merge 1 commit into
mainfrom
pydantic-ai-durable-capability

Conversation

@nvasiu

@nvasiu nvasiu commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Pydantic AI Durable Agent

A Pydantic AI agent as a Lambda durable function. Each model request and tool
call is checkpointed as a durable step.

  • lambda_durability.py — the LambdaDurability capability.
  • pydantic_ai_agent.py — the durable handler and agent.
  • pydantic-ai-agent.zip — prebuilt deployment package (built for the Lambda
    python3.13 / x86_64 runtime).

Requires an execution role with durable execution and bedrock:InvokeModel
permissions, and Bedrock access for PYDANTIC_AI_MODEL
(default us.amazon.nova-micro-v1:0).

Deploy

aws lambda create-function \
  --function-name pydantic-ai-agent \
  --runtime python3.13 \
  --handler pydantic_ai_agent.handler \
  --role <ROLE_ARN> \
  --zip-file fileb://pydantic-ai-agent.zip \
  --timeout 300 --memory-size 1024 \
  --environment '{"Variables":{"PYDANTIC_AI_MODEL":"us.amazon.nova-micro-v1:0"}}' \
  --durable-config '{"ExecutionTimeout":3600,"RetentionPeriodInDays":7}' \
  --region us-east-1

aws lambda wait function-active --function-name pydantic-ai-agent --region us-east-1
aws lambda publish-version --function-name pydantic-ai-agent --region us-east-1

Invoke

aws lambda invoke \
  --function-name pydantic-ai-agent:1 \
  --invocation-type Event \
  --payload '{"prompt":"What is the weather in Vancouver?"}' \
  --region us-east-1 /dev/stdout

aws lambda get-durable-execution \
  --durable-execution-arn "<DurableExecutionArn>" --region us-east-1

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

# is not recorded as failed.
async def call_tool():
try:
return {"retry": False, "value": await handler(args)}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Return value from handler may not be serializable.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To keep this demo simple, tool results are just assumed to be JSON serializable.

In the future we'll add a per-tool config to allow users to specify custom SerDes for each tool.

@nvasiu
nvasiu force-pushed the pydantic-ai-durable-capability branch from 3c5d118 to e8d8d74 Compare July 20, 2026 18:09

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what is this zip for?

rather than check in a zip to github, where it will increase the size of a git clone forever more, can we make this a script instead? A script to make the zip?


from __future__ import annotations

import asyncio

@yaythomas yaythomas Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ruff F401 — unused import asyncio at test/pydantic_ai_agent/test_pydantic_ai_agent.py:5. Auto-fixable.

)
def test_pydantic_ai_agent_runs_in_cloud(durable_runner):
if durable_runner.mode != "cloud":
pytest.skip("Pydantic AI Bedrock example only runs in cloud mode")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the durable_runner fixture calls pytest.fail("handler is required for local mode tests"), so strictly speaking this is dead code, I think?

assert deserialize_operation_payload(result.result)
operation_names = {operation.name for operation in result.operations}
assert any(name and name.startswith("model.request") for name in operation_names)
assert any(name and name.startswith("tool.") for name in operation_names)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

might be handy to have a test where the agent actually suspends mid-run and then succeeds on a later invocation by replaying its checkpointed model and tool steps.

@DouweM

DouweM commented Jul 22, 2026

Copy link
Copy Markdown

Posted by Claude Code on behalf of Douwe.

Thanks for building this, and for asking for the detailed feedback. Consolidating it here rather than as line comments, since the main recommendation changes the design wholesale. I ran your lambda_durability.py verbatim against real Pydantic AI agents, using a local fake DurableContext that models Lambda's step semantics (a raising step body is recorded as a failed operation; a completed step replays its checkpoint without re-running the body), so the two bugs below are reproduced, not theoretical.

Headline: the durability API changed under you

Since this PoC, pydantic/pydantic-ai#4977 merged (shipped in pydantic-ai 2.15.0). It adds a first-class base for durability integrations: pydantic_ai.durable_exec._base.BaseDurabilityCapability, plus shared DurableModel, DurableFunctionToolset / DurableMCPToolset / DurableDynamicToolset, and helpers (wrap_tool_call_result, resolve_tool_durable_config, ...). The three in-tree engines (Temporal, DBOS, Prefect) are built on it, and it is what a new engine should subclass. Building on it gets most of the gaps below right for free.

We went ahead and built the complete version on that base, so you have a concrete reference rather than a checklist: pydantic/pydantic-ai-harness#417 (AWSLambdaDurability). It keeps your sync-to-async bridge design (that part is good, see the bottom) and fills in the rest. Two ways forward: adopt or adapt that, or apply the specific fixes below to this PoC. Where the integration ultimately lives (your SDK repo vs the harness) is a separate conversation Douwe and Jas are already having.

Two reproduced bugs

1. Control-flow exceptions are recorded as step failures. wrap_tool_execute only round-trips ToolRetryError. ApprovalRequired, CallDeferred, and ToolFailed escape the step body, so the step is recorded failed, Lambda retries it per StepConfig, and then the execution fails. ApprovalRequired / CallDeferred are the deferred-tools / human-in-the-loop path, so HITL approval never reaches the caller.

Tool raises Without durability With this capability
ApprovalRequired clean DeferredToolRequests pause step FAILED
CallDeferred clean DeferredToolRequests step FAILED
ToolFailed model sees the failure, run continues step FAILED
ToolRetryError retry prompt handled

Pydantic AI has exactly this as wrap_tool_call_result / unwrap_tool_call_result: they convert the whole family to serialized values that cross the checkpoint and re-raise on the other side.

2. Tool results are not serialization-safe. The model path is correct (TypeAdapter(ModelResponse)), but wrap_tool_execute returns {"retry": False, "value": await handler(args)} with the raw value. Under JSON checkpointing, a tool returning a ToolReturn or BinaryContent (multimodal) raises TypeError. Route tool results through the same TypeAdapter treatment the model path already uses. Worth noting: the example test does not JSON round-trip stored values, so the current suite cannot catch this. Our fake does, which is how it surfaced.

The seam: wrap_tool_execute is not enough for durability

wrap_tool_execute is a real public hook, but it fires per tool call, so it never sees get_tools / get_instructions. For an MCP or dynamic toolset that means the tool-listing network I/O runs outside any step and re-executes on every replay, which Lambda's structural replay validation forbids. None of the in-tree engines use wrap_tool_execute; they wrap at the toolset level (get_wrapper_toolset returning the shared Durable*Toolsets), which checkpoints listing and calls together. This PoC has one plain function tool, so it does not surface yet, but it is the reason the toolset-level seam exists.

Other gaps to be production-complete (all provided by the base)

  • Capability ordering. Set get_ordering() to innermost, so any other capability's contribution to a model request is already applied inside the durable step.
  • Model continuation segments. You checkpoint the whole model request in one step; for models that suspend and continue mid-turn (Anthropic pause_turn, OpenAI background mode) the entire continuation loop becomes one step, which is a real risk against the 15-minute invocation cap. DurableModel + segment executors checkpoint each segment separately.
  • Streaming. Neither captured nor rejected: on replay a streaming consumer silently receives nothing. Lambda has no channel out of a running execution, so the right shape is capture-and-replay events inside the step (the base's capture_event_stream + StreamedActivityResult) and a clear error if someone tries to stream out to a caller.
  • Sequential tools is hardcoded and ungated. parallel_execution_mode('sequential') is applied on every run, including ordinary non-durable ones, which changes their tool-call ordering. Gate it on being inside a durable execution.
  • Step identity. ToolCallPart.tool_call_id is already in wrap_tool_execute and is a stable per-turn id. It is the natural input if the SDK ever offers a caller-assigned step-id variant, which is what would let you drop sequential execution and restore tool parallelism.

Genuinely good, keep it

The two-thread bridge (async agent loop on a background thread, durable steps serviced synchronously on the handler thread), propagating contextvars in both directions, and reusing the background event loop across warm invocations so loop-bound resources stay valid, is the hard part and you got it right. #417 keeps that exact design and hardened a few edges we found by stress-testing it: detecting a genuinely nested step vs. concurrent sibling steps (two MCP servers list their tools concurrently, so a naive "one step at a time" guard rejects a normal configuration), resolving a cancelled step operation instead of hanging the handler thread, and letting the SDK's SuspendExecution propagate out of the handler rather than routing it into the agent loop.

Happy to walk through #417 or the base API on a call. The two files to read first are pydantic_ai_harness/aws_lambda/_capability.py and _bridge.py.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants