diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 30e0901c32a..018ff758d5d 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -5,6 +5,7 @@ from collections.abc import Callable from datetime import datetime, timedelta from functools import cache from typing import Final +from urllib.parse import unquote from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -31,13 +32,19 @@ from litellm.types.utils import StandardLoggingPayload AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default" _ADLS_SAFE_NAME: Final = str.maketrans("/", "_", "=") +_DOT_OR_EMPTY_SEGMENTS: Final = frozenset(("", ".", "..")) def adls_safe_file_name(payload_id: str | None) -> str: - """`=` padding and `/` in a base64 payload id are what the Data Lake service rejects, so the name drops the - padding and maps `/` to `_`. Standard base64 has no `_` and its padding is fixed by the length, so ids from - that alphabet stay distinct; anything else is left as is.""" - return f"{(payload_id or str(uuid.uuid4())).translate(_ADLS_SAFE_NAME)}.json" + """A Responses API id is base64 behind `resp_`, and the Data Lake service rejects its `=` padding and `/`, so + that name drops the padding and maps `/` to `_`. Standard base64 has no `_` and its padding is fixed by the + length, so those ids stay distinct. Every other id, including a caller's `x-litellm-call-id`, is used as is + unless it has an empty, `.` or `..` path segment, which gets the same rewrite so the file keeps its own name in + the log directory""" + name: Final = payload_id or str(uuid.uuid4()) + if not name.startswith("resp_") and _DOT_OR_EMPTY_SEGMENTS.isdisjoint(unquote(name).split("/")): + return f"{name}.json" + return f"{name.translate(_ADLS_SAFE_NAME)}.json" @cache diff --git a/tests/integration/observability/test_azure_storage_file_names.py b/tests/integration/observability/test_azure_storage_file_names.py index 5009dba1d53..b471f55851b 100644 --- a/tests/integration/observability/test_azure_storage_file_names.py +++ b/tests/integration/observability/test_azure_storage_file_names.py @@ -1,3 +1,4 @@ +import json import re import uuid from pathlib import Path @@ -13,9 +14,12 @@ from _s3_v2_support import surface_reply from integration._support.client import Gateway, eventually from integration._support.process import owned_proxy from integration._support.tls import server_context, write_self_signed_cert -from integration._support.wire import wire_server +from integration._support.wire import Reply, Request, wire_server + +from litellm.constants import MAX_LITELLM_CALL_ID_LENGTH ADLS_SAFE_FILE_NAME: Final = re.compile(r"^[A-Za-z0-9._+-]+\.json$") +WORKERS: Final = 2 def _responses_id(candidate: Gateway, model: str, key: str, marker: str) -> str: @@ -37,7 +41,7 @@ def test_responses_ids_with_base64_padding_land_under_adls_safe_names(gateway: G environment: Final = {**azure_storage_environment(store.url, cert), "DEFAULT_FLUSH_INTERVAL_SECONDS": "1"} config: Final = azure_storage_config(tmp_path) with ( - owned_proxy(gateway, tmp_path, environment, config=config, workers=1) as candidate, + owned_proxy(gateway, tmp_path, environment, config=config, workers=WORKERS) as candidate, candidate.scenario() as scenario, ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") @@ -59,3 +63,146 @@ def test_responses_ids_with_base64_padding_land_under_adls_safe_names(gateway: G assert all(ADLS_SAFE_FILE_NAME.match(name) for name in names), names assert len(frozenset(names)) == len(answered), names assert provider.drain() + + +def _embedding_reply(request: Request) -> Reply: + assert request.method == "POST" and request.target.endswith("/embeddings"), request.target + return Reply( + body=json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + } + ).encode() + ) + + +def _failing_chat_reply(request: Request) -> Reply: + assert request.method == "POST" and request.target.endswith("/chat/completions"), request.target + return Reply(status=500, body=json.dumps({"error": {"message": "upstream rejected the request"}}).encode()) + + +def _log_names_by_call_id( + gateway: Gateway, + tmp_path: Path, + call_ids: tuple[str, ...], + *, + inputs: tuple[str, ...] | None = None, + failing: bool = False, +) -> dict[str, str]: + sink: Final = RecordingDataLakeSink() + cert, key = write_self_signed_cert(tmp_path, SINK_HOSTS) + with ( + wire_server(_failing_chat_reply if failing else _embedding_reply) as provider, + wire_server(sink.respond, tls=server_context(cert, key), keep_alive=True) as store, + ): + environment: Final = {**azure_storage_environment(store.url, cert), "DEFAULT_FLUSH_INTERVAL_SECONDS": "1"} + config: Final = azure_storage_config(tmp_path) + with ( + owned_proxy(gateway, tmp_path, environment, config=config, workers=WORKERS) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-4.1-nano" if failing else "openai/text-embedding-3-small", + api_base=provider.url + "/v1", + api_key="synthetic-provider-key", + ) + api_key: Final = scenario.key(models=[model]) + responses: Final = tuple( + candidate.request( + "POST", + "/v1/chat/completions" if failing else "/v1/embeddings", + {"model": model, "messages": [{"role": "user", "content": text}]} + if failing + else {"model": model, "input": text}, + key=api_key, + headers={"x-litellm-call-id": call_id}, + ) + for call_id, text in zip(call_ids, inputs or call_ids, strict=True) + ) + assert all(response.status_code == (500 if failing else 200) for response in responses), tuple( + response.text for response in responses + ) + eventually( + lambda: len(sink.stored()) + len(sink.duplicated()) + len(sink.unauthenticated_targets()), + lambda settled: settled >= len(call_ids), + seconds=60, + ) + assert sink.unauthenticated_targets() == (), sink.unauthenticated_targets() + assert sink.duplicated() == (), f"a later log overwrote an earlier one at {sink.duplicated()}" + assert provider.drain() + return {path.split("/", 3)[3]: str(payload["id"]) for path, payload in sink.payloads().items()} + + +def test_client_call_ids_differing_only_by_slash_or_underscore_land_in_separate_files( + gateway: Gateway, tmp_path: Path +) -> None: + """An embedding response carries no id, so its log is named after the caller's `x-litellm-call-id`. Two + caller ids that differ only by `/` and `_` are two requests and must leave two logs, neither overwriting + the other""" + marker: Final = f"svc-{uuid.uuid4().hex[:8]}" + call_ids: Final = (f"{marker}/req-1", f"{marker}_req-1") + assert _log_names_by_call_id(gateway, tmp_path, call_ids) == {f"{call_id}.json": call_id for call_id in call_ids} + + +def test_client_call_ids_with_parent_segments_stay_inside_the_log_directory(gateway: Gateway, tmp_path: Path) -> None: + """A caller's `x-litellm-call-id` names its log file, so a `..` segment in it must not climb out of the + dated log directory into another day's folder or another filesystem""" + marker: Final = uuid.uuid4().hex[:8] + call_ids: Final = (f"../other-filesystem/{marker}", f"../2026-09-30/{marker}", f"%2e%2e/other-filesystem/{marker}") + assert _log_names_by_call_id(gateway, tmp_path, call_ids) == { + f".._other-filesystem_{marker}.json": call_ids[0], + f".._2026-09-30_{marker}.json": call_ids[1], + f"%2e%2e_other-filesystem_{marker}.json": call_ids[2], + } + + +def test_client_call_ids_with_dot_or_empty_segments_keep_their_own_files(gateway: Gateway, tmp_path: Path) -> None: + """A `.` or empty segment in a caller's `x-litellm-call-id` collapses on the Data Lake path, so `svc/./x` would + overwrite the log of `svc/x` and `svc//x` would fail to upload. Each id must still leave its own log""" + marker: Final = uuid.uuid4().hex[:8] + call_ids: Final = (f"{marker}/x", f"{marker}/./x", f"{marker}//x") + assert _log_names_by_call_id(gateway, tmp_path, call_ids) == { + f"{marker}/x.json": call_ids[0], + f"{marker}_._x.json": call_ids[1], + f"{marker}__x.json": call_ids[2], + } + + +def test_failed_requests_with_look_alike_call_ids_keep_separate_failure_logs(gateway: Gateway, tmp_path: Path) -> None: + """A failed request has no response id, so its failure log is named after the caller's `x-litellm-call-id`. Two + failures whose ids differ only by `/` and `_` must leave two failure logs""" + marker: Final = f"fail-{uuid.uuid4().hex[:8]}" + call_ids: Final = (f"{marker}/req-1", f"{marker}_req-1") + assert _log_names_by_call_id(gateway, tmp_path, call_ids, failing=True) == { + f"{call_id}.json": call_id for call_id in call_ids + } + + +def test_cache_hits_with_look_alike_call_ids_keep_separate_logs(gateway: Gateway, tmp_path: Path) -> None: + """A cached embedding is served without reaching the provider, and its log is named after the caller's call id + plus a cache-hit suffix. Two cache hits whose ids differ only by `/` and `_` must still leave two logs""" + marker: Final = f"hit-{uuid.uuid4().hex[:8]}" + call_ids: Final = (f"{marker}/warm", f"{marker}/req-1", f"{marker}_req-1") + logs: Final = _log_names_by_call_id(gateway, tmp_path, call_ids, inputs=(marker, marker, marker)) + assert {name: payload_id for name, payload_id in logs.items() if "_cache_hit" not in payload_id} == { + f"{call_ids[0]}.json": call_ids[0] + }, logs + cache_hits: Final = {name: payload_id for name, payload_id in logs.items() if "_cache_hit" in payload_id} + assert sorted(payload_id.split("_cache_hit")[0] for payload_id in cache_hits.values()) == sorted(call_ids[1:]), logs + assert all(name == f"{payload_id}.json" for name, payload_id in cache_hits.items()), logs + + +def test_longest_oversized_and_blank_call_ids_each_leave_one_log(gateway: Gateway, tmp_path: Path) -> None: + """The longest accepted caller id keeps its own name, while a 5 KB or blank `x-litellm-call-id` falls back to a + generated id, so none of the three requests loses its log""" + marker: Final = f"edge-{uuid.uuid4().hex[:8]}/" + longest: Final = marker + "x" * (MAX_LITELLM_CALL_ID_LENGTH - len(marker)) + call_ids: Final = (longest, marker + "x" * 5000, "") + logs: Final = _log_names_by_call_id(gateway, tmp_path, call_ids, inputs=(f"{marker}0", f"{marker}1", f"{marker}2")) + assert logs.get(f"{longest}.json") == longest, tuple(logs) + generated: Final = frozenset(payload_id for payload_id in logs.values() if payload_id != longest) + assert len(logs) == 3 and len(generated) == 2 and not generated & frozenset(call_ids), tuple(logs) + assert all(logs[f"{payload_id}.json"] == payload_id for payload_id in generated), tuple(logs) diff --git a/tests/unit/integrations/azure_storage/test_azure_storage.py b/tests/unit/integrations/azure_storage/test_azure_storage.py index 0227906a2dd..01b17edd2cf 100644 --- a/tests/unit/integrations/azure_storage/test_azure_storage.py +++ b/tests/unit/integrations/azure_storage/test_azure_storage.py @@ -447,6 +447,51 @@ def test_adls_safe_file_name_rewrites_base64_padding_and_reserved_characters(pay ) +def test_adls_safe_file_name_rewrites_only_responses_ids(): + ids = ("svc/req-1", "svc_req-1", "trace=7", "trace7", "resp_YWJjZA==", "resp_+/8=") + names = {payload_id: adls_safe_file_name(payload_id) for payload_id in ids} + assert names == { + "svc/req-1": "svc/req-1.json", + "svc_req-1": "svc_req-1.json", + "trace=7": "trace=7.json", + "trace7": "trace7.json", + "resp_YWJjZA==": "resp_YWJjZA.json", + "resp_+/8=": "resp_+_8.json", + }, "caller-chosen ids must keep their own names so none overwrites another, while resp_ ids are rewritten" + + +def test_adls_safe_file_name_rewrites_ids_with_dot_or_empty_path_segments(): + ids = ( + "../other-filesystem/x", + "../2026-09-30/x", + "svc/../../x", + "%2e%2e/other-filesystem/x", + ".%2E/x", + "a..b/c", + "../", + "svc/./x", + "./x", + "svc//x", + "/x", + "x/", + ) + names = {payload_id: adls_safe_file_name(payload_id) for payload_id in ids} + assert names == { + "../other-filesystem/x": ".._other-filesystem_x.json", + "../2026-09-30/x": ".._2026-09-30_x.json", + "svc/../../x": "svc_.._.._x.json", + "%2e%2e/other-filesystem/x": "%2e%2e_other-filesystem_x.json", + ".%2E/x": ".%2E_x.json", + "a..b/c": "a..b/c.json", + "../": ".._.json", + "svc/./x": "svc_._x.json", + "./x": "._x.json", + "svc//x": "svc__x.json", + "/x": "_x.json", + "x/": "x_.json", + }, "a dot or empty segment must never reach the Data Lake path, while ids without one keep their own names" + + def test_adls_safe_file_name_is_deterministic_and_distinct_per_id(): ids = ( "resp_" + base64.b64encode(b"a").decode(),