Address feedback

Author: DouweM

docs/mcp/client.md

@@ -137,7 +137,9 @@ _(This example is complete, it can be run "as is" with Python 3.10+ — you'll n
 The other transport offered by MCP is the [stdio transport](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio) where the server is run as a subprocess and communicates with the client over `stdin` and `stdout`. In this case, you'd use the [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio] class.
 
 !!! note
-    When using [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio] servers, the [`async with agent`][pydantic_ai.Agent.__aenter__] context manager is responsible for starting and stopping the server.
+    When using [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio] servers as `toolsets` on an [`Agent`][pydantic_ai.Agent], you can use the [`async with agent`][pydantic_ai.Agent.__aenter__] context manager to start and stop the server around the context where it'll be used. You can also use [`async with server`][pydantic_ai.mcp.MCPServerStdio.__aenter__] to manage the starting and stopping of a specific server, for example if you'd like to use it with multiple agents.
+
+    If you don't explicitly start the server using one of these context managers, it will automatically be started when it's needed (e.g. to list the available tools or call a specific tool), but it's more efficient to do so around the entire context where you expect it to be used.
 
 ```python {title="mcp_stdio_client.py" py="3.10"}
 from pydantic_ai import Agent

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -639,14 +639,16 @@ async def process_function_tools(  # noqa: C901
     # Then, we handle function tool calls
     calls_to_run: list[_messages.ToolCallPart] = []
     if final_result and ctx.deps.end_strategy == 'early':
-        for call in tool_calls_by_kind['function']:
-            output_parts.append(
+        output_parts.extend(
+            [
                 _messages.ToolReturnPart(
                     tool_name=call.tool_name,
                     content='Tool not executed - a final result was already processed.',
                     tool_call_id=call.tool_call_id,
                 )
-            )
+                for call in tool_calls_by_kind['function']
+            ]
+        )
     else:
         calls_to_run.extend(tool_calls_by_kind['function'])
 
@@ -776,8 +778,8 @@ async def _call_function_tool(
         def process_content(content: Any) -> Any:
             if isinstance(content, _messages.ToolReturn):
                 raise exceptions.UserError(
-                    f"{tool_call.tool_name}'s return contains invalid nested ToolReturn objects. "
-                    f'ToolReturn should be used directly.'
+                    f'The return value of tool {tool_call.tool_name!r} contains invalid nested `ToolReturn` objects. '
+                    f'`ToolReturn` should be used directly.'
                 )
             elif isinstance(content, _messages.MultiModalContentTypes):
                 if isinstance(content, _messages.BinaryContent):
@@ -792,8 +794,8 @@ def process_content(content: Any) -> Any:
                     )
                 )
                 return f'See file {identifier}'
-            else:
-                return content
+
+            return content
 
         if isinstance(tool_result, _messages.ToolReturn):
             if (
@@ -805,7 +807,7 @@ def process_content(content: Any) -> Any:
                 )
             ):
                 raise exceptions.UserError(
-                    f"{tool_call.tool_name}'s `return_value` contains invalid nested MultiModalContentTypes objects. "
+                    f'The `return_value` of tool {tool_call.tool_name!r} contains invalid nested `MultiModalContentTypes` objects. '
                     f'Please use `content` instead.'
                 )
 

pydantic_ai_slim/pydantic_ai/_output.py

@@ -294,7 +294,7 @@ def with_default_mode(self, mode: StructuredOutputMode) -> OutputSchema[OutputDa
 @dataclass(init=False)
 class OutputSchemaWithoutMode(BaseOutputSchema[OutputDataT]):
     processor: ObjectOutputProcessor[OutputDataT] | UnionOutputProcessor[OutputDataT]
-    _toolset: OutputToolset[Any] | None = None
+    _toolset: OutputToolset[Any] | None
 
     def __init__(
         self,
@@ -477,7 +477,7 @@ async def process(
 
 @dataclass(init=False)
 class ToolOutputSchema(OutputSchema[OutputDataT]):
-    _toolset: OutputToolset[Any] | None = None
+    _toolset: OutputToolset[Any] | None
 
     def __init__(self, toolset: OutputToolset[Any] | None, allows_deferred_tool_calls: bool):
         super().__init__(allows_deferred_tool_calls)
@@ -834,8 +834,8 @@ class OutputToolset(CallableToolset[AgentDepsT]):
 
     _tool_defs: list[ToolDefinition]
     processors: dict[str, ObjectOutputProcessor[Any]]
-    max_retries: int = field(default=1)
-    output_validators: list[OutputValidator[AgentDepsT, Any]] = field(default_factory=list)
+    max_retries: int
+    output_validators: list[OutputValidator[AgentDepsT, Any]]
 
     @classmethod
     def build(
@@ -910,12 +910,12 @@ def __init__(
         tool_defs: list[ToolDefinition],
         processors: dict[str, ObjectOutputProcessor[Any]],
         max_retries: int = 1,
-        output_validators: list[OutputValidator[AgentDepsT, Any]] = [],
+        output_validators: list[OutputValidator[AgentDepsT, Any]] | None = None,
     ):
         self.processors = processors
         self._tool_defs = tool_defs
         self.max_retries = max_retries
-        self.output_validators = output_validators
+        self.output_validators = output_validators or []
 
     async def prepare_for_run(self, ctx: RunContext[AgentDepsT]) -> RunToolset[AgentDepsT]:
         return RunToolset(self, ctx)

pydantic_ai_slim/pydantic_ai/agent.py

@@ -4,6 +4,7 @@
 import inspect
 import json
 import warnings
+from asyncio import Lock
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
 from contextvars import ContextVar
@@ -166,6 +167,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     _prepare_output_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
 
+    _enter_lock: Lock = dataclasses.field(repr=False)
     _entered_count: int = dataclasses.field(repr=False)
     _exit_stack: AsyncExitStack | None = dataclasses.field(repr=False)
 
@@ -433,8 +435,9 @@ def __init__(
         self._override_deps: ContextVar[_utils.Option[AgentDepsT]] = ContextVar('_override_deps', default=None)
         self._override_model: ContextVar[_utils.Option[models.Model]] = ContextVar('_override_model', default=None)
 
-        self._exit_stack = None
+        self._enter_lock = Lock()
         self._entered_count = 0
+        self._exit_stack = None
 
     @staticmethod
     def instrument_all(instrument: InstrumentationSettings | bool = True) -> None:
@@ -1795,18 +1798,19 @@ def is_end_node(
         return isinstance(node, End)
 
     async def __aenter__(self) -> Self:
-        """Enter the agent. This will start all [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] registered with the agent so they can be used in a run."""
-        if self._entered_count == 0:
-            self._exit_stack = AsyncExitStack()
-            await self._exit_stack.enter_async_context(self._toolset)
-        self._entered_count += 1
+        async with self._enter_lock:
+            if self._entered_count == 0:
+                self._exit_stack = AsyncExitStack()
+                await self._exit_stack.enter_async_context(self._toolset)
+            self._entered_count += 1
         return self
 
     async def __aexit__(self, *args: Any) -> bool | None:
-        self._entered_count -= 1
-        if self._entered_count <= 0 and self._exit_stack is not None:
-            await self._exit_stack.aclose()
-            self._exit_stack = None
+        async with self._enter_lock:
+            self._entered_count -= 1
+            if self._entered_count == 0 and self._exit_stack is not None:
+                await self._exit_stack.aclose()
+                self._exit_stack = None
 
     def set_mcp_sampling_model(self, model: models.Model | models.KnownModelName | str | None = None) -> None:
         """Set the sampling model on all MCP servers registered with the agent.

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -3,9 +3,10 @@
 import base64
 import functools
 from abc import ABC, abstractmethod
+from asyncio import Lock
 from collections.abc import AsyncIterator, Awaitable, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Callable
 
@@ -60,12 +61,18 @@ class MCPServer(CallableToolset[Any], ABC):
     sampling_model: models.Model | None = None
     # } end of "abstract fields"
 
-    _running_count: int = 0
+    _enter_lock: Lock = field(compare=False)
+    _running_count: int
+    _exit_stack: AsyncExitStack | None
 
     _client: ClientSession
     _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
     _write_stream: MemoryObjectSendStream[SessionMessage]
-    _exit_stack: AsyncExitStack
+
+    def __post_init__(self):
+        self._enter_lock = Lock()
+        self._running_count = 0
+        self._exit_stack = None
 
     @abstractmethod
     @asynccontextmanager
@@ -86,7 +93,7 @@ def name(self) -> str:
         return repr(self)
 
     @property
-    def tool_name_conflict_hint(self) -> str:
+    def _tool_name_conflict_hint(self) -> str:
         return 'Consider setting `tool_prefix` to avoid name conflicts.'
 
     async def list_tools(self) -> list[mcp_types.Tool]:
@@ -188,30 +195,35 @@ def _max_retries_for_tool(self, name: str) -> int:
         return self.max_retries
 
     async def __aenter__(self) -> Self:
-        if self._running_count == 0:
-            self._exit_stack = AsyncExitStack()
-
-            self._read_stream, self._write_stream = await self._exit_stack.enter_async_context(self.client_streams())
-            client = ClientSession(
-                read_stream=self._read_stream,
-                write_stream=self._write_stream,
-                sampling_callback=self._sampling_callback if self.allow_sampling else None,
-                logging_callback=self.log_handler,
-            )
-            self._client = await self._exit_stack.enter_async_context(client)
+        async with self._enter_lock:
+            if self._running_count == 0:
+                self._exit_stack = AsyncExitStack()
+
+                self._read_stream, self._write_stream = await self._exit_stack.enter_async_context(
+                    self.client_streams()
+                )
+                client = ClientSession(
+                    read_stream=self._read_stream,
+                    write_stream=self._write_stream,
+                    sampling_callback=self._sampling_callback if self.allow_sampling else None,
+                    logging_callback=self.log_handler,
+                )
+                self._client = await self._exit_stack.enter_async_context(client)
 
-            with anyio.fail_after(self.timeout):
-                await self._client.initialize()
+                with anyio.fail_after(self.timeout):
+                    await self._client.initialize()
 
-                if log_level := self.log_level:
-                    await self._client.set_logging_level(log_level)
-        self._running_count += 1
+                    if log_level := self.log_level:
+                        await self._client.set_logging_level(log_level)
+            self._running_count += 1
         return self
 
     async def __aexit__(self, *args: Any) -> bool | None:
-        self._running_count -= 1
-        if self._running_count <= 0:
-            await self._exit_stack.aclose()
+        async with self._enter_lock:
+            self._running_count -= 1
+            if self._running_count <= 0 and self._exit_stack is not None:
+                await self._exit_stack.aclose()
+                self._exit_stack = None
 
     @property
     def is_running(self) -> bool:

pydantic_ai_slim/pydantic_ai/tools.py

@@ -364,6 +364,7 @@ class ToolDefinition:
 
     kind: ToolKind = field(default='function')
     """The kind of tool:
+
     - `'function'`: a tool that can be executed by Pydantic AI and has its result returned to the model
     - `'output'`: a tool that passes through an output value that ends the run
     - `'deferred'`: a tool that will be executed not by Pydantic AI, but by the upstream service that called the agent, such as a web application that supports frontend-defined tools provided to Pydantic AI via e.g. [AG-UI](https://docs.ag-ui.com/concepts/tools#frontend-defined-tools).

pydantic_ai_slim/pydantic_ai/toolsets/__init__.py

@@ -27,7 +27,7 @@ def name(self) -> str:
         return self.__class__.__name__.replace('Toolset', ' toolset')
 
     @property
-    def tool_name_conflict_hint(self) -> str:
+    def _tool_name_conflict_hint(self) -> str:
         return 'Consider renaming the tool or wrapping the toolset in a `PrefixedToolset` to avoid name conflicts.'
 
     async def __aenter__(self) -> Self:

pydantic_ai_slim/pydantic_ai/toolsets/_callable.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import replace
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any
 
 from pydantic_core import SchemaValidator
 
@@ -28,7 +28,7 @@ async def _call_tool(self, ctx: RunContext[AgentDepsT], name: str, tool_args: di
     async def call_tool(self, call: ToolCallPart, ctx: RunContext[AgentDepsT], allow_partial: bool = False) -> Any:
         ctx = replace(ctx, tool_name=call.tool_name, tool_call_id=call.tool_call_id)
 
-        pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
+        pyd_allow_partial = 'trailing-strings' if allow_partial else 'off'
         validator = self._get_tool_args_validator(ctx, call.tool_name)
         if isinstance(call.args, str):
             args_dict = validator.validate_json(call.args or '{}', allow_partial=pyd_allow_partial)

pydantic_ai_slim/pydantic_ai/toolsets/combined.py

@@ -3,7 +3,7 @@
 import asyncio
 from collections.abc import Sequence
 from contextlib import AsyncExitStack
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Callable
 
 from typing_extensions import Self
@@ -22,12 +22,16 @@ class CombinedToolset(AbstractToolset[AgentDepsT]):
 
     toolsets: list[AbstractToolset[AgentDepsT]]
     _toolset_per_tool_name: dict[str, AbstractToolset[AgentDepsT]]
-    _exit_stack: AsyncExitStack | None
+
+    _enter_lock: asyncio.Lock = field(compare=False)
     _entered_count: int
+    _exit_stack: AsyncExitStack | None
 
     def __init__(self, toolsets: Sequence[AbstractToolset[AgentDepsT]]):
-        self._exit_stack = None
+        self._enter_lock = asyncio.Lock()
         self._entered_count = 0
+        self._exit_stack = None
+
         self.toolsets = list(toolsets)
 
         self._toolset_per_tool_name = {}
@@ -36,28 +40,30 @@ def __init__(self, toolsets: Sequence[AbstractToolset[AgentDepsT]]):
                 try:
                     existing_toolset = self._toolset_per_tool_name[name]
                     raise UserError(
-                        f'{toolset.name} defines a tool whose name conflicts with existing tool from {existing_toolset.name}: {name!r}. {toolset.tool_name_conflict_hint}'
+                        f'{toolset.name} defines a tool whose name conflicts with existing tool from {existing_toolset.name}: {name!r}. {toolset._tool_name_conflict_hint}'
                     )
                 except KeyError:
                     pass
                 self._toolset_per_tool_name[name] = toolset
 
     async def __aenter__(self) -> Self:
-        if self._entered_count == 0:
-            self._exit_stack = AsyncExitStack()
-            for toolset in self.toolsets:
-                await self._exit_stack.enter_async_context(toolset)
-        self._entered_count += 1
+        async with self._enter_lock:
+            if self._entered_count == 0:
+                self._exit_stack = AsyncExitStack()
+                for toolset in self.toolsets:
+                    await self._exit_stack.enter_async_context(toolset)
+            self._entered_count += 1
         return self
 
     async def __aexit__(self, *args: Any) -> bool | None:
-        self._entered_count -= 1
-        if self._entered_count <= 0 and self._exit_stack is not None:
-            await self._exit_stack.aclose()
-            self._exit_stack = None
+        async with self._enter_lock:
+            self._entered_count -= 1
+            if self._entered_count == 0 and self._exit_stack is not None:
+                await self._exit_stack.aclose()
+                self._exit_stack = None
 
     async def prepare_for_run(self, ctx: RunContext[AgentDepsT]) -> RunToolset[AgentDepsT]:
-        toolsets_for_run = await asyncio.gather(*[toolset.prepare_for_run(ctx) for toolset in self.toolsets])
+        toolsets_for_run = await asyncio.gather(*(toolset.prepare_for_run(ctx) for toolset in self.toolsets))
         combined_for_run = CombinedToolset(toolsets_for_run)
         return RunToolset(combined_for_run, ctx)
 

pydantic_ai_slim/pydantic_ai/toolsets/deferred.py

@@ -11,7 +11,10 @@
 
 
 class DeferredToolset(AbstractToolset[AgentDepsT]):
-    """A toolset that holds deferred tool."""
+    """A toolset that holds deferred tools.
+
+    See [`ToolDefinition.kind`][pydantic_ai.tools.ToolDefinition.kind] for more information about deferred tools.
+    """
 
     _tool_defs: list[ToolDefinition]