Refactor get_http_request and context.session_id (#1242)

This commit is contained in:
hopeful0 2025-07-23 22:05:46 +08:00 committed by GitHub
commit ca74b44d62
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 68 additions and 42 deletions

View file

@ -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:

View file

@ -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):

View file

@ -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: