chore(proxy): use neutral example names for baggage tag mappings (#50011)

This commit is contained in:
Mateo Wang 2026-11-04 14:49:55 +00:00 • committed by GitHub
parent f5e740e69a
commit 80b18d165b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 133 additions and 133 deletions

File diff suppressed because one or more lines are too long

View file

@ -61,10 +61,10 @@ class Rig:
@property
def exact_tag(self) -> str:
return f"usecase-a-{self.marker}"
return f"app-a-{self.marker}"
def usecase_tag(self, value: str) -> str:
return f"usecase-{self.marker}:{value}"
def app_tag(self, value: str) -> str:
return f"app-{self.marker}:{value}"
def tenant_tag(self, value: str) -> str:
return f"tenant-{self.marker}:{value}"
@ -195,7 +195,7 @@ def _routed_deployment(marker: str, provider: Wire, name: str, tag: str) -> dict
def _mapped_models(marker: str, provider: Wire, upstream_url: str) -> list[JsonValue]:
routed: Final = (
_routed_deployment(marker, provider, "tagged", f"usecase-a-{marker}"),
_routed_deployment(marker, provider, "tagged", f"app-a-{marker}"),
_routed_deployment(marker, provider, "default", "default"),
)
realtime: Final[dict[str, JsonValue]] = {
@ -221,8 +221,8 @@ def _mapped_models(marker: str, provider: Wire, upstream_url: str) -> list[JsonV
def _mapped_rules(marker: str, *, tenant_required: bool) -> list[JsonValue]:
return [
{"key": "entrypoint", "value": "a", "tag": f"usecase-a-{marker}"},
{"key": "entrypoint", "tag_prefix": f"usecase-{marker}:"},
{"key": "service", "value": "a", "tag": f"app-a-{marker}"},
{"key": "service", "tag_prefix": f"app-{marker}:"},
{"key": "tenant", "tag_prefix": f"tenant-{marker}:", "required": tenant_required},
{"key": "budgeted", "tag_prefix": f"budgeted-{marker}:", "require_budgeted_tag": True},
{"key": "team+id", "tag_prefix": f"team-{marker}:"},
@ -403,7 +403,7 @@ def test_mapped_baggage_tags_reach_spend_and_callbacks_but_not_requester_metadat
ROUTE_PATHS[route],
_inference_body(route, mapped.chat_model, f"baggage {tenant}", caller_tag),
key=key,
headers={"baggage": f"entrypoint=a, tenant={tenant}"},
headers={"baggage": f"service=a, tenant={tenant}"},
)
assert response.status_code == 200, response.text
assert len(mapped.provider.drain()) == 1
@ -423,7 +423,7 @@ def test_json_string_metadata_gains_exactly_the_mapped_tags(optional: Rig) -> No
key: Final = scenario.key()
optional.provider.drain()
plain: Final = _chat(optional, plain_key, None, fields=fields)
response: Final = _chat(optional, key, f"entrypoint=a,tenant={tenant}", fields=fields)
response: Final = _chat(optional, key, f"service=a,tenant={tenant}", fields=fields)
assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text)
assert len(optional.provider.drain()) == 2
assert _spend_tags(optional.proxy, key) == _spend_tags(optional.proxy, plain_key) | mapped_tags
@ -494,7 +494,7 @@ def test_a2a_spend_row_carries_the_mapped_tags(mapped: Rig) -> None:
},
},
key=key,
headers={"baggage": f"entrypoint=a,tenant={tenant}"},
headers={"baggage": f"service=a,tenant={tenant}"},
)
assert sent.status_code == 200, sent.text
assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key)
@ -514,8 +514,8 @@ def test_spent_mapped_tag_budget_rejects_the_next_request_before_the_provider(
mapped: Rig, key_fields: Mapping[str, JsonValue], fields: Mapping[str, JsonValue], headers: Mapping[str, str]
) -> None:
entry: Final = uuid.uuid4().hex
tag: Final = mapped.usecase_tag(entry)
baggage: Final = f"entrypoint={entry},tenant={entry}"
tag: Final = mapped.app_tag(entry)
baggage: Final = f"service={entry},tenant={entry}"
with mapped.proxy.scenario() as scenario:
key: Final = scenario.key(**key_fields)
mapped.proxy.post("/tag/new", {"name": tag, "max_budget": 0.0001})
@ -528,7 +528,7 @@ def test_spent_mapped_tag_budget_rejects_the_next_request_before_the_provider(
assert blocked.status_code == 422, blocked.text
assert "budget" in blocked.text.lower(), blocked.text
assert mapped.provider.drain() == ()
other: Final = _chat(mapped, key, f"entrypoint=other-{entry},tenant={entry}")
other: Final = _chat(mapped, key, f"service=other-{entry},tenant={entry}")
assert other.status_code == 200, other.text
@ -536,10 +536,10 @@ def test_required_mapping_rejects_a_request_without_its_baggage_value_before_the
with mapped.proxy.scenario() as scenario:
key: Final = scenario.key()
mapped.provider.drain()
rejected: Final = _chat(mapped, key, "entrypoint=a")
rejected: Final = _chat(mapped, key, "service=a")
assert rejected.status_code == 400, rejected.text
assert mapped.provider.drain() == ()
admitted: Final = _chat(mapped, key, f"entrypoint=a,tenant={uuid.uuid4().hex}")
admitted: Final = _chat(mapped, key, f"service=a,tenant={uuid.uuid4().hex}")
assert admitted.status_code == 200, admitted.text
@ -591,7 +591,7 @@ def test_mapped_tags_never_select_a_tag_routed_deployment(mapped: Rig, caller_ta
_chat(
mapped,
key,
f"entrypoint=a,tenant={tenant}",
f"service=a,tenant={tenant}",
model=f"routed-{mapped.marker}",
fields=fields,
**headers,
@ -604,8 +604,8 @@ def test_mapped_tags_never_select_a_tag_routed_deployment(mapped: Rig, caller_ta
def test_mapped_tags_never_trigger_a_tag_scoped_policy_attachment(mapped: Rig) -> None:
entry: Final = uuid.uuid4().hex
tag: Final = mapped.usecase_tag(entry)
baggage: Final = f"entrypoint={entry},tenant={entry}"
tag: Final = mapped.app_tag(entry)
baggage: Final = f"service={entry},tenant={entry}"
proxy: Final = mapped.proxy
with wire_server(lambda _request: Reply(body=GUARDRAIL_VERDICT)) as guardrail, proxy.scenario() as scenario:
key: Final = scenario.key()
@ -662,7 +662,7 @@ def test_required_mapping_refuses_websocket_inference_before_the_upstream(mapped
upstream.get("/__observations").raise_for_status()
status: Final = _handshake_status(
f"ws{proxy_url.removeprefix('http')}/v1/realtime?model=realtime-{mapped.marker}",
{"Authorization": f"Bearer {key}", "baggage": f"entrypoint=a,tenant={uuid.uuid4().hex}"},
{"Authorization": f"Bearer {key}", "baggage": f"service=a,tenant={uuid.uuid4().hex}"},
)
assert status == 403, status
observed: Final = list_value(JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"])
@ -706,7 +706,7 @@ def test_custom_passthrough_spend_row_carries_the_mapped_tags(mapped: Rig) -> No
path,
{"model": mapped.chat_model, "messages": [{"role": "user", "content": tenant}]},
key=key,
headers={"baggage": f"entrypoint=a,tenant={tenant}"},
headers={"baggage": f"service=a,tenant={tenant}"},
)
assert response.status_code == 200, response.text
assert {mapped.exact_tag, mapped.tenant_tag(tenant)} <= _spend_tags(mapped.proxy, key)
@ -777,7 +777,7 @@ def test_passthrough_forwards_the_client_body_unchanged_while_its_spend_row_carr
plain: Final = optional.proxy.request("POST", path, body, key=plain_key)
plain_forwarded: Final = _forwarded_bodies(optional)
response: Final = optional.proxy.request(
"POST", path, body, key=key, headers={"baggage": f"entrypoint=a,tenant={tenant}"}
"POST", path, body, key=key, headers={"baggage": f"service=a,tenant={tenant}"}
)
assert (plain.status_code, response.status_code) == (200, 200), (plain.text, response.text)
assert plain_forwarded == (_upstream_body(entrypoint, body),)
@ -929,7 +929,7 @@ def test_reject_clientside_metadata_tags_with_baggage(
) -> None:
tenant: Final = uuid.uuid4().hex
fields, headers = _rejected_client_tag(client_tag, f"caller-{tenant}")
baggage: Final = f"entrypoint=a,tenant={tenant}" if baggage_maps else f"region={tenant}"
baggage: Final = f"service=a,tenant={tenant}" if baggage_maps else f"region={tenant}"
mapped_tags: Final = frozenset({rejecting.exact_tag, rejecting.tenant_tag(tenant)} if baggage_maps else ())
with rejecting.proxy.scenario() as scenario:
plain_key: Final = scenario.key()
@ -952,7 +952,7 @@ def test_plus_sign_rule_key_matches_either_wire_spelling_but_not_both(
optional: Rig, baggage_keys: tuple[str, ...], maps_team: bool
) -> None:
value: Final = uuid.uuid4().hex
baggage: Final = ",".join(("entrypoint=a", *(f"{baggage_key}={value}" for baggage_key in baggage_keys)))
baggage: Final = ",".join(("service=a", *(f"{baggage_key}={value}" for baggage_key in baggage_keys)))
with optional.proxy.scenario() as scenario:
key: Final = scenario.key()
response: Final = _chat(optional, key, baggage)
@ -991,7 +991,7 @@ def test_baggage_mapping_a_deployment_tag_keeps_the_team_tag_guardrail_mode_firs
@pytest.mark.parametrize("metadata", ["run-123", {"tags": "run-123"}], ids=["non_object", "string_tags"])
@pytest.mark.parametrize(
("header_tag", "baggage"),
[(True, "region=sg"), (True, "entrypoint=a"), (False, "region=sg"), (False, "entrypoint=a")],
[(True, "region=sg"), (True, "service=a"), (False, "region=sg"), (False, "service=a")],
ids=["header_tag_unmapped_baggage", "header_tag_mapped_baggage", "unmapped_baggage", "mapped_baggage"],
)
def test_malformed_metadata_keeps_its_status_under_baggage(

View file

@ -4023,11 +4023,11 @@ async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_t
async def fake_find_many(**kwargs):
if "where" not in kwargs:
return [_tag_registry_row("usecase:spent")]
return [_tag_registry_row("app:spent")]
return [_tag_db_row(name, max_budget=1.0) for name in kwargs["where"]["tag_name"]["in"]]
async def spent_tag_spend(counter_key, fallback_spend, max_budget=None, **kwargs):
return 1.5 if counter_key == "spend:tag:usecase:spent" else fallback_spend
return 1.5 if counter_key == "spend:tag:app:spent" else fallback_spend
async def receive():
return {
@ -4047,7 +4047,7 @@ async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_t
"scheme": "http",
"server": ("localhost", 4000),
"query_string": b"",
"headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=spent")],
"headers": [(b"content-type", b"application/json"), (b"baggage", b"service=spent")],
},
receive,
)
@ -4055,7 +4055,7 @@ async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_t
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
request_body,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
@ -4077,7 +4077,7 @@ async def test_common_checks_rejects_a_spent_mapped_tag_whatever_non_list_tags_t
valid_token=UserAPIKeyAuth(token="test-token", metadata=key_metadata),
request=request,
)
assert exc_info.value.entity_id == "usecase:spent"
assert exc_info.value.entity_id == "app:spent"
@pytest.mark.asyncio
@ -11789,8 +11789,8 @@ def _client_tag_request(client_tags, baggage):
{
"type": "http",
"headers": [
*([(b"baggage", b"entrypoint=a")] if baggage else []),
*([(b"x-litellm-tags", b"usecase:a")] if client_tags == "header" else []),
*([(b"baggage", b"service=a")] if baggage else []),
*([(b"x-litellm-tags", b"app:a")] if client_tags == "header" else []),
],
}
)
@ -11814,7 +11814,7 @@ def _client_tag_rejection(route, data, request):
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/messages", "/v1/responses"])
@pytest.mark.parametrize("client_tags", [None, [], ["usecase:a"], '{"tags":["usecase:a"]}', "header"])
@pytest.mark.parametrize("client_tags", [None, [], ["app:a"], '{"tags":["app:a"]}', "header"])
@pytest.mark.parametrize("client_slot", ["metadata", "litellm_metadata"])
@pytest.mark.parametrize("baggage", [True, False])
async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_it(
@ -11833,7 +11833,7 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
mapped_request,
mapped,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
@ -11843,19 +11843,19 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_
route, unmapped, unmapped_request
)
if client_tags is None and baggage:
assert "usecase:a" in mapped["metadata" if "chat/completions" in route else "litellm_metadata"]["tags"]
assert "app:a" in mapped["metadata" if "chat/completions" in route else "litellm_metadata"]["tags"]
@pytest.mark.asyncio
async def test_baggage_mapping_still_rejects_client_metadata_tags():
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
request = _client_tag_request(["usecase:b"], baggage=True)
data = {"metadata": {"tags": ["usecase:b"]}}
request = _client_tag_request(["app:b"], baggage=True)
data = {"metadata": {"tags": ["app:b"]}}
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)

View file

@ -2203,7 +2203,7 @@ async def test_websocket_inference_is_denied_when_a_required_baggage_mapping_is_
{
"request_tag_mappings": [
{"key": "region", "tag_prefix": "region:"},
{"key": "entrypoint", "tag_prefix": "usecase:", "required": True},
{"key": "service", "tag_prefix": "app:", "required": True},
]
},
),
@ -2227,7 +2227,7 @@ async def test_websocket_invalid_key_is_rejected_before_the_required_baggage_map
with (
patch.dict(
litellm.proxy.proxy_server.general_settings,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "required": True}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "required": True}]},
),
patch(
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth",
@ -2248,7 +2248,7 @@ async def test_websocket_inference_is_allowed_when_no_baggage_mapping_is_require
with (
patch.dict(
litellm.proxy.proxy_server.general_settings,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]},
),
patch(
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth", new=AsyncMock(return_value=authorized)
@ -2273,8 +2273,8 @@ async def _auth_then_serve_chat(monkeypatch, general_settings, body, headers):
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4o", "tags": ["usecase:a"], "mock_response": "hi"},
"model_info": {"id": "usecase-deployment"},
"litellm_params": {"model": "gpt-4o", "tags": ["app:a"], "mock_response": "hi"},
"model_info": {"id": "app-deployment"},
},
{
"model_name": "gpt-4",
@ -2321,30 +2321,30 @@ async def test_auth_header_tag_merge_keeps_json_string_metadata_tags_routable(mo
body = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hi"}],
"metadata": json.dumps({"tags": ["usecase:a"]}),
"metadata": json.dumps({"tags": ["app:a"]}),
}
cached_body, served_by = await _auth_then_serve_chat(
monkeypatch, general_settings, body, [(b"x-litellm-tags", b"other"), (b"baggage", b"region=sg")]
)
assert served_by == "usecase-deployment"
assert cached_body["metadata"] == {"tags": ["usecase:a", "other"]}
assert served_by == "app-deployment"
assert cached_body["metadata"] == {"tags": ["app:a", "other"]}
@pytest.mark.asyncio
@pytest.mark.parametrize("client_metadata", ["run-123", 42, ["x"]])
@pytest.mark.parametrize(
"baggage",
[[], [(b"baggage", b"region=sg")], [(b"baggage", b"entrypoint=a")]],
[[], [(b"baggage", b"region=sg")], [(b"baggage", b"service=a")]],
ids=["no_baggage", "unmapped_baggage", "mapped_baggage"],
)
async def test_non_object_metadata_is_served_like_unset_mappings_whatever_the_baggage_maps(
monkeypatch, client_metadata, baggage
):
body = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "metadata": client_metadata}
headers = [(b"x-litellm-tags", b"usecase:a"), *baggage]
headers = [(b"x-litellm-tags", b"app:a"), *baggage]
unset = await _auth_then_serve_chat(monkeypatch, {}, body, headers)
mapped = await _auth_then_serve_chat(
monkeypatch, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, body, headers
monkeypatch, {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]}, body, headers
)
assert mapped == unset
assert mapped[1] == "usecase-deployment"
assert mapped[1] == "app-deployment"

View file

@ -26,7 +26,7 @@ def test_accounting_enrichment_without_request_state_preserves_legacy_metadata()
@pytest.mark.asyncio
@pytest.mark.parametrize("header", [b"", b"entrypoint=a,entrypoint=b"])
@pytest.mark.parametrize("header", [b"", b"service=a,service=b"])
@pytest.mark.parametrize("client_metadata", [None, '{"requester_metadata":{"guardrails":["qa-guard"]}}'])
async def test_mapping_rejection_keeps_string_metadata_safe_for_auth_failure_logging(header, client_metadata):
from litellm.utils import get_request_guardrails
@ -37,7 +37,7 @@ async def test_mapping_rejection_keeps_string_metadata_safe_for_auth_failure_log
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "required": True}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "required": True}]},
None,
MagicMock(),
)
@ -59,7 +59,7 @@ async def test_invalid_mapper_metadata_maps_like_empty_metadata_and_keeps_the_cl
return {"type": "http.request", "body": json.dumps(body).encode()}
request = Request(
{"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"baggage", b"entrypoint=a")]},
{"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"baggage", b"service=a")]},
receive,
)
data, parse_error = await _read_request_body_deferring_parse_failure(request)
@ -67,11 +67,11 @@ async def test_invalid_mapper_metadata_maps_like_empty_metadata_and_keeps_the_cl
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
assert "usecase:a" in data["litellm_metadata"]["tags"]
assert "app:a" in data["litellm_metadata"]["tags"]
assert "tags" not in data[slot] or slot == "litellm_metadata"
assert get_request_guardrails(data) == []
assert await read_request_body(request) == body
@ -80,60 +80,60 @@ async def test_invalid_mapper_metadata_maps_like_empty_metadata_and_keeps_the_cl
@pytest.mark.parametrize(
"header,value",
[
("entrypoint=a", "a"),
("entrypoint = a", "a"),
("region=sg, entrypoint=customer%20support;version=1", "customer support"),
("service=a", "a"),
("service = a", "a"),
("region=sg, service=customer%20support;version=1", "customer support"),
],
)
def test_variable_mapper_uses_decoded_baggage_value(header, value):
rule = RequestTagMapping(key="entrypoint", tag_prefix="usecase:")
assert resolve_baggage_tag_mappings(header, [rule]) == ((f"usecase:{value}", False),)
rule = RequestTagMapping(key="service", tag_prefix="app:")
assert resolve_baggage_tag_mappings(header, [rule]) == ((f"app:{value}", False),)
@pytest.mark.parametrize(
"header,expected",
[("entrypoint=a", "department:support"), ("entrypoint=b", "department:sales"), ("entrypoint=x", "usecase:x")],
[("service=a", "department:support"), ("service=b", "department:sales"), ("service=x", "app:x")],
)
def test_exact_rules_override_variable_fallback_independent_of_order(header, expected):
rules = [
RequestTagMapping(key="entrypoint", tag_prefix="usecase:"),
RequestTagMapping(key="entrypoint", value="a", tag="department:support"),
RequestTagMapping(key="entrypoint", value="b", tag="department:sales"),
RequestTagMapping(key="service", tag_prefix="app:"),
RequestTagMapping(key="service", value="a", tag="department:support"),
RequestTagMapping(key="service", value="b", tag="department:sales"),
]
assert resolve_baggage_tag_mappings(header, rules) == ((expected, False),)
assert resolve_baggage_tag_mappings(header, list(reversed(rules))) == ((expected, False),)
@pytest.mark.parametrize("header", ["entrypoint=%00", "entrypoint=a%0Ab", "entrypoint=b,entry%70oint=a"])
@pytest.mark.parametrize("header", ["service=%00", "service=a%0Ab", "service=b,serv%69ce=a"])
def test_variable_mapper_skips_control_characters_and_percent_encoded_duplicate_keys(header):
rule = RequestTagMapping(key="entrypoint", tag_prefix="usecase:")
rule = RequestTagMapping(key="service", tag_prefix="app:")
assert resolve_baggage_tag_mappings(header, [rule]) == ()
@pytest.mark.parametrize(
"header,expected",
[
("entry+point=a", (("usecase:a", False),)),
("entry%2Bpoint=a", (("usecase:a", False),)),
("entry+point=a,entry%2Bpoint=b", ()),
("serv+ice=a", (("app:a", False),)),
("serv%2Bice=a", (("app:a", False),)),
("serv+ice=a,serv%2Bice=b", ()),
],
)
def test_key_with_a_plus_matches_both_wire_spellings_and_rejects_both_at_once(header, expected):
rule = RequestTagMapping(key="entry+point", tag_prefix="usecase:")
rule = RequestTagMapping(key="serv+ice", tag_prefix="app:")
assert resolve_baggage_tag_mappings(header, [rule]) == expected
def test_exact_rules_work_without_a_variable_rule():
rule = RequestTagMapping(key="entrypoint", value="a", tag="usecase:manual")
assert resolve_baggage_tag_mappings("entrypoint=a", [rule]) == (("usecase:manual", False),)
assert resolve_baggage_tag_mappings("entrypoint=x", [rule]) == ()
rule = RequestTagMapping(key="service", value="a", tag="app:manual")
assert resolve_baggage_tag_mappings("service=a", [rule]) == (("app:manual", False),)
assert resolve_baggage_tag_mappings("service=x", [rule]) == ()
@pytest.mark.parametrize(
"header", ["", "region=sg", "entrypoint=", "entrypoint=a,entrypoint=b", "entrypoint=b,entry%70oint=a", "entrypoint=x"]
"header", ["", "region=sg", "service=", "service=a,service=b", "service=b,serv%69ce=a", "service=x"]
)
def test_required_mapping_rejects_missing_ambiguous_or_unmapped_values(header):
rule = RequestTagMapping(key="entrypoint", value="a", tag="usecase:a", required=True)
rule = RequestTagMapping(key="service", value="a", tag="app:a", required=True)
with pytest.raises(ProxyException) as error:
resolve_baggage_tag_mappings(header, [rule])
assert error.value.code == "400"
@ -142,10 +142,10 @@ def test_required_mapping_rejects_missing_ambiguous_or_unmapped_values(header):
@pytest.mark.parametrize(
"rule",
[
{"key": "entrypoint", "tag": "usecase:a"},
{"key": "entrypoint", "value": "a", "tag_prefix": "usecase:"},
{"key": "entrypoint", "value": "a", "tag": "usecase:a", "tag_prefix": "usecase:"},
{"key": "", "tag_prefix": "usecase:"},
{"key": "service", "tag": "app:a"},
{"key": "service", "value": "a", "tag_prefix": "app:"},
{"key": "service", "value": "a", "tag": "app:a", "tag_prefix": "app:"},
{"key": "", "tag_prefix": "app:"},
],
)
def test_invalid_rule_is_rejected_during_config_validation(rule):
@ -158,8 +158,8 @@ def test_duplicate_fallbacks_are_rejected():
ConfigGeneralSettings.model_validate(
{
"request_tag_mappings": [
{"key": "entrypoint", "tag_prefix": "one:"},
{"key": "entrypoint", "tag_prefix": "two:"},
{"key": "service", "tag_prefix": "one:"},
{"key": "service", "tag_prefix": "two:"},
]
}
)
@ -171,23 +171,23 @@ def test_duplicate_fallbacks_are_rejected():
[("/chat/completions", "metadata"), ("/v1/messages", "litellm_metadata"), ("/v1/responses", "litellm_metadata")],
)
async def test_mapped_tag_merges_with_existing_tags_in_the_budget_metadata_slot(route, slot):
data = {slot: {"tags": ["existing", "usecase:a"]}}
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
data = {slot: {"tags": ["existing", "app:a"]}}
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(data, route)
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
assert data[slot]["tags"] == ["existing", "usecase:a"]
assert data[slot]["tags"] == ["existing", "app:a"]
@pytest.mark.asyncio
async def test_budgeted_mapper_requires_a_database():
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
config = {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]}
with pytest.raises(ProxyException) as error:
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, {}, config, None, MagicMock()
@ -197,10 +197,10 @@ async def test_budgeted_mapper_requires_a_database():
@pytest.mark.asyncio
async def test_budgeted_mapper_answers_503_when_the_tag_budget_lookup_fails():
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
prisma_client = MagicMock()
prisma_client.db.litellm_tagtable.find_many = AsyncMock(side_effect=Exception("database unreachable"))
config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}
config = {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]}
with pytest.raises(ProxyException, match="could not read tag budgets") as error:
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, {"metadata": {}}, config, prisma_client, UserApiKeyCache()
@ -210,16 +210,16 @@ async def test_budgeted_mapper_answers_503_when_the_tag_budget_lookup_fails():
@pytest.mark.asyncio
async def test_budgeted_mapper_serves_through_a_db_outage_when_requests_are_allowed_on_db_unavailable():
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
prisma_client = MagicMock()
prisma_client.db.litellm_tagtable.find_many = AsyncMock(side_effect=Exception("database unreachable"))
data = {"metadata": {}}
config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}
config = {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]}
with patch.dict("litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": True}):
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, config, prisma_client, UserApiKeyCache()
)
assert data["metadata"]["tags"] == ["usecase:a"]
assert data["metadata"]["tags"] == ["app:a"]
async def _cache_with_tag(tag: LiteLLM_TagTable) -> UserApiKeyCache:
@ -234,11 +234,11 @@ async def _cache_with_tag(tag: LiteLLM_TagTable) -> UserApiKeyCache:
[None, LiteLLM_BudgetTable(budget_id="qa-budget", max_budget=None, tpm_limit=10)],
)
async def test_budgeted_mapper_rejects_a_tag_without_a_max_budget(budget_table):
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a")]})
request = Request({"type": "http", "headers": [(b"baggage", b"service=a")]})
data = {"metadata": {}}
cache = await _cache_with_tag(LiteLLM_TagTable(tag_name="usecase:a", litellm_budget_table=budget_table))
config = {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True}]}
with pytest.raises(ProxyException, match="usecase:a requires an existing tag budget") as error:
cache = await _cache_with_tag(LiteLLM_TagTable(tag_name="app:a", litellm_budget_table=budget_table))
config = {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True}]}
with pytest.raises(ProxyException, match="app:a requires an existing tag budget") as error:
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, config, MagicMock(), cache
)
@ -248,23 +248,23 @@ async def test_budgeted_mapper_rejects_a_tag_without_a_max_budget(budget_table):
@pytest.mark.asyncio
async def test_budgeted_mapper_admits_a_tag_with_a_max_budget_and_unbudgeted_rules_need_none():
request = Request({"type": "http", "headers": [(b"baggage", b"entrypoint=a,region=sg")]})
request = Request({"type": "http", "headers": [(b"baggage", b"service=a,region=sg")]})
data = {"metadata": {}}
cache = await _cache_with_tag(
LiteLLM_TagTable(
tag_name="usecase:a", litellm_budget_table=LiteLLM_BudgetTable(budget_id="qa-budget", max_budget=1.0)
tag_name="app:a", litellm_budget_table=LiteLLM_BudgetTable(budget_id="qa-budget", max_budget=1.0)
)
)
config = {
"request_tag_mappings": [
{"key": "entrypoint", "tag_prefix": "usecase:", "require_budgeted_tag": True},
{"key": "service", "tag_prefix": "app:", "require_budgeted_tag": True},
{"key": "region", "tag_prefix": "region:"},
]
}
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, config, MagicMock(), cache
)
assert data["metadata"]["tags"] == ["usecase:a", "region:sg"]
assert data["metadata"]["tags"] == ["app:a", "region:sg"]
@pytest.mark.parametrize(
@ -325,7 +325,7 @@ async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fiel
"scheme": "http",
"server": ("localhost", 4000),
"query_string": query,
"headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a")],
"headers": [(b"content-type", b"application/json"), (b"baggage", b"service=a")],
},
receive,
)
@ -337,7 +337,7 @@ async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fiel
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request,
data,
{"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]},
{"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]},
None,
MagicMock(),
)
@ -350,7 +350,7 @@ async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fiel
proxy_config=ProxyConfig(),
general_settings={},
)
assert enriched[slot]["tags"] == ["usecase:a"]
assert enriched[slot]["tags"] == ["app:a"]
assert "key" not in enriched
assert "api-version" not in enriched
assert "key" not in enriched["proxy_server_request"]["body"]
@ -360,7 +360,7 @@ async def test_mapping_survives_cached_body_read_without_leaking_auth_query_fiel
@pytest.mark.asyncio
@pytest.mark.parametrize(
"client_metadata",
[{"trace_name": "qa"}, {"trace_name": "qa", "tags": ["client:x"]}, {"trace_name": "qa", "tags": ["usecase:a"]}],
[{"trace_name": "qa"}, {"trace_name": "qa", "tags": ["client:x"]}, {"trace_name": "qa", "tags": ["app:a"]}],
)
async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata):
from litellm.proxy._types import UserAPIKeyAuth
@ -380,14 +380,14 @@ async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata):
"scheme": "http",
"server": ("localhost", 4000),
"query_string": b"",
"headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a")],
"headers": [(b"content-type", b"application/json"), (b"baggage", b"service=a")],
},
receive,
)
data, parse_error = await _read_request_body_deferring_parse_failure(request)
assert parse_error is None
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, None, MagicMock()
request, data, {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]}, None, MagicMock()
)
enriched = await add_litellm_data_to_request(
data=await read_request_body(request),
@ -397,7 +397,7 @@ async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata):
general_settings={},
)
requester_metadata = enriched["metadata"]["requester_metadata"]
assert "usecase:a" in enriched["metadata"]["tags"]
assert "app:a" in enriched["metadata"]["tags"]
assert requester_metadata["trace_name"] == "qa"
assert requester_metadata.get("tags") == client_metadata.get("tags")
assert BAGGAGE_MAPPED_TAGS_METADATA_KEY not in requester_metadata
@ -407,9 +407,9 @@ async def test_mapped_tags_stay_out_of_requester_metadata(client_metadata):
@pytest.mark.parametrize(
"headers,body_fields",
[
([(b"x-litellm-tags", b"usecase:a")], {}),
([], {"tags": ["usecase:a"]}),
([], {"litellm_metadata": {"tags": ["usecase:a"]}}),
([(b"x-litellm-tags", b"app:a")], {}),
([], {"tags": ["app:a"]}),
([], {"litellm_metadata": {"tags": ["app:a"]}}),
],
ids=["header", "body", "litellm_metadata"],
)
@ -432,14 +432,14 @@ async def test_client_sent_tag_still_routes_when_baggage_maps_to_it(headers, bod
"scheme": "http",
"server": ("localhost", 4000),
"query_string": b"",
"headers": [(b"content-type", b"application/json"), (b"baggage", b"entrypoint=a"), *headers],
"headers": [(b"content-type", b"application/json"), (b"baggage", b"service=a"), *headers],
},
receive,
)
data, parse_error = await _read_request_body_deferring_parse_failure(request)
assert parse_error is None
await LiteLLMProxyRequestSetup.apply_baggage_tag_mappings_pre_auth(
request, data, {"request_tag_mappings": [{"key": "entrypoint", "tag_prefix": "usecase:"}]}, None, MagicMock()
request, data, {"request_tag_mappings": [{"key": "service", "tag_prefix": "app:"}]}, None, MagicMock()
)
enriched = await add_litellm_data_to_request(
data=await read_request_body(request),
@ -452,8 +452,8 @@ async def test_client_sent_tag_still_routes_when_baggage_maps_to_it(headers, bod
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4o", "tags": ["usecase:a"]},
"model_info": {"id": "usecase-deployment"},
"litellm_params": {"model": "gpt-4o", "tags": ["app:a"]},
"model_info": {"id": "app-deployment"},
},
{
"model_name": "gpt-4",
@ -470,4 +470,4 @@ async def test_client_sent_tag_still_routes_when_baggage_maps_to_it(headers, bod
metadata=dict(enriched["metadata"]),
mock_response="hi",
)
assert response._hidden_params["model_id"] == "usecase-deployment"
assert response._hidden_params["model_id"] == "app-deployment"

View file

@ -5,8 +5,8 @@ import { describe, expect, it, vi } from "vitest";
import { RequestTagMapping, RequestTagMappingsEditor, requestTagMappingsAreValid } from "./RequestTagMappings";
const exact: RequestTagMapping = { header: "baggage", key: "entrypoint", value: "a", tag: "support" };
const variable: RequestTagMapping = { header: "baggage", key: "entrypoint", tag_prefix: "usecase:" };
const exact: RequestTagMapping = { header: "baggage", key: "service", value: "a", tag: "support" };
const variable: RequestTagMapping = { header: "baggage", key: "service", tag_prefix: "app:" };
function Editor({
initial,
@ -47,20 +47,20 @@ describe("RequestTagMappingsEditor", () => {
expect(screen.getByRole("combobox", { name: "Mapping 2 mode" })).toHaveTextContent("Any value");
expect(screen.getByLabelText("Mapping 1 match value")).toHaveValue("a");
expect(screen.getByLabelText("Mapping 1 tag")).toHaveValue("support");
expect(screen.getByLabelText("Mapping 2 tag prefix")).toHaveValue("usecase:");
expect(screen.getByText("entrypoint=a → support")).toBeInTheDocument();
expect(screen.getByText("entrypoint=x → usecase:x")).toBeInTheDocument();
expect(screen.getByLabelText("Mapping 2 tag prefix")).toHaveValue("app:");
expect(screen.getByText("service=a → support")).toBeInTheDocument();
expect(screen.getByText("service=x → app:x")).toBeInTheDocument();
});
it("adds and edits a new exact value mapping", async () => {
const changed = vi.fn();
render(<Editor initial={[]} changed={changed} />);
await userEvent.click(screen.getByRole("button", { name: "Add mapping" }));
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } });
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "service" } });
fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "b" } });
fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "sales" } });
expect(screen.getByText("entrypoint=b → sales")).toBeInTheDocument();
expect(changed).toHaveBeenLastCalledWith([{ header: "baggage", key: "entrypoint", value: "b", tag: "sales" }]);
expect(screen.getByText("service=b → sales")).toBeInTheDocument();
expect(changed).toHaveBeenLastCalledWith([{ header: "baggage", key: "service", value: "b", tag: "sales" }]);
});
it("changes an exact mapping to any-value mode and removes its exact match fields", async () => {
@ -70,7 +70,7 @@ describe("RequestTagMappingsEditor", () => {
await userEvent.click(await screen.findByRole("option", { name: "Any value" }));
expect(screen.queryByLabelText("Mapping 1 match value")).not.toBeInTheDocument();
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue("");
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveAttribute("placeholder", "usecase:");
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveAttribute("placeholder", "app:");
expect(changed).toHaveBeenLastCalledWith([
expect.objectContaining({ value: undefined, tag: undefined, tag_prefix: "" }),
]);
@ -81,9 +81,9 @@ describe("RequestTagMappingsEditor", () => {
render(<Editor initial={[]} changed={vi.fn()} saved={saved} />);
await userEvent.click(screen.getByRole("button", { name: "Add mapping" }));
expect(screen.getByLabelText("Mapping 1 baggage key")).toHaveValue("");
expect(screen.getByLabelText("Mapping 1 baggage key")).toHaveAttribute("placeholder", "entrypoint");
expect(screen.getByLabelText("Mapping 1 baggage key")).toHaveAttribute("placeholder", "service");
expect(screen.getByText("Enter a baggage key")).toBeInTheDocument();
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } });
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "service" } });
expect(screen.getByText("Enter a match value")).toBeInTheDocument();
await userEvent.click(screen.getByRole("button", { name: "Update" }));
fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "b" } });
@ -93,7 +93,7 @@ describe("RequestTagMappingsEditor", () => {
fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "sales" } });
expect(screen.queryByText(/^Enter a/)).not.toBeInTheDocument();
await userEvent.click(screen.getByRole("button", { name: "Update" }));
expect(saved).toHaveBeenCalledWith([{ header: "baggage", key: "entrypoint", value: "b", tag: "sales" }]);
expect(saved).toHaveBeenCalledWith([{ header: "baggage", key: "service", value: "b", tag: "sales" }]);
});
it("blocks saving an any-value mapping until the admin types a tag prefix", async () => {
@ -105,7 +105,7 @@ describe("RequestTagMappingsEditor", () => {
expect(screen.getByRole("button", { name: "Update" })).toBeDisabled();
fireEvent.change(screen.getByLabelText("Mapping 1 tag prefix"), { target: { value: "team:" } });
await userEvent.click(screen.getByRole("button", { name: "Update" }));
expect(saved).toHaveBeenCalledWith([expect.objectContaining({ key: "entrypoint", tag_prefix: "team:" })]);
expect(saved).toHaveBeenCalledWith([expect.objectContaining({ key: "service", tag_prefix: "team:" })]);
});
it("blocks saving a baggage key that is empty or not a header token", async () => {
@ -139,7 +139,7 @@ describe("RequestTagMappingsEditor", () => {
const changed = vi.fn();
render(<Editor initial={[exact, variable]} changed={changed} />);
await userEvent.click(screen.getByRole("button", { name: "Delete mapping 1" }));
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue("usecase:");
expect(screen.getByLabelText("Mapping 1 tag prefix")).toHaveValue("app:");
expect(changed).toHaveBeenLastCalledWith([variable]);
});
});

View file

@ -111,7 +111,7 @@ export function RequestTagMappingsEditor({ value, onChange, disabled = false }:
{t("routerSettings.requestTagMappings.baggageKey")}
<Input
aria-label={t("routerSettings.requestTagMappings.baggageKeyAriaLabel", { number: index + 1 })}
placeholder="entrypoint"
placeholder="service"
value={rule.key}
onChange={(event) => update(index, { key: event.target.value })}
/>
@ -145,7 +145,7 @@ export function RequestTagMappingsEditor({ value, onChange, disabled = false }:
{t("routerSettings.requestTagMappings.tagPrefix")}
<Input
aria-label={t("routerSettings.requestTagMappings.tagPrefixAriaLabel", { number: index + 1 })}
placeholder="usecase:"
placeholder="app:"
value={rule.tag_prefix ?? ""}
onChange={(event) => update(index, { tag_prefix: event.target.value })}
/>

View file

@ -310,12 +310,12 @@ it("keeps a new baggage mapping unsaved until every field the backend requires i
name: "Update",
});
expect(update).toBeDisabled();
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "entrypoint" } });
fireEvent.change(screen.getByLabelText("Mapping 1 baggage key"), { target: { value: "service" } });
fireEvent.change(screen.getByLabelText("Mapping 1 match value"), { target: { value: "a" } });
expect(update).toBeDisabled();
fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "usecase:support" } });
fireEvent.change(screen.getByLabelText("Mapping 1 tag"), { target: { value: "app:support" } });
await user.click(update);
expect(vi.mocked(updateConfigFieldSetting).mock.calls).toEqual([
["token", "request_tag_mappings", [{ header: "baggage", key: "entrypoint", value: "a", tag: "usecase:support" }]],
["token", "request_tag_mappings", [{ header: "baggage", key: "service", value: "a", tag: "app:support" }]],
]);
});