From ad36c603b669fb403999bc6ee31d940f0126ac42 Mon Sep 17 00:00:00 2001 From: mateo-berri Date: Wed, 12 Aug 2026 17:51:19 -0700 Subject: [PATCH 01/17] feat(model_prices): add meta/muse-spark-1.2 and its contributor tier --- ...odel_prices_and_context_window_backup.json | 72 +++++++++++++++ model_prices_and_context_window.json | 72 +++++++++++++++ .../test_muse_spark_1_2_model_metadata.py | 88 +++++++++++++++++++ 3 files changed, 232 insertions(+) create mode 100644 tests/test_litellm/test_muse_spark_1_2_model_metadata.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b12e4a9fea3..77645f7ba2d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27500,6 +27500,78 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "meta/muse-spark-1.2": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "meta/muse-spark-1.2-contributor": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1e-07, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b12e4a9fea3..77645f7ba2d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27500,6 +27500,78 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "meta/muse-spark-1.2": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "meta/muse-spark-1.2-contributor": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 1e-07, + "litellm_provider": "meta", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/messages" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py new file mode 100644 index 00000000000..81ef96f9b7a --- /dev/null +++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py @@ -0,0 +1,88 @@ +import json +from pathlib import Path + +import pytest + +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + +MUSE_SPARK_STANDARD = "meta/muse-spark-1.2" +MUSE_SPARK_CONTRIBUTOR = "meta/muse-spark-1.2-contributor" + +PRICING = ( + (MUSE_SPARK_STANDARD, 1.25e-06, 1.5e-07, 4.25e-06), + (MUSE_SPARK_CONTRIBUTOR, 1e-07, 2e-09, 2e-07), +) + + +def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> dict: + with open(Path(__file__).parents[2] / filename) as f: + return json.load(f) + + +@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) +def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): + info = _load_cost_map().get(model) + assert info is not None, f"{model} not found in model_prices_and_context_window.json" + + assert info["litellm_provider"] == "meta" + assert info["mode"] == "chat" + + assert info["input_cost_per_token"] == input_cost + assert info["output_cost_per_token"] == output_cost + assert info["cache_read_input_token_cost"] == cached_cost + + assert info["max_input_tokens"] == 1048576 + assert info["max_output_tokens"] == 131072 + assert info["max_tokens"] == 131072 + + assert info["supports_function_calling"] is True + assert info["supports_parallel_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_response_schema"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_pdf_input"] is True + assert info["supports_web_search"] is True + assert info["supports_minimal_reasoning_effort"] is True + assert info["supports_xhigh_reasoning_effort"] is True + + assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + assert info["supported_modalities"] == ["text", "image", "video"] + assert info["supported_output_modalities"] == ["text"] + + +@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) +def test_muse_spark_1_2_cost_per_token(model: str, input_cost: float, cached_cost: float, output_cost: float): + prompt_cost, completion_cost = cost_per_token(model=model, prompt_tokens=1000, completion_tokens=500) + + assert prompt_cost == pytest.approx(1000 * input_cost) + assert completion_cost == pytest.approx(500 * output_cost) + + +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_2_routes_to_meta_model_api(model: str): + routed_model, provider, _, api_base = get_llm_provider(model=model, api_key="sk-test") + + assert routed_model == model.split("/", 1)[1] + assert provider == "meta" + assert api_base == "https://api.meta.ai/v1" + + +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_2_backup_matches_main(model: str): + """Ensure the bundled model cost map stays in sync with the canonical file.""" + main_cost = _load_cost_map() + backup_cost = _load_cost_map("litellm/model_prices_and_context_window_backup.json") + + assert backup_cost.get(model) == main_cost.get(model), f"{model} differs between main and backup model cost maps" + + +def test_muse_spark_contributor_tier_is_cheaper_than_standard(): + cost_map = _load_cost_map() + standard = cost_map[MUSE_SPARK_STANDARD] + contributor = cost_map[MUSE_SPARK_CONTRIBUTOR] + + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + assert contributor[field] < standard[field], f"contributor {field} should undercut the standard tier" From 1369b3e6f3c25446ecb261d5964d5b227fecbcb5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:14:19 -0700 Subject: [PATCH 02/17] test(model_prices): pin muse spark cost assertions to the bundled map cost_per_token reads litellm.model_cost, which CI loads from the model_prices_and_context_window.json on main, so the new entries were missing until merge. Use the same local_model_cost_map fixture the sonnet 5 metadata test uses to force the branch's own backup map. --- .../test_muse_spark_1_2_model_metadata.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py index 81ef96f9b7a..dcce92b0c17 100644 --- a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py @@ -3,6 +3,7 @@ from pathlib import Path import pytest +import litellm from litellm.cost_calculator import cost_per_token from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -20,6 +21,21 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di return json.load(f) +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled backup cost map so assertions don't depend on the + network-fetched ``main`` copy (which lags this branch until merge).""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): info = _load_cost_map().get(model) @@ -54,7 +70,9 @@ def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: f @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) -def test_muse_spark_1_2_cost_per_token(model: str, input_cost: float, cached_cost: float, output_cost: float): +def test_muse_spark_1_2_cost_per_token( + local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float +): prompt_cost, completion_cost = cost_per_token(model=model, prompt_tokens=1000, completion_tokens=500) assert prompt_cost == pytest.approx(1000 * input_cost) From 39a5ede2d9884ebec88de05b4088194394486652 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:44:20 -0700 Subject: [PATCH 03/17] fix(model_prices): bill muse spark web search grounding per query --- .../model_prices_and_context_window_backup.json | 15 +++++++++++++++ model_prices_and_context_window.json | 15 +++++++++++++++ .../test_muse_spark_1_2_model_metadata.py | 15 +++++++++++++++ 3 files changed, 45 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 77645f7ba2d..d3f9c52bd08 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27473,6 +27473,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", @@ -27509,6 +27514,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", @@ -27545,6 +27555,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 77645f7ba2d..d3f9c52bd08 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27473,6 +27473,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", @@ -27509,6 +27514,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 4.25e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", @@ -27545,6 +27555,11 @@ "max_tokens": 131072, "mode": "chat", "output_cost_per_token": 2e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.0025, + "search_context_size_low": 0.0025, + "search_context_size_medium": 0.0025 + }, "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", diff --git a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py index dcce92b0c17..20aa4b11dcd 100644 --- a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py @@ -6,9 +6,11 @@ import pytest import litellm from litellm.cost_calculator import cost_per_token from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import StandardBuiltInToolCostTracking MUSE_SPARK_STANDARD = "meta/muse-spark-1.2" MUSE_SPARK_CONTRIBUTOR = "meta/muse-spark-1.2-contributor" +WEB_SEARCH_COST_PER_QUERY = 0.0025 PRICING = ( (MUSE_SPARK_STANDARD, 1.25e-06, 1.5e-07, 4.25e-06), @@ -68,6 +70,12 @@ def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: f assert info["supported_modalities"] == ["text", "image", "video"] assert info["supported_output_modalities"] == ["text"] + assert info["search_context_cost_per_query"] == { + "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, + "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, + "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, + } + @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) def test_muse_spark_1_2_cost_per_token( @@ -88,6 +96,13 @@ def test_muse_spark_1_2_routes_to_meta_model_api(model: str): assert api_base == "https://api.meta.ai/v1" +@pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) +def test_muse_spark_1_2_web_search_cost_per_query(local_model_cost_map, model: str): + info = litellm.get_model_info(model=model) + + assert StandardBuiltInToolCostTracking.get_cost_for_web_search(model_info=info) == WEB_SEARCH_COST_PER_QUERY + + @pytest.mark.parametrize("model", (MUSE_SPARK_STANDARD, MUSE_SPARK_CONTRIBUTOR)) def test_muse_spark_1_2_backup_matches_main(model: str): """Ensure the bundled model cost map stays in sync with the canonical file.""" From 4bc27f1664d1c742d476c3c751153610d651fa3e Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 13 Aug 2026 16:56:47 -0700 Subject: [PATCH 04/17] fix(auth): carry team grants in lite login session tokens (#36826) CLI session tokens minted by /sso/cli/poll set team_id and team_alias but never team_models or team_model_aliases, so the token carried a team with none of that team's grants. /v1/models bails out to "unrestricted" when both key_models and team_models are empty and listed the whole proxy, and team model aliases never resolved because both can_team_access_model and the pre-call rewrite read team_model_aliases off the token. The team data was not close at hand: _fetch_cli_sso_team_details projected full team rows down to team_id and team_alias before they reached the mint. Widen that projection to include the team's models and its joined alias table, and populate both fields at mint time. Also stop writing the user's personal allowlist into the key models slot when a team is bound, matching virtual-key semantics where a team-bound credential is governed by the team grant. Because an empty team grant is itself a real value meaning unrestricted, a team whose grants cannot be resolved must not be minted as empty: that is the same "unrestricted" bail-out this fix exists to close. The poll now refuses to mint when the selected team has no complete cached detail. That refusal is only safe because a login can no longer be pinned to a team whose grants will never resolve. Deleting an organization drops its team rows but leaves the memberships behind, so the login now offers only teams whose rows still exist, and a lookup that fails outright fails the login rather than caching a session that silently drops every team. --- litellm/proxy/auth/auth_checks.py | 8 +- litellm/proxy/management_endpoints/ui_sso.py | 132 ++++-- .../proxy/auth/test_auth_checks.py | 49 +++ .../proxy/management_endpoints/test_ui_sso.py | 382 +++++++++++++++++- 4 files changed, 541 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b1e444e55d6..17ebb4446e2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2556,6 +2556,8 @@ class ExperimentalUIJWTToken: user_info: LiteLLM_UserTable, team_id: str | None = None, team_alias: str | None = None, + team_models: Sequence[str] | None = None, + team_model_aliases: Mapping[str, str] | None = None, max_budget: float | None = None, ) -> str: """ @@ -2568,6 +2570,8 @@ class ExperimentalUIJWTToken: user_info: User information from the database team_id: Team ID for the user (optional, uses user's team if available) team_alias: Team alias for the selected team, if available + team_models: Model allowlist granted by the selected team + team_model_aliases: Team model aliases for the selected team Returns: Encrypted JWT token string @@ -2606,7 +2610,9 @@ class ExperimentalUIJWTToken: user_id=user_info.user_id, team_id=_team_id, team_alias=team_alias, - models=user_info.models, + team_models=list(team_models) if team_models is not None else [], + team_model_aliases=dict(team_model_aliases) if team_model_aliases is not None else None, + models=[] if _team_id is not None else user_info.models, max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), is_session_token=True, diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index a2c50590dd5..b87ad8597dc 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -21,6 +21,7 @@ from copy import deepcopy from html import escape from typing import ( TYPE_CHECKING, + Annotated, Any, Final, Literal, @@ -40,6 +41,7 @@ if TYPE_CHECKING: import jwt from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status from fastapi.responses import RedirectResponse +from pydantic import BaseModel, BeforeValidator, ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -185,6 +187,7 @@ class _PrismaTableActions(Protocol[_DbRecordT]): async def find_many( self, where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, ) -> Sequence[_DbRecordT]: ... async def update( @@ -241,6 +244,45 @@ def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDe return repo.table +_MODEL_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, str]) + + +def _decode_model_aliases(value: object) -> object: + """``/team/new`` stores team model aliases as a JSON-encoded string in the Json column.""" + if not isinstance(value, str): + return value + try: + return _MODEL_ALIASES_ADAPTER.validate_json(value) + except ValidationError: + return None + + +class _TeamModelAliasTable(BaseModel): + model_config = ConfigDict(protected_namespaces=()) + + model_aliases: Annotated[Mapping[str, str] | None, BeforeValidator(_decode_model_aliases)] = None + + +class _TeamRowGrants(BaseModel): + team_id: str + team_alias: str | None = None + models: tuple[str, ...] = () + litellm_model_table: _TeamModelAliasTable | None = None + + +class _CliSsoTeamDetail(BaseModel): + """The per-team snapshot cached in the CLI SSO flow and echoed to the CLI on poll.""" + + team_id: str | None = None + team_alias: str | None = None + team_models: tuple[str, ...] + team_model_aliases: Mapping[str, str] | None = None + + +_CLI_SSO_TEAM_DETAILS_ADAPTER: Final = TypeAdapter(tuple[_CliSsoTeamDetail, ...]) +_TEAMLESS_CLI_SSO_TEAM_DETAIL: Final = _CliSsoTeamDetail(team_models=()) + + class _CustomSsoCall(Protocol): async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... @@ -2147,27 +2189,55 @@ async def _build_cli_sso_user_defined_values( ) +def _cli_sso_team_detail(team_row: Mapping[str, object]) -> _CliSsoTeamDetail: + team: Final = _TeamRowGrants.model_validate(team_row) + alias_table: Final = team.litellm_model_table + return _CliSsoTeamDetail( + team_id=team.team_id, + team_alias=team.team_alias, + team_models=team.models, + team_model_aliases=alias_table.model_aliases if alias_table is not None else None, + ) + + async def _fetch_cli_sso_team_details( prisma_client: PrismaClient, teams: Sequence[str], -) -> list[dict[str, object]]: - team_details: Final[list[dict[str, object]]] = [] +) -> tuple[_CliSsoTeamDetail, ...] | None: + """``None`` means the lookup itself failed, which is not the same as the user having no teams.""" + if not teams: + return () try: - if teams: - prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many( - where={"team_id": {"in": teams}} - ) - for team_row in prisma_teams: - team_dict = team_row.model_dump() - team_details.append( - { - "team_id": team_dict.get("team_id"), - "team_alias": team_dict.get("team_alias"), - } - ) + prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many( + where={"team_id": {"in": teams}}, + include={"litellm_model_table": True}, + ) except Exception as e: verbose_proxy_logger.error("Error fetching team details for CLI SSO session: %s", e) - return team_details + return None + return tuple(_cli_sso_team_detail(team_row.model_dump()) for team_row in prisma_teams) + + +def _cli_sso_session_teams(team_details: Sequence[_CliSsoTeamDetail]) -> list[str]: + """The teams a login may bind to: only those whose row still exists. + + A team deleted out from under a membership, which is what deleting an organization + leaves behind, can never resolve its grants, so offering it would refuse every + future login for that user with nothing they could do to recover. + """ + return [detail.team_id for detail in team_details if detail.team_id is not None] + + +def _selected_cli_sso_team_detail(team_details: object, team_id: str | None) -> _CliSsoTeamDetail | None: + """``None`` means the team's grants are unknown. An empty grant is a real value meaning unrestricted, + so an unknown one must not be minted as empty.""" + if team_id is None: + return _TEAMLESS_CLI_SSO_TEAM_DETAIL + try: + details: Final = _CLI_SSO_TEAM_DETAILS_ADAPTER.validate_python(team_details) + except ValidationError: + return None + return next((detail for detail in details if detail.team_id == team_id), None) async def _complete_cli_sso_callback_session( @@ -2210,6 +2280,12 @@ async def _complete_cli_sso_callback_session( teams = user_info.teams if isinstance(user_info.teams, list) else [] team_details: Final = await _fetch_cli_sso_team_details(prisma_client=prisma_client, teams=teams) + if team_details is None: + raise HTTPException( + status_code=500, + detail="Could not resolve team model grants for this login. Please try again", + ) + resolved_teams: Final = _cli_sso_session_teams(team_details) attribution_metadata: Final = build_cli_sso_attribution_metadata(result=result) if attribution_metadata: await _persist_cli_sso_user_metadata( @@ -2223,8 +2299,8 @@ async def _complete_cli_sso_callback_session( "user_role": user_info.user_role, "models": user_info.models if hasattr(user_info, "models") else [], "user_email": user_email, - "teams": teams, - "team_details": team_details, + "teams": resolved_teams, + "team_details": [detail.model_dump() for detail in team_details], "attribution_metadata": attribution_metadata, } flow["sso_complete"] = True @@ -2233,7 +2309,10 @@ async def _complete_cli_sso_callback_session( _set_cli_sso_flow(login_id=key, cache=cli_sso_session_cache, flow=flow) verbose_proxy_logger.info( - "Stored CLI SSO session for user: %s, teams: %s, num_teams: %s", user_info.user_id, teams, len(teams) + "Stored CLI SSO session for user: %s, teams: %s, num_teams: %s", + user_info.user_id, + resolved_teams, + len(resolved_teams), ) verify_url: Final = get_custom_url( request_base_url=str(request.base_url), @@ -2401,11 +2480,14 @@ async def cli_poll_key( # If no team_id provided and user has 0 or 1 team, use first team (or None) team_id = user_teams[0] if len(user_teams) > 0 else None - team_alias = None - if team_id and isinstance(user_team_details, list): - team_alias = next( - (team.get("team_alias") for team in user_team_details if team.get("team_id") == team_id), - None, + selected_team: Final = _selected_cli_sso_team_detail( + team_details=user_team_details, + team_id=team_id, + ) + if selected_team is None: + raise HTTPException( + status_code=500, + detail=f"Could not resolve the model grants for team: {team_id}. Please run `lite login` again", ) user_info: Final = LiteLLM_UserTable( @@ -2417,7 +2499,9 @@ async def cli_poll_key( jwt_token: Final = ExperimentalUIJWTToken.get_cli_jwt_auth_token( user_info=user_info, team_id=team_id, - team_alias=team_alias, + team_alias=selected_team.team_alias, + team_models=selected_team.team_models, + team_model_aliases=selected_team.team_model_aliases, max_budget=None, ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 298f8a31b64..28eda6633e8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -155,6 +155,55 @@ def test_get_cli_jwt_auth_token_includes_team_alias(valid_sso_user_defined_value assert token_data["team_alias"] == "test-team" +def test_get_cli_jwt_auth_token_carries_team_grants_not_user_allowlist( + valid_sso_user_defined_values, +): + """A team-bound `lite login` session token must snapshot the team's grants. + + Without team_models the /v1/models bail-out (`not key_models and not team_models`) + treats the session as unrestricted and lists the whole proxy; without + team_model_aliases a team alias never resolves on /chat/completions. The user's + personal allowlist must stay out of the key `models` slot, since a team-bound + credential is governed by the team grant, not by a per-user list. + """ + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + valid_sso_user_defined_values, + team_id="team-123", + team_alias="test-team", + team_models=("claude-sonnet-4-5", "gpt-4.1"), + team_model_aliases={"team-fast": "gpt-4.1-mini"}, + ) + + decrypted_token = decrypt_value_helper( + token, key="ui_hash_key", exception_type="debug" + ) + assert decrypted_token is not None + token_data = json.loads(decrypted_token) + + assert token_data["team_id"] == "team-123" + assert token_data["team_models"] == ["claude-sonnet-4-5", "gpt-4.1"] + assert token_data["team_model_aliases"] == {"team-fast": "gpt-4.1-mini"} + assert valid_sso_user_defined_values.models == ["gpt-3.5-turbo"] + assert token_data["models"] == [] + + +def test_get_cli_jwt_auth_token_keeps_user_allowlist_when_no_team( + valid_sso_user_defined_values, +): + """A session token with no team bound still carries the user's own allowlist.""" + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values) + + decrypted_token = decrypt_value_helper( + token, key="ui_hash_key", exception_type="debug" + ) + assert decrypted_token is not None + token_data = json.loads(decrypted_token) + + assert token_data.get("team_id") is None + assert token_data["models"] == ["gpt-3.5-turbo"] + assert token_data["team_models"] == [] + + def test_get_experimental_ui_login_jwt_auth_token_uses_10_min_expiry( valid_sso_user_defined_values, ): diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 979eb09d7db..da805b864fc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2847,6 +2847,19 @@ class TestCLIKeyRegenerationFlow: "user_code_verified": False, "session_data": None, } + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock( + return_value=[ + MagicMock( + model_dump=lambda team_id=team_id: { + "team_id": team_id, + "team_alias": team_id, + "models": [], + } + ) + for team_id in ("team1", "team2") + ] + ) with ( patch.dict( os.environ, @@ -2859,7 +2872,7 @@ class TestCLIKeyRegenerationFlow: "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", return_value=mock_user_info, ), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), patch( @@ -3156,9 +3169,9 @@ class TestCLIKeyRegenerationFlow: "user_role": "internal_user", "teams": ["team-a", "team-b", "team-c"], "team_details": [ - {"team_id": "team-a", "team_alias": "Team A"}, - {"team_id": "team-b", "team_alias": "Team B"}, - {"team_id": "team-c", "team_alias": "Team C"}, + {"team_id": "team-a", "team_alias": "Team A", "team_models": []}, + {"team_id": "team-b", "team_alias": "Team B", "team_models": []}, + {"team_id": "team-c", "team_alias": "Team C", "team_models": []}, ], "models": ["gpt-4"], "user_email": "test@example.com", @@ -3225,6 +3238,243 @@ class TestCLIKeyRegenerationFlow: # Verify session was deleted after JWT generation mock_cache.delete_cache.assert_called_once() + @pytest.mark.asyncio + async def test_fetch_cli_sso_team_details_projects_team_grants(self): + """The cached team detail must carry the team's model grants. + + The projection used to drop everything except team_id/team_alias, so the + minted CLI token had no team_models and no team_model_aliases to snapshot. + The joined alias table is stored JSON-encoded, so it has to be decoded here + too, otherwise alias lookup at request time is a substring match on a string. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _fetch_cli_sso_team_details, + ) + + team_row = MagicMock() + team_row.model_dump.return_value = { + "team_id": "team-a", + "team_alias": "Team A", + "models": ["claude-sonnet-4-5", "gpt-4.1"], + "litellm_model_table": { + "id": 7, + "model_aliases": json.dumps({"team-fast": "gpt-4.1-mini"}), + "created_by": "admin", + "updated_by": "admin", + }, + } + find_many = AsyncMock(return_value=[team_row]) + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_many = find_many + + details = await _fetch_cli_sso_team_details( + prisma_client=prisma_client, teams=["team-a"] + ) + + assert find_many.await_args.kwargs["include"] == {"litellm_model_table": True} + assert [detail.model_dump() for detail in details] == [ + { + "team_id": "team-a", + "team_alias": "Team A", + "team_models": ("claude-sonnet-4-5", "gpt-4.1"), + "team_model_aliases": {"team-fast": "gpt-4.1-mini"}, + } + ] + + @pytest.mark.asyncio + async def test_fetch_cli_sso_team_details_separates_lookup_failure_from_no_teams(self): + """A failed lookup must not look like a team that resolved to nothing. + + Both used to return [], so a database blip was indistinguishable from a real + answer. The callback needs them apart: a blip has to fail the login, while a + real empty answer means the team rows are genuinely gone. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _fetch_cli_sso_team_details, + ) + + failing_client = MagicMock() + failing_client.db.litellm_teamtable.find_many = AsyncMock( + side_effect=Exception("connection reset") + ) + assert ( + await _fetch_cli_sso_team_details( + prisma_client=failing_client, teams=["team-a"] + ) + is None + ) + + empty_client = MagicMock() + empty_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + assert ( + await _fetch_cli_sso_team_details( + prisma_client=empty_client, teams=["team-a"] + ) + == () + ) + + @pytest.mark.asyncio + async def test_cli_poll_key_mints_jwt_with_selected_team_grants(self): + """The selected team's grants must reach the mint, not just its alias.""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + session_data = { + "user_id": "grants-user", + "user_role": "internal_user", + "teams": ["team-a", "team-b"], + "team_details": [ + { + "team_id": "team-a", + "team_alias": "Team A", + "team_models": ["gpt-4.1"], + "team_model_aliases": {"a-fast": "gpt-4.1-mini"}, + }, + { + "team_id": "team-b", + "team_alias": "Team B", + "team_models": ["claude-sonnet-4-5"], + "team_model_aliases": {"b-fast": "claude-haiku-4-5"}, + }, + ], + "models": ["personal-only"], + "user_email": "grants@example.com", + } + mock_cache = MagicMock(redis_cache=None) + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": session_data, + } + + with ( + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="minted-token", + ) as mock_get_jwt, + ): + result = await cli_poll_key( + key_id="cli-session-grants", + team_id="team-b", + x_litellm_cli_poll_secret="poll-secret", + ) + + assert result["status"] == "ready" + kwargs = mock_get_jwt.call_args.kwargs + assert kwargs["team_id"] == "team-b" + assert kwargs["team_alias"] == "Team B" + assert kwargs["team_models"] == ("claude-sonnet-4-5",) + assert kwargs["team_model_aliases"] == {"b-fast": "claude-haiku-4-5"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "team_details", + [ + pytest.param(None, id="detail_fetch_failed"), + pytest.param( + [{"team_id": "team-other", "team_models": []}], id="selected_team_absent" + ), + pytest.param( + [{"team_id": "team-a", "team_alias": "Team A"}], + id="legacy_detail_without_grants", + ), + ], + ) + async def test_cli_poll_key_refuses_to_mint_when_team_grants_are_unknown( + self, team_details + ): + """An unknown team grant must never be minted as an empty one. + + get_complete_model_list falls through to the whole proxy model list when both + the key allowlist and the team allowlist are empty, and team-bound tokens carry + an empty key allowlist by design. So minting an unresolved team as empty would + hand a team-bound CLI session every model on the proxy. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + mock_cache = MagicMock(redis_cache=None) + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": { + "user_id": "grants-user", + "user_role": "internal_user", + "teams": ["team-a"], + "team_details": team_details, + "models": ["personal-only"], + "user_email": "grants@example.com", + }, + } + + with ( + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="minted-token", + ) as mock_get_jwt, + ): + with pytest.raises(HTTPException) as exc_info: + await cli_poll_key( + key_id="cli-session-grants", + team_id="team-a", + x_litellm_cli_poll_secret="poll-secret", + ) + + assert exc_info.value.status_code == 500 + assert "team-a" in str(exc_info.value.detail) + mock_get_jwt.assert_not_called() + mock_cache.delete_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_cli_poll_key_mints_teamless_session_without_team_grants(self): + """A user with no team still mints, keeping their personal allowlist in the key slot.""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + mock_cache = MagicMock(redis_cache=None) + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": { + "user_id": "teamless-user", + "user_role": "internal_user", + "teams": [], + "team_details": [], + "models": ["personal-only"], + "user_email": "teamless@example.com", + }, + } + + with ( + patch("litellm.proxy.proxy_server.cli_sso_session_cache", mock_cache), + patch( + "litellm.proxy.auth.auth_checks.ExperimentalUIJWTToken.get_cli_jwt_auth_token", + return_value="minted-token", + ) as mock_get_jwt, + ): + result = await cli_poll_key( + key_id="cli-session-teamless", + team_id=None, + x_litellm_cli_poll_secret="poll-secret", + ) + + assert result["status"] == "ready" + kwargs = mock_get_jwt.call_args.kwargs + assert kwargs["team_id"] is None + assert kwargs["team_models"] == () + assert kwargs["user_info"].models == ["personal-only"] + @pytest.mark.asyncio async def test_cli_poll_key_does_not_cap_session_when_user_has_budget(self): """A user with a configured budget must not get the max_ui_session_budget fallback cap.""" @@ -3302,7 +3552,7 @@ class TestCLIKeyRegenerationFlow: "user_id": "unbudgeted-user", "user_role": "internal_user", "teams": ["team-x"], - "team_details": [{"team_id": "team-x", "team_alias": "Team X"}], + "team_details": [{"team_id": "team-x", "team_alias": "Team X", "team_models": []}], "models": ["gpt-4"], "user_email": "unbudgeted@example.com", } @@ -6539,6 +6789,17 @@ class TestCliSsoAttributionMetadata: return_value=MagicMock(metadata={"auth_provider": "generic"}) ) mock_prisma.db.litellm_usertable.update_many = AsyncMock() + mock_prisma.db.litellm_teamtable.find_many = AsyncMock( + return_value=[ + MagicMock( + model_dump=lambda: { + "team_id": "team1", + "team_alias": "team1", + "models": [], + } + ) + ] + ) with ( patch.dict( @@ -7879,6 +8140,117 @@ async def test_cli_completion_persists_assertion_under_db_user_id(): assert response.status_code == 200 +def _cli_callback_kwargs(flow): + return { + "request": _cli_callback_request(), + "key": "cli-login-id", + "flow": flow, + "result": {"sub": "raw-idp-subject"}, + "parsed_openid_result": { + "user_id": "raw-idp-subject", + "user_email": "u@example.com", + "user_role": None, + }, + "user_defined_values": None, + "prisma_client": MagicMock(), + "user_api_key_cache": MagicMock(), + "cli_sso_session_cache": MagicMock(), + "proxy_logging_obj": MagicMock(), + } + + +def _cli_callback_request(): + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + return mock_request + + +def _cli_callback_user_info(teams): + user_info = MagicMock() + user_info.user_id = "cli-user-id" + user_info.user_role = "internal_user" + user_info.models = ["personal-only"] + user_info.teams = teams + return user_info + + +@pytest.mark.asyncio +async def test_cli_completion_drops_teams_whose_rows_no_longer_exist(): + """A membership pointing at a deleted team must not be offered for selection. + + Deleting an organization removes its team rows but leaves the user's membership + behind. If that dead team still reached the session, it would be auto-selected + for a single-team user, its grants could never resolve, and every future login + would be refused with no way for the user to recover. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _CliSsoTeamDetail, + _complete_cli_sso_callback_session, + ) + + live_detail = _CliSsoTeamDetail( + team_id="team-live", team_alias="Live", team_models=("gpt-4.1",) + ) + flow = {} + with ( + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=_cli_callback_user_info(["team-live", "team-deleted"])), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso._fetch_cli_sso_team_details", + AsyncMock(return_value=(live_detail,)), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.build_cli_sso_attribution_metadata", + return_value={}, + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + AsyncMock(), + ), + ): + response = await _complete_cli_sso_callback_session(**_cli_callback_kwargs(flow)) + + assert response.status_code == 200 + assert flow["session_data"]["teams"] == ["team-live"] + assert [d["team_id"] for d in flow["session_data"]["team_details"]] == ["team-live"] + + +@pytest.mark.asyncio +async def test_cli_completion_fails_the_login_when_team_lookup_fails(): + """A lookup failure must fail the login instead of caching a teamless session. + + Silently dropping every team here would hand a team-bound user a session with + their personal allowlist, which is the same "unknown grant treated as a real + grant" bug in a quieter form. + """ + from litellm.proxy.management_endpoints.ui_sso import ( + _complete_cli_sso_callback_session, + ) + + flow = {} + with ( + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + AsyncMock(return_value=_cli_callback_user_info(["team-live"])), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso._fetch_cli_sso_team_details", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.retain_sso_identity_assertion_for_ema", + AsyncMock(), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _complete_cli_sso_callback_session(**_cli_callback_kwargs(flow)) + + assert exc_info.value.status_code == 500 + assert "session_data" not in flow + + class TestSameOriginReturnPath: """The same-origin relative return_to arm added for the MCP gateway DCR authorize round-trip: only strictly relative paths qualify, so login can never redirect the From b72dab8049f3a8c6da5033c866acca0ab96bda95 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 13 Aug 2026 16:58:14 -0700 Subject: [PATCH 05/17] feat(ui): show provider prompt cache tokens in chat response metrics (#36827) The chat metrics bar reported In/Out/Reasoning/Total/cost only, so a playground user had no signal that provider prompt caching worked. The cached-token counts were already visible in the Logs drawer, which meant the answer to "does caching work here" lived on a different page. Adds cacheReadTokens and cacheCreationTokens to TokenUsage and renders them as two chips, reusing the prompt-cache tooltip wording already introduced for the Logs drawer so both surfaces say the same thing. A single helper, extractPromptCacheTokens, normalizes the three usage shapes the playground consumes: Anthropic Messages (cache_read_input_tokens / cache_creation_input_tokens), chat completions (prompt_tokens_details) and the Responses API (input_tokens_details). All three producers call it instead of parsing per surface. Counts that are absent, zero or non-finite are dropped, so providers without prompt caching render exactly what they render today. --- .../llm_calls/anthropic_messages.test.tsx | 59 +++++++++++++++++ .../llm_calls/anthropic_messages.tsx | 2 + .../chat_ui/ResponseMetrics.test.tsx | 36 ++++++++++ .../components/chat_ui/ResponseMetrics.tsx | 44 ++++++++++++- .../src/components/chat_ui/types.ts | 10 +-- .../llm_calls/chat_completion.test.tsx | 65 +++++++++++++++++++ .../components/llm_calls/chat_completion.tsx | 2 + .../llm_calls/responses_api.test.tsx | 56 ++++++++++++++++ .../components/llm_calls/responses_api.tsx | 2 + .../LogDetailsDrawer/LogDetailContent.tsx | 5 +- .../src/utils/promptCacheUsage.test.ts | 39 +++++++++++ .../src/utils/promptCacheUsage.ts | 37 +++++++++++ 12 files changed, 345 insertions(+), 12 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx create mode 100644 ui/litellm-dashboard/src/utils/promptCacheUsage.test.ts create mode 100644 ui/litellm-dashboard/src/utils/promptCacheUsage.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx new file mode 100644 index 00000000000..995104d4b7f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx @@ -0,0 +1,59 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { makeAnthropicMessagesRequest } from "./anthropic_messages"; +import type { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => "https://example.com"), +})); + +const mockMessagesStream = vi.fn(); + +vi.mock("@anthropic-ai/sdk", () => ({ + default: vi.fn(() => ({ messages: { stream: mockMessagesStream } })), +})); + +describe("anthropic_messages prompt cache usage", () => { + const captureUsage = async (usage: Record): Promise => { + async function* mockStream() { + yield { + type: "message_delta", + usage: { input_tokens: 5000, output_tokens: 2, ...usage }, + }; + } + mockMessagesStream.mockReturnValue(mockStream()); + + const onUsageData = vi.fn(); + await makeAnthropicMessagesRequest( + [{ role: "user", content: "Hello" }], + vi.fn(), + "claude-haiku-4-5", + "test-token", + undefined, + undefined, + undefined, + undefined, + onUsageData, + ); + + expect(onUsageData).toHaveBeenCalledTimes(1); + return onUsageData.mock.calls[0][0] as TokenUsage; + }; + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("surfaces read and creation tokens from Anthropic-shape usage", async () => { + await expect( + captureUsage({ cache_read_input_tokens: 4695, cache_creation_input_tokens: 1234 }), + ).resolves.toMatchObject({ cacheReadTokens: 4695, cacheCreationTokens: 1234, promptTokens: 5000 }); + }); + + it("omits cache fields entirely when Anthropic reports no prompt caching", async () => { + const usageData = await captureUsage({}); + + expect(usageData).not.toHaveProperty("cacheReadTokens"); + expect(usageData).not.toHaveProperty("cacheCreationTokens"); + expect(usageData.promptTokens).toBe(5000); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx index 4319315396a..4facf8a2046 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx @@ -5,6 +5,7 @@ import { buildMcpToolBlocks } from "@/components/llm_calls/mcp_tool_blocks"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; +import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; export async function makeAnthropicMessagesRequest( messages: MessageType[], @@ -109,6 +110,7 @@ export async function makeAnthropicMessagesRequest( completionTokens: usage.output_tokens, promptTokens: usage.input_tokens, totalTokens: usage.input_tokens + usage.output_tokens, + ...extractPromptCacheTokens(usage), }; onUsageData(usageData); } diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx new file mode 100644 index 00000000000..5afc94eb043 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx @@ -0,0 +1,36 @@ +import { render, screen } from "@testing-library/react"; +import { describe, it, expect } from "vitest"; +import ResponseMetrics, { type TokenUsage } from "./ResponseMetrics"; + +const baseUsage: TokenUsage = { promptTokens: 5000, completionTokens: 12, totalTokens: 5012 }; + +describe("ResponseMetrics prompt cache chips", () => { + it("renders both cache chips when the provider reports reads and writes", () => { + render(); + + expect(screen.getByText("Cache Read: 4695")).toBeInTheDocument(); + expect(screen.getByText("Cache Write: 1234")).toBeInTheDocument(); + }); + + it("renders only the read chip when the provider reports reads alone", () => { + render(); + + expect(screen.getByText("Cache Read: 4695")).toBeInTheDocument(); + expect(screen.queryByText(/Cache Write/)).not.toBeInTheDocument(); + }); + + it("renders no cache chips for a provider that reports no cache fields", () => { + render(); + + expect(screen.getByText("In: 5000")).toBeInTheDocument(); + expect(screen.queryByText(/Cache Read/)).not.toBeInTheDocument(); + expect(screen.queryByText(/Cache Write/)).not.toBeInTheDocument(); + }); + + it("renders no cache chips when the provider reports zero cache tokens", () => { + render(); + + expect(screen.queryByText(/Cache Read/)).not.toBeInTheDocument(); + expect(screen.queryByText(/Cache Write/)).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx index 58069d43ade..9795201125a 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx @@ -1,12 +1,25 @@ import React from "react"; -import { ArrowDownToLine, ArrowUpFromLine, Clock, DollarSign, Hash, Lightbulb, Wrench } from "lucide-react"; +import { + ArrowDownToLine, + ArrowUpFromLine, + Clock, + Database, + DatabaseBackup, + DollarSign, + Hash, + Lightbulb, + Wrench, +} from "lucide-react"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { PROMPT_CACHE_CREATION_TOOLTIP, PROMPT_CACHE_READ_TOOLTIP } from "@/utils/promptCacheUsage"; export interface TokenUsage { completionTokens?: number; promptTokens?: number; totalTokens?: number; reasoningTokens?: number; + cacheReadTokens?: number; + cacheCreationTokens?: number; cost?: number; } @@ -38,6 +51,33 @@ function MetricItem({ label, tooltip, icon, value }: MetricItemProps) { ); } +function PromptCacheChips({ usage }: { usage?: TokenUsage }) { + const readTokens = usage?.cacheReadTokens ?? 0; + const creationTokens = usage?.cacheCreationTokens ?? 0; + + return ( + <> + {readTokens > 0 && ( +