updates

Author: mattbrandman

tests/test_temporal.py

@@ -71,7 +71,6 @@
  TemporalAgent,
  )
  from pydantic_ai.durable_exec.temporal._function_toolset import TemporalFunctionToolset
- from pydantic_ai.durable_exec.temporal._mcp_server import TemporalMCPServer
  from pydantic_ai.durable_exec.temporal._model import TemporalModel
  from pydantic_ai.durable_exec.temporal._run_context import TemporalRunContext
 except ImportError: # pragma: lax no cover
@@ -1231,32 +1230,14 @@ async def test_temporal_agent():
  assert isinstance(complex_temporal_agent.model, TemporalModel)
  assert complex_temporal_agent.model.wrapped == complex_agent.model
 
- toolsets = complex_temporal_agent.toolsets
- assert len(toolsets) == 4
-
- # Wrapped function toolset for the agent's own tools
- assert isinstance(toolsets[0], TemporalFunctionToolset)
- assert toolsets[0].id == '<agent>'
- assert isinstance(toolsets[0].wrapped, FunctionToolset)
- assert toolsets[0].wrapped.tools.keys() == {'get_weather'}
-
- # Wrapped 'country' toolset
- assert isinstance(toolsets[1], TemporalFunctionToolset)
- assert toolsets[1].id == 'country'
- assert toolsets[1].wrapped == complex_agent.toolsets[1]
- assert isinstance(toolsets[1].wrapped, FunctionToolset)
- assert toolsets[1].wrapped.tools.keys() == {'get_country'}
-
- # Wrapped 'mcp' MCP server
- assert isinstance(toolsets[2], TemporalMCPServer)
- assert toolsets[2].id == 'mcp'
- assert toolsets[2].wrapped == complex_agent.toolsets[2]
-
- # Unwrapped 'external' toolset
- assert isinstance(toolsets[3], ExternalToolset)
- assert toolsets[3].id == 'external'
- assert toolsets[3] == complex_agent.toolsets[3]
+ # Toolsets should match the wrapped agent's toolsets
+ assert len(complex_temporal_agent.toolsets) == 4
+ assert complex_temporal_agent.toolsets[0].id == '<agent>'
+ assert complex_temporal_agent.toolsets[1].id == 'country'
+ assert complex_temporal_agent.toolsets[2].id == 'mcp'
+ assert complex_temporal_agent.toolsets[3].id == 'external'
 
+ # Temporal activities verify that toolsets are properly wrapped for workflow execution
  assert [
  ActivityDefinition.must_from_callable(activity).name # pyright: ignore[reportUnknownMemberType]
  for activity in complex_temporal_agent.temporal_activities