Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .agents/skills/switchyard-lib-core/SKILL.md
Original file line number Diff line number Diff line change
Expand Up @@ -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. |
Expand Down
46 changes: 46 additions & 0 deletions docs/internal/latency_service_routing.md
Original file line number Diff line number Diff line change
Expand Up @@ -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": "<client-facing route id>",
"router_strategy": "latency",
"router_selected_endpoint": "<endpoint id>",
"router_selected_model": "<upstream model>",
"router_selected_provider": "<upstream model's leading path segment>",
"router_correlation_id": "<uuid4>"}
```

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
Expand Down
79 changes: 79 additions & 0 deletions switchyard/lib/backends/latency_service_llm_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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",
Expand All @@ -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, {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -605,21 +639,66 @@ 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,
target_request_type: ChatRequestType,
*,
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,
)
Comment on lines 669 to 703

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔒 Security & Privacy | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Check how the OpenAI SDK / httpx merge extra_headers with default headers,
# and whether case-differing duplicate keys collapse or survive as two headers.
python - <<'EOF'
import inspect
try:
    import httpx
    print(inspect.getsource(httpx.Headers.__init__))
except Exception as e:
    print("httpx introspection failed:", e)
try:
    import openai
    from openai._base_client import BaseClient
    print(inspect.getsource(BaseClient._build_headers))
except Exception as e:
    print("openai SDK introspection failed:", e)
EOF

Repository: NVIDIA-NeMo/Switchyard

Length of output: 1118


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Map the relevant file and inspect the surrounding implementation.
ast-grep outline switchyard/lib/backends/latency_service_llm_backend.py --view expanded || true
sed -n '640,730p' switchyard/lib/backends/latency_service_llm_backend.py

# Find tests and call sites related to extra_headers / spend-logs metadata.
rg -n "extra_headers|spend-logs|SPEND_LOGS_METADATA_HEADER|test_client_body_extra_headers_forwarded_not_collided" switchyard -g '!**/__pycache__/**'

Repository: NVIDIA-NeMo/Switchyard

Length of output: 10424


🌐 Web query:

OpenAI Python SDK BaseClient _build_headers extra_headers case-insensitive httpx.Headers merge behavior

💡 Result:

In the OpenAI Python SDK, the _build_headers method is responsible for constructing the final request headers by merging default client headers with any user-provided extra_headers [1][2][3]. The merge behavior operates as follows: 1. Initial Merging: The SDK first performs a dictionary-based merge of the client's default headers and the provided custom headers using a helper function (_merge_mappings) [1][2][3]. 2. Case-Insensitivity: The result of this dictionary merge is then passed to the httpx.Headers class constructor [1][2][3]. Because httpx.Headers is a case-insensitive multi-dict [4], it normalizes the header keys. 3. Precedence: The extra_headers provided by the user take precedence over client-level default headers when keys clash [5]. This is consistent with the standard behavior of merging dictionaries where the later dictionary overwrites the former, and further enforced by how the httpx.Headers object subsequently handles the provided key-value pairs [4][5]. Essentially, the process relies on the fact that while the preliminary dictionary merge is case-sensitive, the conversion into the httpx.Headers object creates a case-insensitive store that correctly manages the final header set for the request [1][2][3].

Citations:


🏁 Script executed:

python3 - <<'PY'
import inspect
import httpx

print("httpx version:", httpx.__version__)
for name in ["__contains__", "__getitem__", "get_list", "items", "multi_items", "keys", "values", "raw"]:
    try:
        attr = getattr(httpx.Headers, name)
        print(f"\n### Headers.{name}")
        print(inspect.getsource(attr))
    except Exception as e:
        print(f"\n### Headers.{name} introspection failed: {e}")

h = httpx.Headers([
    ("x-litellm-spend-logs-metadata", "safe"),
    ("X-LiteLLM-Spend-Logs-Metadata", "spoof"),
])
print("\nconstructed headers repr:", h)
print("contains lowercase:", "x-litellm-spend-logs-metadata" in h)
print("contains spoof-cased:", "X-LiteLLM-Spend-Logs-Metadata" in h)
print("getitem lowercase:", h["x-litellm-spend-logs-metadata"])
print("get_list lowercase:", h.get_list("x-litellm-spend-logs-metadata"))
print("multi_items:", list(h.multi_items()))
print("raw:", h.raw)
PY

Repository: NVIDIA-NeMo/Switchyard

Length of output: 3990


Normalize the protected header case-insensitively before merging. httpx.Headers treats header names case-insensitively and keeps duplicate raw entries, so a client can still send X-LiteLLM-Spend-Logs-Metadata alongside the protected lowercase key. Drop or rewrite any case variant of the spend-logs header, and add a test for a differently-cased spoof key.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@switchyard/lib/backends/latency_service_llm_backend.py` around lines 669 -
703, Update _call_endpoint’s extra_headers merge to remove or overwrite every
case variant of the protected spend-logs metadata header before applying
extra_headers, ensuring client-provided casing cannot spoof or duplicate it.
Preserve the protected header value from extra_headers and add a test covering a
differently-cased spoof key.


Expand Down
4 changes: 3 additions & 1 deletion switchyard/lib/endpoints/anthropic_messages_endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
15 changes: 12 additions & 3 deletions switchyard/lib/endpoints/dispatch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Comment on lines 111 to +113

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== dispatch.py outline ==\n'
ast-grep outline switchyard/lib/endpoints/dispatch.py --view expanded || true

printf '\n== relevant slice of dispatch.py ==\n'
nl -ba switchyard/lib/endpoints/dispatch.py | sed -n '1,220p'

printf '\n== search for serialize_chain_result usages/tests ==\n'
rg -n "serialize_chain_result|route_selection_headers|CTX_ROUTE_SELECTION|x-switchyard" switchyard -g '!**/__pycache__/**'

printf '\n== tests around route selection headers ==\n'
fd -a "test_route_selection_headers.py|route_selection" switchyard || true

Repository: NVIDIA-NeMo/Switchyard

Length of output: 820


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the endpoint code paths that return Response objects directly.
printf '\n== anthropic_messages_endpoint.py relevant lines ==\n'
nl -ba switchyard/lib/endpoints/anthropic_messages_endpoint.py | sed -n '1,260p'

printf '\n== responses_endpoint.py relevant lines ==\n'
nl -ba switchyard/lib/endpoints/responses_endpoint.py | sed -n '1,260p'

printf '\n== route_selection_headers implementation and tests ==\n'
rg -n "def route_selection_headers|class ProxyContext|CTX_ROUTE_SELECTION|route_selection_headers\\(" switchyard/lib switchyard/tests -g '!**/__pycache__/**'

Repository: NVIDIA-NeMo/Switchyard

Length of output: 253


🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== dispatch.py relevant slice ==\n'
sed -n '1,220p' switchyard/lib/endpoints/dispatch.py

printf '\n== anthropic_messages_endpoint.py relevant slice ==\n'
sed -n '1,260p' switchyard/lib/endpoints/anthropic_messages_endpoint.py

printf '\n== responses_endpoint.py relevant slice ==\n'
sed -n '1,260p' switchyard/lib/endpoints/responses_endpoint.py

printf '\n== search for route selection header tests ==\n'
rg -n "route_selection_headers|CTX_ROUTE_SELECTION|x-switchyard|serialize_chain_result" switchyard -g '!**/__pycache__/**'

Repository: NVIDIA-NeMo/Switchyard

Length of output: 18660


🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== route_selection.py ==\n'
sed -n '1,220p' switchyard/lib/endpoints/route_selection.py

printf '\n== search for chain implementations returning Response ==\n'
rg -n "return .*Response|StreamingResponse\\(|JSONResponse\\(" switchyard/lib -g '!**/__pycache__/**'

printf '\n== tests mentioning serialize_chain_result or route selection headers ==\n'
rg -n "serialize_chain_result|route_selection_headers|CTX_ROUTE_SELECTION|x-switchyard-router-model|x-switchyard-selected-model|x-switchyard-selected-provider" switchyard/tests switchyard -g '!*__pycache__*'

Repository: NVIDIA-NeMo/Switchyard

Length of output: 8990


🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '\n== openai_chat_endpoint.py relevant slice ==\n'
sed -n '1,220p' switchyard/lib/endpoints/openai_chat_endpoint.py

printf '\n== upstream_error.py header stamping slice ==\n'
sed -n '180,260p' switchyard/lib/endpoints/upstream_error.py

Repository: NVIDIA-NeMo/Switchyard

Length of output: 6655


Stamp route-selection headers before returning an existing Response. serialize_chain_result returns Response unchanged before route_selection_headers(ctx) runs, so any pre-built response drops the x-switchyard-* attribution headers even when ctx already has a recorded selection. Merge the headers onto the existing response instead of bypassing them.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@switchyard/lib/endpoints/dispatch.py` around lines 111 - 113, Update the
response handling around serialize_chain_result so route_selection_headers(ctx)
is applied before returning an existing Response. Merge the generated
x-switchyard-* attribution headers onto that response, preserving its existing
headers and body, while retaining the current serialization path for
non-Response results.

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)
2 changes: 1 addition & 1 deletion switchyard/lib/endpoints/openai_chat_endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
4 changes: 3 additions & 1 deletion switchyard/lib/endpoints/responses_endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
68 changes: 68 additions & 0 deletions switchyard/lib/endpoints/route_selection.py
Original file line number Diff line number Diff line change
@@ -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
Loading
Loading