test: cover stale session middleware edges

Author: dialupdisaster

src/godot_ai/asgi.py

@@ -85,9 +85,6 @@ async def _send_response(
         body: bytes,
         send: Send,
     ) -> None:
-        if start_message is None:
-            return
-
         rewritten = self._rewrite_stale_session_body(start_message, body)
         response_body = rewritten if rewritten is not None else body
         headers = start_message.get("headers", [])
@@ -99,8 +96,6 @@ async def _send_response(
         await send({"type": "http.response.body", "body": response_body, "more_body": False})
 
     def _rewrite_stale_session_body(self, start_message: Message, body: bytes) -> bytes | None:
-        if start_message.get("status") != HTTPStatus.NOT_FOUND:
-            return None
         try:
             payload = json.loads(body.decode("utf-8"))
         except (UnicodeDecodeError, json.JSONDecodeError):

tests/unit/test_asgi_session_diagnostics.py

@@ -28,6 +28,19 @@ async def send(message):
     return sent
 
 
+async def _single_asgi_request(app, *, scope):
+    sent = []
+
+    async def receive():
+        return {"type": "websocket.disconnect"}
+
+    async def send(message):
+        sent.append(message)
+
+    await app(scope, receive, send)
+    return sent
+
+
 @pytest.mark.anyio
 async def test_stale_mcp_session_diagnostic_rewrites_sdk_session_not_found():
     async def sdk_stale_session_response(scope, receive, send):
@@ -76,6 +89,36 @@ async def sdk_stale_session_response(scope, receive, send):
     }
 
 
+@pytest.mark.anyio
+async def test_stale_mcp_session_diagnostic_handles_chunked_stale_session_body():
+    async def sdk_stale_session_response(scope, receive, send):
+        await send({"type": "http.response.start", "status": 404, "headers": []})
+        await send(
+            {
+                "type": "http.response.body",
+                "body": b'{"jsonrpc":"2.0","id":"server-error",',
+                "more_body": True,
+            }
+        )
+        await send(
+            {
+                "type": "http.response.body",
+                "body": b'"error":{"code":-32600,"message":"Session not found"}}',
+                "more_body": False,
+            }
+        )
+
+    app = StaleMcpSessionDiagnosticMiddleware(sdk_stale_session_response)
+
+    sent = await _single_http_request(app)
+
+    assert sent[0]["status"] == 404
+    assert (b"content-type", b"application/json") in sent[0]["headers"]
+    body = json.loads(sent[1]["body"])
+    assert body["error"]["message"] == STALE_MCP_SESSION_MESSAGE
+    assert body["error"]["data"]["action"] == "reinitialize_mcp_session"
+
+
 @pytest.mark.anyio
 async def test_stale_mcp_session_diagnostic_leaves_other_responses_unchanged():
     async def ok_response(scope, receive, send):
@@ -130,6 +173,40 @@ async def streaming_response(scope, receive, send):
     ]
 
 
+@pytest.mark.anyio
+async def test_stale_mcp_session_diagnostic_passes_non_http_scopes_through():
+    async def websocket_response(scope, receive, send):
+        await send({"type": "websocket.close", "code": 1000})
+
+    app = StaleMcpSessionDiagnosticMiddleware(websocket_response)
+
+    sent = await _single_asgi_request(app, scope={"type": "websocket", "path": "/ws"})
+
+    assert sent == [{"type": "websocket.close", "code": 1000}]
+
+
+@pytest.mark.anyio
+async def test_stale_mcp_session_diagnostic_passes_unhandled_asgi_messages_through():
+    async def extension_message_response(scope, receive, send):
+        await send(
+            {
+                "type": "http.response.debug",
+                "info": {"note": "kept for downstream middleware"},
+            }
+        )
+
+    app = StaleMcpSessionDiagnosticMiddleware(extension_message_response)
+
+    sent = await _single_http_request(app)
+
+    assert sent == [
+        {
+            "type": "http.response.debug",
+            "info": {"note": "kept for downstream middleware"},
+        }
+    ]
+
+
 @pytest.mark.anyio
 async def test_stale_mcp_session_diagnostic_leaves_other_404_responses_unchanged():
     async def not_found_response(scope, receive, send):
@@ -156,6 +233,34 @@ async def not_found_response(scope, receive, send):
     ]
 
 
+@pytest.mark.anyio
+async def test_stale_mcp_session_diagnostic_leaves_other_json_rpc_404_errors_unchanged():
+    async def json_rpc_not_found_response(scope, receive, send):
+        body = json.dumps(
+            {
+                "jsonrpc": "2.0",
+                "id": "server-error",
+                "error": {"code": -32000, "message": "Tool not found"},
+            }
+        ).encode()
+        await send(
+            {
+                "type": "http.response.start",
+                "status": 404,
+                "headers": [(b"content-type", b"application/json")],
+            }
+        )
+        await send({"type": "http.response.body", "body": body, "more_body": False})
+
+    app = StaleMcpSessionDiagnosticMiddleware(json_rpc_not_found_response)
+
+    sent = await _single_http_request(app)
+
+    body = json.loads(sent[1]["body"])
+    assert body["error"] == {"code": -32000, "message": "Tool not found"}
+    assert "data" not in body["error"]
+
+
 def test_create_server_wraps_streamable_http_app_with_stale_session_diagnostic():
     server = create_server()
 
@@ -164,6 +269,14 @@ def test_create_server_wraps_streamable_http_app_with_stale_session_diagnostic()
     assert isinstance(app, StaleMcpSessionDiagnosticMiddleware)
 
 
+def test_create_server_does_not_wrap_sse_app_with_stale_session_diagnostic():
+    server = create_server()
+
+    app = server.http_app(transport="sse")
+
+    assert not isinstance(app, StaleMcpSessionDiagnosticMiddleware)
+
+
 def test_stale_mcp_session_diagnostic_preserves_fastmcp_app_state():
     server = create_server()