From 3d91b1dcf1e13d245faf978cba4cc2c55b11e697 Mon Sep 17 00:00:00 2001 From: kerry Date: Fri, 2 Oct 2026 21:06:20 +0000 Subject: [PATCH] test(spend): rework audit cells for review findings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend/test_spend_log_request_tags.py | 149 +++------- .../test_spend_log_request_tags_chaos.py | 266 +++++++++++------- 2 files changed, 203 insertions(+), 212 deletions(-) diff --git a/tests/integration/spend/test_spend_log_request_tags.py b/tests/integration/spend/test_spend_log_request_tags.py index 67fee51d597..e064d9f6496 100644 --- a/tests/integration/spend/test_spend_log_request_tags.py +++ b/tests/integration/spend/test_spend_log_request_tags.py @@ -6,7 +6,9 @@ from hashlib import sha256 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 @@ -29,46 +31,18 @@ SENT_HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant- EXPECTED_TAGS: Final = ["User-Agent: claude-cli", "User-Agent: claude-cli/2.0.0", "x-tenant-id: tenant-a"] -def _respond(request: Request) -> Reply: - assert request.method == "POST" and request.target == "/v1/messages", request.target - return Reply( - body=json.dumps( - { - "id": f"msg_{uuid.uuid4().hex}", - "type": "message", - "role": "assistant", - "model": MODEL, - "content": [{"type": "text", "text": "tagged"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 10, "output_tokens": 2}, - } - ).encode() - ) - - -def _request_tags(key: str) -> list[list[str]]: - rows: Final = read_rows( - 'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),) - ) - return [ - json.loads(row["request_tags"]) if isinstance(row["request_tags"], str) else row["request_tags"] for row in rows - ] - - @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(_respond) as wire: - config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - config["litellm_settings"]["extra_spend_tag_headers"] = ["x-tenant-id"] - path: Final = tmp_path / "spend-tag-headers.yaml" - path.write_text(yaml.safe_dump(config)) - environment: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": "synthetic-anthropic-key"} - with owned_proxy(gateway, tmp_path, environment, config=path) as candidate, candidate.scenario() as scenario: + 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" ) @@ -86,30 +60,21 @@ 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] + 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("/") -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"]}}) + 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, @@ -133,12 +98,9 @@ def test_pass_through_anthropic_stream_records_header_tags(gateway: Gateway, tmp assert eventually(lambda: tags_by_key(key), 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"]}}) + 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, @@ -157,12 +119,9 @@ def test_pass_through_anthropic_sdk_records_header_tags(gateway: Gateway, tmp_pa ] -# 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"]}}) + 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, @@ -182,11 +141,10 @@ def test_pass_through_anthropic_sdk_stream_records_header_tags(gateway: Gateway, ] -# 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"]}}) + 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, @@ -209,13 +167,10 @@ def test_pass_through_openai_chat_records_header_tags(gateway: Gateway, tmp_path assert eventually(lambda: tags_by_key(key), 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"]}}) + 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, @@ -240,10 +195,9 @@ def test_pass_through_openai_sdk_records_header_tags(gateway: Gateway, tmp_path: ] -# 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"]}}) + 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, @@ -261,10 +215,9 @@ def test_pass_through_gemini_records_header_tags(gateway: Gateway, tmp_path: Pat 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( + config: Final = write_config( tmp_path, { "litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}, @@ -304,12 +257,9 @@ def test_custom_pass_through_endpoint_records_header_tags(gateway: Gateway, tmp_ 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"]}}) + 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, @@ -338,10 +288,9 @@ def test_pass_through_openai_async_sdk_stream_records_header_tags(gateway: Gatew ] -# 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"]}}) + 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, @@ -363,13 +312,10 @@ def test_unified_chat_completions_records_header_tags(gateway: Gateway, tmp_path ] -# 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"]}}) + 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, @@ -390,10 +336,9 @@ def test_unified_responses_records_header_tags(gateway: Gateway, tmp_path: Path, assert eventually(lambda: tags_by_key(key), 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"]}}) + 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, @@ -412,13 +357,12 @@ def test_unified_cache_hit_twin_records_header_tags(gateway: Gateway, tmp_path: ] -# 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( + config: Final = write_config( tmp_path, { "litellm_settings": { @@ -470,10 +414,9 @@ def test_pass_through_tags_reach_generic_api_sink(gateway: Gateway, tmp_path: Pa ] -# 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"]}}) + 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, @@ -506,7 +449,6 @@ def test_pass_through_tags_accrue_daily_tag_spend(gateway: Gateway, tmp_path: Pa 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")] ) @@ -514,7 +456,7 @@ 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, {}) + 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, @@ -542,13 +484,12 @@ def test_header_tags_without_extra_spend_tag_headers_record_user_agent_only( ] -# 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"]}}) + 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, @@ -583,13 +524,12 @@ def test_routes_without_headers_record_no_tags(gateway: Gateway, tmp_path: Path, assert eventually(lambda: tags_by_id(request_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( + config: Final = write_config( tmp_path, { "litellm_settings": { @@ -624,13 +564,12 @@ def test_disabled_user_agent_keeps_only_extra_header_tags(gateway: Gateway, tmp_ ] -# 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"]}}) + 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, @@ -657,13 +596,12 @@ def test_unsent_configured_header_contributes_no_tag(gateway: Gateway, tmp_path: ] -# 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"]}}) + 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, @@ -691,16 +629,14 @@ def test_default_httpx_user_agent_is_recorded(gateway: Gateway, tmp_path: Path, ] -# 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: +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 = _owned_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}}) + 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, ): - before: Final = _spend_count() anonymous_before: Final = len( read_rows("SELECT request_tags FROM \"LiteLLM_SpendLogs\" WHERE api_key IS NULL OR api_key=''", ()) ) @@ -740,16 +676,14 @@ def test_unauthenticated_request_writes_no_spend_row(gateway: Gateway, tmp_path: assert eventually(lambda: tags_by_id(control.json()["id"]), lambda tags: len(tags) == 1, seconds=70) == [ EXPECTED_TAGS ] - assert _spend_count() == before + 2 -# 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"]}}) + 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, @@ -777,7 +711,6 @@ def test_upstream_failure_still_records_header_tags(gateway: Gateway, tmp_path: assert eventually(lambda: tags_by_key(key), lambda tags: len(tags) == 1, seconds=70) == [expected] -# 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")] @@ -786,7 +719,7 @@ 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}}) + 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, @@ -813,10 +746,9 @@ def test_null_and_empty_extra_spend_tag_headers_record_user_agent_only( ] -# 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"]}}) + 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, @@ -854,14 +786,13 @@ def test_configured_header_case_differs_between_routes(gateway: Gateway, tmp_pat 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"]}}) + 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, @@ -892,10 +823,9 @@ def test_large_header_value_is_stored_verbatim(gateway: Gateway, tmp_path: Path, ] -# 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"]}}) + 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, @@ -926,13 +856,12 @@ def test_duplicate_header_values_follow_carrier_semantics(gateway: Gateway, tmp_ 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"]}}) + 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, @@ -960,13 +889,12 @@ def test_x_litellm_tags_merges_with_header_tags(gateway: Gateway, tmp_path: Path ] -# 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"]}}) + 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, @@ -997,10 +925,9 @@ def test_repeated_requests_each_record_tags(gateway: Gateway, tmp_path: Path, ro ) == [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( + config: Final = write_config( tmp_path, { "litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}, @@ -1048,5 +975,5 @@ def test_guardrail_mode_tag_decider_is_unchanged_on_pass_through(gateway: Gatewa key=key, headers=SENT_HEADERS, ) - assert control.status_code == 500, control.text + assert control.status_code != 200, control.text assert len(wire.drain()) == 0, "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 index 87be663635b..26b33a9c68c 100644 --- a/tests/integration/spend/test_spend_log_request_tags_chaos.py +++ b/tests/integration/spend/test_spend_log_request_tags_chaos.py @@ -1,15 +1,17 @@ import json import threading -from hashlib import sha256 from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 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.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 ( + MODEL, OPENAI_MODEL, T3, provider_env, @@ -17,12 +19,16 @@ from integration.spend._request_tag_helpers import ( 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 +ROUTES: Final = ( + "/anthropic/v1/messages", + "/openai/v1/chat/completions", + "/v1/chat/completions", + "/v1/messages", + "/v1/responses", +) def _ids(response) -> str: @@ -32,77 +38,118 @@ def _ids(response) -> str: 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).""" +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}" - 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, + return ( + candidate.request( + "POST", + "/anthropic/v1/messages", + {"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": marker}], "stream": stream}, + key=key, + headers=ANTHROPIC_HEADERS, + ), + candidate.request( + "POST", + "/openai/v1/chat/completions", + {"model": OPENAI_MODEL, "messages": [{"role": "user", "content": marker}], "stream": stream}, + key=key, + headers=HEADERS, + ), + candidate.request( + "POST", + "/v1/chat/completions", + {"model": openai_model, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers=HEADERS, + ), + candidate.request( + "POST", + "/v1/messages", + { + "model": anthropic_model, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + }, + key=key, + headers=ANTHROPIC_HEADERS, + ), + candidate.request("POST", "/v1/responses", {"model": openai_model, "input": marker}, key=key, headers=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 _deployments(scenario, url: str) -> tuple[str, str]: + return ( + scenario.model(model=f"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, count: int) -> list[dict]: + 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,)), + lambda values: len(values) == count, + seconds=70, + ) + assert len({row["request_id"] for row in landed}) == count + 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(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as candidate, - candidate.scenario() as scenario, + owned_proxy_process(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as owned, + owned.gateway.scenario() as scenario, ): - model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + 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, model, stream=index % 2 == 1, index=index) + return _tagged_requests( + candidate, key, anthropic_model, openai_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 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)) == 30, "duplicate upstream id in burst" - assert len(wire.drain()) == 30 - 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,) - ), - lambda values: len(values) == 30, - seconds=70, - ) - assert len({row["request_id"] for row in landed}) == 30 - for row in landed: - value: Final = row["request_tags"] - assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED + 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(worker.is_running() for worker in psutil.process_iter(pids)) + _landed_tags(key, 50) -# 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() + down: Final = threading.Event() + delivered: Final = [] # mutable-ok: sink thread appends between drains def stoppable_sink(request: Request) -> Reply: - stopped.wait(timeout=30) - return Reply(status=503) + if down.is_set(): + 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( @@ -116,7 +163,7 @@ def test_sink_outage_does_not_lose_spend_log_tags(gateway: Gateway, tmp_path: Pa }, ) with ( - owned_proxy( + owned_proxy_process( gateway, tmp_path, { @@ -126,45 +173,65 @@ def test_sink_outage_does_not_lose_spend_log_tags(gateway: Gateway, tmp_path: Pa }, config=config, workers=2, - ) as candidate, - candidate.scenario() as scenario, + ) as owned, + owned.gateway.scenario() as scenario, ): - model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + 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, model, stream=False, index=index) + return _tagged_requests(candidate, key, anthropic_model, openai_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 + with ThreadPoolExecutor(max_workers=5) as pool: + first: Final = [response for group in pool.map(burst, range(4)) for response in group] + assert all(response.status_code == 200 for response in first), [ + (response.status_code, response.text[:200]) for response in first + ] - def second_burst(index: int) -> tuple: - return _tagged_requests(candidate, key, model, stream=False, index=100 + index) + def call_ids(responses: list) -> set: + return {response.headers["x-litellm-call-id"] for response in responses} - 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] + first_ids: Final = call_ids(first) + + def events_for(ids: set) -> set: + return { + event["litellm_call_id"] + for batch in delivered + for event in json.loads(batch.body) + if event.get("litellm_call_id") in ids + } + + eventually(lambda: events_for(first_ids), lambda found: found == first_ids, seconds=70) + down.set() + with ThreadPoolExecutor(max_workers=5) as pool: + second: Final = [ + response for group in pool.map(lambda i: burst(100 + i), range(4)) for response in group + ] + down.clear() + third: Final = _tagged_requests(candidate, key, anthropic_model, openai_model, False, 200) + responses: Final = [*first, *second, *third] 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)) == len(ids), "duplicate upstream id in burst" - 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,) - ), - lambda values: len(values) == 36, + second_ids: Final = call_ids(second) + third_ids: Final = call_ids(list(third)) + recovery_probe: Final = next(iter(third_ids)) + eventually( + lambda: events_for(third_ids), + lambda found: recovery_probe in found, seconds=70, ) - assert len({row["request_id"] for row in landed}) == 36 - for row in landed: - value: Final = row["request_tags"] - assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED + second_delivered: Final = eventually( + lambda: events_for(second_ids), + lambda found: len(found) == len(second_ids), + seconds=30, + return_last_on_timeout=True, + ) + assert second_delivered == second_ids + _landed_tags(key, len(responses)) -# 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"]}}) @@ -173,41 +240,38 @@ def test_worker_kill_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_path: P owned.gateway.scenario() as scenario, ): candidate: Final = owned.gateway - model: Final = scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{wire.url}/v1") + anthropic_model, openai_model = _deployments(scenario, wire.url) key: Final = scenario.key() def burst(index: int) -> tuple: - return _tagged_requests(candidate, key, model, stream=False, index=index) + return _tagged_requests(candidate, key, anthropic_model, openai_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] + with ThreadPoolExecutor(max_workers=5) as pool: + first: Final = [response for group in pool.map(burst, range(4)) for response in group] workers: Final = [ child for child in psutil.Process(owned.process.pid).children(recursive=True) - if child.status() != psutil.STATUS_ZOMBIE + if any(marker in " ".join(child.cmdline()) for marker in ("spawn_main", "integration._support.proxy")) ] - assert len(workers) >= 2, f"expected two proxy workers, found {[w.pid for w in workers]}" + assert len(workers) == 2, ( + f"expected two uvicorn workers, found {[(w.pid, w.cmdline()[:3]) for w in workers]}" + ) workers[0].kill() + psutil.wait_procs(workers[:1], timeout=10) + assert not workers[0].is_running() - 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] + with ThreadPoolExecutor(max_workers=5) as pool: + second: Final = [ + response for group in pool.map(lambda i: burst(100 + i), range(4)) for response in group + ] 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}" 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" - 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,) - ), - lambda values: len(values) == len(ids), - seconds=70, - ) - assert len({row["request_id"] for row in landed}) == len(ids) - for row in landed: - value: Final = row["request_tags"] - assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED + _landed_tags(key, len(ok))