DX-119753: clean up transport log header parsing

Author: aniket-s-kulkarni

src/dremioai/servers/mcp.py

@@ -23,6 +23,7 @@
 from uuid import uuid4
 from enum import StrEnum, auto
 from functools import reduce, wraps
+from http import HTTPStatus
 from json import dump as jdump
 from json import load
 from operator import ior
@@ -54,6 +55,7 @@
 from pydantic.networks import AnyUrl
 from rich import console, table
 from rich import print as pp
+from starlette.datastructures import Headers
 from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from starlette.requests import Request
@@ -78,6 +80,7 @@
 
 class MCPTransportLoggingMiddleware:
     logger = log.logger("MCPTransportLoggingMiddleware")
+    status_to_log = HTTPStatus.MULTIPLE_CHOICES  # log 300 and above
     max_error_body_log_bytes = 2048
 
     def __init__(self, app: ASGIApp):
@@ -102,7 +105,7 @@ async def send_wrapper(message: Message):
             elif (
                 message["type"] == "http.response.body"
                 and status_code is not None
-                and status_code >= 400
+                and status_code >= self.status_to_log
             ):
                 body = message.get("body", b"")
                 remaining = self.max_error_body_log_bytes - len(captured_body)
@@ -122,7 +125,7 @@ async def send_wrapper(message: Message):
             )
             raise
 
-        if status_code is not None and status_code >= 400:
+        if status_code is not None and status_code >= self.status_to_log:
             response_body = (
                 captured_body.decode("utf-8", errors="replace").strip()
                 if captured_body
@@ -133,20 +136,12 @@ async def send_wrapper(message: Message):
                 status_code=status_code,
                 response_body=response_body,
                 response_body_truncated=response_body_truncated,
-                response_mcp_session_id=self._header(
-                    response_headers, MCP_SESSION_ID_HEADER
+                response_mcp_session_id=Headers(raw=response_headers).get(
+                    MCP_SESSION_ID_HEADER
                 ),
                 **self._request_log_context(request),
             )
 
-    @staticmethod
-    def _header(headers: List[Tuple[bytes, bytes]], name: str) -> str | None:
-        expected = name.lower().encode()
-        for key, value in headers:
-            if key.lower() == expected:
-                return value.decode("utf-8", errors="replace")
-        return None
-
     @staticmethod
     def _request_log_context(request: Request) -> Dict[str, Any]:
         context = {

tests/servers/test_jwks_verifier.py

@@ -331,6 +331,41 @@ async def send(message):
         assert "203.0.113.10" in logged
         assert "secret-token" not in logged
 
+    @pytest.mark.asyncio
+    async def test_logs_redirects(self, caplog):
+        async def app(scope, receive, send):
+            response = Response(status_code=307, headers={"location": "/mcp/new"})
+            await response(scope, receive, send)
+
+        scope = {
+            "type": "http",
+            "method": "GET",
+            "path": "/mcp/project-123",
+            "headers": [],
+            "client": ("203.0.113.10", 4321),
+            "server": ("testserver", 80),
+            "scheme": "http",
+            "query_string": b"",
+        }
+
+        async def receive():
+            return {"type": "http.request", "body": b"", "more_body": False}
+
+        async def send(_message):
+            pass
+
+        middleware = MCPTransportLoggingMiddleware(app)
+        with caplog.at_level(logging.WARNING):
+            await middleware(scope, receive, send)
+
+        warning_messages = [
+            str(r.message) for r in caplog.records if r.levelno >= logging.WARNING
+        ]
+        assert any(
+            "MCP transport request failed" in msg and "307" in msg
+            for msg in warning_messages
+        )
+
 
 class TestStreamableHttpInit:
     def test_streamable_http_transport_is_stateless(self):