Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_redis_breaker_open_silent_miss

This commit is contained in:
mateo-berri 2026-09-10 19:40:52 -07:00
commit fdd423128e
48 changed files with 1416 additions and 236 deletions

View file

@ -59,3 +59,54 @@ tests:
- contains:
path: spec.template.spec.containers[0].args
content: "4"
- it: should give the collector sidecar the same pool env as the proxy container
template: deployment.yaml
set:
collector.enabled: true
db.connectionPool.enabled: true
db.connectionPool.maxDbConnections: 8
db.connectionPool.maxClientConn: 400
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "400"
- it: should give the collector sidecar no pool env when the pool is off
template: deployment.yaml
set:
collector.enabled: true
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: litellm-collector
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true

View file

@ -170,6 +170,9 @@ spec:
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml
{{- end }}
{{- if .Values.database.connectionPool.enabled }}
{{- include "litellm.connectionPoolEnv" $ | nindent 12 }}
{{- end }}
{{- include "litellm.gateway.collectorEnv" . | nindent 12 }}
- name: LITELLM_JOB_ROLE
value: collector

View file

@ -29,16 +29,16 @@ tests:
value: Resource
template: gateway/hpa.yaml
- it: runs the collector as a sidecar sharing env, config and a unix socket emptyDir, and scales on the gateway container only
- it: runs the collector as a sidecar sharing env, config, the pod pool and a unix socket emptyDir, and scales on the gateway container only
set:
gateway.collector.enabled: true
gateway.collector.bufferSize: 250
gateway.collector.onUnavailable: drop
gateway.image.tag: v1.102.0
gateway.numWorkers: 4
gateway.extraEnv:
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
gateway.envSecrets:
- litellm-license
gateway.volumes:
@ -107,6 +107,18 @@ tests:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[1].env
content:

View file

@ -82,9 +82,70 @@ tests:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- it: collector sidecar gets the same pool env as the gateway container, the metrics sidecar none
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
gateway.metricsServer.enabled: true
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: metrics
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- equal:
path: spec.template.spec.containers[2].name
value: collector
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
- it: collector sidecar gets no pool env when the pool is off
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
asserts:
- equal:
path: spec.template.spec.containers[1].name
value: collector
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true
- it: pool with IAM auth renders both the pool and the token auth flag
template: gateway/deployment.yaml
set:
gateway.collector.enabled: true
database.connectionPool.enabled: true
database.writer.useIAMAuth: true
asserts:
@ -98,6 +159,16 @@ tests:
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[1].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- it: pool with Entra auth renders both the pool and the token auth flag
template: gateway/deployment.yaml

View file

@ -234,8 +234,8 @@ database:
# workers run; the workers connect to the pool over loopback, with no extra
# network hop. The chart emits LITELLM_PGBOUNCER_ENABLED /
# LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on
# the gateway container only: the backend runs a single worker and the
# migrations Job must keep a direct connection. With
# the gateway container and its collector sidecar only: the backend runs a
# single worker and the migrations Job must keep a direct connection. With
# `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and
# renews the database token itself, so the workers never see it. Starting profile for
# `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.95"
version = "0.4.96"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.95"
version = "0.4.96"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -295,6 +295,9 @@ def _get_provider_request_id(original_exception: Exception) -> str | None:
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
_MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
_UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset(
("user_api_key_auth", "user_api_key_budget_reservation")
)
sentry_sdk_instance = None
capture_exception = None
@ -5386,23 +5389,23 @@ class StandardLoggingPayloadSetup:
Returns:
dict: Merged metadata with user API key fields taking precedence
"""
merged_metadata: Final[dict] = {}
# Start with metadata (user API key fields) - but skip non-serializable objects
if litellm_params.get("metadata") and isinstance(litellm_params.get("metadata"), dict):
for key, value in litellm_params["metadata"].items():
# Skip non-serializable objects like UserAPIKeyAuth
if key in {"user_api_key_auth", "user_api_key_budget_reservation"}:
continue
merged_metadata[key] = value
# Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys
if litellm_params.get("litellm_metadata") and isinstance(litellm_params.get("litellm_metadata"), dict):
for key, value in litellm_params["litellm_metadata"].items():
if key not in merged_metadata: # Don't overwrite existing keys from metadata
merged_metadata[key] = value
return merged_metadata
metadata: Final = litellm_params.get("metadata")
litellm_metadata: Final = litellm_params.get("litellm_metadata")
user_metadata: Final = MappingProxyType(
{
key: value
for key, value in (metadata.copy().items() if isinstance(metadata, dict) else ())
if key not in _UNSERIALIZABLE_METADATA_KEYS
}
)
model_metadata: Final = MappingProxyType(
{
key: value
for key, value in (litellm_metadata.copy().items() if isinstance(litellm_metadata, dict) else ())
if key not in user_metadata
}
)
return {**user_metadata, **model_metadata} # mutable-ok: function contract returns a plain dict
@staticmethod
def get_standard_logging_metadata(
@ -5660,7 +5663,7 @@ class StandardLoggingPayloadSetup:
additional_logging_headers[key] = additiona_headers[_key]
# Preserve all remaining headers verbatim (e.g. llm_provider-x-request-id)
for k, v in additiona_headers.items():
for k, v in additiona_headers.copy().items():
if k.lower() not in typed_keys:
additional_logging_headers[k] = v

View file

@ -236,9 +236,11 @@ class CustomStreamWrapper:
stream_options=None,
make_call: Callable | None = None,
_response_headers: dict | httpx.Headers | None = None,
count_prompt_tokens: Callable[[], int] | None = None,
):
self.model = model
self.make_call = make_call
self.count_prompt_tokens = count_prompt_tokens
self.custom_llm_provider = custom_llm_provider
self.logging_obj: LiteLLMLoggingObject = logging_obj
self.completion_stream = completion_stream
@ -1641,7 +1643,7 @@ class CustomStreamWrapper:
except Exception:
model_response.choices[0].delta = Delta()
else:
if self.stream_options is not None and self.stream_options["include_usage"] is True:
if self.send_stream_usage is True:
model_response.choices = []
return model_response
self._record_usage_only_chunk(model_response=model_response)
@ -1996,6 +1998,7 @@ class CustomStreamWrapper:
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,
count_prompt_tokens=self.count_prompt_tokens,
)
except Exception as e:
# stream_chunk_builder can re-raise (as APIError) on large agentic
@ -2248,6 +2251,7 @@ class CustomStreamWrapper:
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,
count_prompt_tokens=self.count_prompt_tokens,
)
except Exception as e:
# see sync __next__: a raise from stream_chunk_builder inside this
@ -2371,6 +2375,7 @@ class CustomStreamWrapper:
chunks=self.chunks,
messages=self.messages if isinstance(self.messages, list) else None,
logging_obj=self.logging_obj,
count_prompt_tokens=self.count_prompt_tokens,
)
if partial_response is None:
return

View file

@ -850,6 +850,12 @@ def admission_input_tokens(kwargs: Mapping[str, object]) -> int | None:
)
def admitted_prompt_token_counter(prompt_tokens: int | None) -> Callable[[], int] | None:
if prompt_tokens is None:
return None
return lambda: prompt_tokens
def mock_completion(
model: str,
messages: list,
@ -935,23 +941,26 @@ def mock_completion(
if stream is True:
model_response = ModelResponseStream()
count_prompt_tokens: Final = admitted_prompt_token_counter(prompt_tokens)
# don't try to access stream object,
if kwargs.get("acompletion", False) is True:
return CustomStreamWrapper(
completion_stream=async_mock_completion_streaming_obj(
model_response, mock_response=mock_response, model=model, n=n
model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens
),
model=model,
custom_llm_provider="openai",
logging_obj=logging,
count_prompt_tokens=count_prompt_tokens,
)
return CustomStreamWrapper(
completion_stream=mock_completion_streaming_obj(
model_response, mock_response=mock_response, model=model, n=n
model_response, mock_response=mock_response, model=model, n=n, prompt_tokens=prompt_tokens
),
model=model,
custom_llm_provider="openai",
logging_obj=logging,
count_prompt_tokens=count_prompt_tokens,
)
if isinstance(mock_response, litellm.MockException):
raise mock_response

View file

@ -59,6 +59,10 @@ def _get_priority_settings() -> "PriorityReservationSettings":
return settings
def _is_latin1_encodable(value: object) -> bool:
return all(ord(char) < 256 for char in str(value))
class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"""
Saturation-aware priority-based rate limiter using v3 infrastructure.
@ -666,7 +670,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
if response_has_hidden_params(response):
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers["x-litellm-priority"] = priority or "default"
priority_header: Final = priority or "default"
if _is_latin1_encodable(priority_header):
additional_headers["x-litellm-priority"] = priority_header
else:
verbose_proxy_logger.debug(
"Skipping x-litellm-priority header: priority %r is not Latin-1 encodable", priority
)
additional_headers["x-litellm-rate-limiter-version"] = "v3"
return response

View file

@ -455,6 +455,13 @@ If 2+ reasoning markers are detected in the user message, the request is promote
Reasoning markers in the system prompt do **not** trigger the reasoning override. This prevents system prompts like "Think step by step before answering" from forcing all requests to the reasoning tier.
For requests identified by a `claude-cli/` or `claude-code/` user agent, the LLM classifier omits caller system
text to avoid classifying environment, agent, and skill catalogs. The current ask, configured prior-turn context,
and trajectory signal remain unchanged. The routed completion still receives the original system text. This
also excludes genuine task constraints supplied only in Claude Code system messages. Other clients keep the
existing system-context behavior. The browser routing preview has no client-identity field and retains that
generic behavior; use the real client when checking Claude Code routing.
### Harness Reminder Blocks
Agent harnesses inject their own context into the conversation as ordinary message text. That text is plumbing, not something a human asked for, so the router strips complete reminder blocks before classifying and picking a tier. A turn that is nothing but a reminder block strips to empty and is skipped, and the router falls back to the last real ask instead

View file

@ -48,6 +48,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
request_contains_image_content,
)
from litellm.litellm_core_utils.sensitive_data_masker import mask_credentials_in_payload
from litellm.llms.anthropic.common_utils import is_claude_code_user_agent
from litellm.llms.base_llm.base_utils import type_to_response_format_param
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
from litellm.router_strategy.complexity_router.tier_predictor import (
@ -1993,8 +1994,7 @@ class ComplexityRouter(CustomLogger):
Args:
prompt: The current user ask text (already extracted as the real human ask, not tool results)
system_prompt: The caller's system prompt (task constraints), always included so later
turns never lose it
system_prompt: Caller task constraints, omitted from classification for Claude Code requests
request_kwargs: Request metadata for spend attribution
messages: Full message history for extracting prior turns and the trajectory signal
"""
@ -2027,9 +2027,18 @@ class ComplexityRouter(CustomLogger):
)
encrypted_task: Final = _encrypted_classifier_task(request_kwargs, marker_pairs)
caller_system_prompt: Final = (
None
if any(
is_claude_code_user_agent(user_agent)
for metadata in (self._iter_metadata_dicts(request_kwargs) if request_kwargs is not None else ())
if isinstance(user_agent := metadata.get("user_agent"), str)
)
else system_prompt
)
user_payload: Final = self._build_classifier_user_payload(
prompt="The delegated task in the following agent_message." if encrypted_task is not None else prompt,
system_prompt=system_prompt,
system_prompt=caller_system_prompt,
prior_turns=prior_turns,
messages=messages,
has_prior_conversation=has_prior_conversation,

View file

@ -971,9 +971,11 @@ class ComplexityRouterConfig(BaseModel):
"classified against what it refers to. Counts turns of both roles when "
"classifier_context_include_assistant_turns is enabled. These turns are sent to the classifier "
"model, which may "
"be a different deployment or provider than the routed completion model; that call already "
"carries the current user ask and the caller's system prompt in full. Set to 0 to send neither "
"prior turns nor any conversation context beyond the current ask. Only applies when "
"be a different deployment or provider than the routed completion model; that call carries "
"the current user ask and, except for Claude Code requests, the extracted system-role text in full. "
"Claude Code system text is omitted to avoid classifying harness instructions; the routed "
"completion still receives it. Set to 0 to send neither prior turns nor "
"any conversation context beyond the current ask. Only applies when "
"classifier_type is 'llm'."
),
)
@ -985,9 +987,9 @@ class ComplexityRouterConfig(BaseModel):
"context window, per classification call. Turns are taken newest first and quoted whole "
"while they fit, so a conversation small enough to quote entirely is never cut; once the "
"budget runs out the older turns are dropped whole and only the turn straddling the "
"boundary is truncated, into whatever space is left. The current ask and the caller's "
"system prompt sit outside this budget and are always sent in full, as does the numbering "
"each quoted turn carries. A budget under 120 leaves no room to quote a turn and "
"boundary is truncated, into whatever space is left. The current ask and, except for Claude "
"Code requests, the extracted system-role text sit outside this budget and are sent in full, as does "
"the numbering each quoted turn carries. A budget under 120 leaves no room to quote a turn and "
"suppresses the block; set classifier_context_window_size to 0 to turn context off "
"deliberately. Only applies when classifier_type is 'llm'."
),

View file

@ -66,6 +66,7 @@ from litellm.constants import (
DEFAULT_EMBEDDING_PARAM_VALUES,
DEFAULT_MAX_LRU_CACHE_SIZE,
DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
DEFAULT_TRIM_RATIO,
FUNCTION_DEFINITION_TOKEN_COUNT,
@ -278,7 +279,7 @@ except (ImportError, AttributeError, TypeError):
# Convert to str (if necessary)
claude_json_str = json.dumps(json_data)
import importlib.metadata
from collections.abc import Callable, Iterable, Mapping, Sequence
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args
from litellm import utils as litellm_utils
@ -1195,6 +1196,47 @@ def function_setup(
raise e
def _dispatch_success_logging(
logging_obj: LiteLLMLoggingObject,
result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
is_completion_with_fallbacks: bool,
is_litellm_internal_call: bool,
) -> None:
if not is_litellm_internal_call:
if getattr(logging_obj, "_defer_async_logging", False):
def _enqueue_deferred_logging() -> None:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging
else:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=result,
start_time=start_time,
end_time=end_time,
)
async def _client_async_logging_helper(
logging_obj: LiteLLMLoggingObject,
result,
@ -1662,6 +1704,16 @@ def client(original_function):
kwargs=kwargs,
)
_update_response_metadata: Final = getattr(sys.modules[__name__], "update_response_metadata")
_update_response_metadata(
result=result,
logging_obj=logging_obj,
model=model,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
# LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated
verbose_logger.info("Wrapper: Completed Call, calling success_handler")
# Copy the current context to propagate it to the background thread
@ -1676,15 +1728,6 @@ def client(original_function):
end_time,
)
# RETURN RESULT
update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata")
update_response_metadata(
result=result,
logging_obj=logging_obj,
model=model,
kwargs=kwargs,
start_time=start_time,
end_time=end_time,
)
return result
except Exception as e:
call_type = original_function.__name__
@ -1944,48 +1987,20 @@ def client(original_function):
args=args,
)
# LOG SUCCESS - handle streaming success logging in the _next_ object
# Internal sub-calls (e.g. emulated file-search steps) share the
# parent's logging obj; skip async logging here so only the outer call bills once.
# NOTE: streaming requests return early (before this point) via
# CustomStreamWrapper, so this block is non-streaming only.
if not _is_litellm_internal_call:
if getattr(logging_obj, "_defer_async_logging", False):
def _enqueue_deferred_logging() -> None:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging
else:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=result,
start_time=start_time,
end_time=end_time,
)
# REBUILD EMBEDDING CACHING
if (
isinstance(result, EmbeddingResponse)
and _caching_handler_response is not None
and _caching_handler_response.final_embedding_cached_response is not None
):
_dispatch_success_logging(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
is_litellm_internal_call=_is_litellm_internal_call,
)
return _llm_caching_handler._combine_cached_embedding_response_with_api_result(
_caching_handler_response=_caching_handler_response,
embedding_response=result,
@ -2001,6 +2016,14 @@ def client(original_function):
start_time=start_time,
end_time=end_time,
)
_dispatch_success_logging(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
is_litellm_internal_call=_is_litellm_internal_call,
)
return result
except Exception as e:
@ -7001,7 +7024,26 @@ class TextCompletionStreamWrapper:
raise StopAsyncIteration
def mock_completion_streaming_obj(model_response, mock_response, model, n: int | None = None):
def mock_stream_usage_chunk(model_response: ModelResponseStream, model: str, prompt_tokens: int) -> ModelResponseStream:
return ModelResponseStream(
id=model_response.id,
choices=[], # mutable-ok: ModelResponseStream only treats a list as explicit choices, a tuple gets a default choice
model=model,
usage=Usage(
prompt_tokens=prompt_tokens,
completion_tokens=DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
total_tokens=prompt_tokens + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
),
)
def mock_completion_streaming_obj(
model_response: ModelResponseStream,
mock_response: str | MockException | ModelResponseStream,
model: str,
n: int | None = None,
prompt_tokens: int | None = None,
) -> Iterator[ModelResponseStream]:
if isinstance(mock_response, litellm.MockException):
raise mock_response
if isinstance(mock_response, ModelResponseStream):
@ -7021,14 +7063,17 @@ def mock_completion_streaming_obj(model_response, mock_response, model, n: int |
_all_choices.append(_streaming_choice)
model_response.choices = _all_choices
yield model_response
if prompt_tokens is not None:
yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens)
async def async_mock_completion_streaming_obj(
model_response,
model_response: ModelResponseStream,
mock_response: str | MockException | ModelResponseStream,
model,
model: str,
n: int | None = None,
):
prompt_tokens: int | None = None,
) -> AsyncIterator[ModelResponseStream]:
if isinstance(mock_response, litellm.MockException):
raise mock_response
if isinstance(mock_response, ModelResponseStream):
@ -7048,6 +7093,8 @@ async def async_mock_completion_streaming_obj(
_all_choices.append(_streaming_choice)
model_response.choices = _all_choices
yield model_response
if prompt_tokens is not None:
yield mock_stream_usage_chunk(model_response, model=model, prompt_tokens=prompt_tokens)
########## Reading Config File ############################

View file

@ -67,7 +67,7 @@ proxy = [
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.28.1,<2.0",
"litellm-proxy-extras==0.4.95",
"litellm-proxy-extras==0.4.96",
"litellm-enterprise==0.1.66",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",

View file

@ -362,7 +362,14 @@ The sidecar gets the same database, Redis, master-key, license, proxy
config, and `gateway_extra_env` / `gateway_extra_secrets` values as the
gateway container, runs with `LITELLM_JOB_ROLE=collector`, and is
non-essential with an ECS restart policy, so a sidecar crash restarts it in
place while the gateway falls back to in-process spend tracking.
place while the gateway falls back to in-process spend tracking. With
`gateway_connection_pool_enabled` it also gets the `LITELLM_PGBOUNCER_*` env,
so with a password-authenticated database (`create_database = false`) its
Prisma client goes through the task-local PgBouncer instead of opening a
second pool straight to the database. Under IAM token auth (the module-managed
Aurora cluster) the collector keeps its own direct connection on purpose: the
pooler's auth file only holds the token the gateway container minted, which
the sidecar cannot present, so it mints its own.
```hcl
collector_enabled = true

View file

@ -51,6 +51,43 @@ run "pool_enabled_renders_the_three_vars_with_configured_sizes" {
}
}
run "collector_sidecar_gets_the_same_pool_env_as_the_gateway" {
command = plan
variables {
create_database = false
database_url = "postgresql://litellm:pw@db.internal:5432/litellm"
collector_enabled = true
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 8
gateway_pool_max_client_conn = 250
}
assert {
condition = alltrue([
for env in [local.gateway_environment, local.collector_container[0].environment] : (
{ for e in env : e.name => e.value }["LITELLM_PGBOUNCER_ENABLED"] == "true" &&
{ for e in env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"] == "8" &&
{ for e in env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_CLIENT_CONN"] == "250"
)
])
error_message = "The collector sidecar must carry the same three LITELLM_PGBOUNCER_* vars as the gateway so its Prisma connects to the task-local pool."
}
}
run "collector_sidecar_gets_no_pool_env_when_the_pool_is_off" {
command = plan
variables {
collector_enabled = true
}
assert {
condition = !anytrue([for e in local.collector_container[0].environment : startswith(e.name, "LITELLM_PGBOUNCER_")])
error_message = "The collector sidecar must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set."
}
}
run "pool_enabled_uses_the_module_default_sizes" {
command = plan

View file

@ -339,8 +339,11 @@ instead of the Unix socket helm uses; the proxy rejects any non-loopback
address. The sidecar runs the same Redis CA + `DATABASE_URL` bootstrap as
the gateway container, gets the same database, Redis, master-key, license,
proxy config, and `gateway_extra_env` / `gateway_extra_secrets` values, and
runs with `LITELLM_JOB_ROLE=collector`. When it is unreachable the
gateway falls back to in-process spend tracking.
runs with `LITELLM_JOB_ROLE=collector`. With `gateway_connection_pool_enabled`
it also gets the `LITELLM_PGBOUNCER_*` env, so its Prisma client goes through
the instance-local PgBouncer instead of opening a second pool straight to the
database. When it is unreachable the gateway falls back to in-process spend
tracking.
```hcl
collector_enabled = true

View file

@ -79,6 +79,45 @@ run "pool_enabled_renders_the_three_vars_with_configured_sizes" {
}
}
run "collector_sidecar_gets_the_same_pool_env_as_the_gateway" {
command = plan
variables {
collector_enabled = true
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 8
gateway_pool_max_client_conn = 250
}
assert {
condition = alltrue([
for c in google_cloud_run_v2_service.gateway[0].template[0].containers : (
{ for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_ENABLED"] == "true" &&
{ for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"] == "8" &&
{ for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_CLIENT_CONN"] == "250"
) if c.name == "spend-collector"
]) && length([for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name if c.name == "spend-collector"]) == 1
error_message = "The spend-collector sidecar must carry the same three LITELLM_PGBOUNCER_* vars as the gateway so its Prisma connects to the instance-local pool."
}
}
run "collector_sidecar_gets_no_pool_env_when_the_pool_is_off" {
command = plan
variables {
collector_enabled = true
}
assert {
condition = !anytrue(flatten([
for c in google_cloud_run_v2_service.gateway[0].template[0].containers : [
for e in c.env : startswith(e.name, "LITELLM_PGBOUNCER_")
] if c.name == "spend-collector"
])) && length([for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name if c.name == "spend-collector"]) == 1
error_message = "The spend-collector sidecar must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set."
}
}
run "pool_enabled_uses_the_module_default_sizes" {
command = plan

View file

@ -159,15 +159,23 @@ def test_azure_o_series_routing():
def test_openai_o_series_max_retries_0(mock_get_openai_client):
import litellm
mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {}
mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = (
ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}])
)
litellm.set_verbose = True
response = litellm.completion(
model="azure/o1-preview",
messages=[{"role": "user", "content": "hi"}],
max_retries=0,
api_key="fake-key",
api_base="https://fake-azure.openai.azure.com",
api_version="2024-10-21",
)
mock_get_openai_client.assert_called_once()
assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0
assert response.choices[0].message.content == "Hello"
@pytest.mark.asyncio

View file

@ -335,6 +335,10 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version):
]
with patch.object(client.chat.completions.with_raw_response, "create") as mock_post:
mock_post.return_value.headers = {}
mock_post.return_value.parse.return_value = litellm.ModelResponse(
choices=[{"message": {"role": "assistant", "content": InvestigationOutput().model_dump_json()}}]
)
response = litellm.completion(
model="azure/gpt-4.1-mini",
messages=[
@ -362,6 +366,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version):
assert "response_format" in mock_post.call_args.kwargs
else:
assert "response_format" not in mock_post.call_args.kwargs
assert response.choices[0].message.content == InvestigationOutput().model_dump_json()
def test_map_openai_params():

View file

@ -292,15 +292,21 @@ class TestOpenAIChatCompletion(BaseLLMChatTest):
def test_openai_max_retries_0(mock_get_openai_client):
import litellm
mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {}
mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = (
ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}])
)
litellm.set_verbose = True
response = litellm.completion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
max_retries=0,
api_key="fake-key",
)
mock_get_openai_client.assert_called_once()
assert mock_get_openai_client.call_args.kwargs["max_retries"] == 0
assert response.choices[0].message.content == "Hello"
@patch("litellm.main.openai_chat_completions._get_openai_client")

View file

@ -3999,10 +3999,14 @@ def test_completion_novita_ai():
openai_client = OpenAI(api_key="fake-key")
with patch.object(
openai_client.chat.completions, "create", new=MagicMock()
openai_client.chat.completions.with_raw_response, "create"
) as mock_call:
mock_call.return_value.headers = {}
mock_call.return_value.parse.return_value = litellm.ModelResponse(
choices=[{"message": {"role": "assistant", "content": "Hello"}}]
)
try:
completion(
response = completion(
model="novita/meta-llama/llama-3.3-70b-instruct",
messages=messages,
client=openai_client,
@ -4010,6 +4014,7 @@ def test_completion_novita_ai():
)
mock_call.assert_called_once()
assert response.choices[0].message.content == "Hello"
# Verify model is passed correctly
assert (

View file

@ -1076,7 +1076,7 @@ def test_standard_logging_payload(model, turn_off_message_logging):
)
)
keys_list = list(StandardLoggingPayload.__annotations__.keys())
keys_list = list(StandardLoggingPayload.__required_keys__)
for k in keys_list:
assert (
@ -1190,7 +1190,7 @@ def test_standard_logging_payload_audio(turn_off_message_logging, stream):
)
)
keys_list = list(StandardLoggingPayload.__annotations__.keys())
keys_list = list(StandardLoggingPayload.__required_keys__)
for k in keys_list:
assert (

View file

@ -270,7 +270,7 @@ async def test_datadog_logging_http_request():
message = json.loads(body[0]["message"])
print("logged message", json.dumps(message, indent=4))
expected_message_fields = StandardLoggingPayload.__annotations__.keys()
expected_message_fields = StandardLoggingPayload.__required_keys__
for field in expected_message_fields:
assert field in message, f"Field '{field}' is missing from the message"

View file

@ -3,6 +3,7 @@ import contextlib
import datetime
import os
import sys
from collections.abc import Callable
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@ -6945,3 +6946,61 @@ def test_classifier_audit_is_not_added_to_other_calls(logging_obj, call_type, or
logging_obj.model_call_details["litellm_params"] = {"metadata": {"internal_call_origin": origin}}
logging_obj.pre_call(input=[], api_key=None, additional_args={"complete_input_dict": {"input": "embedding"}})
assert logging_obj.classifier_input is None
def _run_while_a_thread_grows(target: dict, read: Callable[[], None], reads: int) -> None:
import itertools
import threading
stop: Final = threading.Event()
def grow() -> None:
for counter in itertools.count():
if stop.is_set():
return
key: Final = f"late_{counter % 64}"
if key in target:
del target[key]
else:
target[key] = counter
writer: Final = threading.Thread(target=grow, daemon=True)
previous_interval: Final = sys.getswitchinterval()
sys.setswitchinterval(1e-6)
writer.start()
try:
for _ in range(reads):
read()
finally:
stop.set()
writer.join(timeout=5)
sys.setswitchinterval(previous_interval)
def test_merge_litellm_metadata_survives_a_thread_growing_metadata_mid_merge():
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
metadata: Final = {f"key_{i}": i for i in range(2000)}
litellm_params: Final = {"metadata": metadata, "litellm_metadata": {"model_group": "gpt"}}
def read() -> None:
merged: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params)
assert merged["key_1999"] == 1999
assert merged["model_group"] == "gpt"
_run_while_a_thread_grows(metadata, read, reads=300)
def test_get_additional_headers_survives_a_thread_growing_headers_mid_copy():
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
headers: Final = {f"llm_provider-x-custom-{i}": str(i) for i in range(2000)}
headers["x-ratelimit-remaining-requests"] = "7"
def read() -> None:
copied: Final = StandardLoggingPayloadSetup.get_additional_headers(headers)
assert copied is not None
assert copied["x_ratelimit_remaining_requests"] == 7
assert copied["llm_provider-x-custom-1999"] == "1999"
_run_while_a_thread_grows(headers, read, reads=300)

View file

@ -1554,3 +1554,52 @@ def test_stream_chunk_builder_reads_role_from_first_frame_with_choices() -> None
assert response is not None
assert response.choices[0].message.role == "user"
assert response.choices[0].message.content == "Hi"
def _fail_prompt_token_count() -> int:
raise AssertionError("prompt tokens must come from the usage chunk, not the tokenizer")
def test_calculate_usage_reads_prompt_tokens_from_mock_stream_usage_chunk_without_tokenizer_fallback() -> None:
from litellm.utils import mock_completion_streaming_obj
chunks: Final = list(
mock_completion_streaming_obj(
ModelResponseStream(model="gpt-5.4-mini"),
mock_response="ok",
model="gpt-5.4-mini",
prompt_tokens=51234,
)
)
assert chunks[-1].choices == []
usage: Final = ChunkProcessor(chunks=chunks).calculate_usage(
chunks=chunks,
model="gpt-5.4-mini",
completion_output="ok",
count_prompt_tokens=_fail_prompt_token_count,
)
assert usage.prompt_tokens == 51234
assert usage.completion_tokens == chunks[-1].usage.completion_tokens
assert usage.total_tokens == 51234 + usage.completion_tokens
def test_calculate_usage_falls_back_to_prompt_counter_when_mock_stream_has_no_admission_count() -> None:
from litellm.utils import mock_completion_streaming_obj
chunks: Final = list(
mock_completion_streaming_obj(
ModelResponseStream(model="gpt-5.4-mini"), mock_response="ok", model="gpt-5.4-mini"
)
)
assert all(chunk.choices for chunk in chunks)
usage: Final = ChunkProcessor(chunks=chunks).calculate_usage(
chunks=chunks,
model="gpt-5.4-mini",
completion_output="ok",
count_prompt_tokens=lambda: 77,
)
assert usage.prompt_tokens == 77

View file

@ -1918,3 +1918,30 @@ async def test_post_call_success_hook_leaves_raw_provider_dict_untouched():
)
assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("team_metadata", "expected_priority_header"),
[
({"priority": "优先"}, None),
({"priority": "high"}, "high"),
({}, "default"),
],
)
async def test_post_call_success_hook_priority_header_is_always_http_encodable(team_metadata, expected_priority_header):
from starlette.responses import Response
handler = DynamicRateLimitHandler(internal_usage_cache=DualCache())
response = {"id": "msg_123", "type": "message", "role": "assistant", "content": [], "_hidden_params": {}}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(team_id="team-1", team_metadata=team_metadata),
response=response,
)
additional_headers = response["_hidden_params"]["additional_headers"]
http_response = Response(headers={key: str(value) for key, value in additional_headers.items()})
assert http_response.headers.get("x-litellm-priority") == expected_priority_header
assert http_response.headers["x-litellm-rate-limiter-version"] == "v3"

View file

@ -2676,6 +2676,26 @@ class TestEncryptedTaskClassifier:
assert "source-secret" not in json.dumps(call)
assert "originating_request_masked" not in call["proxy_server_request"]["body"]
@pytest.mark.asyncio
async def test_claude_code_encrypted_task_omits_caller_instructions(self):
router, dependency = _native_classifier_router()
task: Final = _encrypted_agent_task()
request: Final = {
"input": [task],
"instructions": "CLAUDE_CODE_SYSTEM",
"litellm_metadata": {"user_agent": "claude-cli/2.1.233"},
}
original: Final = deepcopy(request)
result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs=request)
assert result.routing_decision["cause"] == "llm_classifier"
assert request == original
call: Final = dependency.aresponses.call_args.kwargs
assert call["instructions"] == classification_system_prompt(router.config.classifier_context_window_size)
assert "CLAUDE_CODE_SYSTEM" not in json.dumps(call["input"][:-1])
assert call["input"][-1] == task
@pytest.mark.asyncio
@pytest.mark.parametrize(
"items",
@ -7994,6 +8014,113 @@ _CODEX_ENVELOPES: Final = (
class TestContextAwareClassifier:
"""Test the new classifier context window and trajectory signals."""
@pytest.mark.asyncio
@pytest.mark.parametrize(
"request_metadata,forwards_system",
[
({"metadata": {"user_agent": "claude-cli/2.1.233"}}, False),
({"litellm_metadata": {"user_agent": "claude-code/2.1.233"}}, False),
({"metadata": {"user_agent": "curl/8.7.1"}}, True),
({"litellm_metadata": {}}, True),
(
{"metadata": {"user_agent": "claude-cli/2.1.233"}, "litellm_metadata": {"user_agent": "curl/8.7.1"}},
False,
),
({"metadata": {"user_agent": "Claude-Code/2.1.233"}}, True),
],
)
async def test_claude_code_classifier_omits_harness_system_prompt(
self,
llm_classifier_config: dict[str, object],
request_metadata: dict[str, object],
forwards_system: bool,
) -> None:
dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "COMPLEX"}')))
router: Final = ComplexityRouter(
"test-complexity-router",
dependency,
{
**llm_classifier_config,
"classifier_context_include_assistant_turns": True,
},
)
messages: Final = [
{"role": "user", "content": "Design the retry state machine"},
{"role": "assistant", "content": "The design needs a lease and fencing token"},
{"role": "user", "content": "Now prove it cannot livelock"},
{
"role": "system",
"content": [{"type": "text", "text": "ENVIRONMENT_CATALOG\nAGENT_CATALOG\nSKILL_CATALOG"}],
},
]
top_level_system: Final = [{"type": "text", "text": "TOP_LEVEL_HARNESS_SYSTEM"}]
claude_kwargs: Final = {
"metadata": {"user_agent": "claude-cli/2.1.233"},
"system": top_level_system,
"proxy_server_request": {"body": {"system": top_level_system}},
}
compared_kwargs: Final = {
**request_metadata,
"system": top_level_system,
"proxy_server_request": {"body": {"system": top_level_system}},
}
original_messages: Final = deepcopy(messages)
original_kwargs: Final = deepcopy((claude_kwargs, compared_kwargs))
results: Final = (
await router.async_pre_routing_hook("test-complexity-router", claude_kwargs, messages),
await router.async_pre_routing_hook("test-complexity-router", compared_kwargs, messages),
)
assert all(result is not None and result.routing_decision["cause"] == "llm_classifier" for result in results)
assert all(result is not None and result.messages == original_messages for result in results)
assert messages == original_messages
assert (claude_kwargs, compared_kwargs) == original_kwargs
calls: Final = tuple(call.kwargs["messages"] for call in dependency.acompletion.await_args_list)
assert calls[0][0]["content"] == calls[1][0]["content"] == classification_system_prompt(
router.config.classifier_context_window_size
)
payloads: Final = (calls[0][1]["content"], calls[1][1]["content"])
for payload, expected_system in zip(payloads, (False, forwards_system)):
assert payload.endswith("Classify this message:\nNow prove it cannot livelock")
assert ("ENVIRONMENT_CATALOG" in payload) is expected_system
assert ("AGENT_CATALOG" in payload) is expected_system
assert ("SKILL_CATALOG" in payload) is expected_system
assert "Design the retry state machine" in payload
assert "lease and fencing token" in payload
assert "TOP_LEVEL_HARNESS_SYSTEM" not in payload
assert "Conversation so far: ~35 tokens across the request" in payload
@pytest.mark.asyncio
async def test_claude_code_first_turn_without_context_omits_harness_system_prompt(
self, llm_classifier_config: dict[str, object]
) -> None:
dependency: Final = MagicMock(acompletion=AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')))
router: Final = ComplexityRouter(
"test-complexity-router",
dependency,
{**llm_classifier_config, "classifier_context_window_size": 0},
)
messages: Final = [
{"role": "user", "content": "What is two plus two?"},
{
"role": "system",
"content": [{"type": "text", "text": "ENVIRONMENT_CATALOG\nAGENT_CATALOG\nSKILL_CATALOG"}],
},
]
request_kwargs: Final = {"litellm_metadata": {"user_agent": "claude-code/2.1.233"}}
original: Final = deepcopy((messages, request_kwargs))
result: Final = await router.async_pre_routing_hook("test-complexity-router", request_kwargs, messages)
assert result is not None and result.routing_decision["cause"] == "llm_classifier"
assert result.messages == messages == original[0]
assert request_kwargs == original[1]
classifier_messages: Final = dependency.acompletion.call_args.kwargs["messages"]
assert classifier_messages[0]["content"] == classification_system_prompt(
router.config.classifier_context_window_size
)
assert classifier_messages[1]["content"].strip() == "Classify this message:\nWhat is two plus two?"
@pytest.mark.parametrize(
"tail,expected",
(

View file

@ -2421,6 +2421,253 @@ def test_mock_completion_usage_falls_back_to_default_without_admission_count():
assert response.usage.prompt_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT
_ADMISSION_INPUT_TOKENS: Final = 51234
def _admission_metadata(input_tokens: int) -> dict[str, object]: # mutable-ok: logging writes into metadata
return {"user_api_key_budget_reservation": {"reserved_cost": 1.0, "input_tokens": input_tokens}}
_ADMISSION_METADATA: Final = _admission_metadata(_ADMISSION_INPUT_TOKENS)
_MOCK_STREAM_MESSAGES: Final = [{"role": "user", "content": "hello " * 200}]
_STREAM_CHUNK_BUILDER_TOKEN_COUNTER: Final = "litellm.litellm_core_utils.streaming_chunk_builder_utils.token_counter"
def _prompt_token_counter_calls(token_counter: MagicMock) -> list[object]:
return [call for call in token_counter.call_args_list if call.kwargs.get("messages") is not None]
def _client_usage_chunks(chunks: list[ModelResponseStream]) -> list[Usage]:
return [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None]
@pytest.mark.parametrize("n", (None, 2))
def test_mock_completion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(n: int | None):
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
chunks: Final = list(
litellm.completion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
stream=True,
n=n,
stream_options={"include_usage": True},
metadata=_ADMISSION_METADATA,
)
)
usage_chunks: Final = _client_usage_chunks(chunks)
assert len(usage_chunks) == 1
assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS
assert usage_chunks[0].completion_tokens == litellm_main.DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT
assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens
assert _prompt_token_counter_calls(token_counter) == []
assert all(chunk.choices for chunk in chunks[:-1])
assert {chunk.id for chunk in chunks} == {chunks[0].id}
@pytest.mark.asyncio
@pytest.mark.parametrize("n", (None, 2))
async def test_mock_acompletion_stream_usage_reports_admission_input_tokens_without_tokenizer_fallback(
n: int | None,
):
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
response: Final = await litellm.acompletion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
stream=True,
n=n,
stream_options={"include_usage": True},
litellm_metadata=_ADMISSION_METADATA,
)
chunks: Final = [chunk async for chunk in response]
usage_chunks: Final = _client_usage_chunks(chunks)
assert len(usage_chunks) == 1
assert usage_chunks[0].prompt_tokens == _ADMISSION_INPUT_TOKENS
assert usage_chunks[0].total_tokens == _ADMISSION_INPUT_TOKENS + usage_chunks[0].completion_tokens
assert _prompt_token_counter_calls(token_counter) == []
assert all(chunk.choices for chunk in chunks[:-1])
assert {chunk.id for chunk in chunks} == {chunks[0].id}
def test_mock_completion_stream_without_include_usage_hides_usage_chunk_but_logs_admission_count():
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
chunks: Final = list(
litellm.completion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
stream=True,
metadata=_ADMISSION_METADATA,
)
)
assert _client_usage_chunks(chunks) == []
assert all(len(chunk.choices) == 1 for chunk in chunks)
assert chunks[-1]._hidden_params["usage"].prompt_tokens == _ADMISSION_INPUT_TOKENS
assert _prompt_token_counter_calls(token_counter) == []
def test_mock_completion_stream_with_empty_stream_options_completes_and_logs_admission_count():
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
chunks: Final = list(
litellm.completion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
stream=True,
stream_options={},
metadata=_ADMISSION_METADATA,
)
)
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok"
assert _client_usage_chunks(chunks) == []
assert _prompt_token_counter_calls(token_counter) == []
@pytest.mark.asyncio
async def test_mock_acompletion_stream_with_empty_stream_options_completes_and_logs_admission_count():
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
response: Final = await litellm.acompletion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
stream=True,
stream_options={},
litellm_metadata=_ADMISSION_METADATA,
)
chunks: Final = [chunk async for chunk in response]
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok"
assert _client_usage_chunks(chunks) == []
assert _prompt_token_counter_calls(token_counter) == []
def test_mock_completion_stream_without_admission_count_falls_back_to_tokenizer():
expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES)
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
chunks: Final = list(
litellm.completion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
stream=True,
stream_options={"include_usage": True},
metadata={"user_api_key_budget_reservation": {"reserved_cost": 1.0}},
)
)
usage_chunks: Final = _client_usage_chunks(chunks)
assert len(usage_chunks) == 1
assert usage_chunks[0].prompt_tokens == expected_prompt_tokens
assert usage_chunks[0].total_tokens == expected_prompt_tokens + usage_chunks[0].completion_tokens
assert len(_prompt_token_counter_calls(token_counter)) >= 1
@pytest.mark.asyncio
async def test_mock_acompletion_stream_without_admission_count_falls_back_to_tokenizer():
expected_prompt_tokens: Final = litellm.token_counter(model="openai/gpt-5.4-mini", messages=_MOCK_STREAM_MESSAGES)
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
response: Final = await litellm.acompletion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
stream=True,
stream_options={"include_usage": True},
)
chunks: Final = [chunk async for chunk in response]
usage_chunks: Final = _client_usage_chunks(chunks)
assert len(usage_chunks) == 1
assert usage_chunks[0].prompt_tokens == expected_prompt_tokens
assert len(_prompt_token_counter_calls(token_counter)) >= 1
def _usage_triple(usage: Usage) -> tuple[int, int, int]:
return (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens)
@pytest.mark.parametrize("input_tokens", (_ADMISSION_INPUT_TOKENS, 0))
def test_mock_completion_stream_and_non_stream_report_the_same_admission_usage(input_tokens: int):
metadata: Final = _admission_metadata(input_tokens)
non_stream: Final = litellm.completion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
metadata=metadata,
)
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
chunks: Final = list(
litellm.completion(
model="openai/gpt-5.4-mini",
messages=_MOCK_STREAM_MESSAGES,
mock_response="ok",
api_key="mock",
stream=True,
stream_options={"include_usage": True},
metadata=metadata,
)
)
assert _usage_triple(non_stream.usage) == _usage_triple(_client_usage_chunks(chunks)[0])
assert non_stream.usage.prompt_tokens == input_tokens
assert _prompt_token_counter_calls(token_counter) == []
@pytest.mark.asyncio
async def test_mock_acompletion_stream_reports_zero_admission_input_tokens_without_tokenizer_fallback():
with patch(_STREAM_CHUNK_BUILDER_TOKEN_COUNTER, wraps=litellm.token_counter) as token_counter:
response: Final = await litellm.acompletion(
model="openai/gpt-5.4-mini",
messages=[{"role": "user", "content": ""}],
mock_response="ok",
api_key="mock",
stream=True,
stream_options={"include_usage": True},
litellm_metadata=_admission_metadata(0),
)
chunks: Final = [chunk async for chunk in response]
usage_chunks: Final = _client_usage_chunks(chunks)
assert len(usage_chunks) == 1
assert _usage_triple(usage_chunks[0]) == (0, usage_chunks[0].completion_tokens, usage_chunks[0].completion_tokens)
assert _prompt_token_counter_calls(token_counter) == []
def test_mock_text_completion_stream_and_non_stream_report_the_same_zero_admission_usage():
metadata: Final = _admission_metadata(0)
non_stream: Final = litellm.text_completion(
model="openai/gpt-5.4-mini", prompt="", mock_response="ok", api_key="mock", metadata=metadata
)
chunks: Final = list(
litellm.text_completion(
model="openai/gpt-5.4-mini",
prompt="",
mock_response="ok",
api_key="mock",
stream=True,
stream_options={"include_usage": True},
metadata=metadata,
)
)
stream_usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None)
assert len(stream_usages) == 1
assert _usage_triple(non_stream.usage) == _usage_triple(stream_usages[0])
assert non_stream.usage.prompt_tokens == 0
def test_mock_completion_stream_with_model_response():
"""Test that mock_completion correctly handles stream=True with a ModelResponse as mock_response."""
from litellm import completion

View file

@ -1,9 +1,12 @@
import asyncio
import contextlib
import json
import logging
import os
import queue
import threading
from datetime import datetime, timedelta, timezone
from collections.abc import Iterator
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -15,6 +18,7 @@ from jsonschema import validate
import litellm
from litellm._internal_context import is_internal_call
from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT
from litellm._logging import (
CorrelationContextFilter,
JsonFormatter,
@ -6208,3 +6212,135 @@ def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_with
"api_key": "sk-from-db",
}
assert _credential_warnings(caplog) == []
_MOCK_STREAM_ID: Final = "chatcmpl-mock-stream"
_ChunkSnapshot = tuple[str, tuple[str | None, ...], Usage | None]
def _snapshot(chunk: ModelResponseStream) -> _ChunkSnapshot:
return chunk.id, tuple(choice.delta.content for choice in chunk.choices), getattr(chunk, "usage", None)
def _mock_stream_snapshots(mock_response: object, prompt_tokens: int | None) -> list[_ChunkSnapshot]:
from litellm.utils import mock_completion_streaming_obj
return [
_snapshot(chunk)
for chunk in mock_completion_streaming_obj(
ModelResponseStream(id=_MOCK_STREAM_ID, model="gpt-5.4-mini"),
mock_response=mock_response,
model="gpt-5.4-mini",
prompt_tokens=prompt_tokens,
)
]
async def _async_mock_stream_snapshots(mock_response: object, prompt_tokens: int | None) -> list[_ChunkSnapshot]:
from litellm.utils import async_mock_completion_streaming_obj
return [
_snapshot(chunk)
async for chunk in async_mock_completion_streaming_obj(
ModelResponseStream(id=_MOCK_STREAM_ID, model="gpt-5.4-mini"),
mock_response=mock_response,
model="gpt-5.4-mini",
prompt_tokens=prompt_tokens,
)
]
_CONTENT_SNAPSHOTS: Final = [(_MOCK_STREAM_ID, (content,), None) for content in ("hel", "lo ", "wor", "ld")]
def _assert_trailing_usage_chunk(snapshots: list[_ChunkSnapshot], prompt_tokens: int) -> None:
assert snapshots[:-1] == _CONTENT_SNAPSHOTS
chunk_id, choices, usage = snapshots[-1]
assert chunk_id == _MOCK_STREAM_ID
assert choices == ()
assert usage is not None
assert usage.prompt_tokens == prompt_tokens
assert usage.completion_tokens == DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT
assert usage.total_tokens == prompt_tokens + usage.completion_tokens
@pytest.mark.parametrize("prompt_tokens", (51234, 0))
def test_mock_completion_streaming_obj_emits_usage_chunk_with_admission_prompt_tokens(prompt_tokens: int) -> None:
_assert_trailing_usage_chunk(_mock_stream_snapshots("hello world", prompt_tokens), prompt_tokens)
@pytest.mark.asyncio
@pytest.mark.parametrize("prompt_tokens", (51234, 0))
async def test_async_mock_completion_streaming_obj_emits_usage_chunk_with_admission_prompt_tokens(
prompt_tokens: int,
) -> None:
_assert_trailing_usage_chunk(await _async_mock_stream_snapshots("hello world", prompt_tokens), prompt_tokens)
def test_mock_completion_streaming_obj_emits_no_usage_chunk_without_admission_prompt_tokens() -> None:
assert _mock_stream_snapshots("hello world", None) == _CONTENT_SNAPSHOTS
@pytest.mark.asyncio
async def test_async_mock_completion_streaming_obj_emits_no_usage_chunk_without_admission_prompt_tokens() -> None:
assert await _async_mock_stream_snapshots("hello world", None) == _CONTENT_SNAPSHOTS
def test_mock_completion_streaming_obj_passes_prebuilt_stream_chunk_through_without_usage_chunk() -> None:
prebuilt: Final = ModelResponseStream(
model="gpt-5.4-mini", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="prebuilt"))]
)
assert _mock_stream_snapshots(prebuilt, 51234) == [(prebuilt.id, ("prebuilt",), None)]
@pytest.mark.asyncio
async def test_async_mock_completion_streaming_obj_raises_mock_exception_before_usage_chunk() -> None:
mock_exception: Final = litellm.MockException(
status_code=500, message="boom", llm_provider="openai", model="gpt-5.4-mini"
)
with pytest.raises(litellm.MockException):
await _async_mock_stream_snapshots(mock_exception, 51234)
@contextlib.contextmanager
def _recording_hidden_params_at_submit(submit_target: str) -> "Iterator[queue.SimpleQueue[dict[str, object]]]":
seen: Final = queue.SimpleQueue()
def record_submit(_fn, *args, **_kwargs):
response: Final = next(arg for arg in args if isinstance(arg, litellm.ModelResponse))
seen.put(dict(response._hidden_params))
return MagicMock()
with patch(submit_target, side_effect=record_submit):
yield seen
@pytest.mark.asyncio
async def test_acompletion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread(monkeypatch):
monkeypatch.setattr(litellm, "success_callback", [lambda kwargs, response, start_time, end_time: None])
with _recording_hidden_params_at_submit("litellm.litellm_core_utils.litellm_logging.executor.submit") as seen:
await litellm.acompletion(
model="gpt-5.5",
messages=[{"role": "user", "content": "hi"}],
mock_response="Hello there!",
num_retries=0,
)
snapshot: Final = seen.get_nowait()
assert snapshot["litellm_call_id"]
assert snapshot["response_cost"] is not None
assert snapshot["api_base"]
def test_completion_finishes_response_metadata_before_handing_the_response_to_the_logging_thread():
with _recording_hidden_params_at_submit("litellm.utils.executor.submit") as seen:
litellm.completion(
model="gpt-5.5",
messages=[{"role": "user", "content": "hi"}],
mock_response="Hello there!",
)
snapshot: Final = seen.get_nowait()
assert snapshot["litellm_call_id"]
assert snapshot["response_cost"] is not None
assert snapshot["api_base"]

View file

@ -54,7 +54,7 @@ function ResourceBadge({
fallback,
}: {
resource: AccessGroupResource;
href: string;
href?: string;
fallback: (id: string) => string;
}) {
const badge = (

View file

@ -473,6 +473,92 @@ it("should display 'Default Proxy Admin' for user_id when value is 'default_user
});
});
describe("entity links out of the key rows", () => {
const keyRow = async () => (await screen.findByText("Test Key Alias")).closest("tr") as HTMLElement;
const enableColumn = async (user: ReturnType<typeof userEvent.setup>, title: string) => {
await user.click(screen.getByRole("button", { name: "Columns" }));
await user.click(await screen.findByText(title));
await user.keyboard("{Escape}");
};
const enableCreatedByColumn = (user: ReturnType<typeof userEvent.setup>) => enableColumn(user, "Created By");
it("points the User and Team cells at their detail pages", async () => {
renderWithProviders(<VirtualKeysTable />);
const row = await keyRow();
expect(within(row).getByRole("link", { name: "user@example.com" })).toHaveAttribute(
"href",
"/ui/users?user=user-1",
);
expect(within(row).getByRole("link", { name: "Test Team" })).toHaveAttribute("href", "/ui/teams?team=team-1");
});
it("points the Organization cell at the org's detail page", async () => {
mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, org_id: "org-1" }]));
const user = userEvent.setup();
renderWithProviders(<VirtualKeysTable />);
await enableColumn(user, "Organization");
const row = await keyRow();
expect(within(row).getByRole("link", { name: "Test Organization" })).toHaveAttribute(
"href",
"/ui/organizations?org=org-1",
);
});
it("points the Created By cell at the creator's detail page", async () => {
mockUseKeys.mockReturnValue(
keysResult([
{
...mockKey,
created_by: "creator-1",
created_by_user: { user_id: "creator-1", user_email: "creator@example.com", user_alias: "The Creator" },
},
]),
);
const user = userEvent.setup();
renderWithProviders(<VirtualKeysTable />);
await enableCreatedByColumn(user);
const row = await keyRow();
expect(within(row).getByRole("link", { name: "The Creator" })).toHaveAttribute("href", "/ui/users?user=creator-1");
});
it("leaves the default_user_id placeholder unlinked even once it resolves to a named user", async () => {
const placeholder = { user_id: "default_user_id", user_email: "admin@example.com", user_alias: "Proxy Admin" };
mockUseKeys.mockReturnValue(
keysResult([
{
...mockKey,
user_id: placeholder.user_id,
user_email: placeholder.user_email,
user: placeholder,
created_by: placeholder.user_id,
created_by_user: placeholder,
},
]),
);
const user = userEvent.setup();
renderWithProviders(<VirtualKeysTable />);
await enableCreatedByColumn(user);
const row = await keyRow();
expect(within(row).getAllByText("Proxy Admin")).toHaveLength(2);
expect(within(row).queryByRole("link", { name: "Proxy Admin" })).not.toBeInTheDocument();
});
it("leaves the litellm-dashboard session team unlinked", async () => {
mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, team_id: "litellm-dashboard" }]));
renderWithProviders(<VirtualKeysTable />);
const row = await keyRow();
expect(within(row).getByText("litellm-dashboard")).toBeInTheDocument();
expect(within(row).queryByRole("link", { name: "litellm-dashboard" })).not.toBeInTheDocument();
});
});
it("should render table without crashing when models is null", async () => {
mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, models: null as unknown as string[] }]));

View file

@ -9,15 +9,17 @@ import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/h
import { Skeleton } from "@/components/ui/skeleton";
import {
DateCell,
ENTITY_CELL_TITLE_CLASSES,
IdCell,
IdentityCell,
ModelsCell,
SpendBudgetCell,
StatusBadge,
UserPopoverCell,
type StatusTone,
} from "@/components/shared/table_cells";
import { orgDetailHref, teamDetailHref } from "@/utils/entityLinks";
import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag";
import { KeyResponse, Team } from "../key_team_helpers/key_list";
import { Organization } from "../networking";
@ -62,67 +64,6 @@ const getKeyStatus = (key: KeyResponse): KeyStatus => {
};
};
const UserPopoverCell = ({
userAlias,
userEmail,
userId,
width,
}: {
userAlias: string | null;
userEmail: string | null;
userId: string | null;
width: number;
}) => {
const displayValue = userAlias || userEmail || userId;
const isDefaultAdmin = userId === "default_user_id";
const popoverContent = (
<div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]">
{[
{ label: "User Alias", value: userAlias },
{ label: "User Email", value: userEmail },
{ label: "User ID", value: userId },
].map(({ label, value }) => (
<div key={label} className="flex flex-col min-w-0">
<span className="text-muted-foreground">{label}</span>
{value ? (
<IdCell value={value} variant="plain" copyable className="max-w-full" />
) : (
<span className="font-mono">-</span>
)}
</div>
))}
</div>
);
if (isDefaultAdmin && !userAlias && !userEmail) {
return (
<HoverCard>
<HoverCardTrigger render={<span className="cursor-default" />}>
<DefaultProxyAdminTag userId={userId} />
</HoverCardTrigger>
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
</HoverCard>
);
}
return (
<HoverCard>
<HoverCardTrigger
render={
<span
className="font-mono text-xs truncate block cursor-default"
style={{ maxWidth: width, overflow: "hidden" }}
/>
}
>
{displayValue || "-"}
</HoverCardTrigger>
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
</HoverCard>
);
};
const InfoHeader = ({ label, tooltip }: { label: string; tooltip: string }) => (
<span className="flex items-center gap-1">
{label}
@ -201,12 +142,12 @@ export const getKeyTableColumns = ({
const teamId = info.getValue() as string | null;
if (!teamId) return "-";
const team = allTeams.find((t) => t.team_id === teamId);
const displayValue = team?.team_alias || teamId;
const width = info.cell.column.getSize();
return (
<span className="font-mono text-xs truncate block" style={{ maxWidth: width, overflow: "hidden" }}>
{displayValue}
</span>
<IdentityCell
title={team?.team_alias || teamId}
titleClassName={ENTITY_CELL_TITLE_CLASSES}
href={teamDetailHref(teamId)}
/>
);
},
},
@ -221,12 +162,12 @@ export const getKeyTableColumns = ({
const orgId = info.getValue() as string | null;
if (!orgId) return "-";
const org = organizations.find((o) => o.organization_id === orgId);
const displayValue = org?.organization_alias || orgId;
const width = info.cell.column.getSize();
return (
<span className="font-mono text-xs truncate block" style={{ maxWidth: width, overflow: "hidden" }}>
{displayValue}
</span>
<IdentityCell
title={org?.organization_alias || orgId}
titleClassName={ENTITY_CELL_TITLE_CLASSES}
href={orgDetailHref(orgId)}
/>
);
},
},

View file

@ -1,13 +1,12 @@
import { Badge } from "@/components/ui/badge";
const DEFAULT_USER_ID = "default_user_id";
import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels";
interface DefaultProxyAdminTagProps {
userId: string | null | undefined;
}
export default function DefaultProxyAdminTag({ userId }: DefaultProxyAdminTagProps) {
if (userId === DEFAULT_USER_ID) {
if (userId === DEFAULT_PROXY_ADMIN_USER_ID) {
return <Badge variant="secondary">Default Proxy Admin</Badge>;
}

View file

@ -2,6 +2,7 @@ import React from "react";
import CopyButton from "@/components/shared/CopyButton";
import { EntityLink } from "@/components/shared/EntityLink";
import { cx } from "@/lib/cva.config";
import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels";
import DefaultProxyAdminTag from "./DefaultProxyAdminTag";
interface LabeledFieldProps {
@ -24,7 +25,7 @@ export default function LabeledField({
defaultUserIdCheck = false,
}: LabeledFieldProps) {
const isEmpty = !value;
const isDefaultUser = defaultUserIdCheck && value === "default_user_id";
const isDefaultUser = defaultUserIdCheck && value === DEFAULT_PROXY_ADMIN_USER_ID;
const displayValue = isEmpty ? "-" : value;
const isCopyable = copyable && !isEmpty && !isDefaultUser;
const isLink = href != null && !isEmpty && !isDefaultUser;

View file

@ -25,6 +25,12 @@ describe("EntityLink", () => {
expect(push).toHaveBeenCalledWith("/ui/users?user=u1");
});
it("renders the label as plain text when there is no href to point at", () => {
render(<EntityLink>default_user_id</EntityLink>);
expect(screen.queryByRole("link")).not.toBeInTheDocument();
expect(screen.getByText("default_user_id")).toBeInTheDocument();
});
it("leaves modified clicks to the browser so new-tab shortcuts keep working", async () => {
const user = userEvent.setup();
render(<EntityLink href="/ui/users?user=u1">alice</EntityLink>);

View file

@ -19,12 +19,24 @@ export function useEntityLinkClick(href: string): (e: React.MouseEvent) => void
}
interface EntityLinkProps {
href: string;
href?: string;
className?: string;
children: React.ReactNode;
}
export function EntityLink({ href, className, children }: EntityLinkProps) {
if (!href) {
return <span className={cn("inline-block min-w-0 max-w-full truncate font-semibold", className)}>{children}</span>;
}
return (
<LinkedEntity href={href} className={className}>
{children}
</LinkedEntity>
);
}
function LinkedEntity({ href, className, children }: EntityLinkProps & { href: string }) {
const handleClick = useEntityLinkClick(href);
return (

View file

@ -0,0 +1,62 @@
"use client";
import DefaultProxyAdminTag from "@/components/common_components/DefaultProxyAdminTag";
import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card";
import { userDetailHref } from "@/utils/entityLinks";
import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels";
import { IdCell } from "./id_cell";
import { IdentityCell } from "./identity_cell";
export const ENTITY_CELL_TITLE_CLASSES = "font-mono text-xs font-normal";
interface UserPopoverCellProps {
userAlias: string | null;
userEmail: string | null;
userId: string | null;
width: number;
}
export function UserPopoverCell({ userAlias, userEmail, userId, width }: UserPopoverCellProps) {
const displayValue = userAlias || userEmail || userId;
const isDefaultAdmin = userId === DEFAULT_PROXY_ADMIN_USER_ID;
const popoverContent = (
<div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]">
{[
{ label: "User Alias", value: userAlias },
{ label: "User Email", value: userEmail },
{ label: "User ID", value: userId },
].map(({ label, value }) => (
<div key={label} className="flex flex-col min-w-0">
<span className="text-muted-foreground">{label}</span>
{value ? (
<IdCell value={value} variant="plain" copyable copyLabel={`Copy ${label}`} className="max-w-full" />
) : (
<span className="font-mono">-</span>
)}
</div>
))}
</div>
);
const trigger =
isDefaultAdmin && !userAlias && !userEmail ? (
<DefaultProxyAdminTag userId={userId} />
) : (
<IdentityCell
title={displayValue || "-"}
titleClassName={ENTITY_CELL_TITLE_CLASSES}
href={userId ? userDetailHref(userId) : undefined}
/>
);
return (
<HoverCard>
<HoverCardTrigger render={<span className="block" style={{ maxWidth: width, overflow: "hidden" }} />}>
{trigger}
</HoverCardTrigger>
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
</HoverCard>
);
}

View file

@ -71,6 +71,14 @@ describe("IdCell", () => {
expect(rowClick).not.toHaveBeenCalled();
});
it("names the copy button after the field it copies", async () => {
const user = userEvent.setup();
render(<IdCell value="alice@example.com" copyable copyLabel="Copy User Email" />);
expect(screen.queryByRole("button", { name: "Copy ID" })).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Copy User Email" }));
expect(copyToClipboardMock).toHaveBeenCalledWith("alice@example.com");
});
it("passes dataTestId through to the id element", () => {
render(<IdCell value="k-1" dataTestId="key-id-cell" />);
expect(screen.getByTestId("key-id-cell")).toHaveTextContent("k-1");

View file

@ -15,6 +15,7 @@ interface IdCellProps {
variant?: IdCellVariant;
onClick?: (value: string) => void;
copyable?: boolean;
copyLabel?: string;
truncate?: boolean;
fallback?: string;
tooltip?: React.ReactNode;
@ -39,6 +40,7 @@ export function IdCell({
variant = "pill",
onClick,
copyable = false,
copyLabel = "Copy ID",
truncate = true,
fallback = "-",
tooltip,
@ -80,7 +82,7 @@ export function IdCell({
{withTooltip}
<button
type="button"
aria-label="Copy ID"
aria-label={copyLabel}
className="shrink-0 cursor-pointer text-muted-foreground hover:text-foreground"
onClick={(event) => {
event.stopPropagation();

View file

@ -13,3 +13,4 @@ export { ModelsCell } from "./models_cell";
export { MoneyCell } from "./money_cell";
export { SpendBudgetCell } from "./spend_budget_cell";
export { StatusBadge, type StatusTone } from "./status_badge";
export { UserPopoverCell, ENTITY_CELL_TITLE_CLASSES } from "./UserPopoverCell";

View file

@ -10,6 +10,8 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({
useKeys: vi.fn(),
}));
vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) }));
vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({
getModelDisplayName: vi.fn((model: string) => model),
}));
@ -384,4 +386,66 @@ describe("TeamVirtualKeysTable", () => {
expect(screen.getByText("Key Info View")).toBeInTheDocument();
});
});
describe("entity links out of the key rows", () => {
const renderRow = async (key: KeyResponse, organization: Organization | null = null) => {
mockUseKeys.mockReturnValue({
data: { keys: [key], total_count: 1, current_page: 1, total_pages: 1 } as KeysResponse,
isPending: false,
isFetching: false,
refetch: vi.fn(),
} as any);
renderWithProviders(<TeamVirtualKeysTable {...defaultProps} organization={organization} />);
await screen.findByText(key.key_alias as string);
return screen.getByRole("row", { name: new RegExp(key.key_alias as string) });
};
it("points the Organization ID cell at the org's detail page", async () => {
const row = await renderRow(createMockKey({ organization_id: null }), mockOrganization);
expect(within(row).getByRole("link", { name: "org-123" })).toHaveAttribute(
"href",
"/ui/organizations?org=org-123",
);
});
it("points the User Email and User ID cells at the owning user's detail page", async () => {
const row = await renderRow(
createMockKey({ user_id: "user-1", user: { user_id: "user-1", user_email: "alice@example.com" } }),
);
expect(within(row).getByRole("link", { name: "alice@example.com" })).toHaveAttribute(
"href",
"/ui/users?user=user-1",
);
expect(within(row).getByRole("link", { name: "user-1" })).toHaveAttribute("href", "/ui/users?user=user-1");
});
it("points the Created By cell at the creator's detail page", async () => {
const row = await renderRow(
createMockKey({
created_by: "creator-1",
created_by_user: { user_id: "creator-1", user_email: "creator@example.com", user_alias: "The Creator" },
}),
);
expect(within(row).getByRole("link", { name: "The Creator" })).toHaveAttribute(
"href",
"/ui/users?user=creator-1",
);
});
it("leaves the default_user_id placeholder unlinked in the User ID and Created By cells", async () => {
const placeholder = { user_id: "default_user_id", user_email: "admin@example.com", user_alias: "Proxy Admin" };
const ownedAndCreatedByPlaceholder = {
user_id: placeholder.user_id,
user: placeholder,
created_by: placeholder.user_id,
created_by_user: placeholder,
};
const row = await renderRow(createMockKey(ownedAndCreatedByPlaceholder));
expect(within(row).getByText("Default Proxy Admin")).toBeInTheDocument();
expect(within(row).getByText("Proxy Admin")).toBeInTheDocument();
expect(within(row).queryByRole("link", { name: "Proxy Admin" })).not.toBeInTheDocument();
expect(within(row).queryByRole("link", { name: placeholder.user_email })).not.toBeInTheDocument();
expect(within(row).queryByRole("link", { name: "Default Proxy Admin" })).not.toBeInTheDocument();
});
});
});

View file

@ -1,8 +1,14 @@
"use client";
import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
import { SimpleTooltip } from "@/components/ui/tooltip";
import CopyButton from "@/components/shared/CopyButton";
import { DateCell, IdCell, MoneyCell } from "@/components/shared/table_cells";
import {
DateCell,
ENTITY_CELL_TITLE_CLASSES,
IdCell,
IdentityCell,
MoneyCell,
UserPopoverCell,
} from "@/components/shared/table_cells";
import {
DataTable,
DataTableFilterDrawer,
@ -11,8 +17,9 @@ import {
DataTableToolbar,
} from "@/components/shared/DataTable";
import { Badge } from "@/components/ui/badge";
import { HoverCard, HoverCardContent, HoverCardTrigger } from "@/components/ui/hover-card";
import { Input } from "@/components/ui/input";
import { orgDetailHref, userDetailHref } from "@/utils/entityLinks";
import { DEFAULT_PROXY_ADMIN_USER_ID } from "@/utils/sentinels";
import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
import { ColumnDef, ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
@ -168,7 +175,15 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
header: "Organization ID",
size: 140,
enableSorting: false,
cell: (info) => (info.getValue() ? info.renderValue() : "-"),
cell: (info) => {
const orgId = info.getValue() as string | null;
if (!orgId) return "-";
return (
<SimpleTooltip content={orgId}>
<IdentityCell title={orgId} titleClassName={ENTITY_CELL_TITLE_CLASSES} href={orgDetailHref(orgId)} />
</SimpleTooltip>
);
},
},
{
id: "user_email",
@ -179,9 +194,14 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
cell: (info) => {
const user = info.getValue() as { user_email?: string } | undefined;
const value = user?.user_email;
const userId = info.row.original.user_id;
return (
<SimpleTooltip content={value}>
<span className="block max-w-full truncate font-mono text-xs">{value ?? "-"}</span>
<IdentityCell
title={value ?? "-"}
titleClassName={ENTITY_CELL_TITLE_CLASSES}
href={value && userId ? userDetailHref(userId) : undefined}
/>
</SimpleTooltip>
);
},
@ -194,10 +214,16 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
enableSorting: false,
cell: (info) => {
const userId = info.getValue() as string | null;
const displayValue = userId === "default_user_id" ? "Default Proxy Admin" : userId;
if (userId === DEFAULT_PROXY_ADMIN_USER_ID) {
return <DefaultProxyAdminTag userId={userId} />;
}
return (
<SimpleTooltip content={displayValue}>
<span className="block max-w-full truncate font-mono text-xs">{displayValue ?? "-"}</span>
<SimpleTooltip content={userId}>
<IdentityCell
title={userId ?? "-"}
titleClassName={ENTITY_CELL_TITLE_CLASSES}
href={userId ? userDetailHref(userId) : undefined}
/>
</SimpleTooltip>
);
},
@ -221,53 +247,13 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
const userId = info.getValue() as string | null;
if (!userId) return "-";
const { created_by_user } = info.row.original;
const userAlias = created_by_user?.user_alias ?? null;
const userEmail = created_by_user?.user_email ?? null;
const isDefaultAdmin = userId === "default_user_id";
const displayValue = userAlias || userEmail || userId;
const popoverContent = (
<div className="flex min-w-[200px] max-w-[300px] flex-col gap-2 text-xs">
{[
{ label: "User Alias", value: userAlias },
{ label: "User Email", value: userEmail },
{ label: "User ID", value: userId },
].map(({ label, value }) => (
<div key={label} className="flex flex-col min-w-0">
<span className="text-muted-foreground">{label}</span>
{value ? (
<span className="flex items-center gap-1">
<span className="min-w-0 flex-1 truncate font-mono text-xs">{value}</span>
<CopyButton value={value} label={`Copy ${label}`} />
</span>
) : (
<span className="font-mono">-</span>
)}
</div>
))}
</div>
);
if (isDefaultAdmin && !userAlias && !userEmail) {
return (
<HoverCard>
<HoverCardTrigger render={<span className="cursor-default" />}>
<DefaultProxyAdminTag userId={userId} />
</HoverCardTrigger>
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
</HoverCard>
);
}
return (
<HoverCard>
<HoverCardTrigger
render={<span className="block max-w-full cursor-default truncate font-mono text-xs" />}
>
{displayValue}
</HoverCardTrigger>
<HoverCardContent align="start">{popoverContent}</HoverCardContent>
</HoverCard>
<UserPopoverCell
userAlias={created_by_user?.user_alias ?? null}
userEmail={created_by_user?.user_email ?? null}
userId={userId}
width={130}
/>
);
},
},

View file

@ -2,7 +2,29 @@ import { describe, expect, it, vi } from "vitest";
vi.mock("@/components/networking", () => ({ serverRootPath: "" }));
import { modelGroupHref } from "./entityLinks";
import { modelGroupHref, teamDetailHref, userDetailHref } from "./entityLinks";
describe("userDetailHref", () => {
it("targets the users page filtered to the encoded user id", () => {
expect(userDetailHref("user-1")).toMatch(/\/users\?user=user-1$/);
expect(userDetailHref("a b/c")).toMatch(/\?user=a%20b%2Fc$/);
});
it("returns no href for the proxy admin placeholder, which has no user page", () => {
expect(userDetailHref("default_user_id")).toBeUndefined();
});
});
describe("teamDetailHref", () => {
it("targets the teams page filtered to the encoded team id", () => {
expect(teamDetailHref("team-1")).toMatch(/\/teams\?team=team-1$/);
expect(teamDetailHref("a b/c")).toMatch(/\?team=a%20b%2Fc$/);
});
it("returns no href for the Admin UI session team, which has no team page", () => {
expect(teamDetailHref("litellm-dashboard")).toBeUndefined();
});
});
describe("modelGroupHref", () => {
it("targets the models page filtered to the encoded model group", () => {

View file

@ -1,3 +1,4 @@
import { DEFAULT_PROXY_ADMIN_USER_ID, UI_TEAM_ID } from "@/utils/sentinels";
import { uiHref } from "@/utils/uiHref";
const MODEL_GRANT_SENTINELS: ReadonlySet<string> = new Set([
@ -6,7 +7,8 @@ const MODEL_GRANT_SENTINELS: ReadonlySet<string> = new Set([
"no-default-models",
]);
export function teamDetailHref(teamId: string): string {
export function teamDetailHref(teamId: string): string | undefined {
if (teamId === UI_TEAM_ID) return undefined;
return `${uiHref("teams")}?team=${encodeURIComponent(teamId)}`;
}
@ -14,7 +16,8 @@ export function keyDetailHref(keyToken: string): string {
return `${uiHref("api-keys")}?key=${encodeURIComponent(keyToken)}`;
}
export function userDetailHref(userId: string): string {
export function userDetailHref(userId: string): string | undefined {
if (userId === DEFAULT_PROXY_ADMIN_USER_ID) return undefined;
return `${uiHref("users")}?user=${encodeURIComponent(userId)}`;
}

View file

@ -0,0 +1,3 @@
export const DEFAULT_PROXY_ADMIN_USER_ID = "default_user_id";
export const UI_TEAM_ID = "litellm-dashboard";

4
uv.lock generated
View file

@ -10,7 +10,7 @@ resolution-markers = [
]
[options]
exclude-newer = "2026-09-06T00:40:30.433549Z"
exclude-newer = "2026-09-07T23:09:03.362777Z"
exclude-newer-span = "P3D"
[manifest]
@ -4772,7 +4772,7 @@ source = { editable = "enterprise" }
[[package]]
name = "litellm-proxy-extras"
version = "0.4.95"
version = "0.4.96"
source = { editable = "litellm-proxy-extras" }
[[package]]