Merge 79663c9037 into 6532dcb73b

litellm/litellm_core_utils/litellm_logging.py

@@ -6301,7 +6301,7 @@ class StandardLoggingPayloadSetup:
             return None
         user_agent_tags: list[str] | None = None
         headers: Final = proxy_server_request.get("headers", {})
-        if headers is not None and isinstance(headers, dict):
+        if headers is not None and isinstance(headers, Mapping):
             if "user-agent" in headers:
                 user_agent: Final = headers["user-agent"]
                 if user_agent is not None:
@@ -6326,7 +6326,7 @@ class StandardLoggingPayloadSetup:
             return None
 
         headers: Final = proxy_server_request.get("headers", {})
-        if not isinstance(headers, dict):
+        if not isinstance(headers, Mapping):
             return None
 
         header_tags: Final = []

tests/integration/spend/_request_tag_helpers.py

@@ -0,0 +1,222 @@
+import json
+import uuid
+from hashlib import sha256
+from pathlib import Path
+from typing import Final
+
+import yaml
+from integration._support.database import read_rows
+from integration._support.wire import Reply, Request
+
+ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929"
+OPENAI_MODEL: Final = "gpt-4o-mini"
+GEMINI_MODEL: Final = "gemini-2.5-flash"
+HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"}
+T3: Final = ["User-Agent: claude-cli", "User-Agent: claude-cli/2.0.0", "x-tenant-id: tenant-a"]
+ANTHROPIC_SONNET_BODY: Final = {
+    "id": "msg_synthetic",
+    "type": "message",
+    "role": "assistant",
+    "model": ANTHROPIC_MODEL,
+    "content": [{"type": "text", "text": "tagged"}],
+    "stop_reason": "end_turn",
+    "stop_sequence": None,
+    "usage": {"input_tokens": 10, "output_tokens": 2},
+}
+CHAT_COMPLETION_BODY: Final = {
+    "id": "chatcmpl-synthetic",
+    "object": "chat.completion",
+    "created": 1,
+    "model": OPENAI_MODEL,
+    "choices": [{"index": 0, "message": {"role": "assistant", "content": "tagged"}, "finish_reason": "stop"}],
+    "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
+}
+RESPONSES_BODY: Final = {
+    "id": "resp_synthetic",
+    "object": "response",
+    "created_at": 1,
+    "status": "completed",
+    "model": OPENAI_MODEL,
+    "output": [
+        {
+            "type": "message",
+            "id": "msg_synthetic",
+            "status": "completed",
+            "role": "assistant",
+            "content": [{"type": "output_text", "text": "tagged", "annotations": []}],
+        }
+    ],
+    "usage": {
+        "input_tokens": 10,
+        "output_tokens": 2,
+        "total_tokens": 12,
+        "input_tokens_details": {"cached_tokens": 0},
+        "output_tokens_details": {"reasoning_tokens": 0},
+    },
+}
+GEMINI_BODY: Final = {
+    "candidates": [
+        {
+            "content": {"parts": [{"text": "tagged"}], "role": "model"},
+            "finishReason": "STOP",
+        }
+    ],
+    "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 2, "totalTokenCount": 12},
+}
+ANTHROPIC_STREAM_EVENTS: Final = (
+    {
+        "type": "message_start",
+        "message": {**ANTHROPIC_SONNET_BODY, "content": [], "usage": {"input_tokens": 10, "output_tokens": 1}},
+    },
+    {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+    {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "tagged"}},
+    {"type": "content_block_stop", "index": 0},
+    {
+        "type": "message_delta",
+        "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+        "usage": {"output_tokens": 2},
+    },
+    {"type": "message_stop"},
+)
+OPENAI_STREAM_CHUNKS: Final = (
+    {
+        "id": CHAT_COMPLETION_BODY["id"],
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": OPENAI_MODEL,
+        "choices": [{"index": 0, "delta": {"role": "assistant", "content": "tagged"}, "finish_reason": None}],
+    },
+    {
+        "id": CHAT_COMPLETION_BODY["id"],
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": OPENAI_MODEL,
+        "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+    },
+    {
+        "id": CHAT_COMPLETION_BODY["id"],
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": OPENAI_MODEL,
+        "choices": [],
+        "usage": CHAT_COMPLETION_BODY["usage"],
+    },
+)
+
+
+def _message_id() -> str:
+    return "msg_" + uuid.uuid4().hex
+
+
+def _sse_frames(events: tuple[dict, ...]) -> tuple[bytes, ...]:
+    return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events)
+
+
+def provider_reply(request: Request) -> Reply:
+    """Scripted edge for every route the audit drives: anthropic messages, openai chat completions and
+    responses, gemini generateContent. Error bodies keyed off the sentinel model name."""
+    body: Final = json.loads(request.body) if request.body else {}
+    target: Final = request.target.split("?", 1)[0]
+    if body.get("model") == "claude-nonexistent-model":
+        return Reply(
+            status=400,
+            body=json.dumps(
+                {"error": {"type": "invalid_request_error", "message": "model: claude-nonexistent-model"}}
+            ).encode(),
+        )
+    if target == "/v1/models" or target.startswith("/v1/models/"):
+        return Reply(body=json.dumps({"object": "list", "data": []}).encode())
+    if target == "/v1/messages":
+        identity: Final = _message_id()
+        if body.get("stream") is True:
+            events: Final = tuple(
+                {**event, "message": {**event.get("message", {}), "id": identity}} if "message" in event else event
+                for event in ANTHROPIC_STREAM_EVENTS
+            )
+            return Reply(content_type="text/event-stream", chunks=_sse_frames(events))
+        return Reply(body=json.dumps({**ANTHROPIC_SONNET_BODY, "id": identity}).encode())
+    if target == "/v1/chat/completions":
+        identity = "chatcmpl_" + uuid.uuid4().hex
+        if body.get("stream") is True:
+            frames: Final = tuple(
+                f"data: {json.dumps({**chunk, 'id': identity})}\n\n".encode() for chunk in OPENAI_STREAM_CHUNKS
+            ) + (b"data: [DONE]\n\n",)
+            return Reply(content_type="text/event-stream", chunks=frames)
+        return Reply(body=json.dumps({**CHAT_COMPLETION_BODY, "id": identity}).encode())
+    if target == "/v1/responses":
+        identity = "resp_" + uuid.uuid4().hex
+        message: Final = _message_id()
+        completed_body: Final = {
+            **RESPONSES_BODY,
+            "id": identity,
+            "output": [{**RESPONSES_BODY["output"][0], "id": message}],
+        }
+        if body.get("stream") is True:
+            created: Final = {
+                "type": "response.created",
+                "response": {**completed_body, "status": "in_progress", "output": [], "usage": None},
+            }
+            delta: Final = {
+                "type": "response.output_text.delta",
+                "item_id": message,
+                "output_index": 0,
+                "content_index": 0,
+                "delta": "tagged",
+            }
+            completed: Final = {"type": "response.completed", "response": completed_body}
+            return Reply(content_type="text/event-stream", chunks=_sse_frames((created, delta, completed)))
+        return Reply(body=json.dumps(completed_body).encode())
+    if target.endswith(":generateContent") or target.endswith(":streamGenerateContent"):
+        return Reply(body=json.dumps(GEMINI_BODY).encode())
+    raise AssertionError(f"unexpected upstream target {request.target}")
+
+
+def provider_env(url: str) -> dict[str, str]:
+    return {
+        "ANTHROPIC_API_BASE": url,
+        "ANTHROPIC_API_KEY": "synthetic-anthropic-key",
+        "OPENAI_API_BASE": url,
+        "OPENAI_API_KEY": "synthetic-openai-key",
+        "GEMINI_API_BASE": url,
+        "GEMINI_API_KEY": "synthetic-gemini-key",
+    }
+
+
+def write_config(directory: Path, mutations: dict, name: str = "spend-tag-headers.yaml") -> Path:
+    config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+    for section, values in mutations.items():
+        if isinstance(values, dict) and isinstance(config.get(section), dict):
+            config[section].update(values)
+        else:
+            config[section] = values
+    path: Final = directory / name
+    path.write_text(yaml.safe_dump(config))
+    return path
+
+
+def tags_of(row: dict) -> list:
+    value: Final = row["request_tags"]
+    return json.loads(value) if isinstance(value, str) else value
+
+
+def tags_by_key(key: str) -> list[list]:
+    rows: Final = read_rows(
+        'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)
+    )
+    return [tags_of(row) for row in rows]
+
+
+def tags_by_id(request_id: str) -> list[list]:
+    rows: Final = read_rows('SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,))
+    return [tags_of(row) for row in rows]
+
+
+def row_by_id(request_id: str) -> list[dict]:
+    return read_rows(
+        'SELECT request_id, call_type, request_tags, status, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+        (request_id,),
+    )
+
+
+def unique_marker() -> str:
+    return "tagprobe" + uuid.uuid4().hex[:12]

tests/integration/spend/test_spend_log_request_tags_chaos.py

@@ -0,0 +1,314 @@
+import json
+import threading
+import uuid
+from collections.abc import Callable, Mapping, Sequence, Set
+from concurrent.futures import ThreadPoolExecutor
+from hashlib import sha256
+from itertools import chain
+from pathlib import Path
+from typing import Final
+
+import httpx
+import psutil
+import pytest
+from integration._support.client import Gateway, eventually
+from integration._support.database import read_rows
+from integration._support.process import owned_proxy_process
+from integration._support.wire import Reply, Request, wire_server
+from integration.spend._request_tag_helpers import (
+    ANTHROPIC_MODEL,
+    OPENAI_MODEL,
+    T3,
+    provider_env,
+    provider_reply,
+    write_config,
+)
+
+HEADERS: Final = {"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"}
+ANTHROPIC_HEADERS: Final = {**HEADERS, "anthropic-version": "2023-06-01"}
+EXPECTED: Final = T3
+ROUTES: Final = (
+    "/anthropic/v1/messages",
+    "/openai/v1/chat/completions",
+    "/v1/chat/completions",
+    "/v1/messages",
+    "/v1/responses",
+)
+
+
+def _ids(response) -> str:
+    for prefix in ("msg_", "chatcmpl_", "resp_"):
+        if prefix in response.text:
+            return prefix + response.text.split(prefix)[1].split('"')[0]
+    raise AssertionError(f"no upstream id in {response.text[:200]}")
+
+
+def _requests(
+    candidate: Gateway, key: str, anthropic_model: str, openai_model: str, stream: bool, marker: str
+) -> tuple:
+    """Deferred calls for one request per route in ROUTES order with the same client headers."""
+    return (
+        lambda: candidate.request(
+            "POST",
+            "/anthropic/v1/messages",
+            {
+                "model": ANTHROPIC_MODEL,
+                "max_tokens": 16,
+                "messages": [{"role": "user", "content": marker}],
+                "stream": stream,
+            },
+            key=key,
+            headers=ANTHROPIC_HEADERS,
+        ),
+        lambda: candidate.request(
+            "POST",
+            "/openai/v1/chat/completions",
+            {"model": OPENAI_MODEL, "messages": [{"role": "user", "content": marker}], "stream": stream},
+            key=key,
+            headers=HEADERS,
+        ),
+        lambda: candidate.request(
+            "POST",
+            "/v1/chat/completions",
+            {"model": openai_model, "messages": [{"role": "user", "content": marker}]},
+            key=key,
+            headers=HEADERS,
+        ),
+        lambda: candidate.request(
+            "POST",
+            "/v1/messages",
+            {
+                "model": anthropic_model,
+                "max_tokens": 16,
+                "messages": [{"role": "user", "content": marker}],
+            },
+            key=key,
+            headers=ANTHROPIC_HEADERS,
+        ),
+        lambda: candidate.request(
+            "POST", "/v1/responses", {"model": openai_model, "input": marker}, key=key, headers=HEADERS
+        ),
+    )
+
+
+def _tagged_requests(
+    candidate: Gateway, key: str, anthropic_model: str, openai_model: str, stream: bool, index: int
+) -> tuple:
+    """One call per route in ROUTES order with the same client headers."""
+    marker: Final = f"burst {index} {uuid.uuid4().hex}"
+    return tuple(send() for send in _requests(candidate, key, anthropic_model, openai_model, stream, marker))
+
+
+def _deployments(scenario, url: str) -> tuple[str, str]:
+    return (
+        scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=url, api_key="synthetic-anthropic-key"),
+        scenario.model(model=f"openai/{OPENAI_MODEL}", api_base=f"{url}/v1"),
+    )
+
+
+def _landed_tags(key: str, satisfied) -> Sequence[Mapping]:
+    digest: Final = sha256(key.encode()).hexdigest()
+    landed: Final = eventually(
+        lambda: read_rows('SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)),
+        satisfied,
+        seconds=60,
+    )
+    assert len({row["request_id"] for row in landed}) == len(landed)
+    for row in landed:
+        value: Final = row["request_tags"]
+        assert (json.loads(value) if isinstance(value, str) else value) == EXPECTED
+    return landed
+
+
+def _worker_pids(owned) -> tuple[int, ...]:
+    workers: Final = tuple(
+        child
+        for child in psutil.Process(owned.process.pid).children(recursive=True)
+        if any(marker in " ".join(child.cmdline()) for marker in ("spawn_main", "integration._support.proxy"))
+    )
+    return tuple(worker.pid for worker in workers)
+
+
+def test_burst_across_routes_records_tags_once_per_response(gateway: Gateway, tmp_path: Path) -> None:
+    with wire_server(provider_reply) as wire:
+        config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
+        with (
+            owned_proxy_process(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as owned,
+            owned.gateway.scenario() as scenario,
+        ):
+            candidate: Final = owned.gateway
+            anthropic_model, openai_model = _deployments(scenario, wire.url)
+            key: Final = scenario.key()
+
+            def burst(index: int) -> tuple:
+                return _tagged_requests(
+                    candidate, key, anthropic_model, openai_model, stream=index % 2 == 1, index=index
+                )
+
+            with ThreadPoolExecutor(max_workers=10) as pool:
+                responses: Final = tuple(chain.from_iterable(pool.map(burst, range(10))))
+            assert len(responses) == 50
+            assert all(response.status_code == 200 for response in responses), [
+                (response.status_code, response.text[:200]) for response in responses
+            ]
+            ids: Final = [_ids(response) for response in responses]
+            assert len(set(ids)) == 50, "duplicate upstream id in burst"
+            assert len(wire.drain()) == 50
+            pids: Final = _worker_pids(owned)
+            assert len(set(pids)) == 2, f"expected two uvicorn workers, found {pids}"
+            assert all(psutil.Process(pid).is_running() for pid in pids)
+            _landed_tags(key, lambda values: len(values) == 50)
+
+
+@pytest.mark.timeout(240)
+def test_sink_outage_does_not_lose_spend_log_tags(gateway: Gateway, tmp_path: Path) -> None:
+    down: Final = threading.Event()
+    delivered: Final = []
+    rejected: Final = []
+
+    def stoppable_sink(request: Request) -> Reply:
+        if down.is_set():
+            rejected.append(request)
+            return Reply(status=503)
+        delivered.append(request)
+        return Reply()
+
+    with wire_server(provider_reply) as wire, wire_server(stoppable_sink) as endpoint:
+        config: Final = write_config(
+            tmp_path,
+            {
+                "litellm_settings": {
+                    "extra_spend_tag_headers": ["x-tenant-id"],
+                    "callbacks": ["generic_api"],
+                    "DEFAULT_FLUSH_INTERVAL_SECONDS": 1,
+                }
+            },
+        )
+        with (
+            owned_proxy_process(
+                gateway,
+                tmp_path,
+                {
+                    **provider_env(wire.url),
+                    "GENERIC_LOGGER_ENDPOINT": endpoint.url,
+                    "GENERIC_LOGGER_HEADERS": "Authorization=Bearer synthetic-sink-secret",
+                },
+                config=config,
+                workers=2,
+            ) as owned,
+            owned.gateway.scenario() as scenario,
+        ):
+            candidate: Final = owned.gateway
+            anthropic_model, openai_model = _deployments(scenario, wire.url)
+            key: Final = scenario.key()
+
+            def events_for(batches: Sequence[Request], ids: Set[str]) -> Sequence[Mapping]:
+                events: Final = chain.from_iterable(json.loads(batch.body) for batch in batches)
+                return [event for event in events if event.get("litellm_call_id") in ids]
+
+            def send_and_await(send: Callable[[], httpx.Response], batches: Sequence[Request]) -> httpx.Response:
+                response: Final = send()
+                assert response.status_code == 200, response.text
+                call_id: Final = response.headers["x-litellm-call-id"]
+                eventually(lambda: events_for(batches, {call_id}), lambda found: len(found) >= 1, seconds=30)
+                return response
+
+            def burst_sends(label: str):
+                return chain.from_iterable(
+                    _requests(
+                        candidate,
+                        key,
+                        anthropic_model,
+                        openai_model,
+                        False,
+                        f"{label} {index} {uuid.uuid4().hex}",
+                    )
+                    for index in range(3)
+                )
+
+            first: Final = tuple(send_and_await(send, delivered) for send in burst_sends("burst"))
+            first_ids: Final = {response.headers["x-litellm-call-id"] for response in first}
+            first_events: Final = events_for(delivered, first_ids)
+            first_occurrences: Final = [event["litellm_call_id"] for event in first_events]
+            assert sorted(first_occurrences) == sorted(first_ids), "burst-1 sink delivery is not exactly once per call"
+            for event in first_events:
+                assert event["request_tags"] == EXPECTED
+            down.set()
+            second: Final = tuple(send_and_await(send, rejected) for send in burst_sends("outage"))
+            second_ids: Final = {response.headers["x-litellm-call-id"] for response in second}
+            rejected_ids: Final = {event["litellm_call_id"] for event in events_for(rejected, second_ids)}
+            assert rejected_ids == second_ids, "outage burst was not rejected by the down sink"
+            assert events_for(delivered, second_ids) == [], "burst-2 event delivered to a down sink"
+            down.clear()
+            probe: Final = candidate.request(
+                "POST",
+                "/v1/chat/completions",
+                {
+                    "model": openai_model,
+                    "messages": [{"role": "user", "content": f"recovery probe {uuid.uuid4().hex}"}],
+                },
+                key=key,
+                headers=HEADERS,
+            )
+            assert probe.status_code == 200, probe.text
+            probe_id: Final = probe.headers["x-litellm-call-id"]
+            probe_events: Final = eventually(
+                lambda: events_for(delivered, {probe_id}),
+                lambda found: len(found) >= 1,
+                seconds=70,
+            )
+            assert len(probe_events) == 1, "recovery probe delivered to the sink more than once"
+            assert probe_events[0]["request_tags"] == EXPECTED
+            second_events: Final = events_for(delivered, second_ids)
+            second_occurrences: Final = [event["litellm_call_id"] for event in second_events]
+            assert len(second_occurrences) == len(set(second_occurrences)), (
+                "duplicate burst-2 delivery after the outage"
+            )
+
+            _landed_tags(key, lambda values: len(values) == len(first) + len(second) + 1)
+
+
+def test_worker_kill_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_path: Path) -> None:
+    with wire_server(provider_reply) as wire:
+        config: Final = write_config(tmp_path, {"litellm_settings": {"extra_spend_tag_headers": ["x-tenant-id"]}})
+        with (
+            owned_proxy_process(gateway, tmp_path, provider_env(wire.url), config=config, workers=2) as owned,
+            owned.gateway.scenario() as scenario,
+        ):
+            candidate: Final = owned.gateway
+            anthropic_model, openai_model = _deployments(scenario, wire.url)
+            key: Final = scenario.key()
+
+            def burst(index: int) -> tuple:
+                return _tagged_requests(candidate, key, anthropic_model, openai_model, stream=False, index=index)
+
+            with ThreadPoolExecutor(max_workers=5) as pool:
+                first: Final = tuple(chain.from_iterable(pool.map(burst, range(3))))
+
+            workers: Final = [
+                child
+                for child in psutil.Process(owned.process.pid).children(recursive=True)
+                if any(marker in " ".join(child.cmdline()) for marker in ("spawn_main", "integration._support.proxy"))
+            ]
+            assert len(workers) == 2, (
+                f"expected two uvicorn workers, found {[(w.pid, w.cmdline()[:3]) for w in workers]}"
+            )
+            first_landed: Final = _landed_tags(key, lambda values: len(values) == len(first))
+            workers[0].kill()
+            psutil.wait_procs(workers[:1], timeout=10)
+            assert not workers[0].is_running()
+
+            with ThreadPoolExecutor(max_workers=5) as pool:
+                second: Final = tuple(chain.from_iterable(pool.map(lambda i: burst(100 + i), range(3))))
+            responses: Final = [*first, *second]
+            for position in range(len(ROUTES)):
+                statuses: Final = {
+                    responses[offset + position].status_code for offset in range(0, len(responses), len(ROUTES))
+                }
+                assert 200 in statuses, f"no surviving 200 for route {ROUTES[position]}: {statuses}"
+            second_ok: Final = [response for response in second if response.status_code == 200]
+            assert second_ok, "surviving worker served no second-burst request"
+            ok: Final = [response for response in responses if response.status_code == 200]
+            ids: Final = [_ids(response) for response in ok]
+            assert len(set(ids)) == len(ids), "duplicate upstream id in burst"
+            _landed_tags(key, lambda values: len(values) == len(first_landed) + len(second_ok))

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -2971,6 +2971,30 @@ def test_get_extra_header_tags():
                 delattr(litellm, "extra_spend_tag_headers")
 
 
+def test_get_request_tags_reads_header_tags_from_starlette_headers():
+    from starlette.datastructures import Headers
+
+    import litellm
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    original_extra_headers = getattr(litellm, "extra_spend_tag_headers", None)
+    original_disable_user_agent = litellm.disable_add_user_agent_to_request_tags
+    try:
+        litellm.extra_spend_tag_headers = ["x-tenant-id"]
+        litellm.disable_add_user_agent_to_request_tags = False
+        proxy_server_request = {"headers": Headers({"user-agent": "claude-cli/2.0.0", "x-tenant-id": "tenant-a"})}
+
+        assert StandardLoggingPayloadSetup._get_request_tags(
+            litellm_params={}, proxy_server_request=proxy_server_request
+        ) == ["User-Agent: claude-cli", "User-Agent: claude-cli/2.0.0", "x-tenant-id: tenant-a"]
+    finally:
+        if original_extra_headers is not None:
+            litellm.extra_spend_tag_headers = original_extra_headers
+        elif hasattr(litellm, "extra_spend_tag_headers"):
+            delattr(litellm, "extra_spend_tag_headers")
+        litellm.disable_add_user_agent_to_request_tags = original_disable_user_agent
+
+
 def test_response_cost_calculator_with_response_cost_in_hidden_params(logging_obj):
     from litellm import Router