From f932e292c523f2d238c4b7361c5168776dc0d042 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 2 Oct 2026 19:38:44 +0000 Subject: [PATCH] fix(proxy): carry key, team and project tags into pass-through spend logs (#42662) * fix(proxy): carry key, team and project tags into pass-through spend logs Pass-through endpoints built their request metadata without the key, team and project controls that native routes apply, so spend rows for configured routes and provider pass-throughs like /anthropic dropped the key, team and project tags and the key and team spend_logs_metadata. The native team and project controls now live in a shared helper that both paths call, key spend_logs_metadata is copied instead of aliased from the cached key, client metadata cannot overwrite user_api_key_ fields, and header tags dedupe with the same merge used on native routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): drop covers markers from pass-through tag tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): type the shared team and project control helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): validate pass-through endpoint list instead of suppressing pyright Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover body, streaming, hostile, forged and native cells for pass-through tags Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: mock pass-through request url as httpx.URL after rebase on main * refactor(proxy): merge spend_logs_metadata sources without a stacked comprehension * refactor(proxy): name the team and request spend_logs_metadata merge --------- Co-authored-by: ryan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/litellm_pre_call_utils.py | 83 ++-- .../batch_attribution.py | 4 +- .../pass_through_endpoints.py | 29 +- .../spend/test_passthrough_request_tags.py | 421 ++++++++++++++++++ .../test_pass_through_endpoints.py | 55 +++ 5 files changed, 551 insertions(+), 41 deletions(-) create mode 100644 tests/integration/spend/test_passthrough_request_tags.py diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d6daf6ebe59..a809f53aa85 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1760,7 +1760,7 @@ class LiteLLMProxyRequestSetup: ): # don't override k-v pair sent by request (user request) data[_metadata_variable_name]["spend_logs_metadata"][key] = value else: - data[_metadata_variable_name]["spend_logs_metadata"] = key_metadata["spend_logs_metadata"] + data[_metadata_variable_name]["spend_logs_metadata"] = dict(key_metadata["spend_logs_metadata"]) ## KEY-LEVEL DISABLE FALLBACKS if "disable_fallbacks" in key_metadata and isinstance(key_metadata["disable_fallbacks"], bool): @@ -1777,6 +1777,53 @@ class LiteLLMProxyRequestSetup: ) return data + @staticmethod + def add_team_and_project_level_controls( + user_api_key_dict: UserAPIKeyAuth, metadata: dict[str, object] + ) -> dict[str, object]: + team_metadata: Final = user_api_key_dict.team_metadata or MappingProxyType({}) + project_metadata: Final = user_api_key_dict.project_metadata or MappingProxyType({}) + request_tags: Final = metadata.get("tags") + team_tags: Final = team_metadata.get("tags") + project_tags: Final = project_metadata.get("tags") + disable_global_guardrails: Final = team_metadata.get("disable_global_guardrails") + opted_out_global_guardrails: Final = team_metadata.get("opted_out_global_guardrails") + spend_logs_metadata: Final = LiteLLMProxyRequestSetup._merge_spend_logs_metadata( + team_spend_logs_metadata=team_metadata.get("spend_logs_metadata"), + request_spend_logs_metadata=metadata.get("spend_logs_metadata"), + ) + tags: Final = LiteLLMProxyRequestSetup._merge_tags( + request_tags=LiteLLMProxyRequestSetup._merge_tags( + request_tags=request_tags if isinstance(request_tags, list) else None, + tags_to_add=team_tags if isinstance(team_tags, list) else None, + ), + tags_to_add=project_tags if isinstance(project_tags, list) else None, + ) + controls: Final = ( + ("tags", tags or None), + ("spend_logs_metadata", spend_logs_metadata), + ( + "disable_global_guardrails", + disable_global_guardrails if isinstance(disable_global_guardrails, bool) else None, + ), + ( + "opted_out_global_guardrails", + opted_out_global_guardrails if isinstance(opted_out_global_guardrails, list) else None, + ), + ) + return {**metadata, **{key: value for key, value in controls if value is not None}} + + @staticmethod + def _merge_spend_logs_metadata( + team_spend_logs_metadata: object, request_spend_logs_metadata: object + ) -> dict[str, object] | None: + """Team values as defaults, the request's own values win on the same key. None when neither is a dict""" + team_values: Final = team_spend_logs_metadata if isinstance(team_spend_logs_metadata, dict) else None + request_values: Final = request_spend_logs_metadata if isinstance(request_spend_logs_metadata, dict) else None + if team_values is None and request_values is None: + return None + return {**(team_values or {}), **(request_values or {})} + @staticmethod def _merge_tags(request_tags: list | None, tags_to_add: list | None) -> list: """ @@ -2312,38 +2359,12 @@ async def add_litellm_data_to_request( data=data, _metadata_variable_name=_metadata_variable_name, ) - ## TEAM-LEVEL SPEND LOGS/TAGS + data[_metadata_variable_name] = LiteLLMProxyRequestSetup.add_team_and_project_level_controls( + user_api_key_dict=user_api_key_dict, + metadata=data[_metadata_variable_name], + ) team_metadata: Final = user_api_key_dict.team_metadata or {} - if "tags" in team_metadata and team_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=team_metadata["tags"], - ) - if "disable_global_guardrails" in team_metadata and isinstance(team_metadata["disable_global_guardrails"], bool): - data[_metadata_variable_name]["disable_global_guardrails"] = team_metadata["disable_global_guardrails"] - if "opted_out_global_guardrails" in team_metadata and isinstance( - team_metadata["opted_out_global_guardrails"], list - ): - data[_metadata_variable_name]["opted_out_global_guardrails"] = team_metadata["opted_out_global_guardrails"] - if "spend_logs_metadata" in team_metadata and isinstance(team_metadata["spend_logs_metadata"], dict): - if "spend_logs_metadata" in data[_metadata_variable_name] and isinstance( - data[_metadata_variable_name]["spend_logs_metadata"], dict - ): - for key, value in team_metadata["spend_logs_metadata"].items(): - if ( - key not in data[_metadata_variable_name]["spend_logs_metadata"] - ): # don't override k-v pair sent by request (user request) - data[_metadata_variable_name]["spend_logs_metadata"][key] = value - else: - data[_metadata_variable_name]["spend_logs_metadata"] = team_metadata["spend_logs_metadata"] - - ## PROJECT-LEVEL TAGS project_metadata: Final = user_api_key_dict.project_metadata or {} - if "tags" in project_metadata and project_metadata["tags"] is not None: - data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=project_metadata["tags"], - ) # inherited_tags: every tag key/team/project policy contributed, read # directly from those three sources rather than snapshotted off the shared diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py index e7b608e162e..880fdad92bf 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py @@ -34,9 +34,7 @@ def is_collection_route(url_route: str, collection_suffix: str) -> bool: def request_tags_from_metadata(request_metadata: Mapping[str, object]) -> tuple[str, ...] | None: """Tags for the batch-cost spend row: the request's own tags when it sent any, - otherwise the key's tags, which auth exposes as user_api_key_auth_metadata (a - tagged key does not put its tags in the top-level metadata "tags" on the - passthrough path) + otherwise the key's tags, which auth exposes as user_api_key_auth_metadata """ tags: Final = _sanitized_str_tuple(request_metadata.get("tags")) if tags: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fa592629933..865374a0430 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -619,11 +619,11 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): litellm_metadata: Final = litellm_keys_in_body.get("litellm_metadata") metadata: Final = litellm_keys_in_body.get("metadata") - if litellm_metadata: - _metadata.update(litellm_metadata) - if metadata: - _metadata.update(metadata) + for client_metadata in (litellm_metadata, metadata): + if isinstance(client_metadata, dict): + _metadata.update({k: v for k, v in client_metadata.items() if not k.startswith("user_api_key_")}) + _metadata = _apply_key_team_project_controls(user_api_key_dict=user_api_key_dict, metadata=_metadata) _metadata = _update_metadata_with_tags_in_header( request=request, metadata=_metadata, @@ -1934,6 +1934,20 @@ async def pass_through_request( ) +def _apply_key_team_project_controls( + user_api_key_dict: UserAPIKeyAuth, metadata: dict[str, object] +) -> dict[str, object]: + data: Final = LiteLLMProxyRequestSetup.add_key_level_controls( + key_metadata=user_api_key_dict.metadata, + data={"metadata": metadata}, + _metadata_variable_name="metadata", + ) + return LiteLLMProxyRequestSetup.add_team_and_project_level_controls( + user_api_key_dict=user_api_key_dict, + metadata=data["metadata"], + ) + + def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> dict: """ If tags are in the request headers, add them to the metadata @@ -1954,9 +1968,10 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di # Only add tags key if there are tags to add if tags_to_add: - if "tags" not in metadata: - metadata["tags"] = [] - metadata["tags"].extend(tags_to_add) + metadata["tags"] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=metadata.get("tags"), + tags_to_add=tags_to_add, + ) return metadata diff --git a/tests/integration/spend/test_passthrough_request_tags.py b/tests/integration/spend/test_passthrough_request_tags.py new file mode 100644 index 00000000000..e6f5e15e161 --- /dev/null +++ b/tests/integration/spend/test_passthrough_request_tags.py @@ -0,0 +1,421 @@ +import json +import uuid +from collections.abc import Callable, Mapping +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import TypeAdapter + + +def _chat_reply(marker: str) -> dict[str, JsonValue]: + return { + "id": f"chatcmpl-{marker}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": marker}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + + +def _anthropic_reply(marker: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": marker}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + } + + +def _chat_stream_frames(marker: str) -> tuple[bytes, ...]: + chunk: Final = {"id": f"chatcmpl-{marker}", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"} + return ( + f"data: {json.dumps({**chunk, 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': marker}}]})}\n\n".encode(), + f"data: {json.dumps({**chunk, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}], 'usage': {'prompt_tokens': 5, 'completion_tokens': 3, 'total_tokens': 8}})}\n\n".encode(), + b"data: [DONE]\n\n", + ) + + +def _spend_row(digest: str, call_type: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_tags, metadata, team_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND call_type=%s', + (digest, call_type), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0] + + +def _spend_row_tagged(tag: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_tags, metadata, team_id, api_key FROM "LiteLLM_SpendLogs" WHERE request_tags::text LIKE %s', + (f'%"{tag}"%',), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0] + + +def _policy_tags(row: Mapping[str, JsonValue]) -> list[JsonValue]: + raw: Final = row["request_tags"] + tags: Final = json.loads(raw) if isinstance(raw, str) else raw + assert isinstance(tags, list), row + return [tag for tag in tags if not (isinstance(tag, str) and tag.startswith("User-Agent: "))] + + +def _spend_logs_metadata(row: Mapping[str, JsonValue]) -> JsonValue: + metadata: Final = row["metadata"] + return object_value(json.loads(metadata) if isinstance(metadata, str) else metadata).get("spend_logs_metadata") + + +def _tagged_key(scenario: Scenario, marker: str, **fields: JsonValue) -> tuple[str, str]: + team: Final = scenario.team(metadata={"tags": [f"team-{marker}"], "spend_logs_metadata": {"team_field": marker}}) + project: Final = scenario.project(team, metadata={"tags": [f"project-{marker}"]}) + key: Final = scenario.key( + team_id=team, + project_id=project, + metadata={"tags": [f"key-{marker}"], "spend_logs_metadata": {"cost_center": marker}}, + **fields, + ) + return key, sha256(key.encode()).hexdigest() + + +def _digest(key: str) -> str: + return sha256(key.encode()).hexdigest() + + +def _configured_passthrough(gateway: Gateway, scenario: Scenario, marker: str, target: str, *, auth: bool) -> str: + path: Final = f"/integration-passthrough-{marker}" + created: Final = gateway.post("/config/pass_through_endpoint", {"path": path, "target": target, "auth": auth}) + endpoints: Final = TypeAdapter(list[JsonValue]).validate_python(created["endpoints"]) + endpoint_id: Final = object_value(endpoints[0])["id"] + scenario.cleanups.callback( + lambda: gateway.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": str(endpoint_id)}) + ) + return path + + +def _responses_reply(marker: str, stream: bool) -> Reply: + response: Final[dict[str, JsonValue]] = { + "id": f"resp_{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": marker, "annotations": []}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8}, + } + if not stream: + return Reply(body=json.dumps(response).encode()) + events: Final[tuple[dict[str, JsonValue], ...]] = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{marker}", + "output_index": 0, + "content_index": 0, + "delta": marker, + }, + {"type": "response.completed", "sequence_number": 2, "response": response}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), + ) + + +def _echo_upstream(marker: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + assert request.method == "POST", request + body: Final = object_value(json.loads(request.body)) + assert marker in json.dumps(body), request + if request.target == "/v1/responses": + return _responses_reply(marker, body.get("stream") is True) + assert body["messages"] == [{"role": "user", "content": marker}], request + if body.get("stream") is True: + return Reply(chunks=_chat_stream_frames(marker), content_type="text/event-stream") + return Reply(body=json.dumps(_chat_reply(marker)).encode()) + + return respond + + +def test_configured_passthrough_spend_row_matches_native_route_tags_and_spend_logs_metadata(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=wire.url + "/v1") + path: Final = _configured_passthrough(gateway, scenario, marker, wire.url + "/echo", auth=True) + key, digest = _tagged_key(scenario, marker, models=[model], allowed_passthrough_routes=[path]) + headers: Final = {"x-litellm-tags": f"caller-{marker},key-{marker}", "User-Agent": "integration-tags/1"} + body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "user", "content": marker}]} + + native: Final = gateway.request("POST", "/v1/chat/completions", body, key=key, headers=headers) + assert native.status_code == 200, native.text + passthrough: Final = gateway.request("POST", path, body, key=key, headers=headers) + assert passthrough.status_code == 200, passthrough.text + assert json.loads(passthrough.content) == _chat_reply(marker) + + native_row: Final = _spend_row(digest, "acompletion") + passthrough_row: Final = _spend_row(digest, "pass_through_endpoint") + expected: Final = [f"key-{marker}", f"team-{marker}", f"project-{marker}", f"caller-{marker}"] + assert _policy_tags(native_row) == expected, native_row + assert _policy_tags(passthrough_row) == expected, passthrough_row + assert _spend_logs_metadata(native_row) == {"cost_center": marker, "team_field": marker}, native_row + assert _spend_logs_metadata(passthrough_row) == {"cost_center": marker, "team_field": marker}, passthrough_row + + +@pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"]) +def test_configured_passthrough_body_tags_lead_and_body_spend_logs_metadata_wins_over_key_and_team( + gateway: Gateway, bucket: str +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + path: Final = _configured_passthrough(gateway, scenario, marker, wire.url + "/echo", auth=True) + key, digest = _tagged_key(scenario, marker, allowed_passthrough_routes=[path]) + body: Final[dict[str, JsonValue]] = { + "messages": [{"role": "user", "content": marker}], + bucket: { + "tags": [f"body-{marker}", f"team-{marker}"], + "spend_logs_metadata": {"cost_center": f"body-{marker}"}, + }, + } + response: Final = gateway.request("POST", path, body, key=key) + assert response.status_code == 200, response.text + row: Final = _spend_row(digest, "pass_through_endpoint") + assert _policy_tags(row) == [f"body-{marker}", f"team-{marker}", f"key-{marker}", f"project-{marker}"], row + assert _spend_logs_metadata(row) == {"cost_center": f"body-{marker}", "team_field": marker}, row + + +def test_configured_passthrough_streaming_upstream_row_carries_key_team_project_and_caller_tags( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + path: Final = _configured_passthrough(gateway, scenario, marker, wire.url + "/echo", auth=True) + key, digest = _tagged_key(scenario, marker, allowed_passthrough_routes=[path]) + response: Final = gateway.request( + "POST", + path, + {"stream": True, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers={"x-litellm-tags": f"caller-{marker}"}, + ) + assert response.status_code == 200, response.text + assert response.content == b"".join(_chat_stream_frames(marker)), response.text + row: Final = _spend_row(digest, "pass_through_endpoint") + assert _policy_tags(row) == [f"key-{marker}", f"team-{marker}", f"project-{marker}", f"caller-{marker}"], row + assert _spend_logs_metadata(row) == {"cost_center": marker, "team_field": marker}, row + + +def test_configured_passthrough_key_outside_any_team_carries_its_own_tags_and_spend_logs_metadata( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + path: Final = _configured_passthrough(gateway, scenario, marker, wire.url + "/echo", auth=True) + key: Final = scenario.key( + allowed_passthrough_routes=[path], + metadata={"tags": [f"key-{marker}"], "spend_logs_metadata": {"cost_center": marker}}, + ) + response: Final = gateway.request( + "POST", + path, + {"messages": [{"role": "user", "content": marker}]}, + key=key, + headers={"x-litellm-tags": f"caller-{marker}"}, + ) + assert response.status_code == 200, response.text + row: Final = _spend_row(_digest(key), "pass_through_endpoint") + assert _policy_tags(row) == [f"key-{marker}", f"caller-{marker}"], row + assert _spend_logs_metadata(row) == {"cost_center": marker}, row + + +def test_configured_passthrough_untagged_key_row_keeps_only_caller_tag_and_no_spend_logs_metadata( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + path: Final = _configured_passthrough(gateway, scenario, marker, wire.url + "/echo", auth=True) + team: Final = scenario.team() + key: Final = scenario.key(team_id=team, allowed_passthrough_routes=[path]) + response: Final = gateway.request( + "POST", + path, + {"messages": [{"role": "user", "content": marker}]}, + key=key, + headers={"x-litellm-tags": f"caller-{marker}"}, + ) + assert response.status_code == 200, response.text + row: Final = _spend_row(_digest(key), "pass_through_endpoint") + assert _policy_tags(row) == [f"caller-{marker}"], row + assert _spend_logs_metadata(row) is None, row + assert row["team_id"] == team, row + + +def test_open_passthrough_without_auth_row_carries_only_caller_tag(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + path: Final = _configured_passthrough(gateway, scenario, marker, wire.url + "/echo", auth=False) + response: Final = gateway.client.post( + path, + json={"messages": [{"role": "user", "content": marker}]}, + headers={"x-litellm-tags": f"caller-{marker}"}, + ) + assert response.status_code == 200, response.text + row: Final = _spend_row_tagged(f"caller-{marker}") + assert _policy_tags(row) == [f"caller-{marker}"], row + assert _spend_logs_metadata(row) is None, row + assert row["api_key"] == "", row + + +@pytest.mark.parametrize( + ("metadata", "leading_tags"), + [ + ({"tags": "string-not-list"}, []), + ({"tags": [1, None, "z"]}, [1, None, "z"]), + ({"spend_logs_metadata": "string-not-object"}, []), + ], +) +def test_configured_passthrough_hostile_body_metadata_shapes_still_carry_key_team_project_tags( + gateway: Gateway, metadata: JsonValue, leading_tags: list[JsonValue] +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + path: Final = _configured_passthrough(gateway, scenario, marker, wire.url + "/echo", auth=True) + key, digest = _tagged_key(scenario, marker, allowed_passthrough_routes=[path]) + response: Final = gateway.request( + "POST", path, {"messages": [{"role": "user", "content": marker}], "metadata": metadata}, key=key + ) + assert response.status_code == 200, response.text + row: Final = _spend_row(digest, "pass_through_endpoint") + assert _policy_tags(row) == [*leading_tags, f"key-{marker}", f"team-{marker}", f"project-{marker}"], row + assert _spend_logs_metadata(row) == {"cost_center": marker, "team_field": marker}, row + + +def test_configured_passthrough_body_cannot_forge_user_api_key_attribution_fields(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + path: Final = _configured_passthrough(gateway, scenario, marker, wire.url + "/echo", auth=True) + forged_team: Final = scenario.team() + key, digest = _tagged_key(scenario, marker, allowed_passthrough_routes=[path]) + forged: Final[dict[str, JsonValue]] = { + "user_api_key": "forged-" + marker, + "user_api_key_team_id": forged_team, + "user_api_key_user_id": "forged-" + marker, + "user_api_key_alias": "forged-" + marker, + } + body: Final[dict[str, JsonValue]] = {"messages": [{"role": "user", "content": marker}], "metadata": forged} + response: Final = gateway.request("POST", path, body, key=key) + assert response.status_code == 200, response.text + row: Final = _spend_row(digest, "pass_through_endpoint") + assert row["team_id"] != forged_team, row + assert _policy_tags(row) == [f"key-{marker}", f"team-{marker}", f"project-{marker}"], row + assert read_rows('SELECT api_key FROM "LiteLLM_SpendLogs" WHERE team_id=%s', (forged_team,)) == [], forged_team + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("route", ["/v1/chat/completions", "/v1/messages"]) +def test_native_routes_carry_key_team_project_and_caller_tags_and_key_over_team_spend_logs_metadata( + gateway: Gateway, route: str, stream: bool +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_echo_upstream(marker)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=wire.url + "/v1") + key, digest = _tagged_key(scenario, marker, models=[model]) + body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 16, + "stream": stream, + "messages": [{"role": "user", "content": marker}], + } + response: Final = gateway.request( + "POST", + route, + body, + key=key, + headers={"x-litellm-tags": f"caller-{marker}", "User-Agent": "integration-tags/1"}, + ) + assert response.status_code == 200, response.text + rows: Final = eventually( + lambda: read_rows('SELECT request_tags, metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert _policy_tags(rows[0]) == [f"key-{marker}", f"team-{marker}", f"project-{marker}", f"caller-{marker}"], ( + rows + ) + assert _spend_logs_metadata(rows[0]) == {"cost_center": marker, "team_field": marker}, rows + + +def test_anthropic_passthrough_spend_row_carries_key_team_project_tags_and_spend_logs_metadata( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request + assert request.headers["x-api-key"] == "synthetic-anthropic-key" + return Reply(body=json.dumps(_anthropic_reply(marker)).encode()) + + config: Final = tmp_path / "proxy_config.yaml" + config.write_text( + "model_list: []\n" + "general_settings:\n" + " master_key: os.environ/LITELLM_MASTER_KEY\n" + " database_url: os.environ/DATABASE_URL\n" + " store_model_in_db: true\n" + " disable_spend_logs: false\n" + " proxy_batch_write_at: 1\n" + "router_settings:\n" + " disable_cooldowns: true\n" + ) + with wire_server(respond) as wire: + overrides: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": "synthetic-anthropic-key"} + with owned_proxy(gateway, tmp_path, overrides, config=config) as candidate, candidate.scenario() as scenario: + key, digest = _tagged_key(scenario, marker) + response: Final = candidate.request( + "POST", + "/anthropic/v1/messages", + { + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + }, + key=key, + headers={"x-litellm-tags": f"caller-{marker}", "User-Agent": "integration-tags/1"}, + ) + assert response.status_code == 200, response.text + assert json.loads(response.content) == _anthropic_reply(marker) + row: Final = _spend_row(digest, "pass_through_endpoint") + assert _policy_tags(row) == [f"key-{marker}", f"team-{marker}", f"project-{marker}", f"caller-{marker}"], ( + row + ) + assert _spend_logs_metadata(row) == {"cost_center": marker, "team_field": marker}, row diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 72d76e54a5d..c6c81c14b16 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7592,6 +7592,61 @@ def test_passthrough_sees_the_public_list_rebound_after_import(monkeypatch: pyte assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]} +def test_passthrough_metadata_carries_key_team_project_tags_and_key_spend_logs_metadata(): + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = httpx.URL("http://0.0.0.0:4000/anthropic/v1/messages") + mock_request.headers = Headers({"x-litellm-tags": "caller-tag,key-tag"}) + mock_request.scope = {} + + cached_key = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"tags": ["key-tag", "shared-tag"], "spend_logs_metadata": {"cost_center": "key"}}, + team_metadata={ + "tags": ["team-tag", "shared-tag"], + "spend_logs_metadata": {"cost_center": "team", "team_field": "team"}, + }, + project_metadata={"tags": ["project-tag"]}, + ) + + kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request, + user_api_key_dict=cached_key, + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), + _parsed_body={ + "metadata": { + "tags": ["body-tag"], + "spend_logs_metadata": {"request_id": "body"}, + "user_api_key_auth_metadata": "forged", + } + }, + litellm_call_id="lit-5359-call-id", + ) + second = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request, + user_api_key_dict=cached_key, + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), + _parsed_body={}, + litellm_call_id="lit-5359-second-call-id", + ) + + metadata = kwargs["litellm_params"]["metadata"] + assert metadata["tags"] == ["body-tag", "key-tag", "shared-tag", "team-tag", "project-tag", "caller-tag"] + assert metadata["spend_logs_metadata"] == {"request_id": "body", "cost_center": "key", "team_field": "team"} + assert metadata["user_api_key_auth_metadata"] == { + "tags": ["key-tag", "shared-tag"], + "spend_logs_metadata": {"cost_center": "key"}, + } + assert second["litellm_params"]["metadata"]["spend_logs_metadata"] == {"cost_center": "key", "team_field": "team"} + assert cached_key.metadata == {"tags": ["key-tag", "shared-tag"], "spend_logs_metadata": {"cost_center": "key"}} + assert cached_key.team_metadata == { + "tags": ["team-tag", "shared-tag"], + "spend_logs_metadata": {"cost_center": "team", "team_field": "team"}, + } + + @pytest.mark.asyncio async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_error_for_an_unknown_model( monkeypatch: pytest.MonkeyPatch,