From e1afe2e29cee700710faa85063a9a0f7927104f6 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 16:42:15 -0700 Subject: [PATCH 001/126] test(e2e): bound the post-/model/new servable wait at 40s _await_model_servable used poll_timeout (120s), the spend/log read-back budget. A stuck model reload therefore stalled every suite that creates a deployment for two minutes before failing Give create_model a fixed harness middle ground: model_servable_timeout=40s, polled every 2s, with each /v1/models call capped at 5s and clamped to the remaining deadline so one slow GET cannot overrun the wait. Happy path still returns on the first listing. Not derived from proxy general_settings or env Transport.get accepts an optional per-call timeout for that clamp. Unit tests cover the deadline arithmetic and clamp without a live proxy (cherry picked from commit c082a0e6488f50978bf5255f6b5298ba7e8fd8da) --- tests/e2e/proxy_client.py | 134 ++++++++++-- tests/e2e/test_proxy_client_model_servable.py | 190 ++++++++++++++++++ tests/e2e/transport.py | 13 +- 3 files changed, 313 insertions(+), 24 deletions(-) create mode 100644 tests/e2e/test_proxy_client_model_servable.py diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..87693175e62 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -75,12 +75,95 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] +# After /model/new, poll the data plane until the model is listed (or fail). +# Shorter than poll_timeout (spend/log read-backs ~120s); longer than a single +# request. 40s is the harness middle ground: happy path returns on the first +# poll, a stuck reload fails in under a minute instead of two. +MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_INTERVAL = 2.0 +# Cap each /v1/models poll so one slow request cannot outlast the budget. +# Clamped further to remaining deadline inside await_servable. +MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 + + +@dataclass(frozen=True, slots=True) +class Servable: + """The data plane listed the model within the deadline.""" + + +@dataclass(frozen=True, slots=True) +class NotServable: + """The deadline passed without the data plane listing the model. + + `last_result` is the final /v1/models read, so the caller can tell "the proxy + answered but omitted the model" (propagation) from "the read itself failed" + (network/auth) when reporting.""" + + last_result: Result[ModelsListResponse] | None + + +ServableOutcome = Servable | NotServable + + +def await_servable( + list_models: Callable[[float], Result[ModelsListResponse]], + *, + model_name: str, + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ServableOutcome: + """Poll `list_models` until the data plane lists `model_name` or `timeout` elapses. + + `list_models` receives the per-poll request timeout, clamped to the remaining + deadline so a slow final poll cannot overrun the overall budget. Clock and sleep + are injected so this is exercised without wall-clock waits. Always polls at least + once when the loop starts with a positive budget.""" + deadline = now() + timeout + last_result: Result[ModelsListResponse] | None = None + while True: + remaining = deadline - now() + if remaining <= 0 and last_result is not None: + return NotServable(last_result=last_result) + poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout + last_result = list_models(poll_timeout) + if isinstance(last_result, Success) and any( + entry.id == model_name for entry in last_result.data.data + ): + return Servable() + if now() + interval >= deadline: + return NotServable(last_result=last_result) + sleep(interval) + + +def servable_timeout_message( + *, + model_name: str, + timeout: float, + last_result: Result[ModelsListResponse] | None, +) -> str: + last_error = ( + f"; last /v1/models poll did not succeed: {last_result}" + if last_result is not None and not isinstance(last_result, Success) + else "" + ) + return ( + f"model {model_name!r} was created but never became servable on the data " + f"plane within {timeout}s of /model/new (control/data-plane propagation or " + f"STORE_MODEL_IN_DB reload issue){last_error}" + ) + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport poll_timeout: float = 120.0 poll_interval: float = 5.0 + model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_interval: float = MODEL_SERVABLE_INTERVAL + model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- @@ -167,7 +250,12 @@ class ProxyClient: this returns can race the reload and 400 with "Invalid model name passed". We therefore poll the data-plane /v1/models until the model appears before handing back, so callers can invoke it immediately. In the monolithic case - it is already present on the first poll, so this adds one request.""" + it is already present on the first poll, so this adds one request. + + The wait is bounded by `model_servable_timeout` rather than the much longer + `poll_timeout` used for batched read-backs, so a stuck reload fails in under + a minute instead of two. Happy path still returns as soon as /v1/models lists + the model (usually the first poll).""" model_id = unwrap( self.transport.post( "/model/new", @@ -185,32 +273,34 @@ class ProxyClient: def _await_model_servable(self, model_name: str) -> None: """Block until the data plane lists `model_name`, or fail loudly if it does - not within poll_timeout (a real propagation/config problem, surfaced here - instead of as a downstream "Invalid model name passed").""" - deadline = time.monotonic() + self.poll_timeout - last_result: Result[ModelsListResponse] | None = None - while time.monotonic() < deadline: - last_result = self.transport.get( + not within model_servable_timeout (a real propagation/config problem, + surfaced here instead of as a downstream "Invalid model name passed").""" + outcome = await_servable( + lambda poll_timeout: self.transport.get( "/v1/models", headers=self.transport.master, params=NoBody(), response_type=ModelsListResponse, - ) - if isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ): + timeout=poll_timeout, + ), + model_name=model_name, + timeout=self.model_servable_timeout, + interval=self.model_servable_interval, + request_timeout=self.model_servable_request_timeout, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Servable(): return - time.sleep(self.poll_interval) - last_error = ( - f"; last /v1/models poll did not succeed: {last_result}" - if last_result is not None and not isinstance(last_result, Success) - else "" - ) - raise AssertionError( - f"model {model_name!r} was created but never became servable on the data " - f"plane within {self.poll_timeout}s of /model/new (control/data-plane " - f"propagation or STORE_MODEL_IN_DB reload issue){last_error}" - ) + case NotServable(last_result=last_result): + raise AssertionError( + servable_timeout_message( + model_name=model_name, + timeout=self.model_servable_timeout, + last_result=last_result, + ) + ) def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: """Merge `litellm_params` over the deployment `model_id`'s stored params via diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py new file mode 100644 index 00000000000..0cc63f882e2 --- /dev/null +++ b/tests/e2e/test_proxy_client_model_servable.py @@ -0,0 +1,190 @@ +"""Harness coverage for the bounded wait after /model/new (no live proxy). + +Model propagation is polled to a deadline so a stuck control/data-plane reload fails +fast instead of stalling every test that creates a model. The clock and sleep are +injected, so these assert the deadline arithmetic without waiting. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from e2e_http import NetworkError, Result, Success +from models import ModelListEntry, ModelsListResponse +from proxy_client import ( + MODEL_SERVABLE_REQUEST_TIMEOUT, + MODEL_SERVABLE_TIMEOUT, + NotServable, + Servable, + await_servable, + servable_timeout_message, +) + + +def _listing(*model_names: str) -> Result[ModelsListResponse]: + return Success( + status_code=200, + data=ModelsListResponse(data=tuple(ModelListEntry(id=name) for name in model_names)), + ) + + +@dataclass(slots=True) +class FakeClock: + """A clock that only advances when the code under test sleeps or a slow poll runs.""" + + seconds: float = 0.0 + slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions + + def now(self) -> float: + return self.seconds + + def sleep(self, duration: float) -> None: + self.slept.append(duration) + self.seconds += duration + + +@dataclass(slots=True) +class FakeModelList: + """Returns each queued /v1/models read in turn, repeating the last forever.""" + + responses: tuple[Result[ModelsListResponse], ...] + calls: int = 0 + timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts + + def __call__(self, request_timeout: float) -> Result[ModelsListResponse]: + self.timeouts.append(request_timeout) + response = self.responses[min(self.calls, len(self.responses) - 1)] + self.calls += 1 + return response + + +def test_returns_servable_on_first_listing_without_sleeping() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("my-model"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls == 1 + assert list_models.timeouts == [5.0] + assert clock.slept == [] + + +def test_polls_until_the_model_appears() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls == 3 + assert clock.seconds == 4.0 + + +def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("other"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=10.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds == 8.0 + assert list_models.calls == 5 + assert list_models.timeouts == [5.0, 5.0, 5.0, 4.0, 2.0] + + +def test_does_not_wait_past_the_overall_budget() -> None: + clock = FakeClock() + + outcome = await_servable( + FakeModelList(responses=(_listing("other"),)), + model_name="my-model", + timeout=MODEL_SERVABLE_TIMEOUT, + interval=2.0, + request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds <= MODEL_SERVABLE_TIMEOUT + + +def test_clamps_request_timeout_to_remaining_deadline() -> None: + """A slow final poll must not receive the full request cap when less budget remains. + + Without the clamp, remaining=3 and cap=5 lets the transport block for 5s and the + overall wait overruns model_servable_timeout by up to ~cap seconds. + """ + clock = FakeClock() + timeouts: list[float] = [] + + def list_models(request_timeout: float) -> Result[ModelsListResponse]: + timeouts.append(request_timeout) + clock.seconds += request_timeout + return _listing("other") + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=10.0, + interval=2.0, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert timeouts[0] == 5.0 + assert any(timeout < 5.0 for timeout in timeouts) + assert timeouts[-1] == 3.0 + assert clock.seconds <= 10.0 + + +def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: + clock = FakeClock() + unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") + + outcome = await_servable( + FakeModelList(responses=(unreachable,)), + model_name="my-model", + timeout=1.0, + interval=0.5, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == NotServable(last_result=unreachable) + message = servable_timeout_message(model_name="my-model", timeout=1.0, last_result=unreachable) + assert "connection refused" in message + + listed_without_model = _listing("other") + propagation_message = servable_timeout_message( + model_name="my-model", timeout=1.0, last_result=listed_without_model + ) + assert "did not succeed" not in propagation_message diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index a6adf83ed1f..27b11befc8e 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -58,6 +58,7 @@ class Transport(Protocol): headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: ... def delete[R: BaseModel]( @@ -136,13 +137,16 @@ class HttpTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: + """`timeout` overrides the transport-wide request_timeout for this call, for + pollers whose own deadline is shorter than it.""" return e2e_http.get( self._url(path), headers=headers, params=params, response_type=response_type, - timeout=self.request_timeout, + timeout=self.request_timeout if timeout is None else timeout, ) def delete[R: BaseModel]( @@ -336,9 +340,14 @@ class SplitTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: return self._route(path).get( - path, headers=headers, params=params, response_type=response_type + path, + headers=headers, + params=params, + response_type=response_type, + timeout=timeout, ) def delete[R: BaseModel]( From 5aa66ea33e6c194f66d360e22292f22eddac38e7 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:09:20 -0700 Subject: [PATCH 002/126] fix(e2e): wait one default DB reload interval of continuous listing create_model returned after the first /v1/models hit that listed the model, so chat could still land on a cold gateway worker (numWorkers>1 / peer pod) and 400 Invalid model name. Require continuous listing for the product default add_deployment interval (30s) after first sight so every worker has synced from the DB; first listing still bounded at 40s (cherry picked from commit 7d1ee2ff861b970f6de3f6759ff015947af9d2a1) --- tests/e2e/proxy_client.py | 85 ++++++++---- tests/e2e/test_proxy_client_model_servable.py | 125 ++++++++++++------ 2 files changed, 148 insertions(+), 62 deletions(-) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 87693175e62..dfbb90ac08e 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -75,14 +75,17 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] -# After /model/new, poll the data plane until the model is listed (or fail). -# Shorter than poll_timeout (spend/log read-backs ~120s); longer than a single -# request. 40s is the harness middle ground: happy path returns on the first -# poll, a stuck reload fails in under a minute instead of two. +# After /model/new, the control-plane writer reloads itself immediately, but every +# other gateway worker (and peer pod) only picks the model up on its add_deployment +# job. That job runs every proxy_config_reload_interval_seconds (product default 30). +# A single /v1/models hit can land on a hot worker while the next /chat hits a cold +# one ("Invalid model name"). Wait for first listing within MODEL_SERVABLE_TIMEOUT, +# then require continuous listing for MODEL_SERVABLE_DB_SYNC_SECONDS (the default +# reload interval) so every worker has had a chance to sync from the DB. MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_DB_SYNC_SECONDS = 30.0 MODEL_SERVABLE_INTERVAL = 2.0 -# Cap each /v1/models poll so one slow request cannot outlast the budget. -# Clamped further to remaining deadline inside await_servable. +# Cap each /v1/models poll so one slow request cannot outlast the remaining budget. MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 @@ -112,29 +115,55 @@ def await_servable( timeout: float, interval: float, request_timeout: float, + db_sync_seconds: float, now: Callable[[], float], sleep: Callable[[float], None], ) -> ServableOutcome: - """Poll `list_models` until the data plane lists `model_name` or `timeout` elapses. + """Poll until `model_name` is listed long enough for every worker to DB-sync. - `list_models` receives the per-poll request timeout, clamped to the remaining - deadline so a slow final poll cannot overrun the overall budget. Clock and sleep - are injected so this is exercised without wall-clock waits. Always polls at least - once when the loop starts with a positive budget.""" - deadline = now() + timeout + First listing must happen within `timeout`. After that, the model must stay + listed continuously for `db_sync_seconds` (any miss resets the continuous + window). `db_sync_seconds=0` returns on the first listing. Each poll's request + timeout is clamped to the remaining budget. Clock and sleep are injected.""" + started = now() + first_seen_at: float | None = None last_result: Result[ModelsListResponse] | None = None while True: - remaining = deadline - now() + t = now() + if first_seen_at is None: + deadline = started + timeout + else: + deadline = first_seen_at + db_sync_seconds + remaining = deadline - t if remaining <= 0 and last_result is not None: + if first_seen_at is not None and db_sync_seconds <= 0: + return Servable() + if first_seen_at is not None and t - first_seen_at >= db_sync_seconds: + return Servable() return NotServable(last_result=last_result) poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout last_result = list_models(poll_timeout) - if isinstance(last_result, Success) and any( + listed = isinstance(last_result, Success) and any( entry.id == model_name for entry in last_result.data.data - ): + ) + t = now() + if not listed: + first_seen_at = None + elif first_seen_at is None: + first_seen_at = t + if db_sync_seconds <= 0: + return Servable() + elif t - first_seen_at >= db_sync_seconds: return Servable() - if now() + interval >= deadline: - return NotServable(last_result=last_result) + if first_seen_at is None: + if now() + interval >= started + timeout: + return NotServable(last_result=last_result) + elif now() + interval >= first_seen_at + db_sync_seconds: + # Final stretch: sleep only the remainder of the continuous window. + remainder = first_seen_at + db_sync_seconds - now() + if remainder > 0: + sleep(remainder) + continue sleep(interval) @@ -142,6 +171,7 @@ def servable_timeout_message( *, model_name: str, timeout: float, + db_sync_seconds: float, last_result: Result[ModelsListResponse] | None, ) -> str: last_error = ( @@ -151,7 +181,8 @@ def servable_timeout_message( ) return ( f"model {model_name!r} was created but never became servable on the data " - f"plane within {timeout}s of /model/new (control/data-plane propagation or " + f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous " + f"DB sync) after /model/new (control/data-plane propagation or " f"STORE_MODEL_IN_DB reload issue){last_error}" ) @@ -162,6 +193,7 @@ class ProxyClient: poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT @@ -252,10 +284,10 @@ class ProxyClient: handing back, so callers can invoke it immediately. In the monolithic case it is already present on the first poll, so this adds one request. - The wait is bounded by `model_servable_timeout` rather than the much longer - `poll_timeout` used for batched read-backs, so a stuck reload fails in under - a minute instead of two. Happy path still returns as soon as /v1/models lists - the model (usually the first poll).""" + First listing must arrive within `model_servable_timeout` (not the longer + spend `poll_timeout`). The model must then stay listed for + `model_servable_db_sync_seconds` (product default DB reload interval) so every + gateway worker has run add_deployment before callers use the model.""" model_id = unwrap( self.transport.post( "/model/new", @@ -272,9 +304,10 @@ class ProxyClient: return model_id def _await_model_servable(self, model_name: str) -> None: - """Block until the data plane lists `model_name`, or fail loudly if it does - not within model_servable_timeout (a real propagation/config problem, - surfaced here instead of as a downstream "Invalid model name passed").""" + """Block until the data plane lists `model_name` long enough for DB sync. + + Fails if first listing misses model_servable_timeout, or if continuous listing + for model_servable_db_sync_seconds never holds (multi-worker / peer reload).""" outcome = await_servable( lambda poll_timeout: self.transport.get( "/v1/models", @@ -287,6 +320,7 @@ class ProxyClient: timeout=self.model_servable_timeout, interval=self.model_servable_interval, request_timeout=self.model_servable_request_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, now=time.monotonic, sleep=time.sleep, ) @@ -298,6 +332,7 @@ class ProxyClient: servable_timeout_message( model_name=model_name, timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, last_result=last_result, ) ) diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py index 0cc63f882e2..c9edd148c2b 100644 --- a/tests/e2e/test_proxy_client_model_servable.py +++ b/tests/e2e/test_proxy_client_model_servable.py @@ -1,8 +1,9 @@ """Harness coverage for the bounded wait after /model/new (no live proxy). -Model propagation is polled to a deadline so a stuck control/data-plane reload fails -fast instead of stalling every test that creates a model. The clock and sleep are -injected, so these assert the deadline arithmetic without waiting. +create_model must wait for the product default DB reload interval of continuous +listing so multi-worker gateways finish add_deployment before callers use the +model. Clock and sleep are injected so these assert the deadline arithmetic +without wall-clock waits. """ from __future__ import annotations @@ -12,6 +13,7 @@ from dataclasses import dataclass, field from e2e_http import NetworkError, Result, Success from models import ModelListEntry, ModelsListResponse from proxy_client import ( + MODEL_SERVABLE_DB_SYNC_SECONDS, MODEL_SERVABLE_REQUEST_TIMEOUT, MODEL_SERVABLE_TIMEOUT, NotServable, @@ -30,8 +32,6 @@ def _listing(*model_names: str) -> Result[ModelsListResponse]: @dataclass(slots=True) class FakeClock: - """A clock that only advances when the code under test sleeps or a slow poll runs.""" - seconds: float = 0.0 slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions @@ -45,8 +45,6 @@ class FakeClock: @dataclass(slots=True) class FakeModelList: - """Returns each queued /v1/models read in turn, repeating the last forever.""" - responses: tuple[Result[ModelsListResponse], ...] calls: int = 0 timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts @@ -58,7 +56,7 @@ class FakeModelList: return response -def test_returns_servable_on_first_listing_without_sleeping() -> None: +def test_returns_on_first_listing_when_db_sync_is_zero() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("my-model"),)) @@ -68,17 +66,64 @@ def test_returns_servable_on_first_listing_without_sleeping() -> None: timeout=40.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) assert outcome == Servable() assert list_models.calls == 1 - assert list_models.timeouts == [5.0] assert clock.slept == [] -def test_polls_until_the_model_appears() -> None: +def test_requires_continuous_listing_for_default_db_sync_interval() -> None: + clock = FakeClock() + list_models = FakeModelList(responses=(_listing("my-model"),)) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert clock.seconds >= MODEL_SERVABLE_DB_SYNC_SECONDS + assert list_models.calls >= 2 + + +def test_resets_db_sync_window_when_a_poll_misses() -> None: + clock = FakeClock() + list_models = FakeModelList( + responses=( + _listing("my-model"), + _listing("my-model"), + _listing("other"), + _listing("my-model"), + ) + ) + + outcome = await_servable( + list_models, + model_name="my-model", + timeout=40.0, + interval=2.0, + request_timeout=5.0, + db_sync_seconds=6.0, + now=clock.now, + sleep=clock.sleep, + ) + + assert outcome == Servable() + assert list_models.calls >= 4 + assert clock.seconds >= 6.0 + + +def test_polls_until_the_model_first_appears() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) @@ -88,6 +133,7 @@ def test_polls_until_the_model_appears() -> None: timeout=40.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) @@ -97,7 +143,7 @@ def test_polls_until_the_model_appears() -> None: assert clock.seconds == 4.0 -def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: +def test_gives_up_if_first_listing_never_arrives() -> None: clock = FakeClock() list_models = FakeModelList(responses=(_listing("other"),)) @@ -107,6 +153,7 @@ def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: timeout=10.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=30.0, now=clock.now, sleep=clock.sleep, ) @@ -114,32 +161,9 @@ def test_gives_up_at_the_deadline_rather_than_polling_forever() -> None: assert isinstance(outcome, NotServable) assert clock.seconds == 8.0 assert list_models.calls == 5 - assert list_models.timeouts == [5.0, 5.0, 5.0, 4.0, 2.0] - - -def test_does_not_wait_past_the_overall_budget() -> None: - clock = FakeClock() - - outcome = await_servable( - FakeModelList(responses=(_listing("other"),)), - model_name="my-model", - timeout=MODEL_SERVABLE_TIMEOUT, - interval=2.0, - request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds <= MODEL_SERVABLE_TIMEOUT def test_clamps_request_timeout_to_remaining_deadline() -> None: - """A slow final poll must not receive the full request cap when less budget remains. - - Without the clamp, remaining=3 and cap=5 lets the transport block for 5s and the - overall wait overruns model_servable_timeout by up to ~cap seconds. - """ clock = FakeClock() timeouts: list[float] = [] @@ -154,6 +178,7 @@ def test_clamps_request_timeout_to_remaining_deadline() -> None: timeout=10.0, interval=2.0, request_timeout=5.0, + db_sync_seconds=30.0, now=clock.now, sleep=clock.sleep, ) @@ -161,10 +186,27 @@ def test_clamps_request_timeout_to_remaining_deadline() -> None: assert isinstance(outcome, NotServable) assert timeouts[0] == 5.0 assert any(timeout < 5.0 for timeout in timeouts) - assert timeouts[-1] == 3.0 assert clock.seconds <= 10.0 +def test_does_not_wait_past_first_listing_budget_when_missing() -> None: + clock = FakeClock() + + outcome = await_servable( + FakeModelList(responses=(_listing("other"),)), + model_name="my-model", + timeout=MODEL_SERVABLE_TIMEOUT, + interval=2.0, + request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, + db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, + now=clock.now, + sleep=clock.sleep, + ) + + assert isinstance(outcome, NotServable) + assert clock.seconds <= MODEL_SERVABLE_TIMEOUT + + def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: clock = FakeClock() unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") @@ -175,16 +217,25 @@ def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: timeout=1.0, interval=0.5, request_timeout=5.0, + db_sync_seconds=0.0, now=clock.now, sleep=clock.sleep, ) assert outcome == NotServable(last_result=unreachable) - message = servable_timeout_message(model_name="my-model", timeout=1.0, last_result=unreachable) + message = servable_timeout_message( + model_name="my-model", + timeout=1.0, + db_sync_seconds=0.0, + last_result=unreachable, + ) assert "connection refused" in message listed_without_model = _listing("other") propagation_message = servable_timeout_message( - model_name="my-model", timeout=1.0, last_result=listed_without_model + model_name="my-model", + timeout=1.0, + db_sync_seconds=30.0, + last_result=listed_without_model, ) assert "did not succeed" not in propagation_message From 5953a66eab8212beac24d3265cd118b7f4a51176 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:10:04 -0700 Subject: [PATCH 003/126] test(e2e): drop proxy_client model-servable unit tests Keep the create_model DB-sync wait in the harness; the pure-function unit file is not needed for this PR (cherry picked from commit 89204651d1a4537c6f21550c5ab85448ae0923f8) --- tests/e2e/test_proxy_client_model_servable.py | 241 ------------------ 1 file changed, 241 deletions(-) delete mode 100644 tests/e2e/test_proxy_client_model_servable.py diff --git a/tests/e2e/test_proxy_client_model_servable.py b/tests/e2e/test_proxy_client_model_servable.py deleted file mode 100644 index c9edd148c2b..00000000000 --- a/tests/e2e/test_proxy_client_model_servable.py +++ /dev/null @@ -1,241 +0,0 @@ -"""Harness coverage for the bounded wait after /model/new (no live proxy). - -create_model must wait for the product default DB reload interval of continuous -listing so multi-worker gateways finish add_deployment before callers use the -model. Clock and sleep are injected so these assert the deadline arithmetic -without wall-clock waits. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field - -from e2e_http import NetworkError, Result, Success -from models import ModelListEntry, ModelsListResponse -from proxy_client import ( - MODEL_SERVABLE_DB_SYNC_SECONDS, - MODEL_SERVABLE_REQUEST_TIMEOUT, - MODEL_SERVABLE_TIMEOUT, - NotServable, - Servable, - await_servable, - servable_timeout_message, -) - - -def _listing(*model_names: str) -> Result[ModelsListResponse]: - return Success( - status_code=200, - data=ModelsListResponse(data=tuple(ModelListEntry(id=name) for name in model_names)), - ) - - -@dataclass(slots=True) -class FakeClock: - seconds: float = 0.0 - slept: list[float] = field(default_factory=list) # mutable-ok: records calls for assertions - - def now(self) -> float: - return self.seconds - - def sleep(self, duration: float) -> None: - self.slept.append(duration) - self.seconds += duration - - -@dataclass(slots=True) -class FakeModelList: - responses: tuple[Result[ModelsListResponse], ...] - calls: int = 0 - timeouts: list[float] = field(default_factory=list) # mutable-ok: records call timeouts - - def __call__(self, request_timeout: float) -> Result[ModelsListResponse]: - self.timeouts.append(request_timeout) - response = self.responses[min(self.calls, len(self.responses) - 1)] - self.calls += 1 - return response - - -def test_returns_on_first_listing_when_db_sync_is_zero() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("my-model"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls == 1 - assert clock.slept == [] - - -def test_requires_continuous_listing_for_default_db_sync_interval() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("my-model"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert clock.seconds >= MODEL_SERVABLE_DB_SYNC_SECONDS - assert list_models.calls >= 2 - - -def test_resets_db_sync_window_when_a_poll_misses() -> None: - clock = FakeClock() - list_models = FakeModelList( - responses=( - _listing("my-model"), - _listing("my-model"), - _listing("other"), - _listing("my-model"), - ) - ) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=6.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls >= 4 - assert clock.seconds >= 6.0 - - -def test_polls_until_the_model_first_appears() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("other"), _listing("other"), _listing("other", "my-model"))) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=40.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == Servable() - assert list_models.calls == 3 - assert clock.seconds == 4.0 - - -def test_gives_up_if_first_listing_never_arrives() -> None: - clock = FakeClock() - list_models = FakeModelList(responses=(_listing("other"),)) - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=10.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=30.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds == 8.0 - assert list_models.calls == 5 - - -def test_clamps_request_timeout_to_remaining_deadline() -> None: - clock = FakeClock() - timeouts: list[float] = [] - - def list_models(request_timeout: float) -> Result[ModelsListResponse]: - timeouts.append(request_timeout) - clock.seconds += request_timeout - return _listing("other") - - outcome = await_servable( - list_models, - model_name="my-model", - timeout=10.0, - interval=2.0, - request_timeout=5.0, - db_sync_seconds=30.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert timeouts[0] == 5.0 - assert any(timeout < 5.0 for timeout in timeouts) - assert clock.seconds <= 10.0 - - -def test_does_not_wait_past_first_listing_budget_when_missing() -> None: - clock = FakeClock() - - outcome = await_servable( - FakeModelList(responses=(_listing("other"),)), - model_name="my-model", - timeout=MODEL_SERVABLE_TIMEOUT, - interval=2.0, - request_timeout=MODEL_SERVABLE_REQUEST_TIMEOUT, - db_sync_seconds=MODEL_SERVABLE_DB_SYNC_SECONDS, - now=clock.now, - sleep=clock.sleep, - ) - - assert isinstance(outcome, NotServable) - assert clock.seconds <= MODEL_SERVABLE_TIMEOUT - - -def test_reports_a_failed_read_distinctly_from_a_missing_model() -> None: - clock = FakeClock() - unreachable: Result[ModelsListResponse] = NetworkError(message="connection refused") - - outcome = await_servable( - FakeModelList(responses=(unreachable,)), - model_name="my-model", - timeout=1.0, - interval=0.5, - request_timeout=5.0, - db_sync_seconds=0.0, - now=clock.now, - sleep=clock.sleep, - ) - - assert outcome == NotServable(last_result=unreachable) - message = servable_timeout_message( - model_name="my-model", - timeout=1.0, - db_sync_seconds=0.0, - last_result=unreachable, - ) - assert "connection refused" in message - - listed_without_model = _listing("other") - propagation_message = servable_timeout_message( - model_name="my-model", - timeout=1.0, - db_sync_seconds=30.0, - last_result=listed_without_model, - ) - assert "did not succeed" not in propagation_message From 38d03fd341bf5b7030b2ebbc45c2555a7f1ee1e6 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:16:30 -0700 Subject: [PATCH 004/126] fix(e2e): never skip the final deadline-clamped model-servable poll When less than one full poll interval remained in the first-listing budget, the pre-sleep check returned NotServable without another /v1/models call. Sleep only min(interval, time left) so a model that becomes listable in the last seconds of the timeout still gets a clamped final poll (cherry picked from commit 8439195922c913d118cb146409c75cc081d23e6e) --- tests/e2e/proxy_client.py | 43 ++++++++++++++++++++------------------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index dfbb90ac08e..36dada5770a 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -124,24 +124,28 @@ def await_servable( First listing must happen within `timeout`. After that, the model must stay listed continuously for `db_sync_seconds` (any miss resets the continuous window). `db_sync_seconds=0` returns on the first listing. Each poll's request - timeout is clamped to the remaining budget. Clock and sleep are injected.""" + timeout is clamped to the remaining budget. Sleeps only min(interval, time left) + so a final deadline-clamped poll is never skipped just because a full interval + does not fit. Clock and sleep are injected.""" started = now() first_seen_at: float | None = None last_result: Result[ModelsListResponse] | None = None while True: t = now() - if first_seen_at is None: - deadline = started + timeout - else: - deadline = first_seen_at + db_sync_seconds - remaining = deadline - t - if remaining <= 0 and last_result is not None: - if first_seen_at is not None and db_sync_seconds <= 0: - return Servable() - if first_seen_at is not None and t - first_seen_at >= db_sync_seconds: + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + remaining = phase_deadline - t + if remaining <= 0: + if ( + last_result is not None + and first_seen_at is not None + and (db_sync_seconds <= 0 or t - first_seen_at >= db_sync_seconds) + ): return Servable() return NotServable(last_result=last_result) - poll_timeout = min(request_timeout, remaining) if remaining > 0 else request_timeout + + poll_timeout = min(request_timeout, remaining) last_result = list_models(poll_timeout) listed = isinstance(last_result, Success) and any( entry.id == model_name for entry in last_result.data.data @@ -155,16 +159,13 @@ def await_servable( return Servable() elif t - first_seen_at >= db_sync_seconds: return Servable() - if first_seen_at is None: - if now() + interval >= started + timeout: - return NotServable(last_result=last_result) - elif now() + interval >= first_seen_at + db_sync_seconds: - # Final stretch: sleep only the remainder of the continuous window. - remainder = first_seen_at + db_sync_seconds - now() - if remainder > 0: - sleep(remainder) - continue - sleep(interval) + + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + wait = min(interval, phase_deadline - now()) + if wait > 0: + sleep(wait) def servable_timeout_message( From 87be33f9354f399100c7ededdfba31b0e831d225 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 17:35:13 -0700 Subject: [PATCH 005/126] fix(e2e): reject first listing that returns after the 40s deadline A poll may start with remaining budget and still return after started+timeout if the transport overruns its clamp. Recheck the first-listing deadline after the response so a late listing does not open the continuous DB-sync phase (cherry picked from commit 7ff2bcbf1498ee82f7dbe0b4330c1ab48927ed01) --- tests/e2e/proxy_client.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 36dada5770a..b3fc8538322 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -154,6 +154,8 @@ def await_servable( if not listed: first_seen_at = None elif first_seen_at is None: + if t > started + timeout: + return NotServable(last_result=last_result) first_seen_at = t if db_sync_seconds <= 0: return Servable() From 82fa66908bb72748efc574a8a1a8750155a85c14 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 28 Jul 2026 22:15:21 -0700 Subject: [PATCH 006/126] test(e2e): poll MCP tools across multi-worker lag (#35047) * fix(mcp): resolve call_tool by registry without requiring tool map Multi-worker reloads put MCP servers in the registry from the DB but do not re-run tools/list on every process. Gating call_tool on tool_name_to_mcp_server_name_mapping made cold workers 500 with Tool not found after another worker had already listed the tool. Treat a registry match on server id/name/alias as enough; upstream rejects unknown tools * test(e2e): poll MCP register, tools/list, and tools/call across multi-worker lag Stage multi-worker gateways only load MCP servers and tool maps on the process that handled the request. Poll until the server is listed, the tool appears on tools/list, and tools/call is not a cold-worker 500 so key-access and Datadog MCP e2e stop racing the LB * Revert "fix(mcp): resolve call_tool by registry without requiring tool map" This reverts commit 8b56e51e39b876d13d1112efa4130554ddf5f173. * test(e2e): tighten MCP multi-worker lag classifier Only retry tools/call on gateway shapes Tool not found and server_not_found, not any 500 that mentions tool/server not found, so upstream failures are not retried until the poll deadline * test(e2e): drop unit file for MCP lag classifier The live await_call_tool polls already cover multi-worker lag; a separate string-match unit module is not worth keeping (cherry picked from commit c274cf321c5c35c629220a89bb497d15b56f870f) --- tests/e2e/mcp/mcp_client.py | 91 +++++++++++++++++++++- tests/e2e/mcp/test_mcp_access_group_e2e.py | 1 + tests/e2e/mcp/test_mcp_datadog_e2e.py | 28 ++++--- tests/e2e/mcp/test_mcp_key_access_e2e.py | 15 ++-- 4 files changed, 111 insertions(+), 24 deletions(-) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 6d0f6ddc760..33ec557c339 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -11,13 +11,14 @@ request/response bodies are co-located here because only this suite speaks MCP. from __future__ import annotations +import re import time from collections.abc import Mapping from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field, RootModel -from e2e_http import Headers, NoBody, Result, Success, unwrap +from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap from models import KeyGenerateBody, ObjectPermission from proxy_client import ProxyClient @@ -270,6 +271,60 @@ class McpClient: ) time.sleep(self.proxy.poll_interval) + def await_call_tool( + self, + key: str, + *, + server_id: str, + name: str, + arguments: McpToolArguments, + ) -> McpCallToolResponse: + """Poll tools/call until the result is not a multi-worker registry miss. + + Retries only on the gateway's own cold-worker 500 shapes (Tool + not found / server_not_found). Upstream tool errors and other 500s fail + immediately so non-idempotent calls are not repeated. + """ + deadline = time.monotonic() + self.proxy.poll_timeout + last: Result[McpCallToolResponse] | None = None + while True: + last = self.call_tool(key, server_id=server_id, name=name, arguments=arguments) + if not _is_mcp_not_synced(last, tool_name=name): + return unwrap(last) + if time.monotonic() >= deadline: + raise AssertionError( + f"tools/call for {name!r} on server {server_id} still missing on the " + f"data plane after {self.proxy.poll_timeout}s (multi-worker registry lag); " + f"last result: {last}" + ) + time.sleep(self.proxy.poll_interval) + + def await_call_tool_denied( + self, + key: str, + *, + server_id: str, + name: str, + arguments: McpToolArguments, + ) -> UnknownApiError: + """Poll tools/call until a cold-worker miss clears and the call is 403 access_denied.""" + deadline = time.monotonic() + self.proxy.poll_timeout + last: Result[McpCallToolResponse] | None = None + while True: + last = self.call_tool(key, server_id=server_id, name=name, arguments=arguments) + if isinstance(last, UnknownApiError) and last.status_code == 403: + return last + if not _is_mcp_not_synced(last, tool_name=name): + raise AssertionError( + f"ungranted key's tools/call was not 403 access_denied: {last}" + ) + if time.monotonic() >= deadline: + raise AssertionError( + f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; " + f"last result: {last}" + ) + time.sleep(self.proxy.poll_interval) + def register_mcp_content_filter(self, *, name: str, blocked_keyword: str) -> str: """Register a default-on content-filter guardrail that runs on the MCP tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is @@ -317,5 +372,39 @@ class McpClient: ) +def _is_mcp_not_synced( + result: Result[McpCallToolResponse], + *, + tool_name: str | None = None, +) -> bool: + """True only for gateway multi-worker registry misses, not upstream errors. + + Matches the proxy's own shapes: + - ValueError ``Tool not found`` wrapped as HTTP 500 (cold tool map / + unresolved server on this process) + - REST ``server_not_found`` when this worker has not loaded the MCP server row + + Does not treat arbitrary 500 bodies that merely mention "tool" and "not found" + (e.g. upstream MCP payload text) as lag, so await_call_tool does not retry + real failures or non-idempotent calls. + """ + if not isinstance(result, UnknownApiError) or result.status_code != 500: + return False + body = result.body + body_l = body.lower() + + if "server_not_found" in body_l: + return True + if re.search(r"mcp server ['\"][^'\"]+['\"] was not found", body_l): + return True + + # Gateway: "Tool search_datadog_logs not found" (optionally inside a longer message) + if tool_name is not None: + return ( + re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None + ) + return re.search(r"\btool\s+\S+\s+not found\b", body_l) is not None + + def build_client(proxy: ProxyClient) -> McpClient: return McpClient(proxy=proxy) diff --git a/tests/e2e/mcp/test_mcp_access_group_e2e.py b/tests/e2e/mcp/test_mcp_access_group_e2e.py index 1b53d1ca0b4..f72b75fd43d 100644 --- a/tests/e2e/mcp/test_mcp_access_group_e2e.py +++ b/tests/e2e/mcp/test_mcp_access_group_e2e.py @@ -29,6 +29,7 @@ class TestMcpAccessGroupToolSelection: ) -> None: group = f"e2e-mcp-grp-{unique_marker()}" server_id = register_datadog_mcp(client, resources, mcp_access_groups=[group]) + client.await_registered(server_id) granted = client.generate_key( user_id=f"e2e-mcp-ag-granted-{unique_marker()}", diff --git a/tests/e2e/mcp/test_mcp_datadog_e2e.py b/tests/e2e/mcp/test_mcp_datadog_e2e.py index 8a539b86bff..d093e307f99 100644 --- a/tests/e2e/mcp/test_mcp_datadog_e2e.py +++ b/tests/e2e/mcp/test_mcp_datadog_e2e.py @@ -60,6 +60,7 @@ class TestDatadogMcpRoundTrip: _assert_datadog_logger_active(client.proxy) server_id = register_datadog_mcp(client, resources) + client.await_registered(server_id) marker = f"{MARKER_PREFIX}{unique_marker()}" key = client.generate_key( @@ -78,22 +79,19 @@ class TestDatadogMcpRoundTrip: ) tool_name = client.await_tool(key, server_id, SEARCH_LOGS_TOOL) - - call = unwrap( - client.call_tool( - key, - server_id=server_id, - name=tool_name, - arguments={ - "query": marker, - "from": DD_SEARCH_FROM, - "to": "now", - "max_tokens": 5000, - "telemetry": { - "intent": "e2e assert seeded litellm completion log is searchable via MCP" - }, + call = client.await_call_tool( + key, + server_id=server_id, + name=tool_name, + arguments={ + "query": marker, + "from": DD_SEARCH_FROM, + "to": "now", + "max_tokens": 5000, + "telemetry": { + "intent": "e2e assert seeded litellm completion log is searchable via MCP" }, - ) + }, ) assert call.is_error is not True, f"search_datadog_logs errored: {call}" body = call.all_text diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 35c864c07d8..678424e36d1 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -16,7 +16,7 @@ import pytest from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp from e2e_config import DD_SEARCH_FROM, unique_marker -from e2e_http import UnknownApiError, unwrap +from e2e_http import unwrap from lifecycle import ResourceManager from mcp_client import McpClient @@ -72,13 +72,12 @@ class TestMcpKeyWithoutAccessIsDenied: "max_tokens": 1000, "telemetry": {"intent": "e2e control call proving granted key can invoke Datadog MCP"}, } - permitted_call = unwrap( - client.call_tool(permitted_key, server_id=server_id, name=tool_name, arguments=search_args) + permitted_call = client.await_call_tool( + permitted_key, server_id=server_id, name=tool_name, arguments=search_args ) assert permitted_call.is_error is not True, f"granted key's tool call errored: {permitted_call}" - match client.call_tool(denied_key, server_id=server_id, name=tool_name, arguments=search_args): - case UnknownApiError(status_code=403, body=body): - assert "access_denied" in body, f"403 was not an MCP access denial: {body}" - case other: - pytest.fail(f"ungranted key's tool call was not refused with 403 access_denied: {other}") + denied = client.await_call_tool_denied( + denied_key, server_id=server_id, name=tool_name, arguments=search_args + ) + assert "access_denied" in denied.body, f"403 was not an MCP access denial: {denied.body}" From b93030f84e7a414d2106528114b09f1fca1ad1aa Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:20:25 +0000 Subject: [PATCH 007/126] fix(vertex_ai): surface real error/status on vertex batch create instead of IndexError 500 --- litellm/llms/vertex_ai/batches/handler.py | 44 ++++++++++++---- .../llms/vertex_ai/batches/transformation.py | 46 ++++++++++++++--- .../llms/vertex_ai/batches/test_handler.py | 50 +++++++++++++++---- .../vertex_ai/batches/test_transformation.py | 43 +++++++++++++++- 4 files changed, 152 insertions(+), 31 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index ada1356fb6b..f0fd5480c75 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -13,7 +13,7 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( @@ -98,7 +98,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -130,7 +132,9 @@ class VertexAIBatchPrediction(VertexLLM): ) raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -242,7 +246,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -292,7 +298,9 @@ class VertexAIBatchPrediction(VertexLLM): headers=headers, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -365,7 +373,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = ( @@ -390,7 +400,9 @@ class VertexAIBatchPrediction(VertexLLM): params=params, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response = response.json() vertex_batch_response = ( @@ -475,7 +487,9 @@ class VertexAIBatchPrediction(VertexLLM): raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) # HTTPHandler.get() does not accept a timeout parameter retrieve_response = sync_handler.get( @@ -488,7 +502,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -521,7 +538,9 @@ class VertexAIBatchPrediction(VertexLLM): ) raise if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) # AsyncHTTPHandler.get() does not accept a timeout parameter retrieve_response = await client.get( @@ -534,7 +553,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index df903ba7ef0..e4299bcf2a0 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,7 +1,9 @@ from typing import Any, Dict, Optional +from urllib.parse import unquote from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest @@ -199,16 +201,40 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ - from urllib.parse import unquote - - decoded_uri = unquote(gcs_file_uri) - - model_path = decoded_uri.split("publishers/")[1] - parts = model_path.split("/") - model = f"publishers/{'/'.join(parts[:3])}" + model = cls._parse_model_from_gcs_file(gcs_file_uri) + if model is None: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch creation requires the model to be part of `input_file_id`, but " + f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + "Either upload the input file through LiteLLM (POST /v1/files with " + "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " + "pass a uri of the form " + "gs:////publishers//models//" + ), + ) return model + @classmethod + def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: + """ + Returns the `publishers//models/` path from a gcs uri, or None if the uri + does not contain one. + """ + _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") + if not separator: + return None + + parts = model_path.split("/") + if len(parts) < 3 or parts[1] != "models" or not parts[2]: + return None + + return f"publishers/{'/'.join(parts[:3])}" + @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool: """ @@ -216,7 +242,11 @@ class VertexAIBatchTransformation: LiteLLM-managed unified file id) with a `publishers/` model path that `_get_model_from_gcs_file` can parse. """ - return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id + return ( + input_file_id is not None + and input_file_id.startswith("gs://") + and cls._parse_model_from_gcs_file(input_file_id) is not None + ) @classmethod def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str: diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index cacea234777..b9fb5dfe3c5 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -40,6 +40,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) +from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 HMOD = "litellm.llms.vertex_ai.batches.handler" @@ -184,7 +185,7 @@ def test_create_batch_sync_non_200_raises(): client.post.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500") as exc_info: h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -196,6 +197,32 @@ def test_create_batch_sync_non_200_raises(): max_retries=None, ) + assert exc_info.value.status_code == 500 + assert "error text" in str(exc_info.value) + + +def test_create_batch_input_file_id_without_model_raises_400_before_post(): + """A gs:// uri with no publishers//models/ path is a 400, not a bare 500.""" + h = _make_handler() + client = MagicMock() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data={"input_file_id": "gs://bucket/batch-input.jsonl"}, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + assert "gs://bucket/batch-input.jsonl" in str(exc_info.value) + client.post.assert_not_called() + def test_create_batch_async_non_200_raises(): h = _make_handler() @@ -216,9 +243,12 @@ def test_create_batch_async_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 403"): + with pytest.raises(VertexAIError, match="Error: 403") as exc_info: _run(coro) + assert exc_info.value.status_code == 403 + assert "error text" in str(exc_info.value) + # =========================================================================== # # retrieve_batch @@ -292,7 +322,7 @@ def test_retrieve_batch_sync_non_200_raises(): patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), patch(f"{HMOD}.safe_get", return_value=_http_response(status_code=404)), ): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.retrieve_batch( _is_async=False, batch_id=BATCH_ID, @@ -438,7 +468,7 @@ def test_list_batches_sync_non_200_raises(): client.get.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.list_batches( _is_async=False, after=None, @@ -530,7 +560,7 @@ def test_cancel_batch_sync_cancel_post_non_200_raises(): client.post.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -552,7 +582,7 @@ def test_cancel_batch_sync_retrieve_non_200_raises(): client.get.return_value = _http_response(status_code=404) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -672,7 +702,7 @@ def test_async_retrieve_batch_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -726,7 +756,7 @@ def test_async_list_batches_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -779,7 +809,7 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) async_client_post500.get.assert_not_awaited() @@ -801,5 +831,5 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): _run(coro) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 1b37ade6b30..8352ec16389 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -25,6 +25,7 @@ from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.utils import LiteLLMBatch # noqa: E402 @@ -69,6 +70,24 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +@pytest.mark.parametrize( + "input_file_id", + [ + "gs://bucket/no-model-here.jsonl", + "gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", + "gs://bucket/publishers/google/models", + "gs://bucket/publishers/google/models//file-uuid", + ], +) +def test_transform_openai_request_unparseable_model_raises_400(input_file_id: str): + """An input_file_id with no parseable model path is a client error, not an IndexError -> 500.""" + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": input_file_id}) + + assert exc_info.value.status_code == 400 + assert input_file_id in str(exc_info.value) + + # =========================================================================== # # transform_vertex_ai_batch_response_to_openai_batch_response # =========================================================================== # @@ -299,9 +318,29 @@ def test_get_model_from_gcs_file_url_encoded(): assert T._get_model_from_gcs_file(encoded) == "publishers/google/models/gemini-1.5-flash-001" -def test_get_model_from_gcs_file_no_publishers_raises(): - with pytest.raises(IndexError): +def test_get_model_from_gcs_file_no_publishers_raises_400(): + with pytest.raises(VertexAIError) as exc_info: T._get_model_from_gcs_file("gs://bucket/no-model-here.jsonl") + assert exc_info.value.status_code == 400 + + +# =========================================================================== # +# is_unmanaged_gcs_batch_input_file_id +# =========================================================================== # + + +@pytest.mark.parametrize( + "input_file_id, expected", + [ + (INPUT_FILE, True), + (None, False), + ("file-abc123", False), + ("gs://bucket/no-model-here.jsonl", False), + ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + ], +) +def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): + assert T.is_unmanaged_gcs_batch_input_file_id(input_file_id) is expected # =========================================================================== # From 18d9c7aa21e1308c5ecf05254b29bc0715965bde Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 22:09:49 +0000 Subject: [PATCH 008/126] fix(bedrock): pass SSE-KMS key through to the batch input-file S3 upload --- .../llms/bedrock/batches/transformation.py | 8 ++- litellm/llms/bedrock/common_utils.py | 19 ++++- litellm/llms/bedrock/files/transformation.py | 16 ++++- litellm/types/router.py | 1 + .../bedrock/batches/test_transformation.py | 2 +- .../test_bedrock_files_transformation.py | 72 ++++++++++++++++++- tests/test_litellm/test_router.py | 31 ++++++++ 7 files changed, 141 insertions(+), 8 deletions(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index a4ff1c78467..7500531b81a 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -12,7 +12,6 @@ from litellm.litellm_core_utils.cloud_storage_security import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.bedrock import ( BedrockCreateBatchRequest, BedrockCreateBatchResponse, @@ -29,7 +28,7 @@ from litellm.types.llms.openai import ( from litellm.types.utils import LiteLLMBatch, LlmProviders from ..base_aws_llm import BaseAWSLLM -from ..common_utils import CommonBatchFilesUtils +from ..common_utils import CommonBatchFilesUtils, resolve_s3_encryption_key_id # Bedrock batch input files are uploaded as # s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see @@ -200,7 +199,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) # Add optional KMS encryption key ID if provided - s3_encryption_key_id = litellm_params.get("s3_encryption_key_id") or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + s3_encryption_key_id = resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ) if s3_encryption_key_id: s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5114677ffc0..9d427fa6f12 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -35,7 +35,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( ) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret +from litellm.secret_managers.main import get_secret, get_secret_str if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues @@ -1313,6 +1313,23 @@ def get_anthropic_beta_from_headers(headers: dict) -> List[str]: return [] +def resolve_s3_encryption_key_id( + litellm_params: Mapping[str, Any], + optional_params: Mapping[str, Any] | None = None, +) -> str | None: + """ + Resolve the SSE-KMS key configured for Bedrock batch/file S3 objects. + + Precedence: `s3_encryption_key_id` in litellm_params, then optional_params + (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var. + """ + for source in (litellm_params, optional_params or {}): + value = source.get("s3_encryption_key_id") + if isinstance(value, str) and value: + return value + return get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + + class CommonBatchFilesUtils: """ Common utilities for Bedrock batch and file operations. diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index d4865a1c87a..d1674b260b4 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -53,7 +53,7 @@ from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError +from ..common_utils import BedrockError, resolve_s3_encryption_key_id # litellm_params key used to hand the SigV4-signed GET headers from # `transform_file_content_request` to `validate_environment` (the only hook @@ -741,6 +741,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content=file_content, api_base=api_base, optional_params=optional_params, + s3_encryption_key_id=resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ), ) litellm_params["upload_url"] = api_base @@ -758,6 +762,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content: str, api_base: str, optional_params: dict, + s3_encryption_key_id: str | None = None, ) -> Tuple[dict, str]: """ Sign S3 PUT request using the same proven logic as S3Logger. @@ -790,11 +795,20 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest() # Prepare headers with required S3 headers (same as s3_v2.py) + sse_headers = ( + { + "x-amz-server-side-encryption": "aws:kms", + "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, + } + if s3_encryption_key_id + else {} + ) request_headers = { "Content-Type": "application/json", # JSONL files are JSON content "x-amz-content-sha256": content_hash, # REQUIRED by S3 "Content-Language": "en", "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + **sse_headers, } # Use requests.Request to prepare the request (same pattern as s3_v2.py) diff --git a/litellm/types/router.py b/litellm/types/router.py index 28e4a8272e8..c4d679a2500 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -211,6 +211,7 @@ class CredentialLiteLLMParams(BaseModel): aws_bedrock_runtime_endpoint: Optional[str] = None aws_bedrock_project_id: Optional[str] = None s3_bucket_name: Optional[str] = None + s3_encryption_key_id: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 3681daffe5e..01420eb10df 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -172,7 +172,7 @@ def test_create_request_omits_kms_key_when_absent(config): "generate_unique_job_name", return_value="litellm-batch-1", ), patch.object(config.common_utils, "sign_aws_request") as mock_sign, patch( - "litellm.llms.bedrock.batches.transformation.get_secret_str", + "litellm.llms.bedrock.common_utils.get_secret_str", return_value=None, ): mock_sign.return_value = ({}, b"{}") diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e15..a57e5801327 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -442,7 +442,7 @@ class TestBedrockFilesTransformation: captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -498,7 +498,7 @@ class TestBedrockFilesTransformation: captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -514,6 +514,74 @@ class TestBedrockFilesTransformation: captured_optional_params.get("aws_region_name") == "us-gov-west-1" ), "s3_region_name must override aws_region_name for SigV4 signing" + def _signed_upload_request(self, litellm_params: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + jsonl_content = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + }, + } + ).encode() + + request = config.transform_create_file_request( + model="amazon.nova-pro-v1:0", + create_file_data={ + "file": ("batch.jsonl", jsonl_content, "application/jsonl"), + "purpose": "batch", + }, + optional_params={ + "aws_access_key_id": "test-key-id", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + litellm_params={"s3_bucket_name": "litellm-batch-bucket", **litellm_params}, + ) + assert isinstance(request, dict) + return request + + def test_upload_signs_sse_kms_headers_when_key_configured(self, monkeypatch): + """ + Buckets whose policy requires SSE-KMS reject the batch input-file PutObject + unless the upload carries the aws:kms encryption headers; they must also be + covered by SigV4 SignedHeaders or S3 answers SignatureDoesNotMatch. + """ + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + kms_key = "arn:aws:kms:us-west-2:1234:key/abcd" + + request = self._signed_upload_request({"s3_encryption_key_id": kms_key}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption"] == "aws:kms" + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == kms_key + signed_headers = headers["authorization"].split("SignedHeaders=")[1].split(",")[0] + assert "x-amz-server-side-encryption" in signed_headers + assert "x-amz-server-side-encryption-aws-kms-key-id" in signed_headers + + def test_upload_reads_sse_kms_key_from_env(self, monkeypatch): + monkeypatch.setenv("AWS_S3_ENCRYPTION_KEY_ID", "env-kms-key") + + request = self._signed_upload_request({}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == "env-kms-key" + + def test_upload_omits_sse_headers_when_no_key_configured(self, monkeypatch): + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + + request = self._signed_upload_request({}) + + headers = {key.lower() for key in request["headers"]} + assert "x-amz-server-side-encryption" not in headers + assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers + def test_openai_passthrough_still_works(self): """ Regression test: ensure OpenAI-compatible models (e.g. gpt-oss) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 46b5ce65c3f..fa047d7ee46 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3666,6 +3666,37 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" +def test_get_deployment_credentials_with_provider_includes_s3_encryption_key_id(): + """ + Regression: s3_encryption_key_id must survive the CredentialLiteLLMParams filter, + otherwise the Bedrock batch input-file upload loses the SSE-KMS key and S3 rejects + the PutObject on buckets whose policy requires aws:kms encryption. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-4-20250514-v1:0", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-batch-bucket", + "s3_encryption_key_id": "arn:aws:kms:us-west-2:1234:key/abcd", + }, + } + ], + ) + + credentials = router.get_deployment_credentials_with_provider( + model_id="bedrock-batch" + ) + + assert credentials is not None + assert ( + credentials["s3_encryption_key_id"] + == "arn:aws:kms:us-west-2:1234:key/abcd" + ) + + def test_get_deployment_credentials_with_provider_resolves_credential_name(): """ Test that get_deployment_credentials_with_provider correctly resolves From 0b809cf7d68e368b7a7ee90d7134c4841c944e3c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:21:20 +0000 Subject: [PATCH 009/126] fix(anthropic adapter): stop indexing choices[0] on choiceless streaming chunks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../adapters/streaming_iterator.py | 30 +++++++ .../test_streaming_iterator_empty_choices.py | 87 +++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index d9bcfa19a7f..194cbcc9327 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -329,6 +329,26 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits)) return self._augment_message_delta_usage(merged_chunk) + def _handle_choiceless_chunk(self, chunk: Any) -> bool: + """Consume an OpenAI-compatible chunk that carries no ``choices``. + + ``choices`` is legitimately empty on metadata-only chunks; the final + usage chunk emitted when ``stream_options.include_usage`` is set is the + common case (vLLM and other OpenAI-compatible servers do this). Such a + chunk carries no content-block information, so the caller must not run + the content-block state machine over it. + + Returns True when a merged ``message_delta`` was queued (usage folded + into the held stop-reason chunk); False when the chunk should be + skipped entirely. + """ + if self.holding_stop_reason_chunk is not None and getattr(chunk, "usage", None) is not None: + self.chunk_queue.append(self._merge_usage_into_held_stop_reason_chunk(chunk)) + self.queued_usage_chunk = True + self.holding_stop_reason_chunk = None + return True + return False + def _ensure_context_management_attached(self, message_delta_chunk: Dict[str, Any]) -> Dict[str, Any]: """Attach ``context_management`` to a ``message_delta`` chunk if ``self.applied_edits`` is non-empty and the chunk does not already @@ -490,6 +510,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if chunk == "None" or chunk is None: raise Exception + if not getattr(chunk, "choices", None): + if self._handle_choiceless_chunk(chunk): + return self.chunk_queue.popleft() + continue + should_start_new_block = self._should_start_new_content_block(chunk) is_opening_first_block = self.sent_content_block_start is False if is_opening_first_block and self._is_blank_delta(chunk): @@ -713,6 +738,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if chunk == "None" or chunk is None: raise Exception + if not getattr(chunk, "choices", None): + if self._handle_choiceless_chunk(chunk): + return self.chunk_queue.popleft() + continue + should_start_new_block = self._should_start_new_content_block(chunk) is_opening_first_block = self.sent_content_block_start is False if is_opening_first_block and self._is_blank_delta(chunk): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py new file mode 100644 index 00000000000..3e85872f1e5 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py @@ -0,0 +1,87 @@ +""" +Regression tests for OpenAI-compatible chunks with an empty ``choices`` list. + +``choices: []`` is valid OpenAI-compatible streaming: vLLM (and OpenAI itself, +when ``stream_options.include_usage`` is set) emits a final usage chunk with no +choices, and some gateways emit metadata-only chunks mid-stream. The adapter +used to index ``chunk.choices[0]`` unconditionally, so such a chunk raised +``IndexError: list index out of range`` and killed the ``/v1/messages`` stream. +""" + +import asyncio +import json +from typing import Any, AsyncIterator, Dict, List, Optional + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + +def _text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=text), finish_reason=None)] + ) + + +def _finish_chunk() -> ModelResponseStream: + return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]) + + +def _empty_choices_chunk(usage: Optional[Usage] = None) -> ModelResponseStream: + return ModelResponseStream(choices=[], usage=usage) + + +def _collect_async(wrapper: AnthropicStreamWrapper) -> str: + async def _run() -> str: + return "".join( + [raw.decode() if isinstance(raw, bytes) else raw async for raw in wrapper.async_anthropic_sse_wrapper()] + ) + + return asyncio.run(_run()) + + +def _message_delta(sse: str) -> Dict[str, Any]: + return next( + json.loads(line[len("data: ") :]) + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"message_delta"' in line + ) + + +def test_leading_metadata_chunk_without_choices_does_not_kill_stream(): + """A metadata-only chunk before any content must be skipped, not indexed.""" + chunks: List[ModelResponseStream] = [ + _empty_choices_chunk(), + _text_chunk("Hello"), + _text_chunk(" there"), + _finish_chunk(), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="mock-model") + events = list(wrapper) + + text = "".join( + event["delta"]["text"] for event in events if event.get("type") == "content_block_delta" + ) + assert text == "Hello there" + assert events[-1]["type"] == "message_stop" + + +def test_final_usage_chunk_without_choices_is_merged_into_message_delta(): + """The vLLM/OpenAI final usage chunk carries no choices; its usage must + still land on the Anthropic ``message_delta``.""" + usage = Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13) + + async def _aiter() -> "AsyncIterator[ModelResponseStream]": + for chunk in [_text_chunk("Hi"), _finish_chunk(), _empty_choices_chunk(usage)]: + yield chunk + + sse = _collect_async(AnthropicStreamWrapper(completion_stream=_aiter(), model="mock-model")) + + message_delta = _message_delta(sse) + assert message_delta["delta"]["stop_reason"] == "end_turn" + assert message_delta["usage"]["input_tokens"] == 10 + assert message_delta["usage"]["output_tokens"] == 3 + assert "Hi" in sse + assert "message_stop" in sse From f0ffc6507e1d21daa4f3a13a0245daa55effccd2 Mon Sep 17 00:00:00 2001 From: Rithvik Mysore Suresh Date: Fri, 31 Jul 2026 10:16:35 -0400 Subject: [PATCH 010/126] fix(batches): keep managed files on owner Managed files and batches are provider-owned. Cross-model fallbacks can dispatch creation with credentials that cannot access the input file and replace the owning provider's validation error.\n\nCloses #35359 --- litellm/proxy/batches_endpoints/endpoints.py | 5 ++- .../proxy/batches_endpoints/test_endpoints.py | 3 +- tests/test_litellm/test_router.py | 45 +++++++++++++++++++ 3 files changed, 51 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index a91b29002e3..b7713d388ea 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -262,7 +262,10 @@ async def create_batch( detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) - response = await llm_router.acreate_batch(**_create_batch_data) + response = await llm_router.acreate_batch( + **_create_batch_data, + disable_fallbacks=True, + ) response.input_file_id = input_file_id response._hidden_params["unified_file_id"] = unified_file_id else: diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 6a185988c9b..b382313ea1f 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -469,7 +469,7 @@ async def test_create__fallback_body_custom_llm_provider(harness): @pytest.mark.asyncio -async def test_create__unified_file_id_single_model(harness): +async def test_create__unified_file_id_single_model_disables_cross_model_fallbacks(harness): set_body( harness, { @@ -489,6 +489,7 @@ async def test_create__unified_file_id_single_model(harness): harness.litellm_acreate.assert_not_called() # model injected from the unified id, input_file_id restored, hidden param set assert harness.router_kwargs()["model"] == "gpt-4o-mini" + assert harness.router_kwargs()["disable_fallbacks"] is True assert resp.input_file_id == "litellm_proxy_unified_id" assert resp._hidden_params["unified_file_id"] == "unified-xyz" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 46b5ce65c3f..aa917757bbf 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -6054,6 +6054,51 @@ def test_get_configured_token_limits_coerces_numeric_strings(): assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000) +@pytest.mark.asyncio +async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error(): + router = litellm.Router( + model_list=[ + { + "model_name": "owning-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-owning", + }, + }, + { + "model_name": "fallback-model", + "litellm_params": { + "model": "azure/gpt-4o-mini", + "api_key": "sk-fallback", + "api_base": "https://fallback.openai.azure.com", + "api_version": "2024-08-01-preview", + }, + }, + ], + fallbacks=[{"owning-model": ["fallback-model"]}], + num_retries=0, + ) + owning_provider_error = litellm.BadRequestError( + message="completion_window must be one of: 24h", + model="openai/gpt-4o-mini", + llm_provider="openai", + ) + mock_create = AsyncMock(side_effect=owning_provider_error) + + with patch.object(router, "_acreate_batch", mock_create): + with pytest.raises(litellm.BadRequestError, match="24h"): + await router.acreate_batch( + model="owning-model", + input_file_id="file-owned-by-openai", + endpoint="/v1/chat/completions", + completion_window="5m", + disable_fallbacks=True, + ) + + mock_create.assert_awaited_once() + assert mock_create.call_args.kwargs["model"] == "owning-model" + + @pytest.mark.asyncio async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): import httpx From 55726fc09e979fee39680a465b1e04b95def8c3b Mon Sep 17 00:00:00 2001 From: Rithvik Mysore Suresh Date: Fri, 31 Jul 2026 10:28:52 -0400 Subject: [PATCH 011/126] fix(batches): override existing fallback flag Build one kwargs mapping so managed-file ownership always disables cross-model fallback without duplicating a request-enriched key. --- litellm/proxy/batches_endpoints/endpoints.py | 3 +-- tests/test_litellm/proxy/batches_endpoints/test_endpoints.py | 1 + 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index b7713d388ea..a5a03320f7a 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -263,8 +263,7 @@ async def create_batch( ) response = await llm_router.acreate_batch( - **_create_batch_data, - disable_fallbacks=True, + **{**_create_batch_data, "disable_fallbacks": True}, ) response.input_file_id = input_file_id response._hidden_params["unified_file_id"] = unified_file_id diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index b382313ea1f..f8bc3e10d79 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -476,6 +476,7 @@ async def test_create__unified_file_id_single_model_disables_cross_model_fallbac "input_file_id": "litellm_proxy_unified_id", "endpoint": "/v1/chat/completions", "completion_window": "24h", + "disable_fallbacks": False, }, ) with ( From efb5f74173879660d4a79eed66d882da57980946 Mon Sep 17 00:00:00 2001 From: Rithvik Mysore Suresh Date: Fri, 31 Jul 2026 10:39:05 -0400 Subject: [PATCH 012/126] fix(batches): overwrite fallback flag in place Avoid a fresh mutable kwargs mapping while still replacing any request-enriched value before router dispatch. --- litellm/proxy/batches_endpoints/endpoints.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index a5a03320f7a..f94518b16b6 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -262,9 +262,8 @@ async def create_batch( detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, ) - response = await llm_router.acreate_batch( - **{**_create_batch_data, "disable_fallbacks": True}, - ) + _create_batch_data.update(disable_fallbacks=True) # pyright: ignore[reportCallIssue] # router flag + response = await llm_router.acreate_batch(**_create_batch_data) response.input_file_id = input_file_id response._hidden_params["unified_file_id"] = unified_file_id else: From c5c5a276790529e2de3378654864fd847530c5a0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 00:19:42 +0000 Subject: [PATCH 013/126] fix(files): enforce require_managed_files on file retrieve, content and delete Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../openai_files_endpoints/common_utils.py | 30 +++++ .../openai_files_endpoints/files_endpoints.py | 7 ++ .../test_files_endpoint.py | 116 ++++++++++++++++++ 3 files changed, 153 insertions(+) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 87514b46dbd..3eef3868c94 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -866,6 +866,36 @@ def validate_managed_files_requirement( ) +def validate_managed_file_id_requirement(file_id: str) -> None: + """ + Enforce proxy-level managed files on the file read/delete routes when + ``litellm.require_managed_files`` is enabled. + + Ownership is only recorded for LiteLLM managed files, so a raw provider file id sent to + retrieve/content/delete is forwarded to the provider under shared credentials without any + tenant check; knowing another tenant's provider file id would be enough to read or delete it. + + Raises: + HTTPException: 400 if ``file_id`` is not a LiteLLM managed file id. + """ + import litellm + from fastapi import HTTPException + + if litellm.require_managed_files is not True: + return + + if _is_base64_encoded_unified_file_id(file_id): + return + + raise HTTPException( + status_code=400, + detail=( + "Raw provider file ids cannot be used when require_managed_files is enabled in " + "litellm_settings. Use the LiteLLM managed file id returned when the file was created." + ), + ) + + def _extract_model_param(request: "Request", request_body: dict) -> str | None: """ Extract model parameter from request. diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 4e4718272bd..37f1ced6996 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, + validate_managed_file_id_requirement, validate_managed_files_requirement, ) from litellm.proxy.utils import ProxyLogging, is_known_model @@ -612,6 +613,8 @@ async def get_file_content( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + # Include original request and headers in the data base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -908,6 +911,8 @@ async def get_file( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -1098,6 +1103,8 @@ async def delete_file( data: dict = {"file_id": file_id} try: + validate_managed_file_id_requirement(file_id=file_id) + custom_llm_provider = ( provider or get_custom_llm_provider_from_request_headers(request=request) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ac01c6ae1d1..24b814bae1f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3051,3 +3051,119 @@ def test_list_files_key_allowed_openai_model_still_resolves_team_credentials( mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["team-gpt"] ) assert captured_kwargs.get("api_key") == "team-openai-key" + + +@pytest.mark.parametrize( + "http_method, url, patched_litellm_call", + [ + ("get", "/v1/files/file-victim-abc123", "litellm.afile_retrieve"), + ("get", "/v1/files/file-victim-abc123/content", "litellm.afile_content"), + ("delete", "/v1/files/file-victim-abc123", "litellm.afile_delete"), + ], +) +def test_require_managed_files_rejects_raw_provider_file_id( + mocker: MockerFixture, + monkeypatch, + llm_router: Router, + http_method: str, + url: str, + patched_litellm_call: str, +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", True) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_call = mocker.patch(patched_litellm_call, new=mocker.AsyncMock()) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="attacker-user" + ) + + try: + response = getattr(client, http_method)( + url, headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + monkeypatch.setattr("litellm.require_managed_files", False) + + assert response.status_code == 400, response.text + mock_call.assert_not_called() + + +def _unified_managed_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-3.5-turbo", "file-victim-abc123", "gpt-3.5-turbo-id" + ) + return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + +def test_require_managed_files_allows_unified_managed_file_id(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_file_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", True) + + validate_managed_file_id_requirement(file_id=_unified_managed_file_id()) + + +def test_managed_file_id_requirement_is_opt_in(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_file_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", False) + + validate_managed_file_id_requirement(file_id="file-victim-abc123") + + +def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", False) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_retrieve = mocker.patch( + "litellm.afile_retrieve", + new=mocker.AsyncMock( + return_value=OpenAIFileObject( + id="file-victim-abc123", + object="file", + bytes=3, + created_at=1234567890, + filename="test.txt", + purpose="user_data", + status="uploaded", + ) + ), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="some-user" + ) + + try: + response = client.get( + "/v1/files/file-victim-abc123", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + mock_retrieve.assert_called_once() From a7250f4eeab560299215773b50ba32405f597a1c Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 17:11:50 +0000 Subject: [PATCH 014/126] fix(bedrock): normalize /v1/completions and /v1/responses batch records Bedrock managed-batch file upload read `messages` unconditionally, so a JSONL record shaped for /v1/completions (`prompt`) or /v1/responses (`input`) reached the per-provider transform with an empty message list. Anthropic and Nova rejected it at POST /v1/files, and the passthrough providers shipped an empty conversation to AWS. Classify each record by its OpenAI batch `url`, then normalize the non-embedding shapes to chat completions before the Bedrock transforms: `prompt` wraps into user messages the way litellm.text_completion does in real time, and `input` goes through the existing Responses-to-Chat bridge. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../prompt_templates/common_utils.py | 22 +- litellm/llms/bedrock/files/transformation.py | 194 ++++++++-- litellm/types/llms/bedrock.py | 15 + ...ore_utils_prompt_templates_common_utils.py | 43 +++ .../test_bedrock_files_transformation.py | 347 ++++++++++++++++-- 5 files changed, 559 insertions(+), 62 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 639c93dfb80..777ba398d5a 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,7 +6,7 @@ import io import json import mimetypes import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from os import PathLike from pathlib import Path from typing import ( @@ -1742,3 +1742,23 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: idx = end_idx return results + + +def text_completion_prompt_to_messages(prompt: str | Sequence[str]) -> tuple[AllMessageValues, ...]: + """ + Wrap an OpenAI ``/v1/completions`` ``prompt`` into Chat Completion messages. + + Mirrors what ``litellm.text_completion`` does on the real-time path: a + string becomes a single user message, and a list of strings becomes one + user message per element. Pre-tokenized prompts (``list[int]`` / + ``list[list[int]]``) are only meaningful for the OpenAI-family text + endpoints, so they are rejected here rather than silently forwarded, as is + an empty prompt, which every chat-shaped provider rejects downstream. + """ + if isinstance(prompt, str) and prompt: + return (ChatCompletionUserMessage(role="user", content=prompt),) + if isinstance(prompt, Sequence) and prompt and all(isinstance(entry, str) and entry for entry in prompt): + return tuple(ChatCompletionUserMessage(role="user", content=entry) for entry in prompt) + raise ValueError( + f"`prompt` must be a non-empty string or a non-empty list of strings. Got: {type(prompt).__name__}." + ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 3656088cb9d..0aa832780e5 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -2,7 +2,9 @@ import base64 import json import os import time -from collections.abc import Mapping, MutableMapping +from collections.abc import Iterable, Mapping, MutableMapping +from functools import cache +from itertools import chain from types import MappingProxyType from typing import ( Any, @@ -12,7 +14,7 @@ from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -28,12 +30,16 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) -from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + text_completion_prompt_to_messages, +) from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) +from litellm.types.llms.bedrock import BedrockBatchRecordKind from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -43,6 +49,8 @@ from litellm.types.llms.openai import ( OpenAICreateFileRequestOptionalParams, OpenAIFileObject, PathLike, + ResponseInputParam, + ResponsesAPIOptionalRequestParams, ) from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider @@ -57,6 +65,26 @@ from ..common_utils import BedrockError S3_SIGNED_GET_HEADERS_PARAM = "_s3_signed_get_headers" +def _frozen_mapping(items: Iterable[tuple[str, Any]]) -> Mapping[str, Any]: + return MappingProxyType(dict(items)) + + +# JSONL batch records are untyped json, so the `/v1/responses` fields are +# validated into their concrete Responses API types before being handed to the +# Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't +# define, which is what the bridge would ignore anyway. Built on first use +# rather than at import: `ResponseInputParam` is a deep union and only batch +# files carrying `/v1/responses` records need it. +@cache +def _responses_input_adapter() -> TypeAdapter[str | ResponseInputParam]: + return TypeAdapter(str | ResponseInputParam) + + +@cache +def _responses_request_adapter() -> TypeAdapter[ResponsesAPIOptionalRequestParams]: + return TypeAdapter(ResponsesAPIOptionalRequestParams) + + class _BedrockS3RequestParams(BaseModel): """Typed view of the credential/region params the S3 GetObject path reads.""" @@ -305,41 +333,55 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # example; add others here as they adopt the same schema. CONVERSE_INVOKE_PROVIDERS = ("nova",) - # OpenAI batch URL that signals an embedding request. Per OpenAI Batch API - # spec, every JSONL record carries a `url` field; we use it as the - # authoritative signal to route the line to the embedding code path - # instead of inferring from the presence of `input` vs `messages`. + # OpenAI batch URLs that select which request shape a JSONL line carries. + # Per the OpenAI Batch API spec every record carries a `url`, so we use it + # as the authoritative routing signal instead of inferring from the + # presence of `input` vs `prompt` vs `messages`. OPENAI_EMBEDDINGS_URL = "/v1/embeddings" + OPENAI_TEXT_COMPLETIONS_URL = "/v1/completions" + OPENAI_RESPONSES_URL = "/v1/responses" @staticmethod - def _is_embedding_record(openai_jsonl_record: dict[str, Any]) -> bool: + def _classify_batch_record(openai_jsonl_record: Mapping[str, Any]) -> BedrockBatchRecordKind: """ - Decide whether an OpenAI batch JSONL line is an embedding request. + Decide which OpenAI endpoint shape an OpenAI batch JSONL line carries. - Precedence (strict - any explicit `url` short-circuits): - 1. `url == "/v1/embeddings"` -> embedding. Authoritative per the - OpenAI Batch API spec. - 2. Any other non-empty `url` (e.g. `/v1/chat/completions`) -> NOT - embedding. We trust the caller's explicit signal even if the - body would otherwise suggest embedding; misrouting a chat - record into the embedding transformer would corrupt the - modelInput, while a chat-shaped body sent to the chat path - either succeeds or fails cleanly inside that transformer. - 3. `url` missing/empty -> fall back to body shape. Requires - `input` present AND `messages` absent so a malformed record - carrying both keys routes to the chat path (safer default: - Anthropic transforms ignore unknown top-level keys, whereas - the embedding transformer would silently drop the messages). + Precedence (strict - any recognized `url` short-circuits): + 1. A `url` matching a supported endpoint wins. Authoritative per the + OpenAI Batch API spec, which requires it on every record. + 2. Any other non-empty `url` -> chat. We trust the caller's explicit + signal rather than re-deriving it from the body, and an + unexpectedly-shaped body fails cleanly inside the chat + transformer instead of being silently misrouted. + 3. `url` missing/empty -> fall back to body shape. `messages` wins + over the other keys so a malformed record carrying several of + them keeps its conversation instead of having it dropped, and a + bare `input` stays an embedding for backwards compatibility + (that ambiguity with `/v1/responses` is only resolvable from + `url`). """ - url = openai_jsonl_record.get("url") - if url == BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: - return True - if url: - return False - body = openai_jsonl_record.get("body", {}) - if not isinstance(body, dict): - return False - return "input" in body and "messages" not in body + match openai_jsonl_record.get("url"): + case BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: + return BedrockBatchRecordKind.EMBEDDING + case BedrockFilesConfig.OPENAI_TEXT_COMPLETIONS_URL: + return BedrockBatchRecordKind.TEXT_COMPLETION + case BedrockFilesConfig.OPENAI_RESPONSES_URL: + return BedrockBatchRecordKind.RESPONSES + case None | "": + pass + case _: + return BedrockBatchRecordKind.CHAT + + body = openai_jsonl_record.get("body") + if not isinstance(body, Mapping): + return BedrockBatchRecordKind.CHAT + if "messages" in body: + return BedrockBatchRecordKind.CHAT + if "prompt" in body: + return BedrockBatchRecordKind.TEXT_COMPLETION + if "input" in body: + return BedrockBatchRecordKind.EMBEDDING + return BedrockBatchRecordKind.CHAT # Identifier for the Bedrock Titan v2 InvokeModel body schema as stored # in `model_prices_and_context_window.json`. Centralized so future @@ -546,9 +588,83 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) return dict(titan_config._transform_request(input=input_text, inference_params=inference_params)) + @staticmethod + def _transform_text_completion_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Rewrite an OpenAI `/v1/completions` batch body as a Chat Completions body. + + Bedrock batch `modelInput` is the model's InvokeModel/Converse body, and + no Bedrock batch model takes a bare `prompt`, so the wrapping that + `litellm.text_completion` does in real time has to happen here too. + """ + prompt = openai_request_body.get("prompt") + if prompt is None: + raise ValueError( + "Batch record for /v1/completions is missing required `prompt` field: " + f"model={openai_request_body.get('model', '')}" + ) + return _frozen_mapping( + chain( + ((key, value) for key, value in openai_request_body.items() if key != "prompt"), + (("messages", text_completion_prompt_to_messages(prompt)),), + ) + ) + + @staticmethod + def _transform_responses_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Rewrite an OpenAI `/v1/responses` batch body as a Chat Completions body. + + Delegates to the same Responses-to-Chat bridge the real-time path uses + for providers without a native Responses API (which is every Bedrock + model), so `input`, `instructions`, `max_output_tokens` and the tool + params translate identically in batch and real time. The bridge always + emits a `tools` key; an empty one is dropped rather than shipped as an + empty array inside `modelInput`. + """ + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + responses_input = openai_request_body.get("input") + if responses_input is None: + raise ValueError( + "Batch record for /v1/responses is missing required `input` field: " + f"model={openai_request_body.get('model', '')}" + ) + chat_body = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=openai_request_body.get("model", ""), + input=_responses_input_adapter().validate_python(responses_input), + responses_api_request=_responses_request_adapter().validate_python( + _frozen_mapping( + (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") + ) + ), + ) + return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) + + @staticmethod + def _transform_batch_body_to_chat_body( + openai_request_body: Mapping[str, Any], + record_kind: BedrockBatchRecordKind, + ) -> Mapping[str, Any]: + """ + Normalize a non-embedding batch body to the Chat Completions shape the + per-provider Bedrock transformations expect. + """ + match record_kind: + case BedrockBatchRecordKind.TEXT_COMPLETION: + return BedrockFilesConfig._transform_text_completion_body_to_chat_body(openai_request_body) + case BedrockBatchRecordKind.RESPONSES: + return BedrockFilesConfig._transform_responses_body_to_chat_body(openai_request_body) + case BedrockBatchRecordKind.CHAT: + return openai_request_body + case BedrockBatchRecordKind.EMBEDDING: + raise ValueError("Embedding batch records do not have a chat-completion equivalent") + def _map_openai_to_bedrock_params( self, - openai_request_body: dict[str, Any], + openai_request_body: Mapping[str, Any], provider: str | None = None, ) -> dict[str, Any]: """ @@ -659,14 +775,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): provider = self.get_bedrock_invoke_provider(model) # Route to the embedding transformer when the OpenAI batch line - # targets /v1/embeddings; otherwise fall back to the existing - # chat-completion path. We branch here (rather than inside + # targets /v1/embeddings; every other endpoint shape is normalized + # to chat completions first. We branch here (rather than inside # `_map_openai_to_bedrock_params`) so the chat helper keeps its # narrow contract and the embedding helper can evolve independently. - if self._is_embedding_record(_openai_jsonl_content): + record_kind = self._classify_batch_record(_openai_jsonl_content) + if record_kind is BedrockBatchRecordKind.EMBEDDING: model_input = self._map_openai_embedding_to_bedrock_params(openai_request_body=openai_body) else: - model_input = self._map_openai_to_bedrock_params(openai_request_body=openai_body, provider=provider) + model_input = self._map_openai_to_bedrock_params( + openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), + provider=provider, + ) # Create Bedrock batch record record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index d9f8229dbed..f7a4682cb03 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from enum import Enum from typing import Any, Dict, List, Literal, Optional, Union from typing_extensions import TYPE_CHECKING, Required, TypedDict, override @@ -1100,3 +1101,17 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): # supported subset and strips the field entirely when nothing remains, so # other edit types (e.g. `clear_thinking_20251015`) never reach Bedrock. context_management: dict + + +class BedrockBatchRecordKind(Enum): + """ + Which OpenAI endpoint shape a line of a Bedrock managed-batch JSONL file + carries. Bedrock batch `modelInput` is always the model's InvokeModel / + Converse body, so every non-embedding shape is normalized to Chat + Completions before being handed to the per-provider transformation. + """ + + CHAT = "chat" + TEXT_COMPLETION = "text_completion" + RESPONSES = "responses" + EMBEDDING = "embedding" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1b1db634ed2..d10ccf77703 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -721,3 +721,46 @@ class TestUnpackLegacyDefs: out = unpack_legacy_defs(schema) assert "components" not in out assert out["properties"]["r0"]["properties"]["p0"] == {"type": "string"} + + +class TestTextCompletionPromptToMessages: + """`/v1/completions` prompt wrapping, shared by the real-time and batch paths.""" + + def test_string_prompt_becomes_single_user_message(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + assert text_completion_prompt_to_messages("summarize this") == ( + {"role": "user", "content": "summarize this"}, + ) + + def test_list_of_strings_becomes_one_message_each(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + assert text_completion_prompt_to_messages(["first", "second"]) == ( + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ) + + @pytest.mark.parametrize( + "prompt", + [ + [1, 2, 3], + [[1, 2], [3, 4]], + ["ok", 7], + [], + "", + None, + {"role": "user"}, + ], + ) + def test_unsupported_prompt_shapes_raise(self, prompt): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + with pytest.raises(ValueError, match="non-empty string or a non-empty list of strings"): + text_completion_prompt_to_messages(prompt) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index c548fe53e15..87b03b02e1a 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1072,21 +1072,34 @@ class TestBedrockFilesEmbeddingTransformation: is None ) - def test_is_embedding_record_helper(self): - """Helper detects embeddings via `url` first, then by body shape.""" + def test_classify_batch_record_helper(self): + """Helper classifies by `url` first, then by body shape.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind - assert BedrockFilesConfig._is_embedding_record( - {"url": "/v1/embeddings", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/embeddings", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.EMBEDDING ) # body-only fallback - assert BedrockFilesConfig._is_embedding_record({"body": {"input": "x"}}) - # chat shape - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/chat/completions", "body": {"messages": []}} + assert ( + BedrockFilesConfig._classify_batch_record({"body": {"input": "x"}}) + is BedrockBatchRecordKind.EMBEDDING + ) + # chat shape + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/chat/completions", "body": {"messages": []}} + ) + is BedrockBatchRecordKind.CHAT + ) + # ambiguous body without any recognized key is treated as chat + assert ( + BedrockFilesConfig._classify_batch_record({"body": {}}) + is BedrockBatchRecordKind.CHAT ) - # ambiguous body without `input` is treated as not-embedding - assert not BedrockFilesConfig._is_embedding_record({"body": {}}) def test_explicit_chat_url_with_input_body_short_circuits_to_chat(self): """Explicit url=/v1/chat/completions wins even if body looks like embedding. @@ -1097,15 +1110,20 @@ class TestBedrockFilesEmbeddingTransformation: """ from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind + # Direct helper assertion - assert not BedrockFilesConfig._is_embedding_record( - { - "url": "/v1/chat/completions", - "body": { - "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "input": "this would mis-route under the old precedence", - }, - } + assert ( + BedrockFilesConfig._classify_batch_record( + { + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "input": "this would mis-route under the old precedence", + }, + } + ) + is BedrockBatchRecordKind.CHAT ) # End-to-end: a record like this routes through the chat path. We @@ -1164,20 +1182,301 @@ class TestBedrockFilesEmbeddingTransformation: with pytest.raises(ValueError, match="must be a string"): BedrockFilesConfig._coerce_embedding_input_to_string({"unsupported": True}) - def test_other_non_embedding_urls_route_to_chat(self): - """Any non-/v1/embeddings url short-circuits to chat path.""" + def test_other_non_embedding_urls_do_not_route_to_embeddings(self): + """An `input` body only means "embedding" when the url says so.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind # /v1/completions (legacy completions endpoint) - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/completions", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/completions", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.TEXT_COMPLETION + ) + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/responses", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.RESPONSES ) # Arbitrary unknown url - caller's explicit signal still wins - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/responses", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/moderations", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.CHAT ) +class TestBedrockBatchNonChatEndpointRecords: + """`/v1/completions` and `/v1/responses` JSONL records (issue #35639). + + Bedrock batch `modelInput` is always the model's InvokeModel/Converse body, + so a record shaped for another OpenAI endpoint has to be normalized to chat + completions first. Before this normalization every record below either + raised `BadRequestError` at `POST /v1/files` (Anthropic, Nova) or silently + shipped an empty `messages` list to AWS (passthrough providers). + """ + + ANTHROPIC_MODEL = "bedrock/us.anthropic.claude-sonnet-4-6" + NOVA_MODEL = "bedrock/us.amazon.nova-pro-v1:0" + PASSTHROUGH_MODEL = "bedrock/openai.gpt-oss-120b-1:0" + + def _transform(self, record: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content([record]) + assert len(result) == 1 + assert result[0]["recordId"] == record["custom_id"] + return result[0]["modelInput"] + + def test_anthropic_text_completion_record_wraps_prompt(self): + model_input = self._transform( + { + "custom_id": "1", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "prompt": "Summarize the following call transcript", + "max_tokens": 64, + }, + } + ) + + assert model_input == { + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "Summarize the following call transcript"}], + } + ], + "max_tokens": 64, + "anthropic_version": "bedrock-2023-05-31", + } + + def test_anthropic_text_completion_record_keeps_every_prompt_in_a_list(self): + model_input = self._transform( + { + "custom_id": "2", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "prompt": ["first prompt", "second prompt"], + "max_tokens": 8, + }, + } + ) + + # Consecutive user messages are merged by the Anthropic transform, the + # same way they are on the real-time path. + assert model_input["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "first prompt"}, + {"type": "text", "text": "second prompt"}, + ], + } + ] + assert "prompt" not in model_input + + def test_anthropic_responses_record_wraps_string_input(self): + model_input = self._transform( + { + "custom_id": "3", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "input": "hi", + "max_output_tokens": 16, + }, + } + ) + + assert model_input == { + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + "max_tokens": 16, + "anthropic_version": "bedrock-2023-05-31", + } + assert "tools" not in model_input, "an empty tools array must not be shipped to Bedrock" + + def test_anthropic_responses_record_maps_instructions_and_input_items(self): + """The Responses-specific params go through the same bridge as real time.""" + model_input = self._transform( + { + "custom_id": "4", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "instructions": "be terse", + "input": [ + {"role": "user", "content": "what is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "and 3+3?"}, + ], + "max_output_tokens": 32, + "temperature": 0.2, + }, + } + ) + + assert model_input["system"] == [{"type": "text", "text": "be terse"}] + assert model_input["max_tokens"] == 32 + assert model_input["temperature"] == 0.2 + assert [message["role"] for message in model_input["messages"]] == [ + "user", + "assistant", + "user", + ] + assert model_input["messages"][-1]["content"] == [{"type": "text", "text": "and 3+3?"}] + assert "input" not in model_input + assert "max_output_tokens" not in model_input + + @pytest.mark.parametrize( + "body", + [ + {"prompt": "hi"}, + {"input": "hi"}, + ], + ids=["prompt", "input"], + ) + def test_nova_converse_record_wraps_prompt_and_input(self, body): + url = "/v1/completions" if "prompt" in body else "/v1/responses" + model_input = self._transform( + { + "custom_id": "5", + "method": "POST", + "url": url, + "body": {"model": self.NOVA_MODEL, **body}, + } + ) + + assert model_input["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + + @pytest.mark.parametrize( + "body", + [ + {"prompt": "hi"}, + {"input": "hi"}, + ], + ids=["prompt", "input"], + ) + def test_passthrough_provider_record_no_longer_emits_empty_messages(self, body): + """The passthrough branch used to emit `{"messages": [], "prompt": ...}`. + + That shape is accepted by `POST /v1/files`, so the whole batch job was + submitted to AWS and only failed there. + """ + url = "/v1/completions" if "prompt" in body else "/v1/responses" + model_input = self._transform( + { + "custom_id": "6", + "method": "POST", + "url": url, + "body": {"model": self.PASSTHROUGH_MODEL, **body}, + } + ) + + # Asserted on the serialized form, since the passthrough branch hands + # `messages` straight to S3 without a per-provider transform. + assert json.loads(json.dumps(model_input)) == {"messages": [{"role": "user", "content": "hi"}]} + + def test_mixed_endpoints_in_one_file_keep_their_own_shapes(self): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "chat", + "url": "/v1/chat/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4, + }, + }, + { + "custom_id": "text", + "url": "/v1/completions", + "body": {"model": self.ANTHROPIC_MODEL, "prompt": "hi", "max_tokens": 4}, + }, + { + "custom_id": "responses", + "url": "/v1/responses", + "body": {"model": self.ANTHROPIC_MODEL, "input": "hi", "max_output_tokens": 4}, + }, + { + "custom_id": "embedding", + "url": "/v1/embeddings", + "body": {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": "hi"}, + }, + ] + ) + + assert [record["recordId"] for record in result] == [ + "chat", + "text", + "responses", + "embedding", + ] + for record in result[:3]: + assert record["modelInput"]["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]} + ] + assert result[3]["modelInput"] == {"inputText": "hi"} + + @pytest.mark.parametrize( + ("url", "expected_message"), + [ + ("/v1/completions", "missing required `prompt` field"), + ("/v1/responses", "missing required `input` field"), + ], + ) + def test_missing_required_field_raises_actionable_error(self, url, expected_message): + with pytest.raises(ValueError, match=expected_message): + self._transform( + { + "custom_id": "7", + "method": "POST", + "url": url, + "body": {"model": self.ANTHROPIC_MODEL, "max_tokens": 4}, + } + ) + + def test_prompt_body_without_url_is_still_wrapped(self): + """A record can omit `url`; the body shape then decides.""" + model_input = self._transform( + { + "custom_id": "8", + "body": {"model": self.ANTHROPIC_MODEL, "prompt": "hi", "max_tokens": 4}, + } + ) + + assert model_input["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + + def test_messages_win_over_prompt_when_url_is_absent(self): + model_input = self._transform( + { + "custom_id": "9", + "body": { + "model": self.ANTHROPIC_MODEL, + "messages": [{"role": "user", "content": "from messages"}], + "prompt": "from prompt", + "max_tokens": 4, + }, + } + ) + + assert model_input["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "from messages"}]} + ] + + class TestBedrockFileContentTransformation: """SigV4-signed S3 GetObject retrieval of Bedrock batch output files.""" From 6def61e672d95c4b145613009cd3064d0a133475 Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 17:22:46 +0000 Subject: [PATCH 015/126] fix(bedrock): keep /v1/responses batch metadata through the chat bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/files/transformation.py | 1 + .../files/test_bedrock_files_transformation.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 0aa832780e5..baa2630556a 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -640,6 +640,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") ) ), + metadata=openai_request_body.get("metadata"), ) return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 87b03b02e1a..09aa6b1cf2d 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1337,6 +1337,23 @@ class TestBedrockBatchNonChatEndpointRecords: assert "input" not in model_input assert "max_output_tokens" not in model_input + def test_responses_record_keeps_metadata(self): + """`metadata` reaches the bridge, which reads it as its own kwarg.""" + model_input = self._transform( + { + "custom_id": "4b", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.PASSTHROUGH_MODEL, + "input": "hi", + "metadata": {"tenant": "acct-1"}, + }, + } + ) + + assert model_input["metadata"] == {"tenant": "acct-1"} + @pytest.mark.parametrize( "body", [ From 46751ad83ed3307fdc0e966d7e850660df2446d9 Mon Sep 17 00:00:00 2001 From: Darien Kindlund Date: Tue, 4 Aug 2026 16:17:01 -0400 Subject: [PATCH 016/126] fix(anthropic): coerce explicit additionalProperties to false in output_format schema Anthropic's structured outputs reject any `additionalProperties` value other than `false` ("output_format.schema: For 'object' type, 'additionalProperties: true' is not supported. Please set 'additionalProperties' to false") `filter_anthropic_output_schema` only added the key when it was absent, so an explicit `true` (or a sub-schema) was copied verbatim into output_format.schema and 400'd. Coerce it for object schemas instead, at every recursion depth, matching what the Anthropic Python/TypeScript SDKs do The permissive tool-use path (map_response_format_to_anthropic_tool, used for vertex_ai) is deliberately left alone Fixes #35808 --- litellm/llms/anthropic/chat/transformation.py | 2 +- .../anthropic/test_anthropic_schema_filter.py | 72 ++++++++++++++++++- 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 51b862e79d9..19f5174579b 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -597,7 +597,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Anthropic requires additionalProperties=false for object schemas # See: https://docs.anthropic.com/en/docs/build-with-claude/structured-outputs - if result.get("type") == "object" and "additionalProperties" not in result: + if result.get("type") == "object": result["additionalProperties"] = False return result diff --git a/tests/litellm/llms/anthropic/test_anthropic_schema_filter.py b/tests/litellm/llms/anthropic/test_anthropic_schema_filter.py index c10ac5532a0..71c9cfe8f41 100644 --- a/tests/litellm/llms/anthropic/test_anthropic_schema_filter.py +++ b/tests/litellm/llms/anthropic/test_anthropic_schema_filter.py @@ -281,7 +281,10 @@ class TestFilterAnthropicOutputSchema: "unevaluatedProperties", ): assert field not in result - assert 'properties whose names match each pattern must satisfy: {"^x": {"type": "string"}}' in result["description"] + assert ( + 'properties whose names match each pattern must satisfy: {"^x": {"type": "string"}}' + in result["description"] + ) assert 'property names must satisfy: {"pattern": "^[a-z]+$"}' in result["description"] assert 'dependent required properties: {"first": ["last"]}' in result["description"] assert 'dependent schemas: {"first": {"required": ["last"]}}' in result["description"] @@ -347,3 +350,70 @@ class TestFilterAnthropicOutputSchema: "all array items must be unique, minimum number of matching items: 2, " "maximum number of matching items: 3." ) + + def test_coerces_explicit_additional_properties_true(self): + """An explicit ``additionalProperties: true`` must be coerced to false. + + Anthropic rejects anything other than false with: + "output_format.schema: For 'object' type, 'additionalProperties: true' is + not supported". + """ + schema = { + "type": "object", + "additionalProperties": True, + "properties": {"a": {"type": "string"}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["additionalProperties"] is False + + def test_coerces_additional_properties_true_when_nested(self): + """Nested object schemas are coerced too, at every recursion site.""" + schema = { + "type": "object", + "properties": { + "obj": { + "type": "object", + "additionalProperties": True, + "properties": {"a": {"type": "string"}}, + }, + "rows": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": True, + "properties": {"b": {"type": "string"}}, + }, + }, + }, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["properties"]["obj"]["additionalProperties"] is False + assert result["properties"]["rows"]["items"]["additionalProperties"] is False + + def test_coerces_additional_properties_sub_schema(self): + """A sub-schema value (free-form map) is also rejected by Anthropic.""" + schema = { + "type": "object", + "additionalProperties": {"type": "string"}, + "properties": {"a": {"type": "string"}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["additionalProperties"] is False + + def test_explicit_additional_properties_false_is_preserved(self): + """The already-correct value must survive untouched.""" + schema = { + "type": "object", + "additionalProperties": False, + "properties": {"a": {"type": "string"}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["additionalProperties"] is False From 3d673f9534f961c7f709b0a70063f349ab7cfd2c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:58:20 +0000 Subject: [PATCH 017/126] fix(managed_files): skip unparseable rows when listing managed files get_user_created_file_ids validated every row's file_object without a guard, so a single row failing OpenAIFileObject validation raised ValidationError and turned the whole GET /v1/files response into a 500. #35365 covered the null case only, leaving malformed or partial rows able to take the entire listing down. Rows now parse through a helper that returns None on failure and logs a warning, matching how list_user_batches already tolerates rows it cannot parse, so one bad row costs its own entry instead of the caller's whole listing. Null rows stay silent since the batch cost poller registers those legitimately. Refs #35361 --- .../proxy/hooks/managed_files.py | 23 +++++++++++++++++-- .../proxy/test_managed_files_hook.py | 23 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index ec47b6ac0e6..2349b618a28 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -73,6 +73,20 @@ else: PrismaClient = Any +def _parse_managed_file_object( + raw_file_object: object, unified_file_id: str +) -> Optional[OpenAIFileObject]: + if raw_file_object is None: + return None + try: + return OpenAIFileObject.model_validate(raw_file_object) + except Exception as e: + verbose_logger.warning( + f"Failed to parse managed file object {unified_file_id}: {e}" + ) + return None + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -383,9 +397,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ) return [ - OpenAIFileObject.model_validate(file_object.file_object) + parsed_file_object for file_object in file_ids - if file_object.file_object is not None + if ( + parsed_file_object := _parse_managed_file_object( + file_object.file_object, file_object.unified_file_id + ) + ) + is not None ] async def check_managed_file_id_access( diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 4a4aa7aa5ea..4da6de6353f 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -154,6 +154,29 @@ async def test_get_user_created_file_ids_skips_rows_without_file_object(): assert [file.id for file in files] == ["file-output-abc"] +@pytest.mark.asyncio +async def test_get_user_created_file_ids_skips_unparseable_rows(): + managed_files = _make_managed_files_instance() + managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[ + MagicMock( + file_object={"id": "file-corrupt", "object": "file"}, + unified_file_id="unified-corrupt", + ), + MagicMock( + file_object=_make_file_object().model_dump(), + unified_file_id="unified-valid", + ), + ] + ) + + files = await managed_files.get_user_created_file_ids( + _make_user_api_key_dict(), ["file-output-abc"] + ) + + assert [file.id for file in files] == ["file-output-abc"] + + @pytest.mark.asyncio async def test_should_fallback_when_no_router(): """ From 1b6f3cebf1a4a4804a9bd9a0c3287cfc0d07c971 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:16:05 +0000 Subject: [PATCH 018/126] fix(managed_files): log sanitized validation errors when skipping rows The skip warning interpolated the full pydantic ValidationError, whose string embeds input_value with the rejected row's contents. Managed-file rows carry a caller-supplied filename, so a malformed row copied that into operational logs. Log the error locations, types, and messages via errors() with input, url, and context excluded, keeping the field-level diagnostics without the values. Non-validation failures fall back to the exception type. --- .../proxy/hooks/managed_files.py | 9 ++++++++- .../proxy/test_managed_files_hook.py | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 2349b618a28..688ffb35ff7 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -7,6 +7,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast from fastapi import HTTPException +from pydantic import ValidationError import litellm from litellm import Router, verbose_logger @@ -80,9 +81,15 @@ def _parse_managed_file_object( return None try: return OpenAIFileObject.model_validate(raw_file_object) + except ValidationError as e: + verbose_logger.warning( + f"Failed to parse managed file object {unified_file_id}: " + f"{e.errors(include_input=False, include_url=False, include_context=False)}" + ) + return None except Exception as e: verbose_logger.warning( - f"Failed to parse managed file object {unified_file_id}: {e}" + f"Failed to parse managed file object {unified_file_id}: {type(e).__name__}" ) return None diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 4da6de6353f..6397e0be247 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -6,6 +6,7 @@ async_post_call_success_hook when processing completed batch responses. """ import json +import logging import pytest from typing import Optional @@ -154,6 +155,24 @@ async def test_get_user_created_file_ids_skips_rows_without_file_object(): assert [file.id for file in files] == ["file-output-abc"] +@pytest.mark.asyncio +async def test_parse_managed_file_object_warning_omits_rejected_values(caplog): + from litellm_enterprise.proxy.hooks.managed_files import ( + _parse_managed_file_object, + ) + + with caplog.at_level(logging.WARNING): + parsed = _parse_managed_file_object( + {"id": "file-corrupt", "object": "file", "filename": "confidential.jsonl"}, + "unified-corrupt", + ) + + assert parsed is None + assert "unified-corrupt" in caplog.text + assert "bytes" in caplog.text + assert "confidential.jsonl" not in caplog.text + + @pytest.mark.asyncio async def test_get_user_created_file_ids_skips_unparseable_rows(): managed_files = _make_managed_files_instance() From eef908d4ad542e8003f22d76dd12ad559f25733d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:26:21 -0700 Subject: [PATCH 019/126] fix(batches): register managed output files on batch cancel update_batch_in_database now fetches the batch row by unified_object_id when the caller omits db_batch_object, so the cancel endpoint attributes newly registered output and error files to the batch owner and returns unified managed ids instead of raw provider ids. Idempotent cancels that do not change the stored status also skip the redundant DB write now. Repair two pre-existing mock tests in test_openai_batches_endpoint.py that asserted values inside lazy percent-format log strings, and give the cancel test's prisma mock an awaitable find_first. --- .../openai_files_endpoints/common_utils.py | 17 +++-- .../test_openai_batches_endpoint.py | 7 +- ..._batch_update_db_managed_output_file_id.py | 64 ++++++++++++++++++- 3 files changed, 80 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 37b51e1d3af..290045a5a87 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1141,7 +1141,7 @@ async def update_batch_in_database( managed_files_obj: The managed_files proxy hook object prisma_client: Prisma database client verbose_proxy_logger: Logger instance - db_batch_object: Optional existing database object (for comparison) + db_batch_object: Optional existing database object; fetched by unified_object_id when omitted operation: Description of operation ("update", "cancel", etc.) user_api_key_dict: Optional auth context for creating managed file IDs """ @@ -1154,6 +1154,12 @@ async def update_batch_in_database( if not prisma_client: return + effective_db_batch_object: Final = ( + db_batch_object + if db_batch_object is not None + else await ManagedObjectRepository(prisma_client).table.find_first(where={"unified_object_id": batch_id}) + ) + # Always normalize the response's file IDs to unified managed IDs # (mutates in place) so the caller returns unified IDs to the user # even when we skip the DB update below for an unchanged status. @@ -1163,16 +1169,17 @@ async def update_batch_in_database( prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, user_api_key_dict=user_api_key_dict, - db_batch_object=db_batch_object, + db_batch_object=effective_db_batch_object, + unified_batch_id=unified_batch_id, ) # Only update if status has changed (when db_batch_object is provided) - if db_batch_object and response.status == db_batch_object.status: + if effective_db_batch_object and response.status == effective_db_batch_object.status: return - if db_batch_object: + if effective_db_batch_object: verbose_proxy_logger.info( - "Updating batch %s status from %s to %s", batch_id, db_batch_object.status, response.status + "Updating batch %s status from %s to %s", batch_id, effective_db_batch_object.status, response.status ) else: verbose_proxy_logger.info("Updating batch %s status to %s after %s", batch_id, response.status, operation) diff --git a/tests/openai_endpoints_tests/test_openai_batches_endpoint.py b/tests/openai_endpoints_tests/test_openai_batches_endpoint.py index c6f4128f2c5..db8f75cf640 100644 --- a/tests/openai_endpoints_tests/test_openai_batches_endpoint.py +++ b/tests/openai_endpoints_tests/test_openai_batches_endpoint.py @@ -414,7 +414,7 @@ async def test_batch_status_sync_from_provider_to_database(): # Verify logger was called with status change message mock_logger.info.assert_called() - log_message = mock_logger.info.call_args[0][0] + log_message = mock_logger.info.call_args[0][0] % mock_logger.info.call_args[0][1:] assert "validating" in log_message assert "completed" in log_message @@ -450,6 +450,9 @@ async def test_batch_cancel_updates_database(): # Mock prisma client mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( + return_value=None + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() # Mock managed_files_obj @@ -482,7 +485,7 @@ async def test_batch_cancel_updates_database(): # Verify logger was called mock_logger.info.assert_called() - log_message = mock_logger.info.call_args[0][0] + log_message = mock_logger.info.call_args[0][0] % mock_logger.info.call_args[0][1:] assert "cancel" in log_message.lower() assert "cancelled" in log_message diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index d8669960674..74139fa9238 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -45,9 +45,10 @@ def _build_managed_files_mock(unified_id: str = "file-bWFuYWdlZF9vdXRwdXRfaWQ=") return mock -def _build_prisma_mock(): +def _build_prisma_mock(db_batch_object=None): mock = MagicMock() mock.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) + mock.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=db_batch_object) mock.db.litellm_managedobjecttable.update = AsyncMock() return mock @@ -89,6 +90,67 @@ async def test_update_batch_in_database_stores_unified_output_file_id(): assert stored["output_file_id"] != raw_output_file_id +@pytest.mark.asyncio +async def test_cancel_path_registers_output_file_under_batch_owner(): + unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" + db_batch_object = SimpleNamespace( + created_by="batch-owner", team_id="batch-team", status="in_progress" + ) + response = _build_batch_response( + status="cancelling", + output_file_id="file-raw-output", + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + mock_managed_files = _build_managed_files_mock(unified_id=unified_id) + mock_prisma = _build_prisma_mock(db_batch_object=db_batch_object) + + await update_batch_in_database( + batch_id="batch_managed_ids_test", + unified_batch_id="litellm_proxy;model_id:my-model;llm_batch_id:batch_managed_ids_test", + response=response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + operation="cancel", + ) + + forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[ + "user_api_key_dict" + ] + assert forwarded_auth.user_id == "batch-owner" + assert forwarded_auth.team_id == "batch-team" + stored = json.loads( + mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs["data"][ + "file_object" + ] + ) + assert stored["output_file_id"] == unified_id + + +@pytest.mark.asyncio +async def test_update_batch_derives_model_id_from_unified_batch_id(): + unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" + response = _build_batch_response(output_file_id="file-raw-output", hidden_params={}) + mock_managed_files = _build_managed_files_mock(unified_id=unified_id) + mock_prisma = _build_prisma_mock() + + await update_batch_in_database( + batch_id="batch_managed_ids_test", + unified_batch_id="litellm_proxy;model_id:model-from-batch-id;llm_batch_id:batch_managed_ids_test", + response=response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert ( + mock_managed_files.get_unified_output_file_id.call_args.kwargs["model_id"] + == "model-from-batch-id" + ) + assert response.output_file_id == unified_id + + @pytest.mark.asyncio async def test_ensure_batch_response_normalizes_error_file_id(): """Both output_file_id and error_file_id must be normalized to managed IDs.""" From 6b5c7f92ce7b359ebb09ce69d9837bbaae1e3209 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:34:06 +0000 Subject: [PATCH 020/126] refactor(types): replace implicit and explicit Any across 11 modules Types the values that were flowing through as Any in the highest-density modules, using shapes the code already assumes: TypedDicts for the JSON payloads read by literal key, Protocols for the prisma rows, existing litellm types where they were already modeled, and `object` where a value is only stored and forwarded. Annotation-level only, no runtime behavior change. New annotations use read-only views (Mapping / Sequence / tuple) rather than dict / list, so LIT001 drops alongside the Any counts instead of trading one budget for another. No suppressions, casts, or type guards were added. basedpyright across the touched files: 1547 -> 856 errors, with reportAny down 399 and reportExplicitAny down 134, and no rule increasing. --- litellm/caching/redis_semantic_cache.py | 33 +++--- litellm/integrations/galileo.py | 83 +++++++++------ .../mcp_server/openapi_to_mcp_generator.py | 100 +++++++++++++----- .../claude_code_marketplace.py | 80 ++++++++++---- .../proxy/common_utils/custom_openapi_spec.py | 42 +++++--- .../proxy/container_endpoints/ownership.py | 80 ++++++++++---- .../cato_networks/cato_networks.py | 88 +++++++++++---- .../hiddenlayer/hiddenlayer.py | 45 ++++++-- .../guardrail_hooks/microsoft_purview/base.py | 61 ++++++----- .../workflow_management_endpoints.py | 98 ++++++++++++----- .../policy_engine/attachment_registry.py | 39 ++++--- 11 files changed, 516 insertions(+), 233 deletions(-) diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index b1d298b79bb..604d6395ea1 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -13,6 +13,7 @@ import ast import asyncio import json import os +from collections.abc import Callable, Mapping from typing import Any, Final, cast import litellm @@ -47,7 +48,7 @@ class RedisSemanticCache(BaseCache): similarity_threshold: float | None = None, embedding_model: str = "text-embedding-ada-002", index_name: str | None = None, - **kwargs, + **kwargs: object, ): """ Initialize the Redis Semantic Cache. @@ -150,11 +151,11 @@ class RedisSemanticCache(BaseCache): def _init_semantic_cache( self, - semantic_cache_cls: Any, + semantic_cache_cls: Callable[..., object], index_name: str, redis_url: str, - cache_vectorizer: Any, - ) -> Any: + cache_vectorizer: object, + ) -> object: def _is_schema_mismatch(exc: ValueError) -> bool: error_message: Final = str(exc).lower() return any(phrase in error_message for phrase in ("schema does not match", "index schema")) @@ -206,12 +207,12 @@ class RedisSemanticCache(BaseCache): def _get_cache_filters(self, key: str) -> dict[str, str]: return {self.CACHE_KEY_FIELD_NAME: str(key)} - def _get_cache_key_filter_expression(self, key: str) -> Any: + def _get_cache_key_filter_expression(self, key: str) -> object: from redisvl.query.filter import Tag return Tag(self.CACHE_KEY_FIELD_NAME) == str(key) - def _cache_hit_matches_key(self, cache_hit: dict[str, Any], key: str) -> bool: + def _cache_hit_matches_key(self, cache_hit: Mapping[str, object], key: str) -> bool: # Pre-isolation entries with no ``litellm_cache_key`` field cannot be # safely reassigned to a caller's scope and are treated as misses. cached_key = cache_hit.get(self.CACHE_KEY_FIELD_NAME) @@ -297,7 +298,7 @@ class RedisSemanticCache(BaseCache): return @staticmethod - def _coerce_response_input_value(value: Any) -> Any: + def _coerce_response_input_value(value: object) -> object: model_dump: Final = getattr(value, "model_dump", None) if callable(model_dump): return model_dump() @@ -340,7 +341,7 @@ class RedisSemanticCache(BaseCache): ) return embedding_response["data"][0]["embedding"] - def _get_cache_logic(self, cached_response: Any) -> Any: + def _get_cache_logic(self, cached_response: Any) -> object: """ Process the cached response to prepare it for use. @@ -369,7 +370,7 @@ class RedisSemanticCache(BaseCache): return cached_response - def set_cache(self, key: str, value: Any, **kwargs) -> None: + def set_cache(self, key: str, value: object, **kwargs) -> None: """ Store a value in the semantic cache. @@ -405,7 +406,7 @@ class RedisSemanticCache(BaseCache): except Exception as e: print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {e}") - def get_cache(self, key: str, **kwargs) -> Any: + def get_cache(self, key: str, **kwargs) -> object: """ Retrieve a semantically similar cached response. @@ -428,7 +429,7 @@ class RedisSemanticCache(BaseCache): # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. prompt_embedding: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata")) - check_kwargs: Final[dict[str, Any]] = { + check_kwargs: Final[Mapping[str, object]] = { "prompt": prompt, "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), @@ -508,7 +509,7 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Error generating async embedding: {e}") raise ValueError(f"Failed to generate embedding: {e}") from e - async def async_set_cache(self, key: str, value: Any, **kwargs) -> None: + async def async_set_cache(self, key: str, value: object, **kwargs) -> None: """ Asynchronously store a value in the semantic cache. @@ -548,7 +549,7 @@ class RedisSemanticCache(BaseCache): except Exception as e: print_verbose(f"Error in async_set_cache: {e}") - async def async_get_cache(self, key: str, **kwargs) -> Any: + async def async_get_cache(self, key: str, **kwargs) -> object: """ Asynchronously retrieve a semantically similar cached response. @@ -573,7 +574,7 @@ class RedisSemanticCache(BaseCache): # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - check_kwargs: Final[dict[str, Any]] = { + check_kwargs: Final[Mapping[str, object]] = { "prompt": prompt, "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), @@ -615,7 +616,7 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Error in async_get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def _index_info(self) -> dict[str, Any]: + async def _index_info(self) -> Mapping[str, object]: """ Get information about the Redis index. @@ -625,7 +626,7 @@ class RedisSemanticCache(BaseCache): aindex: Final = await self.llmcache._get_async_index() return await aindex.info() - async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs) -> None: + async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: object) -> None: """ Asynchronously store multiple values in the semantic cache. diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 12c2ac8a53f..f9ec825e922 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -4,7 +4,8 @@ import json import os import re import uuid -from datetime import datetime, timezone +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone, tzinfo from typing import Any, Final, cast import httpx @@ -59,7 +60,7 @@ class LLMResponse(BaseModel): class GalileoObserve(CustomLogger): def __init__(self) -> None: - self.in_memory_records: list[dict] = [] + self.in_memory_records: list[dict[str, Any]] = [] self.batch_size = 1 self.api_key = os.getenv("GALILEO_API_KEY") self.project_id = os.getenv("GALILEO_PROJECT_ID") @@ -176,7 +177,7 @@ class GalileoObserve(CustomLogger): return False @staticmethod - def _galileo_input_messages(messages: Any | None, input_text: str) -> list[dict[str, str]]: + def _galileo_input_messages(messages: object, input_text: str) -> list[dict[str, str]]: if isinstance(messages, dict): messages = messages.get("messages") if not messages: @@ -203,11 +204,11 @@ class GalileoObserve(CustomLogger): return [{"role": "user", "content": input_text}] @staticmethod - def _local_timezone(): + def _local_timezone() -> tzinfo: return datetime.now().astimezone().tzinfo or timezone.utc @staticmethod - def _format_created_at(dt: datetime | Any) -> str: + def _format_created_at(dt: object) -> str: """Serialize timestamps as UTC ISO-8601 for Galileo.""" if not isinstance(dt, datetime): return str(dt) @@ -226,7 +227,7 @@ class GalileoObserve(CustomLogger): return created_at @staticmethod - def _token_metrics_from_record(record: dict[str, Any]) -> dict[str, Any]: + def _token_metrics_from_record(record: Mapping[str, Any]) -> dict[str, Any]: num_input_tokens: Final = int(record.get("num_input_tokens") or 0) num_output_tokens: Final = int(record.get("num_output_tokens") or 0) num_total_tokens = int(record.get("num_total_tokens") or 0) @@ -244,7 +245,7 @@ class GalileoObserve(CustomLogger): @staticmethod def _record_to_v2_span( - record: dict[str, Any], + record: Mapping[str, Any], *, trace_id: str, span_id: str, @@ -275,7 +276,7 @@ class GalileoObserve(CustomLogger): return span @staticmethod - def _record_to_v2_trace(record: dict[str, Any]) -> dict[str, Any]: + def _record_to_v2_trace(record: Mapping[str, Any]) -> dict[str, Any]: trace_id: Final = str(uuid.uuid4()) span_id: Final = str(uuid.uuid4()) created_at: Final = GalileoObserve._normalize_created_at(record.get("created_at", "")) @@ -295,7 +296,7 @@ class GalileoObserve(CustomLogger): "spans": [GalileoObserve._record_to_v2_span(record, trace_id=trace_id, span_id=span_id)], } - def _build_traces_payload(self, records: list[dict]) -> dict[str, Any]: + def _build_traces_payload(self, records: Sequence[Mapping[str, Any]]) -> dict[str, Any]: payload: Final[dict[str, Any]] = { "traces": [self._record_to_v2_trace(record) for record in records], "logging_method": "api_direct", @@ -357,7 +358,7 @@ class GalileoObserve(CustomLogger): @staticmethod def _log_v2_payload_validation(payload: dict[str, Any]) -> None: missing_fields: Final[list[str]] = [] - traces: Final = payload.get("traces", []) + traces: Final[Sequence[object]] = payload.get("traces", []) if not traces: missing_fields.append("traces") @@ -385,7 +386,7 @@ class GalileoObserve(CustomLogger): ) def _log_flush_payload(self, url: str, payload: dict[str, Any]) -> None: - traces: Final = payload.get("traces", []) + traces: Final[Sequence[object]] = payload.get("traces", []) verbose_logger.debug( "Galileo Logger flush URL: %s trace_count=%s", url, @@ -415,8 +416,8 @@ class GalileoObserve(CustomLogger): pass @staticmethod - def _build_prompt(kwargs: dict[str, Any]) -> dict[str, Any]: - optional_params: Final = kwargs.get("optional_params", {}) or {} + def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, Any]: + optional_params: Final[Mapping[str, object]] = kwargs.get("optional_params", {}) or {} prompt: Final[dict[str, Any]] = {"messages": kwargs.get("messages")} if optional_params.get("functions") is not None: prompt["functions"] = optional_params["functions"] @@ -425,13 +426,13 @@ class GalileoObserve(CustomLogger): return prompt @staticmethod - def _serialize_galileo_output(value: Any) -> str: + def _serialize_galileo_output(value: object) -> str: if value is None: return "" if isinstance(value, str): return value - def _json_default(obj: Any) -> Any: + def _json_default(obj: Any) -> object: if hasattr(obj, "model_dump"): return obj.model_dump() return str(obj) @@ -439,8 +440,8 @@ class GalileoObserve(CustomLogger): return json.dumps(value, default=_json_default) @staticmethod - def _prompt_to_input_text(prompt: dict[str, Any]) -> str: - messages: Final = prompt.get("messages") + def _prompt_to_input_text(prompt: Mapping[str, Any]) -> str: + messages: Final[object] = prompt.get("messages") if messages is not None: text: Final = GalileoObserve._input_text_from_messages(messages) if text: @@ -448,7 +449,7 @@ class GalileoObserve(CustomLogger): return json.dumps(prompt, default=str) @staticmethod - def _get_chat_content_for_galileo(response_obj: litellm.ModelResponse) -> Any: + def _get_chat_content_for_galileo(response_obj: litellm.ModelResponse) -> object: if response_obj.choices and len(response_obj.choices) > 0: message: Final = response_obj["choices"][0]["message"] if hasattr(message, "json"): @@ -470,23 +471,23 @@ class GalileoObserve(CustomLogger): @staticmethod def _get_responses_api_content_for_galileo( response_obj: ResponsesAPIResponse, - ) -> Any: + ) -> object: if hasattr(response_obj, "output") and response_obj.output: return response_obj.output return None @staticmethod - def _langfuse_style_rerank_prompt(kwargs: dict[str, Any]) -> dict[str, Any]: + def _langfuse_style_rerank_prompt(kwargs: Mapping[str, object]) -> dict[str, Any]: """Match Langfuse rerank input: prompt = {"messages": kwargs.get("messages")}.""" return {"messages": kwargs.get("messages")} def _get_galileo_input_output_content( self, - kwargs: dict[str, Any], - response_obj: Any, + kwargs: Mapping[str, object], + response_obj: object, level: str = "DEFAULT", status_message: str | None = None, - ) -> tuple[str, str, Any]: + ) -> tuple[str, str, object]: """ Mirror Langfuse _get_langfuse_input_output_content for Galileo ingest. @@ -582,12 +583,12 @@ class GalileoObserve(CustomLogger): return self._prompt_to_input_text(prompt), "", kwargs.get("messages") or [] - def get_output_str_from_response(self, response_obj: Any, kwargs: dict[str, Any]) -> str: + def get_output_str_from_response(self, response_obj: object, kwargs: Mapping[str, object]) -> str: _, output_text, _ = self._get_galileo_input_output_content(kwargs=kwargs, response_obj=response_obj) return output_text @staticmethod - def _input_text_from_messages(messages: Any) -> str: + def _input_text_from_messages(messages: object) -> str: """Return a plain-string summary of the input suitable for the trace-level input field.""" if isinstance(messages, str): return messages @@ -613,7 +614,13 @@ class GalileoObserve(CustomLogger): return str(content) return "" - async def async_log_success_event(self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any): + async def async_log_success_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: verbose_logger.debug("On Async Success") try: await self._async_log_success_event_impl( @@ -625,7 +632,13 @@ class GalileoObserve(CustomLogger): except Exception: verbose_logger.exception("Galileo Logger: unexpected error in async_log_success_event") - async def _async_log_success_event_impl(self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any): + async def _async_log_success_event_impl( + self, + kwargs: Mapping[str, Any], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: if not self._is_configured(): verbose_logger.debug( "Galileo Logger: skipping — GALILEO_PROJECT_ID=%s GALILEO_API_KEY=%s GALILEO_BASE_URL=%s", @@ -635,7 +648,7 @@ class GalileoObserve(CustomLogger): ) return - slo: Final[dict[str, Any] | None] = kwargs.get("standard_logging_object") + slo: Final[Mapping[str, Any] | None] = kwargs.get("standard_logging_object") if slo is None: verbose_logger.debug("Galileo Logger: no standard_logging_object in kwargs, skipping") return @@ -646,8 +659,8 @@ class GalileoObserve(CustomLogger): kwargs=kwargs, response_obj=response_obj ) - raw_start: Final = slo.get("startTime") - raw_end: Final = slo.get("endTime") + raw_start: Final[float | None] = slo.get("startTime") + raw_end: Final[float | None] = slo.get("endTime") if raw_start is None or raw_end is None: verbose_logger.debug( "Galileo Logger: standard_logging_object missing startTime/endTime, " @@ -710,7 +723,7 @@ class GalileoObserve(CustomLogger): if len(self.in_memory_records) >= self.batch_size: await self.flush_in_memory_records() - async def flush_in_memory_records(self): + async def flush_in_memory_records(self) -> None: if not self.in_memory_records: return @@ -774,5 +787,11 @@ class GalileoObserve(CustomLogger): if not self.use_v2_api and response.status_code in (401, 403): self.headers = None - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: verbose_logger.debug("On Async Failure") diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 1794a1b66b8..41057840684 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -7,10 +7,13 @@ import contextvars import json import os import re +from collections.abc import Mapping, Sequence from pathlib import PurePosixPath -from typing import Any, Final +from typing import Any, Final, TypeAlias, TypedDict from urllib.parse import quote +import httpx + # Tool names emitted from OpenAPI specs must work across all major LLM providers. # OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to # ^[a-zA-Z0-9_-]+$ on tool names. Many specs (notably GitHub's REST API) use @@ -44,6 +47,43 @@ from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) +_OpenAPIObject: TypeAlias = Mapping[str, Any] + + +class _OpenAPIParameterSchema(TypedDict, total=False): + type: str + + +class _OpenAPIJSONSchema(TypedDict, total=False): + properties: Mapping[str, object] + + +class _OpenAPIMediaType(TypedDict, total=False): + schema: _OpenAPIJSONSchema + + +class _OpenAPIRequestBody(TypedDict, total=False): + description: str + required: bool + content: Mapping[str, _OpenAPIMediaType] + + +class _OpenAPIOperation(TypedDict, total=False): + operationId: str + summary: str + description: str + parameters: Sequence[_OpenAPIObject] + requestBody: _OpenAPIRequestBody + + +class _OpenAPIPathItem(TypedDict, total=False): + parameters: Sequence[_OpenAPIObject] + + +class _OpenAPIComponents(TypedDict, total=False): + parameters: Mapping[str, _OpenAPIObject] + + # Store the base URL and headers globally BASE_URL: Final = "" HEADERS: Final[dict[str, str]] = {} @@ -69,7 +109,7 @@ _request_resolved_auth_headers: Final[contextvars.ContextVar[dict[str, str] | No ) -def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: +def _sanitize_path_parameter_value(param_value: object, param_name: str) -> str: """Ensure path params cannot introduce directory traversal.""" if param_value is None: return "" @@ -109,7 +149,7 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]: async def load_openapi_spec_async(filepath: str) -> dict[str, Any]: if filepath.startswith("http://") or filepath.startswith("https://"): client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) - r: Final = await async_safe_get(client, filepath) + r: Final[httpx.Response] = await async_safe_get(client, filepath) r.raise_for_status() return r.json() @@ -121,11 +161,11 @@ async def load_openapi_spec_async(filepath: str) -> dict[str, Any]: return json.load(f) -def get_base_url(spec: dict[str, Any], spec_path: str | None = None) -> str: +def get_base_url(spec: _OpenAPIObject, spec_path: str | None = None) -> str: """Extract base URL from OpenAPI spec.""" # OpenAPI 3.x if "servers" in spec and spec["servers"]: - server_url: Final = spec["servers"][0]["url"] + server_url: Final[str] = spec["servers"][0]["url"] # If the server URL is relative (starts with /), derive base from spec_path if server_url.startswith("/") and spec_path: @@ -147,8 +187,8 @@ def get_base_url(spec: dict[str, Any], spec_path: str | None = None) -> str: return server_url # OpenAPI 2.x (Swagger) elif "host" in spec: - scheme: Final = spec.get("schemes", ["https"])[0] - base_path: Final = spec.get("basePath", "") + scheme: Final[str] = spec.get("schemes", ["https"])[0] + base_path: Final[str] = spec.get("basePath", "") return f"{scheme}://{spec['host']}{base_path}" # Fallback: derive base URL from spec_path if it's a URL @@ -172,20 +212,22 @@ def get_base_url(spec: dict[str, Any], spec_path: str | None = None) -> str: return "" -def _resolve_ref(param: dict[str, Any], component_params: dict[str, Any]) -> dict[str, Any] | None: +def _resolve_ref(param: _OpenAPIObject, component_params: Mapping[str, _OpenAPIObject]) -> _OpenAPIObject | None: """Resolve a single parameter, following a $ref if present. Returns the resolved param dict, or None if the $ref target is absent from components (so callers can skip/filter it rather than propagating a stub with name=None that would corrupt deduplication). """ - ref: Final = param.get("$ref", "") + ref: Final[str] = param.get("$ref", "") if not ref.startswith("#/components/parameters/"): return param return component_params.get(ref.split("/")[-1]) -def _resolve_param_list(raw: list[dict[str, Any]], component_params: dict[str, Any]) -> list[dict[str, Any]]: +def _resolve_param_list( + raw: Sequence[_OpenAPIObject], component_params: Mapping[str, _OpenAPIObject] +) -> list[_OpenAPIObject]: """Resolve $refs in a parameter list, dropping any unresolvable entries.""" result: Final = [] for p in raw: @@ -196,9 +238,9 @@ def _resolve_param_list(raw: list[dict[str, Any]], component_params: dict[str, A def resolve_operation_params( - operation: dict[str, Any], - path_item: dict[str, Any], - components: dict[str, Any], + operation: _OpenAPIOperation, + path_item: _OpenAPIPathItem, + components: _OpenAPIComponents, ) -> dict[str, Any]: """Return a copy of *operation* with fully-resolved, merged parameters. @@ -214,7 +256,7 @@ def resolve_operation_params( merged with the operation-level params; operation-level wins when the same ``name`` + ``in`` combination appears in both. """ - component_params: Final = components.get("parameters", {}) + component_params: Final[Mapping[str, _OpenAPIObject]] = components.get("parameters", {}) path_level: Final = _resolve_param_list(path_item.get("parameters", []), component_params) op_level: Final = _resolve_param_list(operation.get("parameters", []), component_params) op_keys: Final = {(p["name"], p.get("in")) for p in op_level} @@ -224,8 +266,9 @@ def resolve_operation_params( return result -def extract_parameters(operation: dict[str, Any]) -> tuple: +def extract_parameters(operation: _OpenAPIObject) -> tuple[Sequence[str], Sequence[str], Sequence[str]]: """Extract parameter names from OpenAPI operation.""" + param: _OpenAPIObject path_params: Final = [] query_params: Final = [] body_params: Final = [] @@ -235,7 +278,7 @@ def extract_parameters(operation: dict[str, Any]) -> tuple: for param in operation["parameters"]: if "name" not in param: continue - param_name = param["name"] + param_name: str = param["name"] if param.get("in") == "path": path_params.append(param_name) elif param.get("in") == "query": @@ -250,8 +293,9 @@ def extract_parameters(operation: dict[str, Any]) -> tuple: return path_params, query_params, body_params -def build_input_schema(operation: dict[str, Any]) -> dict[str, Any]: +def build_input_schema(operation: _OpenAPIObject) -> dict[str, Any]: """Build MCP input schema from OpenAPI operation.""" + param: _OpenAPIObject properties: Final = {} required: Final = [] @@ -260,9 +304,9 @@ def build_input_schema(operation: dict[str, Any]) -> dict[str, Any]: for param in operation["parameters"]: if "name" not in param: continue - param_name = param["name"] - param_schema = param.get("schema", {}) - param_type = param_schema.get("type", "string") + param_name: str = param["name"] + param_schema: _OpenAPIParameterSchema = param.get("schema", {}) + param_type: str = param_schema.get("type", "string") properties[param_name] = { "type": param_type, @@ -274,12 +318,12 @@ def build_input_schema(operation: dict[str, Any]) -> dict[str, Any]: # Process requestBody (OpenAPI 3.x) if "requestBody" in operation: - request_body: Final = operation["requestBody"] - content: Final = request_body.get("content", {}) + request_body: Final[_OpenAPIRequestBody] = operation["requestBody"] + content: Final[Mapping[str, _OpenAPIMediaType]] = request_body.get("content", {}) # Try to get JSON schema if "application/json" in content: - schema: Final = content["application/json"].get("schema", {}) + schema: Final[_OpenAPIJSONSchema] = content["application/json"].get("schema", {}) properties["body"] = { "type": "object", "description": request_body.get("description", "Request body"), @@ -347,7 +391,7 @@ def _merge_openapi_tool_request_headers( def create_tool_function( path: str, method: str, - operation: dict[str, Any], + operation: _OpenAPIObject, base_url: str, headers: dict[str, str] | None = None, ): @@ -373,7 +417,7 @@ def create_tool_function( path_params, query_params, body_params = extract_parameters(operation) original_method: Final = method.lower() - async def tool_function(**kwargs: Any) -> str: + async def tool_function(**kwargs: object) -> str: """ Dynamically generated tool function. @@ -448,10 +492,10 @@ def create_tool_function( return tool_function -def register_tools_from_openapi(spec: dict[str, Any], base_url: str): +def register_tools_from_openapi(spec: _OpenAPIObject, base_url: str) -> None: """Register MCP tools from OpenAPI specification.""" - paths: Final = spec.get("paths", {}) - used_names: Final[set] = set() + paths: Final[Mapping[str, Mapping[str, _OpenAPIOperation]]] = spec.get("paths", {}) + used_names: Final = set() for path, path_item in paths.items(): for method in ["get", "post", "put", "delete", "patch"]: diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 579c735b180..d340ca2ced3 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -18,8 +18,9 @@ Endpoints: import json import re +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final +from typing import Final, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse @@ -41,7 +42,30 @@ from litellm.types.proxy.claude_code_endpoints import ( router: Final = APIRouter() -async def _get_prisma_client(): +class _PluginRecord(Protocol): + id: str + name: str + version: str | None + description: str | None + manifest_json: str + enabled: bool + created_at: datetime | None + updated_at: datetime | None + created_by: str | None + + +class _MarketplaceEntry(TypedDict, total=False): + name: str + source: object + version: str + description: str + author: object + homepage: object + keywords: object + category: object + + +async def _get_prisma_client() -> object: """Get the prisma client from proxy_server.""" from litellm.proxy.proxy_server import prisma_client @@ -77,12 +101,14 @@ async def get_marketplace(): try: prisma_client: Final = await _get_prisma_client() - plugins: Final = await ClaudeCodePluginRepository(prisma_client).table.find_many(where={"enabled": True}) + plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( + where={"enabled": True} + ) plugin_list: Final = [] for plugin in plugins: try: - manifest = json.loads(plugin.manifest_json) + manifest: Mapping[str, object] = json.loads(plugin.manifest_json) except json.JSONDecodeError: verbose_proxy_logger.warning("Plugin %s has invalid manifest JSON, skipping", plugin.name) continue @@ -92,7 +118,7 @@ async def get_marketplace(): verbose_proxy_logger.warning("Plugin %s has no source field, skipping", plugin.name) continue - entry: dict[str, Any] = { + entry: _MarketplaceEntry = { "name": plugin.name, "source": manifest["source"], } @@ -137,7 +163,7 @@ async def get_marketplace(): _VALID_GIT_SUBDIR_PATH_RE: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]*(/[a-zA-Z0-9][a-zA-Z0-9._-]*)*$") -def _validate_plugin_source(source: dict[str, Any]) -> None: +def _validate_plugin_source(source: Mapping[str, str]) -> None: """Validate plugin source format, raising HTTPException on invalid input.""" source_type: Final = source.get("source") if source_type == "github": @@ -179,9 +205,9 @@ def _validate_plugin_source(source: dict[str, Any]) -> None: ) -def _build_plugin_manifest(name: str, spec: PluginSpec) -> dict[str, Any]: +def _build_plugin_manifest(name: str, spec: PluginSpec) -> Mapping[str, object]: """Build the stored manifest dict shared by plugin create and update.""" - dumped = spec.model_dump(exclude_none=True) + dumped: Final[Mapping[str, object]] = spec.model_dump(exclude_none=True) return {"name": name, **{key: value for key, value in dumped.items() if value and key != "name"}} @@ -255,14 +281,16 @@ async def register_plugin( _validate_plugin_source(request.source) - existing = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": request.name}) + existing: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( + where={"name": request.name} + ) if existing: raise _name_conflict_error(request.name) - manifest = _build_plugin_manifest(request.name, request) + manifest: Final[Mapping[str, object]] = _build_plugin_manifest(request.name, request) try: - plugin = await ClaudeCodePluginRepository(prisma_client).table.create( + plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.create( data={ "name": request.name, "version": request.version, @@ -326,7 +354,9 @@ async def list_plugins( prisma_client: Final = await _get_prisma_client() where: Final = {"enabled": True} if enabled_only else {} - plugins: Final = await ClaudeCodePluginRepository(prisma_client).table.find_many(where=where) + plugins: Final[Sequence[_PluginRecord]] = await ClaudeCodePluginRepository(prisma_client).table.find_many( + where=where + ) plugin_list: Final = [] for p in plugins: @@ -391,7 +421,9 @@ async def get_plugin( try: prisma_client: Final = await _get_prisma_client() - plugin: Final = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name}) + plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( + where={"name": plugin_name} + ) if not plugin: raise HTTPException( @@ -399,7 +431,7 @@ async def get_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - manifest: Final = json.loads(plugin.manifest_json) if plugin.manifest_json else {} + manifest: Final[Mapping[str, object]] = json.loads(plugin.manifest_json) if plugin.manifest_json else {} return { "id": plugin.id, @@ -477,19 +509,19 @@ async def update_plugin( from prisma.errors import PrismaError try: - prisma_client = await _get_prisma_client() + prisma_client: Final = await _get_prisma_client() _validate_plugin_source(request.source) - existing = await ClaudeCodePluginRepository(prisma_client).table.find_unique( + existing: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} # mutable-ok: prisma query arguments must be plain dicts ) if not existing: raise _error_response(404, f"Plugin '{plugin_name}' not found") - manifest = _build_plugin_manifest(plugin_name, request) + manifest: Final[Mapping[str, object]] = _build_plugin_manifest(plugin_name, request) - plugin = await ClaudeCodePluginRepository(prisma_client).table.update( + plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.update( where={"name": plugin_name}, # mutable-ok: prisma query arguments must be plain dicts data={ # mutable-ok: prisma query arguments must be plain dicts "version": request.version, @@ -540,7 +572,9 @@ async def enable_plugin( try: prisma_client: Final = await _get_prisma_client() - plugin: Final = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name}) + plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( + where={"name": plugin_name} + ) if not plugin: raise HTTPException( status_code=404, @@ -583,7 +617,9 @@ async def disable_plugin( try: prisma_client: Final = await _get_prisma_client() - plugin: Final = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name}) + plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( + where={"name": plugin_name} + ) if not plugin: raise HTTPException( status_code=404, @@ -626,7 +662,9 @@ async def delete_plugin( try: prisma_client: Final = await _get_prisma_client() - plugin: Final = await ClaudeCodePluginRepository(prisma_client).table.find_unique(where={"name": plugin_name}) + plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( + where={"name": plugin_name} + ) if not plugin: raise HTTPException( status_code=404, diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index cb91805aafc..3bedb6f720c 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -1,8 +1,14 @@ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, TypedDict from litellm._logging import verbose_proxy_logger +class _FieldSchema(TypedDict, total=False): + type: str + anyOf: Sequence["_FieldSchema"] + + class CustomOpenAPISpec: """ Handler for customizing OpenAPI specifications with Pydantic models @@ -26,7 +32,7 @@ class CustomOpenAPISpec: RESPONSES_API_PATHS = ["/v1/responses", "/responses"] @staticmethod - def get_pydantic_schema(model_class) -> dict[str, Any] | None: + def get_pydantic_schema(model_class) -> Mapping[str, object] | None: """ Get JSON schema from a Pydantic model, handling both v1 and v2 APIs. @@ -53,7 +59,9 @@ class CustomOpenAPISpec: return None @staticmethod - def add_schema_to_components(openapi_schema: dict[str, Any], schema_name: str, schema_def: dict[str, Any]) -> None: + def add_schema_to_components( + openapi_schema: dict[str, Any], schema_name: str, schema_def: Mapping[str, object] + ) -> None: """ Add a schema definition to the OpenAPI components/schemas section. @@ -72,7 +80,7 @@ class CustomOpenAPISpec: CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def}) @staticmethod - def add_request_body_to_paths(openapi_schema: dict[str, Any], paths: list[str], schema_ref: str) -> None: + def add_request_body_to_paths(openapi_schema: dict[str, Any], paths: Sequence[str], schema_ref: str) -> None: """ Add request body with expanded form fields for better Swagger UI display. This keeps the request body but expands it to show individual fields in the UI. @@ -130,7 +138,7 @@ class CustomOpenAPISpec: openapi_schema["paths"][path]["post"]["parameters"] = filtered_params @staticmethod - def _move_defs_to_components(openapi_schema: dict[str, Any], defs: dict[str, Any]) -> None: + def _move_defs_to_components(openapi_schema: dict[str, Any], defs: Mapping[str, Mapping[str, Any]]) -> None: """ Move $defs from Pydantic v2 schema to OpenAPI components/schemas. This makes the definitions resolvable in Swagger/OpenAPI viewers. @@ -190,7 +198,7 @@ class CustomOpenAPISpec: return schema @staticmethod - def _extract_field_schema(field_def: dict[str, Any]) -> dict[str, Any]: + def _extract_field_schema(field_def: _FieldSchema) -> _FieldSchema: """ Extract a simple schema from a Pydantic field definition for parameter display. @@ -218,7 +226,7 @@ class CustomOpenAPISpec: return {"type": "string"} @staticmethod - def _expand_field_definition(field_def: dict[str, Any]) -> dict[str, Any]: + def _expand_field_definition(field_def: dict[str, object]) -> dict[str, object]: """ Expand a Pydantic field definition for inline use in OpenAPI schema. This creates a full field definition that Swagger UI can render as individual form fields. @@ -234,12 +242,12 @@ class CustomOpenAPISpec: @staticmethod def add_request_schema( - openapi_schema: dict[str, Any], + openapi_schema: dict[str, object], model_class: type, schema_name: str, - paths: list[str], + paths: Sequence[str], operation_name: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Generic method to add a request schema to OpenAPI specification. @@ -279,8 +287,8 @@ class CustomOpenAPISpec: @staticmethod def add_chat_completion_request_schema( - openapi_schema: dict[str, Any], - ) -> dict[str, Any]: + openapi_schema: dict[str, object], + ) -> dict[str, object]: """ Add ProxyChatCompletionRequest schema to chat completion endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -306,7 +314,7 @@ class CustomOpenAPISpec: return openapi_schema @staticmethod - def add_embedding_request_schema(openapi_schema: dict[str, Any]) -> dict[str, Any]: + def add_embedding_request_schema(openapi_schema: dict[str, object]) -> dict[str, object]: """ Add EmbeddingRequest schema to embedding endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -333,8 +341,8 @@ class CustomOpenAPISpec: @staticmethod def add_responses_api_request_schema( - openapi_schema: dict[str, Any], - ) -> dict[str, Any]: + openapi_schema: dict[str, object], + ) -> dict[str, object]: """ Add ResponsesAPIRequestParams schema to responses API endpoints for documentation. This shows the request body in Swagger without runtime validation. @@ -361,8 +369,8 @@ class CustomOpenAPISpec: @staticmethod def add_llm_api_request_schema_body( - openapi_schema: dict[str, Any], - ) -> dict[str, Any]: + openapi_schema: dict[str, object], + ) -> dict[str, object]: """ Add LLM API request schema bodies to OpenAPI specification for documentation. diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index e582ddb36c7..c40f4d4fef2 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,5 +1,7 @@ import json -from typing import Any, Final +from collections.abc import Mapping, Sequence +from collections.abc import Set as AbstractSet +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import HTTPException @@ -15,6 +17,36 @@ from litellm.proxy.common_utils.resource_ownership import ( from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.responses.utils import ResponsesAPIRequestUtils +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +class _ManagedObjectRow(Protocol): + model_object_id: str + unified_object_id: str | None + file_purpose: str | None + created_by: str | None + + +class _ManagedObjectTable(Protocol): + async def find_unique(self, *, where: Mapping[str, str]) -> _ManagedObjectRow | None: ... + + async def find_first(self, *, where: Mapping[str, str]) -> _ManagedObjectRow | None: ... + + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ManagedObjectRow]: ... + + async def create(self, *, data: Mapping[str, str]) -> _ManagedObjectRow: ... + + async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> _ManagedObjectRow | None: ... + + +class _ContainerListResponse(Protocol): + data: Sequence[object] + first_id: str | None + last_id: str | None + has_more: bool + + CONTAINER_OBJECT_PURPOSE: Final = "container" # 60s LRU/TTL cache absorbs every container access check before it reaches @@ -39,7 +71,7 @@ _CONTAINER_STORED_ID_CACHE: Final = InMemoryCache(max_size_in_memory=10000, defa _ALLOWED_CONTAINER_IDS_CACHE: Final = InMemoryCache(max_size_in_memory=2048, default_ttl=60) -def _allowed_container_ids_cache_key(owner_scopes: list[str]) -> str: +def _allowed_container_ids_cache_key(owner_scopes: Sequence[str]) -> str: """JSON-encode the sorted scope list — using a separator like ``|`` would collide for any tenant whose user_id / team_id / org_id / api_key happens to contain the separator. JSON quoting escapes @@ -86,7 +118,7 @@ async def get_container_forwarding_params( return params -def _get_response_id(response: Any) -> str | None: +def _get_response_id(response: object) -> str | None: if response is None: return None if isinstance(response, dict): @@ -96,7 +128,7 @@ def _get_response_id(response: Any) -> str | None: return value if isinstance(value, str) else None -def _dump_response(response: Any) -> dict[str, Any]: +def _dump_response(response: Any) -> dict[str, object]: if isinstance(response, dict): return dict(response) if hasattr(response, "model_dump"): @@ -106,17 +138,17 @@ def _dump_response(response: Any) -> dict[str, Any]: return {"id": _get_response_id(response)} -async def _get_prisma_client(): +async def _get_prisma_client() -> "PrismaClient | None": from litellm.proxy.proxy_server import prisma_client return prisma_client def _custom_llm_provider_from_responses_response( - response: Any, + response: object, default: str = "openai", ) -> str: - hidden_params: dict[str, Any] = {} + hidden_params: Mapping[str, object] = {} if isinstance(response, dict): hidden_params = response.get("_hidden_params") or {} else: @@ -129,7 +161,7 @@ def _custom_llm_provider_from_responses_response( async def record_container_owners_from_responses_response( - response: Any, + response: object, user_api_key_dict: UserAPIKeyAuth, custom_llm_provider: str | None = None, ) -> None: @@ -160,10 +192,10 @@ async def record_container_owners_from_responses_response( async def record_container_owner( - response: Any, + response: object, user_api_key_dict: UserAPIKeyAuth, custom_llm_provider: str, -) -> Any: +) -> object: container_id: Final = _get_response_id(response) if container_id is None: verbose_proxy_logger.warning("Skipping container ownership tracking because provider response has no id") @@ -195,7 +227,7 @@ async def record_container_owner( verbose_proxy_logger.warning("Skipping container ownership tracking because prisma_client is None") return response - table: Final = ManagedObjectRepository(prisma_client).table + table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table existing: Final = await table.find_unique(where={"model_object_id": model_object_id}) if existing is not None: if getattr(existing, "file_purpose", None) != CONTAINER_OBJECT_PURPOSE: @@ -247,15 +279,15 @@ async def _get_container_owner(original_container_id: str, custom_llm_provider: if prisma_client is None: return None - row: Final = await ManagedObjectRepository(prisma_client).table.find_first( + row: Final[_ManagedObjectRow | None] = await ManagedObjectRepository(prisma_client).table.find_first( where={ "model_object_id": model_object_id, "file_purpose": CONTAINER_OBJECT_PURPOSE, } ) - owner: Final = getattr(row, "created_by", None) if row is not None else None + owner: Final[str | None] = getattr(row, "created_by", None) if row is not None else None _CONTAINER_OWNER_CACHE.set_cache(model_object_id, owner if owner is not None else _NEGATIVE_OWNER_SENTINEL) - stored_id: Final = getattr(row, "unified_object_id", None) if row is not None else None + stored_id: Final[str | None] = getattr(row, "unified_object_id", None) if row is not None else None _CONTAINER_STORED_ID_CACHE.set_cache( model_object_id, (stored_id if isinstance(stored_id, str) and stored_id else _NEGATIVE_STORED_ID_SENTINEL), @@ -283,13 +315,13 @@ async def _get_stored_container_id(original_container_id: str, custom_llm_provid if prisma_client is None: return None - row: Final = await ManagedObjectRepository(prisma_client).table.find_first( + row: Final[_ManagedObjectRow | None] = await ManagedObjectRepository(prisma_client).table.find_first( where={ "model_object_id": model_object_id, "file_purpose": CONTAINER_OBJECT_PURPOSE, } ) - stored_id: Final = getattr(row, "unified_object_id", None) if row is not None else None + stored_id: Final[str | None] = getattr(row, "unified_object_id", None) if row is not None else None _CONTAINER_STORED_ID_CACHE.set_cache( model_object_id, (stored_id if isinstance(stored_id, str) and stored_id else _NEGATIVE_STORED_ID_SENTINEL), @@ -317,7 +349,7 @@ async def assert_user_can_access_container( return original_container_id, resolved_provider -def _get_container_list_data(response: Any) -> list[Any] | None: +def _get_container_list_data(response: object) -> Sequence[object] | None: if response is None: return None if isinstance(response, dict): @@ -327,7 +359,9 @@ def _get_container_list_data(response: Any) -> list[Any] | None: return data if isinstance(data, list) else None -def _set_container_list_data(response: Any, data: list[Any], removed_filtered_items: bool = False) -> Any: +def _set_container_list_data( + response: _ContainerListResponse, data: list[object], removed_filtered_items: bool = False +) -> _ContainerListResponse: if isinstance(response, dict): response["data"] = data if data: @@ -353,7 +387,7 @@ def _set_container_list_data(response: Any, data: list[Any], removed_filtered_it async def _get_allowed_container_ids( user_api_key_dict: UserAPIKeyAuth, -) -> set[str]: +) -> AbstractSet[str]: owner_scopes: Final = get_resource_owner_scopes(user_api_key_dict) if not owner_scopes: return set() @@ -367,7 +401,7 @@ async def _get_allowed_container_ids( if prisma_client is None: return set() - rows: Final = await ManagedObjectRepository(prisma_client).table.find_many( + rows: Final[Sequence[_ManagedObjectRow]] = await ManagedObjectRepository(prisma_client).table.find_many( where={ "file_purpose": CONTAINER_OBJECT_PURPOSE, "created_by": {"in": owner_scopes}, @@ -382,10 +416,10 @@ async def _get_allowed_container_ids( async def filter_container_list_response( - response: Any, + response: _ContainerListResponse, user_api_key_dict: UserAPIKeyAuth, custom_llm_provider: str, -) -> Any: +) -> _ContainerListResponse: if is_proxy_admin(user_api_key_dict): return response @@ -394,7 +428,7 @@ async def filter_container_list_response( return response allowed_container_ids: Final = await _get_allowed_container_ids(user_api_key_dict) - filtered: Final[list[Any]] = [] + filtered: Final[list[object]] = [] for item in data: container_id = _get_response_id(item) if container_id is None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index c5481c9ce63..79867d492a1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -9,11 +9,13 @@ import contextlib import json import os import ssl -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping, Sequence +from ssl import SSLContext from typing import TYPE_CHECKING, Any, Final from fastapi import HTTPException from pydantic import BaseModel +from typing_extensions import NotRequired, TypedDict from websockets.asyncio.client import ClientConnection, connect from websockets.exceptions import ConnectionClosed @@ -35,8 +37,8 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( CallTypesLiteral, Choices, - EmbeddingResponse, - ImageResponse, + LLMResponseTypes, + Message, ModelResponse, ModelResponseStream, ResponsesAPIResponse, @@ -50,6 +52,40 @@ class CatoNetworksGuardrailMissingSecrets(Exception): pass +class _WsSslKwargs(TypedDict, total=False): + ssl: bool | str | SSLContext + + +class _CatoRequiredAction(TypedDict, total=False): + action_type: str + detection_message: str + + +class _CatoRedactedMessage(TypedDict): + role: NotRequired[str] + content: str | None + + +class _CatoRedactedChat(TypedDict, total=False): + all_redacted_messages: Sequence[_CatoRedactedMessage] + + +class _CatoAnalyzeResponse(TypedDict): + required_action: _CatoRequiredAction + analysis_result: NotRequired[Mapping[str, Mapping[str, object]]] + redacted_chat: NotRequired[_CatoRedactedChat] + + +class _CatoOutputRedaction(TypedDict): + redacted_output: str + + +class _CatoStreamMessage(TypedDict, total=False): + verified_chunk: Mapping[str, object] + done: bool + blocking_message: str + + class CatoNetworksGuardrail(CustomGuardrail): @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -80,7 +116,7 @@ class CatoNetworksGuardrail(CustomGuardrail): super().__init__(**kwargs) @staticmethod - def _build_ws_ssl_kwargs(ssl_verify: bool | str | None, ws_api_base: str) -> dict: + def _build_ws_ssl_kwargs(ssl_verify: bool | str | None, ws_api_base: str) -> _WsSslKwargs: """Resolve the ``ssl`` argument for ``websockets.connect``. Mirrors the ``ssl_verify`` handling applied to the HTTP handler so a custom Cato instance behind TLS honours the same verification settings for streaming.""" @@ -156,7 +192,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return flattened @staticmethod - def _prompt_inspection_messages(prompt: Any) -> list: + def _prompt_inspection_messages(prompt: object) -> Sequence[Mapping[str, str]]: """Synthetic user messages for a legacy completion ``prompt`` (a string or a list of string prompts).""" if isinstance(prompt, str): @@ -208,7 +244,7 @@ class CatoNetworksGuardrail(CustomGuardrail): stack.extend(reversed(node)) @classmethod - def _extra_inspection_sources(cls, data: dict) -> list: + def _extra_inspection_sources(cls, data: dict) -> Sequence[tuple[str, Sequence[Mapping[str, str]]]]: """Text the proxy forwards to the model outside chat ``messages``: Responses-API ``input`` and ``instructions``, legacy completion ``prompt`` and tool/function/``response_format`` schema strings. Returned @@ -251,7 +287,7 @@ class CatoNetworksGuardrail(CustomGuardrail): json={"messages": self._inspection_messages(data)}, ) response.raise_for_status() - res: Final = response.json() + res: Final[_CatoAnalyzeResponse] = response.json() required_action: Final = res.get("required_action") action_type: Final = required_action and required_action.get("action_type", None) if action_type is None: @@ -267,7 +303,11 @@ class CatoNetworksGuardrail(CustomGuardrail): verbose_proxy_logger.error("Cato: %s action", action_type) return data - def _handle_block_action(self, analysis_result: Any, required_action: Any) -> None: + def _handle_block_action( + self, + analysis_result: Mapping[str, Mapping[str, object]], + required_action: _CatoRequiredAction, + ) -> None: detection_message: Final = required_action.get("detection_message", None) verbose_proxy_logger.info( "Cato: Violation detected enabled policies: {policies}".format( @@ -348,7 +388,7 @@ class CatoNetworksGuardrail(CustomGuardrail): hook: str, key_alias: str | None, user_email: str | None = None, - ) -> dict | None: + ) -> _CatoOutputRedaction | None: call_id: Final = request_data.get("litellm_call_id") inspection_messages: Final = self._inspection_messages(request_data) assistant_index: Final = len(inspection_messages) @@ -363,7 +403,7 @@ class CatoNetworksGuardrail(CustomGuardrail): json={"messages": inspection_messages + [{"role": "assistant", "content": output}]}, ) response.raise_for_status() - res: Final = response.json() + res: Final[_CatoAnalyzeResponse] = response.json() required_action: Final = res.get("required_action") action_type: Final = required_action and required_action.get("action_type", None) if action_type and action_type == "block_action": @@ -378,7 +418,11 @@ class CatoNetworksGuardrail(CustomGuardrail): return {"redacted_output": redacted_output} return None - def _handle_block_action_on_output(self, analysis_result: Any, required_action: Any) -> None: + def _handle_block_action_on_output( + self, + analysis_result: Mapping[str, Mapping[str, object]], + required_action: _CatoRequiredAction, + ) -> None: detection_message: Final = required_action.get("detection_message", None) verbose_proxy_logger.info( "Cato: detected: {detected}, enabled policies: {policies}".format( @@ -422,7 +466,7 @@ class CatoNetworksGuardrail(CustomGuardrail): ) @staticmethod - def _output_fragments(message: Any) -> list: + def _output_fragments(message: Message) -> Sequence[tuple[tuple[str, int | None], str]]: """Assistant text the proxy returns to the caller: ``content`` plus every ``tool_calls[].function.arguments`` string, each tagged with where a redaction must be written back. ``content`` is only included when present @@ -439,7 +483,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return fragments @staticmethod - def _apply_output_fragment(message: Any, target: tuple, redacted: str) -> None: + def _apply_output_fragment(message: Any, target: tuple[str, int | None], redacted: str) -> None: kind, idx = target if kind == "content": message.content = redacted @@ -447,11 +491,11 @@ class CatoNetworksGuardrail(CustomGuardrail): message.tool_calls[idx].function.arguments = redacted @staticmethod - def _responses_output_field(item: Any, key: str) -> Any: + def _responses_output_field(item: object, key: str) -> str | Sequence[object] | None: return item.get(key) if isinstance(item, dict) else getattr(item, key, None) @classmethod - def _responses_output_fragments(cls, response: ResponsesAPIResponse) -> list: + def _responses_output_fragments(cls, response: ResponsesAPIResponse) -> Sequence[tuple[object, str, str]]: """Assistant text the Responses API returns to the caller: every ``output_text`` content block plus every function-call ``arguments`` string, each paired with the ``(container, key)`` a Cato redaction is @@ -474,7 +518,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return fragments @staticmethod - def _apply_responses_output_fragment(container: Any, key: str, redacted: str) -> None: + def _apply_responses_output_fragment(container: object, key: str, redacted: str) -> None: if isinstance(container, dict): container[key] = redacted else: @@ -505,8 +549,8 @@ class CatoNetworksGuardrail(CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - response: Any | ModelResponse | EmbeddingResponse | ImageResponse, - ) -> Any: + response: LLMResponseTypes, + ) -> LLMResponseTypes: user_email: Final = self._resolve_cato_user_email(user_api_key_dict) if isinstance(response, ModelResponse) and response.choices: for choice in response.choices: @@ -526,7 +570,7 @@ class CatoNetworksGuardrail(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response, + response: AsyncGenerator[object, None], request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: from litellm.proxy.proxy_server import StreamingCallbackError @@ -547,7 +591,7 @@ class CatoNetworksGuardrail(CustomGuardrail): try: while True: raw_message = await self._await_cato_message(websocket, sender) - result = json.loads(raw_message) + result: _CatoStreamMessage = json.loads(raw_message) if verified_chunk := result.get("verified_chunk"): yield ModelResponseStream.model_validate(verified_chunk) continue @@ -560,7 +604,7 @@ class CatoNetworksGuardrail(CustomGuardrail): finally: await self._cancel_background_task(sender) - async def _await_cato_message(self, websocket: ClientConnection, sender: asyncio.Task) -> Any: + async def _await_cato_message(self, websocket: ClientConnection, sender: asyncio.Task[None]) -> str | bytes: """Wait for the next Cato message, surfacing a dead forwarding task instead of blocking.""" from litellm.proxy.proxy_server import StreamingCallbackError @@ -578,7 +622,7 @@ class CatoNetworksGuardrail(CustomGuardrail): async def forward_the_stream_to_cato( self, websocket: ClientConnection, - response_iter: AsyncGenerator[Any, None], + response_iter: AsyncGenerator[object, None], ) -> None: async for chunk in response_iter: if isinstance(chunk, BaseModel): diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index 131ce5f9392..ca84ff47884 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -1,7 +1,8 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Any, Final, Literal +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict from urllib.parse import urlparse from uuid import uuid4 @@ -29,9 +30,31 @@ from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from pydantic import BaseModel + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +class _HiddenlayerEvaluation(TypedDict, total=False): + action: str + threat_level: str + + +class _HiddenlayerAnalysisEntry(TypedDict, total=False): + name: str + detected: bool + + +class _HiddenlayerModifiedSide(TypedDict): + messages: Any + + +class _HiddenlayerResponse(TypedDict, total=False): + evaluation: _HiddenlayerEvaluation + analysis: Sequence[_HiddenlayerAnalysisEntry] + modified_data: Mapping[str, _HiddenlayerModifiedSide] + + def is_saas(host: str) -> bool: """Checks whether the connection is to the SaaS platform""" @@ -43,7 +66,7 @@ def is_saas(host: str) -> bool: return False -def _get_jwt(auth_url, api_id, api_key): +def _get_jwt(auth_url, api_id, api_key) -> str: token_url: Final = f"{auth_url}/oauth2/token?grant_type=client_credentials" resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key)) @@ -139,7 +162,7 @@ class HiddenlayerGuardrail(CustomGuardrail): if scan_params := inputs.get("structured_messages"): last_msg: Final = scan_params[-1] - result = await self._call_hiddenlayer( + result: _HiddenlayerResponse = await self._call_hiddenlayer( project_id, hl_request_metadata, { @@ -205,11 +228,11 @@ class HiddenlayerGuardrail(CustomGuardrail): async def _call_hiddenlayer( self, project_id: str | None, - metadata: dict[str, str], - payload: dict[str, Any], + metadata: Mapping[str, str], + payload: Mapping[str, Sequence[Mapping[str, str]]], input_type: Literal["request", "response"], - ) -> dict[str, Any]: - data: Final[dict[str, Any]] = {"metadata": metadata} + ) -> _HiddenlayerResponse: + data: Final[dict[str, object]] = {"metadata": metadata} if input_type == "request": data["input"] = payload @@ -235,7 +258,7 @@ class HiddenlayerGuardrail(CustomGuardrail): headers=headers, ) response.raise_for_status() - result = response.json() + result: _HiddenlayerResponse = response.json() verbose_proxy_logger.debug("Hiddenlayer reponse: %s", result) @@ -265,7 +288,7 @@ class HiddenlayerGuardrail(CustomGuardrail): return result @staticmethod - def get_config_model() -> type[GuardrailConfigModel] | None: + def get_config_model() -> type[GuardrailConfigModel[BaseModel]] | None: from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( HiddenlayerGuardrailConfigModel, ) @@ -343,7 +366,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): if "hl-requester-id" not in hl_headers: hl_headers["hl-requester-id"] = "LiteLLM" - payload: Any + payload: object if input_type == "request": payload = { "messages": inputs.get("structured_messages"), @@ -461,7 +484,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): return response @staticmethod - def get_config_model() -> type[GuardrailConfigModel] | None: + def get_config_model() -> type[GuardrailConfigModel[BaseModel]] | None: from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( HiddenlayerGuardrailConfigModel, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index 4ae29ade0d7..80beb90cf27 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -2,7 +2,10 @@ import threading import time import uuid from collections import OrderedDict -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from typing_extensions import NotRequired, TypedDict from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -15,6 +18,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import AllMessageValues GRAPH_API_BASE: Final = "https://graph.microsoft.com/v1.0" @@ -25,6 +29,11 @@ GRAPH_SCOPE: Final = "https://graph.microsoft.com/.default" SCOPE_CACHE_TTL_SECONDS: Final = 3600.0 +class GraphTokenResponse(TypedDict): + access_token: str + expires_in: NotRequired[int] + + class PurviewGuardrailBase: """ Base class for Microsoft Purview guardrails. @@ -41,8 +50,8 @@ class PurviewGuardrailBase: client_secret: str, purview_app_name: str = "LiteLLM", user_id_field: str = "user_id", - **kwargs: Any, - ): + **kwargs: object, + ) -> None: # Forward remaining kwargs to the next class in the MRO # (typically CustomGuardrail). super().__init__(**kwargs) @@ -59,7 +68,7 @@ class PurviewGuardrailBase: # Protection scope cache: user_id -> (etag, scope_response, fetched_at) # Capped at 1000 entries (LRU eviction) to avoid unbounded growth. - self._scope_cache: OrderedDict[str, tuple[str, dict[str, Any], float]] = OrderedDict() + self._scope_cache: OrderedDict[str, tuple[str, Mapping[str, object], float]] = OrderedDict() self._scope_cache_maxsize = 1000 # Use a threading.Lock (not asyncio.Lock) because this lock is acquired # from both the proxy's main asyncio event loop and from short-lived @@ -100,7 +109,7 @@ class PurviewGuardrailBase: headers={"Content-Type": "application/x-www-form-urlencoded"}, ) response.raise_for_status() - token_data: Final = response.json() + token_data: Final[GraphTokenResponse] = response.json() access_token: Final = token_data["access_token"] expires_in: Final = int(token_data.get("expires_in", 3599)) # Recompute ``now`` after the await so the expiry reflects when the @@ -117,9 +126,9 @@ class PurviewGuardrailBase: async def _graph_post( self, url: str, - json_body: dict[str, Any], - extra_headers: dict[str, str] | None = None, - ) -> tuple[dict[str, Any], dict[str, str]]: + json_body: dict[str, object], + extra_headers: Mapping[str, str] | None = None, + ) -> tuple[dict[str, object], dict[str, str]]: """POST to Graph API with bearer auth. Returns: @@ -136,7 +145,7 @@ class PurviewGuardrailBase: verbose_proxy_logger.debug("Purview Graph POST %s", url) response: Final = await self.async_handler.post(url=url, headers=headers, json=json_body) response.raise_for_status() - response_json: Final[dict[str, Any]] = response.json() + response_json: Final[dict[str, object]] = response.json() response_headers: Final = dict(response.headers) verbose_proxy_logger.debug("Purview Graph response: %s", response_json) return response_json, response_headers @@ -145,7 +154,7 @@ class PurviewGuardrailBase: # Protection scopes # ------------------------------------------------------------------ - async def _compute_protection_scopes(self, user_id: str) -> tuple[str, dict[str, Any]]: + async def _compute_protection_scopes(self, user_id: str) -> tuple[str, Mapping[str, object]]: """Call protectionScopes/compute and cache with ETag. Returns: @@ -161,7 +170,7 @@ class PurviewGuardrailBase: return cached[0], cached[1] url: Final = f"{GRAPH_API_BASE}/users/{encoded_user_id}/dataSecurityAndGovernance/protectionScopes/compute" - body: Final[dict[str, Any]] = { + body: Final[dict[str, object]] = { "activities": "uploadText,downloadText", "locations": [ { @@ -199,7 +208,7 @@ class PurviewGuardrailBase: activity: str, etag: str, correlation_id: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Call processContent for DLP policy evaluation. Args: @@ -211,7 +220,7 @@ class PurviewGuardrailBase: """ encoded_user_id: Final = self._encode_graph_user_id(user_id) url: Final = f"{GRAPH_API_BASE}/users/{encoded_user_id}/dataSecurityAndGovernance/processContent" - body: Final[dict[str, Any]] = { + body: Final[dict[str, object]] = { "contentToProcess": { "contentEntries": [ { @@ -261,7 +270,9 @@ class PurviewGuardrailBase: # User ID resolution # ------------------------------------------------------------------ - def _resolve_user_id(self, data: dict[str, Any], user_api_key_dict: Any) -> str | None: + def _resolve_user_id( + self, data: Mapping[str, Mapping[str, object]], user_api_key_dict: "UserAPIKeyAuth" + ) -> str | None: """Resolve the Entra user object ID from request data or auth context. Returns the strongest available identity walking down four sources, in @@ -284,7 +295,7 @@ class PurviewGuardrailBase: if hasattr(user_api_key_dict, "end_user_id") and user_api_key_dict.end_user_id: return str(user_api_key_dict.end_user_id) - metadata: Final = data.get("metadata") or data.get("litellm_metadata") or {} + metadata: Final[Mapping[str, object]] = data.get("metadata") or data.get("litellm_metadata") or {} uid = metadata.get("user_api_key_user_id") if uid: return str(uid) @@ -296,15 +307,15 @@ class PurviewGuardrailBase: return None @staticmethod - def _logging_kwargs_metadata(kwargs: dict[str, Any]) -> dict[str, Any]: + def _logging_kwargs_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: """Metadata dict from ``model_call_details`` / logging kwargs.""" - litellm_params: Final = kwargs.get("litellm_params") or {} + litellm_params: Final[object] = kwargs.get("litellm_params") or {} if not isinstance(litellm_params, dict): return {} md: Final = litellm_params.get("metadata") return md if isinstance(md, dict) else {} - def _resolve_trusted_user_id(self, data: dict[str, Any], user_api_key_dict: Any) -> str | None: + def _resolve_trusted_user_id(self, data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth") -> str | None: """Resolve user ID from API-key/JWT-bound identity for blocking DLP. Uses only ``UserAPIKeyAuth.user_id`` (bound on the LiteLLM key or JWT). @@ -325,7 +336,7 @@ class PurviewGuardrailBase: return None - def _resolve_user_id_from_logging_kwargs(self, kwargs: dict[str, Any]) -> str | None: + def _resolve_user_id_from_logging_kwargs(self, kwargs: Mapping[str, object]) -> str | None: """Trusted-identity-only resolver for logging-only hooks. Uses only the proxy-injected ``user_api_key_user_id`` (populated from @@ -348,7 +359,7 @@ class PurviewGuardrailBase: # ------------------------------------------------------------------ @staticmethod - def _should_block(response: dict[str, Any]) -> bool: + def _should_block(response: Mapping[str, Sequence[Mapping[str, str]]]) -> bool: """Return True if any policyAction requires blocking.""" for action in response.get("policyActions", []): odata_type = action.get("@odata.type", "") @@ -365,7 +376,7 @@ class PurviewGuardrailBase: # ------------------------------------------------------------------ @staticmethod - def is_token_id_prompt(prompt: Any) -> bool: + def is_token_id_prompt(prompt: str | Sequence[object] | None) -> bool: """Return True if ``prompt`` carries OpenAI completions token ids. Covers every list shape that ``completion_prompt_to_str`` cannot decode @@ -383,7 +394,7 @@ class PurviewGuardrailBase: return False @staticmethod - def completion_prompt_to_str(prompt: Any) -> str | None: + def completion_prompt_to_str(prompt: str | Sequence[object] | None) -> str | None: """Normalize OpenAI ``/v1/completions`` ``prompt`` for text DLP. Supports string prompts and list-of-string prompts. List-of-token-id prompts @@ -408,7 +419,7 @@ class PurviewGuardrailBase: return None @staticmethod - def _extract_tool_call_args_from_message(message: Any) -> list[str]: + def _extract_tool_call_args_from_message(message: object) -> list[str]: """Return plaintext arguments strings from tool_calls and function_call fields. Covers both the request path (assistant messages in chat histories that @@ -419,7 +430,9 @@ class PurviewGuardrailBase: args: Final[list[str]] = [] # tool_calls: [{"function": {"arguments": "..."}}] - tool_calls = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) + tool_calls: Final[Sequence[object] | None] = ( + message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) + ) if tool_calls: for tc in tool_calls: fn = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None) diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index 70a6cc507f5..eeb64ab2773 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -14,7 +14,8 @@ GET /v1/workflows/runs/{run_id}/messages - Fetch conversation history """ import json -from typing import Any, Final, Literal +from collections.abc import Mapping, Sequence +from typing import Final, Literal, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query @@ -43,7 +44,7 @@ router: Final = APIRouter() _MAX_SEQUENCE_RETRIES: Final = 5 -def _json(value: Any) -> str: +def _json(value: object) -> str: """Serialize a Python value for prisma-client-py Json fields (must be a string).""" return json.dumps(value) @@ -62,7 +63,7 @@ def _caller_key(user_api_key_dict: UserAPIKeyAuth) -> str | None: # Status transitions driven by event_type -_EVENT_STATUS_MAP: Final[dict[str, str]] = { +_EVENT_STATUS_MAP: Final[Mapping[str, str]] = { "step.started": "running", "step.failed": "failed", "hook.waiting": "paused", @@ -77,8 +78,8 @@ _EVENT_STATUS_MAP: Final[dict[str, str]] = { class WorkflowRunCreateRequest(BaseModel): workflow_type: str - input: dict[str, Any] | None = None - metadata: dict[str, Any] | None = None + input: Mapping[str, object] | None = None + metadata: Mapping[str, object] | None = None WorkflowRunStatus = Literal["pending", "running", "paused", "completed", "failed"] @@ -86,14 +87,14 @@ WorkflowRunStatus = Literal["pending", "running", "paused", "completed", "failed class WorkflowRunUpdateRequest(BaseModel): status: WorkflowRunStatus | None = None - output: dict[str, Any] | None = None - metadata: dict[str, Any] | None = None + output: Mapping[str, object] | None = None + metadata: Mapping[str, object] | None = None class WorkflowEventCreateRequest(BaseModel): event_type: str step_name: str - data: dict[str, Any] | None = None + data: Mapping[str, object] | None = None class WorkflowMessageCreateRequest(BaseModel): @@ -102,15 +103,60 @@ class WorkflowMessageCreateRequest(BaseModel): session_id: str | None = None +class _RunRow(Protocol): + @property + def created_by(self) -> str | None: ... + + +class _SeqRow(Protocol): + @property + def sequence_number(self) -> int: ... + + +class _RunCreateData(TypedDict, total=False): + workflow_type: str + created_by: str | None + input: str + metadata: str + + +class _RunWhere(TypedDict, total=False): + workflow_type: str + status: str | Mapping[str, Sequence[str]] + created_by: str + + +class _RunUpdateData(TypedDict, total=False): + status: WorkflowRunStatus + output: str + metadata: str + + +class _EventCreateData(TypedDict, total=False): + run_id: str + event_type: str + step_name: str + sequence_number: int + data: str + + +class _MessageCreateData(TypedDict, total=False): + run_id: str + role: str + content: str + sequence_number: int + session_id: str + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) -> int: +async def _get_next_sequence_number(prisma_client: object, run_id: str, table: str) -> int: """Return MAX(sequence_number) + 1 for the given run, for either events or messages.""" if table == "events": - rows = await WorkflowEventRepository(prisma_client).table.find_many( + rows: Sequence[_SeqRow] = await WorkflowEventRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "desc"}, take=1, @@ -125,12 +171,12 @@ async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) async def _require_run( - prisma_client: Any, + prisma_client: object, run_id: str, user_api_key_dict: UserAPIKeyAuth | None = None, -) -> Any: +) -> _RunRow: """Return the run or raise 404. For non-admin callers, also enforce key ownership.""" - run: Final = await WorkflowRunRepository(prisma_client).table.find_unique(where={"run_id": run_id}) + run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.find_unique(where={"run_id": run_id}) if run is None: raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") if user_api_key_dict is not None and not _is_admin(user_api_key_dict): @@ -165,7 +211,7 @@ async def create_workflow_run( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - create_data: Final[dict[str, Any]] = { + create_data: Final[_RunCreateData] = { "workflow_type": data.workflow_type, "created_by": _caller_key(user_api_key_dict), } @@ -173,7 +219,7 @@ async def create_workflow_run( create_data["input"] = _json(data.input) if data.metadata is not None: create_data["metadata"] = _json(data.metadata) - run: Final = await WorkflowRunRepository(prisma_client).table.create(data=create_data) + run: Final[_RunRow] = await WorkflowRunRepository(prisma_client).table.create(data=create_data) return run except Exception as e: verbose_proxy_logger.exception("Error creating workflow run: %s", e) @@ -200,7 +246,7 @@ async def list_workflow_runs( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - where: Final[dict[str, Any]] = {} + where: Final[_RunWhere] = {} if workflow_type: where["workflow_type"] = workflow_type if status: @@ -214,7 +260,7 @@ async def list_workflow_runs( where["created_by"] = caller try: - runs: Final = await WorkflowRunRepository(prisma_client).table.find_many( + runs: Final[Sequence[object]] = await WorkflowRunRepository(prisma_client).table.find_many( where=where, order={"created_at": "desc"}, take=limit, @@ -241,7 +287,7 @@ async def get_workflow_run( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - run: Final = await WorkflowRunRepository(prisma_client).table.find_unique( + run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.find_unique( where={"run_id": run_id}, include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}}, ) @@ -275,7 +321,7 @@ async def update_workflow_run( if prisma_client is None: raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) - update: Final[dict[str, Any]] = {} + update: Final[_RunUpdateData] = {} if data.status is not None: update["status"] = data.status if data.output is not None: @@ -290,7 +336,7 @@ async def update_workflow_run( await _require_run(prisma_client, run_id, user_api_key_dict) try: - run: Final = await WorkflowRunRepository(prisma_client).table.update( + run: Final[_RunRow | None] = await WorkflowRunRepository(prisma_client).table.update( where={"run_id": run_id}, data=update, ) @@ -332,7 +378,7 @@ async def append_workflow_event( for attempt in range(_MAX_SEQUENCE_RETRIES): try: seq = await _get_next_sequence_number(prisma_client, run_id, "events") - event_data: dict[str, Any] = { + event_data: _EventCreateData = { "run_id": run_id, "event_type": data.event_type, "step_name": data.step_name, @@ -342,7 +388,7 @@ async def append_workflow_event( event_data["data"] = _json(data.data) async with prisma_client.db.tx() as tx: - event = await tx.litellm_workflowevent.create(data=event_data) + event: object = await tx.litellm_workflowevent.create(data=event_data) if new_status: await tx.litellm_workflowrun.update( where={"run_id": run_id}, @@ -389,7 +435,7 @@ async def list_workflow_events( await _require_run(prisma_client, run_id, _read_scope_caller(user_api_key_dict)) try: - events: Final = await WorkflowEventRepository(prisma_client).table.find_many( + events: Final[Sequence[object]] = await WorkflowEventRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "asc"}, take=limit, @@ -424,7 +470,7 @@ async def append_workflow_message( for attempt in range(_MAX_SEQUENCE_RETRIES): try: seq = await _get_next_sequence_number(prisma_client, run_id, "messages") - msg_data: dict[str, Any] = { + msg_data: _MessageCreateData = { "run_id": run_id, "role": data.role, "content": data.content, @@ -432,7 +478,7 @@ async def append_workflow_message( } if data.session_id is not None: msg_data["session_id"] = data.session_id - msg = await WorkflowMessageRepository(prisma_client).table.create(data=msg_data) + msg: object = await WorkflowMessageRepository(prisma_client).table.create(data=msg_data) return msg except Exception as e: @@ -473,7 +519,7 @@ async def list_workflow_messages( await _require_run(prisma_client, run_id, _read_scope_caller(user_api_key_dict)) try: - messages: Final = await WorkflowMessageRepository(prisma_client).table.find_many( + messages: Final[Sequence[object]] = await WorkflowMessageRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "asc"}, take=limit, diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 8ee1dd268e6..05df44242aa 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -6,7 +6,7 @@ This allows the same policy to be attached to multiple scopes. """ from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypedDict from litellm._logging import verbose_proxy_logger from litellm.repositories.table_repositories import PolicyAttachmentRepository @@ -18,9 +18,18 @@ from litellm.types.proxy.policy_engine import ( ) if TYPE_CHECKING: + from collections.abc import Sequence + + from prisma.models import LiteLLM_PolicyAttachmentTable + from litellm.proxy.utils import PrismaClient +class PolicyAttachmentMatch(TypedDict): + policy_name: str + matched_via: str + + class AttachmentRegistry: """ In-memory registry for storing and managing policy attachments. @@ -40,7 +49,7 @@ class AttachmentRegistry: ``` """ - def __init__(self): + def __init__(self) -> None: self._attachments: list[PolicyAttachment] = [] self._config_attachments: tuple[PolicyAttachment, ...] = () self._initialized: bool = False @@ -98,7 +107,7 @@ class AttachmentRegistry: """ return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context)] - def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[dict[str, Any]]: + def get_attached_policies_with_reasons(self, context: PolicyMatchContext) -> list[PolicyAttachmentMatch]: """ Get list of policy names and match reasons for the given context. @@ -107,8 +116,8 @@ class AttachmentRegistry: """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - results: Final[list[dict[str, Any]]] = [] - seen_policies: Final[set] = set() + results: Final[list[PolicyAttachmentMatch]] = [] + seen_policies: Final[set[str]] = set() for attachment in self._attachments: scope = attachment.to_policy_scope() @@ -280,7 +289,9 @@ class AttachmentRegistry: PolicyAttachmentDBResponse with the created attachment """ try: - created_attachment: Final = await PolicyAttachmentRepository(prisma_client).table.create( + created_attachment: Final[LiteLLM_PolicyAttachmentTable] = await PolicyAttachmentRepository( + prisma_client + ).table.create( data={ "policy_name": attachment_request.policy_name, "scope": attachment_request.scope, @@ -340,9 +351,9 @@ class AttachmentRegistry: """ try: # Get attachment before deleting - attachment: Final = await PolicyAttachmentRepository(prisma_client).table.find_unique( - where={"attachment_id": attachment_id} - ) + attachment: Final[LiteLLM_PolicyAttachmentTable | None] = await PolicyAttachmentRepository( + prisma_client + ).table.find_unique(where={"attachment_id": attachment_id}) if attachment is None: raise Exception(f"Attachment with ID {attachment_id} not found") @@ -375,9 +386,9 @@ class AttachmentRegistry: PolicyAttachmentDBResponse if found, None otherwise """ try: - attachment: Final = await PolicyAttachmentRepository(prisma_client).table.find_unique( - where={"attachment_id": attachment_id} - ) + attachment: Final[LiteLLM_PolicyAttachmentTable | None] = await PolicyAttachmentRepository( + prisma_client + ).table.find_unique(where={"attachment_id": attachment_id}) if attachment is None: return None @@ -413,7 +424,9 @@ class AttachmentRegistry: List of PolicyAttachmentDBResponse objects """ try: - attachments: Final = await PolicyAttachmentRepository(prisma_client).table.find_many( + attachments: Final[Sequence[LiteLLM_PolicyAttachmentTable]] = await PolicyAttachmentRepository( + prisma_client + ).table.find_many( order={"created_at": "desc"}, ) From 5339ec50e788a6bb9090380d1e43060e13bcdb97 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:42:20 -0700 Subject: [PATCH 021/126] fix(batches): persist managed file ids for cancelled/failed/expired batches When the batch cost poller found a batch in a terminal failed, expired, or cancelled state it wrote the provider response straight to the managed object table, so the stored blob kept raw provider file ids and a raw batch id. Since the row is final after batch_processed=True and the read paths only resolve existing managed ids, every later GET /batches/{id} and GET /batches leaked raw provider output and error file ids that clients cannot fetch through the proxy. The terminal branch now normalizes the response with ensure_batch_response_managed_file_ids before persisting, minting managed ids under the batch owner's identity POST /batches/{id}/cancel had the same gap: it called update_batch_in_database without the caller's auth context, so a cancel response that already carried provider file ids could never mint managed ids. The endpoint now forwards user_api_key_dict --- .../proxy/common_utils/check_batch_cost.py | 14 ++ litellm/proxy/batches_endpoints/endpoints.py | 1 + .../proxy_unit_tests/test_check_batch_cost.py | 151 +++++++++++++++--- .../proxy/batches_endpoints/test_endpoints.py | 11 ++ 4 files changed, 158 insertions(+), 19 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 3ed63b0d9ee..a12b0bb7170 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -656,6 +656,20 @@ class CheckBatchCost: elif response.status in ("failed", "expired", "cancelled"): try: + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ensure_batch_response_managed_file_ids, + ) + + response.id = job.unified_object_id + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=self.proxy_logging_obj.get_proxy_hook("managed_files"), + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + db_batch_object=job, + unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id), + ) update_data = { "status": response.status, "file_object": response.model_dump_json(), diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index deed665d3f2..7b5af6c1068 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -961,6 +961,7 @@ async def cancel_batch( prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, operation="cancel", + user_api_key_dict=user_api_key_dict, ) ### CALL HOOKS ### - modify outgoing data diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index a15abd023d8..d6c5d8fc809 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -499,7 +499,7 @@ class TestCheckBatchCost: must be written back with that status and batch_processed=True so it stops being polled forever. """ - from unittest.mock import patch + import base64 mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( return_value=0 @@ -511,7 +511,9 @@ class TestCheckBatchCost: mock_job = MagicMock() mock_job.id = "job-terminal-1" - mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.unified_object_id = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True @@ -527,23 +529,7 @@ class TestCheckBatchCost: mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" - - with ( - patch( - "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", - side_effect=[decoded_id, None], - ), - patch( - "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", - return_value="model-123", - ), - patch( - "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", - return_value="batch-456", - ), - ): - await check_batch_cost_instance.check_batch_cost() + await check_batch_cost_instance.check_batch_cost() assert ( mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 @@ -556,6 +542,133 @@ class TestCheckBatchCost: update_data["batch_processed"] is True ), "terminal-status update() must set batch_processed=True so polling stops" + @pytest.mark.asyncio + @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) + async def test_terminal_status_persists_managed_output_file_ids( + self, + check_batch_cost_instance, + mock_prisma_client, + mock_llm_router, + terminal_status, + ): + """A cancelled/failed/expired batch with provider output files must be persisted + with unified managed file IDs, never raw provider IDs. Raw IDs written here leak + to every later GET /batches/{id} and GET /batches because the terminal row is + final (batch_processed=True) and read paths only resolve, never mint. + """ + import base64 + import json + + from litellm.types.utils import LiteLLMBatch + + unified_batch_uid = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() + raw_output_file_id = "file-terminal-out-abc" + raw_error_file_id = "file-terminal-err-xyz" + raw_input_file_id = "file-terminal-in-123" + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch" + ).decode() + unified_output_file_id = base64.urlsafe_b64encode( + f"litellm_proxy:application/octet-stream;unified_id,u-1;llm_output_file_id,{raw_output_file_id}".encode() + ).decode() + unified_error_file_id = base64.urlsafe_b64encode( + f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() + ).decode() + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) + + input_file_row = MagicMock() + input_file_row.unified_file_id = unified_input_file_id + + def find_managed_file(where): + if where["flat_model_file_ids"]["has"] == raw_input_file_id: + return input_file_row + return None + + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) + + mock_job = MagicMock() + mock_job.id = "job-terminal-mint-1" + mock_job.unified_object_id = unified_batch_uid + mock_job.created_by = "user-1" + mock_job.team_id = "team-1" + + check_batch_cost_instance._has_batch_processed_column = True + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + response = LiteLLMBatch( + id="batch-456", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id=raw_input_file_id, + object="batch", + status=terminal_status, + output_file_id=raw_output_file_id, + error_file_id=raw_error_file_id, + ) + mock_llm_router.aretrieve_batch = AsyncMock(return_value=response) + + mock_hook = MagicMock() + mock_hook.get_unified_output_file_id.side_effect = [ + unified_output_file_id, + unified_error_file_id, + ] + mock_hook.store_unified_file_id = AsyncMock() + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) + + await check_batch_cost_instance.check_batch_cost() + + mock_hook.get_unified_output_file_id.assert_any_call( + output_file_id=raw_output_file_id, + model_id="model-123", + model_name="gpt-5-batch", + ) + mock_hook.get_unified_output_file_id.assert_any_call( + output_file_id=raw_error_file_id, + model_id="model-123", + model_name="gpt-5-batch", + ) + stored = { + next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_id"] + for c in mock_hook.store_unified_file_id.call_args_list + } + assert stored == { + raw_output_file_id: unified_output_file_id, + raw_error_file_id: unified_error_file_id, + } + for store_call in mock_hook.store_unified_file_id.call_args_list: + assert store_call.kwargs["user_api_key_dict"].user_id == "user-1" + assert store_call.kwargs["user_api_key_dict"].team_id == "team-1" + + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args + assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1"} + update_data = update_call.kwargs["data"] + assert update_data["status"] == terminal_status + assert update_data["batch_processed"] is True + persisted = json.loads(update_data["file_object"]) + assert persisted["id"] == unified_batch_uid + assert persisted["input_file_id"] == unified_input_file_id + assert persisted["output_file_id"] == unified_output_file_id + assert persisted["error_file_id"] == unified_error_file_id + assert raw_output_file_id not in update_data["file_object"] + assert raw_error_file_id not in update_data["file_object"] + @pytest.mark.asyncio async def test_raw_output_file_id_converted_to_managed_id( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index f2d37fbe842..64e1dcda5c9 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -1896,6 +1896,17 @@ async def test_cancel__unified_batch_id_routes_to_router(cancel_harness): assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel" +@pytest.mark.asyncio +async def test_cancel__db_write_receives_caller_auth(cancel_harness): + """update_batch_in_database can only mint managed IDs for a cancelled batch's + output files when it has an auth context, so cancel must forward the caller's.""" + caller = UserAPIKeyAuth(api_key="sk-test", user_id="user-cancel-1") + with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): + await call_cancel(cancel_harness, "batch-unified-blob", user=caller) + + assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller + + @pytest.mark.asyncio async def test_cancel__unified_missing_model_id_400(cancel_harness): # unified id with no model_id segment -> get_model_id returns None -> 400. From 0991692e684fb00601c9ad7fe0b7d9dbddb602e5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:44:25 -0700 Subject: [PATCH 022/126] test: roll back runtime model registrations between tests Since #35491, register_model records every registration in the process-global _runtime_registered_model_cost ledger, and every cost map swap replays that ledger on top of the freshly adopted map. Under pytest-xdist, any earlier test in the same worker that registered gpt-3.5-turbo leaked into TestPriceDataReloadIntegration::test_distributed_reload_check_function: the replay ballooned its sparse mocked entry into a full ModelInfo dict and failed the exact-equality assert, breaking the proxy-infra shard whenever loadscope happened to co-schedule such a test first (reruns cannot help since the pollution is process-wide) The autouse isolate_litellm_state fixture now snapshots the ledger before each test and restores it in place on teardown, so no test's registrations outlive it. A regression pair in test_conftest_isolation.py asserts the rollback --- tests/test_litellm/conftest.py | 9 +++++++++ tests/test_litellm/test_conftest_isolation.py | 13 +++++++++++++ 2 files changed, 22 insertions(+) create mode 100644 tests/test_litellm/test_conftest_isolation.py diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index f4aa1926d21..c0993051d33 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -19,6 +19,7 @@ sys.path.insert( import asyncio import litellm +from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.prompt_templates import ( image_handling as image_handling_module, @@ -238,6 +239,11 @@ def isolate_litellm_state(): if hasattr(litellm, _attr): original_state[_attr] = getattr(litellm, _attr) + original_runtime_registered_model_cost = { + model_key: dict(model_value) + for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() + } + # Store LiteLLM logger state. Some tests reconfigure handlers/propagation for # JSON logging and do not restore them, which breaks later caplog-based tests. logger_state = {} @@ -304,6 +310,9 @@ def isolate_litellm_state(): if hasattr(litellm, attr_name): setattr(litellm, attr_name, original_value) + litellm_utils_module._runtime_registered_model_cost.clear() + litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + # Restore logger configuration mutated by logging-focused tests. for logger in ALL_LOGGERS: original_logger_state = logger_state.get(logger.name) diff --git a/tests/test_litellm/test_conftest_isolation.py b/tests/test_litellm/test_conftest_isolation.py new file mode 100644 index 00000000000..88889ad7740 --- /dev/null +++ b/tests/test_litellm/test_conftest_isolation.py @@ -0,0 +1,13 @@ +import litellm +from litellm import utils as litellm_utils_module + +CANARY_MODEL = "conftest-isolation-canary-model" + + +def test_register_model_ledger_entry_is_scoped_to_this_test(): + litellm.register_model({CANARY_MODEL: {"litellm_provider": "openai", "input_cost_per_token": 0.001}}) + assert CANARY_MODEL in litellm_utils_module._runtime_registered_model_cost + + +def test_register_model_ledger_entry_was_rolled_back(): + assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost From 7d00f9d019f84be709a7515094fed4ce7bbee900 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:46:00 -0700 Subject: [PATCH 023/126] fix(managed_files): return unified output file ids from GET /batches list_user_batches parsed each stored batch blob and returned it as-is, so any row whose blob still carried raw provider file ids (for example a batch that reached a terminal state through the cost poller, or rows written before output registration existed) leaked raw output_file_id and error_file_id values that clients cannot fetch through the proxy. The list path now runs each row through ensure_batch_response_managed_file_ids, which swaps in existing managed ids and registers missing ones under the batch owner's identity, matching what GET /batches/{id} already does --- .../proxy/hooks/managed_files.py | 12 ++ .../proxy/hooks/test_managed_files.py | 142 ++++++++++++++++++ 2 files changed, 154 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 0036603bcd1..07a1f959940 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -31,6 +31,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + ensure_batch_response_managed_file_ids, get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, @@ -352,6 +353,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) batch_obj = LiteLLMBatch.model_validate(batch_data) batch_obj.id = batch.unified_object_id + await ensure_batch_response_managed_file_ids( + response=batch_obj, + managed_files_obj=self, + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=batch, + unified_batch_id=_is_base64_encoded_unified_file_id( + batch.unified_object_id + ), + ) batch_objects.append(batch_obj) except Exception as e: diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 50af6465d06..fc10a1257e1 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1813,6 +1813,148 @@ def _create_unified_batch_id(model_id: str, batch_id: str) -> str: return base64.urlsafe_b64encode(unified_str.encode()).decode().rstrip("=") +def _decode_unified_id(b64_id: str) -> str: + return base64.urlsafe_b64decode(b64_id + "=" * (-len(b64_id) % 4)).decode() + + +def _terminal_batch_record( + unified_batch_uid: str, + raw_input_file_id: str, + raw_output_file_id: str, + raw_error_file_id: str, +): + record = MagicMock() + record.unified_object_id = unified_batch_uid + record.created_by = "owner-user" + record.team_id = "owner-team" + record.status = "cancelled" + record.file_object = json.dumps( + { + "id": "batch-raw-456", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "cancelled", + "created_at": 1234567890, + "input_file_id": raw_input_file_id, + "output_file_id": raw_output_file_id, + "error_file_id": raw_error_file_id, + } + ) + return record + + +@pytest.mark.asyncio +async def test_list_batches_registers_and_returns_unified_output_file_ids(): + """A stored batch blob with raw provider file IDs (e.g. persisted by the cost + poller for a cancelled batch) must be listed with unified managed IDs, and the + output/error files must be registered in the managed file table so GET + /files/{id}/content can route them.""" + from litellm.proxy._types import UserAPIKeyAuth + + unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") + raw_input_file_id = "file-list-in-1" + raw_output_file_id = "file-list-out-1" + raw_error_file_id = "file-list-err-1" + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch" + ).decode() + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [ + _terminal_batch_record( + unified_batch_uid, raw_input_file_id, raw_output_file_id, raw_error_file_id + ) + ] + + input_file_row = MagicMock() + input_file_row.unified_file_id = unified_input_file_id + + def find_managed_file(where): + if where["flat_model_file_ids"]["has"] == raw_input_file_id: + return input_file_row + return None + + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=10, + ) + + listed = result["data"][0] + assert listed.id == unified_batch_uid + assert listed.input_file_id == unified_input_file_id + + decoded_output = _decode_unified_id(listed.output_file_id) + assert decoded_output.startswith("litellm_proxy") + assert f"llm_output_file_id,{raw_output_file_id}" in decoded_output + assert "llm_output_file_model_id,model-123" in decoded_output + assert "target_model_names,gpt-5-batch" in decoded_output + + decoded_error = _decode_unified_id(listed.error_file_id) + assert f"llm_output_file_id,{raw_error_file_id}" in decoded_error + + upsert_calls = prisma_client.db.litellm_managedfiletable.upsert.await_args_list + stored_raw_ids = { + c.kwargs["data"]["create"]["flat_model_file_ids"][0] for c in upsert_calls + } + assert stored_raw_ids == {raw_output_file_id, raw_error_file_id} + for c in upsert_calls: + assert c.kwargs["data"]["create"]["created_by"] == "owner-user" + assert c.kwargs["data"]["create"]["team_id"] == "owner-team" + + +@pytest.mark.asyncio +async def test_list_batches_resolves_existing_managed_rows_without_minting(): + """When the raw provider file IDs already have managed file rows, listing must + swap in the existing unified IDs and must not upsert duplicate rows.""" + from litellm.proxy._types import UserAPIKeyAuth + + unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") + raw_output_file_id = "file-list-out-existing" + existing_unified_output_id = base64.urlsafe_b64encode( + f"litellm_proxy:application/json;unified_id,u-9;llm_output_file_id,{raw_output_file_id}".encode() + ).decode() + + record = _terminal_batch_record( + unified_batch_uid, "file-list-in-9", raw_output_file_id, "" + ) + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [record] + + existing_row = MagicMock() + existing_row.unified_file_id = existing_unified_output_id + + def find_managed_file(where): + if where["flat_model_file_ids"]["has"] == raw_output_file_id: + return existing_row + return None + + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=10, + ) + + assert result["data"][0].output_file_id == existing_unified_output_id + prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited() + + @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception(): from litellm.proxy._types import UserAPIKeyAuth From fa47c47020cb702cd28bf98cde3b924ae4b3d491 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 21:33:24 -0700 Subject: [PATCH 024/126] fix(lint): measure the basedpyright budget gate in a gate-owned venv The gate previously measured whatever environment the caller happened to have. Locally that is the fat bootstrap venv (--extra proxy pulls in fastapi-sso, whose type info flips a reportUnnecessaryIsInstance diagnostic in ui_sso.py), while CI's publisher venv only has the proxy-dev and e2e-dev groups, so identical trees measured 866 locally vs 865 in CI and every local gate run breached by a phantom +1 scripts/type_check_gate.py now provisions .venv-typecheck itself: a frozen uv sync of the canonical proxy-dev and e2e-dev groups, the interpreter pinned to pyrightconfig.json's pythonVersion, plus the generated Prisma client. Every measurement pass is pinned to that env with --pythonpath, because basedpyright auto-detects a .venv in the project root and that auto-detection beats both PATH order and VIRTUAL_ENV, so the CLI flag is the only pin that actually works. The dependency-group set is folded into the environment fingerprint, so artifacts or caches recorded under a different group set never match and the gate falls back to computing base counts locally instead of comparing mismatched environments The publisher workflow drops its own install and prisma steps and lets the script build the measurement env, and the node heap for the full-tree pass drops from 12GB to 8GB (peak RSS measured at 5.4GB) --- .../publish-basedpyright-base-counts.yml | 18 +-- .gitignore | 1 + Makefile | 8 +- scripts/type_check_gate.py | 118 +++++++++++++++--- tests/test_litellm/test_type_check_gate.py | 99 +++++++++++++-- 5 files changed, 203 insertions(+), 41 deletions(-) diff --git a/.github/workflows/publish-basedpyright-base-counts.yml b/.github/workflows/publish-basedpyright-base-counts.yml index d9b034684a4..c85d30df0ce 100644 --- a/.github/workflows/publish-basedpyright-base-counts.yml +++ b/.github/workflows/publish-basedpyright-base-counts.yml @@ -43,22 +43,14 @@ jobs: with: version: "0.10.9" - - name: Install dependencies - run: | - uv sync --frozen --group proxy-dev --group e2e-dev - - # Mirrors test-linting.yml's lint job: basedpyright resolves Prisma's - # generated client only after `prisma generate`, and the published counts - # must match what that job would measure for the same tree. - - name: Generate Prisma client + # The gate provisions its own measurement env (.venv-typecheck: a frozen + # uv sync of its canonical dependency groups plus a generated Prisma + # client), so no install step here can drift from what local runs measure. + - name: Emit basedpyright counts for HEAD env: PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Emit basedpyright counts for HEAD - run: | - uv run --no-sync python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts" + python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts" counts_file=$(ls "$RUNNER_TEMP"/basedpyright-counts/basedpyright-counts-*.json) echo "COUNTS_ARTIFACT_NAME=$(basename "$counts_file" .json)" >> "$GITHUB_ENV" diff --git a/.gitignore b/.gitignore index 13f2202305d..3329f39ca10 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ .python-version .venv +.venv-typecheck .venv_policy_test .env .claude diff --git a/Makefile b/Makefile index 3e82e141c77..493828571b7 100644 --- a/Makefile +++ b/Makefile @@ -124,10 +124,10 @@ lint-fetch-base: git fetch origin litellm_internal_staging # Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated -# Prisma client, so basedpyright resolves the same modules CI does (without the generated -# client the DB wrappers typed against it degrade to Unknown, drifting the budget from -# CI's). --inexact tops up the venv instead of pruning the proxy extras gen:api and the -# running proxy need. +# Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The +# budget gate itself no longer measures here (scripts/type_check_gate.py provisions its +# own .venv-typecheck). --inexact tops up the venv instead of pruning the proxy extras +# gen:api and the running proxy need. lint-install: $(UV) sync --inexact --frozen --group proxy-dev --group e2e-dev $(UV_RUN) python scripts/prisma_generate_if_needed.py diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py index b8c6cb29a4e..f9d7f2912fd 100644 --- a/scripts/type_check_gate.py +++ b/scripts/type_check_gate.py @@ -12,6 +12,18 @@ a red once two PRs each land near the limit and their sum crosses it: the bystander's count equals its base, so it is spared, while any PR that actually grows the rule past its limit still fails. +Installed packages are part of the measurement: a typed dependency that is +present changes what basedpyright can prove (and therefore which diagnostics +fire) versus when it is absent, so counts from two differently provisioned +venvs are not comparable and their comparison produces phantom breaches no +diff hunk explains. The gate therefore provisions its own environment at +``.venv-typecheck`` (a frozen ``uv sync`` of one canonical dependency-group +set, plus a generated Prisma client) and runs every basedpyright pass from it, +so pre-commit, the CI lint job, and the artifact publisher measure one package +set by construction; re-syncs of an up-to-date env are a near-instant no-op. +The group set is folded into the cache and artifact fingerprint, so counts +recorded under a different set are never matched, only recomputed. + The gate runs basedpyright itself, for both the head and the base pass, with ``NODE_OPTIONS`` raised to the heap this repo needs: basedpyright's node process OOMs at the ~4 GB default, and when callers had to remember the flag, @@ -21,8 +33,8 @@ matters once some rule is over its limit, so when none is the base pass is skipped outright. When it is needed, it is a second basedpyright pass over a detached worktree at the merge-base, run under the same environment so import resolution matches, and its per-rule counts are cached under the repo's git -common dir keyed by merge-base commit, -``pyrightconfig.json``, and ``uv.lock``, so re-runs against the same branch +common dir keyed by merge-base commit, ``pyrightconfig.json``, ``uv.lock``, +and the dependency-group set, so re-runs against the same branch point pay for it once. A CI workflow publishes every staging commit's counts as an artifact (``--emit-counts-dir`` is its entry point), and on a disk-cache miss the gate first tries to download the merge-base's artifact through the ``gh`` @@ -64,10 +76,18 @@ CACHE_FILE_PREFIX = "basedpyright-base-" ARTIFACT_NAME_PREFIX = "basedpyright-counts-" GH_TIMEOUT_SECONDS = 10 +# The one environment every basedpyright pass measures in. The group set is +# the slim one the CI publisher has always installed (not bootstrap's fatter +# --extra proxy env), so the committed budgets stay valid; changing it re-keys +# every cache and artifact fingerprint, so stale counts can never be matched. +TYPECHECK_ENV_DIR = REPO_ROOT / ".venv-typecheck" +TYPECHECK_DEP_GROUPS = ("proxy-dev", "e2e-dev") +PRISMA_GENERATE_SCRIPT = REPO_ROOT / "scripts" / "prisma_generate_if_needed.py" + # basedpyright's node process needs more than the ~4 GB default heap on this # repo; appended last so it wins node's last-flag-wins resolution over any # caller-set value while preserving the caller's other NODE_OPTIONS flags. -NODE_HEAP_OPTION = "--max-old-space-size=12288" +NODE_HEAP_OPTION = "--max-old-space-size=8192" # Bucket for a basedpyright diagnostic with no `rule`. Counted so it's gated. UNCODED = "" @@ -129,14 +149,78 @@ def node_options_with_heap(base_env: Mapping[str, str]) -> str: return f"{base_env.get('NODE_OPTIONS', '')} {NODE_HEAP_OPTION}".strip() -def run_basedpyright(cwd: Path = REPO_ROOT) -> str: - """One basedpyright pass over `cwd` with the raised node heap exported. +def typecheck_python_version() -> str | None: + """The interpreter version to build the owned env with, read from + pyrightconfig's `pythonVersion` so the packages installed for basedpyright + to see always come from the same version it type-checks against.""" + try: + config = json.loads(PYRIGHT_CONFIG.read_text()) + except (OSError, ValueError): + return None + version: Final = config.get("pythonVersion") if isinstance(config, dict) else None + return version if isinstance(version, str) else None - Exit 0 (clean) and 1 (errors found) are both output-bearing runs; anything - else is a crash and fails loudly instead of reading as zero errors.""" - exe = shutil.which("basedpyright") or "basedpyright" + +def typecheck_env_commands(env_dir: Path = TYPECHECK_ENV_DIR) -> tuple[tuple[str, ...], ...]: + python_pin: Final = typecheck_python_version() + sync: Final = ( + "uv", + "sync", + "--frozen", + *(("--python", python_pin) if python_pin else ()), + *(flag for group in TYPECHECK_DEP_GROUPS for flag in ("--group", group)), + ) + generate: Final = (str(env_dir / "bin" / "python"), str(PRISMA_GENERATE_SCRIPT)) + return (sync, generate) + + +def _run_provision_step(cmd: tuple[str, ...], env: Mapping[str, str]) -> int: proc = subprocess.run( - [exe, "--outputjson"], + list(cmd), cwd=REPO_ROOT, env=dict(env), capture_output=True, text=True + ) + if proc.returncode != 0: + sys.stderr.write(proc.stdout) + sys.stderr.write(proc.stderr) + return proc.returncode + + +def ensure_typecheck_env( + env_dir: Path = TYPECHECK_ENV_DIR, + run: Callable[[tuple[str, ...], Mapping[str, str]], int] = _run_provision_step, +) -> Path: + """Sync the gate-owned venv (and its generated Prisma client) before a + measurement pass. Unconditional on purpose: an up-to-date env makes both + steps near-instant no-ops, and skipping them on a heuristic is how the + measured environment and the fingerprinted one drift apart.""" + env: Final = {**os.environ, "UV_PROJECT_ENVIRONMENT": str(env_dir)} + for cmd in typecheck_env_commands(env_dir): + if run(cmd, env) != 0: + raise SystemExit( + f"could not provision the type-check environment at {env_dir}: " + f"`{' '.join(cmd)}` failed" + ) + return env_dir + + +def run_basedpyright(cwd: Path = REPO_ROOT, env_dir: Path = TYPECHECK_ENV_DIR) -> str: + """One basedpyright pass over `cwd` from the gate-owned venv, with the + raised node heap exported. + + `--pythonpath` pins import resolution to the owned env's interpreter; it is + the only pin that works, because basedpyright auto-detects a `.venv` in the + project root and that beats both PATH order and VIRTUAL_ENV, silently + measuring the caller's fatter venv (whose extra typed packages flip + diagnostics) whenever the repo has one. Exit 0 (clean) and 1 (errors + found) are both output-bearing runs; anything else is a crash and fails + loudly instead of reading as zero errors.""" + bin_dir: Final = env_dir / "bin" + proc = subprocess.run( + [ + str(bin_dir / "basedpyright"), + "--outputjson", + "--pythonpath", + str(bin_dir / "python"), + ], cwd=cwd, capture_output=True, text=True, @@ -208,11 +292,16 @@ def over_ceiling( ) -def environment_fingerprints() -> tuple[str, ...]: - return tuple( - hashlib.sha256(path.read_bytes()).hexdigest() - for path in (PYRIGHT_CONFIG, UV_LOCK) - if path.exists() +def environment_fingerprints( + dep_groups: tuple[str, ...] = TYPECHECK_DEP_GROUPS, +) -> tuple[str, ...]: + return ( + *( + hashlib.sha256(path.read_bytes()).hexdigest() + for path in (PYRIGHT_CONFIG, UV_LOCK) + if path.exists() + ), + "groups:" + ",".join(dep_groups), ) @@ -560,6 +649,7 @@ def main() -> None: parser.add_argument("--update", action="store_true") parser.add_argument("--emit-counts-dir", type=Path) args = parser.parse_args() + ensure_typecheck_env() head = count_basedpyright(run_basedpyright()) if args.emit_counts_dir is not None: cmd_emit_counts( diff --git a/tests/test_litellm/test_type_check_gate.py b/tests/test_litellm/test_type_check_gate.py index e381f787d78..2d9731db53f 100644 --- a/tests/test_litellm/test_type_check_gate.py +++ b/tests/test_litellm/test_type_check_gate.py @@ -1,6 +1,5 @@ import importlib.util import json -import os import subprocess from pathlib import Path @@ -84,33 +83,51 @@ def test_node_options_with_heap_appends_after_caller_flags_so_it_wins(): assert merged == f"--max-old-space-size=4096 --no-warnings {gate.NODE_HEAP_OPTION}" -def _stub_basedpyright(tmp_path, monkeypatch, script_body): - stub = tmp_path / "basedpyright" +def _stub_env(tmp_path, script_body): + bin_dir = tmp_path / "bin" + bin_dir.mkdir(exist_ok=True) + stub = bin_dir / "basedpyright" stub.write_text(f"#!/bin/sh\n{script_body}\n") stub.chmod(0o755) - monkeypatch.setenv("PATH", str(tmp_path), prepend=os.pathsep) + return tmp_path def test_run_basedpyright_exports_the_raised_heap_to_the_child(tmp_path, monkeypatch): captured = tmp_path / "node_options.txt" - _stub_basedpyright( + env_dir = _stub_env( tmp_path, - monkeypatch, f'echo "$NODE_OPTIONS" > "{captured}"\necho \'{{"generalDiagnostics": []}}\'', ) monkeypatch.delenv("NODE_OPTIONS", raising=False) - assert json.loads(gate.run_basedpyright(cwd=tmp_path)) == {"generalDiagnostics": []} + assert json.loads(gate.run_basedpyright(cwd=tmp_path, env_dir=env_dir)) == { + "generalDiagnostics": [] + } assert captured.read_text().strip() == gate.NODE_HEAP_OPTION -def test_run_basedpyright_fails_loudly_on_a_crash_exit_code(tmp_path, monkeypatch): +def test_run_basedpyright_pins_import_resolution_to_the_owned_env(tmp_path): + # basedpyright auto-detects a `.venv` in the project root, and that beats + # PATH order and VIRTUAL_ENV; only an explicit --pythonpath keeps the + # caller's fatter venv (whose extra typed packages flip diagnostics vs CI) + # out of the measurement. + captured = tmp_path / "argv.txt" + env_dir = _stub_env( + tmp_path, + f'echo "$@" > "{captured}"\necho \'{{"generalDiagnostics": []}}\'', + ) + gate.run_basedpyright(cwd=tmp_path, env_dir=env_dir) + argv = captured.read_text().split() + assert argv[argv.index("--pythonpath") + 1] == str(env_dir / "bin" / "python") + + +def test_run_basedpyright_fails_loudly_on_a_crash_exit_code(tmp_path): import pytest # 134 is SIGABRT, what node dies with on a heap OOM; it must never read as a # clean zero-error run. - _stub_basedpyright(tmp_path, monkeypatch, "exit 134") + env_dir = _stub_env(tmp_path, "exit 134") with pytest.raises(SystemExit): - gate.run_basedpyright(cwd=tmp_path) + gate.run_basedpyright(cwd=tmp_path, env_dir=env_dir) def test_at_or_under_ceiling_passes(): @@ -248,6 +265,68 @@ def test_cache_key_changes_with_base_point_and_each_fingerprint(): assert gate.cache_key("abc", ("cfg", "lock2")) != key +def test_fingerprints_carry_the_dependency_group_set(): + # Counts measured under one group set must never be compared against + # another's: the fingerprint difference re-keys every cache entry and + # artifact name, so a changed canonical set falls back to recompute. + assert gate.environment_fingerprints() == gate.environment_fingerprints() + assert gate.environment_fingerprints( + dep_groups=("proxy-dev",) + ) != gate.environment_fingerprints(dep_groups=("proxy-dev", "e2e-dev")) + assert gate.environment_fingerprints()[-1] == "groups:" + ",".join( + gate.TYPECHECK_DEP_GROUPS + ) + + +def test_env_commands_sync_the_canonical_groups_then_generate_prisma(): + sync, generate = gate.typecheck_env_commands(Path("/envdir")) + assert sync[:3] == ("uv", "sync", "--frozen") + adjacent = list(zip(sync, sync[1:])) + for group in gate.TYPECHECK_DEP_GROUPS: + assert ("--group", group) in adjacent + assert generate == ( + str(Path("/envdir") / "bin" / "python"), + str(gate.PRISMA_GENERATE_SCRIPT), + ) + + +def test_env_interpreter_pin_tracks_pyrightconfigs_python_version(): + configured = json.loads((ROOT / "pyrightconfig.json").read_text())[ + "pythonVersion" + ] + assert gate.typecheck_python_version() == configured + sync = gate.typecheck_env_commands()[0] + assert sync[sync.index("--python") + 1] == configured + + +def test_ensure_env_targets_the_owned_dir_and_runs_sync_then_generate(tmp_path): + calls = [] + + def runner(cmd, env): + calls.append((cmd[:2], env["UV_PROJECT_ENVIRONMENT"])) + return 0 + + assert gate.ensure_typecheck_env(env_dir=tmp_path, run=runner) == tmp_path + assert calls == [ + (("uv", "sync"), str(tmp_path)), + ((str(tmp_path / "bin" / "python"), str(gate.PRISMA_GENERATE_SCRIPT)), str(tmp_path)), + ] + + +def test_ensure_env_fails_loudly_and_stops_at_the_first_failed_step(tmp_path): + import pytest + + calls = [] + + def failing(cmd, env): + calls.append(cmd) + return 2 + + with pytest.raises(SystemExit): + gate.ensure_typecheck_env(env_dir=tmp_path, run=failing) + assert len(calls) == 1 + + def test_cached_counts_round_trip(tmp_path): path = gate.cache_path(tmp_path, "abc123", ("f1", "f2")) gate.store_counts(tmp_path, path, "abc123", {"reportAny": 3, "reportCall": 1}) From e2cb01c87cda40b95ac86fe02bec92e7db892112 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 04:37:48 +0000 Subject: [PATCH 025/126] fix(types): correct annotations that were false about their runtime values An adversarial review of the previous commit found annotations that described what the code wished were true rather than what flows through. A false annotation is worse than the Any it replaced, since it launders a wrong assumption past the type checker. - purview: `_resolve_user_id` claimed every request-body value was a Mapping, contradicting `_resolve_trusted_user_id` one method over, which types the same argument `Mapping[str, object]`. `_should_block` claimed every Graph response value was a sequence of str->str mappings and was not assignable from its own producer's return type. - cato: `_CatoAnalyzeResponse.required_action` was required and non-nullable while the API returns null, as seven fixtures in the guardrail's own suite assert. `analysis_result` had the same problem. The streaming hook narrowed an override parameter below what `ProxyLogging` actually passes it. - marketplace: `_PluginRecord.manifest_json` was `str` against a nullable column. Making it honest surfaced a latent crash, covered below. - ownership: two functions took an attribute Protocol while their own bodies branch on `isinstance(response, dict)`, which no Protocol can satisfy. - openapi generator: `paths` claimed every path-item value was an operation, though path items also carry `parameters`, `summary` and `$ref`. - custom openapi spec: a TypedDict asserted a shape that the function returns raw Pydantic sub-schemas out of. Reverted to Any, which is imprecise but not false. `get_marketplace` did an unguarded `json.loads` on the nullable `manifest_json` inside an `except json.JSONDecodeError`, which cannot catch the TypeError a NULL raises, so one NULL row 500s the endpoint. It now skips the plugin like the file's other two read sites already do, with a regression test that fails without the guard. Where honesty cost precision, precision lost. `_should_block` went back to its original signature entirely: the narrowing needed to type it turned a fail-closed DLP control fail-open, because the TypeError it used to raise on a malformed response reached `except Exception` and became a 400. --- litellm/integrations/galileo.py | 17 +++++-- .../mcp_server/openapi_to_mcp_generator.py | 46 +++++++++---------- .../claude_code_marketplace.py | 6 +-- .../proxy/common_utils/custom_openapi_spec.py | 9 +--- .../proxy/container_endpoints/ownership.py | 24 ++++------ .../cato_networks/cato_networks.py | 26 ++++++----- .../guardrail_hooks/microsoft_purview/base.py | 13 +++--- .../test_claude_code_marketplace.py | 27 ++++++++++- 8 files changed, 98 insertions(+), 70 deletions(-) diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index f9ec825e922..2c9ac63941c 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -6,7 +6,7 @@ import re import uuid from collections.abc import Mapping, Sequence from datetime import datetime, timezone, tzinfo -from typing import Any, Final, cast +from typing import Any, Final, TypedDict, cast import httpx from pydantic import BaseModel, Field @@ -35,6 +35,17 @@ GALILEO_CLOUD_API_BASE_URL: Final = "https://api.galileo.ai" GALILEO_MAX_IN_MEMORY_RECORDS: Final = 1000 +class GalileoStandardLoggingFields(TypedDict, total=False): + call_type: str + model: str + prompt_tokens: int + completion_tokens: int + total_tokens: int + response_cost: float + startTime: float + endTime: float + + class LLMResponse(BaseModel): latency_ms: int status_code: int @@ -60,7 +71,7 @@ class LLMResponse(BaseModel): class GalileoObserve(CustomLogger): def __init__(self) -> None: - self.in_memory_records: list[dict[str, Any]] = [] + self.in_memory_records: list[Mapping[str, object]] = [] self.batch_size = 1 self.api_key = os.getenv("GALILEO_API_KEY") self.project_id = os.getenv("GALILEO_PROJECT_ID") @@ -648,7 +659,7 @@ class GalileoObserve(CustomLogger): ) return - slo: Final[Mapping[str, Any] | None] = kwargs.get("standard_logging_object") + slo: Final[GalileoStandardLoggingFields | None] = kwargs.get("standard_logging_object") if slo is None: verbose_logger.debug("Galileo Logger: no standard_logging_object in kwargs, skipping") return diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 41057840684..2cc761f99ed 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -47,11 +47,7 @@ from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) -_OpenAPIObject: TypeAlias = Mapping[str, Any] - - -class _OpenAPIParameterSchema(TypedDict, total=False): - type: str +_OpenAPIParameter: TypeAlias = Mapping[str, Any] class _OpenAPIJSONSchema(TypedDict, total=False): @@ -72,16 +68,18 @@ class _OpenAPIOperation(TypedDict, total=False): operationId: str summary: str description: str - parameters: Sequence[_OpenAPIObject] + parameters: Sequence[_OpenAPIParameter] requestBody: _OpenAPIRequestBody class _OpenAPIPathItem(TypedDict, total=False): - parameters: Sequence[_OpenAPIObject] + summary: str + description: str + parameters: Sequence[_OpenAPIParameter] class _OpenAPIComponents(TypedDict, total=False): - parameters: Mapping[str, _OpenAPIObject] + parameters: Mapping[str, _OpenAPIParameter] # Store the base URL and headers globally @@ -161,7 +159,7 @@ async def load_openapi_spec_async(filepath: str) -> dict[str, Any]: return json.load(f) -def get_base_url(spec: _OpenAPIObject, spec_path: str | None = None) -> str: +def get_base_url(spec: Mapping[str, Any], spec_path: str | None = None) -> str: """Extract base URL from OpenAPI spec.""" # OpenAPI 3.x if "servers" in spec and spec["servers"]: @@ -212,7 +210,9 @@ def get_base_url(spec: _OpenAPIObject, spec_path: str | None = None) -> str: return "" -def _resolve_ref(param: _OpenAPIObject, component_params: Mapping[str, _OpenAPIObject]) -> _OpenAPIObject | None: +def _resolve_ref( + param: _OpenAPIParameter, component_params: Mapping[str, _OpenAPIParameter] +) -> _OpenAPIParameter | None: """Resolve a single parameter, following a $ref if present. Returns the resolved param dict, or None if the $ref target is absent from @@ -226,8 +226,8 @@ def _resolve_ref(param: _OpenAPIObject, component_params: Mapping[str, _OpenAPIO def _resolve_param_list( - raw: Sequence[_OpenAPIObject], component_params: Mapping[str, _OpenAPIObject] -) -> list[_OpenAPIObject]: + raw: Sequence[_OpenAPIParameter], component_params: Mapping[str, _OpenAPIParameter] +) -> list[_OpenAPIParameter]: """Resolve $refs in a parameter list, dropping any unresolvable entries.""" result: Final = [] for p in raw: @@ -256,7 +256,7 @@ def resolve_operation_params( merged with the operation-level params; operation-level wins when the same ``name`` + ``in`` combination appears in both. """ - component_params: Final[Mapping[str, _OpenAPIObject]] = components.get("parameters", {}) + component_params: Final[Mapping[str, _OpenAPIParameter]] = components.get("parameters", {}) path_level: Final = _resolve_param_list(path_item.get("parameters", []), component_params) op_level: Final = _resolve_param_list(operation.get("parameters", []), component_params) op_keys: Final = {(p["name"], p.get("in")) for p in op_level} @@ -266,9 +266,8 @@ def resolve_operation_params( return result -def extract_parameters(operation: _OpenAPIObject) -> tuple[Sequence[str], Sequence[str], Sequence[str]]: +def extract_parameters(operation: Mapping[str, Any]) -> tuple[Sequence[str], Sequence[str], Sequence[str]]: """Extract parameter names from OpenAPI operation.""" - param: _OpenAPIObject path_params: Final = [] query_params: Final = [] body_params: Final = [] @@ -278,7 +277,7 @@ def extract_parameters(operation: _OpenAPIObject) -> tuple[Sequence[str], Sequen for param in operation["parameters"]: if "name" not in param: continue - param_name: str = param["name"] + param_name = param["name"] if param.get("in") == "path": path_params.append(param_name) elif param.get("in") == "query": @@ -293,9 +292,8 @@ def extract_parameters(operation: _OpenAPIObject) -> tuple[Sequence[str], Sequen return path_params, query_params, body_params -def build_input_schema(operation: _OpenAPIObject) -> dict[str, Any]: +def build_input_schema(operation: Mapping[str, Any]) -> dict[str, Any]: """Build MCP input schema from OpenAPI operation.""" - param: _OpenAPIObject properties: Final = {} required: Final = [] @@ -304,9 +302,9 @@ def build_input_schema(operation: _OpenAPIObject) -> dict[str, Any]: for param in operation["parameters"]: if "name" not in param: continue - param_name: str = param["name"] - param_schema: _OpenAPIParameterSchema = param.get("schema", {}) - param_type: str = param_schema.get("type", "string") + param_name = param["name"] + param_schema = param.get("schema", {}) + param_type = param_schema.get("type", "string") properties[param_name] = { "type": param_type, @@ -391,7 +389,7 @@ def _merge_openapi_tool_request_headers( def create_tool_function( path: str, method: str, - operation: _OpenAPIObject, + operation: Mapping[str, Any], base_url: str, headers: dict[str, str] | None = None, ): @@ -492,9 +490,9 @@ def create_tool_function( return tool_function -def register_tools_from_openapi(spec: _OpenAPIObject, base_url: str) -> None: +def register_tools_from_openapi(spec: Mapping[str, Any], base_url: str) -> None: """Register MCP tools from OpenAPI specification.""" - paths: Final[Mapping[str, Mapping[str, _OpenAPIOperation]]] = spec.get("paths", {}) + paths: Final[Mapping[str, Mapping[str, Any]]] = spec.get("paths", {}) used_names: Final = set() for path, path_item in paths.items(): diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index d340ca2ced3..46ee9b0911d 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -47,7 +47,7 @@ class _PluginRecord(Protocol): name: str version: str | None description: str | None - manifest_json: str + manifest_json: str | None enabled: bool created_at: datetime | None updated_at: datetime | None @@ -108,7 +108,7 @@ async def get_marketplace(): plugin_list: Final = [] for plugin in plugins: try: - manifest: Mapping[str, object] = json.loads(plugin.manifest_json) + manifest: Mapping[str, object] = json.loads(plugin.manifest_json or "{}") except json.JSONDecodeError: verbose_proxy_logger.warning("Plugin %s has invalid manifest JSON, skipping", plugin.name) continue @@ -431,7 +431,7 @@ async def get_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - manifest: Final[Mapping[str, object]] = json.loads(plugin.manifest_json) if plugin.manifest_json else {} + manifest: Final[Mapping[str, object]] = json.loads(plugin.manifest_json or "{}") if plugin.manifest_json else {} return { "id": plugin.id, diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index 3bedb6f720c..bc7b80801fe 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -1,14 +1,9 @@ from collections.abc import Mapping, Sequence -from typing import Any, Final, TypedDict +from typing import Any, Final from litellm._logging import verbose_proxy_logger -class _FieldSchema(TypedDict, total=False): - type: str - anyOf: Sequence["_FieldSchema"] - - class CustomOpenAPISpec: """ Handler for customizing OpenAPI specifications with Pydantic models @@ -198,7 +193,7 @@ class CustomOpenAPISpec: return schema @staticmethod - def _extract_field_schema(field_def: _FieldSchema) -> _FieldSchema: + def _extract_field_schema(field_def: dict[str, Any]) -> dict[str, Any]: """ Extract a simple schema from a Pydantic field definition for parameter display. diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index c40f4d4fef2..a559ab49cfa 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -40,13 +40,6 @@ class _ManagedObjectTable(Protocol): async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> _ManagedObjectRow | None: ... -class _ContainerListResponse(Protocol): - data: Sequence[object] - first_id: str | None - last_id: str | None - has_more: bool - - CONTAINER_OBJECT_PURPOSE: Final = "container" # 60s LRU/TTL cache absorbs every container access check before it reaches @@ -279,7 +272,8 @@ async def _get_container_owner(original_container_id: str, custom_llm_provider: if prisma_client is None: return None - row: Final[_ManagedObjectRow | None] = await ManagedObjectRepository(prisma_client).table.find_first( + table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table + row: Final[_ManagedObjectRow | None] = await table.find_first( where={ "model_object_id": model_object_id, "file_purpose": CONTAINER_OBJECT_PURPOSE, @@ -315,7 +309,8 @@ async def _get_stored_container_id(original_container_id: str, custom_llm_provid if prisma_client is None: return None - row: Final[_ManagedObjectRow | None] = await ManagedObjectRepository(prisma_client).table.find_first( + table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table + row: Final[_ManagedObjectRow | None] = await table.find_first( where={ "model_object_id": model_object_id, "file_purpose": CONTAINER_OBJECT_PURPOSE, @@ -359,9 +354,7 @@ def _get_container_list_data(response: object) -> Sequence[object] | None: return data if isinstance(data, list) else None -def _set_container_list_data( - response: _ContainerListResponse, data: list[object], removed_filtered_items: bool = False -) -> _ContainerListResponse: +def _set_container_list_data(response: Any, data: list[object], removed_filtered_items: bool = False) -> object: if isinstance(response, dict): response["data"] = data if data: @@ -401,7 +394,8 @@ async def _get_allowed_container_ids( if prisma_client is None: return set() - rows: Final[Sequence[_ManagedObjectRow]] = await ManagedObjectRepository(prisma_client).table.find_many( + table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table + rows: Final[Sequence[_ManagedObjectRow]] = await table.find_many( where={ "file_purpose": CONTAINER_OBJECT_PURPOSE, "created_by": {"in": owner_scopes}, @@ -416,10 +410,10 @@ async def _get_allowed_container_ids( async def filter_container_list_response( - response: _ContainerListResponse, + response: object, user_api_key_dict: UserAPIKeyAuth, custom_llm_provider: str, -) -> _ContainerListResponse: +) -> object: if is_proxy_admin(user_api_key_dict): return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 79867d492a1..958f84e18de 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -9,7 +9,7 @@ import contextlib import json import os import ssl -from collections.abc import AsyncGenerator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence from ssl import SSLContext from typing import TYPE_CHECKING, Any, Final @@ -70,9 +70,13 @@ class _CatoRedactedChat(TypedDict, total=False): all_redacted_messages: Sequence[_CatoRedactedMessage] +class _CatoAnalysisResult(TypedDict, total=False): + policy_drill_down: Mapping[str, object] + + class _CatoAnalyzeResponse(TypedDict): - required_action: _CatoRequiredAction - analysis_result: NotRequired[Mapping[str, Mapping[str, object]]] + required_action: NotRequired[_CatoRequiredAction | None] + analysis_result: NotRequired[_CatoAnalysisResult] redacted_chat: NotRequired[_CatoRedactedChat] @@ -202,7 +206,7 @@ class CatoNetworksGuardrail(CustomGuardrail): return [] @staticmethod - def _iter_schema_string_refs(data: dict): + def _iter_schema_string_refs(data: Mapping[str, Any]): """Yield ``(container, key)`` for every non-empty schema string the proxy forwards to the model inside tool/function and structured-output schemas: each ``tools[].function`` and legacy ``functions[]`` entry plus the @@ -244,7 +248,7 @@ class CatoNetworksGuardrail(CustomGuardrail): stack.extend(reversed(node)) @classmethod - def _extra_inspection_sources(cls, data: dict) -> Sequence[tuple[str, Sequence[Mapping[str, str]]]]: + def _extra_inspection_sources(cls, data: Mapping[str, Any]) -> Sequence[tuple[str, Sequence[Mapping[str, str]]]]: """Text the proxy forwards to the model outside chat ``messages``: Responses-API ``input`` and ``instructions``, legacy completion ``prompt`` and tool/function/``response_format`` schema strings. Returned @@ -305,8 +309,8 @@ class CatoNetworksGuardrail(CustomGuardrail): def _handle_block_action( self, - analysis_result: Mapping[str, Mapping[str, object]], - required_action: _CatoRequiredAction, + analysis_result: _CatoAnalysisResult, + required_action: Any, ) -> None: detection_message: Final = required_action.get("detection_message", None) verbose_proxy_logger.info( @@ -420,8 +424,8 @@ class CatoNetworksGuardrail(CustomGuardrail): def _handle_block_action_on_output( self, - analysis_result: Mapping[str, Mapping[str, object]], - required_action: _CatoRequiredAction, + analysis_result: _CatoAnalysisResult, + required_action: Any, ) -> None: detection_message: Final = required_action.get("detection_message", None) verbose_proxy_logger.info( @@ -570,7 +574,7 @@ class CatoNetworksGuardrail(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: AsyncGenerator[object, None], + response: AsyncIterable[object], request_data: dict, ) -> AsyncGenerator[ModelResponseStream, None]: from litellm.proxy.proxy_server import StreamingCallbackError @@ -622,7 +626,7 @@ class CatoNetworksGuardrail(CustomGuardrail): async def forward_the_stream_to_cato( self, websocket: ClientConnection, - response_iter: AsyncGenerator[object, None], + response_iter: AsyncIterable[object], ) -> None: async for chunk in response_iter: if isinstance(chunk, BaseModel): diff --git a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py index 80beb90cf27..3f666178970 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/base.py @@ -3,7 +3,7 @@ import time import uuid from collections import OrderedDict from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Any, Final from typing_extensions import NotRequired, TypedDict @@ -270,9 +270,7 @@ class PurviewGuardrailBase: # User ID resolution # ------------------------------------------------------------------ - def _resolve_user_id( - self, data: Mapping[str, Mapping[str, object]], user_api_key_dict: "UserAPIKeyAuth" - ) -> str | None: + def _resolve_user_id(self, data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth") -> str | None: """Resolve the Entra user object ID from request data or auth context. Returns the strongest available identity walking down four sources, in @@ -295,7 +293,10 @@ class PurviewGuardrailBase: if hasattr(user_api_key_dict, "end_user_id") and user_api_key_dict.end_user_id: return str(user_api_key_dict.end_user_id) - metadata: Final[Mapping[str, object]] = data.get("metadata") or data.get("litellm_metadata") or {} + metadata_value: Final[object] = data.get("metadata") or data.get("litellm_metadata") or {} + if not isinstance(metadata_value, Mapping): + return None + metadata: Final[Mapping[str, object]] = metadata_value uid = metadata.get("user_api_key_user_id") if uid: return str(uid) @@ -359,7 +360,7 @@ class PurviewGuardrailBase: # ------------------------------------------------------------------ @staticmethod - def _should_block(response: Mapping[str, Sequence[Mapping[str, str]]]) -> bool: + def _should_block(response: dict[str, Any]) -> bool: """Return True if any policyAction requires blocking.""" for action in response.get("policyActions", []): odata_type = action.get("@odata.type", "") diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py index f94cd471a01..18e0f2cb559 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -18,13 +18,14 @@ from litellm.types.proxy.claude_code_endpoints import ( UpdatePluginRequest, ) from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( + get_marketplace, register_plugin, update_plugin, ) def _make_mock_prisma(): - """Stateful prisma mock that supports find_unique, create, and update.""" + """Stateful prisma mock that supports find_unique, find_many, create, and update.""" store: dict = {} mock_client = MagicMock() @@ -34,6 +35,12 @@ def _make_mock_prisma(): async def _find_unique(where): return store.get(where.get("name")) + async def _find_many(where=None): + records = list(store.values()) + if where and "enabled" in where: + return [r for r in records if r.enabled == where["enabled"]] + return records + async def _create(data): record = MagicMock() record.id = "test-id" @@ -52,6 +59,7 @@ def _make_mock_prisma(): return record mock_table.find_unique = AsyncMock(side_effect=_find_unique) + mock_table.find_many = AsyncMock(side_effect=_find_many) mock_table.create = AsyncMock(side_effect=_create) mock_table.update = AsyncMock(side_effect=_update) mock_client.db.litellm_claudecodeplugintable = mock_table @@ -211,6 +219,23 @@ async def test_update_plugin_db_error_maps_to_structured_500(): assert "connection lost" in exc_info.value.detail["error"] +@pytest.mark.asyncio +async def test_get_marketplace_skips_plugin_with_null_manifest(): + await register_plugin( + request=RegisterPluginRequest(name="good-plugin", source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + + table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable + await table.create(data={"name": "null-manifest-plugin", "manifest_json": None, "enabled": True}) + + response = await get_marketplace() + + assert response.status_code == 200 + body = json.loads(response.body) + assert [plugin["name"] for plugin in body["plugins"]] == ["good-plugin"] + + @pytest.mark.asyncio async def test_register_plugin_git_subdir_missing_url(): """git-subdir without url field raises HTTP 400.""" From 4ab7a33d2c296923b0a2486fc8e3783197385238 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 04:50:01 +0000 Subject: [PATCH 026/126] chore(lint): ratchet lint budgets down by what this branch fixed Lowers the committed ceilings so the headroom shrinks by exactly what was cleared instead of leaving stale slack for the next change to spend. basedpyright -653 errors across 48 rules, with reportAny 29204 -> 28842 and reportExplicitAny 9227 -> 9105. Strict ruff -80 violations, led by ANN401 -59. LIT rules -85, led by LIT001 -76. --- basedpyright-code-budget.json | 16 ++++++++-------- ruff-strict-budget.json | 18 +++++++++--------- type-discipline-budget.json | 4 ++-- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 27d96e415fd..62e2ef2c1df 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 29204 + "limit": 28842 }, "reportArgumentType": { "limit": 2635 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9227 + "limit": 9105 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5850 + "limit": 5843 }, "reportMissingTypeArgument": { - "limit": 15833 + "limit": 15816 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45242 + "limit": 45207 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40340 + "limit": 40297 }, "reportUnknownParameterType": { - "limit": 20293 + "limit": 20272 }, "reportUnknownVariableType": { - "limit": 31796 + "limit": 31750 }, "reportUnnecessaryCast": { "limit": 122 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 421b424757b..81f4a8b97fb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,30 +1,30 @@ { "ANN001": { - "limit": 3126 + "limit": 3121 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 836 + "limit": 834 }, "ANN201": { - "limit": 2037 + "limit": 2033 }, "ANN202": { - "limit": 869 + "limit": 865 }, "ANN204": { - "limit": 715 + "limit": 713 }, "ANN205": { - "limit": 115 + "limit": 114 }, "ANN206": { "limit": 133 }, "ANN401": { - "limit": 1689 + "limit": 1630 }, "ASYNC230": { "limit": 11 @@ -222,7 +222,7 @@ "limit": 0 }, "RET504": { - "limit": 178 + "limit": 177 }, "RUF010": { "limit": 0 @@ -306,7 +306,7 @@ "limit": 0 }, "TID251": { - "limit": 1242 + "limit": 1240 }, "TRY002": { "limit": 528 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e26ce54ede7..4db4c61ef4b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23343 + "limit": 23267 }, "LIT002": { "limit": 27213 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16802 + "limit": 16793 }, "LIT011": { "limit": 5602 From e34b47f682d2d0183685e5d44058a3cf62d09b6e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:11:21 -0700 Subject: [PATCH 027/126] docs: cap all GitHub comments at 15-25 words, curb semicolon splices The 15-25 word cap previously applied only to replies/rebuttals to AI PR review bots; it now covers every GitHub comment (issue comments and PR discussion comments included). The public-writing punctuation bullet also gains a warning that a word cap is not a one-sentence cap, so tight budgets should be met with period splits or conjunctions rather than ";" splices, at most one ";" per message --- CLAUDE.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 1bc4d108da7..926d3e8698a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -29,7 +29,7 @@ Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We pref If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y: - don't use emojis -- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y +- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap is not a one-sentence cap: when writing under tight word budgets (tldrs, 15-25 word review replies), prefer a period split or a conjunction over ";", and keep to at most one ";" per message - don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. - don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose - don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "." @@ -59,7 +59,7 @@ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages When working on a PR, keep the PR description in sync with new commits being made -Replies/rebuttals to AI PR review bots must be 15-25 word human-readable replies +All GitHub comments (issue comments, PR discussion comments, and replies/rebuttals to AI PR review bots) must be 15-25 word human-readable messages, 25 words max Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in From 3e5dee7317631d0a24b0e799485ab5b2103166a3 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:21:18 -0700 Subject: [PATCH 028/126] chore: make it clearer to Claude that GitHub comments must be concise and to use fewer semicolons --- CLAUDE.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 926d3e8698a..4d66297e0a0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -29,7 +29,7 @@ Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We pref If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y: - don't use emojis -- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap is not a one-sentence cap: when writing under tight word budgets (tldrs, 15-25 word review replies), prefer a period split or a conjunction over ";", and keep to at most one ";" per message +- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message - don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc. - don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose - don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "." @@ -41,7 +41,7 @@ Python max line length is 120, not 88 When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing -`make pre-commit` always saves its complete output to a per-worktree log file and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice, and re-run only after the working tree actually changed +`make pre-commit` saves its complete output to a log file in .git (overwriting previous pre-commit logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in @@ -59,7 +59,7 @@ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages When working on a PR, keep the PR description in sync with new commits being made -All GitHub comments (issue comments, PR discussion comments, and replies/rebuttals to AI PR review bots) must be 15-25 word human-readable messages, 25 words max +All GitHub comments must be human-readable and 15-25 word max Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in From 54e9964eb829b242db45d99d3f6214262b7ae56f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 5 Aug 2026 23:22:11 -0700 Subject: [PATCH 029/126] fix(proxy): stop resolving the UI session sentinel team on /search_tools/list Every Admin UI session key is stamped with the reserved team id `litellm-dashboard`, which never has a row in LiteLLM_TeamTable, so the team lookup in _filter_visible_search_tools raised 404 and the endpoint returned 500 for every non-admin dashboard session. Skip the lookup for that sentinel and scope the caller by its key-level allowlist alone, matching how MCP and agent permission checks already treat it. A real team id is still resolved, and a genuine lookup failure now surfaces with its own status instead of being masked as a 500. --- .../search_tool_management.py | 66 ++++-- .../test_search_tool_management.py | 194 ++++++++++++++++++ 2 files changed, 241 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 381a5e14ca2..69edf681e4d 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -2,13 +2,15 @@ CRUD ENDPOINTS FOR SEARCH TOOLS """ +from collections.abc import Awaitable, Callable from datetime import datetime -from typing import Any, Final +from typing import Any, Final, TypeAlias from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel from litellm._logging import verbose_proxy_logger +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import ( LiteLLM_TeamTable, LitellmUserRoles, @@ -46,9 +48,46 @@ def _convert_datetime_to_str(value: datetime | str | None) -> str | None: return value +TeamObjectLookup: TypeAlias = Callable[[str, UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable]] + + +async def _team_object_from_db(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + + +def _allowlist_team_id(user_api_key_dict: UserAPIKeyAuth) -> str | None: + """ + The team whose object_permission allowlist scopes this caller, or None when there is none. + + Every Admin UI session key is stamped with UI_SESSION_TOKEN_TEAM_ID, a reserved sentinel that + never has a row in LiteLLM_TeamTable (`/team/new` rejects it as a real team id), so looking it + up would raise 404 instead of resolving a team. It carries no allowlist of its own, so the + caller is scoped by its key-level allowlist alone. Any other team id is looked up for real and + a failed lookup still surfaces. + """ + team_id: Final = user_api_key_dict.team_id + if not team_id or team_id == UI_SESSION_TOKEN_TEAM_ID: + return None + return team_id + + async def _filter_visible_search_tools( search_tools: list[SearchToolInfoResponse], user_api_key_dict: UserAPIKeyAuth, + lookup_team_object: TeamObjectLookup = _team_object_from_db, ) -> list[SearchToolInfoResponse]: """ Drop search tools the caller is not authorized to invoke, applying the same @@ -60,25 +99,12 @@ async def _filter_visible_search_tools( ): return search_tools - from litellm.proxy.auth.auth_checks import ( - can_user_view_search_tool, - get_team_object, - ) - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy.auth.auth_checks import can_user_view_search_tool - team_object: LiteLLM_TeamTable | None = None - if user_api_key_dict.team_id: - team_object = await get_team_object( - team_id=user_api_key_dict.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_dict.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) + allowlist_team_id: Final = _allowlist_team_id(user_api_key_dict) + team_object: Final[LiteLLM_TeamTable | None] = ( + await lookup_team_object(allowlist_team_id, user_api_key_dict) if allowlist_team_id else None + ) visible: Final[list[SearchToolInfoResponse]] = [] for tool in search_tools: @@ -213,6 +239,8 @@ async def list_search_tools( visible_search_tools: Final = await _filter_visible_search_tools(search_tool_configs, user_api_key_dict) return ListSearchToolsResponse(search_tools=visible_search_tools) + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception("Error getting search tools: %s", e) raise HTTPException(status_code=500, detail=str(e)) diff --git a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index f2ccfcd0155..158dee7cb78 100644 --- a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -5,6 +5,7 @@ from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient sys.path.insert( @@ -815,3 +816,196 @@ async def test_list_search_tools_admin_with_restricted_key_still_sees_all(): assert response.status_code == 200 names = {t["search_tool_name"] for t in response.json()["search_tools"]} assert names == {"db-tool-1", "db-tool-2", "db-tool-3"} + + +def _search_tool_responses(*names): + from litellm.types.search import SearchToolInfoResponse + + return [ + SearchToolInfoResponse( + search_tool_id=f"id-{name}", + search_tool_name=name, + litellm_params={"search_provider": "perplexity"}, + search_tool_info=None, + created_at=None, + updated_at=None, + is_from_config=False, + ) + for name in names + ] + + +def _recording_team_lookup(team_object=None, raises=None): + """A `TeamObjectLookup` double that records the team ids it was asked to resolve.""" + asked_for = [] + + async def _lookup(team_id, user_api_key_dict): + asked_for.append(team_id) + if raises is not None: + raise raises + return team_object + + return _lookup, asked_for + + +@pytest.mark.asyncio +async def test_list_search_tools_dashboard_session_key_does_not_look_up_the_ui_team(): + """ + Regression: the Admin UI session key is stamped with the reserved team id + ``litellm-dashboard``, which has no row in LiteLLM_TeamTable. Resolving it as a real team + raised 404, which the endpoint reported as a 500, so the Search Tools page was broken for + every non-admin browsing the dashboard. + """ + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + dashboard_session_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id=UI_SESSION_TOKEN_TEAM_ID, + ) + ui_team_is_not_a_real_team = AsyncMock( + side_effect=HTTPException( + status_code=404, + detail={"error": f"Team doesn't exist in db. Team={UI_SESSION_TOKEN_TEAM_ID}."}, + ) + ) + + with ( + _mock_search_tool_backend(_scoping_db_tools()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + ui_team_is_not_a_real_team, + ), + _override_auth(dashboard_session_user), + ): + response = TestClient(app).get("/search_tools/list") + + assert response.status_code == 200 + names = {t["search_tool_name"] for t in response.json()["search_tools"]} + assert names == {"db-tool-1", "db-tool-2", "db-tool-3"} + ui_team_is_not_a_real_team.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_filter_visible_search_tools_dashboard_session_still_honors_key_allowlist(): + """ + Skipping the synthetic team must not widen visibility: a dashboard session whose key + carries a search_tools allowlist stays scoped to it. + """ + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy.search_endpoints.search_tool_management import ( + _filter_visible_search_tools, + ) + + restricted_dashboard_session = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id=UI_SESSION_TOKEN_TEAM_ID, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", + search_tools=["db-tool-3"], + ), + ) + lookup, asked_for = _recording_team_lookup() + + visible = await _filter_visible_search_tools( + _search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"), + restricted_dashboard_session, + lookup, + ) + + assert [t["search_tool_name"] for t in visible] == ["db-tool-3"] + assert asked_for == [] + + +@pytest.mark.asyncio +async def test_filter_visible_search_tools_still_applies_a_real_team_allowlist(): + """A caller with a real team is still resolved and scoped by that team's allowlist.""" + from litellm.proxy.search_endpoints.search_tool_management import ( + _filter_visible_search_tools, + ) + + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="team-1", + ) + lookup, asked_for = _recording_team_lookup( + team_object=LiteLLM_TeamTable( + team_id="team-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-team", + search_tools=["db-tool-2"], + ), + ) + ) + + visible = await _filter_visible_search_tools( + _search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"), + team_member, + lookup, + ) + + assert [t["search_tool_name"] for t in visible] == ["db-tool-2"] + assert asked_for == ["team-1"] + + +@pytest.mark.asyncio +async def test_filter_visible_search_tools_propagates_a_real_team_lookup_failure(): + """ + A caller whose real team cannot be resolved must not fall through to "no team", which + would drop that team's allowlist and show tools the caller may not call. + """ + from litellm.proxy.search_endpoints.search_tool_management import ( + _filter_visible_search_tools, + ) + + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="deleted-team", + ) + lookup, asked_for = _recording_team_lookup( + raises=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."}) + ) + + with pytest.raises(HTTPException) as exc_info: + await _filter_visible_search_tools( + _search_tool_responses("db-tool-1", "db-tool-2"), + team_member, + lookup, + ) + + assert exc_info.value.status_code == 404 + assert asked_for == ["deleted-team"] + + +@pytest.mark.asyncio +async def test_list_search_tools_reports_a_missing_real_team_as_404(): + """ + The endpoint surfaces a genuine team lookup failure with its own status instead of + masking it as a 500 or quietly returning an unscoped list. + """ + team_member = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="internal_user", + team_id="deleted-team", + ) + + with ( + _mock_search_tool_backend(_scoping_db_tools()), + patch( + "litellm.proxy.auth.auth_checks.get_team_object", + AsyncMock( + side_effect=HTTPException( + status_code=404, + detail={"error": "Team doesn't exist in db. Team=deleted-team."}, + ) + ), + ), + _override_auth(team_member), + ): + response = TestClient(app).get("/search_tools/list") + + assert response.status_code == 404 + assert "search_tools" not in response.json() From eea292abbabed6a9b4bb86630c60efa74fd7e9a5 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 5 Aug 2026 23:24:37 -0700 Subject: [PATCH 030/126] fix(proxy): allow non-admins to reach /user/daily/activity/aggregated The aggregated route was missing from LiteLLMRoutes.self_managed_routes while its paginated sibling /user/daily/activity was listed, so auth rejected every internal user with a 401 before the handler ran. That route backs the default "Your Usage" view in the dashboard, which left the main Usage page broken for non-admin users. The handler already self-scopes: it checks admin view first, then falls back to require_caller_user_id_for_non_admin, defaults a missing user_id to the caller's own, and returns 403 when a non-admin asks for someone else's data. Listing the route restores reachability without widening what a caller can read. check_route_access matches exactly (plus explicit wildcards), so the parent entry never covered the /aggregated sub-path. --- litellm/proxy/_types.py | 1 + .../proxy/auth/test_route_checks.py | 52 ++++++++++++++ .../test_internal_user_endpoints.py | 69 +++++++++++++++++++ 3 files changed, 122 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7bc8ed59a6a..5b1134650f2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -781,6 +781,7 @@ class LiteLLMRoutes(enum.Enum): "/model/update", "/model/delete", "/user/daily/activity", + "/user/daily/activity/aggregated", "/user/available_roles", # read-only role metadata; any authenticated user may read "/user/list", # org admins checked in endpoint; non-admins get 403 "/model/{model_id}/update", diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 9285b997efc..0bfb10320f7 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3246,3 +3246,55 @@ def test_internal_user_still_blocked_from_another_users_info(): assert exc_info.value.status_code == 403 assert "key not allowed to access this user's info" in str(exc_info.value.detail) + + +@pytest.mark.parametrize( + "route", + [ + "/user/daily/activity", + "/user/daily/activity/aggregated", + ], +) +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_user_daily_activity_routes_reachable_by_non_admin(route, user_role): + """Both /user/daily/activity and its /aggregated sibling power the default + "Your Usage" dashboard view, and both handlers self-scope to the caller + (_user_has_admin_view -> require_caller_user_id_for_non_admin -> 403 on a + user_id mismatch). self_managed_routes is the ONLY list that grants either + route to a non-admin, so dropping one from it 401s every internal user's + main Usage page before the handler ever runs. + """ + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=user_role, + ) + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) + request = MagicMock(spec=Request) + request.query_params = {} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + + +def test_user_daily_activity_aggregated_not_covered_by_prefix_match(): + """check_route_access is exact-match plus explicit wildcards, so listing the + parent /user/daily/activity does not implicitly cover the /aggregated + sub-path. Pins the reason the sibling needs its own entry. + """ + assert not RouteChecks.check_route_access( + route="/user/daily/activity/aggregated", + allowed_routes=["/user/daily/activity"], + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index aab9a0b4fd0..056c2d3657a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2294,6 +2294,75 @@ async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch) ) +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_users( + monkeypatch, +): + """ + Same scoping contract as + test_get_user_daily_activity_non_admin_cannot_view_other_users, on the + aggregated route. Non-admins reach this handler now that the route is in + self_managed_routes, so the 403-on-mismatch and default-to-self behaviour + has to hold here too: opening the route must not widen access. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity_aggregated, + ) + + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + non_admin_key_dict = UserAPIKeyAuth( + user_id="regular-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + # Case 1: Non-admin targets another user's data — 403, helper never reached + with patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_get_daily_agg: + with pytest.raises(HTTPException) as exc_info: + await get_user_daily_activity_aggregated( + start_date="2025-01-01", + end_date="2025-01-31", + model=None, + api_key=None, + user_id="other-user-456", + timezone=None, + user_api_key_dict=non_admin_key_dict, + ) + + assert exc_info.value.status_code == 403 + assert "Non-admin users can only view their own spend data" in str(exc_info.value.detail) + mock_get_daily_agg.assert_not_called() + + # Case 2: Non-admin omits user_id — scoped to their own user_id, not global + mock_response = MagicMock() + with patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_get_daily_agg: + result = await get_user_daily_activity_aggregated( + start_date="2025-01-01", + end_date="2025-01-31", + model=None, + api_key=None, + user_id=None, + timezone=None, + user_api_key_dict=non_admin_key_dict, + ) + + assert result is mock_response + mock_get_daily_agg.assert_called_once() + assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "regular-user-123" + + @pytest.mark.asyncio async def test_delete_user_cleans_up_created_by_invitation_links(mocker): """ From d5a471b2a7be5b5eb923d156ff8327309313895d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 5 Aug 2026 23:40:16 -0700 Subject: [PATCH 031/126] test(proxy): type the search tool test helpers and record lookups with AsyncMock Replaces the hand-rolled recording double with AsyncMock so the awaited team ids come from await_args_list instead of a mutated list, and annotates the response factory now that SearchToolInfoResponse is imported at module level. --- .../test_search_tool_management.py | 34 ++++++------------- 1 file changed, 11 insertions(+), 23 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index 158dee7cb78..a64397d9818 100644 --- a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -24,6 +24,7 @@ import litellm.proxy.proxy_server as ps # Now we can safely import app from litellm.proxy.proxy_server import app +from litellm.types.search import SearchToolInfoResponse client = TestClient(app) @@ -818,9 +819,7 @@ async def test_list_search_tools_admin_with_restricted_key_still_sees_all(): assert names == {"db-tool-1", "db-tool-2", "db-tool-3"} -def _search_tool_responses(*names): - from litellm.types.search import SearchToolInfoResponse - +def _search_tool_responses(*names: str) -> list[SearchToolInfoResponse]: return [ SearchToolInfoResponse( search_tool_id=f"id-{name}", @@ -835,17 +834,8 @@ def _search_tool_responses(*names): ] -def _recording_team_lookup(team_object=None, raises=None): - """A `TeamObjectLookup` double that records the team ids it was asked to resolve.""" - asked_for = [] - - async def _lookup(team_id, user_api_key_dict): - asked_for.append(team_id) - if raises is not None: - raise raises - return team_object - - return _lookup, asked_for +def _team_ids_looked_up(lookup: AsyncMock) -> list[str]: + return [awaited.args[0] for awaited in lookup.await_args_list] @pytest.mark.asyncio @@ -906,7 +896,7 @@ async def test_filter_visible_search_tools_dashboard_session_still_honors_key_al search_tools=["db-tool-3"], ), ) - lookup, asked_for = _recording_team_lookup() + lookup = AsyncMock() visible = await _filter_visible_search_tools( _search_tool_responses("db-tool-1", "db-tool-2", "db-tool-3"), @@ -915,7 +905,7 @@ async def test_filter_visible_search_tools_dashboard_session_still_honors_key_al ) assert [t["search_tool_name"] for t in visible] == ["db-tool-3"] - assert asked_for == [] + lookup.assert_not_awaited() @pytest.mark.asyncio @@ -930,8 +920,8 @@ async def test_filter_visible_search_tools_still_applies_a_real_team_allowlist() user_id="internal_user", team_id="team-1", ) - lookup, asked_for = _recording_team_lookup( - team_object=LiteLLM_TeamTable( + lookup = AsyncMock( + return_value=LiteLLM_TeamTable( team_id="team-1", object_permission=LiteLLM_ObjectPermissionTable( object_permission_id="op-team", @@ -947,7 +937,7 @@ async def test_filter_visible_search_tools_still_applies_a_real_team_allowlist() ) assert [t["search_tool_name"] for t in visible] == ["db-tool-2"] - assert asked_for == ["team-1"] + assert _team_ids_looked_up(lookup) == ["team-1"] @pytest.mark.asyncio @@ -965,9 +955,7 @@ async def test_filter_visible_search_tools_propagates_a_real_team_lookup_failure user_id="internal_user", team_id="deleted-team", ) - lookup, asked_for = _recording_team_lookup( - raises=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."}) - ) + lookup = AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})) with pytest.raises(HTTPException) as exc_info: await _filter_visible_search_tools( @@ -977,7 +965,7 @@ async def test_filter_visible_search_tools_propagates_a_real_team_lookup_failure ) assert exc_info.value.status_code == 404 - assert asked_for == ["deleted-team"] + assert _team_ids_looked_up(lookup) == ["deleted-team"] @pytest.mark.asyncio From e8245107654a28dc4021de4712503230633e37df Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 01:54:26 -0700 Subject: [PATCH 032/126] fix(lint): generate the prisma client into the gate-owned venv prisma resolves its prisma-client-py generator through a plain /bin/sh PATH lookup, never through the interpreter that ran prisma generate, so the gate's generate step landed the client in whatever venv the caller had on PATH: the owned env never received one, every gate run regenerated, the caller's venv was mutated instead, and any invocation without a venv on PATH (the rewritten publisher workflow) failed outright The generate now runs with the target interpreter's bin directory pinned to the front of the child PATH. The prisma schema joins the environment fingerprint so clientless counts recorded before this commit can never be compared against clientful ones, a cold provision announces itself on stderr instead of sitting silent for two minutes, and the CI gate step reuses the job's prisma binary cache --- .github/workflows/test-linting.yml | 1 + scripts/prisma_generate_if_needed.py | 36 +++++++++++++++---- scripts/type_check_gate.py | 12 +++++-- .../test_prisma_generate_if_needed.py | 29 +++++++++++++++ tests/test_litellm/test_type_check_gate.py | 22 ++++++++++++ 5 files changed, 91 insertions(+), 9 deletions(-) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 5125dd0a354..5e333f2a3ca 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -115,6 +115,7 @@ jobs: - name: Check basedpyright budget (delta vs base) env: GH_TOKEN: ${{ github.token }} + PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA" diff --git a/scripts/prisma_generate_if_needed.py b/scripts/prisma_generate_if_needed.py index d2c40adf820..f2d2232b23f 100644 --- a/scripts/prisma_generate_if_needed.py +++ b/scripts/prisma_generate_if_needed.py @@ -9,13 +9,21 @@ client (a fresh or reinstalled prisma package) forces a regenerate even when the stamp matches. The prisma package itself is never imported here: once generated it re-exports the whole client on import, which costs more than the generate this script exists to skip. + +prisma resolves its generator command (``prisma-client-py``) through a plain +PATH lookup, never through the interpreter that invoked ``prisma generate``, +so the generate runs with this interpreter's own bin directory pinned to the +front of PATH; without that pin the client lands in whichever venv the caller +happened to have on PATH (or the generate fails outright when none is). """ import hashlib import importlib.metadata import importlib.util +import os import subprocess import sys +from collections.abc import Callable, Mapping from pathlib import Path REPO_ROOT = Path(__file__).resolve().parent.parent @@ -46,6 +54,25 @@ def client_is_generated() -> bool: ) +def env_with_own_bin_first(base_env: Mapping[str, str]) -> dict[str, str]: + bin_dir = str(Path(sys.executable).parent) + inherited = base_env.get("PATH") + path = os.pathsep.join((bin_dir, inherited)) if inherited else bin_dir + return {**base_env, "PATH": path} + + +def _run_command(cmd: list[str], cwd: Path, env: dict[str, str]) -> int: + return subprocess.run(cmd, cwd=cwd, env=env).returncode + + +def run_generate(run: Callable[[list[str], Path, dict[str, str]], int] = _run_command) -> int: + return run( + [sys.executable, "-m", "prisma", "generate", "--schema", str(SCHEMA)], + REPO_ROOT, + env_with_own_bin_first(os.environ), + ) + + def main() -> int: version = importlib.metadata.version("prisma") expected = stamp_value(SCHEMA.read_bytes(), version) @@ -55,12 +82,9 @@ def main() -> int: f"(prisma {version}); skipping prisma generate" ) return 0 - result = subprocess.run( - [sys.executable, "-m", "prisma", "generate", "--schema", str(SCHEMA)], - cwd=REPO_ROOT, - ) - if result.returncode != 0: - return result.returncode + returncode = run_generate() + if returncode != 0: + return returncode STAMP.write_text(expected) return 0 diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py index f9d7f2912fd..2fab7f8363a 100644 --- a/scripts/type_check_gate.py +++ b/scripts/type_check_gate.py @@ -34,8 +34,8 @@ skipped outright. When it is needed, it is a second basedpyright pass over a detached worktree at the merge-base, run under the same environment so import resolution matches, and its per-rule counts are cached under the repo's git common dir keyed by merge-base commit, ``pyrightconfig.json``, ``uv.lock``, -and the dependency-group set, so re-runs against the same branch -point pay for it once. A CI workflow publishes every staging commit's counts as +the Prisma schema, and the dependency-group set, so re-runs against the same +branch point pay for it once. A CI workflow publishes every staging commit's counts as an artifact (``--emit-counts-dir`` is its entry point), and on a disk-cache miss the gate first tries to download the merge-base's artifact through the ``gh`` CLI; any fetch failure falls back silently to the local base pass, so the gate @@ -83,6 +83,7 @@ GH_TIMEOUT_SECONDS = 10 TYPECHECK_ENV_DIR = REPO_ROOT / ".venv-typecheck" TYPECHECK_DEP_GROUPS = ("proxy-dev", "e2e-dev") PRISMA_GENERATE_SCRIPT = REPO_ROOT / "scripts" / "prisma_generate_if_needed.py" +PRISMA_SCHEMA = REPO_ROOT / "litellm" / "proxy" / "schema.prisma" # basedpyright's node process needs more than the ~4 GB default heap on this # repo; appended last so it wins node's last-flag-wins resolution over any @@ -192,6 +193,11 @@ def ensure_typecheck_env( measurement pass. Unconditional on purpose: an up-to-date env makes both steps near-instant no-ops, and skipping them on a heuristic is how the measured environment and the fingerprinted one drift apart.""" + if not env_dir.exists(): + sys.stderr.write( + f"provisioning {env_dir.name} (first run installs packages and " + "generates the Prisma client; re-runs are near-instant no-ops)\n" + ) env: Final = {**os.environ, "UV_PROJECT_ENVIRONMENT": str(env_dir)} for cmd in typecheck_env_commands(env_dir): if run(cmd, env) != 0: @@ -298,7 +304,7 @@ def environment_fingerprints( return ( *( hashlib.sha256(path.read_bytes()).hexdigest() - for path in (PYRIGHT_CONFIG, UV_LOCK) + for path in (PYRIGHT_CONFIG, UV_LOCK, PRISMA_SCHEMA) if path.exists() ), "groups:" + ",".join(dep_groups), diff --git a/tests/test_litellm/test_prisma_generate_if_needed.py b/tests/test_litellm/test_prisma_generate_if_needed.py index 39b9fcc4202..c36f580e24f 100644 --- a/tests/test_litellm/test_prisma_generate_if_needed.py +++ b/tests/test_litellm/test_prisma_generate_if_needed.py @@ -1,4 +1,6 @@ import importlib.util +import os +import sys from pathlib import Path _MODULE_PATH = ( @@ -33,3 +35,30 @@ def test_skip_requires_a_generated_client_even_with_a_matching_stamp(tmp_path): expected = mod.stamp_value(b"schema", "0.11.0") stamp.write_text(expected) assert mod.should_skip(stamp, expected, client_generated=False) is False + + +def test_env_puts_this_interpreters_bin_dir_first_on_path(): + env = mod.env_with_own_bin_first({"PATH": "/usr/bin", "HOME": "/home"}) + bin_dir = str(Path(sys.executable).parent) + assert env["PATH"].split(os.pathsep) == [bin_dir, "/usr/bin"] + assert env["HOME"] == "/home" + + +def test_env_without_an_inherited_path_is_just_the_bin_dir(): + env = mod.env_with_own_bin_first({}) + assert env["PATH"] == str(Path(sys.executable).parent) + + +def test_generate_runs_prisma_with_its_own_bin_dir_leading_the_childs_path(): + seen = {} + + def recorder(cmd, cwd, env): + seen["cmd"] = cmd + seen["cwd"] = cwd + seen["env"] = env + return 0 + + assert mod.run_generate(run=recorder) == 0 + assert seen["cmd"][:4] == [sys.executable, "-m", "prisma", "generate"] + assert seen["cwd"] == mod.REPO_ROOT + assert seen["env"]["PATH"].split(os.pathsep)[0] == str(Path(sys.executable).parent) diff --git a/tests/test_litellm/test_type_check_gate.py b/tests/test_litellm/test_type_check_gate.py index 2d9731db53f..6cf21707629 100644 --- a/tests/test_litellm/test_type_check_gate.py +++ b/tests/test_litellm/test_type_check_gate.py @@ -1,3 +1,4 @@ +import hashlib import importlib.util import json import subprocess @@ -278,6 +279,11 @@ def test_fingerprints_carry_the_dependency_group_set(): ) +def test_fingerprints_cover_the_prisma_schema(): + schema_hash = hashlib.sha256(gate.PRISMA_SCHEMA.read_bytes()).hexdigest() + assert schema_hash in gate.environment_fingerprints() + + def test_env_commands_sync_the_canonical_groups_then_generate_prisma(): sync, generate = gate.typecheck_env_commands(Path("/envdir")) assert sync[:3] == ("uv", "sync", "--frozen") @@ -327,6 +333,22 @@ def test_ensure_env_fails_loudly_and_stops_at_the_first_failed_step(tmp_path): assert len(calls) == 1 +def test_ensure_env_announces_a_cold_provision(tmp_path, capsys): + def runner(cmd, env): + return 0 + + gate.ensure_typecheck_env(env_dir=tmp_path / "fresh", run=runner) + assert "provisioning" in capsys.readouterr().err + + +def test_ensure_env_is_silent_when_the_env_already_exists(tmp_path, capsys): + def runner(cmd, env): + return 0 + + gate.ensure_typecheck_env(env_dir=tmp_path, run=runner) + assert capsys.readouterr().err == "" + + def test_cached_counts_round_trip(tmp_path): path = gate.cache_path(tmp_path, "abc123", ("f1", "f2")) gate.store_counts(tmp_path, path, "abc123", {"reportAny": 3, "reportCall": 1}) From 59041240f036fe80776b297b36757c48d85f7978 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:23:32 -0700 Subject: [PATCH 033/126] fix(managed_files): cap batch list page size at 100 and bulk-resolve raw file ids in one query --- .../proxy/hooks/managed_files.py | 104 +++++++++++++----- .../openai_files_endpoints/common_utils.py | 30 +++++ .../proxy/hooks/test_managed_files.py | 94 +++++++++++----- 3 files changed, 172 insertions(+), 56 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 07a1f959940..6fa6ef46ad4 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,6 +3,7 @@ import base64 import json +from collections.abc import Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast from uuid import NAMESPACE_URL, uuid5 @@ -31,10 +32,12 @@ from litellm.proxy._types import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + apply_unified_file_ids, ensure_batch_response_managed_file_ids, get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, + map_raw_file_ids_to_unified, normalize_mime_type_for_provider, resolve_managed_output_file_model_name, ) @@ -62,6 +65,9 @@ if TYPE_CHECKING: if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from prisma.models import ( + LiteLLM_ManagedObjectTable as PrismaManagedObjectRow, + ) from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient @@ -75,6 +81,20 @@ else: PrismaClient = Any +def _decode_json_blob(blob: object) -> object: + return json.loads(blob) if isinstance(blob, str) else blob + + +def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMBatch]: + try: + batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object)) + except Exception as e: + verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {e}") + return None + batch_obj.id = row.unified_object_id + return batch_obj + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes def __init__( @@ -329,7 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"Invalid 'after' cursor: no batch found with id '{after}'.", ) - page_size = limit or 20 + page_size: Final = min(limit or 20, 100) cursor_args: Dict[str, Any] = ( {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} ) @@ -343,36 +363,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): has_more = len(batches) > page_size - batch_objects: List[LiteLLMBatch] = [] - for batch in batches[:page_size]: - try: - batch_data = ( - json.loads(batch.file_object) - if isinstance(batch.file_object, str) - else batch.file_object - ) - batch_obj = LiteLLMBatch.model_validate(batch_data) - batch_obj.id = batch.unified_object_id - await ensure_batch_response_managed_file_ids( - response=batch_obj, - managed_files_obj=self, - prisma_client=self.prisma_client, - verbose_proxy_logger=verbose_logger, - user_api_key_dict=user_api_key_dict, - db_batch_object=batch, - unified_batch_id=_is_base64_encoded_unified_file_id( - batch.unified_object_id - ), - ) - batch_objects.append(batch_obj) + parsed_rows: Final = tuple( + (row, batch_obj) + for row in batches[:page_size] + if (batch_obj := _parse_managed_batch_row(row)) is not None + ) + unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified( + raw_file_ids=frozenset( + file_id + for _, batch_obj in parsed_rows + for file_id in (batch_obj.input_file_id, batch_obj.output_file_id, batch_obj.error_file_id) + if file_id and not _is_base64_encoded_unified_file_id(file_id) + ), + prisma_client=self.prisma_client, + ) + resolved_batches: Final = [ + await self._resolve_listed_batch( + row=row, + batch_obj=batch_obj, + unified_id_by_raw_id=unified_id_by_raw_id, + user_api_key_dict=user_api_key_dict, + ) + for row, batch_obj in parsed_rows + ] + return build_list_page( + [batch_obj for batch_obj in resolved_batches if batch_obj is not None], + has_more=has_more, + ) - except Exception as e: - verbose_logger.warning( - f"Failed to parse batch object {batch.unified_object_id}: {e}" - ) - continue - - return build_list_page(batch_objects, has_more=has_more) + async def _resolve_listed_batch( + self, + row: "PrismaManagedObjectRow", + batch_obj: LiteLLMBatch, + unified_id_by_raw_id: Mapping[str, str], + user_api_key_dict: UserAPIKeyAuth, + ) -> Optional[LiteLLMBatch]: + apply_unified_file_ids(batch_obj, unified_id_by_raw_id) + try: + await ensure_batch_response_managed_file_ids( + response=batch_obj, + managed_files_obj=self, + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=row, + unified_batch_id=_is_base64_encoded_unified_file_id( + row.unified_object_id + ), + ) + except Exception as e: + verbose_logger.warning( + f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}" + ) + return None + return batch_obj async def get_user_created_file_ids( self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str] diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 080b8b80ae4..bf83a7cf25c 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1,6 +1,7 @@ import base64 import mimetypes import re +from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional @@ -16,6 +17,7 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedObjectTable from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.types.utils import LiteLLMBatch @@ -1002,6 +1004,34 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None: pass +async def map_raw_file_ids_to_unified( + raw_file_ids: frozenset[str], prisma_client: "PrismaClient | None" +) -> Mapping[str, str]: + if not raw_file_ids or not prisma_client: + return MappingProxyType({}) + managed_files: Final = await ManagedFileRepository(prisma_client).table.find_many( + where={"flat_model_file_ids": {"hasSome": sorted(raw_file_ids)}} # mutable-ok: prisma where is a plain dict + ) + return MappingProxyType( + { + raw_id: managed_file.unified_file_id + for managed_file in managed_files + for raw_id in managed_file.flat_model_file_ids + if raw_id in raw_file_ids + } + ) + + +def apply_unified_file_ids(response: "LiteLLMBatch", unified_id_by_raw_id: Mapping[str, str]) -> None: + for file_attr, raw_id in ( + ("input_file_id", getattr(response, "input_file_id", None)), + ("output_file_id", getattr(response, "output_file_id", None)), + ("error_file_id", getattr(response, "error_file_id", None)), + ): + if isinstance(raw_id, str) and raw_id in unified_id_by_raw_id: + setattr(response, file_attr, unified_id_by_raw_id[raw_id]) + + async def ensure_batch_response_managed_file_ids( response, managed_files_obj, diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index fc10a1257e1..e1e5cc6c532 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1869,15 +1869,12 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id + input_file_row.flat_model_file_ids = [raw_input_file_id] - def find_managed_file(where): - if where["flat_model_file_ids"]["has"] == raw_input_file_id: - return input_file_row - return None - - prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[input_file_row] ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client @@ -1892,6 +1889,13 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): assert listed.id == unified_batch_uid assert listed.input_file_id == unified_input_file_id + bulk_lookup = prisma_client.db.litellm_managedfiletable.find_many.await_args + assert set(bulk_lookup.kwargs["where"]["flat_model_file_ids"]["hasSome"]) == { + raw_input_file_id, + raw_output_file_id, + raw_error_file_id, + } + decoded_output = _decode_unified_id(listed.output_file_id) assert decoded_output.startswith("litellm_proxy") assert f"llm_output_file_id,{raw_output_file_id}" in decoded_output @@ -1914,33 +1918,43 @@ async def test_list_batches_registers_and_returns_unified_output_file_ids(): @pytest.mark.asyncio async def test_list_batches_resolves_existing_managed_rows_without_minting(): """When the raw provider file IDs already have managed file rows, listing must - swap in the existing unified IDs and must not upsert duplicate rows.""" + swap in the existing unified IDs via one bulk lookup for the whole page, with + no per-row queries and no duplicate upserts.""" from litellm.proxy._types import UserAPIKeyAuth - unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") - raw_output_file_id = "file-list-out-existing" - existing_unified_output_id = base64.urlsafe_b64encode( - f"litellm_proxy:application/json;unified_id,u-9;llm_output_file_id,{raw_output_file_id}".encode() + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-9;target_model_names,gpt-5-batch" ).decode() + raw_output_file_ids = ["file-list-out-existing-1", "file-list-out-existing-2"] + existing_unified_output_ids = [ + base64.urlsafe_b64encode( + f"litellm_proxy:application/json;unified_id,u-{i};llm_output_file_id,{raw_id}".encode() + ).decode() + for i, raw_id in enumerate(raw_output_file_ids) + ] - record = _terminal_batch_record( - unified_batch_uid, "file-list-in-9", raw_output_file_id, "" - ) + records = [ + _terminal_batch_record( + _create_unified_batch_id("model-123", f"batch-{i}"), + unified_input_file_id, + raw_id, + "", + ) + for i, raw_id in enumerate(raw_output_file_ids) + ] prisma_client = AsyncMock() - prisma_client.db.litellm_managedobjecttable.find_many.return_value = [record] + prisma_client.db.litellm_managedobjecttable.find_many.return_value = records - existing_row = MagicMock() - existing_row.unified_file_id = existing_unified_output_id + existing_rows = [ + MagicMock(unified_file_id=unified_id, flat_model_file_ids=[raw_id]) + for raw_id, unified_id in zip(raw_output_file_ids, existing_unified_output_ids) + ] - def find_managed_file(where): - if where["flat_model_file_ids"]["has"] == raw_output_file_id: - return existing_row - return None - - prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=existing_rows ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock() proxy_managed_files = _PROXY_LiteLLMManagedFiles( DualCache(), prisma_client=prisma_client @@ -1951,10 +1965,38 @@ async def test_list_batches_resolves_existing_managed_rows_without_minting(): limit=10, ) - assert result["data"][0].output_file_id == existing_unified_output_id + assert [b.output_file_id for b in result["data"]] == existing_unified_output_ids + prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once() + prisma_client.db.litellm_managedfiletable.find_first.assert_not_awaited() prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited() +@pytest.mark.asyncio +async def test_list_batches_caps_page_size_at_100(): + """The list page size must be capped at 100 rows (matching OpenAI's limit) + even when the caller asks for more, so one request cannot fan out into an + unbounded scan.""" + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=100000, + ) + + assert ( + prisma_client.db.litellm_managedobjecttable.find_many.await_args.kwargs["take"] + == 101 + ) + assert result["data"] == [] + + @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception(): from litellm.proxy._types import UserAPIKeyAuth From 526f6d793ec96d9ba1ce93a854f6486499bff0e5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:23:52 -0700 Subject: [PATCH 034/126] fix(lint): retire the single-slot base-counts cache Storing a baseline used to prune every other cache entry, so gate runs in concurrent worktrees kept evicting each other's baselines and forcing full recomputes: this bit six times across two nights of benchmarking. The store now writes alongside existing entries and evicts only the oldest beyond eight, keyed as before by merge-base and environment fingerprint, so parallel worktrees' baselines simply coexist --- scripts/type_check_gate.py | 23 +++++++++++--- tests/test_litellm/test_type_check_gate.py | 35 ++++++++++++++++++++-- 2 files changed, 52 insertions(+), 6 deletions(-) diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py index 2fab7f8363a..c9f774c6113 100644 --- a/scripts/type_check_gate.py +++ b/scripts/type_check_gate.py @@ -63,7 +63,7 @@ import sys import tempfile import zipfile from collections import Counter -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence from pathlib import Path from typing import Final, NamedTuple @@ -73,6 +73,7 @@ PYRIGHT_CONFIG = REPO_ROOT / "pyrightconfig.json" UV_LOCK = REPO_ROOT / "uv.lock" DEFAULT_BASE = "origin/litellm_internal_staging" CACHE_FILE_PREFIX = "basedpyright-base-" +CACHE_KEEP_ENTRIES = 8 ARTIFACT_NAME_PREFIX = "basedpyright-counts-" GH_TIMEOUT_SECONDS = 10 @@ -367,16 +368,30 @@ def counts_payload(base_point: str, counts: Mapping[str, int]) -> str: ) +def entry_recency(path: Path) -> float: + try: + return path.stat().st_mtime + except OSError: + return 0.0 + + +def evicted_beyond_cap(entries: Sequence[Path], keep: int) -> tuple[Path, ...]: + newest_first: Final = sorted(entries, key=entry_recency, reverse=True) + return tuple(newest_first[keep:]) + + def store_counts( directory: Path, path: Path, base_point: str, counts: Mapping[str, int] ) -> None: directory.mkdir(parents=True, exist_ok=True) - for stale in directory.glob(f"{CACHE_FILE_PREFIX}*.json"): - if stale != path: - stale.unlink(missing_ok=True) scratch = scratch_path(path) scratch.write_text(counts_payload(base_point, counts)) scratch.replace(path) + siblings: Final = tuple( + entry for entry in directory.glob(f"{CACHE_FILE_PREFIX}*.json") if entry != path + ) + for stale in evicted_beyond_cap(siblings, CACHE_KEEP_ENTRIES - 1): + stale.unlink(missing_ok=True) def parse_origin_slug(url: str) -> str | None: diff --git a/tests/test_litellm/test_type_check_gate.py b/tests/test_litellm/test_type_check_gate.py index 6cf21707629..d104e4ca0c8 100644 --- a/tests/test_litellm/test_type_check_gate.py +++ b/tests/test_litellm/test_type_check_gate.py @@ -1,6 +1,7 @@ import hashlib import importlib.util import json +import os import subprocess from pathlib import Path @@ -387,15 +388,45 @@ def test_store_prune_spares_a_concurrent_runs_in_flight_scratch(tmp_path): assert gate.load_cached_counts(mine) == {"reportAny": 1} -def test_store_prunes_entries_for_other_branch_points(tmp_path): +def test_store_keeps_a_concurrent_worktrees_entry_for_another_branch_point(tmp_path): old = gate.cache_path(tmp_path, "old", ("f",)) gate.store_counts(tmp_path, old, "old", {"reportAny": 1}) new = gate.cache_path(tmp_path, "new", ("f",)) gate.store_counts(tmp_path, new, "new", {"reportAny": 2}) - assert not old.exists() + assert gate.load_cached_counts(old) == {"reportAny": 1} assert gate.load_cached_counts(new) == {"reportAny": 2} +def test_store_evicts_only_the_oldest_entries_beyond_the_cap(tmp_path): + aged = [ + gate.cache_path(tmp_path, f"base{i}", ("f",)) + for i in range(gate.CACHE_KEEP_ENTRIES) + ] + for age, path in enumerate(aged): + gate.store_counts(tmp_path, path, f"base{age}", {"reportAny": age}) + os.utime(path, (age, age)) + newest = gate.cache_path(tmp_path, "newest", ("f",)) + gate.store_counts(tmp_path, newest, "newest", {"reportAny": 99}) + assert not aged[0].exists() + assert all(path.exists() for path in aged[1:]) + assert gate.load_cached_counts(newest) == {"reportAny": 99} + + +def test_store_never_evicts_the_entry_it_just_wrote_even_on_mtime_ties(tmp_path): + others = [ + gate.cache_path(tmp_path, f"base{i}", ("f",)) + for i in range(gate.CACHE_KEEP_ENTRIES + 2) + ] + for path in others: + gate.store_counts(tmp_path, path, path.name, {"reportAny": 1}) + os.utime(path, (9_999_999_999, 9_999_999_999)) + mine = gate.cache_path(tmp_path, "mine", ("f",)) + gate.store_counts(tmp_path, mine, "mine", {"reportAny": 2}) + assert gate.load_cached_counts(mine) == {"reportAny": 2} + survivors = list(tmp_path.glob(f"{gate.CACHE_FILE_PREFIX}*.json")) + assert len(survivors) == gate.CACHE_KEEP_ENTRIES + + def _no_fetch(ref): return None From 470ebc208991c0883386bba4a7bd4b9c1d7d2c09 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:27:59 -0700 Subject: [PATCH 035/126] test(batches): cover caller-supplied db row skipping the cancel-path lookup --- ..._batch_update_db_managed_output_file_id.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index 74139fa9238..1b60c97b510 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -127,6 +127,42 @@ async def test_cancel_path_registers_output_file_under_batch_owner(): assert stored["output_file_id"] == unified_id +@pytest.mark.asyncio +async def test_update_batch_skips_lookup_when_db_batch_object_supplied(): + unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" + caller_row = SimpleNamespace( + created_by="caller-owner", team_id="caller-team", status="in_progress" + ) + decoy_row = SimpleNamespace( + created_by="decoy-owner", team_id="decoy-team", status="in_progress" + ) + response = _build_batch_response( + status="cancelling", + output_file_id="file-raw-output", + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + mock_managed_files = _build_managed_files_mock(unified_id=unified_id) + mock_prisma = _build_prisma_mock(db_batch_object=decoy_row) + + await update_batch_in_database( + batch_id="batch_managed_ids_test", + unified_batch_id="litellm_proxy;model_id:my-model;llm_batch_id:batch_managed_ids_test", + response=response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + db_batch_object=caller_row, + operation="retrieve", + ) + + mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_called() + forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[ + "user_api_key_dict" + ] + assert forwarded_auth.user_id == "caller-owner" + assert forwarded_auth.team_id == "caller-team" + + @pytest.mark.asyncio async def test_update_batch_derives_model_id_from_unified_batch_id(): unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" From 920e05c484337350af5724d46b7233d4ab8eaacc Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:45:44 -0700 Subject: [PATCH 036/126] chore: fix typo --- CLAUDE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index 4d66297e0a0..fa891bc7ae9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -59,7 +59,7 @@ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages When working on a PR, keep the PR description in sync with new commits being made -All GitHub comments must be human-readable and 15-25 word max +All GitHub comments must be human-readable and 15-25 words max Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in From 0c1a4b127dbfdf8aca292e2344f6c76b678d7d75 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:31:40 -0700 Subject: [PATCH 037/126] chore(lint): name MappingProxyType in the mutable-collection fix messages LIT001/LIT002 and the typing.Dict ban all steered dict-shaped values to frozen dataclasses or suppression even though the checker already accepts MappingProxyType as a freezing wrapper; the messages now name it so the dict-shaped freeze path is actually discoverable at fix time. --- CLAUDE.md | 4 ++-- ruff-strict.toml | 2 +- scripts/check_type_discipline.py | 16 ++++++++++------ tests/test_litellm/test_check_type_discipline.py | 8 ++++++++ 4 files changed, 21 insertions(+), 9 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 1bc4d108da7..81e560af849 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -45,7 +45,7 @@ When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-bud If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in -If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason +If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()`, and freeze dict-shaped values with `types.MappingProxyType({...})` (annotated as `Mapping[...]`), instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # `. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing @@ -72,7 +72,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega - Composition over inheritance - Never-nester: early returns over deep nesting - Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) -- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), etc. +- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `types.MappingProxyType`, etc. - Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: ` explaining why - Use dependency injection - Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed diff --git a/ruff-strict.toml b/ruff-strict.toml index d58885fe848..66d8f281fce 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -15,7 +15,7 @@ max-args = 5 "typing.Any".msg = "Use a concrete type. Frozen slots=True dataclass (preferred) / NamedTuple / ReadOnly TypedDict for payloads." "typing_extensions.Any".msg = "Same as typing.Any." "typing.List".msg = "tuple[X, ...] for state, Sequence[X] for params." -"typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; create a Mapping alias with concrete value types if truly dynamic." +"typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; if truly dynamic, a Mapping alias with concrete value types, frozen at runtime with types.MappingProxyType({...})." "typing.Set".msg = "frozenset[X] or AbstractSet[X]." "typing.MutableSequence".msg = "Sequence[X]." "typing.MutableMapping".msg = "See typing.Dict." diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index c303fbaffce..d73f1095895 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -11,14 +11,16 @@ LIT001 Mutable collection in a type annotation, anywhere it appears: function collection lets whoever holds it grow or rewrite it after the fact; annotate a read-only view instead (Mapping/Sequence/AbstractSet/tuple[X, ...]/ frozenset[X], or a frozen dataclass / NamedTuple / ReadOnly TypedDict) and - build it functionally (comprehension / map, not append-in-a-loop). + build it functionally (comprehension / map / MappingProxyType({...}), not + append-in-a-loop). Suppress with `# mutable-ok: ` on the offending line. LIT002 Mutable-collection *construction*: a list/dict/set literal or comprehension, or a call to a mutable constructor (list/dict/set/deque/defaultdict/Counter/...). Catches the unannotated seed-then-mutate pattern LIT001 cannot see (`acc = []`). Build the value in one shot and freeze it: a `tuple`/`frozenset` wrapping a - generator (`tuple(f(x) for x in xs)`), a tuple literal, or a frozen dataclass / - NamedTuple / ReadOnly TypedDict. Generator expressions and `tuple`/`frozenset` + generator (`tuple(f(x) for x in xs)`), a tuple literal, `MappingProxyType({...})` + for a dict-shaped value, or a frozen dataclass / NamedTuple / ReadOnly TypedDict. + Generator expressions and `tuple`/`frozenset` calls are not construction and pass. Annotation-internal lists (`Callable[[int], str]`) are exempt, as is a value passed directly to a freezing wrapper (`tuple(...)`, `frozenset(...)`, `MappingProxyType(...)`): it is frozen before @@ -279,7 +281,8 @@ def _mutable_ann(path: Path, line: int, name: str, where: str) -> Violation: f"mutable `{name}` in {where}: a mutable collection can be grown or rewritten " f"by whoever holds it. Annotate a read-only view -- Mapping[...], Sequence[...], " f"AbstractSet[...], tuple[X, ...], frozenset[X], or a frozen dataclass / " - f"NamedTuple / ReadOnly TypedDict -- and build it functionally, not by " + f"NamedTuple / ReadOnly TypedDict -- and build it functionally " + f"(comprehension / map / `MappingProxyType({{...}})`), not by " f"append-in-a-loop (suppress: `# mutable-ok: `)", ) @@ -488,8 +491,9 @@ def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) path, node.lineno, "LIT002", f"mutable {kind}: this builds a collection that can be grown or rewritten. " f"Build it in one shot and freeze it -- a tuple/frozenset wrapping a generator " - f"(`tuple(f(x) for x in xs)`), a tuple literal, or a frozen dataclass / NamedTuple " - f"/ ReadOnly TypedDict (suppress: `# mutable-ok: `)", + f"(`tuple(f(x) for x in xs)`), a tuple literal, `MappingProxyType({{...}})` for a " + f"dict-shaped value, or a frozen dataclass / NamedTuple / ReadOnly TypedDict " + f"(suppress: `# mutable-ok: `)", ) diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index f2bfc637095..81efbf0a4f8 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -175,6 +175,14 @@ def test_unfrozen_literal_still_counts(tmp_path): assert "LIT002" in _codes(tmp_path, "from types import MappingProxyType\nd = {'a': 1}\nm = MappingProxyType(d)\n") +def test_fix_messages_name_mappingproxytype(tmp_path): + f = tmp_path / "snippet.py" + f.write_text("x: dict[str, int] = {}\n", encoding="utf-8") + messages = {v.code: v.message for v in checker.check_file(f)} + assert "MappingProxyType" in messages["LIT001"] + assert "MappingProxyType" in messages["LIT002"] + + def test_mutable_ok_with_reason_suppresses_both_rules(tmp_path): codes = _codes(tmp_path, "x: dict[str, int] = {} # mutable-ok: in-place buffer mutated hot path\n") assert "LIT001" not in codes From d7ca4dc77f04e1dbd903fd15de26e508504ba64f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:34:34 -0700 Subject: [PATCH 038/126] chore(lint): say annotate Mapping[...] instead of Mapping alias in typing.Dict ban --- ruff-strict.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ruff-strict.toml b/ruff-strict.toml index 66d8f281fce..5e5192aa864 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -15,7 +15,7 @@ max-args = 5 "typing.Any".msg = "Use a concrete type. Frozen slots=True dataclass (preferred) / NamedTuple / ReadOnly TypedDict for payloads." "typing_extensions.Any".msg = "Same as typing.Any." "typing.List".msg = "tuple[X, ...] for state, Sequence[X] for params." -"typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; if truly dynamic, a Mapping alias with concrete value types, frozen at runtime with types.MappingProxyType({...})." +"typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; if truly dynamic, annotate Mapping[...] with concrete value types and freeze at runtime with types.MappingProxyType({...})." "typing.Set".msg = "frozenset[X] or AbstractSet[X]." "typing.MutableSequence".msg = "Sequence[X]." "typing.MutableMapping".msg = "See typing.Dict." From d0712e1a5e8b116f603fad6160b3bee298863ecc Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:37:46 -0700 Subject: [PATCH 039/126] chore: make CLAUDE.md more concise --- CLAUDE.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 81e560af849..a66e36e75cf 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -45,7 +45,7 @@ When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-bud If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in -If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()`, and freeze dict-shaped values with `types.MappingProxyType({...})` (annotated as `Mapping[...]`), instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason +If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # `. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing @@ -72,7 +72,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega - Composition over inheritance - Never-nester: early returns over deep nesting - Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) -- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `types.MappingProxyType`, etc. +- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc. - Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: ` explaining why - Use dependency injection - Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed From 818e319b536fb986294ca5ae8ab75389569137e1 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:38:28 -0700 Subject: [PATCH 040/126] chore: make it more concise --- ruff-strict.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ruff-strict.toml b/ruff-strict.toml index 5e5192aa864..01faf04805f 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -15,7 +15,7 @@ max-args = 5 "typing.Any".msg = "Use a concrete type. Frozen slots=True dataclass (preferred) / NamedTuple / ReadOnly TypedDict for payloads." "typing_extensions.Any".msg = "Same as typing.Any." "typing.List".msg = "tuple[X, ...] for state, Sequence[X] for params." -"typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; if truly dynamic, annotate Mapping[...] with concrete value types and freeze at runtime with types.MappingProxyType({...})." +"typing.Dict".msg = "Frozen dataclass / NamedTuple / ReadOnly TypedDict; if truly dynamic, use MappingProxyType." "typing.Set".msg = "frozenset[X] or AbstractSet[X]." "typing.MutableSequence".msg = "Sequence[X]." "typing.MutableMapping".msg = "See typing.Dict." From 20a94a91001aa7f1ff01de3b8dc3774d4db3bbae Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:57:44 -0700 Subject: [PATCH 041/126] chore(lint): move MappingProxyType to the dynamic tail of the LIT002 freeze menu Revert the LIT001 build-clause inserts, phrase the LIT002 menu as 'or (if it really must be dynamic) a MappingProxyType wrapping a dict literal or comprehension', and fold the two freezing-wrapper exemption sentences into one that names MappingProxyType beside tuple/frozenset. --- scripts/check_type_discipline.py | 28 +++++++++---------- .../test_check_type_discipline.py | 9 +++--- 2 files changed, 17 insertions(+), 20 deletions(-) diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index d73f1095895..65d0424fb5a 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -11,21 +11,20 @@ LIT001 Mutable collection in a type annotation, anywhere it appears: function collection lets whoever holds it grow or rewrite it after the fact; annotate a read-only view instead (Mapping/Sequence/AbstractSet/tuple[X, ...]/ frozenset[X], or a frozen dataclass / NamedTuple / ReadOnly TypedDict) and - build it functionally (comprehension / map / MappingProxyType({...}), not - append-in-a-loop). + build it functionally (comprehension / map, not append-in-a-loop). Suppress with `# mutable-ok: ` on the offending line. LIT002 Mutable-collection *construction*: a list/dict/set literal or comprehension, or a call to a mutable constructor (list/dict/set/deque/defaultdict/Counter/...). Catches the unannotated seed-then-mutate pattern LIT001 cannot see (`acc = []`). Build the value in one shot and freeze it: a `tuple`/`frozenset` wrapping a - generator (`tuple(f(x) for x in xs)`), a tuple literal, `MappingProxyType({...})` - for a dict-shaped value, or a frozen dataclass / NamedTuple / ReadOnly TypedDict. - Generator expressions and `tuple`/`frozenset` - calls are not construction and pass. Annotation-internal lists (`Callable[[int], - str]`) are exempt, as is a value passed directly to a freezing wrapper - (`tuple(...)`, `frozenset(...)`, `MappingProxyType(...)`): it is frozen before - it can escape, though anything mutable nested inside it still counts. - Suppress with `# mutable-ok: `. + generator (`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / + NamedTuple / ReadOnly TypedDict, or (if it really must be dynamic) a + MappingProxyType wrapping a dict literal or comprehension. Generator expressions + and freezing-wrapper calls (`tuple(...)`, `frozenset(...)`, + `MappingProxyType(...)`) are not construction and pass, as does the value passed + directly to a wrapper: it is frozen before it can escape, though anything + mutable nested inside it still counts. Annotation-internal lists + (`Callable[[int], str]`) are exempt. Suppress with `# mutable-ok: `. LIT003 noqa suppression without rule codes or without a reason. Required shape: `# noqa: TID251 # ` LIT004 pyright/mypy ignore without bracketed codes or without a reason. @@ -281,8 +280,7 @@ def _mutable_ann(path: Path, line: int, name: str, where: str) -> Violation: f"mutable `{name}` in {where}: a mutable collection can be grown or rewritten " f"by whoever holds it. Annotate a read-only view -- Mapping[...], Sequence[...], " f"AbstractSet[...], tuple[X, ...], frozenset[X], or a frozen dataclass / " - f"NamedTuple / ReadOnly TypedDict -- and build it functionally " - f"(comprehension / map / `MappingProxyType({{...}})`), not by " + f"NamedTuple / ReadOnly TypedDict -- and build it functionally, not by " f"append-in-a-loop (suppress: `# mutable-ok: `)", ) @@ -491,9 +489,9 @@ def iter_construction_violations(path: Path, tree: ast.AST, comments: Comments) path, node.lineno, "LIT002", f"mutable {kind}: this builds a collection that can be grown or rewritten. " f"Build it in one shot and freeze it -- a tuple/frozenset wrapping a generator " - f"(`tuple(f(x) for x in xs)`), a tuple literal, `MappingProxyType({{...}})` for a " - f"dict-shaped value, or a frozen dataclass / NamedTuple / ReadOnly TypedDict " - f"(suppress: `# mutable-ok: `)", + f"(`tuple(f(x) for x in xs)`), a tuple literal, a frozen dataclass / NamedTuple " + f"/ ReadOnly TypedDict, or (if it really must be dynamic) a MappingProxyType " + f"wrapping a dict literal or comprehension (suppress: `# mutable-ok: `)", ) diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 81efbf0a4f8..4b8533df604 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -175,12 +175,11 @@ def test_unfrozen_literal_still_counts(tmp_path): assert "LIT002" in _codes(tmp_path, "from types import MappingProxyType\nd = {'a': 1}\nm = MappingProxyType(d)\n") -def test_fix_messages_name_mappingproxytype(tmp_path): +def test_lit002_fix_message_names_mappingproxytype(tmp_path): f = tmp_path / "snippet.py" - f.write_text("x: dict[str, int] = {}\n", encoding="utf-8") - messages = {v.code: v.message for v in checker.check_file(f)} - assert "MappingProxyType" in messages["LIT001"] - assert "MappingProxyType" in messages["LIT002"] + f.write_text("x = {'a': 1}\n", encoding="utf-8") + messages = [v.message for v in checker.check_file(f) if v.code == "LIT002"] + assert "MappingProxyType" in messages[0] def test_mutable_ok_with_reason_suppresses_both_rules(tmp_path): From 2c91166d32d5253ae7434c796c2365598c18b8b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 11:39:34 +0000 Subject: [PATCH 042/126] chore: ignore the mechanical lint and typing sweeps in git blame Seven wide-reaching but semantically neutral commits landed since the last entry, together rewriting roughly 162k lines across ~4,700 file touches. Blame on any line they reflowed points at the sweep instead of the commit that wrote the logic. They cover the safe ruff autofix pass, the collections.abc import move, the f-string !s cleanup, lazy log message construction, the LIT010 and LIT011 Final and frozen-parameter rollout, ruff coverage for litellm/types, and the inert type: ignore strip. Smaller ratchet commits are left out on purpose: each touches a few hundred lines at most, so listing them would grow the file faster than it buys back blame accuracy --- .git-blame-ignore-revs | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs index 2527239b904..a0943cff53d 100644 --- a/.git-blame-ignore-revs +++ b/.git-blame-ignore-revs @@ -17,3 +17,24 @@ # style: unify ruff format width on 120 (#31518) 48b5a5a0cc5a694a11219416ee0b6eb6e620e74e + +# refactor(imports): move collections.abc names out of typing (#35495) +397e8e4918777e4e60a7f5e88699e0a9a7dabb3d + +# refactor(lint): apply every safe ruff autofix and zero 28 strict-rule budgets (#35495) +b604e2b20c6db2099085a2f0e59b7e99e87eed6f + +# refactor(logging): drop redundant !s conversion flags from f-strings (#35546) +7b2d3440cba3160277470f7a0180098ae9b87864 + +# perf: build log messages lazily so filtered-out log records cost nothing (#35703) +c9887a1f94bc1e7e4bdfe64d640f0509a0bc19dd + +# feat(lint): enforce Final on locals and freeze function parameters (#35807) +2708620d6a599cc73c1950a942d26ac26a7ed3d4 + +# chore(lint): remove litellm/types from the ruff lint exclusion (#35926) +4e32a8bf6a1e1af1e04b67c759841ccef44b2235 + +# chore(lint): strip inert type: ignore comments and zero LIT009/LIT010/LIT011 headroom (#35928) +338e411103ad5d7003e97f34f04fa36bca542dbe From ca7453bc6922f707761ac212b0c5872f777cd3fd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 12:04:38 +0000 Subject: [PATCH 043/126] chore(lint): recompute budget ceilings after merging base The base branch ratcheted the same limits in 28a277e9, so the conflicting files were reset to base and the ratchet re-run against the new merge-base rather than resolved by hand. Each limit is now the base value minus this branch's own delta, so both ratchets survive: basedpyright -653 across 48 rules, strict ruff -80, LIT -85. --- basedpyright-code-budget.json | 16 ++++++++-------- ruff-strict-budget.json | 18 +++++++++--------- type-discipline-budget.json | 4 ++-- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d98e6c6c911..632743236c4 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 29204 + "limit": 28842 }, "reportArgumentType": { "limit": 2634 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9225 + "limit": 9103 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5850 + "limit": 5843 }, "reportMissingTypeArgument": { - "limit": 15833 + "limit": 15816 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45145 + "limit": 45110 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39881 + "limit": 39838 }, "reportUnknownParameterType": { - "limit": 20258 + "limit": 20237 }, "reportUnknownVariableType": { - "limit": 31429 + "limit": 31383 }, "reportUnnecessaryCast": { "limit": 122 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 65c98f6aab3..60356eda05b 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,30 +1,30 @@ { "ANN001": { - "limit": 3126 + "limit": 3121 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 836 + "limit": 834 }, "ANN201": { - "limit": 2037 + "limit": 2033 }, "ANN202": { - "limit": 869 + "limit": 865 }, "ANN204": { - "limit": 715 + "limit": 713 }, "ANN205": { - "limit": 115 + "limit": 114 }, "ANN206": { "limit": 133 }, "ANN401": { - "limit": 1689 + "limit": 1630 }, "ASYNC230": { "limit": 11 @@ -222,7 +222,7 @@ "limit": 0 }, "RET504": { - "limit": 178 + "limit": 177 }, "RUF010": { "limit": 0 @@ -306,7 +306,7 @@ "limit": 0 }, "TID251": { - "limit": 1242 + "limit": 1240 }, "TRY002": { "limit": 528 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 8064e63f1aa..ab8198304bb 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23332 + "limit": 23256 }, "LIT002": { "limit": 27213 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16792 + "limit": 16783 }, "LIT011": { "limit": 5602 From 2d2994c9e9f45132c4ebb401a7fbf4a63d35a45f Mon Sep 17 00:00:00 2001 From: Praveena Mundolimoole <103165192+Praveena-617@users.noreply.github.com> Date: Thu, 6 Aug 2026 22:02:44 +0530 Subject: [PATCH 044/126] fix(proxy): yaml store_prompts_in_spend_logs should take precedence over DB cached value (#35769) When store_model_in_db is true, general_settings are persisted to the LiteLLM_Config DB table. On subsequent startups and periodic reloads, _add_general_settings_from_db_config() unconditionally overwrites the in-memory general_settings with DB-cached values, including store_prompts_in_spend_logs. This means a YAML config change (e.g. store_prompts_in_spend_logs: false) deployed via CI/CD has no effect because the stale DB value (true) always wins. The admin must manually update via /config/update API after every deploy, defeating config-as-code. Fix: track which general_settings keys were explicitly set in YAML at startup (_yaml_general_settings_keys). During DB config merge, prefer the YAML value for tracked keys. The DB value is only used as fallback when YAML does not set the key, preserving the admin UI's ability to change settings at runtime. Steps to reproduce: 1. Start proxy with store_model_in_db: true, store_prompts_in_spend_logs: true 2. Change YAML to store_prompts_in_spend_logs: false, restart 3. Send a request, query LiteLLM_SpendLogs - prompts still stored 4. Check LiteLLM_Config table - DB still has true, overriding YAML Slack thread: https://dataset-jsonhackathon.slack.com/archives/C0ACUS7LM29/p1785835131860139 --- litellm/proxy/proxy_server.py | 19 ++- .../test_proxy_config_unit_test.py | 122 ++++++++++++------ 2 files changed, 99 insertions(+), 42 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 539b68c1aee..07eaed9fe45 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3909,6 +3909,10 @@ class ProxyConfig: # whether an existing request predates the prices it just fetched, and re-serving one # costs a single fetch where skipping one leaves it priced wrong indefinitely self.model_cost_map_applied_revision: int = 0 + # Keys explicitly set in the YAML config file. Used to give YAML + # precedence over stale DB-cached values for these specific keys + # during periodic config reloads (_update_general_settings). + self._yaml_general_settings_keys: set[str] = set() # mutable-ok: populated once at startup, read-only thereafter # fmt: skip def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -4839,6 +4843,11 @@ class ProxyConfig: _hc_staleness = None _hc_ignore_transient = False if general_settings: + # Record which keys were explicitly set in the YAML config file. + # These keys take precedence over DB-cached values during periodic + # reloads (see _update_general_settings). + self._yaml_general_settings_keys = set(general_settings.keys()) # mutable-ok: snapshot of YAML keys at load time # fmt: skip + ### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ### key_management_settings: Final = general_settings.get("key_management_settings", None) if key_management_settings is not None: @@ -6049,7 +6058,15 @@ class ProxyConfig: ## STORE PROMPTS IN SPEND LOGS ## if "store_prompts_in_spend_logs" in _general_settings: - value = _general_settings["store_prompts_in_spend_logs"] + # If the YAML config explicitly set this key, prefer the YAML value + # over the DB-cached value. This ensures config changes deployed via + # CI/CD take effect without requiring a manual /config/update call. + # When YAML does not set this key, the DB value is used (preserving + # admin UI runtime changes). + if "store_prompts_in_spend_logs" in self._yaml_general_settings_keys: + value = general_settings.get("store_prompts_in_spend_logs") + else: + value = _general_settings["store_prompts_in_spend_logs"] # Normalize case: handle True/true/TRUE, False/false/FALSE, None/null if value is None: general_settings["store_prompts_in_spend_logs"] = None diff --git a/tests/proxy_unit_tests/test_proxy_config_unit_test.py b/tests/proxy_unit_tests/test_proxy_config_unit_test.py index 36cd08fa5f5..e6b38f31b48 100644 --- a/tests/proxy_unit_tests/test_proxy_config_unit_test.py +++ b/tests/proxy_unit_tests/test_proxy_config_unit_test.py @@ -15,9 +15,7 @@ import os # this file is to test litellm/proxy -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import asyncio import logging @@ -88,25 +86,14 @@ async def test_read_config_file_with_os_environ_vars(): # Read config proxy_config_instance = ProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) - config_path = os.path.join( - current_path, "example_config_yaml", "config_with_env_vars.yaml" - ) + config_path = os.path.join(current_path, "example_config_yaml", "config_with_env_vars.yaml") config = await proxy_config_instance.get_config(config_file_path=config_path) print(config) # Add assertions - assert ( - config["litellm_settings"]["default_internal_user_params"]["user_role"] - == "admin" - ) - assert ( - config["litellm_settings"]["s3_callback_params"]["s3_aws_access_key_id"] - == "1234567890" - ) - assert ( - config["litellm_settings"]["s3_callback_params"]["s3_aws_secret_access_key"] - == "1234567890" - ) + assert config["litellm_settings"]["default_internal_user_params"]["user_role"] == "admin" + assert config["litellm_settings"]["s3_callback_params"]["s3_aws_access_key_id"] == "1234567890" + assert config["litellm_settings"]["s3_callback_params"]["s3_aws_secret_access_key"] == "1234567890" for model in config["model_list"]: if "azure" in model["litellm_params"]["model"]: @@ -129,17 +116,13 @@ async def test_basic_include_directive(): """ proxy_config_instance = ProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) - config_path = os.path.join( - current_path, "example_config_yaml", "config_with_include.yaml" - ) + config_path = os.path.join(current_path, "example_config_yaml", "config_with_include.yaml") config = await proxy_config_instance.get_config(config_file_path=config_path) # Verify the included model list was merged assert len(config["model_list"]) > 0 - assert any( - model["model_name"] == "included-model" for model in config["model_list"] - ) + assert any(model["model_name"] == "included-model" for model in config["model_list"]) # Verify original config settings remain assert config["litellm_settings"]["callbacks"] == ["prometheus"] @@ -152,9 +135,7 @@ async def test_missing_include_file(): """ proxy_config_instance = ProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) - config_path = os.path.join( - current_path, "example_config_yaml", "config_with_missing_include.yaml" - ) + config_path = os.path.join(current_path, "example_config_yaml", "config_with_missing_include.yaml") with pytest.raises(FileNotFoundError): await proxy_config_instance.get_config(config_file_path=config_path) @@ -167,20 +148,14 @@ async def test_multiple_includes(): """ proxy_config_instance = ProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) - config_path = os.path.join( - current_path, "example_config_yaml", "config_with_multiple_includes.yaml" - ) + config_path = os.path.join(current_path, "example_config_yaml", "config_with_multiple_includes.yaml") config = await proxy_config_instance.get_config(config_file_path=config_path) # Verify models from both included files are present assert len(config["model_list"]) == 2 - assert any( - model["model_name"] == "included-model-1" for model in config["model_list"] - ) - assert any( - model["model_name"] == "included-model-2" for model in config["model_list"] - ) + assert any(model["model_name"] == "included-model-1" for model in config["model_list"]) + assert any(model["model_name"] == "included-model-2" for model in config["model_list"]) # Verify original config settings remain assert config["litellm_settings"]["callbacks"] == ["prometheus"] @@ -211,8 +186,7 @@ def test_add_callbacks_from_db_config(): # 1 instance of LangfusePromptManagement should exist in litellm.success_callback num_langfuse_instances = sum( - isinstance(callback, LangfusePromptManagement) - for callback in litellm.success_callback + isinstance(callback, LangfusePromptManagement) for callback in litellm.success_callback ) assert num_langfuse_instances == 1 assert len(litellm.success_callback) == 2 @@ -290,9 +264,7 @@ async def test_json_logs_calls_turn_on_json(): "litellm_settings": {"json_logs": True}, } - with tempfile.NamedTemporaryFile( - mode="w", suffix=".yaml", delete=False - ) as temp_file: + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as temp_file: yaml.dump(config_content, temp_file) temp_file_path = temp_file.name @@ -316,3 +288,71 @@ async def test_json_logs_calls_turn_on_json(): # Cleanup os.unlink(temp_file_path) litellm.json_logs = False + + +class TestYamlStorePromptsDbOverride: + """ + Test that YAML store_prompts_in_spend_logs takes precedence over DB-cached value. + + When store_model_in_db=true, LiteLLM persists general_settings to the DB. + On periodic reloads, _update_general_settings() must NOT override + YAML-explicit values with stale DB values. + """ + + def _make_proxy_config_with_yaml_keys(self, yaml_keys: set) -> "ProxyConfig": + """Helper: create ProxyConfig with pre-populated _yaml_general_settings_keys.""" + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = yaml_keys + return proxy_config + + @pytest.mark.asyncio + async def test_yaml_value_takes_precedence_over_db(self): + """When YAML sets store_prompts_in_spend_logs=false, DB value (true) should be ignored.""" + proxy_config = self._make_proxy_config_with_yaml_keys({"store_prompts_in_spend_logs"}) + + test_general_settings = {"store_prompts_in_spend_logs": False} + + with mock.patch("litellm.proxy.proxy_server.general_settings", test_general_settings): + await proxy_config._update_general_settings( + db_general_settings={"store_prompts_in_spend_logs": True}, + ) + + assert test_general_settings["store_prompts_in_spend_logs"] is False + + @pytest.mark.asyncio + async def test_db_value_used_when_yaml_does_not_set_key(self): + """When YAML does NOT set store_prompts_in_spend_logs, DB value should be used.""" + proxy_config = self._make_proxy_config_with_yaml_keys({"master_key", "database_url"}) + + test_general_settings = {"master_key": "sk-test"} + + with mock.patch("litellm.proxy.proxy_server.general_settings", test_general_settings): + await proxy_config._update_general_settings( + db_general_settings={"store_prompts_in_spend_logs": True}, + ) + + assert test_general_settings["store_prompts_in_spend_logs"] is True + + @pytest.mark.asyncio + async def test_admin_ui_change_works_when_yaml_omits_key(self): + """Admin UI change (DB update) should work when YAML doesn't set the key.""" + proxy_config = self._make_proxy_config_with_yaml_keys({"master_key"}) + + test_general_settings = {"master_key": "sk-test"} + + with mock.patch("litellm.proxy.proxy_server.general_settings", test_general_settings): + await proxy_config._update_general_settings( + db_general_settings={"store_prompts_in_spend_logs": True}, + ) + assert test_general_settings["store_prompts_in_spend_logs"] is True + + await proxy_config._update_general_settings( + db_general_settings={"store_prompts_in_spend_logs": False}, + ) + + assert test_general_settings["store_prompts_in_spend_logs"] is False + + def test_yaml_general_settings_keys_populated_on_load(self): + """_yaml_general_settings_keys should be empty on init.""" + proxy_config = ProxyConfig() + assert proxy_config._yaml_general_settings_keys == set() From b5823d5894d28130b1a8748c9edea898d4055452 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 09:49:13 -0700 Subject: [PATCH 045/126] feat(terraform): sync provider 0.3.0 from mirror and cut 0.4.0 The provider's release gate in project-releaser publishes only when the topmost released heading in terraform/provider/CHANGELOG.md moves past the tag the mirror already carries. That heading has been 0.2.2 since 2026-05-13, so every stable release since has correctly decided there was nothing to publish and the registry has gone stale. Two things were blocking a release: 1. The mirror shipped 0.3.0 out-of-band on 2026-07-13 (pricing_base_model, BerriAI/terraform-provider-litellm#47) after the source move, so that code exists only in the mirror. The publish rsyncs monorepo -> mirror with --delete, so publishing without this port would have deleted a released feature from the registry. 2. Nothing here declared a new version. Port #47 verbatim (resource_model.go and resource_model_crud.go are now byte-identical to the mirror's released files), backfill the 0.3.0 changelog entry it shipped under, and cut 0.4.0 covering the changes made here since the source move. 0.3.0 is not reusable as the next version -- the mirror holds that tag and the publish workflow's tag guard rejects it. --- terraform/provider/CHANGELOG.md | 13 ++++++++++++ terraform/provider/docs/resources/model.md | 2 ++ terraform/provider/litellm/resource_model.go | 8 +++++++ .../provider/litellm/resource_model_crud.go | 21 +++++++++++++++++-- 4 files changed, 42 insertions(+), 2 deletions(-) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 101519c0b08..7c744f04064 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -7,13 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.4.0] - 2026-08-06 + ### Fixed - **organization**: Send `PATCH` instead of `POST` to `/organization/update` and `/organization/member_update`, matching the methods the LiteLLM proxy serves; organization and organization member updates previously failed with a 405 +- **team_member**: Include `role` in the update payload so a role change on an existing `litellm_team_member` is applied instead of being silently dropped ### Changed - The provider source of truth moved to `terraform/provider/` in [BerriAI/litellm](https://github.com/BerriAI/litellm); this repository is now a release mirror. CI in the monorepo statically audits every endpoint the provider calls against the proxy's OpenAPI schema on every change +- **mcp_server**, **vector_store**: `env` and `litellm_params` are now marked sensitive, so they are redacted from plan/apply output, and they are no longer read back from the API into state — the configured value is authoritative. If the proxy returns values that differ from the configuration, that drift is no longer surfaced on refresh +- Dependency updates: `grpc` and `golang.org/x` modules + +## [0.3.0] - 2026-07-13 + +Released from the mirror repository before the source move was complete; this entry backfills it in the monorepo changelog. + +### Added + +- **model**: Add optional `pricing_base_model` attribute that sets `model_info.base_model` (the cost-map lookup key) independently of routing. Deployments whose routing name differs from the pricing key (for example Azure Data Zone, routed as `azure/gpt-4.1` but priced via `us/gpt-4.1-2025-04-14`) can now be billed correctly without breaking routing. When unset, behavior is unchanged and `base_model` continues to drive both routing and pricing (#47) ## [0.2.2] - 2026-05-13 diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index 5a46fe2f073..0409b48b391 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -118,6 +118,8 @@ The following arguments are supported: * `base_model` - (Required) string. The actual model identifier from the provider (e.g., "gpt-4", "claude-2"). +* `pricing_base_model` - (Optional) string. A pricing key fed to `model_info.base_model` **independently of routing**. When set, `litellm_params.model` still routes via `base_model`, but LiteLLM looks up cost against this key. Useful when the routing/deployment name differs from the cost-map key — e.g. an Azure deployment routed as `azure/gpt-4.1` whose real tier is Data Zone: set `pricing_base_model = "us/gpt-4.1-2025-04-14"` so it is billed at the Data Zone rate. When unset, `base_model` drives pricing as before. + * `litellm_credential_name` - (Optional) string. Name of a LiteLLM credential to use for this model. * `tier` - (Optional) string. The usage tier for this model. Valid values are `"free"` or `"paid"`. Default: `"free"`. diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index 2858b6e763d..4bad057871d 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -73,6 +73,14 @@ func resourceLiteLLMModel() *schema.Resource { Type: schema.TypeString, Required: true, }, + "pricing_base_model": { + // Optional pricing key fed to model_info.base_model, DECOUPLED + // from routing. When set, litellm_params.model still routes via + // base_model, but cost is looked up against this key (e.g. + // "us/gpt-4.1-2025-04-14" for Azure Data Zone pricing). + Type: schema.TypeString, + Optional: true, + }, "tier": { Type: schema.TypeString, Optional: true, diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 40766c8e312..fc5d5b09dd5 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -68,6 +68,14 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e baseModel := d.Get("base_model").(string) modelName := fmt.Sprintf("%s/%s", customLLMProvider, baseModel) + // Pricing base_model, decoupled from routing. When pricing_base_model is + // set it feeds model_info.base_model (the cost-lookup key) WITHOUT changing + // the routing string above; otherwise base_model drives pricing as before. + pricingBaseModel := baseModel + if v, ok := d.GetOk("pricing_base_model"); ok && v.(string) != "" { + pricingBaseModel = v.(string) + } + // Generate a UUID for new models modelID := d.Id() if !isUpdate { @@ -240,7 +248,7 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e ModelInfo: ModelInfo{ ID: modelID, DBModel: true, - BaseModel: baseModel, + BaseModel: pricingBaseModel, Tier: d.Get("tier").(string), Mode: d.Get("mode").(string), TeamID: d.Get("team_id").(string), @@ -306,7 +314,16 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { d.Set("rpm", GetIntValue(modelResp.LiteLLMParams.RPM, d.Get("rpm").(int))) d.Set("model_api_base", GetStringValue(modelResp.LiteLLMParams.APIBase, d.Get("model_api_base").(string))) d.Set("api_version", GetStringValue(modelResp.LiteLLMParams.APIVersion, d.Get("api_version").(string))) - d.Set("base_model", GetStringValue(modelResp.ModelInfo.BaseModel, d.Get("base_model").(string))) + // base_model / pricing_base_model read-back. When pricing_base_model is + // configured, model_info.base_model holds the PRICING key, so recover the + // routing base_model from state (not returned by the API) and read + // pricing_base_model from model_info. + if pbm, ok := d.GetOk("pricing_base_model"); ok && pbm.(string) != "" { + d.Set("base_model", d.Get("base_model").(string)) + d.Set("pricing_base_model", GetStringValue(modelResp.ModelInfo.BaseModel, pbm.(string))) + } else { + d.Set("base_model", GetStringValue(modelResp.ModelInfo.BaseModel, d.Get("base_model").(string))) + } d.Set("tier", GetStringValue(modelResp.ModelInfo.Tier, d.Get("tier").(string))) d.Set("mode", GetStringValue(modelResp.ModelInfo.Mode, d.Get("mode").(string))) d.Set("team_id", GetStringValue(modelResp.ModelInfo.TeamID, d.Get("team_id").(string))) From 845680ed1dc1e2f4b6c4493a00289e2f9422bbf0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:50:09 -0700 Subject: [PATCH 046/126] test(proxy): unit test batch file id mapping helpers directly --- .../test_common_utils.py | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py new file mode 100644 index 00000000000..4a021627c3e --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_common_utils.py @@ -0,0 +1,97 @@ +import os +import sys +from types import MappingProxyType +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.openai_files_endpoints.common_utils import ( + apply_unified_file_ids, + map_raw_file_ids_to_unified, +) +from litellm.types.utils import LiteLLMBatch + + +def _batch(input_file_id, output_file_id, error_file_id) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + object="batch", + status="cancelled", + output_file_id=output_file_id, + error_file_id=error_file_id, + ) + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_empty_ids_skips_db(): + prisma_client = MagicMock() + + assert await map_raw_file_ids_to_unified(frozenset(), prisma_client) == {} + + prisma_client.db.litellm_managedfiletable.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_no_prisma_client_returns_empty(): + assert await map_raw_file_ids_to_unified(frozenset({"file-raw-1"}), None) == {} + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_bulk_queries_and_filters_to_requested_ids(): + row_a = MagicMock( + unified_file_id="unified-a", + flat_model_file_ids=["file-raw-a", "file-raw-other"], + ) + row_b = MagicMock(unified_file_id="unified-b", flat_model_file_ids=["file-raw-b"]) + prisma_client = MagicMock() + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[row_a, row_b]) + + mapping = await map_raw_file_ids_to_unified( + frozenset({"file-raw-b", "file-raw-a", "file-raw-missing"}), prisma_client + ) + + prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once_with( + where={"flat_model_file_ids": {"hasSome": ["file-raw-a", "file-raw-b", "file-raw-missing"]}} + ) + assert dict(mapping) == {"file-raw-a": "unified-a", "file-raw-b": "unified-b"} + + +def test_apply_unified_file_ids_swaps_only_mapped_ids(): + batch = _batch(input_file_id="file-raw-in", output_file_id="file-raw-out", error_file_id=None) + + apply_unified_file_ids(batch, MappingProxyType({"file-raw-out": "unified-out"})) + + assert batch.input_file_id == "file-raw-in" + assert batch.output_file_id == "unified-out" + assert batch.error_file_id is None + + +def test_apply_unified_file_ids_swaps_all_three_ids(): + batch = _batch( + input_file_id="file-raw-in", + output_file_id="file-raw-out", + error_file_id="file-raw-err", + ) + + apply_unified_file_ids( + batch, + MappingProxyType( + { + "file-raw-in": "unified-in", + "file-raw-out": "unified-out", + "file-raw-err": "unified-err", + } + ), + ) + + assert (batch.input_file_id, batch.output_file_id, batch.error_file_id) == ( + "unified-in", + "unified-out", + "unified-err", + ) From 495eb7e7f428a64ebfb9b57004026dc7739dcbc1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 11:26:32 -0700 Subject: [PATCH 047/126] test(router): assert the auto-router max_input_chars kwarg PR #35956 added the max_input_chars passthrough to the AutoRouter constructor but left this mock assertion in tests/router_unit_tests unchanged, so test_init_auto_router_deployment_success has been failing on litellm_internal_staging ever since. The passthrough itself is intentional and its behaviour is already covered by TestAutoRouterMaxInputCharsWiring in tests/test_litellm, so only the stale expected kwargs need updating. Assert the shared constant rather than the literal 2000 so tuning the default does not break this test again. --- tests/router_unit_tests/test_router_helper_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index bcc70fae67c..0655763d41b 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -15,6 +15,7 @@ from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload from litellm.types.utils import StandardLoggingPayload from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS @pytest.fixture @@ -1816,6 +1817,7 @@ def test_init_auto_router_deployment_success(mock_auto_router, model_list): default_model="gpt-5-mini", embedding_model="text-embedding-3-small", litellm_router_instance=router, + max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ) # Verify the auto-router was added to the router's auto_routers dict From b7749f67f172fa21176f6d96991ee2ddbecf0bb6 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:27:41 -0700 Subject: [PATCH 048/126] fix(proxy): warn at startup when max_budget is set but no database is connected (#36041) * warn at startup when a proxy-wide budget is set but no DB is connected litellm.max_budget is only enforced via DB-loaded global spend, so a DB-less proxy silently ignores it. Log a one-time startup warning. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): inject max_budget into DB-less budget warning Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover DB-less budget warning startup call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): pin DB-less budget warning call site Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): stabilize budget warning call-site pin Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: tin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 17 +++++++ .../proxy/proxy_server/test_lifecycle.py | 49 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 07eaed9fe45..2e24a2d4f3c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1111,6 +1111,10 @@ async def proxy_startup_event(app: FastAPI): prisma_client=prisma_client, ) ) + ProxyStartupEvent._warn_budget_without_db( + max_budget=litellm.max_budget, + prisma_client=prisma_client, + ) ### START BATCH WRITING DB + CHECKING NEW MODELS### if prisma_client is not None: @@ -7825,6 +7829,19 @@ def giveup(e): class ProxyStartupEvent: + @staticmethod + def _warn_budget_without_db(max_budget: float | None, prisma_client: PrismaClient | None) -> None: + if prisma_client is not None or not max_budget or max_budget <= 0: + return + + verbose_proxy_logger.warning( + "A proxy-wide budget (litellm.max_budget=%s) is configured but no database is connected, " + "so the budget will NOT be enforced and requests will never be blocked. Set DATABASE_URL or " + "general_settings.database_url and restart. Redis and fail_closed_budget_enforcement do not " + "cover the proxy-wide budget because there is no global spend counter; Redis alone is not a substitute.", + max_budget, + ) + @classmethod def _initialize_startup_logging( cls, diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index cf83300ab3b..6ac1e15e7b5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -20,6 +20,7 @@ from __future__ import annotations import asyncio import inspect import json +import logging import os from typing import List, Optional, Union from unittest.mock import AsyncMock, MagicMock, patch @@ -31,6 +32,7 @@ from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps from litellm.proxy.proxy_server import ( + ProxyStartupEvent, _initialize_shared_aiohttp_session, _resolve_pydantic_type, _resolve_typed_dict_type, @@ -728,3 +730,50 @@ def test_otel_global_provider_published_after_callback_init(): "preset logger will not exist yet and a second generic logger will own " "the global provider, orphaning gen-ai spans" ) + + +def test_startup_warns_for_global_budget_without_database(caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=None) + + assert "litellm.max_budget=100.0" in caplog.text + assert "will NOT be enforced" in caplog.text + assert "requests will never be blocked" in caplog.text + + +def test_startup_does_not_warn_for_global_budget_with_database(caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=MagicMock()) + + assert "litellm.max_budget" not in caplog.text + + +@pytest.mark.parametrize("max_budget", [0, None]) +def test_startup_does_not_warn_without_global_budget(caplog, max_budget): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyStartupEvent._warn_budget_without_db(max_budget=max_budget, prisma_client=None) + + assert "litellm.max_budget" not in caplog.text + + +def test_proxy_startup_event_warns_for_global_budget_without_database(): + """Pin the lifespan call that prevents silent DB-less budgets. + + The call must follow Prisma setup so DB-backed deployments do not false-positive. + Direct ``_warn_budget_without_db`` tests cover the warning behavior itself. + """ + wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + source = inspect.getsource(wrapped) + budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:") + warn_pos = source.find("_warn_budget_without_db(") + next_startup_section_pos = source.find( + "await ProxyStartupEvent.initialize_scheduled_background_jobs(", + budget_check_pos, + ) + + assert budget_check_pos != -1, "global budget startup block not found" + assert warn_pos != -1, "DB-less budget warning call not found" + assert next_startup_section_pos != -1, "startup section after budget block not found" + assert budget_check_pos < warn_pos < next_startup_section_pos, ( + "DB-less budget warning must run after Prisma setup and the DB-backed budget block" + ) From 1d2e8b4c29a74a730b68d71d771048867084952d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 17:01:15 -0700 Subject: [PATCH 049/126] bump: litellm-enterprise 0.1.53 -> 0.1.54, litellm-proxy-extras 0.4.83 -> 0.4.84 --- enterprise/pyproject.toml | 4 ++-- litellm-proxy-extras/pyproject.toml | 4 ++-- pyproject.toml | 4 ++-- uv.lock | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 5489eba1494..a069bd81eca 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.53" +version = "0.1.54" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.53" +version = "0.1.54" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index beddd899472..fc58ff68b4d 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.83" +version = "0.4.84" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.83" +version = "0.4.84" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index 414b09eb3b4..35fd949c2e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,8 +66,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.83", - "litellm-enterprise==0.1.53", + "litellm-proxy-extras==0.4.84", + "litellm-enterprise==0.1.54", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/uv.lock b/uv.lock index 9c2897b5e4f..a42a164e5f0 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-02T02:14:05.876141Z" +exclude-newer = "2026-08-04T00:00:57.623181Z" exclude-newer-span = "P3D" [manifest] @@ -4583,12 +4583,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.53" +version = "0.1.54" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.83" +version = "0.4.84" source = { editable = "litellm-proxy-extras" } [[package]] From 988ee8b85ddeaaadc98875777e12c380fb3c618a Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 6 Aug 2026 17:07:43 -0700 Subject: [PATCH 050/126] fix(proxy): promote caller metadata trace fields into litellm_metadata (#35866) * fix(proxy): promote caller metadata trace fields into litellm_metadata Routes in LITELLM_METADATA_ROUTES keep the caller's metadata as a provider passthrough field and track proxy state in litellm_metadata, which is the dict the logging integrations read. The caller's trace_id, session_id, trace_user_id and trace_metadata therefore never reached any callback on /v1/responses, /v1/messages, /v1/batches or /v1/files, and mask_input / mask_output were dropped with them so a caller asking for redaction had their prompt logged in full. Promote an explicit allow-list of those fields from the requester_metadata snapshot into litellm_metadata, never overwriting a value already set so header-derived ids keep precedence. Trace-mutation controls (existing_trace_id, update_trace_keys) and trace_public are deliberately excluded: langfuse applies them to an arbitrary caller-chosen trace with no ownership check. tags is excluded because per-tag budget enforcement runs earlier, at auth time. This covers providers with a native Responses API config. Providers reaching /v1/responses through the chat-completions bridge need the companion change to get_litellm_params. * ci: retrigger workflows --- litellm/proxy/litellm_pre_call_utils.py | 33 ++++ .../proxy/test_litellm_pre_call_utils.py | 182 +++++++++++++++++- 2 files changed, 214 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b00ba35b14e..83ae59ef050 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -151,6 +151,20 @@ LITELLM_METADATA_ROUTES: Final = ( "files", ) +LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( + { + "mask_input", + "mask_output", + "session_id", + "trace_id", + "trace_metadata", + "trace_name", + "trace_release", + "trace_user_id", + "trace_version", + } +) + _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "proxy_server_request", "standard_logging_object", @@ -458,6 +472,18 @@ def _get_metadata_variable_name(request: Request) -> str: return "metadata" +def _promoted_trace_control_fields( + requester_metadata: Mapping[str, Any], + litellm_metadata: Mapping[str, Any], +) -> tuple[tuple[str, Any], ...]: + """Return the caller's trace-control fields that ``litellm_metadata`` does not already set.""" + return tuple( + (key, value) + for key, value in requester_metadata.items() + if key in LITELLM_TRACE_CONTROL_METADATA_FIELDS and key not in litellm_metadata + ) + + def _extract_generic_session_id_from_headers( normalized: dict[str, str], ) -> str | None: @@ -1670,6 +1696,13 @@ async def add_litellm_data_to_request( # paths may read from it. if "metadata" in data and isinstance(data["metadata"], dict): data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"]) + if _metadata_variable_name == "litellm_metadata": + data[_metadata_variable_name].update( + _promoted_trace_control_fields( + requester_metadata=data[_metadata_variable_name]["requester_metadata"], + litellm_metadata=data[_metadata_variable_name], + ) + ) # Merge litellm_metadata into the metadata variable (preserving existing # values). Runs after the user_api_key_* / _pipeline_managed_guardrails diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 0e9aac7bf85..6d6fd2e5507 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -20,6 +20,7 @@ from litellm.proxy.litellm_pre_call_utils import ( _get_dynamic_logging_metadata, _get_enforced_params, _get_metadata_variable_name, + _promoted_trace_control_fields, _resolve_credential_from_model_config, _resolve_provider_from_deployment, _update_model_if_key_alias_exists, @@ -5869,4 +5870,183 @@ async def test_key_level_callback_vars_survive_the_strip(): ) assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "key-dd-key", "dd_site": "us5.datadoghq.com"} - assert updated["dd_site"] == "us5.datadoghq.com" \ No newline at end of file + assert updated["dd_site"] == "us5.datadoghq.com" + + +class TestPromotedTraceControlFields: + """LIT-5137: caller metadata trace fields must reach litellm_metadata.""" + + def _make_request(self, path: str) -> MagicMock: + request = MagicMock(spec=Request) + request.url = MagicMock() + request.url.path = path + request.url.__str__.return_value = f"http://localhost{path}" + request.method = "POST" + request.query_params = {} + request.headers = {"Content-Type": "application/json"} + request.client = MagicMock() + request.client.host = "127.0.0.1" + return request + + async def _run(self, path: str, data: dict, headers: dict | None = None) -> dict: + request = self._make_request(path) + if headers is not None: + request.headers = {"Content-Type": "application/json", **headers} + return await add_litellm_data_to_request( + data=data, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + def test_returns_litellm_metadata_for_responses_route(self): + assert _get_metadata_variable_name(self._make_request("/v1/responses")) == "litellm_metadata" + + def test_promotes_trace_prefixed_and_allow_listed_fields(self): + requester_metadata = { + "trace_id": "trace-1", + "trace_name": "name-1", + "trace_user_id": "user-1", + "trace_metadata": {"tenant_id": "tenant-1"}, + "trace_version": "v1", + "trace_release": "r1", + "session_id": "session-1", + "mask_input": True, + "mask_output": True, + } + + promoted = _promoted_trace_control_fields( + requester_metadata=requester_metadata, + litellm_metadata={}, + ) + + assert dict(promoted) == requester_metadata + + def test_does_not_promote_unlisted_trace_prefixed_fields(self): + """trace_public flips a trace to publicly readable, so the allow-list is explicit.""" + promoted = _promoted_trace_control_fields( + requester_metadata={"trace_id": "trace-1", "trace_public": True, "trace_tags": ["a"]}, + litellm_metadata={}, + ) + + assert dict(promoted) == {"trace_id": "trace-1"} + + def test_does_not_promote_non_trace_fields(self): + promoted = _promoted_trace_control_fields( + requester_metadata={ + "trace_id": "trace-1", + "tags": ["free-tier"], + "user_api_key": "forged", + "user_api_key_user_id": "forged-user", + "spend_logs_metadata": {"forged": True}, + "guardrails": ["disabled"], + "debug_langfuse": True, + "session": "not-session-id", + "existing_trace_id": "victim-trace", + "update_trace_keys": ["input", "output"], + }, + litellm_metadata={}, + ) + + assert dict(promoted) == {"trace_id": "trace-1"} + + def test_does_not_promote_trace_mutation_controls(self): + """existing_trace_id + update_trace_keys let a caller overwrite any trace in the project.""" + promoted = _promoted_trace_control_fields( + requester_metadata={ + "trace_id": "trace-1", + "existing_trace_id": "someone-elses-trace", + "update_trace_keys": ["input", "output"], + }, + litellm_metadata={}, + ) + + assert dict(promoted) == {"trace_id": "trace-1"} + + def test_existing_litellm_metadata_value_wins(self): + promoted = _promoted_trace_control_fields( + requester_metadata={"trace_id": "from-body", "session_id": "from-body", "trace_name": "from-body"}, + litellm_metadata={"trace_id": "from-header", "session_id": "from-header"}, + ) + + assert dict(promoted) == {"trace_name": "from-body"} + + def test_empty_requester_metadata_promotes_nothing(self): + assert _promoted_trace_control_fields(requester_metadata={}, litellm_metadata={}) == () + + @pytest.mark.asyncio + async def test_responses_route_end_to_end(self): + caller_metadata = { + "trace_id": "22662678-30c1-41a1-a24b-216d6e5fb83d", + "session_id": "218af06c-28a2-4705-8a0a-5f9970d39326", + "trace_user_id": "user-123", + "trace_metadata": {"tenant_id": "tenant-1"}, + "mask_input": True, + } + + updated = await self._run( + "/v1/responses", + {"model": "gpt-4.1-mini", "input": "say resp", "metadata": copy.deepcopy(caller_metadata)}, + ) + + litellm_metadata = updated["litellm_metadata"] + for key, value in caller_metadata.items(): + assert litellm_metadata[key] == value + assert updated["metadata"] == caller_metadata + + @pytest.mark.asyncio + async def test_messages_route_end_to_end(self): + updated = await self._run( + "/v1/messages", + { + "model": "claude-sonnet-4-5", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"trace_id": "msg-trace-1", "session_id": "msg-session-1"}, + }, + ) + + assert updated["litellm_metadata"]["trace_id"] == "msg-trace-1" + assert updated["litellm_metadata"]["session_id"] == "msg-session-1" + + @pytest.mark.asyncio + async def test_session_id_header_beats_body_metadata(self): + updated = await self._run( + "/v1/responses", + {"model": "gpt-4.1-mini", "input": "say resp", "metadata": {"session_id": "from-body"}}, + headers={"x-litellm-session-id": "from-header-12345678"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "from-header-12345678" + + @pytest.mark.asyncio + async def test_forged_user_api_key_fields_are_not_promoted(self): + updated = await self._run( + "/v1/responses", + { + "model": "gpt-4.1-mini", + "input": "say resp", + "metadata": {"trace_id": "trace-1", "user_api_key_user_id": "forged", "spend_logs_metadata": {"a": 1}}, + }, + ) + + litellm_metadata = updated["litellm_metadata"] + assert litellm_metadata["trace_id"] == "trace-1" + assert litellm_metadata.get("user_api_key_user_id") != "forged" + + @pytest.mark.asyncio + async def test_chat_completions_route_is_untouched(self): + updated = await self._run( + "/v1/chat/completions", + { + "model": "gpt-4.1-mini", + "messages": [{"role": "user", "content": "hello"}], + "metadata": {"trace_id": "trace-1", "session_id": "session-1"}, + }, + ) + + assert "litellm_metadata" not in updated + assert updated["metadata"]["trace_id"] == "trace-1" + assert updated["metadata"]["session_id"] == "session-1" From f4f59ec4c35ff1b22d54e4f5f517f5922d1ea733 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 6 Aug 2026 17:25:52 -0700 Subject: [PATCH 051/126] fix(guardrails): honor configured timeout in Zscaler AI Guard (#36110) The shared `timeout` guardrail param already parsed into LitellmParams, but the Zscaler initializer never forwarded it and _send_request hardcoded a 5 second constant, so a configured value was silently ignored and slow scans failed with `Timeout passed=5` regardless of config. Forward litellm_params.timeout through to the HTTP call, keep 5 seconds as the default, fall back to it for non-positive values, and declare the field on the config model so the dashboard renders it. --- .../zscaler_ai_guard/__init__.py | 1 + .../zscaler_ai_guard/zscaler_ai_guard.py | 30 ++++- .../guardrail_hooks/zscaler_ai_guard.py | 9 ++ .../guardrails_tests/test_zscaler_ai_guard.py | 119 ++++++++++++++++++ 4 files changed, 157 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py index 408260d8483..270c28b094e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py @@ -18,6 +18,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" send_user_api_key_alias=litellm_params.send_user_api_key_alias, send_user_api_key_user_id=litellm_params.send_user_api_key_user_id, send_user_api_key_team_id=litellm_params.send_user_api_key_team_id, + timeout=litellm_params.timeout, guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, default_on=litellm_params.default_on, diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index c5c66988cb4..1aefa38ecf8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -22,9 +22,10 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.guardrails import LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel -GUARDRAIL_TIMEOUT: Final = 5 +DEFAULT_GUARDRAIL_TIMEOUT: Final = 5.0 class ZscalerAIGuard(CustomGuardrail): @@ -43,6 +44,7 @@ class ZscalerAIGuard(CustomGuardrail): send_user_api_key_alias: bool | None = None, send_user_api_key_user_id: bool | None = None, send_user_api_key_team_id: bool | None = None, + timeout: float | None = None, **kwargs, ): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) @@ -68,6 +70,7 @@ class ZscalerAIGuard(CustomGuardrail): if send_user_api_key_team_id is not None else os.getenv("SEND_USER_API_KEY_TEAM_ID", "False").lower() in ("true", "1") ) + self.timeout = self._resolve_timeout(timeout) verbose_proxy_logger.debug( "send_user_api_key_alias: %s, \n send_user_api_key_user_id:%s, \n send_user_api_key_team_id:%s", @@ -80,6 +83,29 @@ class ZscalerAIGuard(CustomGuardrail): verbose_proxy_logger.debug("ZscalerAIGuard Initializing ...") + @staticmethod + def _resolve_timeout(timeout: float | None) -> float: + """ + Resolve the effective per-request timeout, falling back to the default + when it is unset or non-positive. + """ + if timeout is None: + return DEFAULT_GUARDRAIL_TIMEOUT + + if timeout <= 0: + verbose_proxy_logger.warning( + "Ignoring non-positive Zscaler AI Guard timeout %s, using %s seconds", + timeout, + DEFAULT_GUARDRAIL_TIMEOUT, + ) + return DEFAULT_GUARDRAIL_TIMEOUT + + return timeout + + def update_in_memory_litellm_params(self, litellm_params: "LitellmParams") -> None: + super().update_in_memory_litellm_params(litellm_params) + self.timeout = self._resolve_timeout(litellm_params.timeout) + @staticmethod def _resolve_metadata_value(request_data: dict | None, key: str) -> str | None: """ @@ -267,7 +293,7 @@ class ZscalerAIGuard(CustomGuardrail): f"{url}", headers=headers, json=data, - timeout=GUARDRAIL_TIMEOUT, + timeout=self.timeout, ) response.raise_for_status() return response diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py index 3991cee8548..37125c4d583 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py @@ -79,6 +79,15 @@ class ZscalerAIGuardConfigModel(GuardrailConfigModel): json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL}, ) + timeout: float | None = Field( + default=None, + description=( + "Timeout for each Zscaler AI Guard API call, in seconds. Must be positive. " + "Raise it if scans fail under load with 'Connection timed out'. " + "Defaults to 5 seconds." + ), + ) + @model_validator(mode="after") def validate_endpoint_configuration(self) -> "ZscalerAIGuardConfigModel": """ diff --git a/tests/guardrails_tests/test_zscaler_ai_guard.py b/tests/guardrails_tests/test_zscaler_ai_guard.py index 51c86c15dcb..76e498673b0 100644 --- a/tests/guardrails_tests/test_zscaler_ai_guard.py +++ b/tests/guardrails_tests/test_zscaler_ai_guard.py @@ -396,3 +396,122 @@ async def test_apply_guardrail_block_does_not_log_error(mock_api_call): mock_logger.error.assert_not_called() assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_send_request_uses_default_timeout_when_unconfigured(): + """ + Regression: unconfigured guardrails must keep the historical 5s timeout. + """ + guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1) + + assert guardrail.timeout == 5.0 + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.zscaler_ai_guard.get_async_httpx_client" + ) as mock_get_client: + mock_client = Mock() + mock_client.post = AsyncMock(return_value=Mock(status_code=200)) + mock_get_client.return_value = mock_client + + await guardrail._send_request("http://example.com", {}, {}) + + assert mock_client.post.call_args.kwargs["timeout"] == 5.0 + + +@pytest.mark.asyncio +async def test_send_request_uses_configured_timeout(): + """ + Regression for LIT-5222: a configured timeout must reach the HTTP call. + + Before the fix _send_request passed a module-level constant, so a slow + upstream failed at 5s with `Timeout passed=5` no matter what was configured. + """ + guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=30) + + assert guardrail.timeout == 30 + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.zscaler_ai_guard.get_async_httpx_client" + ) as mock_get_client: + mock_client = Mock() + mock_client.post = AsyncMock(return_value=Mock(status_code=200)) + mock_get_client.return_value = mock_client + + await guardrail._send_request("http://example.com", {}, {}) + + assert mock_client.post.call_args.kwargs["timeout"] == 30 + + +def test_initialize_guardrail_forwards_configured_timeout(): + """ + Regression for LIT-5222: the `timeout` key from config.yaml must survive + initialization. It reaches LitellmParams already, but the initializer used + to drop it before it could reach the guardrail instance. + """ + from litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="zscaler_ai_guard", + mode="pre_call", + api_key="test_key", + api_base="http://example.com", + policy_id=1, + timeout="30", + ) + + guardrail = initialize_guardrail( + litellm_params, {"guardrail_name": "zscaler-configured-timeout"} + ) + + assert guardrail.timeout == 30.0 + + +def test_config_model_exposes_timeout_to_dashboard(): + """ + The dashboard guardrail form is built from get_config_model(), so the field + has to be declared there for the setting to be reachable outside config.yaml. + """ + config_model = ZscalerAIGuard.get_config_model() + + assert config_model is not None + assert "timeout" in config_model.model_fields + + +@pytest.mark.parametrize("bad_timeout", [0, -1]) +def test_non_positive_timeout_falls_back_to_default(bad_timeout): + """ + Regression: httpx rejects a negative timeout and treats 0 as "fail + immediately", so a non-positive value would break every scan instead of + relaxing the limit the operator was trying to raise. + """ + guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=bad_timeout) + + assert guardrail.timeout == 5.0 + + +def test_update_in_memory_litellm_params_keeps_timeout_resolved(): + """ + Regression: the base implementation copies every LitellmParams attribute + onto the guardrail, so an unset timeout would overwrite the resolved value + with None and silently fall back to the shared client's 600s default. + """ + from litellm.types.guardrails import LitellmParams + + guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=30) + assert guardrail.timeout == 30 + + guardrail.update_in_memory_litellm_params( + LitellmParams(guardrail="zscaler_ai_guard", mode="pre_call", api_key="test_key") + ) + assert guardrail.timeout == 5.0 + + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="zscaler_ai_guard", mode="pre_call", api_key="test_key", timeout=45 + ) + ) + assert guardrail.timeout == 45.0 From f3f72c4574f37ff4403ea08da6cc36cfd39b0500 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 6 Aug 2026 17:32:20 -0700 Subject: [PATCH 052/126] fix(logging): fall back to litellm_metadata when metadata is empty (#36105) get_litellm_params returned metadata=None whenever only litellm_metadata was supplied, which overwrote the fallback function_setup had already applied and left litellm_params["metadata"] empty. On the /v1/responses completion-transformation bridge, used by every provider without a native Responses API config, and on /v1/messages, that discarded the caller's trace fields a second time after the proxy had promoted them. Resolve metadata to a copy of litellm_metadata when metadata is empty, guarding on isinstance because the proxy leaves an unparseable litellm_metadata string in place and a null metadata would otherwise suppress the backfill and break the merge. update_from_kwargs copies rather than aliases for the same reason: on these routes it is handed the caller's provider-bound dict and would otherwise write user_api_key_auth into it. --- .../litellm_core_utils/get_litellm_params.py | 7 ++- litellm/litellm_core_utils/litellm_logging.py | 4 +- .../test_get_litellm_params.py | 53 +++++++++++++++++++ .../test_litellm_logging.py | 20 +++++++ 4 files changed, 80 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index d6433ad3332..f251ab4d74a 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -115,8 +115,11 @@ def get_litellm_params( litellm_request_debug: bool | None = None, **kwargs, ) -> dict: + _litellm_metadata_dict: Final = litellm_metadata if isinstance(litellm_metadata, dict) else None + resolved_metadata: Final = _litellm_metadata_dict.copy() if not metadata and _litellm_metadata_dict else metadata + # Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining) - _meta: Final = metadata or {} + _meta: Final = resolved_metadata or {} if litellm_session_id is None: litellm_session_id = _meta.get("session_id") or _meta.get("trace_id") if litellm_trace_id is None: @@ -139,7 +142,7 @@ def get_litellm_params( "model_alias_map": model_alias_map, "completion_call_id": completion_call_id, "aembedding": aembedding, - "metadata": metadata, + "metadata": resolved_metadata, "model_info": model_info, "proxy_server_request": proxy_server_request, "preset_cache_key": preset_cache_key, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9475441e214..99721c3ffa2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -585,8 +585,8 @@ class Logging(LiteLLMLoggingBaseClass): """ base_litellm_params: Final[dict[str, Any]] = {} - if "metadata" in kwargs: - base_litellm_params["metadata"] = kwargs["metadata"] + if isinstance(kwargs.get("metadata"), dict): + base_litellm_params["metadata"] = kwargs["metadata"].copy() if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict): base_litellm_params["litellm_metadata"] = kwargs["litellm_metadata"] if "metadata" not in base_litellm_params: diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index 55db31efd2c..fb4cb494bee 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -162,3 +162,56 @@ class TestGetLitellmParamsDataResidency: api_base="https://eu.api.openai.com/v1", ) assert result["data_residency"] is None + + +class TestMetadataFallsBackToLitellmMetadata: + def test_metadata_falls_back_to_litellm_metadata_when_absent(self): + result = get_litellm_params(litellm_metadata={"trace_id": "trace-1"}) + assert result["metadata"] == {"trace_id": "trace-1"} + assert result["litellm_metadata"] == {"trace_id": "trace-1"} + + def test_empty_metadata_falls_back_to_litellm_metadata(self): + result = get_litellm_params(metadata={}, litellm_metadata={"trace_id": "trace-1"}) + assert result["metadata"] == {"trace_id": "trace-1"} + + def test_metadata_wins_when_both_present(self): + result = get_litellm_params( + metadata={"trace_id": "from-metadata"}, + litellm_metadata={"trace_id": "from-litellm-metadata"}, + ) + assert result["metadata"] == {"trace_id": "from-metadata"} + + @pytest.mark.parametrize("bad_value", ["not-json-a-string", 12345, ["a"], True]) + def test_non_dict_litellm_metadata_is_ignored(self, bad_value): + result = get_litellm_params(litellm_metadata=bad_value) + assert result["metadata"] is None + + def test_metadata_stays_none_without_litellm_metadata(self): + result = get_litellm_params(api_key="test-key") + assert result["metadata"] is None + + def test_session_and_trace_id_derived_from_litellm_metadata(self): + result = get_litellm_params( + litellm_metadata={"trace_id": "trace-1", "session_id": "session-1"}, + ) + assert result["litellm_session_id"] == "session-1" + assert result["litellm_trace_id"] == "trace-1" + + def test_explicit_session_and_trace_id_are_not_overridden(self): + result = get_litellm_params( + litellm_session_id="explicit-session", + litellm_trace_id="explicit-trace", + litellm_metadata={"trace_id": "trace-1", "session_id": "session-1"}, + ) + assert result["litellm_session_id"] == "explicit-session" + assert result["litellm_trace_id"] == "explicit-trace" + + def test_litellm_metadata_fallback_is_copied_not_aliased(self): + litellm_metadata = {"trace_id": "trace-1"} + + result = get_litellm_params(litellm_metadata=litellm_metadata) + + assert result["metadata"] == litellm_metadata + assert result["metadata"] is not litellm_metadata + result["metadata"].pop("trace_id") + assert litellm_metadata == {"trace_id": "trace-1"} diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index a09c45cb141..23e0975cd08 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -526,6 +526,26 @@ class TestUpdateFromKwargs: ) assert logging_obj.litellm_params["litellm_call_id"] == "call-empty" + @pytest.mark.parametrize("caller_metadata", [None, "not-a-dict", 42]) + def test_non_dict_caller_metadata_does_not_break_the_merge(self, logging_obj, caller_metadata): + logging_obj.update_from_kwargs( + kwargs={"metadata": caller_metadata, "litellm_metadata": {"user_api_key_hash": "hashed"}}, + litellm_params={"metadata": {"user_api_key_hash": "hashed", "litellm_api_version": "1.0"}}, + ) + + assert logging_obj.litellm_params["metadata"]["user_api_key_hash"] == "hashed" + + def test_does_not_mutate_caller_metadata_dict(self, logging_obj): + caller_metadata: dict = {} + + logging_obj.update_from_kwargs( + kwargs={"metadata": caller_metadata, "litellm_metadata": {"user_api_key_hash": "hashed"}}, + litellm_params={"metadata": {"user_api_key_hash": "hashed", "litellm_api_version": "1.0"}}, + ) + + assert caller_metadata == {} + assert logging_obj.litellm_params["metadata"]["user_api_key_hash"] == "hashed" + def test_logging_prevent_double_logging(logging_obj): """ From 210ffe65fea9ad9404352f5fdb48e6067889b3df Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 6 Aug 2026 17:41:29 -0700 Subject: [PATCH 053/126] fix(proxy): re-assert the authenticated identity on passthrough requests (#36121) * fix(proxy): re-assert the authenticated identity on passthrough requests The passthrough merges the client's litellm_metadata into the request metadata and then re-asserts only user_api_key and the parent span. Every other identity field the spend and budget pipeline reads stays whatever the request body set, so a body carrying user_api_key_user_id, user_api_key_team_id, user_api_key_org_id or user_api_key_end_user_id charges that user, team, org or end user instead of the caller. Re-assert the whole sanitized identity after the merge, so the client's copy of any of those fields is overwritten by the authenticated key's own values. * test(passthrough): assert no authenticated identity field is client settable The existing regression names seven fields; the re-assertion covers every field get_sanitized_user_information_from_key returns, which is twenty today. Derive the set from the helper so a field added to StandardLoggingUserAPIKeyMetadata is covered without touching the test. Two of the twenty were not covered before, including user_api_key_hash, which is distinct from user_api_key and was client settable. --- .../pass_through_endpoints.py | 3 + .../test_pass_through_unit_tests.py | 117 ++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 64e52d252ca..8a526fcd6cb 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -565,6 +565,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): # real parent span. _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span + _metadata.update( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) + ) kwargs: Final = { "litellm_params": { diff --git a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py index 65448c6281e..c263b8ce381 100644 --- a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py +++ b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py @@ -30,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( ) from fastapi import Request from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( _update_metadata_with_tags_in_header, HttpPassThroughEndpointHelpers, @@ -652,3 +653,119 @@ def test_custom_pricing_used_in_cost_calculation(): print(f"Cache-aware cost: {cache_cost}") print("✅ Custom pricing parameters are correctly used in cost calculation") + + +def test_init_kwargs_client_metadata_cannot_spoof_authenticated_identity( + mock_request, mock_user_api_key_dict +): + request = mock_request() + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://test.com", + request_body={}, + ) + authenticated_key = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-user", + key_alias="real-key", + team_alias="Real Team", + user_email="real@example.com", + org_id="real-org", + ) + + result = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=request, + user_api_key_dict=authenticated_key, + passthrough_logging_payload=passthrough_payload, + litellm_call_id="test-call-id", + logging_obj=LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=False, + call_type="test-call-type", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ), + _parsed_body={ + "litellm_metadata": { + "user_api_key_org_id": "victim-org", + "user_api_key_end_user_id": "victim-end-user", + "user_api_key_user_id": "victim-user", + "user_api_key_team_id": "victim-team", + "user_api_key_team_alias": "Victim Team", + "user_api_key_alias": "victim-key", + "user_api_key_user_email": "victim@example.com", + } + }, + ) + + metadata = result["litellm_params"]["metadata"] + assert metadata["user_api_key_user_id"] == "test-user" + assert metadata["user_api_key_team_id"] == "test-team" + assert metadata["user_api_key_team_alias"] == "Real Team" + assert metadata["user_api_key_alias"] == "real-key" + assert metadata["user_api_key_user_email"] == "real@example.com" + assert metadata["user_api_key_org_id"] == "real-org" + assert metadata["user_api_key_end_user_id"] == "test-user" + + +def test_init_kwargs_no_authenticated_identity_field_is_client_settable( + mock_request, mock_user_api_key_dict +): + authenticated_key = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-end-user", + key_alias="real-key", + team_alias="Real Team", + user_email="real@example.com", + org_id="real-org", + organization_alias="Real Org", + project_id="real-project", + project_alias="Real Project", + spend=1.5, + max_budget=10.0, + user_spend=2.5, + user_max_budget=20.0, + team_spend=3.5, + team_max_budget=30.0, + metadata={"real": "auth-metadata"}, + ) + expected = dict( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=authenticated_key + ) + ) + assert len(expected) >= 20 + + spoofed = {key: f"SPOOFED-{key}" for key in expected} + + result = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request(), + user_api_key_dict=authenticated_key, + passthrough_logging_payload=PassthroughStandardLoggingPayload( + url="https://test.com", request_body={} + ), + litellm_call_id="test-call-id", + logging_obj=LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=False, + call_type="test-call-type", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ), + _parsed_body={"litellm_metadata": dict(spoofed), "metadata": dict(spoofed)}, + ) + + metadata = result["litellm_params"]["metadata"] + survived = { + key: metadata.get(key) + for key in expected + if metadata.get(key) != expected[key] + } + assert survived == {}, f"client-supplied values survived for: {sorted(survived)}" From 7da891a42a7604697f06ddcfe4e12d7e55a79d29 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 6 Aug 2026 17:47:05 -0700 Subject: [PATCH 054/126] fix(ui): match auto-router preset models against wildcard-expanded model groups (#36111) --- .../add_model/add_auto_router_tab.test.tsx | 62 +++++++- .../src/lib/autorouter_presets.test.ts | 142 ++++++++++++++++++ .../src/lib/autorouter_presets.ts | 36 ++++- 3 files changed, 237 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 57afdbbd28b..6a5a1e0f159 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -636,8 +636,11 @@ describe("AddAutoRouterTab", () => { expect(labels).toEqual(["Anthropic Family", "OpenAI Family", "Custom Configuration"]); }); - it("never lets a wildcard deployment satisfy a preset", async () => { - const wildcard = [{ model_name: "openai-wild", litellm_params: { model: "openai/*" } }]; + it.each([ + ["a wildcard group", "openai/*"], + ["a plain group over a wildcard underlying model", "openai-wild"], + ])("never lets %s satisfy a preset when the hub lists no expansions", async (_label, modelName) => { + const wildcard = [{ model_name: modelName, litellm_params: { model: "openai/*" } }]; mockFetchAvailableModels.mockResolvedValue(groupsFor(wildcard)); mockFetchAllModelDeployments.mockResolvedValue(wildcard); @@ -650,4 +653,59 @@ describe("AddAutoRouterTab", () => { expect(isOptionDisabled(optionByLabel("OpenAI Family")!)).toBe(true); }); }); + + describe("wildcard-matched presets", () => { + const WILDCARD_DEPLOYMENTS = [{ model_name: "someprovider/*", litellm_params: { model: "someprovider/*" } }]; + + const expandedGroupFor = (model: string): string => `someprovider/${model}`; + + const EXPANDED_HUB_GROUPS: ModelGroup[] = [ + { model_group: "someprovider/*", mode: "chat" }, + ...[...new Set(getAllPresets().flatMap((preset) => [...getRequiredModelsInPreset(preset)]))].map((model) => ({ + model_group: expandedGroupFor(model), + mode: "chat", + })), + ]; + + it("enables a preset whose models exist only as wildcard-expanded groups, labeling the match", async () => { + mockFetchAvailableModels.mockResolvedValue(EXPANDED_HUB_GROUPS); + mockFetchAllModelDeployments.mockResolvedValue(WILDCARD_DEPLOYMENTS); + + renderWithProviders(); + openTemplateDropdown(); + + await waitFor(() => { + expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(false); + }); + expect(optionByLabel("Anthropic Family")!.textContent).toContain("Matches your deployments"); + }); + + it("prefills the expanded group names and submits them", async () => { + const user = userEvent.setup(); + mockFetchAvailableModels.mockResolvedValue(EXPANDED_HUB_GROUPS); + mockFetchAllModelDeployments.mockResolvedValue(WILDCARD_DEPLOYMENTS); + + renderWithProviders(); + openTemplateDropdown(); + await waitFor(() => { + expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(false); + }); + fireEvent.click(optionByLabel("Anthropic Family")!); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "wildcard-router"); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]).toMatchObject({ + complexity_router_config: { + tiers: { + SIMPLE: ANTHROPIC_TIERS.SIMPLE.map(expandedGroupFor), + MEDIUM: ANTHROPIC_TIERS.MEDIUM.map(expandedGroupFor), + COMPLEX: ANTHROPIC_TIERS.COMPLEX.map(expandedGroupFor), + REASONING: ANTHROPIC_TIERS.REASONING.map(expandedGroupFor), + }, + }, + }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index 0d965be054b..fca8420966f 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -191,6 +191,148 @@ describe("autorouter_presets", () => { ); }); + describe("wildcard deployment matching (expanded model groups)", () => { + const wildcardDeployment = (pattern: string) => ({ modelGroup: pattern, underlyingModels: [pattern] }); + + const simpleTierConfig = (presetModel: string) => ({ + tiers: { SIMPLE: [presetModel], MEDIUM: [], COMPLEX: [], REASONING: [] }, + classifier_type: "heuristic" as const, + session_affinity: false, + }); + + it("resolves a preset model to a group expanded from a wildcard deployment", () => { + const availability = buildModelAvailability( + ["anthropic/*", "anthropic/claude-opus-5", "bedrock/anthropic.claude-opus-5"], + [wildcardDeployment("anthropic/*")], + ); + const config = simpleTierConfig("claude-opus-5"); + expect(getMissingModels(config, availability)).toEqual([]); + expect(buildPresetPrefill(config, availability).complexityRouterConfig.tiers.SIMPLE).toEqual([ + "anthropic/claude-opus-5", + ]); + }); + + it("normalizes an expanded group's namespaced own name the same way as a deployment's", () => { + const availability = buildModelAvailability( + ["bedrock/*", "bedrock/us.anthropic.claude-sonnet-5"], + [wildcardDeployment("bedrock/*")], + ); + expect(getMissingModels(simpleTierConfig("claude-sonnet-5"), availability)).toEqual([]); + }); + + it("anchors a partial wildcard pattern and treats its dots literally", () => { + const availability = buildModelAvailability( + ["bedrock/us.anthropic.claude-opus-5", "bedrock/usXanthropic.claude-fable-5"], + [wildcardDeployment("bedrock/us.*")], + ); + expect(getMissingModels(simpleTierConfig("claude-opus-5"), availability)).toEqual([]); + expect(getMissingModels(simpleTierConfig("claude-fable-5"), availability)).toEqual(["claude-fable-5"]); + }); + + it.each([ + ["gpt-5.4", "openai/gpt-5.4-mini"], + ["gpt-5.4-mini", "openai/gpt-5.4"], + ["o3", "openai/o3-mini"], + ])("never lets %s be satisfied by the expanded group %s", (presetModel, expandedGroup) => { + const availability = buildModelAvailability(["openai/*", expandedGroup], [wildcardDeployment("openai/*")]); + expect(getMissingModels(simpleTierConfig(presetModel), availability)).toEqual([presetModel]); + }); + + it("anchors the pattern's suffix and keeps middle segments in order", () => { + const availability = buildModelAvailability( + ["bedrock/us.anthropic.claude-opus-5", "bedrock/anthropic.us.claude-sonnet-5"], + [wildcardDeployment("bedrock/*.anthropic.*")], + ); + expect(getMissingModels(simpleTierConfig("claude-opus-5"), availability)).toEqual([]); + expect(getMissingModels(simpleTierConfig("claude-sonnet-5"), availability)).toEqual(["claude-sonnet-5"]); + }); + + it("matches a pathological many-star pattern in linear time instead of backtracking", () => { + const hostile = `prov/a*${"a*".repeat(30)}b`; + const nonMatching = `prov/${"a".repeat(120)}`; + const availability = buildModelAvailability([nonMatching], [wildcardDeployment(hostile)]); + expect(availability.underlyingIndex.size).toBe(0); + }); + + it("expands a bare-star model_name through its underlying wildcard, not as match-all", () => { + const availability = buildModelAvailability( + ["openai/gpt-5.4", "team-a/claude-opus-5"], + [{ modelGroup: "*", underlyingModels: ["openai/*"] }], + ); + expect(getMissingModels(simpleTierConfig("gpt-5.4"), availability)).toEqual([]); + expect(getMissingModels(simpleTierConfig("claude-opus-5"), availability)).toEqual(["claude-opus-5"]); + }); + + it.each([ + ["a bare-star underlying", "*"], + ["a non-wildcard underlying", "openai/gpt-4o"], + ["a slashless wildcard underlying", "gpt*"], + ])("derives no pattern from a bare-star model_name with %s", (_label, underlying) => { + const availability = buildModelAvailability( + ["openai/gpt-5.4"], + [{ modelGroup: "*", underlyingModels: [underlying] }], + ); + expect(availability.underlyingIndex.size).toBe(0); + }); + + it("derives no pattern from a slashless wildcard model_name", () => { + const availability = buildModelAvailability(["gpt-5.4"], [wildcardDeployment("gpt*")]); + expect(availability.underlyingIndex.size).toBe(0); + }); + + it("does not trust a group's name when no wildcard deployment covers it", () => { + const availability = buildModelAvailability( + ["team-a/claude-opus-5", "openai/*"], + [wildcardDeployment("openai/*")], + ); + expect(getMissingModels(simpleTierConfig("claude-opus-5"), availability)).toEqual(["claude-opus-5"]); + }); + + it("never resolves to the wildcard group itself when the hub lists no expansions", () => { + const availability = buildModelAvailability(["openai/*"], [wildcardDeployment("openai/*")]); + expect(getMissingModels(simpleTierConfig("gpt-5.4"), availability)).toEqual(["gpt-5.4"]); + expect(availability.underlyingIndex.size).toBe(0); + }); + + it("applies a wildcard deployment's pattern even when the wildcard group is not itself listed", () => { + const availability = buildModelAvailability(["anthropic/claude-opus-5"], [wildcardDeployment("anthropic/*")]); + expect(getMissingModels(simpleTierConfig("claude-opus-5"), availability)).toEqual([]); + }); + + it("keeps the groups-only availability strict even when expanded groups are listed", () => { + const availability = groupsOnly(["anthropic/*", "anthropic/claude-opus-5"]); + expect(getMissingModels(simpleTierConfig("claude-opus-5"), availability)).toEqual(["claude-opus-5"]); + }); + + it("prefers the alphabetically first covered group when several expansions serve the model", () => { + const availability = buildModelAvailability( + ["bedrock/us.anthropic.claude-opus-5", "anthropic/claude-opus-5", "bedrock/anthropic.claude-opus-5"], + [wildcardDeployment("anthropic/*"), wildcardDeployment("bedrock/*")], + ); + const config = simpleTierConfig("claude-opus-5"); + expect(buildPresetPrefill(config, availability).complexityRouterConfig.tiers.SIMPLE).toEqual([ + "anthropic/claude-opus-5", + ]); + }); + + it.each(getAllPresets().map((preset) => [preset.key, preset] as const))( + "fully resolves the %s preset through wildcard-expanded groups only", + (_key, preset) => { + const required = [...getRequiredModelsInPreset(preset)]; + const expandedGroups = required.map((model) => `someprovider/${model}`); + const availability = buildModelAvailability( + ["someprovider/*", ...expandedGroups], + [wildcardDeployment("someprovider/*")], + ); + expect(getMissingModelsInPreset(preset, availability)).toEqual([]); + const prefilled = buildPresetPrefill(preset.complexity_router_config, availability); + const prefilledModels = Object.values(prefilled.complexityRouterConfig.tiers).flat(); + expect(prefilledModels.length).toBeGreaterThan(0); + for (const model of prefilledModels) expect(expandedGroups).toContain(model); + }, + ); + }); + describe("deploymentRefsFromModelInfo", () => { it("keeps litellm_params.model and model_info.base_model, drops rows with neither or no name", () => { const refs = deploymentRefsFromModelInfo([ diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts index 5b6e3dc6f20..ae3f30c90f5 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts @@ -77,12 +77,30 @@ const normalizeUnderlyingModel = (model: string): string | null => { return stripped.toLowerCase() || null; }; +// A linear glob scan rather than a RegExp: patterns are admin-controlled model_name values, and a +// backtracking regex built from one ("a*a*a*...") can freeze another admin's dashboard. +const matchesWildcard = (pattern: string, name: string): boolean => { + const parts = pattern.split("*"); + if (parts.length === 1) return pattern === name; + const head = parts[0]; + const tail = parts[parts.length - 1]; + if (!name.startsWith(head) || !name.endsWith(tail)) return false; + if (name.length < head.length + tail.length) return false; + const scanEnd = name.length - tail.length; + const scanResult = parts.slice(1, -1).reduce((searchFrom: number, part: string) => { + if (searchFrom < 0) return -1; + const found = name.indexOf(part, searchFrom); + return found === -1 || found + part.length > scanEnd ? -1 : found + part.length; + }, head.length); + return scanResult >= 0; +}; + export const buildModelAvailability = ( modelGroups: Iterable, deployments: readonly DeploymentModelRef[], ): ModelAvailability => { const groups = new Set(modelGroups); - const entries = deployments + const literalEntries = deployments .filter((deployment) => groups.has(deployment.modelGroup)) .flatMap((deployment) => deployment.underlyingModels @@ -90,6 +108,22 @@ export const buildModelAvailability = ( .filter((key): key is string => key !== null) .map((key) => ({ key, modelGroup: deployment.modelGroup })), ); + // Mirrors get_known_models_from_wildcard: a bare "*" model_name expands via its underlying + // wildcard (or not at all), and a wildcard without a "/" expands to nothing. + const wildcardPatterns = Array.from( + new Set( + deployments + .flatMap((deployment) => + deployment.modelGroup === "*" ? deployment.underlyingModels : [deployment.modelGroup], + ) + .filter((pattern) => pattern !== "*" && pattern.includes("*") && pattern.includes("/")), + ), + ); + const wildcardEntries = Array.from(groups) + .filter((group) => !group.includes("*") && wildcardPatterns.some((pattern) => matchesWildcard(pattern, group))) + .map((group) => ({ key: normalizeUnderlyingModel(group), modelGroup: group })) + .filter((entry): entry is { key: string; modelGroup: string } => entry.key !== null); + const entries = [...literalEntries, ...wildcardEntries]; const grouped = new Map>(); for (const entry of entries) { const groupsForKey = grouped.get(entry.key) ?? new Set(); From 1ef019437c071d83a0e5ed573013c4b76347bd5f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:04:50 -0700 Subject: [PATCH 055/126] chore: rerun ci From 0253154780f81bd19d084eef5a23246097bc8ce5 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 18:14:26 -0700 Subject: [PATCH 056/126] build(deps-dev): bump js-yaml to 4.3.1 Closes GHSA-5p4m-2wfm-xmqj (CVSS 7.5), flagged by osv-scan against ui/litellm-dashboard/package-lock.json. js-yaml is pinned by an exact npm override, so the override and the lock move together. Dev-only dependency: js-yaml reaches the tree through eslintrc, knip and @redocly/openapi-core, none of which ship in the built dashboard. 4.3.1 published 2026-07-31, clear of the 3-day min-release-age cooldown. --- ui/litellm-dashboard/package-lock.json | 6 +++--- ui/litellm-dashboard/package.json | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index ce7cc9a13db..34dea2a39f5 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -8613,9 +8613,9 @@ "license": "MIT" }, "node_modules/js-yaml": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.0.tgz", - "integrity": "sha512-1td788aAnnZ5qs7V2QIRl1owjtYpbKt749Y3xauqQgwIIGF/xXWz1wMTEBx5O3LK3lXLVuqXPdPxj2BoFHaW9Q==", + "version": "4.3.1", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.1.tgz", + "integrity": "sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==", "dev": true, "funding": [ { diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index a9ac92cd023..3953b41a2c2 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -90,7 +90,7 @@ }, "overrides": { "prismjs": "1.30.0", - "js-yaml": "4.3.0", + "js-yaml": "4.3.1", "brace-expansion": "5.0.9", "glob": "13.0.0", "minimatch": "10.2.4", From 3a81f90ba278d16aa75b15ff580ecd1fd4a0c675 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 6 Aug 2026 18:44:07 -0700 Subject: [PATCH 057/126] build(deps): bump h2 to 4.4.1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes GHSA-6hr6-w5qg-qmwg (CVSS 5.3), the second finding from the same osv-scan run as the js-yaml bump. Bundled here so the scan goes green in one merge instead of two PRs that each stay red on the other's finding. Re-derived with `uv lock --upgrade-package h2` rather than taking the Dependabot lock wholesale: that keeps the diff to the two packages that actually move (h2, plus hpack 4.2.0 which h2 4.4.1 requires) and leaves the `exclude-newer` snapshot a real timestamp. h2 4.4.1 published 2026-08-03, hpack 4.2.0 on 2026-06-23 — both clear of the 3-day exclude-newer window. --- uv.lock | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/uv.lock b/uv.lock index a42a164e5f0..5ac80a92568 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-04T00:00:57.623181Z" +exclude-newer = "2026-08-04T01:43:29.894567Z" exclude-newer-span = "P3D" [manifest] @@ -3102,15 +3102,15 @@ wheels = [ [[package]] name = "h2" -version = "4.3.0" +version = "4.4.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "hpack" }, { name = "hyperframe" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1d/17/afa56379f94ad0fe8defd37d6eb3f89a25404ffc71d4d848893d270325fc/h2-4.3.0.tar.gz", hash = "sha256:6c59efe4323fa18b47a632221a1888bd7fde6249819beda254aeca909f221bf1", size = 2152026, upload-time = "2025-08-23T18:12:19.778Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/85/7c366e69d84c17bb778fe41419e1fbcce3033d5b7ce29bbffff0a98b859f/h2-4.4.1.tar.gz", hash = "sha256:4e866ffb1a869ae14dd9b5e6beb5c24a13da0495ad72b65925ded182521c1516", size = 2157281, upload-time = "2026-08-03T11:45:09.509Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/69/b2/119f6e6dcbd96f9069ce9a2665e0146588dc9f88f29549711853645e736a/h2-4.3.0-py3-none-any.whl", hash = "sha256:c438f029a25f7945c69e0ccf0fb951dc3f73a5f6412981daee861431b70e2bdd", size = 61779, upload-time = "2025-08-23T18:12:17.779Z" }, + { url = "https://files.pythonhosted.org/packages/7e/22/e85faf23bd72a92d1921e37d674ca56eb298a3c8be31fdecef0ff2b3aaac/h2-4.4.1-py3-none-any.whl", hash = "sha256:0e25f1462b23c9cb82d9eb02e28bc706dac2a68cb457c6a0d74d63c8a2a5d0e6", size = 62636, upload-time = "2026-08-03T11:44:59.164Z" }, ] [[package]] @@ -3239,11 +3239,11 @@ wheels = [ [[package]] name = "hpack" -version = "4.1.0" +version = "4.2.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/2c/48/71de9ed269fdae9c8057e5a4c0aa7402e8bb16f2c6e90b3aa53327b113f8/hpack-4.1.0.tar.gz", hash = "sha256:ec5eca154f7056aa06f196a557655c5b009b382873ac8d1e66e79e87535f1dca", size = 51276, upload-time = "2025-01-22T21:44:58.347Z" } +sdist = { url = "https://files.pythonhosted.org/packages/26/5b/fcabf6028144a8723726318b07a32c2f3314acdff6265743cf08a344b18e/hpack-4.2.0.tar.gz", hash = "sha256:0895cfa3b5531fc65fe439c05eb65144f123bf7a394fcaa56aa423548d8e45c0", size = 51300, upload-time = "2026-06-23T18:34:46.667Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/07/c6/80c95b1b2b94682a72cbdbfb85b81ae2daffa4291fbfa1b1464502ede10d/hpack-4.1.0-py3-none-any.whl", hash = "sha256:157ac792668d995c657d93111f46b4535ed114f0c9c8d672271bbec7eae1b496", size = 34357, upload-time = "2025-01-22T21:44:56.92Z" }, + { url = "https://files.pythonhosted.org/packages/71/b4/4a9fcfb2aef6ba44d9073ecd301443aa00b3dac95de5619f2a7de7ec8a91/hpack-4.2.0-py3-none-any.whl", hash = "sha256:858ac0b02280fa582b5080d68db0899c62a80375e0e5413a74970c5e518b6986", size = 34246, upload-time = "2026-06-23T18:34:45.472Z" }, ] [[package]] From 429a5dc430857735bec2d2a31d26312d68151f4e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 6 Aug 2026 18:46:28 -0700 Subject: [PATCH 058/126] fix(ui): allow clearing a key's budget reset from the Edit Key form (#36140) --- .../src/components/TeamSSOSettings.tsx | 2 +- .../budget_duration_dropdown.tsx | 6 +- .../organisms/create_key_button.tsx | 5 +- .../KeyInfoView.handleKeyUpdate.test.tsx | 62 +++++++++++++++++++ .../templates/key_edit_view.test.tsx | 59 ++++++++++++++++++ .../components/templates/key_edit_view.tsx | 6 +- 6 files changed, 135 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx b/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx index f6d0f38e2db..6691ef6abdb 100644 --- a/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx @@ -236,7 +236,7 @@ const TeamSSOSettings: React.FC = ({ accessToken }) => { editContent={ update("budget_duration", v)} + onChange={(v) => update("budget_duration", v ?? null)} style={{ maxWidth: 320 }} /> } diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index aa74bc60aa1..847a6ca1949 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -5,9 +5,10 @@ const { Option } = Select; interface BudgetDurationDropdownProps { value?: string | null; - onChange?: (value: string) => void; + onChange?: (value: string | undefined) => void; className?: string; style?: React.CSSProperties; + placeholder?: string; } const BudgetDurationDropdown: React.FC = ({ @@ -15,6 +16,7 @@ const BudgetDurationDropdown: React.FC = ({ onChange, className = "", style = {}, + placeholder = "n/a", }) => { return (