Merge remote-tracking branch 'origin/main' into litellm_model_access_denied_message

This commit is contained in:
yassin 2026-09-15 23:36:30 +00:00
commit 7267c6bed7
93 changed files with 6232 additions and 402 deletions

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.67"
version = "0.1.68"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.67"
version = "0.1.68"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -0,0 +1,18 @@
-- DropIndex
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx";
-- DropIndex
DROP INDEX IF EXISTS "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key";
-- AlterTable
-- NOT NULL DEFAULT '' (not nullable): Postgres unique constraints treat every
-- NULL as distinct, so a nullable column would let multiple unscoped mappings
-- collide on the same claim without a constraint violation. The constant
-- default is a fast, metadata-only backfill for existing rows, not a rewrite.
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD COLUMN IF NOT EXISTS "jwt_issuer" TEXT NOT NULL DEFAULT '';
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_idx" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value", "is_active");
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_JWTKeyMapping_jwt_issuer_jwt_claim_name_jwt_claim_v_key" ON "LiteLLM_JWTKeyMapping"("jwt_issuer", "jwt_claim_name", "jwt_claim_value");

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.97"
version = "0.4.98"
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.97"
version = "0.4.98"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -344,6 +344,7 @@ _anthropic_prompt_caching_ttl_env: Optional[str] = os.getenv("LITELLM_ANTHROPIC_
anthropic_prompt_caching_ttl: Optional[Literal["5m", "1h"]] = (
"1h" if _anthropic_prompt_caching_ttl_env == "1h" else "5m" if _anthropic_prompt_caching_ttl_env == "5m" else None
)
openai_system_messages_first: bool = False
disable_vertex_batch_output_transformation: bool = False
extra_spend_tag_headers: Optional[List[str]] = None
in_memory_llm_clients_cache: "LLMClientCache"

View file

@ -35,6 +35,7 @@ from litellm.litellm_core_utils.logging_utils import (
_assemble_complete_response_from_streaming_chunks,
)
from litellm.types.caching import CachedEmbedding
from litellm.types.integrations.custom_logger import converted_stream_requested
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.rerank import RerankResponse
from litellm.types.utils import (
@ -107,17 +108,31 @@ def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
return "choices" in cached_result
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, object]) -> bool:
def _stream_replay_requested(kwargs: Mapping[str, object]) -> bool:
if kwargs.get("stream", False) is True:
return True
return converted_stream_requested(kwargs) and not kwargs.get("_agentic_loop_depth")
def _should_defer_streaming_cache_hit_callbacks(*, cached_result: object) -> bool:
"""
When stream=True, do not run success callbacks at cache-hit time.
When the cache hit is replayed as a stream, do not run success callbacks at cache-hit time.
Cached chat/text completion replay uses CustomStreamWrapper; cached Responses
replay uses CachedResponsesAPIStreamingIterator; cached Anthropic Messages
replay uses CachedAnthropicMessagesStreamIterator. All invoke logging success
handlers when the stream finishes; firing them here too would double-count
spend and callback records.
spend and callback records. A plain (non-stream) replay logs here, since nothing
else will.
"""
return kwargs.get("stream", False) is True
from litellm.llms.anthropic.experimental_pass_through.messages.response_cache import (
CachedAnthropicMessagesStreamIterator,
)
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
return isinstance(
cached_result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator, CachedAnthropicMessagesStreamIterator)
)
def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") -> Mapping[str, object]:
@ -267,7 +282,7 @@ class LLMCachingHandler:
custom_llm_provider=kwargs.get("custom_llm_provider", None),
args=args,
)
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
# LOG SUCCESS
self._async_log_cache_hit_on_callbacks(
logging_obj=logging_obj,
@ -383,7 +398,7 @@ class LLMCachingHandler:
is_async=False,
)
if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=cached_result,
start_time=start_time,
@ -823,7 +838,7 @@ class LLMCachingHandler:
if (call_type == CallTypes.acompletion.value or call_type == CallTypes.completion.value) and isinstance(
cached_result, dict
):
if kwargs.get("stream", False) is True:
if _stream_replay_requested(kwargs):
cached_result = self._convert_cached_stream_response(
cached_result=cached_result,
call_type=call_type,
@ -838,7 +853,7 @@ class LLMCachingHandler:
if (
call_type == CallTypes.atext_completion.value or call_type == CallTypes.text_completion.value
) and isinstance(cached_result, dict):
if kwargs.get("stream", False) is True:
if _stream_replay_requested(kwargs):
cached_result = self._convert_cached_stream_response(
cached_result=cached_result,
call_type=call_type,
@ -893,7 +908,7 @@ class LLMCachingHandler:
elif (call_type == "aresponses" or call_type == "responses") and isinstance(cached_result, dict):
use_chat_completion_cache: Final = _is_chat_completion_cached_dict(cached_result)
if use_chat_completion_cache:
if kwargs.get("stream", False) is True:
if _stream_replay_requested(kwargs):
bridge_call_type: Final = (
CallTypes.acompletion.value if call_type == "aresponses" else CallTypes.completion.value
)
@ -921,7 +936,7 @@ class LLMCachingHandler:
):
response_obj._hidden_params["cache_hit"] = True
if kwargs.get("stream", False) is True:
if _stream_replay_requested(kwargs):
cached_result = CachedResponsesAPIStreamingIterator(
response=response_obj,
logging_obj=logging_obj,

View file

@ -1777,6 +1777,7 @@ DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_
LENGTH_OF_LITELLM_GENERATED_KEY: Final = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16))
MINIMUM_CUSTOM_KEY_LENGTH: Final = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16))
SECRET_MANAGER_REFRESH_INTERVAL: Final = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400))
OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS: Final = frozenset({"openai", "azure"})
LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"default_internal_user_params",
"default_team_params",
@ -1794,6 +1795,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
# test_general_settings_ui_fields_are_db_overridable enforces that pairing.
"enable_anthropic_prompt_caching",
"anthropic_prompt_caching_ttl",
"openai_system_messages_first",
"max_ui_session_budget",
"budget_rollover",
"mcp_tool_search",

View file

@ -39,6 +39,8 @@ def is_serializable(value):
class LangsmithLogger(CustomBatchLogger):
preserve_events_added_during_flush = True
def __init__(
self,
langsmith_api_key: str | None = None,

View file

@ -238,6 +238,8 @@ def get_llm_provider(
if dynamic_api_key is not None and not isinstance(dynamic_api_key, str):
raise Exception(f"dynamic_api_key needs to be a string. Got type={type(dynamic_api_key).__name__}")
return model, custom_llm_provider, dynamic_api_key, api_base
if "/" in model and is_registered_custom_provider(provider_prefix):
return model.split("/", 1)[1], provider_prefix, dynamic_api_key, api_base
# check if api base is a known openai compatible endpoint
if api_base:
for endpoint in litellm.openai_compatible_endpoints:
@ -536,6 +538,10 @@ def get_llm_provider(
)
def is_registered_custom_provider(custom_llm_provider: str | None) -> bool:
return any(item["provider"] == custom_llm_provider for item in litellm.custom_provider_map)
def _dashscope_family_chat_config(custom_llm_provider: str) -> "litellm.DashScopeChatConfig":
if custom_llm_provider == "qwencloud":
return litellm.QwenCloudChatConfig()

View file

@ -2256,6 +2256,22 @@ def drop_tool_reference_parts_from_tool_messages(
return [_drop_tool_reference_parts(message) for message in messages] # mutable-ok: pipelines mutate message lists
INSTRUCTION_MESSAGE_ROLES: Final = frozenset({"system", "developer"})
def _is_instruction_message(message: AllMessageValues) -> bool:
return message.get("role") in INSTRUCTION_MESSAGE_ROLES
def system_messages_first(
messages: list[AllMessageValues], # mutable-ok: message pipelines type messages as mutable lists
) -> list[AllMessageValues]: # mutable-ok: message pipelines type messages as mutable lists
return [ # mutable-ok: pipelines mutate message lists
*(message for message in messages if _is_instruction_message(message)),
*(message for message in messages if not _is_instruction_message(message)),
]
def _attempt_json_repair(s: str) -> object | None:
"""
Attempt to repair truncated JSON produced by LLM tool calls.

View file

@ -1500,19 +1500,33 @@ def convert_to_gemini_tool_call_result(
return _part
def _sanitize_anthropic_tool_use_id(tool_use_id: str) -> str:
"""
Sanitize tool_use_id to match Anthropic's required pattern: ^[a-zA-Z0-9_-]+$
_TOOL_USE_ID_FALLBACK: Final = "tool_use_id"
_ANTHROPIC_TOOL_USE_ID_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]")
_BEDROCK_TOOL_USE_ID_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_.:-]")
_BEDROCK_TOOL_USE_ID_MAX_LEN: Final = 64
_BEDROCK_TOOL_USE_ID_HASH_LEN: Final = 8
Anthropic requires tool_use_id to only contain alphanumeric characters, underscores, and hyphens.
This function replaces any invalid characters with underscores.
def _replace_invalid_tool_use_id_chars(tool_use_id: str, invalid_chars: re.Pattern[str]) -> str:
return invalid_chars.sub("_", tool_use_id) or _TOOL_USE_ID_FALLBACK
def _sanitize_anthropic_tool_use_id(tool_use_id: str) -> str:
"""Anthropic requires tool_use_id to match ^[a-zA-Z0-9_-]+$."""
return _replace_invalid_tool_use_id_chars(tool_use_id, _ANTHROPIC_TOOL_USE_ID_INVALID_CHARS)
def _sanitize_bedrock_tool_use_id(tool_use_id: str) -> str:
"""
# Replace any character that's not alphanumeric, underscore, or hyphen with underscore
sanitized = re.sub(r"[^a-zA-Z0-9_-]", "_", tool_use_id)
# Ensure it's not empty (fallback to a default if needed)
if not sanitized:
sanitized = "tool_use_id"
return sanitized
Bedrock Converse requires toolUseId to match [a-zA-Z0-9_.:-]+ and be at most 64 chars.
Ids that need rewriting get a short hash of the original appended so two ids that only
differ in a replaced char or past the cut still map to distinct values.
"""
sanitized: Final = _replace_invalid_tool_use_id_chars(tool_use_id, _BEDROCK_TOOL_USE_ID_INVALID_CHARS)
if sanitized == tool_use_id and len(sanitized) <= _BEDROCK_TOOL_USE_ID_MAX_LEN:
return sanitized
digest: Final = hashlib.sha256(tool_use_id.encode()).hexdigest()[:_BEDROCK_TOOL_USE_ID_HASH_LEN]
return f"{sanitized[: _BEDROCK_TOOL_USE_ID_MAX_LEN - _BEDROCK_TOOL_USE_ID_HASH_LEN - 1]}_{digest}"
_ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES: Final = {"application/pdf", "text/plain"}
@ -3661,7 +3675,9 @@ def _convert_to_bedrock_tool_call_invoke(
if parsed_objects:
# First object keeps the original tool id.
for obj_idx, obj in enumerate(parsed_objects):
block_id = tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}"
block_id = _sanitize_bedrock_tool_use_id(
tool_id if obj_idx == 0 else f"{tool_id}_{obj_idx}"
)
bedrock_tool = BedrockToolUseBlock(input=obj, name=name, toolUseId=block_id)
_parts_list.append(BedrockContentBlock(toolUse=bedrock_tool))
# cache_control applies to the whole original
@ -3678,7 +3694,9 @@ def _convert_to_bedrock_tool_call_invoke(
# Fallback: no objects extracted — use empty dict.
arguments_dict = {}
bedrock_tool = BedrockToolUseBlock(input=arguments_dict, name=name, toolUseId=tool_id)
bedrock_tool = BedrockToolUseBlock(
input=arguments_dict, name=name, toolUseId=_sanitize_bedrock_tool_use_id(tool_id)
)
bedrock_content_block = BedrockContentBlock(toolUse=bedrock_tool)
_parts_list.append(bedrock_content_block)
@ -3849,7 +3867,7 @@ def _convert_to_bedrock_tool_call_result(
tool_result_content_blocks, used_search_results = _build_bedrock_tool_result_content_blocks(message)
message.get("name", "")
id: Final = str(message.get("tool_call_id", str(uuid.uuid4())))
id: Final = _sanitize_bedrock_tool_use_id(str(message.get("tool_call_id", str(uuid.uuid4()))))
tool_result: Final = BedrockToolResultBlock(content=tool_result_content_blocks, toolUseId=id)
if used_search_results:

View file

@ -539,6 +539,13 @@ class AnthropicModelInfo(BaseLLMModelInfo):
value: Final = litellm.model_cost.get(model, {}).get(key)
return value if isinstance(value, bool) else None
@staticmethod
def supports_fast_mode(model: str, custom_llm_provider: str) -> bool:
return (
custom_llm_provider == "anthropic"
and AnthropicModelInfo._get_exact_model_capability(model, "supports_fast_mode") is True
)
@staticmethod
def _get_provider_resolved_capability(model: str, key: str, custom_llm_provider: str) -> bool | None:
"""Resolve boolean capability ``key`` for ``model`` under the caller's provider.

View file

@ -9,6 +9,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
drop_tool_reference_parts_from_tool_messages,
flatten_combinators_and_drop_non_python_regex_patterns,
hoist_images_from_tool_messages,
system_messages_first,
tool_with_sanitized_parameters,
)
from litellm.litellm_core_utils.prompt_templates.factory import (
@ -276,7 +277,8 @@ class AzureOpenAIConfig(BaseConfig):
litellm_params: dict,
headers: dict,
) -> dict:
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(messages)
ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages)
azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages))
return {
"model": model,

View file

@ -12,6 +12,7 @@ from urllib.parse import urlparse
import httpx
import litellm
from litellm.constants import OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_extract_reasoning_content,
@ -24,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
flatten_combinators_and_drop_non_python_regex_patterns,
get_tool_call_names,
hoist_images_from_tool_messages,
system_messages_first,
tool_with_sanitized_parameters,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import (
@ -463,6 +465,15 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
]
return MappingProxyType({"tools": sanitized})
def _prompt_cache_ordered_messages(
self, messages: list[AllMessageValues], litellm_params: Mapping[str, object]
) -> list[AllMessageValues]:
if not litellm.openai_system_messages_first:
return messages
if litellm_params.get("custom_llm_provider") not in OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS:
return messages
return system_messages_first(messages)
def transform_request(
self,
model: str,
@ -477,7 +488,9 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
Returns:
dict: The transformed request. Sent as the body of the API call.
"""
messages = self._transform_messages(messages=messages, model=model)
messages = self._transform_messages(
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model
)
if not self._should_preserve_cache_control_for_endpoint(
litellm_params.get("custom_llm_provider"), litellm_params.get("api_base")
):
@ -506,7 +519,9 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
litellm_params: dict,
headers: dict,
) -> dict:
transformed_messages = await self._transform_messages(messages=messages, model=model, is_async=True)
transformed_messages = await self._transform_messages(
messages=self._prompt_cache_ordered_messages(messages, litellm_params), model=model, is_async=True
)
if not self._should_preserve_cache_control_for_endpoint(
litellm_params.get("custom_llm_provider"), litellm_params.get("api_base")
):

File diff suppressed because it is too large Load diff

View file

@ -15226,6 +15226,17 @@
"title": "Jwt Claim Value",
"type": "string"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"key": {
"title": "Key",
"type": "string"
@ -15310,6 +15321,17 @@
"title": "Jwt Claim Value",
"type": "string"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"updated_at": {
"format": "date-time",
"title": "Updated At",
@ -15366,6 +15388,17 @@
],
"title": "Is Active"
},
"jwt_issuer": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Jwt Issuer"
},
"key": {
"anyOf": [
{

View file

@ -4499,12 +4499,14 @@ class CreateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
jwt_claim_name: str
jwt_claim_value: str
key: str
jwt_issuer: str | None = None
description: str | None = None
class UpdateJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
id: str
key: str | None = None
jwt_issuer: str | None = None
description: str | None = None
is_active: bool | None = None
@ -4515,6 +4517,7 @@ class DeleteJWTKeyMappingRequest(LiteLLMPydanticObjectBase):
class JWTKeyMappingResponse(LiteLLMPydanticObjectBase):
id: str
jwt_issuer: str | None = None
jwt_claim_name: str
jwt_claim_value: str
description: str | None = None

View file

@ -150,6 +150,7 @@ class _PrismaDictableRow(Protocol):
class _PrismaJWTKeyMappingRow(Protocol):
token: str
jwt_issuer: str
jwt_claim_name: str
jwt_claim_value: str
@ -3603,9 +3604,18 @@ async def _fetch_key_object_from_db_with_reconnect(
raise
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str) -> str:
"""Cache key under which ``_resolve_jwt_to_virtual_key`` stores a JWT-claim-to-key mapping."""
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
def jwt_key_mapping_cache_key(jwt_claim_name: str, jwt_claim_value: str, jwt_issuer: str | None = None) -> str:
"""Cache key under which a JWT-claim-to-key mapping is stored, scoped to one
issuer (or the issuer-agnostic/global scope when ``jwt_issuer`` is falsy).
Scoped by issuer (when one is configured) so a cached hit or ``__NO_MAPPING__`` miss
for one issuer's claim value can never be served to a different issuer whose claim
value happens to collide. Unchanged for the global scope, keeping the single-issuer
(no ``litellm_jwtauth.issuers`` configured) cache key format stable across this fix.
"""
if not jwt_issuer:
return f"jwt_key_mapping:{jwt_claim_name}:{jwt_claim_value}"
return f"jwt_key_mapping:{jwt_issuer}:{jwt_claim_name}:{jwt_claim_value}"
@log_db_metrics
@ -3617,7 +3627,7 @@ async def get_jwt_key_mapping_cache_keys_for_token(
mappings: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_many(
where={"token": hashed_token}
)
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value) for m in mappings)
return tuple(jwt_key_mapping_cache_key(m.jwt_claim_name, m.jwt_claim_value, m.jwt_issuer) for m in mappings)
@log_db_metrics
@ -3625,9 +3635,14 @@ async def get_jwt_key_mapping_object(
jwt_claim_name: str,
jwt_claim_value: str,
prisma_client: PrismaClient,
jwt_issuer: str | None = None,
) -> str | None:
"""
Lookup a JWT-to-virtual-key mapping from the database.
Lookup a JWT-to-virtual-key mapping from the database for one exact scope:
``jwt_issuer`` (or the global/issuer-agnostic scope when falsy). Does not fall
back to the global scope itself -- a caller that wants "issuer-scoped mapping,
else the global one" queries both scopes itself, so each result can be cached
under its own scope's key (see ``_resolve_jwt_to_virtual_key``).
Returns the hashed token (str) if a matching active mapping is found, else None.
"""
@ -3635,6 +3650,7 @@ async def get_jwt_key_mapping_object(
where={
"jwt_claim_name": jwt_claim_name,
"jwt_claim_value": jwt_claim_value,
"jwt_issuer": jwt_issuer or "",
"is_active": True,
}
)

View file

@ -269,6 +269,17 @@ class _TokenTeamModels(Protocol):
def team_models(self) -> list[str]: ...
class _RawCacheRead(Protocol):
async def async_get_cache(self, *, key: str) -> object: ...
def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead:
"""View an untyped cache object's ``async_get_cache`` as returning ``object``
instead of ``Any``, so a caller can ``isinstance``-narrow it without paying
the ``reportAny`` cost of the underlying (unannotated) cache implementation."""
return cache
def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
return valid_token.team_models
@ -842,6 +853,7 @@ class _PendingAutoRegister(NamedTuple):
claim_field: str
claim_value: str
cache_key: str
jwt_issuer: str | None = None
async def _auto_register_jwt_mapping(
@ -853,6 +865,7 @@ async def _auto_register_jwt_mapping(
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging,
cache_key: str,
jwt_issuer: str | None = None,
team_id: str | None = None,
user_id: str | None = None,
org_id: str | None = None,
@ -905,6 +918,7 @@ async def _auto_register_jwt_mapping(
try:
await prisma_client.db.litellm_jwtkeymapping.create(
data={
"jwt_issuer": jwt_issuer or "",
"jwt_claim_name": virtual_key_claim_field,
"jwt_claim_value": claim_value,
"token": token_hash,
@ -939,6 +953,7 @@ async def _auto_register_jwt_mapping(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=jwt_issuer,
)
if token_hash is None:
# The winner's mapping vanished between the unique-constraint
@ -983,6 +998,43 @@ async def _auto_register_jwt_mapping(
return auto_registered_key
async def _lookup_jwt_mapping_token_hash(
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
virtual_key_claim_field: str,
claim_value: str,
normalized_issuer: str | None,
cache_key: str,
ttl: float,
) -> str | None:
issuer_scoped: Final = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=normalized_issuer,
)
if issuer_scoped is not None:
await user_api_key_cache.async_set_cache(key=cache_key, value=issuer_scoped, ttl=ttl)
return issuer_scoped
if normalized_issuer is None:
return None
# Another issuer may have already resolved (and cached) this same
# global mapping -- check its cache entry before re-querying the DB.
global_cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, claim_value)
cached_global: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=global_cache_key)
if isinstance(cached_global, str) and cached_global != "__NO_MAPPING__":
return cached_global
global_row: Final = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=claim_value,
prisma_client=prisma_client,
jwt_issuer=None,
)
if global_row is not None:
await user_api_key_cache.async_set_cache(key=global_cache_key, value=global_row, ttl=ttl)
return global_row
async def _resolve_jwt_to_virtual_key(
jwt_claims: dict,
jwt_handler: JWTHandler,
@ -1041,7 +1093,7 @@ async def _resolve_jwt_to_virtual_key(
)
return None
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value))
cache_key: Final = jwt_key_mapping_cache_key(virtual_key_claim_field, str(claim_value), normalized_issuer)
raw_cached_mapping: Final = await user_api_key_cache.async_get_cache(cache_key)
sentinel_written_by_this_policy: Final = behavior == UnregisteredJWTClientBehavior.AUTO_REGISTER
cached_mapping: Final = (
@ -1081,6 +1133,7 @@ async def _resolve_jwt_to_virtual_key(
claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
cache_key=cache_key,
jwt_issuer=normalized_issuer,
)
return None
elif cached_mapping is not None:
@ -1094,21 +1147,30 @@ async def _resolve_jwt_to_virtual_key(
)
# Resolve the mapping from DB, or treat prisma_client=None as a definitive
# miss (no DB → no mapping can exist → apply no-match policy below).
token_hash: str | None = None
if prisma_client is not None:
token_hash = await get_jwt_key_mapping_object(
jwt_claim_name=virtual_key_claim_field,
jwt_claim_value=str(claim_value),
# miss (no DB → no mapping can exist → apply no-match policy below). An
# issuer-scoped row wins; falling back to the global (no-issuer) row keeps
# mappings created before issuer scoping existed working for every issuer.
# Each tier is cached under ITS OWN key (the global tier under the
# issuer-less cache key, not under `cache_key`/this issuer's key) so that
# updating or deleting either row invalidates exactly the cache entries it
# can affect. Caching a global-row hit under the requesting issuer's key
# would leave every OTHER issuer that had fallen back to that same global
# mapping serving its stale token until TTL after the row changes.
token_hash: Final = (
await _lookup_jwt_mapping_token_hash(
prisma_client=prisma_client,
)
if token_hash is not None:
await user_api_key_cache.async_set_cache(
key=cache_key,
value=token_hash,
user_api_key_cache=user_api_key_cache,
virtual_key_claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
normalized_issuer=normalized_issuer,
cache_key=cache_key,
ttl=jwt_handler.litellm_jwtauth.virtual_key_mapping_cache_ttl,
)
if prisma_client is not None
else None
)
if token_hash is not None:
return IdentityStore.key_from_principal(
await IdentityStore(
prisma_client,
@ -1149,6 +1211,7 @@ async def _resolve_jwt_to_virtual_key(
claim_field=virtual_key_claim_field,
claim_value=str(claim_value),
cache_key=cache_key,
jwt_issuer=normalized_issuer,
)
# FALLBACK_TEAM_MAPPING (default): cache the miss and return None so the
@ -1641,6 +1704,7 @@ async def _user_api_key_auth_builder(
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
cache_key=pending_auto_register.cache_key,
jwt_issuer=pending_auto_register.jwt_issuer,
team_id=team_id,
user_id=user_id,
org_id=org_id,

View file

@ -490,7 +490,7 @@ lite codex exec "summarize the repo"
Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process.
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-<UTF-8 hex of the group name>` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway.
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-<UTF-8 hex of the group name>` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. Codex gets the same list as a catalog file, `$CODEX_HOME/litellm-models.json` (default `~/.codex/`), passed as `-c model_catalog_json=<path>` so `/model` lists exactly the proxy's chat models; a proxy model the installed Codex already knows (`gpt-5.5`, say) keeps that Codex's own entry, reasoning levels and prompt included, and only its place in the picker comes from the proxy, while a model Codex does not know gets the plain entry Codex uses for an unknown `-m` slug. Before launching, `lite codex` asks the installed Codex for its own list and then has it read the written file back (both through `codex debug models`), and when the fetch, either of those or the write fails (Codex releases older than 0.130 have no such command) it says so on stderr and launches with Codex's built-in catalog, leaving a rejected file in place.
pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/<id>`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/<other-id>` wins, and inside the TUI the `/model` picker lists every synced litellm model.

View file

@ -4,15 +4,16 @@ import re
import shutil
import subprocess
import sys
import tempfile
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final, TypeAlias
from typing import Final, Literal, TypeAlias
import click
import requests
from pydantic import BaseModel, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login
from .claude_settings import ClaudeSettingsError, install_statusline_script
@ -65,6 +66,10 @@ _INSTALL_DOCS: Final[dict[str, str]] = {
_HIDDEN_AGENTS: Final = frozenset({"pi"})
CODEX_PROXY_PROVIDER: Final = "litellm"
CODEX_HOME_ENV: Final = "CODEX_HOME"
CODEX_MODEL_CATALOG_FILENAME: Final = "litellm-models.json"
_CODEX_BASE_INSTRUCTIONS_PATH: Final = Path(__file__).with_name("codex_base_instructions.md")
_CODEX_PREFLIGHT_TIMEOUT_SECONDS: Final = 10.0
class AgentRunError(Exception):
@ -252,7 +257,7 @@ def agent_launch_args(command: str, base_url: str) -> list[str]:
class ListedModel(BaseModel):
"""The fields of a /v1/models entry that an OpenCode model entry is built from."""
"""The fields of a /v1/models entry that an OpenCode or Codex model entry is built from."""
id: str
mode: str | None = None
@ -265,7 +270,7 @@ class _ModelListing(BaseModel):
_MODEL_LISTING: Final = TypeAdapter(_ModelListing)
_OPENCODE_CHAT_MODES: Final[frozenset[str]] = frozenset({"chat", "responses"})
_CHAT_MODES: Final[frozenset[str]] = frozenset({"chat", "responses"})
_NO_EXTRA_ENV: Final[Mapping[str, str]] = MappingProxyType({})
@ -274,6 +279,40 @@ class ModelSyncSkipped:
reason: str
@dataclass(frozen=True, slots=True)
class ModelSyncArgs:
"""CLI args, placed before the user's own, that hand an agent the synced model list."""
args: tuple[str, ...]
ModelSyncResult: TypeAlias = Mapping[str, str] | ModelSyncArgs | ModelSyncSkipped
def _chat_models(models: Sequence[ListedModel]) -> tuple[ListedModel, ...]:
return tuple(m for m in models if m.mode is None or m.mode in _CHAT_MODES)
def _fetch_model_listing(
base_url: str,
api_key: str,
*,
get: Callable[..., requests.Response],
) -> tuple[ListedModel, ...] | ModelSyncSkipped:
url: Final = base_url.rstrip("/") + "/v1/models"
try:
resp: Final = get(url, headers=MappingProxyType({"Authorization": f"Bearer {api_key}"}), timeout=10)
except requests.RequestException as e:
return ModelSyncSkipped(f"could not reach {url}: {e}")
if resp.status_code != 200:
return ModelSyncSkipped(f"{url} returned HTTP {resp.status_code}")
try:
listing: Final = _MODEL_LISTING.validate_json(resp.content)
except ValidationError:
return ModelSyncSkipped(f"{url} returned an unexpected body")
return listing.data
class _OpenCodeLimit(BaseModel):
context: int
output: int
@ -317,7 +356,7 @@ def opencode_provider_config(base_url: str, models: Sequence[ListedModel]) -> st
it never lands in the config text. OpenCode merges this inline config over
the user's own files, leaving unrelated keys and providers untouched.
"""
chat_models: Final = tuple(m for m in models if m.mode is None or m.mode in _OPENCODE_CHAT_MODES)
chat_models: Final = _chat_models(models)
provider: Final = _OpenCodeProvider(
npm=OPENCODE_PROVIDER_NPM,
name=OPENCODE_PROVIDER_NAME,
@ -347,40 +386,269 @@ def opencode_model_sync_env(
"""
if OPENCODE_CONFIG_CONTENT_ENV in base_env:
return ModelSyncSkipped(f"{OPENCODE_CONFIG_CONTENT_ENV} is already set")
url: Final = base_url.rstrip("/") + "/v1/models"
listing: Final = _fetch_model_listing(base_url, api_key, get=get)
if isinstance(listing, ModelSyncSkipped):
return listing
return MappingProxyType({OPENCODE_CONFIG_CONTENT_ENV: opencode_provider_config(base_url, listing)})
class _CodexTruncationPolicy(BaseModel):
mode: Literal["bytes"] = "bytes"
limit: int = 10_000
class _CodexModel(BaseModel):
"""One `ModelInfo` entry of a Codex model catalog for a model the installed Codex does not know.
Every field that some Codex release since `model_catalog_json` appeared
(0.105.0) deserializes without a default is spelled out here, so one catalog
parses on all of them; the values match the fallback metadata Codex uses for
a model slug it does not know, so picking such a proxy model behaves the
same as `codex -m` did.
"""
slug: str
display_name: str
description: None = None
supported_reasoning_levels: tuple[()] = ()
shell_type: Literal["unified_exec"] = "unified_exec"
visibility: Literal["list"] = "list"
supported_in_api: Literal[True] = True
priority: int
availability_nux: None = None
upgrade: None = None
support_verbosity: Literal[False] = False
supports_reasoning_summaries: Literal[False] = False
supports_parallel_tool_calls: Literal[False] = False
default_verbosity: None = None
apply_patch_tool_type: None = None
truncation_policy: _CodexTruncationPolicy = _CodexTruncationPolicy()
experimental_supported_tools: tuple[()] = ()
context_window: int | None
base_instructions: str
class _StockCodexUpgrade(BaseModel):
model_config = ConfigDict(extra="allow")
model: str
class _StockCodexModel(BaseModel):
"""One `ModelInfo` entry as the installed Codex prints it from `codex debug models`.
Only the fields the sync rewrites are named; everything else that release
knows about the model (its reasoning levels, prompt, tool support) rides
along untouched, whatever the release's schema.
"""
model_config = ConfigDict(extra="allow")
slug: str
priority: int
visibility: str
supported_in_api: bool = True
upgrade: _StockCodexUpgrade | None = None
class _StockCodexCatalog(BaseModel):
models: tuple[_StockCodexModel, ...]
class _CodexCatalog(BaseModel):
models: tuple[_CodexModel | _StockCodexModel, ...]
def _codex_catalog_entry(
priority: int,
listed: ListedModel,
stock: _StockCodexModel | None,
served: frozenset[str],
instructions: str,
) -> _CodexModel | _StockCodexModel:
if stock is None:
return _CodexModel(
slug=listed.id,
display_name=listed.id,
priority=priority,
context_window=listed.max_input_tokens,
base_instructions=instructions,
)
upgrade: Final = stock.upgrade if stock.upgrade is not None and stock.upgrade.model in served else None
return stock.model_copy(
update={"priority": priority, "visibility": "list", "supported_in_api": True, "upgrade": upgrade}
)
def codex_model_catalog(
models: Sequence[ListedModel], stock: Sequence[_StockCodexModel], instructions: str
) -> str | None:
"""The `model_catalog_json` body listing the proxy's chat models, or None if there are none.
Codex refuses an empty catalog, hence None instead of `{"models": []}`.
Passing a catalog replaces Codex's built-in one, so a proxy model the
installed Codex knows keeps that Codex's own entry and the proxy only
decides its place in the picker: the listing orders it, lists it even when
Codex hides it or keeps it off the API, and keeps Codex's upgrade nudge only
when the model it points at is served too. A model Codex does not know gets the fallback
entry, with the same base instructions Codex itself uses so the agent never
runs without a system prompt.
"""
chat_models: Final = _chat_models(models)
if not chat_models:
return None
served: Final = frozenset(m.id for m in chat_models)
known: Final = MappingProxyType({m.slug: m for m in stock})
catalog: Final = _CodexCatalog(
models=tuple(
_codex_catalog_entry(index, m, known.get(m.id), served, instructions) for index, m in enumerate(chat_models)
)
)
return catalog.model_dump_json()
def codex_model_catalog_path(env: Mapping[str, str], *, home: Callable[[], Path] = Path.home) -> Path:
override: Final = env.get(CODEX_HOME_ENV)
root: Final = Path(override) if override else home() / ".codex"
return root / CODEX_MODEL_CATALOG_FILENAME
def _replace_file(path: Path, text: str) -> None:
with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as tmp:
_ = tmp.write(text)
try:
resp: Final = get(url, headers=MappingProxyType({"Authorization": f"Bearer {api_key}"}), timeout=10)
except requests.RequestException as e:
return ModelSyncSkipped(f"could not reach {url}: {e}")
if resp.status_code != 200:
return ModelSyncSkipped(f"{url} returned HTTP {resp.status_code}")
os.replace(tmp.name, path)
except OSError:
Path(tmp.name).unlink(missing_ok=True)
raise
def _codex_debug_models(
binary: str,
args: Sequence[str],
env: Mapping[str, str],
*,
run: Callable[..., subprocess.CompletedProcess[str]],
) -> str | ModelSyncSkipped:
"""What `codex debug models` prints with `args` in front, or why the installed Codex could not run it.
The command prints the catalog Codex would launch with, without touching
the network, so it lists the installed Codex's own models and parses a
catalog override the way a launch does. Releases before 0.130.0 have no
such command and are reported the same way. A batch shim goes through
cmd.exe exactly as the launch will.
"""
name: Final = os.path.basename(binary)
command: Final = _windows_command(binary, (binary, *args, "debug", "models"))
try:
listing: Final = _MODEL_LISTING.validate_json(resp.content)
except ValidationError:
return ModelSyncSkipped(f"{url} returned an unexpected body")
return MappingProxyType({OPENCODE_CONFIG_CONTENT_ENV: opencode_provider_config(base_url, listing.data)})
completed: Final = run(
command,
env=dict(env),
stdin=subprocess.DEVNULL,
capture_output=True,
encoding="utf-8",
timeout=_CODEX_PREFLIGHT_TIMEOUT_SECONDS,
)
except (OSError, subprocess.TimeoutExpired) as e:
return ModelSyncSkipped(f"`{name} debug models` failed: {e}")
if completed.returncode == 0:
return completed.stdout
lines: Final = completed.stderr.strip().splitlines()
detail: Final = lines[0] if lines else "no output"
return ModelSyncSkipped(f"`{name} debug models` exited {completed.returncode}: {detail}")
def _stock_codex_models(
binary: str, env: Mapping[str, str], *, run: Callable[..., subprocess.CompletedProcess[str]]
) -> tuple[_StockCodexModel, ...] | ModelSyncSkipped:
printed: Final = _codex_debug_models(binary, (), env, run=run)
if isinstance(printed, ModelSyncSkipped):
return printed
try:
return _StockCodexCatalog.model_validate_json(printed).models
except ValidationError as e:
name: Final = os.path.basename(binary)
return ModelSyncSkipped(f"`{name} debug models` printed no model catalog: {e.errors()[0]['msg']}")
def codex_model_sync_args(
base_env: Mapping[str, str],
base_url: str,
api_key: str,
*,
binary: str = "codex",
get: Callable[..., requests.Response] = requests.get,
run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
home: Callable[[], Path] = Path.home,
instructions_path: Path = _CODEX_BASE_INSTRUCTIONS_PATH,
) -> ModelSyncArgs | ModelSyncSkipped:
"""`-c model_catalog_json=...` pointing Codex at the proxy's model list, or why it was skipped.
Codex has no env or inline equivalent of OPENCODE_CONFIG_CONTENT: the catalog
must be a file, so it is written under $CODEX_HOME (default ~/.codex) and
atomically replaced on every launch. The Codex at `binary` first lists its
own models, so the ones the proxy serves keep that Codex's entries, and then
reads the file back once before it is handed over. The key never lands in
the file. A failed fetch, read, listing, write or read-back is reported
rather than raised: Codex still launches with its built-in catalog and takes
a proxy model by name via -m, and a rejected file stays on disk to be looked
at.
"""
listing: Final = _fetch_model_listing(base_url, api_key, get=get)
if isinstance(listing, ModelSyncSkipped):
return listing
try:
instructions: Final = instructions_path.read_text(encoding="utf-8")
except OSError as e:
return ModelSyncSkipped(f"could not read {instructions_path}: {e}")
path: Final = codex_model_catalog_path(base_env, home=home)
try:
path.parent.mkdir(parents=True, exist_ok=True)
except OSError as e:
return ModelSyncSkipped(f"could not write {path}: {e}")
stock: Final = _stock_codex_models(binary, base_env, run=run)
if isinstance(stock, ModelSyncSkipped):
return stock
catalog: Final = codex_model_catalog(listing, stock, instructions)
if catalog is None:
return ModelSyncSkipped(f"{base_url.rstrip('/')}/v1/models lists no chat models")
try:
_replace_file(path, catalog)
except OSError as e:
return ModelSyncSkipped(f"could not write {path}: {e}")
override: Final = f"model_catalog_json={json.dumps(str(path))}"
read_back: Final = _codex_debug_models(binary, ("-c", override), base_env, run=run)
if isinstance(read_back, ModelSyncSkipped):
return read_back
return ModelSyncArgs(("-c", override))
def agent_model_sync_env(
command: str,
binary: str,
base_env: Mapping[str, str],
base_url: str,
api_key: str,
skip_verify: bool,
*,
get: Callable[..., requests.Response] = requests.get,
) -> Mapping[str, str] | ModelSyncSkipped:
"""Extra env an agent needs to see the proxy's model list.
run: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
) -> ModelSyncResult:
"""Extra env or args an agent needs to see the proxy's model list.
Only OpenCode needs one: Claude Code discovers models through
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY and Codex takes the model by name.
skip_verify means the caller wants no pre-launch proxy call at all, so the
listing is skipped too rather than hanging on an offline proxy.
binary is the resolved path the launch will run (`codex.cmd` on a Windows
npm install). OpenCode takes the list as env, Codex as a `-c` override that
binary has read back first; Claude Code discovers models itself through
CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY. skip_verify means the caller
wants no pre-launch proxy call at all, so the listing is skipped too rather
than hanging on an offline proxy.
"""
if os.path.basename(command) != "opencode":
agent: Final = os.path.splitext(os.path.basename(binary))[0]
if agent not in ("opencode", "codex"):
return _NO_EXTRA_ENV
if skip_verify:
return ModelSyncSkipped(f"{_SKIP_VERIFY_FLAG} was passed")
if agent == "codex":
return codex_model_sync_args(base_env, base_url, api_key, binary=binary, get=get, run=run)
return opencode_model_sync_env(base_env, base_url, api_key, get=get)
@ -508,9 +776,7 @@ def run_agent(
base_env: Mapping[str, str] | None = None,
which: Callable[[str], str | None] = shutil.which,
verify: Callable[[str, str], None] = verify_proxy_key,
sync_models: Callable[[str, Mapping[str, str], str, str, bool], Mapping[str, str] | ModelSyncSkipped] = (
agent_model_sync_env
),
sync_models: Callable[[str, Mapping[str, str], str, str, bool], ModelSyncResult] = agent_model_sync_env,
warn: Callable[[str], None] = _warn,
launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off,
reattach_terminal: Callable[[], None] | None = None,
@ -537,7 +803,7 @@ def run_agent(
verify(base_url, api_key)
env_before_sync: Final = base_env if base_env is not None else os.environ
synced: Final = sync_models(command[0], env_before_sync, base_url, api_key, skip_verify)
synced: Final = sync_models(binary, env_before_sync, base_url, api_key, skip_verify)
if isinstance(synced, ModelSyncSkipped):
warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}")
@ -547,10 +813,11 @@ def run_agent(
env: Final = MappingProxyType(
{
**build_agent_env(env_before_sync, base_url, api_key, profiles),
**(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced),
**(synced if isinstance(synced, Mapping) else _NO_EXTRA_ENV),
}
)
extra_args: Final = (*agent_launch_args(command[0], base_url), *prepared_args)
synced_args: Final = synced.args if isinstance(synced, ModelSyncArgs) else ()
extra_args: Final = (*agent_launch_args(command[0], base_url), *synced_args, *prepared_args)
if reattach_terminal is not None:
reattach_terminal()
launcher(binary, [command[0], *extra_args, *command[1:]], env)

View file

@ -0,0 +1,275 @@
You are a coding agent running in the Codex CLI, a terminal-based coding assistant. Codex CLI is an open source project led by OpenAI. You are expected to be precise, safe, and helpful.
Your capabilities:
- Receive user prompts and other context provided by the harness, such as files in the workspace.
- Communicate with the user by streaming thinking & responses, and by making & updating plans.
- Emit function calls to run terminal commands and apply patches. Depending on how this specific run is configured, you can request that these function calls be escalated to the user for approval before running. More on this in the "Sandbox and approvals" section.
Within this context, Codex refers to the open-source agentic coding interface (not the old Codex language model built by OpenAI).
# How you work
## Personality
Your default personality and tone is concise, direct, and friendly. You communicate efficiently, always keeping the user clearly informed about ongoing actions without unnecessary detail. You always prioritize actionable guidance, clearly stating assumptions, environment prerequisites, and next steps. Unless explicitly asked, you avoid excessively verbose explanations about your work.
# AGENTS.md spec
- Repos often contain AGENTS.md files. These files can appear anywhere within the repository.
- These files are a way for humans to give you (the agent) instructions or tips for working within the container.
- Some examples might be: coding conventions, info about how code is organized, or instructions for how to run or test code.
- Instructions in AGENTS.md files:
- The scope of an AGENTS.md file is the entire directory tree rooted at the folder that contains it.
- For every file you touch in the final patch, you must obey instructions in any AGENTS.md file whose scope includes that file.
- Instructions about code style, structure, naming, etc. apply only to code within the AGENTS.md file's scope, unless the file states otherwise.
- More-deeply-nested AGENTS.md files take precedence in the case of conflicting instructions.
- Direct system/developer/user instructions (as part of a prompt) take precedence over AGENTS.md instructions.
- The contents of the AGENTS.md file at the root of the repo and any directories from the CWD up to the root are included with the developer message and don't need to be re-read. When working in a subdirectory of CWD, or a directory outside the CWD, check for any AGENTS.md files that may be applicable.
## Responsiveness
### Preamble messages
Before making tool calls, send a brief preamble to the user explaining what youre about to do. When sending preamble messages, follow these principles and examples:
- **Logically group related actions**: if youre about to run several related commands, describe them together in one preamble rather than sending a separate note for each.
- **Keep it concise**: be no more than 1-2 sentences, focused on immediate, tangible next steps. (812 words for quick updates).
- **Build on prior context**: if this is not your first tool call, use the preamble message to connect the dots with whats been done so far and create a sense of momentum and clarity for the user to understand your next actions.
- **Keep your tone light, friendly and curious**: add small touches of personality in preambles feel collaborative and engaging.
- **Exception**: Avoid adding a preamble for every trivial read (e.g., `cat` a single file) unless its part of a larger grouped action.
**Examples:**
- “Ive explored the repo; now checking the API route definitions.”
- “Next, Ill patch the config and update the related tests.”
- “Im about to scaffold the CLI commands and helper functions.”
- “Ok cool, so Ive wrapped my head around the repo. Now digging into the API routes.”
- “Configs looking tidy. Next up is patching helpers to keep things in sync.”
- “Finished poking at the DB gateway. I will now chase down error handling.”
- “Alright, build pipeline order is interesting. Checking how it reports failures.”
- “Spotted a clever caching util; now hunting where it gets used.”
## Planning
You have access to an `update_plan` tool which tracks steps and progress and renders them to the user. Using the tool helps demonstrate that you've understood the task and convey how you're approaching it. Plans can help to make complex, ambiguous, or multi-phase work clearer and more collaborative for the user. A good plan should break the task into meaningful, logically ordered steps that are easy to verify as you go.
Note that plans are not for padding out simple work with filler steps or stating the obvious. The content of your plan should not involve doing anything that you aren't capable of doing (i.e. don't try to test things that you can't test). Do not use plans for simple or single-step queries that you can just do or answer immediately.
Do not repeat the full contents of the plan after an `update_plan` call — the harness already displays it. Instead, summarize the change made and highlight any important context or next step.
Before running a command, consider whether or not you have completed the previous step, and make sure to mark it as completed before moving on to the next step. It may be the case that you complete all steps in your plan after a single pass of implementation. If this is the case, you can simply mark all the planned steps as completed. Sometimes, you may need to change plans in the middle of a task: call `update_plan` with the updated plan and make sure to provide an `explanation` of the rationale when doing so.
Use a plan when:
- The task is non-trivial and will require multiple actions over a long time horizon.
- There are logical phases or dependencies where sequencing matters.
- The work has ambiguity that benefits from outlining high-level goals.
- You want intermediate checkpoints for feedback and validation.
- When the user asked you to do more than one thing in a single prompt
- The user has asked you to use the plan tool (aka "TODOs")
- You generate additional steps while working, and plan to do them before yielding to the user
### Examples
**High-quality plans**
Example 1:
1. Add CLI entry with file args
2. Parse Markdown via CommonMark library
3. Apply semantic HTML template
4. Handle code blocks, images, links
5. Add error handling for invalid files
Example 2:
1. Define CSS variables for colors
2. Add toggle with localStorage state
3. Refactor components to use variables
4. Verify all views for readability
5. Add smooth theme-change transition
Example 3:
1. Set up Node.js + WebSocket server
2. Add join/leave broadcast events
3. Implement messaging with timestamps
4. Add usernames + mention highlighting
5. Persist messages in lightweight DB
6. Add typing indicators + unread count
**Low-quality plans**
Example 1:
1. Create CLI tool
2. Add Markdown parser
3. Convert to HTML
Example 2:
1. Add dark mode toggle
2. Save preference
3. Make styles look good
Example 3:
1. Create single-file HTML game
2. Run quick sanity check
3. Summarize usage instructions
If you need to write a plan, only write high quality plans, not low quality ones.
## Task execution
You are a coding agent. Please keep going until the query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved. Autonomously resolve the query to the best of your ability, using the tools available to you, before coming back to the user. Do NOT guess or make up an answer.
You MUST adhere to the following criteria when solving queries:
- Working on the repo(s) in the current environment is allowed, even if they are proprietary.
- Analyzing code for vulnerabilities is allowed.
- Showing user code and tool call details is allowed.
- Use the `apply_patch` tool to edit files (NEVER try `applypatch` or `apply-patch`, only `apply_patch`): {"command":["apply_patch","*** Begin Patch\\n*** Update File: path/to/file.py\\n@@ def example():\\n- pass\\n+ return 123\\n*** End Patch"]}
If completing the user's task requires writing or modifying files, your code and final answer should follow these coding guidelines, though user instructions (i.e. AGENTS.md) may override these guidelines:
- Fix the problem at the root cause rather than applying surface-level patches, when possible.
- Avoid unneeded complexity in your solution.
- Do not attempt to fix unrelated bugs or broken tests. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
- Update documentation as necessary.
- Keep changes consistent with the style of the existing codebase. Changes should be minimal and focused on the task.
- Use `git log` and `git blame` to search the history of the codebase if additional context is required.
- NEVER add copyright or license headers unless specifically requested.
- Do not waste tokens by re-reading files after calling `apply_patch` on them. The tool call will fail if it didn't work. The same goes for making folders, deleting folders, etc.
- Do not `git commit` your changes or create new git branches unless explicitly requested.
- Do not add inline comments within code unless explicitly requested.
- Do not use one-letter variable names unless explicitly requested.
- NEVER output inline citations like "【F:README.md†L5-L14】" in your outputs. The CLI is not able to render these so they will just be broken in the UI. Instead, if you output valid filepaths, users will be able to click on them to open the files in their editor.
## Validating your work
If the codebase has tests or the ability to build or run, consider using them to verify that your work is complete.
When testing, your philosophy should be to start as specific as possible to the code you changed so that you can catch issues efficiently, then make your way to broader tests as you build confidence. If there's no test for the code you changed, and if the adjacent patterns in the codebases show that there's a logical place for you to add a test, you may do so. However, do not add tests to codebases with no tests.
Similarly, once you're confident in correctness, you can suggest or use formatting commands to ensure that your code is well formatted. If there are issues you can iterate up to 3 times to get formatting right, but if you still can't manage it's better to save the user time and present them a correct solution where you call out the formatting in your final message. If the codebase does not have a formatter configured, do not add one.
For all of testing, running, building, and formatting, do not attempt to fix unrelated bugs. It is not your responsibility to fix them. (You may mention them to the user in your final message though.)
Be mindful of whether to run validation commands proactively. In the absence of behavioral guidance:
- When running in the non-interactive approval mode **never**, proactively run tests, lint and do whatever you need to ensure you've completed the task.
- When working in interactive approval modes like **untrusted**, or **on-request**, hold off on running tests or lint commands until the user is ready for you to finalize your output, because these commands take time to run and slow down iteration. Instead suggest what you want to do next, and let the user confirm first.
- When working on test-related tasks, such as adding tests, fixing tests, or reproducing a bug to verify behavior, you may proactively run tests regardless of approval mode. Use your judgement to decide whether this is a test-related task.
## Ambition vs. precision
For tasks that have no prior context (i.e. the user is starting something brand new), you should feel free to be ambitious and demonstrate creativity with your implementation.
If you're operating in an existing codebase, you should make sure you do exactly what the user asks with surgical precision. Treat the surrounding codebase with respect, and don't overstep (i.e. changing filenames or variables unnecessarily). You should balance being sufficiently ambitious and proactive when completing tasks of this nature.
You should use judicious initiative to decide on the right level of detail and complexity to deliver based on the user's needs. This means showing good judgment that you're capable of doing the right extras without gold-plating. This might be demonstrated by high-value, creative touches when scope of the task is vague; while being surgical and targeted when scope is tightly specified.
## Sharing progress updates
For especially longer tasks that you work on (i.e. requiring many tool calls, or a plan with multiple steps), you should provide progress updates back to the user at reasonable intervals. These updates should be structured as a concise sentence or two (no more than 8-10 words long) recapping progress so far in plain language: this update demonstrates your understanding of what needs to be done, progress so far (i.e. files explores, subtasks complete), and where you're going next.
Before doing large chunks of work that may incur latency as experienced by the user (i.e. writing a new file), you should send a concise message to the user with an update indicating what you're about to do to ensure they know what you're spending time on. Don't start editing or writing large files before informing the user what you are doing and why.
The messages you send before tool calls should describe what is immediately about to be done next in very concise language. If there was previous work done, this preamble message should also include a note about the work done so far to bring the user along.
## Presenting your work and final message
Your final message should read naturally, like an update from a concise teammate. For casual conversation, brainstorming tasks, or quick questions from the user, respond in a friendly, conversational tone. You should ask questions, suggest ideas, and adapt to the users style. If you've finished a large amount of work, when describing what you've done to the user, you should follow the final answer formatting guidelines to communicate substantive changes. You don't need to add structured formatting for one-word answers, greetings, or purely conversational exchanges.
You can skip heavy formatting for single, simple actions or confirmations. In these cases, respond in plain sentences with any relevant next step or quick option. Reserve multi-section structured responses for results that need grouping or explanation.
The user is working on the same computer as you, and has access to your work. As such there's no need to show the full contents of large files you have already written unless the user explicitly asks for them. Similarly, if you've created or modified files using `apply_patch`, there's no need to tell users to "save the file" or "copy the code into a file"—just reference the file path.
If there's something that you think you could help with as a logical next step, concisely ask the user if they want you to do so. Good examples of this are running tests, committing changes, or building out the next logical component. If theres something that you couldn't do (even with approval) but that the user might want to do (such as verifying changes by running the app), include those instructions succinctly.
Brevity is very important as a default. You should be very concise (i.e. no more than 10 lines), but can relax this requirement for tasks where additional detail and comprehensiveness is important for the user's understanding.
### Final answer structure and style guidelines
You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value.
**Section Headers**
- Use only when they improve clarity — they are not mandatory for every answer.
- Choose descriptive names that fit the content
- Keep headers short (13 words) and in `**Title Case**`. Always start headers with `**` and end with `**`
- Leave no blank line before the first bullet under a header.
- Section headers should only be used where they genuinely improve scanability; avoid fragmenting the answer.
**Bullets**
- Use `-` followed by a space for every bullet.
- Merge related points when possible; avoid a bullet for every trivial detail.
- Keep bullets to one line unless breaking for clarity is unavoidable.
- Group into short lists (46 bullets) ordered by importance.
- Use consistent keyword phrasing and formatting across sections.
**Monospace**
- Wrap all commands, file paths, env vars, and code identifiers in backticks (`` `...` ``).
- Apply to inline examples and to bullet keywords if the keyword itself is a literal file/command.
- Never mix monospace and bold markers; choose one based on whether its a keyword (`**`) or inline code/path (`` ` ``).
**File References**
When referencing files in your response, make sure to include the relevant start line and always follow the below rules:
* Use inline code to make file paths clickable.
* Each reference should have a stand alone path. Even if it's the same file.
* Accepted: absolute, workspacerelative, a/ or b/ diff prefixes, or bare filename/suffix.
* Line/column (1based, optional): :line[:column] or #Lline[Ccolumn] (column defaults to 1).
* Do not use URIs like file://, vscode://, or https://.
* Do not provide range of lines
* Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\repo\project\main.rs:12:5
**Structure**
- Place related bullets together; dont mix unrelated concepts in the same section.
- Order sections from general → specific → supporting info.
- For subsections (e.g., “Binaries” under “Rust Workspace”), introduce with a bolded keyword bullet, then list items under it.
- Match structure to complexity:
- Multi-part or detailed results → use clear headers and grouped bullets.
- Simple results → minimal headers, possibly just a short list or paragraph.
**Tone**
- Keep the voice collaborative and natural, like a coding partner handing off work.
- Be concise and factual — no filler or conversational commentary and avoid unnecessary repetition
- Use present tense and active voice (e.g., “Runs tests” not “This will run tests”).
- Keep descriptions self-contained; dont refer to “above” or “below”.
- Use parallel structure in lists for consistency.
**Dont**
- Dont use literal words “bold” or “monospace” in the content.
- Dont nest bullets or create deep hierarchies.
- Dont output ANSI escape codes directly — the CLI renderer applies them.
- Dont cram unrelated keywords into a single bullet; split for clarity.
- Dont let keyword lists run long — wrap or reformat for scanability.
Generally, ensure your final answers adapt their shape and depth to the request. For example, answers to code explanations should have a precise, structured explanation with code references that answer the question directly. For tasks with a simple implementation, lead with the outcome and supplement only with whats needed for clarity. Larger changes can be presented as a logical walkthrough of your approach, grouping related steps, explaining rationale where it adds value, and highlighting next actions to accelerate the user. Your answers should provide the right level of detail while being easily scannable.
For casual greetings, acknowledgements, or other one-off conversational messages that are not delivering substantive information or structured results, respond naturally without section headers or bullet formatting.
# Tool Guidelines
## Shell commands
When using the shell, you must adhere to the following guidelines:
- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.)
- Do not use python scripts to attempt to output larger chunks of a file.
## `update_plan`
A tool named `update_plan` is available to you. You can use it to keep an uptodate, stepbystep plan for the task.
To create a new plan, call `update_plan` with a short list of 1sentence steps (no more than 5-7 words each) with a `status` for each step (`pending`, `in_progress`, or `completed`).
When steps have been completed, use `update_plan` to mark each finished step as `completed` and the next step you are working on as `in_progress`. There should always be exactly one `in_progress` step until everything is done. You can mark multiple items as complete in a single `update_plan` call.
If all steps are complete, ensure you call `update_plan` to mark all steps as `completed`.

View file

@ -189,8 +189,9 @@ async def _read_request_body(request: Request | None) -> dict:
try:
parsed_body = json.loads(body_str)
except json.JSONDecodeError:
# If both orjson and json.loads fail, throw a proper error
json.dumps(parsed_body, ensure_ascii=False).encode("utf-8")
except (json.JSONDecodeError, UnicodeEncodeError):
# json.loads accepts lone surrogate escapes that no provider can encode
verbose_proxy_logger.error("Invalid JSON payload received: %s", e)
raise ProxyException(
message=f"Invalid JSON payload: {e}",

View file

@ -28,6 +28,9 @@ class _JWTKeyMappingRecord(Protocol):
@property
def id(self) -> str: ...
@property
def jwt_issuer(self) -> str: ...
@property
def jwt_claim_name(self) -> str: ...
@ -78,6 +81,7 @@ def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse:
"""Convert a Prisma mapping object to a safe response (no hashed token)."""
return JWTKeyMappingResponse(
id=mapping.id,
jwt_issuer=mapping.jwt_issuer or None,
jwt_claim_name=mapping.jwt_claim_name,
jwt_claim_value=mapping.jwt_claim_value,
description=mapping.description,
@ -109,6 +113,7 @@ async def create_jwt_key_mapping(
try:
hashed_key: Final = hash_token(data.key)
create_data: Final = {
"jwt_issuer": data.jwt_issuer or "",
"jwt_claim_name": data.jwt_claim_name,
"jwt_claim_value": data.jwt_claim_value,
"token": hashed_key,
@ -120,7 +125,7 @@ async def create_jwt_key_mapping(
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value)
cache_key: Final = jwt_key_mapping_cache_key(data.jwt_claim_name, data.jwt_claim_value, data.jwt_issuer)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return _to_response(new_mapping)
@ -131,7 +136,10 @@ async def create_jwt_key_mapping(
if "unique" in error_str or "p2002" in error_str:
raise HTTPException(
status_code=409,
detail=f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' already exists.",
detail=(
f"A mapping for claim '{data.jwt_claim_name}' = '{data.jwt_claim_value}' "
f"already exists for issuer '{data.jwt_issuer}'."
),
)
if "foreign" in error_str or "p2003" in error_str:
raise HTTPException(
@ -161,6 +169,9 @@ async def update_jwt_key_mapping(
update_data: Final = data.model_dump(exclude_unset=True, exclude={"id", "key"})
if data.key is not None:
update_data["token"] = hash_token(data.key)
if "jwt_issuer" in update_data:
# DB column is NOT NULL (see schema.prisma); "" is the global/unscoped sentinel.
update_data["jwt_issuer"] = update_data["jwt_issuer"] or ""
update_data["updated_by"] = user_api_key_dict.user_id
try:
@ -178,9 +189,11 @@ async def update_jwt_key_mapping(
# Evict only after the write commits: a concurrent request between an
# early eviction and the commit would re-cache the old mapping and keep
# it authorized until TTL.
old_cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
old_cache_key: Final = jwt_key_mapping_cache_key(
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
)
new_cache_key: Final = jwt_key_mapping_cache_key(
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value
updated_mapping.jwt_claim_name, updated_mapping.jwt_claim_value, updated_mapping.jwt_issuer
)
cache_keys: Final = (old_cache_key,) if old_cache_key == new_cache_key else (old_cache_key, new_cache_key)
await evict_and_broadcast(cache_keys=cache_keys, user_api_key_cache=user_api_key_cache)
@ -227,7 +240,9 @@ async def delete_jwt_key_mapping(
# Evict only after the row is gone, else a concurrent request can
# re-cache the deleted mapping and keep it authorized until TTL.
cache_key: Final = jwt_key_mapping_cache_key(old_mapping.jwt_claim_name, old_mapping.jwt_claim_value)
cache_key: Final = jwt_key_mapping_cache_key(
old_mapping.jwt_claim_name, old_mapping.jwt_claim_value, old_mapping.jwt_issuer
)
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
return {"status": "success"}
except HTTPException:

View file

@ -77,6 +77,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
from litellm.secret_managers.main import get_secret_str, str_to_bool
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
@ -1322,7 +1323,7 @@ def _resolve_vertex_model_from_router(
endpoint: str,
vertex_project: str | None,
vertex_location: str | None,
) -> tuple[str, str, str | None, str | None]:
) -> tuple[str, str, str | None, str | None, Mapping[str, object] | None]:
"""
Resolve Vertex AI model configuration from router.
@ -1335,18 +1336,21 @@ def _resolve_vertex_model_from_router(
vertex_location: Current vertex location (may be from URL)
Returns:
tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location)
with resolved values from router config
tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location, deployment_model_info)
with resolved values from router config; deployment_model_info is the resolved
deployment's `model_info`, or None when no deployment matched
"""
if not llm_router:
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
try:
deployment: Final = llm_router.get_available_deployment_for_pass_through(model=model_id)
if not deployment:
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
litellm_params: Final = deployment.get("litellm_params", {})
model_info: Final = deployment.get("model_info")
deployment_model_info: Final = model_info if isinstance(model_info, Mapping) else None
# Always override with router config values (they take precedence over URL values)
config_vertex_project: Final = litellm_params.get("vertex_project")
@ -1387,10 +1391,11 @@ def _resolve_vertex_model_from_router(
encoded_endpoint = encoded_endpoint.replace(model_id, actual_model)
endpoint = endpoint.replace(model_id, actual_model)
return encoded_endpoint, endpoint, vertex_project, vertex_location, deployment_model_info
except Exception as e:
verbose_proxy_logger.debug("Error resolving vertex model from router for model %s: %s", model_id, e)
return encoded_endpoint, endpoint, vertex_project, vertex_location
return encoded_endpoint, endpoint, vertex_project, vertex_location, None
def _is_bedrock_agent_runtime_route(endpoint: str) -> bool:
@ -2134,6 +2139,7 @@ async def _base_vertex_proxy_route(
endpoint,
vertex_project,
vertex_location,
deployment_model_info,
) = _resolve_vertex_model_from_router(
model_id=model_id,
llm_router=llm_router,
@ -2142,6 +2148,8 @@ async def _base_vertex_proxy_route(
vertex_project=vertex_project,
vertex_location=vertex_location,
)
if deployment_model_info:
setattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, deployment_model_info)
vertex_credentials: Final = passthrough_endpoint_router.get_vertex_credentials(
project_id=vertex_project,

View file

@ -96,6 +96,7 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
EndpointType,
@ -613,6 +614,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
_metadata.update(
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict)
)
_request_state: Final = getattr(request, "state", None)
deployment_model_info: Final = getattr(
_request_state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, None
)
if isinstance(deployment_model_info, Mapping):
_metadata["model_info"] = dict(deployment_model_info)
kwargs: Final = {
"litellm_params": {
@ -2002,6 +2009,8 @@ def create_pass_through_route(
delattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY)
if hasattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY):
delattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY)
if hasattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY):
delattr(request.state, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY)
# The upstream withholds its response headers until its first token, so
# the whole time-to-first-token is spent inside _relay with nothing on

View file

@ -17467,6 +17467,16 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie
"tab": "prompt_caching",
"description": "Empty uses Anthropic's 5m default. 1h suits long sessions but doubles the cache write cost.",
},
"openai_system_messages_first": {
"type": "Boolean",
"tab": "prompt_caching",
"description": (
"Moves system and developer messages to the front of the messages array on OpenAI and "
"Azure OpenAI chat completions requests, keeping their relative order. OpenAI's prompt cache "
"matches on the exact prefix, so a system message that arrives mid-conversation otherwise "
"breaks the cached prefix on every turn."
),
},
"budget_rollover": { # mutable-ok: registry literal, frozen with its siblings below
"type": "Boolean",
"description": (

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -33,6 +33,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils
from litellm.types.integrations.custom_logger import converted_stream_requested
from litellm.types.llms.openai import (
PART_UNION_TYPES,
ResponseAPIUsage,
@ -626,7 +627,9 @@ class BaseResponsesAPIStreamingIterator:
return
request_kwargs = getattr(caching_handler, "request_kwargs", None)
if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True:
if not _is_json_object(request_kwargs):
return
if request_kwargs.get("stream") is not True and not converted_stream_requested(request_kwargs):
return
request_kwargs = request_kwargs.copy()
preset_cache_key = getattr(caching_handler, "preset_cache_key", None)

View file

@ -79,7 +79,10 @@ from litellm.litellm_core_utils.core_helpers import (
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
from litellm.litellm_core_utils.get_llm_provider_logic import (
declared_authenticating_provider,
is_registered_custom_provider,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.ptu_pricing import (
PTU_COST_ATTRIBUTION_ENV_VAR,
@ -99,6 +102,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
mask_sensitive_structure,
)
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
from litellm.llms.base_llm.passthrough.transformation import replace_path_segment
from litellm.llms.base_llm.vector_store.transformation import (
RouterVectorStoreEmbeddingExecutor,
@ -3774,7 +3778,16 @@ class Router:
self, deployment: dict, kwargs: dict, function_name: str | None = None
) -> Deployment:
"""
Handle clientside credential
Build a per-request Deployment carrying the caller-supplied api_key/api_base,
with its own stable id for cooldown, logging, and cost-map identity.
This deployment is deliberately never registered with the router (no
upsert_deployment/add_deployment call): doing so used to add it to
self.model_list under the shared model_name, which made a request-scoped,
caller-supplied provider credential a permanent, load-balanced deployment
that every other caller of that model group could be routed onto. Its
pricing is still registered directly, so a custom price configured on the
underlying deployment still applies to this call.
"""
model_info: Final = deployment.get("model_info", {}).copy()
litellm_params: Final = deployment["litellm_params"].copy()
@ -3793,7 +3806,7 @@ class Router:
litellm_params=LiteLLM_Params(**dynamic_litellm_params),
model_info=model_info,
)
self.upsert_deployment(deployment=deployment_pydantic_obj) # add new deployment to router
Router._register_deployment_pricing(deployment=deployment_pydantic_obj)
return deployment_pydantic_obj
@staticmethod
@ -9554,8 +9567,10 @@ class Router:
)
# done reading model["litellm_params"]
# Check if provider is supported: either in enum or JSON-configured
if custom_llm_provider not in litellm.provider_list and not JSONProviderRegistry.exists(
custom_llm_provider
if (
custom_llm_provider not in litellm.provider_list
and not JSONProviderRegistry.exists(custom_llm_provider)
and not is_registered_custom_provider(custom_llm_provider)
):
raise Exception(f"Unsupported provider - {custom_llm_provider}")
@ -9701,40 +9716,7 @@ class Router:
# initialize client
self._add_deployment(deployment=deployment)
_model_info_dict: Final[dict] = deployment.model_info.model_dump(exclude_none=True)
for field in CustomPricingLiteLLMParams.model_fields:
field_value = deployment.litellm_params.get(field)
if field_value is not None:
_model_info_dict[field] = field_value
Router._inherit_builtin_base_rates_for_off_peak(
model_info=_model_info_dict,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
if _model_info_dict.get("input_cost_per_token") is not None:
Router._inherit_builtin_cache_pricing(
model_info=_model_info_dict,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._inherit_builtin_tiered_output_rate(
model_info=_model_info_dict,
backend_model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
# Register custom pricing in litellm.model_cost.
# Mirrors _create_deployment() logic to ensure dynamically-added deployments
# (e.g., loaded from DB) also have their custom pricing registered.
# Without this, _is_model_cost_zero() cannot detect explicitly-configured
# zero-cost models, causing budget checks to block free models.
Router._register_deployment_in_model_cost(
model_id=deployment.model_info.id,
model_info=_model_info_dict,
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
Router._register_deployment_pricing(deployment=deployment)
# add to model names
self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id)
@ -9996,6 +9978,21 @@ class Router:
)
return model_info
@staticmethod
def _register_deployment_pricing(deployment: Deployment) -> None:
"""Register a deployment's custom/inherited pricing in ``litellm.model_cost``.
Takes only a ``Deployment``, so it registers pricing for a deployment that
is never added to ``self.model_list`` (a per-request client-side-credential
deployment) just as readily as one that is.
"""
Router._register_deployment_in_model_cost(
model_id=deployment.model_info.id,
model_info=Router._deployment_model_cost_payload(deployment),
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
)
@staticmethod
def _register_deployment_in_model_cost(
*,
@ -10822,6 +10819,7 @@ class Router:
"model_group": user_facing_model_group_name,
"providers": [llm_provider],
**model_info,
"supports_fast_mode": True,
"supported_reasoning_efforts": None,
}
)
@ -10900,6 +10898,9 @@ class Router:
if model_info.get("rpm", None) is not None and _deployment_rpm is None:
_deployment_rpm = model_info.get("rpm")
model_group_info.supports_fast_mode = model_group_info.supports_fast_mode and (
AnthropicModelInfo.supports_fast_mode(litellm_model, llm_provider)
)
deployment_reasoning_efforts = (
resolve_supported_reasoning_efforts( # rebind-ok: recalculated per deployment
model_info, deployment_is_mapped=deployment_is_mapped

View file

@ -11,6 +11,9 @@ LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY: Final = "litellm_pass_through_custom
# exact byte/string body, such as AWS SigV4-signed requests.
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY: Final = "litellm_pass_through_raw_body"
# `model_info` of the router deployment a provider route (e.g. Vertex) resolved for this request.
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: Final = "litellm_pass_through_deployment_model_info"
# Attribute set on the FastAPI endpoint function of every user-defined pass-through
# route. Auth reads it off the dispatched endpoint (``request.scope["endpoint"]``) to
# decide whether a request body ``model`` names an upstream model rather than a

View file

@ -722,6 +722,7 @@ class ModelGroupInfo(BaseModel):
supports_url_context: bool = Field(default=False)
supports_reasoning: bool = Field(default=False)
supports_function_calling: bool = Field(default=False)
supports_fast_mode: bool = Field(default=False)
supported_reasoning_efforts: tuple[str, ...] | None = Field(default=None)
supported_openai_params: list[str] | None = Field(default=[])
configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None

View file

@ -846,6 +846,13 @@ def _is_streaming_response_for_correlation(result: object) -> bool:
return isinstance(result, CustomStreamWrapper)
def _is_converted_stream_result(result: object) -> bool:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
return isinstance(result, (CustomStreamWrapper, BaseResponsesAPIStreamingIterator))
# Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
def function_setup(
original_function: str,
@ -1889,6 +1896,9 @@ def client(original_function):
_caching_handler_response.cached_result is not None
and _caching_handler_response.final_embedding_cached_response is None
):
if _is_converted_stream_result(_caching_handler_response.cached_result):
logging_obj.stream = True
logging_obj.model_call_details["stream"] = True
return _caching_handler_response.cached_result
elif _caching_handler_response.embedding_all_elements_cache_hit is True:
@ -1946,10 +1956,9 @@ def client(original_function):
raise
end_time = datetime.datetime.now()
if _is_streaming_request(
kwargs=kwargs,
call_type=call_type,
):
if _is_streaming_request(kwargs=kwargs, call_type=call_type) or _is_converted_stream_result(result):
logging_obj.stream = True
logging_obj.model_call_details["stream"] = True
if "complete_response" in kwargs and kwargs["complete_response"] is True:
chunks: Final = []
for idx, chunk in enumerate(result):

File diff suppressed because it is too large Load diff

View file

@ -716,6 +716,9 @@
"supports_embedding_image_input": {
"type": "boolean"
},
"supports_fast_mode": {
"type": "boolean"
},
"supports_forced_tool_use": {
"type": "boolean"
},

View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.102.0"
version = "1.103.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -67,8 +67,8 @@ proxy = [
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.28.1,<2.0",
"litellm-proxy-extras==0.4.97",
"litellm-enterprise==0.1.67",
"litellm-proxy-extras==0.4.98",
"litellm-enterprise==0.1.68",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",
@ -290,6 +290,7 @@ editable-profile = "dev"
include = [
"litellm/proxy/_experimental/out/**",
"litellm/router_strategy/complexity_router/artifacts/*.json",
"litellm/proxy/client/cli/commands/codex_base_instructions.md",
]
exclude = [
"litellm/proxy/enterprise",
@ -331,7 +332,7 @@ members = ["enterprise", "litellm-proxy-extras"]
profile = "black"
[tool.commitizen]
version = "1.102.0"
version = "1.103.0"
version_files = [
"pyproject.toml:^version",
]

View file

@ -487,6 +487,10 @@ model LiteLLM_VerificationToken {
model LiteLLM_JWTKeyMapping {
id String @id @default(uuid())
jwt_issuer String @default("") // Scopes the mapping to one configured issuer; "" matches any issuer.
// Not nullable: Postgres unique constraints treat every NULL as
// distinct, so a nullable column would let multiple unscoped
// mappings collide on the same claim without a constraint violation.
jwt_claim_name String // e.g. "sub", "email"
jwt_claim_value String // The claim value to match
token String // Hashed virtual key (FK)
@ -499,8 +503,8 @@ model LiteLLM_JWTKeyMapping {
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@@unique([jwt_issuer, jwt_claim_name, jwt_claim_value])
@@index([jwt_issuer, jwt_claim_name, jwt_claim_value, is_active])
}
// Deprecated keys during grace period - allows old key to work until revoke_at

View file

@ -927,24 +927,14 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request():
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={"stream": True},
)
is True
)
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={"stream": False},
)
is False
)
assert (
_should_defer_streaming_cache_hit_callbacks(
kwargs={},
)
is False
logging_obj = MagicMock()
logging_obj.model_call_details = {}
stream_replay = CustomStreamWrapper(
completion_stream=iter(()), model="gpt-4o", logging_obj=logging_obj
)
assert _should_defer_streaming_cache_hit_callbacks(cached_result=stream_replay) is True
assert _should_defer_streaming_cache_hit_callbacks(cached_result=ModelResponse()) is False
assert _should_defer_streaming_cache_hit_callbacks(cached_result={"id": "msg_1"}) is False
@pytest.mark.asyncio

View file

@ -3,6 +3,7 @@
import sys, os, time
import traceback, asyncio
import httpx
import pytest
import litellm
@ -402,6 +403,10 @@ def test_router_redis_cache():
def test_router_handle_clientside_credential():
"""A caller-supplied credential must stay scoped to the current call: it must
never be registered as a router deployment, or a later caller with no override
of their own can be load-balanced onto it and reach the provider with someone
else's credential (see LIT-7811)."""
deployment = {
"model_name": "gemini/*",
"litellm_params": {"model": "gemini/*"},
@ -421,7 +426,67 @@ def test_router_handle_clientside_credential():
)
assert new_deployment.litellm_params.api_key == "123"
assert len(router.get_model_list()) == 2
assert len(router.get_model_list()) == 1
assert router.get_deployment(model_id=new_deployment.model_info.id) is None
async def test_router_clientside_credential_not_reused_by_other_callers(
respx_mock, monkeypatch: pytest.MonkeyPatch
):
"""End-to-end regression test for LIT-7811.
One caller's request-scoped api_key must never leak into a later, unrelated
caller's request. Before the fix, the router registered the caller-supplied
credential as a second, permanent deployment for the shared model group, so
plain follow-up calls with no override of their own could be load-balanced
onto it and reach the provider with the first caller's key.
"""
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 0,
"model": "gpt-4o",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
)
)
router = Router(
model_list=[
{
"model_name": "shared-model",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "configured-key"},
"model_info": {"id": "configured-deployment"},
}
]
)
await router.acompletion(
model="shared-model",
messages=[{"role": "user", "content": "hi"}],
api_key="alternate-tenant-key",
)
assert route.calls[-1].request.headers["authorization"] == "Bearer alternate-tenant-key"
# The forwarded credential must never become a routable deployment for the
# model group other callers share.
assert [d["model_info"]["id"] for d in router.get_model_list(model_name="shared-model")] == [
"configured-deployment"
]
for _ in range(20):
await router.acompletion(
model="shared-model",
messages=[{"role": "user", "content": "hi"}],
)
used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]}
assert used_auth_headers == {"Bearer configured-key"}
def test_router_get_async_openai_model_client():

View file

@ -91,6 +91,154 @@ async def test_jwt_to_virtual_key_mapping_resolution():
prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
@pytest.mark.asyncio
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
DIFFERENT issuer whose claim value happens to collide, even though both issuers
map the same claim field (``sub``) to a virtual key."""
issuer_a = "https://issuer-a.example.com"
issuer_b = "https://issuer-b.example.com"
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
virtual_key_claim_field="sub", virtual_key_mapping_cache_ttl=3600
)
rows = [
{
"jwt_issuer": issuer_b,
"jwt_claim_name": "sub",
"jwt_claim_value": "dev-alice",
"token": "hashed-issuer-b-key",
"is_active": True,
}
]
async def fake_find_first(where):
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return MagicMock(**row)
return None
prisma_client = MagicMock()
prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(side_effect=fake_find_first)
# Dependency-inject the resolved key via the cache (IdentityStore._resolve_key
# reads it from here) instead of monkeypatching IdentityStore itself.
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-issuer-b-key",
value=UserAPIKeyAuth(token="hashed-issuer-b-key", team_id="issuer-b-team"),
)
# The rightful owner: issuer-b's own claim resolves to its mapping.
owner_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert isinstance(owner_result, UserAPIKeyAuth)
assert owner_result.token == "hashed-issuer-b-key"
# A validly-signed token from issuer-a carrying the SAME claim value must not
# inherit issuer-b's mapping. Default behavior is fallback_team_mapping, so a
# correctly-scoped miss returns None instead of resolving to issuer-b's key.
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert colliding_result is None
@pytest.mark.asyncio
async def test_global_mapping_resolution_is_cached_under_the_global_key_not_the_requesting_issuer():
"""LIT-7417: caching a global (unscoped) mapping's hit under the REQUESTING
issuer's key would leave every issuer that falls back to it holding its own
stale copy after the row is updated/deleted -- CRUD only evicts the cache key
computed from the row's own scope (global), so a copy cached under some other
issuer's key would keep resolving to the old token until TTL. Caching it under
the global key instead means every issuer shares (and CRUD correctly evicts)
the exact same entry."""
issuer_a = "https://issuer-a.example.com"
issuer_b = "https://issuer-b.example.com"
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
issuers=[
{
"issuer": issuer_a,
"jwks_url": f"{issuer_a}/jwks",
"virtual_key_claim_field": "sub",
"disable_audience_validation": True,
},
{
"issuer": issuer_b,
"jwks_url": f"{issuer_b}/jwks",
"virtual_key_claim_field": "sub",
"disable_audience_validation": True,
},
]
)
rows = [
{
"jwt_issuer": "",
"jwt_claim_name": "sub",
"jwt_claim_value": "legacy-user",
"token": "hashed-legacy-key",
"is_active": True,
}
]
async def fake_find_first(where):
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return MagicMock(**row)
return None
prisma_client = MagicMock()
find_first = AsyncMock(side_effect=fake_find_first)
prisma_client.db.litellm_jwtkeymapping.find_first = find_first
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-legacy-key",
value=UserAPIKeyAuth(token="hashed-legacy-key", team_id="legacy-team"),
)
resolved_a = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_a, "sub": "legacy-user"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert isinstance(resolved_a, UserAPIKeyAuth)
assert find_first.await_count == 2 # issuer-a-scoped miss, then global hit
# issuer-b resolving the SAME global mapping must hit the cache issuer-a's
# resolution populated, not issue a fresh DB query for the global row again.
resolved_b = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer_b, "sub": "legacy-user"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert isinstance(resolved_b, UserAPIKeyAuth)
assert resolved_b.token == "hashed-legacy-key"
assert find_first.await_count == 3 # +1 for issuer-b's own issuer-scoped miss; global tier served from cache
@pytest.mark.asyncio
async def test_jwt_to_virtual_key_mapping_no_mapping():
"""
@ -223,6 +371,7 @@ def test_to_response_excludes_token():
now = datetime.now(timezone.utc)
mock_mapping = MagicMock()
mock_mapping.id = "mapping-1"
mock_mapping.jwt_issuer = None
mock_mapping.jwt_claim_name = "email"
mock_mapping.jwt_claim_value = "user@example.com"
mock_mapping.token = "hashed_secret_value"
@ -275,10 +424,12 @@ def _mock_mapping(
id="mapping-1",
claim_name="email",
claim_value="user@example.com",
issuer=None,
):
now = datetime.now(timezone.utc)
m = MagicMock()
m.id = id
m.jwt_issuer = issuer
m.jwt_claim_name = claim_name
m.jwt_claim_value = claim_value
m.token = "hashed_token"
@ -485,6 +636,35 @@ async def test_create_success_returns_response_without_token():
assert result.jwt_claim_name == "email"
@pytest.mark.asyncio
async def test_create_without_issuer_stores_empty_string_not_null():
"""LIT-7417: the DB column is NOT NULL (see schema.prisma). Storing a real NULL
for an unscoped mapping would let Postgres accept unlimited duplicate unscoped
rows for the same claim (NULL is never equal to NULL in a unique constraint),
so two mappings for the same claim value could point at two different keys with
no conflict, and resolution would pick whichever one Postgres returns first."""
from litellm.proxy._types import CreateJWTKeyMappingRequest
mock_prisma = _mock_prisma()
mock_prisma.db.litellm_jwtkeymapping.create.return_value = _mock_mapping()
mock_cache = AsyncMock()
data = CreateJWTKeyMappingRequest(jwt_claim_name="sub", jwt_claim_value="dev-alice", key="sk-test-key")
with (
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
"litellm.proxy.proxy_server.prisma_client", mock_prisma
),
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
"litellm.proxy.proxy_server.user_api_key_cache", mock_cache
),
):
await create_jwt_key_mapping(data=data, user_api_key_dict=_make_admin_auth())
sent_data = mock_prisma.db.litellm_jwtkeymapping.create.call_args.kwargs["data"]
assert sent_data["jwt_issuer"] == ""
# ──────────────────────────────────────────────
# Tests: unregistered_jwt_client_behavior
# ──────────────────────────────────────────────

View file

@ -2099,8 +2099,12 @@ def test_handle_clientside_credential_metadata_loading(
assert result_deployment.model_info.id != "original-id-123"
assert result_deployment.model_info.original_model_id == "original-id-123"
# Verify the deployment was added to the router
assert len(router.model_list) == len(model_list) + 1
# The caller-supplied credential must stay scoped to this call: it must never be
# registered as a router deployment, or a later caller with no override of their
# own could be load-balanced onto it and reach the provider with this credential
# (see LIT-7811).
assert len(router.model_list) == len(model_list)
assert router.get_deployment(model_id=result_deployment.model_info.id) is None
# Test that the function correctly uses the right metadata key
# For acompletion, it should use "metadata"
@ -2260,14 +2264,63 @@ def test_handle_clientside_credential_with_responses_function(model_list):
assert result_deployment.model_info.id != "original-id-responses"
assert result_deployment.model_info.original_model_id == "original-id-responses"
# Verify the deployment was added to the router
assert len(router.model_list) == len(model_list) + 1
# The caller-supplied credential must stay scoped to this call: it must never be
# registered as a router deployment (see LIT-7811).
assert len(router.model_list) == len(model_list)
assert router.get_deployment(model_id=result_deployment.model_info.id) is None
print(
"✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata"
)
def test_handle_clientside_credential_still_registers_custom_pricing(model_list):
"""A clientside-credential call must still price against the deployment's own
custom rate, even though the call's ephemeral deployment is never added to the
router (see LIT-7811): losing that registration would silently fall back to
public catalog pricing for every clientside-credential call on a deployment
with a custom rate configured."""
router = Router(model_list=model_list)
deployment = {
"model_name": "gpt-4.1",
"litellm_params": {
"model": "gpt-4.1",
"api_key": "test_key",
"input_cost_per_token": 0.0001234,
"output_cost_per_token": 0.0005678,
},
"model_info": {"id": "original-id-pricing"},
}
kwargs = {"api_key": "client_side_key", "metadata": {"model_group": "gpt-4.1"}}
result_deployment = router._handle_clientside_credential(
deployment=deployment, kwargs=kwargs, function_name="acompletion"
)
registered = litellm.model_cost.get(result_deployment.model_info.id)
assert registered is not None
assert registered["input_cost_per_token"] == 0.0001234
assert registered["output_cost_per_token"] == 0.0005678
def test_register_deployment_pricing_direct_call():
"""Direct-call unit test for the pricing-registration helper `_handle_clientside_credential`
relies on, so it prices a deployment that is deliberately never added to `self.model_list`."""
deployment = Deployment(
model_name="gpt-4.1",
litellm_params=LiteLLM_Params(
model="gpt-4.1",
api_key="test_key",
input_cost_per_token=0.0009999,
),
model_info=ModelInfo(id="direct-call-pricing-id"),
)
Router._register_deployment_pricing(deployment=deployment)
assert litellm.model_cost["direct-call-pricing-id"]["input_cost_per_token"] == 0.0009999
def test_get_metadata_variable_name_from_kwargs(model_list):
"""
Test _get_metadata_variable_name_from_kwargs method returns correct metadata variable name based on kwargs content.

View file

@ -693,3 +693,90 @@ async def test_cache_hit_records_the_looked_up_key_as_the_preset_cache_key(monke
assert handler.preset_cache_key is not None
assert logging_obj.litellm_params["preset_cache_key"] == handler.preset_cache_key
assert hit.cached_result._hidden_params["cache_key"] == handler.preset_cache_key
@pytest.mark.asyncio
async def test_converted_stream_cache_hit_replayed_as_plain_object_logs_at_hit_time(monkeypatch):
import litellm
from litellm.caching.caching import Cache
from litellm.types.utils import CallTypes
async def aanthropic_messages(**kwargs):
return None
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
kwargs = {
"model": "claude-sonnet-5",
"messages": [{"role": "user", "content": "hello"}],
"max_tokens": 16,
"caching": True,
"stream": False,
"_websearch_interception_converted_stream": True,
}
cached_message = {
"id": "msg_1",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "hi"}],
}
await litellm.cache.async_add_cache(cached_message, **kwargs)
handler = LLMCachingHandler(original_function=aanthropic_messages, request_kwargs=kwargs, start_time=datetime.now())
logging_obj = _build_logging_obj(CallTypes.aanthropic_messages.value, stream=False)
logging_obj.async_success_handler = AsyncMock()
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
hit = await handler._async_get_cache(
model="claude-sonnet-5",
original_function=aanthropic_messages,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.aanthropic_messages.value,
kwargs=kwargs,
args=(),
)
assert hit is not None and hit.cached_result == cached_message
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True
@pytest.mark.asyncio
async def test_agentic_loop_followup_cache_hit_with_converted_stream_marker_replays_as_plain_object(monkeypatch):
import litellm
from litellm.caching.caching import Cache
from litellm.types.utils import CallTypes
async def acompletion(**kwargs):
return None
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
kwargs = {
"model": "gpt-5.6",
"messages": [{"role": "user", "content": "run the code"}],
"caching": True,
"stream": False,
"_code_interpreter_interception_converted_stream": True,
"_agentic_loop_depth": 1,
}
await litellm.cache.async_add_cache(
litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "done"}}]), **kwargs
)
handler = LLMCachingHandler(original_function=acompletion, request_kwargs=kwargs, start_time=datetime.now())
logging_obj = _build_logging_obj(CallTypes.acompletion.value, stream=False)
logging_obj.async_success_handler = AsyncMock()
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
hit = await handler._async_get_cache(
model="gpt-5.6",
original_function=acompletion,
logging_obj=logging_obj,
start_time=datetime.now(),
call_type=CallTypes.acompletion.value,
kwargs=kwargs,
args=(),
)
assert hit is not None and isinstance(hit.cached_result, litellm.ModelResponse)
assert hit.cached_result.choices[0].message.content == "done"
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["cache_hit"] is True

View file

@ -1,5 +1,6 @@
import asyncio
import os
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -7,6 +8,7 @@ import pytest
import litellm
from litellm.integrations.langsmith import LangsmithLogger
from litellm.types.integrations.langsmith import LangsmithQueueObject
@pytest.fixture
@ -531,3 +533,44 @@ class TestLangsmithRootRunIdConsistency:
assert data["trace_id"] == "trace-1"
assert data["dotted_order"] == dotted
@pytest.mark.asyncio
async def test_events_appended_during_flush_are_not_dropped():
logger = LangsmithLogger(langsmith_api_key="test-key", langsmith_project="test-project")
try:
sent_batches: Final[list[list[dict[str, str]]]] = []
late_event: Final = LangsmithQueueObject(
credentials=logger.default_credentials, data={"id": "late"}
)
async def fake_post(
url: str, json: dict[str, list[dict[str, str]]], headers: dict[str, str]
) -> MagicMock:
if not sent_batches:
logger.log_queue.append(late_event)
sent_batches.append(json["post"])
response = MagicMock()
response.status_code = 200
response.raise_for_status = MagicMock()
return response
logger.async_httpx_client = MagicMock(post=AsyncMock(side_effect=fake_post))
logger.log_queue = [
LangsmithQueueObject(credentials=logger.default_credentials, data={"id": "a"}),
LangsmithQueueObject(credentials=logger.default_credentials, data={"id": "b"}),
]
await logger.flush_queue()
assert [e["id"] for e in sent_batches[0]] == ["a", "b"]
assert logger.log_queue == [late_event]
await logger.flush_queue()
assert [e["id"] for e in sent_batches[1]] == ["late"]
assert logger.log_queue == []
finally:
if logger._flush_task is not None:
logger._flush_task.cancel()
await asyncio.gather(logger._flush_task, return_exceptions=True)

View file

@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
responses_reasoning_items_from_thinking_blocks,
split_concatenated_json_objects,
strip_encrypted_reasoning_from_messages,
system_messages_first,
update_messages_with_model_file_ids,
)
@ -1107,6 +1108,38 @@ def test_drop_tool_reference_parts_leaves_non_tool_messages_alone():
assert result[2]["content"] == ""
class TestSystemMessagesFirst:
def test_stable_partition_keeps_order_within_each_group(self):
messages = [
{"role": "user", "content": "u1"},
{"role": "system", "content": "s1"},
{"role": "assistant", "content": "a1"},
{"role": "developer", "content": "d1"},
{"role": "tool", "tool_call_id": "c1", "content": "t1"},
{"role": "system", "content": "s2"},
]
result = system_messages_first(messages)
assert [m["content"] for m in result] == ["s1", "d1", "s2", "u1", "a1", "t1"]
assert [m["content"] for m in messages] == ["u1", "s1", "a1", "d1", "t1", "s2"]
assert all(
result_message is original for result_message, original in zip(result[3:], messages[::2], strict=True)
)
@pytest.mark.parametrize(
"messages",
[
[],
[{"role": "user", "content": "u1"}, {"role": "assistant", "content": "a1"}],
[{"role": "system", "content": "s1"}, {"role": "user", "content": "u1"}],
[{"role": "system", "content": "s1"}, {"role": "system", "content": "s2"}],
],
)
def test_already_ordered_messages_come_back_unchanged(self, messages):
assert system_messages_first(messages) == messages
class TestFlattenTopLevelSchemaCombinators:
def _customer_anyof_schema(self):
return {

View file

@ -2,6 +2,7 @@ import base64
import json
import logging
import os
import re
from typing import Final
from unittest.mock import MagicMock, patch
@ -19,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
_convert_to_bedrock_tool_call_invoke,
_convert_to_bedrock_tool_call_result,
anthropic_messages_pt,
convert_to_anthropic_tool_result,
convert_to_gemini_tool_call_result,
make_valid_bedrock_tool_name,
ollama_pt,
@ -2208,6 +2210,104 @@ def test_bedrock_tool_call_invoke_empty_arguments():
assert result[0]["toolUse"]["input"] == {}
_BEDROCK_TOOL_USE_ID_RE = re.compile(r"^[a-zA-Z0-9_.:-]{1,64}$")
@pytest.mark.parametrize(
"tool_call_id",
[
"call_" + "x" * 100,
"call|with|pipes",
"call_" + "y" * 60 + "|end",
"call:ok.dots-and_under",
"",
],
)
def test_bedrock_tool_use_id_is_sanitized_consistently_for_invoke_and_result(tool_call_id):
"""
Regression test for https://github.com/BerriAI/litellm/issues/34239: client-minted
tool_call ids longer than 64 chars or with chars outside [a-zA-Z0-9_.:-] made Bedrock
return a 400. The invoke and result paths must produce the same valid toolUseId so the
toolUse/toolResult pair still correlates.
"""
invoke = _convert_to_bedrock_tool_call_invoke(
[
{
"id": tool_call_id,
"type": "function",
"function": {"name": "get_weather", "arguments": '{"location": "Boston"}'},
}
]
)
result = _convert_to_bedrock_tool_call_result(
{"tool_call_id": tool_call_id, "role": "tool", "name": "get_weather", "content": "sunny"}
)
tool_use_id = invoke[0]["toolUse"]["toolUseId"]
assert _BEDROCK_TOOL_USE_ID_RE.match(tool_use_id)
assert result["toolResult"]["toolUseId"] == tool_use_id
def test_bedrock_tool_use_id_valid_ids_pass_through_unchanged():
result = _convert_to_bedrock_tool_call_result(
{"tool_call_id": "tooluse_Ab.c:1-2_3", "role": "tool", "name": "f", "content": "ok"}
)
assert result["toolResult"]["toolUseId"] == "tooluse_Ab.c:1-2_3"
def test_bedrock_tool_use_id_truncation_keeps_distinct_ids_distinct():
prefix = "call_" + "z" * 70
ids = {
_convert_to_bedrock_tool_call_result(
{"tool_call_id": f"{prefix}{suffix}", "role": "tool", "name": "f", "content": "ok"}
)["toolResult"]["toolUseId"]
for suffix in ("a", "b")
}
assert len(ids) == 2
assert all(len(i) == 64 for i in ids)
def test_bedrock_tool_use_id_replaced_chars_do_not_collide_with_existing_ids():
ids = {
_convert_to_bedrock_tool_call_result({"tool_call_id": i, "role": "tool", "name": "f", "content": "ok"})[
"toolResult"
]["toolUseId"]
for i in ("call|x", "call_x")
}
assert len(ids) == 2
def test_bedrock_tool_call_invoke_concatenated_json_long_id_stays_within_limit():
long_id = "call_" + "q" * 62
result = _convert_to_bedrock_tool_call_invoke(
[
{
"id": long_id,
"type": "function",
"function": {"name": "run", "arguments": '{"cmd":"a"}{"cmd":"b"}'},
}
]
)
ids = [block["toolUse"]["toolUseId"] for block in result]
assert len(ids) == 2
assert len(set(ids)) == 2
assert all(_BEDROCK_TOOL_USE_ID_RE.match(i) for i in ids)
@pytest.mark.parametrize(
("tool_call_id", "expected"),
[
("call|with|pipes", "call_with_pipes"),
("call:ok.dots", "call_ok_dots"),
("call_" + "x" * 100, "call_" + "x" * 100),
("toolu_01AbC-xyz", "toolu_01AbC-xyz"),
("", "tool_use_id"),
],
)
def test_anthropic_tool_use_id_keeps_pattern_only_rewrite_with_no_cap_or_hash(tool_call_id, expected):
result = convert_to_anthropic_tool_result({"role": "tool", "tool_call_id": tool_call_id, "content": "ok"})
assert result["tool_use_id"] == expected
def test_bedrock_tool_call_invoke_concatenated_json():
"""
Tool call whose arguments contain multiple concatenated JSON objects

View file

@ -0,0 +1,55 @@
from typing import Final
import pytest
import litellm
from litellm import CustomLLM
from litellm.litellm_core_utils.get_llm_provider_logic import (
get_llm_provider,
is_registered_custom_provider,
)
CUSTOM_PROVIDER: Final = "test-onprem-llm"
@pytest.fixture
def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str:
monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": CUSTOM_PROVIDER, "custom_handler": CustomLLM()}])
monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list))
monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers))
return CUSTOM_PROVIDER
def test_get_llm_provider_resolves_custom_provider_map_prefix_before_first_completion(
registered_custom_provider: str,
) -> None:
assert registered_custom_provider not in litellm.provider_list
model, provider, dynamic_api_key, api_base = get_llm_provider(model=f"{registered_custom_provider}/my-model")
assert (model, provider, dynamic_api_key, api_base) == ("my-model", registered_custom_provider, None, None)
def test_get_llm_provider_strips_prefix_when_custom_provider_passed_explicitly(
registered_custom_provider: str,
) -> None:
model, provider, _, api_base = get_llm_provider(
model="my-model",
custom_llm_provider=registered_custom_provider,
api_base="http://onprem.internal:8080",
)
assert (model, provider, api_base) == ("my-model", registered_custom_provider, "http://onprem.internal:8080")
def test_get_llm_provider_still_rejects_unregistered_prefix(registered_custom_provider: str) -> None:
with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"):
get_llm_provider(model="not-registered-llm/my-model")
@pytest.mark.parametrize(
("candidate", "expected"),
[(CUSTOM_PROVIDER, True), ("not-registered-llm", False), (None, False), ("", False)],
)
def test_is_registered_custom_provider(registered_custom_provider: str, candidate: str | None, expected: bool) -> None:
assert is_registered_custom_provider(candidate) is expected

View file

@ -132,6 +132,32 @@ def test_transform_request_drops_tool_reference_parts():
assert request["messages"][2]["content"] == ""
@pytest.mark.parametrize(
"enabled, expected", [(False, ("hi", "sys", "reply", "more")), (True, ("sys", "hi", "reply", "more"))]
)
def test_transform_request_system_messages_first_follows_global_flag(monkeypatch, enabled, expected):
"""Azure OpenAI shares OpenAI's prefix-matched prompt cache, so the same flag moves
system messages ahead of the conversation on the Azure request body."""
monkeypatch.setattr(litellm, "openai_system_messages_first", enabled)
messages = [
{"role": "user", "content": "hi"},
{"role": "system", "content": "sys"},
{"role": "assistant", "content": "reply"},
{"role": "user", "content": "more"},
]
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=messages,
optional_params={},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert tuple(m["content"] for m in request["messages"]) == expected
assert [m["content"] for m in messages] == ["hi", "sys", "reply", "more"]
@pytest.mark.parametrize(
"model, emitted_key, absent_key",
[

View file

@ -68,3 +68,24 @@ def test_azure_o_series_transform_request_flattens_top_level_anyof():
assert parameters["required"] == ["id"]
assert "anyOf" in tool["function"]["parameters"]
assert optional_params["tools"][0] is tool
def test_azure_o_series_transform_request_moves_system_messages_first(monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
messages = [
{"role": "user", "content": "hi"},
{"role": "developer", "content": "dev"},
{"role": "assistant", "content": "reply"},
{"role": "user", "content": "more"},
]
request = AzureOpenAIO1Config().transform_request(
model="o3-mini",
messages=messages,
optional_params={},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert [m["content"] for m in request["messages"]] == ["dev", "hi", "reply", "more"]
assert [m["content"] for m in messages] == ["hi", "dev", "reply", "more"]

View file

@ -1124,6 +1124,80 @@ class TestToolReferenceStripping:
assert request["messages"][2]["content"] == ""
class TestSystemMessagesFirst:
"""With litellm.openai_system_messages_first on, requests bound for OpenAI put system and
developer messages ahead of the conversation, keeping each group's order, so the instruction
prefix stays byte-stable for OpenAI's prefix-matched prompt cache."""
MESSAGES: Final = (
{"role": "user", "content": "first turn"},
{"role": "system", "content": "sys 1"},
{"role": "assistant", "content": "reply"},
{"role": "developer", "content": "dev"},
{"role": "user", "content": "second turn"},
{"role": "system", "content": "sys 2"},
)
ORIGINAL_ORDER: Final = ("first turn", "sys 1", "reply", "dev", "second turn", "sys 2")
ORDERED: Final = ("sys 1", "dev", "sys 2", "first turn", "reply", "second turn")
def setup_method(self):
self.config = OpenAIGPTConfig()
def _messages(self):
return [dict(m) for m in self.MESSAGES]
def _transform(self, provider):
return self.config.transform_request(
model="gpt-4.1",
messages=self._messages(),
optional_params={},
litellm_params={"custom_llm_provider": provider},
headers={},
)
def test_default_off_keeps_caller_order(self, monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", False)
assert tuple(m["content"] for m in self._transform("openai")["messages"]) == self.ORIGINAL_ORDER
def test_moves_system_and_developer_messages_first_for_openai(self, monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
assert tuple(m["content"] for m in self._transform("openai")["messages"]) == self.ORDERED
def test_leaves_openai_compatible_providers_alone(self, monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
assert tuple(m["content"] for m in self._transform("deepseek")["messages"]) == self.ORIGINAL_ORDER
def test_does_not_mutate_caller_messages(self, monkeypatch):
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
messages = self._messages()
self.config.transform_request(
model="gpt-4.1",
messages=messages,
optional_params={},
litellm_params={"custom_llm_provider": "openai"},
headers={},
)
assert tuple(m["content"] for m in messages) == self.ORIGINAL_ORDER
@pytest.mark.asyncio
async def test_async_transform_request_moves_system_messages_first(self, monkeypatch):
class UninstantiatedOpenAIGPTConfig(OpenAIGPTConfig):
_is_base_class = True
def __init__(self) -> None:
pass
monkeypatch.setattr(litellm, "openai_system_messages_first", True)
request = await UninstantiatedOpenAIGPTConfig().async_transform_request(
model="gpt-4.1",
messages=self._messages(),
optional_params={},
litellm_params={"custom_llm_provider": "openai"},
headers={},
)
assert tuple(m["content"] for m in request["messages"]) == self.ORDERED
class TestOpenAIPromptCacheBreakpointChatPath:
"""Chat-path shape for OpenAI explicit prompt caching (#37509)."""

View file

@ -32,7 +32,13 @@ from litellm.proxy._types import (
JWTRoutingOverride,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object
from litellm.proxy.auth.auth_checks import (
TeamNotFoundError,
UserNotFoundError,
get_key_object,
_cache_key_object,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import (
_check_key_model_budget_with_fallback,
@ -7948,13 +7954,38 @@ def _per_issuer_virtual_key_jwt_handler(
def _fake_prisma_with_jwt_key_mapping(hashed_token: str | None) -> tuple[SimpleNamespace, AsyncMock]:
"""Every ``find_first`` call (issuer-scoped or global fallback) resolves the same way."""
find_first = AsyncMock(return_value=None if hashed_token is None else SimpleNamespace(token=hashed_token))
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
return prisma_client, find_first
def _mapping_where(claim_name: str, claim_value: str) -> dict[str, str | bool]:
return {"jwt_claim_name": claim_name, "jwt_claim_value": claim_value, "is_active": True}
def _fake_prisma_jwt_key_mapping_table(rows: list[dict[str, object]]) -> tuple[SimpleNamespace, AsyncMock]:
"""A ``find_first`` whose result depends on the ``where`` clause, like a real table.
Matches a row when every key present in ``where`` equals that key on the row --
a key ``get_jwt_key_mapping_object`` omits (e.g. old, issuer-blind code never
sending ``jwt_issuer``) does not constrain the match, exactly like Prisma.
"""
async def _find_first(where: dict[str, object]) -> SimpleNamespace | None:
for row in rows:
if all(row.get(k) == v for k, v in where.items()):
return SimpleNamespace(**row)
return None
find_first = AsyncMock(side_effect=_find_first)
prisma_client = SimpleNamespace(db=SimpleNamespace(litellm_jwtkeymapping=SimpleNamespace(find_first=find_first)))
return prisma_client, find_first
def _mapping_where(claim_name: str, claim_value: str, jwt_issuer: str | None) -> dict[str, str | bool]:
return {
"jwt_claim_name": claim_name,
"jwt_claim_value": claim_value,
"jwt_issuer": jwt_issuer or "",
"is_active": True,
}
@pytest.mark.asyncio
@ -7978,11 +8009,13 @@ async def test_per_issuer_virtual_key_claim_field_selects_the_issuer_mapping_for
proxy_logging_obj=MagicMock(),
)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7"))
# Issuer-scoped lookup hits on the first query, so no global fallback query runs.
find_first.assert_awaited_once_with(where=_mapping_where("sub", "svc-account-7", ISSUER_TWO))
assert isinstance(resolved, UserAPIKeyAuth)
assert resolved.token == "hashed-mapped-key"
assert resolved.team_id == "svc-team"
assert await user_api_key_cache.async_get_cache("jwt_key_mapping:sub:svc-account-7") == "hashed-mapped-key"
cache_key = jwt_key_mapping_cache_key("sub", "svc-account-7", ISSUER_TWO)
assert await user_api_key_cache.async_get_cache(cache_key) == "hashed-mapped-key"
@pytest.mark.asyncio
@ -8015,7 +8048,11 @@ async def test_per_issuer_reject_behavior_does_not_leak_into_the_team_issuer():
assert exc.value.status_code == 403
assert "No registered mapping for sub='unknown-svc'" in str(exc.value.detail)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "unknown-svc"))
# REJECT checks the issuer-scoped row first, then falls back to a global (NULL-issuer) row.
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("sub", "unknown-svc", ISSUER_TWO),
_mapping_where("sub", "unknown-svc", None),
]
@pytest.mark.asyncio
@ -8025,7 +8062,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub", global_behavior="auto_register")
prisma_client, find_first = _fake_prisma_with_jwt_key_mapping(None)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(key="jwt_key_mapping:sub:admin-7", value=_JWT_PROXY_ADMIN_SENTINEL)
# Sentinel cached under issuer-one's own key -- must never answer issuer-two's lookup.
await user_api_key_cache.async_set_cache(
key=jwt_key_mapping_cache_key("sub", "admin-7", ISSUER_ONE), value=_JWT_PROXY_ADMIN_SENTINEL
)
auto_register_issuer_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "admin-7"},
@ -8050,7 +8090,10 @@ async def test_proxy_admin_sentinel_cached_by_another_issuer_does_not_bypass_rej
assert exc.value.status_code == 403
assert "No registered mapping for sub='admin-7'" in str(exc.value.detail)
find_first.assert_awaited_once_with(where=_mapping_where("sub", "admin-7"))
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("sub", "admin-7", ISSUER_TWO),
_mapping_where("sub", "admin-7", None),
]
@pytest.mark.asyncio
@ -8079,7 +8122,125 @@ async def test_issuer_without_virtual_key_claim_field_falls_back_to_the_global_f
assert with_claim is None
assert without_claim is None
find_first.assert_awaited_once_with(where=_mapping_where("client_id", "app-9"))
# without_claim has no claim value and returns before ever reaching the DB.
assert [c.kwargs["where"] for c in find_first.await_args_list] == [
_mapping_where("client_id", "app-9", ISSUER_ONE),
_mapping_where("client_id", "app-9", None),
]
@pytest.mark.asyncio
async def test_colliding_claim_value_from_another_issuer_does_not_resolve_to_the_wrong_virtual_key():
"""LIT-7417: a mapping registered for one issuer must not answer a lookup from a
DIFFERENT issuer whose claim value happens to collide, even though both issuers
use the same claim field (``sub``) for their virtual-key mapping."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table(
[
{
"jwt_issuer": ISSUER_TWO,
"jwt_claim_name": "sub",
"jwt_claim_value": "dev-alice",
"token": "hashed-issuer-b-key",
"is_active": True,
}
]
)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-issuer-b-key",
value=UserAPIKeyAuth(token="hashed-issuer-b-key", api_key="hashed-issuer-b-key", team_id="issuer-b-team"),
)
owner_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_TWO, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert isinstance(owner_result, UserAPIKeyAuth)
assert owner_result.token == "hashed-issuer-b-key"
# issuer-one's behavior is fallback_team_mapping: a correctly-scoped miss must
# return None (fall through to team-based JWT auth), never issuer-two's key.
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert colliding_result is None
assert find_first.await_count == 3 # owner hit (1 call) + colliding miss (issuer-scoped + global fallback)
@pytest.mark.asyncio
async def test_cached_resolution_for_one_issuer_does_not_leak_to_a_colliding_issuer():
"""A cached positive resolution must be keyed by issuer too, or a colliding
claim value from another issuer could be served straight from cache without
ever reaching the (correctly issuer-scoped) DB lookup."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, find_first = _fake_prisma_jwt_key_mapping_table([])
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key=jwt_key_mapping_cache_key("sub", "dev-alice", ISSUER_TWO), value="hashed-issuer-b-key"
)
colliding_result = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: ISSUER_ONE, "sub": "dev-alice"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert colliding_result is None
# Must have gone to the DB rather than serving issuer-two's cached token.
assert find_first.await_count == 2
@pytest.mark.asyncio
async def test_issuer_agnostic_mapping_matches_every_issuer():
"""A mapping created before issuer scoping existed (``jwt_issuer`` is NULL) keeps
matching any issuer, so existing global mappings are not broken by this fix."""
from litellm.proxy.auth.user_api_key_auth import _resolve_jwt_to_virtual_key
jwt_handler = _per_issuer_virtual_key_jwt_handler(global_claim_field="sub")
prisma_client, _find_first = _fake_prisma_jwt_key_mapping_table(
[
{
"jwt_issuer": "",
"jwt_claim_name": "sub",
"jwt_claim_value": "legacy-user",
"token": "hashed-legacy-key",
"is_active": True,
}
]
)
user_api_key_cache = DualCache()
await user_api_key_cache.async_set_cache(
key="hashed-legacy-key",
value=UserAPIKeyAuth(token="hashed-legacy-key", api_key="hashed-legacy-key", team_id="legacy-team"),
)
for issuer in (ISSUER_ONE, ISSUER_TWO):
resolved = await _resolve_jwt_to_virtual_key(
jwt_claims={JWTHandler.LITELLM_JWT_ISSUER_CLAIM: issuer, "sub": "legacy-user"},
jwt_handler=jwt_handler,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
)
assert isinstance(resolved, UserAPIKeyAuth)
assert resolved.token == "hashed-legacy-key"
@pytest.mark.asyncio

View file

@ -1,7 +1,9 @@
import inspect
import json
import os
import subprocess
import sys
from pathlib import Path
from unittest.mock import patch
import click
@ -9,10 +11,9 @@ import pytest
import requests
from click.testing import CliRunner
from litellm.proxy.client.cli.commands.agents import (
AgentRunError,
ModelSyncArgs,
ModelSyncSkipped,
_hand_off,
_replace_process,
@ -22,6 +23,7 @@ from litellm.proxy.client.cli.commands.agents import (
agent_model_sync_env,
agent_profile,
build_agent_env,
codex_model_sync_args,
opencode_model_sync_env,
run_agent,
verify_proxy_key,
@ -55,6 +57,112 @@ class _Recorder:
return self.returns
_STOCK_REASONING_LEVELS = [
{"effort": "low", "description": "Fast responses with lighter reasoning"},
{"effort": "medium", "description": "Balances speed and reasoning depth for everyday tasks"},
{"effort": "high", "description": "Greater reasoning depth for complex problems"},
]
_STOCK_MODELS = {
"gpt-5.6-terra": {
"slug": "gpt-5.6-terra",
"display_name": "GPT-5.6 Terra",
"description": "Balanced agentic coding model for everyday work.",
"default_reasoning_level": "medium",
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
"shell_type": "unified_exec",
"visibility": "list",
"supported_in_api": True,
"priority": 7,
"availability_nux": None,
"upgrade": None,
"base_instructions": "You are Codex, a coding agent based on GPT-5.6.",
"apply_patch_tool_type": "freeform",
"supports_parallel_tool_calls": True,
"context_window": 272000,
"comp_hash": "terra-hash",
},
"gpt-5.5": {
"slug": "gpt-5.5",
"display_name": "GPT-5.5",
"description": "Frontier model for complex coding, research, and real-world work.",
"default_reasoning_level": "medium",
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
"shell_type": "unified_exec",
"visibility": "list",
"supported_in_api": True,
"priority": 12,
"availability_nux": None,
"upgrade": None,
"base_instructions": "You are Codex, a coding agent based on GPT-5.",
"apply_patch_tool_type": "freeform",
"supports_parallel_tool_calls": True,
"context_window": 272000,
"comp_hash": "gpt-5.5-hash",
},
"gpt-5.4": {
"slug": "gpt-5.4",
"display_name": "GPT-5.4",
"description": "Strong model for everyday coding.",
"default_reasoning_level": "medium",
"supported_reasoning_levels": _STOCK_REASONING_LEVELS,
"shell_type": "unified_exec",
"visibility": "hide",
"supported_in_api": True,
"priority": 16,
"availability_nux": None,
"upgrade": {
"model": "gpt-5.6-terra",
"migration_markdown": "GPT-5.4 is no longer available. Switch to GPT-5.6 Terra to continue.",
"retirement_at": "2026-08-31T19:00:00Z",
},
"base_instructions": "You are Codex, a coding agent based on GPT-5.",
"apply_patch_tool_type": "freeform",
"supports_parallel_tool_calls": True,
"context_window": 272000,
"comp_hash": "gpt-5.4-hash",
},
"codex-auto-review": {
"slug": "codex-auto-review",
"display_name": "Codex Auto Review",
"description": None,
"supported_reasoning_levels": [],
"shell_type": "unified_exec",
"visibility": "hide",
"supported_in_api": False,
"priority": 43,
"availability_nux": None,
"upgrade": None,
"base_instructions": "You are Codex, reviewing a change.",
"apply_patch_tool_type": None,
"supports_parallel_tool_calls": True,
"context_window": 272000,
"comp_hash": "review-hash",
},
}
_STOCK_CATALOG = json.dumps({"models": list(_STOCK_MODELS.values())})
class _FakeRun:
"""A `codex` that prints `stock` from a bare `debug models` and answers a catalog override with `returncode`.
`stock=None` is a Codex with no `debug models` at all: every call answers with `returncode` and `stderr`.
"""
def __init__(self, returncode=0, stderr="", stock=_STOCK_CATALOG):
self.returncode = returncode
self.stderr = stderr
self.stock = stock
self.calls = []
def __call__(self, args, **kwargs):
self.calls.append((args, kwargs))
if self.stock is not None and "model_catalog_json=" not in str(args):
return subprocess.CompletedProcess(args, 0, self.stock, "")
return subprocess.CompletedProcess(args, self.returncode, "", self.stderr)
class _FakeJsonResponse:
def __init__(self, status_code, payload=None):
self.status_code = status_code
@ -90,9 +198,7 @@ class TestAgentProfile:
class TestBuildAgentEnv:
def test_anthropic_profile_uses_bare_root_and_bearer(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"})
)
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"}))
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
assert env["ENABLE_TOOL_SEARCH"] == "true"
@ -128,9 +234,7 @@ class TestBuildAgentEnv:
assert "ANTHROPIC_API_KEY" not in env
def test_openai_profile_appends_v1(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"openai"})
)
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"openai"}))
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
assert env["OPENAI_API_KEY"] == "sk-key"
assert "ANTHROPIC_BASE_URL" not in env
@ -138,9 +242,7 @@ class TestBuildAgentEnv:
assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in env
def test_both_profiles_set_everything(self):
env = build_agent_env(
{}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"})
)
env = build_agent_env({}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"}))
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
@ -148,9 +250,7 @@ class TestBuildAgentEnv:
assert env["ENABLE_TOOL_SEARCH"] == "true"
def test_litellm_profile_exports_only_the_proxy_key(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"litellm"})
)
env = build_agent_env({}, "http://localhost:4000/", "sk-key", frozenset({"litellm"}))
assert env["LITELLM_PROXY_API_KEY"] == "sk-key"
assert "ANTHROPIC_BASE_URL" not in env
assert "OPENAI_BASE_URL" not in env
@ -158,9 +258,7 @@ class TestBuildAgentEnv:
def test_preserves_unrelated_env_and_does_not_mutate_input(self):
base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
env = build_agent_env(
base, "http://localhost:4000", "sk-key", frozenset({"anthropic"})
)
env = build_agent_env(base, "http://localhost:4000", "sk-key", frozenset({"anthropic"}))
assert env["PATH"] == "/usr/bin"
assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
@ -322,9 +420,7 @@ class TestOpencodeModelSync:
assert "refused" in result.reason
def test_non_200_is_reported(self):
result = opencode_model_sync_env(
{}, "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)
)
result = opencode_model_sync_env({}, "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500))
assert isinstance(result, ModelSyncSkipped)
assert "HTTP 500" in result.reason
@ -335,10 +431,10 @@ class TestOpencodeModelSync:
assert isinstance(result, ModelSyncSkipped)
assert "unexpected body" in result.reason
@pytest.mark.parametrize("command", ["claude", "codex", "/usr/bin/claude"])
def test_only_opencode_syncs(self, command):
@pytest.mark.parametrize("command", ["claude", "pi", "/usr/bin/claude"])
def test_only_opencode_and_codex_sync(self, command):
def boom(*a, **k):
raise AssertionError("no agent other than opencode should call the proxy")
raise AssertionError("no agent other than opencode or codex should call the proxy")
assert agent_model_sync_env(command, {}, "http://localhost:4000", "sk-key", False, get=boom) == {}
@ -367,7 +463,369 @@ class TestOpencodeModelSync:
assert _default_of(opencode_model_sync_env, "get") is requests.get
class TestCodexModelSync:
@staticmethod
def _listing(*models):
return {"object": "list", "data": list(models)}
@staticmethod
def _row(model_id, **extra):
return {"id": model_id, "object": "model", "created": 1, "owned_by": "openai", **extra}
def _sync(self, listing, codex_home, base_url="http://localhost:4000/", run=None):
captured = {}
def fake_get(url, headers, timeout):
captured["url"] = url
captured["headers"] = headers
return _FakeResponse(200, listing)
result = codex_model_sync_args(
{"CODEX_HOME": str(codex_home)},
base_url,
"sk-key",
get=fake_get,
run=_FakeRun() if run is None else run,
)
return captured, result
@staticmethod
def _catalog_path(result):
assert isinstance(result, ModelSyncArgs)
flag, override = result.args
assert flag == "-c"
key, _, value = override.partition("=")
assert key == "model_catalog_json"
return json.loads(value)
def test_writes_catalog_under_codex_home_and_points_codex_at_it(self, tmp_path):
listing = self._listing(self._row("gpt-5.5", mode="chat"), self._row("claude-opus-4-7"))
captured, result = self._sync(listing, tmp_path / "codex")
assert captured["url"] == "http://localhost:4000/v1/models"
assert captured["headers"] == {"Authorization": "Bearer sk-key"}
path = self._catalog_path(result)
assert path == str(tmp_path / "codex" / "litellm-models.json")
text = (tmp_path / "codex" / "litellm-models.json").read_text()
assert "sk-key" not in text
catalog = json.loads(text)
assert [m["slug"] for m in catalog["models"]] == ["gpt-5.5", "claude-opus-4-7"]
assert [m["display_name"] for m in catalog["models"]] == ["GPT-5.5", "claude-opus-4-7"]
assert [m["priority"] for m in catalog["models"]] == [0, 1]
def _entries(self, codex_home):
return {m["slug"]: m for m in json.loads((codex_home / "litellm-models.json").read_text())["models"]}
def test_known_model_keeps_the_installed_codex_entry(self, tmp_path):
self._sync(self._listing(self._row("gpt-5.5", mode="chat")), tmp_path)
assert self._entries(tmp_path)["gpt-5.5"] == {**_STOCK_MODELS["gpt-5.5"], "priority": 0}
def test_hidden_stock_model_is_listed_when_the_proxy_serves_it(self, tmp_path):
self._sync(self._listing(self._row("gpt-5.4")), tmp_path)
entry = self._entries(tmp_path)["gpt-5.4"]
assert entry["visibility"] == "list"
assert entry["upgrade"] is None
assert entry["supported_reasoning_levels"] == _STOCK_REASONING_LEVELS
def test_api_disabled_stock_model_is_selectable_when_the_proxy_serves_it(self, tmp_path):
self._sync(self._listing(self._row("codex-auto-review")), tmp_path)
entry = self._entries(tmp_path)["codex-auto-review"]
assert entry["supported_in_api"] is True
assert entry["visibility"] == "list"
assert entry["base_instructions"] == _STOCK_MODELS["codex-auto-review"]["base_instructions"]
def test_stock_upgrade_nudge_survives_when_its_target_is_listed(self, tmp_path):
self._sync(self._listing(self._row("gpt-5.4"), self._row("gpt-5.6-terra")), tmp_path)
entries = self._entries(tmp_path)
assert entries["gpt-5.4"]["upgrade"] == _STOCK_MODELS["gpt-5.4"]["upgrade"]
assert [entries["gpt-5.4"]["priority"], entries["gpt-5.6-terra"]["priority"]] == [0, 1]
def test_stock_catalog_is_decoded_as_utf8_regardless_of_locale(self, tmp_path):
description = "Modelo equilibrado para el trabajo diario, con acentos y ñ."
catalog = {"models": [{**_STOCK_MODELS["gpt-5.5"], "description": description}]}
stock = json.dumps(catalog, ensure_ascii=False).encode("utf-8")
def locale_bound_run(args, **kwargs):
if "model_catalog_json=" in str(args):
return subprocess.CompletedProcess(args, 0, "", "")
return subprocess.CompletedProcess(args, 0, stock.decode(kwargs.get("encoding") or "ascii"), "")
_, result = self._sync(self._listing(self._row("gpt-5.5")), tmp_path, run=locale_bound_run)
assert isinstance(result, ModelSyncArgs)
written = json.loads((tmp_path / "litellm-models.json").read_text(encoding="utf-8"))["models"]
assert [m["description"] for m in written] == [description]
def test_unparseable_stock_catalog_is_reported(self, tmp_path):
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(stock="not json"))
assert isinstance(result, ModelSyncSkipped)
assert result.reason.startswith("`codex debug models` printed no model catalog: ")
assert not (tmp_path / "litellm-models.json").exists()
def test_unknown_model_gets_the_fields_codex_requires(self, tmp_path):
_, result = self._sync(self._listing(self._row("m")), tmp_path)
entry = json.loads((tmp_path / "litellm-models.json").read_text())["models"][0]
assert entry["visibility"] == "list"
assert entry["supported_in_api"] is True
assert entry["shell_type"] == "unified_exec"
assert entry["supported_reasoning_levels"] == []
assert entry["truncation_policy"] == {"mode": "bytes", "limit": 10000}
assert entry["experimental_supported_tools"] == []
assert entry["support_verbosity"] is False
assert entry["supports_reasoning_summaries"] is False
assert entry["supports_parallel_tool_calls"] is False
for nullable in ("description", "availability_nux", "upgrade", "default_verbosity", "apply_patch_tool_type"):
assert nullable in entry and entry[nullable] is None
assert entry["base_instructions"].startswith("You are a coding agent running in the Codex CLI")
def test_context_window_comes_from_max_input_tokens_for_unknown_models_only(self, tmp_path):
listing = self._listing(
self._row("big", max_input_tokens=400000),
self._row("unknown"),
self._row("gpt-5.5", max_input_tokens=400000),
)
self._sync(listing, tmp_path)
models = self._entries(tmp_path)
assert models["big"]["context_window"] == 400000
assert models["unknown"]["context_window"] is None
assert models["gpt-5.5"]["context_window"] == 272000
def test_non_chat_models_are_left_out(self, tmp_path):
listing = self._listing(
self._row("chat", mode="chat"),
self._row("resp", mode="responses"),
self._row("embed", mode="embedding"),
self._row("img", mode="image_generation"),
)
self._sync(listing, tmp_path)
slugs = {m["slug"] for m in json.loads((tmp_path / "litellm-models.json").read_text())["models"]}
assert slugs == {"chat", "resp"}
def test_listing_without_chat_models_is_skipped_and_writes_nothing(self, tmp_path):
_, result = self._sync(self._listing(self._row("embed", mode="embedding")), tmp_path)
assert isinstance(result, ModelSyncSkipped)
assert "no chat models" in result.reason
assert not (tmp_path / "litellm-models.json").exists()
def test_catalog_is_rewritten_on_every_launch(self, tmp_path):
self._sync(self._listing(self._row("old")), tmp_path)
self._sync(self._listing(self._row("new")), tmp_path)
slugs = [m["slug"] for m in json.loads((tmp_path / "litellm-models.json").read_text())["models"]]
assert slugs == ["new"]
def test_catalog_is_replaced_whole_and_leaves_no_temp_files(self, tmp_path):
self._sync(self._listing(*(self._row(f"m{i}") for i in range(50))), tmp_path)
self._sync(self._listing(self._row("new")), tmp_path)
assert [p.name for p in tmp_path.iterdir()] == ["litellm-models.json"]
assert json.loads((tmp_path / "litellm-models.json").read_text())["models"][0]["slug"] == "new"
def test_defaults_to_dot_codex_in_home(self, tmp_path):
result = codex_model_sync_args(
{},
"http://localhost:4000",
"sk-key",
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
run=_FakeRun(),
home=lambda: tmp_path,
)
assert self._catalog_path(result) == str(tmp_path / ".codex" / "litellm-models.json")
def test_default_home_is_the_users(self):
assert _default_of(codex_model_sync_args, "home") == Path.home
def test_missing_base_instructions_is_reported_not_raised(self, tmp_path):
result = codex_model_sync_args(
{"CODEX_HOME": str(tmp_path)},
"http://localhost:4000",
"sk-key",
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
instructions_path=tmp_path / "missing.md",
)
assert isinstance(result, ModelSyncSkipped)
assert "could not read" in result.reason
assert not (tmp_path / "litellm-models.json").exists()
def test_unwritable_catalog_path_is_reported_not_raised(self, tmp_path):
blocker = tmp_path / "file"
blocker.write_text("")
_, result = self._sync(self._listing(self._row("m")), blocker / "codex")
assert isinstance(result, ModelSyncSkipped)
assert "could not write" in result.reason
def test_failed_replace_is_reported_and_leaves_no_temp_file(self, tmp_path):
(tmp_path / "litellm-models.json").mkdir()
_, result = self._sync(self._listing(self._row("m")), tmp_path)
assert isinstance(result, ModelSyncSkipped)
assert "could not write" in result.reason
assert [p.name for p in tmp_path.iterdir()] == ["litellm-models.json"]
def test_unreachable_proxy_is_reported_not_raised(self, tmp_path):
def boom(*a, **k):
raise requests.ConnectionError("refused")
result = codex_model_sync_args({"CODEX_HOME": str(tmp_path)}, "http://localhost:4000", "sk-key", get=boom)
assert isinstance(result, ModelSyncSkipped)
assert "refused" in result.reason
assert not (tmp_path / "litellm-models.json").exists()
@pytest.mark.parametrize(
("response", "reason"),
[(_FakeResponse(500), "HTTP 500"), (_FakeResponse(200, {"data": "nope"}), "unexpected body")],
)
def test_bad_response_is_reported(self, tmp_path, response, reason):
result = codex_model_sync_args(
{"CODEX_HOME": str(tmp_path)}, "http://localhost:4000", "sk-key", get=lambda *a, **k: response
)
assert isinstance(result, ModelSyncSkipped)
assert reason in result.reason
@pytest.mark.parametrize("binary", ["codex", "/opt/bin/codex", "codex.cmd", "/c/npm/codex.CMD"])
def test_codex_syncs_through_the_agent_dispatch_with_the_binary_it_will_run(self, tmp_path, binary):
run = _FakeRun()
result = agent_model_sync_env(
binary,
{"CODEX_HOME": str(tmp_path)},
"http://localhost:4000",
"sk-key",
False,
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
run=run,
)
assert self._catalog_path(result) == str(tmp_path / "litellm-models.json")
assert len(run.calls) == 2
assert all(binary in command for command, _ in run.calls)
def test_opencode_dispatch_never_runs_codex(self):
def boom(*a, **k):
raise AssertionError("only the Codex sync reads its catalog back")
result = agent_model_sync_env(
"opencode",
{},
"http://localhost:4000",
"sk-key",
False,
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
run=boom,
)
assert "OPENCODE_CONFIG_CONTENT" in result
def test_codex_lists_its_own_models_then_reads_the_catalog_back_before_launch(self, tmp_path):
run = _FakeRun()
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=run)
path = self._catalog_path(result)
assert [command for command, _ in run.calls] == [
("codex", "debug", "models"),
("codex", "-c", f"model_catalog_json={json.dumps(path)}", "debug", "models"),
]
for _, options in run.calls:
assert options["env"] == {"CODEX_HOME": str(tmp_path)}
assert options["stdin"] is subprocess.DEVNULL
assert options["capture_output"] is True
assert options["encoding"] == "utf-8"
assert options["timeout"] == 10
def test_codex_rejecting_the_catalog_skips_the_sync_and_keeps_the_file(self, tmp_path):
stderr = (
"Error: failed to parse model_catalog_json path `/home/me/.codex/litellm-models.json` as JSON: "
"missing field `supports_parallel_tool_calls` at line 1 column 21648\n"
)
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(1, stderr))
assert isinstance(result, ModelSyncSkipped)
assert result.reason == (
"`codex debug models` exited 1: Error: failed to parse model_catalog_json path "
"`/home/me/.codex/litellm-models.json` as JSON: missing field `supports_parallel_tool_calls` "
"at line 1 column 21648"
)
assert (tmp_path / "litellm-models.json").exists()
def test_codex_without_debug_models_skips_the_sync(self, tmp_path):
stderr = "error: unrecognized subcommand 'models'\n\nUsage: codex debug [OPTIONS] <COMMAND>\n"
run = _FakeRun(2, stderr, stock=None)
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=run)
assert isinstance(result, ModelSyncSkipped)
assert result.reason == "`codex debug models` exited 2: error: unrecognized subcommand 'models'"
assert len(run.calls) == 1
assert not (tmp_path / "litellm-models.json").exists()
def test_codex_failing_silently_is_reported(self, tmp_path):
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=_FakeRun(1))
assert isinstance(result, ModelSyncSkipped)
assert result.reason == "`codex debug models` exited 1: no output"
@pytest.mark.parametrize(
"error", [OSError("codex vanished"), subprocess.TimeoutExpired("codex", 10)], ids=["oserror", "timeout"]
)
def test_unrunnable_preflight_is_reported_not_raised(self, tmp_path, error):
def failing_run(*a, **k):
raise error
_, result = self._sync(self._listing(self._row("m")), tmp_path, run=failing_run)
assert isinstance(result, ModelSyncSkipped)
assert result.reason.startswith("`codex debug models` failed: ")
assert str(error) in result.reason
def test_windows_shim_preflight_goes_through_cmd_exe(self, tmp_path):
shim = _WINDOWS_CLAUDE_CMD.replace("claude", "codex")
run = _FakeRun()
result = codex_model_sync_args(
{"CODEX_HOME": str(tmp_path)},
"http://localhost:4000",
"sk-key",
binary=shim,
get=lambda *a, **k: _FakeResponse(200, self._listing(self._row("m"))),
run=run,
)
override = f"model_catalog_json={json.dumps(self._catalog_path(result))}"
doubled = override.replace('"', '""')
assert [command for command, _ in run.calls] == [
f'{_CMD_PREFIX}""{shim}" "debug" "models""',
f'{_CMD_PREFIX}""{shim}" "-c" "{doubled}" "debug" "models""',
]
def test_default_binary_is_codex_on_path(self):
assert _default_of(codex_model_sync_args, "binary") == "codex"
def test_default_runner_is_subprocess_run(self):
assert _default_of(codex_model_sync_args, "run") is subprocess.run
assert _default_of(agent_model_sync_env, "run") is subprocess.run
def test_skip_verify_keeps_the_launch_offline(self):
def boom(*a, **k):
raise AssertionError("--skip-verify must not touch the proxy")
result = agent_model_sync_env("codex", {}, "http://localhost:4000", "sk-key", True, get=boom)
assert isinstance(result, ModelSyncSkipped)
assert "--skip-verify" in result.reason
def test_default_http_client_is_requests_get(self):
assert _default_of(codex_model_sync_args, "get") is requests.get
class TestRunAgent:
def test_synced_args_precede_user_args_and_follow_provider_overrides(self):
calls = {}
run_agent(
"http://localhost:4000",
"sk-key",
["codex", "exec", "hi"],
base_env={},
sync_models=lambda *a: ModelSyncArgs(("-c", 'model_catalog_json="/tmp/c.json"')),
which=lambda name: "/usr/local/bin/codex",
verify=lambda *a: None,
launcher=lambda p, a, e: calls.update(args=tuple(a), env=dict(e)),
)
args = calls["args"]
assert args[-2:] == ("exec", "hi")
assert args[args.index('model_catalog_json="/tmp/c.json"') - 1] == "-c"
assert (
args.index('model_provider="litellm"') < args.index('model_catalog_json="/tmp/c.json"') < args.index("exec")
)
assert calls["env"]["OPENAI_API_KEY"] == "sk-key"
assert "model_catalog_json" not in json.dumps(calls["env"])
def test_synced_model_config_reaches_the_agent_alongside_profile_env(self):
calls = {}
run_agent(
@ -405,7 +863,13 @@ class TestRunAgent:
launcher=lambda p, a, e: order.append("launch"),
)
assert order == ["verify", "sync", "launch"]
assert calls["args"] == ("opencode", {"HOME": "/home/me"}, "http://localhost:4000", "sk-key", False)
assert calls["args"] == (
"/usr/local/bin/opencode",
{"HOME": "/home/me"},
"http://localhost:4000",
"sk-key",
False,
)
def test_unreachable_proxy_is_not_asked_for_models(self):
def failing_verify(*a):
@ -1050,10 +1514,7 @@ class TestAgentCommands:
assert captured["api_key"] == "sk-key"
assert captured["command"] == ["claude", "--resume", "-p", "hi"]
assert captured["skip_verify"] is False
assert (
"routing Claude Code through proxy at http://localhost:4000"
in result.output
)
assert "routing Claude Code through proxy at http://localhost:4000" in result.output
def test_codex_shows_friendly_name(self):
captured = {}
@ -1126,14 +1587,10 @@ class TestAgentCommands:
with (
patch(f"{AGENTS_MODULE}._is_interactive", return_value=True),
patch(f"{AGENTS_MODULE}.login", fake_login),
patch(
f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login"
) as mock_get,
patch(f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login") as mock_get,
patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda base_url, api_key, command, **k: captured.update(
api_key=api_key
),
side_effect=lambda base_url, api_key, command, **k: captured.update(api_key=api_key),
),
):
result = self.runner.invoke(

View file

@ -512,8 +512,8 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch):
the repair must be skipped and the existing 400 raised immediately, while bodies
at or below the limit still get repaired.
`\\ud83d` is a lone high-surrogate escape: orjson rejects it, the json fallback
accepts it, so a body containing it is only salvaged when the repair path runs.
`NaN` is rejected by orjson and accepted by the json fallback, so a body containing
it is only salvaged when the repair path runs.
"""
import litellm.proxy.common_utils.http_parsing_utils as http_parsing_utils
@ -522,14 +522,14 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch):
http_parsing_utils, "MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB", 100 / (1024 * 1024)
)
small_body = b'{"model":"gpt-4o","x":"\\ud83d"}'
small_body = b'{"model":"gpt-4o","x":NaN}'
assert len(small_body) <= 100
repaired = await _read_request_body(_make_json_request(small_body))
assert repaired["model"] == "gpt-4o"
padding = "a" * 200
large_body = (
b'{"model":"gpt-4o","pad":"' + padding.encode() + b'","x":"\\ud83d"}'
b'{"model":"gpt-4o","pad":"' + padding.encode() + b'","x":NaN}'
)
assert len(large_body) > 100
with pytest.raises(ProxyException) as exc_info:
@ -546,6 +546,33 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch):
assert repaired_large["model"] == "gpt-4o"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"content",
[
pytest.param(b"say ok \\ud83d", id="lone-high-surrogate"),
pytest.param(b"say ok \\ude00", id="lone-low-surrogate"),
pytest.param(b"\\ud83d\\ud83d\\ude00", id="lone-high-before-valid-pair"),
],
)
async def test_lone_surrogate_escape_is_rejected_with_400(content: bytes):
"""
orjson rejects a lone surrogate escape, and the json fallback accepts it, so the
parsed body used to carry a code point no provider request can UTF-8 encode. That
surfaced as a 500 from the provider handler instead of a 400 for the bad input.
"""
body = b'{"model":"gpt-4o","messages":[{"role":"user","content":"' + content + b'"}]}'
with pytest.raises(ProxyException) as exc_info:
await _read_request_body(_make_json_request(body))
assert exc_info.value.code == "400"
assert exc_info.value.type == "invalid_request_error"
assert "Invalid JSON payload" in exc_info.value.message
paired = body.replace(content, b"say ok \\ud83d\\ude00")
parsed = await _read_request_body(_make_json_request(paired))
assert parsed["messages"][0]["content"] == "say ok \U0001F600"
@pytest.mark.asyncio
async def test_get_form_data():
"""

View file

@ -36,7 +36,11 @@ from litellm.proxy._types import (
UpdateKeyRequest,
)
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.proxy.auth.auth_checks import _delete_cache_key_object, _project_cache_key
from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
_project_cache_key,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
@ -5132,10 +5136,11 @@ async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
class _JWTMappingRow:
def __init__(self, token, jwt_claim_name, jwt_claim_value):
def __init__(self, token, jwt_claim_name, jwt_claim_value, jwt_issuer=None):
self.token = token
self.jwt_claim_name = jwt_claim_name
self.jwt_claim_value = jwt_claim_value
self.jwt_issuer = jwt_issuer
class _CascadingJWTMappingTable:
@ -5226,7 +5231,7 @@ async def test_delete_verification_tokens_evicts_jwt_key_mapping_cache(monkeypat
),
)
assert recording_evict.cache_keys == ("jwt_key_mapping:email:user@example.com",)
assert recording_evict.cache_keys == (jwt_key_mapping_cache_key("email", "user@example.com", None),)
@pytest.mark.asyncio
@ -13131,11 +13136,11 @@ async def test_regenerate_evicts_jwt_key_mapping_cache_so_next_jwt_call_gets_new
_execute_virtual_key_regeneration,
)
stale_cache_key = "jwt_key_mapping:sub:user1"
stale_cache_key = jwt_key_mapping_cache_key("sub", "user1", None)
existing_key = _make_regenerate_existing_key()
mock_prisma_client = _make_regenerate_mock_prisma()
mock_prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1")]
return_value=[MagicMock(jwt_claim_name="sub", jwt_claim_value="user1", jwt_issuer=None)]
)
mock_prisma_client.db.litellm_jwtkeymapping.find_first = AsyncMock(
return_value=MagicMock(token="new-hashed-token")

View file

@ -34,6 +34,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
from litellm.proxy.pass_through_endpoints.success_handler import (
@ -5934,6 +5935,32 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
)
@pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"])
def test_passthrough_logs_the_resolved_deployment_model_info_over_the_request_body(client_metadata_key: str):
"""A provider route that resolved a router deployment stashes its model_info on request.state. That
deployment, not a model_info the client put in its own body, is what spend logs and metrics attribute
the call to (LIT-1761: passthrough successes carried model_id="")."""
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
mock_request.headers = Headers({})
mock_request.scope = {}
mock_request.state = SimpleNamespace(
**{LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY: {"id": "vertex-gemini-38-flash-dep"}}
)
kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=mock_request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
passthrough_logging_payload=MagicMock(),
logging_obj=MagicMock(),
_parsed_body={client_metadata_key: {"model_info": {"id": "client-forged-id"}}},
litellm_call_id="lit-1761-call-id",
)
assert kwargs["litellm_params"]["metadata"]["model_info"] == {"id": "vertex-gemini-38-flash-dep"}
@pytest.mark.asyncio
async def test_chat_completion_pass_through_endpoint_answers_an_openai_typed_error_for_an_unknown_model(
monkeypatch: pytest.MonkeyPatch,

View file

@ -1,13 +1,19 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import Request
from starlette.datastructures import Headers, State
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
VertexAIPassThroughHandler,
_base_vertex_proxy_route,
_resolve_vertex_model_from_router,
_upstream_headers_for_vertex_route,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
)
from litellm.types.router import DeploymentTypedDict
@ -758,3 +764,115 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
assert (
"gemini-3-pro" in target_url
), f"Actual Vertex AI model name should be in target URL. Got: {target_url}"
@pytest.mark.asyncio
async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment():
"""The router deployment that rewrote the upstream URL is the one the logging kwargs must name, so
the Prometheus model_id label (and SpendLogs.model_id) on a Vertex passthrough success reads the
deployment's id instead of "" (LIT-1761)."""
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://0.0.0.0:4000/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent"
mock_request.headers = Headers({})
mock_request.scope = {}
mock_request.state = State()
mock_handler = MagicMock()
mock_handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
mock_router = MagicMock()
mock_router.get_available_deployment_for_pass_through.return_value = {
"model_name": "gemini-3.8-flash",
"litellm_params": {
"model": "vertex_ai/gemini-3.8-flash",
"vertex_project": "p",
"vertex_location": "global",
"use_in_pass_through": True,
},
"model_info": {"id": "vertex-gemini-38-flash-dep"},
}
async def relay_returning_logging_kwargs(
request: Request, fastapi_response: object, user_api_key_dict: UserAPIKeyAuth
) -> dict:
return HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint(
request=request,
user_api_key_dict=user_api_key_dict,
passthrough_logging_payload=MagicMock(),
logging_obj=MagicMock(),
_parsed_body={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
litellm_call_id="lit-1761-call-id",
)
with (
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
"litellm.proxy.proxy_server.llm_router", mock_router
),
patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
) as mock_pt_router,
patch( # test-quality-ok: the route offers no injection point for its header preparation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
new_callable=AsyncMock,
return_value=({}, False, "p", "global"),
),
patch( # test-quality-ok: the relay is captured here to read the logging kwargs, the route offers no seam
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
return_value=relay_returning_logging_kwargs,
),
patch( # test-quality-ok: the route calls auth directly rather than through Depends
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
new_callable=AsyncMock,
return_value=UserAPIKeyAuth(api_key="hashed-key"),
),
):
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
logging_kwargs = await _base_vertex_proxy_route(
endpoint="v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
request=mock_request,
fastapi_response=MagicMock(),
get_vertex_pass_through_handler=mock_handler,
)
assert logging_kwargs["litellm_params"]["metadata"]["model_info"]["id"] == "vertex-gemini-38-flash-dep"
def _router_without_deployment() -> MagicMock:
router = MagicMock()
router.get_available_deployment_for_pass_through.return_value = None
return router
def _router_raising_on_lookup() -> MagicMock:
router = MagicMock()
router.get_available_deployment_for_pass_through.side_effect = ValueError("no healthy deployment")
return router
@pytest.mark.parametrize(
"llm_router",
[None, _router_without_deployment(), _router_raising_on_lookup()],
ids=["no-router", "no-matching-deployment", "lookup-raises"],
)
def test_vertex_passthrough_without_a_resolved_deployment_keeps_the_url_and_reports_no_model_info(
llm_router: MagicMock | None,
):
"""A Vertex passthrough call that no router deployment serves must keep the URL-derived values and carry no
deployment model_info, so logging cannot attribute it to a deployment that never handled it."""
resolved = _resolve_vertex_model_from_router(
model_id="gemini-3.8-flash",
llm_router=llm_router,
encoded_endpoint="/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
endpoint="v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
vertex_project="url-project",
vertex_location="url-location",
)
assert resolved == (
"/v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
"v1/projects/p/locations/global/publishers/google/models/gemini-3.8-flash:generateContent",
"url-project",
"url-location",
None,
)

View file

@ -10750,6 +10750,7 @@ def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch):
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", "1h")
monkeypatch.setattr(litellm, "openai_system_messages_first", False)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
@ -10771,6 +10772,10 @@ def test_get_config_list_includes_anthropic_prompt_caching_fields(monkeypatch):
assert fields["enable_anthropic_prompt_caching"]["field_tab"] == "prompt_caching"
assert fields["anthropic_prompt_caching_ttl"]["field_tab"] == "prompt_caching"
assert fields["budget_exceeded_throttle_percentage"]["field_tab"] is None
assert fields["openai_system_messages_first"]["field_type"] == "Boolean"
assert fields["openai_system_messages_first"]["field_value"] is False
assert fields["openai_system_messages_first"]["field_tab"] == "prompt_caching"
finally:
app.dependency_overrides.clear()
@ -10975,6 +10980,7 @@ def test_general_settings_ui_defaults_unchanged_for_existing_fields():
[
("enable_anthropic_prompt_caching", True),
("anthropic_prompt_caching_ttl", "1h"),
("openai_system_messages_first", True),
],
)
def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_name, db_value):
@ -11033,6 +11039,8 @@ def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypa
("enable_anthropic_prompt_caching", False),
("anthropic_prompt_caching_ttl", "5m"),
("anthropic_prompt_caching_ttl", "1h"),
("openai_system_messages_first", True),
("openai_system_messages_first", False),
],
)
@pytest.mark.asyncio
@ -11081,6 +11089,8 @@ async def test_update_config_field_prompt_caching_persists_to_litellm_settings(m
("anthropic_prompt_caching_ttl", "10m"),
("anthropic_prompt_caching_ttl", "1H"),
("anthropic_prompt_caching_ttl", 3600),
("openai_system_messages_first", "yes"),
("openai_system_messages_first", 1),
],
)
@pytest.mark.asyncio
@ -11120,6 +11130,7 @@ async def test_update_config_field_prompt_caching_rejects_invalid(monkeypatch, f
[
("enable_anthropic_prompt_caching", False),
("anthropic_prompt_caching_ttl", None),
("openai_system_messages_first", False),
("budget_exceeded_throttle_percentage", None),
],
)

View file

@ -1360,6 +1360,56 @@ def test_add_invalid_provider_to_router():
assert router.pattern_router.patterns == {}
@pytest.fixture
def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str:
from litellm import CustomLLM
from litellm.types.utils import ModelResponse
class OnPremLLM(CustomLLM):
def completion(self, *args, **kwargs) -> ModelResponse:
return litellm.completion(
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="served by onprem handler"
)
monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}])
monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list))
monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers))
return "test-onprem-llm"
def test_router_init_accepts_custom_provider_map_prefix_before_first_completion(registered_custom_provider: str):
assert registered_custom_provider not in litellm.provider_list
router = litellm.Router(
model_list=[
{"model_name": "onprem", "litellm_params": {"model": f"{registered_custom_provider}/my-model"}},
],
)
assert router.get_model_list(model_name="onprem")[0]["litellm_params"]["model"] == (
f"{registered_custom_provider}/my-model"
)
response = router.completion(model="onprem", messages=[{"role": "user", "content": "hi"}])
assert response.choices[0].message.content == "served by onprem handler"
def test_router_add_deployment_accepts_explicit_custom_provider_from_custom_provider_map(
registered_custom_provider: str,
):
from litellm.types.router import Deployment
router = litellm.Router(model_list=[])
router.add_deployment(
Deployment(
model_name="onprem",
litellm_params={"model": "my-model", "custom_llm_provider": registered_custom_provider},
)
)
assert router.get_model_list(model_name="onprem")[0]["litellm_params"]["model"] == "my-model"
@pytest.mark.asyncio
async def test_router_ageneric_api_call_with_fallbacks_helper():
"""
@ -12156,6 +12206,83 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o
@pytest.mark.parametrize(
"model,provider,expected",
[
("anthropic/claude-opus-5", None, True),
("claude-opus-4-8", None, True),
("anthropic/claude-opus-4-7", None, False),
("anthropic/claude-opus-4-6", None, False),
("anthropic/claude-sonnet-5", None, False),
("anthropic/off-map-opus", None, False),
("vertex_ai/claude-opus-5", None, False),
("bedrock/claude-opus-5", None, False),
("claude-opus-5", "vertex_ai", False),
("claude-opus-5", "bedrock", False),
],
)
@pytest.mark.parametrize("operator_flag", [True, False])
def test_model_group_info_fast_mode_uses_exact_provider_catalog(
local_model_cost_map: None, model: str, provider: str | None, expected: bool, operator_flag: bool
) -> None:
router: Final = Router(model_list=[{
"model_name": "fast-group",
"litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"},
"model_info": {"supports_fast_mode": operator_flag},
}])
result: Final = router.get_model_group_info("fast-group")
assert result is not None
assert result.supports_fast_mode is expected
@pytest.mark.parametrize("flag", [None, False, "true", 1])
def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean(
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, flag: object
) -> None:
entry: Final = {key: value for key, value in litellm.model_cost["claude-opus-5"].items()
if key != "supports_fast_mode"}
if flag is not None:
entry["supports_fast_mode"] = flag
monkeypatch.setitem(litellm.model_cost, "claude-opus-5", entry)
router: Final = Router(model_list=[{
"model_name": "fast-group",
"litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"},
"model_info": {"supports_fast_mode": True},
}])
result: Final = router.get_model_group_info("fast-group")
assert result is not None
assert result.supports_fast_mode is False
@pytest.mark.parametrize("other_model,expected", [
("anthropic/claude-opus-4-8", True),
("anthropic/claude-opus-4-7", False),
("anthropic/off-map-opus", False),
("vertex_ai/claude-opus-5", False),
("bedrock/claude-opus-5", False),
])
@pytest.mark.parametrize("reverse", [True, False])
def test_model_group_info_fast_mode_requires_every_deployment(
local_model_cost_map: None, other_model: str, expected: bool, reverse: bool
) -> None:
models: Final = (other_model, "anthropic/claude-opus-5") if reverse else (
"anthropic/claude-opus-5", other_model
)
router: Final = Router(model_list=[{
"model_name": "fast-group",
"litellm_params": {"model": model, "api_key": "fake-key"},
} for model in models])
result: Final = router.get_model_group_info("fast-group")
assert result is not None
assert result.supports_fast_mode is expected
def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map):
"""``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose
registry entry declares parallel function calling must flip the group to True instead of False."""

View file

@ -20,6 +20,8 @@ from jsonschema import validate
import litellm
from litellm._internal_context import is_internal_call
from litellm.caching.caching import Cache
from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
from litellm.constants import DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT
from litellm._logging import (
CorrelationContextFilter,
@ -32,6 +34,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor
from litellm.proxy.utils import is_valid_api_key
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY
from litellm.types.utils import (
CallTypes,
Delta,
@ -44,6 +47,7 @@ from litellm.types.utils import (
from litellm.types.utils import all_litellm_params, bedrock_batch_litellm_params
from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams
from litellm.utils import (
CustomStreamWrapper,
ProviderConfigManager,
TextCompletionStreamWrapper,
_check_provider_match,
@ -1099,6 +1103,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"supports_sampling_params": {"type": "boolean"},
"supports_output_config": {"type": "boolean"},
"supports_speed": {"type": "boolean"},
"supports_fast_mode": {"type": "boolean"},
"supported_audio_formats": {
"type": "array",
"items": {
@ -4953,6 +4958,208 @@ async def test_wrapper_async_restores_originating_task_context_after_success(mon
session_id_var.set("")
class _ConvertStreamDeploymentHook(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict[str, object] | None:
if not kwargs.get("stream"):
return None
return {**kwargs, "stream": False, HEADROOM_CONVERTED_STREAM_KEY: True}
class _SuccessKwargsCapture(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.success_kwargs: list[dict[str, object]] = []
self.stream_event_responses: list[object] = []
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.success_kwargs.append(kwargs)
async def async_log_stream_event(
self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
self.stream_event_responses.append(response_obj)
def _install_converted_stream_callbacks(monkeypatch: pytest.MonkeyPatch) -> _SuccessKwargsCapture:
capture: Final = _SuccessKwargsCapture()
monkeypatch.setattr(litellm, "callbacks", [_ConvertStreamDeploymentHook(), capture])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
return capture
async def _wait_for_success_kwargs(capture: _SuccessKwargsCapture, count: int = 1) -> dict[str, object]:
for _ in range(50):
if len(capture.success_kwargs) >= count and not _PENDING_CACHE_WRITES:
break
await asyncio.sleep(0.05)
await asyncio.sleep(0.2)
assert len(capture.success_kwargs) == count
return capture.success_kwargs[-1]
def _assert_cache_hit_logged_as_stream(capture: _SuccessKwargsCapture, success_kwargs: dict[str, object]) -> None:
standard_logging_object: Final = success_kwargs["standard_logging_object"]
assert isinstance(standard_logging_object, dict)
assert standard_logging_object["cache_hit"] is True
assert standard_logging_object["stream"] is True
assert success_kwargs["stream"] is True
assert capture.stream_event_responses == []
@pytest.mark.asyncio
async def test_wrapper_async_logs_converted_chat_stream_with_standard_logging_object(
monkeypatch: pytest.MonkeyPatch,
) -> None:
capture: Final = _install_converted_stream_callbacks(monkeypatch)
response: Final = await litellm.acompletion(
model="gpt-5.6",
messages=[{"role": "user", "content": "hi"}],
stream=True,
mock_response="converted stream body",
num_retries=0,
)
assert isinstance(response, CustomStreamWrapper)
chunks: Final = [chunk async for chunk in response]
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "converted stream body"
success_kwargs: Final = await _wait_for_success_kwargs(capture)
standard_logging_object: Final = success_kwargs["standard_logging_object"]
assert isinstance(standard_logging_object, dict)
assert standard_logging_object["response_cost"] > 0
assert standard_logging_object["stream"] is True
assert success_kwargs["stream"] is True
@pytest.mark.asyncio
@respx.mock
async def test_wrapper_async_logs_converted_responses_stream_with_standard_logging_object(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
capture: Final = _install_converted_stream_callbacks(monkeypatch)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
respx.post("https://api.openai.com/v1/responses").respond(
json={
"id": "resp_converted",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-5.6",
"output": [
{
"type": "message",
"id": "msg_converted",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "converted stream body", "annotations": []}],
}
],
"usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7},
}
)
response: Final = await litellm.aresponses(
model="openai/gpt-5.6", input="hi", stream=True, api_key="sk-test", num_retries=0
)
assert isinstance(response, BaseResponsesAPIStreamingIterator)
events: Final = [event async for event in response]
assert events[-1].type == "response.completed"
success_kwargs: Final = await _wait_for_success_kwargs(capture)
standard_logging_object: Final = success_kwargs["standard_logging_object"]
assert isinstance(standard_logging_object, dict)
assert standard_logging_object["response_cost"] > 0
assert standard_logging_object["stream"] is True
assert success_kwargs["stream"] is True
@pytest.mark.asyncio
async def test_wrapper_async_replays_cached_converted_chat_stream_as_stream(
monkeypatch: pytest.MonkeyPatch,
) -> None:
capture: Final = _install_converted_stream_callbacks(monkeypatch)
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
request: Final = {
"model": "gpt-5.6",
"messages": [{"role": "user", "content": "replay me from cache"}],
"stream": True,
"mock_response": "converted stream body",
"num_retries": 0,
}
first: Final = await litellm.acompletion(**request)
first_chunks: Final = [chunk async for chunk in first]
assert "".join(chunk.choices[0].delta.content or "" for chunk in first_chunks) == "converted stream body"
await _wait_for_success_kwargs(capture)
replay: Final = await litellm.acompletion(**request)
assert isinstance(replay, CustomStreamWrapper)
replay_chunks: Final = [chunk async for chunk in replay]
assert "".join(chunk.choices[0].delta.content or "" for chunk in replay_chunks) == "converted stream body"
_assert_cache_hit_logged_as_stream(capture, await _wait_for_success_kwargs(capture, count=2))
@pytest.mark.asyncio
@respx.mock
async def test_wrapper_async_replays_cached_converted_responses_stream_as_stream(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
capture: Final = _install_converted_stream_callbacks(monkeypatch)
monkeypatch.setattr(litellm, "cache", Cache(type="local"))
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
route: Final = respx.post("https://api.openai.com/v1/responses").respond(
json={
"id": "resp_cached_converted",
"object": "response",
"created_at": 1,
"status": "completed",
"model": "gpt-5.6",
"output": [
{
"type": "message",
"id": "msg_cached_converted",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "converted stream body", "annotations": []}],
}
],
"usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7},
}
)
request: Final = {
"model": "openai/gpt-5.6",
"input": "replay me from cache",
"stream": True,
"api_key": "sk-test",
"num_retries": 0,
}
first: Final = await litellm.aresponses(**request)
assert [event async for event in first][-1].type == "response.completed"
await _wait_for_success_kwargs(capture)
replay: Final = await litellm.aresponses(**request)
assert isinstance(replay, BaseResponsesAPIStreamingIterator)
assert [event async for event in replay][-1].type == "response.completed"
assert route.call_count == 1
_assert_cache_hit_logged_as_stream(capture, await _wait_for_success_kwargs(capture, count=2))
def test_function_setup_failure_after_logging_construction_restores_context(monkeypatch):
"""If function_setup() constructs Logging() (which already mutated
trace_id_var/session_id_var in __init__) but then raises before returning,

View file

@ -1,8 +1,10 @@
import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized";
import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
import { fireEvent, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import type { OnUrlUpdateFunction } from "nuqs/adapters/testing";
import { beforeEach, describe, expect, it, Mock, vi } from "vitest";
import { renderWithProviders } from "../../../../../tests/test-utils";
import AllModelsTab from "./AllModelsTab";
import { STATUS_COLUMN_ID, toServerSortField } from "./ModelsTableColumns";
@ -111,6 +113,9 @@ const setModelsInfo = (rows: Record<string, unknown>[], totalCount = rows.length
const lastModelsInfoCall = (): ModelsInfoArgs => modelsInfoCalls[modelsInfoCalls.length - 1];
const lastUrlParams = (onUrlUpdate: Mock<OnUrlUpdateFunction>): URLSearchParams | undefined =>
onUrlUpdate.mock.calls.at(-1)?.[0].searchParams;
const SEARCH_SETTLE_MS = 400;
const MOCK_AUTHORIZED = {
@ -121,6 +126,8 @@ const MOCK_AUTHORIZED = {
userId: "user-123",
userEmail: "test@example.com",
userRole: "Admin",
userRoleLabel: "Admin",
isViewOnly: false,
premiumUser: true,
disabledPersonalKeyCreation: false,
showSSOBanner: false,
@ -149,14 +156,14 @@ describe("AllModelsTab", () => {
it("renders the fetched models and the server row count", async () => {
setModelsInfo([makeRow()], 137);
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(await screen.findByText("gpt-4")).toBeInTheDocument();
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-50 of 137");
});
it("does not re-query after the mount-time debounced search settles unchanged", async () => {
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
const callsAfterMount = modelsInfoCalls.length;
await new Promise((resolve) => setTimeout(resolve, SEARCH_SETTLE_MS));
@ -166,14 +173,14 @@ describe("AllModelsTab", () => {
it("shows the empty state when the proxy returns no models", () => {
setModelsInfo([], 0);
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.getByText("No models found")).toBeInTheDocument();
});
it("shows the loading skeleton while the first page is in flight", () => {
setModelsInfo([], 0, true);
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0);
expect(screen.queryByText("No models found")).not.toBeInTheDocument();
@ -197,7 +204,7 @@ describe("AllModelsTab", () => {
it.each(cases)("sorts %s using the server field %s", async (_label, columnId, serverField, firstDirection) => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(sortHeader(columnId));
await expectIndicator(columnId, firstDirection);
@ -212,7 +219,7 @@ describe("AllModelsTab", () => {
it("cycles a sorted column back to unsorted", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(sortHeader("model_info_updated_at"));
await expectIndicator("model_info_updated_at", "asc");
@ -230,7 +237,7 @@ describe("AllModelsTab", () => {
it("queries the selected team and resets to the first page", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(lastModelsInfoCall().teamId).toBeUndefined();
@ -244,8 +251,7 @@ describe("AllModelsTab", () => {
});
it("debounces the model name search into the server query", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "claude" } });
@ -254,9 +260,138 @@ describe("AllModelsTab", () => {
});
});
describe("URL persistence", () => {
it("writes the typed search to the URL and drops the page so a reload keeps the search", async () => {
setModelsInfo([makeRow()], 200);
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "3" }, onUrlUpdate });
expect(lastModelsInfoCall().page).toBe(3);
fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "claude" } });
await waitFor(() => {
expect(lastUrlParams(onUrlUpdate)?.get("model_search")).toBe("claude");
});
expect(lastUrlParams(onUrlUpdate)?.get("page")).toBeNull();
await waitFor(() => {
expect(lastModelsInfoCall().page).toBe(1);
});
});
it("restores the search box and server query from ?model_search= on mount", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { model_search: "haiku" } });
expect(screen.getByTestId("datatable-search")).toHaveValue("haiku");
expect(lastModelsInfoCall().search).toBe("haiku");
});
it("restores team, sort, page and page size from the URL into the server query", () => {
setModelsInfo([makeRow()], 200);
renderWithProviders(<AllModelsTab {...defaultProps} />, {
searchParams: {
filter_team: "team-1",
sort_by: "model_info_updated_at",
sort_order: "desc",
page: "2",
page_size: "25",
},
});
const expectedQuery: ModelsInfoArgs = {
teamId: "team-1",
sortBy: "updated_at",
sortOrder: "desc",
page: 2,
size: 25,
};
expect(lastModelsInfoCall()).toMatchObject(expectedQuery);
expect(screen.getByTestId("models-team-select")).toHaveTextContent("Engineering");
});
it("restores the access group and view mode from the URL", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, {
searchParams: { access_group: "sales-team", view_mode: "all" },
});
expect(lastModelsInfoCall().accessGroup).toBe("sales-team");
expect(screen.queryByText(/To access these models/)).not.toBeInTheDocument();
});
it("clamps a hand-edited page and page size into the range the table supports", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "0", page_size: "5000" } });
expect(lastModelsInfoCall().page).toBe(1);
expect(lastModelsInfoCall().size).toBe(100);
});
it("keeps the default page size when the URL value is not a number", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page_size: "lots" } });
expect(lastModelsInfoCall().size).toBe(50);
});
it("ignores a sort_by the table cannot sort by instead of forwarding it to the server", () => {
renderWithProviders(<AllModelsTab {...defaultProps} />, {
searchParams: { sort_by: "litellm_credential_name", sort_order: "desc" },
});
expect(lastModelsInfoCall().sortBy).toBeUndefined();
expect(lastModelsInfoCall().sortOrder).toBeUndefined();
});
it("writes sort changes to the URL with the page cleared", async () => {
setModelsInfo([makeRow()], 200);
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<AllModelsTab {...defaultProps} />, { searchParams: { page: "2" }, onUrlUpdate });
await user.click(screen.getByTestId("sort-header-model_info_updated_at"));
await waitFor(() => {
expect(lastUrlParams(onUrlUpdate)?.get("sort_by")).toBe("model_info_updated_at");
});
expect(lastUrlParams(onUrlUpdate)?.get("sort_order")).toBeNull();
expect(lastUrlParams(onUrlUpdate)?.get("page")).toBeNull();
await user.click(screen.getByTestId("sort-header-model_info_updated_at"));
await waitFor(() => {
expect(lastUrlParams(onUrlUpdate)?.get("sort_order")).toBe("desc");
});
});
it("clears every table param from the URL on drawer reset", async () => {
setModelsInfo([makeRow()], 200);
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<AllModelsTab {...defaultProps} />, {
searchParams: {
model_search: "haiku",
filter_team: "team-1",
sort_by: "model_name",
page: "2",
view_mode: "all",
},
onUrlUpdate,
});
await user.click(screen.getByTestId("datatable-filters-trigger"));
await user.click(await screen.findByTestId("filter-drawer-reset"));
await waitFor(() => {
expect(lastUrlParams(onUrlUpdate)?.toString()).toBe("");
});
expect(screen.getByTestId("datatable-search")).toHaveValue("");
const defaultQuery: ModelsInfoArgs = { search: undefined, teamId: undefined, sortBy: undefined, page: 1 };
await waitFor(() => {
expect(lastModelsInfoCall()).toMatchObject(defaultQuery);
});
});
});
it("applies a public model name filter through the drawer", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(screen.getByTestId("datatable-filters-trigger"));
await user.click(await screen.findByPlaceholderText("Filter by Public Model Name"));
@ -270,7 +405,7 @@ describe("AllModelsTab", () => {
it("renders every row the server returned for the selected model group so rows match the footer total", () => {
setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "claude-opus" }], 2);
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
const table = screen.getByRole("table");
expect(within(table).getByText("claude-opus")).toBeInTheDocument();
@ -280,7 +415,7 @@ describe("AllModelsTab", () => {
it("asks the server for wildcard deployments instead of hiding rows client-side", () => {
setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "openai/*" }], 2);
render(<AllModelsTab {...defaultProps} selectedModelGroup="wildcard" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="wildcard" />);
expect(lastModelsInfoCall().wildcardOnly).toBe(true);
expect(within(screen.getByRole("table")).getByText("gpt-4")).toBeInTheDocument();
@ -289,7 +424,7 @@ describe("AllModelsTab", () => {
it("asks the server for the selected access group instead of hiding rows client-side", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(lastModelsInfoCall().wildcardOnly).toBe(false);
await user.click(screen.getByTestId("datatable-filters-trigger"));
@ -303,20 +438,20 @@ describe("AllModelsTab", () => {
});
it("asks the server for the exact selected model group so deployments beyond the first page are found", () => {
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
expect(lastModelsInfoCall().modelName).toBe("claude-opus");
expect(lastModelsInfoCall().search).toBeUndefined();
});
it.each(["all", "wildcard"])("sends no exact model name for the %s pseudo group", (group) => {
render(<AllModelsTab {...defaultProps} selectedModelGroup={group} />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup={group} />);
expect(lastModelsInfoCall().modelName).toBeUndefined();
});
it("keeps the exact model group alongside a typed search", async () => {
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
fireEvent.change(screen.getByPlaceholderText("Search model names…"), { target: { value: "opus" } });
@ -326,7 +461,7 @@ describe("AllModelsTab", () => {
it("resets search, filters, team and sorting from the drawer reset button", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} selectedModelGroup="gpt-4" />);
renderWithProviders(<AllModelsTab {...defaultProps} selectedModelGroup="gpt-4" />);
await user.click(screen.getByTestId("models-team-select"));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
@ -343,7 +478,7 @@ describe("AllModelsTab", () => {
it("opens the delete modal from the row and deletes the model", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(await screen.findByTestId("model-delete-model-1"));
expect(await screen.findByText("Delete Model")).toBeInTheDocument();
@ -357,7 +492,7 @@ describe("AllModelsTab", () => {
it("pauses a model through the row toggle", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(await screen.findByTestId("model-pause-toggle-model-1"));
@ -368,7 +503,7 @@ describe("AllModelsTab", () => {
it("opens the model settings modal from the toolbar", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.queryByTestId("model-settings-modal")).not.toBeInTheDocument();
await user.click(screen.getByTestId("models-settings-trigger"));
@ -377,7 +512,7 @@ describe("AllModelsTab", () => {
it("opens the model detail view from the model ID cell", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(await screen.findByTestId("model-id-model-1"));
@ -386,7 +521,7 @@ describe("AllModelsTab", () => {
it("opens the team detail view from the team ID cell", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(await screen.findByTestId("model-team-id-model-1"));
@ -395,20 +530,20 @@ describe("AllModelsTab", () => {
describe("virtual key hint", () => {
it("explains personal key creation while viewing current team models", () => {
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.getByText(/create a Virtual Key without selecting a team/i)).toBeInTheDocument();
});
it("links the Virtual Keys page through the migrated /ui route", () => {
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
expect(screen.getByRole("link", { name: "Virtual Keys page" })).toHaveAttribute("href", "/ui/api-keys");
});
it("links the team hint's Virtual Keys page through the migrated /ui route", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(screen.getByTestId("models-team-select"));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
@ -419,7 +554,7 @@ describe("AllModelsTab", () => {
it("names the selected team in the hint", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(screen.getByTestId("models-team-select"));
await user.click(await screen.findByRole("option", { name: "Engineering" }));
@ -429,7 +564,7 @@ describe("AllModelsTab", () => {
it("hides the hint when viewing all available models", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
renderWithProviders(<AllModelsTab {...defaultProps} />);
await user.click(screen.getByTestId("models-view-select"));
await user.click(await screen.findByRole("option", { name: "All Available Models" }));

View file

@ -10,10 +10,11 @@ import { toast } from "@/lib/toast";
import { uiHref } from "@/utils/uiHref";
import { modelDeleteCall, modelPatchUpdateCall } from "@/components/networking";
import { useQueryClient } from "@tanstack/react-query";
import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
import { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
import { Info } from "lucide-react";
import { useCallback, useEffect, useMemo, useState } from "react";
import { createParser, parseAsInteger, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs";
import { useCallback, useMemo, useState } from "react";
import { useModelsInfo } from "../../hooks/models/useModels";
import { transformModelData } from "../utils/modelDataTransformer";
@ -24,11 +25,40 @@ import {
PERSONAL_TEAM_VALUE,
WILDCARD_MODEL_GROUP_VALUE,
} from "./AllModelsTable";
import { ACCESS_GROUPS_COLUMN_ID, MODEL_NAME_COLUMN_ID, toServerSortField } from "./ModelsTableColumns";
import {
ACCESS_GROUPS_COLUMN_ID,
isModelTableSortColumnId,
MODEL_NAME_COLUMN_ID,
MODEL_TABLE_SORT_COLUMN_IDS,
toServerSortField,
} from "./ModelsTableColumns";
const SEARCH_DEBOUNCE_WAIT_MS = 200;
const DEFAULT_PAGE_SIZE = 50;
const DEFAULT_PAGINATION: PaginationState = { pageIndex: 0, pageSize: DEFAULT_PAGE_SIZE };
const MAX_PAGE_SIZE = 100;
const MAX_PAGE = 100_000;
const MODEL_VIEW_MODES = ["current_team", "all"] as const satisfies readonly ModelViewMode[];
const boundedInteger = (min: number, max: number, fallback: number) =>
createParser({
parse: (value: string) => {
const parsed = parseAsInteger.parse(value);
return parsed === null ? null : Math.min(Math.max(parsed, min), max);
},
serialize: String,
}).withDefault(fallback);
const TABLE_STATE = {
model_search: parseAsString.withDefault(""),
view_mode: parseAsStringLiteral(MODEL_VIEW_MODES).withDefault("current_team"),
filter_team: parseAsString.withDefault(PERSONAL_TEAM_VALUE),
access_group: parseAsString.withDefault(""),
sort_by: parseAsStringLiteral(MODEL_TABLE_SORT_COLUMN_IDS),
sort_order: parseAsStringLiteral(["asc", "desc"] as const).withDefault("asc"),
page: boundedInteger(1, MAX_PAGE, 1),
page_size: boundedInteger(1, MAX_PAGE_SIZE, DEFAULT_PAGE_SIZE),
};
interface AllModelsTabProps {
selectedModelGroup: string | null;
@ -52,34 +82,25 @@ const AllModelsTab = ({
const { data: teams, isLoading: isLoadingTeams } = useTeams();
const queryClient = useQueryClient();
const [modelNameSearch, setModelNameSearch] = useState<string>("");
const [debouncedSearch, setDebouncedSearch] = useState<string>("");
const [modelViewMode, setModelViewMode] = useState<ModelViewMode>("current_team");
const [selectedTeamValue, setSelectedTeamValue] = useState<string>(PERSONAL_TEAM_VALUE);
const [selectedModelAccessGroupFilter, setSelectedModelAccessGroupFilter] = useState<string | null>(null);
const [pagination, setPagination] = useState<PaginationState>(DEFAULT_PAGINATION);
const [sorting, setSorting] = useState<SortingState>([]);
const [tableState, setTableState] = useQueryStates(TABLE_STATE);
const modelNameSearch = tableState.model_search;
const [debouncedSearch] = useDebouncedValue(modelNameSearch, { wait: SEARCH_DEBOUNCE_WAIT_MS });
const modelViewMode = tableState.view_mode;
const selectedTeamValue = tableState.filter_team;
const selectedModelAccessGroupFilter = tableState.access_group || null;
const pagination = useMemo<PaginationState>(
() => ({ pageIndex: tableState.page - 1, pageSize: tableState.page_size }),
[tableState.page, tableState.page_size],
);
const sorting = useMemo<SortingState>(
() => (tableState.sort_by ? [{ id: tableState.sort_by, desc: tableState.sort_order === "desc" }] : []),
[tableState.sort_by, tableState.sort_order],
);
const [isModelSettingsModalVisible, setIsModelSettingsModalVisible] = useState(false);
const [deleteModalModelId, setDeleteModalModelId] = useState<string | null>(null);
const [deleteLoading, setDeleteLoading] = useState(false);
const [pausingModelId, setPausingModelId] = useState<string | null>(null);
const resetToFirstPage = useCallback(() => {
setPagination((previous) => (previous.pageIndex === 0 ? previous : { ...previous, pageIndex: 0 }));
}, []);
const debouncedUpdateSearch = useDebouncedCallback(
(value: string) => {
setDebouncedSearch(value);
resetToFirstPage();
},
{ wait: SEARCH_DEBOUNCE_WAIT_MS },
);
useEffect(() => {
debouncedUpdateSearch(modelNameSearch);
}, [modelNameSearch, debouncedUpdateSearch]);
const teamIdForQuery = selectedTeamValue === PERSONAL_TEAM_VALUE ? undefined : selectedTeamValue;
const isConcreteModelGroup =
Boolean(selectedModelGroup) &&
@ -152,33 +173,49 @@ const AllModelsTab = ({
[selectedModelGroup, selectedModelAccessGroupFilter],
);
const handleSearchChange = useCallback(
(value: string) => {
void setTableState({ model_search: value || null, page: null });
},
[setTableState],
);
const handleColumnFiltersChange: OnChangeFn<ColumnFiltersState> = (updater) => {
const next = typeof updater === "function" ? updater(columnFilters) : updater;
const next = functionalUpdate(updater, columnFilters);
const modelGroup = next.find((entry) => entry.id === MODEL_NAME_COLUMN_ID)?.value;
const accessGroup = next.find((entry) => entry.id === ACCESS_GROUPS_COLUMN_ID)?.value;
setSelectedModelGroup(typeof modelGroup === "string" ? modelGroup : ALL_MODEL_GROUPS_VALUE);
setSelectedModelAccessGroupFilter(typeof accessGroup === "string" ? accessGroup : null);
resetToFirstPage();
void setTableState({ access_group: typeof accessGroup === "string" ? accessGroup : null, page: null });
};
const handleSortingChange: OnChangeFn<SortingState> = (updater) => {
setSorting(typeof updater === "function" ? updater(sorting) : updater);
resetToFirstPage();
const active = functionalUpdate(updater, sorting)[0];
void setTableState({
sort_by: active && isModelTableSortColumnId(active.id) ? active.id : null,
sort_order: active?.desc ? "desc" : null,
page: null,
});
};
const handlePaginationChange = useCallback<OnChangeFn<PaginationState>>(
(updater) => {
const next = functionalUpdate(updater, pagination);
void setTableState({ page: next.pageIndex + 1, page_size: next.pageSize });
},
[pagination, setTableState],
);
const handleTeamChange = (value: string) => {
setSelectedTeamValue(value);
resetToFirstPage();
void setTableState({ filter_team: value, page: null });
};
const handleViewModeChange = (value: ModelViewMode) => {
void setTableState({ view_mode: value });
};
const resetFilters = () => {
setModelNameSearch("");
setSelectedModelGroup(ALL_MODEL_GROUPS_VALUE);
setSelectedModelAccessGroupFilter(null);
setSelectedTeamValue(PERSONAL_TEAM_VALUE);
setModelViewMode("current_team");
setPagination(DEFAULT_PAGINATION);
setSorting([]);
void setTableState(null);
};
const teamOptions = useMemo(
@ -264,18 +301,18 @@ const AllModelsTab = ({
sorting={sorting}
onSortingChange={handleSortingChange}
pagination={pagination}
onPaginationChange={setPagination}
onPaginationChange={handlePaginationChange}
columnFilters={columnFilters}
onColumnFiltersChange={handleColumnFiltersChange}
onResetFilters={resetFilters}
searchValue={modelNameSearch}
onSearchChange={setModelNameSearch}
onSearchChange={handleSearchChange}
teamOptions={teamOptions}
selectedTeamValue={selectedTeamValue}
onTeamChange={handleTeamChange}
isLoadingTeams={isLoadingTeams}
viewMode={modelViewMode}
onViewModeChange={setModelViewMode}
onViewModeChange={handleViewModeChange}
onOpenModelSettings={handleOpenModelSettings}
availableModelGroups={availableModelGroups}
availableModelAccessGroups={availableModelAccessGroups}

View file

@ -24,6 +24,19 @@ export const TEAM_ID_COLUMN_ID = "model_info_team_id";
export const ACCESS_GROUPS_COLUMN_ID = "model_info_access_groups";
export const STATUS_COLUMN_ID = "model_info_db_model";
export const MODEL_TABLE_SORT_COLUMN_IDS = [
MODEL_NAME_COLUMN_ID,
CREATED_BY_COLUMN_ID,
UPDATED_AT_COLUMN_ID,
COSTS_COLUMN_ID,
STATUS_COLUMN_ID,
] as const;
export type ModelTableSortColumnId = (typeof MODEL_TABLE_SORT_COLUMN_IDS)[number];
export const isModelTableSortColumnId = (columnId: string): columnId is ModelTableSortColumnId =>
(MODEL_TABLE_SORT_COLUMN_IDS as readonly string[]).includes(columnId);
const COLUMN_ID_TO_SERVER_SORT_FIELD: Record<string, string> = {
[COSTS_COLUMN_ID]: "costs",
[STATUS_COLUMN_ID]: "status",

View file

@ -35,10 +35,12 @@ beforeEach(() => {
Element.prototype.scrollIntoView = () => {};
});
const CHAT_REQUEST_ARG_COUNT = 26;
const CHAT_REQUEST_ARG_COUNT = 27;
const STREAMING_ENABLED_ARG_INDEX = 25;
const MESSAGES_REQUEST_ARG_COUNT = 19;
const CHAT_CUSTOM_HEADERS_ARG_INDEX = 26;
const MESSAGES_REQUEST_ARG_COUNT = 20;
const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18;
const MESSAGES_CUSTOM_HEADERS_ARG_INDEX = 19;
async function openComboboxByPlaceholder(placeholder: string) {
const user = userEvent.setup();
@ -447,6 +449,63 @@ describe("ChatUI", () => {
expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false);
});
it("should send custom headers entered in the sidebar with /v1/chat/completions and /v1/messages requests", async () => {
const user = userEvent.setup();
render(
<ChatUI
accessToken="1234567890"
token="1234567890"
userRole="user"
userID="1234567890"
disabledPersonalKeyCreation={false}
/>,
);
await waitFor(() => {
expect(screen.getByText("Test Key")).toBeInTheDocument();
});
await selectComboboxOption("Select a Model", "Model 1");
await user.click(screen.getByRole("button", { name: "Add Header" }));
await user.click(screen.getByRole("button", { name: "Add Header" }));
const [firstName] = screen.getAllByPlaceholderText("Header Name");
const [firstValue, secondValue] = screen.getAllByPlaceholderText("Header Value");
fireEvent.change(firstName, { target: { value: "anthropic-beta" } });
fireEvent.change(firstValue, { target: { value: "context-1m-2025-08-07" } });
fireEvent.change(secondValue, { target: { value: "ignored because the name is blank" } });
const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)");
await act(async () => {
fireEvent.change(messageInput, { target: { value: "hello" } });
});
await act(async () => {
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
});
await waitFor(() => {
expect(makeOpenAIChatCompletionRequest).toHaveBeenCalledTimes(1);
});
const chatArgs = vi.mocked(makeOpenAIChatCompletionRequest).mock.calls[0];
expect(chatArgs).toHaveLength(CHAT_REQUEST_ARG_COUNT);
expect(chatArgs[CHAT_CUSTOM_HEADERS_ARG_INDEX]).toEqual({ "anthropic-beta": "context-1m-2025-08-07" });
await selectComboboxOption("Select an endpoint", "/v1/messages");
await selectComboboxOption("Select a Model", "Model 1");
await act(async () => {
fireEvent.change(messageInput, { target: { value: "hello again" } });
});
await act(async () => {
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
});
await waitFor(() => {
expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1);
});
const messagesArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0];
expect(messagesArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT);
expect(messagesArgs[MESSAGES_CUSTOM_HEADERS_ARG_INDEX]).toEqual({ "anthropic-beta": "context-1m-2025-08-07" });
});
it("should force streaming in simplified mode even when the playground setting is off", async () => {
sessionStorage.setItem("streamingEnabled", "false");

View file

@ -9,6 +9,7 @@ import {
Info,
Key,
Link2,
ListPlus,
Loader2,
Settings,
Shield,
@ -40,6 +41,8 @@ import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages
import { makeOpenAIAudioSpeechRequest } from "../../llm_calls/audio_speech";
import { makeOpenAIAudioTranscriptionRequest } from "../../llm_calls/audio_transcriptions";
import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion";
import { customHeadersFromPairs, parseStoredHeaderPairs } from "@/components/llm_calls/request_headers";
import KeyValueInput, { type KeyValuePair } from "@/components/key_value_input";
import { makeOpenAIEmbeddingsRequest } from "../../llm_calls/embeddings_api";
import { Agent, fetchAvailableAgents } from "../../llm_calls/fetch_agents";
import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models";
@ -220,6 +223,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
return [];
}
});
const [customHeaderPairs, setCustomHeaderPairs] = useState<readonly KeyValuePair[]>(() =>
parseStoredHeaderPairs(getSecureItem("customHeaders")),
);
const customHeaders = useMemo(() => customHeadersFromPairs(customHeaderPairs), [customHeaderPairs]);
const [selectedVoice, setSelectedVoice] = useState<OpenAIVoice>(() => {
const saved = sessionStorage.getItem("selectedVoice");
if (!saved) return "alloy";
@ -346,6 +353,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedSdk,
selectedVoice,
proxySettings,
customHeaders,
});
setGeneratedCode(code);
}
@ -367,12 +375,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
endpointType,
selectedModel,
proxySettings,
customHeaders,
]);
useEffect(() => {
try {
setSecureItem("apiKeySource", JSON.stringify(apiKeySource));
setSecureItem("apiKey", apiKey);
setSecureItem("customHeaders", JSON.stringify(customHeaderPairs));
} catch {
// Storage full or unavailable — non-critical, skip persisting.
}
@ -410,6 +420,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mcpServerToolRestrictions,
selectedVoice,
streamingEnabled,
customHeaderPairs,
]);
useEffect(() => {
@ -921,6 +932,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mockTestFallbacks,
mcpToolsets,
streamingEnabled,
customHeaders,
);
} else if (endpointType === EndpointType.IMAGE) {
// For image generation
@ -932,6 +944,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedTags,
signal,
customProxyBaseUrl || undefined,
customHeaders,
);
} else if (endpointType === EndpointType.SPEECH) {
// For audio speech
@ -946,6 +959,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
undefined, // responseFormat
undefined, // speed
customProxyBaseUrl || undefined,
customHeaders,
);
} else if (endpointType === EndpointType.IMAGE_EDITS) {
// For image edits
@ -959,6 +973,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedTags,
signal,
customProxyBaseUrl || undefined,
customHeaders,
);
}
} else if (endpointType === EndpointType.RESPONSES) {
@ -1004,6 +1019,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mcpToolsets,
streamingEnabled,
updateTotalLatency,
customHeaders,
);
} else if (endpointType === EndpointType.ANTHROPIC_MESSAGES) {
const apiChatHistory = [
@ -1033,6 +1049,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mcpServerToolRestrictions,
mcpToolsets,
streamingEnabled,
customHeaders,
);
} else if (endpointType === EndpointType.EMBEDDINGS) {
await makeOpenAIEmbeddingsRequest(
@ -1042,6 +1059,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
effectiveApiKey,
selectedTags,
customProxyBaseUrl || undefined,
customHeaders,
);
} else if (endpointType === EndpointType.TRANSCRIPTION) {
// For audio transcriptions
@ -1058,6 +1076,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
undefined, // responseFormat
undefined, // temperature
customProxyBaseUrl || undefined,
customHeaders,
);
}
} else if (endpointType === EndpointType.INTERACTIONS) {
@ -1069,6 +1088,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
selectedTags,
signal,
customProxyBaseUrl || undefined,
undefined,
customHeaders,
);
}
}
@ -1086,13 +1107,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
resolvedServerId = toolEntry?.server_id ?? rawSelected;
}
if (resolvedServerId && !resolvedServerId.startsWith("toolset:") && selectedMCPDirectTool) {
const result = await callMCPTool(
effectiveApiKey,
resolvedServerId,
selectedMCPDirectTool,
mcpToolArguments,
selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : undefined,
);
const result = await callMCPTool(effectiveApiKey, resolvedServerId, selectedMCPDirectTool, mcpToolArguments, {
...(selectedGuardrails.length > 0 ? { guardrails: selectedGuardrails } : {}),
customHeaders,
});
const resultText =
result?.content?.length > 0
? JSON.stringify(
@ -1118,6 +1136,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
updateA2AMetadata,
customProxyBaseUrl || undefined,
selectedGuardrails.length > 0 ? selectedGuardrails : undefined,
customHeaders,
);
}
} catch (error) {
@ -1485,6 +1504,18 @@ const ChatUI: React.FC<ChatUIProps> = ({
/>
</div>
{endpointType !== EndpointType.REALTIME && (
<div>
<label className="mb-2 flex items-center text-sm font-medium text-foreground">
<ListPlus className="mr-2 size-4" aria-hidden="true" /> Custom Headers
</label>
<KeyValueInput value={customHeaderPairs} onChange={setCustomHeaderPairs} />
<p className="mt-2 text-xs text-muted-foreground">
Sent with every playground request, e.g. provider-specific headers like anthropic-beta.
</p>
</div>
)}
<div>
<div className="mb-2 flex items-center gap-1 text-sm font-medium text-foreground">
<Wrench className="mr-1 size-4" aria-hidden="true" />

View file

@ -3,6 +3,7 @@
import { v4 as uuidv4 } from "uuid";
import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking";
import { type CustomHeaders, withRequiredHeaders } from "@/components/llm_calls/request_headers";
import { A2ATaskMetadata } from "@/components/chat_ui/types";
interface A2AMessagePart {
@ -116,6 +117,7 @@ export const makeA2ASendMessageRequest = async (
onA2AMetadata?: (metadata: A2ATaskMetadata) => void,
customBaseUrl?: string,
guardrails?: string[],
customHeaders?: CustomHeaders,
): Promise<void> => {
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
const url = proxyBaseUrl ? `${proxyBaseUrl}/a2a/${agentId}/message/send` : `/a2a/${agentId}/message/send`;
@ -146,10 +148,10 @@ export const makeA2ASendMessageRequest = async (
try {
const response = await fetch(url, {
method: "POST",
headers: {
headers: withRequiredHeaders(customHeaders ?? {}, {
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
}),
body: JSON.stringify(jsonRpcRequest),
signal,
});

View file

@ -1,4 +1,5 @@
import { afterEach, describe, expect, it, vi } from "vitest";
import Anthropic from "@anthropic-ai/sdk";
import { makeAnthropicMessagesRequest } from "./anthropic_messages";
import type { TokenUsage } from "@/components/chat_ui/ResponseMetrics";
@ -122,4 +123,26 @@ describe("anthropic_messages non-streaming", () => {
expect(mockMessagesCreate).not.toHaveBeenCalled();
expect(mockMessagesStream.mock.calls[0][0]).toMatchObject({ stream: true });
});
it("sends custom headers alongside the tags header on the Anthropic client", async () => {
mockMessagesCreate.mockResolvedValue({ content: [{ type: "text", text: "OK" }], usage: {} });
await makeAnthropicMessagesRequest(
[{ role: "user", content: "Hello" }],
vi.fn(),
"claude-haiku-4-5",
"test-token",
["team-a"],
undefined,
undefined,
undefined,
undefined,
...NON_STREAMING_ARGS,
{ "anthropic-beta": "context-1m-2025-08-07" },
);
expect(vi.mocked(Anthropic).mock.calls[0][0]).toMatchObject({
defaultHeaders: { "x-litellm-tags": "team-a", "anthropic-beta": "context-1m-2025-08-07" },
});
});
});

View file

@ -2,6 +2,7 @@ import Anthropic from "@anthropic-ai/sdk";
import { MessageType } from "@/components/chat_ui/types";
import { TokenUsage } from "@/components/chat_ui/ResponseMetrics";
import { buildMcpToolBlocks } from "@/components/llm_calls/mcp_tool_blocks";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { MCPServer, MCPToolset } from "@/components/mcp_tools/types";
import { getProxyBaseUrl } from "@/components/networking";
import { toast } from "@/lib/toast";
@ -34,6 +35,7 @@ export async function makeAnthropicMessagesRequest(
mcpServerToolRestrictions?: Record<string, string[]>,
mcpToolsets?: MCPToolset[],
streamingEnabled: boolean = true,
customHeaders?: CustomHeaders,
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -46,11 +48,7 @@ export async function makeAnthropicMessagesRequest(
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
// Prepare headers with tags and trace ID
const headers: Record<string, string> = {};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const headers = buildPlaygroundHeaders(tags, customHeaders);
const client = new Anthropic({
apiKey: accessToken,

View file

@ -1,5 +1,6 @@
import openai from "openai";
import { getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { toast } from "@/lib/toast";
import type { OpenAIVoice } from "../components/chat_ui/chatConstants";
@ -14,6 +15,7 @@ export async function makeOpenAIAudioSpeechRequest(
responseFormat?: string,
speed?: number,
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -25,7 +27,7 @@ export async function makeOpenAIAudioSpeechRequest(
apiKey: accessToken,
baseURL: proxyBaseUrl,
dangerouslyAllowBrowser: true,
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
});
try {

View file

@ -1,5 +1,6 @@
import openai from "openai";
import { getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { toast } from "@/lib/toast";
export async function makeOpenAIAudioTranscriptionRequest(
@ -14,6 +15,7 @@ export async function makeOpenAIAudioTranscriptionRequest(
responseFormat?: string,
temperature?: number,
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -26,7 +28,7 @@ export async function makeOpenAIAudioTranscriptionRequest(
apiKey: accessToken,
baseURL: proxyBaseUrl,
dangerouslyAllowBrowser: true,
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
});
try {

View file

@ -76,4 +76,42 @@ describe("embeddings_api", () => {
input: "Sample text",
});
});
it("sends custom headers on the fetch request, letting them override the tags header", async () => {
await makeOpenAIEmbeddingsRequest(
"Sample text",
mockUpdateEmbeddingsUI,
"text-embedding-3-small",
"abcdef",
["team-a"],
undefined,
{ "x-litellm-tags": "team-b", "x-request-source": "playground" },
);
expect(mockFetch.mock.calls[0][1]).toMatchObject({
headers: {
Authorization: "Bearer abcdef",
"x-litellm-tags": "team-b",
"x-request-source": "playground",
},
});
});
it("does not let custom headers replace the gateway auth or content-type headers", async () => {
await makeOpenAIEmbeddingsRequest(
"Sample text",
mockUpdateEmbeddingsUI,
"text-embedding-3-small",
"abcdef",
undefined,
undefined,
{ authorization: "Bearer stolen", "Content-Type": "text/plain", "x-request-source": "playground" },
);
expect(mockFetch.mock.calls[0][1].headers).toEqual({
Authorization: "Bearer abcdef",
"Content-Type": "application/json",
"x-request-source": "playground",
});
});
});

View file

@ -1,5 +1,10 @@
import { toast } from "@/lib/toast";
import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking";
import {
buildPlaygroundHeaders,
type CustomHeaders,
withRequiredHeaders,
} from "@/components/llm_calls/request_headers";
export async function makeOpenAIEmbeddingsRequest(
input: string,
@ -8,6 +13,7 @@ export async function makeOpenAIEmbeddingsRequest(
accessToken: string,
tags?: string[],
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -20,11 +26,10 @@ export async function makeOpenAIEmbeddingsRequest(
}
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
// Prepare headers with tags and trace ID
const headers: Record<string, string> = {};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const headers = withRequiredHeaders(buildPlaygroundHeaders(tags, customHeaders), {
"Content-Type": "application/json",
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
});
try {
const normalizedBaseUrl = proxyBaseUrl.endsWith("/") ? proxyBaseUrl.slice(0, -1) : proxyBaseUrl;
@ -32,11 +37,7 @@ export async function makeOpenAIEmbeddingsRequest(
const response = await fetch(requestUrl, {
method: "POST",
headers: {
"Content-Type": "application/json",
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
...headers,
},
headers,
body: JSON.stringify({
model: selectedModel,
input,

View file

@ -1,5 +1,6 @@
import openai from "openai";
import { getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { toast } from "@/lib/toast";
export async function makeOpenAIImageEditsRequest(
@ -11,6 +12,7 @@ export async function makeOpenAIImageEditsRequest(
tags?: string[],
signal?: AbortSignal,
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -23,7 +25,7 @@ export async function makeOpenAIImageEditsRequest(
apiKey: accessToken,
baseURL: proxyBaseUrl,
dangerouslyAllowBrowser: true,
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
});
try {

View file

@ -1,5 +1,6 @@
import openai from "openai";
import { getProxyBaseUrl } from "@/components/networking";
import { buildPlaygroundHeaders, type CustomHeaders } from "@/components/llm_calls/request_headers";
import { toast } from "@/lib/toast";
export async function makeOpenAIImageGenerationRequest(
@ -10,6 +11,7 @@ export async function makeOpenAIImageGenerationRequest(
tags?: string[],
signal?: AbortSignal,
customBaseUrl?: string,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -21,7 +23,7 @@ export async function makeOpenAIImageGenerationRequest(
apiKey: accessToken,
baseURL: proxyBaseUrl,
dangerouslyAllowBrowser: true,
defaultHeaders: tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : undefined,
defaultHeaders: buildPlaygroundHeaders(tags, customHeaders),
});
try {

View file

@ -1,5 +1,10 @@
import { toast } from "@/lib/toast";
import { getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking";
import {
buildPlaygroundHeaders,
type CustomHeaders,
withRequiredHeaders,
} from "@/components/llm_calls/request_headers";
export async function makeInteractionsRequest(
input: string,
@ -10,6 +15,7 @@ export async function makeInteractionsRequest(
signal?: AbortSignal,
customBaseUrl?: string,
previousInteractionId?: string,
customHeaders?: CustomHeaders,
): Promise<void> {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -24,13 +30,10 @@ export async function makeInteractionsRequest(
const normalizedBaseUrl = proxyBaseUrl.endsWith("/") ? proxyBaseUrl.slice(0, -1) : proxyBaseUrl;
const requestUrl = `${normalizedBaseUrl}/v1beta/interactions`;
const headers: Record<string, string> = {
const headers: Record<string, string> = withRequiredHeaders(buildPlaygroundHeaders(tags, customHeaders), {
"Content-Type": "application/json",
[getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`,
};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
});
const body: Record<string, unknown> = {
model: selectedModel,

View file

@ -45,6 +45,15 @@ const SETTINGS_FIXTURE = [
field_tab: "prompt_caching",
field_default_value: null,
},
{
field_name: "openai_system_messages_first",
field_type: "Boolean",
field_value: false,
field_description: "openai system first toggle",
stored_in_db: null,
field_tab: "prompt_caching",
field_default_value: false,
},
{
field_name: "max_ui_session_budget",
field_type: "Dollar",
@ -157,6 +166,39 @@ describe("GeneralSettings General tab", () => {
});
});
describe("GeneralSettings Prompt Caching tab", () => {
beforeEach(() => {
vi.mocked(getGeneralSettingsCall).mockResolvedValue([...SETTINGS_FIXTURE.map((s) => ({ ...s }))]);
vi.mocked(updateConfigFieldSetting).mockClear();
vi.mocked(deleteConfigFieldSetting).mockClear();
});
it("persists openai_system_messages_first when its switch is turned on", async () => {
const user = userEvent.setup();
renderWithProviders(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);
await user.click(await screen.findByRole("tab", { name: "Prompt Caching" }));
const toggle = await screen.findByRole("switch", { name: "System messages first for OpenAI" });
expect(toggle).not.toBeChecked();
await user.click(toggle);
expect(toggle).toBeChecked();
expect(updateConfigFieldSetting).toHaveBeenCalledWith("token", "openai_system_messages_first", true);
expect(deleteConfigFieldSetting).not.toHaveBeenCalled();
});
it("keeps the prompt caching rows off the General tab table", async () => {
const user = userEvent.setup();
renderWithProviders(<GeneralSettings accessToken="token" userRole="Admin" userID="user" />);
await user.click(screen.getByText("General"));
await settingsRow("max_ui_session_budget");
expect(screen.queryByText("openai_system_messages_first")).not.toBeInTheDocument();
});
});
// The five tabs here are proxy-wide settings. Auto-routers moved to Models + Endpoints.
describe("GeneralSettings tabs", () => {
beforeEach(() => {

View file

@ -18,6 +18,9 @@ import RoutingGroups from "@/components/routing_groups";
const PROMPT_CACHING_TAB = "prompt_caching";
const ENABLE_ANTHROPIC_PROMPT_CACHING = "enable_anthropic_prompt_caching";
const ANTHROPIC_PROMPT_CACHING_TTL = "anthropic_prompt_caching_ttl";
const OPENAI_SYSTEM_MESSAGES_FIRST = "openai_system_messages_first";
const isOn = (value: unknown) => value === true || value === "true";
interface GeneralSettingsPageProps {
accessToken: string | null;
@ -129,14 +132,15 @@ export const PromptCachingPanel: React.FC<{
}> = ({ accessToken, settings, onChange }) => {
const enableSetting = settings.find((s) => s.field_name === ENABLE_ANTHROPIC_PROMPT_CACHING);
const ttlSetting = settings.find((s) => s.field_name === ANTHROPIC_PROMPT_CACHING_TTL);
const systemFirstSetting = settings.find((s) => s.field_name === OPENAI_SYSTEM_MESSAGES_FIRST);
// The two rows come from the same registry the General tab reads; if they
// The rows come from the same registry the General tab reads; if they
// are not loaded yet there is nothing to render.
if (!enableSetting) {
return null;
}
const enabled = enableSetting.field_value === true || enableSetting.field_value === "true";
const enabled = isOn(enableSetting.field_value);
// Apply immediately: a toggle and a dropdown are direct controls, so there is
// no separate Update button. Clearing the ttl resets it to the provider default.
@ -187,6 +191,20 @@ export const PromptCachingPanel: React.FC<{
</Select>
</div>
)}
{systemFirstSetting && (
<div className="mt-6 flex items-start justify-between gap-8">
<div className="min-w-0 max-w-2xl">
<p className="font-medium">System messages first for OpenAI</p>
<p className="mt-1 break-words text-xs text-muted-foreground">{systemFirstSetting.field_description}</p>
</div>
<Switch
aria-label="System messages first for OpenAI"
checked={isOn(systemFirstSetting.field_value)}
onCheckedChange={(checked) => persist(OPENAI_SYSTEM_MESSAGES_FIRST, checked)}
/>
</div>
)}
</CardContent>
</Card>
);

View file

@ -42,9 +42,10 @@ import { Restricted, restrictedBy } from "./TierRestrictions";
import { type TierSetAction, applyTierSetAction, setFallbackTier } from "./tier_set_actions";
import {
ReasoningEffort,
TierModelParamChange,
TierModelParamsByTier,
classifierEffortOptionsForModels,
setTierModelReasoningEffort,
setTierModelParam,
tierEffortOptionsForModels,
tierRowLabel,
} from "./complexity_router_tiers";
@ -621,6 +622,9 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
const exitToBuiltInTiers = () => dispatch({ kind: "restore" });
const tierEffortOptionsByModel = tierEffortOptionsForModels(modelInfo);
const fastModeByModel = Object.fromEntries(
modelInfo.map((model) => [model.model_group, model.supports_fast_mode === true]),
);
const classifierEffortOptionsByModel = classifierEffortOptionsForModels(modelInfo);
// Embedding models can't serve a chat-completion role, so they're excluded here.
@ -631,12 +635,11 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
label: model.model_group,
}));
const handleTierModelEffortChange = (tier: string, model: string, effort: ReasoningEffort | undefined) => {
const handleTierModelParamChange = (tier: string, model: string, change: TierModelParamChange) =>
onChange({
...value,
tier_model_params: setTierModelReasoningEffort(value.tier_model_params, tier, model, effort),
tier_model_params: setTierModelParam(value.tier_model_params, tier, model, change),
});
};
// Clearing the select drops the key entirely rather than storing "", so an emptied pin reads as
// "track the tiers" everywhere downstream instead of as a blank model name.
@ -734,7 +737,13 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
models={row.models}
effortOptionsByModel={tierEffortOptionsByModel}
paramsByModel={row.params}
onEffortChange={(model, effort) => handleTierModelEffortChange(row.id, model, effort)}
fastModeByModel={fastModeByModel}
onEffortChange={(model, effort) =>
handleTierModelParamChange(row.id, model, ["reasoning_effort", effort])
}
onFastModeChange={(model, enabled) =>
handleTierModelParamChange(row.id, model, ["speed", enabled ? "fast" : undefined])
}
/>
{row.models.length > 1 && (
<span className="text-xs text-muted-foreground">

View file

@ -0,0 +1,143 @@
import userEvent from "@testing-library/user-event";
import React from "react";
import { describe, expect, it, vi } from "vitest";
import { renderWithProviders, screen } from "../../../tests/test-utils";
import {
buildUpdatedComplexityRouterConfig,
hydrateComplexityRouterConfig,
} from "../edit_auto_router/edit_auto_router_modal";
import type { ModelGroup } from "../llm_calls/fetch_models";
import ComplexityRouterConfig, { type ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
const modelInfo: ModelGroup[] = [
{ model_group: "primary", supported_reasoning_efforts: ["low", "high"], supports_fast_mode: true },
{ model_group: "secondary", supports_fast_mode: true },
{ model_group: "blocked", supported_reasoning_efforts: ["low"], supports_fast_mode: false },
{ model_group: "missing", supported_reasoning_efforts: ["low"] },
];
it.each([false, true])("edits and round-trips independent model settings with custom tiers=%s", async (custom) => {
const user = userEvent.setup();
const tier = custom ? "custom-a" : "COMPLEX";
const otherTier = custom ? "custom-b" : "REASONING";
const label = custom ? "Interactive" : "Complex";
const models = ["primary", "secondary", "blocked", "missing"];
const initial: ComplexityRouterConfigValue = {
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: models, REASONING: ["primary"] },
classifier_type: "heuristic",
...(custom && {
custom_tier_set: {
tiers: [
{ id: tier, name: label, definition: "Interactive requests", models },
{ id: otherTier, name: "Deliberate", definition: "Careful requests", models: ["primary"] },
],
fallback_tier_id: tier,
},
}),
tier_model_params: {
[tier]: {
primary: { reasoning_effort: "high", max_tokens: 1024 },
secondary: { speed: "fast" },
blocked: { speed: "fast" },
},
[otherTier]: { primary: { speed: "fast", reasoning_effort: "low" } },
},
};
const onChange = vi.fn<(value: ComplexityRouterConfigValue) => void>();
const editor = (value: ComplexityRouterConfigValue) => (
<ComplexityRouterConfig modelInfo={modelInfo} value={value} onChange={onChange} />
);
const view = renderWithProviders(editor(initial));
const fast = () => screen.getByRole("switch", { name: `Fast mode for primary in the ${label} tier` });
expect(screen.getAllByRole("switch", { name: /^Fast mode for/ })).toHaveLength(3);
expect(screen.queryByRole("switch", { name: /^Fast mode for (blocked|missing)/ })).not.toBeInTheDocument();
expect(screen.queryByRole("combobox", { name: /^Reasoning effort for secondary/ })).not.toBeInTheDocument();
expect(screen.getByRole("switch", { name: `Fast mode for secondary in the ${label} tier` })).toBeChecked();
expect(fast()).not.toBeChecked();
expect(onChange).not.toHaveBeenCalled();
await user.click(fast());
const enabled = onChange.mock.lastCall![0];
expect(enabled.tier_model_params).toEqual({
...initial.tier_model_params,
[tier]: {
...initial.tier_model_params![tier],
primary: { reasoning_effort: "high", max_tokens: 1024, speed: "fast" },
},
});
const saved = buildUpdatedComplexityRouterConfig({}, enabled);
expect(saved.tier_model_configs).toEqual({
[custom ? label : tier]: [
{ model_name: "primary", litellm_params: { reasoning_effort: "high", max_tokens: 1024, speed: "fast" } },
{ model_name: "secondary", litellm_params: { speed: "fast" } },
{ model_name: "blocked", litellm_params: { speed: "fast" } },
],
[custom ? "Deliberate" : otherTier]: [
{ model_name: "primary", litellm_params: { speed: "fast", reasoning_effort: "low" } },
],
});
const reopened = hydrateComplexityRouterConfig(saved, undefined);
const reopenedTier = custom ? reopened.custom_tier_set!.tiers[0].id : tier;
view.rerender(editor(reopened));
expect(fast()).toBeChecked();
await user.click(screen.getByRole("combobox", { name: `Reasoning effort for primary in the ${label} tier` }));
await user.click(await screen.findByRole("option", { name: "low" }));
const effortChanged = onChange.mock.lastCall![0];
expect(effortChanged.tier_model_params?.[reopenedTier].primary).toEqual({
reasoning_effort: "low",
max_tokens: 1024,
speed: "fast",
});
view.rerender(editor(effortChanged));
await user.click(fast());
const disabled = onChange.mock.lastCall![0];
expect(disabled.tier_model_params).toEqual({
...effortChanged.tier_model_params,
[reopenedTier]: {
...effortChanged.tier_model_params![reopenedTier],
primary: { reasoning_effort: "low", max_tokens: 1024 },
},
});
view.rerender(editor(disabled));
expect(fast()).not.toBeChecked();
const picker = () => screen.getByRole("combobox", { name: `Select model(s) for ${label.toLowerCase()} queries` });
await user.click(picker());
await user.click(await screen.findByRole("option", { name: "primary" }));
await user.keyboard("{Escape}");
const deselected = onChange.mock.lastCall![0];
expect(deselected.tier_model_params?.[reopenedTier]).toEqual({
secondary: { speed: "fast" },
blocked: { speed: "fast" },
});
view.rerender(editor(deselected));
expect(screen.queryByRole("switch", { name: `Fast mode for primary in the ${label} tier` })).not.toBeInTheDocument();
await user.click(picker());
await user.click(await screen.findByRole("option", { name: "primary" }));
await user.keyboard("{Escape}");
const reselected = onChange.mock.lastCall![0];
view.rerender(editor(reselected));
expect(fast()).not.toBeChecked();
expect(screen.getByRole("combobox", { name: `Reasoning effort for primary in the ${label} tier` })).toHaveTextContent(
"Default",
);
});
describe("Fast mode metadata", () => {
it("offers nothing before model capabilities load and leaves stored speed untouched", () => {
const value: ComplexityRouterConfigValue = {
tiers: { SIMPLE: ["primary"], MEDIUM: [], COMPLEX: [], REASONING: [] },
classifier_type: "heuristic",
tier_model_params: { SIMPLE: { primary: { speed: "fast" } } },
};
const onChange = vi.fn();
renderWithProviders(<ComplexityRouterConfig modelInfo={[]} value={value} onChange={onChange} />);
expect(screen.queryByRole("switch", { name: /^Fast mode for/ })).not.toBeInTheDocument();
expect(onChange).not.toHaveBeenCalled();
expect(buildUpdatedComplexityRouterConfig({}, value).tier_model_configs).toEqual({
SIMPLE: [{ model_name: "primary", litellm_params: { speed: "fast" } }],
});
});
});

View file

@ -1,5 +1,6 @@
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { SimpleTooltip } from "@/components/ui/tooltip";
import { Switch } from "@/components/ui/switch";
import { Info } from "lucide-react";
import React from "react";
import { ReasoningEffort, TierModelParams } from "./complexity_router_tiers";
@ -18,6 +19,8 @@ interface TierModelEffortRowsProps {
effortOptionsByModel: Record<string, string[]>;
paramsByModel: Record<string, TierModelParams> | undefined;
onEffortChange: (model: string, effort: ReasoningEffort | undefined) => void;
fastModeByModel?: Record<string, boolean>;
onFastModeChange: (model: string, enabled: boolean) => void;
}
export interface TierEffortRow {
@ -29,13 +32,16 @@ export interface TierEffortRow {
/**
* A stored effort outside the model's supported set (hand-authored, or capabilities changed since
* it was saved) is listed anyway, so the row renders with its value selected and can be cleared.
* Only a model with no supported level and nothing stored drops out.
*/
export const tierEffortRows = ({
models,
effortOptionsByModel,
paramsByModel,
}: Pick<TierModelEffortRowsProps, "models" | "effortOptionsByModel" | "paramsByModel">): TierEffortRow[] =>
fastModeByModel,
}: Pick<
TierModelEffortRowsProps,
"models" | "effortOptionsByModel" | "paramsByModel" | "fastModeByModel"
>): TierEffortRow[] =>
models
.map((model) => {
const effort = storedEffort(paramsByModel?.[model]);
@ -43,56 +49,74 @@ export const tierEffortRows = ({
const listed = effort !== undefined && !supported.includes(effort) ? [...supported, effort] : supported;
return { model, effort, options: Array.from(new Set(listed)) };
})
.filter(({ options }) => options.length > 0);
.filter(({ model, options }) => options.length > 0 || fastModeByModel?.[model] === true);
const TierModelEffortRows: React.FC<TierModelEffortRowsProps> = ({
tierLabel,
models,
effortOptionsByModel,
paramsByModel,
onEffortChange,
}) => {
const rows = tierEffortRows({ models, effortOptionsByModel, paramsByModel });
const TierModelEffortRows: React.FC<TierModelEffortRowsProps> = (props) => {
const { tierLabel, paramsByModel, onEffortChange, fastModeByModel, onFastModeChange } = props;
const rows = tierEffortRows(props);
if (rows.length === 0) return null;
return (
<div className="mt-2 space-y-1">
<div className="flex items-center gap-1">
<span className="text-xs font-medium text-muted-foreground">Reasoning effort</span>
<SimpleTooltip
content={`Sent as reasoning_effort on requests this tier routes to the model, overriding the caller's value. Default leaves the request untouched.`}
>
<Info className="size-3 text-muted-foreground/70" />
</SimpleTooltip>
</div>
{rows.map(({ model, effort, options }) => (
<div key={model} className="flex items-center justify-between gap-2">
<span className="truncate text-xs">{model}</span>
<Select
items={[
{ value: PROVIDER_DEFAULT, label: "Default" },
...options.map((option) => ({ value: option, label: option })),
]}
value={effort ?? PROVIDER_DEFAULT}
onValueChange={(selected: string | null) =>
selected !== null && onEffortChange(model, selected === PROVIDER_DEFAULT ? undefined : selected)
}
{rows.some(({ options }) => options.length > 0) && (
<div className="flex items-center gap-1">
<span className="text-xs font-medium text-muted-foreground">Reasoning effort</span>
<SimpleTooltip
content={`Sent as reasoning_effort on requests this tier routes to the model, overriding the caller's value. Default leaves the request untouched.`}
>
<SelectTrigger
size="sm"
className="w-36"
aria-label={`Reasoning effort for ${model} in the ${tierLabel} tier`}
>
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value={PROVIDER_DEFAULT}>Default</SelectItem>
{options.map((option) => (
<SelectItem key={option} value={option}>
{option}
</SelectItem>
))}
</SelectContent>
</Select>
<Info className="size-3 text-muted-foreground/70" />
</SimpleTooltip>
</div>
)}
{rows.map(({ model, effort, options }) => (
<div key={model} className="flex flex-wrap items-center justify-between gap-2">
<span className="min-w-0 flex-1 basis-32 truncate text-xs" title={model}>
{model}
</span>
<div className="flex flex-wrap items-center gap-3">
{options.length > 0 && (
<Select
items={[
{ value: PROVIDER_DEFAULT, label: "Default" },
...options.map((option) => ({ value: option, label: option })),
]}
value={effort ?? PROVIDER_DEFAULT}
onValueChange={(selected: string | null) =>
selected !== null && onEffortChange(model, selected === PROVIDER_DEFAULT ? undefined : selected)
}
>
<SelectTrigger
size="sm"
className="w-36"
aria-label={`Reasoning effort for ${model} in the ${tierLabel} tier`}
>
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value={PROVIDER_DEFAULT}>Default</SelectItem>
{options.map((option) => (
<SelectItem key={option} value={option}>
{option}
</SelectItem>
))}
</SelectContent>
</Select>
)}
{fastModeByModel?.[model] === true && (
<SimpleTooltip content="Fast mode has higher pricing and requires an eligible provider account. Off removes this tier's speed override and inherits the request or provider default">
<label
className="flex items-center gap-2 text-xs"
aria-label={`Fast mode for ${model} in the ${tierLabel} tier`}
>
<Switch
size="sm"
checked={paramsByModel?.[model]?.speed === "fast"}
onCheckedChange={(enabled) => onFastModeChange(model, enabled)}
/>
Fast mode
</label>
</SimpleTooltip>
)}
</div>
</div>
))}
</div>

View file

@ -48,6 +48,19 @@ const baseParams: BuildComplexityRouterConfigParams = {
};
describe("buildComplexityRouterConfig", () => {
it("carries Fast and reasoning overrides independently into a new router payload", () => {
const params = { speed: "fast", reasoning_effort: "high", max_tokens: 1024 };
const config = buildComplexityRouterConfig({
...baseParams,
tiers: { ...tiers, COMPLEX: ["primary"], REASONING: ["secondary"] },
tierModelParams: { COMPLEX: { primary: params }, REASONING: { secondary: { speed: "fast" } } },
});
expect(config.tier_model_configs).toEqual({
COMPLEX: [{ model_name: "primary", litellm_params: params }],
REASONING: [{ model_name: "secondary", litellm_params: { speed: "fast" } }],
});
});
it("emits tiers, classifier_type, and escalation_keywords when nothing else is configured", () => {
const config = buildComplexityRouterConfig(baseParams);
const expected = {

View file

@ -7,6 +7,7 @@ import {
serializeTierModelConfigs,
tierRowLabel,
setTierModelReasoningEffort,
setTierModelParam,
} from "./complexity_router_tiers";
import { resolveComplexityDefaultModel } from "./tier_rows";
@ -218,6 +219,28 @@ describe("setTierModelReasoningEffort", () => {
});
});
describe("setTierModelParam", () => {
it.each(["reasoning_effort", "speed"] as const)("clears only %s and preserves the input", (key) => {
const params = { reasoning_effort: "high", speed: "fast", max_tokens: 512 };
const current = { COMPLEX: { primary: params, secondary: { speed: "fast" } }, REASONING: { primary: params } };
const cleared = setTierModelParam(current, "COMPLEX", "primary", [key, undefined]);
expect(cleared).toEqual({
...current,
COMPLEX: {
...current.COMPLEX,
primary: key === "speed" ? { reasoning_effort: "high", max_tokens: 512 } : { speed: "fast", max_tokens: 512 },
},
});
expect(current.COMPLEX.primary).toEqual({ reasoning_effort: "high", speed: "fast", max_tokens: 512 });
});
it("removes empty records when the only override is Fast", () => {
const enabled = setTierModelParam(undefined, "COMPLEX", "primary", ["speed", "fast"]);
expect(enabled).toEqual({ COMPLEX: { primary: { speed: "fast" } } });
expect(setTierModelParam(enabled, "COMPLEX", "primary", ["speed", undefined])).toBeUndefined();
});
});
describe("pruneTierModelParams", () => {
it("drops params for models deselected from the tier", () => {
expect(

View file

@ -114,14 +114,16 @@ export const serializeTierModelConfigs = (
return serialized.length > 0 ? Object.fromEntries(serialized) : undefined;
};
export const setTierModelReasoningEffort = (
export type TierModelParamChange = ["reasoning_effort", ReasoningEffort | undefined] | ["speed", "fast" | undefined];
export const setTierModelParam = (
current: TierModelParamsByTier | undefined,
tier: string,
model: string,
effort: ReasoningEffort | undefined,
[key, value]: TierModelParamChange,
): TierModelParamsByTier | undefined => {
const { reasoning_effort: _dropped, ...rest } = current?.[tier]?.[model] ?? {};
const params = effort === undefined ? rest : { ...rest, reasoning_effort: effort };
const { [key]: _dropped, ...rest } = current?.[tier]?.[model] ?? {};
const params = value === undefined ? rest : { ...rest, [key]: value };
const byModel = Object.fromEntries(
Object.entries({ ...current?.[tier], [model]: params }).filter(([, value]) => Object.keys(value).length > 0),
);
@ -131,6 +133,13 @@ export const setTierModelReasoningEffort = (
return Object.keys(next).length > 0 ? next : undefined;
};
export const setTierModelReasoningEffort = (
current: TierModelParamsByTier | undefined,
tier: string,
model: string,
effort: ReasoningEffort | undefined,
): TierModelParamsByTier | undefined => setTierModelParam(current, tier, model, ["reasoning_effort", effort]);
export const pruneTierModelParams = (
current: TierModelParamsByTier | undefined,
tier: string,

View file

@ -48,6 +48,28 @@ describe("CodeSnippets", () => {
expect(code).toContain("print(response.data[0].embedding)");
});
describe("custom headers", () => {
const customHeaders = { "anthropic-beta": "context-1m-2025-08-07", "x-request-source": "playground" };
it("passes configured headers as default_headers on the OpenAI client", () => {
const code = generateCodeSnippet({ ...baseParams, endpointType: EndpointType.CHAT, customHeaders });
expect(code).toContain('base_url="http://localhost:4000",\n\tdefault_headers={');
expect(code).toContain('"anthropic-beta": "context-1m-2025-08-07"');
expect(code).toContain('"x-request-source": "playground"');
});
it("passes configured headers as default_headers on the Azure client", () => {
const code = generateCodeSnippet({ ...baseParams, selectedSdk: "azure", customHeaders });
expect(code).toContain('api_version="2024-02-01",\n\tdefault_headers={');
expect(code).toContain('"anthropic-beta": "context-1m-2025-08-07"');
});
it("omits default_headers when no custom headers are configured", () => {
expect(generateCodeSnippet(baseParams)).not.toContain("default_headers");
expect(generateCodeSnippet({ ...baseParams, customHeaders: {} })).not.toContain("default_headers");
});
});
describe("base URL selection", () => {
it("should use LITELLM_UI_API_DOC_BASE_URL when provided", () => {
const customBaseUrl = "https://custom-doc.example.com";

View file

@ -1,6 +1,7 @@
import { MessageType } from "./types";
import { EndpointType } from "./mode_endpoint_mapping";
import { MCPServer } from "@/components/mcp_tools/types";
import type { CustomHeaders } from "@/components/llm_calls/request_headers";
interface CodeGenMetadata {
tags?: string[];
@ -30,6 +31,7 @@ interface GenerateCodeParams {
PROXY_BASE_URL?: string;
LITELLM_UI_API_DOC_BASE_URL?: string | null;
};
customHeaders?: CustomHeaders;
}
export const generateCodeSnippet = (params: GenerateCodeParams): string => {
@ -48,6 +50,7 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
selectedModel,
selectedSdk,
proxySettings,
customHeaders,
} = params;
const effectiveApiKey = apiKeySource === "session" ? accessToken : apiKey;
@ -76,6 +79,11 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
const modelNameForCode = selectedModel || "your-model-name";
const defaultHeadersCode =
customHeaders && Object.keys(customHeaders).length > 0
? `,\n\tdefault_headers=${JSON.stringify(customHeaders, null, 2).replace(/\n/g, "\n\t")}`
: "";
const clientInitialization =
selectedSdk === "azure"
? `import openai
@ -83,13 +91,13 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => {
client = openai.AzureOpenAI(
api_key="${effectiveApiKey || "YOUR_LITELLM_API_KEY"}",
azure_endpoint="${apiBase}",
api_version="2024-02-01"
api_version="2024-02-01"${defaultHeadersCode}
)`
: `import openai
client = openai.OpenAI(
api_key="${effectiveApiKey || "YOUR_LITELLM_API_KEY"}",
base_url="${apiBase}"
base_url="${apiBase}"${defaultHeadersCode}
)`;
let endpointSpecificCode;

View file

@ -1,4 +1,5 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import openai from "openai";
import { makeOpenAIChatCompletionRequest } from "./chat_completion";
import type { TokenUsage } from "../chat_ui/ResponseMetrics";
@ -615,3 +616,47 @@ describe("chat_completion response cache", () => {
expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }));
});
});
describe("chat_completion custom headers", () => {
afterEach(() => {
vi.clearAllMocks();
});
it("sends custom headers alongside the tags header on the OpenAI client", async () => {
mockCreate.mockReturnValueOnce(nonStreamingResponse({ choices: [{ message: { content: "Hi" } }] }));
await makeOpenAIChatCompletionRequest(
[{ role: "user", content: "Hello" }],
vi.fn(),
"gpt-4",
"test-token",
["team-a"],
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
false,
{ "anthropic-beta": "context-1m-2025-08-07", "x-litellm-tags": "overridden" },
);
expect(vi.mocked(openai.OpenAI).mock.calls[0][0]).toMatchObject({
defaultHeaders: { "anthropic-beta": "context-1m-2025-08-07", "x-litellm-tags": "overridden" },
});
});
});

View file

@ -6,6 +6,7 @@ import { getProxyBaseUrl } from "@/components/networking";
import { MCPServer, MCPToolset, type MCPEvent } from "@/components/mcp_tools/types";
import { extractPromptCacheTokens } from "@/utils/promptCacheUsage";
import { parseUsageCost } from "./usage_cost";
import { buildPlaygroundHeaders, type CustomHeaders } from "./request_headers";
const completionAsSingleChunk = (completion: ChatCompletion): ChatCompletionChunk =>
({
@ -50,6 +51,7 @@ export async function makeOpenAIChatCompletionRequest(
mockTestFallbacks?: boolean,
mcpToolsets?: MCPToolset[],
streamingEnabled: boolean = true,
customHeaders?: CustomHeaders,
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
@ -57,11 +59,7 @@ export async function makeOpenAIChatCompletionRequest(
console.log = function () {};
}
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
// Prepare headers with tags and trace ID
const headers: Record<string, string> = {};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const headers = buildPlaygroundHeaders(tags, customHeaders);
const client = new openai.OpenAI({
apiKey: accessToken,

View file

@ -52,6 +52,23 @@ describe("fetchAvailableModels", () => {
]);
});
it("carries only explicitly supported Fast capabilities, not accepted speed parameters", async () => {
modelHubCallMock.mockResolvedValue({
data: [
{ model_group: "fast", supports_fast_mode: true },
{ model_group: "blocked", supports_fast_mode: false },
{ model_group: "missing", supports_speed: true },
{ model_group: "unknown", supports_fast_mode: null },
],
});
expect(await fetchAvailableModels("token")).toEqual([
{ model_group: "blocked" },
{ model_group: "fast", supports_fast_mode: true },
{ model_group: "missing" },
{ model_group: "unknown" },
]);
});
it("preserves absent, unknown, empty, and explicit effort capability states", async () => {
modelHubCallMock.mockResolvedValue({
data: [

View file

@ -7,6 +7,7 @@ export interface ModelGroup {
model_group: string;
mode?: string;
supports_reasoning?: boolean;
supports_fast_mode?: boolean;
supported_reasoning_efforts?: string[] | null;
}
@ -16,6 +17,7 @@ interface AvailableModel {
id?: string | null;
mode?: string | null;
supports_reasoning?: boolean | null;
supports_fast_mode?: boolean | null;
supported_reasoning_efforts?: string[] | null;
}
@ -25,6 +27,7 @@ const toModelGroup = (item: AvailableModel): ModelGroup => {
model_group: groupName,
...(item.mode && { mode: item.mode }),
...(item.supports_reasoning === true && { supports_reasoning: true }),
...(item.supports_fast_mode === true && { supports_fast_mode: true }),
...(item.supported_reasoning_efforts !== undefined && {
supported_reasoning_efforts: item.supported_reasoning_efforts,
}),

View file

@ -0,0 +1,60 @@
import { describe, expect, it } from "vitest";
import {
buildPlaygroundHeaders,
customHeadersFromPairs,
parseStoredHeaderPairs,
withRequiredHeaders,
} from "./request_headers";
describe("customHeadersFromPairs", () => {
it("trims header names and drops rows without a name", () => {
expect(
customHeadersFromPairs([
[" anthropic-beta ", "context-1m-2025-08-07"],
["", "orphan value"],
[" ", "whitespace name"],
["x-empty", ""],
]),
).toEqual({ "anthropic-beta": "context-1m-2025-08-07", "x-empty": "" });
});
});
describe("parseStoredHeaderPairs", () => {
it("round-trips pairs persisted as JSON", () => {
const pairs = [["anthropic-beta", "context-1m-2025-08-07"]] as const;
expect(parseStoredHeaderPairs(JSON.stringify(pairs))).toEqual(pairs);
});
it("returns no pairs for missing, malformed, or wrongly shaped storage", () => {
expect(parseStoredHeaderPairs(null)).toEqual([]);
expect(parseStoredHeaderPairs("not json")).toEqual([]);
expect(parseStoredHeaderPairs(JSON.stringify({ "anthropic-beta": "x" }))).toEqual([]);
expect(parseStoredHeaderPairs(JSON.stringify([["ok", "pair"], ["one"], [1, 2], "str"]))).toEqual([["ok", "pair"]]);
});
});
describe("buildPlaygroundHeaders", () => {
it("joins tags into x-litellm-tags and lets custom headers override it", () => {
expect(buildPlaygroundHeaders(["a", "b"], { "x-custom": "1" })).toEqual({
"x-litellm-tags": "a,b",
"x-custom": "1",
});
expect(buildPlaygroundHeaders(["a"], { "x-litellm-tags": "b" })).toEqual({ "x-litellm-tags": "b" });
});
it("omits x-litellm-tags when there are no tags", () => {
expect(buildPlaygroundHeaders([], { "x-custom": "1" })).toEqual({ "x-custom": "1" });
expect(buildPlaygroundHeaders(undefined, undefined)).toEqual({});
});
});
describe("withRequiredHeaders", () => {
it("keeps required headers regardless of custom header name casing", () => {
expect(
withRequiredHeaders(
{ authorization: "Bearer stolen", "content-type": "text/plain", "x-custom": "1" },
{ Authorization: "Bearer real", "Content-Type": "application/json" },
),
).toEqual({ Authorization: "Bearer real", "Content-Type": "application/json", "x-custom": "1" });
});
});

View file

@ -0,0 +1,38 @@
import type { KeyValuePair } from "@/components/key_value_input";
export type CustomHeaders = Readonly<Record<string, string>>;
export const customHeadersFromPairs = (pairs: readonly KeyValuePair[]): CustomHeaders =>
Object.fromEntries(pairs.map(([name, value]) => [name.trim(), value]).filter(([name]) => name !== ""));
const isHeaderPair = (entry: unknown): entry is KeyValuePair =>
Array.isArray(entry) && entry.length === 2 && entry.every((part) => typeof part === "string");
export const parseStoredHeaderPairs = (raw: string | null): readonly KeyValuePair[] => {
if (!raw) return [];
try {
const parsed: unknown = JSON.parse(raw);
return Array.isArray(parsed) ? parsed.filter(isHeaderPair) : [];
} catch {
return [];
}
};
export const buildPlaygroundHeaders = (
tags?: readonly string[],
customHeaders?: CustomHeaders,
): Record<string, string> => ({
...(tags && tags.length > 0 ? { "x-litellm-tags": tags.join(",") } : {}),
...customHeaders,
});
export const withRequiredHeaders = (
headers: Readonly<Record<string, string>>,
required: Readonly<Record<string, string>>,
): Record<string, string> => {
const reserved = new Set(Object.keys(required).map((name) => name.toLowerCase()));
return {
...Object.fromEntries(Object.entries(headers).filter(([name]) => !reserved.has(name.toLowerCase()))),
...required,
};
};

View file

@ -1,4 +1,5 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import openai from "openai";
import { makeOpenAIResponsesRequest } from "./responses_api";
import { MessageType } from "../chat_ui/types";
import type { TokenUsage } from "../chat_ui/ResponseMetrics";
@ -611,3 +612,46 @@ describe("responses_api response cache", () => {
expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }), "");
});
});
describe("responses_api custom headers", () => {
afterEach(() => {
vi.clearAllMocks();
});
it("sends custom headers alongside the tags header on the OpenAI client", async () => {
mockResponsesCreate.mockReturnValueOnce(nonStreamingResponse({ id: "resp_1", output: [] }));
await makeOpenAIResponsesRequest(
[{ role: "user", content: "Hello" }],
vi.fn(),
"gpt-4",
"test-token",
["team-a"],
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
undefined,
false,
undefined,
{ "anthropic-beta": "context-1m-2025-08-07" },
);
expect(vi.mocked(openai.OpenAI).mock.calls[0][0]).toMatchObject({
defaultHeaders: { "x-litellm-tags": "team-a", "anthropic-beta": "context-1m-2025-08-07" },
});
});
});

View file

@ -5,6 +5,7 @@ import { getProxyBaseUrl } from "@/components/networking";
import { toast } from "@/lib/toast";
import { extractPromptCacheTokens } from "@/utils/promptCacheUsage";
import { parseUsageCost } from "./usage_cost";
import { buildPlaygroundHeaders, type CustomHeaders } from "./request_headers";
import type { MCPEvent } from "@/components/mcp_tools/types";
import { MCPServer, MCPToolset } from "@/components/mcp_tools/types";
import {
@ -85,6 +86,7 @@ export async function makeOpenAIResponsesRequest(
mcpToolsets?: MCPToolset[],
streamingEnabled: boolean = true,
onTotalLatency?: (latency: number) => void,
customHeaders?: CustomHeaders,
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -101,11 +103,7 @@ export async function makeOpenAIResponsesRequest(
}
const proxyBaseUrl = customBaseUrl || getProxyBaseUrl();
// Prepare headers with tags and trace ID
const headers: Record<string, string> = {};
if (tags && tags.length > 0) {
headers["x-litellm-tags"] = tags.join(",");
}
const headers = buildPlaygroundHeaders(tags, customHeaders);
const client = new openai.OpenAI({
apiKey: accessToken,

View file

@ -27297,6 +27297,8 @@ export interface components {
jwt_claim_name: string;
/** Jwt Claim Value */
jwt_claim_value: string;
/** Jwt Issuer */
jwt_issuer?: string | null;
/** Key */
key: string;
};
@ -28910,6 +28912,8 @@ export interface components {
jwt_claim_name: string;
/** Jwt Claim Value */
jwt_claim_value: string;
/** Jwt Issuer */
jwt_issuer?: string | null;
/**
* Updated At
* Format: date-time
@ -32704,6 +32708,11 @@ export interface components {
supported_openai_params: string[] | null;
/** Supported Reasoning Efforts */
supported_reasoning_efforts?: string[] | null;
/**
* Supports Fast Mode
* @default false
*/
supports_fast_mode: boolean;
/**
* Supports Function Calling
* @default false
@ -38643,6 +38652,8 @@ export interface components {
id: string;
/** Is Active */
is_active?: boolean | null;
/** Jwt Issuer */
jwt_issuer?: string | null;
/** Key */
key?: string | null;
};

8
uv.lock generated
View file

@ -10,7 +10,7 @@ resolution-markers = [
]
[options]
exclude-newer = "2026-09-12T03:51:49.261499386Z"
exclude-newer = "2026-09-12T22:48:38.53978Z"
exclude-newer-span = "P3D"
[manifest]
@ -4457,7 +4457,7 @@ wheels = [
[[package]]
name = "litellm"
version = "1.102.0"
version = "1.103.0"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },
@ -4874,12 +4874,12 @@ proxy-dev = [
[[package]]
name = "litellm-enterprise"
version = "0.1.67"
version = "0.1.68"
source = { editable = "enterprise" }
[[package]]
name = "litellm-proxy-extras"
version = "0.4.97"
version = "0.4.98"
source = { editable = "litellm-proxy-extras" }
[[package]]