vinez/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-10 14:27:00 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Per-request backend selection for routing extensions (#2809)

Browse source 
## The gap

Headroom picks its egress backend **once**, at startup:
`create_proxy_backend` returns a single `Backend` (or `None` for the
direct Anthropic path) and every request goes through it. That is the
right shape for *"run this whole proxy against Bedrock instead of
Anthropic"* and the wrong shape for *"this request is cheaper on a
different provider than the last one."*

`ModelRouter` already lets an extension change `body["model"]` per
request — but only within the protocol the request arrived in, because a
model id alone cannot move a request to another provider.

So an extension can currently **decide** something Headroom has no way
to **carry out**. This adds the missing half.

## The seam

An extension publishes a decision on the request state:

```python
request.state.headroom_route = SimpleNamespace(
    model="moonshot/kimi-k2",   # required
    provider="moonshot",        # optional; inferred from the model id if absent
    reason="cheaper at this prefix length",
)
```

Headroom resolves a `LiteLLMBackend` for that provider — which is where
translation already lives — and serves **that one request** from it.
Nothing in core names any particular extension; the field is duck-typed,
so an extension does not import Headroom to talk to Headroom.

## Absent means unchanged

This is the property the tests are built around, and the reason this
should be safe to merge.

With nothing published, every path is what it was before. Advice that is
**absent, malformed, names an unknown provider, names a native provider,
or fails to build** all resolve to `self.anthropic_backend` — including
when that is `None`, which is the direct-API path and must survive. A
routing preference can never take traffic down.

## Coverage

| path | |
|---|---|
| `/v1/messages` | non-streaming + streaming |
| `/v1/chat/completions` | non-streaming + streaming |
| Responses API | untouched — does not use the backend abstraction |

Streaming is the one that matters. The resolver rewrites
`body["model"]`, so had `_stream_response_bedrock` kept reading
`self.anthropic_backend`, every streamed routed request would have sent
a foreign model id to Anthropic. Both streaming helpers now take an
optional `backend`, defaulting to the configured one.

## Details worth review

- **Validate the provider name before building.** `LiteLLMBackend`
accepts *any* provider string — the registry falls through to a generic
pass-through config — so a typo silently builds a backend that only
fails later, at request time, with an error pointing nowhere near the
typo. `_known_provider()` checks against `litellm.provider_list` first.
- **Cache per provider, and cache the failures too**, or a broken
provider name costs a construction attempt on every request. (Bedrock
construction calls out to AWS to enumerate inference profiles — it is
not free.)
- **`backend_owns_translation` now asks the per-request backend.** It
decides whether Headroom or the backend owns the `max_tokens` /
`max_completion_tokens` spelling; asking `self.anthropic_backend` would
answer "Headroom does" for a request about to be served by a translating
backend that does.
- **`_route_resolver` lives in `route_advice.py`, not on a handler
mixin.** Two mixins need it, and reaching across sibling mixins only
works by accident of how `HeadroomProxy` composes them.

## Tests

`tests/test_route_advice.py` — 20 tests, most of them asserting the
absent-means-unchanged property from a different angle.

Local runs: 20/20 on the new file; **1102 passed, 1 failed** on `-k
"openai or chat_completions or ccr"`, and **414 passed, 0 failed** on
`-k "stream or bedrock or route_advice"`. The single failure is
`test_realignment_live_multi_turn::test_ccr_marker_round_trip_live`,
which fails identically on this branch's merge-base — verified by
checking out `59314cff~1` and re-running it.

Note for anyone reproducing: `pytest-asyncio` is a declared dev
dependency but was missing from my venv, which made every `async def
test_` in the repo fail. Worth checking before diagnosing a large
failure count.

## Docs

`docs/content/docs/pipeline-extensions.mdx` gains a section on the
contract, next to the existing `x-headroom-base-url` one.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Tejas Chopra 2026-08-05 17:01:32 -07:00 • committed by GitHub
parent 0237cbffbb
commit c07da992dd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 572 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

23
docs/content/docs/pipeline-extensions.mdx
View file

@ -77,3 +77,26 @@ curl http://localhost:8787/v1/chat/completions \
```
Internal `x-headroom-*` headers (including this one) are stripped before the request is forwarded upstream by default — see `HEADROOM_STRIP_INTERNAL_HEADERS` in [Configuration](/docs/configuration).
## Per-request model routing with `request.state.headroom_route`
`x-headroom-base-url` is client-driven and points at one OpenAI-compatible base. When the choice of model belongs to an extension instead of the caller — a router that picks a cheaper model per turn, say — publish it on the request state and Headroom serves that one request from a backend that speaks the target provider:
```python
# middleware or an extension holding the request
request.state.headroom_route = SimpleNamespace(
model="moonshot/kimi-k2", # required
provider="moonshot", # optional; inferred from the model id if absent
reason="cheaper at this prefix length",
)
```
The contract, in `headroom/proxy/route_advice.py`:
- **Absent means unchanged.** No advice — or advice that is malformed, names an unknown provider, or fails to build a backend — and the request takes exactly the path it took before. A routing preference can never take traffic down.
- **Duck-typed**, so an extension does not import Headroom to publish one.
- A **native** provider (`anthropic`) needs no backend switch — rewrite `body["model"]` yourself. A foreign one is translated by a `LiteLLMBackend` built for it, and Headroom writes the model id.
- Backends are **built once per provider** and cached; a provider that fails to build is not retried per request.
- Honored on `/v1/messages` and `/v1/chat/completions`, streaming and non-streaming alike. (Not the Responses API, which does not use the backend abstraction.)
`routemegood` is the reference consumer of this seam: it decides, Headroom routes.

21
headroom/proxy/handlers/anthropic.py
View file

@ -2644,7 +2644,17 @@ class AnthropicHandlerMixin:
headers["anthropic-beta"] = _client_beta_value
# Forward request - use Bedrock backend if configured, otherwise direct API
if self.anthropic_backend is not None:
#
# An extension may have published a per-request routing decision on
# `request.state.headroom_route` (see proxy/route_advice.py). When
# it names a provider we do not already speak, serve this ONE
# request from a translating backend for it. Absent, unresolvable,
# or same-protocol -> `self.anthropic_backend`, i.e. exactly the
# behaviour this line had before.
from headroom.proxy.route_advice import resolver_for
request_backend = resolver_for(self).for_request(request, body=body)
if request_backend is not None:
# Route through Bedrock backend
try:
if stream:
@ -2675,12 +2685,11 @@ class AnthropicHandlerMixin:
original_messages=original_client_messages,
prefix_tracker=prefix_tracker,
optimized_messages=optimized_messages,
backend=request_backend,
)
else:
async with stage_timer.measure("upstream_connect"):
backend_response = await self.anthropic_backend.send_message(
body, headers
)
backend_response = await request_backend.send_message(body, headers)
self.pipeline_extensions.emit(
PipelineStage.POST_SEND,
operation="proxy.request",
@ -2732,9 +2741,7 @@ class AnthropicHandlerMixin:
usage = backend_response.body.get("usage", {})
output_tokens = usage.get("output_tokens", 0)
_backend_name = (
self.anthropic_backend.name if self.anthropic_backend else "anthropic"
)
_backend_name = request_backend.name if request_backend else "anthropic"
# Eligible-only denominator for the active
# compression ratio: tokens in the live zone we
# actually attempted to compress. Frozen prefix

34
headroom/proxy/handlers/openai.py
View file

@ -3702,8 +3702,15 @@ class OpenAIHandlerMixin:
# translate it here — the proxy already owns the outbound body — and
# those requests work unchanged. No-op when the caller already set
# `max_completion_tokens`.
# Resolved without `body=` so nothing is rewritten yet -- we only need
# to know WHICH backend will serve this request, because a translating
# one owns the max_tokens spelling. Cached, so the dispatch-site call
# below is a dict lookup.
from headroom.proxy.route_advice import resolver_for
_normalize_openai_max_tokens(
body, backend_owns_translation=self.anthropic_backend is not None
body,
backend_owns_translation=resolver_for(self).for_request(request) is not None,
)
# Output shaping (opt-in via HEADROOM_OUTPUT_SHAPER): verbosity steering
@ -3763,8 +3770,12 @@ class OpenAIHandlerMixin:
f"{_shape_result.labels}"
)
# Route through LiteLLM/any-llm backend if configured
if self.anthropic_backend is not None:
# Route through LiteLLM/any-llm backend if configured -- or through a
# per-request one an extension asked for (see proxy/route_advice.py).
# No advice resolves to `self.anthropic_backend`, so this is the same
# condition it has always been.
request_backend = resolver_for(self).for_request(request, body=body)
if request_backend is not None:
try:
if stream:
self.pipeline_extensions.emit(
@ -3794,12 +3805,11 @@ class OpenAIHandlerMixin:
waste_signals=waste_signals_dict,
prefix_tracker=openai_prefix_tracker,
optimized_messages=optimized_messages,
backend=request_backend,
)
else:
# Non-streaming: use send_openai_message() → JSON
backend_response = await self.anthropic_backend.send_openai_message(
body, headers
)
backend_response = await request_backend.send_openai_message(body, headers)
self.pipeline_extensions.emit(
PipelineStage.POST_SEND,
operation="proxy.request",
@ -3858,7 +3868,7 @@ class OpenAIHandlerMixin:
):
logger.info(
f"[{request_id}] CCR: Detected retrieval tool call "
f"on backend path, handling via {self.anthropic_backend.name}"
f"on backend path, handling via {request_backend.name}"
)
# Continuation closure — delegates transport to
@ -3885,13 +3895,13 @@ class OpenAIHandlerMixin:
)
}
assert self.anthropic_backend is not None
assert request_backend is not None
logger.info(
f"[{request_id}] CCR: Issuing continuation via "
f"{self.anthropic_backend.name} backend "
f"{request_backend.name} backend "
f"({len(msgs)} messages)"
)
cont_resp = await self.anthropic_backend.send_openai_message(
cont_resp = await request_backend.send_openai_message(
continuation_body, continuation_headers
)
return cont_resp.body
@ -4017,7 +4027,7 @@ class OpenAIHandlerMixin:
await self._record_request_outcome(
RequestOutcome(
request_id=request_id,
provider=self.anthropic_backend.name,
provider=request_backend.name,
model=model,
original_tokens=original_tokens,
# Local count, same tokenizer as original_tokens, so
@ -4049,7 +4059,7 @@ class OpenAIHandlerMixin:
if tokens_saved > 0:
logger.info(
f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} "
f"(saved {tokens_saved:,} tokens) via {self.anthropic_backend.name}"
f"(saved {tokens_saved:,} tokens) via {request_backend.name}"
)
return JSONResponse(

27
headroom/proxy/handlers/streaming.py
View file

@ -1679,6 +1679,7 @@ class StreamingMixin:
original_messages: list[dict] | None = None,
prefix_tracker: Any | None = None,
optimized_messages: list[dict] | None = None,
backend: Any | None = None,
) -> StreamingResponse:
"""Stream response from Bedrock backend with metrics tracking.
@ -1697,6 +1698,11 @@ class StreamingMixin:
from headroom.proxy.outcome import RequestOutcome
# ``backend`` lets the caller serve this one request from somewhere
# other than the configured backend (see proxy/route_advice.py). None
# means "the configured one", i.e. what this method always did.
backend = backend if backend is not None else self.anthropic_backend
client = classify_client(headers)
start_time = time.time()
@ -1721,7 +1727,7 @@ class StreamingMixin:
async def generate():
try:
assert self.anthropic_backend is not None
assert backend is not None
# Emit a synthetic ping before the first message_start so that
# downstream clients (e.g. Claude Code) arm their mid-turn
@ -1731,7 +1737,7 @@ class StreamingMixin:
# (issue #902).
yield b"event: ping\ndata: {}\n\n"
async for event in self.anthropic_backend.stream_message(body, headers):
async for event in backend.stream_message(body, headers):
# Record TTFB on first event
if stream_state["ttfb_ms"] is None:
stream_state["ttfb_ms"] = (time.time() - start_time) * 1000
@ -1807,9 +1813,7 @@ class StreamingMixin:
finally:
total_latency = (time.time() - start_time) * 1000
_backend_name = (
self.anthropic_backend.name if self.anthropic_backend else "anthropic"
)
_backend_name = backend.name if backend else "anthropic"
# Update prefix cache tracker for the next turn — mirrors
# _finalize_stream_response (direct-API streaming path)
@ -1908,6 +1912,7 @@ class StreamingMixin:
waste_signals: dict[str, int] | None = None,
prefix_tracker: Any | None = None,
optimized_messages: list[dict] | None = None,
backend: Any | None = None,
) -> StreamingResponse:
"""Stream OpenAI chat completion response from backend.
@ -1942,7 +1947,11 @@ class StreamingMixin:
from headroom.proxy.handlers.openai import _infer_openai_cache_write_tokens
from headroom.proxy.outcome import RequestOutcome
assert self.anthropic_backend is not None
# ``backend`` lets the caller serve this one request from somewhere
# other than the configured backend (see proxy/route_advice.py). None
# means "the configured one", i.e. what this method always did.
backend = backend if backend is not None else self.anthropic_backend
assert backend is not None
client = classify_client(headers)
async def generate():
@ -1972,7 +1981,7 @@ class StreamingMixin:
stream_state[key] = usage[key]
try:
async for sse_chunk in self.anthropic_backend.stream_openai_message(body, headers):
async for sse_chunk in backend.stream_openai_message(body, headers):
chunk_bytes = sse_chunk.encode() if isinstance(sse_chunk, str) else sse_chunk
stream_state["sse_buffer"].extend(chunk_bytes)
full_sse_bytes.extend(chunk_bytes)
@ -2074,7 +2083,7 @@ class StreamingMixin:
# instead of collapsing the dashboard headline to 0%.
outcome = RequestOutcome.from_stream(
body=body,
provider=self.anthropic_backend.name,
provider=backend.name,
model=model,
request_id=request_id,
original_tokens=original_tokens,
@ -2099,7 +2108,7 @@ class StreamingMixin:
if tokens_saved > 0:
logger.info(
f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} "
f"(saved {tokens_saved:,} tokens) via {self.anthropic_backend.name} [stream]"
f"(saved {tokens_saved:,} tokens) via {backend.name} [stream]"
)
return StreamingResponse(

206
headroom/proxy/route_advice.py Normal file
View file

@ -0,0 +1,206 @@
"""Per-request backend selection, published by an extension.
WHY THIS EXISTS
Headroom picks its egress backend ONCE, at startup: `create_proxy_backend`
returns a single `Backend` (or None for the direct Anthropic path) and
every request goes through it. That is right for "run this whole proxy
against Bedrock instead of Anthropic", and it is the wrong shape for "this
request is cheaper on a different provider than the last one".
Meanwhile `ModelRouter` already lets an extension change `body["model"]`
per request -- but only within whatever protocol the request arrived in,
because a model id alone cannot move a request to another provider.
So a router extension can currently decide something Headroom has no way to
carry out. This module is the missing half: a routing decision that names a
PROVIDER as well as a model, and a per-request backend to serve it.
THE CONTRACT, AND WHY IT IS VENDOR-NEUTRAL
An extension sets `request.state.headroom_route` to a `RouteAdvice`.
Headroom reads it, resolves a backend for `advice.provider`, and dispatches
there. Nothing in core names any particular extension.
ABSENT MEANS UNCHANGED
This is the property that matters most. With no extension installed, no
extension setting the field, an unresolvable provider, or a backend that
fails to construct, `backend_for()` returns exactly what
`self.anthropic_backend` would have been -- including None, which keeps the
direct-API path. Nothing about the default deployment moves.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Any
log = logging.getLogger(__name__)
STATE_ATTR = "headroom_route"
# Providers that speak the Anthropic Messages API natively, i.e. the shape the
# proxy already holds. Everything else needs a translating backend.
_NATIVE = ("anthropic",)
@dataclass(frozen=True)
class RouteAdvice:
"""A routing decision an extension wants Headroom to carry out.
Only `model` and `provider` are required; the rest is explanatory and is
logged rather than acted on, so an extension can be as terse as it likes.
"""
model: str
provider: str = ""
reason: str = ""
def __post_init__(self) -> None:
if not self.model:
raise ValueError("RouteAdvice needs a model")
def advice_from(request: Any) -> RouteAdvice | None:
"""The advice an extension published, if any. Never raises.
Duck-typed on purpose: an extension should not have to import Headroom to
talk to Headroom. Anything carrying `.model` (and optionally `.provider`)
works, which also keeps extensions from breaking when this dataclass gains
a field.
"""
obj = getattr(getattr(request, "state", None), STATE_ATTR, None)
if obj is None:
return None
model = getattr(obj, "model", None)
if not isinstance(model, str) or not model:
return None
provider = getattr(obj, "provider", "") or ""
return RouteAdvice(
model=model,
provider=provider if isinstance(provider, str) else "",
reason=str(getattr(obj, "reason", "") or "")[:400],
)
class BackendResolver:
"""Lazily builds and caches one translating backend per provider.
Construction is not free -- the Bedrock path calls out to AWS to enumerate
inference profiles -- so backends are built on first use for a provider and
kept. A provider that fails to construct is remembered as a failure and not
retried on every request; the request falls back to the default backend
rather than erroring, because a routing preference must never be able to
take traffic down.
"""
def __init__(self, default: Any = None) -> None:
self._default = default
self._cache: dict[str, Any] = {}
self._failed: set[str] = set()
@property
def default(self) -> Any:
return self._default
def for_request(self, request: Any, *, body: dict | None = None) -> Any:
"""The backend to serve this request. Falls back to the default."""
advice = advice_from(request)
if advice is None:
return self._default
provider = advice.provider or _provider_of(advice.model)
if not provider or provider in _NATIVE:
# Same protocol the proxy already speaks -- a `body["model"]`
# rewrite is enough and the extension has already done it. Nothing
# to switch.
return self._default
if provider in self._failed:
return self._default
backend = self._cache.get(provider)
if backend is None:
try:
backend = self._build(provider)
except Exception as exc: # noqa: BLE001 - see the class docstring
log.warning("route advice: building %r raised (%s)", provider, exc)
backend = None
if backend is None:
self._failed.add(provider)
return self._default
self._cache[provider] = backend
if body is not None:
# The extension chose the model; make sure the body agrees, since
# it could not safely write a foreign model id itself. Log here
# too -- a caller asking without a body is only checking WHICH
# backend will serve the request, and would log the same decision
# a second time.
body["model"] = advice.model
log.info(
"route advice: %s via %s (%s)",
advice.model,
provider,
advice.reason or "no reason given",
)
return backend
def _build(self, provider: str) -> Any:
# Validate the name FIRST. `LiteLLMBackend` accepts an unknown provider
# happily -- the registry falls through to a generic pass-through
# config -- so a typo builds a backend that only fails later, at
# request time, with an error that points nowhere near the typo.
if not _known_provider(provider):
log.warning("route advice: %r is not a litellm provider; ignoring", provider)
return None
try:
from headroom.backends.litellm import LiteLLMBackend
return LiteLLMBackend(provider=provider)
except Exception as exc: # noqa: BLE001 - never fail a request for this
log.warning(
"route advice: cannot build a backend for %r (%s); "
"falling back to the configured backend",
provider,
exc,
)
return None
def resolver_for(handler: Any) -> BackendResolver:
"""The handler's resolver, built once and reused.
Lives here rather than on a handler mixin so every protocol handler can
reach it without depending on a sibling mixin. Rebuilt if
`anthropic_backend` is reassigned (tests do this), so the resolver can
never serve a stale default.
"""
default = getattr(handler, "anthropic_backend", None)
cached = getattr(handler, "_route_resolver_cache", None)
if cached is None or cached.default is not default:
cached = BackendResolver(default)
handler._route_resolver_cache = cached
return cached
def _known_provider(provider: str) -> bool:
"""Is this a provider litellm actually knows about?"""
try:
import litellm
names = {getattr(p, "value", None) or str(p) for p in getattr(litellm, "provider_list", [])}
return provider in names
except Exception: # noqa: BLE001
return False
def _provider_of(model: str) -> str:
"""Provider from a `provider/model` spec, else from litellm's table."""
if "/" in model:
return model.split("/", 1)[0]
try:
import litellm
entry = litellm.model_cost.get(model) or {}
return str(entry.get("litellm_provider") or "")
except Exception: # noqa: BLE001
return ""

289
tests/test_route_advice.py Normal file
View file

@ -0,0 +1,289 @@
"""Per-request backend selection published by an extension.
The property under test throughout is the one that makes this safe to merge:
with nothing published, every path is exactly what it was before.
"""
from __future__ import annotations
import asyncio
from types import SimpleNamespace
import pytest
from headroom.proxy.route_advice import (
BackendResolver,
RouteAdvice,
advice_from,
)
class _Req:
def __init__(self, **state):
self.state = SimpleNamespace(**state)
DEFAULT = object() # stands in for the configured backend
# --- reading what an extension published ------------------------------------
def test_no_extension_means_no_advice():
assert advice_from(_Req()) is None
assert advice_from(object()) is None # no .state at all
assert advice_from(None) is None
def test_advice_is_duck_typed_so_extensions_need_not_import_us():
a = advice_from(
_Req(
headroom_route=SimpleNamespace(
model="moonshot/kimi-k2", provider="moonshot", reason="cheaper"
)
)
)
assert a == RouteAdvice("moonshot/kimi-k2", "moonshot", "cheaper")
def test_an_extension_may_omit_everything_but_the_model():
a = advice_from(_Req(headroom_route=SimpleNamespace(model="gpt-5-mini")))
assert a.model == "gpt-5-mini" and a.provider == ""
def test_malformed_advice_is_ignored_rather_than_raised():
for bad in (
SimpleNamespace(),
SimpleNamespace(model=""),
SimpleNamespace(model=123),
"not an object",
):
assert advice_from(_Req(headroom_route=bad)) is None
def test_advice_needs_a_model():
with pytest.raises(ValueError):
RouteAdvice("")
# --- absent means unchanged, which is the whole safety argument -------------
def test_no_advice_returns_the_configured_backend():
r = BackendResolver(DEFAULT)
assert r.for_request(_Req()) is DEFAULT
def test_no_advice_and_no_configured_backend_stays_none():
"""None is not "no backend", it is the direct-API path. It must survive."""
assert BackendResolver(None).for_request(_Req()) is None
def test_a_native_provider_does_not_switch_backends():
"""Anthropic is the shape the proxy already holds, so a model rewrite is
enough and the extension has already done it."""
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="claude-haiku-4-5", provider="anthropic"))
assert r.for_request(req) is DEFAULT
def test_a_bare_anthropic_model_resolves_its_provider_and_stays_put():
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="claude-haiku-4-5"))
assert r.for_request(req) is DEFAULT
# --- switching, and refusing to switch --------------------------------------
def test_a_foreign_provider_gets_its_own_backend(monkeypatch):
built = []
class FakeBackend:
def __init__(self, provider):
built.append(provider)
self.provider = provider
monkeypatch.setattr(BackendResolver, "_build", lambda self, p: FakeBackend(p))
r = BackendResolver(DEFAULT)
body = {"model": "claude-opus-5"}
req = _Req(headroom_route=SimpleNamespace(model="moonshot/kimi-k2", provider="moonshot"))
got = r.for_request(req, body=body)
assert isinstance(got, FakeBackend) and got.provider == "moonshot"
# The extension could not safely write a foreign model id; we do it.
assert body["model"] == "moonshot/kimi-k2"
def test_backends_are_built_once_per_provider(monkeypatch):
built = []
monkeypatch.setattr(BackendResolver, "_build", lambda self, p: built.append(p) or object())
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="x", provider="moonshot"))
for _ in range(5):
r.for_request(req)
assert built == ["moonshot"], "construction is expensive; cache it"
def test_a_backend_that_will_not_build_falls_back_and_stops_retrying(monkeypatch):
calls = []
monkeypatch.setattr(BackendResolver, "_build", lambda self, p: calls.append(p) or None)
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="x", provider="nope"))
for _ in range(5):
assert r.for_request(req) is DEFAULT
assert calls == ["nope"], "a broken provider must not be retried per request"
def test_a_routing_preference_can_never_take_traffic_down(monkeypatch):
"""Missing credentials, missing optional dependency -- whatever the reason,
the request still has to be served."""
def boom(self, provider):
raise RuntimeError("no credentials")
monkeypatch.setattr(BackendResolver, "_build", boom)
r = BackendResolver(DEFAULT)
req = _Req(headroom_route=SimpleNamespace(model="x", provider="moonshot"))
assert r.for_request(req) is DEFAULT
def test_an_unknown_provider_is_rejected_at_resolve_time():
"""LiteLLMBackend accepts ANY provider string -- the registry falls through
to a generic pass-through -- so a typo silently builds a backend that only
fails later, at request time, with an error pointing nowhere near the typo.
Validate the name up front instead."""
r = BackendResolver(DEFAULT)
assert r._build("definitely-not-a-provider-name") is None
req = _Req(headroom_route=SimpleNamespace(model="x", provider="definitely-not-a-provider-name"))
assert r.for_request(req) is DEFAULT
def test_a_real_provider_name_is_accepted():
assert BackendResolver(DEFAULT)._build("moonshot") is not None
# --- streaming, which is the path agents actually take ----------------------
class _Backend:
"""Records that it, and not some other backend, served the request."""
def __init__(self, name):
self.name = name
self.served = False
async def stream_message(self, body, headers):
self.served = True
return
yield # pragma: no cover -- makes this an async generator
async def stream_openai_message(self, body, headers):
self.served = True
return
yield # pragma: no cover
async def _drive(handler, **kw):
from headroom.proxy.handlers.streaming import StreamingMixin
resp = await StreamingMixin._stream_response_bedrock(
handler,
{"messages": []},
{},
"anthropic",
"m",
"rid",
0,
0,
0,
[],
{},
0.0,
**kw,
)
async for _ in resp.body_iterator:
pass
class _Config:
"""Every proxy flag off. Named individually the list would rot; the test
is about which backend served the request, not about config."""
def __getattr__(self, name):
return False
def _handler(default):
return SimpleNamespace(
anthropic_backend=default,
config=_Config(),
_record_request_outcome=lambda outcome: _noop(),
)
async def _noop():
return None
def test_streaming_honours_the_routed_backend():
"""The non-streaming branch was the easy half. body["model"] has already
been rewritten to a foreign id by the time we get here, so streaming to
the configured backend would send e.g. moonshot/kimi-k2 to Anthropic."""
configured, routed = _Backend("anthropic"), _Backend("moonshot")
asyncio.run(_drive(_handler(configured), backend=routed))
assert routed.served and not configured.served
def test_streaming_without_a_route_uses_the_configured_backend():
configured = _Backend("anthropic")
asyncio.run(_drive(_handler(configured)))
assert configured.served
async def _drive_openai(handler, **kw):
from headroom.proxy.handlers.streaming import StreamingMixin
resp = await StreamingMixin._stream_openai_via_backend(
handler,
{"messages": []},
{},
"m",
"rid",
0.0,
0,
0,
0,
[],
{},
0.0,
**kw,
)
async for _ in resp.body_iterator:
pass
def test_openai_streaming_honours_the_routed_backend():
"""opencode and pi can speak either protocol, so the OpenAI chat path
needs the same treatment as the Anthropic one."""
configured, routed = _Backend("openai"), _Backend("moonshot")
asyncio.run(_drive_openai(_handler(configured), backend=routed))
assert routed.served and not configured.served
def test_openai_streaming_without_a_route_uses_the_configured_backend():
configured = _Backend("openai")
asyncio.run(_drive_openai(_handler(configured)))
assert configured.served
def test_the_resolver_follows_a_reassigned_default():
class H:
anthropic_backend = None
h = H()
from headroom.proxy.route_advice import BackendResolver as BR
first = BR(h.anthropic_backend)
assert first.default is None
h.anthropic_backend = DEFAULT
assert BR(h.anthropic_backend).default is DEFAULT