import asyncio import base64 import json import threading import uuid from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path from queue import SimpleQueue from typing import Final import anthropic import httpx import openai import pytest import yaml from _s3_v2_support import _chat_stream_frames, _responses_stream_frames from integration._support.client import ( Gateway, Scenario, eventually, gateway_from_environment, object_value, string_value, ) from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy, owned_proxy_process from integration._support.timing import assert_cut_short, assert_within_factor_of_control, timed from integration._support.wire import Reply, Request, Wire, wire_server from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest from opentelemetry.proto.common.v1.common_pb2 import KeyValue from opentelemetry.proto.trace.v1.trace_pb2 import Span from pydantic import BaseModel, JsonValue, TypeAdapter pytestmark: Final = pytest.mark.shares_worker_state PUBLIC_KEY: Final = "pk-lf-integration" SECRET_KEY: Final = "sk-lf-integration" PROJECTS_PATH: Final = "/api/public/projects" TRACES_PATH: Final = "/api/public/otel/v1/traces" PROMPTS_PATH: Final = "/api/public/v2/prompts/" STOCK_CONFIG: Final = Path("tests/integration/proxy_config.yaml") CONFIG_SECTIONS: Final = ("litellm_settings", "environment_variables") LANGFUSE_ENVIRONMENT: Final = ("LANGFUSE_HOST", "LANGFUSE_PUBLIC_KEY", "LANGFUSE_SECRET_KEY") INHERITED_ENVIRONMENT: Final = (*LANGFUSE_ENVIRONMENT, "DATABASE_URL_READ_REPLICA") _PROXY_CONFIG: Final = TypeAdapter(dict[str, object]) _SETTINGS: Final = TypeAdapter(dict[str, object]) class _ProviderBody(BaseModel): messages: list[object] def _completion(text: str) -> Reply: return Reply( body=json.dumps( { "id": "chatcmpl-" + text, "object": "chat.completion", "created": 1, "model": "gpt-4o-mini", "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, } ).encode() ) def _projects() -> Reply: return Reply( body=json.dumps( { "data": [ { "id": "integration-project", "name": "integration", "organization": {"id": "integration-org", "name": "integration"}, "metadata": {}, } ] } ).encode() ) def _text_prompt(name: str) -> Reply: return Reply( body=json.dumps( { "type": "text", "name": name, "version": 1, "prompt": "Say {{word}}", "config": {}, "labels": ["production"], "tags": [], } ).encode() ) def _langfuse_config(tmp_path: Path, general_settings: Mapping[str, JsonValue] | None = None) -> Path: config: Final = _PROXY_CONFIG.validate_python(yaml.safe_load(STOCK_CONFIG.read_text())) settings: Final = { **_SETTINGS.validate_python(config["litellm_settings"]), "success_callback": ["langfuse"], "failure_callback": ["langfuse"], } general: Final = { **_SETTINGS.validate_python(config["general_settings"]), **(general_settings or {}), } name: Final = "langfuse.yaml" if general_settings is None else "langfuse-merged.yaml" path: Final = tmp_path / name path.write_text(yaml.safe_dump({**config, "litellm_settings": settings, "general_settings": general})) return path def _langfuse_environment(langfuse: Wire) -> dict[str, str]: return { "LANGFUSE_HOST": langfuse.url, "LANGFUSE_PUBLIC_KEY": PUBLIC_KEY, "LANGFUSE_SECRET_KEY": SECRET_KEY, "LANGFUSE_FLUSH_INTERVAL": "1", } def _config_rows(database_url: str) -> list[dict[str, JsonValue]]: return read_rows( 'SELECT param_name, param_value FROM "LiteLLM_Config" WHERE param_name IN (%s, %s) ORDER BY param_name', CONFIG_SECTIONS, database_url=database_url, ) def _attribute(entries: Sequence[KeyValue], key: str) -> str | list[str] | None: for entry in entries: if entry.key != key: continue if entry.value.HasField("array_value"): return [item.string_value for item in entry.value.array_value.values] return entry.value.string_value return None def _spans(batches: Sequence[Request]) -> tuple[Span, ...]: return tuple( span for batch in batches if batch.target == TRACES_PATH and batch.headers.get("content-type") == "application/x-protobuf" for resource_spans in ExportTraceServiceRequest.FromString(batch.body).resource_spans for scope_spans in resource_spans.scope_spans for span in scope_spans.spans ) def test_langfuse_callback_delivers_the_generation_over_otlp_v4_with_the_caller_trace_fields( gateway: Gateway, tmp_path: Path ) -> None: marker: Final = "langfuse" + uuid.uuid4().hex trace_id: Final = uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker def upstream(request: Request) -> Reply: assert request.headers["authorization"] == f"Bearer {provider_secret}" return _completion(marker + "-answer") def langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() return Reply(body=b"", content_type="application/x-protobuf") with ( wire_server(upstream) as provider, wire_server(langfuse) as destination, owned_proxy( gateway, tmp_path, _langfuse_environment(destination), config=_langfuse_config(tmp_path) ) as candidate, candidate.scenario() as scenario, ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = candidate.request( "POST", "/v1/chat/completions", { "model": model, "messages": [{"role": "user", "content": marker + "-question"}], "metadata": { "trace_id": trace_id, "trace_name": marker + "-trace", "generation_name": marker, "trace_user_id": marker + "-user", "session_id": marker + "-session", "tags": [marker], }, "cache": {"no-cache": True}, }, ) assert response.status_code == 200, response.text received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones def exported() -> tuple[Span, ...]: received.extend(destination.drain()) return tuple(span for span in _spans(received) if span.name == marker) spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20) span: Final = spans[0] posts: Final = tuple(request for request in received if request.method == "POST") assert {request.target for request in posts} == {TRACES_PATH}, [request.target for request in received] basic: Final = "Basic " + base64.b64encode(f"{PUBLIC_KEY}:{SECRET_KEY}".encode()).decode() for request in posts: assert request.headers["authorization"] == basic assert request.headers["content-type"] == "application/x-protobuf" assert request.headers["x-langfuse-ingestion-version"] == "4" assert provider_secret.encode() not in request.body assert candidate.key.encode() not in request.body assert span.trace_id.hex() == trace_id assert span.parent_span_id == b"" attributes: Final = span.attributes assert _attribute(attributes, "langfuse.observation.type") == "generation" assert _attribute(attributes, "langfuse.trace.name") == marker + "-trace" assert _attribute(attributes, "user.id") == marker + "-user" assert _attribute(attributes, "session.id") == marker + "-session" assert marker in (_attribute(attributes, "langfuse.trace.tags") or ()) assert _attribute(attributes, "langfuse.observation.model.name") == "openai/gpt-4o-mini" assert json.loads(str(_attribute(attributes, "langfuse.observation.usage_details"))) == { "input": 11, "output": 4, "total": 15, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, } assert marker + "-question" in str(_attribute(attributes, "langfuse.observation.input")) assert marker + "-answer" in str(_attribute(attributes, "langfuse.observation.output")) assert ( _attribute(attributes, "langfuse.observation.metadata.litellm_call_id") == response.headers["x-litellm-call-id"] ) def test_langfuse_callback_stored_in_the_db_through_config_update_delivers_the_generation_over_otlp_v4( gateway: Gateway, tmp_path: Path ) -> None: marker: Final = "langfusedb" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker public_key: Final = "pk-lf-db-" + marker secret_key: Final = "sk-lf-db-" + marker stock_settings: Final = _SETTINGS.validate_python( _PROXY_CONFIG.validate_python(yaml.safe_load(STOCK_CONFIG.read_text()))["litellm_settings"] ) assert "langfuse" not in json.dumps( [stock_settings.get(key) for key in ("callbacks", "success_callback", "failure_callback")] ) def upstream(request: Request) -> Reply: assert request.headers["authorization"] == f"Bearer {provider_secret}" return _completion(marker + "-answer") def langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() return Reply(body=b"", content_type="application/x-protobuf") with ( scratch_database() as scratch_url, wire_server(upstream) as provider, wire_server(langfuse) as destination, owned_proxy( gateway, tmp_path, {"DATABASE_URL": scratch_url, "LANGFUSE_FLUSH_INTERVAL": "1"}, remove_environment=INHERITED_ENVIRONMENT, ) as candidate, candidate.scenario() as scenario, ): candidate.post( "/config/update", { "litellm_settings": {"success_callback": ["langfuse"]}, "environment_variables": { "LANGFUSE_HOST": destination.url, "LANGFUSE_PUBLIC_KEY": public_key, "LANGFUSE_SECRET_KEY": secret_key, }, }, ) model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) body: Final = candidate.post( "/v1/chat/completions", { "model": model, "messages": [{"role": "user", "content": marker + "-question"}], "metadata": {"generation_name": marker}, "cache": {"no-cache": True}, }, ) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones def exported() -> tuple[Span, ...]: received.extend(destination.drain()) return tuple(span for span in _spans(received) if span.name == marker) spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20) posts: Final = tuple(request for request in received if request.method == "POST") assert {request.target for request in posts} == {TRACES_PATH}, [request.target for request in received] basic: Final = "Basic " + base64.b64encode(f"{public_key}:{secret_key}".encode()).decode() for request in posts: assert request.headers["authorization"] == basic assert request.headers["content-type"] == "application/x-protobuf" assert request.headers["x-langfuse-ingestion-version"] == "4" assert provider_secret.encode() not in request.body assert candidate.key.encode() not in request.body attributes: Final = spans[0].attributes assert _attribute(attributes, "langfuse.observation.type") == "generation" assert _attribute(attributes, "langfuse.observation.metadata.response_id") == string_value(body["id"]) assert marker + "-question" in str(_attribute(attributes, "langfuse.observation.input")) assert marker + "-answer" in str(_attribute(attributes, "langfuse.observation.output")) stored: Final = {string_value(row["param_name"]): row["param_value"] for row in _config_rows(scratch_url)} callbacks: Final = TypeAdapter(list[str]).validate_python( object_value(stored["litellm_settings"]).get("success_callback") or [] ) assert "langfuse" in callbacks, stored assert set(object_value(stored["environment_variables"])) >= set(LANGFUSE_ENVIRONMENT), stored assert secret_key not in json.dumps(stored["environment_variables"]), stored def test_prompt_fetch_encodes_the_name_retries_a_5xx_once_and_keeps_langfuse_headers_off_the_client( gateway: Gateway, tmp_path: Path ) -> None: marker: Final = "prompt" + uuid.uuid4().hex leak: Final = "leak-" + marker flaky_prompt: Final = f"{marker}/what?" encoded_flaky_prompt: Final = f"{marker}%2Fwhat%3F" missing_prompt: Final = marker + "-missing" steady_prompt: Final = marker + "-steady" retry_after: Final = 30 retried_within_control_factor: Final = 3 seen_prompt_gets: Final[list[str]] = [] # mutable-ok: the double counts attempts across requests def upstream(request: Request) -> Reply: return _completion(marker + "-answer") def langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() if request.method == "POST": return Reply(body=b"", content_type="application/x-protobuf") assert request.target.startswith(PROMPTS_PATH), request.target assert request.headers["authorization"].startswith("Basic ") if request.target.startswith(PROMPTS_PATH + encoded_flaky_prompt): prior: Final = sum(1 for seen in seen_prompt_gets if seen.startswith(PROMPTS_PATH + encoded_flaky_prompt)) seen_prompt_gets.append(request.target) if prior == 0: return Reply(status=503, body=b'{"message":"try later"}', headers={"retry-after": str(retry_after)}) return _text_prompt(flaky_prompt) if request.target.startswith(PROMPTS_PATH + steady_prompt): return _text_prompt(steady_prompt) seen_prompt_gets.append(request.target) return Reply( status=404, body=b'{"message":"Prompt not found","error":"LangfuseNotFoundError"}', headers={"set-cookie": f"session={leak}; Path=/", "x-upstream-internal": leak, "server": leak}, ) with ( wire_server(upstream) as provider, wire_server(langfuse) as destination, owned_proxy( gateway, tmp_path, _langfuse_environment(destination), config=_langfuse_config(tmp_path) ) as candidate, candidate.scenario() as scenario, ): flaky: Final = scenario.model( model="langfuse/gpt-4o-mini", prompt_id=flaky_prompt, api_base=provider.url + "/v1", api_key="synthetic" ) missing: Final = scenario.model( model="langfuse/gpt-4o-mini", prompt_id=missing_prompt, api_base=provider.url + "/v1", api_key="synthetic" ) steady: Final = scenario.model( model="langfuse/gpt-4o-mini", prompt_id=steady_prompt, api_base=provider.url + "/v1", api_key="synthetic" ) def cold_prompt_chat(model: str) -> httpx.Response: return candidate.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}], "prompt_variables": {"word": marker}}, ) control: Final = timed(lambda: cold_prompt_chat(steady)) assert control.value.status_code == 200, control.value.text call: Final = timed(lambda: cold_prompt_chat(flaky)) response: Final = call.value assert response.status_code == 200, response.text assert_cut_short(call.seconds, injected_delay=retry_after) assert_within_factor_of_control(call.seconds, control=control.seconds, factor=retried_within_control_factor) attempts: Final = tuple( target for target in seen_prompt_gets if target.startswith(PROMPTS_PATH + encoded_flaky_prompt) ) assert len(attempts) == 2, seen_prompt_gets assert all(target.split("?", 1)[0] == PROMPTS_PATH + encoded_flaky_prompt for target in attempts), attempts sent: Final = _ProviderBody.model_validate_json(provider.drain()[-1].body).messages assert any("Say " + marker in json.dumps(message) for message in sent), sent failure: Final = candidate.request( "POST", "/v1/chat/completions", {"model": missing, "messages": [{"role": "user", "content": marker}], "prompt_variables": {"word": marker}}, ) assert failure.status_code == 404, failure.text assert "Prompt not found" in failure.text assert leak not in failure.text assert leak not in json.dumps(dict(failure.headers)) assert "set-cookie" not in failure.headers and "x-upstream-internal" not in failure.headers assert sum(1 for target in seen_prompt_gets if target.startswith(PROMPTS_PATH + missing_prompt)) == 1 def _responses_result(identity: str) -> Reply: return Reply( body=json.dumps( { "id": identity, "object": "response", "created_at": 1, "status": "completed", "model": "gpt-4o-mini", "output": [ { "id": "msg_" + identity, "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "integration answer", "annotations": []}], } ], "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}, } ).encode() ) def _trace_body(kind: str, model: str, marker: str, metadata: Mapping[str, JsonValue] | None) -> dict[str, JsonValue]: metadata_field: Final[dict[str, JsonValue]] = {} if metadata is None else {"metadata": dict(metadata)} if kind == "responses": return {"model": model, "input": marker + "-question", **metadata_field} if kind == "messages": return { "model": model, "max_tokens": 16, "messages": [{"role": "user", "content": marker + "-question"}], **metadata_field, } return { "model": model, "messages": [{"role": "user", "content": marker + "-question"}], "cache": {"no-cache": True}, **metadata_field, } def _w3c_headers(header_trace: str, baggage_session: str | None) -> dict[str, str]: baggage_field: Final = {} if baggage_session is None else {"baggage": f"session.id={baggage_session}"} return {"traceparent": f"00-{header_trace}-00f067aa0ba902b7-01", **baggage_field} def _await_span(received: list[Request], destination: Wire, call_id: str) -> Span: def exported() -> tuple[Span, ...]: received.extend(destination.drain()) return tuple( span for span in _spans(received) if _attribute(span.attributes, "langfuse.observation.metadata.litellm_call_id") == call_id ) spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=20) return spans[0] @pytest.mark.parametrize( ("endpoint", "kind", "metadata_mode", "expected_trace", "expected_session", "expected_target"), ( pytest.param( "/v1/chat/completions", "chat", "both", "caller", "caller", "/v1/chat/completions", id="chat_caller_ids" ), pytest.param( "/v1/responses", "responses", "both", "caller", "caller", "/v1/responses", id="responses_caller_ids" ), pytest.param( "/v1/messages", "messages", "both", "caller", "caller", "/v1/responses", id="messages_caller_ids" ), pytest.param( "/v1/chat/completions", "chat", "none", "header", "baggage", "/v1/chat/completions", id="chat_header_ids" ), pytest.param( "/v1/chat/completions", "chat", "trace", "caller", "baggage", "/v1/chat/completions", id="chat_caller_trace_header_session", ), pytest.param( "/v1/chat/completions", "chat", "empty_session", "header", "baggage", "/v1/chat/completions", id="chat_empty_session_header_ids", ), ), ) def test_langfuse_trace_and_session_prefer_caller_metadata_over_w3c_headers( gateway: Gateway, tmp_path: Path, endpoint: str, kind: str, metadata_mode: str, expected_trace: str, expected_session: str, expected_target: str, ) -> None: marker: Final = "w3c" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}" metadata: Final = { "both": {"trace_id": caller_trace, "session_id": caller_session}, "trace": {"trace_id": caller_trace}, "empty_session": {"session_id": ""}, "none": None, }[metadata_mode] expected_trace_value: Final = {"caller": caller_trace, "header": header_trace}[expected_trace] expected_session_value: Final = {"caller": caller_session, "baggage": baggage_session}[expected_session] upstream_targets: Final[list[str]] = [] # mutable-ok: records which upstream endpoint each call hit def upstream(request: Request) -> Reply: if not request.body: return Reply(status=404) assert request.headers["authorization"] == f"Bearer {provider_secret}" upstream_targets.append(request.target) if request.target == "/v1/responses": return _responses_result("resp-" + marker) assert request.target == "/v1/chat/completions", request.target return _completion(marker + "-answer") def langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() return Reply(body=b"", content_type="application/x-protobuf") with ( wire_server(upstream) as provider, wire_server(langfuse) as destination, owned_proxy( gateway, tmp_path, _langfuse_environment(destination), config=_langfuse_config(tmp_path) ) as candidate, candidate.scenario() as scenario, ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = candidate.request( "POST", endpoint, _trace_body(kind, model, marker, metadata), headers=_w3c_headers(header_trace, baggage_session), ) assert response.status_code == 200, response.text assert upstream_targets == [expected_target], upstream_targets received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones span: Final = _await_span(received, destination, response.headers["x-litellm-call-id"]) assert span.trace_id.hex() == expected_trace_value assert _attribute(span.attributes, "session.id") == expected_session_value def test_missing_session_id_reject_accepts_caller_metadata_and_baggage_fallback( gateway: Gateway, tmp_path: Path ) -> None: marker: Final = "reject" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker upstream_targets: Final[list[str]] = [] # mutable-ok: records which upstream endpoint each call hit def upstream(request: Request) -> Reply: if not request.body: return Reply(status=404) assert request.headers["authorization"] == f"Bearer {provider_secret}" upstream_targets.append(request.target) if request.target == "/v1/responses": return _responses_result("resp-" + marker) assert request.target == "/v1/chat/completions", request.target return _completion(marker + "-answer") def langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() return Reply(body=b"", content_type="application/x-protobuf") with ( wire_server(upstream) as provider, wire_server(langfuse) as destination, owned_proxy( gateway, tmp_path, _langfuse_environment(destination), config=_langfuse_config(tmp_path, {"missing_session_id": "reject"}), ) as candidate, candidate.scenario() as scenario, ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones caller_session: Final = f"my-session-id-{marker}-r1" header_trace: Final = uuid.uuid4().hex first: Final = candidate.request( "POST", "/v1/responses", {"model": model, "input": marker + "-r1", "metadata": {"session_id": caller_session}}, headers=_w3c_headers(header_trace, None), ) assert first.status_code == 200, first.text first_span: Final = _await_span(received, destination, first.headers["x-litellm-call-id"]) assert _attribute(first_span.attributes, "session.id") == caller_session baggage_session: Final = "baggage-" + marker + "-r2" second: Final = candidate.request( "POST", "/v1/chat/completions", { "model": model, "messages": [{"role": "user", "content": marker + "-r2"}], "metadata": {"session_id": ""}, }, headers=_w3c_headers(uuid.uuid4().hex, baggage_session), ) assert second.status_code == 200, second.text second_span: Final = _await_span(received, destination, second.headers["x-litellm-call-id"]) assert _attribute(second_span.attributes, "session.id") == baggage_session third: Final = candidate.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker + "-r3"}]}, headers=_w3c_headers(uuid.uuid4().hex, None), ) assert third.status_code == 400, third.text assert upstream_targets == ["/v1/responses", "/v1/chat/completions"], upstream_targets @pytest.mark.parametrize( ("endpoint", "kind", "expected_target"), ( pytest.param("/v1/responses", "responses", "/v1/responses", id="responses_caller_trace"), pytest.param("/v1/messages", "messages", "/v1/responses", id="messages_caller_trace"), ), ) def test_missing_session_id_generate_derives_session_from_caller_trace( gateway: Gateway, tmp_path: Path, endpoint: str, kind: str, expected_target: str ) -> None: marker: Final = "gen" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker caller_trace: Final = uuid.uuid4().hex upstream_targets: Final[list[str]] = [] # mutable-ok: records which upstream endpoint each call hit def upstream(request: Request) -> Reply: if not request.body: return Reply(status=404) assert request.headers["authorization"] == f"Bearer {provider_secret}" upstream_targets.append(request.target) if request.target == "/v1/responses": return _responses_result("resp-" + marker) assert request.target == "/v1/chat/completions", request.target return _completion(marker + "-answer") def langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() return Reply(body=b"", content_type="application/x-protobuf") with ( wire_server(upstream) as provider, wire_server(langfuse) as destination, owned_proxy( gateway, tmp_path, _langfuse_environment(destination), config=_langfuse_config(tmp_path, {"missing_session_id": "generate"}), ) as candidate, candidate.scenario() as scenario, ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = candidate.request( "POST", endpoint, _trace_body(kind, model, marker, {"trace_id": caller_trace}), headers=_w3c_headers(uuid.uuid4().hex, None), ) assert response.status_code == 200, response.text assert upstream_targets == [expected_target], upstream_targets received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones span: Final = _await_span(received, destination, response.headers["x-litellm-call-id"]) assert _attribute(span.attributes, "session.id") == caller_trace _AUDIT_ENDPOINTS: Final = ( pytest.param("/v1/chat/completions", "chat", "/v1/chat/completions", id="chat"), pytest.param("/v1/responses", "responses", "/v1/responses", id="responses"), pytest.param("/v1/messages", "messages", "/v1/responses", id="messages"), ) def _audit_upstream(provider_secret: str, marker: str, targets: list[str]): def upstream(request: Request) -> Reply: if not request.body: return Reply(status=404) assert request.headers["authorization"] == f"Bearer {provider_secret}" targets.append(request.target) body: Final = TypeAdapter(dict[str, JsonValue]).validate_json(request.body) reply_id: Final = uuid.uuid4().hex if request.target == "/v1/responses": if body.get("stream"): return Reply( content_type="text/event-stream", chunks=_responses_stream_frames(f"resp-{marker}-{reply_id}") ) return _responses_result(f"resp-{marker}-{reply_id}") assert request.target == "/v1/chat/completions", request.target if body.get("stream"): return Reply(content_type="text/event-stream", chunks=_chat_stream_frames(f"chatcmpl-{marker}-{reply_id}")) return _completion(f"{marker}-{reply_id}-answer") return upstream def _audit_sink(): def langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() return Reply(body=b"", content_type="application/x-protobuf") return langfuse @dataclass(frozen=True, slots=True) class _AuditRig: candidate: Gateway scenario: Scenario destination: Wire @pytest.fixture(scope="module") def audit_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_AuditRig]: directory: Final = tmp_path_factory.mktemp("audit") with ( gateway_from_environment() as gateway, wire_server(_audit_sink()) as destination, owned_proxy( gateway, directory, _langfuse_environment(destination), config=_langfuse_config(directory) ) as candidate, candidate.scenario() as scenario, ): yield _AuditRig(candidate, scenario, destination) @pytest.fixture(scope="module") def audit_generate_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_AuditRig]: directory: Final = tmp_path_factory.mktemp("audit-generate") with ( gateway_from_environment() as gateway, wire_server(_audit_sink()) as destination, owned_proxy( gateway, directory, _langfuse_environment(destination), config=_langfuse_config(directory, {"missing_session_id": "generate"}), ) as candidate, candidate.scenario() as scenario, ): yield _AuditRig(candidate, scenario, destination) @pytest.fixture(scope="module") def audit_reject_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_AuditRig]: directory: Final = tmp_path_factory.mktemp("audit-reject") with ( gateway_from_environment() as gateway, wire_server(_audit_sink()) as destination, owned_proxy( gateway, directory, _langfuse_environment(destination), config=_langfuse_config(directory, {"missing_session_id": "reject"}), ) as candidate, candidate.scenario() as scenario, ): yield _AuditRig(candidate, scenario, destination) @pytest.fixture(scope="module") def audit_omit_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_AuditRig]: directory: Final = tmp_path_factory.mktemp("audit-omit") with ( gateway_from_environment() as gateway, wire_server(_audit_sink()) as destination, owned_proxy( gateway, directory, _langfuse_environment(destination), config=_langfuse_config(directory, {"missing_session_id": "omit"}), ) as candidate, candidate.scenario() as scenario, ): yield _AuditRig(candidate, scenario, destination) @pytest.fixture(scope="module") def audit_otel_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_AuditRig]: directory: Final = tmp_path_factory.mktemp("audit-otel") def otlp(request: Request) -> Reply: return Reply() with ( gateway_from_environment() as gateway, wire_server(otlp) as destination, owned_proxy( gateway, directory, {"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"}, config=_otel_config(directory, destination.url), ) as candidate, candidate.scenario() as scenario, ): yield _AuditRig(candidate, scenario, destination) def _await_spend_row(call_id: str) -> dict[str, JsonValue]: rows: Final = eventually( lambda: read_rows( 'SELECT session_id, status, metadata, cache_hit FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,), ), lambda values: len(values) == 1, seconds=240, ) return rows[0] def _assert_call( response: httpx.Response, received: list[Request], destination: Wire, targets: list[str], expected_target: str, expected_trace: str | None, expected_session: str | None, ) -> None: assert targets == [expected_target], targets _assert_span_spend(response, received, destination, expected_trace, expected_session) def _assert_span_spend( response: httpx.Response, received: list[Request], destination: Wire, expected_trace: str | None, expected_session: str | None, ) -> None: call_id: Final = response.headers["x-litellm-call-id"] span: Final = _await_span(received, destination, call_id) if expected_trace is not None: assert span.trace_id.hex() == expected_trace, f"call {call_id}: trace id" assert _attribute(span.attributes, "session.id") == expected_session, f"call {call_id}: session.id" row: Final = _await_spend_row(call_id) assert row["session_id"] == expected_session, f"call {call_id}: spend session {row}" @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) @pytest.mark.parametrize("metadata_mode", ("both", "none", "trace", "session")) def test_audit_caller_metadata_wins_over_w3c_per_field( audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str, metadata_mode: str ) -> None: marker: Final = "audit" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}" metadata: Final = { "both": {"trace_id": caller_trace, "session_id": caller_session}, "trace": {"trace_id": caller_trace}, "session": {"session_id": caller_session}, "none": None, }[metadata_mode] expected_trace: Final = caller_trace if metadata_mode in ("both", "trace") else header_trace expected_session: Final = caller_session if metadata_mode in ("both", "session") else baggage_session targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, metadata), headers=_w3c_headers(header_trace, baggage_session), ) assert response.status_code == 200, response.text received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones _assert_call( response, received, audit_rig.destination, targets, expected_target, expected_trace, expected_session ) @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) def test_audit_caller_metadata_wins_on_streamed_calls( audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str ) -> None: marker: Final = "auditstream" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}" targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_rig.candidate.request( "POST", endpoint, { **_trace_body(kind, model, marker, {"trace_id": caller_trace, "session_id": caller_session}), "stream": True, }, headers=_w3c_headers(header_trace, baggage_session), ) assert response.status_code == 200, response.text received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones _assert_call(response, received, audit_rig.destination, targets, expected_target, caller_trace, caller_session) def test_audit_caller_metadata_wins_through_official_sdk_clients(audit_rig: _AuditRig) -> None: marker: Final = "auditsdk" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: candidate: Final = audit_rig.candidate destination: Final = audit_rig.destination model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) base_url: Final = str(candidate.client.base_url).rstrip("/") received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones def drive(client_kind: str) -> tuple[str, str, httpx.Headers]: caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"{client_kind}-session-{marker}" body_metadata: Final = { "trace_id": caller_trace, "session_id": caller_session, "generation_name": marker + "-" + client_kind, } headers: Final = _w3c_headers(uuid.uuid4().hex, "baggage-" + marker + "-" + client_kind) if client_kind == "chat_openai_sync": return caller_trace, caller_session, ( openai.OpenAI(base_url=f"{base_url}/v1", api_key=candidate.key) .chat.completions.with_raw_response.create( model=model, messages=[{"role": "user", "content": marker + "-chat"}], extra_body={"metadata": body_metadata, "cache": {"no-cache": True}}, extra_headers=headers, ) .headers ) if client_kind == "responses_openai_async": async def responses_call() -> httpx.Headers: answer: Final = await openai.AsyncOpenAI( base_url=f"{base_url}/v1", api_key=candidate.key ).responses.with_raw_response.create( model=model, input=marker + "-responses", extra_body={"metadata": body_metadata, "cache": {"no-cache": True}}, extra_headers=headers, ) return answer.headers return caller_trace, caller_session, asyncio.run(responses_call()) if client_kind == "messages_anthropic_sync": return caller_trace, caller_session, ( anthropic.Anthropic(base_url=base_url, api_key=candidate.key) .messages.with_raw_response.create( model=model, max_tokens=16, messages=[{"role": "user", "content": marker + "-messages"}], extra_body={"metadata": body_metadata}, extra_headers=headers, ) .headers ) async def messages_call() -> httpx.Headers: answer: Final = await anthropic.AsyncAnthropic( base_url=base_url, api_key=candidate.key ).messages.with_raw_response.create( model=model, max_tokens=16, messages=[{"role": "user", "content": marker + "-messages-async"}], extra_body={"metadata": body_metadata}, extra_headers=headers, ) return answer.headers return caller_trace, caller_session, asyncio.run(messages_call()) expected: Final = { call_headers["x-litellm-call-id"]: (caller_trace, caller_session, client_kind) for client_kind, (caller_trace, caller_session, call_headers) in ( (kind, drive(kind)) for kind in ( "chat_openai_sync", "responses_openai_async", "messages_anthropic_sync", "messages_anthropic_async", ) ) } def spans_named() -> tuple[Span, ...]: received.extend(destination.drain()) return tuple( span for span in _spans(received) if _attribute(span.attributes, "langfuse.observation.metadata.litellm_call_id") in expected ) spans: Final = eventually(spans_named, lambda values: len(values) == len(expected), seconds=60) for span in spans: call_id: Final = str( _attribute(span.attributes, "langfuse.observation.metadata.litellm_call_id") ) caller_trace, caller_session, client_kind = expected[call_id] assert span.trace_id.hex() == caller_trace, f"{client_kind}: trace id" assert _attribute(span.attributes, "session.id") == caller_session, f"{client_kind}: session.id" row: Final = _await_spend_row(call_id) assert row["session_id"] == caller_session, f"{client_kind} {call_id}: spend session {row}" def _otel_config(tmp_path: Path, sink_url: str) -> Path: config: Final = _PROXY_CONFIG.validate_python(yaml.safe_load(STOCK_CONFIG.read_text())) settings: Final = {**_SETTINGS.validate_python(config["litellm_settings"]), "callbacks": ["otel"]} path: Final = tmp_path / "otel.yaml" path.write_text( yaml.safe_dump( { **config, "litellm_settings": settings, "callback_settings": { "otel": {"exporter": "http/json", "endpoint": sink_url, "mapper_names": ["genai"]} }, } ) ) return path _OTEL_SPAN: Final = TypeAdapter(dict[str, JsonValue]) def _otel_spans(batches: Sequence[Request]) -> tuple[dict[str, JsonValue], ...]: spans: list[dict[str, JsonValue]] = [] # mutable-ok: flattens nested OTLP batches into a tuple for batch in batches: if not batch.target.endswith("/v1/traces"): continue payload: Final = TypeAdapter(JsonValue).validate_json(batch.body) envelopes: Final = payload if isinstance(payload, list) else [payload] for envelope in envelopes: for resource in TypeAdapter(list[JsonValue]).validate_python( object_value(envelope)["resourceSpans"] ): for scope in object_value(resource)["scopeSpans"]: spans.extend(TypeAdapter(list[JsonValue]).validate_python(object_value(scope)["spans"])) return tuple(_OTEL_SPAN.validate_python(span) for span in spans) def _otel_attribute(span: Mapping[str, JsonValue], key: str) -> str | None: for attribute in TypeAdapter(list[JsonValue]).validate_python(span.get("attributes") or []): entry: Final = object_value(attribute) if entry["key"] == key: value: Final = object_value(entry["value"]) raw: Final = value.get("stringValue") return str(raw) if raw is not None else None return None @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS[:2]) def test_audit_otel_span_carries_caller_ids( audit_otel_rig: _AuditRig, endpoint: str, kind: str, expected_target: str ) -> None: marker: Final = "auditotel" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}" targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: destination: Final = audit_otel_rig.destination model: Final = audit_otel_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_otel_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, {"trace_id": caller_trace, "session_id": caller_session}), headers=_w3c_headers(header_trace, baggage_session), ) assert response.status_code == 200, response.text assert targets == [expected_target], targets response_id: Final = string_value(object_value(response.json())["id"]) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones def exported() -> tuple[dict[str, JsonValue], ...]: received.extend(destination.drain()) return tuple( span for span in _otel_spans(received) if _otel_attribute(span, "gen_ai.response.id") == response_id ) spans: Final = eventually(exported, lambda values: len(values) == 1, seconds=60) span: Final = spans[0] print( f"H7 record: otel span trace={span['traceId']} header={header_trace} " f"session.id={_otel_attribute(span, 'session.id')} " f"gen_ai.conversation.id={_otel_attribute(span, 'gen_ai.conversation.id')}" ) assert span["traceId"] == header_trace, ( f"otel span for {response_id}: trace id must be the ambient W3C header trace" ) assert _otel_attribute(span, "gen_ai.conversation.id") == caller_session, ( f"otel span conversation id for {response_id}" ) @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) @pytest.mark.parametrize( ("case", "metadata_mode", "headers_mode", "expected_session"), ( pytest.param("g1", "trace", "traceparent", "caller_trace", id="g1_caller_trace_and_header"), pytest.param("g3", "none", "traceparent", "header_trace", id="g3_header_only"), pytest.param("g5", "session", "traceparent", "caller_session", id="g5_caller_session"), ), ) def test_audit_generate_policy_derives_session_from_caller_trace( audit_generate_rig: _AuditRig, endpoint: str, kind: str, expected_target: str, case: str, metadata_mode: str, headers_mode: str, expected_session: str, ) -> None: marker: Final = "auditgen" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}" metadata: Final = {"trace": {"trace_id": caller_trace}, "session": {"session_id": caller_session}, "none": None}[ metadata_mode ] expected: Final = {"caller_trace": caller_trace, "header_trace": header_trace, "caller_session": caller_session}[ expected_session ] targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_generate_rig.scenario.model( api_base=provider.url + "/v1", api_key=provider_secret ) response: Final = audit_generate_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, metadata), headers={"traceparent": f"00-{header_trace}-00f067aa0ba902b7-01"}, ) assert response.status_code == 200, response.text received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones call_id: Final = response.headers["x-litellm-call-id"] assert targets == [expected_target], targets span: Final = _await_span(received, audit_generate_rig.destination, call_id) assert _attribute(span.attributes, "session.id") == expected, f"call {call_id}: session.id" row: Final = _await_spend_row(call_id) assert row["session_id"] == expected, f"call {call_id}: spend session {row}" @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) def test_audit_generate_policy_is_stable_across_repeated_caller_trace( audit_generate_rig: _AuditRig, endpoint: str, kind: str, expected_target: str ) -> None: marker: Final = "auditgen2" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker caller_trace: Final = uuid.uuid4().hex targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_generate_rig.scenario.model( api_base=provider.url + "/v1", api_key=provider_secret ) responses: Final = tuple( audit_generate_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, f"{marker}-{attempt}", {"trace_id": caller_trace}), ) for attempt in ("first", "second") ) sessions: Final[list[str | None]] = [] # mutable-ok: collects the two observed sessions in order for attempt, response in zip(("first", "second"), responses): assert response.status_code == 200, f"{attempt}: {response.text}" row: Final = _await_spend_row(response.headers["x-litellm-call-id"]) sessions.append(string_value(row["session_id"])) assert sessions == [caller_trace, caller_trace], ( f"generate must derive both sessions from the caller trace {caller_trace}" ) @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) def test_audit_generate_policy_fresh_session_without_any_ids( audit_generate_rig: _AuditRig, endpoint: str, kind: str, expected_target: str ) -> None: marker: Final = "auditgen4" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_generate_rig.scenario.model( api_base=provider.url + "/v1", api_key=provider_secret ) sessions: Final[list[str | None]] = [] # mutable-ok: collects the two observed sessions in order for attempt in ("first", "second"): response: Final = audit_generate_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, f"{marker}-{attempt}", None) ) assert response.status_code == 200, response.text session: Final = _await_spend_row(response.headers["x-litellm-call-id"])["session_id"] assert session, f"{attempt}: generated session id must be non-empty" sessions.append(string_value(session)) assert sessions[0] != sessions[1], f"two id-less calls must not share a session: {sessions}" assert targets and set(targets) == {expected_target}, targets @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) @pytest.mark.parametrize( ("case", "metadata", "baggage", "expected_status", "expected_session"), ( pytest.param( "r1", "caller_session", None, 200, "caller_session", id="r1_caller_session_no_baggage" ), pytest.param("r2", "empty_session", "baggage", 200, "baggage", id="r2_empty_session_baggage"), pytest.param("r3", "none", None, 400, None, id="r3_nothing_rejected"), ), ) def test_audit_reject_policy( audit_reject_rig: _AuditRig, endpoint: str, kind: str, expected_target: str, case: str, metadata: str, baggage: str, expected_status: int, expected_session: str | None, ) -> None: marker: Final = "auditreject" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker baggage_session: Final = "baggage-" + marker caller_session: Final = f"my-session-id-{marker}" body_metadata: Final = { "caller_session": {"session_id": caller_session}, "empty_session": {"session_id": ""}, "none": None, }[metadata] expected: Final = {"caller_session": caller_session, "baggage": baggage_session}[expected_session] if expected_session else None targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_reject_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_reject_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, body_metadata), headers=_w3c_headers(uuid.uuid4().hex, baggage_session if baggage else None), ) assert response.status_code == expected_status, response.text if expected_status != 200: assert targets == [], targets rejected_rows: Final = eventually( lambda: read_rows( 'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (response.headers["x-litellm-call-id"],), ), lambda values: len(values) == 1, seconds=240, ) assert rejected_rows[0]["status"] == "failure", ( f"rejected call {response.headers['x-litellm-call-id']} must write exactly one failure spend row: {rejected_rows}" ) return received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones _assert_call(response, received, audit_reject_rig.destination, targets, expected_target, None, expected) @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) def test_audit_omit_policy_records_no_session(audit_omit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str) -> None: marker: Final = "auditomit" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_omit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_omit_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, None) ) assert response.status_code == 200, response.text call_id: Final = response.headers["x-litellm-call-id"] assert targets == [expected_target], targets received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones span: Final = _await_span(received, audit_omit_rig.destination, call_id) row: Final = _await_spend_row(call_id) span_session: Final = _attribute(span.attributes, "session.id") assert row["session_id"] == span_session, ( f"call {call_id}: spend session {row['session_id']!r} must match span session {span_session!r}" ) @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) @pytest.mark.parametrize("bad_value", (123, ["x"]), ids=["int", "list"]) def test_audit_non_string_caller_ids_fall_back_to_w3c( audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str, bad_value: JsonValue ) -> None: marker: Final = "auditbad" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, {"trace_id": bad_value, "session_id": bad_value}), headers=_w3c_headers(header_trace, baggage_session), ) assert response.status_code == 200, response.text received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones _assert_call(response, received, audit_rig.destination, targets, expected_target, header_trace, baggage_session) @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) @pytest.mark.parametrize( "metadata", ({"trace_id": "", "session_id": ""}, {"trace_id": None, "session_id": None}), ids=["empty", "null"], ) def test_audit_empty_and_null_caller_ids_fall_back_to_w3c( audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str, metadata: dict[str, object] ) -> None: marker: Final = "auditempty" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, dict(metadata)), headers=_w3c_headers(header_trace, baggage_session), ) assert response.status_code == 200, response.text received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones _assert_call(response, received, audit_rig.destination, targets, expected_target, header_trace, baggage_session) @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS[:2]) def test_audit_five_kilobyte_caller_ids_win_verbatim( audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str ) -> None: marker: Final = "auditbig" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker caller_trace: Final = "T" * 5120 caller_session: Final = "S" * 5120 targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, {"trace_id": caller_trace, "session_id": caller_session}), headers=_w3c_headers(header_trace, baggage_session), ) assert response.status_code == 200, response.text call_id: Final = response.headers["x-litellm-call-id"] assert targets == [expected_target], targets received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones span: Final = _await_span(received, audit_rig.destination, call_id) assert span.trace_id.hex() != header_trace, f"call {call_id}: caller trace must beat the W3C header" assert _attribute(span.attributes, "session.id") == caller_session, f"call {call_id}: session.id" assert _await_spend_row(call_id)["session_id"] == caller_session, f"call {call_id}: spend session" @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS) def test_audit_identical_requests_twice_log_per_call( audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str ) -> None: marker: Final = "auditdup" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}" targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones for attempt in ("first", "second"): response: Final = audit_rig.candidate.request( "POST", endpoint, _trace_body( kind, model, marker, {"trace_id": caller_trace, "session_id": caller_session} ), headers=_w3c_headers(header_trace, baggage_session), ) assert response.status_code == 200, f"{attempt}: {response.text}" call_id: Final = response.headers["x-litellm-call-id"] span: Final = _await_span(received, audit_rig.destination, call_id) assert span.trace_id.hex() == caller_trace, f"{attempt} {call_id}: trace id" assert _attribute(span.attributes, "session.id") == caller_session, f"{attempt} {call_id}: session.id" assert _await_spend_row(call_id)["session_id"] == caller_session, f"{attempt} {call_id}: spend" assert targets and set(targets) == {expected_target}, targets def test_audit_malformed_w3c_headers_are_ignored(audit_rig: _AuditRig) -> None: marker: Final = "auditmal" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_rig.candidate.request( "POST", "/v1/chat/completions", _trace_body("chat", model, marker, None), headers={"traceparent": "00-zz-00f067aa0ba902b7-01", "baggage": "not-a-session-key"}, ) assert response.status_code == 200, response.text assert targets == ["/v1/chat/completions"], targets received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones span: Final = _await_span(received, audit_rig.destination, response.headers["x-litellm-call-id"]) assert len(span.trace_id.hex()) == 32 and "zz" not in span.trace_id.hex(), span.trace_id.hex() row: Final = _await_spend_row(response.headers["x-litellm-call-id"]) span_session: Final = _attribute(span.attributes, "session.id") assert span_session in (None, row["session_id"]), ( f"span session {span_session!r} diverges from spend session {row['session_id']!r}" ) assert row["session_id"], row def test_audit_unauthenticated_call_leaves_no_spend_row(audit_rig: _AuditRig) -> None: marker: Final = "auditunauth" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: candidate: Final = audit_rig.candidate model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) denied: Final = candidate.request( "POST", "/v1/chat/completions", _trace_body( "chat", model, marker, {"trace_id": uuid.uuid4().hex, "session_id": f"my-session-id-{marker}"}, ), headers=_w3c_headers(uuid.uuid4().hex, "baggage-" + marker), key="sk-wrong-key", ) assert denied.status_code == 401, denied.text assert targets == [], targets control: Final = candidate.request( "POST", "/v1/chat/completions", _trace_body("chat", model, marker + "-control", None) ) assert control.status_code == 200, control.text _await_spend_row(control.headers["x-litellm-call-id"]) assert ( read_rows( 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (denied.headers.get("x-litellm-call-id") or "",), ) == [] ), "an unauthenticated call must not write a spend row" @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS[:2]) @pytest.mark.parametrize("upstream_status", (500, 401), ids=["upstream_500", "upstream_401"]) def test_audit_upstream_error_still_logs_caller_session( audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str, upstream_status: int, ) -> None: marker: Final = "auditerr" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}" targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit def upstream(request: Request) -> Reply: if request.body: targets.append(request.target) return Reply(status=upstream_status, body=b'{"error": {"message": "scripted upstream failure"}}') with wire_server(upstream) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) response: Final = audit_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, {"trace_id": caller_trace, "session_id": caller_session}), headers=_w3c_headers(uuid.uuid4().hex, "baggage-" + marker), ) assert response.status_code == upstream_status, response.text assert targets and set(targets) == {expected_target}, targets call_id: Final = response.headers["x-litellm-call-id"] row: Final = _await_spend_row(call_id) assert row["session_id"] == caller_session, f"call {call_id}: failure spend session {row}" def test_audit_sink_rejection_does_not_break_the_caller(gateway: Gateway, tmp_path: Path) -> None: marker: Final = "auditreject" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit attempts: Final[list[int]] = [] # mutable-ok: sink status sequence counter def langfuse(request: Request) -> Reply: if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() if request.target == TRACES_PATH: attempts.append(1) if len(attempts) == 1: return Reply(status=403) if len(attempts) == 2: return Reply(status=404) return Reply(body=b"", content_type="application/x-protobuf") with ( wire_server(_audit_upstream(provider_secret, marker, targets)) as provider, wire_server(langfuse) as destination, owned_proxy( gateway, tmp_path, _langfuse_environment(destination), config=_langfuse_config(tmp_path) ) as candidate, candidate.scenario() as scenario, ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones for attempt in ("first", "second", "third"): response: Final = candidate.request( "POST", "/v1/chat/completions", _trace_body( "chat", model, f"{marker}-{attempt}", {"trace_id": uuid.uuid4().hex, "session_id": f"my-session-id-{marker}-{attempt}"}, ), ) assert response.status_code == 200, f"{attempt}: {response.text}" _await_spend_row(response.headers["x-litellm-call-id"]) assert targets == ["/v1/chat/completions"] * 3, targets @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS[:2]) def test_audit_string_metadata_body_does_not_crash(audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str) -> None: marker: Final = "auditstr" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: candidate: Final = audit_rig.candidate model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) body: Final = {**_trace_body(kind, model, marker, None), "metadata": "x"} response: Final = candidate.request( "POST", endpoint, body, headers=_w3c_headers(uuid.uuid4().hex, "baggage-" + marker) ) assert response.status_code < 500, f"metadata string must not crash the proxy: {response.status_code} {response.text}" follow_up: Final = candidate.request( "POST", "/v1/chat/completions", _trace_body("chat", model, marker + "-follow", None) ) assert follow_up.status_code == 200, follow_up.text @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS[:2]) def test_audit_key_metadata_session_still_honoured(audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str) -> None: marker: Final = "auditkey" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker key_session: Final = f"key-session-{marker}" targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) token: Final = audit_rig.scenario.key(metadata={"session_id": key_session}) response: Final = audit_rig.candidate.request( "POST", endpoint, _trace_body(kind, model, marker, None), headers=_w3c_headers(uuid.uuid4().hex, "baggage-" + marker), key=token, ) assert response.status_code == 200, response.text call_id: Final = response.headers["x-litellm-call-id"] assert targets == [expected_target], targets received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones span: Final = _await_span(received, audit_rig.destination, call_id) row: Final = _await_spend_row(call_id) assert _attribute(span.attributes, "session.id") == row["session_id"], ( f"call {call_id}: spend session {row['session_id']!r} must match span session" ) @pytest.mark.parametrize(("endpoint", "kind", "expected_target"), _AUDIT_ENDPOINTS[:2]) @pytest.mark.parametrize("metadata_mode", ("both", "none"), ids=["caller_ids", "no_metadata"]) def test_audit_cache_hit_call_keeps_winning_ids( audit_rig: _AuditRig, endpoint: str, kind: str, expected_target: str, metadata_mode: str ) -> None: marker: Final = "auditcache" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}" metadata: Final = ( {"trace_id": caller_trace, "session_id": caller_session} if metadata_mode == "both" else None ) expected_trace: Final = caller_trace if metadata_mode == "both" else header_trace expected_session: Final = caller_session if metadata_mode == "both" else baggage_session targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: candidate: Final = audit_rig.candidate model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones body: Final = {key: value for key, value in _trace_body(kind, model, marker, metadata).items() if key != "cache"} headers: Final = _w3c_headers(header_trace, baggage_session) first: Final = candidate.request("POST", endpoint, body, headers=headers) assert first.status_code == 200, first.text second: Final = candidate.request("POST", endpoint, body, headers=headers) assert second.status_code == 200, second.text assert second.headers.get("x-litellm-cache-key"), ( f"second identical call must be a cache hit: {dict(second.headers)}" ) call_id: Final = second.headers["x-litellm-call-id"] assert targets == [expected_target], targets span: Final = _await_span(received, audit_rig.destination, call_id) if metadata_mode == "both": assert span.trace_id.hex() == expected_trace, f"cache hit {call_id}: trace id" assert _attribute(span.attributes, "session.id") == expected_session, ( f"cache hit {call_id}: session.id" ) else: span_session: Final = _attribute(span.attributes, "session.id") print( f"E2 record: cache-hit span trace={span.trace_id.hex()} session={span_session!r} " f"header trace={header_trace} baggage session={baggage_session!r}" ) assert len(span.trace_id.hex()) == 32, span.trace_id.hex() assert _await_spend_row(call_id)["session_id"] == expected_session, f"cache hit {call_id}: spend" def test_audit_concurrent_requests_each_keep_their_caller_ids(audit_rig: _AuditRig) -> None: marker: Final = "auditconc" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with wire_server(_audit_upstream(provider_secret, marker, targets)) as provider: candidate: Final = audit_rig.candidate model: Final = audit_rig.scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) jobs: Final = tuple( (index, endpoint, kind, uuid.uuid4().hex, f"my-session-id-{marker}-{index}") for index, (endpoint, kind, _) in tuple( enumerate(tuple(row.values for row in _AUDIT_ENDPOINTS) * 4) )[:10] ) def fire(job: tuple[int, str, str, str, str]) -> tuple[str, str, httpx.Response]: index, endpoint, kind, caller_trace, caller_session = job response: Final = candidate.request( "POST", endpoint, _trace_body( kind, model, f"{marker}-{index}", {"trace_id": caller_trace, "session_id": caller_session} ), headers=_w3c_headers(header_trace, baggage_session), ) return caller_trace, caller_session, response with ThreadPoolExecutor(max_workers=10) as pool: answered: Final = tuple(pool.map(fire, jobs)) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones assert len(targets) == 10, targets for caller_trace, caller_session, response in answered: assert response.status_code == 200, response.text _assert_span_spend(response, received, audit_rig.destination, caller_trace, caller_session) def test_audit_sink_outage_mid_burst_loses_no_spend_rows(gateway: Gateway, tmp_path: Path) -> None: marker: Final = "auditoutage" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit outage: Final = threading.Event() rejected_export: Final = threading.Event() accepted: Final[SimpleQueue[Request]] = SimpleQueue() def langfuse(request: Request) -> Reply: if outage.is_set(): if request.method == "POST" and request.target == TRACES_PATH: rejected_export.set() return Reply(status=503) if request.method == "GET" and request.target.startswith(PROJECTS_PATH): return _projects() accepted.put(request) return Reply(body=b"", content_type="application/x-protobuf") with ( wire_server(_audit_upstream(provider_secret, marker, targets)) as provider, wire_server(langfuse) as destination, owned_proxy( gateway, tmp_path, _langfuse_environment(destination), config=_langfuse_config(tmp_path) ) as candidate, candidate.scenario() as scenario, ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) jobs: Final = tuple( (index, endpoint, kind, uuid.uuid4().hex, f"my-session-id-{marker}-{index}") for index, (endpoint, kind, _) in enumerate(tuple(row.values for row in _AUDIT_ENDPOINTS) * 10) ) def fire(job: tuple[int, str, str, str, str]) -> tuple[str, str, httpx.Response]: index, endpoint, kind, caller_trace, caller_session = job stream: Final = index % 3 == 0 response: Final = candidate.request( "POST", endpoint, { **_trace_body( kind, model, f"{marker}-{index}", {"trace_id": caller_trace, "session_id": caller_session}, ), "stream": stream, }, headers=_w3c_headers(header_trace, baggage_session), ) return caller_trace, caller_session, response outage.set() with ThreadPoolExecutor(max_workers=30) as pool: first_wave: Final = tuple(pool.map(fire, jobs[:10])) assert rejected_export.wait(timeout=10), "the exporter did not reach the scripted outage" unhealthy: Final = candidate.request("GET", "/health/services?service=langfuse") assert unhealthy.status_code != 200 or "unhealthy" in unhealthy.text, ( f"langfuse must report unhealthy while the sink 503s: {unhealthy.status_code} {unhealthy.text}" ) outage.clear() healthy: Final = candidate.request("GET", "/health/services?service=langfuse") assert healthy.status_code == 200, healthy.text with ThreadPoolExecutor(max_workers=30) as pool: answered: Final = first_wave + tuple(pool.map(fire, jobs[10:])) for _, _, response in answered: assert response.status_code == 200, response.text expected_by_call: Final = { response.headers["x-litellm-call-id"]: caller_session for caller_session, response in ((session, res) for _, session, res in answered) } call_ids: Final = sorted(expected_by_call) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones delivered: Final = eventually( lambda: ( received.extend(accepted.get_nowait() for _ in range(accepted.qsize())) or tuple( span for span in _spans(received) if _attribute(span.attributes, "langfuse.observation.metadata.litellm_call_id") in expected_by_call ) ), lambda spans: len(spans) >= len(answered), seconds=60, ) spans_per_call: Final = { call_id: sum( 1 for span in delivered if _attribute(span.attributes, "langfuse.observation.metadata.litellm_call_id") == call_id ) for call_id in call_ids } assert sorted(spans_per_call.values()) == [1] * len(answered), ( f"each call id must arrive on exactly one span: {spans_per_call}" ) for span in delivered: span_call: Final = str( _attribute(span.attributes, "langfuse.observation.metadata.litellm_call_id") ) assert _attribute(span.attributes, "session.id") == expected_by_call[span_call], ( f"call {span_call}: session.id" ) spend_rows: Final = eventually( lambda: read_rows( 'SELECT session_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = ANY(%s)', (call_ids,), ), lambda values: len(values) == len(answered), seconds=150, ) for row in spend_rows: row_call: Final = string_value(row["litellm_call_id"]) assert row["session_id"] == expected_by_call[row_call], f"call {row_call}: spend {row}" def test_audit_surviving_worker_keeps_serving_after_kill(gateway: Gateway, tmp_path: Path) -> None: marker: Final = "auditworker" + uuid.uuid4().hex provider_secret: Final = "synthetic-provider-secret-" + marker header_trace: Final = uuid.uuid4().hex baggage_session: Final = "baggage-" + marker targets: Final[list[str]] = [] # mutable-ok: the upstream records every hit with ( wire_server(_audit_upstream(provider_secret, marker, targets)) as provider, wire_server(_audit_sink()) as destination, owned_proxy_process( gateway, tmp_path, _langfuse_environment(destination), config=_langfuse_config(tmp_path), workers=2, ) as owned, owned.gateway.scenario() as scenario, ): candidate: Final = owned.gateway model: Final = scenario.model(api_base=provider.url + "/v1", api_key=provider_secret) workers: Final = owned.worker_pids() assert workers, "expected worker processes under the owned proxy" def fire(index: int) -> tuple[str, httpx.Response]: caller_trace: Final = uuid.uuid4().hex caller_session: Final = f"my-session-id-{marker}-{index}" response: Final = candidate.request( "POST", "/v1/chat/completions", _trace_body( "chat", model, f"{marker}-{index}", {"trace_id": caller_trace, "session_id": caller_session}, ), headers=_w3c_headers(header_trace, baggage_session), ) return caller_session, response with owned.replacing_worker(min(workers)) as kill, ThreadPoolExecutor(max_workers=10) as pool: kill() answered: Final = tuple(pool.map(fire, range(10))) received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier ones for caller_session, response in answered: assert response.status_code == 200, response.text call_id: Final = response.headers["x-litellm-call-id"] span: Final = _await_span(received, destination, call_id) assert _attribute(span.attributes, "session.id") == caller_session, f"{call_id}: session.id" expected_by_call: Final = { response.headers["x-litellm-call-id"]: caller_session for caller_session, response in answered } spend_rows: Final = eventually( lambda: read_rows( 'SELECT session_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = ANY(%s)', (list(expected_by_call),), ), lambda values: len(values) == len(answered), seconds=150, ) for row in spend_rows: row_call: Final = string_value(row["litellm_call_id"]) assert row["session_id"] == expected_by_call[row_call], f"call {row_call}: spend {row}"