Address latest review feedback

- Fix docstring indentation: use 4-space continuation indent filled to
  120 cols consistently across all Args parameters
- Change stateless+idle_timeout error from ValueError to RuntimeError
- Remove unnecessary None guard on session_id in idle timeout cleanup
- Replace while+sleep(0) polling with anyio.Event in test

Github-Issue: #1283

Author: felixweinberger

src/mcp/server/streamable_http_manager.py

@@ -47,25 +47,20 @@ class StreamableHTTPSessionManager:
 
     Args:
         app: The MCP server instance
-        event_store: Optional event store for resumability support.
-                     If provided, enables resumable connections where clients
-                     can reconnect and receive missed events.
-                     If None, sessions are still tracked but not resumable.
+        event_store: Optional event store for resumability support. If provided, enables resumable connections
+            where clients can reconnect and receive missed events. If None, sessions are still tracked but not
+            resumable.
         json_response: Whether to use JSON responses instead of SSE streams
-        stateless: If True, creates a completely fresh transport for each request
-                   with no session tracking or state persistence between requests.
+        stateless: If True, creates a completely fresh transport for each request with no session tracking or
+            state persistence between requests.
         security_settings: Optional transport security settings.
-        retry_interval: Retry interval in milliseconds to suggest to clients in SSE
-                       retry field. Used for SSE 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.
+        retry_interval: Retry interval in milliseconds to suggest to clients in SSE retry field. Used for SSE
+            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__(
@@ -81,7 +76,7 @@ def __init__(
         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 ValueError("session_idle_timeout is not supported in stateless mode")
+            raise RuntimeError("session_idle_timeout is not supported in stateless mode")
 
         self.app = app
         self.event_store = event_store
@@ -248,10 +243,10 @@ async def run_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORE
                                 )
 
                             if idle_scope.cancelled_caught:
+                                assert http_transport.mcp_session_id is not None
                                 session_id = http_transport.mcp_session_id
                                 logger.info(f"Session {session_id} idle timeout")
-                                if session_id is not None:  # pragma: no branch
-                                    self._server_instances.pop(session_id, None)
+                                self._server_instances.pop(session_id, None)
                                 await http_transport.terminate()
                         except Exception:
                             logger.exception(f"Session {http_transport.mcp_session_id} crashed")

tests/server/test_streamable_http_manager.py

@@ -337,12 +337,24 @@ async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestP
 
 @pytest.mark.anyio
 async def test_idle_session_is_reaped():
-    """Idle timeout sets a cancel scope deadline and reaps the session when it fires."""
-    idle_timeout = 300
+    """After idle timeout fires, the session returns 404."""
     app = Server("test-idle-reap")
-    manager = StreamableHTTPSessionManager(app=app, session_idle_timeout=idle_timeout)
+
+    run_finished = anyio.Event()
+    original_run = app.run
+
+    async def tracked_run(*args: Any, **kwargs: Any) -> None:
+        try:
+            await original_run(*args, **kwargs)
+        finally:
+            run_finished.set()
+
+    app.run = tracked_run  # type: ignore[assignment]
+
+    manager = StreamableHTTPSessionManager(app=app, session_idle_timeout=300)
 
     async with manager.run():
+        # Create a session
         sent_messages: list[Message] = []
 
         async def mock_send(message: Message):
@@ -358,7 +370,6 @@ async def mock_send(message: Message):
         async def mock_receive():  # pragma: no cover
             return {"type": "http.request", "body": b"", "more_body": False}
 
-        before = anyio.current_time()
         await manager.handle_request(scope, mock_receive, mock_send)
 
         session_id = None
@@ -372,35 +383,54 @@ async def mock_receive():  # pragma: no cover
                     break
 
         assert session_id is not None, "Session ID not found in response headers"
-        assert session_id in manager._server_instances
 
-        # Verify the idle deadline was set correctly
+        # Force the idle deadline to expire
         transport = manager._server_instances[session_id]
         assert transport.idle_scope is not None
-        assert transport.idle_scope.deadline >= before + idle_timeout
-
-        # Simulate time passing by expiring the deadline
         transport.idle_scope.deadline = anyio.current_time()
 
+        # Wait for app.run to exit via the cancel scope, then one checkpoint for cleanup
         with anyio.fail_after(5):
-            while session_id in manager._server_instances:
-                await anyio.sleep(0)
-
-        assert session_id not in manager._server_instances
-
-        # Verify terminate() is idempotent
-        await transport.terminate()
-        assert transport.is_terminated
-
-
-@pytest.mark.parametrize(
-    "kwargs,match",
-    [
-        ({"session_idle_timeout": -1}, "positive number"),
-        ({"session_idle_timeout": 0}, "positive number"),
-        ({"session_idle_timeout": 30, "stateless": True}, "not supported in stateless"),
-    ],
-)
-def test_session_idle_timeout_validation(kwargs: dict[str, Any], match: str):
-    with pytest.raises(ValueError, match=match):
-        StreamableHTTPSessionManager(app=Server("test"), **kwargs)
+            await run_finished.wait()
+        await anyio.sleep(0)
+
+        # Verify session is gone via public API: request with old session ID returns 404
+        response_messages: list[Message] = []
+        response_body = b""
+
+        async def capture_send(message: Message):
+            nonlocal response_body
+            response_messages.append(message)
+            if message["type"] == "http.response.body":
+                response_body += message.get("body", b"")
+
+        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)