Give inference requests an X-Request-ID

src/common/middleware.py

@@ -3,13 +3,54 @@
 """
 
 import logging
+import re
 import uuid
 from typing import List, Optional
 from fastapi import Request, status
 from common.http_client import request_id_ctx
 
 logger = logging.getLogger(__name__)
 
+# A caller's id is logged and echoed back, so only short, plain ids are kept.
+_REQUEST_ID = re.compile(r"[A-Za-z0-9._:-]{1,128}")
+
+
+class RequestIDMiddleware:
+    """Gives each HTTP request an X-Request-ID, the caller's or a new one.
+
+    Plain ASGI rather than BaseHTTPMiddleware, so streamed responses pass
+    through untouched. The id is set in request_id_ctx for logs and outgoing
+    internal calls, in request.state.request_id, and on the response.
+    """
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        sent = dict(scope["headers"]).get(b"x-request-id", b"").decode("latin-1")
+        request_id = sent if _REQUEST_ID.fullmatch(sent) else str(uuid.uuid4())
+        scope.setdefault("state", {})["request_id"] = request_id
+
+        async def send_with_id(message):
+            if message["type"] == "http.response.start":
+                headers = [
+                    (k, v) for k, v in message.get("headers", [])
+                    if k.lower() != b"x-request-id"
+                ]
+                headers.append((b"x-request-id", request_id.encode("latin-1")))
+                message = {**message, "headers": headers}
+            await send(message)
+
+        token = request_id_ctx.set(request_id)
+        try:
+            await self.app(scope, receive, send_with_id)
+        finally:
+            request_id_ctx.reset(token)
+
 
 def _route_path(request: Request) -> str:
     """Path relative to the ASGI root_path. Under a sub-app mount Starlette sets

src/inference/app.py

@@ -21,6 +21,7 @@
 from common.exception_handlers import register_exception_handlers
 from common.logger import setup_logging
 from common.app_setup import setup_cors, add_standard_health_routes
+from common.middleware import RequestIDMiddleware
 
 # Configure logging
 logger = setup_logging(
@@ -44,6 +45,9 @@
 
 setup_cors(app, os.getenv("ALLOWED_ORIGINS", ""), settings.is_development)
 
+# Added last so it runs first: every request, CORS preflights included, gets an id.
+app.add_middleware(RequestIDMiddleware)
+
 
 # Add standard / and /health routes
 add_standard_health_routes(

src/inference/client.py

@@ -8,6 +8,7 @@
 from typing import Any, Dict, List, Optional
 
 import httpx
+from common.http_client import request_id_ctx
 from inference.config import settings
 from fastapi import HTTPException
 from fastapi import status as http_status
@@ -94,6 +95,11 @@ def _get_headers(self, auth_token: Optional[str] = None) -> Dict[str, str]:
         if auth_token:
             headers["Authorization"] = auth_token
 
+        # Lets the gateway's logs be matched to the inference request.
+        request_id = request_id_ctx.get()
+        if request_id:
+            headers["X-Request-ID"] = request_id
+
         return headers
 
     async def check_quota(self, user_id: str, model: str) -> None:

src/tests/inference/test_request_id.py

@@ -0,0 +1,84 @@
+"""Every inference request carries an X-Request-ID through to logs and the gateway."""
+import uuid
+
+import httpx
+import pytest
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+
+from common.http_client import request_id_ctx
+from common.middleware import RequestIDMiddleware
+from inference.app import app as inference_app
+from inference.client import ApiGatewayClient
+
+
+def _app():
+    app = FastAPI()
+    app.add_middleware(RequestIDMiddleware)
+
+    @app.get("/seen")
+    async def seen(request: Request):
+        return {"ctx": request_id_ctx.get(), "state": request.state.request_id}
+
+    @app.get("/stream")
+    async def stream():
+        async def chunks():
+            for part in ("a", "b", "c"):
+                yield f"data: {part}\n\n"
+        return StreamingResponse(chunks(), media_type="text/event-stream")
+
+    return app
+
+
+async def _get(app, path, headers=None):
+    transport = httpx.ASGITransport(app=app)
+    async with httpx.AsyncClient(transport=transport, base_url="http://inf") as c:
+        return await c.get(path, headers=headers)
+
+
+@pytest.mark.asyncio
+async def test_a_request_without_an_id_gets_a_new_one():
+    resp = await _get(_app(), "/seen")
+    rid = resp.headers["x-request-id"]
+    uuid.UUID(rid)
+    assert resp.json() == {"ctx": rid, "state": rid}
+    assert request_id_ctx.get() is None
+
+
+@pytest.mark.asyncio
+async def test_the_callers_id_is_kept():
+    resp = await _get(_app(), "/seen", {"X-Request-ID": "trace-123"})
+    assert resp.headers["x-request-id"] == "trace-123"
+    assert resp.json()["ctx"] == "trace-123"
+
+
+@pytest.mark.parametrize("sent", ["x" * 129, "bad id", "a\x7fb"])
+@pytest.mark.asyncio
+async def test_an_unusable_callers_id_is_replaced(sent):
+    resp = await _get(_app(), "/seen", {"X-Request-ID": sent})
+    rid = resp.headers["x-request-id"]
+    assert rid != sent
+    uuid.UUID(rid)
+
+
+@pytest.mark.asyncio
+async def test_a_streamed_response_passes_through_with_the_id():
+    resp = await _get(_app(), "/stream", {"X-Request-ID": "trace-9"})
+    assert resp.headers["x-request-id"] == "trace-9"
+    assert resp.text == "data: a\n\ndata: b\n\ndata: c\n\n"
+
+
+@pytest.mark.asyncio
+async def test_the_inference_app_sets_the_id():
+    resp = await _get(inference_app, "/health", {"X-Request-ID": "trace-health"})
+    assert resp.headers["x-request-id"] == "trace-health"
+
+
+def test_calls_to_the_gateway_forward_the_id():
+    client = ApiGatewayClient()
+    assert "X-Request-ID" not in client._get_headers()
+    token = request_id_ctx.set("trace-7")
+    try:
+        assert client._get_headers()["X-Request-ID"] == "trace-7"
+    finally:
+        request_id_ctx.reset(token)