From c01623b7656748acb92e80a424b157c639ec2da5 Mon Sep 17 00:00:00 2001 From: kerry Date: Fri, 2 Oct 2026 17:47:09 +0000 Subject: [PATCH] test(spend): audit header-derived request tags across routes, modes and outages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/spend/_request_tag_helpers.py | 219 +++++ .../spend/test_spend_log_request_tags.py | 924 ++++++++++++++++++ .../test_spend_log_request_tags_chaos.py | 209 ++++ 3 files changed, 1352 insertions(+) create mode 100644 tests/integration/spend/_request_tag_helpers.py create mode 100644 tests/integration/spend/test_spend_log_request_tags_chaos.py diff --git a/tests/integration/spend/_request_tag_helpers.py b/tests/integration/spend/_request_tag_helpers.py new file mode 100644 index 00000000000..6bf72f594c5 --- /dev/null +++ b/tests/integration/spend/_request_tag_helpers.py @@ -0,0 +1,219 @@ +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 {} + 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 request.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 request.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 request.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 request.target.endswith(":generateContent") or request.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 index 9d89a866cc8..f905af8258b 100644 --- a/tests/integration/spend/test_spend_log_request_tags.py +++ b/tests/integration/spend/test_spend_log_request_tags.py @@ -1,15 +1,26 @@ +import asyncio import json import uuid from hashlib import sha256 from pathlib import Path from typing import Final +import httpx 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 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, + write_config, +) MODEL: Final = "claude-sonnet-4-5-20250929" SENT_HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"} @@ -74,3 +85,916 @@ def test_header_derived_spend_tags_are_recorded_on_anthropic_messages_routes( assert response.status_code == 200, response.text assert len(wire.drain()) == 1 assert eventually(lambda: _request_tags(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS] + + +def _base_url(candidate: Gateway) -> str: + return str(candidate.client.base_url).rstrip("/") + + +def _spend_count() -> int: + return read_rows('SELECT count(*) AS n FROM "LiteLLM_SpendLogs"')[0]["n"] + + +def _owned_config(tmp_path: Path, mutations: dict) -> Path: + return write_config(tmp_path, mutations) + + +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" + + +# H2: pass-through streaming anthropic request records the header tags +def test_pass_through_anthropic_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}], + "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 + request_id: Final = f"msg_{response.text.split('msg_')[1].split(chr(34))[0]}" + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +# H3: Anthropic SDK non-stream call through the pass-through route +def test_pass_through_anthropic_sdk_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + import anthropic + + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}] + ) + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(message.id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +# H4: Anthropic SDK streaming call through the pass-through route +def test_pass_through_anthropic_sdk_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + import anthropic + + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}] + ) as stream: + message: Final = stream.get_final_message() + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(message.id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +# H5: openai pass-through route records the header tags +@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 = _owned_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": "tag me"}], + **({"stream": True} if stream else {}), + }, + key=key, + headers=SENT_HEADERS, + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + request_id: Final = f"chatcmpl_{response.text.split('chatcmpl_')[1].split(chr(34))[0]}" + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +# H6: OpenAI SDK sync and stream calls through the pass-through route +@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: + import openai + + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}], stream=stream + ) + if stream: + request_id: Final = next(chunk.id for chunk in completion) + for _ in completion: + pass + else: + request_id = completion.id + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +# H7: gemini pass-through route records the header tags +def test_pass_through_gemini_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}]}]}, + 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] + + +# H8: user-defined pass-through endpoint records the header 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 = _owned_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": "tag me"}]}, + 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] + + +# H9: async OpenAI SDK streaming call through the pass-through route +def test_pass_through_openai_async_sdk_stream_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + import openai + + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}], stream=True + ) + first: Final = await completion.__anext__() + async for _ in completion: + pass + return first.id + + request_id: Final = asyncio.run(call()) + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +# H10: unified chat completions control records the header tags +def test_unified_chat_completions_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}]}, + 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 + ] + + +# H11: unified /v1/responses via the OpenAI SDK, sync and stream +@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: + import openai + + with wire_server(provider_reply) as wire: + config: Final = _owned_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="tag me", 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 len(wire.drain()) == 1 + assert eventually(lambda: tags_by_id(request_id), lambda tags: len(tags) == 1, seconds=70) == [ + EXPECTED_TAGS + ] + + +# H12: a unified cache-hit twin records the same tags on both spend rows +def test_unified_cache_hit_twin_records_header_tags(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = _owned_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, + ] + + +# H13: generic_api callback sink sees the same request_tags as the spend row +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 = _owned_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": "tag me"}]}, + 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() -> list[dict]: + batches.extend(endpoint.drain()) + return [event for batch in batches for event in json.loads(batch.body) 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 + ] + + +# H14: LiteLLM_DailyTagSpend accrues each header-derived tag +def test_pass_through_tags_accrue_daily_tag_spend(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}]}, + 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) + + +# S1: extra_spend_tag_headers unset records only the user-agent tags on both routes +@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 = _owned_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": "tag me"}], + }, + 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 + ] + + +# S2: pass-through request without the headers records no tags +@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 = _owned_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": "tag me"}], + }, + key=key, + 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) == [ + [] + ] + + +# S3: disable_add_user_agent_to_request_tags keeps only the extra header tags +@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 = _owned_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": "tag me"}], + }, + 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] + ] + + +# S4: a configured header the client never sends contributes no 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 = _owned_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": "tag me"}], + }, + 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] + ] + + +# S5: httpx default user-agent is recorded when the client sends none +@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 = _owned_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": "tag me"}], + }, + 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 + ] + + +# S6: unauthenticated requests return 401 and write no spend row +@pytest.mark.parametrize("route", [pytest.param("/anthropic/v1/messages", id="passthrough")]) +def test_unauthenticated_request_writes_no_spend_row(gateway: Gateway, tmp_path: Path, route: str) -> None: + with wire_server(provider_reply) as wire: + config: Final = _owned_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, + ): + before: Final = _spend_count() + response: Final = candidate.client.post( + route, + json={"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "tag me"}]}, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 401, response.text + key: Final = scenario.key() + control: Final = candidate.request( + "POST", + route, + {"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "tag me"}]}, + 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 + ] + assert _spend_count() == before + 1 + + +# S7: an upstream 400 surfaces the same status and its spend row records the 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 = _owned_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": "claude-nonexistent-model" if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "tag me"}], + }, + key=key, + headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 400, response.text + assert len(wire.drain()) == 1 + assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [EXPECTED_TAGS] + + +# S8: null and empty extra_spend_tag_headers behave like unset +@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 = _owned_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": "tag me"}], + }, + 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] + ] + + +# S9: pass-through matches the configured header name case-insensitively, unified is case-sensitive +def test_configured_header_case_differs_between_routes(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}]} + 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 + + +# E1: a 5KB header value is stored verbatim +@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 = _owned_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": "tag me"}], + }, + 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}"] + ] + + +# E2: duplicate configured headers record first-value on pass-through, last-value on unified +def test_duplicate_header_values_follow_carrier_semantics(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(provider_reply) as wire: + config: Final = _owned_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": "tag me"}]} + 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"] + + +# E4: x-litellm-tags and header-derived tags land together +@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 = _owned_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": "tag me"}], + }, + 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] + if route == "/v1/messages" + else [UA_FAMILY_TAG, UA_TAG, TENANT_TAG, "team-x"] + ) + assert eventually(lambda: tags_by_id(response.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ + expected + ] + + +# E5: three identical requests write three spend rows, each with the tags +@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 = _owned_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 if route == "/anthropic/v1/messages" else model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"repeat {uuid.uuid4().hex}"}], + } + ids: Final = [] + for _ in range(3): + response: Final = candidate.request( + "POST", route, body, key=key, headers={**SENT_HEADERS, "anthropic-version": "2023-06-01"} + ) + assert response.status_code == 200, response.text + ids.append(response.json()["id"]) + 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] + + +# E6: guardrail mode-by-tag decider behaves identically with and without the fix +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 = _owned_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()) == 1, "tag-matched guardrail should have blocked before the upstream" 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..cd4f8e16691 --- /dev/null +++ b/tests/integration/spend/test_spend_log_request_tags_chaos.py @@ -0,0 +1,209 @@ +import json +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import psutil +from integration._support.client import Gateway, eventually +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 ( + OPENAI_MODEL, + T3, + provider_env, + provider_reply, + write_config, +) + +from tests.integration._support.database import read_rows + +MODEL: Final = "claude-sonnet-4-5-20250929" +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 + + +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 _tagged_requests(candidate: Gateway, key: str, model: str, stream: bool, index: int) -> tuple: + """One call per route shape, all with the same client headers; returns (response, request_id).""" + marker: Final = f"burst {index}" + anthropic: Final = candidate.request( + "POST", + "/anthropic/v1/messages", + {"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": marker}], "stream": stream}, + key=key, + headers=ANTHROPIC_HEADERS, + ) + openai: Final = candidate.request( + "POST", + "/openai/v1/chat/completions", + {"model": OPENAI_MODEL, "messages": [{"role": "user", "content": marker}], "stream": stream}, + key=key, + headers=HEADERS, + ) + unified: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers=HEADERS, + ) + return anthropic, openai, unified + + +# C1: 10 concurrent bursts x 3 routes; each response id lands exactly one spend row with the tags +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(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() + + def burst(index: int) -> tuple: + return _tagged_requests(candidate, key, model, stream=index % 2 == 1, index=index) + + with ThreadPoolExecutor(max_workers=10) as pool: + responses: Final = [response for group in pool.map(burst, range(10)) for response in group] + assert len(responses) == 30 + 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)) == 30, "duplicate upstream id in burst" + assert len(wire.drain()) == 30 + landed: Final = eventually( + lambda: read_rows( + 'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', + (ids,), + ), + lambda values: len(values) == 30, + seconds=70, + ) + assert sorted(row["request_id"] for row in landed) == sorted(ids) + for row in landed: + value: Final = row["request_tags"] + assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED + + +# C2: generic_api sink down mid burst; spend rows still land exactly once with the tags +def test_sink_outage_does_not_lose_spend_log_tags(gateway: Gateway, tmp_path: Path) -> None: + stopped: Final = threading.Event() + + def stoppable_sink(request: Request) -> Reply: + stopped.wait(timeout=30) + return Reply(status=503) + + 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( + 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, + ): + model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + key: Final = scenario.key() + + def burst(index: int) -> tuple: + return _tagged_requests(candidate, key, model, stream=False, index=index) + + with ThreadPoolExecutor(max_workers=6) as pool: + first: Final = [response for group in pool.map(burst, range(6)) for response in group] + stopped.set() # sink goes down: the peer now returns 503 to every flush + + def second_burst(index: int) -> tuple: + return _tagged_requests(candidate, key, model, stream=False, index=100 + index) + + with ThreadPoolExecutor(max_workers=6) as pool: + second: Final = [response for group in pool.map(second_burst, range(6)) for response in group] + responses: Final = [*first, *second] + 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] + landed: Final = eventually( + lambda: read_rows( + 'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', + (ids,), + ), + lambda values: len(values) == 36, + seconds=70, + ) + for row in landed: + value: Final = row["request_tags"] + assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED + + +# C3: killing one proxy worker mid burst loses no spend row +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 + model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + key: Final = scenario.key() + + def burst(index: int) -> tuple: + return _tagged_requests(candidate, key, model, stream=False, index=index) + + with ThreadPoolExecutor(max_workers=6) as pool: + first: Final = [response for group in pool.map(burst, range(6)) for response in group] + + workers: Final = [ + child + for child in psutil.Process(owned.process.pid).children(recursive=True) + if child.status() != psutil.STATUS_ZOMBIE + ] + assert len(workers) >= 2, f"expected two proxy workers, found {[w.pid for w in workers]}" + workers[0].kill() + + def second_burst(index: int) -> tuple: + return _tagged_requests(candidate, key, model, stream=False, index=100 + index) + + with ThreadPoolExecutor(max_workers=6) as pool: + second: Final = [response for group in pool.map(second_burst, range(6)) for response in group] + responses: Final = [*first, *second] + ok: Final = [response for response in responses if response.status_code == 200] + ids: Final = [_ids(response) for response in ok] + landed: Final = eventually( + lambda: read_rows( + 'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s)', + (ids,), + ), + lambda values: len(values) == len(ids), + seconds=70, + ) + assert sorted(row["request_id"] for row in landed) == sorted(ids) + for row in landed: + value: Final = row["request_tags"] + assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED