From 461d3407bf669cb2751502ae50215c4e2994961b Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Tue, 10 Nov 2026 09:58:46 -0800 Subject: [PATCH 01/11] fix(router): hold the registry lock across router-core writes and mirror reads (#50933) * fix(router): hold the registry lock across router-core writes and the mirror reads at their positions A registry lookup returns positions in router-core's list and the host reads its mirror at them, while an upsert wrote router-core and read the mirror back unlocked. A read on another thread in between got an IndexError past the mirror, or after a removal the neighbouring model's entries * test(router): cover the registry lock's read and write paths Move the concurrent-read regression into the registry writes characterization, drop the explanatory lock comments and the unused litellm-model lookup --- litellm/router_rust/core.py | 15 +- litellm/router_rust/surface/model_list.py | 63 ++-- litellm/router_rust/surface/registry.py | 278 ++++++++++-------- .../test_registry_characterization_lookups.py | 15 + .../test_registry_characterization_writes.py | 81 +++++ 5 files changed, 293 insertions(+), 159 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 5015795657f..b7e67fffd57 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -1815,13 +1815,14 @@ class RustCore( def _rebuild(self) -> None: """A new Rust core for changed static settings, over the same deployments and the providers the old one's builds resolved for them.""" - old: Final = self._handle - model_list: Final = list(self._hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # see __init__ - self._handle = self._new_handle(model_list) - self._handle.adopt_providers(old) - old.retire() - # The live list and indexes read the new handle's (the same entries, kept as the same objects). - self._hosted.refresh() + with self._hosted._deployment_index_lock: # pyright: ignore[reportPrivateUsage] # the registry's lock, held across the swap and the read back + old: Final = self._handle + model_list: Final = list(self._hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # see __init__ + self._handle = self._new_handle(model_list) + self._handle.adopt_providers(old) + old.retire() + # The live list and indexes read the new handle's (the same entries, kept as the same objects). + self._hosted.refresh() # The new handle builds its indexes from the deployments: written ones are handed again. self._index_sent = None diff --git a/litellm/router_rust/surface/model_list.py b/litellm/router_rust/surface/model_list.py index bb7fb53acea..32ab8f1cb76 100644 --- a/litellm/router_rust/surface/model_list.py +++ b/litellm/router_rust/surface/model_list.py @@ -26,7 +26,9 @@ model info caches) mean what they did: the same object until a write changes tha import copy import datetime +import threading from collections.abc import Callable, Iterable, Mapping, Sequence +from contextlib import AbstractContextManager from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -189,9 +191,12 @@ class HostedModelList(list[Entry]): apply: Callable[[Mapping[str, object]], None], live: Callable[[object], Mapping[str, object]] = lambda _: {}, live_info: Callable[[object], Mapping[str, object]] = lambda _: {}, + *, + lock: AbstractContextManager[object] | None = None, ) -> None: super().__init__() self._handle: Final = handle + self._lock: Final[AbstractContextManager[object]] = threading.RLock() if lock is None else lock self._apply: Final = apply # A deployment's litellm_params values its record carries as text, by id (``live_params``). self._live: Final = live @@ -209,6 +214,10 @@ class HostedModelList(list[Entry]): def refresh(self) -> None: """Read router-core's deployments again when its registry moved since the last read.""" + with self._lock: + self._refresh() + + def _refresh(self) -> None: handle: Final = self._handle() if handle is None or self._building is not None: # While a build runs, the list shows what it accepted until router-core serves it. @@ -260,28 +269,31 @@ class HostedModelList(list[Entry]): """A build starts: router-core serves what it adds once it ends, so until then the list shows each deployment it accepts here (``built``), after the entries it keeps, or alone when it replaces them, as the Python router's list grows while it adds them.""" - self.refresh() - if replacing: - self._show([]) - self._id_counts = {} - self._building = (len(self), self._id_counts, self._read_at) + with self._lock: + self._refresh() + if replacing: + self._show([]) + self._id_counts = {} + self._building = (len(self), self._id_counts, self._read_at) def built(self, entry: Mapping[str, object]) -> None: """A deployment the build accepted, shown until the build ends.""" - if self._building is not None: - super().extend(self._kept([entry], appended=True)) + with self._lock: + if self._building is not None: + super().extend(self._kept([entry], appended=True)) def end_build(self) -> None: """The build ended (or failed): router-core serves what it added, read back in place of the entries shown meanwhile, each kept as the object it was shown as while unchanged.""" - if self._building is None: - return - start, counts, read_at = self._building - self._building = None - super().__delitem__(slice(start, None)) - self._id_counts = counts - self._read_at = read_at - self.refresh() + with self._lock: + if self._building is None: + return + start, counts, read_at = self._building + self._building = None + super().__delitem__(slice(start, None)) + self._id_counts = counts + self._read_at = read_at + self._refresh() def _seed(self, entries: Sequence[object]) -> None: """Keep ``entries``, the list router-core now holds, as the objects its records for them @@ -311,19 +323,21 @@ class HostedModelList(list[Entry]): def _replace(self, entries: list[Entry]) -> None: # mutable-ok: the list router-core takes # Seeded first: router-core's records for them read back as these objects however soon # the list is read again. - self._seed(entries) - self._apply({"replace": entries}) - self.refresh() + with self._lock: + self._seed(entries) + self._apply({"replace": entries}) + self._refresh() def replace(self, entries: Iterable[object]) -> None: """Every deployment at once (``router.model_list = [...]``).""" self._replace(cast(list[Entry], list(entries))) # cast-ok: the entries as assigned, which router-core validates def append(self, value: Entry) -> None: - self.refresh() - self._seed_appended([value]) - self._apply({"append": [value]}) - self.refresh() + with self._lock: + self._refresh() + self._seed_appended([value]) + self._apply({"append": [value]}) + self._refresh() def extend(self, values: Iterable[Entry]) -> None: for value in tuple(values): @@ -362,8 +376,9 @@ class HostedModelList(list[Entry]): # A delete removes every deployment with the id; the list removes only this one. self._replace(changed) return - self._apply({"delete": [removed_id]}) - self.refresh() + with self._lock: + self._apply({"delete": [removed_id]}) + self._refresh() def pop(self, index: SupportsIndex = -1) -> Entry: entry: Final = self[index] diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index e48569c4c9d..83008bdbb94 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -743,7 +743,7 @@ class HostRegistry(HostModelInfo): self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = model_group_alias # mutable-ok: the live alias map # fmt: skip # The live model_list: router-core's deployments read back, every write a changeset. self._core_model_list: Final[HostedModelList] = HostedModelList( - self._live_handle, self._hand_changes, self.live_params, self.live_info + self._live_handle, self._hand_changes, self.live_params, self.live_info, lock=self._deployment_index_lock ) # The litellm_params values a deployment's JSON cannot carry (an Exception or model object as # ``mock_response``, a callable ``azure_ad_token_provider``), by deployment id: an attempt is called with them, as Python's list held them. @@ -1472,9 +1472,10 @@ class HostRegistry(HostModelInfo): def _hand_changes(self, changes: Mapping[str, object]) -> None: if self.changes_runner is None: raise RuntimeError("the host registry writes deployments through its Rust core") - self.changes_runner(changes) - self._core_model_list.refresh() - self._drop_removed_live_values() + with self._deployment_index_lock: + self.changes_runner(changes) + self._core_model_list.refresh() + self._drop_removed_live_values() def _drop_removed_live_values(self) -> None: """Forget the live values of ids no deployment carries any more, so a router that rotates @@ -1678,10 +1679,11 @@ class HostRegistry(HostModelInfo): return handle is not None and key is not None and handle.has_model_id(key) def _core_get_model_info(self, id: str) -> UntypedEntry | None: - handle: Final = self._live_handle() - key: Final = _index_key(id) - idx: Final[object] = None if handle is None or key is None else handle.id_position(key) - return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(id) + idx: Final[object] = None if handle is None or key is None else handle.id_position(key) + return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it def live_params(self, model_id: object) -> Mapping[str, object]: """The litellm_params values of deployment ``model_id`` that its JSON record cannot carry.""" @@ -1696,21 +1698,23 @@ class HostRegistry(HostModelInfo): return None if entry is None else Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry def _core_get_deployment_by_model_group_name(self, model_group_name: str) -> Deployment | None: - handle: Final = self._live_handle() - key: Final = _index_key(model_group_name) - indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) - model: Final = self._core_model_list[indices[0]] if indices else None # pyright: ignore[reportIndexIssue, reportArgumentType, reportCallIssue] # positions, or what a caller wrote, read as Python reads them - return None if model is None else Deployment(**model) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(model_group_name) + indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) + model: Final = self._core_model_list[indices[0]] if indices else None # pyright: ignore[reportIndexIssue, reportArgumentType, reportCallIssue] # positions, or what a caller wrote, read as Python reads them + return None if model is None else Deployment(**model) # pyright: ignore[reportArgumentType] # a stored entry def _at(self, indices: object) -> tuple[UntypedEntry, ...]: """The entries at ``indices`` (what an index holds), read as Python's list reads them.""" return tuple(self._core_model_list[idx] for idx in indices) # pyright: ignore[reportGeneralTypeIssues, reportArgumentType, reportCallIssue] # positions, or what a caller wrote def _core_deployments_named(self, model_name: str) -> tuple[UntypedEntry, ...] | None: - handle: Final = self._live_handle() - key: Final = _index_key(model_name) - indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) - return None if indices is None else self._at(indices) + with self._deployment_index_lock: + handle: Final = self._live_handle() + key: Final = _index_key(model_name) + indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) + return None if indices is None else self._at(indices) def team_public_model_names_of(self, team_id: object) -> frozenset[str]: handle: Final = self._live_handle() @@ -1719,13 +1723,14 @@ class HostRegistry(HostModelInfo): return frozenset(handle.team_public_names_of(team_id)) def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[UntypedEntry, ...] | None: - handle: Final = self._live_handle() - team: Final = _index_key(team_id) - public: Final = _index_key(public_model_name) - indices: Final[object] = ( - None if handle is None or team is None or public is None else handle.team_positions(team, public) - ) - return None if indices is None else self._at(indices) + with self._deployment_index_lock: + handle: Final = self._live_handle() + team: Final = _index_key(team_id) + public: Final = _index_key(public_model_name) + indices: Final[object] = ( + None if handle is None or team is None or public is None else handle.team_positions(team, public) + ) + return None if indices is None else self._at(indices) # -- lookups (router-core `lookup.rs`): rows read at the stored entries -------------------- @@ -1841,24 +1846,26 @@ class HostRegistry(HostModelInfo): def _team_deployments_across_teams(self, model: str) -> list[UntypedEntry]: # mutable-ok: a fresh list """Every team's deployments under public name ``model``, for a proxy admin calling without a team.""" - rows, teams = self._lookup().lookup_across_teams(model) - team_deployments: Final = self._rows(rows) - if teams > 1: - raise litellm_exceptions.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) - return team_deployments + with self._deployment_index_lock: + rows, teams = self._lookup().lookup_across_teams(model) + team_deployments: Final = self._rows(rows) + if teams > 1: + raise litellm_exceptions.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments def _core_get_all_deployments( self, model_name: str, model_alias: str | None = None, team_id: str | None = None ) -> list[UntypedEntry]: # mutable-ok: a fresh list, as the core returns it - key: Final = _index_key(model_name) - return [] if key is None else self._rows(self._lookup().lookup_all(key, team_id, model_alias)) + with self._deployment_index_lock: + key: Final = _index_key(model_name) + return [] if key is None else self._rows(self._lookup().lookup_all(key, team_id, model_alias)) def get_routing_group(self, model_name: str) -> RoutingGroup | None: """The routing group callable as ``model_name``: a deployment model_name or an alias of the @@ -1877,61 +1884,70 @@ class HostRegistry(HostModelInfo): def _core_routing_group_has_alternatives(self, model_group: str | None) -> bool: """A callable routing group whose member union spans more than one deployment.""" - if model_group is None: - return False - resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group - group: Final = self.get_routing_group(resolved) - if group is None: - return False - handle: Final = self._lookup() + with self._deployment_index_lock: + if model_group is None: + return False + resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group + group: Final = self.get_routing_group(resolved) + if group is None: + return False + handle: Final = self._lookup() - def held(member: str) -> Sequence[object]: - return cast(Sequence[object], handle.named_positions(member) or ()) # cast-ok: positions, or junk + def held(member: str) -> Sequence[object]: + return cast(Sequence[object], handle.named_positions(member) or ()) # cast-ok: positions, or junk - return sum(len(held(member)) for member in group.models) > 1 + return sum(len(held(member)) for member in group.models) > 1 def _core_get_model_list_from_model_alias( self, model_name: str | None = None ) -> list[UntypedEntry]: # mutable-ok: a fresh list - return self._rows( - self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="aliases")) - ) + with self._deployment_index_lock: + return self._rows( + self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="aliases")) + ) def _core_get_model_list_from_routing_groups(self, model_name: str | None = None) -> Sequence[UntypedEntry]: """Callable routing groups as model-list rows: each member deployment under the group's name.""" - if model_name is None and self._routing_group_rows is not None: - return self._routing_group_rows - rows: Final = tuple( - self._rows(self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="groups"))) - ) - if model_name is None: - self._routing_group_rows = rows - return rows + with self._deployment_index_lock: + if model_name is None and self._routing_group_rows is not None: + return self._routing_group_rows + rows: Final = tuple( + self._rows( + self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="groups")) + ) + ) + if model_name is None: + self._routing_group_rows = rows + return rows def _core_get_model_list( self, model_name: str | None = None, team_id: str | None = None ) -> list[UntypedEntry] | None: # mutable-ok: a fresh list, as get_model_list returns - if model_name is not None and _index_key(model_name) is None: - # No index holds it, so `_get_wildcard_deployments` reads it: what the pattern routers raise. - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - return [ - *self._core_pattern_router.get_deployments_by_pattern(model=model_name), - *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), - ] - return self._rows(self._lookup().lookup_model_list(model_name, team_id, self._list_scope(model_name))) + with self._deployment_index_lock: + if model_name is not None and _index_key(model_name) is None: + # No index holds it, so `_get_wildcard_deployments` reads it: what the pattern routers raise. + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + return [ + *self._core_pattern_router.get_deployments_by_pattern(model=model_name), + *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), + ] + return self._rows(self._lookup().lookup_model_list(model_name, team_id, self._list_scope(model_name))) def _core_wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: - from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + with self._deployment_index_lock: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider - if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model reads the base routers - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_owned: Final = team_router.get_deployments_by_pattern(model=model) if team_router is not None else () - return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) - rows: Final = self._lookup().lookup_wildcard_route( - model, team_id, [model, f"{inferred_provider(model)}/{model}"] - ) - routed: Final = cast(Sequence[Mapping[str, object]], rows) # cast-ok: router-core's wildcard rows - return [cast(Mapping[str, object], row["pattern"]) for row in routed] # cast-ok: each carries its routed copy + if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model reads the base routers + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + team_owned: Final = ( + team_router.get_deployments_by_pattern(model=model) if team_router is not None else () + ) + return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) + rows: Final = self._lookup().lookup_wildcard_route( + model, team_id, [model, f"{inferred_provider(model)}/{model}"] + ) + routed: Final = cast(Sequence[Mapping[str, object]], rows) # cast-ok: router-core's wildcard rows + return [cast(Mapping[str, object], row["pattern"]) for row in routed] # cast-ok: its routed copy def _core_wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None @@ -1942,23 +1958,27 @@ class HostRegistry(HostModelInfo): def _core_config_deployments(self) -> Iterator[Mapping[str, object]]: """A generator, as the Python core's: the rows are read when iteration starts.""" - yield from self._rows(self._lookup().lookup_config()) + with self._deployment_index_lock: + rows: Final = self._rows(self._lookup().lookup_config()) + yield from rows # -- names ------------------------------------------------------------------------------------ def _core_get_model_names(self, team_id: str | None = None) -> list[str]: # mutable-ok: a fresh list - names, raw = self._lookup().lookup_model_names(team_id, self._list_scope(None)) - self._raise_at(raw) - return names + with self._deployment_index_lock: + names, raw = self._lookup().lookup_model_names(team_id, self._list_scope(None)) + self._raise_at(raw) + return names def _core_get_model_ids( self, model_name: str | None = None, exclude_team_models: bool = False ) -> list[str]: # mutable-ok: a fresh list - if model_name is not None and _index_key(model_name) is None: - return [] - ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models) - self._raise_at(raw) - return ids + with self._deployment_index_lock: + if model_name is not None and _index_key(model_name) is None: + return [] + ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models) + self._raise_at(raw) + return ids # -- resolving a name ------------------------------------------------------------------------- @@ -2017,12 +2037,6 @@ class HostRegistry(HostModelInfo): return provider_prefixed return None - def _get_deployment_by_litellm_model(self, model: object) -> list[UntypedEntry]: # mutable-ok: a fresh list - handle: Final = self._live_handle() - if handle is None or not isinstance(model, str): - return [] - return list(self._at(handle.litellm_model_positions(model))) - # -- credentials ------------------------------------------------------------------------------ @staticmethod @@ -2041,18 +2055,20 @@ class HostRegistry(HostModelInfo): def _get_model_group_deployment_usable_by_team( self, model_group_name: str, team_id: str | None ) -> Deployment | None: - if _index_key(model_group_name) is None: - return None - position, raw = self._lookup().lookup_group_usable_by_team(model_group_name, team_id) - self._raise_at(raw) - return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + if _index_key(model_group_name) is None: + return None + position, raw = self._lookup().lookup_group_usable_by_team(model_group_name, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: - if team_id is None or _index_key(model_id) is None: - return None - position, raw = self._lookup().lookup_team_public_name(model_id, team_id) - self._raise_at(raw) - return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + with self._deployment_index_lock: + if team_id is None or _index_key(model_id) is None: + return None + position, raw = self._lookup().lookup_team_public_name(model_id, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None @@ -2105,7 +2121,10 @@ class HostRegistry(HostModelInfo): # -- routable names --------------------------------------------------------------------------- def _routable_deployments(self, model_name: str, team_id: str | None) -> tuple[UntypedEntry, ...]: - return tuple(self._rows(self._lookup().lookup_routable(self._core_routable_model_group(model_name), team_id))) + with self._deployment_index_lock: + return tuple( + self._rows(self._lookup().lookup_routable(self._core_routable_model_group(model_name), team_id)) + ) def _core_team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, @@ -2169,24 +2188,26 @@ class HostRegistry(HostModelInfo): """The deployment ids that could serve ``model`` for ``team_id``: the alias target, then a routing group, then a team public name or wildcard route for a name that is no model group, then the model group (``get_candidate_model_ids_for_route``). Read-only.""" - from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + with self._deployment_index_lock: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider - resolved: Final = self._core_get_model_from_alias(model=model) or model - group: Final = self.get_routing_group(resolved) if self._core_routing_groups else None - ids, raw = self._lookup().lookup_candidate_ids( - resolved, - team_id, - None if group is None else list(group.models), - [resolved, f"{inferred_provider(resolved)}/{resolved}"], - ) - self._raise_at(raw) - return frozenset(ids) + resolved: Final = self._core_get_model_from_alias(model=model) or model + group: Final = self.get_routing_group(resolved) if self._core_routing_groups else None + ids, raw = self._lookup().lookup_candidate_ids( + resolved, + team_id, + None if group is None else list(group.models), + [resolved, f"{inferred_provider(resolved)}/{resolved}"], + ) + self._raise_at(raw) + return frozenset(ids) def _core_team_model_has_alternatives(self, deployment_id: str) -> bool: """Whether a team deployment's public name has another unpaused deployment of that team.""" - found, raw = self._lookup().lookup_team_model_has_alternatives(deployment_id) - self._raise_at(raw) - return found + with self._deployment_index_lock: + found, raw = self._lookup().lookup_team_model_has_alternatives(deployment_id) + self._raise_at(raw) + return found def _core_resolve_served_model_name(self, model_name: str, served_model: str) -> str: """``served_model`` when ``model_name`` routes to it (the alias target, a routing group @@ -2315,17 +2336,18 @@ class HostRegistry(HostModelInfo): def _core_get_model_access_groups( self, model_name: str | None = None, model_access_group: str | None = None, team_id: str | None = None ) -> Mapping[str, Sequence[str]]: - use_cache: Final = model_name is None and model_access_group is None and team_id is None - if use_cache and self._access_groups_cache is not None: - return self._access_groups_cache - pairs, raw = self._lookup().lookup_access_groups( - model_name, model_access_group, team_id, self._list_scope(model_name) - ) - self._raise_at(raw) - groups: Final[dict[str, list[str]]] = {g: list(n) for g, n in pairs} # mutable-ok: as a defaultdict - if use_cache: - self._access_groups_cache = groups - return groups + with self._deployment_index_lock: + use_cache: Final = model_name is None and model_access_group is None and team_id is None + if use_cache and self._access_groups_cache is not None: + return self._access_groups_cache + pairs, raw = self._lookup().lookup_access_groups( + model_name, model_access_group, team_id, self._list_scope(model_name) + ) + self._raise_at(raw) + groups: Final[dict[str, list[str]]] = {g: list(n) for g, n in pairs} # mutable-ok: as a defaultdict + if use_cache: + self._access_groups_cache = groups + return groups def _core_is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: groups: Final = self._core_get_model_access_groups(model_access_group=model_access_group) diff --git a/tests/unit/test_router/core/test_registry_characterization_lookups.py b/tests/unit/test_router/core/test_registry_characterization_lookups.py index 19d58ff10d4..94c88de73fe 100644 --- a/tests/unit/test_router/core/test_registry_characterization_lookups.py +++ b/tests/unit/test_router/core/test_registry_characterization_lookups.py @@ -4,6 +4,7 @@ access groups; blocked and unhealthy groups; credentials lookups. """ import time +from typing import Final import pytest @@ -680,3 +681,17 @@ def test_assigning_model_list_swaps_the_list_without_reindexing(make_core: MakeC found = core.registry().deployment("a") assert found is not None assert found.model_name == "n" + + +def test_the_routing_groups_source_reads_its_rows_once(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + routing_groups=[{"group_name": "grp", "models": ["g"], "routing_strategy": "simple-shuffle"}], + ) + first: Final = core.registry().deployments(source="routing_groups") + assert ids(first) == ["a", "b"] + assert core.registry().deployments(source="routing_groups") is first + + +def test_the_credential_deployment_of_a_non_string_id_is_none(make_core: MakeCore) -> None: + assert _core(make_core).registry().credentials(123, form="deployment") is None # pyright: ignore[reportArgumentType] # a caller's non-string id diff --git a/tests/unit/test_router/core/test_registry_characterization_writes.py b/tests/unit/test_router/core/test_registry_characterization_writes.py index 0558a16fe9d..23fa2a8e0e9 100644 --- a/tests/unit/test_router/core/test_registry_characterization_writes.py +++ b/tests/unit/test_router/core/test_registry_characterization_writes.py @@ -5,7 +5,11 @@ what ``add_deployment``, ``set_model_list``, upserts and deletes leave behind in and ``registry()`` views. """ +import threading import warnings +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from typing import Final import pytest @@ -514,3 +518,80 @@ async def test_the_deprecated_arefresh_model_info_warns(make_core: MakeCore) -> await core.arefresh_model_info() # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public assert any("arefresh_model_info is deprecated" in str(w.message) for w in caught) + + +# -- reads racing writes ---------------------------------------------------------------------------- + +_GROUPS: Final = 4 + + +def _two_per_group(make_core: MakeCore) -> RouterCore: + return make_core(model_list=[deployment(f"m{index % _GROUPS}", f"d{index}") for index in range(2 * _GROUPS)]) + + +def _group_ids(core: RouterCore) -> Mapping[int, Sequence[str]]: + return {group: ids(core.registry().deployments(f"m{group}", source="all")) for group in range(_GROUPS)} + + +@pytest.mark.parametrize("upserted", [0, 3, _GROUPS + 1]) +def test_a_read_on_another_thread_waits_for_an_upsert_to_read_the_list_back(make_core: MakeCore, upserted: int) -> None: + """A lookup's positions are read against the mirror of the same registry version: a read that + starts once router-core holds a write waits until the mirror is read back, so it never reads + past the list nor, after a removal, the neighbouring group's entries.""" + core: Final = _two_per_group(make_core) + hosted: Final = core._hosted # pyright: ignore[reportAttributeAccessIssue] # the registry under test + write: Final[Callable[[Mapping[str, object]], None]] = hosted.changes_runner + started: Final = threading.Event() + pool: Final = ThreadPoolExecutor(max_workers=1) + reads: Final[list[object]] = [] # mutable-ok: one per write + + def read() -> Mapping[int, Sequence[str]]: + started.set() + return _group_ids(core) + + def write_with_a_read_started(changes: Mapping[str, object]) -> None: + write(changes) + started.clear() + reads.append(pool.submit(read)) + assert started.wait(timeout=5) + # The reader runs while router-core holds the write and the mirror does not yet. + _ = threading.Event().wait(0.02) + + hosted.changes_runner = write_with_a_read_started + try: + _ = core.apply(ChangeSet(upsert=[typed(f"m{upserted % _GROUPS}", f"d{upserted}", timeout=11)])) + finally: + hosted.changes_runner = write + + assert reads + for future in reads: + seen = future.result(timeout=5) # pyright: ignore[reportAttributeAccessIssue] # a Future + for group in range(_GROUPS): + # Between the upsert's removal and its add, its group is served by the other deployment. + assert seen[group], (group, seen) + assert set(seen[group]) <= {f"d{group}", f"d{group + _GROUPS}"}, (group, seen) + pool.shutdown() + assert sorted(_group_ids(core)[upserted % _GROUPS]) == sorted( + [f"d{upserted % _GROUPS}", f"d{upserted % _GROUPS + _GROUPS}"] + ) + + +def test_appending_to_the_live_model_list_reads_the_entry_back(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + live.append(deployment("m0", "d8")) + assert ids(live) == [*(f"d{index}" for index in range(2 * _GROUPS)), "d8"] + + +def test_deleting_a_live_model_list_entry_reads_the_list_back(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + del live[0] + assert ids(live) == [f"d{index}" for index in range(1, 2 * _GROUPS)] + + +def test_ending_a_build_that_never_began_leaves_the_list(make_core: MakeCore) -> None: + core: Final = _two_per_group(make_core) + live: Final = core._hosted._core_model_list # pyright: ignore[reportAttributeAccessIssue] # the live list + live.end_build() + assert ids(live) == [f"d{index}" for index in range(2 * _GROUPS)] From 3e3842046d8b2b4b70fefe7e2819333c770aaa9c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 17:59:48 +0000 Subject: [PATCH 02/11] test(budgets): expect per-model tpm and rpm limits in budget info (#50967) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/management/test_budget_updates.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/integration/management/test_budget_updates.py b/tests/integration/management/test_budget_updates.py index 6939b447866..aa351078ff9 100644 --- a/tests/integration/management/test_budget_updates.py +++ b/tests/integration/management/test_budget_updates.py @@ -36,7 +36,8 @@ def test_shortening_budget_duration_moves_reset_at_onto_the_new_schedule(gateway BUDGET_COLUMNS: Final = ( "budget_id, max_budget, soft_budget, max_parallel_requests, tpm_limit, rpm_limit, tpd_limit, model_max_budget, " - "budget_duration, allowed_models, temp_budget_increase, temp_budget_expiry, created_by, updated_by" + "model_tpm_limit, model_rpm_limit, budget_duration, allowed_models, temp_budget_increase, temp_budget_expiry, " + "created_by, updated_by" ) VOLATILE_INFO_FIELDS: Final = frozenset({"created_at", "updated_at", "budget_reset_at"}) UNLOADED_RELATIONS: Final[dict[str, JsonValue]] = dict.fromkeys( @@ -128,6 +129,8 @@ def test_partial_update_keeps_sibling_budget_fields(gateway: Gateway) -> None: "rpm_limit": 60, "tpd_limit": 100000, "model_max_budget": {"openai/gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}}, + "model_tpm_limit": None, + "model_rpm_limit": None, "budget_duration": "30d", "allowed_models": [], "temp_budget_increase": None, From c229851924c4187718d9b102ca32e4a804319a94 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 10 Nov 2026 10:00:20 -0800 Subject: [PATCH 03/11] feat(ui): customer portal (beta) surfacing the existing customer endpoints and customer usage (#50807) * feat(ui): Customer Portal (Beta) surfacing the existing customer endpoints and customer usage Adds a Customer Portal mode to the dashboard (product switcher entry, Customers and Chargebacks nav) built only on endpoints and data the proxy already serves. Customers: list ranked by this month's spend from the customer daily-activity aggregate, status (active, over budget, blocked), search, Create Customer through /customer/new (id, alias, monthly max_budget, tpm_limit, rpm_limit), Edit alias and budget through /customer/update, Block and Unblock. Customer detail shows the keys that sent traffic as that customer this month from the api_key_breakdown of the same aggregate, read only. Chargebacks: per month, grouped by customer, customer and key, or customer and model, with CSV export. All strings go through i18n. Pure helpers (customerForms, customerPortalData, chargebackRows) carry unit tests, and the ViewSwitcher integration test covers entering and leaving the portal. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): refresh the customers list after editing or blocking from customer detail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): name the expected objects in the customer portal tests so the inline-object lint budget holds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): mock next/navigation in the navbar test now that ViewSwitcher routes with useRouter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): format customer portal money and limits through the format keys instead of inline formatters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): note that customer blocking needs blocked_user_check and neutralise formula prefixes in the chargeback CSV Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chargebacks/chargebackRows.test.ts | 106 +++++++++++ .../(dashboard)/chargebacks/chargebackRows.ts | 87 +++++++++ .../src/app/(dashboard)/chargebacks/page.tsx | 175 ++++++++++++++++++ .../_components/CreateCustomerDialog.tsx | 122 ++++++++++++ .../customers/_components/CustomerDetail.tsx | 156 ++++++++++++++++ .../_components/EditCustomerDialog.tsx | 96 ++++++++++ .../customers/_components/StatusBadge.tsx | 15 ++ .../_components/customerForms.test.ts | 49 +++++ .../customers/_components/customerForms.ts | 46 +++++ .../_components/customerPortalData.test.ts | 110 +++++++++++ .../_components/customerPortalData.ts | 89 +++++++++ .../src/app/(dashboard)/customers/page.tsx | 153 +++++++++++++++ .../src/app/(dashboard)/layout.tsx | 10 +- .../src/app/(dashboard)/legacyPageRoutes.ts | 2 + .../Navbar/ViewSwitcher.integration.test.tsx | 28 ++- .../src/components/Navbar/ViewSwitcher.tsx | 40 +++- .../src/components/leftnav.tsx | 32 +++- .../src/components/navbar.test.tsx | 2 + .../src/contexts/PluginModeContext.tsx | 6 +- ui/litellm-dashboard/src/i18n/locales/en.json | 118 +++++++++++- 20 files changed, 1429 insertions(+), 13 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/chargebacks/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerDetail.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/EditCustomerDialog.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/StatusBadge.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.test.ts new file mode 100644 index 00000000000..041d683803c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.test.ts @@ -0,0 +1,106 @@ +import { describe, expect, test } from "vitest"; +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import type { CustomerUsage } from "../customers/_components/customerPortalData"; +import { rowsByCustomer, rowsByCustomerAndKey, rowsByCustomerAndModel, sumRows, toCsv } from "./chargebackRows"; + +const totals = (spend: number, requests = 1) => ({ + spend, + requests, + inputTokens: requests * 10, + outputTokens: requests * 5, +}); + +const usage: ReadonlyMap = new Map([ + [ + "acme", + { + ...totals(6, 3), + keys: [ + { hash: "sk-hash-aaaaaaaaaaaa", alias: "prod", ...totals(5, 2) }, + { hash: "sk-hash-bbbbbbbbbbbb", alias: null, ...totals(1, 1) }, + ], + }, + ], + ["idle", { ...totals(0, 0), keys: [] }], + ["globex", { ...totals(9, 1), keys: [{ hash: "h3", alias: "main", ...totals(9, 1) }] }], +]); + +describe("rowsByCustomer", () => { + test("one row per customer with traffic, highest spend first", () => { + expect(rowsByCustomer(usage).map((r) => [r.customer, r.dimension, r.spend])).toEqual([ + ["globex", null, 9], + ["acme", null, 6], + ]); + }); +}); + +describe("rowsByCustomerAndKey", () => { + test("one row per customer and key, labelled by alias or hash prefix", () => { + expect(rowsByCustomerAndKey(usage).map((r) => [r.customer, r.dimension, r.spend])).toEqual([ + ["globex", "main", 9], + ["acme", "prod", 5], + ["acme", "sk-hash-bbbb", 1], + ]); + }); +}); + +describe("rowsByCustomerAndModel", () => { + const metrics = (spend: number): SpendMetrics => ({ + spend, + prompt_tokens: 1, + completion_tokens: 1, + total_tokens: 2, + api_requests: 1, + successful_requests: 1, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }); + const day = (models: Record): DailyData => ({ + date: "2026-11-01", + metrics: metrics(0), + breakdown: { + models: Object.fromEntries( + Object.entries(models).map(([m, s]) => [m, { metrics: metrics(s), metadata: {}, api_key_breakdown: {} }]), + ), + model_groups: {}, + mcp_servers: {}, + providers: {}, + api_keys: {}, + entities: {}, + }, + }); + test("sums each customer's models across days", () => { + const rows = rowsByCustomerAndModel(new Map([["acme", [day({ "gpt-5": 1, claude: 2 }), day({ "gpt-5": 4 })]]])); + expect(rows.map((r) => [r.customer, r.dimension, r.spend, r.requests])).toEqual([ + ["acme", "gpt-5", 5, 2], + ["acme", "claude", 2, 1], + ]); + }); +}); + +describe("sumRows and toCsv", () => { + test("total row matches the sum of the rows", () => { + const expected = { spend: 15, requests: 4, inputTokens: 40, outputTokens: 20 }; + expect(sumRows(rowsByCustomer(usage))).toEqual(expected); + }); + test("csv carries alias, optional dimension column and escaped quotes", () => { + const rows = rowsByCustomerAndKey(usage).slice(0, 1); + const csv = toCsv(rows, new Map([["globex", 'Globex "Inc"']]), "key"); + expect(csv.split("\n")).toEqual([ + '"customer_id","customer","key","requests","input_tokens","output_tokens","spend_usd"', + '"globex","Globex ""Inc""","main","1","10","5","9.000000"', + ]); + }); + test("csv neutralises ids and aliases that a spreadsheet would run as a formula", () => { + const rows = rowsByCustomer(new Map([['=HYPERLINK("x")', usage.get("acme")!]])); + const line = toCsv(rows, new Map([['=HYPERLINK("x")', "-1+1"]]), null).split("\n")[1]; + expect(line.startsWith(`"'=HYPERLINK(""x"")","'-1+1",`)).toBe(true); + expect(line.endsWith('"3","30","15","6.000000"')).toBe(true); + }); + test("csv without a dimension has no dimension column", () => { + expect(toCsv(rowsByCustomer(usage), new Map(), null).split("\n")[0]).toBe( + '"customer_id","customer","requests","input_tokens","output_tokens","spend_usd"', + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.ts b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.ts new file mode 100644 index 00000000000..7f86b7cf6d6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/chargebackRows.ts @@ -0,0 +1,87 @@ +import type { DailyData } from "@/components/UsagePage/types"; +import { addMetrics, type CustomerUsage, type Totals } from "../customers/_components/customerPortalData"; + +export type GroupBy = "customer" | "key" | "model"; + +export interface ChargebackRow extends Totals { + customer: string; + dimension: string | null; +} + +const ZERO: Totals = { spend: 0, requests: 0, inputTokens: 0, outputTokens: 0 }; + +const hasTraffic = (r: Totals) => r.requests > 0 || r.spend > 0; +const bySpend = (a: Totals, b: Totals) => b.spend - a.spend; + +export const rowsByCustomer = (usage: ReadonlyMap): ChargebackRow[] => + [...usage.entries()] + .map(([customer, u]) => ({ + customer, + dimension: null, + spend: u.spend, + requests: u.requests, + inputTokens: u.inputTokens, + outputTokens: u.outputTokens, + })) + .filter(hasTraffic) + .sort(bySpend); + +export const rowsByCustomerAndKey = (usage: ReadonlyMap): ChargebackRow[] => + [...usage.entries()] + .flatMap(([customer, u]) => u.keys.map((k) => ({ ...k, customer, dimension: k.alias ?? k.hash.slice(0, 12) }))) + .filter(hasTraffic) + .sort(bySpend); + +export const rowsByCustomerAndModel = (perCustomer: ReadonlyMap): ChargebackRow[] => + [...perCustomer.entries()] + .flatMap(([customer, days]) => { + const byModel = new Map(); + for (const day of days) { + for (const [model, m] of Object.entries(day.breakdown.models)) + byModel.set(model, addMetrics(byModel.get(model) ?? ZERO, m.metrics)); + } + return [...byModel.entries()].map(([model, t]) => ({ ...t, customer, dimension: model })); + }) + .filter(hasTraffic) + .sort(bySpend); + +export const sumRows = (rows: readonly Totals[]): Totals => + rows.reduce( + (acc, r) => ({ + spend: acc.spend + r.spend, + requests: acc.requests + r.requests, + inputTokens: acc.inputTokens + r.inputTokens, + outputTokens: acc.outputTokens + r.outputTokens, + }), + ZERO, + ); + +const FORMULA_PREFIX = /^[=+\-@\t\r]/; +const neutralise = (v: string) => (FORMULA_PREFIX.test(v) ? `'${v}` : v); +const csvCell = (v: string | number) => `"${(typeof v === "string" ? neutralise(v) : String(v)).replace(/"/g, '""')}"`; + +export const toCsv = ( + rows: readonly ChargebackRow[], + aliasOf: ReadonlyMap, + dimension: string | null, +): string => { + const header = [ + "customer_id", + "customer", + ...(dimension ? [dimension] : []), + "requests", + "input_tokens", + "output_tokens", + "spend_usd", + ]; + const lines = rows.map((r) => [ + r.customer, + aliasOf.get(r.customer) ?? r.customer, + ...(dimension ? [r.dimension ?? ""] : []), + r.requests, + r.inputTokens, + r.outputTokens, + r.spend.toFixed(6), + ]); + return [header, ...lines].map((l) => l.map(csvCell).join(",")).join("\n"); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/page.tsx new file mode 100644 index 00000000000..262b918bde5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/chargebacks/page.tsx @@ -0,0 +1,175 @@ +"use client"; + +import { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; +import { useQuery } from "@tanstack/react-query"; +import { Download } from "lucide-react"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { dailyActivityAggregatedCall } from "@/components/networking"; +import { toDailyData } from "@/components/UsagePage/dailyActivityApi"; +import { Button } from "@/components/ui/button"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { monthRange, useCustomerUsage } from "../customers/_components/customerPortalData"; +import { + rowsByCustomer, + rowsByCustomerAndKey, + rowsByCustomerAndModel, + sumRows, + toCsv, + type GroupBy, +} from "./chargebackRows"; + +const GROUPS: readonly GroupBy[] = ["customer", "key", "model"]; +const MONTHS = [0, 1, 2] as const; + +const downloadCsv = (csv: string, filename: string) => { + const url = URL.createObjectURL(new Blob([csv], { type: "text/csv" })); + Object.assign(document.createElement("a"), { href: url, download: filename }).click(); + URL.revokeObjectURL(url); +}; + +export default function ChargebacksPage() { + const { t, i18n } = useTranslation(); + const { accessToken } = useAuthorized(); + const [monthsBack, setMonthsBack] = useState(0); + const [groupBy, setGroupBy] = useState("customer"); + const customersQuery = useCustomers(100); + const customers = useMemo(() => customersQuery.data?.pages.flat() ?? [], [customersQuery.data]); + const aliasOf = useMemo(() => new Map(customers.map((c) => [c.user_id, c.alias || c.user_id])), [customers]); + const usageQuery = useCustomerUsage(accessToken, monthsBack); + const customerIds = useMemo(() => [...(usageQuery.data?.keys() ?? [])].sort(), [usageQuery.data]); + + const perModel = useQuery({ + queryKey: ["customer-portal", "usage-by-model", monthsBack, customerIds], + enabled: Boolean(accessToken) && groupBy === "model" && usageQuery.isSuccess, + queryFn: async () => { + const { start, end } = monthRange(monthsBack); + const entries = await Promise.all( + customerIds.map(async (id) => { + const request = { accessToken: accessToken!, startTime: start, endTime: end, entityIds: [id] }; + const res = await dailyActivityAggregatedCall("customer", request); + return [id, toDailyData(res)] as const; + }), + ); + return new Map(entries); + }, + }); + + const rows = useMemo(() => { + if (groupBy === "model") return perModel.data ? rowsByCustomerAndModel(perModel.data) : []; + if (!usageQuery.data) return []; + return groupBy === "key" ? rowsByCustomerAndKey(usageQuery.data) : rowsByCustomer(usageQuery.data); + }, [groupBy, usageQuery.data, perModel.data]); + const total = sumRows(rows); + const isLoading = usageQuery.isLoading || (groupBy === "model" && perModel.isLoading); + const dimensionLabel = groupBy === "customer" ? null : t(`customerPortal.chargebacks.dimension.${groupBy}`); + const monthLabel = (back: number) => + new Intl.DateTimeFormat(i18n.language, { month: "long", year: "numeric", timeZone: "UTC" }).format( + monthRange(back).start, + ); + const exportCsv = () => + downloadCsv( + toCsv(rows, aliasOf, dimensionLabel?.toLowerCase() ?? null), + `chargebacks-${monthRange(monthsBack).start.toISOString().slice(0, 7)}-${groupBy}.csv`, + ); + + return ( +
+
+
+

{t("customerPortal.chargebacks.title")}

+

{t("customerPortal.chargebacks.subtitle")}

+
+
+ + + +
+
+
+ + + + {t("customerPortal.chargebacks.columns.customer")} + {dimensionLabel && {dimensionLabel}} + {t("customerPortal.chargebacks.columns.requests")} + {t("customerPortal.chargebacks.columns.inputTokens")} + {t("customerPortal.chargebacks.columns.outputTokens")} + {t("customerPortal.chargebacks.columns.spend")} + + + + {isLoading && ( + + + {t("customerPortal.common.loading")} + + + )} + {!isLoading && rows.length === 0 && ( + + + {t("customerPortal.chargebacks.empty", { month: monthLabel(monthsBack) })} + + + )} + {rows.map((r) => ( + + +

{aliasOf.get(r.customer) ?? r.customer}

+

{r.customer}

+
+ {dimensionLabel && {r.dimension}} + {t("format.number", { val: r.requests })} + {t("format.number", { val: r.inputTokens })} + {t("format.number", { val: r.outputTokens })} + {t("format.usd", { val: r.spend })} +
+ ))} + {rows.length > 0 && ( + + + {t("customerPortal.common.total")} + + {t("format.number", { val: total.requests })} + + {t("format.number", { val: total.inputTokens })} + + + {t("format.number", { val: total.outputTokens })} + + {t("format.usd", { val: total.spend })} + + )} +
+
+
+
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx new file mode 100644 index 00000000000..c645c180e53 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerDialog.tsx @@ -0,0 +1,122 @@ +"use client"; + +import { useState } from "react"; +import { useTranslation } from "react-i18next"; +import { toast } from "sonner"; +import { $api } from "@/lib/http/api"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { buildCreateCustomerBody, digitsOnly, decimalOnly } from "./customerForms"; + +interface Props { + open: boolean; + onOpenChange: (o: boolean) => void; + onCreated: (userId: string) => void; +} + +const errorText = (e: unknown) => (e instanceof Error ? e.message : String(e)); + +export default function CreateCustomerDialog({ open, onOpenChange, onCreated }: Props) { + const { t } = useTranslation(); + const [userId, setUserId] = useState(""); + const [alias, setAlias] = useState(""); + const [budget, setBudget] = useState(""); + const [tpm, setTpm] = useState(""); + const [rpm, setRpm] = useState(""); + const create = $api.useMutation("post", "/customer/new", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.created", { name: alias.trim() || userId.trim() })); + onCreated(userId.trim()); + onOpenChange(false); + }, + onError: (e: unknown) => toast.error(t("customerPortal.toasts.createFailed", { error: errorText(e) })), + }); + const submit = (e: React.FormEvent) => { + e.preventDefault(); + const fields = { userId, alias, budget, tpm, rpm }; + const body = buildCreateCustomerBody(fields); + if (body) create.mutate({ body }); + }; + return ( + + +
+ + {t("customerPortal.createDialog.title")} + {t("customerPortal.createDialog.description")} + +
+ + setUserId(e.target.value)} + placeholder={t("customerPortal.createDialog.customerIdPlaceholder")} + autoFocus + required + /> +
+
+ + setAlias(e.target.value)} + placeholder={t("customerPortal.createDialog.aliasPlaceholder")} + /> +
+
+ + setBudget(decimalOnly(e.target.value))} + placeholder={t("customerPortal.createDialog.maxBudgetPlaceholder")} + /> +
+
+
+ + setTpm(digitsOnly(e.target.value))} + placeholder={t("customerPortal.createDialog.noLimit")} + /> +
+
+ + setRpm(digitsOnly(e.target.value))} + placeholder={t("customerPortal.createDialog.noLimit")} + /> +
+
+

{t("customerPortal.createDialog.limitsHelp")}

+ + + + +
+
+
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerDetail.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerDetail.tsx new file mode 100644 index 00000000000..6e9e94dedf2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerDetail.tsx @@ -0,0 +1,156 @@ +"use client"; + +import { useState } from "react"; +import { useTranslation } from "react-i18next"; +import { useQueryClient } from "@tanstack/react-query"; +import { Ban, Pencil } from "lucide-react"; +import { toast } from "sonner"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { $api } from "@/lib/http/api"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { customerStatus, useCustomerUsage } from "./customerPortalData"; +import EditCustomerDialog from "./EditCustomerDialog"; +import StatusBadge from "./StatusBadge"; + +function UsageSummary({ spend, budget }: { spend: number; budget: EndUser["litellm_budget_table"] | undefined }) { + const { t } = useTranslation(); + const max = budget?.max_budget ?? null; + const tpmLimit = budget?.tpm_limit ?? null; + const rpmLimit = budget?.rpm_limit ?? null; + return ( +

+ + {t("customerPortal.detail.spendThisMonth")}{" "} + {t("format.usd", { val: spend })} + + + {max === null + ? t("customerPortal.detail.noBudget") + : t("customerPortal.detail.ofBudget", { max: t("format.usd", { val: max }) })} + + {tpmLimit !== null && ( + + {t("customerPortal.detail.tpm", { val: t("format.number", { val: tpmLimit }) })} + + )} + {rpmLimit !== null && ( + + {t("customerPortal.detail.rpm", { val: t("format.number", { val: rpmLimit }) })} + + )} +

+ ); +} + +export default function CustomerDetail({ customerId, onBack }: { customerId: string; onBack: () => void }) { + const { t } = useTranslation(); + const { accessToken } = useAuthorized(); + const info = $api.useQuery( + "get", + "/customer/info", + { params: { query: { end_user_id: customerId } } }, + { enabled: Boolean(accessToken) }, + ); + const usageQuery = useCustomerUsage(accessToken); + const [editOpen, setEditOpen] = useState(false); + const queryClient = useQueryClient(); + const refresh = () => { + info.refetch(); + queryClient.invalidateQueries({ queryKey: ["get", "/customer/list"] }); + }; + const block = $api.useMutation("post", "/customer/block", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.blocked")); + refresh(); + }, + }); + const unblock = $api.useMutation("post", "/customer/update", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.unblocked")); + refresh(); + }, + }); + + const c = info.data; + const usage = usageQuery.data?.get(customerId); + const spend = usage?.spend ?? 0; + const keys = usage?.keys ?? []; + const label = c?.alias || customerId; + + return ( +
+
+
+
+

{label}

+ {c && } + {customerId} +
+ +
+
+ + {c?.blocked ? ( + + ) : ( + + )} +
+
+

{t("customerPortal.common.blockHint")}

+ +
+
+

{t("customerPortal.detail.keysTitle")}

+

{t("customerPortal.detail.keysSubtitle")}

+
+ {keys.length === 0 ? ( +

{t("customerPortal.detail.keysEmpty")}

+ ) : ( + + + + {t("customerPortal.detail.columns.key")} + {t("customerPortal.detail.columns.requests")} + {t("customerPortal.detail.columns.tokens")} + {t("customerPortal.detail.columns.spend")} + + + + {keys.map((k) => ( + + +

{k.alias ?? t("customerPortal.detail.unnamedKey")}

+

{k.hash.slice(0, 12)}

+
+ {t("format.number", { val: k.requests })} + + {t("format.number", { val: k.inputTokens + k.outputTokens })} + + {t("format.usd", { val: k.spend })} +
+ ))} +
+
+ )} +
+ + {c && } +
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/EditCustomerDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/EditCustomerDialog.tsx new file mode 100644 index 00000000000..aebc8d60053 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/EditCustomerDialog.tsx @@ -0,0 +1,96 @@ +"use client"; + +import { useState } from "react"; +import { useTranslation } from "react-i18next"; +import { toast } from "sonner"; +import { $api } from "@/lib/http/api"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { buildUpdateCustomerBody, decimalOnly } from "./customerForms"; + +interface Props { + open: boolean; + onOpenChange: (o: boolean) => void; + customer: EndUser; + onSaved: () => void; +} + +const asText = (n: number | null | undefined) => (n === null || n === undefined ? "" : String(n)); + +function EditCustomerForm({ onOpenChange, customer, onSaved }: Omit) { + const { t } = useTranslation(); + const [alias, setAlias] = useState(customer.alias ?? ""); + const [budget, setBudget] = useState(asText(customer.litellm_budget_table?.max_budget)); + const budgetId = customer.litellm_budget_table?.budget_id ?? null; + const update = $api.useMutation("post", "/customer/update", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.updated")); + onSaved(); + onOpenChange(false); + }, + onError: (e: unknown) => + toast.error(t("customerPortal.toasts.updateFailed", { error: e instanceof Error ? e.message : String(e) })), + }); + const submit = (e: React.FormEvent) => { + e.preventDefault(); + const fields = { userId: customer.user_id, alias, budget, blocked: Boolean(customer.blocked) }; + update.mutate({ body: buildUpdateCustomerBody(fields) }); + }; + return ( + +
+ + {t("customerPortal.editDialog.title")} + + {customer.user_id} + + +
+ + setAlias(e.target.value)} autoFocus /> +
+
+ + setBudget(decimalOnly(e.target.value))} + placeholder={t("customerPortal.editDialog.maxBudgetPlaceholder")} + /> +
+ {budgetId && ( +

+ {t("customerPortal.editDialog.sharedBudgetWarning", { budgetId })} +

+ )} + + + + +
+
+ ); +} + +export default function EditCustomerDialog({ open, ...form }: Props) { + return ( + + {open && } + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/StatusBadge.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/StatusBadge.tsx new file mode 100644 index 00000000000..005171abc53 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/StatusBadge.tsx @@ -0,0 +1,15 @@ +import { useTranslation } from "react-i18next"; +import { Badge } from "@/components/ui/badge"; +import type { CustomerStatus } from "./customerPortalData"; + +const VARIANT = { active: "secondary", over_budget: "destructive", blocked: "outline" } as const; +const LABEL_KEY = { + active: "customerPortal.status.active", + over_budget: "customerPortal.status.overBudget", + blocked: "customerPortal.status.blocked", +} as const; + +export default function StatusBadge({ status }: { status: CustomerStatus }) { + const { t } = useTranslation(); + return {t(LABEL_KEY[status])}; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.test.ts new file mode 100644 index 00000000000..3c7ecc97484 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.test.ts @@ -0,0 +1,49 @@ +import { describe, expect, test } from "vitest"; +import { buildCreateCustomerBody, buildUpdateCustomerBody } from "./customerForms"; + +describe("buildCreateCustomerBody", () => { + const fields = { userId: " acme ", alias: " Acme Corp ", budget: "25", tpm: "1000", rpm: "60" }; + const emptyFields = { userId: "acme", alias: "", budget: "", tpm: "", rpm: "" }; + + test("trims ids and aliases and sends numeric limits with a monthly budget reset", () => { + const expected = { + user_id: "acme", + alias: "Acme Corp", + blocked: false, + max_budget: 25, + budget_duration: "1mo", + tpm_limit: 1000, + rpm_limit: 60, + }; + expect(buildCreateCustomerBody(fields)).toEqual(expected); + }); + + test("empty optional fields become null and no budget means no budget_duration", () => { + const expected = { + user_id: "acme", + alias: null, + blocked: false, + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + }; + expect(buildCreateCustomerBody(emptyFields)).toEqual(expected); + }); + + test("a blank customer id produces no request", () => { + expect(buildCreateCustomerBody({ ...fields, userId: " " })).toBeNull(); + }); +}); + +describe("buildUpdateCustomerBody", () => { + test("preserves the blocked flag and clears the budget when emptied", () => { + const fields = { userId: "acme", alias: "", budget: "", blocked: true }; + const expected = { user_id: "acme", alias: null, blocked: true, max_budget: null }; + expect(buildUpdateCustomerBody(fields)).toEqual(expected); + }); + test("parses a decimal budget", () => { + const fields = { userId: "acme", alias: "Acme", budget: "12.50", blocked: false }; + expect(buildUpdateCustomerBody(fields).max_budget).toBe(12.5); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.ts new file mode 100644 index 00000000000..f1fafc0c718 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerForms.ts @@ -0,0 +1,46 @@ +import type { components } from "@/lib/http/schema"; + +export type NewCustomerBody = components["schemas"]["NewCustomerRequest"]; +export type UpdateCustomerBody = components["schemas"]["UpdateCustomerRequest"]; + +export const digitsOnly = (s: string) => s.replace(/[^0-9]/g, ""); +export const decimalOnly = (s: string) => s.replace(/[^0-9.]/g, ""); + +const numberOrNull = (s: string): number | null => (s.trim() === "" ? null : Number(s)); + +export interface CreateCustomerFields { + userId: string; + alias: string; + budget: string; + tpm: string; + rpm: string; +} + +export const buildCreateCustomerBody = (f: CreateCustomerFields): NewCustomerBody | null => { + const userId = f.userId.trim(); + if (!userId) return null; + const maxBudget = numberOrNull(f.budget); + return { + user_id: userId, + alias: f.alias.trim() || null, + blocked: false, + max_budget: maxBudget, + budget_duration: maxBudget === null ? null : "1mo", + tpm_limit: numberOrNull(f.tpm), + rpm_limit: numberOrNull(f.rpm), + }; +}; + +export interface EditCustomerFields { + userId: string; + alias: string; + budget: string; + blocked: boolean; +} + +export const buildUpdateCustomerBody = (f: EditCustomerFields): UpdateCustomerBody => ({ + user_id: f.userId, + alias: f.alias.trim() || null, + blocked: f.blocked, + max_budget: numberOrNull(f.budget), +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.test.ts new file mode 100644 index 00000000000..cf48c81db5b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.test.ts @@ -0,0 +1,110 @@ +import { describe, expect, test } from "vitest"; +import type { DailyData, KeyMetricWithMetadata, MetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; +import { customerStatus, monthRange, rankBySpend, usageByCustomer } from "./customerPortalData"; + +const metrics = (spend: number, requests = 1, prompt = 10, completion = 5): SpendMetrics => ({ + spend, + prompt_tokens: prompt, + completion_tokens: completion, + total_tokens: prompt + completion, + api_requests: requests, + successful_requests: requests, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}); + +const key = (alias: string | null, spend: number, requests = 1): KeyMetricWithMetadata => ({ + metrics: metrics(spend, requests), + metadata: { key_alias: alias, team_id: null }, +}); + +const entity = (spend: number, keys: Record, requests = 1): MetricWithMetadata => ({ + metrics: metrics(spend, requests), + metadata: {}, + api_key_breakdown: keys, +}); + +const day = (date: string, entities: Record): DailyData => ({ + date, + metrics: metrics(0, 0, 0, 0), + breakdown: { models: {}, model_groups: {}, mcp_servers: {}, providers: {}, api_keys: {}, entities }, +}); + +describe("usageByCustomer", () => { + test("sums a customer's spend, requests and tokens across days", () => { + const usage = usageByCustomer([ + day("2026-11-01", { acme: entity(1.5, {}, 2) }), + day("2026-11-02", { acme: entity(2.25, {}, 3) }), + ]); + const expected = { spend: 3.75, requests: 5, inputTokens: 20, outputTokens: 10 }; + expect(usage.get("acme")).toMatchObject(expected); + }); + + test("keeps customers separate", () => { + const usage = usageByCustomer([day("2026-11-01", { acme: entity(1, {}), globex: entity(4, {}) })]); + expect(usage.get("acme")?.spend).toBe(1); + expect(usage.get("globex")?.spend).toBe(4); + }); + + test("merges the same key across days and sorts keys by spend", () => { + const usage = usageByCustomer([ + day("2026-11-01", { acme: entity(3, { h1: key("prod", 1), h2: key("staging", 2) }) }), + day("2026-11-02", { acme: entity(5, { h1: key("prod", 5) }) }), + ]); + const prod = { hash: "h1", alias: "prod", spend: 6, requests: 2 }; + const staging = { hash: "h2", alias: "staging", spend: 2, requests: 1 }; + expect(usage.get("acme")?.keys).toEqual([expect.objectContaining(prod), expect.objectContaining(staging)]); + }); + + test("keeps the first non-null alias seen for a key", () => { + const usage = usageByCustomer([ + day("2026-11-01", { acme: entity(1, { h1: key(null, 1) }) }), + day("2026-11-02", { acme: entity(1, { h1: key("prod", 1) }) }), + ]); + expect(usage.get("acme")?.keys[0]?.alias).toBe("prod"); + }); +}); + +describe("customerStatus", () => { + test("blocked wins over budget", () => { + expect(customerStatus({ blocked: true, litellm_budget_table: { max_budget: 1 } }, 5)).toBe("blocked"); + }); + test("over budget once spend reaches max_budget", () => { + expect(customerStatus({ blocked: false, litellm_budget_table: { max_budget: 10 } }, 10)).toBe("over_budget"); + expect(customerStatus({ blocked: false, litellm_budget_table: { max_budget: 10 } }, 9.99)).toBe("active"); + }); + test("no budget is active regardless of spend", () => { + expect(customerStatus({ blocked: false, litellm_budget_table: null }, 1000)).toBe("active"); + }); +}); + +describe("rankBySpend", () => { + const customer = (user_id: string) => ({ user_id, blocked: false, spend: 0, litellm_budget_table: null }); + test("orders customers by this month's spend, highest first, zero for customers without traffic", () => { + const usage = usageByCustomer([day("2026-11-01", { b: entity(2, {}), c: entity(9, { h1: key("k", 9) }) })]); + const rows = rankBySpend([customer("a"), customer("b"), customer("c")], usage); + expect(rows.map((r) => [r.c.user_id, r.spend, r.keys])).toEqual([ + ["c", 9, 1], + ["b", 2, 0], + ["a", 0, 0], + ]); + }); +}); + +describe("monthRange", () => { + test("covers the whole UTC month, monthsBack months ago", () => { + const now = new Date(Date.UTC(2026, 10, 9, 12)); + expect(monthRange(0, now)).toEqual({ + start: new Date(Date.UTC(2026, 10, 1)), + end: new Date(Date.UTC(2026, 10, 30, 23, 59, 59)), + }); + expect(monthRange(1, now)).toEqual({ + start: new Date(Date.UTC(2026, 9, 1)), + end: new Date(Date.UTC(2026, 9, 31, 23, 59, 59)), + }); + }); + test("crosses a year boundary", () => { + expect(monthRange(2, new Date(Date.UTC(2026, 0, 15))).start).toEqual(new Date(Date.UTC(2025, 10, 1))); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.ts new file mode 100644 index 00000000000..5bc1c2ec437 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPortalData.ts @@ -0,0 +1,89 @@ +import { useQuery } from "@tanstack/react-query"; +import { dailyActivityAggregatedCall } from "@/components/networking"; +import { toDailyData } from "@/components/UsagePage/dailyActivityApi"; +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; + +export type CustomerStatus = "active" | "over_budget" | "blocked"; + +export const monthRange = (monthsBack = 0, now = new Date()): { start: Date; end: Date } => ({ + start: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() - monthsBack, 1)), + end: new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth() - monthsBack + 1, 0, 23, 59, 59)), +}); + +export interface Totals { + spend: number; + requests: number; + inputTokens: number; + outputTokens: number; +} + +export interface CustomerKeyUsage extends Totals { + hash: string; + alias: string | null; +} + +export interface CustomerUsage extends Totals { + keys: readonly CustomerKeyUsage[]; +} + +const ZERO: Totals = { spend: 0, requests: 0, inputTokens: 0, outputTokens: 0 }; + +export const addMetrics = (t: Totals, m: Partial | undefined): Totals => ({ + spend: t.spend + (m?.spend ?? 0), + requests: t.requests + (m?.api_requests ?? 0), + inputTokens: t.inputTokens + (m?.prompt_tokens ?? 0), + outputTokens: t.outputTokens + (m?.completion_tokens ?? 0), +}); + +export const usageByCustomer = (days: readonly DailyData[]): Map => { + const totals = new Map(); + const keys = new Map>(); + for (const day of days) { + for (const [id, entity] of Object.entries(day.breakdown.entities)) { + totals.set(id, addMetrics(totals.get(id) ?? ZERO, entity.metrics)); + const byKey = keys.get(id) ?? new Map(); + for (const [hash, key] of Object.entries(entity.api_key_breakdown)) { + const prev = byKey.get(hash) ?? { ...ZERO, hash, alias: key.metadata.key_alias }; + byKey.set(hash, { ...prev, ...addMetrics(prev, key.metrics), alias: prev.alias ?? key.metadata.key_alias }); + } + keys.set(id, byKey); + } + } + return new Map( + [...totals.entries()].map(([id, t]) => [ + id, + { ...t, keys: [...(keys.get(id)?.values() ?? [])].sort((a, b) => b.spend - a.spend) }, + ]), + ); +}; + +export const useCustomerUsage = (accessToken: string | null, monthsBack = 0) => + useQuery({ + queryKey: ["customer-portal", "usage", monthsBack], + enabled: Boolean(accessToken), + queryFn: async () => { + const { start, end } = monthRange(monthsBack); + const res = await dailyActivityAggregatedCall("customer", { + accessToken: accessToken!, + startTime: start, + endTime: end, + }); + return usageByCustomer(toDailyData(res)); + }, + }); + +export const customerStatus = (c: Pick, spend: number): CustomerStatus => { + if (c.blocked) return "blocked"; + const max = c.litellm_budget_table?.max_budget; + return max !== null && max !== undefined && spend >= max ? "over_budget" : "active"; +}; + +export const rankBySpend = (customers: readonly EndUser[], usage: ReadonlyMap | undefined) => + customers + .map((c) => { + const u = usage?.get(c.user_id); + const spend = u?.spend ?? 0; + return { c, spend, keys: u?.keys.length ?? 0, status: customerStatus(c, spend) }; + }) + .sort((a, b) => b.spend - a.spend); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx new file mode 100644 index 00000000000..f80d15f5acf --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx @@ -0,0 +1,153 @@ +"use client"; + +import { useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; +import { parseAsString, useQueryState } from "nuqs"; +import { Plus } from "lucide-react"; +import { toast } from "sonner"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useCustomers, type EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { $api } from "@/lib/http/api"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import CustomerDetail from "./_components/CustomerDetail"; +import CreateCustomerDialog from "./_components/CreateCustomerDialog"; +import StatusBadge from "./_components/StatusBadge"; +import { rankBySpend, useCustomerUsage } from "./_components/customerPortalData"; + +export default function CustomersPage() { + const { t } = useTranslation(); + const { accessToken } = useAuthorized(); + const [selectedId, setSelectedId] = useQueryState("id", parseAsString); + const [search, setSearch] = useState(""); + const [createOpen, setCreateOpen] = useState(false); + const customersQuery = useCustomers(100, search || undefined); + const usageQuery = useCustomerUsage(accessToken); + const customers: EndUser[] = useMemo(() => customersQuery.data?.pages.flat() ?? [], [customersQuery.data]); + const rows = useMemo(() => rankBySpend(customers, usageQuery.data), [customers, usageQuery.data]); + const block = $api.useMutation("post", "/customer/block", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.blocked")); + customersQuery.refetch(); + }, + }); + const unblock = $api.useMutation("post", "/customer/update", { + onSuccess: () => { + toast.success(t("customerPortal.toasts.unblocked")); + customersQuery.refetch(); + }, + }); + + if (selectedId) return setSelectedId(null)} />; + + return ( +
+
+
+

{t("customerPortal.list.title")}

+

{t("customerPortal.list.subtitle")}

+
+ +
+

{t("customerPortal.common.blockHint")}

+
+
+ setSearch(e.target.value)} + /> +
+ + + + {t("customerPortal.list.columns.customer")} + {t("customerPortal.list.columns.spend")} + {t("customerPortal.list.columns.budget")} + {t("customerPortal.list.columns.keys")} + {t("customerPortal.list.columns.status")} + + + + + {customersQuery.isLoading && ( + + + {t("customerPortal.list.loading")} + + + )} + {!customersQuery.isLoading && rows.length === 0 && ( + + + {t("customerPortal.list.empty")} + + + )} + {rows.map(({ c, spend, keys, status }) => { + const max = c.litellm_budget_table?.max_budget ?? null; + return ( + setSelectedId(c.user_id, { history: "push" })} + > + +

{c.alias || c.user_id}

+

{c.user_id}

+
+ + {t("format.usd", { val: spend })} + + + {max === null + ? t("customerPortal.common.noBudget") + : t("customerPortal.list.budgetUsage", { + spend: t("format.usd", { val: spend }), + max: t("format.usd", { val: max }), + })} + + {keys} + + + + e.stopPropagation()}> + {c.blocked ? ( + + ) : ( + + )} + +
+ ); + })} +
+
+
+ { + customersQuery.refetch(); + setSelectedId(id, { history: "push" }); + }} + /> +
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 5c93d214812..d898f1b0373 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -116,13 +116,19 @@ export function AgentControlPlaneView() { } const FULL_BLEED_SEGMENTS = new Set(["logs"]); +const CUSTOMER_PORTAL_SEGMENTS = new Set(["customers", "chargebacks"]); function DashboardShell({ children }: { children: React.ReactNode }) { const { t } = useTranslation(); const { accessToken } = useAuth(); - const { mode } = usePluginMode(); + const { mode, setMode } = usePluginMode(); const pathname = usePathname(); const routeSegment = routeSegmentForPathname(pathname); + const isPortalRoute = CUSTOMER_PORTAL_SEGMENTS.has(routeSegment); + useEffect(() => { + if (isPortalRoute && mode !== "customer-portal") setMode("customer-portal"); + if (!isPortalRoute && mode === "customer-portal") setMode("ai-gateway"); + }, [isPortalRoute, mode, setMode]); useEffect(() => { void recordUiEvent({ page: routeSegment || "home", action: "view" }); }, [routeSegment]); @@ -142,7 +148,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) { const sidebarCollapsed = sidebarOverride?.segment === routeSegment ? sidebarOverride.collapsed : isFullBleed; const toggleSidebar = () => setSidebarOverride({ segment: routeSegment, collapsed: !sidebarCollapsed }); - const isGateway = mode === "ai-gateway"; + const isGateway = mode === "ai-gateway" || mode === "customer-portal"; // Non-gateway (agent control plane) mode keeps the original full-width Navbar, // which carries the account menu; the redesigned sidebar + header shell is diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts index 19dac805948..90b49e56c5a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts @@ -12,6 +12,8 @@ const LEGACY_PAGE_ROUTES: ReadonlyMap = new Map( chat: "chat", "access-groups": "access-groups", budgets: "budgets", + customers: "customers", + chargebacks: "chargebacks", workflows: "workflows", "guardrails-monitor": "guardrails-monitor", "mcp-servers": "mcp-servers", diff --git a/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx b/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx index 64bb8f15585..a0f6e59d029 100644 --- a/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx @@ -8,13 +8,14 @@ import { setServerRootPath } from "@/lib/serverRootPath"; import ViewSwitcher from "./ViewSwitcher"; -const { state, getUiSettingsMock, getUiConfigMock } = vi.hoisted(() => ({ +const { state, getUiSettingsMock, getUiConfigMock, pushMock } = vi.hoisted(() => ({ state: { pathname: "/ui/", plugins: [] as Plugin[] }, getUiSettingsMock: vi.fn(), getUiConfigMock: vi.fn(), + pushMock: vi.fn(), })); -vi.mock("next/navigation", () => ({ usePathname: () => state.pathname })); +vi.mock("next/navigation", () => ({ usePathname: () => state.pathname, useRouter: () => ({ push: pushMock }) })); vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "", getUiSettings: getUiSettingsMock, @@ -56,6 +57,8 @@ describe("ViewSwitcher", () => { }); afterEach(() => { + pushMock.mockClear(); + assignSpy.mockClear(); Object.defineProperty(window, "location", { configurable: true, value: originalLocation }); testQueryClient.clear(); localStorage.clear(); @@ -103,6 +106,27 @@ describe("ViewSwitcher", () => { expect(assignSpy).not.toHaveBeenCalled(); }); + it("should open the Customer Portal at the customers page and come back to the gateway", async () => { + const user = userEvent.setup(); + renderSwitcher(); + + await user.click(screen.getByRole("button", { name: "Switch workspace, current: AI Gateway" })); + await user.click(await screen.findByRole("menuitem", { name: /Customer Portal/ })); + expect(pushMock).toHaveBeenCalledWith("/ui/customers"); + expect(assignSpy).not.toHaveBeenCalled(); + expect(localStorage.getItem("litellm_plugin_mode")).toBeNull(); + }); + + it("should route back to the gateway from Customer Portal mode instead of only flipping the mode", async () => { + localStorage.setItem("litellm_plugin_mode", "customer-portal"); + const user = userEvent.setup(); + renderSwitcher(); + + await user.click(await screen.findByRole("button", { name: "Switch workspace, current: Customer Portal" })); + await user.click(await screen.findByRole("menuitem", { name: /^AI Gateway/ })); + expect(pushMock).toHaveBeenCalledWith("/ui/api-keys"); + }); + it("should navigate from the gateway to enabled Chat without changing plugin mode", async () => { getUiSettingsMock.mockResolvedValue({ values: { enable_chat_ui: true } }); const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx b/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx index 5c8d4a7a850..38b051515fb 100644 --- a/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx +++ b/ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx @@ -1,7 +1,7 @@ import { useTerminology } from "@/contexts/TerminologyContext"; import React from "react"; import { useTranslation } from "react-i18next"; -import { usePathname } from "next/navigation"; +import { usePathname, useRouter } from "next/navigation"; import { Button } from "@/components/ui/button"; import { DropdownMenu, @@ -11,6 +11,7 @@ import { } from "@/components/ui/dropdown-menu"; import { Check, ChevronDown, ChevronsUpDown } from "lucide-react"; import { useSidebar } from "@/components/shared/Sidebar"; +import BetaBadge from "@/components/BetaBadge"; import { usePluginMode } from "@/contexts/PluginModeContext"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import { uiHref } from "@/utils/uiHref"; @@ -19,6 +20,7 @@ import i18n from "@/i18n/i18n"; import { productNames } from "@/i18n/untranslated/productNames"; const GATEWAY = "ai-gateway"; +const CUSTOMER_PORTAL = "customer-portal"; const CHAT = "chat"; const MOYAI = "moyai"; @@ -51,6 +53,9 @@ interface ViewSwitcherProps { contentClassName?: string; } +const pluginLabel = (mode: string, plugins: readonly { name: string; display_name: string }[], portalLabel: string) => + mode === CUSTOMER_PORTAL ? portalLabel : plugins.find((p) => p.name === mode)?.display_name; + export default function ViewSwitcher({ className, contentClassName }: ViewSwitcherProps) { const { t } = useTranslation(); const term = useTerminology(); @@ -58,6 +63,7 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch const { mode, setMode, plugins } = usePluginMode(); const { data: uiSettings } = useUISettings(); const pathname = usePathname(); + const router = useRouter(); const chatEnabled = Boolean(uiSettings?.values?.enable_chat_ui); const moyaiUrl = (uiSettings?.values?.moyai_url as string | undefined) ?? null; @@ -71,7 +77,7 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch const activeLabel = activeWorkspaceLabel( isChatRoute, isMoyaiRoute, - plugins.find((p) => p.name === mode)?.display_name, + pluginLabel(mode, plugins, t("customerPortal.switcher.label")), gatewayLabel, ); @@ -81,6 +87,14 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch ]; const selectMode = (key: string) => { + if (key === CUSTOMER_PORTAL) { + router.push(uiHref("customers")); + return; + } + if (mode === CUSTOMER_PORTAL && key === GATEWAY) { + router.push(uiHref("api-keys")); + return; + } setMode(key); if (isStandaloneRoute) { window.location.assign(uiHref("")); @@ -107,6 +121,21 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch onClick: chatEnabled ? () => window.location.assign(uiHref(CHAT)) : undefined, }; + const customerPortalItem: ViewSwitcherItem = { + key: CUSTOMER_PORTAL, + label: ( +
+ + + {t("customerPortal.switcher.label")} + {" "} + {t("customerPortal.switcher.description")} + + {!isStandaloneRoute && mode === CUSTOMER_PORTAL && } +
+ ), + onClick: () => selectMode(CUSTOMER_PORTAL), + }; const items: ViewSwitcherItem[] = [ ...modeEntries.map((e) => ({ key: e.key, @@ -118,6 +147,7 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch ), onClick: () => selectMode(e.key), })), + customerPortalItem, chatItem, { key: MOYAI, @@ -161,6 +191,12 @@ export default function ViewSwitcher({ className, contentClassName }: ViewSwitch <> {activeLabel} + {!isStandaloneRoute && mode === CUSTOMER_PORTAL && ( + <> + {" "} + + + )} diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index f48271a872c..6040da32f6c 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -55,12 +55,14 @@ import { Trophy, User, Users, + Receipt, Wallet, Wrench, Workflow, } from "lucide-react"; import Link from "next/link"; import ViewSwitcher from "@/components/Navbar/ViewSwitcher"; +import { usePluginMode } from "@/contexts/PluginModeContext"; import { AriaSidebarLink } from "@/components/liteadmin/LiteAdmin"; import { usePathname } from "next/navigation"; import { useMemo, useState } from "react"; @@ -197,6 +199,28 @@ export const visibleMenuGroups = (groups: readonly MenuGroup[], context: MenuVis // Menu groups organized by category - defined outside component for export. // Shape (key/page/label/roles/children) is consumed by page_utils.ts; only the // icons changed to lucide as part of the sidebar redesign. +const customerPortalGroups: MenuGroup[] = [ + { + groupLabel: "CUSTOMER PORTAL", + items: [ + { + key: "customers", + page: "customers", + label: "shell.nav.items.customers", + icon: , + roles: all_admin_roles, + }, + { + key: "chargebacks", + page: "chargebacks", + label: "shell.nav.items.chargebacks", + icon: , + roles: all_admin_roles, + }, + ], + }, +]; + const menuGroups: MenuGroup[] = [ { groupLabel: "AI GATEWAY", @@ -542,6 +566,7 @@ const findMenuItemKey = (route: string): string => { const GROUP_LABELS: Record = { "AI GATEWAY": "shell.nav.groups.aiGateway", + "CUSTOMER PORTAL": "shell.nav.groups.customerPortal", OBSERVABILITY: "shell.nav.groups.observability", "ACCESS CONTROL": "shell.nav.groups.accessControl", "DEVELOPER TOOLS": "shell.nav.groups.developerTools", @@ -553,6 +578,7 @@ export const groupText = (groupLabel: string): string => const SECTION_DISPLAY: Record = { "AI GATEWAY": "shell.nav.sections.aiGateway", + "CUSTOMER PORTAL": "shell.nav.sections.customerPortal", OBSERVABILITY: "shell.nav.sections.observability", "ACCESS CONTROL": "shell.nav.sections.accessControl", "DEVELOPER TOOLS": "shell.nav.sections.developerTools", @@ -575,7 +601,7 @@ export const labelText = (item: MenuItem): string => export const getBreadcrumb = (pathname: string): { section: string | null; title: string } => { const route = routeForPathname(pathname); if (route === "aria") return { section: null, title: "Aria Agent" }; - for (const group of menuGroups) { + for (const group of [...menuGroups, ...customerPortalGroups]) { for (const item of group.items) { const section = sectionText(group.groupLabel); if (routeOf(item) === route) return { section, title: labelText(item) }; @@ -597,6 +623,7 @@ const Sidebar_: React.FC = ({ allowVectorStoresForTeamAdmins, }) => { const { userId, accessToken, userRole, isViewOnly } = useAuthorized(); + const { mode: pluginMode } = usePluginMode(); const { t } = useTranslation(); const term = useTerminology(); const labelOf = (item: MenuItem) => (typeof item.label === "string" ? term(labelText(item)) : item.label); @@ -639,8 +666,9 @@ const Sidebar_: React.FC = ({ disableVectorStoresForInternalUsers, allowVectorStoresForTeamAdmins, }; - return visibleMenuGroups(menuGroups, context); + return visibleMenuGroups(pluginMode === "customer-portal" ? customerPortalGroups : menuGroups, context); }, [ + pluginMode, allowAgentsForTeamAdmins, allowVectorStoresForTeamAdmins, disableAgentsForInternalUsers, diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index cc90e0b8e99..38b539c51b9 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -4,6 +4,8 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor, within } from "../../tests/test-utils"; import Navbar from "./navbar"; +vi.mock("next/navigation", () => ({ usePathname: () => "/api-keys", useRouter: () => ({ push: vi.fn() }) })); + // Mock the hooks and utilities vi.mock("@/components/networking", () => ({ getProxyBaseUrl: vi.fn(() => "http://localhost:4000"), diff --git a/ui/litellm-dashboard/src/contexts/PluginModeContext.tsx b/ui/litellm-dashboard/src/contexts/PluginModeContext.tsx index 380a9479829..247cae44235 100644 --- a/ui/litellm-dashboard/src/contexts/PluginModeContext.tsx +++ b/ui/litellm-dashboard/src/contexts/PluginModeContext.tsx @@ -4,7 +4,8 @@ import React, { createContext, useContext, useState, useEffect } from "react"; import { createApiClient } from "@/lib/http/client"; import { getProxyBaseUrl } from "@/components/networking"; -export type PluginMode = "ai-gateway" | string; // "ai-gateway" or a registered plugin name +export type PluginMode = "ai-gateway" | "customer-portal" | string; // built-in modes or a registered plugin name +export const BUILT_IN_MODES: ReadonlySet = new Set(["ai-gateway", "customer-portal"]); export interface PluginNavItem { key: string; @@ -84,7 +85,8 @@ export function PluginModeProvider({ children, accessToken }: PluginModeProvider // Once plugins have loaded, fall back to ai-gateway if the persisted mode is // no longer registered — including when the list came back empty (all plugins // removed). Derived rather than setState-in-effect to avoid cascading renders. - const effectiveMode = mode !== "ai-gateway" && loaded && !plugins.some((p) => p.name === mode) ? "ai-gateway" : mode; + const effectiveMode = + !BUILT_IN_MODES.has(mode) && loaded && !plugins.some((p) => p.name === mode) ? "ai-gateway" : mode; const setMode = (m: PluginMode) => { setModeState(m); diff --git a/ui/litellm-dashboard/src/i18n/locales/en.json b/ui/litellm-dashboard/src/i18n/locales/en.json index d4ba9c14998..113c77e1916 100644 --- a/ui/litellm-dashboard/src/i18n/locales/en.json +++ b/ui/litellm-dashboard/src/i18n/locales/en.json @@ -7994,14 +7994,16 @@ "observability": "OBSERVABILITY", "accessControl": "ACCESS CONTROL", "developerTools": "DEVELOPER TOOLS", - "settings": "SETTINGS" + "settings": "SETTINGS", + "customerPortal": "CUSTOMER PORTAL" }, "sections": { "aiGateway": "AI Gateway", "observability": "Observability", "accessControl": "Access Control", "developerTools": "Developer Tools", - "settings": "Settings" + "settings": "Settings", + "customerPortal": "Customer Portal" }, "items": { "home": "Home", @@ -8047,7 +8049,9 @@ "adminSettings": "Admin Settings", "costTracking": "Cost Tracking", "uiTheme": "UI Theme", - "telemetry": "Telemetry" + "telemetry": "Telemetry", + "customers": "Customers", + "chargebacks": "Chargebacks" }, "homeLink": "LiteLLM home" }, @@ -20493,5 +20497,113 @@ "removeGroup": "Remove weights for {{group}}", "summary": "Deployment Weights:", "loadFailed": "Could not load this model group's deployments. Saved weights are kept." + }, + "customerPortal": { + "switcher": { + "label": "Customer Portal", + "description": "End customers, budgets and chargebacks" + }, + "status": { + "active": "Active", + "overBudget": "Over budget", + "blocked": "Blocked" + }, + "common": { + "noBudget": "No budget", + "cancel": "Cancel", + "save": "Save", + "block": "Block", + "blockHint": "Block and Unblock change the customer record. The proxy only rejects a blocked customer's requests when the blocked_user_check callback is enabled in litellm_settings.callbacks.", + "unblock": "Unblock", + "edit": "Edit", + "total": "Total", + "loading": "Loading..." + }, + "toasts": { + "blocked": "Customer blocked", + "unblocked": "Customer unblocked", + "created": "Created {{name}}", + "updated": "Customer updated", + "createFailed": "Failed to create customer: {{error}}", + "updateFailed": "Failed to update customer: {{error}}" + }, + "list": { + "title": "Customers", + "subtitle": "End users of your product, identified by customer id on every request.", + "create": "Create Customer", + "searchPlaceholder": "Search by customer id or alias", + "columns": { + "customer": "Customer", + "spend": "Spend this month", + "budget": "Budget", + "keys": "Keys", + "status": "Status" + }, + "loading": "Loading customers...", + "empty": "No customers yet. Create one to start tracking spend.", + "budgetUsage": "{{spend}} / {{max}}" + }, + "detail": { + "spendThisMonth": "Spend this month", + "ofBudget": "of {{max}} budget", + "noBudget": "no budget", + "tpm": "{{val}} TPM", + "rpm": "{{val}} RPM", + "keysTitle": "Keys", + "keysSubtitle": "Keys that sent requests as this customer this month.", + "keysEmpty": "No traffic from this customer this month.", + "columns": { + "key": "Key", + "requests": "Requests", + "tokens": "Tokens", + "spend": "Spend" + }, + "unnamedKey": "Unnamed key", + "back": "Back to customers" + }, + "createDialog": { + "title": "Create customer", + "description": "A customer is an end user of your product. Pass its id as `user` on every request to track spend and apply its limits.", + "customerId": "Customer id", + "customerIdPlaceholder": "e.g. acme-corp", + "alias": "Alias", + "aliasPlaceholder": "Optional display name", + "maxBudget": "Monthly budget (USD)", + "maxBudgetPlaceholder": "Leave empty for no budget", + "tpm": "TPM limit", + "rpm": "RPM limit", + "noLimit": "No limit", + "limitsHelp": "Budget and limits apply to this customer across every key it uses.", + "submit": "Create" + }, + "editDialog": { + "title": "Edit customer", + "alias": "Alias", + "maxBudget": "Max budget (USD)", + "maxBudgetPlaceholder": "Leave empty for no budget", + "sharedBudgetWarning": "This customer is on budget {{budgetId}}. Changing the budget here changes it for every customer on that budget." + }, + "chargebacks": { + "title": "Chargebacks", + "subtitle": "Spend per customer for a month, from the daily spend rollup.", + "groupBy": { + "customer": "Group by customer", + "key": "Customer and key", + "model": "Customer and model" + }, + "dimension": { + "key": "Key", + "model": "Model" + }, + "exportCsv": "Export CSV", + "columns": { + "customer": "Customer", + "requests": "Requests", + "inputTokens": "Input tokens", + "outputTokens": "Output tokens", + "spend": "Spend" + }, + "empty": "No customer traffic in {{month}}." + } } } From c1abd490b26758df7ea8b426d2a1ee4efa55c264 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 10 Nov 2026 10:08:21 -0800 Subject: [PATCH 04/11] feat(customers): reject unknown fields on /customer/new and /customer/update instead of dropping them (#50965) * feat(customers): reject unknown fields on customer create and update instead of dropping them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(customers): derive the budget hint from the budget request fields and fix the create example Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 25 +++++++++ .../customer_endpoints.py | 3 +- .../test_customer_unknown_fields.py | 55 +++++++++++++++++++ .../test_customer_budget.py | 3 - tests/unit/proxy/test__types.py | 50 +++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 +- 6 files changed, 133 insertions(+), 9 deletions(-) create mode 100644 tests/integration/management/test_customer_unknown_fields.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index abf61a909b2..228e80da4bd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1860,6 +1860,21 @@ class CustomerBase(LiteLLMPydanticObjectBase): blocked: bool = False +def reject_unknown_customer_fields(request: type[BaseModel], values: object) -> object: + if not isinstance(values, dict): + return values + unknown: Final = tuple(sorted(str(key) for key in values if key not in request.model_fields)) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # raw request body + if not unknown: + return values + on_budget: Final = tuple(field for field in unknown if field in BudgetNewRequest.model_fields) + hint: Final = ( + f" {', '.join(on_budget)}: these live on the customer's budget object, update it with /budget/update." + if on_budget + else "" + ) + raise ValueError(f"Unknown fields: {', '.join(unknown)}.{hint}") + + class NewCustomerRequest(BudgetNewRequest): """ Create a new customer, allocate a budget to them @@ -1877,6 +1892,11 @@ class NewCustomerRequest(BudgetNewRequest): models: list[str] | None = None object_permission: LiteLLM_ObjectPermissionBase | None = None + @model_validator(mode="before") + @classmethod + def reject_unknown_fields(cls, values: object) -> object: + return reject_unknown_customer_fields(cls, values) + @model_validator(mode="before") @classmethod def check_user_info(cls, values): @@ -1904,6 +1924,11 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): models: list[str] | None = None object_permission: LiteLLM_ObjectPermissionBase | None = None + @model_validator(mode="before") + @classmethod + def reject_unknown_fields(cls, values: object) -> object: + return reject_unknown_customer_fields(cls, values) + class DeleteCustomerRequest(LiteLLMPydanticObjectBase): """ diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index aec5e6840a5..bd202b85f73 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -353,7 +353,6 @@ async def new_end_user( - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. - models: Optional[list[str]] - Restrict this customer's access to the listed models. - - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) @@ -387,7 +386,7 @@ async def new_end_user( --header 'Content-Type: application/json' \ --data '{ "user_id" : "ishaan-jaff-3", - "allowed_region": "eu", + "allowed_model_region": "eu", "budget_id": "free_tier", "models": ["gpt-4o-mini"], "default_model": "azure/gpt-3.5-turbo-eu" diff --git a/tests/integration/management/test_customer_unknown_fields.py b/tests/integration/management/test_customer_unknown_fields.py new file mode 100644 index 00000000000..235d5745791 --- /dev/null +++ b/tests/integration/management/test_customer_unknown_fields.py @@ -0,0 +1,55 @@ +import uuid +from typing import Final + +import httpx +import pytest + +from tests.integration._support.client import Gateway, Scenario + +pytestmark: Final = pytest.mark.shares_worker_state + + +def _customer(scenario: Scenario, **fields: object) -> str: + customer: Final = f"integration-customer-{uuid.uuid4().hex}" + scenario.gateway.post("/customer/new", {"user_id": customer, **fields}) # pyright: ignore[reportArgumentType] # JsonValue kwargs + scenario.cleanups.callback(scenario.gateway.post, "/customer/delete", {"user_ids": [customer]}) + return customer + + +def _post(gateway: Gateway, route: str, body: dict[str, object]) -> httpx.Response: + return gateway.request("POST", route, body) # pyright: ignore[reportArgumentType] # JsonValue body + + +def _info(gateway: Gateway, customer: str) -> dict[str, object]: + return gateway.get("/customer/info", {"end_user_id": customer}) + + +@pytest.mark.parametrize( + ("body", "unknown_field"), + [ + pytest.param({"tpm_limit": 5}, "tpm_limit", id="budget_field_on_update"), + pytest.param({"quota": {"opus": 1}}, "quota", id="made_up_field"), + ], +) +def test_customer_update_rejects_unknown_fields_and_leaves_the_customer_unchanged( + gateway: Gateway, body: dict[str, object], unknown_field: str +) -> None: + with gateway.scenario() as scenario: + customer: Final = _customer(scenario, alias="before") + rejected: Final = _post(gateway, "/customer/update", {"user_id": customer, **body}) + assert rejected.status_code == 422, rejected.text + assert unknown_field in rejected.text, rejected.text + assert _info(gateway, customer)["alias"] == "before" + + accepted: Final = _post(gateway, "/customer/update", {"user_id": customer, "alias": "after"}) + assert accepted.status_code == 200, accepted.text + assert _info(gateway, customer)["alias"] == "after" + + +def test_customer_new_rejects_a_misspelled_field_and_creates_nothing(gateway: Gateway) -> None: + customer: Final = f"integration-customer-{uuid.uuid4().hex}" + rejected: Final = _post(gateway, "/customer/new", {"user_id": customer, "tpm_limti": 5}) + assert rejected.status_code == 422, rejected.text + assert "tpm_limti" in rejected.text, rejected.text + missing: Final = gateway.request("GET", "/customer/info", params={"end_user_id": customer}) + assert missing.status_code == 404, missing.text diff --git a/tests/unit/proxy/management_endpoints/test_customer_budget.py b/tests/unit/proxy/management_endpoints/test_customer_budget.py index 0beca0c15e8..ac172ac2ffb 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_budget.py +++ b/tests/unit/proxy/management_endpoints/test_customer_budget.py @@ -144,7 +144,6 @@ async def test_update_customer_creates_budget_with_proper_relations( update_request = UpdateCustomerRequest( user_id="test-user", max_budget=100.0, # This triggers budget creation - rpm_limit=200, # Use valid budget field ) # Act @@ -267,7 +266,6 @@ async def test_update_customer_budget_creation_with_fallback_admin( update_request = UpdateCustomerRequest( user_id="test-user", max_budget=150.0, - tpm_limit=1000, # Add another budget field ) # Act @@ -325,7 +323,6 @@ async def test_update_customer_with_budget_id_and_creation_fields( user_id="test-user", budget_id="existing-budget-link", # For linking to existing budget max_budget=300.0, # This should trigger new budget creation - rpm_limit=500, # Use valid budget field ) # Act diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 40e66f1da65..1c97e9054ae 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -26,6 +26,7 @@ from litellm.proxy._types import ( OrganizationMemberUpdateRequest, RegenerateKeyRequest, ResetSpendRequest, + UpdateCustomerRequest, UpdateKeyRequest, UpdateMCPServerRequest, UpdateTeamRequest, @@ -634,3 +635,52 @@ def test_batch_file_caps_accept_only_positive_limits(setting): with pytest.raises(ValidationError): ConfigGeneralSettings.model_validate({setting: invalid}) assert getattr(ConfigGeneralSettings.model_validate({setting: 3}), setting) == 3 + + +@pytest.mark.parametrize( + ("request_type", "body", "message"), + [ + pytest.param( + UpdateCustomerRequest, + {"user_id": "acme", "tpm_limit": 5, "quota": {"opus": 1}}, + "Unknown fields: quota, tpm_limit. tpm_limit: these live on the customer's budget object, " + "update it with /budget/update.", + id="update_with_a_budget_field_and_a_made_up_one", + ), + pytest.param( + UpdateCustomerRequest, + {"user_id": "acme", "alias": "Acme", "quota": 1}, + "Unknown fields: quota.", + id="update_with_a_made_up_field_has_no_budget_hint", + ), + pytest.param( + NewCustomerRequest, + {"user_id": "acme", "tpm_limti": 5}, + "Unknown fields: tpm_limti.", + id="create_with_a_misspelled_budget_field", + ), + ], +) +def test_a_customer_request_with_unknown_fields_is_refused_naming_them( + request_type: type[BaseModel], body: dict[str, object], message: str +) -> None: + with pytest.raises(ValidationError) as exc: + request_type.model_validate(body) + assert [error["msg"] for error in exc.value.errors()] == [f"Value error, {message}"] + + +@pytest.mark.parametrize( + ("request_type", "body"), + [ + pytest.param(UpdateCustomerRequest, {"user_id": "acme", "alias": "Acme", "blocked": True}, id="update"), + pytest.param( + NewCustomerRequest, + {"user_id": "acme", "tpm_limit": 5, "model_tpm_limit": {"opus": 60}}, + id="create_with_inline_budget_fields", + ), + ], +) +def test_a_customer_request_with_only_known_fields_is_accepted( + request_type: type[BaseModel], body: dict[str, object] +) -> None: + assert request_type.model_validate(body).model_dump(exclude_unset=True) == body diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b229461e981..e1b0104e089 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4549,7 +4549,6 @@ export interface paths { * - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. * - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. * - models: Optional[list[str]] - Restrict this customer's access to the listed models. - * - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) @@ -4580,7 +4579,7 @@ export interface paths { * ``` * curl --location 'http://0.0.0.0:4000/customer/new' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{ * "user_id" : "ishaan-jaff-3", - * "allowed_region": "eu", + * "allowed_model_region": "eu", * "budget_id": "free_tier", * "models": ["gpt-4o-mini"], * "default_model": "azure/gpt-3.5-turbo-eu" @@ -5229,7 +5228,6 @@ export interface paths { * - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. * - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. * - models: Optional[list[str]] - Restrict this customer's access to the listed models. - * - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) @@ -5260,7 +5258,7 @@ export interface paths { * ``` * curl --location 'http://0.0.0.0:4000/customer/new' --header "Authorization: Bearer $LITELLM_MASTER_KEY" --header 'Content-Type: application/json' --data '{ * "user_id" : "ishaan-jaff-3", - * "allowed_region": "eu", + * "allowed_model_region": "eu", * "budget_id": "free_tier", * "models": ["gpt-4o-mini"], * "default_model": "azure/gpt-3.5-turbo-eu" From 90f37e113ca80b593e5c25ef2550dac9ce48af0a Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Tue, 10 Nov 2026 10:09:19 -0800 Subject: [PATCH 05/11] fix(proxy): send a prompt-only call_endpoint token count to every provider count API (#50783) * fix(proxy): count a prompt-only call_endpoint request with the provider on every counter POST /utils/token_counter?call_endpoint=true with only a prompt sent the prompt to the provider count API on Gemini and Vertex only. Anthropic, OpenAI, Bedrock and Azure AI (Claude) deployments answered a local estimate instead. Every counter now gets the prompt as one user message, so a provider error surfaces the same as the messages form, and the framed stand-in after a failed provider count frames the prompt as one user turn, matching the messages form * fix(proxy): count a call_endpoint prompt like a user message on providers without a count API * fix(proxy): count a call_endpoint prompt as a turn only on chat deployments A prompt-only /utils/token_counter request with call_endpoint=true now goes to the count API as one user message only when the deployment serves chat (mode chat, responses, or unknown). Embedding rows and models with no deployment keep the raw local count. The Anthropic, Azure AI, and OpenAI count calls now honor the deployment's timeout, and a count error with an empty body names its status code. * fix(proxy): fall back to router_settings.timeout for count calls and answer 408 on a count timeout Count calls to Anthropic, Azure AI, and OpenAI now pick their timeout the way a chat call to the same deployment does: the deployment's timeout, else router_settings.timeout, else request_timeout. A count call that times out answers 408 instead of a 500 processing error. * fix(proxy): count a prompt as a turn only where a count API exists and give count calls chat's timeout order * fix(proxy): read the count call timeout through the router settings snapshot The token counter now asks llm_router.settings().timeout_for() for the chat timeout order instead of a new Router member, so the router surface does not grow * fix(router): answer timeout_for on every router settings snapshot The proxy reads the count call timeout through llm_router.settings().timeout_for(), but only the Rust core's snapshot had that method, so a plain RouterSettings (the unit tests' router stub) raised AttributeError. timeout_for now lives on RouterSettings and reads the core's _core_deployment_call_timeout, like the snapshot's other members --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/anthropic/count_tokens/handler.py | 2 + .../anthropic/count_tokens/token_counter.py | 3 +- .../anthropic/count_tokens/handler.py | 2 + .../anthropic/count_tokens/token_counter.py | 3 +- litellm/llms/base_llm/base_utils.py | 15 +- .../openai/responses/count_tokens/handler.py | 2 + .../responses/count_tokens/token_counter.py | 13 +- litellm/proxy/proxy_server.py | 40 ++- litellm/router_api.py | 3 + litellm/router_rust/core.py | 5 +- ...anthropic_count_tokens_system_lift_wire.py | 16 + .../test_bedrock_mantle_count_tokens_wire.py | 26 ++ tests/unit/proxy/router_entry_point_stubs.py | 1 + tests/unit/proxy/test_proxy_token_counter.py | 332 ++++++++++++++++-- 14 files changed, 427 insertions(+), 36 deletions(-) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 7a7d1d776d3..46c90295198 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -125,6 +125,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise AnthropicError(status_code=e.status_code, message=str(e)) from e except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise AnthropicError( diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 99b1cd7f241..755ccd129d8 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -18,7 +18,7 @@ from litellm.llms.anthropic.count_tokens.transformation import ( ChatCountTranslationError, anthropic_count_fields, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests @@ -134,6 +134,7 @@ class AnthropicTokenCounter(BaseTokenCounter): optional_params=count_options, context_management=context_management.spec if context_management is not None else None, drop_params=drops_unsupported_params(litellm_params.get("drop_params")), + timeout=deployment_timeout(litellm_params), ) if result is not None: diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 5fcb84593d2..443a3a5edad 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -147,6 +147,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise AnthropicError(status_code=e.status_code, message=str(e)) from e except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise AnthropicError( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 8657fcddab4..5031b69f83b 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -20,7 +20,7 @@ from litellm.llms.anthropic.count_tokens.transformation import ( from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests @@ -137,6 +137,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): system=system, optional_params=count_options, context_management=None if context_management is None else context_management.spec, + timeout=deployment_timeout(litellm_params), ) if result is not None: diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 45f623730a8..8b1ccda3a61 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -15,8 +15,9 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Any, Final, TypeAlias +import httpx from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, JsonValue, PositiveFloat, TypeAdapter, ValidationError from typing_extensions import TypeIs # noqa: TID251 # narrows a message to its system member by the role literal from litellm._logging import verbose_logger @@ -38,6 +39,18 @@ from litellm.types.utils import ( TokenCountResponse, ) +_SECONDS: Final[TypeAdapter[float]] = TypeAdapter(PositiveFloat) + + +def deployment_timeout(litellm_params: Mapping[str, object]) -> float | httpx.Timeout | None: + timeout: Final = litellm_params.get("timeout") + if timeout is None or isinstance(timeout, httpx.Timeout): + return timeout + try: + return _SECONDS.validate_python(timeout) + except ValidationError: + return None + class BaseTokenCounter(ABC): @abstractmethod diff --git a/litellm/llms/openai/responses/count_tokens/handler.py b/litellm/llms/openai/responses/count_tokens/handler.py index 2cf48a75524..ce9ed6c0940 100644 --- a/litellm/llms/openai/responses/count_tokens/handler.py +++ b/litellm/llms/openai/responses/count_tokens/handler.py @@ -111,6 +111,8 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig): status_code=e.response.status_code, message=e.response.text, ) + except litellm.Timeout as e: + raise OpenAIError(status_code=e.status_code, message=str(e)) from e except (httpx.RequestError, json.JSONDecodeError, ValueError) as e: verbose_logger.error("Error in CountTokens handler: %s", e) raise OpenAIError( diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index c82a2982a54..cd60bb1327f 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Any, Final +import httpx import openai from pydantic import JsonValue from typing_extensions import assert_never @@ -21,7 +22,7 @@ from litellm.llms.anthropic.messages.responses_adapters.handler import count_vie from litellm.llms.anthropic.messages.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) -from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.llms.base_llm.base_utils import BaseTokenCounter, deployment_timeout from litellm.llms.openai.common_utils import OpenAIError, decoded_json_object, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, @@ -61,6 +62,7 @@ def _workload_identity_config( class _CountCredential: api_base: str | None credential: str | OpenAIWorkloadIdentityConfig + timeout: float | httpx.Timeout | None def _count_credential(deployment: Mapping[str, object] | None, model_label: str) -> _CountCredential | OpenAIError: @@ -82,7 +84,9 @@ def _count_credential(deployment: Mapping[str, object] | None, model_label: str) "workload identity on the deployment, or OPENAI_API_KEY in the proxy environment." ), ) - return _CountCredential(api_base=credentials.api_base, credential=credential) + return _CountCredential( + api_base=credentials.api_base, credential=credential, timeout=deployment_timeout(litellm_params) + ) def _count_error(model_to_use: str, request_model: str, message: str, status_code: int) -> TokenCountResponse: @@ -245,6 +249,7 @@ class OpenAITokenCounter(BaseTokenCounter): tools=counted_tools, instructions=instructions, text=text, + timeout=resolved.timeout, ), model_to_use, request_model, @@ -264,7 +269,9 @@ class OpenAITokenCounter(BaseTokenCounter): body: Final = native_count_request(request, model_to_use, policy) return await self._counted( resolved, - lambda api_key, api_base: openai_count_tokens_handler.post_count_request(body, api_key, api_base), + lambda api_key, api_base: openai_count_tokens_handler.post_count_request( + body, api_key, api_base, resolved.timeout + ), model_to_use, request_model, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e6447754436..3fddca6f4d8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15325,7 +15325,12 @@ def _provider_count_outcome( if result is not None and result.error is True: return ProviderCountFailure( status_code=result.status_code or 500, - message=result.error_message or "Token counting failed", + message=result.error_message + or ( + "Token counting failed" + if result.status_code is None + else f"The provider's count API answered {result.status_code} with no error message" + ), rejected_before_provider_call=result.rejected_before_provider_call, ) return result @@ -15794,6 +15799,14 @@ def _token_count_settings_model(deployment: _TokenCountDeployment, request_model _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} ) +_CHAT_MODES: Final = frozenset({"chat", "responses"}) + + +def _serves_chat(deployment: _TokenCountDeployment, cost_map_row: Mapping[str, object]) -> bool: + mode: Final = _non_empty_str(deployment.get("model_info", {}).get("mode")) or _non_empty_str( + cost_map_row.get("mode") + ) + return mode is None or mode in _CHAT_MODES _ESTIMABLE_GOOGLE_PART_KEYS: Final = frozenset({"text", "functionCall", "functionResponse", "thoughtSignature"}) @@ -16104,6 +16117,12 @@ async def _routed_count_deployment( return None +def _with_chat_timeout(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + if llm_router is None: + return litellm_params + return {**litellm_params, "timeout": llm_router.settings().timeout_for(litellm_params)} + + async def _count_request_tokens( request: TokenCountRequest, routed: Mapping[str, object] | None, @@ -16128,8 +16147,8 @@ async def _count_request_tokens( litellm_model_name = deployment.get("litellm_params", {}).get("model") deployment = { # rebind-ok: merge the named credential into a copy **deployment, - "litellm_params": with_named_credential( - _LITELLM_PARAMS_COPY.validate_python(deployment.get("litellm_params", {})) + "litellm_params": _with_chat_timeout( + with_named_credential(_LITELLM_PARAMS_COPY.validate_python(deployment.get("litellm_params", {}))) ), } # remove the custom_llm_provider_prefix in the litellm_model_name @@ -16142,9 +16161,10 @@ async def _count_request_tokens( deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python( deployment or _model_group_deployment(llm_router, request.model) or {} ) + cost_map_row: Final = await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)) settings: Final = await asyncify(resolve_token_count_settings)( configured=deployment_view.get("model_info", {}), - default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)), + default=cost_map_row, deployment_region=partial(_deployment_aws_region, deployment_view), ) @@ -16169,6 +16189,15 @@ async def _count_request_tokens( code=400, ) prompt_only: Final = bool(prompt) and conversation_empty + prompt_as_turn: Final = ( + prompt_only + and provider_counter is not None + and provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider) + and _serves_chat(deployment_view, cost_map_row) + ) + prompt_as_message: Final = prompt_as_turn or ( + prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES + ) request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment) local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment) malformed_messages: Final = ( @@ -16205,7 +16234,6 @@ async def _count_request_tokens( ) inference_system: Final = system if edited_request is None else edited_request.system contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents)) - prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment) routed_request: Final = ( None if responses_request is None else with_provider_file_ids(responses_request, routed_deployment) @@ -16306,7 +16334,7 @@ async def _count_request_tokens( _with_get(model_info).get("custom_tokenizer", None), ) system_message: Final = _system_message(local_system) - prompt_turn: Final = prompt_only and not (system_message is None and not tools) + prompt_turn: Final = prompt_as_turn or (prompt_only and not (system_message is None and not tools)) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes Sequence[AllMessageValues] | None, [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else local_messages), diff --git a/litellm/router_api.py b/litellm/router_api.py index d07b135e9a5..c85a764151f 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1308,6 +1308,9 @@ class RouterSettings: """The guardrail configured under this name (was ``get_available_guardrail``).""" return self._owner("_core_get_available_guardrail")._core_get_available_guardrail(name) + def timeout_for(self, litellm_params: Mapping[str, object]) -> object: + return self._owner("_core_deployment_call_timeout")._core_deployment_call_timeout(litellm_params) + class RouterHealth: """Cooldowns, blocked deployments, health checks and usage (``health()``).""" diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index b7e67fffd57..682aa557007 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -113,7 +113,7 @@ from litellm.router_rust.invoke.endpoints import ( stamp_error, unrouted_fields, ) -from litellm.router_rust.invoke.shaping import Shaping +from litellm.router_rust.invoke.shaping import Shaping, non_stream_timeout from litellm.router_rust.invoke.slots import enter_deployment_slot, max_parallel_requests_limit from litellm.router_rust.invoke.streaming import ( StreamAttempt, @@ -1592,6 +1592,9 @@ class RustCore( def _shaping(self) -> Shaping: return self._endpoints.shaping(self._resolved) + def _core_deployment_call_timeout(self, litellm_params: Mapping[str, object]) -> object: + return non_stream_timeout(self._shaping, MappingProxyType({}), litellm_params) + # -- building -------------------------------------------------------------------------------- def _new_handle(self, model_list: Sequence[object]) -> "RouterCoreHandle": diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index 380881a6796..86a8cb3df4f 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -332,6 +332,22 @@ def test_utils_token_counter_local_mode_never_calls_the_peer_for_a_leading_syste assert peer.drain() == () +def test_utils_token_counter_call_endpoint_counts_a_prompt_as_one_user_message( + deployment: _Deployment, gateway: Gateway +) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/utils/token_counter", + {"model": deployment.model, "prompt": USER_TEXT}, + params={"call_endpoint": "true"}, + ) + payload: Final = _payload(response) + assert (payload["total_tokens"], payload["count_source"]) == (_COUNT, "provider"), response.text + assert payload["tokenizer_type"] == deployment.provider.tokenizer, response.text + assert _bodies(peer, deployment.provider) == ({"model": _MODEL, "messages": [USER]},) + + def test_responses_input_tokens_lifts_instructions(deployment: _Deployment, gateway: Gateway) -> None: with _listening(deployment) as peer: response: Final = gateway.request( diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index 2a0508402ae..07394c03f0c 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -550,6 +550,32 @@ def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: Owned assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) +def test_utils_token_counter_prompt_answers_the_count_model_404_like_the_messages_form( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _metadata_deployment(scenario, runtime.url) + from_prompt: Final = gateway.request( + "POST", + "/utils/token_counter", + {"model": model, "prompt": "Count this message"}, + params={"call_endpoint": "true"}, + ) + from_messages: Final = gateway.request( + "POST", "/utils/token_counter", _bare(model), params={"call_endpoint": "true"} + ) + assert from_prompt.status_code == from_messages.status_code == 404, from_prompt.text + assert _REJECTION in from_prompt.text, from_prompt.text + assert _JSON_OBJECT.validate_json(from_prompt.content) == _JSON_OBJECT.validate_json(from_messages.content) + assert runtime.drain() == () + assert _mantle_bodies(mantle) == (_MANTLE_BARE, _MANTLE_BARE) + + def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with ( diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 306740489d7..10176accf9f 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -184,6 +184,7 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: _core_get_settings=lambda: router.get_settings(), _core_get_available_guardrail=lambda name: router.get_available_guardrail(guardrail_name=name), _core_get_allowed_fails_from_policy=lambda exception: router.get_allowed_fails_from_policy(exception=exception), + _core_deployment_call_timeout=lambda litellm_params: router.deployment_call_timeout(litellm_params), ) router.settings.side_effect = lambda: RouterSettings( **{name: getattr(router, name) for name in SETTING_FIELDS}, _router=settings_backend diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 694b257b1f9..34868606012 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -8,10 +8,12 @@ import hashlib import itertools import json import logging +import socket import threading import traceback import uuid -from collections.abc import Callable, Coroutine, Mapping +from collections.abc import Callable, Coroutine, Iterator, Mapping +from contextlib import contextmanager from dataclasses import replace from functools import partial from pathlib import Path @@ -44,8 +46,11 @@ import litellm from litellm import Router from litellm._logging import JsonFormatter, verbose_proxy_logger from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler +from litellm.llms.azure_ai.anthropic.count_tokens import handler as azure_ai_count_tokens_http from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_http +from litellm.llms.bedrock.count_tokens import mantle_handler as bedrock_mantle_count_tokens_http from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.kimi_k3 import ( @@ -4833,36 +4838,317 @@ async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens +_PROVIDER_COUNTED_DEPLOYMENTS: Final = ( + pytest.param( + {"model": "anthropic/claude-sonnet-5-5", "api_key": "test-key", "api_base": "https://count-tokens.test"}, + id="anthropic", + ), + pytest.param( + {"model": "azure_ai/claude-sonnet-5-5", "api_key": "test-key", "api_base": "https://count-tokens.test"}, + id="azure_ai", + ), + pytest.param( + {"model": "openai/gpt-5.6", "api_key": "test-key", "api_base": "https://count-tokens.test/v1"}, id="openai" + ), + pytest.param( + { + "model": "bedrock/us.anthropic.claude-sonnet-5-5", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + id="bedrock", + ), +) + + +@contextmanager +def _count_apis_answering(peer: Callable[[httpx.Request], httpx.Response]) -> Iterator[None]: + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + with ( + patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client), + patch.object(azure_ai_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(bedrock_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(bedrock_mantle_count_tokens_http, "get_async_httpx_client", return_value=client), + patch.object(openai_count_tokens_http, "get_async_httpx_client", return_value=client), + ): + yield + + +def _route_to_one_deployment(litellm_params: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "provider-counted", "litellm_params": litellm_params}]), + ) + + +_FOX_AS_PROMPT: Final = TokenCountRequest(model="provider-counted", prompt=_FOX) +_FOX_AS_MESSAGE: Final = TokenCountRequest(model="provider-counted", messages=[{"role": "user", "content": _FOX}]) + + @pytest.mark.asyncio -async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages( +@pytest.mark.parametrize("litellm_params", _PROVIDER_COUNTED_DEPLOYMENTS) +async def test_call_endpoint_sends_a_prompt_to_the_provider_as_one_user_message( + litellm_params: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + _route_to_one_deployment(litellm_params, monkeypatch) + sent_bodies: Final[list[JsonValue]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, json={"object": "response.input_tokens", "input_tokens": 41, "inputTokens": 41}, request=request + ) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + + assert (from_prompt.total_tokens, from_prompt.error) == (41, False) + assert from_prompt.count_source != "local_estimate" + assert from_prompt == from_message + assert len(sent_bodies) == 2 and sent_bodies[0] == sent_bodies[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [404, 503]) +@pytest.mark.parametrize("litellm_params", _PROVIDER_COUNTED_DEPLOYMENTS) +async def test_call_endpoint_answers_a_provider_error_on_a_prompt_like_on_the_same_text_sent_as_a_user_message( + litellm_params: dict[str, JsonValue], status: int, monkeypatch: pytest.MonkeyPatch +) -> None: + _route_to_one_deployment(litellm_params, monkeypatch) + + def count_api(request: httpx.Request) -> httpx.Response: + return httpx.Response(status, json={"error": {"message": f"scripted {status}"}}, request=request) + + with _count_apis_answering(count_api): + with pytest.raises(ProxyException) as from_prompt: + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + with pytest.raises(ProxyException) as from_message: + await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + + assert from_prompt.value.code == str(status) + assert f"scripted {status}" in from_prompt.value.message + assert (from_prompt.value.code, from_prompt.value.type, from_prompt.value.message) == ( + from_message.value.code, + from_message.value.type, + from_message.value.message, + ) + + +@pytest.mark.asyncio +async def test_framed_stand_in_counts_a_prompt_like_the_same_text_sent_as_a_user_message( monkeypatch: pytest.MonkeyPatch, ) -> None: - router: Final = Router( - model_list=[ - { - "model_name": "anthropic-counted", - "litellm_params": { - "model": "anthropic/claude-sonnet-5-5", - "api_key": "test-key", - "api_base": "https://count-tokens.test", - }, - } - ] - ) - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_unavailable_peer)) - request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.") + _framed_kimi(monkeypatch) - with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client): - with_endpoint: Final = await token_counter(request=request, call_endpoint=True) + from_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX), call_endpoint=True) + from_message: Final = await token_counter( + request=TokenCountRequest(model="kimi", messages=[{"role": "user", "content": _FOX}]), call_endpoint=True + ) + local_prompt: Final = await token_counter(request=TokenCountRequest(model="kimi", prompt=_FOX)) + + assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False) + assert from_prompt.total_tokens == from_message.total_tokens > local_prompt.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", ["fireworks_ai/accounts/fireworks/models/glm-5", "deepseek/deepseek-v4-flash"], ids=["plain", "framed"] +) +async def test_call_endpoint_counts_a_prompt_on_a_provider_without_a_count_api_as_raw_text( + model: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _route_to_one_deployment({"model": model, "api_key": "test-key"}, monkeypatch) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + from_message: Final = await token_counter(request=_FOX_AS_MESSAGE, call_endpoint=True) + local_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False) + + assert sent == [] + assert (from_prompt.count_source, from_prompt.error) == ("local_estimate", False) + assert from_prompt.total_tokens == local_prompt.total_tokens < from_message.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment", + [ + pytest.param( + {"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "test-key"}}, id="cost map mode" + ), + pytest.param( + { + "litellm_params": {"model": "openai/in-house-embedder", "api_key": "test-key"}, + "model_info": {"mode": "embedding"}, + }, + id="configured mode", + ), + ], +) +async def test_call_endpoint_counts_a_prompt_on_a_non_chat_deployment_as_raw_text_without_the_provider( + deployment: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, "llm_router", Router(model_list=[{"model_name": "provider-counted", **deployment}]) + ) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_endpoint: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + local: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=False) + + assert sent == [] + assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["chat", "responses"]) +async def test_call_endpoint_sends_a_prompt_on_a_chat_mode_deployment_to_the_provider( + mode: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "provider-counted", + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "test-key", + "api_base": "https://count-tokens.test/v1", + }, + "model_info": {"mode": mode}, + } + ] + ), + ) + sent: Final[list[httpx.Request]] = [] # mutable-ok: collects what the mock peer receives + + def count_api(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"object": "response.input_tokens", "input_tokens": 41}, request=request) + + with _count_apis_answering(count_api): + from_prompt: Final = await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert (from_prompt.total_tokens, from_prompt.error) == (41, False) + assert len(sent) == 1 + + +@pytest.mark.asyncio +async def test_call_endpoint_counts_a_prompt_for_a_model_without_a_deployment_as_raw_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router(model_list=[{"model_name": "other", "litellm_params": {"model": "openai/gpt-5.6"}}]), + ) + request: Final = TokenCountRequest(model="deepseek/deepseek-v4-flash", prompt=_FOX) + + from_endpoint: Final = await token_counter(request=request, call_endpoint=True) local: Final = await token_counter(request=request, call_endpoint=False) - assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == ( - False, - local.total_tokens, - local.tokenizer_type, + assert (from_endpoint.total_tokens, from_endpoint.error) == (local.total_tokens, False) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "path"), + [("anthropic/claude-sonnet-5-5", ""), ("azure_ai/claude-sonnet-5-5", ""), ("openai/gpt-5.6", "/v1")], + ids=["anthropic", "azure_ai", "openai"], +) +@pytest.mark.parametrize( + ("row_timeout", "router_timeout"), + [({"timeout": 0.5}, None), ({"request_timeout": 0.5}, 30.0), ({}, 0.5)], + ids=["deployment timeout", "deployment request_timeout", "router_settings timeout"], +) +async def test_call_endpoint_answers_a_timeout_when_a_stalled_count_api_outlasts_the_chat_timeout( + model: str, + path: str, + row_timeout: dict[str, JsonValue], + router_timeout: float | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with socket.create_server(("127.0.0.1", 0)) as never_answers: + port: Final = never_answers.getsockname()[1] + litellm_params: Final[dict[str, JsonValue]] = { + "model": model, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{port}{path}", + } + deployment: Final[dict[str, JsonValue]] = { + "model_name": "provider-counted", + "litellm_params": {**litellm_params, **row_timeout}, + } + monkeypatch.setattr( + litellm.proxy.proxy_server, "llm_router", Router(model_list=[deployment], timeout=router_timeout) + ) + + with pytest.raises(ProxyException) as error: + async with asyncio.timeout(10): + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert error.value.code == "408" + + +@pytest.mark.asyncio +async def test_responses_count_answers_a_timeout_when_a_stalled_count_api_outlasts_the_deployment_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with socket.create_server(("127.0.0.1", 0)) as never_answers: + port: Final = never_answers.getsockname()[1] + _route_to_one_deployment( + { + "model": "openai/gpt-5.6", + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{port}/v1", + "timeout": 0.5, + }, + monkeypatch, + ) + + with pytest.raises(ProxyException) as error: + async with asyncio.timeout(10): + await count_request_tokens( + request=_FOX_AS_MESSAGE, + call_endpoint=True, + responses_request={"model": "provider-counted", "input": _FOX}, + ) + + assert error.value.code == "408" + + +@pytest.mark.asyncio +async def test_call_endpoint_names_the_status_of_a_count_api_error_with_an_empty_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _route_to_one_deployment( + {"model": "openai/gpt-5.6", "api_key": "test-key", "api_base": "https://count-tokens.test/v1"}, monkeypatch ) + with _count_apis_answering(lambda request: httpx.Response(404, request=request)): + with pytest.raises(ProxyException) as error: + await token_counter(request=_FOX_AS_PROMPT, call_endpoint=True) + + assert error.value.code == "404" + assert "404" in error.value.message + _PINNED_COUNTS: Final = MappingProxyType({"deployment-a.test": 11, "deployment-b.test": 22}) From 3539c0afc999c47eec8e07b027c746f49ed2a38d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 10 Nov 2026 10:11:16 -0800 Subject: [PATCH 06/11] test(integration): cap the OAuth metadata lookup in the mixed capabilities subscription case (#50867) The second server in this case is a plain MCP SDK server registered as oauth_delegate. The gateway's first catalog listing waits for that server's OAuth metadata lookup, which never resolves and runs for the full MCP_METADATA_TIMEOUT of 10 seconds. The MCP SDK client also gives server/discover 10 seconds, so the case passed or failed on about a tenth of a second of head start. Setting LITELLM_MCP_METADATA_TIMEOUT to 5 for this case keeps the lookup well inside the client limit. The case still asserts that subscription changes are delivered and tools are refreshed across both servers. --- tests/integration/mcp/test_interactions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/mcp/test_interactions.py b/tests/integration/mcp/test_interactions.py index 5577cf16067..b77deb5236a 100644 --- a/tests/integration/mcp/test_interactions.py +++ b/tests/integration/mcp/test_interactions.py @@ -1227,6 +1227,7 @@ def test_modern_subscription_delivers_changes_between_calls_and_refreshes_tools( "PYTHONPATH": os.pathsep.join((str(source), str(source / "tests"))), "LITELLM_SALT_KEY": "subscription-pagination-key", **({"LITELLM_MCP_TOOL_LISTING_TIMEOUT": "2"} if access == "probe-timeout" else {}), + **({"LITELLM_MCP_METADATA_TIMEOUT": "5"} if access == "mixed-capabilities" else {}), "STORE_MODEL_IN_DB": "False", "DISABLE_SCHEMA_UPDATE": "true", "LITELLM_LOCAL_MODEL_COST_MAP": "True", From e6db6206d19de2c8d33aeaccdd00049fdb296603 Mon Sep 17 00:00:00 2001 From: shrey Date: Tue, 10 Nov 2026 10:23:50 -0800 Subject: [PATCH 07/11] fix(mcp): accept comma-separated vendor scope grants on the connect check (#50717) --- .../mcp_server/discoverable_endpoints.py | 12 +- litellm/proxy/mcp_server/oauth_utils.py | 31 ++++- .../mcp_server/test_discoverable_endpoints.py | 110 +++++++++++++++++- 3 files changed, 150 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/mcp_server/discoverable_endpoints.py b/litellm/proxy/mcp_server/discoverable_endpoints.py index 07a267d6caa..ecc40b6e98f 100644 --- a/litellm/proxy/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/mcp_server/discoverable_endpoints.py @@ -105,9 +105,11 @@ from litellm.proxy.mcp_server.oauth_utils import ( get_cimd_client_id, get_cimd_document_url, get_request_base_url, + grant_covers_scopes, needs_cimd_discovery, oauth_client_registration_matches, resolve_upstream_resource, + server_scope_names, validate_trusted_redirect_uri, well_known_root_suffix, ) @@ -2537,6 +2539,9 @@ async def _vendor_credential_state( get_user_oauth_credential, oauth_grant_state, ) + from litellm.proxy.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # circular import at module load + global_mcp_server_manager, + ) from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # circular import at module load if prisma_client is None: @@ -2548,7 +2553,12 @@ async def _vendor_credential_state( if oauth_grant_state(credential) == "absent": return "absent" granted_scopes: Final = credential.get("scopes") if credential else None - return "present" if set(requested_scopes).issubset(granted_scopes or ()) else "absent" + covered: Final = grant_covers_scopes( + tuple(granted_scopes or ()), + requested_scopes, + lambda: server_scope_names(global_mcp_server_manager.get_mcp_server_by_id(server_id)), + ) + return "present" if covered else "absent" @router.get("/authorize/flow") diff --git a/litellm/proxy/mcp_server/oauth_utils.py b/litellm/proxy/mcp_server/oauth_utils.py index 76343f77e32..3171c476ede 100644 --- a/litellm/proxy/mcp_server/oauth_utils.py +++ b/litellm/proxy/mcp_server/oauth_utils.py @@ -3,7 +3,7 @@ import os import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from ipaddress import ip_address from types import MappingProxyType @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Final, NoReturn from urllib.parse import ParseResult, parse_qsl, urlencode, urlparse, urlsplit, urlunparse, urlunsplit from fastapi import HTTPException, Request +from pydantic import TypeAdapter from starlette.types import Scope from litellm._logging import verbose_logger @@ -131,6 +132,34 @@ def build_upstream_authorize_url(authorization_url: str, params: Mapping[str, st return urlunparse(parsed._replace(query=urlencode(merged))) +def grant_covers_scopes( + granted: tuple[str, ...], requested: tuple[str, ...], server_scopes: Callable[[], frozenset[str]] +) -> bool: + """Whether a stored grant holds every requested scope. A grant stored as one comma-joined token, the form + Slack and GitHub reply with, counts as its parts only when every part is a requested or server scope.""" + if all(scope in granted for scope in requested): + return True + if len(granted) != 1 or "," not in granted[0]: + return False + parts: Final = frozenset(granted[0].split(",")) + known: Final = frozenset(requested) | server_scopes() + return parts <= known and all(scope in granted or scope in parts for scope in requested) + + +_PUBLISHED_SCOPES: Final = TypeAdapter[tuple[str, ...] | None](tuple[str, ...] | None) + + +def server_scope_names(server: "MCPServer | None") -> frozenset[str]: + if server is None: + return frozenset() + published: Final = ( + _PUBLISHED_SCOPES.validate_python(server.authorization_server_metadata.get("scopes_supported")) + if server.authorization_server_metadata is not None + else None + ) + return frozenset((*(server.scopes or ()), *(published or ()))) + + def proxy_base_url_from_env() -> str | None: """The gateway's fixed public URL from ``PROXY_BASE_URL``, ``None`` when unset or malformed.""" return _resolve_proxy_base_url_env() diff --git a/tests/unit/proxy/mcp_server/test_discoverable_endpoints.py b/tests/unit/proxy/mcp_server/test_discoverable_endpoints.py index c17f4b392cd..47d24a9626f 100644 --- a/tests/unit/proxy/mcp_server/test_discoverable_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_discoverable_endpoints.py @@ -16,7 +16,7 @@ from litellm.proxy._types import LiteLLM_MCPServerTable from litellm.types.mcp import MCPAuth if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterator, Mapping import httpx from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey @@ -15278,6 +15278,9 @@ async def test_managed_resource_metadata_excludes_refresh_permission(monkeypatch (["tools.read", "tools.write"], ("tools.write",), "present"), (None, ("tools.write",), "absent"), (["tools.read"], (), "present"), + (["channels:read,chat:write"], ("channels:read", "chat:write"), "present"), + (["channels:read,chat:write"], ("channels:read", "files:write"), "absent"), + (["channels:read"], ("channels:read", "chat:write"), "absent"), ], ) async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, required, expected): @@ -15293,6 +15296,111 @@ async def test_connect_requires_the_requested_vendor_grant(monkeypatch, granted, read.assert_awaited_once_with(prisma, "owner", "server") +_SLACK_SCOPES: Final = frozenset(("channels:read", "chat:write", "files:write")) +_COMMA_NAMED_SCOPES: Final = frozenset(("files:read", "files:read,write")) + + +@pytest.mark.parametrize( + ("granted", "requested", "known", "expected"), + [ + (("channels:read,chat:write",), ("channels:read", "chat:write"), _SLACK_SCOPES, True), + (("channels:read,chat:write",), ("channels:read", "files:write"), _SLACK_SCOPES, False), + (("channels:read,chat:write,files:read",), ("channels:read",), _SLACK_SCOPES, False), + (("x:unknown,chat:write",), ("chat:write",), _SLACK_SCOPES, False), + (("files:read,write",), ("files:read",), _COMMA_NAMED_SCOPES, False), + (("files:read,write",), ("files:read,write",), _COMMA_NAMED_SCOPES, True), + (("channels:read,chat:write",), ("read",), _SLACK_SCOPES | {"read"}, False), + (("read", "write,admin"), ("admin",), frozenset(("read", "write", "admin")), False), + (("write,admin", "read"), ("admin",), frozenset(("read", "write", "admin")), False), + (("read", "write"), ("read", "write"), frozenset[str](), True), + ((), ("read",), frozenset(("read",)), False), + ((), (), frozenset[str](), True), + ], +) +def test_grant_covers_scopes_splits_a_lone_token_only_into_known_scopes( + granted: tuple[str, ...], requested: tuple[str, ...], known: frozenset[str], expected: bool +) -> None: + from litellm.proxy.mcp_server.oauth_utils import grant_covers_scopes + + assert grant_covers_scopes(granted, requested, lambda: known) is expected + + +@pytest.mark.parametrize( + ("granted", "requested", "expected"), + [ + (("channels:read",), ("channels:read",), True), + (("read", "write"), ("read", "write"), True), + (("read", "write"), ("admin",), False), + (("read",), ("admin",), False), + ((), (), True), + ], +) +def test_grant_covers_scopes_answers_like_a_subset_check_without_reading_server_scopes( + granted: tuple[str, ...], requested: tuple[str, ...], expected: bool +) -> None: + from litellm.proxy.mcp_server.oauth_utils import grant_covers_scopes + + def unreachable_server_scopes() -> frozenset[str]: + raise AssertionError("server scopes read for a grant with no comma token") + + assert grant_covers_scopes(granted, requested, unreachable_server_scopes) is expected + + +def test_server_scope_names_joins_configured_and_published_scopes() -> None: + from litellm.proxy.mcp_server.oauth_utils import server_scope_names + + server: Final = _cimd_oauth_server().model_copy( + update={ + "scopes": ["channels:read"], + "authorization_server_metadata": { + "issuer": "https://idp.example.com", + "authorization_endpoint": "https://idp.example.com/authorize", + "token_endpoint": "https://idp.example.com/token", + "response_types_supported": ["code"], + "scopes_supported": ["chat:write", "files:write"], + }, + } + ) + assert server_scope_names(server) == frozenset(("channels:read", "chat:write", "files:write")) + assert server_scope_names(_cimd_oauth_server()) == frozenset[str]() + assert server_scope_names(None) == frozenset[str]() + + +_PUBLISHED_SLACK_METADATA: Final = { + "issuer": "https://idp.example.com", + "authorization_endpoint": "https://idp.example.com/authorize", + "token_endpoint": "https://idp.example.com/token", + "response_types_supported": ["code"], + "scopes_supported": ["channels:read", "chat:write", "files:read"], +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server_update", + [ + {"scopes": ["channels:read", "chat:write", "files:read"]}, + {"authorization_server_metadata": _PUBLISHED_SLACK_METADATA}, + ], +) +async def test_connect_reads_a_comma_grant_against_the_server_scopes( + monkeypatch: pytest.MonkeyPatch, server_update: "Mapping[str, object]" +) -> None: + from litellm.proxy.mcp_server import db + from litellm.proxy.mcp_server import discoverable_endpoints as endpoints + from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager + + server: Final = _cimd_oauth_server().model_copy(update=server_update) + credential: Final = {**_stored_grant(), "scopes": ["channels:read,chat:write,files:read"]} + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr(db, "get_user_oauth_credential", AsyncMock(return_value=credential)) + assert await endpoints._vendor_credential_state("owner", server.server_id, ("channels:read",)) == "absent" + monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) + assert await endpoints._vendor_credential_state("owner", server.server_id, ("channels:read",)) == "present" + assert await endpoints._vendor_credential_state("owner", server.server_id, ("files:write",)) == "absent" + assert await endpoints._vendor_credential_state("owner", server.server_id, ()) == "present" + + @pytest.mark.parametrize( ("header", "expected"), [ From 43a4b465e94a849c64816e1bfaceb1eed6e7a200 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 10 Nov 2026 10:30:00 -0800 Subject: [PATCH 08/11] fix(cost-map): give the Gemini 4 Argon rows the shared Gemini 3 search free tier (#50959) --- litellm/model_prices_and_context_window_backup.json | 3 +++ model_prices_and_context_window.json | 3 +++ tests/unit/test_model_prices_and_context_window.py | 11 +++++++---- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3196b3c0c08..3145ca4a181 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -78910,6 +78910,7 @@ "prompt_cache_min_tokens": 4096.0, "rpm": 2000.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -78970,6 +78971,7 @@ "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -79029,6 +79031,7 @@ "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3196b3c0c08..3145ca4a181 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -78910,6 +78910,7 @@ "prompt_cache_min_tokens": 4096.0, "rpm": 2000.0, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -78970,6 +78971,7 @@ "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -79029,6 +79031,7 @@ "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 62c1fab244c..3408a95935d 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -958,23 +958,26 @@ def test_gemini_4_argon_google_rows_carry_supports_continuation_in_both_cost_map assert _supports_continuation_rows(BACKUP_PRICES_PATH) == main_rows -def test_gemini_4_carbon_rows_share_the_gemini_3_search_allowance_in_both_cost_map_copies() -> None: +def test_gemini_4_rows_share_the_gemini_3_search_allowance_in_both_cost_map_copies() -> None: main_rows: Final = _search_free_requests_rows(PRICES_PATH) # 5,000 free grounding requests per month shared across Gemini 3 and newer: # https://ai.google.dev/gemini-api/docs/pricing#gemini-4-carbon and # https://cloud.google.com/vertex-ai/generative-ai/pricing (2026-11-10) - carbon_rows: Final = { + gemini_4_rows: Final = { + "gemini/gemini-4-argon", + "vertex_ai/gemini-4-argon", + "vertex_ai/gemini-4-argon-cyber-permissive", "gemini/gemini-4-carbon", "vertex_ai/gemini-4-carbon", "vertex_ai/gemini-4-carbon-cyber-permissive", } - assert carbon_rows <= main_rows.keys(), sorted(main_rows) + assert gemini_4_rows <= main_rows.keys(), sorted(main_rows) gemini_3_allowances: Final = { json.dumps(main_rows[key]["search_free_requests"], sort_keys=True) for key in ("gemini/gemini-3.8-flash", "vertex_ai/gemini-3.8-flash") } assert { - json.dumps(main_rows[key]["search_free_requests"], sort_keys=True) for key in carbon_rows + json.dumps(main_rows[key]["search_free_requests"], sort_keys=True) for key in gemini_4_rows } == gemini_3_allowances assert _search_free_requests_rows(BACKUP_PRICES_PATH) == main_rows From 52999149640659344155cbfe6ded6634c3511e9e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:31:57 +0000 Subject: [PATCH 09/11] test(http2): drain client traffic after sending GOAWAY (#50972) Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/sdk/test_http2_wire.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index a263472a45f..30a0635f3b8 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -531,6 +531,9 @@ def _serve_h2(context: ssl.SSLContext, raw: socket.socket, drain: bool) -> None: tls.sendall(connection.data_to_send()) while data := tls.recv(65535): if _reply_to_ended_streams(connection, tls, data, drain) and drain: + tls.shutdown(socket.SHUT_WR) + while tls.recv(65535): + pass return From d47c03175bba41a39c4c3f8a947e28938708eef2 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:33:27 +0000 Subject: [PATCH 10/11] fix(responses): keep unset-effort gpt-5.4/5.5 tool calls on chat completions (#50968) * fix(responses): keep unset-effort gpt-5.4/5.5 tool calls on chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(responses): explicitly route bridge fixtures to responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kenneth Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 4 + litellm/responses/bridge_check.py | 43 ++---- model_prices_and_context_window.json | 4 + .../sdk/test_unset_effort_tool_bridge_wire.py | 146 ++++++++++++++++++ ...hropic_messages_adapters_transformation.py | 26 ++-- .../test_deferred_guardrail_logging.py | 3 +- .../test_responses_api_bridge_flag.py | 65 +++++++- tests/unit/test_main.py | 23 +-- 8 files changed, 249 insertions(+), 65 deletions(-) create mode 100644 tests/integration/sdk/test_unset_effort_tool_bridge_wire.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3145ca4a181..01a32ee7cdb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -72516,6 +72516,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-mini": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 8.25e-08, @@ -72537,6 +72538,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-nano": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 2.2e-08, @@ -73140,6 +73142,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-mini": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 8.25e-08, @@ -73161,6 +73164,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-nano": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 2.2e-08, diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 683a51fc6b6..6a55302027d 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -192,35 +192,11 @@ def responses_api_bridge_check( model_info["mode"] = "responses" model = model.replace("responses/", "") - # OpenAI/Azure GPT-5 chat-completions that need Responses-only fields (e.g. - # ``reasoningSummary`` in ``extra_body``) must be bridged; Chat Completions rejects - # those keys. - # - # - gpt-5.4+: FUNCTION tools with reasoning active must be bridged. OpenAI enables - # reasoning by default for these models (unset reasoning_effort means medium - # server-side), and Chat Completions rejects function tools whenever reasoning is - # on ("Function tools with reasoning_effort are not supported ... use - # /v1/responses or set reasoning_effort to 'none'"), so only an explicit - # ``"none"`` the chat gate forwards keeps the request chat-servable: a ``none`` - # the model's map row refuses is dropped (or refused) before the call, which - # leaves the provider's default effort in force. Custom (grammar) tools are served - # natively by Chat Completions with reasoning on, so custom-only requests stay on - # chat and keep their native custom tool_call response shape. - # - The UNSET-effort arm only fires against endpoints known to enforce that - # constraint (any api.openai.com host, or Azure OpenAI where api_base is - # always set): chat-only OpenAI-compatible backends registered under the openai - # provider with a custom api_base and gpt-5.4+ model names serve tools without - # reasoning fine and have no /responses route, so they keep pre-existing - # behavior (bridge only on an explicit reasoning_effort). - # - Azure AI Foundry's OpenAI v1 hosts (azure_ai provider) enforce it later in the series: - # an explicit effort with function tools is rejected from gpt-5.6 on, and the unset - # effort only from gpt-6 on (gpt-5.6 serves tools with reasoning silently off), so the - # azure_ai gate keys on those measured boundaries instead of gpt-5.4+. - # - Bedrock Mantle serves the same API and enforces the gpt-5.4+ rule even with the effort - # unset (measured on gpt-5.6-terra), so a Mantle row that also serves Responses follows - # these rules, which keeps a deployment pinned to ``mode: chat`` from a 400 on them. - # - Older GPT-5 names (e.g. ``gpt-5``, ``gpt-5.1``): bridge only when a reasoning - # summary alias is present with ``reasoning_effort`` (tools alone stay on chat). + # Function tools with explicit active reasoning bridge on GPT-5.4+. With effort unset, + # OpenAI/Azure and Responses-capable Mantle endpoints bridge only GPT-5.6+ whose map + # default is not "none". Custom OpenAI-compatible endpoints keep unset effort on chat. + # Foundry uses its separate measured boundaries. Custom-only tools remain chat-native; + # reasoning summaries and tool_search retain their independent Responses requirements. has_function_tool: Final = any( ( tool.get("type") == "function" and (isinstance(tool.get("function"), dict) or "name" in tool) @@ -274,7 +250,14 @@ def responses_api_bridge_check( if on_foundry_openai_endpoint else ( OpenAIGPT5Config.is_model_gpt_5_4_plus_model(gpt_series_name) - and (reasoning_effort is not None or on_constraint_enforcing_endpoint) + and ( + reasoning_effort is not None + or ( + on_constraint_enforcing_endpoint + and OpenAIGPT5Config.is_model_gpt_5_6_plus_model(gpt_series_name) + and model_info.get("default_reasoning_effort") != "none" + ) + ) ) ) ) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3145ca4a181..01a32ee7cdb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -72516,6 +72516,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-mini": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 8.25e-08, @@ -72537,6 +72538,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-5.4-nano": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 2.2e-08, @@ -73140,6 +73142,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-mini": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 8.25e-08, @@ -73161,6 +73164,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-5.4-nano": { + "default_reasoning_effort": "none", "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/model-retirement-schedule", "cache_read_input_token_cost": 2.2e-08, diff --git a/tests/integration/sdk/test_unset_effort_tool_bridge_wire.py b/tests/integration/sdk/test_unset_effort_tool_bridge_wire.py new file mode 100644 index 00000000000..11740f1d163 --- /dev/null +++ b/tests/integration/sdk/test_unset_effort_tool_bridge_wire.py @@ -0,0 +1,146 @@ +import json +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm import Router +from tests.integration._support.wire import Reply, Request, wire_server + +_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +_ARGUMENTS: Final = '{"city":"Paris"}' + + +def _tool_reply(request: Request) -> Reply: + body: Final = _OBJECT.validate_json(request.body) + if "/responses" in request.target: + return Reply( + body=json.dumps( + { + "id": "resp_weather", + "object": "response", + "created_at": 1789788253, + "model": "gpt-5.6", + "status": "completed", + "output": [ + { + "type": "function_call", + "id": "fc_weather", + "call_id": "call_weather", + "name": "get_weather", + "arguments": _ARGUMENTS, + "status": "completed", + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-weather", + "object": "chat.completion", + "created": 1789788253, + "model": body.get("model"), + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": _ARGUMENTS}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + +@pytest.mark.parametrize( + ("model", "placement", "endpoint"), + ( + ("gpt-5.4-mini", "unset", "chat/completions"), + ("gpt-5.5", "unset", "chat/completions"), + ("gpt-5.6", "unset", "responses"), + ("gpt-5.4-mini", "extra-body", "responses"), + ("gpt-5.4-mini", "none", "chat/completions"), + ("gpt-5.4-mini", "deployment", "responses"), + ), +) +def test_function_tools_reach_the_chat_servable_endpoint_and_keep_tool_call_shape( + model: str, placement: str, endpoint: str +) -> None: + with wire_server(_tool_reply) as wire: + parameters: Final = { + "model": f"azure/{model}", + "api_base": wire.url, + "api_key": "synthetic-azure-key", + "api_version": "2025-04-01-preview", + } + fields: Final = ( + {"extra_body": {"reasoning_effort": "medium"}} + if placement == "extra-body" + else {"reasoning_effort": "none"} + if placement == "none" + else {} + ) + messages: Final = [{"role": "user", "content": "What's the weather in Paris? Use the tool."}] + if placement == "deployment": + router: Final = Router( + model_list=[ + { + "model_name": "weather", + "litellm_params": {**parameters, "reasoning_effort": "medium"}, + } + ], + num_retries=0, + ) + response: Final = router.completion(model="weather", messages=messages, tools=[_TOOL]) + else: + response_direct: Final = litellm.completion( + **parameters, + **fields, + messages=messages, + tools=[_TOOL], + num_retries=0, + ) + answer: Final = response if placement == "deployment" else response_direct + (sent,) = tuple(request for request in wire.drain() if request.method == "POST") + expected_path: Final = ( + "/openai/responses" if endpoint == "responses" else f"/openai/deployments/{model}/chat/completions" + ) + assert sent.target.split("?")[0] == expected_path, sent.target + assert answer.choices[0].finish_reason == "tool_calls" + assert answer.choices[0].message.tool_calls[0].function.name == "get_weather" + assert json.loads(answer.choices[0].message.tool_calls[0].function.arguments) == {"city": "Paris"} + outbound: Final = _OBJECT.validate_json(sent.body) + if endpoint == "responses": + assert outbound["tools"] == [ + { + "type": "function", + "name": "get_weather", + "parameters": _TOOL["function"]["parameters"], + "description": None, + "strict": None, + } + ] + else: + assert outbound["tools"] == [_TOOL] diff --git a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py index bad7b3f1750..a76d5e82a37 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py @@ -759,7 +759,7 @@ def test_translate_openai_reasoning_items_are_not_duplicated_when_thinking_block def _replayed_reasoning_request(tools: list[dict[str, object]]) -> AnthropicMessagesRequest: return AnthropicMessagesRequest( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", max_tokens=64, messages=[ AnthropicMessagesUserMessageParam(role="user", content="Which city is colder?"), @@ -5820,7 +5820,7 @@ _RETRY_REPLY: Final = { "object": "response", "created_at": 1700000000, "status": "completed", - "model": "gpt-5.4-nano", + "model": "gpt-5.6", "output": [ { "type": "message", @@ -5868,7 +5868,7 @@ async def test_async_messages_retry_once_without_reasoning_the_backend_cannot_de with caplog.at_level(logging.WARNING, logger="LiteLLM"): response: Final = await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -5890,7 +5890,7 @@ async def test_async_messages_retry_once_without_reasoning_the_backend_cannot_de assert retry_warnings == ( "Backend rejected replayed encrypted reasoning, retrying once without it. " "Enable the encrypted_content_affinity pre-call check to drop such blocks before the call. " - "provider=azure model=azure/gpt-5.4-nano status=400 code=invalid_encrypted_content", + "provider=azure model=azure/gpt-5.6 status=400 code=invalid_encrypted_content", ) assert all( "gAAA_foreign" not in message and "litellm_encrypted_reasoning:" not in message for message in warnings @@ -5906,7 +5906,7 @@ def test_sync_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( ) response: Final = litellm.anthropic.messages.create( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -5930,7 +5930,7 @@ async def test_async_messages_with_num_retries_send_the_rejected_reasoning_only_ ) await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -5953,7 +5953,7 @@ def test_sync_messages_with_global_num_retries_send_the_rejected_reasoning_once_ ) litellm.anthropic.messages.create( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -5981,7 +5981,7 @@ async def test_async_messages_retry_a_server_error_with_the_reasoning_and_drop_i ) response: Final = await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6004,7 +6004,7 @@ def test_sync_messages_retry_a_server_error_with_the_reasoning_and_drop_it_only_ ) response: Final = litellm.anthropic.messages.create( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6063,7 +6063,7 @@ async def test_async_messages_log_only_the_outcome_of_the_turn_sent_again( with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6089,7 +6089,7 @@ def test_sync_messages_log_only_the_outcome_of_the_turn_sent_again( with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): litellm.anthropic.messages.create( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6124,7 +6124,7 @@ async def test_messages_do_not_retry_other_rejections( with pytest.raises(litellm.BadRequestError): await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", @@ -6277,7 +6277,7 @@ async def test_streamed_reasoning_summaries_stay_with_their_own_encrypted_item( ) stream: Final = await litellm.anthropic.messages.acreate( - model="azure/gpt-5.4-nano", + model="azure/gpt-5.6", api_base=_RETRY_API_BASE, api_key="fake-key", api_version="2025-04-01-preview", diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index d2161fddeb0..bfbf7318d52 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -330,7 +330,7 @@ async def test_deferred_slot_keeps_the_innermost_wrapper_result(): async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the_provider_usage( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): - """/v1/messages on an Azure gpt-5.4+ deployment with function tools runs three nested + """/v1/messages on an Azure deployment with function tools and explicit reasoning runs three nested wrappers: anthropic_messages, the chat adapter's acompletion, and the Responses bridge acompletion hands the call to, which retags the call as ``responses``. With logging deferred for a post-call guardrail the stored closure must carry the innermost provider @@ -382,6 +382,7 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the response: Final = await litellm.anthropic_messages( model="azure/gpt-5.4-nano", + reasoning_effort="medium", messages=[{"role": "user", "content": "hi"}], max_tokens=16, tools=[ diff --git a/tests/unit/responses/test_responses_api_bridge_flag.py b/tests/unit/responses/test_responses_api_bridge_flag.py index 210488f0888..f6084529ed6 100644 --- a/tests/unit/responses/test_responses_api_bridge_flag.py +++ b/tests/unit/responses/test_responses_api_bridge_flag.py @@ -786,17 +786,74 @@ _FUNCTION_TOOL: Final[dict[str, JsonValue]] = { @pytest.mark.parametrize("custom_llm_provider", ["openai", "azure"]) -def test_a_chat_pin_cannot_pull_gpt_5_4_function_tool_calls_off_the_responses_api(custom_llm_provider: str) -> None: - """Chat Completions rejects function tools while reasoning is on for gpt-5.4+, so the capability rule - must still bridge the request even when the deployment pins ``mode: chat``.""" +@pytest.mark.parametrize(("effort", "mode"), ((None, "chat"), ("medium", "responses"))) +def test_a_chat_pin_routes_gpt_5_4_function_tools_by_explicit_effort( + custom_llm_provider: str, effort: str | None, mode: str +) -> None: info, model = responses_api_bridge_check( model="gpt-5.4", custom_llm_provider=custom_llm_provider, tools=[_FUNCTION_TOOL], + reasoning_effort=effort, api_base="https://example.openai.azure.com" if custom_llm_provider == "azure" else None, deployment_model_info=_CHAT, ) - assert (info["mode"], model) == ("responses", "gpt-5.4") + assert (info["mode"], model) == (mode, "gpt-5.4") + + +@pytest.mark.parametrize("provider", ("openai", "azure")) +@pytest.mark.parametrize( + ("model", "effort", "extra_body", "expected"), + ( + ("gpt-5.4-mini", None, None, "chat"), + ("gpt-5.4-nano", None, None, "chat"), + ("gpt-5.5", None, None, "chat"), + ("gpt-5.6", None, None, "responses"), + ("gpt-6.1-sol", None, None, "responses"), + ("gpt-5.4-mini", "none", None, "chat"), + ("gpt-5.4-mini", "medium", None, "responses"), + ("gpt-5.4-mini", {"effort": "medium"}, None, "responses"), + ("gpt-5.4-mini", None, {"reasoning_effort": "medium"}, "responses"), + ), +) +def test_function_tools_bridge_only_when_chat_rejects_the_effective_effort( + provider: str, + model: str, + effort: str | Mapping[str, object] | None, + extra_body: Mapping[str, object] | None, + expected: str, +) -> None: + info, _ = responses_api_bridge_check( + model=model, + custom_llm_provider=provider, + tools=[_FUNCTION_TOOL], + reasoning_effort=bridge_reasoning_effort(effort, extra_body), + api_base="https://api.openai.com/v1", + ) + assert info.get("mode") == expected + + +@pytest.mark.parametrize("region", ("us", "eu")) +@pytest.mark.parametrize("model", ("gpt-5.4-mini", "gpt-5.4-nano")) +def test_regional_azure_default_none_function_tools_stay_on_chat(region: str, model: str) -> None: + info, _ = responses_api_bridge_check( + model=f"{region}/{model}", custom_llm_provider="azure", tools=[_FUNCTION_TOOL] + ) + assert info.get("mode") == "chat" + assert info.get("default_reasoning_effort") == litellm.model_cost[f"azure/{model}"]["default_reasoning_effort"] + + +def test_registered_gpt_5_6_default_none_function_tools_stay_on_chat(monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "gpt-5.6-test-default-none" + monkeypatch.setitem( + litellm.model_cost, + model, + {**litellm.model_cost["gpt-5.6"], "default_reasoning_effort": "none"}, + ) + info, _ = responses_api_bridge_check( + model=model, custom_llm_provider="openai", tools=[_FUNCTION_TOOL], api_base="https://api.openai.com/v1" + ) + assert info.get("mode") == "chat" @pytest.mark.parametrize("model", ["openai.gpt-5.6-terra", "us-east-1/openai.gpt-5.6-terra"]) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index dbd9ba8a32e..d2035b6a682 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1418,12 +1418,7 @@ def test_responses_api_bridge_check_azure_gpt_5_4_tools_plus_reasoning_routes_to assert model_info.get("mode") == "responses" -def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): - """ - Azure gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables - reasoning by default for gpt-5.4+, and Chat Completions rejects function tools - whenever reasoning is on. - """ +def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_stays_chat(): from litellm.main import responses_api_bridge_check with patch("litellm.responses.bridge_check.get_model_info_helper") as mock_get_model_info: @@ -1436,15 +1431,10 @@ def test_responses_api_bridge_check_azure_gpt_5_4_tools_with_default_reasoning_r ) assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" + assert model_info.get("mode") != "responses" -def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_to_responses(): - """ - gpt-5.4 with tools and UNSET reasoning_effort must bridge: OpenAI enables reasoning - by default for gpt-5.4+, and Chat Completions rejects function tools whenever - reasoning is on ("use /v1/responses or set reasoning_effort to 'none'"). - """ +def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_stays_chat(): from litellm.main import responses_api_bridge_check with patch("litellm.responses.bridge_check.get_model_info_helper") as mock_get_model_info: @@ -1457,7 +1447,7 @@ def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_ ) assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" + assert model_info.get("mode") != "responses" @pytest.mark.parametrize("region", ("us", "eu")) @@ -1972,8 +1962,7 @@ def test_responses_api_bridge_check_custom_api_base_with_explicit_effort_still_r assert model_info.get("mode") == "responses" -def test_responses_api_bridge_check_azure_with_api_base_and_unset_effort_routes(): - """Azure OpenAI always sets api_base and does enforce the constraint; keep bridging.""" +def test_responses_api_bridge_check_azure_gpt_5_4_with_api_base_and_unset_effort_stays_chat(): from litellm.main import responses_api_bridge_check with patch("litellm.responses.bridge_check.get_model_info_helper") as mock_get_model_info: @@ -1987,7 +1976,7 @@ def test_responses_api_bridge_check_azure_with_api_base_and_unset_effort_routes( ) assert model == "gpt-5.4" - assert model_info.get("mode") == "responses" + assert model_info.get("mode") != "responses" _FOUNDRY_API_BASE: Final = "https://myproject.services.ai.azure.com" From 2624f7caead93e7e84e85a282a22650bb3d98ae9 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 10 Nov 2026 10:42:31 -0800 Subject: [PATCH 11/11] fix(cost-map): mark the azure_ai Claude rows that bill web search as supporting it (#50960) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 15 +++++++++++++++ model_prices_and_context_window.json | 15 +++++++++++++++ .../unit/test_model_prices_and_context_window.py | 13 +++++++++++++ 3 files changed, 43 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 01a32ee7cdb..a4c16fbeaa3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3965,6 +3965,7 @@ "supports_vision": true, "supports_output_config": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 4096 }, "azure_ai/claude-opus-4-7": { @@ -3997,6 +3998,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 2048 }, "azure_ai/claude-fable-5": { @@ -4031,6 +4033,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, + "supports_web_search": true, "prompt_cache_min_tokens": 512 }, "azure_ai/claude-fable-5-1": { @@ -4065,6 +4068,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, + "supports_web_search": true, "prompt_cache_min_tokens": 512, "deprecation_date": "2027-12-05", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" @@ -4136,6 +4140,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-5-5": { @@ -4171,6 +4176,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" }, @@ -4208,6 +4214,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-5": { @@ -4267,6 +4274,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-6": { @@ -76637,6 +76645,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, "supports_forced_tool_use": false, + "supports_web_search": true, "thinking_always_on": true, "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" @@ -78208,6 +78217,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, "supports_forced_tool_use": true, + "supports_web_search": true, "thinking_always_on": false, "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/models/haiku-5-5/migration-guide" @@ -79277,6 +79287,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "baseten/meta-models/muse-spark-open": { @@ -88134,6 +88145,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-mythos-5-2-fast": { @@ -88169,6 +88181,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-fable-5-2": { @@ -88200,6 +88213,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-fable-5-2-fast": { @@ -88231,6 +88245,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 01a32ee7cdb..a4c16fbeaa3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3965,6 +3965,7 @@ "supports_vision": true, "supports_output_config": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 4096 }, "azure_ai/claude-opus-4-7": { @@ -3997,6 +3998,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 2048 }, "azure_ai/claude-fable-5": { @@ -4031,6 +4033,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, + "supports_web_search": true, "prompt_cache_min_tokens": 512 }, "azure_ai/claude-fable-5-1": { @@ -4065,6 +4068,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": false, + "supports_web_search": true, "prompt_cache_min_tokens": 512, "deprecation_date": "2027-12-05", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" @@ -4136,6 +4140,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 512 }, "azure_ai/claude-opus-5-5": { @@ -4171,6 +4176,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" }, @@ -4208,6 +4214,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-5": { @@ -4267,6 +4274,7 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, + "supports_web_search": true, "prompt_cache_min_tokens": 1024 }, "azure_ai/claude-sonnet-4-6": { @@ -76637,6 +76645,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, "supports_forced_tool_use": false, + "supports_web_search": true, "thinking_always_on": true, "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" @@ -78208,6 +78217,7 @@ "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, "supports_forced_tool_use": true, + "supports_web_search": true, "thinking_always_on": false, "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/models/haiku-5-5/migration-guide" @@ -79277,6 +79287,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "baseten/meta-models/muse-spark-open": { @@ -88134,6 +88145,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-mythos-5-2-fast": { @@ -88169,6 +88181,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-fable-5-2": { @@ -88200,6 +88213,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "azure_ai/claude-fable-5-2-fast": { @@ -88231,6 +88245,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, + "supports_web_search": true, "thinking_always_on": true }, "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 3408a95935d..b270d385693 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -1048,3 +1048,16 @@ def test_bedrock_gpt_oss_rows_declare_vision_support_like_their_safeguard_siblin def test_bedrock_gpt_oss_vision_is_explicitly_disabled_for_the_router_gate() -> None: assert litellm.utils.is_vision_explicitly_disabled("bedrock/openai.gpt-oss-120b-1:0") is True assert litellm.utils.is_vision_explicitly_disabled("bedrock_mantle/openai.gpt-oss-20b") is True + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_azure_ai_claude_rows_priced_for_web_search_say_they_support_it(path: Path) -> None: + prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + priced_rows: Final = { + key: row + for key, row in prices.items() + if key.startswith("azure_ai/claude") and isinstance(row, dict) and row.get("search_context_cost_per_query") + } + assert priced_rows + unflagged: Final = tuple(key for key, row in priced_rows.items() if row.get("supports_web_search") is not True) + assert unflagged == (), unflagged