Use contextvars for tracking the MCP sampling model (#2117)

Author: dmontagu

pydantic_ai_slim/pydantic_ai/agent.py

@@ -1740,7 +1740,7 @@ async def run_mcp_servers(
         try:
             for mcp_server in self._mcp_servers:
                 if sampling_model is not None:  # pragma: no branch
-                    mcp_server.sampling_model = sampling_model
+                    exit_stack.enter_context(mcp_server.override_sampling_model(sampling_model))
                 await exit_stack.enter_async_context(mcp_server)
             yield
         finally:

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -3,8 +3,9 @@
 import base64
 import functools
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Awaitable, Sequence
-from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
+from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
+from contextvars import ContextVar
 from dataclasses import dataclass
 from pathlib import Path
 from types import TracebackType
@@ -60,6 +61,22 @@ class MCPServer(ABC):
     _exit_stack: AsyncExitStack
     sampling_model: models.Model | None = None
 
+    def __post_init__(self):
+        self._override_sampling_model: ContextVar[models.Model | None] = ContextVar(
+            '_override_sampling_model', default=None
+        )
+
+    @contextmanager
+    def override_sampling_model(
+        self,
+        model: models.Model,
+    ) -> Iterator[None]:
+        token = self._override_sampling_model.set(model)
+        try:
+            yield
+        finally:
+            self._override_sampling_model.reset(token)
+
     @abstractmethod
     @asynccontextmanager
     async def client_streams(
@@ -184,7 +201,8 @@ async def _sampling_callback(
         self, context: RequestContext[ClientSession, Any], params: mcp_types.CreateMessageRequestParams
     ) -> mcp_types.CreateMessageResult | mcp_types.ErrorData:
         """MCP sampling callback."""
-        if self.sampling_model is None:
+        sampling_model = self._override_sampling_model.get() or self.sampling_model
+        if sampling_model is None:
             raise ValueError('Sampling model is not set')  # pragma: no cover
 
         pai_messages = _mcp.map_from_mcp_params(params)
@@ -196,15 +214,15 @@ async def _sampling_callback(
         if stop_sequences := params.stopSequences:  # pragma: no branch
             model_settings['stop_sequences'] = stop_sequences
 
-        model_response = await self.sampling_model.request(
+        model_response = await sampling_model.request(
             pai_messages,
             model_settings,
             models.ModelRequestParameters(),
         )
         return mcp_types.CreateMessageResult(
             role='assistant',
             content=_mcp.map_from_model_response(model_response),
-            model=self.sampling_model.model_name,
+            model=sampling_model.model_name,
         )
 
     def _map_tool_result_part(

tests/test_examples.py

@@ -6,6 +6,7 @@
 import shutil
 import sys
 from collections.abc import AsyncIterator, Iterable, Sequence
+from contextlib import nullcontext
 from dataclasses import dataclass
 from inspect import FrameInfo
 from io import StringIO
@@ -258,6 +259,7 @@ def rich_prompt_ask(prompt: str, *_args: Any, **_kwargs: Any) -> str:
 
 class MockMCPServer:
     is_running = True
+    override_sampling_model = nullcontext
 
     async def __aenter__(self) -> MockMCPServer:
         return self