From b1e29a85af26f3830ffd2550c10c41ac3ea559f0 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 13 Jan 2026 21:51:21 -0500 Subject: [PATCH 1/9] Add session-scoped state persistence State now persists across tool calls within an MCP session via ctx.get_state() and ctx.set_state(). Uses pykeyvalue with configurable backends (defaults to in-memory). State keys are prefixed with session_id for isolation between clients. Includes 1-day TTL to prevent memory leaks. --- examples/persistent_state/README.md | 51 +++++ examples/persistent_state/client.py | 85 ++++++++ examples/persistent_state/client_stdio.py | 79 ++++++++ examples/persistent_state/server.py | 42 ++++ src/fastmcp/server/context.py | 86 ++++++-- src/fastmcp/server/low_level.py | 2 +- src/fastmcp/server/server.py | 20 +- .../test_initialization_middleware.py | 139 +++++++++---- tests/server/test_context.py | 188 ++++++++++++++---- 9 files changed, 585 insertions(+), 107 deletions(-) create mode 100644 examples/persistent_state/README.md create mode 100644 examples/persistent_state/client.py create mode 100644 examples/persistent_state/client_stdio.py create mode 100644 examples/persistent_state/server.py diff --git a/examples/persistent_state/README.md b/examples/persistent_state/README.md new file mode 100644 index 0000000000..04b412bb7a --- /dev/null +++ b/examples/persistent_state/README.md @@ -0,0 +1,51 @@ +# Persistent Session State + +This example demonstrates session-scoped state that persists across tool calls within the same MCP session. + +## What it shows + +- State set in one tool call is readable in subsequent calls +- Different clients have isolated state (same keys, different values) +- Reconnecting creates a new session with fresh state + +## Running + +**HTTP transport:** + +```bash +# Terminal 1: Start the server +uv run python server.py + +# Terminal 2: Run the client +uv run python client.py +``` + +**STDIO transport (in-process):** + +```bash +uv run python client_stdio.py +``` + +## Example output + +``` +Each line below is a separate tool call + +Alice connects + session a9f6eaa3 + set user = Alice + set secret = alice-password + get user → Alice + get secret → alice-password + +Bob connects (different session) + session 0c3bffc5 + get user → not found + get secret → not found + set user = Bob + get user → Bob + +Alice reconnects (new session) + session e39640e3 + get user → not found +``` diff --git a/examples/persistent_state/client.py b/examples/persistent_state/client.py new file mode 100644 index 0000000000..ea4c3b9fa8 --- /dev/null +++ b/examples/persistent_state/client.py @@ -0,0 +1,85 @@ +"""Client for testing persistent state. + +Run the server first: + uv run python examples/persistent_state/server.py + +Then run this client: + uv run python examples/persistent_state/client.py +""" + +import asyncio + +from rich.console import Console + +from fastmcp import Client +from fastmcp.client.transports import StreamableHttpTransport + +URL = "http://127.0.0.1:8000/mcp" +console = Console() + + +async def main(): + console.print() + console.print("[dim italic]Each line below is a separate tool call[/dim italic]") + console.print() + + # --- Alice's session --- + console.print("[dim]Alice connects[/dim]") + + transport1 = StreamableHttpTransport(url=URL) + async with Client(transport=transport1) as alice: + result = await alice.call_tool("list_session_info", {}) + console.print(f" session [cyan]{result.data['session_id'][:8]}[/cyan]") + + await alice.call_tool("set_value", {"key": "user", "value": "Alice"}) + console.print(" set [white]user[/white] = [green]Alice[/green]") + + await alice.call_tool("set_value", {"key": "secret", "value": "alice-password"}) + console.print(" set [white]secret[/white] = [green]alice-password[/green]") + + result = await alice.call_tool("get_value", {"key": "user"}) + console.print(" get [white]user[/white] → [green]Alice[/green]") + + result = await alice.call_tool("get_value", {"key": "secret"}) + console.print(" get [white]secret[/white] → [green]alice-password[/green]") + + console.print() + + # --- Bob's session --- + console.print("[dim]Bob connects (different session)[/dim]") + + transport2 = StreamableHttpTransport(url=URL) + async with Client(transport=transport2) as bob: + result = await bob.call_tool("list_session_info", {}) + console.print(f" session [cyan]{result.data['session_id'][:8]}[/cyan]") + + await bob.call_tool("get_value", {"key": "user"}) + console.print(" get [white]user[/white] → [dim]not found[/dim]") + + await bob.call_tool("get_value", {"key": "secret"}) + console.print(" get [white]secret[/white] → [dim]not found[/dim]") + + await bob.call_tool("set_value", {"key": "user", "value": "Bob"}) + console.print(" set [white]user[/white] = [green]Bob[/green]") + + await bob.call_tool("get_value", {"key": "user"}) + console.print(" get [white]user[/white] → [green]Bob[/green]") + + console.print() + + # --- Alice reconnects --- + console.print("[dim]Alice reconnects (new session)[/dim]") + + transport3 = StreamableHttpTransport(url=URL) + async with Client(transport=transport3) as alice_again: + result = await alice_again.call_tool("list_session_info", {}) + console.print(f" session [cyan]{result.data['session_id'][:8]}[/cyan]") + + await alice_again.call_tool("get_value", {"key": "user"}) + console.print(" get [white]user[/white] → [dim]not found[/dim]") + + console.print() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/persistent_state/client_stdio.py b/examples/persistent_state/client_stdio.py new file mode 100644 index 0000000000..544a6245dd --- /dev/null +++ b/examples/persistent_state/client_stdio.py @@ -0,0 +1,79 @@ +"""Client for testing persistent state over STDIO. + +Run directly: + uv run python examples/persistent_state/client_stdio.py +""" + +import asyncio + +from rich.console import Console + +from fastmcp import Client + +from server import server + +console = Console() + + +async def main(): + console.print() + console.print("[dim italic]Each line below is a separate tool call[/dim italic]") + console.print() + + # --- Alice's session --- + console.print("[dim]Alice connects[/dim]") + + async with Client(server) as alice: + result = await alice.call_tool("list_session_info", {}) + console.print(f" session [cyan]{result.data['session_id'][:8]}[/cyan]") + + await alice.call_tool("set_value", {"key": "user", "value": "Alice"}) + console.print(" set [white]user[/white] = [green]Alice[/green]") + + await alice.call_tool("set_value", {"key": "secret", "value": "alice-password"}) + console.print(" set [white]secret[/white] = [green]alice-password[/green]") + + await alice.call_tool("get_value", {"key": "user"}) + console.print(" get [white]user[/white] → [green]Alice[/green]") + + await alice.call_tool("get_value", {"key": "secret"}) + console.print(" get [white]secret[/white] → [green]alice-password[/green]") + + console.print() + + # --- Bob's session --- + console.print("[dim]Bob connects (different session)[/dim]") + + async with Client(server) as bob: + result = await bob.call_tool("list_session_info", {}) + console.print(f" session [cyan]{result.data['session_id'][:8]}[/cyan]") + + await bob.call_tool("get_value", {"key": "user"}) + console.print(" get [white]user[/white] → [dim]not found[/dim]") + + await bob.call_tool("get_value", {"key": "secret"}) + console.print(" get [white]secret[/white] → [dim]not found[/dim]") + + await bob.call_tool("set_value", {"key": "user", "value": "Bob"}) + console.print(" set [white]user[/white] = [green]Bob[/green]") + + await bob.call_tool("get_value", {"key": "user"}) + console.print(" get [white]user[/white] → [green]Bob[/green]") + + console.print() + + # --- Alice reconnects --- + console.print("[dim]Alice reconnects (new session)[/dim]") + + async with Client(server) as alice_again: + result = await alice_again.call_tool("list_session_info", {}) + console.print(f" session [cyan]{result.data['session_id'][:8]}[/cyan]") + + await alice_again.call_tool("get_value", {"key": "user"}) + console.print(" get [white]user[/white] → [dim]not found[/dim]") + + console.print() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/persistent_state/server.py b/examples/persistent_state/server.py new file mode 100644 index 0000000000..f212d5b93c --- /dev/null +++ b/examples/persistent_state/server.py @@ -0,0 +1,42 @@ +"""Example: Persistent session-scoped state. + +This demonstrates using Context.get_state() and set_state() to store +data that persists across tool calls within the same MCP session. + +Run with: + uv run python examples/persistent_state/server.py +""" + +from fastmcp import FastMCP +from fastmcp.server.context import Context + +server = FastMCP("StateExample") + + +@server.tool +async def set_value(key: str, value: str, ctx: Context) -> str: + """Store a value in session state.""" + await ctx.set_state(key, value) + return f"Stored '{key}' = '{value}'" + + +@server.tool +async def get_value(key: str, ctx: Context) -> str: + """Retrieve a value from session state.""" + value = await ctx.get_state(key) + if value is None: + return f"Key '{key}' not found" + return f"'{key}' = '{value}'" + + +@server.tool +async def list_session_info(ctx: Context) -> dict: + """Get information about the current session.""" + return { + "session_id": ctx.session_id, + "transport": ctx.transport, + } + + +if __name__ == "__main__": + server.run(transport="streamable-http") diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 9ed73119a6..dc5c8d9431 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -1,6 +1,5 @@ from __future__ import annotations -import copy import json import logging import weakref @@ -55,7 +54,7 @@ from fastmcp.server.sampling.run import ( execute_tools as run_sampling_tools, ) -from fastmcp.server.server import FastMCP +from fastmcp.server.server import FastMCP, StateValue from fastmcp.utilities.json_schema import compress_schema from fastmcp.utilities.logging import _clamp_logger, get_logger from fastmcp.utilities.types import get_cached_typeadapter @@ -156,29 +155,33 @@ async def my_tool(x: int, ctx: Context) -> str: request_id = ctx.request_id client_id = ctx.client_id - # Manage state across the request - ctx.set_state("key", "value") - value = ctx.get_state("key") + # Manage state across the session (persists across requests) + await ctx.set_state("key", "value") + value = await ctx.get_state("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. + Context provides session-scoped state that persists across requests within + the same MCP session. State is automatically keyed by session, ensuring + isolation between different clients. + + For STDIO and SSE transports, state set during `on_initialize` middleware + will persist to tool calls. For StreamableHTTP, state during `on_initialize` + is isolated because the session ID comes from request headers (not available + during init). 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. """ - def __init__(self, fastmcp: FastMCP): + def __init__(self, fastmcp: FastMCP, session: ServerSession | None = None): self._fastmcp: weakref.ref[FastMCP] = weakref.ref(fastmcp) + self._session: ServerSession | None = session # For state ops during init self._tokens: list[Token] = [] self._notification_queue: list[mcp.types.ServerNotificationType] = [] - self._state: dict[str, Any] = {} self._exit_stack: AsyncExitStack | None = None self._cancel_scope: anyio.CancelScope | None = None @@ -192,11 +195,6 @@ def fastmcp(self) -> FastMCP: 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) @@ -1112,13 +1110,57 @@ async def elicit( else: raise ValueError(f"Unexpected elicitation action: {result.action}") - def set_state(self, key: str, value: Any) -> None: - """Set a value in the context state.""" - self._state[key] = value + def _get_state_prefix(self) -> str: + """Get the prefix for state keys. + + Uses session_id when available (consistent with the public API). + Falls back to id(session) during on_initialize when session_id + isn't available yet. + """ + # When request_context is available, use session_id for consistency + if self.request_context is not None: + return self.session_id + + # During on_initialize, fall back to id(session) + if self._session is not None: + return str(id(self._session)) + + raise RuntimeError("No session available for state operations") + + def _make_state_key(self, key: str) -> str: + """Create session-prefixed key for state storage.""" + return f"{self._get_state_prefix()}:{key}" - 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) + # Default TTL for session state: 1 day in seconds + _STATE_TTL_SECONDS: int = 86400 + + async def set_state(self, key: str, value: Any) -> None: + """Set a value in the session-scoped state store. + + Values persist across requests within the same MCP session. + The key is automatically prefixed with the session identifier. + State expires after 1 day to prevent unbounded memory growth. + """ + prefixed_key = self._make_state_key(key) + await self.fastmcp._state_store.put( + key=prefixed_key, + value=StateValue(value=value), + ttl=self._STATE_TTL_SECONDS, + ) + + async def get_state(self, key: str) -> Any: + """Get a value from the session-scoped state store. + + Returns None if the key is not found. + """ + prefixed_key = self._make_state_key(key) + result = await self.fastmcp._state_store.get(key=prefixed_key) + return result.value if result is not None else None + + async def delete_state(self, key: str) -> None: + """Delete a value from the session-scoped state store.""" + prefixed_key = self._make_state_key(key) + await self.fastmcp._state_store.delete(key=prefixed_key) async def _periodic_flush(self) -> None: """Background task that flushes the notification queue every second.""" diff --git a/src/fastmcp/server/low_level.py b/src/fastmcp/server/low_level.py index f6b579e25d..355f554b79 100644 --- a/src/fastmcp/server/low_level.py +++ b/src/fastmcp/server/low_level.py @@ -96,7 +96,7 @@ async def call_original_handler( return None async with fastmcp.server.context.Context( - fastmcp=self.fastmcp + fastmcp=self.fastmcp, session=self ) as fastmcp_ctx: # Create the middleware context. mw_context = MiddlewareContext( diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 4f013e6fef..84af887e7f 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -30,6 +30,9 @@ import httpx import mcp.types import uvicorn +from key_value.aio.adapters.pydantic import PydanticAdapter +from key_value.aio.protocols import AsyncKeyValue +from key_value.aio.stores.memory import MemoryStore from mcp.server.lowlevel.server import LifespanResultT, NotificationOptions from mcp.server.stdio import stdio_server from mcp.shared.exceptions import McpError @@ -99,7 +102,7 @@ from fastmcp.utilities.cli import log_server_banner from fastmcp.utilities.components import FastMCPComponent from fastmcp.utilities.logging import get_logger, temporary_log_level -from fastmcp.utilities.types import NotSet, NotSetT +from fastmcp.utilities.types import FastMCPBaseModel, NotSet, NotSetT if TYPE_CHECKING: from docket import Docket @@ -219,6 +222,12 @@ async def wrap( return wrap +class StateValue(FastMCPBaseModel): + """Wrapper for stored context state values.""" + + value: Any + + class FastMCP(Generic[LifespanResultT]): def __init__( self, @@ -241,6 +250,7 @@ def __init__( on_duplicate: DuplicateBehavior | None = None, strict_input_validation: bool | None = None, tasks: bool | None = None, + session_state_store: AsyncKeyValue | None = None, # --- # --- DEPRECATED parameters --- # --- @@ -277,6 +287,14 @@ def __init__( self._additional_http_routes: list[BaseRoute] = [] + # Session-scoped state store (shared across all requests) + self._state_storage: AsyncKeyValue = session_state_store or MemoryStore() + self._state_store: PydanticAdapter[StateValue] = PydanticAdapter[StateValue]( + key_value=self._state_storage, + pydantic_model=StateValue, + default_collection="fastmcp_state", + ) + # Create LocalProvider for local components self._local_provider: LocalProvider = LocalProvider( on_duplicate=self._on_duplicate diff --git a/tests/server/middleware/test_initialization_middleware.py b/tests/server/middleware/test_initialization_middleware.py index ae5aba5d38..4716af5f07 100644 --- a/tests/server/middleware/test_initialization_middleware.py +++ b/tests/server/middleware/test_initialization_middleware.py @@ -12,7 +12,12 @@ class InitializationMiddleware(Middleware): - """Middleware that captures initialization details.""" + """Middleware that captures initialization details. + + Note: Session state is NOT available during on_initialize because + the MCP session has not been established yet. Use instance variables + to store data that needs to persist across the session. + """ def __init__(self): super().__init__() @@ -25,7 +30,7 @@ async def on_initialize( context: MiddlewareContext[mt.InitializeRequest], call_next: CallNext[mt.InitializeRequest, None], ) -> None: - """Capture initialization details and store session data.""" + """Capture initialization details.""" self.initialized = True # Extract client info from the initialize params @@ -34,13 +39,13 @@ async def on_initialize( ): self.client_info = context.message.params.clientInfo - # Store data in the context state for cross-request access - if context.fastmcp_context: - context.fastmcp_context.set_state("client_initialized", True) - if self.client_info: - context.fastmcp_context.set_state( - "client_name", getattr(self.client_info, "name", "unknown") - ) + # Store in instance for cross-request access + # (session state is not available during on_initialize) + self.session_data["client_initialized"] = True + if self.client_info: + self.session_data["client_name"] = getattr( + self.client_info, "name", "unknown" + ) return await call_next(context) @@ -194,41 +199,36 @@ def test_tool() -> str: assert detect_mw.tools_modified is True -async def test_initialization_middleware_with_state_sharing(): - """Test that state set during initialization is available in later requests.""" +async def test_session_state_persists_across_tool_calls(): + """Test that session-scoped state persists across multiple tool calls. + + Session state is only available after the session is established, + so it can't be set during on_initialize. This test shows state set + during one tool call is accessible in subsequent tool calls. + """ server = FastMCP("TestServer") class StateTrackingMiddleware(Middleware): def __init__(self): super().__init__() - self.init_state = {} - self.tool_state = {} - - async def on_initialize( - self, - context: MiddlewareContext[mt.InitializeRequest], - call_next: CallNext[mt.InitializeRequest, None], - ) -> None: - # Store some state during initialization - if context.fastmcp_context: - context.fastmcp_context.set_state("init_timestamp", "2024-01-01") - context.fastmcp_context.set_state("client_id", "test-123") - self.init_state["timestamp"] = "2024-01-01" - self.init_state["client_id"] = "test-123" - - return await call_next(context) + self.call_count = 0 + self.state_values = [] async def on_call_tool( self, context: MiddlewareContext[mt.CallToolRequestParams], call_next: CallNext[mt.CallToolRequestParams, Any], ) -> Any: - # Try to access state from initialization + self.call_count += 1 + if context.fastmcp_context: - timestamp = context.fastmcp_context.get_state("init_timestamp") - client_id = context.fastmcp_context.get_state("client_id") - self.tool_state["timestamp"] = timestamp - self.tool_state["client_id"] = client_id + # Read existing state + counter = await context.fastmcp_context.get_state("call_counter") + self.state_values.append(counter) + + # Increment and save + new_counter = (counter or 0) + 1 + await context.fastmcp_context.set_state("call_counter", new_counter) return await call_next(context) @@ -240,20 +240,23 @@ def test_tool() -> str: return "success" async with Client(server) as client: - # Initialization should have set state - assert middleware.init_state["timestamp"] == "2024-01-01" - assert middleware.init_state["client_id"] == "test-123" - - # Call a tool - state should be accessible + # First call - state should be None initially result = await client.call_tool("test_tool", {}) assert isinstance(result.content[0], TextContent) assert result.content[0].text == "success" - # State should have been accessible during tool call - # Note: State is request-scoped, so it won't persist across requests - # This test shows the pattern, but actual cross-request state would need - # external storage (Redis, DB, etc.) - # The middleware.tool_state might be None if state doesn't persist + # Second call - state should show previous value (1) + result = await client.call_tool("test_tool", {}) + assert isinstance(result.content[0], TextContent) + + # Third call - state should show previous value (2) + result = await client.call_tool("test_tool", {}) + assert isinstance(result.content[0], TextContent) + + # Verify state persisted across calls within the session + assert middleware.call_count == 3 + # First call saw None, second saw 1, third saw 2 + assert middleware.state_values == [None, 1, 2] async def test_middleware_can_access_initialize_result(): @@ -375,3 +378,55 @@ async def on_initialize( pass assert middleware.error_raised is True + + +async def test_state_isolation_between_streamable_http_clients(): + """Test that different HTTP clients have isolated session state. + + Each client should have its own session ID and isolated state. + """ + from fastmcp.client.transports import StreamableHttpTransport + from fastmcp.server.context import Context + from fastmcp.utilities.tests import run_server_async + + server = FastMCP("TestServer") + + @server.tool + async def store_and_read(value: str, ctx: Context) -> dict: + """Store a value and return session info.""" + existing = await ctx.get_state("client_value") + await ctx.set_state("client_value", value) + return { + "existing": existing, + "stored": value, + "session_id": ctx.session_id, + } + + async with run_server_async(server, transport="streamable-http") as url: + import json + + # Client 1 stores its value + transport1 = StreamableHttpTransport(url=url) + async with Client(transport=transport1) as client1: + result1 = await client1.call_tool( + "store_and_read", {"value": "client1-value"} + ) + data1 = json.loads(result1.content[0].text) + assert data1["existing"] is None + assert data1["stored"] == "client1-value" + session_id_1 = data1["session_id"] + + # Client 2 should have completely isolated state + transport2 = StreamableHttpTransport(url=url) + async with Client(transport=transport2) as client2: + result2 = await client2.call_tool( + "store_and_read", {"value": "client2-value"} + ) + data2 = json.loads(result2.content[0].text) + # Should NOT see client1's value + assert data2["existing"] is None + assert data2["stored"] == "client2-value" + session_id_2 = data2["session_id"] + + # Session IDs should be different + assert session_id_1 != session_id_2 diff --git a/tests/server/test_context.py b/tests/server/test_context.py index 310dc4e9b9..3013696480 100644 --- a/tests/server/test_context.py +++ b/tests/server/test_context.py @@ -86,47 +86,153 @@ def test_session_id_without_http_headers(self, context): class TestContextState: """Test suite for Context state functionality.""" - 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" - - 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 + async def test_context_state_basic(self): + """Test basic get/set/delete state operations.""" + server = FastMCP("test") + mock_session = MagicMock() # Use same session for consistent id() + + async with Context(fastmcp=server, session=mock_session) as context: + # Initially empty + assert await context.get_state("test1") is None + assert await context.get_state("test2") is None + + # Set values + await context.set_state("test1", "value") + await context.set_state("test2", 2) + + # Retrieve values + assert await context.get_state("test1") == "value" + assert await context.get_state("test2") == 2 + + # Update value + await context.set_state("test1", "new_value") + assert await context.get_state("test1") == "new_value" + + # Delete value + await context.delete_state("test1") + assert await context.get_state("test1") is None + + async def test_context_state_session_isolation(self): + """Test that different sessions have isolated state.""" + server = FastMCP("test") + session_a = MagicMock() + session_b = MagicMock() + + async with Context(fastmcp=server, session=session_a) as context1: + await context1.set_state("key", "value-from-A") + + async with Context(fastmcp=server, session=session_b) as context2: + # Session B should not see session A's state + assert await context2.get_state("key") is None + await context2.set_state("key", "value-from-B") + assert await context2.get_state("key") == "value-from-B" + + # Verify session A's state is still intact + async with Context(fastmcp=server, session=session_a) as context3: + assert await context3.get_state("key") == "value-from-A" + + async def test_context_state_persists_across_requests(self): + """Test that state persists across multiple context instances (requests).""" + server = FastMCP("test") + mock_session = MagicMock() # Same session = same id() + + # First request sets state + async with Context(fastmcp=server, session=mock_session) as context1: + await context1.set_state("counter", 1) + + # Second request in same session sees the state + async with Context(fastmcp=server, session=mock_session) as context2: + counter = await context2.get_state("counter") + assert counter == 1 + await context2.set_state("counter", counter + 1) + + # Third request sees updated state + async with Context(fastmcp=server, session=mock_session) as context3: + assert await context3.get_state("counter") == 2 + + async def test_context_state_nested_contexts_share_state(self): + """Test that nested contexts within the same session share state.""" + server = FastMCP("test") + mock_session = MagicMock() + + async with Context(fastmcp=server, session=mock_session) as context1: + await context1.set_state("key", "outer-value") + + async with Context(fastmcp=server, session=mock_session) as context2: + # Nested context sees same state (same session) + assert await context2.get_state("key") == "outer-value" + + # Nested context can modify shared state + await context2.set_state("key", "inner-value") + + # Outer context sees the modification + assert await context1.get_state("key") == "inner-value" + + async def test_two_clients_same_key_isolated_by_session(self): + """Test that two different clients can store the same key independently. + + Each client gets an auto-generated session ID, and their state is isolated. + """ + import json + + from fastmcp import Client + + server = FastMCP("test") + stored_session_ids: list[str] = [] + + @server.tool + async def store_and_read(value: str, ctx: Context) -> dict: + """Store a value and return all state info.""" + stored_session_ids.append(ctx.session_id) + existing = await ctx.get_state("shared_key") + await ctx.set_state("shared_key", value) + new_value = await ctx.get_state("shared_key") + return { + "session_id": ctx.session_id, + "existing_value": existing, + "new_value": new_value, + } + + # Client 1 stores "value-from-client-1" + async with Client(server) as client1: + result1 = await client1.call_tool( + "store_and_read", {"value": "value-from-client-1"} + ) + data1 = json.loads(result1.content[0].text) + assert data1["existing_value"] is None # First write + assert data1["new_value"] == "value-from-client-1" + session_id_1 = data1["session_id"] + + # Client 2 stores "value-from-client-2" with the SAME key + async with Client(server) as client2: + result2 = await client2.call_tool( + "store_and_read", {"value": "value-from-client-2"} + ) + data2 = json.loads(result2.content[0].text) + # Client 2 should NOT see client 1's value (different session) + assert data2["existing_value"] is None + assert data2["new_value"] == "value-from-client-2" + session_id_2 = data2["session_id"] + + # Verify session IDs were auto-generated and are different + assert session_id_1 is not None + assert session_id_2 is not None + assert session_id_1 != session_id_2 + + # Verify they look like UUIDs (36 chars with dashes) + assert len(session_id_1) == 36 + assert len(session_id_2) == 36 + + # Client 1 reconnects and should still see their value + async with Client(server) as client1_again: + # But this is a NEW session (new connection = new session ID) + result3 = await client1_again.call_tool( + "store_and_read", {"value": "value-from-client-1-again"} + ) + data3 = json.loads(result3.content[0].text) + # New session, so existing value is None + assert data3["existing_value"] is None + assert data3["session_id"] != session_id_1 # Different session class TestContextMeta: From e5c5d7f5630697496d23c5e58d0e723207827a1a Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Tue, 13 Jan 2026 23:29:34 -0500 Subject: [PATCH 2/9] Fix import ordering using sys.path pattern from tasks example --- examples/persistent_state/client_stdio.py | 11 ++++++++++- pyproject.toml | 2 ++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/examples/persistent_state/client_stdio.py b/examples/persistent_state/client_stdio.py index 544a6245dd..b447bc16a7 100644 --- a/examples/persistent_state/client_stdio.py +++ b/examples/persistent_state/client_stdio.py @@ -5,12 +5,21 @@ """ import asyncio +import sys +from pathlib import Path from rich.console import Console from fastmcp import Client -from server import server +# Add parent directory to path for importing the server module +examples_dir = Path(__file__).parent.parent.parent +if str(examples_dir) not in sys.path: + sys.path.insert(0, str(examples_dir)) + +import examples.persistent_state.server as server_module # noqa: E402 + +server = server_module.server console = Console() diff --git a/pyproject.toml b/pyproject.toml index 189a6fdc80..d030c1486e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -169,6 +169,8 @@ extend-select = [ "UP", # flake8-unused-imports: Catches unused imports ] +[tool.ruff.lint.isort] +known-first-party = ["fastmcp"] [tool.ruff.lint.per-file-ignores] "__init__.py" = ["F401", "I001", "RUF013"] From 72f31135be0590de0b4bc3c1c5c0b0bab873caea Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Wed, 14 Jan 2026 18:08:36 -0500 Subject: [PATCH 3/9] Document session state feature in context docs --- docs/servers/context.mdx | 72 +++++++++++++++++++++------------------- 1 file changed, 37 insertions(+), 35 deletions(-) diff --git a/docs/servers/context.mdx b/docs/servers/context.mdx index 3dc1b85e71..3c0f6161a2 100644 --- a/docs/servers/context.mdx +++ b/docs/servers/context.mdx @@ -22,7 +22,7 @@ The `Context` object provides a clean interface to access MCP features within yo - **Prompt Access**: List and retrieve prompts registered with the server - **LLM Sampling**: Request the client's LLM to generate text based on provided messages - **User Elicitation**: Request structured input from users during tool execution -- **State Management**: Store and share data between middleware and the handler within a single request +- **Session State**: Store data that persists across requests within an MCP session - **Request Information**: Access metadata about the current request - **Server Access**: When needed, access the underlying FastMCP server instance @@ -209,55 +209,57 @@ messages = result.messages - **`ctx.list_prompts() -> list[MCPPrompt]`**: Returns list of all available prompts - **`ctx.get_prompt(name: str, arguments: dict[str, Any] | None = None) -> GetPromptResult`**: Get a specific prompt with optional arguments -### State Management +### Session State - + -Store and share data between middleware and handlers within a single MCP request. Each MCP request (such as calling a tool, reading a resource, listing tools, or listing resources) receives its own context object with isolated state. Context state is particularly useful for passing information from [middleware](/servers/middleware) to your handlers. +Store data that persists across multiple requests within the same MCP session. Session state is automatically keyed by the client's session, ensuring isolation between different clients. -To store a value in the context state, use `ctx.set_state(key, value)`. To retrieve a value, use `ctx.get_state(key)`. +```python +from fastmcp import FastMCP, Context - -Context state is scoped to a single MCP request. Each operation (tool call, resource read, list operation, etc.) receives a new context object. State set during one request will not be available in subsequent requests. For persistent data storage across requests, use external storage mechanisms like databases, files, or in-memory caches. - +mcp = FastMCP("stateful-app") -This simplified example shows how to use MCP middleware to store user info in the context state, and how to access that state in a tool: +@mcp.tool +async def increment_counter(ctx: Context) -> int: + """Increment a counter that persists across tool calls.""" + count = await ctx.get_state("counter") or 0 + await ctx.set_state("counter", count + 1) + return count + 1 -```python {7-8, 16-17} -from fastmcp.server.middleware import Middleware, MiddlewareContext +@mcp.tool +async def get_counter(ctx: Context) -> int: + """Get the current counter value.""" + return await ctx.get_state("counter") or 0 +``` -class UserAuthMiddleware(Middleware): - async def on_call_tool(self, context: MiddlewareContext, call_next): +Each client session has its own isolated state—two different clients calling `increment_counter` will each have their own counter. - # Middleware stores user info in context state - context.fastmcp_context.set_state("user_id", "user_123") - context.fastmcp_context.set_state("permissions", ["read", "write"]) +**Method signatures:** +- **`await ctx.set_state(key: str, value: Any) -> None`**: Store a value in session state +- **`await ctx.get_state(key: str) -> Any`**: Retrieve a value (returns None if not found) +- **`await ctx.delete_state(key: str) -> None`**: Remove a value from session state - return await call_next(context) + +State methods are async and require `await`. State expires after 1 day to prevent unbounded memory growth. + -@mcp.tool -async def secure_operation(data: str, ctx: Context) -> str: - """Tool can access state set by middleware.""" +#### Custom Storage Backends - user_id = ctx.get_state("user_id") # "user_123" - permissions = ctx.get_state("permissions") # ["read", "write"] - - if "write" not in permissions: - return "Access denied" - - return f"Processing {data} for user {user_id}" +By default, session state uses an in-memory store suitable for single-server deployments. For distributed or serverless deployments, provide a custom storage backend: + +```python +from key_value.aio.stores.redis import RedisStore + +# Use Redis for distributed state +mcp = FastMCP("distributed-app", session_state_store=RedisStore(...)) ``` -**Method signatures:** -- **`ctx.set_state(key: str, value: Any) -> None`**: Store a value in the context state -- **`ctx.get_state(key: str) -> Any`**: Retrieve a value from the context state (returns None if not found) +Any backend compatible with the [py-key-value-aio](https://github.com/strawgate/py-key-value) `AsyncKeyValue` protocol works. See [Storage Backends](/servers/storage-backends) for more options including Redis, DynamoDB, and MongoDB. -**State Inheritance:** -When a new context is created (nested contexts), it inherits a copy of its parent's state. This ensures that: -- State set on a child context never affects the parent context -- State set on a parent context after the child context is initialized is not propagated to the child context +#### State During Initialization -This makes state management predictable and prevents unexpected side effects between nested operations. +For STDIO and SSE transports, state set during `on_initialize` middleware persists to subsequent tool calls (same session object). For StreamableHTTP, state during initialization is isolated because the session ID comes from request headers which aren't available during init. ### Change Notifications From 17d5265d5c28f770dbef7e755f2a901c7d8196ee Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 16 Jan 2026 11:33:07 -0500 Subject: [PATCH 4/9] Address review feedback: fix init state bug and add type annotations Fix prefix mismatch that made init state unreachable: session_id now uses str(id(session)) for STDIO/SSE instead of UUIDs, ensuring state set during on_initialize persists to tool calls. Also add type annotations to examples and fix test isolation issues. --- examples/persistent_state/README.md | 2 +- examples/persistent_state/client.py | 2 +- examples/persistent_state/client_stdio.py | 10 ++++---- examples/persistent_state/server.py | 2 +- src/fastmcp/server/context.py | 8 +++---- tests/server/test_context.py | 29 +++++++++++++---------- 6 files changed, 28 insertions(+), 25 deletions(-) diff --git a/examples/persistent_state/README.md b/examples/persistent_state/README.md index 04b412bb7a..efe5579c01 100644 --- a/examples/persistent_state/README.md +++ b/examples/persistent_state/README.md @@ -28,7 +28,7 @@ uv run python client_stdio.py ## Example output -``` +```text Each line below is a separate tool call Alice connects diff --git a/examples/persistent_state/client.py b/examples/persistent_state/client.py index ea4c3b9fa8..a5bb6368f8 100644 --- a/examples/persistent_state/client.py +++ b/examples/persistent_state/client.py @@ -18,7 +18,7 @@ console = Console() -async def main(): +async def main() -> None: console.print() console.print("[dim italic]Each line below is a separate tool call[/dim italic]") console.print() diff --git a/examples/persistent_state/client_stdio.py b/examples/persistent_state/client_stdio.py index b447bc16a7..5264a80cc8 100644 --- a/examples/persistent_state/client_stdio.py +++ b/examples/persistent_state/client_stdio.py @@ -10,21 +10,21 @@ from rich.console import Console -from fastmcp import Client +from fastmcp import Client, FastMCP # Add parent directory to path for importing the server module -examples_dir = Path(__file__).parent.parent.parent +examples_dir: Path = Path(__file__).parent.parent.parent if str(examples_dir) not in sys.path: sys.path.insert(0, str(examples_dir)) import examples.persistent_state.server as server_module # noqa: E402 -server = server_module.server +server: FastMCP = server_module.server -console = Console() +console: Console = Console() -async def main(): +async def main() -> None: console.print() console.print("[dim italic]Each line below is a separate tool call[/dim italic]") console.print() diff --git a/examples/persistent_state/server.py b/examples/persistent_state/server.py index f212d5b93c..df44cd94c3 100644 --- a/examples/persistent_state/server.py +++ b/examples/persistent_state/server.py @@ -30,7 +30,7 @@ async def get_value(key: str, ctx: Context) -> str: @server.tool -async def list_session_info(ctx: Context) -> dict: +async def list_session_info(ctx: Context) -> dict[str, str | None]: """Get information about the current session.""" return { "session_id": ctx.session_id, diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index dc5c8d9431..dd97ec6ec3 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -489,11 +489,11 @@ def store_data(data: dict, ctx: Context) -> str: if request: session_id = request.headers.get("mcp-session-id") - # Generate a session ID if it doesn't exist. + # For STDIO/SSE (no mcp-session-id header), use id(session). + # This ensures consistency with state set during on_initialize + # (which also uses id(session) since request_context isn't available). if session_id is None: - from uuid import uuid4 - - session_id = str(uuid4()) + session_id = str(id(session)) # Save the session id to the session attributes session._fastmcp_id = session_id # type: ignore[attr-defined] diff --git a/tests/server/test_context.py b/tests/server/test_context.py index 3013696480..5901f1c00a 100644 --- a/tests/server/test_context.py +++ b/tests/server/test_context.py @@ -58,29 +58,36 @@ def test_session_id_with_http_headers(self, context): ) ) - assert context.session_id == "test-session-123" - - request_ctx.reset(token) + try: + assert context.session_id == "test-session-123" + finally: + request_ctx.reset(token) def test_session_id_without_http_headers(self, context): - """Test that session_id returns a UUID string when no HTTP headers are available.""" - import uuid + """Test that session_id returns id(session) when no HTTP headers are available. + For STDIO/SSE transports, we use id(session) to ensure consistency with + state set during on_initialize (which also uses id(session)). + """ from mcp.server.lowlevel.server import request_ctx from mcp.shared.context import RequestContext + mock_session = MagicMock(wraps={}) token = request_ctx.set( RequestContext( request_id=0, meta=None, - session=MagicMock(wraps={}), + session=mock_session, lifespan_context=MagicMock(), ) ) - assert uuid.UUID(context.session_id) - - request_ctx.reset(token) + try: + # session_id should be str(id(session)) for non-HTTP transports + session_id = context.session_id + assert session_id == str(id(mock_session)) + finally: + request_ctx.reset(token) class TestContextState: @@ -219,10 +226,6 @@ async def store_and_read(value: str, ctx: Context) -> dict: assert session_id_2 is not None assert session_id_1 != session_id_2 - # Verify they look like UUIDs (36 chars with dashes) - assert len(session_id_1) == 36 - assert len(session_id_2) == 36 - # Client 1 reconnects and should still see their value async with Client(server) as client1_again: # But this is a NEW session (new connection = new session ID) From cb4cd0d8da0a74bfb2b41127b35e725ad3e441f3 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:26:15 -0500 Subject: [PATCH 5/9] Fix session ID to use UUID instead of id(session) Using id(session) for session IDs caused issues because: 1. Memory can be reused after garbage collection 2. For in-memory transport, sessions might be reused Now both session_id and _get_state_prefix() generate and cache a UUID on the session object (_fastmcp_state_prefix), ensuring unique IDs per logical session. --- src/fastmcp/server/context.py | 53 ++++++++++++++++++++++++----------- tests/server/test_context.py | 14 +++++---- 2 files changed, 45 insertions(+), 22 deletions(-) diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index dd97ec6ec3..62356e82bf 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -479,8 +479,8 @@ def store_data(data: dict, ctx: Context) -> str: ) session = request_ctx.session - # Try to get the session ID from the session attributes - session_id = getattr(session, "_fastmcp_id", None) + # Check for cached session ID (shared with _get_state_prefix for consistency) + session_id = getattr(session, "_fastmcp_state_prefix", None) if session_id is not None: return session_id @@ -489,14 +489,14 @@ def store_data(data: dict, ctx: Context) -> str: if request: session_id = request.headers.get("mcp-session-id") - # For STDIO/SSE (no mcp-session-id header), use id(session). - # This ensures consistency with state set during on_initialize - # (which also uses id(session) since request_context isn't available). + # For STDIO/SSE/in-memory, generate a UUID if session_id is None: - session_id = str(id(session)) + from uuid import uuid4 - # Save the session id to the session attributes - session._fastmcp_id = session_id # type: ignore[attr-defined] + session_id = str(uuid4()) + + # Cache for consistency with state prefix + session._fastmcp_state_prefix = session_id # type: ignore[attr-defined] return session_id @property @@ -1113,19 +1113,38 @@ async def elicit( def _get_state_prefix(self) -> str: """Get the prefix for state keys. - Uses session_id when available (consistent with the public API). - Falls back to id(session) during on_initialize when session_id - isn't available yet. + Generates a unique prefix per session and caches it. This works during + both on_initialize (when request_context is None) and tool calls. + For HTTP, uses the mcp-session-id header if available. """ - # When request_context is available, use session_id for consistency + from uuid import uuid4 + + # Get session from either source if self.request_context is not None: - return self.session_id + session = self.request_context.session + elif self._session is not None: + session = self._session + else: + raise RuntimeError("No session available for state operations") - # During on_initialize, fall back to id(session) - if self._session is not None: - return str(id(self._session)) + # Check for cached prefix (set during init or previous call) + prefix = getattr(session, "_fastmcp_state_prefix", None) + if prefix is not None: + return prefix - raise RuntimeError("No session available for state operations") + # For HTTP, try to get from header + if self.request_context is not None: + request = self.request_context.request + if request: + header_id = request.headers.get("mcp-session-id") + if header_id: + session._fastmcp_state_prefix = header_id # type: ignore[attr-defined] + return header_id + + # Generate new prefix (UUID) for STDIO/SSE/in-memory + prefix = str(uuid4()) + session._fastmcp_state_prefix = prefix # type: ignore[attr-defined] + return prefix def _make_state_key(self, key: str) -> str: """Create session-prefixed key for state storage.""" diff --git a/tests/server/test_context.py b/tests/server/test_context.py index 5901f1c00a..89c55fdc33 100644 --- a/tests/server/test_context.py +++ b/tests/server/test_context.py @@ -64,11 +64,13 @@ def test_session_id_with_http_headers(self, context): request_ctx.reset(token) def test_session_id_without_http_headers(self, context): - """Test that session_id returns id(session) when no HTTP headers are available. + """Test that session_id returns a UUID when no HTTP headers are available. - For STDIO/SSE transports, we use id(session) to ensure consistency with - state set during on_initialize (which also uses id(session)). + For STDIO/SSE/in-memory transports, we generate a UUID and cache it + on the session for consistency with state operations. """ + import uuid + from mcp.server.lowlevel.server import request_ctx from mcp.shared.context import RequestContext @@ -83,9 +85,11 @@ def test_session_id_without_http_headers(self, context): ) try: - # session_id should be str(id(session)) for non-HTTP transports + # session_id should be a valid UUID for non-HTTP transports session_id = context.session_id - assert session_id == str(id(mock_session)) + assert uuid.UUID(session_id) # Valid UUID format + # Should be cached on session + assert mock_session._fastmcp_state_prefix == session_id finally: request_ctx.reset(token) From 4b11ad7df4a330d5546fc7f9ba4cf508e0950d31 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:27:06 -0500 Subject: [PATCH 6/9] Update docstring and docs for init state behavior --- docs/servers/context.mdx | 2 +- src/fastmcp/server/context.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/servers/context.mdx b/docs/servers/context.mdx index 3c0f6161a2..fb1fbdb171 100644 --- a/docs/servers/context.mdx +++ b/docs/servers/context.mdx @@ -259,7 +259,7 @@ Any backend compatible with the [py-key-value-aio](https://github.com/strawgate/ #### State During Initialization -For STDIO and SSE transports, state set during `on_initialize` middleware persists to subsequent tool calls (same session object). For StreamableHTTP, state during initialization is isolated because the session ID comes from request headers which aren't available during init. +State set during `on_initialize` middleware persists to subsequent tool calls when using the same session object (STDIO, SSE, single-server HTTP). For distributed/serverless HTTP deployments where different machines handle init and tool calls, state is isolated by the `mcp-session-id` header. ### Change Notifications diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 62356e82bf..d2fc4f9294 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -167,10 +167,10 @@ async def my_tool(x: int, ctx: Context) -> str: the same MCP session. State is automatically keyed by session, ensuring isolation between different clients. - For STDIO and SSE transports, state set during `on_initialize` middleware - will persist to tool calls. For StreamableHTTP, state during `on_initialize` - is isolated because the session ID comes from request headers (not available - during init). + State set during `on_initialize` middleware will persist to subsequent tool + calls when using the same session object (STDIO, SSE, single-server HTTP). + For distributed/serverless HTTP deployments where different machines handle + the init and tool calls, state is isolated by the mcp-session-id header. 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. From eb76cac516b4bf38bf8fb076af93d44c0426c750 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:37:46 -0500 Subject: [PATCH 7/9] Fix session_id return type in docs (raises RuntimeError, not None) --- docs/servers/context.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/servers/context.mdx b/docs/servers/context.mdx index fb1fbdb171..6e78111813 100644 --- a/docs/servers/context.mdx +++ b/docs/servers/context.mdx @@ -347,7 +347,7 @@ async def request_info(ctx: Context) -> dict: - **`ctx.request_id -> str`**: Get the unique ID for the current MCP request - **`ctx.client_id -> str | None`**: Get the ID of the client making the request, if provided during initialization -- **`ctx.session_id -> str | None`**: Get the MCP session ID for session-based data sharing (HTTP transports only) +- **`ctx.session_id -> str`**: Get the MCP session ID for session-based data sharing. Raises `RuntimeError` if the MCP session is not yet established. #### Request Context Availability From b494d23fd1668e11103f94951ab109242039a4c9 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:42:24 -0500 Subject: [PATCH 8/9] Add session-scoped state to v3-features and upgrade-guide --- docs/development/upgrade-guide.mdx | 44 +++++++++++++++++++++ docs/development/v3-notes/v3-features.mdx | 47 +++++++++++++++++++++++ 2 files changed, 91 insertions(+) diff --git a/docs/development/upgrade-guide.mdx b/docs/development/upgrade-guide.mdx index c4cdfddb2e..a920385ec3 100644 --- a/docs/development/upgrade-guide.mdx +++ b/docs/development/upgrade-guide.mdx @@ -256,6 +256,50 @@ The environment variable for controlling the server banner has been renamed: This change reflects that the setting now applies to all server startup methods, not just the CLI. The banner is now suppressed when running `python server.py` directly, not just when using `fastmcp run`. +### Context State Methods Are Async + + +**Breaking Change:** `ctx.set_state()` and `ctx.get_state()` are now async methods. Synchronous calls will fail. + + +Context state has changed from request-scoped to session-scoped, persisting across multiple tool calls within the same MCP session. The methods are now async because they interact with a pluggable storage backend. + + +```python Before +@mcp.tool +def my_tool(ctx: Context) -> str: + ctx.set_state("key", "value") + value = ctx.get_state("key") + return value +``` + +```python After +@mcp.tool +async def my_tool(ctx: Context) -> str: + await ctx.set_state("key", "value") + value = await ctx.get_state("key") + return value +``` + + +**What changed:** +- State now persists across requests within a session (not just within a single request) +- Different clients have isolated state (keyed by session ID) +- State expires after 1 day to prevent unbounded memory growth +- New method: `await ctx.delete_state(key)` + +**Custom storage backends:** + +By default, state uses an in-memory store. For distributed deployments, provide a custom backend: + +```python +from key_value.aio.stores.redis import RedisStore + +mcp = FastMCP("server", session_state_store=RedisStore(...)) +``` + +See [Session State](/servers/context#session-state) for full documentation. + ## v2.14.0 ### OpenAPI Parser Promotion diff --git a/docs/development/v3-notes/v3-features.mdx b/docs/development/v3-notes/v3-features.mdx index ff39d8b3a8..c6257e6dbc 100644 --- a/docs/development/v3-notes/v3-features.mdx +++ b/docs/development/v3-notes/v3-features.mdx @@ -150,6 +150,37 @@ Documentation: `docs/servers/providers/transforms.mdx`, `docs/servers/visibility --- +## Session-Scoped State + +v3.0 changes context state from request-scoped to session-scoped. State now persists across multiple tool calls within the same MCP session. + +```python +@mcp.tool +async def increment_counter(ctx: Context) -> int: + count = await ctx.get_state("counter") or 0 + await ctx.set_state("counter", count + 1) + return count + 1 +``` + +State is automatically keyed by session ID, ensuring isolation between different clients. The implementation uses [pykeyvalue](https://github.com/strawgate/py-key-value) for pluggable storage backends: + +```python +from key_value.aio.stores.redis import RedisStore + +# Use Redis for distributed deployments +mcp = FastMCP("server", session_state_store=RedisStore(...)) +``` + +**Key details:** +- Methods are now async: `await ctx.get_state()`, `await ctx.set_state()`, `await ctx.delete_state()` +- State expires after 1 day (TTL) to prevent unbounded memory growth +- Works during `on_initialize` middleware when using the same session object +- For distributed HTTP, session identity comes from the `mcp-session-id` header + +Documentation: `docs/servers/context.mdx` + +--- + ## Visibility System Components can be dynamically enabled/disabled at runtime using the visibility system (`src/fastmcp/server/transforms/visibility.py`). @@ -537,3 +568,19 @@ See `docs/development/v3-notes/auth-provider-env-vars.mdx` for rationale. `FASTMCP_SHOW_CLI_BANNER` → `FASTMCP_SHOW_SERVER_BANNER` Now applies to all server startup methods, not just the CLI. + +### Context State Methods Are Async + +`ctx.set_state()` and `ctx.get_state()` are now async and session-scoped: + +```python +# v2.x +ctx.set_state("key", "value") +value = ctx.get_state("key") + +# v3.0 +await ctx.set_state("key", "value") +value = await ctx.get_state("key") +``` + +State now persists across requests within a session. See "Session-Scoped State" above. From 6ef6c740e8b9ecf79c908f2b57ea473aa4d083a0 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:58:14 -0500 Subject: [PATCH 9/9] Update context.py --- src/fastmcp/server/context.py | 75 ++++++++++------------------------- 1 file changed, 22 insertions(+), 53 deletions(-) diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index d2fc4f9294..16e240ff5a 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -177,6 +177,9 @@ async def my_tool(x: int, ctx: Context) -> str: """ + # Default TTL for session state: 1 day in seconds + _STATE_TTL_SECONDS: int = 86400 + def __init__(self, fastmcp: FastMCP, session: ServerSession | None = None): self._fastmcp: weakref.ref[FastMCP] = weakref.ref(fastmcp) self._session: ServerSession | None = session # For state ops during init @@ -460,7 +463,7 @@ def session_id(self) -> str: for other transports. Raises: - RuntimeError if MCP request context is not available. + RuntimeError if no session is available. Example: ```python @@ -471,31 +474,36 @@ def store_data(data: dict, ctx: Context) -> str: return f"Data stored for session {session_id}" ``` """ + from uuid import uuid4 + + # Get session from request context or _session (for on_initialize) request_ctx = self.request_context - if request_ctx is None: + if request_ctx is not None: + session = request_ctx.session + elif self._session is not None: + session = self._session + else: raise RuntimeError( - "session_id is not available because the MCP session has not been established yet. " - "Check `context.request_context` for None before accessing this attribute." + "session_id is not available because no session exists. " + "This typically means you're outside a request context." ) - session = request_ctx.session - # Check for cached session ID (shared with _get_state_prefix for consistency) + # Check for cached session ID session_id = getattr(session, "_fastmcp_state_prefix", None) if session_id is not None: return session_id - # Try to get the session ID from the http request headers - request = request_ctx.request - if request: - session_id = request.headers.get("mcp-session-id") + # For HTTP, try to get from header + if request_ctx is not None: + request = request_ctx.request + if request: + session_id = request.headers.get("mcp-session-id") # For STDIO/SSE/in-memory, generate a UUID if session_id is None: - from uuid import uuid4 - session_id = str(uuid4()) - # Cache for consistency with state prefix + # Cache on session for consistency session._fastmcp_state_prefix = session_id # type: ignore[attr-defined] return session_id @@ -1110,48 +1118,9 @@ async def elicit( else: raise ValueError(f"Unexpected elicitation action: {result.action}") - def _get_state_prefix(self) -> str: - """Get the prefix for state keys. - - Generates a unique prefix per session and caches it. This works during - both on_initialize (when request_context is None) and tool calls. - For HTTP, uses the mcp-session-id header if available. - """ - from uuid import uuid4 - - # Get session from either source - if self.request_context is not None: - session = self.request_context.session - elif self._session is not None: - session = self._session - else: - raise RuntimeError("No session available for state operations") - - # Check for cached prefix (set during init or previous call) - prefix = getattr(session, "_fastmcp_state_prefix", None) - if prefix is not None: - return prefix - - # For HTTP, try to get from header - if self.request_context is not None: - request = self.request_context.request - if request: - header_id = request.headers.get("mcp-session-id") - if header_id: - session._fastmcp_state_prefix = header_id # type: ignore[attr-defined] - return header_id - - # Generate new prefix (UUID) for STDIO/SSE/in-memory - prefix = str(uuid4()) - session._fastmcp_state_prefix = prefix # type: ignore[attr-defined] - return prefix - def _make_state_key(self, key: str) -> str: """Create session-prefixed key for state storage.""" - return f"{self._get_state_prefix()}:{key}" - - # Default TTL for session state: 1 day in seconds - _STATE_TTL_SECONDS: int = 86400 + return f"{self.session_id}:{key}" async def set_state(self, key: str, value: Any) -> None: """Set a value in the session-scoped state store.