diff --git a/.agents/skills/switchyard-lib-core/SKILL.md b/.agents/skills/switchyard-lib-core/SKILL.md index de8b3b05..17f7409d 100644 --- a/.agents/skills/switchyard-lib-core/SKILL.md +++ b/.agents/skills/switchyard-lib-core/SKILL.md @@ -43,6 +43,7 @@ right validation set. If the change is driven by a launcher need, also read | Stats / telemetry | Reuse `StatsRequestProcessor`, `StatsResponseProcessor`, `StatsLlmBackend`, and `StatsAccumulator`. A profile config should thread one accumulator through all three when stats are enabled. Do not write a parallel collector. | | A fixed-path endpoint contributed by per-route components | Set `Endpoint.register_once = True`; `build_switchyard_app(...)` mounts the first instance while still running every component's lifecycle. Leave the default `False` for configurable endpoint classes that may mount distinct instances. | | Per-endpoint attribution on `/metrics` for a Python backend that can't be wrapped by `StatsLlmBackend` | Set `ctx.selected_model = endpoint_id` before returning the response. Also set `ctx.backend_call_latency_ms = upstream_call_ms` so the response processor can compute routing overhead. `LatencyServiceLLMBackend.call` is the reference. | +| Spend/tokenomics attribution when a front proxy (LiteLLM) sits above Switchyard | Stamp each upstream attempt with the `x-litellm-spend-logs-metadata` request header (selection re-stamped per failover attempt, one `router_correlation_id` uuid per client request) and record the successful attempt's selection in `ctx.metadata[CTX_ROUTE_SELECTION]`; the endpoint layer maps it to the `x-switchyard-*` response headers via `route_selection_headers` in `lib/endpoints/route_selection.py` — stamped by `dispatch.serialize_chain_result` on success and by `upstream_error.handle_chain_exception` when a failure follows a billed upstream success. Header values are sanitized there (client-controlled `router_model` must stay legal header material). `LatencyServiceLLMBackend.call` is the reference; the payload contract lives on `CTX_ROUTE_SELECTION` in `switchyard/lib/proxy_context.py`. | | State metrics on `/metrics` | Register a `PrometheusEmitter` via `switchyard.lib.endpoints.prometheus_emitter.register(...)` and unregister on `shutdown()`. This is for backend-owned state, not request-flow counters. | | Error-rate / retry-recovery counters | Use `switchyard.lib.endpoints.outcome_metrics`. FastAPI middleware records client outcomes. A retrying Python backend records each upstream attempt itself (and `record_retry_recovered()`) and sets `CTX_UPSTREAM_ATTEMPTS_RECORDED` so the endpoint skips its fallback — `LatencyServiceLLMBackend.call` is the reference. Single-attempt backends (Rust native / passthrough / multi) record nothing themselves; the endpoint fallback (`record_upstream_attempt_success` / `record_upstream_attempt_failure` in `upstream_error.py`, called from `dispatch_chat_request` and `handle_chain_exception`) counts their one attempt. Don't add a `model` label — these counters are layer-aggregate. Keep labels bounded. For a per-model error breakdown, emit a backend-owned counter via the `PrometheusEmitter` instead (labels bounded to config-derived ids: the route id `config.route_model` + endpoint ids, else the `other` sentinel) — `LatencyServiceLLMBackend._render_prometheus_lines` exposes `switchyard_latency_upstream_attempts_total{requested_model,upstream_model,outcome,code}` next to the aggregate, leaving the shared counter model-free. | | Per-event error log | Use `switchyard.lib.endpoints.upstream_error_log.log_upstream_attempt_failure(...)` on the failure path. Events belong in logs/traces, not Prometheus sample timestamps. | diff --git a/docs/internal/latency_service_routing.md b/docs/internal/latency_service_routing.md index 9df19dfb..9df6ffac 100644 --- a/docs/internal/latency_service_routing.md +++ b/docs/internal/latency_service_routing.md @@ -431,6 +431,52 @@ front of Switchyard (e.g. LiteLLM) should tag its own failures the same way and propagate these headers from below. Mid-stream failures after HTTP 200 has been committed cannot carry headers and are not annotated today. +## Route-selection spend attribution + +For tokenomics reporting, a LiteLLM front proxy needs to tie its provider spend-log +rows back to the Switchyard logical route that selected them. The backend stamps +**every upstream attempt** with a spend-logs metadata request header (LiteLLM copies +its JSON value into the provider spend-log DB row): + +```text +x-litellm-spend-logs-metadata: {"router_model": "", + "router_strategy": "latency", + "router_selected_endpoint": "", + "router_selected_model": "", + "router_selected_provider": "", + "router_correlation_id": ""} +``` + +One correlation id is generated per client request and shared by every failover +attempt; the selection fields are re-stamped per attempt, so each provider row +records the endpoint it actually hit. `router_model` is the model the client asked +for, falling back to the configured `route_model` when the body carried no model. + +**Successful** responses then return the selection that served the request, so the +front proxy can enrich its own (parent) spend-log row with the same correlation id +the provider row received: + +| Header | Value | +|---|---| +| `x-switchyard-router-model` | `router_model` from the selection | +| `x-switchyard-selected-model` | `router_selected_model` | +| `x-switchyard-selected-provider` | `router_selected_provider` | +| `x-switchyard-router-correlation-id` | `router_correlation_id` | + +Streaming responses carry the same headers (the upstream call completes before the +SSE response is committed). Error responses carry the failure-source headers above; +the selection headers join them only when an upstream call already **succeeded** +before the failure (e.g. a response-translation error after a billed 200), so the +provider row's correlation id stays joinable — a request that never reached a +successful upstream claims no selection. Because the client controls `router_model`, +header values are re-validated as header material: a value that is not +latin-1-encodable or contains control characters is omitted rather than breaking +(or splitting) the response; the outbound JSON header is immune (`json.dumps` +escapes it) and still records the raw value. The cross-component contract is +`CTX_ROUTE_SELECTION` in `switchyard/lib/proxy_context.py` (written by the backend, +read by the endpoint layer via `switchyard/lib/endpoints/route_selection.py`); +wire-level proofs live in `tests/test_route_selection_headers.py`. + ## Observability When `enable_stats` is `true` (the default), `LatencyServiceProfileConfig` wires a diff --git a/switchyard/lib/backends/latency_service_llm_backend.py b/switchyard/lib/backends/latency_service_llm_backend.py index fa8f573c..04c59b5a 100644 --- a/switchyard/lib/backends/latency_service_llm_backend.py +++ b/switchyard/lib/backends/latency_service_llm_backend.py @@ -20,10 +20,13 @@ Responses-mode endpoints receive the OpenAI Responses API natively. """ +import json import logging import random import threading import time +import uuid +from collections.abc import Mapping from dataclasses import dataclass from openai import APIStatusError, AsyncStream @@ -46,6 +49,7 @@ from switchyard.lib.proxy_context import ( CTX_CALLER_API_KEY, CTX_ERROR_SOURCE, + CTX_ROUTE_SELECTION, CTX_UPSTREAM_ATTEMPTS_RECORDED, CTX_UPSTREAM_HTTP_BODY, CTX_UPSTREAM_HTTP_STATUS, @@ -63,6 +67,14 @@ log = logging.getLogger(__name__) +#: Request header read by a LiteLLM front proxy; its JSON value is copied into +#: the provider spend-log DB row, tying that row back to the Switchyard route +#: that selected it (see :data:`CTX_ROUTE_SELECTION` for the payload contract). +SPEND_LOGS_METADATA_HEADER = "x-litellm-spend-logs-metadata" + +#: ``router_strategy`` value identifying this backend's selection algorithm. +ROUTER_STRATEGY_LATENCY = "latency" + @dataclass(frozen=True) class _RouteDecision: @@ -380,6 +392,10 @@ async def call(self, ctx: ProxyContext, request: ChatRequest) -> ChatResponse: # Captured before the per-attempt ``body["model"]`` override so the span # records the model the client asked for, not the selected endpoint. incoming_model = request.model + # One correlation id per client request, shared by every failover + # attempt's spend-logs header and by the response headers — the join + # key between a front proxy's spend-log row and the provider rows. + correlation_id = str(uuid.uuid4()) # Resolve the session-affinity pin once (keyed on the stable conversation # prefix). ``None`` when affinity is disabled or the conversation isn't # pinned yet, so every attempt routes purely by health + latency. @@ -426,6 +442,15 @@ async def call(self, ctx: ProxyContext, request: ChatRequest) -> ChatResponse: target_request_type = self._request_types[model_id] body = self._body_for_endpoint_request_type(ctx, request, target_request_type) body["model"] = upstream_model + # Per-attempt route-selection record: each upstream call is stamped + # with the endpoint it actually hits, so a failover retry's + # spend-log row doesn't inherit the failed attempt's selection. + route_selection = self._route_selection( + incoming_model=incoming_model, + model_id=model_id, + upstream_model=upstream_model, + correlation_id=correlation_id, + ) log.debug( "LatencyServiceLLMBackend: attempt=%d model=%s upstream=%s " "request_type=%s stream=%s", @@ -450,6 +475,9 @@ async def call(self, ctx: ProxyContext, request: ChatRequest) -> ChatResponse: target_request_type, api_key=api_key_override, body=body, + extra_headers={ + SPEND_LOGS_METADATA_HEADER: json.dumps(route_selection), + }, ) except APIStatusError as exc: set_tags(attempt_span, { @@ -547,6 +575,12 @@ async def call(self, ctx: ProxyContext, request: ChatRequest) -> ChatResponse: # Rust ``StatsLlmBackend`` that normally publishes this signal. ctx.backend_call_latency_ms = backend_latency_ms + # The selection that actually served the request, surfaced to + # the endpoint layer as ``x-switchyard-*`` response headers so + # a front proxy can enrich its own spend-log row with the same + # correlation id the provider row received. + ctx.metadata[CTX_ROUTE_SELECTION] = route_selection + # Pin this conversation to the endpoint that served it so later # turns reuse it (warm cache). Re-pinning on every success also # follows a recovery: if the previous pin degraded and we @@ -605,6 +639,33 @@ def _body_for_endpoint_request_type( ) return dict(normalized.body) + def _route_selection( + self, + *, + incoming_model: str | None, + model_id: str, + upstream_model: str, + correlation_id: str, + ) -> dict[str, str | None]: + """Build the route-selection payload for one upstream attempt. + + This is the JSON value of the outbound ``x-litellm-spend-logs-metadata`` + header and, for the attempt that succeeds, the + :data:`CTX_ROUTE_SELECTION` record behind the ``x-switchyard-*`` + response headers. ``router_model`` falls back to the configured route id + when the client body carried no model; ``router_selected_provider`` is + the upstream model's leading path segment (IH/LiteLLM naming, e.g. + ``"openai/openai/gpt-5.4"`` → ``"openai"``). + """ + return { + "router_model": incoming_model or self._config.route_model, + "router_strategy": ROUTER_STRATEGY_LATENCY, + "router_selected_endpoint": model_id, + "router_selected_model": upstream_model, + "router_selected_provider": upstream_model.split("/", 1)[0], + "router_correlation_id": correlation_id, + } + async def _call_endpoint( self, model_id: str, @@ -612,14 +673,32 @@ async def _call_endpoint( *, api_key: str | None, 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. + client_extra = body.pop("extra_headers", None) + headers: dict[str, object] = ( + {**client_extra, **extra_headers} + if isinstance(client_extra, Mapping) + else dict(extra_headers) + ) if target_request_type == ChatRequestType.OPENAI_RESPONSES: return await self._clients[model_id].aresponses( api_key=api_key, + extra_headers=headers, **body, ) return await self._clients[model_id].acompletion( api_key=api_key, + extra_headers=headers, **body, ) diff --git a/switchyard/lib/endpoints/anthropic_messages_endpoint.py b/switchyard/lib/endpoints/anthropic_messages_endpoint.py index beb3a858..f484e401 100644 --- a/switchyard/lib/endpoints/anthropic_messages_endpoint.py +++ b/switchyard/lib/endpoints/anthropic_messages_endpoint.py @@ -106,7 +106,9 @@ async def anthropic_messages( stream, type(result).__name__, ) - return serialize_chain_result(result, stream=stream, sse_iter=iter_anthropic_sse) + return serialize_chain_result( + result, stream=stream, sse_iter=iter_anthropic_sse, ctx=ctx + ) except (SwitchyardContextPoolExhaustedError, SwitchyardContextWindowExceededError) as exc: return context_exhausted_response(exc, inbound="anthropic") except Exception as exc: diff --git a/switchyard/lib/endpoints/dispatch.py b/switchyard/lib/endpoints/dispatch.py index bdac8cd5..ce348660 100644 --- a/switchyard/lib/endpoints/dispatch.py +++ b/switchyard/lib/endpoints/dispatch.py @@ -9,6 +9,7 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse from switchyard.lib.endpoints.error_envelope import error_response +from switchyard.lib.endpoints.route_selection import route_selection_headers from switchyard.lib.endpoints.upstream_error import record_upstream_attempt_success from switchyard.lib.proxy_context import ProxyContext from switchyard.lib.roles import TranslatedResponse @@ -96,16 +97,24 @@ def serialize_chain_result( *, stream: bool, sse_iter: Callable[[Any], AsyncIterator[str]], + ctx: ProxyContext, ) -> 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 ``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. """ if isinstance(result, Response): return result + headers = route_selection_headers(ctx) if stream and hasattr(result, "__aiter__"): - return StreamingResponse(sse_iter(result), media_type="text/event-stream") + return StreamingResponse( + sse_iter(result), media_type="text/event-stream", headers=headers + ) if hasattr(result, "model_dump"): - return JSONResponse(content=result.model_dump()) - return JSONResponse(content=result) + return JSONResponse(content=result.model_dump(), headers=headers) + return JSONResponse(content=result, headers=headers) diff --git a/switchyard/lib/endpoints/openai_chat_endpoint.py b/switchyard/lib/endpoints/openai_chat_endpoint.py index 18cef631..e6f75eec 100644 --- a/switchyard/lib/endpoints/openai_chat_endpoint.py +++ b/switchyard/lib/endpoints/openai_chat_endpoint.py @@ -89,7 +89,7 @@ async def chat_completions( try: result: Any = await dispatch_chat_request(obj, chat_request, ctx) return serialize_chain_result( - result, stream=stream, sse_iter=iter_chat_completion_sse + result, stream=stream, sse_iter=iter_chat_completion_sse, ctx=ctx ) except (SwitchyardContextPoolExhaustedError, SwitchyardContextWindowExceededError) as exc: return context_exhausted_response(exc, inbound="openai") diff --git a/switchyard/lib/endpoints/responses_endpoint.py b/switchyard/lib/endpoints/responses_endpoint.py index b9a9a3be..cdc23519 100644 --- a/switchyard/lib/endpoints/responses_endpoint.py +++ b/switchyard/lib/endpoints/responses_endpoint.py @@ -84,7 +84,9 @@ async def responses( stream, type(result).__name__, ) - return serialize_chain_result(result, stream=stream, sse_iter=iter_preframed_sse) + return serialize_chain_result( + result, stream=stream, sse_iter=iter_preframed_sse, ctx=ctx + ) except (SwitchyardContextPoolExhaustedError, SwitchyardContextWindowExceededError) as exc: return context_exhausted_response(exc, inbound="openai-responses") except Exception as exc: diff --git a/switchyard/lib/endpoints/route_selection.py b/switchyard/lib/endpoints/route_selection.py new file mode 100644 index 00000000..20911053 --- /dev/null +++ b/switchyard/lib/endpoints/route_selection.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Route-selection response headers for spend/tokenomics attribution. + +Maps the :data:`CTX_ROUTE_SELECTION` record a routing backend stored on +``ctx`` (see :mod:`switchyard.lib.proxy_context`) to the ``x-switchyard-*`` +response headers a front proxy such as LiteLLM copies into its parent +spend-log row. Shared by the success serializer (``dispatch``) and the +error path (``upstream_error``) — a failure that happens *after* a billed +upstream success must still expose the selection, or the provider spend-log +row's correlation id becomes unjoinable. +""" + +from collections.abc import Mapping + +from switchyard.lib.proxy_context import CTX_ROUTE_SELECTION, ProxyContext + +#: Response headers exposing the route selection behind an upstream call, +#: carrying the same correlation id the provider row received via the +#: outbound ``x-litellm-spend-logs-metadata`` header. +ROUTER_MODEL_HEADER = "x-switchyard-router-model" +SELECTED_MODEL_HEADER = "x-switchyard-selected-model" +SELECTED_PROVIDER_HEADER = "x-switchyard-selected-provider" +ROUTER_CORRELATION_ID_HEADER = "x-switchyard-router-correlation-id" + +_ROUTE_SELECTION_RESPONSE_HEADERS = ( + (ROUTER_MODEL_HEADER, "router_model"), + (SELECTED_MODEL_HEADER, "router_selected_model"), + (SELECTED_PROVIDER_HEADER, "router_selected_provider"), + (ROUTER_CORRELATION_ID_HEADER, "router_correlation_id"), +) + + +def _is_header_value_safe(value: str) -> bool: + """Whether *value* can be emitted as an HTTP/1.1 response-header value. + + ``router_model`` echoes the client-supplied model string, so it must be + re-validated as header material: Starlette encodes response-header values + as latin-1 (a non-encodable value would fail response construction after + the upstream call already succeeded and was billed), and CTL characters — + CR/LF above all — would be a response-splitting vector on permissive + ASGI stacks. + """ + try: + value.encode("latin-1") + except UnicodeEncodeError: + return False + return not any(ord(ch) < 0x20 or ord(ch) == 0x7F for ch in value) + + +def route_selection_headers(ctx: ProxyContext) -> dict[str, str]: + """Response headers for the route selection recorded on *ctx*, if any. + + Empty when no routing backend recorded a selection (passthrough chains, + failures before any upstream success). A recorded field that is absent or + not emittable as a header value is skipped — headers never carry + placeholder or unsafe values. + """ + selection = ctx.metadata.get(CTX_ROUTE_SELECTION) + if not isinstance(selection, Mapping): + return {} + headers: dict[str, str] = {} + for header_name, selection_key in _ROUTE_SELECTION_RESPONSE_HEADERS: + value = selection.get(selection_key) + if isinstance(value, str) and value and _is_header_value_safe(value): + headers[header_name] = value + return headers diff --git a/switchyard/lib/endpoints/upstream_error.py b/switchyard/lib/endpoints/upstream_error.py index d51b39d4..285eda79 100644 --- a/switchyard/lib/endpoints/upstream_error.py +++ b/switchyard/lib/endpoints/upstream_error.py @@ -23,6 +23,7 @@ error_response, upstream_error_response, ) +from switchyard.lib.endpoints.route_selection import route_selection_headers from switchyard.lib.proxy_context import ( CTX_ERROR_SOURCE, CTX_UPSTREAM_ATTEMPTS_RECORDED, @@ -197,14 +198,23 @@ def handle_chain_exception( record_upstream_attempt_failure(ctx, exc) upstream = upstream_response_from_ctx(ctx, inbound=inbound, exc=exc) if upstream is not None: - return upstream - _log.error(log_msg, exc_info=exc) - return internal_chain_error_response( - exc, - inbound=inbound, - error_source=_ctx_error_source(ctx, default=ERROR_SOURCE_SWITCHYARD), - upstream_model=_ctx_upstream_model(ctx), - ) + response = upstream + else: + _log.error(log_msg, exc_info=exc) + response = internal_chain_error_response( + exc, + inbound=inbound, + error_source=_ctx_error_source(ctx, default=ERROR_SOURCE_SWITCHYARD), + upstream_model=_ctx_upstream_model(ctx), + ) + # A selection on ctx means an upstream call already SUCCEEDED (and was + # billed, with the spend-logs header stamped) before this failure — e.g. + # response-side translation rejected the 200 payload. Surface the + # selection headers on the error response too, so the front proxy can + # still join its (failed) parent spend-log row to the billed provider row. + for header_name, value in route_selection_headers(ctx).items(): + response.headers[header_name] = value + return response def context_exhausted_response(exc: BaseException, inbound: Inbound) -> JSONResponse: diff --git a/switchyard/lib/proxy_context.py b/switchyard/lib/proxy_context.py index 34bd5c29..19ae2779 100644 --- a/switchyard/lib/proxy_context.py +++ b/switchyard/lib/proxy_context.py @@ -89,6 +89,19 @@ #: ``x-switchyard-upstream-model`` response header. CTX_UPSTREAM_MODEL = "_upstream_model" +#: Route-selection record for spend/tokenomics attribution, written by a +#: routing backend after a successful upstream call. A dict with the keys +#: ``router_model`` (client-facing route id), ``router_strategy``, +#: ``router_selected_endpoint``, ``router_selected_model``, +#: ``router_selected_provider``, and ``router_correlation_id``. The same +#: payload is stamped on the outbound upstream call as the +#: ``x-litellm-spend-logs-metadata`` header (one per attempt, so provider +#: spend-log rows record the endpoint they actually hit); the endpoint layer +#: reads this key to return the ``x-switchyard-*`` route-selection response +#: headers, letting a front proxy enrich its own spend-log row with the same +#: correlation id. +CTX_ROUTE_SELECTION = "_route_selection" + __all__ = [ "CTX_CALLER_API_KEY", @@ -97,6 +110,7 @@ "CTX_ORIGINAL_MODEL", "CTX_ORIGINAL_REQUEST", "CTX_PROXY_ACTUAL_MODEL", + "CTX_ROUTE_SELECTION", "CTX_ROUTING", "CTX_TARGET_FORMAT", "CTX_UPSTREAM_ATTEMPTS_RECORDED", diff --git a/tests/_chain_test_helpers.py b/tests/_chain_test_helpers.py index 3e91687b..55202e93 100644 --- a/tests/_chain_test_helpers.py +++ b/tests/_chain_test_helpers.py @@ -155,7 +155,13 @@ def do_POST(self) -> None: raw = self.rfile.read(length) body = json.loads(raw.decode("utf-8")) with owner._lock: - owner._requests.append({"path": self.path, "body": body}) + # Header names lower-cased so tests can look them up + # without caring how the client cased them on the wire. + owner._requests.append({ + "path": self.path, + "body": body, + "headers": {k.lower(): v for k, v in self.headers.items()}, + }) if owner._responses: status, content, content_type = owner._responses.pop(0) else: @@ -174,7 +180,13 @@ def log_message(self, _format: str, *args: object) -> None: return None self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + # ``shutdown()`` blocks up to serve_forever's poll interval; the 0.5s + # default adds half a second of idle teardown to every test using the + # stub, so poll frequently. + server = self._server + self._thread = threading.Thread( + target=lambda: server.serve_forever(poll_interval=0.05), daemon=True + ) self._thread.start() return self diff --git a/tests/test_latency_service_llm_backend.py b/tests/test_latency_service_llm_backend.py index 6b6d0c40..230263bb 100644 --- a/tests/test_latency_service_llm_backend.py +++ b/tests/test_latency_service_llm_backend.py @@ -3,9 +3,11 @@ """Unit tests for :class:`LatencyServiceLLMBackend` (usage case).""" +import json import random import threading import time +import uuid from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -22,6 +24,7 @@ HealthPoller, ) from switchyard.lib.backends.latency_service_llm_backend import ( + SPEND_LOGS_METADATA_HEADER, LatencyServiceLLMBackend, ) from switchyard.lib.config.latency_service_backend_config import ( @@ -30,6 +33,7 @@ ) from switchyard.lib.proxy_context import ( CTX_ERROR_SOURCE, + CTX_ROUTE_SELECTION, CTX_UPSTREAM_HTTP_STATUS, CTX_UPSTREAM_MODEL, ProxyContext, @@ -983,6 +987,154 @@ async def test_upstream_model_default_falls_back_to_model(self): assert call_kwargs["model"] == "openai/gpt-5.5" +# --------------------------------------------------------------------------- +# Route-selection spend-logs metadata (tokenomics attribution) +# --------------------------------------------------------------------------- + + +def _spend_logs_payload(mock: AsyncMock, call_index: int = -1) -> dict[str, object]: + """Parse the spend-logs metadata header stamped on one recorded SDK call.""" + call_kwargs = mock.call_args_list[call_index].kwargs + header_value = call_kwargs["extra_headers"][SPEND_LOGS_METADATA_HEADER] + payload = json.loads(header_value) + assert isinstance(payload, dict) + return payload + + +class TestRouteSelectionSpendLogs: + """Every upstream attempt is stamped with route-selection metadata. + + The ``x-litellm-spend-logs-metadata`` request header ties a LiteLLM + provider spend-log row back to the Switchyard route that selected it, and + the successful attempt's selection is recorded on + ``ctx.metadata[CTX_ROUTE_SELECTION]`` so the endpoint layer can return the + matching ``x-switchyard-*`` response headers. + """ + + async def test_outbound_header_carries_route_selection(self): + config = LatencyServiceBackendConfig( + latency_service_url=LATENCY_SERVICE_URL, + endpoints=[ + LatencyServiceEndpoint( + model="openai/gpt-5.5", + upstream_model="openai/openai/gpt-5.5", + base_url="https://inference-api.test/v1", + api_key="k", + ), + ], + ) + backend = _make_backend(config) + backend._clients["openai/gpt-5.5"].acompletion = AsyncMock( + return_value=_make_completion() + ) + + await backend.call(ProxyContext(), _openai_request(model="nvidia/switchyard/gpt-5.5")) + + payload = _spend_logs_payload(backend._clients["openai/gpt-5.5"].acompletion) + assert payload["router_model"] == "nvidia/switchyard/gpt-5.5" + assert payload["router_strategy"] == "latency" + assert payload["router_selected_endpoint"] == "openai/gpt-5.5" + assert payload["router_selected_model"] == "openai/openai/gpt-5.5" + assert payload["router_selected_provider"] == "openai" + # Must parse as a UUID — the join key between spend-log rows. + uuid.UUID(str(payload["router_correlation_id"])) + + async def test_ctx_records_selection_matching_outbound_header(self): + backend = _make_backend(_config("model-A")) + backend._clients["model-A"].acompletion = AsyncMock( + return_value=_make_completion() + ) + + ctx = ProxyContext() + await backend.call(ctx, _openai_request()) + + payload = _spend_logs_payload(backend._clients["model-A"].acompletion) + assert ctx.metadata[CTX_ROUTE_SELECTION] == payload + # Endpoint id without a provider prefix: the provider segment + # degenerates to the whole id (split on the first "/"). + assert payload["router_selected_provider"] == "model-A" + + async def test_failover_restamps_selection_and_keeps_correlation_id(self): + backend = _make_backend(_config("model-A", "model-B")) + _set_health( + backend, + { + "model-A": EndpointHealthStatus.HEALTHY, + "model-B": EndpointHealthStatus.DEGRADED, + }, + ) + failing = AsyncMock(side_effect=_api_status_error(500)) + succeeding = AsyncMock(return_value=_make_completion()) + backend._clients["model-A"].acompletion = failing + backend._clients["model-B"].acompletion = succeeding + + ctx = ProxyContext() + await backend.call(ctx, _openai_request()) + + first = _spend_logs_payload(failing) + second = _spend_logs_payload(succeeding) + # Each attempt records the endpoint it actually hit… + assert first["router_selected_endpoint"] == "model-A" + assert second["router_selected_endpoint"] == "model-B" + # …while sharing one correlation id per client request. + assert first["router_correlation_id"] == second["router_correlation_id"] + # ctx carries the selection that served the request, not the failure. + assert ctx.metadata[CTX_ROUTE_SELECTION] == second + + async def test_correlation_id_is_fresh_per_request(self): + backend = _make_backend(_config("model-A")) + mock = AsyncMock(return_value=_make_completion()) + backend._clients["model-A"].acompletion = mock + + await backend.call(ProxyContext(), _openai_request()) + await backend.call(ProxyContext(), _openai_request()) + + first = _spend_logs_payload(mock, call_index=0) + second = _spend_logs_payload(mock, call_index=1) + assert first["router_correlation_id"] != second["router_correlation_id"] + + async def test_router_model_falls_back_to_configured_route_model(self): + backend = _make_backend( + _config("model-A", route_model="nvidia/switchyard/gpt-5.5") + ) + backend._clients["model-A"].acompletion = AsyncMock( + return_value=_make_completion() + ) + + body: dict = {"messages": [{"role": "user", "content": "hi"}]} + await backend.call(ProxyContext(), ChatRequest.openai_chat(body)) + + payload = _spend_logs_payload(backend._clients["model-A"].acompletion) + assert payload["router_model"] == "nvidia/switchyard/gpt-5.5" + + async def test_no_selection_recorded_when_all_attempts_fail(self): + backend = _make_backend(_config("model-A")) + backend._clients["model-A"].acompletion = AsyncMock( + side_effect=_api_status_error(401), + ) + + ctx = ProxyContext() + with pytest.raises(openai.APIStatusError): + await backend.call(ctx, _openai_request()) + + assert CTX_ROUTE_SELECTION not in ctx.metadata + + async def test_responses_surface_is_stamped_too(self): + backend = _make_backend(_config("model-A", request_type="openai_responses")) + backend._clients["model-A"].aresponses = AsyncMock( + return_value={"id": "resp-test", "object": "response", "output": []} + ) + + await backend.call( + ProxyContext(), + ChatRequest.openai_responses({"model": "incoming-model", "input": "hi"}), + ) + + payload = _spend_logs_payload(backend._clients["model-A"].aresponses) + assert payload["router_selected_endpoint"] == "model-A" + assert payload["router_strategy"] == "latency" + + # --------------------------------------------------------------------------- # Credential policy # --------------------------------------------------------------------------- diff --git a/tests/test_route_selection_headers.py b/tests/test_route_selection_headers.py new file mode 100644 index 00000000..10b0ac7a --- /dev/null +++ b/tests/test_route_selection_headers.py @@ -0,0 +1,442 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""End-to-end gate for route-selection spend-attribution headers. + +For tokenomics reporting, a front proxy (e.g. LiteLLM) must be able to tie +its provider spend-log rows back to the Switchyard logical route that +selected them: + +* outbound — every upstream attempt carries an ``x-litellm-spend-logs-metadata`` + request header whose JSON payload records the route selection plus a + per-request correlation id; +* inbound — the Switchyard response returns ``x-switchyard-*`` headers with + the successful attempt's selection and the same correlation id, so the + parent spend-log row can be enriched to match the provider row. + +The app-level tests run the real ``OpenAILLMClient`` against a loopback HTTP +stub, so the outbound header is asserted on the wire — not on a mock. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from typing import Any +from unittest.mock import patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient +from pytest import fixture + +from switchyard.lib.backends.health_poller import HealthPoller +from switchyard.lib.backends.latency_service_llm_backend import ( + SPEND_LOGS_METADATA_HEADER, +) +from switchyard.lib.config.latency_service_backend_config import ( + LatencyServiceBackendConfig, + LatencyServiceEndpoint, +) +from switchyard.lib.endpoints.dispatch import serialize_chain_result +from switchyard.lib.endpoints.route_selection import ( + ROUTER_CORRELATION_ID_HEADER, + ROUTER_MODEL_HEADER, + SELECTED_MODEL_HEADER, + SELECTED_PROVIDER_HEADER, + route_selection_headers, +) +from switchyard.lib.endpoints.upstream_error import handle_chain_exception +from switchyard.lib.profiles import LatencyServiceProfileConfig, ProfileSwitchyard +from switchyard.lib.proxy_context import CTX_ROUTE_SELECTION, ProxyContext +from switchyard.server.switchyard_app import build_switchyard_app +from tests._chain_test_helpers import _backend_payload, _OpenAICompatStub, _sse_body, _stream_chunk + +ROUTE_MODEL = "nvidia/switchyard/test-route" +ENDPOINT_ID = "openai/test-model" +UPSTREAM_MODEL = "openai/openai/test-model" + + +def _selection(**overrides: object) -> dict[str, object]: + payload: dict[str, object] = { + "router_model": ROUTE_MODEL, + "router_strategy": "latency", + "router_selected_endpoint": ENDPOINT_ID, + "router_selected_model": UPSTREAM_MODEL, + "router_selected_provider": "openai", + "router_correlation_id": "11111111-2222-3333-4444-555555555555", + } + payload.update(overrides) + return payload + + +# --------------------------------------------------------------------------- +# Unit: ctx → response-header mapping +# --------------------------------------------------------------------------- + + +class TestRouteSelectionHeaders: + def test_empty_without_selection(self) -> None: + assert route_selection_headers(ProxyContext()) == {} + + def test_maps_selection_to_response_headers(self) -> None: + ctx = ProxyContext() + ctx.metadata[CTX_ROUTE_SELECTION] = _selection() + + assert route_selection_headers(ctx) == { + ROUTER_MODEL_HEADER: ROUTE_MODEL, + SELECTED_MODEL_HEADER: UPSTREAM_MODEL, + SELECTED_PROVIDER_HEADER: "openai", + ROUTER_CORRELATION_ID_HEADER: "11111111-2222-3333-4444-555555555555", + } + + def test_skips_absent_fields_instead_of_stamping_placeholders(self) -> None: + ctx = ProxyContext() + ctx.metadata[CTX_ROUTE_SELECTION] = _selection(router_model=None) + + headers = route_selection_headers(ctx) + + assert ROUTER_MODEL_HEADER not in headers + assert headers[SELECTED_MODEL_HEADER] == UPSTREAM_MODEL + + def test_ignores_non_mapping_value(self) -> None: + ctx = ProxyContext() + ctx.metadata[CTX_ROUTE_SELECTION] = "bogus" + + assert route_selection_headers(ctx) == {} + + def test_skips_values_unsafe_as_header_material(self) -> None: + """The client-controlled router_model must be re-validated as a header. + + A CRLF-bearing value would be a response-splitting vector, and a + non-latin-1 value would crash Starlette response construction after + the upstream call already succeeded and was billed. + """ + ctx = ProxyContext() + ctx.metadata[CTX_ROUTE_SELECTION] = _selection(router_model="gpt\r\nx-evil: 1") + headers = route_selection_headers(ctx) + assert ROUTER_MODEL_HEADER not in headers + assert headers[SELECTED_MODEL_HEADER] == UPSTREAM_MODEL + + ctx = ProxyContext() + ctx.metadata[CTX_ROUTE_SELECTION] = _selection(router_model="gpt-4中文") + headers = route_selection_headers(ctx) + assert ROUTER_MODEL_HEADER not in headers + assert headers[ROUTER_CORRELATION_ID_HEADER] == ( + "11111111-2222-3333-4444-555555555555" + ) + + +class TestSerializeChainResultHeaders: + @staticmethod + async def _sse_iter(_result: Any) -> Any: + yield "data: {}\n\n" + + def test_json_response_carries_selection_headers(self) -> None: + ctx = ProxyContext() + ctx.metadata[CTX_ROUTE_SELECTION] = _selection() + + response = serialize_chain_result( + {"ok": True}, stream=False, sse_iter=self._sse_iter, ctx=ctx + ) + + assert response.headers[SELECTED_MODEL_HEADER] == UPSTREAM_MODEL + assert ( + response.headers[ROUTER_CORRELATION_ID_HEADER] + == "11111111-2222-3333-4444-555555555555" + ) + + def test_streaming_response_carries_selection_headers(self) -> None: + # The backend call completes before the StreamingResponse is built, so + # the selection is final by the time SSE headers are committed. + class _EmptyStream: + def __aiter__(self) -> _EmptyStream: + return self + + async def __anext__(self) -> str: + raise StopAsyncIteration + + ctx = ProxyContext() + ctx.metadata[CTX_ROUTE_SELECTION] = _selection() + + response = serialize_chain_result( + _EmptyStream(), stream=True, sse_iter=self._sse_iter, ctx=ctx + ) + + assert response.media_type == "text/event-stream" + assert response.headers[ROUTER_MODEL_HEADER] == ROUTE_MODEL + + def test_selection_free_ctx_no_selection_headers(self) -> None: + response = serialize_chain_result( + {"ok": True}, stream=False, sse_iter=self._sse_iter, ctx=ProxyContext() + ) + + assert ROUTER_CORRELATION_ID_HEADER not in response.headers + + +class TestErrorPathSelectionHeaders: + def test_failure_after_billed_success_keeps_selection_headers(self) -> None: + """A post-backend failure must still expose the billed selection. + + The upstream call succeeded (provider spend-log row written with the + stamped correlation id) before e.g. response translation raised; the + error response must carry the selection headers or that provider row + becomes unjoinable. + """ + ctx = ProxyContext() + ctx.metadata[CTX_ROUTE_SELECTION] = _selection() + + response = handle_chain_exception( + RuntimeError("response translation failed"), + ctx, + inbound="openai", + log_msg="test", + ) + + assert response.status_code == 500 + assert response.headers[ROUTER_CORRELATION_ID_HEADER] == ( + "11111111-2222-3333-4444-555555555555" + ) + assert response.headers[SELECTED_MODEL_HEADER] == UPSTREAM_MODEL + + def test_failure_without_selection_has_no_selection_headers(self) -> None: + response = handle_chain_exception( + RuntimeError("backend never succeeded"), + ProxyContext(), + inbound="openai", + log_msg="test", + ) + + assert response.status_code == 500 + assert ROUTER_CORRELATION_ID_HEADER not in response.headers + + +# --------------------------------------------------------------------------- +# End-to-end: HTTP in → wire header out → response headers back +# --------------------------------------------------------------------------- + + +@fixture +def latency_app() -> Iterator[tuple[FastAPI, _OpenAICompatStub]]: + """Latency-service app whose single endpoint targets a loopback stub. + + The backend uses the real ``OpenAILLMClient``/OpenAI SDK, so the stub + records the actual HTTP headers the upstream would receive. Only the + health poller is stubbed out (no Latency Service in tests). + """ + with _OpenAICompatStub() as stub: + config = LatencyServiceBackendConfig( + latency_service_url="http://latency-service.test:8080", + route_model=ROUTE_MODEL, + endpoints=[ + LatencyServiceEndpoint( + model=ENDPOINT_ID, + upstream_model=UPSTREAM_MODEL, + base_url=stub.base_url, + api_key="test-key", + ), + ], + ) + with patch.object(HealthPoller, "start"), patch.object(HealthPoller, "stop"): + switchyard = ProfileSwitchyard( + LatencyServiceProfileConfig.from_config(config) + .build() + .with_runtime_components(enable_stats=config.enable_stats) + ) + yield build_switchyard_app(switchyard), stub + + +def _wire_spend_logs_payload(stub: _OpenAICompatStub) -> dict[str, object]: + """Parse the spend-logs metadata header the stub received on the wire.""" + headers = stub.requests[-1]["headers"] + assert isinstance(headers, dict) + payload = json.loads(str(headers[SPEND_LOGS_METADATA_HEADER])) + assert isinstance(payload, dict) + return payload + + +def test_chat_completion_roundtrips_selection_and_correlation_id( + latency_app: tuple[FastAPI, _OpenAICompatStub], +) -> None: + app, stub = latency_app + stub.respond_json(_backend_payload(content="hello", model=UPSTREAM_MODEL)) + + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json={ + "model": ROUTE_MODEL, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + + assert response.status_code == 200 + payload = _wire_spend_logs_payload(stub) + assert payload["router_model"] == ROUTE_MODEL + assert payload["router_strategy"] == "latency" + assert payload["router_selected_endpoint"] == ENDPOINT_ID + assert payload["router_selected_model"] == UPSTREAM_MODEL + assert payload["router_selected_provider"] == "openai" + assert response.headers[ROUTER_MODEL_HEADER] == ROUTE_MODEL + assert response.headers[SELECTED_MODEL_HEADER] == UPSTREAM_MODEL + assert response.headers[SELECTED_PROVIDER_HEADER] == "openai" + # The join key: the response header must equal the id the provider row got. + assert response.headers[ROUTER_CORRELATION_ID_HEADER] == payload["router_correlation_id"] + + +def test_streaming_chat_completion_carries_selection_headers( + latency_app: tuple[FastAPI, _OpenAICompatStub], +) -> None: + app, stub = latency_app + stub.respond_sse(_sse_body([_stream_chunk(content="hel"), _stream_chunk(finish="stop")])) + + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json={ + "model": ROUTE_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + }, + ) + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert "data:" in response.text + payload = _wire_spend_logs_payload(stub) + assert response.headers[ROUTER_CORRELATION_ID_HEADER] == payload["router_correlation_id"] + assert response.headers[SELECTED_MODEL_HEADER] == UPSTREAM_MODEL + + +def test_anthropic_messages_endpoint_carries_selection_headers( + latency_app: tuple[FastAPI, _OpenAICompatStub], +) -> None: + """Anthropic inbound rides the same dispatch path and gets the same headers.""" + app, stub = latency_app + stub.respond_json(_backend_payload(content="hello", model=UPSTREAM_MODEL)) + + with TestClient(app) as client: + response = client.post( + "/v1/messages", + json={ + "model": ROUTE_MODEL, + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + + assert response.status_code == 200 + payload = _wire_spend_logs_payload(stub) + assert response.headers[ROUTER_CORRELATION_ID_HEADER] == payload["router_correlation_id"] + assert response.headers[SELECTED_PROVIDER_HEADER] == "openai" + + +def test_responses_endpoint_carries_selection_headers( + latency_app: tuple[FastAPI, _OpenAICompatStub], +) -> None: + """OpenAI Responses inbound rides the same dispatch path and gets the same headers.""" + app, stub = latency_app + stub.respond_json(_backend_payload(content="hello", model=UPSTREAM_MODEL)) + + with TestClient(app) as client: + response = client.post( + "/v1/responses", + json={"model": ROUTE_MODEL, "input": "hi"}, + ) + + assert response.status_code == 200 + payload = _wire_spend_logs_payload(stub) + assert response.headers[ROUTER_CORRELATION_ID_HEADER] == payload["router_correlation_id"] + assert response.headers[SELECTED_MODEL_HEADER] == UPSTREAM_MODEL + + +def test_upstream_error_response_has_no_selection_headers( + latency_app: tuple[FastAPI, _OpenAICompatStub], +) -> None: + """No upstream call succeeded — the error response must not claim a selection. + + The failed attempt itself is still stamped on the wire (the provider row + for the failure keeps its attribution); the client-facing headers are + reserved for a recorded selection, i.e. a billed upstream success — which + never happened here. (A failure *after* a billed success does carry them; + see TestErrorPathSelectionHeaders.) + """ + app, stub = latency_app + stub.respond_json({"error": {"message": "bad key"}}, status=401) + + with TestClient(app, raise_server_exceptions=False) as client: + response = client.post( + "/v1/chat/completions", + json={ + "model": ROUTE_MODEL, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + + assert response.status_code == 401 + assert ROUTER_CORRELATION_ID_HEADER not in response.headers + assert SPEND_LOGS_METADATA_HEADER in stub.requests[-1]["headers"] # type: ignore[operator] + + +def test_hostile_model_string_never_breaks_the_response( + latency_app: tuple[FastAPI, _OpenAICompatStub], +) -> None: + """A model string that is not legal header material must not 500 or inject. + + Single-chain profile serving forwards any client model string, so + router_model is client-controlled; the billed 200 must survive, with the + unsafe header skipped and the config-derived headers intact. + """ + app, stub = latency_app + + with TestClient(app) as client: + for hostile_model in ("gpt-4中文", "gpt\r\nx-evil: 1"): + stub.respond_json(_backend_payload(content="hello", model=UPSTREAM_MODEL)) + response = client.post( + "/v1/chat/completions", + json={ + "model": hostile_model, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + + assert response.status_code == 200 + assert ROUTER_MODEL_HEADER not in response.headers + assert "x-evil" not in response.headers + assert response.headers[SELECTED_MODEL_HEADER] == UPSTREAM_MODEL + # The outbound header is JSON (ensure_ascii escapes control and + # non-ASCII chars), so the wire payload still records the model. + assert _wire_spend_logs_payload(stub)["router_model"] == hostile_model + + +def test_client_body_extra_headers_forwarded_not_collided( + latency_app: tuple[FastAPI, _OpenAICompatStub], +) -> None: + """A passthrough body carrying an SDK-style extra_headers field keeps working. + + Pre-spend-logs it rode ``**body`` into the SDK's header kwarg; it must + still reach the wire — and must not be able to spoof the spend-logs + header, which wins the name conflict. + """ + app, stub = latency_app + stub.respond_json(_backend_payload(content="hello", model=UPSTREAM_MODEL)) + + with TestClient(app) as client: + response = client.post( + "/v1/chat/completions", + json={ + "model": ROUTE_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "extra_headers": { + "x-client-tag": "42", + SPEND_LOGS_METADATA_HEADER: "spoofed", + }, + }, + ) + + assert response.status_code == 200 + headers = stub.requests[-1]["headers"] + assert isinstance(headers, dict) + assert headers["x-client-tag"] == "42" + payload = _wire_spend_logs_payload(stub) + assert payload["router_selected_endpoint"] == ENDPOINT_ID # not "spoofed" + assert response.headers[ROUTER_CORRELATION_ID_HEADER] == payload["router_correlation_id"]