from __future__ import annotations from typing import TYPE_CHECKING, Any, overload from pydantic import TypeAdapter from typing_extensions import TypeVar from pydantic_ai._run_context import AnchoredEvidence, CapabilityEventT, CustomEventT from pydantic_ai.durable_exec._toolset import EnqueueGuard, enqueue_not_supported_message from pydantic_ai.exceptions import UserError from pydantic_ai.messages import CapabilityEvent, CustomEvent from pydantic_ai.tools import RunContext from pydantic_ai.usage import RunUsage, UsageLimits if TYPE_CHECKING: from pydantic_ai.agent.abstract import AbstractAgent AgentDepsT = TypeVar('AgentDepsT', default=object, covariant=True) """Type variable for the agent dependencies in `RunContext`.""" # The serialized run context crosses the activity boundary as untyped JSON (`Any`, so # `TemporalRunContext` subclasses can add their own fields), which means a value whose type isn't # JSON-native arrives back in a different shape than it went in: a model or content object as a # plain dict, a set as a list. Every carried field with such a type is listed here with the shape # it arrives in and the adapter that turns it back into the real thing, so activity-side code gets # the objects it would get in a non-durable run: `usage`/`usage_limits` drive the mid-chain # continuation usage check, while the history-derived sets and dispatch-only supplements feed tool # availability. Activity-side mutations are discarded; reveal deltas are applied to workflow state # after the activity result returns. _str_set_ta: TypeAdapter[set[str]] = TypeAdapter(set[str]) _REHYDRATORS: tuple[tuple[str, type[Any], TypeAdapter[Any]], ...] = ( ('usage', dict, TypeAdapter(RunUsage)), ('usage_limits', dict, TypeAdapter(UsageLimits)), ('loaded_capability_ids', list, _str_set_ta), ('discovered_tool_names', list, _str_set_ta), ('available_tool_names', list, _str_set_ta), ('active_capability_ids', list, _str_set_ta), ('_deferred_capability_ids', list, _str_set_ta), ('_anchored_evidence', dict, TypeAdapter(AnchoredEvidence)), ) # Fields that `serialize_run_context` doesn't carry but that are still readable inside an activity, # as `None` unless the framework attaches something: `agent` and `root_capability` are re-attached # from the worker's agent instance, `pending_messages` is replaced by an `EnqueueGuard` so # `ctx.enqueue()` raises an explanation, and `tool_manager` is documented to be `None` inside # activities — `available_tool_names` then reads the snapshot serialized at dispatch time (see # the property override below), or falls back to `discovered_tool_names` without one. # `realtime_session` is a live session object that cannot cross the boundary, and its contract # already makes `None` mean "not available here". `_run_held_toolsets` is the same: an activity may # run in another process, so it enters its own toolset instead of reusing the one the run holds. _NONE_UNLESS_ATTACHED = ( 'agent', 'root_capability', 'pending_messages', 'tool_manager', 'realtime_session', '_durable_operations', '_run_capabilities_by_id', '_run_held_toolsets', ) # Defaulted rather than guarded when a payload doesn't carry it. Unlike the guarded fields, the # dataclass default can't be mistaken for real run state here: empty means "no anchored evidence", # which is exactly what `is_tool_available` reads when the serving response has no provenance. A # custom `serialize_run_context` written before this field existed therefore keeps answering — with # the history-derived window — instead of raising for a field it never knew to carry. _DEFAULTED_UNLESS_CARRIED: tuple[tuple[str, Any], ...] = (('_anchored_evidence', AnchoredEvidence()),) # Payloads written by a worker running an older version, or by a custom `serialize_run_context` that # still spells the old name. An activity can be dispatched by one worker version and replayed by # another, so the old key has to keep arriving at the renamed field — otherwise the field lands in # `__dict__` under a name nothing reads, and `capability_active` raises as "not carried". _RENAMED_FIELDS: tuple[tuple[str, str], ...] = ( ('capability_loaded', 'capability_active'), ('available_capability_ids', 'active_capability_ids'), ) # Reading any other omitted field raises instead of returning the `RunContext` dataclass default, # which would silently pass for real run state (e.g. `instrumentation_version` reading as the # default version rather than the run's, or `prompt` as `None` for a subclass that drops it). _GUARDED_FIELDS = frozenset(RunContext.__dataclass_fields__) - {'deps', *_NONE_UNLESS_ATTACHED} 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`, `conversation_id`, `metadata`, `retries`, `tool_call_id`, `tool_name`, `tool_call_approved`, `tool_call_metadata`, `retry`, `max_retries`, `run_step`, `usage`, `usage_limits`, `partial_output`, `trace_include_content`, `instrumentation_version`, `loaded_capability_ids`, `discovered_tool_names`, the private dispatch-only availability supplements, and `capability_active` attributes will be available. Reading any other attribute raises a `UserError` explaining how to make it available, rather than returning its default value, so a field that didn't cross the boundary can't be mistaken for real run state. `agent` and `root_capability` are re-attached from the worker's agent instance, `pending_messages` holds a guard that makes [`enqueue`][pydantic_ai.tools.RunContext.enqueue] raise inside an activity, and `tool_manager` and `realtime_session` are `None`: they hold live run state that isn't serializable (for `tool_manager`, `available_tool_names` returns the resolved snapshot serialized at activity dispatch time, falling back to `discovered_tool_names` if a custom subclass doesn't carry it; for `realtime_session`, `None` already means "not available here"). The `capabilities` registry is excluded for the same reason — it holds live capability objects (toolsets, hooks, callables) — so `active_capability_ids` likewise returns a snapshot serialized at dispatch time, which is what lets [`is_tool_available`][pydantic_ai.tools.RunContext.is_tool_available] answer for a capability-owned tool inside an activity; reading `capabilities` itself still raises. `model` and `tracer` are excluded as live objects too. `messages` is excluded because the full history would be duplicated into every activity payload, and `prompt` is excluded because a multi-modal prompt can carry large `BinaryContent` that would likewise ride in every activity payload, risking Temporal's 2 MB limit. `model_settings` is excluded because it's only set for model requests, which receive it as their own activity parameter, and `validation_context` because it's an arbitrary user object with no serialization contract. 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 as the `run_context_type` argument to [`TemporalDurability`][pydantic_ai.durable_exec.temporal.TemporalDurability]. A subclass can use this escape hatch to opt in to carrying `prompt` if it knows its prompts are text-only. """ def __init__(self, deps: AgentDepsT, **kwargs: Any): self.__dict__ = {**kwargs, 'deps': deps} for old_name, new_name in _RENAMED_FIELDS: # Keyed on presence, not truthiness: `capability_active` is `None` for every activity # dispatched outside capability dispatch — the common case — and a value-based guard # would drop the key there, leaving the renamed field absent and the guard below # reporting it as one that never crossed the boundary. if old_name in self.__dict__: self.__dict__.setdefault(new_name, self.__dict__.pop(old_name)) for name in _NONE_UNLESS_ATTACHED: self.__dict__.setdefault(name, None) for name, default in _DEFAULTED_UNLESS_CARRIED: self.__dict__.setdefault(name, default) for name, wire_type, adapter in _REHYDRATORS: if isinstance(value := self.__dict__.get(name), wire_type): self.__dict__[name] = adapter.validate_python(value) setattr( self, '__dataclass_fields__', {name: field for name, field in RunContext.__dataclass_fields__.items() if name in self.__dict__}, ) def __getattribute__(self, name: str) -> Any: if name in _GUARDED_FIELDS and name not in object.__getattribute__(self, '__dataclass_fields__'): raise UserError( f'{name!r} is not available on {self.__class__.__name__!r} inside a Temporal activity. ' 'To make the 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 as the `run_context_type` argument to `TemporalDurability`.' ) return super().__getattribute__(name) def _expose_field(self, name: str) -> None: """Mark a framework-attached field as readable after deserialization.""" instance_fields = object.__getattribute__(self, '__dataclass_fields__') instance_fields[name] = RunContext.__dataclass_fields__[name] @property def available_tool_names(self) -> set[str]: """The availability snapshot serialized at activity dispatch time. Live tool state doesn't cross the activity boundary, but availability was already resolved when the activity was dispatched — so the name form of [`is_tool_available`][pydantic_ai.tools.RunContext.is_tool_available] answers correctly for always-visible tools too, instead of degrading to the `discovered_tool_names` fallback. Custom subclasses whose `serialize_run_context` doesn't carry the snapshot keep the base fallback behavior. """ if (snapshot := self.__dict__.get('available_tool_names')) is not None: return snapshot return super().available_tool_names @property def active_capability_ids(self) -> set[str]: """The set of active capability ids serialized at activity dispatch time. The `capabilities` registry itself can't cross the boundary, but the ids it resolves to can, so [`is_tool_available`][pydantic_ai.tools.RunContext.is_tool_available] still answers for a capability-owned tool instead of raising. Custom subclasses whose `serialize_run_context` doesn't carry the snapshot fall back to the base property, which reads the registry and raises inside an activity. """ if (snapshot := self.__dict__.get('active_capability_ids')) is not None: return snapshot return super().active_capability_ids @property def _deferred_capability_ids(self) -> set[str]: """The set of on-demand capability ids serialized at activity dispatch time. `is_tool_available` needs the *configured* shape of a capability, not just what history says was loaded, and reads it from the registry — which cannot cross the boundary. Carrying the ids keeps a loaded capability's own tools answering as available inside an activity instead of falling back to a reveal marker that, for these tools, nothing can regenerate. Custom subclasses whose `serialize_run_context` omits the snapshot fall back to the base property, which reads the registry and raises inside an activity. """ if (snapshot := self.__dict__.get('_deferred_capability_ids')) is not None: return snapshot return super()._deferred_capability_ids @overload async def emit(self, event: CustomEventT, /) -> CustomEventT: ... @overload async def emit(self, event: CapabilityEventT, /) -> CapabilityEventT: ... async def emit(self, event: CustomEvent | CapabilityEvent, /) -> CustomEvent | CapabilityEvent: """Reject `emit` from inside a Temporal activity. Tools and event stream handlers run inside activities where the run's event stream isn't reachable, so events emitted there can't currently flow back into the stream; raising beats silently dropping them. This covers a capability's own tools emitting a `CapabilityEvent`, which run in activities like any other tool. Events emitted workflow-side (e.g. from capability hooks) work as usual. Lifting this needs a transport from the activity back to the workflow and a decision on what a retried attempt's events mean; tracked in https://github.com/pydantic/pydantic-ai/issues/7971. """ raise UserError( 'Emitting events from a tool or event stream handler is not supported under Temporal yet, as ' 'they run inside activities that cannot reach the run event stream. This includes a capability ' 'emitting a `CapabilityEvent` from one of its own tools. Emit events from capability hooks, ' 'which run in the workflow, instead. Tracked in ' 'https://github.com/pydantic/pydantic-ai/issues/7971.' ) @classmethod def serialize_run_context(cls, ctx: RunContext[Any]) -> dict[str, Any]: """Serialize the run context to a `dict[str, Any]`.""" return { 'run_id': ctx.run_id, 'conversation_id': ctx.conversation_id, 'metadata': ctx.metadata, 'retries': ctx.retries, 'tool_call_id': ctx.tool_call_id, 'tool_name': ctx.tool_name, 'tool_call_approved': ctx.tool_call_approved, 'tool_call_metadata': ctx.tool_call_metadata, 'retry': ctx.retry, 'max_retries': ctx.max_retries, 'run_step': ctx.run_step, 'partial_output': ctx.partial_output, 'trace_include_content': ctx.trace_include_content, 'instrumentation_version': ctx.instrumentation_version, 'usage': ctx.usage, 'usage_limits': ctx.usage_limits, 'loaded_capability_ids': ctx.loaded_capability_ids, 'discovered_tool_names': ctx.discovered_tool_names, # The dispatch-time widening of the two sets above, which `is_tool_available` reads for # a call the model has already made. Carried so a tool asking whether it may run gets # the same answer inside an activity as it would in-process. '_anchored_evidence': ctx._anchored_evidence, # A resolved snapshot: at dispatch time live tool state exists, so this carries the # always-visible tools that the in-activity `discovered_tool_names` fallback misses. 'available_tool_names': ctx.available_tool_names, # Likewise a snapshot rather than the registry: these ids are plain strings, while the # capability objects they key are not serializable. `is_tool_available` consults this # for any capability-owned tool, so without it the definition form — the form the docs # send toolset authors to — raises inside an activity instead of answering. 'active_capability_ids': ctx.active_capability_ids, # The configured on-demand set, which `is_tool_available` consults to tell a loaded # deferred capability (whose load is itself the reveal for its tools) from one that has # since been reconfigured always-on. Derived from the registry, so it must travel too. '_deferred_capability_ids': ctx._deferred_capability_ids, 'capability_active': ctx.capability_active, } @classmethod def deserialize_run_context(cls, ctx: dict[str, Any], deps: Any) -> TemporalRunContext[Any]: """Deserialize the run context from a `dict[str, Any]`.""" return cls(**ctx, deps=deps) def deserialize_run_context( run_context_type: type[TemporalRunContext[Any]], serialized: dict[str, Any], *, deps: Any, agent: AbstractAgent[Any, Any] | None, ) -> RunContext[Any]: """Deserialize a run context and attach the agent instance. This is a helper used internally by the Temporal wrappers. It calls the (potentially user-overridden) `TemporalRunContext.deserialize_run_context` and then sets `agent` and `root_capability` on the result so custom subclasses don't need to know about either parameter. Setting `root_capability` lets the durability capability fire the capability chain against the live model stream inside the activity, which is required for capabilities like `ProcessEventStream` to see real (non-replayed) events. """ ctx = run_context_type.deserialize_run_context(serialized, deps=deps) if agent is not None: ctx.__dict__['agent'] = agent ctx.__dict__['root_capability'] = agent.root_capability # `pending_messages` isn't serialized across the activity boundary, and any code running inside # an activity (a tool, a `process_tool_call` hook, an `event_stream_handler`) is in a durable # unit whose result is replayed without re-running it, so an enqueue would be dropped. Install # the same guard the in-process engines use so `ctx.enqueue()` raises the shared explanation. ctx.__dict__['pending_messages'] = EnqueueGuard(enqueue_not_supported_message('activity', 'workflow')) return ctx