Add idle session timeout to StreamableHTTPSessionManager (#2022)
This commit is contained in:
parent
be5bb7c4f2
commit
92140e5086
4 changed files with 129 additions and 15 deletions
|
|
@ -31,6 +31,11 @@ This document contains critical information about working with this codebase. Fo
|
||||||
- IMPORTANT: Before pushing, verify 100% branch coverage on changed files by running
|
- IMPORTANT: Before pushing, verify 100% branch coverage on changed files by running
|
||||||
`uv run --frozen pytest -x` (coverage is configured in `pyproject.toml` with `fail_under = 100`
|
`uv run --frozen pytest -x` (coverage is configured in `pyproject.toml` with `fail_under = 100`
|
||||||
and `branch = true`). If any branch is uncovered, add a test for it before pushing.
|
and `branch = true`). If any branch is uncovered, add a test for it before pushing.
|
||||||
|
- Avoid `anyio.sleep()` with a fixed duration to wait for async operations. Instead:
|
||||||
|
- Use `anyio.Event` — set it in the callback/handler, `await event.wait()` in the test
|
||||||
|
- For stream messages, use `await stream.receive()` instead of `sleep()` + `receive_nowait()`
|
||||||
|
- Exception: `sleep()` is appropriate when testing time-based features (e.g., timeouts)
|
||||||
|
- Wrap indefinite waits (`event.wait()`, `stream.receive()`) in `anyio.fail_after(5)` to prevent hangs
|
||||||
|
|
||||||
Test files mirror the source tree: `src/mcp/client/streamable_http.py` → `tests/client/test_streamable_http.py`
|
Test files mirror the source tree: `src/mcp/client/streamable_http.py` → `tests/client/test_streamable_http.py`
|
||||||
Add tests to the existing file for that module.
|
Add tests to the existing file for that module.
|
||||||
|
|
|
||||||
|
|
@ -169,6 +169,8 @@ class StreamableHTTPServerTransport:
|
||||||
] = {}
|
] = {}
|
||||||
self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[dict[str, str]]] = {}
|
self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[dict[str, str]]] = {}
|
||||||
self._terminated = False
|
self._terminated = False
|
||||||
|
# Idle timeout cancel scope; managed by the session manager.
|
||||||
|
self.idle_scope: anyio.CancelScope | None = None
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_terminated(self) -> bool:
|
def is_terminated(self) -> bool:
|
||||||
|
|
|
||||||
|
|
@ -39,6 +39,7 @@ class StreamableHTTPSessionManager:
|
||||||
2. Resumability via an optional event store
|
2. Resumability via an optional event store
|
||||||
3. Connection management and lifecycle
|
3. Connection management and lifecycle
|
||||||
4. Request handling and transport setup
|
4. Request handling and transport setup
|
||||||
|
5. Idle session cleanup via optional timeout
|
||||||
|
|
||||||
Important: Only one StreamableHTTPSessionManager instance should be created
|
Important: Only one StreamableHTTPSessionManager instance should be created
|
||||||
per application. The instance cannot be reused after its run() context has
|
per application. The instance cannot be reused after its run() context has
|
||||||
|
|
@ -46,16 +47,20 @@ class StreamableHTTPSessionManager:
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
app: The MCP server instance
|
app: The MCP server instance
|
||||||
event_store: Optional event store for resumability support.
|
event_store: Optional event store for resumability support. If provided, enables resumable connections
|
||||||
If provided, enables resumable connections where clients
|
where clients can reconnect and receive missed events. If None, sessions are still tracked but not
|
||||||
can reconnect and receive missed events.
|
resumable.
|
||||||
If None, sessions are still tracked but not resumable.
|
|
||||||
json_response: Whether to use JSON responses instead of SSE streams
|
json_response: Whether to use JSON responses instead of SSE streams
|
||||||
stateless: If True, creates a completely fresh transport for each request
|
stateless: If True, creates a completely fresh transport for each request with no session tracking or
|
||||||
with no session tracking or state persistence between requests.
|
state persistence between requests.
|
||||||
security_settings: Optional transport security settings.
|
security_settings: Optional transport security settings.
|
||||||
retry_interval: Retry interval in milliseconds to suggest to clients in SSE
|
retry_interval: Retry interval in milliseconds to suggest to clients in SSE retry field. Used for SSE
|
||||||
retry field. Used for SSE polling behavior.
|
polling behavior.
|
||||||
|
session_idle_timeout: Optional idle timeout in seconds for stateful sessions. If set, sessions that
|
||||||
|
receive no HTTP requests for this duration will be automatically terminated and removed. When
|
||||||
|
retry_interval is also configured, ensure the idle timeout comfortably exceeds the retry interval to
|
||||||
|
avoid reaping sessions during normal SSE polling gaps. Default is None (no timeout). A value of 1800
|
||||||
|
(30 minutes) is recommended for most deployments.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
|
|
@ -66,13 +71,20 @@ class StreamableHTTPSessionManager:
|
||||||
stateless: bool = False,
|
stateless: bool = False,
|
||||||
security_settings: TransportSecuritySettings | None = None,
|
security_settings: TransportSecuritySettings | None = None,
|
||||||
retry_interval: int | None = None,
|
retry_interval: int | None = None,
|
||||||
|
session_idle_timeout: float | None = None,
|
||||||
):
|
):
|
||||||
|
if session_idle_timeout is not None and session_idle_timeout <= 0:
|
||||||
|
raise ValueError("session_idle_timeout must be a positive number of seconds")
|
||||||
|
if stateless and session_idle_timeout is not None:
|
||||||
|
raise RuntimeError("session_idle_timeout is not supported in stateless mode")
|
||||||
|
|
||||||
self.app = app
|
self.app = app
|
||||||
self.event_store = event_store
|
self.event_store = event_store
|
||||||
self.json_response = json_response
|
self.json_response = json_response
|
||||||
self.stateless = stateless
|
self.stateless = stateless
|
||||||
self.security_settings = security_settings
|
self.security_settings = security_settings
|
||||||
self.retry_interval = retry_interval
|
self.retry_interval = retry_interval
|
||||||
|
self.session_idle_timeout = session_idle_timeout
|
||||||
|
|
||||||
# Session tracking (only used if not stateless)
|
# Session tracking (only used if not stateless)
|
||||||
self._session_creation_lock = anyio.Lock()
|
self._session_creation_lock = anyio.Lock()
|
||||||
|
|
@ -184,6 +196,9 @@ class StreamableHTTPSessionManager:
|
||||||
if request_mcp_session_id is not None and request_mcp_session_id in self._server_instances:
|
if request_mcp_session_id is not None and request_mcp_session_id in self._server_instances:
|
||||||
transport = self._server_instances[request_mcp_session_id]
|
transport = self._server_instances[request_mcp_session_id]
|
||||||
logger.debug("Session already exists, handling request directly")
|
logger.debug("Session already exists, handling request directly")
|
||||||
|
# Push back idle deadline on activity
|
||||||
|
if transport.idle_scope is not None and self.session_idle_timeout is not None:
|
||||||
|
transport.idle_scope.deadline = anyio.current_time() + self.session_idle_timeout # pragma: no cover
|
||||||
await transport.handle_request(scope, receive, send)
|
await transport.handle_request(scope, receive, send)
|
||||||
return
|
return
|
||||||
|
|
||||||
|
|
@ -210,16 +225,31 @@ class StreamableHTTPSessionManager:
|
||||||
read_stream, write_stream = streams
|
read_stream, write_stream = streams
|
||||||
task_status.started()
|
task_status.started()
|
||||||
try:
|
try:
|
||||||
await self.app.run(
|
# Use a cancel scope for idle timeout — when the
|
||||||
read_stream,
|
# deadline passes the scope cancels app.run() and
|
||||||
write_stream,
|
# execution continues after the ``with`` block.
|
||||||
self.app.create_initialization_options(),
|
# Incoming requests push the deadline forward.
|
||||||
stateless=False, # Stateful mode
|
idle_scope = anyio.CancelScope()
|
||||||
)
|
if self.session_idle_timeout is not None:
|
||||||
|
idle_scope.deadline = anyio.current_time() + self.session_idle_timeout
|
||||||
|
http_transport.idle_scope = idle_scope
|
||||||
|
|
||||||
|
with idle_scope:
|
||||||
|
await self.app.run(
|
||||||
|
read_stream,
|
||||||
|
write_stream,
|
||||||
|
self.app.create_initialization_options(),
|
||||||
|
stateless=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
if idle_scope.cancelled_caught:
|
||||||
|
assert http_transport.mcp_session_id is not None
|
||||||
|
logger.info(f"Session {http_transport.mcp_session_id} idle timeout")
|
||||||
|
self._server_instances.pop(http_transport.mcp_session_id, None)
|
||||||
|
await http_transport.terminate()
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception(f"Session {http_transport.mcp_session_id} crashed")
|
logger.exception(f"Session {http_transport.mcp_session_id} crashed")
|
||||||
finally:
|
finally:
|
||||||
# Only remove from instances if not terminated
|
|
||||||
if ( # pragma: no branch
|
if ( # pragma: no branch
|
||||||
http_transport.mcp_session_id
|
http_transport.mcp_session_id
|
||||||
and http_transport.mcp_session_id in self._server_instances
|
and http_transport.mcp_session_id in self._server_instances
|
||||||
|
|
|
||||||
|
|
@ -333,3 +333,80 @@ async def test_e2e_streamable_http_server_cleanup():
|
||||||
Client(streamable_http_client(f"http://{host}/mcp", http_client=http_client)) as client,
|
Client(streamable_http_client(f"http://{host}/mcp", http_client=http_client)) as client,
|
||||||
):
|
):
|
||||||
await client.list_tools()
|
await client.list_tools()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_idle_session_is_reaped():
|
||||||
|
"""After idle timeout fires, the session returns 404."""
|
||||||
|
app = Server("test-idle-reap")
|
||||||
|
manager = StreamableHTTPSessionManager(app=app, session_idle_timeout=0.05)
|
||||||
|
|
||||||
|
async with manager.run():
|
||||||
|
sent_messages: list[Message] = []
|
||||||
|
|
||||||
|
async def mock_send(message: Message):
|
||||||
|
sent_messages.append(message)
|
||||||
|
|
||||||
|
scope = {
|
||||||
|
"type": "http",
|
||||||
|
"method": "POST",
|
||||||
|
"path": "/mcp",
|
||||||
|
"headers": [(b"content-type", b"application/json")],
|
||||||
|
}
|
||||||
|
|
||||||
|
async def mock_receive(): # pragma: no cover
|
||||||
|
return {"type": "http.request", "body": b"", "more_body": False}
|
||||||
|
|
||||||
|
await manager.handle_request(scope, mock_receive, mock_send)
|
||||||
|
|
||||||
|
session_id = None
|
||||||
|
for msg in sent_messages: # pragma: no branch
|
||||||
|
if msg["type"] == "http.response.start": # pragma: no branch
|
||||||
|
for header_name, header_value in msg.get("headers", []): # pragma: no branch
|
||||||
|
if header_name.decode().lower() == MCP_SESSION_ID_HEADER.lower():
|
||||||
|
session_id = header_value.decode()
|
||||||
|
break
|
||||||
|
if session_id: # pragma: no branch
|
||||||
|
break
|
||||||
|
|
||||||
|
assert session_id is not None, "Session ID not found in response headers"
|
||||||
|
|
||||||
|
# Wait for the 50ms idle timeout to fire and cleanup to complete
|
||||||
|
await anyio.sleep(0.1)
|
||||||
|
|
||||||
|
# Verify via public API: old session ID now returns 404
|
||||||
|
response_messages: list[Message] = []
|
||||||
|
|
||||||
|
async def capture_send(message: Message):
|
||||||
|
response_messages.append(message)
|
||||||
|
|
||||||
|
scope_with_session = {
|
||||||
|
"type": "http",
|
||||||
|
"method": "POST",
|
||||||
|
"path": "/mcp",
|
||||||
|
"headers": [
|
||||||
|
(b"content-type", b"application/json"),
|
||||||
|
(b"mcp-session-id", session_id.encode()),
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
await manager.handle_request(scope_with_session, mock_receive, capture_send)
|
||||||
|
|
||||||
|
response_start = next(
|
||||||
|
(msg for msg in response_messages if msg["type"] == "http.response.start"),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
assert response_start is not None
|
||||||
|
assert response_start["status"] == 404
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_idle_timeout_rejects_non_positive():
|
||||||
|
with pytest.raises(ValueError, match="positive number"):
|
||||||
|
StreamableHTTPSessionManager(app=Server("test"), session_idle_timeout=-1)
|
||||||
|
with pytest.raises(ValueError, match="positive number"):
|
||||||
|
StreamableHTTPSessionManager(app=Server("test"), session_idle_timeout=0)
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_idle_timeout_rejects_stateless():
|
||||||
|
with pytest.raises(RuntimeError, match="not supported in stateless"):
|
||||||
|
StreamableHTTPSessionManager(app=Server("test"), session_idle_timeout=30, stateless=True)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue