[minimcp] Add ContextManager for handler contexts

sreenaths · sreenaths · commit 1209d9624808 · 2025-12-10T00:46:14.000-08:00
- Implement ContextManager for tracking active handler contexts
- Add Context dataclass for holding request metadata (message, time_limiter, scope, responder)
- Support thread-safe and async-safe context isolation using contextvars
- Add helper methods (get_scope, get_responder) for common access patterns
- Add comprehensive unit test suite
diff --git a/src/mcp/server/minimcp/managers/context_manager.py b/src/mcp/server/minimcp/managers/context_manager.py
@@ -0,0 +1,137 @@
+import logging
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+from mcp.server.minimcp.exceptions import ContextError
+from mcp.server.minimcp.limiter import TimeLimiter
+from mcp.server.minimcp.responder import Responder
+from mcp.types import JSONRPCMessage
+
+logger = logging.getLogger(__name__)
+
+
+ScopeT = TypeVar("ScopeT", bound=object)
+
+
+@dataclass(slots=True)  # Use NamedTuple once MCP drops support for Python 3.10
+class Context(Generic[ScopeT]):
+    """
+    Context object holds request metadata available to message handlers.
+
+    The Context provides access to the current message, time limiter for managing
+    idle timeout, optional scope object for passing custom data (auth, user info,
+    session data, database handles, etc.), and optional responder for sending
+    notifications back to the client.
+
+    Attributes:
+        message: The parsed JSON-RPC request message being handled.
+        time_limiter: TimeLimiter for managing handler idle timeout. Call
+            time_limiter.reset() to extend the deadline during active processing.
+        scope: Optional scope object passed to mcp.handle(). Use this to pass
+            authentication details, user info, session data, or database handles
+            to your handlers.
+        responder: Optional responder for sending notifications to the client.
+            Available when bidirectional communication is supported by the transport.
+    """
+
+    message: JSONRPCMessage
+    time_limiter: TimeLimiter
+    scope: ScopeT | None = None
+    responder: Responder | None = None
+
+
+class ContextManager(Generic[ScopeT]):
+    """
+    ContextManager tracks the currently active handler context.
+
+    The ContextManager provides access to request metadata (such as the message,
+    scope, responder, and timeout) directly inside handlers. It uses contextvars
+    to maintain thread-safe and async-safe context isolation across concurrent
+    handler executions.
+
+    You can retrieve the current context using mcp.context.get(). If called
+    outside of a handler, this method raises a ContextError.
+
+    For common use cases, helper methods (get_scope, get_responder) are provided
+    to avoid null checks when accessing optional context attributes.
+    """
+
+    _ctx: ContextVar[Context[ScopeT]] = ContextVar("ctx")
+
+    @contextmanager
+    def active(self, context: Context[ScopeT]):
+        """Set the active context for the current handler execution.
+
+        This context manager sets the context at the start of handler execution
+        and clears it when the handler completes, ensuring proper cleanup.
+
+        Args:
+            context: The context to make active during handler execution.
+
+        Yields:
+            None. The context is accessible via get() during execution.
+        """
+        # Set context
+        token: Token[Context[ScopeT]] = self._ctx.set(context)
+        try:
+            yield
+        finally:
+            # Clear context
+            self._ctx.reset(token)
+
+    def get(self) -> Context[ScopeT]:
+        """Get the current handler context.
+
+        Returns:
+            The active Context object containing message, time_limiter, scope,
+            and responder.
+
+        Raises:
+            ContextError: If called outside of an active handler context.
+        """
+        try:
+            return self._ctx.get()
+        except LookupError as e:
+            msg = "No Context: Called mcp.context.get() outside of an active handler context"
+            logger.error(msg)
+            raise ContextError(msg) from e
+
+    def get_scope(self) -> ScopeT:
+        """Get the scope object from the current context.
+
+        This helper method retrieves the scope and raises an error if it's not
+        available, avoiding the need for null checks in your code.
+
+        Returns:
+            The scope object passed to mcp.handle().
+
+        Raises:
+            ContextError: If called outside of an active handler context or if
+                no scope was provided to mcp.handle().
+        """
+        scope = self.get().scope
+        if scope is None:
+            raise ContextError("ContextError: Scope is not available in current context")
+        return scope
+
+    def get_responder(self) -> Responder:
+        """Get the responder from the current context.
+
+        This helper method retrieves the responder and raises an error if it's not
+        available, avoiding the need for null checks in your code. The responder
+        is only available when using transports that support bidirectional
+        communication (stdio, Streamable HTTP).
+
+        Returns:
+            The Responder for sending notifications to the client.
+
+        Raises:
+            ContextError: If called outside of an active handler context or if
+                the responder is not available (e.g., when using HTTP transport).
+        """
+        responder = self.get().responder
+        if responder is None:
+            raise ContextError("ContextError: Responder is not available in current context")
+        return responder
diff --git a/tests/server/minimcp/unit/managers/test_context_manager.py b/tests/server/minimcp/unit/managers/test_context_manager.py
@@ -0,0 +1,233 @@
+from unittest.mock import Mock
+
+import pytest
+
+import mcp.types as types
+from mcp.server.minimcp.exceptions import ContextError
+from mcp.server.minimcp.limiter import TimeLimiter
+from mcp.server.minimcp.managers.context_manager import Context, ContextManager
+from mcp.server.minimcp.responder import Responder
+
+pytestmark = pytest.mark.anyio
+
+
+class TestContext:
+    """Test suite for Context dataclass."""
+
+    def test_context_creation_minimal(self):
+        """Test creating a Context with minimal required fields."""
+        message = types.JSONRPCMessage(types.JSONRPCRequest(method="test", id=1, jsonrpc="2.0"))
+        time_limiter = Mock(spec=TimeLimiter)
+
+        context = Context[None](message=message, time_limiter=time_limiter)
+
+        assert context.message == message
+        assert context.time_limiter == time_limiter
+        assert context.scope is None
+        assert context.responder is None
+
+    def test_context_creation_with_all_fields(self):
+        """Test creating a Context with all fields."""
+        message = types.JSONRPCMessage(types.JSONRPCRequest(method="test", id=1, jsonrpc="2.0"))
+        time_limiter = Mock(spec=TimeLimiter)
+        scope = {"test": "scope"}
+        responder = Mock(spec=Responder)
+
+        context = Context(message=message, time_limiter=time_limiter, scope=scope, responder=responder)
+
+        assert context.message == message
+        assert context.time_limiter == time_limiter
+        assert context.scope == scope
+        assert context.responder == responder
+
+
+class TestContextManager:
+    """Test suite for ContextManager class."""
+
+    @pytest.fixture
+    def sample_context(self) -> Context[dict[str, str]]:
+        """Create a sample Context for testing."""
+        message = types.JSONRPCMessage(types.JSONRPCRequest(method="test", id=1, jsonrpc="2.0"))
+        time_limiter = Mock(spec=TimeLimiter)
+        scope = {"test": "scope"}
+        responder = Mock(spec=Responder)
+
+        return Context[dict[str, str]](message=message, time_limiter=time_limiter, scope=scope, responder=responder)
+
+    def test_get_without_active_context_raises_error(self):
+        """Test that get() raises ContextError when no context is active."""
+        context_manager = ContextManager[None]()
+        with pytest.raises(ContextError, match="outside of an active handler context"):
+            context_manager.get()
+
+    def test_get_scope_without_active_context_raises_error(self):
+        """Test that get_scope() raises ContextError when no context is active."""
+        context_manager = ContextManager[None]()
+        with pytest.raises(ContextError, match="outside of an active handler context"):
+            context_manager.get_scope()
+
+    def test_get_responder_without_active_context_raises_error(self):
+        """Test that get_responder() raises ContextError when no context is active."""
+        context_manager = ContextManager[None]()
+        with pytest.raises(ContextError, match="outside of an active handler context"):
+            context_manager.get_responder()
+
+    def test_active_context_manager(self, sample_context: Context[dict[str, str]]):
+        """Test the active context manager sets and clears context properly."""
+        context_manager = ContextManager[dict[str, str]]()
+        # Verify no context initially
+        with pytest.raises(ContextError):
+            context_manager.get()
+
+        # Use context manager
+        with context_manager.active(sample_context):
+            # Context should be available
+            retrieved_context = context_manager.get()
+            assert retrieved_context == sample_context
+
+        # Context should be cleared after exiting
+        with pytest.raises(ContextError):
+            context_manager.get()
+
+    def test_get_within_active_context(self, sample_context: Context[dict[str, str]]):
+        """Test that get() returns the active context."""
+        context_manager = ContextManager[dict[str, str]]()
+        with context_manager.active(sample_context):
+            retrieved_context = context_manager.get()
+            assert retrieved_context == sample_context
+            assert retrieved_context.message == sample_context.message
+            assert retrieved_context.time_limiter == sample_context.time_limiter
+            assert retrieved_context.scope == sample_context.scope
+            assert retrieved_context.responder == sample_context.responder
+
+    def test_get_scope_within_active_context(self, sample_context: Context[dict[str, str]]):
+        """Test that get_scope() returns the scope from active context."""
+        context_manager = ContextManager[dict[str, str]]()
+        with context_manager.active(sample_context):
+            scope = context_manager.get_scope()
+            assert scope == sample_context.scope
+
+    def test_get_scope_when_scope_is_none_raises_error(self):
+        """Test that get_scope() raises ContextError when scope is None."""
+
+        context_manager = ContextManager[None]()
+        message = types.JSONRPCMessage(types.JSONRPCRequest(method="test", id=1, jsonrpc="2.0"))
+        time_limiter = Mock(spec=TimeLimiter)
+
+        context_without_scope = Context[None](message=message, time_limiter=time_limiter, scope=None)
+
+        with context_manager.active(context_without_scope):
+            with pytest.raises(ContextError, match="Scope is not available in current context"):
+                context_manager.get_scope()
+
+    def test_get_responder_within_active_context(self, sample_context: Context[dict[str, str]]):
+        """Test that get_responder() returns the responder from active context."""
+        context_manager = ContextManager[dict[str, str]]()
+        with context_manager.active(sample_context):
+            responder = context_manager.get_responder()
+            assert responder == sample_context.responder
+
+    def test_get_responder_when_responder_is_none_raises_error(self):
+        """Test that get_responder() raises ContextError when responder is None."""
+        context_manager = ContextManager[None]()
+        message = types.JSONRPCMessage(types.JSONRPCRequest(method="test", id=1, jsonrpc="2.0"))
+        time_limiter = Mock(spec=TimeLimiter)
+        context_without_responder = Context[None](message=message, time_limiter=time_limiter, responder=None)
+
+        with context_manager.active(context_without_responder):
+            with pytest.raises(ContextError, match="Responder is not available in current context"):
+                context_manager.get_responder()
+
+    def test_nested_context_managers(self):
+        """Test nested context managers work correctly."""
+        context_manager = ContextManager[str]()
+        message1 = types.JSONRPCMessage(types.JSONRPCRequest(method="test1", id=1, jsonrpc="2.0"))
+        message2 = types.JSONRPCMessage(types.JSONRPCRequest(method="test2", id=2, jsonrpc="2.0"))
+        time_limiter = Mock(spec=TimeLimiter)
+
+        context1 = Context[str](message=message1, time_limiter=time_limiter, scope="scope1")
+        context2 = Context[str](message=message2, time_limiter=time_limiter, scope="scope2")
+
+        with context_manager.active(context1):
+            assert context_manager.get_scope() == "scope1"
+
+            with context_manager.active(context2):
+                assert context_manager.get_scope() == "scope2"
+
+            # Should return to outer context
+            assert context_manager.get_scope() == "scope1"
+
+    def test_context_manager_exception_handling(self, sample_context: Context[dict[str, str]]):
+        """Test that context is properly cleared even when exception occurs."""
+        context_manager = ContextManager[dict[str, str]]()
+        with pytest.raises(ValueError):
+            with context_manager.active(sample_context):
+                # Context should be available
+                assert context_manager.get() == sample_context
+                # Raise exception
+                raise ValueError("Test exception")
+
+        # Context should be cleared even after exception
+        with pytest.raises(ContextError):
+            context_manager.get()
+
+    def test_multiple_context_managers_share_context_var(self):
+        """Test that multiple ContextManager instances share the same ContextVar (by design)."""
+        manager1 = ContextManager[str]()
+        manager2 = ContextManager[str]()
+
+        message1 = types.JSONRPCMessage(types.JSONRPCRequest(method="test1", id=1, jsonrpc="2.0"))
+        message2 = types.JSONRPCMessage(types.JSONRPCRequest(method="test2", id=2, jsonrpc="2.0"))
+        time_limiter = Mock(spec=TimeLimiter)
+
+        context1 = Context[str](message=message1, time_limiter=time_limiter, scope="scope1")
+        context2 = Context[str](message=message2, time_limiter=time_limiter, scope="scope2")
+
+        # ContextVar is shared across instances, so the last set context wins
+        with manager1.active(context1):
+            with manager2.active(context2):
+                # Both managers see the same context (context2) since they share the ContextVar
+                assert manager1.get_scope() == "scope2"
+                assert manager2.get_scope() == "scope2"
+
+    async def test_context_var_isolation(self):
+        """Test that ContextVar properly isolates contexts across different execution contexts."""
+        import anyio
+
+        context_manager = ContextManager[str]()
+
+        message1 = types.JSONRPCMessage(types.JSONRPCRequest(method="test1", id=1, jsonrpc="2.0"))
+        message2 = types.JSONRPCMessage(types.JSONRPCRequest(method="test2", id=2, jsonrpc="2.0"))
+        time_limiter = Mock(spec=TimeLimiter)
+
+        context1 = Context[str](message=message1, time_limiter=time_limiter, scope="scope1")
+        context2 = Context[str](message=message2, time_limiter=time_limiter, scope="scope2")
+
+        async def task1():
+            with context_manager.active(context1):
+                await anyio.sleep(0.01)
+                return context_manager.get_scope()
+
+        async def task2():
+            with context_manager.active(context2):
+                await anyio.sleep(0.01)
+                return context_manager.get_scope()
+
+        results: list[str] = []
+
+        async def collect_task1():
+            result = await task1()
+            results.append(result)
+
+        async def collect_task2():
+            result = await task2()
+            results.append(result)
+
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(collect_task1)
+            tg.start_soon(collect_task2)
+
+        # Results order may vary, so check both are present
+        assert "scope1" in results
+        assert "scope2" in results
+        assert len(results) == 2
