feat: add a global timeout

Author: DEENUU1

pydantic_ai_slim/pydantic_ai/agent/__init__.py

@@ -147,6 +147,7 @@ class Agent(AbstractAgent[AgentDepsT, OutputDataT]):
     _prepare_output_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
     _max_tool_retries: int = dataclasses.field(repr=False)
+    _tool_timeout: float | None = dataclasses.field(repr=False)
     _validation_context: Any | Callable[[RunContext[AgentDepsT]], Any] = dataclasses.field(repr=False)
 
     _event_stream_handler: EventStreamHandler[AgentDepsT] | None = dataclasses.field(repr=False)
@@ -179,6 +180,7 @@ def __init__(
         instrument: InstrumentationSettings | bool | None = None,
         history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        tool_timeout: float | None = None,
     ) -> None: ...
 
     @overload
@@ -206,6 +208,7 @@ def __init__(
         instrument: InstrumentationSettings | bool | None = None,
         history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        tool_timeout: float | None = None,
     ) -> None: ...
 
     def __init__(
@@ -231,6 +234,7 @@ def __init__(
         instrument: InstrumentationSettings | bool | None = None,
         history_processors: Sequence[HistoryProcessor[AgentDepsT]] | None = None,
         event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+        tool_timeout: float | None = None,
         **_deprecated_kwargs: Any,
     ):
         """Create an agent.
@@ -285,6 +289,9 @@ def __init__(
                 Each processor takes a list of messages and returns a modified list of messages.
                 Processors can be sync or async and are applied in sequence.
             event_stream_handler: Optional handler for events from the model's streaming response and the agent's execution of tools.
+            tool_timeout: Default timeout in seconds for tool execution. If a tool takes longer than this,
+                a retry prompt is returned to the model. Individual tools can override this with their own timeout.
+                Defaults to None (no timeout).
         """
         if model is None or defer_model_check:
             self._model = model
@@ -318,6 +325,7 @@ def __init__(
 
         self._max_result_retries = output_retries if output_retries is not None else retries
         self._max_tool_retries = retries
+        self._tool_timeout = tool_timeout
 
         self._validation_context = validation_context
 
@@ -331,7 +339,10 @@ def __init__(
             self._output_toolset.max_retries = self._max_result_retries
 
         self._function_toolset = _AgentFunctionToolset(
-            tools, max_retries=self._max_tool_retries, output_schema=self._output_schema
+            tools,
+            max_retries=self._max_tool_retries,
+            default_timeout=self._tool_timeout,
+            output_schema=self._output_schema,
         )
         self._dynamic_toolsets = [
             DynamicToolset[AgentDepsT](toolset_func=toolset)
@@ -1101,7 +1112,7 @@ async def spam(ctx: RunContext[str], y: float) -> float:
                 See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
             metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization.
             timeout: Timeout in seconds for tool execution. If the tool takes longer, a retry prompt is returned to the model.
-                Defaults to None (no timeout).
+                Overrides the agent-level `tool_timeout` if set. Defaults to None (no timeout).
         """
 
         def tool_decorator(
@@ -1217,7 +1228,7 @@ async def spam(ctx: RunContext[str]) -> float:
                 See the [tools documentation](../deferred-tools.md#human-in-the-loop-tool-approval) for more info.
             metadata: Optional metadata for the tool. This is not sent to the model but can be used for filtering and tool behavior customization.
             timeout: Timeout in seconds for tool execution. If the tool takes longer, a retry prompt is returned to the model.
-                Defaults to None (no timeout).
+                Overrides the agent-level `tool_timeout` if set. Defaults to None (no timeout).
         """
 
         def tool_decorator(func_: ToolFuncPlain[ToolParams]) -> ToolFuncPlain[ToolParams]:
@@ -1414,7 +1425,10 @@ def toolsets(self) -> Sequence[AbstractToolset[AgentDepsT]]:
 
         if some_tools := self._override_tools.get():
             function_toolset = _AgentFunctionToolset(
-                some_tools.value, max_retries=self._max_tool_retries, output_schema=self._output_schema
+                some_tools.value,
+                max_retries=self._max_tool_retries,
+                default_timeout=self._tool_timeout,
+                output_schema=self._output_schema,
             )
         else:
             function_toolset = self._function_toolset
@@ -1521,11 +1535,12 @@ def __init__(
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [],
         *,
         max_retries: int = 1,
+        default_timeout: float | None = None,
         id: str | None = None,
         output_schema: _output.OutputSchema[Any],
     ):
         self.output_schema = output_schema
-        super().__init__(tools, max_retries=max_retries, id=id)
+        super().__init__(tools, max_retries=max_retries, default_timeout=default_timeout, id=id)
 
     @property
     def id(self) -> str:

pydantic_ai_slim/pydantic_ai/toolsets/function.py

@@ -37,6 +37,7 @@ class FunctionToolset(AbstractToolset[AgentDepsT]):
 
     tools: dict[str, Tool[Any]]
     max_retries: int
+    default_timeout: float | None
     _id: str | None
     docstring_format: DocstringFormat
     require_parameter_descriptions: bool
@@ -47,6 +48,7 @@ def __init__(
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [],
         *,
         max_retries: int = 1,
+        default_timeout: float | None = None,
         docstring_format: DocstringFormat = 'auto',
         require_parameter_descriptions: bool = False,
         schema_generator: type[GenerateJsonSchema] = GenerateToolJsonSchema,
@@ -62,6 +64,9 @@ def __init__(
             tools: The tools to add to the toolset.
             max_retries: The maximum number of retries for each tool during a run.
                 Applies to all tools, unless overridden when adding a tool.
+            default_timeout: Default timeout in seconds for tool execution. If a tool takes longer than this,
+                a retry prompt is returned to the model. Individual tools can override this with their own timeout.
+                Defaults to None (no timeout).
             docstring_format: Format of tool docstring, see [`DocstringFormat`][pydantic_ai.tools.DocstringFormat].
                 Defaults to `'auto'`, such that the format is inferred from the structure of the docstring.
                 Applies to all tools, unless overridden when adding a tool.
@@ -82,6 +87,7 @@ def __init__(
                 in which case the ID will be used to identify the toolset's activities within the workflow.
         """
         self.max_retries = max_retries
+        self.default_timeout = default_timeout
         self._id = id
         self.docstring_format = docstring_format
         self.require_parameter_descriptions = require_parameter_descriptions
@@ -360,7 +366,8 @@ async def call_tool(
     ) -> Any:
         assert isinstance(tool, FunctionToolsetTool)
 
-        timeout = tool.timeout
+        # Per-tool timeout takes precedence over default timeout
+        timeout = tool.timeout if tool.timeout is not None else self.default_timeout
         if timeout is not None:
             try:
                 with anyio.fail_after(timeout):

tests/test_tools.py

@@ -2677,3 +2677,83 @@ async def always_slow_tool() -> str:
 
     with pytest.raises(UnexpectedModelBehavior, match='exceeded max retries'):
         await agent.run('call always_slow_tool')
+
+
+@pytest.mark.anyio
+async def test_agent_level_tool_timeout():
+    """Test that agent-level tool_timeout applies to all tools."""
+    import asyncio
+
+    call_count = 0
+
+    async def model_logic(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            return ModelResponse(parts=[ToolCallPart(tool_name='slow_tool', args={}, tool_call_id='call-1')])
+        return ModelResponse(parts=[TextPart(content='done')])
+
+    # Set global tool_timeout on Agent
+    agent = Agent(FunctionModel(model_logic), tool_timeout=0.1)
+
+    @agent.tool_plain
+    async def slow_tool() -> str:
+        await asyncio.sleep(1.0)  # 1 second, but agent timeout is 0.1s
+        return 'done'  # pragma: no cover
+
+    result = await agent.run('call slow_tool')
+
+    # Check that retry prompt was sent
+    retry_parts = [
+        part
+        for msg in result.all_messages()
+        if isinstance(msg, ModelRequest)
+        for part in msg.parts
+        if isinstance(part, RetryPromptPart) and 'Timed out' in str(part.content)
+    ]
+    assert len(retry_parts) == 1
+    assert 'Timed out after 0.1 seconds' in retry_parts[0].content
+
+
+@pytest.mark.anyio
+async def test_per_tool_timeout_overrides_agent_timeout():
+    """Test that per-tool timeout overrides agent-level timeout."""
+    import asyncio
+
+    call_count = 0
+
+    async def model_logic(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            return ModelResponse(parts=[ToolCallPart(tool_name='fast_timeout_tool', args={}, tool_call_id='call-1')])
+        return ModelResponse(parts=[TextPart(content='done')])
+
+    # Agent has generous 10s timeout, but per-tool timeout is only 0.1s
+    agent = Agent(FunctionModel(model_logic), tool_timeout=10.0)
+
+    @agent.tool_plain(timeout=0.1)  # Per-tool timeout overrides agent timeout
+    async def fast_timeout_tool() -> str:
+        await asyncio.sleep(1.0)  # 1 second, per-tool timeout is 0.1s
+        return 'done'  # pragma: no cover
+
+    result = await agent.run('call fast_timeout_tool')
+
+    # Should timeout because per-tool timeout (0.1s) is applied, not agent timeout (10s)
+    retry_parts = [
+        part
+        for msg in result.all_messages()
+        if isinstance(msg, ModelRequest)
+        for part in msg.parts
+        if isinstance(part, RetryPromptPart) and 'Timed out' in str(part.content)
+    ]
+    assert len(retry_parts) == 1
+    assert 'Timed out after 0.1 seconds' in retry_parts[0].content
+
+
+def test_agent_tool_timeout_passed_to_toolset():
+    """Test that agent-level tool_timeout is passed to FunctionToolset as default_timeout."""
+    agent = Agent(TestModel(), tool_timeout=30.0)
+
+    # The agent's tool_timeout should be passed to the toolset as default_timeout
+    assert agent._function_toolset.default_timeout == 30.0