mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-26 23:44:17 +02:00
Refactor get_http_request and context.session_id (#1242)
This commit is contained in:
parent
653fe0f89f
commit
ca74b44d62
3 changed files with 68 additions and 42 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue