From ca74b44d6277fe7c2617a6d32ba62f1680574153 Mon Sep 17 00:00:00 2001 From: hopeful0 Date: Wed, 23 Jul 2025 22:05:46 +0800 Subject: [PATCH] Refactor `get_http_request` and `context.session_id` (#1242) --- src/fastmcp/server/context.py | 45 ++++++++++++++++--------- src/fastmcp/server/dependencies.py | 11 ++++-- tests/server/test_context.py | 54 +++++++++++++++++------------- 3 files changed, 68 insertions(+), 42 deletions(-) diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 42ff32e56..5af3597b7 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -130,7 +130,7 @@ class Context: _current_context.reset(token) @property - def request_context(self) -> RequestContext: + def request_context(self) -> RequestContext[ServerSession, Any, Request]: """Access to the underlying request context. If called outside of a request context, this will raise a ValueError. @@ -217,35 +217,48 @@ class Context: return str(self.request_context.request_id) @property - def session_id(self) -> str | None: - """Get the MCP session ID for HTTP transports. + def session_id(self) -> str: + """Get the MCP session ID for ALL transports. Returns the session ID that can be used as a key for session-based data storage (e.g., Redis) to share data between tool calls within the same client session. Returns: - The session ID for HTTP transports (SSE, StreamableHTTP), or None - for stdio and in-memory transports which don't use session IDs. + The session ID for StreamableHTTP transports, or a generated ID + for other transports. Example: ```python @server.tool def store_data(data: dict, ctx: Context) -> str: - if session_id := ctx.session_id: - redis_client.set(f"session:{session_id}:data", json.dumps(data)) - return f"Data stored for session {session_id}" - return "No session ID available (stdio/memory transport)" + session_id = ctx.session_id + redis_client.set(f"session:{session_id}:data", json.dumps(data)) + return f"Data stored for session {session_id}" ``` """ - try: - from fastmcp.server.dependencies import get_http_headers + request_ctx = self.request_context + session = request_ctx.session - headers = get_http_headers(include_all=True) - return headers.get("mcp-session-id") - except RuntimeError: - # No HTTP context available (stdio/in-memory transport) - return None + # Try to get the session ID from the session attributes + session_id = getattr(session, "_fastmcp_id", 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") + + # Generate a session ID if it doesn't exist. + if session_id is None: + from uuid import uuid4 + + session_id = str(uuid4()) + + # Save the session id to the session attributes + setattr(session, "_fastmcp_id", session_id) + return session_id @property def session(self) -> ServerSession: diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py index 08775225e..7b288573b 100644 --- a/src/fastmcp/server/dependencies.py +++ b/src/fastmcp/server/dependencies.py @@ -37,9 +37,14 @@ def get_context() -> Context: def get_http_request() -> Request: - from fastmcp.server.http import _current_http_request + from mcp.server.lowlevel.server import request_ctx + + request = None + try: + request = request_ctx.get().request + except LookupError: + pass - request = _current_http_request.get() if request is None: raise RuntimeError("No active HTTP request found.") return request @@ -72,6 +77,8 @@ def get_http_headers(include_all: bool = False) -> dict[str, str]: "proxy-authenticate", "proxy-authorization", "proxy-connection", + # MCP-related headers + "mcp-session-id", } # (just in case) if not all(h.lower() == h for h in exclude_headers): diff --git a/tests/server/test_context.py b/tests/server/test_context.py index 6a8f4d1c9..17fc13f83 100644 --- a/tests/server/test_context.py +++ b/tests/server/test_context.py @@ -91,38 +91,44 @@ class TestParseModelPreferences: class TestSessionId: def test_session_id_with_http_headers(self, context): """Test that session_id returns the value from mcp-session-id header.""" + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext + mock_headers = {"mcp-session-id": "test-session-123"} - with patch( - "fastmcp.server.dependencies.get_http_headers", return_value=mock_headers - ): - assert context.session_id == "test-session-123" + token = request_ctx.set( + RequestContext( + request_id=0, + meta=None, + session=MagicMock(wraps={}), + lifespan_context=MagicMock(), + request=MagicMock(headers=mock_headers), + ) + ) + + assert context.session_id == "test-session-123" + + request_ctx.reset(token) def test_session_id_without_http_headers(self, context): - """Test that session_id returns None when no HTTP headers are available.""" - with patch( - "fastmcp.server.dependencies.get_http_headers", - side_effect=RuntimeError("No active HTTP request found."), - ): - assert context.session_id is None + """Test that session_id returns a UUID string when no HTTP headers are available.""" + import uuid - def test_session_id_with_missing_header(self, context): - """Test that session_id returns None when mcp-session-id header is missing.""" - mock_headers = {"other-header": "value"} + from mcp.server.lowlevel.server import request_ctx + from mcp.shared.context import RequestContext - with patch( - "fastmcp.server.dependencies.get_http_headers", return_value=mock_headers - ): - assert context.session_id is None + token = request_ctx.set( + RequestContext( + request_id=0, + meta=None, + session=MagicMock(wraps={}), + lifespan_context=MagicMock(), + ) + ) - def test_session_id_with_empty_header(self, context): - """Test that session_id returns None when mcp-session-id header is empty.""" - mock_headers = {"mcp-session-id": ""} + assert uuid.UUID(context.session_id) - with patch( - "fastmcp.server.dependencies.get_http_headers", return_value=mock_headers - ): - assert context.session_id == "" # Empty string is still returned as-is + request_ctx.reset(token) class TestContextState: