Revert "Use contextvars for tracking the MCP sampling model" (#2132)

Author: DouweM

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
-                    exit_stack.enter_context(mcp_server.override_sampling_model(sampling_model))
+                    mcp_server.sampling_model = sampling_model
                 await exit_stack.enter_async_context(mcp_server)
             yield
         finally:

pydantic_ai_slim/pydantic_ai/mcp.py

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

tests/test_examples.py

@@ -6,7 +6,6 @@
 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
@@ -259,7 +258,6 @@ 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