From 03d316cfe832d47ac219a8b031791b09f8bd13c3 Mon Sep 17 00:00:00 2001 From: Priyanshu-u07 Date: Sun, 25 Oct 2026 17:28:06 +0530 Subject: [PATCH] Give inference requests an X-Request-ID The inference app never set a request id, so its logs and error responses had none and its calls to the gateway could not be matched to a request. Add a plain ASGI request-id middleware to common (it leaves streamed responses untouched), install it on the inference app, and forward the id in ApiGatewayClient. A caller's id is kept only if it is short and plain. Signed-off-by: Priyanshu-u07 --- src/common/middleware.py | 41 +++++++++++++ src/inference/app.py | 4 ++ src/inference/client.py | 6 ++ src/tests/inference/test_request_id.py | 84 ++++++++++++++++++++++++++ 4 files changed, 135 insertions(+) create mode 100644 src/tests/inference/test_request_id.py diff --git a/src/common/middleware.py b/src/common/middleware.py index 974bd9ce..67164384 100644 --- a/src/common/middleware.py +++ b/src/common/middleware.py @@ -3,6 +3,7 @@ """ import logging +import re import uuid from typing import List, Optional from fastapi import Request, status @@ -10,6 +11,46 @@ 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 diff --git a/src/inference/app.py b/src/inference/app.py index 6b72aeb8..b6e8f665 100644 --- a/src/inference/app.py +++ b/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( diff --git a/src/inference/client.py b/src/inference/client.py index 1a16e131..851a8b07 100644 --- a/src/inference/client.py +++ b/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: diff --git a/src/tests/inference/test_request_id.py b/src/tests/inference/test_request_id.py new file mode 100644 index 00000000..aad8ced3 --- /dev/null +++ b/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)