forked from retoor/devplacepy
Trust only the upstream X-Gateway-Model header in the AI gateway
An upstream the gateway forwards to may itself emit X-Gateway-* headers (e.g. another DevPlace-style gateway), which can collide with the ones about to be built for the response. Only X-Gateway-Model is ever trusted from upstream and relayed as-is - it is the one field an upstream can legitimately know better than we do (it may have resolved an alias or served a different pinned version). Every other header (cost, tokens, latency, context, app-reference) is always our own measurement and is never overwritten, since blending in an upstream's own accounting would corrupt the usage ledger's per-model rollups and the quota math built on top of it. usage.upstream_reported_model() extracts that one header defensively (case-insensitive lookup, rejects anything oversized or containing a control character) and gateway._apply_served_model() applies it, display- only, at the tail of every response-header build across chat, streaming, embeddings, images, and passthrough. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BWJy6PrMMt5hwWxQwia2rd
This commit is contained in:
@@ -11,6 +11,7 @@ from devplacepy.services.openai_gateway.usage import (
|
||||
new_usage_totals,
|
||||
normalize_usage,
|
||||
parse_usage_headers,
|
||||
upstream_reported_model,
|
||||
usage_metric_cards,
|
||||
)
|
||||
|
||||
@@ -248,3 +249,79 @@ def test_validate_app_reference_strips_whitespace():
|
||||
|
||||
assert _validate_app_reference(" my-app ") == "my-app"
|
||||
assert _validate_app_reference("\tdevplace\t") == "devplace"
|
||||
|
||||
|
||||
class DictOnlyHeaders(dict):
|
||||
"""A plain mapping with no case-insensitive .get, like a hand-rolled
|
||||
fake response's headers dict in a test double - exercises the
|
||||
manual case-insensitive fallback scan."""
|
||||
|
||||
|
||||
def test_upstream_reported_model_reads_exact_case_via_get():
|
||||
assert upstream_reported_model({"X-Gateway-Model": "vendor/served-model"}) == (
|
||||
"vendor/served-model"
|
||||
)
|
||||
|
||||
|
||||
def test_upstream_reported_model_is_case_insensitive_via_real_headers():
|
||||
import httpx
|
||||
|
||||
headers = httpx.Headers({"x-gateway-model": "vendor/lowercase-model"})
|
||||
assert upstream_reported_model(headers) == "vendor/lowercase-model"
|
||||
|
||||
|
||||
def test_upstream_reported_model_falls_back_to_manual_scan_for_plain_dicts():
|
||||
headers = DictOnlyHeaders({"x-GATEWAY-model": "vendor/mixed-case"})
|
||||
assert upstream_reported_model(headers) == "vendor/mixed-case"
|
||||
|
||||
|
||||
def test_upstream_reported_model_strips_surrounding_whitespace():
|
||||
assert upstream_reported_model({"X-Gateway-Model": " vendor/spaced "}) == (
|
||||
"vendor/spaced"
|
||||
)
|
||||
|
||||
|
||||
def test_upstream_reported_model_missing_header_returns_none():
|
||||
assert upstream_reported_model({}) is None
|
||||
assert upstream_reported_model({"Content-Type": "application/json"}) is None
|
||||
|
||||
|
||||
def test_upstream_reported_model_handles_falsy_and_malformed_containers():
|
||||
assert upstream_reported_model(None) is None
|
||||
assert upstream_reported_model({}) is None
|
||||
assert upstream_reported_model(object()) is None
|
||||
assert upstream_reported_model(["not", "a", "mapping"]) is None
|
||||
|
||||
|
||||
def test_upstream_reported_model_rejects_empty_or_whitespace_only():
|
||||
assert upstream_reported_model({"X-Gateway-Model": ""}) is None
|
||||
assert upstream_reported_model({"X-Gateway-Model": " "}) is None
|
||||
|
||||
|
||||
def test_upstream_reported_model_rejects_non_string_value():
|
||||
assert upstream_reported_model({"X-Gateway-Model": 12345}) is None
|
||||
assert upstream_reported_model({"X-Gateway-Model": None}) is None
|
||||
|
||||
|
||||
def test_upstream_reported_model_rejects_oversized_value():
|
||||
huge = "x" * 201
|
||||
assert upstream_reported_model({"X-Gateway-Model": huge}) is None
|
||||
exactly_max = "x" * 200
|
||||
assert upstream_reported_model({"X-Gateway-Model": exactly_max}) == exactly_max
|
||||
|
||||
|
||||
def test_upstream_reported_model_rejects_control_characters():
|
||||
assert upstream_reported_model({"X-Gateway-Model": "vendor\r\nInjected: header"}) is None
|
||||
assert upstream_reported_model({"X-Gateway-Model": "vendor\nmodel"}) is None
|
||||
assert upstream_reported_model({"X-Gateway-Model": "vendor\x00model"}) is None
|
||||
|
||||
|
||||
def test_upstream_reported_model_get_raising_never_propagates():
|
||||
class HostileHeaders:
|
||||
def get(self, key, default=None):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
def items(self):
|
||||
raise RuntimeError("boom too")
|
||||
|
||||
assert upstream_reported_model(HostileHeaders()) is None
|
||||
|
||||
Reference in New Issue
Block a user