diff --git a/docs/durable_execution/temporal.md b/docs/durable_execution/temporal.md index 01872b9186..f52626caff 100644 --- a/docs/durable_execution/temporal.md +++ b/docs/durable_execution/temporal.md @@ -172,7 +172,7 @@ As workflows and activities run in separate processes, any values passed between To account for these limitations, tool functions and the [event stream handler](#streaming) running inside activities receive a limited version of the agent's [`RunContext`][pydantic_ai.tools.RunContext], and it's your responsibility to make sure that the [dependencies](../dependencies.md) object provided to [`TemporalAgent.run()`][pydantic_ai.durable_exec.temporal.TemporalAgent.run] can be serialized using Pydantic. -Specifically, only the `deps`, `run_id`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry`, `max_retries`, `run_step` and `partial_output` fields are available by default, and trying to access `model`, `usage`, `prompt`, `messages`, or `tracer` will raise an error. +Specifically, only the `deps`, `run_id`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry`, `max_retries`, `run_step`, `usage`, and `partial_output` fields are available by default, and trying to access `model`, `prompt`, `messages`, or `tracer` will raise an error. If you need one or more of these attributes to be available inside activities, you can create a [`TemporalRunContext`][pydantic_ai.durable_exec.temporal.TemporalRunContext] subclass with custom `serialize_run_context` and `deserialize_run_context` class methods and pass it to [`TemporalAgent`][pydantic_ai.durable_exec.temporal.TemporalAgent] as `run_context_type`. ### Streaming diff --git a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_run_context.py b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_run_context.py index aab04ea093..645b20cafb 100644 --- a/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_run_context.py +++ b/pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_run_context.py @@ -14,7 +14,7 @@ class TemporalRunContext(RunContext[AgentDepsT]): """The [`RunContext`][pydantic_ai.tools.RunContext] subclass to use to serialize and deserialize the run context for use inside a Temporal activity. - By default, only the `deps`, `run_id`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry`, `max_retries`, `run_step` and `partial_output` attributes will be available. + By default, only the `deps`, `run_id`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `retry`, `max_retries`, `run_step`, `usage`, and `partial_output` attributes will be available. To make another attribute available, create a `TemporalRunContext` subclass with a custom `serialize_run_context` class method that returns a dictionary that includes the attribute and pass it to [`TemporalAgent`][pydantic_ai.durable_exec.temporal.TemporalAgent]. """ @@ -51,6 +51,7 @@ def serialize_run_context(cls, ctx: RunContext[Any]) -> dict[str, Any]: 'max_retries': ctx.max_retries, 'run_step': ctx.run_step, 'partial_output': ctx.partial_output, + 'usage': ctx.usage, } @classmethod diff --git a/tests/test_temporal.py b/tests/test_temporal.py index 9a315a82b9..f072561997 100644 --- a/tests/test_temporal.py +++ b/tests/test_temporal.py @@ -2193,6 +2193,27 @@ def test_temporal_run_context_preserves_run_id(): assert reconstructed.run_id == 'run-123' +def test_temporal_run_context_serializes_usage(): + ctx = RunContext( + deps=None, + model=TestModel(), + usage=RunUsage( + requests=2, + tool_calls=1, + input_tokens=123, + output_tokens=456, + details={'foo': 1}, + ), + run_id='run-123', + ) + + serialized = TemporalRunContext.serialize_run_context(ctx) + assert serialized['usage'] == ctx.usage + + reconstructed = TemporalRunContext.deserialize_run_context(serialized, deps=None) + assert reconstructed.usage == ctx.usage + + fastmcp_agent = Agent( model, name='fastmcp_agent',