Add state dict to Context (#1118) (#1160)

Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>

src/fastmcp/server/context.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import copy
 import warnings
 from collections.abc import Generator
 from contextlib import contextmanager
@@ -83,9 +84,19 @@ class Context:
         request_id = ctx.request_id
         client_id = ctx.client_id
 
+        # Manage state across the request
+        ctx.set_state_value("key", "value")
+        value = ctx.get_state_value("key")
+
         return str(x)
     ```
 
+    State Management:
+    Context objects maintain a state dictionary that can be used to store and share
+    data across middleware and tool calls within a request. When a new context
+    is created (nested contexts), it inherits a copy of its parent's state, ensuring
+    that modifications in child contexts don't affect parent contexts.
+
     The context parameter name can be anything as long as it's annotated with Context.
     The context is optional - tools that don't need it can omit the parameter.
 
@@ -95,9 +106,15 @@ class Context:
         self.fastmcp = fastmcp
         self._tokens: list[Token] = []
         self._notification_queue: set[str] = set()  # Dedupe notifications
+        self._state: dict[str, Any] = {}
 
     async def __aenter__(self) -> Context:
         """Enter the context manager and set this context as the current context."""
+        parent_context = _current_context.get(None)
+        if parent_context is not None:
+            # Inherit state from parent context
+            self._state = copy.deepcopy(parent_context._state)
+
         # Always set this context and save the token
         token = _current_context.set(self)
         self._tokens.append(token)
@@ -455,6 +472,14 @@ class Context:
 
         return fastmcp.server.dependencies.get_http_request()
 
+    def set_state(self, key: str, value: Any) -> None:
+        """Set a value in the context state."""
+        self._state[key] = value
+
+    def get_state(self, key: str) -> Any:
+        """Get a value from the context state. Returns None if the key is not found."""
+        return self._state.get(key)
+
     def _queue_tool_list_changed(self) -> None:
         """Queue a tool list changed notification."""
         self._notification_queue.add("notifications/tools/list_changed")

tests/server/test_context.py

@@ -123,3 +123,51 @@ class TestSessionId:
             "fastmcp.server.dependencies.get_http_headers", return_value=mock_headers
         ):
             assert context.session_id == ""  # Empty string is still returned as-is
+
+
+class TestContextState:
+    """Test suite for Context state functionality."""
+
+    @pytest.mark.asyncio
+    async def test_context_state(self):
+        """Test that state modifications in child contexts don't affect parent."""
+        mock_fastmcp = MagicMock()
+
+        async with Context(fastmcp=mock_fastmcp) as context:
+            assert context.get_state("test1") is None
+            assert context.get_state("test2") is None
+            context.set_state("test1", "value")
+            context.set_state("test2", 2)
+            assert context.get_state("test1") == "value"
+            assert context.get_state("test2") == 2
+            context.set_state("test1", "new_value")
+            assert context.get_state("test1") == "new_value"
+
+    @pytest.mark.asyncio
+    async def test_context_state_inheritance(self):
+        """Test that child contexts inherit parent state."""
+        mock_fastmcp = MagicMock()
+
+        async with Context(fastmcp=mock_fastmcp) as context1:
+            context1.set_state("key1", "key1-context1")
+            context1.set_state("key2", "key2-context1")
+            async with Context(fastmcp=mock_fastmcp) as context2:
+                # Override one key
+                context2.set_state("key1", "key1-context2")
+                assert context2.get_state("key1") == "key1-context2"
+                assert context1.get_state("key1") == "key1-context1"
+                assert context2.get_state("key2") == "key2-context1"
+
+                async with Context(fastmcp=mock_fastmcp) as context3:
+                    # Verify state was inherited
+                    assert context3.get_state("key1") == "key1-context2"
+                    assert context3.get_state("key2") == "key2-context1"
+
+                    # Add a new key and verify parents were not affected
+                    context3.set_state("key-context3-only", 1)
+                    assert context1.get_state("key-context3-only") is None
+                    assert context2.get_state("key-context3-only") is None
+                    assert context3.get_state("key-context3-only") == 1
+
+            assert context1.get_state("key1") == "key1-context1"
+            assert context1.get_state("key-context3-only") is None