diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5a3a17f338c..5c77b076bed 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6301,7 +6301,7 @@ class StandardLoggingPayloadSetup: return None user_agent_tags: list[str] | None = None headers: Final = proxy_server_request.get("headers", {}) - if headers is not None and isinstance(headers, dict): + if headers is not None and isinstance(headers, Mapping): if "user-agent" in headers: user_agent: Final = headers["user-agent"] if user_agent is not None: @@ -6326,7 +6326,7 @@ class StandardLoggingPayloadSetup: return None headers: Final = proxy_server_request.get("headers", {}) - if not isinstance(headers, dict): + if not isinstance(headers, Mapping): return None header_tags: Final = [] diff --git a/tests/integration/spend/_request_tag_helpers.py b/tests/integration/spend/_request_tag_helpers.py new file mode 100644 index 00000000000..4cf6b16d475 --- /dev/null +++ b/tests/integration/spend/_request_tag_helpers.py @@ -0,0 +1,222 @@ +import json +import uuid +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import yaml +from integration._support.database import read_rows +from integration._support.wire import Reply, Request + +ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929" +OPENAI_MODEL: Final = "gpt-4o-mini" +GEMINI_MODEL: Final = "gemini-2.5-flash" +HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"} +T3: Final = ["User-Agent: claude-cli", "User-Agent: claude-cli/2.0.0", "x-tenant-id: tenant-a"] +ANTHROPIC_SONNET_BODY: Final = { + "id": "msg_synthetic", + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": "tagged"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 2}, +} +CHAT_COMPLETION_BODY: Final = { + "id": "chatcmpl-synthetic", + "object": "chat.completion", + "created": 1, + "model": OPENAI_MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "tagged"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, +} +RESPONSES_BODY: Final = { + "id": "resp_synthetic", + "object": "response", + "created_at": 1, + "status": "completed", + "model": OPENAI_MODEL, + "output": [ + { + "type": "message", + "id": "msg_synthetic", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "tagged", "annotations": []}], + } + ], + "usage": { + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, +} +GEMINI_BODY: Final = { + "candidates": [ + { + "content": {"parts": [{"text": "tagged"}], "role": "model"}, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 2, "totalTokenCount": 12}, +} +ANTHROPIC_STREAM_EVENTS: Final = ( + { + "type": "message_start", + "message": {**ANTHROPIC_SONNET_BODY, "content": [], "usage": {"input_tokens": 10, "output_tokens": 1}}, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "tagged"}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, +) +OPENAI_STREAM_CHUNKS: Final = ( + { + "id": CHAT_COMPLETION_BODY["id"], + "object": "chat.completion.chunk", + "created": 1, + "model": OPENAI_MODEL, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "tagged"}, "finish_reason": None}], + }, + { + "id": CHAT_COMPLETION_BODY["id"], + "object": "chat.completion.chunk", + "created": 1, + "model": OPENAI_MODEL, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + { + "id": CHAT_COMPLETION_BODY["id"], + "object": "chat.completion.chunk", + "created": 1, + "model": OPENAI_MODEL, + "choices": [], + "usage": CHAT_COMPLETION_BODY["usage"], + }, +) + + +def _message_id() -> str: + return "msg_" + uuid.uuid4().hex + + +def _sse_frames(events: tuple[dict, ...]) -> tuple[bytes, ...]: + return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) + + +def provider_reply(request: Request) -> Reply: + """Scripted edge for every route the audit drives: anthropic messages, openai chat completions and + responses, gemini generateContent. Error bodies keyed off the sentinel model name.""" + body: Final = json.loads(request.body) if request.body else {} + target: Final = request.target.split("?", 1)[0] + if body.get("model") == "claude-nonexistent-model": + return Reply( + status=400, + body=json.dumps( + {"error": {"type": "invalid_request_error", "message": "model: claude-nonexistent-model"}} + ).encode(), + ) + if target == "/v1/models" or target.startswith("/v1/models/"): + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + if target == "/v1/messages": + identity: Final = _message_id() + if body.get("stream") is True: + events: Final = tuple( + {**event, "message": {**event.get("message", {}), "id": identity}} if "message" in event else event + for event in ANTHROPIC_STREAM_EVENTS + ) + return Reply(content_type="text/event-stream", chunks=_sse_frames(events)) + return Reply(body=json.dumps({**ANTHROPIC_SONNET_BODY, "id": identity}).encode()) + if target == "/v1/chat/completions": + identity = "chatcmpl_" + uuid.uuid4().hex + if body.get("stream") is True: + frames: Final = tuple( + f"data: {json.dumps({**chunk, 'id': identity})}\n\n".encode() for chunk in OPENAI_STREAM_CHUNKS + ) + (b"data: [DONE]\n\n",) + return Reply(content_type="text/event-stream", chunks=frames) + return Reply(body=json.dumps({**CHAT_COMPLETION_BODY, "id": identity}).encode()) + if target == "/v1/responses": + identity = "resp_" + uuid.uuid4().hex + message: Final = _message_id() + completed_body: Final = { + **RESPONSES_BODY, + "id": identity, + "output": [{**RESPONSES_BODY["output"][0], "id": message}], + } + if body.get("stream") is True: + created: Final = { + "type": "response.created", + "response": {**completed_body, "status": "in_progress", "output": [], "usage": None}, + } + delta: Final = { + "type": "response.output_text.delta", + "item_id": message, + "output_index": 0, + "content_index": 0, + "delta": "tagged", + } + completed: Final = {"type": "response.completed", "response": completed_body} + return Reply(content_type="text/event-stream", chunks=_sse_frames((created, delta, completed))) + return Reply(body=json.dumps(completed_body).encode()) + if target.endswith(":generateContent") or target.endswith(":streamGenerateContent"): + return Reply(body=json.dumps(GEMINI_BODY).encode()) + raise AssertionError(f"unexpected upstream target {request.target}") + + +def provider_env(url: str) -> dict[str, str]: + return { + "ANTHROPIC_API_BASE": url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + "OPENAI_API_BASE": url, + "OPENAI_API_KEY": "synthetic-openai-key", + "GEMINI_API_BASE": url, + "GEMINI_API_KEY": "synthetic-gemini-key", + } + + +def write_config(directory: Path, mutations: dict, name: str = "spend-tag-headers.yaml") -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + for section, values in mutations.items(): + if isinstance(values, dict) and isinstance(config.get(section), dict): + config[section].update(values) + else: + config[section] = values + path: Final = directory / name + path.write_text(yaml.safe_dump(config)) + return path + + +def tags_of(row: dict) -> list: + value: Final = row["request_tags"] + return json.loads(value) if isinstance(value, str) else value + + +def tags_by_key(key: str) -> list[list]: + rows: Final = read_rows( + 'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),) + ) + return [tags_of(row) for row in rows] + + +def tags_by_id(request_id: str) -> list[list]: + rows: Final = read_rows('SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)) + return [tags_of(row) for row in rows] + + +def row_by_id(request_id: str) -> list[dict]: + return read_rows( + 'SELECT request_id, call_type, request_tags, status, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (request_id,), + ) + + +def unique_marker() -> str: + return "tagprobe" + uuid.uuid4().hex[:12] diff --git a/tests/integration/spend/test_spend_log_request_tags.py b/tests/integration/spend/test_spend_log_request_tags.py new file mode 100644 index 00000000000..d0311d0e0e0 --- /dev/null +++ b/tests/integration/spend/test_spend_log_request_tags.py @@ -0,0 +1,1002 @@ +import asyncio +import http.client +import json +import uuid +from collections.abc import Mapping, Sequence +from hashlib import sha256 +from itertools import chain +from pathlib import Path +from typing import Final + +import anthropic +import httpx +import openai +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from integration.spend._request_tag_helpers import ( + GEMINI_MODEL, + OPENAI_MODEL, + provider_env, + provider_reply, + tags_by_id, + tags_by_key, + tags_of, + write_config, +) + +MODEL: Final = "claude-sonnet-4-5-20250929" +SENT_HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"} +EXPECTED_TAGS: Final = ["User-Agent: claude-cli", "User-Agent: claude-cli/2.0.0", "x-tenant-id: tenant-a"] + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_header_derived_spend_tags_are_recorded_on_anthropic_messages_routes( + gateway: Gateway, tmp_path: Path, route: str +) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS] + + +def _base_url(candidate: Gateway) -> str: + return str(candidate.client.base_url).rstrip("/") + + +UA_TAG: Final = "User-Agent: claude-cli/2.0.0" +UA_FAMILY_TAG: Final = "User-Agent: claude-cli" +TENANT_TAG: Final = "x-tenant-id: tenant-a" + + +def test_pass_through_anthropic_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + "/anthropic/v1/messages", + { + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + "stream": True, + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert '"type":"message_start"' in response.text.replace(" ", ""), response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS] + + +def test_pass_through_anthropic_sdk_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + client: Final = anthropic.Anthropic( + base_url=f"{_base_url(candidate)}/anthropic", auth_token=key, default_headers=SENT_HEADERS + ) + message: Final = client.messages.create( + model=MODEL, max_tokens=16, messages=[{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}] + ) + assert message.id.startswith("msg_") + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [ + ["User-Agent: Anthropic", f"User-Agent: Anthropic/Python {anthropic.__version__}", TENANT_TAG] + ] + + +def test_pass_through_anthropic_sdk_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + client: Final = anthropic.Anthropic( + base_url=f"{_base_url(candidate)}/anthropic", auth_token=key, default_headers=SENT_HEADERS + ) + with client.messages.stream( + model=MODEL, max_tokens=16, messages=[{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}] + ) as stream: + message: Final = stream.get_final_message() + assert message.id.startswith("msg_") + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [ + ["User-Agent: Anthropic", f"User-Agent: Anthropic/Python {anthropic.__version__}", TENANT_TAG] + ] + + +@pytest.mark.parametrize("stream", [pytest.param(False, id="sync"), pytest.param(True, id="stream")]) +def test_pass_through_openai_chat_records_header_tags(gateway: Gateway, tmp_path: Path, stream: bool) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + "/openai/v1/chat/completions", + { + "model": OPENAI_MODEL, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + **({"stream": True} if stream else {}), + }, + key=key, + headers=SENT_HEADERS, + ) + assert response.status_code == 200, response.text + assert "chatcmpl_" in response.text, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS] + + +@pytest.mark.parametrize("stream", [pytest.param(False, id="sync"), pytest.param(True, id="stream")]) +def test_pass_through_openai_sdk_records_header_tags(gateway: Gateway, tmp_path: Path, stream: bool) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + client: Final = openai.OpenAI( + api_key=key, base_url=f"{_base_url(candidate)}/openai/v1", default_headers=SENT_HEADERS + ) + completion: Final = client.chat.completions.create( + model=OPENAI_MODEL, messages=[{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], stream=stream + ) + if stream: + request_id: Final = next(chunk.id for chunk in completion) + for _ in completion: + pass + else: + request_id = completion.id + assert request_id.startswith("chatcmpl_") + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [ + ["User-Agent: OpenAI", f"User-Agent: OpenAI/Python {openai.__version__}", TENANT_TAG] + ] + + +def test_pass_through_gemini_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + f"/gemini/v1beta/models/{GEMINI_MODEL}:generateContent", + {"contents": [{"parts": [{"text": f"tag me {uuid.uuid4().hex}"}]}]}, + key=key, + headers={**SENT_HEADERS, "x-goog-api-key": key}, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS] + + +def test_custom_pass_through_endpoint_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config( + tmp_path, + { + "litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}, + "general_settings": { + "pass_through_endpoints": [ + { + "path": "/custom-anthropic", + "target": f"{wire.url}/v1/messages", + "auth": True, + "headers": { + "x-api-key": "synthetic-anthropic-key", + "anthropic-version": "2023-06-01", + }, + } + ] + }, + }, + ) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key(allowed_passthrough_routes=["/custom-anthropic"]) + response: Final = candidate.request( + "POST", + "/custom-anthropic", + { + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers=SENT_HEADERS, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS] + + +def test_pass_through_openai_async_sdk_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + client: Final = openai.AsyncOpenAI( + api_key=key, base_url=f"{_base_url(candidate)}/openai/v1", default_headers=SENT_HEADERS + ) + + async def call() -> str: + completion: Final = await client.chat.completions.create( + model=OPENAI_MODEL, + messages=[{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + stream=True, + ) + first: Final = await completion.__anext__() + async for _ in completion: + pass + return first.id + + request_id: Final = asyncio.run(call()) + assert request_id.startswith("chatcmpl_") + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [ + ["User-Agent: AsyncOpenAI", f"User-Agent: AsyncOpenAI/Python {openai.__version__}", TENANT_TAG] + ] + + +def test_unified_chat_completions_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}]}, + key=key, + headers=SENT_HEADERS, + ) + assert response.status_code == 200, response.text + request_id: Final = response.json()["id"] + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +@pytest.mark.parametrize("stream", [pytest.param(False, id="sync"), pytest.param(True, id="stream")]) +def test_unified_responses_records_header_tags(gateway: Gateway, tmp_path: Path, stream: bool) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + key: Final = scenario.key() + client: Final = openai.OpenAI( + api_key=key, base_url=f"{_base_url(candidate)}/v1", default_headers=SENT_HEADERS + ) + created: Final = client.responses.create(model=model, input=f"tag me {uuid.uuid4().hex}", stream=stream) + if stream: + frames: Final = list(created) + request_id: Final = next(event.response.id for event in frames if event.type == "response.completed") + else: + request_id = created.id + assert request_id.startswith("resp_") + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS] + + +def test_unified_cache_hit_twin_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + key: Final = scenario.key() + body: Final = {"model": model, "messages": [{"role": "user", "content": f"cache {uuid.uuid4().hex}"}]} + first: Final = candidate.request("POST", "/v1/chat/completions", body, key=key, headers=SENT_HEADERS) + assert first.status_code == 200, first.text + second: Final = candidate.request("POST", "/v1/chat/completions", body, key=key, headers=SENT_HEADERS) + assert second.status_code == 200, second.text + assert len(wire.drain()) == 1, "identical second call should have hit the response cache" + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 2, seconds=70) == [ + EXPECTED_TAGS, + EXPECTED_TAGS, + ] + + +def test_pass_through_tags_reach_generic_api_sink(gateway: Gateway, tmp_path: Path) -> None: + def sink(request: Request) -> Reply: + return Reply() + + with wire_server(provider_reply) as wire, wire_server(sink) as endpoint: + config: Final = write_config( + tmp_path, + { + "litellm_settings": { + "extra_spend_tag_headers": ["x-tenant-id"], + "callbacks": ["generic_api"], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + } + }, + ) + with ( + owned_proxy( + gateway, + tmp_path, + { + **provider_env(wire.url), + "GENERIC_LOGGER_ENDPOINT": endpoint.url, + "GENERIC_LOGGER_HEADERS": "Authorization=Bearer synthetic-sink-secret", + }, + config=config, + workers=2, + ) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + "/anthropic/v1/messages", + { + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + request_id: Final = response.json()["id"] + assert len(wire.drain()) == 1 + batches: Final[list[Request]] = [] # mutable-ok: drain consumes the queue between polls + + def delivered() -> Sequence[Mapping]: + batches.extend(endpoint.drain()) + events: Final = chain.from_iterable(json.loads(batch.body) for batch in batches) + return [event for event in events if event.get("id") == request_id] + + events: Final = eventually(delivered, lambda values: len(values) == 1, seconds=30) + assert events[0]["request_tags"] == EXPECTED_TAGS + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +def test_pass_through_tags_accrue_daily_tag_spend(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + "/anthropic/v1/messages", + { + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + request_id: Final = response.json()["id"] + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + rows: Final = eventually( + lambda: read_rows( + 'SELECT tag FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda values: len(values) == 3, + seconds=70, + ) + assert {row["tag"] for row in rows} == set(EXPECTED_TAGS) + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_header_tags_without_extra_spend_tag_headers_record_user_agent_only( + gateway: Gateway, tmp_path: Path, route: str +) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + expected: Final = [UA_FAMILY_TAG, UA_TAG] + assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + expected + ] + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_routes_without_headers_record_no_tags(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + connection: Final = http.client.HTTPConnection("127.0.0.1", candidate.client.base_url.port) + connection.request( + "POST", + route, + body=json.dumps( + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + } + ), + headers={ + "authorization": f"Bearer {key}", + "content-type": "application/json", + "anthropic-version": "2023-06-01", + }, + ) + raw: Final = connection.getresponse() + payload: Final = raw.read() + connection.close() + assert raw.status == 200, payload + request_id: Final = json.loads(payload)["id"] + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [[]] + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_disabled_user_agent_keeps_only_extra_header_tags(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config( + tmp_path, + { + "litellm_settings": { + "extra_spend_tag_headers": ["x-tenant-id"], + "disable_add_user_agent_to_request_tags": True, + } + }, + ) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + [TENANT_TAG] + ] + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_unsent_configured_header_contributes_no_tag(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-never-sent"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + [UA_FAMILY_TAG, UA_TAG] + ] + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_default_httpx_user_agent_is_recorded(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={"anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + expected: Final = ["User-Agent: python-httpx", f"User-Agent: python-httpx/{httpx.__version__}"] + assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + expected + ] + + +@pytest.mark.parametrize("route", [pytest.param("/anthropic/v1/messages", id="passthrough")]) +def test_unauthenticated_pass_through_writes_untagged_spend_row(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + anonymous_before: Final = len( + read_rows("SELECT request_tags FROM \"LiteLLM_SpendLogs\" WHERE api_key IS NULL OR api_key=''", ()) + ) + response: Final = candidate.client.post( + route, + json={ + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 401, response.text + anonymous: Final = eventually( + lambda: read_rows( + 'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key IS NULL OR api_key=\'\' ORDER BY "startTime" DESC', + (), + ), + lambda rows: len(rows) == anonymous_before + 1, + seconds=70, + ) + assert tags_of(anonymous[0]) == [] + key: Final = scenario.key() + control: Final = candidate.request( + "POST", + route, + { + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert control.status_code == 200, control.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(control.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_upstream_failure_still_records_header_tags(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model="anthropic/claude-nonexistent-model", + api_base=wire.url, + api_key="synthetic-anthropic-key", + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": "claude-nonexistent-model" if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 400, response.text + assert len(wire.drain()) == 1 + expected: Final = [] if route == "/anthropic/v1/messages" else EXPECTED_TAGS + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [expected] + + +@pytest.mark.parametrize("extra", [pytest.param(None, id="null"), pytest.param([], id="empty")]) +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_null_and_empty_extra_spend_tag_headers_record_user_agent_only( + gateway: Gateway, tmp_path: Path, route: str, extra: object +) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": extra}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + [UA_FAMILY_TAG, UA_TAG] + ] + + +def test_configured_header_case_differs_between_routes(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["X-Tenant-Id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + body: Final = { + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + } + passthrough: Final = candidate.request( + "POST", + "/anthropic/v1/messages", + body, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert passthrough.status_code == 200, passthrough.text + unified: Final = candidate.request( + "POST", + "/v1/messages", + {**body, "model": model}, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert unified.status_code == 200, unified.text + assert len(wire.drain()) == 2 + rows: Final = eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 2, seconds=70) + passthrough_tags: Final = tags_by_id(passthrough.json()["id"])[0] + unified_tags: Final = tags_by_id(unified.json()["id"])[0] + assert passthrough_tags == [UA_FAMILY_TAG, UA_TAG, "X-Tenant-Id: tenant-a"], rows + assert unified_tags == [UA_FAMILY_TAG, UA_TAG], rows + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_large_header_value_is_stored_verbatim(gateway: Gateway, tmp_path: Path, route: str) -> None: + big: Final = "x" * 5000 + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={ + "user-agent": "claude-cli/2.0.0", + "x-tenant-id": big, + "anthropic-version": "2023-06-01", + }, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + [UA_FAMILY_TAG, UA_TAG, f"x-tenant-id: {big}"] + ] + + +def test_duplicate_header_values_follow_carrier_semantics(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + body: Final = { + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + } + duplicated: Final = [ + ("Authorization", f"Bearer {key}"), + ("user-agent", "claude-cli/2.0.0"), + ("anthropic-version", "2023-06-01"), + ("x-tenant-id", "t1"), + ("x-tenant-id", "t2"), + ] + passthrough: Final = candidate.client.post("/anthropic/v1/messages", json=body, headers=duplicated) + assert passthrough.status_code == 200, passthrough.text + unified: Final = candidate.client.post("/v1/messages", json={**body, "model": model}, headers=duplicated) + assert unified.status_code == 200, unified.text + assert len(wire.drain()) == 2 + eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 2, seconds=70) + assert tags_by_id(passthrough.json()["id"])[0] == [UA_FAMILY_TAG, UA_TAG, "x-tenant-id: t1"] + assert tags_by_id(unified.json()["id"])[0] == [UA_FAMILY_TAG, UA_TAG, "x-tenant-id: t2"] + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_x_litellm_tags_merges_with_header_tags(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + response: Final = candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"tag me {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01", "x-litellm-tags": "team-x"}, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + expected: Final = ["team-x", UA_FAMILY_TAG, UA_TAG, TENANT_TAG] + assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + expected + ] + + +@pytest.mark.parametrize( + "route", [pytest.param("/anthropic/v1/messages", id="passthrough"), pytest.param("/v1/messages", id="unified")] +) +def test_repeated_requests_each_record_tags(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"anthropic/{MODEL}", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + key: Final = scenario.key() + + def repeat(_: int) -> httpx.Response: + return candidate.request( + "POST", + route, + { + "model": MODEL if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"repeat {uuid.uuid4().hex}"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + + responses: Final = tuple(repeat(index) for index in range(3)) + assert all(response.status_code == 200 for response in responses), [ + (response.status_code, response.text) for response in responses + ] + ids: Final = tuple(response.json()["id"] for response in responses) + assert len(wire.drain()) == 3 + for identity in ids: + assert eventually( + lambda identity=identity: tags_by_id(identity), lambda tags: len(tags) == 1, seconds=70 + ) == [EXPECTED_TAGS] + + +def test_guardrail_mode_tag_decider_is_unchanged_on_pass_through(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config( + tmp_path, + { + "litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}, + "general_settings": { + "pass_through_endpoints": [ + { + "path": "/custom-anthropic", + "target": f"{wire.url}/v1/messages", + "auth": True, + "guardrails": ["tag-blocker"], + "headers": { + "x-api-key": "synthetic-anthropic-key", + "anthropic-version": "2023-06-01", + }, + } + ] + }, + "guardrails": [ + { + "guardrail_name": "tag-blocker", + "litellm_params": { + "guardrail": "litellm_content_filter", + "blocked_words": [{"keyword": "bananablock", "action": "BLOCK"}], + "mode": {"tags": {UA_FAMILY_TAG: "pre_call"}, "default": "post_call"}, + "default_on": True, + }, + } + ], + }, + ) + with ( + owned_proxy(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + key: Final = scenario.key(allowed_passthrough_routes=["/custom-anthropic"]) + body: Final = {"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "bananablock"}]} + passthrough: Final = candidate.request("POST", "/custom-anthropic", body, key=key, headers=SENT_HEADERS) + assert passthrough.status_code == 200, passthrough.text + assert len(wire.drain()) == 1 + control: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "bananablock"}]}, + key=key, + headers=SENT_HEADERS, + ) + assert control.status_code != 200, control.text + assert len(wire.drain()) == 0, "tag-matched guardrail should have blocked before the upstream" + digest: Final = sha256(key.encode()).hexdigest() + + def blocked_rows() -> Sequence[Mapping]: + return [ + row + for row in read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (digest,), + ) + if "Content blocked: keyword 'bananablock' detected" in json.dumps(row["metadata"]) + and row["metadata"].get("status") == "failure" + ] + + assert eventually(blocked_rows, lambda rows: len(rows) == 1, seconds=70), ( + "guardrail block was not recorded on the key's spend row" + ) diff --git a/tests/integration/spend/test_spend_log_request_tags_chaos.py b/tests/integration/spend/test_spend_log_request_tags_chaos.py new file mode 100644 index 00000000000..068e2917e1a --- /dev/null +++ b/tests/integration/spend/test_spend_log_request_tags_chaos.py @@ -0,0 +1,314 @@ +import json +import threading +import uuid +from collections.abc import Callable, Mapping, Sequence, Set +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from itertools import chain +from pathlib import Path +from typing import Final + +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from integration.spend._request_tag_helpers import ( + ANTHROPIC_MODEL, + OPENAI_MODEL, + T3, + provider_env, + provider_reply, + write_config, +) + +HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"} +ANTHROPIC_HEADERS: Final = {**HEADERS, "anthropic-version": "2023-06-01"} +EXPECTED: Final = T3 +ROUTES: Final = ( + "/anthropic/v1/messages", + "/openai/v1/chat/completions", + "/v1/chat/completions", + "/v1/messages", + "/v1/responses", +) + + +def _ids(response) -> str: + for prefix in ("msg_", "chatcmpl_", "resp_"): + if prefix in response.text: + return prefix + response.text.split(prefix)[1].split('"')[0] + raise AssertionError(f"no upstream id in {response.text[:200]}") + + +def _requests( + candidate: Gateway, key: str, anthropic_model: str, openai_model: str, stream: bool, marker: str +) -> tuple: + """Deferred calls for one request per route in ROUTES order with the same client headers.""" + return ( + lambda: candidate.request( + "POST", + "/anthropic/v1/messages", + { + "model": ANTHROPIC_MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + "stream": stream, + }, + key=key, + headers=ANTHROPIC_HEADERS, + ), + lambda: candidate.request( + "POST", + "/openai/v1/chat/completions", + {"model": OPENAI_MODEL, "messages": [{"role": "user", "content": marker}], "stream": stream}, + key=key, + headers=HEADERS, + ), + lambda: candidate.request( + "POST", + "/v1/chat/completions", + {"model": openai_model, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers=HEADERS, + ), + lambda: candidate.request( + "POST", + "/v1/messages", + { + "model": anthropic_model, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + }, + key=key, + headers=ANTHROPIC_HEADERS, + ), + lambda: candidate.request( + "POST", "/v1/responses", {"model": openai_model, "input": marker}, key=key, headers=HEADERS + ), + ) + + +def _tagged_requests( + candidate: Gateway, key: str, anthropic_model: str, openai_model: str, stream: bool, index: int +) -> tuple: + """One call per route in ROUTES order with the same client headers.""" + marker: Final = f"burst {index} {uuid.uuid4().hex}" + return tuple(send() for send in _requests(candidate, key, anthropic_model, openai_model, stream, marker)) + + +def _deployments(scenario, url: str) -> tuple[str, str]: + return ( + scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=url, api_key="synthetic-anthropic-key"), + scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{url}/v1"), + ) + + +def _landed_tags(key: str, satisfied) -> Sequence[Mapping]: + digest: Final = sha256(key.encode()).hexdigest() + landed: Final = eventually( + lambda: read_rows('SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)), + satisfied, + seconds=60, + ) + assert len({row["request_id"] for row in landed}) == len(landed) + for row in landed: + value: Final = row["request_tags"] + assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED + return landed + + +def _worker_pids(owned) -> tuple[int, ...]: + workers: Final = tuple( + child + for child in psutil.Process(owned.process.pid).children(recursive=True) + if any(marker in " ".join(child.cmdline()) for marker in ("spawn_main", "integration._support.proxy")) + ) + return tuple(worker.pid for worker in workers) + + +def test_burst_across_routes_records_tags_once_per_response(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy_process(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + anthropic_model, openai_model = _deployments(scenario, wire.url) + key: Final = scenario.key() + + def burst(index: int) -> tuple: + return _tagged_requests( + candidate, key, anthropic_model, openai_model, stream=index % 2 == 1, index=index + ) + + with ThreadPoolExecutor(max_workers=10) as pool: + responses: Final = tuple(chain.from_iterable(pool.map(burst, range(10)))) + assert len(responses) == 50 + assert all(response.status_code == 200 for response in responses), [ + (response.status_code, response.text[:200]) for response in responses + ] + ids: Final = [_ids(response) for response in responses] + assert len(set(ids)) == 50, "duplicate upstream id in burst" + assert len(wire.drain()) == 50 + pids: Final = _worker_pids(owned) + assert len(set(pids)) == 2, f"expected two uvicorn workers, found {pids}" + assert all(psutil.Process(pid).is_running() for pid in pids) + _landed_tags(key, lambda values: len(values) == 50) + + +@pytest.mark.timeout(240) +def test_sink_outage_does_not_lose_spend_log_tags(gateway: Gateway, tmp_path: Path) -> None: + down: Final = threading.Event() + delivered: Final = [] + rejected: Final = [] + + def stoppable_sink(request: Request) -> Reply: + if down.is_set(): + rejected.append(request) + return Reply(status=503) + delivered.append(request) + return Reply() + + with wire_server(provider_reply) as wire, wire_server(stoppable_sink) as endpoint: + config: Final = write_config( + tmp_path, + { + "litellm_settings": { + "extra_spend_tag_headers": ["x-tenant-id"], + "callbacks": ["generic_api"], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + } + }, + ) + with ( + owned_proxy_process( + gateway, + tmp_path, + { + **provider_env(wire.url), + "GENERIC_LOGGER_ENDPOINT": endpoint.url, + "GENERIC_LOGGER_HEADERS": "Authorization=Bearer synthetic-sink-secret", + }, + config=config, + workers=2, + ) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + anthropic_model, openai_model = _deployments(scenario, wire.url) + key: Final = scenario.key() + + def events_for(batches: Sequence[Request], ids: Set[str]) -> Sequence[Mapping]: + events: Final = chain.from_iterable(json.loads(batch.body) for batch in batches) + return [event for event in events if event.get("litellm_call_id") in ids] + + def send_and_await(send: Callable[[], httpx.Response], batches: Sequence[Request]) -> httpx.Response: + response: Final = send() + assert response.status_code == 200, response.text + call_id: Final = response.headers["x-litellm-call-id"] + eventually(lambda: events_for(batches, {call_id}), lambda found: len(found) >= 1, seconds=30) + return response + + def burst_sends(label: str): + return chain.from_iterable( + _requests( + candidate, + key, + anthropic_model, + openai_model, + False, + f"{label} {index} {uuid.uuid4().hex}", + ) + for index in range(3) + ) + + first: Final = tuple(send_and_await(send, delivered) for send in burst_sends("burst")) + first_ids: Final = {response.headers["x-litellm-call-id"] for response in first} + first_events: Final = events_for(delivered, first_ids) + first_occurrences: Final = [event["litellm_call_id"] for event in first_events] + assert sorted(first_occurrences) == sorted(first_ids), "burst-1 sink delivery is not exactly once per call" + for event in first_events: + assert event["request_tags"] == EXPECTED + down.set() + second: Final = tuple(send_and_await(send, rejected) for send in burst_sends("outage")) + second_ids: Final = {response.headers["x-litellm-call-id"] for response in second} + rejected_ids: Final = {event["litellm_call_id"] for event in events_for(rejected, second_ids)} + assert rejected_ids == second_ids, "outage burst was not rejected by the down sink" + assert events_for(delivered, second_ids) == [], "burst-2 event delivered to a down sink" + down.clear() + probe: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": openai_model, + "messages": [{"role": "user", "content": f"recovery probe {uuid.uuid4().hex}"}], + }, + key=key, + headers=HEADERS, + ) + assert probe.status_code == 200, probe.text + probe_id: Final = probe.headers["x-litellm-call-id"] + probe_events: Final = eventually( + lambda: events_for(delivered, {probe_id}), + lambda found: len(found) >= 1, + seconds=70, + ) + assert len(probe_events) == 1, "recovery probe delivered to the sink more than once" + assert probe_events[0]["request_tags"] == EXPECTED + second_events: Final = events_for(delivered, second_ids) + second_occurrences: Final = [event["litellm_call_id"] for event in second_events] + assert len(second_occurrences) == len(set(second_occurrences)), ( + "duplicate burst-2 delivery after the outage" + ) + + _landed_tags(key, lambda values: len(values) == len(first) + len(second) + 1) + + +def test_worker_kill_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + with ( + owned_proxy_process(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + anthropic_model, openai_model = _deployments(scenario, wire.url) + key: Final = scenario.key() + + def burst(index: int) -> tuple: + return _tagged_requests(candidate, key, anthropic_model, openai_model, stream=False, index=index) + + with ThreadPoolExecutor(max_workers=5) as pool: + first: Final = tuple(chain.from_iterable(pool.map(burst, range(3)))) + + workers: Final = [ + child + for child in psutil.Process(owned.process.pid).children(recursive=True) + if any(marker in " ".join(child.cmdline()) for marker in ("spawn_main", "integration._support.proxy")) + ] + assert len(workers) == 2, ( + f"expected two uvicorn workers, found {[(w.pid, w.cmdline()[:3]) for w in workers]}" + ) + first_landed: Final = _landed_tags(key, lambda values: len(values) == len(first)) + workers[0].kill() + psutil.wait_procs(workers[:1], timeout=10) + assert not workers[0].is_running() + + with ThreadPoolExecutor(max_workers=5) as pool: + second: Final = tuple(chain.from_iterable(pool.map(lambda i: burst(100 + i), range(3)))) + responses: Final = [*first, *second] + for position in range(len(ROUTES)): + statuses: Final = { + responses[offset + position].status_code for offset in range(0, len(responses), len(ROUTES)) + } + assert 200 in statuses, f"no surviving 200 for route {ROUTES[position]}: {statuses}" + second_ok: Final = [response for response in second if response.status_code == 200] + assert second_ok, "surviving worker served no second-burst request" + ok: Final = [response for response in responses if response.status_code == 200] + ids: Final = [_ids(response) for response in ok] + assert len(set(ids)) == len(ids), "duplicate upstream id in burst" + _landed_tags(key, lambda values: len(values) == len(first_landed) + len(second_ok)) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 4b25da2ff79..662482d34da 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2971,6 +2971,30 @@ def test_get_extra_header_tags(): delattr(litellm, "extra_spend_tag_headers") +def test_get_request_tags_reads_header_tags_from_starlette_headers(): + from starlette.datastructures import Headers + + import litellm + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + original_extra_headers = getattr(litellm, "extra_spend_tag_headers", None) + original_disable_user_agent = litellm.disable_add_user_agent_to_request_tags + try: + litellm.extra_spend_tag_headers = ["x-tenant-id"] + litellm.disable_add_user_agent_to_request_tags = False + proxy_server_request = {"headers": Headers({"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"})} + + assert StandardLoggingPayloadSetup._get_request_tags( + litellm_params={}, proxy_server_request=proxy_server_request + ) == ["User-Agent: claude-cli", "User-Agent: claude-cli/2.0.0", "x-tenant-id: tenant-a"] + finally: + if original_extra_headers is not None: + litellm.extra_spend_tag_headers = original_extra_headers + elif hasattr(litellm, "extra_spend_tag_headers"): + delattr(litellm, "extra_spend_tag_headers") + litellm.disable_add_user_agent_to_request_tags = original_disable_user_agent + + def test_response_cost_calculator_with_response_cost_in_hidden_params(logging_obj): from litellm import Router