fix(auto-router compression): honor the policy on the SDK path and on re-save

The router reused the model hop's compression for routing whenever both hops named
the same guardrail, on the premise that arm_pre_call had already run it. Only the
proxy calls arm_pre_call, so through the SDK nothing armed the guardrail and nothing
had compressed anything: the shortcut skipped routing compression too and served the
request with no compression on either hop. The reuse is now conditional on the model
hop actually having been armed.

The Admin UI hydrated an absent auto_router_model_compression as same-as-routing,
while the backend reads it as no model-hop compression. Opening a router configured
with only auto_router_routing_compression and saving any unrelated edit wrote the
routing guardrail onto the model hop, silently starting to compress the model call.

Both carry a regression test that fails when the fix is reverted.
This commit is contained in:
moe-berri 2026-09-05 09:31:18 -07:00
parent f03f82381e
commit 00b49ccc8b
5 changed files with 347 additions and 666 deletions

View file

@ -45,6 +45,19 @@ def suppressed_compression_guardrails() -> frozenset[str]:
return _suppressed_compression_guardrails.get()
# Whether `arm_pre_call` actually armed a model-side compression guardrail for this
# request. Only the proxy calls `arm_pre_call`, so on the SDK path nothing arms and
# nothing compresses; the router must not assume the model hop already ran.
_model_hop_armed: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar(
"litellm_auto_router_model_hop_armed", default=False
)
def model_hop_compression_armed() -> bool:
"""True when this request's model-side compression guardrail was actually armed."""
return _model_hop_armed.get()
@dataclass(frozen=True, slots=True)
class AutoRouterCompressionPolicy:
"""An auto router's compression choice for each hop. ``None`` means no compression."""
@ -147,6 +160,7 @@ async def arm_pre_call(
guardrail the policy names (if any) even when it isn't ``default_on``.
"""
_suppressed_compression_guardrails.set(frozenset())
_model_hop_armed.set(False)
if llm_router is None:
return
@ -179,6 +193,7 @@ async def arm_pre_call(
)
if policy.model is not None:
_model_hop_armed.set(True)
_, metadata = get_or_create_metadata_bucket(data)
requested: Final = metadata.get("guardrails")
existing: Final = tuple(requested) if isinstance(requested, (list, tuple)) else ()

View file

@ -13039,6 +13039,7 @@ class Router:
from litellm.proxy.guardrails.auto_router_compression import (
messages_for_routing,
model_hop_compression_armed,
policy_for_model,
team_id_from_request,
)
@ -13057,8 +13058,13 @@ class Router:
# (arm_pre_call armed it whether or not it is `default_on`); reuse that result
# for routing too instead of paying for a second compression call against the
# same content.
#
# Only the proxy calls arm_pre_call, so that reuse is conditional on it having
# actually run: on the SDK path nothing arms the model hop and nothing has
# compressed anything, and taking the shortcut there would skip both hops and
# silently serve the request with no compression at all.
needs_independent_routing_compression: Final = compression_policy is not None and not (
compression_policy.is_same and compression_policy.model is not None
compression_policy.is_same and compression_policy.model is not None and model_hop_compression_armed()
)
routing_messages: Final = (
await messages_for_routing(policy=compression_policy, messages=messages, request_kwargs=request_kwargs)

File diff suppressed because it is too large Load diff

View file

@ -80,9 +80,20 @@ describe("hydrateAutoRouterCompression", () => {
expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "headroom-b" });
});
it("treats a missing model key as same-as-routing", () => {
it("treats a missing model key as no model-hop compression, not same-as-routing", () => {
const state = hydrateAutoRouterCompression({ auto_router_routing_compression: "headroom-a" });
expect(state).toEqual({ routing: "headroom-a", sameAsRouting: true, model: undefined });
expect(state).toEqual({ routing: "headroom-a", sameAsRouting: false, model: "none" });
});
it("re-saving a routing-only config leaves the model hop uncompressed", () => {
// Regression: the backend reads an absent model key as no model-hop compression.
// Hydrating it as same-as-routing made opening the router and saving any unrelated
// edit write the routing guardrail onto the model hop, so the model call silently
// started receiving compressed messages.
const stored = { auto_router_routing_compression: "headroom-a" };
const rebuilt = buildAutoRouterCompressionParams(hydrateAutoRouterCompression(stored));
expect(rebuilt.auto_router_model_compression).toBe("none");
expect(rebuilt.auto_router_model_compression).not.toBe("headroom-a");
});
it("round-trips through buildAutoRouterCompressionParams", () => {

View file

@ -53,7 +53,11 @@ export const hydrateAutoRouterCompression = (litellmParams: {
const routing = litellmParams.auto_router_routing_compression ?? undefined;
if (routing === undefined) return DEFAULT_AUTO_ROUTER_COMPRESSION;
const model = litellmParams.auto_router_model_compression ?? undefined;
const sameAsRouting = model === undefined || model === routing;
// An absent model key is no model-hop compression, not same-as-routing: the backend
// reads it as None (policy_from_litellm_params). Hydrating it as same-as-routing
// would make re-saving an unrelated edit write the routing guardrail onto the model
// hop and silently start compressing the model call.
const model = litellmParams.auto_router_model_compression ?? NO_COMPRESSION;
const sameAsRouting = model === routing;
return { routing, sameAsRouting, model: sameAsRouting ? undefined : model };
};