fix(latency): strip case-variant spend-logs spoofs, stamp selection on pre-built responses

Author: eric-liu-nvidia

switchyard/lib/backends/latency_service_llm_backend.py

@@ -368,6 +368,15 @@ def _pick_by_latency(
     # -- Request processing (hot path — no Latency Service call) ------------
 
     async def call(self, ctx: ProxyContext, request: ChatRequest) -> ChatResponse:
+        """Serve *request* on the best healthy endpoint, failing over on error.
+
+        Candidates are ordered by health and observed latency (with optional
+        session affinity); every upstream attempt is stamped with the
+        spend-logs metadata header, and the attempt that succeeds records its
+        route selection on *ctx* (see :data:`CTX_ROUTE_SELECTION`). Exhausting
+        all candidates re-raises the last upstream error, annotated for the
+        client-facing error envelope.
+        """
         # This backend records its own per-attempt ``outcome_metrics`` counters
         # below (one per failover attempt). Claim attempt accounting for this
         # request so the endpoint-layer fallback in ``dispatch`` /
@@ -675,21 +684,33 @@ async def _call_endpoint(
         body: dict[str, object],
         extra_headers: dict[str, str],
     ) -> object:
-        # ``extra_headers`` rides the client wrapper's ``**kwargs`` into the
-        # OpenAI SDK's per-request header merge (over ``default_headers``).
-        # A passthrough body may itself carry an SDK-style ``extra_headers``
-        # field (shape-preserving translation keeps unknown keys); merge it
-        # under ours rather than letting it collide with the keyword — it used
-        # to ride ``**body`` into the same SDK parameter, and the spend-logs
-        # header must win a name conflict so callers can't spoof it. A
-        # non-mapping value is dropped: it could only ever have broken the
-        # SDK call.
+        """Make one upstream SDK call with the protected *extra_headers* stamped on.
+
+        ``extra_headers`` rides the client wrapper's ``**kwargs`` into the
+        OpenAI SDK's per-request header merge (over ``default_headers``). A
+        passthrough body may itself carry an SDK-style ``extra_headers`` field
+        (shape-preserving translation keeps unknown keys); it is still
+        forwarded — it used to ride ``**body`` into the same SDK parameter —
+        but the protected headers must win the merge so callers can't spoof
+        them, and header names are case-insensitive on the wire while dict
+        keys are not, so every case variant of a protected name is stripped
+        from the client mapping (a client-cased duplicate would otherwise
+        reach the wire alongside ours and win first-match parsing on the
+        receiving proxy). A non-mapping value is dropped: it could only ever
+        have broken the SDK call.
+        """
         client_extra = body.pop("extra_headers", None)
+        protected = {name.lower() for name in extra_headers}
         headers: dict[str, object] = (
-            {**client_extra, **extra_headers}
+            {
+                name: value
+                for name, value in client_extra.items()
+                if not (isinstance(name, str) and name.lower() in protected)
+            }
             if isinstance(client_extra, Mapping)
-            else dict(extra_headers)
+            else {}
         )
+        headers.update(extra_headers)
         if target_request_type == ChatRequestType.OPENAI_RESPONSES:
             return await self._clients[model_id].aresponses(
                 api_key=api_key,

switchyard/lib/endpoints/anthropic_messages_endpoint.py

@@ -74,6 +74,7 @@ async def anthropic_messages(
             request: Request,
             body: Annotated[dict[str, Any], Body(...)],
         ) -> Response:
+            """Anthropic-compatible Messages endpoint."""
             obj = request.app.state.switchyard
             _strip_unsupported_output_config(body)
             model = str(body.get("model", "<none>"))

switchyard/lib/endpoints/dispatch.py

@@ -101,16 +101,22 @@ def serialize_chain_result(
 ) -> Response:
     """Serialize a chain result to the appropriate HTTP response.
 
-    Returns the result as-is if it is already a ``Response``, wraps it in a
+    Returns the result itself if it is already a ``Response``, wraps it in a
     ``StreamingResponse`` when streaming is requested, or JSON-serializes it.
     Any route selection recorded on *ctx* is stamped as ``x-switchyard-*``
-    response headers (streaming included — the backend call completed before
-    the response object is built, so the selection is final). ``ctx`` is
-    required so a new endpoint cannot silently opt out of spend attribution.
+    response headers on every branch, pre-built responses included (streaming
+    too — the backend call completed before the response object is built, so
+    the selection is final). ``ctx`` is required so a new endpoint cannot
+    silently opt out of spend attribution.
     """
+    headers = route_selection_headers(ctx)
     if isinstance(result, Response):
+        # Merge rather than pass through untouched: no current chain path
+        # yields a pre-built Response after a billed upstream success, but if
+        # one ever does, dropping the recorded selection here would silently
+        # break spend attribution.
+        result.headers.update(headers)
         return result
-    headers = route_selection_headers(ctx)
     if stream and hasattr(result, "__aiter__"):
         return StreamingResponse(
             sse_iter(result), media_type="text/event-stream", headers=headers

switchyard/lib/endpoints/openai_chat_endpoint.py

@@ -73,6 +73,7 @@ async def chat_completions(
             request: Request,
             body: Annotated[dict[str, Any], Body(...)],
         ) -> Response:
+            """OpenAI-compatible Chat Completions endpoint."""
             obj = request.app.state.switchyard
             chat_request = ChatRequest.openai_chat(body)
             # Reject semantically invalid input (e.g. empty messages) at the

switchyard/lib/endpoints/responses_endpoint.py

@@ -57,6 +57,7 @@ async def responses(
             request: Request,
             body: Annotated[dict[str, Any], Body(...)],
         ) -> Response:
+            """OpenAI-compatible Responses endpoint."""
             obj = request.app.state.switchyard
             model = str(body.get("model", "<none>"))
             stream = bool(body.get("stream"))

tests/_chain_test_helpers.py

@@ -151,6 +151,7 @@ class Handler(BaseHTTPRequestHandler):
             protocol_version = "HTTP/1.1"
 
             def do_POST(self) -> None:
+                """Record the request (path, body, headers) and pop the queued reply."""
                 length = int(self.headers.get("content-length", "0"))
                 raw = self.rfile.read(length)
                 body = json.loads(raw.decode("utf-8"))

tests/test_latency_service_llm_backend.py

@@ -1012,6 +1012,7 @@ class TestRouteSelectionSpendLogs:
     """
 
     async def test_outbound_header_carries_route_selection(self):
+        """The outbound header records the full selection for the chosen endpoint."""
         config = LatencyServiceBackendConfig(
             latency_service_url=LATENCY_SERVICE_URL,
             endpoints=[
@@ -1040,6 +1041,7 @@ async def test_outbound_header_carries_route_selection(self):
         uuid.UUID(str(payload["router_correlation_id"]))
 
     async def test_ctx_records_selection_matching_outbound_header(self):
+        """ctx records exactly the payload the wire header carried."""
         backend = _make_backend(_config("model-A"))
         backend._clients["model-A"].acompletion = AsyncMock(
             return_value=_make_completion()
@@ -1055,6 +1057,7 @@ async def test_ctx_records_selection_matching_outbound_header(self):
         assert payload["router_selected_provider"] == "model-A"
 
     async def test_failover_restamps_selection_and_keeps_correlation_id(self):
+        """Each attempt is stamped with its own endpoint under one correlation id."""
         backend = _make_backend(_config("model-A", "model-B"))
         _set_health(
             backend,
@@ -1082,6 +1085,7 @@ async def test_failover_restamps_selection_and_keeps_correlation_id(self):
         assert ctx.metadata[CTX_ROUTE_SELECTION] == second
 
     async def test_correlation_id_is_fresh_per_request(self):
+        """Two client requests never share a correlation id."""
         backend = _make_backend(_config("model-A"))
         mock = AsyncMock(return_value=_make_completion())
         backend._clients["model-A"].acompletion = mock
@@ -1094,6 +1098,7 @@ async def test_correlation_id_is_fresh_per_request(self):
         assert first["router_correlation_id"] != second["router_correlation_id"]
 
     async def test_router_model_falls_back_to_configured_route_model(self):
+        """A model-less client body attributes to the configured route id."""
         backend = _make_backend(
             _config("model-A", route_model="nvidia/switchyard/gpt-5.5")
         )
@@ -1108,6 +1113,7 @@ async def test_router_model_falls_back_to_configured_route_model(self):
         assert payload["router_model"] == "nvidia/switchyard/gpt-5.5"
 
     async def test_no_selection_recorded_when_all_attempts_fail(self):
+        """No billed success → no selection recorded on ctx."""
         backend = _make_backend(_config("model-A"))
         backend._clients["model-A"].acompletion = AsyncMock(
             side_effect=_api_status_error(401),
@@ -1120,6 +1126,7 @@ async def test_no_selection_recorded_when_all_attempts_fail(self):
         assert CTX_ROUTE_SELECTION not in ctx.metadata
 
     async def test_responses_surface_is_stamped_too(self):
+        """The Responses-API surface is stamped like the Chat surface."""
         backend = _make_backend(_config("model-A", request_type="openai_responses"))
         backend._clients["model-A"].aresponses = AsyncMock(
             return_value={"id": "resp-test", "object": "response", "output": []}
@@ -1134,6 +1141,34 @@ async def test_responses_surface_is_stamped_too(self):
         assert payload["router_selected_endpoint"] == "model-A"
         assert payload["router_strategy"] == "latency"
 
+    async def test_cased_spoof_of_spend_logs_header_is_stripped(self):
+        """A differently-cased spoof key cannot ride client extra_headers to the wire.
+
+        Header names are case-insensitive on the wire while dict merges are
+        not: exactly one instance of the protected header — ours — may reach
+        the SDK, while benign client headers still pass through.
+        """
+        backend = _make_backend(_config("model-A"))
+        mock = AsyncMock(return_value=_make_completion())
+        backend._clients["model-A"].acompletion = mock
+
+        await backend.call(
+            ProxyContext(),
+            _openai_request(
+                extra_headers={
+                    "X-LiteLLM-Spend-Logs-Metadata": "spoofed",
+                    "x-client-tag": "42",
+                },
+            ),
+        )
+
+        headers = mock.call_args.kwargs["extra_headers"]
+        spoof_keys = [k for k in headers if k.lower() == SPEND_LOGS_METADATA_HEADER]
+        assert spoof_keys == [SPEND_LOGS_METADATA_HEADER]
+        payload = _spend_logs_payload(mock)
+        assert payload["router_selected_endpoint"] == "model-A"
+        assert headers["x-client-tag"] == "42"
+
 
 # ---------------------------------------------------------------------------
 # Credential policy