Fix typevar variance for agent deps

Author: dmontagu

pydantic_ai_slim/pydantic_ai/_run_context.py

@@ -19,12 +19,15 @@
 AgentDepsT = TypeVar('AgentDepsT', default=None, contravariant=True)
 """Type variable for agent dependencies."""
 
+RunContextAgentDepsT = TypeVar('RunContextAgentDepsT', default=None, covariant=True)
+"""Type variable for the agent dependencies in `RunContext`."""
+
 
 @dataclasses.dataclass(repr=False, kw_only=True)
-class RunContext(Generic[AgentDepsT]):
+class RunContext(Generic[RunContextAgentDepsT]):
     """Information about the current call."""
 
-    deps: AgentDepsT
+    deps: RunContextAgentDepsT
     """Dependencies for the agent."""
     model: Model
     """The model used in this run."""

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_run_context.py

@@ -2,8 +2,15 @@
 
 from typing import Any
 
+from typing_extensions import TypeVar
+
 from pydantic_ai.exceptions import UserError
-from pydantic_ai.tools import AgentDepsT, RunContext
+from pydantic_ai.tools import RunContext
+
+AgentDepsT = TypeVar('AgentDepsT', default=None, covariant=True)
+"""Type variable for the agent dependencies in `RunContext`."""
+
+T = TypeVar('T', default=None)
 
 
 class TemporalRunContext(RunContext[AgentDepsT]):
@@ -46,7 +53,7 @@ def serialize_run_context(cls, ctx: RunContext[Any]) -> dict[str, Any]:
             'run_step': ctx.run_step,
         }
 
-    @classmethod
-    def deserialize_run_context(cls, ctx: dict[str, Any], deps: AgentDepsT) -> TemporalRunContext[AgentDepsT]:
+    @staticmethod
+    def deserialize_run_context(ctx: dict[str, Any], deps: T) -> TemporalRunContext[T]:
         """Deserialize the run context from a `dict[str, Any]`."""
-        return cls(**ctx, deps=deps)
+        return TemporalRunContext(**ctx, deps=deps)

tests/typed_deps.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+
+from pydantic_ai import Agent, RunContext
+
+
+@dataclass
+class DepsA:
+    a: int
+
+
+@dataclass
+class DepsB:
+    b: str
+
+
+@dataclass
+class AgentDeps(DepsA, DepsB):
+    pass
+
+
+agent = Agent(
+    instructions='...',
+    model='...',
+    deps_type=AgentDeps,
+)
+
+
+@agent.tool
+def tool_1(ctx: RunContext[DepsA]) -> int:
+    return ctx.deps.a
+
+
+@agent.tool
+def tool_2(ctx: RunContext[DepsB]) -> str:
+    return ctx.deps.b
+
+
+# Ensure that you can use tools with deps that are supertypes of the agent's deps
+agent.run_sync('...', deps=AgentDeps(a=0, b='test'))