From 9cb102b72159d4809256d4f87b3a31b3551e3364 Mon Sep 17 00:00:00 2001 From: shrey-berri Date: Mon, 26 Oct 2026 13:41:15 -0700 Subject: [PATCH 01/19] chore(deps): remove expired boto3 botocore and grpcio overrides (#48447) --- pyproject.toml | 5 ----- uv.lock | 3 --- 2 files changed, 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b15eb042959..88bb0b13d91 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -379,11 +379,6 @@ exclude-newer = "3 days" [tool.uv.exclude-newer-package] # TODO: ryan, remove after 2026-10-23, once certifi 2026.10.19 (drops the compromised eMudhra root) clears the 3-day window certifi = "2026-10-20T03:00:00Z" -# TODO(shrey): remove after 2026-10-25, once grpcio 1.84.1 (drops the compromised eMudhra roots from roots.pem) clears the 3-day window -grpcio = "2026-10-22T15:00:00Z" -# TODO(shrey): remove after 2026-10-24, once boto3/botocore 1.43.117 (drops the compromised eMudhra roots from cacert.pem) clear the 3-day window -boto3 = "2026-10-21T19:31:00Z" -botocore = "2026-10-21T19:31:00Z" # TODO: remove after 2026-10-25, once brotli 1.2.1 (CVE-2026-111497) clears the 3-day window brotli = "2026-10-22T14:01:00Z" # TODO: ryan, remove after 2026-10-24, once mlflow 3.17.2 (GHSA-wqmh-qccw-pqwf) clears the 3-day window diff --git a/uv.lock b/uv.lock index 60bcb367e0f..fc4155524d7 100644 --- a/uv.lock +++ b/uv.lock @@ -15,11 +15,8 @@ exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for exclude-newer-span = "P3D" [options.exclude-newer-package] -boto3 = "2026-10-21T19:31:00Z" -botocore = "2026-10-21T19:31:00Z" brotli = "2026-10-22T14:01:00Z" certifi = "2026-10-20T03:00:00Z" -grpcio = "2026-10-22T15:00:00Z" mlflow = "2026-10-21T03:00:00Z" mlflow-skinny = "2026-10-21T03:00:00Z" mlflow-tracing = "2026-10-21T03:00:00Z" From 768960610bc7e77e854bdb79e14a47268f1c04fb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:41:58 +0000 Subject: [PATCH 02/19] feat(cost-map): add baseten elephant-1 and mammoth-1 decision models (#48453) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++++++++ model_prices_and_context_window.json | 18 ++++++++++++++++++ tests/unit/test_utils.py | 1 + 3 files changed, 37 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3bb8b9f8029..4ee3bcb831d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -78851,5 +78851,23 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true + }, + "baseten/decision-models/baseten-elephant-9b-latest": { + "input_cost_per_token": 0.0, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true + }, + "baseten/decision-models/baseten-mammoth-27b-latest": { + "input_cost_per_token": 0.0, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3bb8b9f8029..4ee3bcb831d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -78851,5 +78851,23 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true + }, + "baseten/decision-models/baseten-elephant-9b-latest": { + "input_cost_per_token": 0.0, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true + }, + "baseten/decision-models/baseten-mammoth-27b-latest": { + "input_cost_per_token": 0.0, + "litellm_provider": "baseten", + "max_input_tokens": 16384, + "mode": "decisions", + "output_cost_per_token": 0.0, + "source": "https://inference.baseten.co/v1/models", + "supports_vision": true } } diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 55b7bc78a11..b9baef4056a 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1086,6 +1086,7 @@ _COST_MAP_ROW_SCHEMA: Final = { "chat", "completion", "container", + "decisions", "image_edit", "embedding", "evaluation", From 3beee3cfe97f3b3c04458ba5ff2612c8bd9358a7 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 26 Oct 2026 13:42:01 -0700 Subject: [PATCH 03/19] fix(elevenlabs): map wav and pcm TTS formats to 24kHz output formats (#48455) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../text_to_speech/transformation.py | 4 +-- .../test_transformation.py | 2 +- ...levenlabs_text_to_speech_transformation.py | 28 +++++++++++++++++++ 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index f910ef8c057..23ccf41c9da 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -56,9 +56,9 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): # Response format mappings from OpenAI to ElevenLabs FORMAT_MAPPINGS = { "mp3": "mp3_44100_128", - "pcm": "pcm_44100", + "pcm": "pcm_24000", "opus": "opus_48000_128", - # ElevenLabs does not support WAV, AAC, or FLAC formats. + "wav": "wav_24000", } ELEVENLABS_QUERY_PARAMS_KEY = "__elevenlabs_query_params__" diff --git a/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py b/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py index 86252a1a95d..f55644eccae 100644 --- a/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py +++ b/tests/unit/llms/elevenlabs/audio_transcription/test_transformation.py @@ -314,4 +314,4 @@ class TestElevenLabsTextToSpeechTransformation: ) assert voice_id in url - assert "output_format=pcm_44100" in url + assert "output_format=pcm_24000" in url diff --git a/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py index d05371d7df9..88f0ae1d524 100644 --- a/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py +++ b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py @@ -55,3 +55,31 @@ def test_speech_keeps_an_internal_prefixed_kwarg_out_of_the_elevenlabs_request(r sent: Final = json.loads(respx_mock.calls[0].request.content) assert "_litellm_undeclared_sentinel" not in sent, sent assert sent["text"] == "hi" + + +@pytest.mark.parametrize( + ("response_format", "expected_output_format"), + [("wav", "wav_24000"), ("pcm", "pcm_24000")], +) +def test_speech_maps_wav_and_pcm_to_elevenlabs_output_formats( + respx_mock: respx.MockRouter, + response_format: str, + expected_output_format: str, +) -> None: + api_base: Final = "http://localhost:12346" + mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/v1/text-to-speech/.*").mock( + return_value=httpx.Response(status_code=200, content=b"audio", headers={"content-type": "audio/mpeg"}) + ) + + litellm.speech( + model="elevenlabs/eleven_multilingual_v2", + input="hi", + voice="21m00Tcm4TlvDq8ikWAM", + response_format=response_format, + api_base=api_base, + api_key="fake_elevenlabs_api_key", + ) + + assert mock_route.called + request_url: Final = httpx.URL(str(respx_mock.calls[0].request.url)) + assert request_url.params["output_format"] == expected_output_format From c843cee062e08113f5afa101b0b6879824dd9025 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:55:51 +0000 Subject: [PATCH 04/19] feat(pricing): add github_copilot/gpt-6.1-terra (#48450) * feat(pricing): add github_copilot/gpt-6.1-terra * fix(pricing): raise github_copilot gpt-6 max_input_tokens to 922000 --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 32 ++++++++++++++++--- model_prices_and_context_window.json | 32 ++++++++++++++++--- 2 files changed, 56 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4ee3bcb831d..b7e0c881455 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -33698,7 +33698,7 @@ "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33722,7 +33722,7 @@ "input_cost_per_token": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33746,7 +33746,7 @@ "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33770,7 +33770,7 @@ "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -78869,5 +78869,29 @@ "output_cost_per_token": 0.0, "source": "https://inference.baseten.co/v1/models", "supports_vision": true + }, + "github_copilot/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4ee3bcb831d..b7e0c881455 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33698,7 +33698,7 @@ "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33722,7 +33722,7 @@ "input_cost_per_token": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33746,7 +33746,7 @@ "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -33770,7 +33770,7 @@ "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "github_copilot", - "max_input_tokens": 272000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -78869,5 +78869,29 @@ "output_cost_per_token": 0.0, "source": "https://inference.baseten.co/v1/models", "supports_vision": true + }, + "github_copilot/gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "source": "https://github.blog/changelog/2026-10-26-gpt-6-1-terra-in-github-copilot", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true } } From b7a2226f438499509730c9fe649beb026b7cb4fe Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 16:58:51 -0400 Subject: [PATCH 05/19] test(proxy): make tokenizer offload check deterministic (#48458) --- ...test_proxy_server_endpoints_and_startup.py | 37 ++++++++++--------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 14196d4a1ac..a859ce96c70 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -15207,21 +15207,23 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun assert_loop_stayed_free(took, lags) -async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch): - from litellm.rust_bridge._native import Tokenizer +async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None: + from queue import SimpleQueue + from threading import Thread, current_thread + + from tokenizers import Tokenizer from litellm import Router - from tests.unit.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags - claude_tokenizer: Final = litellm.utils.select_tokenizer("claude-fable-5")["tokenizer"] + claude_tokenizer: Final = Tokenizer.from_str(litellm.utils.claude_json_str) + loop_thread: Final = current_thread() + loader_threads: Final[SimpleQueue[Thread]] = SimpleQueue() - class SlowHubTokenizer: - @staticmethod - def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: - time.sleep(0.3) - return claude_tokenizer + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + loader_threads.put(current_thread()) + return claude_tokenizer - monkeypatch.setattr("litellm.rust_bridge.tokenizer.from_pretrained", SlowHubTokenizer.from_pretrained) + monkeypatch.setattr("tokenizers.Tokenizer", types.SimpleNamespace(from_pretrained=from_pretrained)) monkeypatch.setattr( "litellm.proxy.proxy_server.llm_router", Router( @@ -15237,15 +15239,16 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp ), ) - response, took, lags = await timed_with_loop_lags( - lambda: proxy_server_module.token_counter( + litellm.utils._select_custom_tokenizer_helper.cache_clear() + try: + response: Final = await proxy_server_module.token_counter( TokenCountRequest(model="self-hosted", prompt="count me off the loop") ) - ) - - assert response.tokenizer_type == "huggingface_tokenizer" - assert response.total_tokens > 0 - assert_loop_stayed_free(took, lags) + assert loader_threads.get_nowait() is not loop_thread + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens > 0 + finally: + litellm.utils._select_custom_tokenizer_helper.cache_clear() async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch): From 6d9344535093b6ad5a1b246af449c20d1db46004 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 17:09:19 -0400 Subject: [PATCH 06/19] fix(proxy): preserve limiter state across SSE keepalive tasks (#48452) --- litellm/proxy/common_request_processing.py | 3 + .../hooks/test_parallel_request_limiter_v3.py | 136 +++++++++++++++++- 2 files changed, 138 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index cec176ac559..61cf4d696a1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -119,6 +119,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( ) from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression +from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( route_request, @@ -1416,6 +1417,8 @@ async def open_sse_before_first_byte( if interval is None: return await produce_response + # Share admission state with stream cleanup before the producer copies this context. + get_or_create_request_stash() produce_task: Final = asyncio.ensure_future(produce_response) await asyncio.wait((produce_task,), timeout=interval) if produce_task.done(): diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 1ec505362dd..345f37f1db4 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7,7 +7,7 @@ import contextvars import logging import sys import time -from collections.abc import Iterator, Sequence +from collections.abc import AsyncGenerator, Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone from typing import Any, Dict, Final, List, Optional @@ -4644,6 +4644,140 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): ) +_SSE_KEEPALIVE_CALLBACK_TIME: Final = datetime(2026, 1, 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("after_ping", [False, True], ids=["fast", "pinged"]) +@pytest.mark.parametrize("fails", [False, True], ids=["success", "failure"]) +@pytest.mark.parametrize("preadmitted", [False, True], ids=["producer-admission", "existing-admission"]) +async def test_sse_keepalive_releases_parallel_capacity( + after_ping: bool, fails: bool, preadmitted: bool +) -> None: + from fastapi.responses import StreamingResponse + + from litellm.proxy.common_request_processing import open_sse_before_first_byte + from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES + + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + api_key: Final = hash_token("sk-keepalive-release") + auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + counter: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + + async def request(call_id: str) -> None: + data: Final = {"model": "test-model", "litellm_call_id": call_id} + upstream_answers: Final = asyncio.Event() + if not after_ping: + upstream_answers.set() + + async def admit() -> None: + await handler.async_pre_call_hook( + user_api_key_dict=auth, cache=cache, data=data, call_type="acompletion" + ) + + async def failure(exc: Exception) -> None: + await handler.async_post_call_failure_hook( + request_data=data, original_exception=exc, user_api_key_dict=auth + ) + + async def stream() -> AsyncGenerator[bytes, None]: + yield b"data: [DONE]\n\n" + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": call_id, + "standard_logging_object": {"metadata": {"user_api_key_hash": api_key}}, + }, + response_obj=ModelResponse(usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2)), + start_time=_SSE_KEEPALIVE_CALLBACK_TIME, + end_time=_SSE_KEEPALIVE_CALLBACK_TIME, + ) + + async def produce() -> StreamingResponse: + if not preadmitted: + await admit() + await upstream_answers.wait() + if fails: + raise HTTPException(status_code=502, detail="upstream failed") + return StreamingResponse(stream()) + + if preadmitted: + await admit() + response_call: Final = open_sse_before_first_byte( + produce(), 0.001 if after_ping else 60.0, on_late_failure=failure + ) + if fails and not after_ping: + with pytest.raises(HTTPException) as rejected: + await response_call + assert rejected.value.status_code == 502 + await failure(rejected.value) + return + response: Final = await response_call + first: Final = await anext(response.body_iterator) + if after_ping: + upstream_answers.set() + rest: Final = tuple([chunk async for chunk in response.body_iterator]) + assert first == (SSE_COMMENT_PING_BYTES if after_ping else b"data: [DONE]\n\n") + if fails: + error_frame: Final = rest[-2] + assert isinstance(error_frame, bytes) and b"upstream failed" in error_frame + assert rest[-1] == b"data: [DONE]\n\n" + elif after_ping: + assert rest[-1] == b"data: [DONE]\n\n" + else: + assert rest == () + + for index in range(3): + await contextvars.Context().run(asyncio.create_task, request(f"keepalive-{index}")) + assert handler._gauge_in_flight_from_cache_value(await cache.async_get_cache(counter)) == 0 + + +@pytest.mark.asyncio +async def test_sse_keepalive_rejects_overlap_and_reuses_completed_capacity() -> None: + from fastapi.responses import StreamingResponse + + from litellm.proxy.common_request_processing import open_sse_before_first_byte + + cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + api_key: Final = hash_token("sk-keepalive-overlap") + auth: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + admitted: Final = asyncio.Event() + finish: Final = asyncio.Event() + + async def request(call_id: str) -> None: + async def stream() -> AsyncGenerator[bytes, None]: + await finish.wait() + yield b"data: [DONE]\n\n" + await handler.async_post_call_success_hook( + data={"litellm_call_id": call_id}, user_api_key_dict=auth, response=ModelResponse() + ) + + async def produce() -> StreamingResponse: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=cache, + data={"model": "test-model", "litellm_call_id": call_id}, + call_type="acompletion", + ) + admitted.set() + return StreamingResponse(stream()) + + response: Final = await open_sse_before_first_byte(produce(), 60.0) + assert tuple([chunk async for chunk in response.body_iterator]) == (b"data: [DONE]\n\n",) + + first: Final = contextvars.Context().run(asyncio.create_task, request("first")) + await admitted.wait() + try: + with pytest.raises(HTTPException) as rejected: + await contextvars.Context().run(asyncio.create_task, request("overlapping")) + assert rejected.value.status_code == 429 + finally: + finish.set() + await first + await contextvars.Context().run(asyncio.create_task, request("after-completion")) + + @pytest.mark.asyncio async def test_read_only_gauge_check_counts_without_acquiring_v3(): """ From ccd783d907d8b8a88add48c35732371e8701febf Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 17:16:15 -0400 Subject: [PATCH 07/19] fix(proxy): reject negative RPM and TPM request limits (#48444) --- litellm/proxy/_types.py | 12 ++- .../management_endpoints/key_requests.py | 4 + .../management_endpoints/user_requests.py | 4 + tests/unit/proxy/test__types.py | 83 ++++++++++++++++++- 4 files changed, 98 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2b01e02bb06..db433f07f3b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1766,8 +1766,10 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase): max_parallel_requests: int | None = Field( default=None, description="Max concurrent requests allowed for this budget id." ) - tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.") - rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.") + tpm_limit: int | None = Field(default=None, ge=0, description="Max tokens per minute, allowed for this budget id.") + rpm_limit: int | None = Field( + default=None, ge=0, description="Max requests per minute, allowed for this budget id." + ) tpd_limit: int | None = Field( default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id." ) @@ -1878,6 +1880,8 @@ RouterSettingsDict = Annotated[ class NewTeamRequest(TeamBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) router_settings: RouterSettingsDict | None = None model_aliases: dict | None = None model_max_budget: GenericBudgetConfigType | None = Field( @@ -1957,8 +1961,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): team_alias: str | None = None organization_id: str | None = None metadata: dict | None = None - tpm_limit: int | None = None - rpm_limit: int | None = None + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index baa02b68d98..172357d8ac4 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -55,6 +55,8 @@ class LiteLLMKeyType(str, enum.Enum): class GenerateKeyRequest(KeyRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) soft_budget: float | None = None send_invite_email: bool | None = None key_type: LiteLLMKeyType | None = Field( @@ -82,6 +84,8 @@ def _is_mapping(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: class UpdateKeyRequest(KeyRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) # Note: the defaults of all Params here MUST BE NONE # else they will get overwritten duration: str | None = None diff --git a/litellm/types/proxy/management_endpoints/user_requests.py b/litellm/types/proxy/management_endpoints/user_requests.py index 49dda5995bc..e53dc4152c3 100644 --- a/litellm/types/proxy/management_endpoints/user_requests.py +++ b/litellm/types/proxy/management_endpoints/user_requests.py @@ -16,6 +16,8 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): class NewUserRequest(GenerateRequestBase): + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) max_budget: float | None = None user_email: str | None = None user_alias: str | None = None @@ -47,6 +49,8 @@ class NewUserRequest(GenerateRequestBase): class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest + tpm_limit: int | None = Field(default=None, ge=0) + rpm_limit: int | None = Field(default=None, ge=0) # repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model password: str | None = Field(default=None, repr=False) spend: float | None = None diff --git a/tests/unit/proxy/test__types.py b/tests/unit/proxy/test__types.py index 68bfe4b34e5..76ca49362b4 100644 --- a/tests/unit/proxy/test__types.py +++ b/tests/unit/proxy/test__types.py @@ -2,23 +2,35 @@ import json from typing import Final import pytest -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from litellm.proxy._types import ( ROLES_WITHIN_ORG, + BudgetNewRequest, ChangePasswordRequest, GenerateKeyRequest, + GenerateKeyResponse, KeyRequest, LiteLLM_AuditLogs, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, + NewCustomerRequest, NewMCPServerRequest, + NewTeamRequest, NewUserRequest, + NewUserResponse, OrganizationMemberUpdateRequest, + RegenerateKeyRequest, ResetSpendRequest, UpdateKeyRequest, UpdateMCPServerRequest, + UpdateTeamRequest, UpdateUserRequest, + UpdateUserRequestNoUserIDorEmail, UserAPIKeyAuth, ) @@ -213,6 +225,75 @@ def test_a_max_budget_that_was_actually_sent_is_kept(sent): assert GenerateKeyRequest(max_budget=sent).max_budget == sent +RATE_LIMIT_REQUESTS: Final = ( + (GenerateKeyRequest, {}), + (RegenerateKeyRequest, {}), + (UpdateKeyRequest, {"key": "sk-test"}), + (NewUserRequest, {}), + (UpdateUserRequestNoUserIDorEmail, {}), + (UpdateUserRequest, {"user_id": "user-test"}), + (NewTeamRequest, {}), + (UpdateTeamRequest, {"team_id": "team-test"}), + (BudgetNewRequest, {}), + (NewCustomerRequest, {"user_id": "customer-test"}), +) + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +@pytest.mark.parametrize("field", ("rpm_limit", "tpm_limit")) +@pytest.mark.parametrize("value", (-1, -100)) +def test_management_requests_reject_negative_rate_limits( + request_model: type[BaseModel], identifiers: dict[str, str], field: str, value: int +) -> None: + with pytest.raises(ValidationError) as error: + request_model.model_validate({**identifiers, field: value}) + assert error.value.errors()[0]["loc"] == (field,) + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +@pytest.mark.parametrize("value", (None, 0, 100)) +def test_management_requests_preserve_explicit_rate_limits( + request_model: type[BaseModel], identifiers: dict[str, str], value: int | None +) -> None: + limits: Final = {"rpm_limit": value, "tpm_limit": value} + request: Final = request_model.model_validate({**identifiers, **limits}) + assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == limits + + +@pytest.mark.parametrize(("request_model", "identifiers"), RATE_LIMIT_REQUESTS) +def test_management_requests_leave_omitted_rate_limits_unset( + request_model: type[BaseModel], identifiers: dict[str, str] +) -> None: + request: Final = request_model.model_validate({**identifiers}) + assert request.model_dump(include={"rpm_limit", "tpm_limit"}) == {"rpm_limit": None, "tpm_limit": None} + assert request.model_dump(include={"rpm_limit", "tpm_limit"}, exclude_unset=True) == {} + + +@pytest.mark.parametrize(("response_model", "identifiers"), ( + (GenerateKeyResponse, {"key": "sk-test"}), + (NewUserResponse, {"key": "sk-test"}), + (LiteLLM_TeamTable, {"team_id": "team-test"}), + (LiteLLM_UserTable, {"user_id": "user-test"}), + (LiteLLM_BudgetTable, {"budget_id": "budget-test"}), +)) +def test_existing_negative_rate_limits_remain_readable( + response_model: type[BaseModel], identifiers: dict[str, str] +) -> None: + limits: Final = {"rpm_limit": -1, "tpm_limit": -100} + response: Final = response_model.model_validate({**identifiers, **limits}) + assert response.model_dump(include={"rpm_limit", "tpm_limit"}) == limits + + +def test_existing_customer_negative_rate_limits_remain_readable() -> None: + customer: Final = LiteLLM_EndUserTable.model_validate({ + "user_id": "customer-test", "blocked": False, + "litellm_budget_table": {"rpm_limit": -1, "tpm_limit": -100}, + }) + assert customer.litellm_budget_table is not None + assert customer.litellm_budget_table.rpm_limit == -1 + assert customer.litellm_budget_table.tpm_limit == -100 + + USER_IDENTIFIER_CASES = ( ({"user_id": "user-1"}, True), ({"user_email": "user@example.com"}, True), From dbd55d32023157fa1e13db994645878f8f870fd8 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 26 Oct 2026 14:55:08 -0700 Subject: [PATCH 08/19] fix(ui): keep Aria accessible in the dashboard header (#48454) * fix(ui): keep Aria accessible in the dashboard header * test(ui): wait for every pagination fixture log * test(ui): update assistant setup for Aria branding --- tests/e2e/ui/globalSetup.ts | 4 +- tests/e2e/ui/helpers/navigation.ts | 23 ++++-- .../e2e/ui/tests/logs/logsPagination.spec.ts | 61 ++++++++++++++- ui/litellm-dashboard/AGENTS.md | 2 +- .../src/components/DashboardHeader.test.tsx | 3 +- .../src/components/DashboardHeader.tsx | 75 ++++++++++--------- .../liteadmin/LiteAdmin.integration.test.tsx | 67 ++++++++--------- .../src/components/liteadmin/LiteAdmin.tsx | 37 +++------ 8 files changed, 161 insertions(+), 111 deletions(-) diff --git a/tests/e2e/ui/globalSetup.ts b/tests/e2e/ui/globalSetup.ts index d97789c227e..e18d1be390c 100644 --- a/tests/e2e/ui/globalSetup.ts +++ b/tests/e2e/ui/globalSetup.ts @@ -2,7 +2,7 @@ import { chromium, expect, request } from "@playwright/test"; import { users, Role, STORAGE_PATHS } from "./fixtures/users"; import { ARTIFACT_DIR, UI_BASE_URL } from "./constants"; import { expectUnrestrictedDashboard, setInvitedUserPassword } from "./helpers/userOnboarding"; -import { hideLiteAdmin } from "./helpers/navigation"; +import { setAriaHidden } from "./helpers/navigation"; import { masterKey as getMasterKey } from "./helpers/traffic"; import * as fs from "fs"; import * as path from "path"; @@ -78,7 +78,7 @@ async function globalSetup() { await dismiss.click(); } if (role === Role.ProxyAdmin) { - await hideLiteAdmin(page); + await setAriaHidden(page, true); } // The login flow stores a post-login return URL in the litellm_return_url // cookie. If the snapshot captures it before the app consumes it, every diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index 78288fab743..be46fcf0395 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -62,18 +62,29 @@ export async function dismissFeedbackPopup(page: PlaywrightPage): Promise } } -export async function hideLiteAdmin(page: PlaywrightPage): Promise { +export async function setAriaHidden(page: PlaywrightPage, hidden: boolean): Promise { await page.getByRole("button", { name: /Account menu/i }).click(); const panel = page.getByTestId("sidebar-account-menu-panel"); await expect(panel).toBeVisible({ timeout: 5_000 }); - const toggle = panel.getByRole("switch", { name: "Toggle hide LiteAdmin" }); - if ((await toggle.getAttribute("aria-checked")) !== "true") { - await toggle.click(); + const premium = await panel.getByText("Premium", { exact: true }).isVisible(); + const toggle = panel.getByRole("switch", { name: "Toggle hide Aria Agent", exact: true }); + if (premium) { + await expect(toggle).toBeVisible(); + if ((await toggle.getAttribute("aria-checked")) !== String(hidden)) { + await toggle.click(); + } + await expect(toggle).toHaveAttribute("aria-checked", String(hidden)); + } else { + await expect(toggle).toHaveCount(0); } - await expect(toggle).toHaveAttribute("aria-checked", "true"); await page.keyboard.press("Escape"); await expect(panel).toBeHidden(); - await expect(page.getByRole("button", { name: "LiteAdmin", exact: true })).toBeHidden(); + const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + if (premium && hidden) { + await expect(aria).toBeHidden(); + } else { + await expect(aria).toBeVisible(); + } } /** diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts index 416e1c3171f..838a829a3e2 100644 --- a/tests/e2e/ui/tests/logs/logsPagination.spec.ts +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -1,8 +1,14 @@ import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test"; import { ADMIN_STORAGE_PATH } from "../../constants"; -import { navigateToPage, dismissFeedbackPopup } from "../../helpers/navigation"; +import { navigateToPage, dismissFeedbackPopup, setAriaHidden } from "../../helpers/navigation"; import { Page } from "../../fixtures/pages"; -import { CHAT_MODEL_A, createVirtualKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic"; +import { + CHAT_MODEL_A, + createVirtualKey, + deleteVirtualKey, + sendChatCompletion, + waitForSpendLog, +} from "../../helpers/traffic"; /** * Session-grouped pagination (#38060): a page of N rows must render exactly N session rows, a @@ -45,6 +51,57 @@ async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): await page.getByRole("option", { name: size, exact: true }).click(); } +test.describe("Logs pagination with Aria in the header", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + for (const viewport of [ + { width: 1440, height: 900 }, + { width: 390, height: 844 }, + ]) { + test(`keeps pagination reachable and Aria available at ${viewport.width}px`, async ({ page, request }) => { + await page.setViewportSize(viewport); + const alias = `e2e-aria-header-${uniqueSuffix()}`; + const key = await createVirtualKey(request, { key_alias: alias }); + try { + const ids = await Promise.all( + Array.from({ length: 26 }, (_, index) => + sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt: `${alias}-${index}`, + apiKey: key.key, + }), + ), + ); + await Promise.all(ids.map((id) => waitForSpendLog(request, id))); + await openLogs(page); + await setAriaHidden(page, false); + await applyKeyAliasFilter(page, await openFilterDrawer(page), alias); + await setRowsPerPage(page, "25"); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + await visibleTestId(page, "pagination-next").click(); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 26-26 of 26"); + + const aria = page.getByRole("banner").getByRole("button", { name: "Aria Agent", exact: true }); + await aria.click(); + await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); + await page.keyboard.press("Escape"); + await expect(aria).toBeFocused(); + await visibleTestId(page, "pagination-first").click(); + await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); + for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { + await navigateToPage(page, destination); + await aria.click(); + await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); + await page.keyboard.press("Escape"); + await expect(aria).toBeFocused(); + } + } finally { + await deleteVirtualKey(request, key.token); + } + }); + } +}); + test.describe("Logs page session-grouped pagination", () => { test.use({ storageState: ADMIN_STORAGE_PATH }); diff --git a/ui/litellm-dashboard/AGENTS.md b/ui/litellm-dashboard/AGENTS.md index bbf2e49aad7..b6fe5219022 100644 --- a/ui/litellm-dashboard/AGENTS.md +++ b/ui/litellm-dashboard/AGENTS.md @@ -30,7 +30,7 @@ Type tests are `*.test-d.ts` files run by the `types` vitest project (`npm run t -## This is NOT the Next.js you know +# This is NOT the Next.js you know This version has breaking changes — APIs, conventions, and file structure may all differ from your training data. Read the relevant guide in `node_modules/next/dist/docs/` (resolved from this file's directory; in monorepos the `next` package may not be visible from the repo root) before writing any code. Heed deprecation notices. diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx index 770315af49f..202a111ea97 100644 --- a/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/DashboardHeader.test.tsx @@ -124,7 +124,8 @@ describe("DashboardHeader breadcrumb", () => { const tools = await screen.findByRole("dialog", { name: "Gateway tools" }); expect(within(tools).getByRole("button", { name: "Switch workspace, current: AI Gateway" })).toBeInTheDocument(); expect(within(tools).getByRole("link", { name: "Docs" })).toBeInTheDocument(); - expect(within(tools).getByRole("button", { name: "Aria Agent" })).toBeInTheDocument(); + expect(within(screen.getByRole("banner")).getByRole("button", { name: "Aria Agent" })).toBeInTheDocument(); + expect(screen.getAllByRole("button", { name: "Aria Agent" })).toHaveLength(1); }); it("closes mobile tools when switching to desktop and keeps them closed when returning", async () => { diff --git a/ui/litellm-dashboard/src/components/DashboardHeader.tsx b/ui/litellm-dashboard/src/components/DashboardHeader.tsx index b31fc794f79..55f163c2495 100644 --- a/ui/litellm-dashboard/src/components/DashboardHeader.tsx +++ b/ui/litellm-dashboard/src/components/DashboardHeader.tsx @@ -73,43 +73,46 @@ export function DashboardHeader({ -
- {showWorkerSwitch && ( - <> - - - - )} +
- - - {!hideCommunityLinks && } - - - {deploymentHealth} - -
-
- - {deploymentHealth} - - - }> - - - - Gateway tools - - {showWorkerSwitch && } -
- - - - - {!hideCommunityLinks && } -
-
-
+
+ {showWorkerSwitch && ( + <> + + + + )} + + + {!hideCommunityLinks && } + + + {deploymentHealth} + +
+
+ + {deploymentHealth} + + + } + > + + + + Gateway tools + + {showWorkerSwitch && } +
+ + + + {!hideCommunityLinks && } +
+
+
+
); diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx index f32c99bedf3..ed66f7ce2ef 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.integration.test.tsx @@ -240,7 +240,7 @@ describe("Aria Agent in the gateway", () => { session(role); const requests = gateway([answer("Here is your weekly spend.")]); renderWidget(); - await userEvent.setup().click(await screen.findByRole("button", { name: "Launch Aria Agent" })); + await userEvent.setup().click(await screen.findByRole("button", { name: "Aria Agent" })); const question = "How much have I spent this week?"; const suggestion = await screen.findByRole("button", { name: question }); await waitFor(() => expect(suggestion).toBeEnabled()); @@ -268,7 +268,7 @@ describe("Aria Agent in the gateway", () => { it("should disable suggested questions when no chat model is available", async () => { const requests = gateway([], { models: [] }); renderWidget(); - await userEvent.setup().click(await screen.findByRole("button", { name: "Launch Aria Agent" })); + await userEvent.setup().click(await screen.findByRole("button", { name: "Aria Agent" })); const picker = await screen.findByRole("combobox", { name: "Aria Agent model" }); expect(picker).toBeEnabled(); expect(screen.queryByText("No chat models are available for your account.")).not.toBeInTheDocument(); @@ -285,7 +285,7 @@ describe("Aria Agent in the gateway", () => { expect(requests.filter((request) => request.url.endsWith("/chat/completions"))).toHaveLength(0); }); - it("should restore launcher focus and preserve the conversation when the popup closes with Escape", async () => { + it("should restore header trigger focus and preserve the conversation when the popup closes with Escape", async () => { const requests = gateway([answer("Your current usage is $12.34."), answer("Last month was $10.")]); const user = userEvent.setup(); const backgroundAction = vi.fn(); @@ -293,7 +293,7 @@ describe("Aria Agent in the gateway", () => { return ; } renderWidget(PageAction); - await user.click(await screen.findByRole("button", { name: "Launch Aria Agent" })); + await user.click(await screen.findByRole("button", { name: "Aria Agent" })); const popup = await screen.findByRole("dialog", { name: "Aria Agent" }); expect(popup).toBeVisible(); await waitFor(() => expect(screen.getByPlaceholderText("Ask Aria…")).toBeEnabled()); @@ -306,8 +306,8 @@ describe("Aria Agent in the gateway", () => { await user.click(screen.getByPlaceholderText("Ask Aria…")); await user.keyboard("{Escape}"); await waitFor(() => expect(screen.queryByRole("dialog", { name: "Aria Agent" })).not.toBeInTheDocument()); - const launcher = screen.getByRole("button", { name: "Launch Aria Agent" }); - await waitFor(() => expect(launcher).toHaveFocus()); + const trigger = screen.getByRole("button", { name: "Aria Agent" }); + await waitFor(() => expect(trigger).toHaveFocus()); await user.keyboard("{Enter}"); expect(await screen.findByRole("dialog", { name: "Aria Agent" })).toBeVisible(); expect(screen.getByText("Your current usage is $12.34.")).toBeInTheDocument(); @@ -324,26 +324,25 @@ describe("Aria Agent in the gateway", () => { ]); await user.click(screen.getByRole("button", { name: "Close Aria Agent" })); await waitFor(() => expect(screen.queryByRole("dialog", { name: "Aria Agent" })).not.toBeInTheDocument()); - await waitFor(() => expect(screen.getByRole("button", { name: "Launch Aria Agent" })).toHaveFocus()); + await waitFor(() => expect(screen.getByRole("button", { name: "Aria Agent" })).toHaveFocus()); }); - it("should show the same product name on Home, Usage, the launcher and the popup", async () => { + it("should show the same product name on Home, Usage, the header trigger and the popup", async () => { gateway([]); const view = renderWidget(undefined, true); const home = await screen.findByRole("region", { name: "Ask Aria Agent" }); expect(within(home).getByRole("heading", { name: "Aria Agent" })).toHaveTextContent("Aria Agent"); expect(screen.getByRole("button", { name: "Aria Agent" })).toHaveTextContent("Aria Agent"); - expect(screen.getByRole("button", { name: "Launch Aria Agent" })).toHaveTextContent("Aria Agent"); view.refresh(false, true); const usage = await screen.findByRole("region", { name: "Ask Aria about usage" }); expect(within(usage).getByRole("heading", { name: "Aria Agent" })).toHaveTextContent("Aria Agent"); expect(screen.getByRole("button", { name: "Open Aria Agent" })).toHaveTextContent("Aria Agent"); - act(() => fireEvent.click(screen.getByRole("button", { name: "Launch Aria Agent" }))); + act(() => fireEvent.click(screen.getByRole("button", { name: "Aria Agent" }))); const popup = await screen.findByRole("dialog", { name: "Aria Agent" }); expect(within(popup).getByRole("heading", { name: "Aria Agent" })).toHaveTextContent("Aria Agent"); }); - it("should keep the same scoped conversation when leaving Usage and reopening the launcher", async () => { + it("should keep the same scoped conversation when leaving Usage and reopening the header trigger", async () => { session("internal_user"); const requests = gateway( [answer("Your usage is within budget."), answer("The same usage conversation continues.")], @@ -355,7 +354,7 @@ describe("Aria Agent in the gateway", () => { }, ); const view = renderWidget(undefined, false, true); - expect(screen.queryByRole("button", { name: "Launch Aria Agent" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); const usage = await screen.findByRole("region", { name: "Ask Aria about usage" }); await selectModel("other-chat-model"); act(() => { @@ -366,14 +365,13 @@ describe("Aria Agent in the gateway", () => { view.refresh(false, false); expect(screen.queryByRole("region", { name: "Ask Aria about usage" })).not.toBeInTheDocument(); - expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); - const launcher = screen.getByRole("button", { name: "Launch Aria Agent" }); - expect(launcher).toHaveTextContent("Aria Agent"); - act(() => launcher.focus()); + const trigger = screen.getByRole("button", { name: "Aria Agent" }); + expect(trigger).toHaveTextContent("Aria Agent"); + act(() => trigger.focus()); await userEvent.setup().keyboard("{Enter}"); expect(await screen.findByPlaceholderText("Ask Aria…")).toHaveValue("Check my usage"); expect(screen.getByRole("combobox", { name: "Aria Agent model" })).toHaveValue("other-chat-model"); - expect(screen.queryByRole("button", { name: "Launch Aria Agent" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toHaveAttribute("aria-expanded", "true"); act(() => fireEvent.click(screen.getByRole("button", { name: "Send message" }))); await screen.findByText("Your usage is within budget."); act(() => { @@ -385,7 +383,7 @@ describe("Aria Agent in the gateway", () => { const restoredUsage = await screen.findByRole("region", { name: "Ask Aria about usage" }); expect(within(restoredUsage).getByRole("textbox", { name: "Message Aria Agent" })).toHaveValue("And last month?"); view.refresh(false, false); - act(() => fireEvent.click(screen.getByRole("button", { name: "Launch Aria Agent" }))); + act(() => fireEvent.click(screen.getByRole("button", { name: "Aria Agent" }))); expect(screen.getAllByRole("dialog", { name: "Aria Agent" })).toHaveLength(1); expect(screen.getByText("Your usage is within budget.")).toBeInTheDocument(); expect(screen.getByPlaceholderText("Ask Aria…")).toHaveValue("And last month?"); @@ -405,13 +403,13 @@ describe("Aria Agent in the gateway", () => { ]); }); - it("should keep the launcher unavailable without a signed-in session", async () => { + it("should keep the header trigger unavailable without a signed-in session", async () => { document.cookie = "token=; Max-Age=0; Path=/"; const requests = gateway([]); const { client } = renderWidget(); await screen.findByText("Session ready"); await waitFor(() => expect(client.isFetching()).toBe(0)); - expect(screen.queryByRole("button", { name: "Launch Aria Agent" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); act(() => fireEvent.keyDown(document, { key: "j", ctrlKey: true })); expect(screen.queryByRole("dialog", { name: "Aria Agent" })).not.toBeInTheDocument(); expect(requests.every((request) => request.url.endsWith("/litellm-ui-config"))).toBe(true); @@ -494,7 +492,7 @@ describe("Aria Agent in the gateway", () => { expect(await screen.findByText("Your usage totals $12.34.")).toBeInTheDocument(); expect(screen.getAllByRole("dialog", { name: "Aria Agent" })).toHaveLength(1); expect(usage).not.toBeVisible(); - if (!writable) expect(screen.queryByRole("button", { name: "Aria Agent", hidden: true })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toHaveAttribute("aria-expanded", "true"); const read = requests.find((request) => request.url.includes("/user/daily/activity/aggregated")); expect(read).toBeDefined(); expect(new URL(read!.url).searchParams.get("user_id")).toBe(scope === "user" ? "first-admin" : null); @@ -553,7 +551,7 @@ describe("Aria Agent in the gateway", () => { await waitFor(() => expect(usageInput).toBeEnabled()); expect(usageInput).toHaveValue(""); expect(screen.queryByRole("button", { name: "Confirm change" })).not.toBeInTheDocument(); - expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toHaveAttribute("aria-expanded", "false"); expect(screen.queryByRole("dialog", { name: "Aria Agent" })).not.toBeInTheDocument(); act(() => send("Show my usage")); expect(await screen.findByText("Only your own usage is available.")).toBeInTheDocument(); @@ -812,7 +810,6 @@ describe("Aria Agent in the gateway", () => { await user.click(toggle); expect(toggle).toBeChecked(); expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); - expect(screen.queryByRole("button", { name: "Launch Aria Agent" })).not.toBeInTheDocument(); expect(screen.queryByRole("region", { name: "Ask Aria Agent", hidden: true })).not.toBeInTheDocument(); view.unmount(); @@ -820,13 +817,12 @@ describe("Aria Agent in the gateway", () => { await screen.findByText("Session ready"); await waitFor(() => expect(restored.client.isFetching()).toBe(0)); expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); - expect(screen.queryByRole("button", { name: "Launch Aria Agent" })).not.toBeInTheDocument(); await user.click(screen.getByRole("button", { name: /account menu/i })); const savedToggle = await screen.findByRole("switch", { name: "Toggle hide Aria Agent" }); expect(savedToggle).toBeChecked(); await user.click(savedToggle); expect(await screen.findByRole("button", { name: "Aria Agent" })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Launch Aria Agent" })).toBeVisible(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toBeVisible(); expect(await screen.findByRole("region", { name: "Ask Aria Agent" })).toBeVisible(); for (const premiumUser of [false, null, undefined, "true"]) { @@ -838,27 +834,26 @@ describe("Aria Agent in the gateway", () => { }); it.each([false, null, undefined, "true"])( - "should offer the scoped launcher without admin Home or header access when premium_user is %s", + "should offer the scoped header trigger without admin Home access when premium_user is %s", async (premiumUser) => { session("proxy_admin", "first-admin", { premium_user: premiumUser }); const requests = gateway([]); const { client } = renderWidget(undefined, true); await screen.findByText("Session ready"); await waitFor(() => expect(client.isFetching()).toBe(0)); - expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toBeVisible(); expect(screen.queryByRole("region", { name: "Ask Aria Agent", hidden: true })).not.toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Launch Aria Agent" })).toBeVisible(); expect(requests.every((request) => request.url.endsWith("/litellm-ui-config"))).toBe(true); act(() => fireEvent.keyDown(document, { key: "j", metaKey: true })); expect(await screen.findByRole("dialog", { name: "Aria Agent" })).toBeInTheDocument(); - expect(screen.queryByRole("button", { name: "Launch Aria Agent" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toHaveAttribute("aria-expanded", "true"); act(() => fireEvent.keyDown(document, { key: "j", ctrlKey: true })); expect(screen.queryByRole("dialog", { name: "Aria Agent" })).not.toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Launch Aria Agent" })).toBeVisible(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toBeVisible(); }, ); - it("should discard admin state and preserve scoped launcher access when enterprise access is lost", async () => { + it("should discard admin state and preserve scoped header trigger access when enterprise access is lost", async () => { const requests = gateway([answer("Your usage remains available.")]); const view = renderWidget(); await openWidget(); @@ -866,9 +861,9 @@ describe("Aria Agent in the gateway", () => { await waitFor(() => expect(view.client.isFetching()).toBe(0)); session("proxy_admin", "first-admin", { premium_user: false }); view.refresh(); - expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toHaveAttribute("aria-expanded", "false"); expect(screen.queryByRole("dialog", { name: "Aria Agent" })).not.toBeInTheDocument(); - act(() => fireEvent.click(screen.getByRole("button", { name: "Launch Aria Agent" }))); + act(() => fireEvent.click(screen.getByRole("button", { name: "Aria Agent" }))); await waitFor(() => expect(screen.getByPlaceholderText("Ask Aria…")).toBeEnabled()); expect(screen.getByPlaceholderText("Ask Aria…")).toHaveValue(""); act(() => send("Show my usage")); @@ -997,17 +992,17 @@ describe("Aria Agent in the gateway", () => { }); it.each(["proxy_admin_viewer", "internal_user", "internal_user_viewer", "org_admin"])( - "should expose only scoped usage tools from the dashboard launcher to %s", + "should expose only scoped usage tools from the dashboard header trigger to %s", async (role) => { session(role); const requests = gateway([answer("Usage is available.")]); const { client } = renderWidget(undefined, true); await screen.findByText("Session ready"); await waitFor(() => expect(client.isFetching()).toBe(0)); - expect(screen.queryByRole("button", { name: "Aria Agent" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Aria Agent" })).toBeVisible(); expect(screen.queryByRole("region", { name: "Ask Aria Agent", hidden: true })).not.toBeInTheDocument(); expect(requests.every((request) => request.url.endsWith("/litellm-ui-config"))).toBe(true); - act(() => fireEvent.click(screen.getByRole("button", { name: "Launch Aria Agent" }))); + act(() => fireEvent.click(screen.getByRole("button", { name: "Aria Agent" }))); await waitFor(() => expect(screen.getByPlaceholderText("Ask Aria…")).toBeEnabled()); act(() => send("Show my usage")); await screen.findByText("Usage is available."); diff --git a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx index 66c1ffbc9b7..02194aabe7b 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/LiteAdmin.tsx @@ -45,7 +45,7 @@ type LiteAdminState = { toggle: () => void; show: () => void; access: LiteAdminSession["access"]; - launcherRef: RefObject; + triggerRef: RefObject; homeContainer: HTMLDivElement | null; setHomeContainer: (container: HTMLDivElement | null) => void; usageContainer: HTMLDivElement | null; @@ -78,12 +78,11 @@ function useLiteAdminSession() { }; } -/** Keeps one Aria conversation available from the dashboard and its floating launcher. */ export function LiteAdminFrame({ children }: { children: ReactNode }) { const configured = useLiteAdminSession(); const sessionKey = configured?.key ?? null; const [open, setOpen] = useState(false); - const launcherRef = useRef(null); + const triggerRef = useRef(null); const [started, setStarted] = useState(false); const [homeContainer, setHomeContainer] = useState(null); const [usageContainer, updateUsageContainer] = useState(null); @@ -128,7 +127,7 @@ export function LiteAdminFrame({ children }: { children: ReactNode }) { toggle, show, access: configured.session.access, - launcherRef, + triggerRef, homeContainer, setHomeContainer, usageContainer, @@ -138,7 +137,6 @@ export function LiteAdminFrame({ children }: { children: ReactNode }) { } > {children} - {configured && !open && } {configured && showPanel && ( setOpen(false)} /> )} @@ -146,23 +144,6 @@ export function LiteAdminFrame({ children }: { children: ReactNode }) { ); } -function AriaLauncher({ show, buttonRef }: { show: () => void; buttonRef: RefObject }) { - return ( - - ); -} - export function LiteAdminHome() { const state = useContext(LiteAdminContext); if (!state || state.access.mode !== "admin") return null; @@ -189,18 +170,20 @@ export function LiteAdminHome() { export default function LiteAdminTrigger() { const state = useContext(LiteAdminContext); - if (!state || state.access.mode !== "admin") return null; + if (!state) return null; + const { triggerRef, open, toggle } = state; return ( ); @@ -358,7 +341,7 @@ function Panel({ initialFocus={() => initialFocus?.current ?? ref.current?.querySelector("textarea:enabled") ?? ref.current } - finalFocus={state?.launcherRef} + finalFocus={state?.triggerRef} className="flex h-[min(48rem,calc(100dvh-2rem))] flex-col gap-0 overflow-hidden p-0 shadow-xl sm:max-w-4xl" > {children} From a6fdc7c5a3fa8952f69623d56a52ada73dc5c893 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 26 Oct 2026 15:01:40 -0700 Subject: [PATCH 09/19] fix(proxy): scope dashboard search grants to user (#48464) Co-authored-by: Jodok Batlogg Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 35 +++++- .../auth/test_search_tool_deny_by_default.py | 114 +++++++++++++++++- 2 files changed, 139 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 79e385c376a..448643ef59b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -6461,8 +6461,29 @@ def _search_tool_deny_by_default(general_settings: Mapping[str, object]) -> bool return True +def _is_dashboard_session(valid_token: UserAPIKeyAuth | None) -> bool: + return valid_token is not None and valid_token.team_id == UI_TEAM_ID + + +def _is_search_tool_strict_identity(valid_token: UserAPIKeyAuth) -> bool: + if not _is_dashboard_session(valid_token): + return _is_strict_grant_identity(valid_token) + return valid_token.user_role not in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ) + + +def _search_tool_strict_layers( + valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None +) -> frozenset[GrantLayer]: + if _is_dashboard_session(valid_token): + return frozenset(("user",)) + return _strict_grant_layers(True, valid_token, team_object) + + def is_search_tool_deny_by_default_applied(valid_token: UserAPIKeyAuth, general_settings: Mapping[str, object]) -> bool: - return _search_tool_deny_by_default(general_settings) and _is_strict_grant_identity(valid_token) + return _search_tool_deny_by_default(general_settings) and _is_search_tool_strict_identity(valid_token) def _search_tool_denied(object_type: GrantLayer, message: str) -> ProxyException: @@ -6506,12 +6527,12 @@ async def resolve_search_tool_grants( load_team_object: TeamObjectLoader, ) -> SearchToolGrants: """ - Without `general_settings.search_tool_deny_by_default`, or for the master key and dashboard sessions, + Without `general_settings.search_tool_deny_by_default`, or for the master key and admin dashboard sessions, the key and team lists only restrict when nonempty. With it, the identities `_strict_grant_layers` names must each list the tool, and a missing record, `null`, `[]`, or a team or user that fails to load grants nothing """ if not is_search_tool_deny_by_default_applied(valid_token, general_settings): - team_object: Final = await load_team_object() + team_object: Final = await load_team_object() if not _is_dashboard_session(valid_token) else None return SearchToolGrants( grants=MappingProxyType( {"key": valid_token.object_permission, "team": team_object.object_permission if team_object else None} @@ -6521,8 +6542,10 @@ async def resolve_search_tool_grants( from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - strict_team_object: Final = await _strict_team_object(load_team_object) - strict_layers: Final = _strict_grant_layers(True, valid_token, strict_team_object) + strict_team_object: Final = ( + await _strict_team_object(load_team_object) if not _is_dashboard_session(valid_token) else None + ) + strict_layers: Final = _search_tool_strict_layers(valid_token, strict_team_object) if prisma_client is None: return SearchToolGrants(grants=MappingProxyType({}), strict_layers=strict_layers) user_object: Final = ( @@ -6607,7 +6630,7 @@ def check_unregistered_search_fallback( """The unregistered provider fallback names no search tool that a grant could list, so deny-by-default denies it""" if not is_search_tool_deny_by_default_applied(valid_token, general_settings): return True - layer: Final = min(_strict_grant_layers(True, valid_token, None), key=("key", "team", "user").index) + layer: Final = min(_search_tool_strict_layers(valid_token, None), key=("key", "team", "user").index) raise _search_tool_denied( layer, "No registered search tool is available and search_tool_deny_by_default is enabled", diff --git a/tests/unit/proxy/auth/test_search_tool_deny_by_default.py b/tests/unit/proxy/auth/test_search_tool_deny_by_default.py index cd285e0386f..da036795889 100644 --- a/tests/unit/proxy/auth/test_search_tool_deny_by_default.py +++ b/tests/unit/proxy/auth/test_search_tool_deny_by_default.py @@ -1,8 +1,9 @@ from collections.abc import Mapping -from typing import Final, TypedDict +from typing import Final, TypedDict, cast from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy import proxy_server @@ -78,6 +79,14 @@ def _team_loader(team_tools: list[str] | None | str) -> TeamObjectLoader: return load +def _missing_dashboard_team_loader() -> tuple[TeamObjectLoader, AsyncMock]: + missing_team: Final = HTTPException( + status_code=404, detail="Team doesn't exist in db. Team=litellm-dashboard" + ) + load_team: Final = AsyncMock(side_effect=missing_team) + return cast(TeamObjectLoader, load_team), load_team + + async def _no_team() -> None: return None @@ -194,9 +203,6 @@ async def test_search_tool_access_is_unchanged_without_search_tool_deny_by_defau id="proxy admin virtual key not exempt", ), pytest.param(CallerFields(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS), None, None, None, id="master key exempt"), - pytest.param( - CallerFields(team_id=UI_TEAM_ID, virtual_key=False), None, None, None, id="dashboard session exempt" - ), ], ) @pytest.mark.asyncio @@ -248,6 +254,79 @@ async def test_search_tool_deny_by_default_does_not_load_the_user_when_off( get_user_object.assert_not_awaited() +@pytest.mark.parametrize( + "general_settings", + [{}, {"search_tool_deny_by_default": False}], + ids=["omitted", "false"], +) +@pytest.mark.parametrize( + "user_role", + [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.INTERNAL_USER], + ids=["admin", "internal-user"], +) +@pytest.mark.parametrize( + "key_tools, expected", + [ + pytest.param(None, None, id="unrestricted-key"), + pytest.param(["search-a"], None, id="key-grants-search-a"), + pytest.param(["search-b"], _LEGACY, id="key-excludes-search-a"), + ], +) +@pytest.mark.asyncio +async def test_dashboard_search_tool_grants_without_deny_by_default( + general_settings: Mapping[str, object], + user_role: LitellmUserRoles, + key_tools: list[str] | None, + expected: ProxyErrorTypes | None, +) -> None: + caller: Final = _caller( + virtual_key=False, + key_tools=key_tools, + team_id=UI_TEAM_ID, + user_role=user_role, + ) + load_team, load_team_mock = _missing_dashboard_team_loader() + assert await _denied_by(general_settings, caller, load_team) == expected + load_team_mock.assert_not_awaited() + + +@pytest.mark.parametrize( + "user_role", + [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], + ids=["admin", "admin-view-only"], +) +@pytest.mark.asyncio +async def test_admin_dashboard_search_tool_grants_with_deny_by_default(user_role: LitellmUserRoles) -> None: + caller: Final = _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=user_role) + load_team, load_team_mock = _missing_dashboard_team_loader() + assert await _denied_by(_DENY_ON, caller, load_team) is None + load_team_mock.assert_not_awaited() + + +@pytest.mark.parametrize( + "user_id, user_tools, expected", + [ + pytest.param("user-1", ["search-a"], None, id="user-grants-search-a"), + pytest.param("user-1", ["search-b"], _USER, id="user-grants-another-tool"), + pytest.param("user-1", None, _USER, id="user-grant-null"), + pytest.param("user-2", None, _USER, id="no-user-row"), + ], +) +@pytest.mark.asyncio +async def test_internal_dashboard_search_tool_requires_the_user_grant_with_deny_by_default( + cache: UserApiKeyCache, + user_id: str, + user_tools: list[str] | None, + expected: ProxyErrorTypes | None, +) -> None: + if user_id == "user-1": + _cache_user(cache, user_tools) + caller: Final = _caller(virtual_key=False, team_id=UI_TEAM_ID, user_id=user_id) + load_team, load_team_mock = _missing_dashboard_team_loader() + assert await _denied_by(_DENY_ON, caller, load_team) == expected + load_team_mock.assert_not_awaited() + + @pytest.mark.parametrize( "general_settings, caller, expected", [ @@ -257,6 +336,24 @@ async def test_search_tool_deny_by_default_does_not_load_the_user_when_off( pytest.param(_DENY_ON, _caller(virtual_key=False), _USER, id="keyless user"), pytest.param(_DENY_ON, _caller(user_role=LitellmUserRoles.PROXY_ADMIN), _KEY, id="proxy admin key"), pytest.param(_DENY_ON, _caller(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS), None, id="master key"), + pytest.param( + _DENY_ON, + _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.INTERNAL_USER), + _USER, + id="internal dashboard user", + ), + pytest.param( + _DENY_ON, + _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.PROXY_ADMIN), + None, + id="admin dashboard", + ), + pytest.param( + _DENY_ON, + _caller(virtual_key=False, team_id=UI_TEAM_ID, user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), + None, + id="admin view-only dashboard", + ), ], ) def test_unregistered_search_fallback_follows_search_tool_deny_by_default( @@ -268,3 +365,12 @@ def test_unregistered_search_fallback_follows_search_tool_deny_by_default( with pytest.raises(ProxyException) as exc_info: check_unregistered_search_fallback(caller, general_settings) assert (exc_info.value.type, exc_info.value.code) == (expected, "403") + + +@pytest.mark.asyncio +async def test_search_preserves_missing_real_team_error_without_deny_by_default() -> None: + missing_team: Final = HTTPException(status_code=404, detail="Team doesn't exist in db. Team=team-1") + load_team: Final = AsyncMock(side_effect=missing_team) + with pytest.raises(HTTPException) as exc_info: + await can_caller_call_search_tool("search-a", _caller(team_id="team-1"), {}, load_team) + assert exc_info.value is missing_team From 9a38ff4cb59fd60a5ca2bed9b0410783149a3e68 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 18:27:13 -0400 Subject: [PATCH 10/19] fix(gemini): preserve tool media for older models (#48428) --- .../prompt_templates/factory.py | 5 +- .../llms/vertex_ai/gemini/transformation.py | 4 + .../vertex_and_google_ai_studio_gemini.py | 4 + .../test_vertex_ai_gemini_transformation.py | 73 +++++++++++++++---- 4 files changed, 71 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0aed2956b0a..3bc8dc4e23e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1490,6 +1490,7 @@ def convert_to_gemini_tool_call_result( message: ChatCompletionToolMessage | ChatCompletionFunctionMessage, last_message_with_tool_calls: dict | None, forward_function_call_id: bool = False, + supports_multimodal_function_responses: bool = True, ) -> VertexPartType | list[VertexPartType]: """ OpenAI message with a tool result looks like: @@ -1652,9 +1653,9 @@ def convert_to_gemini_tool_call_result( _part: Final[VertexPartType] = {"function_response": _function_response} - # For multimodal function responses, Gemini expects media parts nested - # inside functionResponse.parts instead of sibling content parts. if inline_data_list: + if not supports_multimodal_function_responses: + return [_part, *(VertexPartType(inline_data=inline_data) for inline_data in inline_data_list)] _function_response["parts"] = [{"inline_data": inline_data} for inline_data in inline_data_list] return [_part] diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 71bb05ee155..ed86934db63 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -726,6 +726,9 @@ def gemini_convert_messages_with_history( from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig forward_function_call_id: Final = VertexGeminiConfig.forward_gemini_function_call_id(model or "") + supports_multimodal_function_responses: Final = model is None or ( + VertexGeminiConfig.supports_multimodal_function_responses(model) + ) try: while msg_i < len(messages): @@ -1029,6 +1032,7 @@ def gemini_convert_messages_with_history( messages[msg_i], last_message_with_tool_calls, forward_function_call_id=forward_function_call_id, + supports_multimodal_function_responses=supports_multimodal_function_responses, ) msg_i += 1 # Handle both single part and list of parts (for Computer Use with images) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d788ccb0f1d..9acbc9c459c 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -689,6 +689,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) -> bool: return cls._forward_gemini_function_call_id(model) + @classmethod + def supports_multimodal_function_responses(cls, model: str) -> bool: + return cls.is_gemini_3_or_newer(model) + def _supports_penalty_parameters(self, model: str) -> bool: # Gemini 3 models do not support penalty parameters if VertexGeminiConfig._is_gemini_3_or_newer(model): diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index b97724319c7..75449eacfd3 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -2,7 +2,7 @@ import base64 from collections.abc import Sequence from itertools import chain from pathlib import Path -from typing import Final +from typing import Final, Literal import httpx import pytest @@ -1650,10 +1650,28 @@ def test_convert_tool_response_with_base64_image(): assert inline_data["data"] == test_image_base64 -def test_gemini_history_nests_multimodal_tool_response_parts(): - """Full history conversion should not emit sibling inline_data tool result parts.""" - test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" - messages = [ +@pytest.mark.parametrize("provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) +@pytest.mark.parametrize( + "model,nested", + [ + (None, True), + ("gemini-2.5-pro", False), + ("gemini/gemini-2.5-flash", False), + ("vertex_ai/gemini-2.5-flash-lite", False), + ("123456789", False), + ("gemini-3.8-flash", True), + ("gemini/gemini-3.8-flash", True), + ("vertex_ai/gemini-3.8-flash", True), + ("gemini-flash-latest", True), + ], +) +def test_gemini_history_places_tool_media_for_model( + model: str | None, + nested: bool, + provider: Literal["gemini", "vertex_ai", "vertex_ai_beta"], +) -> None: + test_image_base64: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + messages: Final = [ {"role": "user", "content": "Get me an image"}, { "role": "assistant", @@ -1679,24 +1697,53 @@ def test_gemini_history_nests_multimodal_tool_response_parts(): "data": test_image_base64, }, }, + { + "type": "file", + "file": {"file_data": "data:application/pdf;base64,JVBERi0xLjQ="}, + }, ], }, + {"role": "user", "content": "Describe the image and file again"}, ] - contents = gemini_convert_messages_with_history(messages=messages) - - tool_response_parts = contents[-1]["parts"] - assert len(tool_response_parts) == 1 - assert "inline_data" not in tool_response_parts[0] - function_response = tool_response_parts[0]["function_response"] - assert function_response["parts"] == [ + contents: Final = ( + gemini_convert_messages_with_history(messages=messages) + if model is None + else transform_request_body( + messages=messages, + model=model, + custom_llm_provider=provider, + optional_params={}, + litellm_params={}, + cached_content=None, + )["contents"] + ) + media_parts: Final = [ { "inline_data": { "data": test_image_base64, "mime_type": "image/png", } - } + }, + {"inline_data": {"data": "JVBERi0xLjQ=", "mime_type": "application/pdf"}}, ] + function_response: Final = { + "name": "get_image", + "response": {"image_ref": "inline"}, + **({"id": "call_get_image"} if model is not None and nested else {}), + **({"parts": media_parts} if nested else {}), + } + assert contents[-2] == { + "role": "user", + "parts": [ + {"function_response": function_response}, + *([] if nested else media_parts), + ], + } + assert contents[-1] == { + "role": "user", + "parts": [{"text": "Describe the image and file again"}], + } def test_convert_tool_response_text_only(): From 57ac10731149fbc2a5612a811c1dc15f6780dac1 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:28:53 -0700 Subject: [PATCH 11/19] ci: remove the Together AI and weekly cost map sync workflows (#48473) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...to_update_price_and_context_window_file.py | 319 ---------- .../auto_update_price_and_context_window.yml | 36 -- .github/workflows/sync-together-ai-models.yml | 65 --- scripts/sync_together_ai_models.py | 547 ------------------ ...to_update_price_and_context_window_file.py | 222 ------- tests/unit/test_sync_together_ai_models.py | 378 ------------ 6 files changed, 1567 deletions(-) delete mode 100644 .github/scripts/auto_update_price_and_context_window_file.py delete mode 100644 .github/workflows/auto_update_price_and_context_window.yml delete mode 100644 .github/workflows/sync-together-ai-models.yml delete mode 100644 scripts/sync_together_ai_models.py delete mode 100644 tests/unit/test_auto_update_price_and_context_window_file.py delete mode 100644 tests/unit/test_sync_together_ai_models.py diff --git a/.github/scripts/auto_update_price_and_context_window_file.py b/.github/scripts/auto_update_price_and_context_window_file.py deleted file mode 100644 index 4c9a05bd4f0..00000000000 --- a/.github/scripts/auto_update_price_and_context_window_file.py +++ /dev/null @@ -1,319 +0,0 @@ -import asyncio -import aiohttp -import json -import math -from typing import Any - -# Asynchronously fetch data from a given URL -async def fetch_data(url): - try: - # Create an asynchronous session - async with aiohttp.ClientSession() as session: - # Send a GET request to the URL - async with session.get(url) as resp: - # Raise an error if the response status is not OK - resp.raise_for_status() - # Parse the response JSON - resp_json = await resp.json() - print("Fetch the data from URL.") - # Return the 'data' field from the JSON response - return resp_json['data'] - except Exception as e: - # Print an error message if fetching data fails - print("Error fetching data from URL:", e) - return None - - -FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models" -FRIENDLI_PROVIDER = "friendliai" - -INHERITABLE_BASE_KEYS = ( - "supports_pdf_input", - "supports_assistant_prefill", - "supports_adaptive_thinking", - "supports_output_config", -) - -REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max") - - -def _find_base_model_entry(base_model: str, local_data: dict) -> str | None: - if not base_model: - return None - bm_tail = base_model.split("/")[-1].lower() - if base_model in local_data: - return base_model - for key in local_data: - if key.startswith("sample_spec") or key == "fallback_generalizations": - continue - if key.split("/")[-1].lower() == bm_tail: - return key - return None - - -def _reasoning_effort_levels(reasoning_options: list) -> list: - offered = { - val - for opt in reasoning_options or [] - if opt.get("type") == "effort" - for val in opt.get("values", []) - } - return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered] - - -def _valid_token_price(value: object) -> bool: - try: - price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except - except (TypeError, ValueError): - return False - return math.isfinite(price) and price >= 0 - - -def _has_valid_token_prices(pricing: dict | None) -> bool: - prices = pricing or {} - return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output")) - - -def _pricing(pricing: dict) -> dict: - out: dict[str, Any] = {} - if not pricing: - return out - if "input" in pricing: - out["input_cost_per_token"] = float(pricing["input"]) - if "output" in pricing: - out["output_cost_per_token"] = float(pricing["output"]) - if "input_cache_read" in pricing and pricing["input_cache_read"] is not None: - out["cache_read_input_token_cost"] = float(pricing["input_cache_read"]) - return out - - -def _modality_flags(input_mods: list) -> dict: - mods = input_mods or [] - has_image = "image" in mods - return { - "supports_vision": has_image, - "supports_image_input": has_image, - "supports_video_input": "video" in mods, - } - - -def transform_friendli_data(data: list, local_data: dict) -> dict: - transformed: dict[str, dict] = {} - if not data: - return transformed - for model in data: - # An unpriced row must never wholesale-replace an already priced local entry: - # missing prices cost-calculate as zero, silently zeroing tracked spend - if not _has_valid_token_prices(model.get("pricing")): - continue - model_id = model["id"] - base_model = model.get("base_model") or "" - entry: dict[str, Any] = { - "litellm_provider": FRIENDLI_PROVIDER, - } - - base_key = _find_base_model_entry(base_model, local_data) - if base_key: - base_entry = local_data[base_key] - for k in INHERITABLE_BASE_KEYS: - if k in base_entry: - entry[k] = base_entry[k] - - ctx = model.get("context_length") - if ctx is not None: - entry["max_input_tokens"] = int(ctx) - max_out = model.get("max_completion_tokens") - if max_out is not None: - entry["max_output_tokens"] = int(max_out) - entry["max_tokens"] = int(max_out) - - pricing = _pricing(model.get("pricing", {})) - entry.update(pricing) - entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing - - reasoning = model.get("reasoning") is True - entry["supports_reasoning"] = reasoning - if reasoning: - entry["reasoning_effort_levels"] = _reasoning_effort_levels( - model.get("reasoning_options", []) - ) - - func = model.get("functionality", {}) - entry["supports_function_calling"] = func.get("tool_call") is True - entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True - is_struct = func.get("structured_output") is True - entry["supports_response_schema"] = is_struct - entry["supports_native_structured_output"] = is_struct - entry["supports_system_messages"] = func.get("system_messages") is True - entry["supports_tool_choice"] = func.get("tool_choice") is True - - entry.update(_modality_flags(model.get("input_modalities", []))) - - entry["mode"] = model.get("mode", "chat") - - desc = model.get("description") - if desc: - entry["comment"] = desc - - dep = model.get("deprecation_date") - if dep: - entry["deprecation_date"] = dep.split("T")[0] - - entry["source"] = FRIENDLI_API_URL - - transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry - return transformed - -# Synchronize local data with remote data -def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()): - # Update existing keys in local_data with values from remote_data - # (replace_keys entries are swapped wholesale so a field the remote catalog - # dropped, e.g. cache pricing, cannot survive as a stale value) - for key in (set(local_data) & set(remote_data)): - if key in replace_keys: - local_data[key] = remote_data[key] - else: - local_data[key].update(remote_data[key]) - - # Add new keys from remote_data to local_data - for key in (set(remote_data) - set(local_data)): - local_data[key] = remote_data[key] - -# Write data to the json file -def write_to_file(file_path, data): - try: - # Open the file in write mode - with open(file_path, "w") as file: - # Dump the data as JSON into the file - json.dump(data, file, indent=4) - print("Values updated successfully.") - except Exception as e: - # Print an error message if writing to file fails - print("Error updating JSON file:", e) - -# Update the existing models and add the missing models for OpenRouter -def transform_openrouter_data(data): - transformed = {} - if not data: - return transformed - for row in data: - # Add the fields 'max_tokens' and 'input_cost_per_token' - obj = { - "max_tokens": row["context_length"], - "input_cost_per_token": float(row["pricing"]["prompt"]), - } - - # Add 'max_output_tokens' as a field if it is not None - if "top_provider" in row and "max_completion_tokens" in row["top_provider"] and row["top_provider"]["max_completion_tokens"] is not None: - obj['max_output_tokens'] = int(row["top_provider"]["max_completion_tokens"]) - - # Add the field 'output_cost_per_token' - obj.update({ - "output_cost_per_token": float(row["pricing"]["completion"]), - }) - - # Add field 'input_cost_per_image' if it exists and is non-zero - if "pricing" in row and "image" in row["pricing"] and float(row["pricing"]["image"]) != 0.0: - obj['input_cost_per_image'] = float(row["pricing"]["image"]) - - # Add the fields 'litellm_provider' and 'mode' - obj.update({ - "litellm_provider": "openrouter", - "mode": "chat" - }) - - # Add the 'supports_vision' field if the modality is 'multimodal' - if row.get('architecture', {}).get('modality') == 'multimodal': - obj['supports_vision'] = True - - # Use a composite key to store the transformed object - transformed[f'openrouter/{row["id"]}'] = obj - - return transformed - -# Update the existing models and add the missing models for Vercel AI Gateway -def transform_vercel_ai_gateway_data(data): - transformed = {} - if not data: - return transformed - for row in data: - # Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync - if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any( - row.get("pricing", {}).get(k) is None for k in ("input", "output") - ): - continue - obj = { - "max_tokens": row["context_window"], - "input_cost_per_token": float(row["pricing"]["input"]), - "output_cost_per_token": float(row["pricing"]["output"]), - 'max_output_tokens': row['max_tokens'], - 'max_input_tokens': row["context_window"], - } - - # Handle cache pricing if available - if "pricing" in row: - if "input_cache_read" in row["pricing"] and row["pricing"]["input_cache_read"] is not None: - obj['cache_read_input_token_cost'] = float(f"{float(row['pricing']['input_cache_read']):e}") - - if "input_cache_write" in row["pricing"] and row["pricing"]["input_cache_write"] is not None: - obj['cache_creation_input_token_cost'] = float(f"{float(row['pricing']['input_cache_write']):e}") - - mode = "embedding" if "embedding" in row["id"].lower() else "chat" - - obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode}) - - transformed[f'vercel_ai_gateway/{row["id"]}'] = obj - - return transformed - - -# Load local data from a specified file -def load_local_data(file_path): - try: - # Open the file in read mode - with open(file_path, "r") as file: - # Load and return the JSON data - return json.load(file) - except FileNotFoundError: - # Print an error message if the file is not found - print("File not found:", file_path) - return None - except json.JSONDecodeError as e: - # Print an error message if JSON decoding fails - print("Error decoding JSON:", e) - return None - -def main(): - local_file_path = "model_prices_and_context_window.json" # Path to the local data file - openrouter_url = "https://openrouter.ai/api/v1/models" # URL to fetch OpenRouter data - vercel_ai_gateway_url = "https://ai-gateway.vercel.sh/v1/models" # URL to fetch Vercel AI Gateway data - - # Load local data from file - local_data = load_local_data(local_file_path) - - # Fetch OpenRouter data - openrouter_data = asyncio.run(fetch_data(openrouter_url)) - # Transform the fetched OpenRouter data - openrouter_data = transform_openrouter_data(openrouter_data) - - # Fetch Vercel AI Gateway data - vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url)) - # Transform the fetched Vercel AI Gateway data - vercel_data = transform_vercel_ai_gateway_data(vercel_data) - - friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL)) - friendli_data = transform_friendli_data(friendli_data, local_data) - - # Combine both datasets - all_remote_data = {**openrouter_data, **vercel_data, **friendli_data} - - # If both local and openrouter data are available, synchronize and save - if local_data and all_remote_data: - sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data)) - write_to_file(local_file_path, local_data) - else: - print("Failed to fetch model data from either local file or URL.") - -# Entry point of the script -if __name__ == "__main__": - main() diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml deleted file mode 100644 index 838ba35f831..00000000000 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ /dev/null @@ -1,36 +0,0 @@ -name: Updates model_prices_and_context_window.json and Create Pull Request - -on: - schedule: - - cron: "0 0 * * 0" # Run every Sundays at midnight - #- cron: "0 0 * * *" # Run daily at midnight - -permissions: - contents: write - pull-requests: write - -jobs: - auto_update_price_and_context_window: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - name: Update JSON Data - run: | - uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py" - - name: Create Pull Request - run: | - git add model_prices_and_context_window.json - git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')" - gh pr create --title "Update model_prices_and_context_window.json file" \ - --body "Automated update for model_prices_and_context_window.json" \ - --head auto-update-price-and-context-window-$(date +'%Y-%m-%d') \ - --base main - env: - GH_TOKEN: ${{ secrets.GH_TOKEN }} diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml deleted file mode 100644 index 5325510e45e..00000000000 --- a/.github/workflows/sync-together-ai-models.yml +++ /dev/null @@ -1,65 +0,0 @@ -name: Sync Together AI model registry - -on: - schedule: - - cron: "30 6 * * *" - workflow_dispatch: - -permissions: - contents: write - pull-requests: write - -jobs: - sync_together_ai_models: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - env: - BASE_BRANCH: ${{ github.event.repository.default_branch }} - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - ref: ${{ env.BASE_BRANCH }} - persist-credentials: false - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - name: Look for an already-open sync PR - id: existing - run: | - open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \ - --jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')" - echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT" - if [ -n "$open_pr" ]; then - echo "An open sync PR already exists on branch $open_pr; skipping this run." - fi - env: - GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} - - name: Run the sync - if: steps.existing.outputs.open_pr == '' - run: | - uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md" - env: - TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }} - - name: Create a pull request when the registry changed - if: steps.existing.outputs.open_pr == '' - run: | - if git diff --quiet; then - echo "Registry already in sync; no PR needed." - exit 0 - fi - branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')" - git config user.name "github-actions[bot]" - git config user.email "41898282+github-actions[bot]@users.noreply.github.com" - git checkout -b "$branch" - git add model_prices_and_context_window.json \ - litellm/model_prices_and_context_window_backup.json - git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')" - gh auth setup-git - git push origin "$branch" - gh pr create --title "feat(models): sync together_ai model registry" \ - --body-file "$RUNNER_TEMP/pr_body.md" \ - --head "$branch" \ - --base "$BASE_BRANCH" - env: - GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} diff --git a/scripts/sync_together_ai_models.py b/scripts/sync_together_ai_models.py deleted file mode 100644 index 12b128890f1..00000000000 --- a/scripts/sync_together_ai_models.py +++ /dev/null @@ -1,547 +0,0 @@ -"""Sync the together_ai entries of model_prices_and_context_window.json with Together's live serverless catalog. - -Pulls ``GET https://api.together.ai/v1/models?serverless`` plus the deprecations doc, maps API fields onto -registry fields, merges the reviewed capability rules below for everything the API cannot express, and diffs -the result against the registry. Dry run (the default) prints the diff summary and the generated PR body; -``--write`` applies the changes to the root cost map and its ``litellm/`` backup copy. - -Policy highlights: -- Prices arrive per 1M tokens with float artifacts and are normalized to clean per-token values. -- A registry entry absent from the serverless catalog is marked with ``deprecation_date`` from the docs - deprecation table, never deleted; absences with no docs date are surfaced for a human call. -- Availability comes from the API: a model the docs list as removed but the API still serves stays live, - with the conflict surfaced as a warning. -- Manually curated values the API cannot express (``metadata.successor``, ``max_output_tokens`` on existing - entries, capability flags no rule covers) are never overwritten; conflicts are surfaced instead. -""" - -import argparse -import json -import os -import re -import sys -from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field -from pathlib import Path -from types import MappingProxyType -from typing import Final - -import httpx -from pydantic import BaseModel, TypeAdapter, ValidationError - -MODELS_URL: Final = "https://api.together.ai/v1/models?serverless" -DEPRECATIONS_URL: Final = "https://docs.together.ai/docs/deprecations.md" -PROVIDER: Final = "together_ai" -PREFIX: Final = "together_ai/" -SOURCE_URL: Final = "https://docs.together.ai/docs/serverless-models" -COST_MAP_RELPATHS: Final = ( - "model_prices_and_context_window.json", - "litellm/model_prices_and_context_window_backup.json", -) - -TYPE_TO_MODE: Final = MappingProxyType({"chat": "chat", "embedding": "embedding", "moderation": "chat"}) - - -class SyncError(RuntimeError): - pass - - -class CatalogPricing(BaseModel): - input: float - output: float - cached_input: float | None = None - - -class CatalogModel(BaseModel): - id: str - type: str - context_length: int | None = None - pricing: CatalogPricing - - -CATALOG_ADAPTER: Final = TypeAdapter(list[CatalogModel]) - -RegistryEntry = dict[str, object] -CostMap = dict[str, object] - - -@dataclass(frozen=True, slots=True) -class CapabilityRule: - model_id: str - fields: Mapping[str, bool | int] - provenance: str - - -def _rule(model_id: str, provenance: str, **fields: bool | int) -> CapabilityRule: - return CapabilityRule(model_id=model_id, fields=MappingProxyType(dict(fields)), provenance=provenance) - - -_TOOLS: Final = MappingProxyType( - { - "supports_function_calling": True, - "supports_parallel_function_calling": True, - "supports_response_schema": True, - "supports_tool_choice": True, - } -) - -CAPABILITY_RULES: Final = ( - _rule( - "MiniMaxAI/MiniMax-M3", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/minimax-m3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule("Prism-ML/Ternary-Bonsai-27B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "Qwen/Qwen3.5-9B", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/qwen3-5-9b", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule( - "Qwen/Qwen3.6-Plus", - "reviewed for the LIT-5968 backfill; hybrid reasoning model without a documented tools contract", - supports_reasoning=True, - ), - _rule("Qwen/Qwen3.7-Max", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("Qwen/Qwen3.7-Plus", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("Qwen/Qwen3.8-2.4T-A95B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule("arize-ai/qwen-2-1.5b-instruct", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "deepseek-ai/DeepSeek-V4-Flash-0731", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-flash", - **_TOOLS, - ), - _rule( - "deepseek-ai/DeepSeek-V4-Pro", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "deepseek-ai/DeepSeek-V4-Pro-0813", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro", - **_TOOLS, - ), - _rule("google/gemma-3n-E4B-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "google/gemma-4-31B-it", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gemma-4-31b-it", - **_TOOLS, - supports_vision=True, - ), - _rule( - "intfloat/multilingual-e5-large-instruct", - "embedding dims per https://huggingface.co/intfloat/multilingual-e5-large-instruct", - output_vector_size=1024, - ), - _rule( - "meta-llama/Llama-3.3-70B-Instruct-Turbo", - "reviewed for the LIT-5968 backfill against https://docs.together.ai/docs/function-calling", - **_TOOLS, - ), - _rule( - "meta-llama/Llama-Guard-4-12B", - "moderation classifier with a chat-shaped API; no tools per the LIT-5968 backfill review", - ), - _rule("meta-models/Muse-Glimmer-30B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "moonshotai/Kimi-K2.7-Code", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k2-7-code", - **_TOOLS, - supports_vision=True, - ), - _rule( - "moonshotai/Kimi-K3", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - ), - _rule( - "nvidia/nemotron-3-ultra-550b-a55b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/nemotron-3-ultra", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "openai/gpt-oss-120b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-120b", - **_TOOLS, - supports_reasoning=True, - ), - _rule( - "openai/gpt-oss-20b", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-20b", - **_TOOLS, - ), - _rule("pearl-ai/gemma-4-31b-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "thinkingmachines/Inkling", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/inkling", - **_TOOLS, - ), - _rule("thinkingmachines/Inkling-Small", "reviewed for the LIT-5968 backfill; no tool or vision support documented"), - _rule( - "zai-org/GLM-5.2", - "reviewed for the LIT-5968 backfill against https://www.together.ai/models/glm-5-2;" - " 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.2", - **_TOOLS, - supports_reasoning=True, - max_output_tokens=128000, - max_tokens=128000, - ), - _rule( - "zai-org/GLM-5.3-Flash", - "reviewed for LIT-6489 against https://www.together.ai/models/glm-5-3-flash;" - " 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.3", - **_TOOLS, - supports_reasoning=True, - supports_vision=True, - max_output_tokens=128000, - max_tokens=128000, - ), -) - -RULES_BY_ID: Final = MappingProxyType({rule.model_id: rule for rule in CAPABILITY_RULES}) - - -@dataclass(frozen=True, slots=True) -class DeprecationDoc: - removal_dates: Mapping[str, str] - redirects: Mapping[str, str] - - -_REDIRECT_ROW: Final = re.compile(r"^\|\s*`([^`]+)`\s*\|\s*`([^`]+)`\s*\|") -_REMOVAL_ROW: Final = re.compile(r"^\|\s*(\d{4}-\d{2}-\d{2})\s*\|\s*`([^`]+)`\s*\|") - - -def _section(markdown: str, heading: str) -> str: - level: Final = heading.split(" ", 1)[0] - start: Final = markdown.find(f"\n{heading}\n") - if start < 0: - return "" - body: Final = markdown[start + 1 + len(heading) :] - next_heading: Final = re.search(rf"^{re.escape(level)} ", body, flags=re.MULTILINE) - return body[: next_heading.start()] if next_heading else body - - -def parse_deprecations(markdown: str) -> DeprecationDoc: - redirect_rows: Final = tuple( - m.groups() - for m in (_REDIRECT_ROW.match(line) for line in _section(markdown, "## Active model redirects").splitlines()) - if m - ) - inference: Final = _section(_section(markdown, "## Deprecation history"), "### Inference") - removal_rows: Final = tuple(m.groups() for m in (_REMOVAL_ROW.match(line) for line in inference.splitlines()) if m) - if not redirect_rows or not removal_rows: - raise SyncError( - "deprecations doc parsed to zero redirect or removal rows; the table format at " - f"{DEPRECATIONS_URL} changed and the parser needs updating" - ) - removal_dates: Final = {model: date for date, model in reversed(removal_rows)} - return DeprecationDoc( - removal_dates=MappingProxyType(dict(reversed(removal_dates.items()))), - redirects=MappingProxyType({original: target for original, target in redirect_rows}), - ) - - -def per_token(price_per_million: float) -> float: - return float(f"{price_per_million / 1e6:.6g}") - - -def _resolve_name(name: str, universe: frozenset[str]) -> str | None: - if name in universe: - return name - suffix_matches: Final = tuple(candidate for candidate in universe if candidate.endswith(f"/{name}")) - return suffix_matches[0] if len(suffix_matches) == 1 else None - - -def resolve_successor(model_id: str, doc: DeprecationDoc, live_ids: frozenset[str]) -> str | None: - canonical: Final = live_ids | frozenset(doc.removal_dates) - redirects: Final = { - (_resolve_name(raw_source, canonical) or raw_source): (_resolve_name(raw_target, canonical) or raw_target) - for raw_source, raw_target in doc.redirects.items() - } - seen: Final = set() - current = model_id # rebind-ok: walks the redirect chain - while current in redirects and current not in seen: - seen.add(current) - current = redirects[current] # rebind-ok: walks the redirect chain - return current if current != model_id and current in live_ids else None - - -@dataclass(frozen=True, slots=True) -class SyncOutcome: - cost_map: CostMap - added: tuple[str, ...] = () - updated: tuple[str, ...] = () - deprecated: tuple[str, ...] = () - reappeared: tuple[str, ...] = () - warnings: tuple[str, ...] = () - skipped_types: Mapping[str, int] = field(default_factory=dict) - - @property - def has_changes(self) -> bool: - return bool(self.added or self.updated or self.deprecated or self.reappeared) - - -def _api_fields(model: CatalogModel) -> RegistryEntry: - cached: Final = model.pricing.cached_input - return { - "input_cost_per_token": per_token(model.pricing.input), - "output_cost_per_token": per_token(model.pricing.output), - **({"cache_read_input_token_cost": per_token(cached), "supports_prompt_caching": True} if cached else {}), - **({"max_input_tokens": model.context_length} if model.context_length is not None else {}), - } - - -def _new_entry(model: CatalogModel, mode: str) -> RegistryEntry: - rule: Final = RULES_BY_ID.get(model.id) - legacy_ceiling: Final = {} if model.context_length is None else {"max_tokens": model.context_length} - merged: Final = { - **_api_fields(model), - **legacy_ceiling, - "litellm_provider": PROVIDER, - "mode": mode, - "source": SOURCE_URL, - **(dict(rule.fields) if rule else {}), - } - return dict(sorted(merged.items())) - - -def _updated_entry(entry: RegistryEntry, model: CatalogModel) -> tuple[RegistryEntry, tuple[str, ...]]: - rule: Final = RULES_BY_ID.get(model.id) - desired: Final = {**_api_fields(model), **(dict(rule.fields) if rule else {})} - dropped: Final = () if model.pricing.cached_input else ("cache_read_input_token_cost", "supports_prompt_caching") - changes: Final = tuple( - f"{name}: {entry.get(name)!r} -> {value!r}" for name, value in desired.items() if entry.get(name) != value - ) + tuple( - f"{name}: {entry[name]!r} removed (no longer in the catalog pricing)" for name in dropped if name in entry - ) - merged: Final = {name: value for name, value in {**entry, **desired}.items() if name not in dropped} - return dict(sorted(merged.items())), changes - - -def _with_new_keys_in_block(original: CostMap, result: CostMap, new_keys: Sequence[str]) -> CostMap: - provider_keys: Final = tuple(key for key in original if key.startswith(PREFIX)) - if not new_keys or not provider_keys: - return result - block_end: Final = provider_keys[-1] - return { - key: value - for existing in original - for key, value in ( - (existing, result[existing]), - *((new, result[new]) for new in sorted(new_keys) if existing == block_end), - ) - } - - -def compute_sync(cost_map: CostMap, catalog: Sequence[CatalogModel], doc: DeprecationDoc) -> SyncOutcome: - live_ids: Final = frozenset(model.id for model in catalog) - token_models: Final = {model.id: model for model in catalog if model.type in TYPE_TO_MODE} - skipped: Final = { - model.type: sum(1 for m in catalog if m.type == model.type) - for model in catalog - if model.type not in TYPE_TO_MODE - } - registry_ids: Final = {key.removeprefix(PREFIX): key for key in cost_map if key.startswith(PREFIX)} - - added: Final[list[str]] = [] - updated: Final[list[str]] = [] - deprecated: Final[list[str]] = [] - reappeared: Final[list[str]] = [] - warnings: Final[list[str]] = [] - result: Final[CostMap] = dict(cost_map) - - for model_id, model in sorted(token_models.items()): - mode: Final = TYPE_TO_MODE[model.type] - key: Final = f"{PREFIX}{model_id}" - if model_id in doc.removal_dates: - warnings.append( - f"`{key}` is listed as removed on {doc.removal_dates[model_id]} in the docs but the serverless " - "catalog still serves it; availability kept from the API" - ) - entry = result.get(key) - if not isinstance(entry, dict): - result[key] = _new_entry(model, mode) - added.append(key) - if model.type == "chat" and model_id not in RULES_BY_ID: - warnings.append( - f"`{key}` added without a capability rule; review its tools/vision/reasoning support and add one" - ) - continue - if entry.get("mode") != mode: - warnings.append( - f"`{key}` has curated mode {entry.get('mode')!r} but the catalog maps to {mode!r}; left unchanged" - ) - new_entry, changes = _updated_entry(entry, model) - if "deprecation_date" in new_entry: - new_entry.pop("deprecation_date") - reappeared.append(key) - if changes: - updated.append(f"{key}: " + "; ".join(changes)) - if changes or key in reappeared: - result[key] = new_entry - - for model_id, key in sorted(registry_ids.items()): - if model_id in token_models: - continue - entry = result.get(key) - if not isinstance(entry, dict): - continue - removal_date: Final = doc.removal_dates.get(model_id) - successor: Final = resolve_successor(model_id, doc, live_ids) - metadata = entry.get("metadata") - curated_successor: Final = metadata.get("successor") if isinstance(metadata, dict) else None - new_entry = dict(entry) - if removal_date is not None and entry.get("deprecation_date") != removal_date: - if "deprecation_date" in entry: - warnings.append( - f"`{key}` has curated deprecation_date {entry.get('deprecation_date')!r} but the docs list " - f"{removal_date!r}; left unchanged" - ) - else: - new_entry["deprecation_date"] = removal_date - if removal_date is None and "deprecation_date" not in entry: - warnings.append( - f"`{key}` is absent from the serverless catalog with no removal date in the docs; " - "needs a human deprecation call" - ) - if successor is not None: - desired_successor: Final = f"{PREFIX}{successor}" - if curated_successor is None: - new_entry["metadata"] = dict( - sorted({**(metadata if isinstance(metadata, dict) else {}), "successor": desired_successor}.items()) - ) - elif curated_successor != desired_successor: - warnings.append( - f"`{key}` has curated successor {curated_successor!r} but the docs redirects resolve to " - f"{desired_successor!r}; left unchanged" - ) - if new_entry != entry: - result[key] = dict(sorted(new_entry.items())) - deprecated.append(f"{key}: " + ", ".join(sorted(set(new_entry) - set(entry)) or ["updated"])) - - return SyncOutcome( - cost_map=_with_new_keys_in_block(cost_map, result, tuple(added)), - added=tuple(added), - updated=tuple(updated), - deprecated=tuple(deprecated), - reappeared=tuple(reappeared), - warnings=tuple(warnings), - skipped_types=MappingProxyType(skipped), - ) - - -def _section_block(title: str, lines: Sequence[str], backtick: bool) -> str: - bullets: Final = "\n".join(f"- `{line}`" if backtick else f"- {line}" for line in lines) or "- none" - return f"### {title} ({len(lines)})\n{bullets}\n" - - -def render_pr_body(outcome: SyncOutcome) -> str: - skipped: Final = ", ".join(f"{kind} ({count})" for kind, count in sorted(outcome.skipped_types.items())) or "none" - return ( - "Automated daily sync of the together_ai entries in model_prices_and_context_window.json against " - f"`GET {MODELS_URL}` and {DEPRECATIONS_URL} by scripts/sync_together_ai_models.py.\n" - "\n" - f"{_section_block('Added', outcome.added, backtick=True)}" - "\n" - f"{_section_block('Updated', outcome.updated, backtick=True)}" - "\n" - f"{_section_block('Marked deprecated', outcome.deprecated, backtick=True)}" - "\n" - f"{_section_block('Returned to the catalog', outcome.reappeared, backtick=True)}" - "\n" - f"{_section_block('Warnings needing a human call', outcome.warnings, backtick=False)}" - "\n" - f"Catalog model types outside the sync's token-pricing scope, skipped: {skipped}\n" - ) - - -def render_summary(outcome: SyncOutcome) -> str: - return ( - f"added={len(outcome.added)} updated={len(outcome.updated)} deprecated={len(outcome.deprecated)} " - f"reappeared={len(outcome.reappeared)} warnings={len(outcome.warnings)}" - ) - - -def load_catalog(raw: bytes) -> list[CatalogModel]: - parsed: Final = json.loads(raw) - entries: Final = parsed.get("data") if isinstance(parsed, dict) else parsed - try: - catalog: Final = CATALOG_ADAPTER.validate_python(entries) - except ValidationError as error: - raise SyncError(f"the catalog response no longer matches the expected shape: {error}") from error - if not any(model.type in TYPE_TO_MODE for model in catalog): - raise SyncError( - "the catalog response contains no token-priced models; refusing to mark the whole registry deprecated" - ) - return catalog - - -def _fetch(url: str, headers: Mapping[str, str]) -> bytes: - response: Final = httpx.get(url, headers=dict(headers), timeout=30, follow_redirects=True) - if response.status_code != 200: - raise SyncError(f"GET {url} returned {response.status_code}") - return response.content - - -def _serialize(cost_map: CostMap) -> str: - return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n" - - -def main(argv: Sequence[str]) -> int: - parser: Final = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--write", action="store_true", help="apply the sync to the cost map files (default: dry run)") - parser.add_argument("--models-json", type=Path, help="recorded catalog response to use instead of the live API") - parser.add_argument( - "--deprecations-md", type=Path, help="recorded deprecations doc to use instead of the live docs" - ) - parser.add_argument("--pr-body-file", type=Path, help="write the generated PR body to this path") - parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parent.parent) - args: Final = parser.parse_args(argv) - - if args.models_json is not None: - catalog_raw: Final = args.models_json.read_bytes() - else: - api_key: Final = os.environ.get("TOGETHER_API_KEY") - if not api_key: - raise SyncError("TOGETHER_API_KEY is not set and --models-json was not given") - catalog_raw = _fetch(MODELS_URL, {"Authorization": f"Bearer {api_key}"}) # rebind-ok: branch-dependent source - catalog: Final = load_catalog(catalog_raw) - markdown: Final = ( - args.deprecations_md.read_text() if args.deprecations_md is not None else _fetch(DEPRECATIONS_URL, {}).decode() - ) - doc: Final = parse_deprecations(markdown) - - cost_map_path: Final = args.repo_root / COST_MAP_RELPATHS[0] - cost_map: Final = json.loads(cost_map_path.read_text()) - outcome: Final = compute_sync(cost_map, catalog, doc) - body: Final = render_pr_body(outcome) - - if args.pr_body_file is not None: - args.pr_body_file.write_text(body) - if args.write and outcome.has_changes: - for relpath in COST_MAP_RELPATHS: - (args.repo_root / relpath).write_text(_serialize(outcome.cost_map)) - print(render_summary(outcome)) - print() - print(body) - if not args.write: - print("dry run: no files were touched") - elif not outcome.has_changes: - print("registry already in sync: no files were touched") - return 0 - - -if __name__ == "__main__": - try: - raise SystemExit(main(sys.argv[1:])) - except SyncError as error: - print(f"SYNC FAILED: {error}", file=sys.stderr) - raise SystemExit(1) from error diff --git a/tests/unit/test_auto_update_price_and_context_window_file.py b/tests/unit/test_auto_update_price_and_context_window_file.py deleted file mode 100644 index 435747e9a09..00000000000 --- a/tests/unit/test_auto_update_price_and_context_window_file.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Unit tests for the Friendli transform in -`.github/scripts/auto_update_price_and_context_window_file.py`.""" - -from __future__ import annotations - -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] - / ".github" - / "scripts" - / "auto_update_price_and_context_window_file.py" -) - - -@pytest.fixture(scope="module") -def sync_module(): - spec = importlib.util.spec_from_file_location( - "auto_update_price_and_context_window_file", SCRIPT_PATH - ) - assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" - module = importlib.util.module_from_spec(spec) - sys.modules["auto_update_price_and_context_window_file"] = module - spec.loader.exec_module(module) - return module - - -def _reasoning_model(**overrides: object) -> dict: - model = { - "id": "zai-org/GLM-Test", - "base_model": "zhipuai/glm-test", - "context_length": 1048576, - "max_completion_tokens": 131072, - "pricing": {"input": "0.00000015", "output": "0.0000005", "input_cache_read": "0.00000003"}, - "reasoning": True, - "reasoning_options": [{"type": "effort", "values": ["max", "high", "low"]}], - "functionality": { - "tool_call": True, - "parallel_tool_call": True, - "structured_output": True, - "system_messages": True, - "tool_choice": True, - }, - "input_modalities": ["text", "image", "video"], - "mode": "chat", - } - model.update(overrides) - return model - - -def test_transform_emits_declared_effort_levels_in_canonical_order(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_reasoning"] is True - assert entry["reasoning_effort_levels"] == ["low", "high", "max"] - assert not any(k.endswith("_reasoning_effort") for k in entry) - - -def test_transform_reasoning_model_without_effort_options_declares_empty_levels(sync_module): - model = _reasoning_model(reasoning_options=[{"type": "budget_tokens", "values": []}]) - entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"] - assert entry["reasoning_effort_levels"] == [] - - -def test_transform_non_reasoning_model_declares_no_levels(sync_module): - model = _reasoning_model(reasoning=False, reasoning_options=[]) - entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"] - assert entry["supports_reasoning"] is False - assert "reasoning_effort_levels" not in entry - - -def test_transform_max_tokens_mirrors_output_cap_not_context(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["max_input_tokens"] == 1048576 - assert entry["max_output_tokens"] == 131072 - assert entry["max_tokens"] == entry["max_output_tokens"] - - -def test_transform_prompt_caching_follows_cache_pricing(sync_module): - cached = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert cached["supports_prompt_caching"] is True - assert cached["cache_read_input_token_cost"] == 3e-08 - - uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"}) - uncached = sync_module.transform_friendli_data([uncached_model], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert uncached["supports_prompt_caching"] is False - assert "cache_read_input_token_cost" not in uncached - - -def test_transform_modalities_set_vision_image_and_video_flags(sync_module): - entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_vision"] is True - assert entry["supports_image_input"] is True - assert entry["supports_video_input"] is True - - text_only = _reasoning_model(input_modalities=["text"]) - entry_text = sync_module.transform_friendli_data([text_only], {})[ - "friendliai/zai-org/GLM-Test" - ] - assert entry_text["supports_vision"] is False - assert entry_text["supports_image_input"] is False - assert entry_text["supports_video_input"] is False - - -def test_transform_skips_rows_without_valid_token_prices_so_priced_local_entries_survive(sync_module): - local = { - "friendliai/zai-org/GLM-Test": { - "litellm_provider": "friendliai", - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - } - } - unpriced_rows = [ - _reasoning_model(pricing={}), - _reasoning_model(pricing=None), - _reasoning_model(pricing={"input": "0.00000015"}), - _reasoning_model(pricing={"output": "0.0000005"}), - _reasoning_model(pricing={"input": "not-a-number", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "-0.00000015", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "inf", "output": "0.0000005"}), - _reasoning_model(pricing={"input": "nan", "output": "0.0000005"}), - ] - remote = sync_module.transform_friendli_data(unpriced_rows, local) - assert remote == {} - sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote)) - assert local["friendliai/zai-org/GLM-Test"]["input_cost_per_token"] == 1.5e-07 - assert local["friendliai/zai-org/GLM-Test"]["output_cost_per_token"] == 5e-07 - - -def test_transform_keeps_zero_priced_rows(sync_module): - free_model = _reasoning_model(pricing={"input": "0", "output": "0"}) - entry = sync_module.transform_friendli_data([free_model], {})["friendliai/zai-org/GLM-Test"] - assert entry["input_cost_per_token"] == 0.0 - assert entry["output_cost_per_token"] == 0.0 - - -def test_transforms_survive_failed_fetch(sync_module): - assert sync_module.transform_friendli_data(None, {}) == {} - assert sync_module.transform_friendli_data([], {}) == {} - assert sync_module.transform_openrouter_data(None) == {} - assert sync_module.transform_vercel_ai_gateway_data(None) == {} - - -def test_vercel_transform_skips_rows_without_token_pricing_or_limits(sync_module): - rows = [ - { - "id": "wan-video", - "pricing": {"video_duration_pricing": [{"resolution": "720p", "cost_per_second": "0.1"}]}, - }, - { - "id": "qwen3-embedding", - "context_window": 32768, - "max_tokens": 32768, - "pricing": {"input": "0.00000001"}, - }, - { - "id": "no-limits-chat", - "pricing": {"input": "0.000001", "output": "0.000002"}, - }, - { - "id": "good-chat", - "context_window": 128000, - "max_tokens": 8192, - "pricing": {"input": "0.000001", "output": "0.000002"}, - }, - ] - transformed = sync_module.transform_vercel_ai_gateway_data(rows) - assert list(transformed) == ["vercel_ai_gateway/good-chat"] - assert transformed["vercel_ai_gateway/good-chat"]["input_cost_per_token"] == 1e-06 - assert transformed["vercel_ai_gateway/good-chat"]["output_cost_per_token"] == 2e-06 - - -def test_sync_replaces_friendli_entries_so_dropped_cache_pricing_does_not_survive(sync_module): - local = { - "friendliai/zai-org/GLM-Test": { - "litellm_provider": "friendliai", - "cache_read_input_token_cost": 3e-08, - "supports_prompt_caching": True, - } - } - uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"}) - remote = sync_module.transform_friendli_data([uncached_model], local) - sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote)) - synced = local["friendliai/zai-org/GLM-Test"] - assert "cache_read_input_token_cost" not in synced - assert synced["supports_prompt_caching"] is False - - -def test_sync_still_merges_entries_outside_replace_keys(sync_module): - local = {"openrouter/some-model": {"input_cost_per_token": 1e-06, "supports_vision": True}} - remote = {"openrouter/some-model": {"input_cost_per_token": 2e-06}} - sync_module.sync_local_data_with_remote(local, remote) - assert local["openrouter/some-model"] == {"input_cost_per_token": 2e-06, "supports_vision": True} - - -def test_transform_inherits_allowlisted_keys_from_base_model_entry(sync_module): - local = { - "zhipuai/glm-test": { - "supports_pdf_input": True, - "supports_assistant_prefill": True, - "input_cost_per_token": 9e-06, - } - } - entry = sync_module.transform_friendli_data([_reasoning_model()], local)[ - "friendliai/zai-org/GLM-Test" - ] - assert entry["supports_pdf_input"] is True - assert entry["supports_assistant_prefill"] is True - assert entry["input_cost_per_token"] == 1.5e-07 diff --git a/tests/unit/test_sync_together_ai_models.py b/tests/unit/test_sync_together_ai_models.py deleted file mode 100644 index b8a85bcfbdc..00000000000 --- a/tests/unit/test_sync_together_ai_models.py +++ /dev/null @@ -1,378 +0,0 @@ -import importlib.util -import json -from pathlib import Path -from types import MappingProxyType - -import pytest - -ROOT = Path(__file__).resolve().parents[2] -SCRIPT = ROOT / "scripts" / "sync_together_ai_models.py" -FIXTURES = Path(__file__).resolve().parent / "fixtures" / "together_ai_sync" - -_spec = importlib.util.spec_from_file_location("sync_together_ai_models", SCRIPT) -assert _spec is not None and _spec.loader is not None -sync = importlib.util.module_from_spec(_spec) -_spec.loader.exec_module(sync) - -RECORDED_CATALOG = sync.load_catalog(FIXTURES.joinpath("models_serverless.json").read_bytes()) -RECORDED_DOC = sync.parse_deprecations(FIXTURES.joinpath("deprecations.md").read_text()) - - -def _doc(removal_dates: dict[str, str], redirects: dict[str, str] | None = None) -> object: - return sync.DeprecationDoc( - removal_dates=MappingProxyType(removal_dates), - redirects=MappingProxyType(redirects or {}), - ) - - -def _chat_model(model_id: str, ctx: int = 4096, price: float = 1.0, cached: float | None = None) -> object: - return sync.CatalogModel( - id=model_id, - type="chat", - context_length=ctx, - pricing=sync.CatalogPricing(input=price, output=price, cached_input=cached), - ) - - -@pytest.mark.parametrize( - ("per_million", "expected"), - [ - (3, 3e-06), - (15, 1.5e-05), - (1.4, 1.4e-06), - (0.25999999999999995, 2.6e-07), - (0.060000000000000005, 6e-08), - (1.0399999999999998, 1.04e-06), - (0, 0.0), - ], -) -def test_per_token_normalizes_float_artifacts(per_million: float, expected: float) -> None: - assert sync.per_token(per_million) == expected - - -def test_parse_deprecations_recorded_fixture() -> None: - assert dict(RECORDED_DOC.redirects) == { - "mistralai/Mistral-7B-Instruct-v0.3": "mistralai/Ministral-3-14B-Instruct-2512", - "Kimi-K2": "Kimi-K2-0905", - "DeepSeek-V3": "DeepSeek-V3.1", - "DeepSeek-V3-0324": "DeepSeek-V3.1", - "DeepSeek-R1": "DeepSeek-R1-0528", - } - assert len(RECORDED_DOC.removal_dates) == 208 - assert RECORDED_DOC.removal_dates["google/gemma-3n-E4B-it"] == "2026-08-04" - - -def test_parse_deprecations_duplicate_rows_keep_most_recent_date() -> None: - assert RECORDED_DOC.removal_dates["Qwen/Qwen3-235B-A22B-Thinking-2507"] == "2026-04-16" - - -@pytest.mark.parametrize( - "markdown", - [ - "# Deprecations\n\nNothing here anymore.\n", - "\n## Active model redirects\n\n| A | B |\n| --- | --- |\n| `x` | `y` |\n\n## Something else\n", - "\n## Deprecation history\n\n### Inference\n\n| Date | Model | R |\n| --- | --- | --- |\n| 2026-01-01 | `m` | No |\n", - ], -) -def test_parse_deprecations_raises_when_a_table_parses_empty(markdown: str) -> None: - with pytest.raises(sync.SyncError): - sync.parse_deprecations(markdown) - - -def test_load_catalog_raises_on_shape_change() -> None: - with pytest.raises(sync.SyncError): - sync.load_catalog(b'[{"id": "x", "type": "chat"}]') - - -def test_load_catalog_raises_when_no_token_models_remain() -> None: - only_video = json.dumps([{"id": "v", "type": "video", "pricing": {"input": 0, "output": 0}}]).encode() - with pytest.raises(sync.SyncError): - sync.load_catalog(only_video) - - -def test_recorded_catalog_counts() -> None: - assert len(RECORDED_CATALOG) == 102 - assert sum(1 for model in RECORDED_CATALOG if model.type in sync.TYPE_TO_MODE) == 26 - assert sum(1 for model in RECORDED_CATALOG if model.pricing.cached_input) == 13 - - -def test_added_chat_model_matches_reviewed_registry_shape() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - assert len(outcome.added) == 26 - assert not outcome.deprecated - assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"] == { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 3e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048576, - "max_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - "source": "https://docs.together.ai/docs/serverless-models", - "supports_function_calling": True, - "supports_parallel_function_calling": True, - "supports_prompt_caching": True, - "supports_reasoning": True, - "supports_response_schema": True, - "supports_tool_choice": True, - "supports_vision": True, - } - - -def test_added_embedding_model_has_no_output_token_cap() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - assert outcome.cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] == { - "input_cost_per_token": 2e-08, - "litellm_provider": "together_ai", - "max_input_tokens": 514, - "max_tokens": 514, - "mode": "embedding", - "output_cost_per_token": 2e-08, - "output_vector_size": 1024, - "source": "https://docs.together.ai/docs/serverless-models", - } - - -def test_moderation_type_maps_to_chat_mode() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - guard = outcome.cost_map["together_ai/meta-llama/Llama-Guard-4-12B"] - assert guard["mode"] == "chat" - assert "max_output_tokens" not in guard - - -def test_output_ceiling_comes_from_the_rule_never_from_context_length() -> None: - glm = next(model for model in RECORDED_CATALOG if model.id == "zai-org/GLM-5.2") - fresh = sync.compute_sync({}, [_chat_model("acme/unreviewed", ctx=1048576), glm], _doc({"x": "2026-01-01"})) - unreviewed = fresh.cost_map["together_ai/acme/unreviewed"] - assert "max_output_tokens" not in unreviewed - assert (unreviewed["max_input_tokens"], unreviewed["max_tokens"]) == (1048576, 1048576) - reviewed = fresh.cost_map["together_ai/zai-org/GLM-5.2"] - assert (reviewed["max_input_tokens"], reviewed["max_output_tokens"], reviewed["max_tokens"]) == ( - 1048575, - 128000, - 128000, - ) - inflated = { - "together_ai/zai-org/GLM-5.2": { - "input_cost_per_token": 1.4e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048575, - "max_output_tokens": 1048575, - "max_tokens": 1048575, - "mode": "chat", - "output_cost_per_token": 4.4e-06, - } - } - corrected = sync.compute_sync(inflated, [glm], _doc({"x": "2026-01-01"})) - assert corrected.cost_map["together_ai/zai-org/GLM-5.2"]["max_output_tokens"] == 128000 - assert any("max_output_tokens: 1048575 -> 128000" in line for line in corrected.updated) - - -def test_docs_removed_but_live_model_stays_live_with_warning() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - gemma = outcome.cost_map["together_ai/google/gemma-3n-E4B-it"] - assert "deprecation_date" not in gemma - assert any("gemma-3n-E4B-it" in warning and "2026-08-04" in warning for warning in outcome.warnings) - - -def test_price_change_updates_api_fields_and_keeps_curated_ones() -> None: - registry = { - "together_ai/acme/chat-1": { - "input_cost_per_token": 9e-07, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "max_output_tokens": 2048, - "max_tokens": 2048, - "mode": "chat", - "output_cost_per_token": 9e-07, - "supports_audio_input": True, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/chat-1", ctx=8192, price=2.0)], _doc({"x": "2026-01-01"})) - entry = outcome.cost_map["together_ai/acme/chat-1"] - assert entry["input_cost_per_token"] == 2e-06 - assert entry["max_input_tokens"] == 8192 - assert entry["max_output_tokens"] == 2048 - assert entry["supports_audio_input"] is True - assert len(outcome.updated) == 1 - assert "input_cost_per_token" in outcome.updated[0] - - -def test_cached_input_appearing_and_disappearing() -> None: - registry = { - "together_ai/acme/chat-1": { - "cache_read_input_token_cost": 1e-07, - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - "supports_prompt_caching": True, - }, - "together_ai/acme/chat-2": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - } - catalog = [_chat_model("acme/chat-1"), _chat_model("acme/chat-2", cached=0.25999999999999995)] - outcome = sync.compute_sync(registry, catalog, _doc({"x": "2026-01-01"})) - assert "cache_read_input_token_cost" not in outcome.cost_map["together_ai/acme/chat-1"] - assert "supports_prompt_caching" not in outcome.cost_map["together_ai/acme/chat-1"] - assert outcome.cost_map["together_ai/acme/chat-2"]["cache_read_input_token_cost"] == 2.6e-07 - assert outcome.cost_map["together_ai/acme/chat-2"]["supports_prompt_caching"] is True - - -def test_capability_rule_backfills_existing_entry() -> None: - registry = { - "together_ai/moonshotai/Kimi-K3": { - "input_cost_per_token": 3e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 1048576, - "mode": "chat", - "output_cost_per_token": 1.5e-05, - } - } - kimi = next(model for model in RECORDED_CATALOG if model.id == "moonshotai/Kimi-K3") - outcome = sync.compute_sync(registry, [kimi], _doc({"x": "2026-01-01"})) - assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"]["supports_reasoning"] is True - assert any("supports_reasoning" in line for line in outcome.updated) - - -def test_disappeared_model_gets_docs_date_and_is_never_deleted() -> None: - registry = { - "together_ai/acme/gone": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"})) - assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-07-01" - assert outcome.deprecated == ("together_ai/acme/gone: deprecation_date",) - - -def test_disappeared_model_without_docs_date_warns_instead() -> None: - registry = { - "together_ai/acme/gone": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"other": "2026-07-01"})) - assert "deprecation_date" not in outcome.cost_map["together_ai/acme/gone"] - assert not outcome.deprecated - assert any("acme/gone" in warning and "human" in warning for warning in outcome.warnings) - - -def test_curated_deprecation_date_is_never_overwritten() -> None: - registry = { - "together_ai/acme/gone": { - "deprecation_date": "2026-06-15", - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"})) - assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-06-15" - assert any("2026-06-15" in warning and "2026-07-01" in warning for warning in outcome.warnings) - - -def test_redirect_chain_resolves_to_final_live_model() -> None: - doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b", "acme/b": "acme/c"}) - live = frozenset({"acme/c"}) - assert sync.resolve_successor("acme/a", doc, live) == "acme/c" - - -def test_redirect_dead_end_yields_no_successor() -> None: - doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b"}) - assert sync.resolve_successor("acme/a", doc, frozenset({"acme/other"})) is None - - -def test_redirect_short_names_resolve_by_unique_suffix() -> None: - doc = _doc({"moonshotai/Kimi-K2": "2026-01-01"}, redirects={"Kimi-K2": "Kimi-K2-0905"}) - live = frozenset({"moonshotai/Kimi-K2-0905"}) - assert sync.resolve_successor("moonshotai/Kimi-K2", doc, live) == "moonshotai/Kimi-K2-0905" - - -def test_successor_written_only_when_not_curated() -> None: - registry = { - "together_ai/acme/a": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - "together_ai/acme/b": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "metadata": {"successor": "together_ai/acme/curated"}, - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - } - doc = _doc({"acme/a": "2026-01-01", "acme/b": "2026-01-01"}, redirects={"acme/a": "acme/c", "acme/b": "acme/c"}) - outcome = sync.compute_sync(registry, [_chat_model("acme/c")], doc) - assert outcome.cost_map["together_ai/acme/a"]["metadata"] == {"successor": "together_ai/acme/c"} - assert outcome.cost_map["together_ai/acme/b"]["metadata"] == {"successor": "together_ai/acme/curated"} - assert any("acme/curated" in warning for warning in outcome.warnings) - - -def test_reappearance_clears_deprecation_date() -> None: - registry = { - "together_ai/acme/back": { - "deprecation_date": "2026-05-01", - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "max_input_tokens": 4096, - "mode": "chat", - "output_cost_per_token": 1e-06, - } - } - outcome = sync.compute_sync(registry, [_chat_model("acme/back")], _doc({"x": "2026-01-01"})) - assert "deprecation_date" not in outcome.cost_map["together_ai/acme/back"] - assert outcome.reappeared == ("together_ai/acme/back",) - - -def test_new_chat_model_without_rule_is_flagged() -> None: - outcome = sync.compute_sync({}, [_chat_model("acme/unreviewed")], _doc({"x": "2026-01-01"})) - assert any("acme/unreviewed" in warning and "capability rule" in warning for warning in outcome.warnings) - - -def test_new_keys_land_at_the_end_of_the_provider_block() -> None: - registry = { - "aaa": {"mode": "chat"}, - "together_ai/acme/old": { - "input_cost_per_token": 1e-06, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1e-06, - }, - "zzz": {"mode": "chat"}, - } - outcome = sync.compute_sync(registry, [_chat_model("acme/old"), _chat_model("acme/new")], _doc({"x": "2026-01-01"})) - assert list(outcome.cost_map) == ["aaa", "together_ai/acme/old", "together_ai/acme/new", "zzz"] - - -def test_sync_is_idempotent_over_the_repo_cost_map() -> None: - cost_map = json.loads((ROOT / "model_prices_and_context_window.json").read_text()) - first = sync.compute_sync(cost_map, RECORDED_CATALOG, RECORDED_DOC) - second = sync.compute_sync(first.cost_map, RECORDED_CATALOG, RECORDED_DOC) - assert not second.has_changes - assert second.cost_map == first.cost_map - - -def test_pr_body_lists_every_section_and_the_skipped_types() -> None: - outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC) - body = sync.render_pr_body(outcome) - assert "### Added (26)" in body - assert "### Warnings needing a human call" in body - assert "image (29)" in body - assert "video (38)" in body From 8b1910c95ae7f9c47135cb41c1e7bbb2700f36e5 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 18:30:05 -0400 Subject: [PATCH 12/19] fix(mcp): advertise valid OAuth resources for toolset URLs (#48474) --- .../mcp_server/auth/user_api_key_auth_mcp.py | 8 +- litellm/proxy/proxy_server.py | 3 +- tests/integration/mcp/test_mcp_oauth_flows.py | 92 ++++++++++++++++--- .../auth/test_user_api_key_auth_mcp.py | 41 +++++++++ 4 files changed, 129 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 824633bdf5a..f883031a529 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -27,6 +27,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_passthrough_resource_metadata_url, get_passthrough_www_authenticate, get_request_base_url, + get_route_relative_request_path, well_known_root_suffix, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( @@ -387,7 +388,7 @@ def _gateway_dcr_challenge( """The RFC 9728 challenge pointing the client at the protected-resource metadata matching the scope it requested: the per-server document (same URL spelling the request arrived on) when the single target advertises gateway-owned sign-in, - else the gateway's aggregate document. Either way the client discovers the gateway + else a gateway document containing the public resource. Either way the client discovers the gateway as its authorization server and starts the same sign-in flow. ``invalid_token`` adds the RFC 6750 error code for a request that DID @@ -397,10 +398,13 @@ def _gateway_dcr_challenge( client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) target: Final = _gateway_dcr_challenge_target(route, mcp_servers, client_ip) server_scope: Final = gateway_dcr_challenge_scope(route, mcp_servers, client_ip) if target is None else None + public_path: Final = get_route_relative_request_path(request.scope).rstrip("/") + aggregate_resource: Final = public_path == "/mcp" or public_path.startswith("/mcp/") + metadata_suffix: Final = f"{well_known_root_suffix()}/mcp" if aggregate_resource else "" resource_metadata_url: Final = ( get_passthrough_resource_metadata_url(request.scope, target) if target is not None - else f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{well_known_root_suffix()}/mcp" + else f"{get_request_base_url(request)}/.well-known/oauth-protected-resource{metadata_suffix}" ) error_attr: Final = 'error="invalid_token", ' if invalid_token else "" scope_attr: Final = f', scope="{server_scope}"' if server_scope is not None else "" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index df948b69a67..129cd7e1589 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -21795,8 +21795,7 @@ async def toolset_mcp_route(toolset_name: str, request: Request): detail=f"Toolset '{toolset_name}' not found", ) - scope: Final = dict(request.scope) - scope["path"] = "/mcp" + scope: Final = {**request.scope, "_original_path": request.scope.get("path", ""), "path": "/mcp"} token: Final = mcp_active_toolset_id.set(toolset.toolset_id) try: diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 327b053d333..90a462472b2 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -792,17 +792,23 @@ def _ui_session_cookie(gateway: Gateway, user_id: str) -> dict[str, str]: def _gateway_session_sign_in( - gateway: Gateway, user_id: str, resource: str | None = None + gateway: Gateway, + user_id: str, + resource: str | None = None, + *, + registration_endpoint: str = "/register", + authorization_endpoint: str = "/authorize", + token_endpoint: str = "/token", ) -> tuple[dict[str, object], str]: registered: Final = gateway.client.post( - "/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"} + registration_endpoint, json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"} ) assert registered.status_code in (200, 201), registered.text client_id: Final = str(_response_object(registered)["client_id"]) pkce: Final = _Pkce(secrets.token_urlsafe(48)) cookies: Final = _ui_session_cookie(gateway, user_id) started: Final = gateway.client.get( - "/authorize", + authorization_endpoint, params={ "client_id": client_id, "redirect_uri": CLIENT_REDIRECT, @@ -823,13 +829,14 @@ def _gateway_session_sign_in( callback: Final = parse_qs(urlsplit(completed.headers["location"]).query) assert "code" in callback, completed.headers["location"] issued: Final = gateway.client.post( - "/token", + token_endpoint, data={ "grant_type": "authorization_code", "code": callback["code"][0], "redirect_uri": CLIENT_REDIRECT, "client_id": client_id, "code_verifier": pkce.verifier, + **({} if resource is None else {"resource": resource}), }, ) assert issued.status_code == 200, issued.text @@ -1292,10 +1299,18 @@ def test_lite_logout_revokes_the_cli_refresh_token_sent_the_way_the_cli_sends_it assert completion.status_code == 200, completion.text -def _toolset_rpc(gateway: Gateway, bearer: str, name: str, method: str, params: dict[str, object]) -> Outcome: +def _toolset_rpc( + gateway: Gateway, + bearer: str, + name: str, + method: str, + params: dict[str, object], + *, + route: str = "/toolset/{name}/mcp", +) -> Outcome: def post(rpc_method: str, rpc_params: dict[str, object]) -> httpx.Response: return gateway.client.post( - f"/toolset/{name}/mcp", + route.format(name=name), json={"jsonrpc": "2.0", "id": 1, "method": rpc_method, "params": rpc_params}, headers={"Authorization": f"Bearer {bearer}", "Accept": "application/json, text/event-stream"}, ) @@ -1306,7 +1321,12 @@ def _toolset_rpc(gateway: Gateway, bearer: str, name: str, method: str, params: return _outcome_from_rpc(post(method, params)) -def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_its_route(gateway: Gateway) -> None: +@pytest.mark.parametrize("route", ["/toolset/{name}/mcp", "/{name}/mcp"]) +def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_its_route( + gateway: Gateway, route: str +) -> None: + from mcp.shared.auth_utils import check_resource_allowed + with mcp_peer() as peer, gateway.scenario() as scenario: alias: Final = "lit6029sess" + uuid.uuid4().hex[:6] server_id: Final = register_mcp(scenario, peer, alias) @@ -1315,19 +1335,69 @@ def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_i granted_id: Final = create_toolset(scenario, ((server_id, "add"),), toolset_name=granted_name) create_toolset(scenario, ((server_id, "multiply"),), toolset_name=withheld_name) member: Final = scenario.member(scenario.team(object_permission={"mcp_toolsets": [granted_id]})) - bearer: Final = _gateway_session_bearer(gateway, member) + path: Final = route.format(name=granted_name) + anonymous: Final = gateway.client.post( + path, + headers=ACCEPT, + json={"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": dict(INITIALIZE)}, + ) + assert anonymous.status_code == 401, anonymous.text + metadata_url: Final = anonymous.headers["www-authenticate"].split('resource_metadata="', 1)[1].split('"', 1)[0] + protected: Final = gateway.client.get(metadata_url) + assert protected.status_code == 200, protected.text + protected_body: Final = _response_object(protected) + resource: Final = str(protected_body["resource"]) + assert check_resource_allowed( + requested_resource=f"{_base(gateway)}{path}", configured_resource=resource + ), protected.text + issuers: Final = TypeAdapter(tuple[str, ...]).validate_python(protected_body["authorization_servers"]) + assert len(issuers) == 1, protected.text + issuer: Final = urlsplit(issuers[0]) + authorization: Final = gateway.client.get( + f"{issuer.scheme}://{issuer.netloc}/.well-known/oauth-authorization-server{issuer.path}" + ) + assert authorization.status_code == 200, authorization.text + authorization_body: Final = _response_object(authorization) + assert authorization_body["issuer"] == issuers[0], authorization.text + token_endpoint: Final = str(authorization_body["token_endpoint"]) + issued, client_id = _gateway_session_sign_in( + gateway, + member, + resource, + registration_endpoint=str(authorization_body["registration_endpoint"]), + authorization_endpoint=str(authorization_body["authorization_endpoint"]), + token_endpoint=token_endpoint, + ) + bearer: Final = _issued_token(issued) assert bearer.startswith("llm_session_"), bearer[:16] - listed: Final = _toolset_rpc(gateway, bearer, granted_name, "tools/list", {}) + listed: Final = _toolset_rpc(gateway, bearer, granted_name, "tools/list", {}, route=route) assert listed.tools == (f"{alias}-add",), listed.raw peer.drain() called: Final = _toolset_rpc( - gateway, bearer, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": {"a": 4, "b": 5}} + gateway, bearer, granted_name, "tools/call", + {"name": f"{alias}-add", "arguments": {"a": 4, "b": 5}}, route=route ) assert called.ok and called.text == "9", called.raw assert len(tool_calls(peer.drain())) == 1 - denied: Final = _toolset_rpc(gateway, bearer, withheld_name, "tools/list", {}) + denied: Final = _toolset_rpc(gateway, bearer, withheld_name, "tools/list", {}, route=route) assert denied.status == 403, denied.raw assert tool_calls(peer.drain()) == () + rotated: Final = gateway.client.post( + token_endpoint, + data={ + "grant_type": "refresh_token", + "refresh_token": str(issued["refresh_token"]), + "client_id": client_id, + "resource": resource, + }, + ) + assert rotated.status_code == 200, rotated.text + rotated_body: Final = _response_object(rotated) + assert rotated_body["refresh_token"] != issued["refresh_token"], rotated.text + refreshed: Final = _toolset_rpc( + gateway, _issued_token(rotated_body), granted_name, "tools/list", {}, route=route + ) + assert refreshed.tools == listed.tools, refreshed.raw def test_resource_scoped_session_bearer_opens_a_team_toolset_inside_its_server_and_none_outside( diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index e8689b2ba89..16130364364 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7693,6 +7693,47 @@ class TestMCPDcrBridgeDualCredential: assert "invalid_token" in str(exc_info.value.headers) +@pytest.mark.parametrize("public_path", ["/toolset/analysis/mcp", "/analysis/mcp", "/mcp", "/mcp/", "/mcp/proxy"]) +@pytest.mark.parametrize("root_path", ["", "/litellm"]) +@pytest.mark.parametrize("invalid_token", [False, True]) +def test_gateway_challenge_metadata_contains_the_public_resource( + public_path: str, root_path: str, invalid_token: bool, +) -> None: + from mcp.shared.auth_utils import check_resource_allowed + from starlette.requests import Request + from starlette.types import Scope + + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import _gateway_dcr_challenge + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_aggregate_protected_resource_response, + oauth_protected_resource_root, + ) + + scope: Final[Scope] = { + "type": "http", "method": "POST", "scheme": "https", "path": "/mcp", + "root_path": root_path, "_original_path": f"{root_path}{public_path}", + "headers": [(b"host", b"gateway.example.com")], + } + request: Final = Request(scope) + challenge: Final = _gateway_dcr_challenge(request, "/mcp", None, invalid_token) + header: Final = (challenge.headers or {})["WWW-Authenticate"] + metadata_url: Final = header.split('resource_metadata="', 1)[1].split('"', 1)[0] + toolset: Final = public_path in ("/toolset/analysis/mcp", "/analysis/mcp") + assert metadata_url == ( + f"https://gateway.example.com{root_path}/.well-known/oauth-protected-resource" + + ("" if toolset else "/mcp") + ) + metadata: Final = ( + oauth_protected_resource_root(request) if toolset else _build_aggregate_protected_resource_response(request) + ) + assert check_resource_allowed( + requested_resource=f"https://gateway.example.com{root_path}{public_path}", + configured_resource=metadata["resource"], + ) + assert challenge.status_code == 401 + assert ('error="invalid_token"' in header) is invalid_token + + @pytest.mark.asyncio class TestAggregateGatewayDcrChallenge: """The mcp_gateway_dcr front door: a 401 on the aggregate /mcp scope must From 48a893aefc1f4cee13f41d2b6f5bde8b382aa6b9 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 26 Oct 2026 15:38:17 -0700 Subject: [PATCH 13/19] build(deps): drop expired exclude-newer pins for certifi, mlflow and pypdf (#48461) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- pyproject.toml | 8 -------- uv.lock | 5 ----- 2 files changed, 13 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 88bb0b13d91..fc860a8fdf4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -377,16 +377,8 @@ required-version = ">=0.10.9" exclude-newer = "3 days" [tool.uv.exclude-newer-package] -# TODO: ryan, remove after 2026-10-23, once certifi 2026.10.19 (drops the compromised eMudhra root) clears the 3-day window -certifi = "2026-10-20T03:00:00Z" # TODO: remove after 2026-10-25, once brotli 1.2.1 (CVE-2026-111497) clears the 3-day window brotli = "2026-10-22T14:01:00Z" -# TODO: ryan, remove after 2026-10-24, once mlflow 3.17.2 (GHSA-wqmh-qccw-pqwf) clears the 3-day window -mlflow = "2026-10-21T03:00:00Z" -mlflow-skinny = "2026-10-21T03:00:00Z" -mlflow-tracing = "2026-10-21T03:00:00Z" -# TODO: ryan, remove after 2026-10-25, once pypdf 6.21.2 (GHSA-6c7f-9265-g8fw) clears the 3-day window -pypdf = "2026-10-21T17:00:00Z" [tool.uv.sources] litellm-proxy-extras = { workspace = true } diff --git a/uv.lock b/uv.lock index fc4155524d7..f86614046e1 100644 --- a/uv.lock +++ b/uv.lock @@ -16,11 +16,6 @@ exclude-newer-span = "P3D" [options.exclude-newer-package] brotli = "2026-10-22T14:01:00Z" -certifi = "2026-10-20T03:00:00Z" -mlflow = "2026-10-21T03:00:00Z" -mlflow-skinny = "2026-10-21T03:00:00Z" -mlflow-tracing = "2026-10-21T03:00:00Z" -pypdf = "2026-10-21T17:00:00Z" [manifest] members = [ From 9a35ed93f9c9dbf7778c28fb3250a9df7cec4e70 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 26 Oct 2026 15:40:03 -0700 Subject: [PATCH 14/19] feat(ui): simplify auto router setup with deployed model defaults (#48451) * feat(ui): simplify auto router setup with deployed model defaults * fix(ui): recognize catalog Opus deployment formats * fix(ui): address auto router setup review gaps * fix(ui): preserve Gemini auto router suggestions * fix(ui): prefer current aliases and wait for team selection * fix(ui): default auto routers to classify each user turn * fix(ui): parse Grok snapshot dates separately from releases * refactor(ui): use non-mutating auto router sorting * feat(ui): allow custom models in simple auto router setup --- .../app/(dashboard)/hooks/models/useModels.ts | 1 + .../AutoRouters/AutoRoutersPanel.tsx | 5 +- .../add_model/AutoRouterPresetSelect.tsx | 125 ++++ .../add_model/AutoRouterSimpleSetup.tsx | 190 +++++++ .../add_auto_router_tab.integration.test.tsx | 458 ++++++++++----- .../add_model/add_auto_router_tab.tsx | 432 +++++++------- .../components/add_model/auto_setup.test.ts | 537 ++++++++++-------- .../src/components/add_model/auto_setup.ts | 347 +++++++---- 8 files changed, 1413 insertions(+), 682 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/AutoRouterPresetSelect.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/AutoRouterSimpleSetup.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index 24f76065f34..744db87ad11 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -121,6 +121,7 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment { /** False for config.yaml-defined deployments, which the update and delete routes refuse. */ db_model?: boolean | null; base_model?: string | null; + blocked?: boolean | null; mode?: string | null; created_at?: string | null; updated_at?: string | null; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx index 6f7e04f92c6..2d1c24685f0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx @@ -106,10 +106,7 @@ export function AutoRoutersPanel({ > Add Auto Router - - Choose a classifier to route each request to a model. Called like any other model, so clients keep using a - single model name. - + Choose models for your auto router. { + switch (availability.kind) { + case "available": + return null; + case "loading": + return "Checking model availability..."; + case "unverifiable": + return "Cannot verify these models are available"; + case "missing_models": + return `Missing: ${availability.models.join(", ")}`; + } +}; + +// "loading"/"unverifiable" are transient system states, not a gap specific to this preset; only a +// caller-specific missing-model reason gets the alarming red treatment. +const isPresetHintAlarming = (availability: PresetAvailability): boolean => availability.kind === "missing_models"; + +type Props = { + selectedPreset?: string; + presets: AutoRouterPreset[]; + presetAvailability: (preset: AutoRouterPreset) => PresetAvailability; + onPresetChange: (key: string | undefined) => void; + pending: boolean; + unavailable: boolean; + onRetry: () => void; +}; + +export default function AutoRouterPresetSelect({ + selectedPreset, + presets, + presetAvailability, + onPresetChange, + pending, + unavailable, + onRetry, +}: Props) { + const sortedPresetOptions = React.useMemo( + () => + presets + .map((preset) => ({ preset, availability: presetAvailability(preset) })) + .sort((a, b) => Number(b.availability.kind === "available") - Number(a.availability.kind === "available")), + [presets, presetAvailability], + ); + + const templateItems = React.useMemo( + () => [ + ...sortedPresetOptions.map(({ preset }) => ({ value: preset.key, label: preset.label })), + { value: "custom", label: "Start from scratch" }, + ], + [sortedPresetOptions], + ); + + return ( +
+
+ +

Replaces classifier and model choices

+ {pending &&
Loading templates...
} + {unavailable && ( +
+ Could not load templates, so only Start from scratch is shown.{" "} + +
+ )} +
+
+ ); +} diff --git a/ui/litellm-dashboard/src/components/add_model/AutoRouterSimpleSetup.tsx b/ui/litellm-dashboard/src/components/add_model/AutoRouterSimpleSetup.tsx new file mode 100644 index 00000000000..64c3d4de63f --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AutoRouterSimpleSetup.tsx @@ -0,0 +1,190 @@ +import { Button } from "@/components/ui/button"; +import { Label } from "@/components/ui/label"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { ROUTER_FAMILIES, ROUTER_GOALS, type FamilyPreference, type RouterGoal, type SolverModel } from "./auto_setup"; + +type Props = { + customized: boolean; + hasSuggestions: boolean; + onReset: () => void; + goal: RouterGoal; + family: FamilyPreference; + models: readonly SolverModel[]; + everyday?: SolverModel; + hard?: SolverModel; + waitingForTeam: boolean; + loading: boolean; + error: boolean; + notice: string | null; + onGoalChange: (goal: RouterGoal) => void; + onFamilyChange: (family: FamilyPreference) => void; + onModelChange: (role: "everyday" | "hard", model: string) => void; + onAdvanced: () => void; + onRetry: () => void; +}; + +const sharedFamily = (everyday?: SolverModel, hard?: SolverModel) => + everyday?.family && everyday.family === hard?.family ? ROUTER_FAMILIES[everyday.family] : null; + +const modelDescription = (model: SolverModel) => + [model.model !== model.name ? model.model : null, model.provider].filter(Boolean).join(" · "); + +export default function AutoRouterSimpleSetup({ + customized, + hasSuggestions, + onReset, + goal, + family, + models, + everyday, + hard, + waitingForTeam, + loading, + error, + notice, + onGoalChange, + onFamilyChange, + onModelChange, + onAdvanced, + onRetry, +}: Props) { + const inventoryReady = !waitingForTeam && !loading; + const showPreferences = !inventoryReady || error || hasSuggestions; + const sameFamily = sharedFamily(everyday, hard); + const familyItems = [ + { value: "auto", label: sameFamily ? `Auto (${sameFamily})` : "Auto" }, + ...Object.entries(ROUTER_FAMILIES).map(([value, label]) => ({ value, label })), + ]; + return ( +
+ {showPreferences && ( + <> +
+ +
+ {Object.entries(ROUTER_GOALS).map(([key, label]) => ( + + ))} +
+
+
+ + +
+ + )} + {waitingForTeam &&

Select a team to see available models.

} + {!waitingForTeam && loading && ( +

+ + Loading deployed models +

+ )} + {inventoryReady && error && ( +

+ Could not load deployed models.{" "} + +

+ )} + {inventoryReady && !error && ( + <> +
+
+ + {customized && hasSuggestions && ( + + )} +
+
+ {(["everyday", "hard"] as const).map((role) => { + const chosen = role === "everyday" ? everyday : hard; + const other = role === "everyday" ? hard : everyday; + const choices = models.filter((model) => model.name !== other?.name); + const items = choices.map((model) => ({ value: model.name, label: model.name })); + const label = role === "everyday" ? "Everyday" : "Hard tasks"; + return ( +
+
+ +

+ {role === "everyday" ? "Your cheaper model" : "Your stronger model"} +

+
+ + {chosen &&

{modelDescription(chosen)}

} +
+ ); + })} +
+
+ {(!everyday || !hard) && ( +

+ {models.length < 2 + ? "Two deployed models are needed. Add a model or use Advanced configuration." + : "Choose two models to get started."} +

+ )} + {notice && !customized && ( +

+ {notice} +

+ )} + + )} + +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx index ee9cbe3c480..52c4e9093a4 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.integration.test.tsx @@ -1,3 +1,4 @@ +import { autoRouterListKey } from "@/app/(dashboard)/hooks/models/useModels"; import { normalizeTierModels } from "./complexity_router_tiers"; import { openAutoRouterAdvanced, @@ -174,6 +175,12 @@ vi.mock("../common_components/team_dropdown", () => ({ ), })); +const renderAdvanced: typeof renderWithProviders = (...args) => { + const result = renderWithProviders(...args); + fireEvent.click(screen.getByRole("button", { name: "Advanced configuration", exact: true })); + return result; +}; + const Harness = () => ; describe("AddAutoRouterTab", () => { @@ -192,12 +199,185 @@ describe("AddAutoRouterTab", () => { vi.mocked(modelCostMap).mockResolvedValue({}); }); + it("creates the goal-selected latest OpenAI pair using its Bedrock deployment aliases", async () => { + mockFetchAvailableModels.mockResolvedValue( + ["small", "middle-old", "middle", "large"].map((model_group) => ({ model_group, mode: "chat" })), + ); + mockFetchAllModelDeployments.mockResolvedValue([ + { model_name: "small", litellm_params: { model: "bedrock/us.openai.gpt-6-luna" } }, + { model_name: "middle-old", litellm_params: { model: "openai/gpt-6-sol" } }, + { model_name: "middle", litellm_params: { model: "bedrock/us.openai.gpt-6.1-sol" } }, + { model_name: "large", litellm_params: { model: "bedrock/global.openai.gpt-6-astra" } }, + ]); + renderWithProviders(); + expect(await screen.findByRole("combobox", { name: "Everyday model" })).toHaveTextContent("small"); + expect(screen.queryByText("Define your tiers")).not.toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Best quality", exact: true })); + expect(screen.getByRole("combobox", { name: "Everyday model" })).toHaveTextContent("middle"); + expect(screen.getAllByText("gpt-6.1-sol · Bedrock").some((element) => element.tagName === "P")).toBe(true); + await waitFor(() => expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled()); + fireEvent.click(screen.getByRole("button", { name: "Add Auto Router" })); + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledOnce()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls[0][0]).toMatchObject({ + auto_router_name: "auto-router", + complexity_router_config: { + classifier_type: "heuristic", + classification_mode: "user_turn", + tiers: { SIMPLE: ["middle"], MEDIUM: ["middle"], COMPLEX: ["large"], REASONING: ["large"] }, + }, + }); + }); + + it("creates a router from two custom models without a goal or family decision", async () => { + const customModels = ["fast-custom", "quality-custom"]; + mockFetchAvailableModels.mockResolvedValue(customModels.map((model_group) => ({ model_group, mode: "chat" }))); + mockFetchAllModelDeployments.mockResolvedValue( + customModels.map((model_name) => ({ + model_name, + litellm_params: { model: `openai/${model_name}` }, + })), + ); + renderWithProviders(); + expect(await screen.findByText("Choose two models to get started.")).toBeVisible(); + expect(screen.queryByRole("group", { name: "Routing goal" })).not.toBeInTheDocument(); + expect(screen.queryByRole("combobox", { name: "Model family" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); + await chooseSelectOption( + userEvent.setup(), + screen.getByRole("combobox", { name: "Everyday model" }), + /fast-custom/, + ); + await chooseSelectOption( + userEvent.setup(), + screen.getByRole("combobox", { name: "Hard tasks model" }), + /quality-custom/, + ); + expect(screen.queryByRole("button", { name: "Use suggested models" })).not.toBeInTheDocument(); + await waitFor(() => expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled()); + fireEvent.click(screen.getByRole("button", { name: "Add Auto Router" })); + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledOnce()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls[0][0].complexity_router_config).toMatchObject({ + classifier_type: "heuristic", + classification_mode: "user_turn", + tiers: { + SIMPLE: ["fast-custom"], + MEDIUM: ["fast-custom"], + COMPLEX: ["quality-custom"], + REASONING: ["quality-custom"], + }, + }); + }); + + it("saves concrete solver names expanded by the backend from an openai wildcard route", async () => { + mockFetchAvailableModels.mockResolvedValue( + ["openai/gpt-6-luna", "openai/gpt-6.1-sol", "openai/gpt-6-astra"].map((model_group) => ({ + model_group, + mode: "chat", + providers: ["openai"], + })), + ); + mockFetchAllModelDeployments.mockResolvedValue([{ model_name: "openai/*", litellm_params: { model: "openai/*" } }]); + renderWithProviders(); + await waitFor(() => expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled()); + fireEvent.click(screen.getByRole("button", { name: "Add Auto Router" })); + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledOnce()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls[0][0]).toMatchObject({ + complexity_router_config: { + classifier_type: "heuristic", + classification_mode: "user_turn", + tiers: { + SIMPLE: ["openai/gpt-6-luna"], + MEDIUM: ["openai/gpt-6-luna"], + COMPLEX: ["openai/gpt-6-astra"], + REASONING: ["openai/gpt-6-astra"], + }, + }, + }); + }); + + it.each(["groups", "deployments"])( + "keeps chosen models and saving available after a failed %s background refresh", + async (source) => { + mockFetchAvailableModels.mockResolvedValue( + ["claude-haiku-5-5", "claude-opus-5-5"].map((model_group) => ({ model_group, mode: "chat" })), + ); + renderWithProviders(); + await waitFor(() => expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled()); + const fetch = source === "groups" ? mockFetchAvailableModels : mockFetchAllModelDeployments; + fetch.mockRejectedValue(new Error("Background refresh unavailable")); + const queryKey = + source === "groups" ? ["availableModels", "autoRouter", "token"] : autoRouterListKey("", "Admin"); + await act(async () => { + await testQueryClient.refetchQueries({ queryKey }); + }); + await waitFor(() => expect(testQueryClient.getQueryState(queryKey)?.status).toBe("error")); + expect(screen.getByRole("combobox", { name: "Everyday model" })).toHaveTextContent("claude-haiku-5-5"); + expect(screen.queryByText("Could not load deployed models.")).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled(); + fireEvent.click(screen.getByRole("button", { name: "Add Auto Router" })); + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalledOnce()); + }, + ); + + it("keeps manual solver choices on goal changes and carries them into Advanced", async () => { + mockFetchAvailableModels.mockResolvedValue( + ["claude-haiku-5-5", "custom-fast", "claude-opus-5-5"].map((model_group) => ({ + model_group, + mode: "chat", + })), + ); + renderWithProviders(); + const everyday = await screen.findByRole("combobox", { name: "Everyday model" }); + await userEvent.click(everyday); + await userEvent.click(await screen.findByRole("option", { name: /custom-fast/ })); + fireEvent.click(screen.getByRole("button", { name: "Cheapest", exact: true })); + expect(screen.getByRole("combobox", { name: "Everyday model" })).toHaveTextContent("custom-fast"); + expect(screen.getByRole("combobox", { name: "Hard tasks model" })).toHaveTextContent("claude-opus-5-5"); + fireEvent.click(screen.getByRole("button", { name: "Advanced configuration", exact: true })); + expectTierModel("Simple", "custom-fast"); + expectTierModel("Reasoning", "claude-opus-5-5"); + }); + + it("keeps standard setup usable when the Heuristic v2 allowance is exhausted", async () => { + mockFetchAvailableModels.mockResolvedValue( + ["claude-haiku-5-5", "claude-opus-5-5"].map((model_group) => ({ model_group, mode: "chat" })), + ); + vi.mocked(apiClient.post).mockResolvedValue({ + allowances: [{ key: "heuristic_v2", limit: 1, remaining: 0, available: true }], + error: null, + }); + renderWithProviders(); + await waitFor(() => expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled()); + expect(screen.queryByText(/allowance/i)).not.toBeInTheDocument(); + }); + + it("blocks creation when deployment identity cannot be verified, then recovers on retry", async () => { + mockFetchAvailableModels.mockResolvedValue( + ["claude-haiku-5-5", "claude-opus-5-5"].map((model_group) => ({ model_group, mode: "chat" })), + ); + mockFetchAllModelDeployments.mockRejectedValueOnce(new Error("Unavailable")); + renderWithProviders(); + expect(await screen.findByRole("alert")).toHaveTextContent("Could not load deployed models"); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); + fireEvent.click(screen.getByRole("button", { name: "Retry", exact: true })); + await waitFor(() => expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeEnabled()); + }); + + it("offers Advanced instead of inventing a second solver for a one-model deployment", async () => { + mockFetchAvailableModels.mockResolvedValue([{ model_group: "claude-opus-5-5", mode: "chat" }]); + renderWithProviders(); + expect(await screen.findByText(/Two deployed models are needed/)).toBeVisible(); + expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); + fireEvent.click(screen.getByRole("button", { name: "Advanced configuration", exact: true })); + expect(screen.getByText("Define your tiers")).toBeVisible(); + }); + it("shows Jev essentials before tiers and saves tuning from its own disclosure", async () => { mockFetchAvailableModels.mockResolvedValue([ ...ALL_FAMILY_MODELS, { model_group: "my-decision", mode: "evaluation" }, ]); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); await userEvent.click(screen.getByRole("radio", { name: "Decisions Model" })); expect(screen.queryByLabelText("Classifier Model")).not.toBeInTheDocument(); @@ -250,7 +430,7 @@ describe("AddAutoRouterTab", () => { { model_name: "ordinary-chat", litellm_params: { model: "mixed-provider/chat-model" } }, { model_name: "unavailable-decider", litellm_params: { model: underlyingModel } }, ]); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); await userEvent.click(screen.getByRole("radio", { name: "Decisions Model" })); await userEvent.click(screen.getByRole("combobox", { name: "Decision model" })); @@ -269,7 +449,7 @@ describe("AddAutoRouterTab", () => { it("uses the judge picker empty state when no decision deployments exist", async () => { mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); await userEvent.click(screen.getByRole("radio", { name: "Decisions Model" })); await userEvent.click(screen.getByRole("combobox", { name: "Decision model" })); @@ -280,7 +460,7 @@ describe("AddAutoRouterTab", () => { it("does not refill a tier cleared by the user when models refresh", async () => { mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); await userEvent.click(within(tierChips("Simple")).getByRole("button", { name: "" })); await act(async () => { @@ -294,7 +474,7 @@ describe("AddAutoRouterTab", () => { allowances: [{ key: "heuristic_v2", limit: 1, remaining, available: true }], error: null, }); - renderWithProviders(); + renderAdvanced(); await waitFor(() => expect(screen.getByRole("button", { name: "Heuristic" })).toHaveTextContent("Rule-based")); }); @@ -305,21 +485,21 @@ describe("AddAutoRouterTab", () => { ? "This change needs an available rule-based tuning allowance" : null, })); - renderWithProviders(); + renderAdvanced(); await waitFor(() => expect(screen.getByRole("button", { name: "Heuristic" })).toHaveTextContent("Rule-based")); expect(screen.queryByRole("alert")).not.toBeInTheDocument(); expect(screen.getByRole("button", { name: "Add Auto Router" })).toBeDisabled(); }); it("keeps the Rule-based default when the form reopens with a cached allowance", async () => { - const first = renderWithProviders(); + const first = renderAdvanced(); await waitFor(() => expect(apiClient.post).toHaveBeenCalled()); first.unmount(); vi.mocked(apiClient.post).mockResolvedValueOnce({ allowances: [{ key: "heuristic_v2", limit: 1, remaining: 0, available: true }], error: null, }); - renderWithProviders(); + renderAdvanced(); await waitFor(() => expect(screen.getByRole("button", { name: "Heuristic" })).toHaveTextContent("Rule-based")); }); @@ -330,7 +510,7 @@ describe("AddAutoRouterTab", () => { allowances: [{ key: "heuristic_v2", limit: 1, remaining: 0, available: true }], error: null, }); - renderWithProviders(); + renderAdvanced(); fireEvent.click(screen.getByRole("button", { name: "Heuristic" })); const option = screen.getByRole("menuitemradio", { name: /^Heuristic v2/ }); await waitFor(() => expect(option).toHaveTextContent("0 of 1 available")); @@ -383,7 +563,7 @@ describe("AddAutoRouterTab", () => { ...ALL_FAMILY_MODELS, { model_group: "my-decision", mode: "evaluation" }, ]); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); await user.click(screen.getByRole("radio", { name: "Decisions Model" })); await userEvent.click(screen.getByRole("combobox", { name: "Decision model" })); @@ -431,7 +611,7 @@ describe("AddAutoRouterTab", () => { ...ALL_FAMILY_MODELS, { model_group: "my-decision", mode: "evaluation" }, ]); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); await user.click(screen.getByRole("radio", { name: "Decisions Model" })); await userEvent.click(screen.getByRole("combobox", { name: "Decision model" })); @@ -468,7 +648,7 @@ describe("AddAutoRouterTab", () => { complete = resolve; }), ); - renderWithProviders(); + renderAdvanced(); await userEvent.click(screen.getByRole("radio", { name: "LLM" })); complete?.({ allowances: [{ key: "heuristic_v2", limit: 1, remaining: 0, available: true }], error: null }); await waitFor(() => expect(screen.getByRole("radio", { name: "LLM" })).toBeChecked()); @@ -484,7 +664,7 @@ describe("AddAutoRouterTab", () => { complete = resolve; }), ); - renderWithProviders(); + renderAdvanced(); await userEvent.click(screen.getByRole("radio", { name: family })); await selectAutoRouterOption("How often to classify", "Every new user message"); await act(async () => { @@ -508,8 +688,8 @@ describe("AddAutoRouterTab", () => { { model_group: "capable", mode: "chat" }, { model_group: "judge", mode: "chat" }, ]); - renderWithProviders(); - await user.type(screen.getByLabelText("Auto Router Name"), "forecast-router"); + renderAdvanced(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "forecast-router" } }); await selectAutoRouterApproach(label); expect(screen.getByLabelText("Auto Router Name")).toHaveValue("forecast-router"); expect(screen.queryByTestId("template-selector")).not.toBeInTheDocument(); @@ -575,7 +755,7 @@ describe("AddAutoRouterTab", () => { it("restores the automatic/template/detail flow on the Complexity tab", async () => { const user = userEvent.setup(); mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); await selectAutoRouterApproach("Capability"); expect(screen.queryByTestId("configure-automatically-button")).not.toBeInTheDocument(); @@ -603,9 +783,9 @@ describe("AddAutoRouterTab", () => { { model_group: "capable", mode: "chat" }, { model_group: "judge", mode: "chat" }, ]); - renderWithProviders(); + renderAdvanced(); await selectAutoRouterApproach(label); - await user.type(screen.getByLabelText("Auto Router Name"), "forecast-retry"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "forecast-retry" } }); const capability = label === "Capability"; const policyField = capability ? "Solve probability threshold" : "Efficient solver profile"; fireEvent.change(screen.getByLabelText(policyField), { @@ -627,7 +807,7 @@ describe("AddAutoRouterTab", () => { // Detailed Configuration starts collapsed so the modal opens onto just Name + Template; a caller // opts into the full tier/classifier form rather than always seeing it up front. it("shows models immediately and keeps advanced settings collapsed", async () => { - renderWithProviders(); + renderAdvanced(); expect(screen.getByRole("button", { name: "Advanced settings" })).toHaveAttribute("aria-expanded", "false"); @@ -641,14 +821,14 @@ describe("AddAutoRouterTab", () => { { model_group: "unknown-model-a", mode: "chat" }, { model_group: "unknown-model-b", mode: "chat" }, ]); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await waitFor(() => expect(optionByLabel("Anthropic Family")).toHaveTextContent("Missing:")); expect(screen.queryByTestId("configure-automatically-button")).not.toBeInTheDocument(); }); - it("mixes preferred tier models even when one complete preset is available", async () => { + it("prefers a same-family pair when a complete family is available", async () => { const anthropicPreset = getPresetByKey("anthropic_family")!; mockFetchAvailableModels.mockResolvedValue( [...getRequiredModelsInPreset(anthropicPreset), "gpt-5.6-luna"].map((model_group) => ({ @@ -657,18 +837,18 @@ describe("AddAutoRouterTab", () => { })), ); mockFetchAllModelDeployments.mockResolvedValue([]); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); - expectTierModel("Simple", "gpt-5.6-luna"); - expectTierModel("Medium", "claude-sonnet-5"); + expectTierModel("Simple", ANTHROPIC_TIERS.SIMPLE[0]); + expectTierModel("Medium", ANTHROPIC_TIERS.SIMPLE[0]); expectTierModel("Complex", "claude-opus-5-5"); expectTierModel("Reasoning", "claude-opus-5-5"); expect(toast.success).not.toHaveBeenCalledWith(expect.stringContaining("Configured with")); }); - it("mixes available models from the preferred tier catalog when no complete template fits", async () => { + it("uses two models within the available family without a complete template", async () => { mockFetchAvailableModels.mockResolvedValue( ["gpt-5.6-luna", "claude-sonnet-5", "gpt-5.6-sol"].map((model_group) => ({ model_group, @@ -676,12 +856,12 @@ describe("AddAutoRouterTab", () => { })), ); mockFetchAllModelDeployments.mockResolvedValue([]); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); expectTierModel("Simple", "gpt-5.6-luna"); - expectTierModel("Medium", "claude-sonnet-5"); + expectTierModel("Medium", "gpt-5.6-luna"); expectTierModel("Complex", "gpt-5.6-sol"); expectTierModel("Reasoning", "gpt-5.6-sol"); }); @@ -689,20 +869,20 @@ describe("AddAutoRouterTab", () => { it("opens Detailed Configuration on the tiers automatic setup just filled in", async () => { const simpleModel = "gpt-5.6-luna"; mockFetchAvailableModels.mockResolvedValue([...ALL_FAMILY_MODELS, { model_group: simpleModel, mode: "chat" }]); - renderWithProviders(); + renderAdvanced(); expect(screen.getByRole("button", { name: "Advanced settings" })).toHaveAttribute("aria-expanded", "false"); await waitForAutomaticTiers(); expect(screen.getByText("Define your tiers")).toBeInTheDocument(); - expectTierModel("Simple", simpleModel); + expectTierModel("Simple", ANTHROPIC_TIERS.SIMPLE[0]); }); // Nothing is filled in, so there is nothing to submit. The button reports that itself instead of // accepting a click and answering with a toast. it("offers no submit at all until every tier has a model", async () => { - renderWithProviders(); + renderAdvanced(); expect(screen.getByRole("button", { name: /add auto router/i })).toBeDisabled(); }); @@ -710,7 +890,8 @@ describe("AddAutoRouterTab", () => { it("still flags the router name once the config no longer blocks the submit", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -720,15 +901,27 @@ describe("AddAutoRouterTab", () => { }); it("offers no team selector to a proxy admin, who may create an unscoped router", async () => { - renderWithProviders(); + renderAdvanced(); expect(screen.queryByTestId("team-dropdown")).not.toBeInTheDocument(); }); - it("requires a team admin to pick a team", async () => { + it("waits for a team before showing model inventory in simple setup", async () => { renderWithProviders( , ); + expect(screen.getByText("Select a team to see available models.")).toBeInTheDocument(); + expect(screen.queryByText(/Two deployed models are needed/)).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: /add auto router/i })).toBeDisabled(); + await userEvent.selectOptions(screen.getByTestId("team-dropdown"), "team-1"); + expect(await screen.findByText("Chosen models")).toBeInTheDocument(); + expect(screen.queryByText("Select a team to see available models.")).not.toBeInTheDocument(); + }); + + it("requires a team admin to pick a team", async () => { + renderAdvanced( + , + ); expect(screen.getByTestId("team-dropdown")).toBeInTheDocument(); expect(screen.getByText("Select Team")).toBeInTheDocument(); @@ -741,11 +934,11 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders( + renderAdvanced( , ); - await user.type(screen.getByPlaceholderText(/smart_router/i), "team-scoped-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "team-scoped-router" } }); await user.selectOptions(screen.getByTestId("team-dropdown"), "team-1"); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -763,7 +956,7 @@ describe("AddAutoRouterTab", () => { ); vi.mocked(handleAddAutoRouterSubmit).mockImplementationOnce(actualSubmit.handleAddAutoRouterSubmit); const user = userEvent.setup(); - renderWithProviders( + renderAdvanced( { error: "session_affinity cannot be combined with tier_definitions", }); - renderWithProviders(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "rejected-router"); + renderAdvanced(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "rejected-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -838,7 +1031,7 @@ describe("AddAutoRouterTab", () => { }), ); - const { container } = renderWithProviders(); + const { container } = renderAdvanced(); fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "double-submit-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); @@ -858,8 +1051,8 @@ describe("AddAutoRouterTab", () => { vi.mocked(getMissingTiersError).mockReturnValue(null); validateAutoRouterConfig.mockResolvedValueOnce({ valid: true }); - renderWithProviders(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "accepted-router"); + renderAdvanced(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "accepted-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -873,9 +1066,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "keyword-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Keyword/Semantic Matching"); await user.click(screen.getByRole("button", { name: /add keyword rule/i })); @@ -890,9 +1083,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "keyword-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Keyword/Semantic Matching"); await user.click(screen.getByRole("button", { name: /add keyword rule/i })); @@ -908,9 +1101,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "orphan-rule-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "orphan-rule-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Keyword/Semantic Matching"); await user.click(screen.getByRole("button", { name: /add keyword rule/i })); @@ -927,9 +1120,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "keyword-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Keyword/Semantic Matching"); await user.click(screen.getByRole("button", { name: /add keyword rule/i })); @@ -944,9 +1137,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "keyword-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Keyword/Semantic Matching"); await user.click(screen.getByRole("button", { name: /add keyword rule/i })); @@ -965,11 +1158,11 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders( + renderAdvanced( , ); - await user.type(screen.getByPlaceholderText(/smart_router/i), "team-scoped-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "team-scoped-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -983,11 +1176,11 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders( + renderAdvanced( , ); - await user.type(screen.getByPlaceholderText(/smart_router/i), "team-scoped-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "team-scoped-router" } }); await user.selectOptions(screen.getByTestId("team-dropdown"), "team-1"); await user.click(screen.getByTestId("team-dropdown-clear")); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); @@ -1001,9 +1194,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "affinity-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Classification Method"); expect(await screen.findByRole("combobox", { name: "How often to classify" })).not.toHaveTextContent( @@ -1022,7 +1215,8 @@ describe("AddAutoRouterTab", () => { it("blocks invalid success thresholds and creates a heuristic v2 router with explicit zero", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "" } }); fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "threshold-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Classification Method"); @@ -1059,7 +1253,7 @@ describe("AddAutoRouterTab", () => { "./handle_add_auto_router_submit", ); vi.mocked(handleAddAutoRouterSubmit).mockImplementationOnce(actualSubmit.handleAddAutoRouterSubmit); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "chained-router" } }); await userEvent.click(screen.getByRole("radio", { name: "LLM" })); @@ -1087,7 +1281,7 @@ describe("AddAutoRouterTab", () => { it("preserves classifier tuning when choosing models automatically", async () => { const user = userEvent.setup(); mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "reset-threshold-router" } }); openAutoRouterAdvanced("Classification Method"); @@ -1112,7 +1306,8 @@ describe("AddAutoRouterTab", () => { it("clears an invalid inactive threshold before creating the router", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "" } }); fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "clear-threshold-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Classification Method"); @@ -1133,13 +1328,15 @@ describe("AddAutoRouterTab", () => { it.each([false, true])("creates a router with cache routing only after an explicit opt-in: %s", async (enabled) => { const user = userEvent.setup(); mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForAutomaticTiers(); fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "cache-aware-router" } }); openAutoRouterAdvanced("Cache-aware routing"); const toggle = screen.getByRole("switch", { name: "Cache-aware routing" }); expect(toggle).not.toBeChecked(); if (enabled) { + expect(toggle).toHaveAttribute("aria-disabled", "true"); + await selectAutoRouterOption("How often to classify", "Every request"); await user.click(toggle); fireEvent.change(screen.getByLabelText("Expected output tokens"), { target: { value: "512" } }); fireEvent.change(screen.getByLabelText("Prediction timeout (ms)"), { target: { value: "750" } }); @@ -1158,9 +1355,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-window-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "ctx-window-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Context Window Escalation"); const toggle = await screen.findByRole("switch", { name: "Escalate oversized prompts to a tier that fits" }); @@ -1181,9 +1378,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-buffer-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "ctx-buffer-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Context Window Escalation"); await user.click(screen.getByRole("switch", { name: "Escalate oversized prompts to a tier that fits" })); @@ -1204,9 +1401,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-clear-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "ctx-clear-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Context Window Escalation"); await user.click(screen.getByRole("switch", { name: "Escalate oversized prompts to a tier that fits" })); @@ -1230,9 +1427,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "no-compression-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -1246,9 +1443,11 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "no-compression-explicit-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { + target: { value: "no-compression-explicit-router" }, + }); expandDetailedConfiguration(); openAutoRouterAdvanced("Compression"); await chooseSelectOption( @@ -1270,9 +1469,11 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "different-compression-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { + target: { value: "different-compression-router" }, + }); expandDetailedConfiguration(); openAutoRouterAdvanced("Compression"); await chooseSelectOption( @@ -1299,9 +1500,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "override-floor-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "override-floor-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Classification Method"); await selectAutoRouterOption("Heuristic", "Rule-based"); @@ -1321,9 +1522,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "affinity-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Classification Method"); await selectAutoRouterOption("How often to classify", "Once per session"); @@ -1346,9 +1547,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "user-turn-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "user-turn-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Classification Method"); await selectAutoRouterOption("How often to classify", "Every new user message"); @@ -1362,16 +1563,19 @@ describe("AddAutoRouterTab", () => { }); }); - it("writes every_request into the create payload when the default frequency stays selected", async () => { + it("allows every request in Advanced while defaulting to every new user message", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "default-timing-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "default-timing-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Classification Method"); - expect(await screen.findByRole("combobox", { name: "How often to classify" })).toHaveTextContent("Every request"); + expect(await screen.findByRole("combobox", { name: "How often to classify" })).toHaveTextContent( + "Every new user message", + ); + await selectAutoRouterOption("How often to classify", "Every request"); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -1386,9 +1590,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "affinity-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Affinity"); expect(await screen.findByRole("switch", { name: "Pin one model deployment per tier" })).toBeChecked(); @@ -1406,9 +1610,9 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "affinity-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Affinity"); await user.click(await screen.findByRole("switch", { name: "Pin one model deployment per tier" })); @@ -1426,7 +1630,7 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "modality-router" } }); @@ -1444,7 +1648,7 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); - renderWithProviders(); + renderAdvanced(); fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "modality-router" } }); expandDetailedConfiguration(); @@ -1465,7 +1669,7 @@ describe("AddAutoRouterTab", () => { // Custom is the escape hatch, not the headline choice, so it's listed after every bundled preset // rather than first. it("lists Start from scratch after the bundled presets", async () => { - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); const labels = visibleOptions().map((option) => option.querySelector(".font-medium")?.textContent); @@ -1487,7 +1691,7 @@ describe("AddAutoRouterTab", () => { ); vi.mocked(getMissingTiersError).mockImplementation(actual.getMissingTiersError); - renderWithProviders(); + renderAdvanced(); expect(screen.getByTestId("auto-router-test-routing-btn")).toBeDisabled(); }); @@ -1504,8 +1708,8 @@ describe("AddAutoRouterTab", () => { }, }); - renderWithProviders(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "keyword-router"); + renderAdvanced(); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "keyword-router" } }); expandDetailedConfiguration(); openAutoRouterAdvanced("Keyword/Semantic Matching"); await user.click(screen.getByRole("button", { name: /add keyword rule/i })); @@ -1540,7 +1744,7 @@ describe("AddAutoRouterTab", () => { }, }); - renderWithProviders(); + renderAdvanced(); await user.click(screen.getByTestId("auto-router-test-routing-btn")); await user.type(await screen.findByTestId("auto-router-routing-test-prompt"), "reconcile this invoice"); await user.click(screen.getByTestId("auto-router-routing-test-send")); @@ -1564,7 +1768,7 @@ describe("AddAutoRouterTab", () => { }), ); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); const anthropicOption = optionByLabel("Anthropic Family")!; @@ -1582,7 +1786,7 @@ describe("AddAutoRouterTab", () => { it("disables every preset and offers a retry when the model list fails to load", async () => { mockFetchAvailableModels.mockRejectedValue(new Error("network error")); - renderWithProviders(); + renderAdvanced(); expect(await screen.findByText("Could not load available models.")).toBeInTheDocument(); openTemplateDropdown(); @@ -1595,7 +1799,7 @@ describe("AddAutoRouterTab", () => { mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); mockFetchAllModelDeployments.mockRejectedValue(new Error("network error")); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await waitForPresetEnabled("OpenAI Family"); @@ -1606,7 +1810,7 @@ describe("AddAutoRouterTab", () => { ALL_FAMILY_MODELS.filter((m) => m.model_group !== ANTHROPIC_ONLY_MODEL), ); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await waitFor(() => { @@ -1618,7 +1822,7 @@ describe("AddAutoRouterTab", () => { it("enables a preset once every model it needs is available", async () => { mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await waitForPresetEnabled("OpenAI Family"); @@ -1626,7 +1830,7 @@ describe("AddAutoRouterTab", () => { it("shows the template's model choices while advanced settings stay collapsed", async () => { mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await selectTemplate("Anthropic Family"); expect(screen.getByRole("button", { name: "Advanced settings" })).toHaveAttribute("aria-expanded", "false"); @@ -1635,7 +1839,7 @@ describe("AddAutoRouterTab", () => { }); it("expands detailed configuration when Start from scratch is chosen", async () => { - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await selectTemplate("Start from scratch"); @@ -1647,7 +1851,7 @@ describe("AddAutoRouterTab", () => { it("lets a caller manually re-expand a detailed configuration a preset just collapsed", async () => { mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await selectTemplate("Anthropic Family"); expect(screen.queryByText("Keyword/Semantic Matching")).not.toBeInTheDocument(); @@ -1666,11 +1870,11 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await selectTemplate("Anthropic Family"); - await user.type(screen.getByPlaceholderText(/smart_router/i), "anthropic-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "anthropic-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -1689,7 +1893,7 @@ describe("AddAutoRouterTab", () => { it("blocks a form submit when a referenced model disappears after the tiers are filled in", async () => { mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - const { container } = renderWithProviders(); + const { container } = renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await selectTemplate("Anthropic Family"); fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "stale-model-router" } }); @@ -1713,11 +1917,11 @@ describe("AddAutoRouterTab", () => { const user = userEvent.setup(); mockFetchAvailableModels.mockResolvedValue(ALL_FAMILY_MODELS); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await selectTemplate("Anthropic Family"); - await user.type(screen.getByPlaceholderText(/smart_router/i), "anthropic-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "anthropic-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -1751,10 +1955,10 @@ describe("AddAutoRouterTab", () => { it("submits the pinned model in place of the one the tiers derive", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderAdvanced(); await applyPresetAndPin(user); - await user.type(screen.getByPlaceholderText(/smart_router/i), "pinned-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "pinned-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -1771,7 +1975,7 @@ describe("AddAutoRouterTab", () => { it("blocks a submit whose pinned model is no longer available", async () => { const user = userEvent.setup(); - const { container } = renderWithProviders(); + const { container } = renderAdvanced(); await applyPresetAndPin(user); fireEvent.change(screen.getByPlaceholderText(/smart_router/i), { target: { value: "stale-pin-router" } }); @@ -1796,11 +2000,11 @@ describe("AddAutoRouterTab", () => { it("omits plan_mode_min_tier from the payload when never touched", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await selectTemplate("Anthropic Family"); - await user.type(screen.getByPlaceholderText(/smart_router/i), "no-plan-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "no-plan-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -1812,7 +2016,7 @@ describe("AddAutoRouterTab", () => { it("carries the enabled override through to the create payload", async () => { const user = userEvent.setup(); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Anthropic Family"); await selectTemplate("Anthropic Family"); @@ -1820,7 +2024,7 @@ describe("AddAutoRouterTab", () => { openAutoRouterAdvanced("Plan-Mode Override"); await user.click(await screen.findByRole("switch", { name: "Route plan-mode requests to a minimum tier" })); - await user.type(screen.getByPlaceholderText(/smart_router/i), "plan-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "plan-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -1847,7 +2051,7 @@ describe("AddAutoRouterTab", () => { mockFetchAvailableModels.mockResolvedValue(groupsFor(ALL_RENAMED_DEPLOYMENTS)); mockFetchAllModelDeployments.mockResolvedValue(ALL_RENAMED_DEPLOYMENTS); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await waitFor(() => { @@ -1874,7 +2078,7 @@ describe("AddAutoRouterTab", () => { mockFetchAvailableModels.mockResolvedValue(groupsFor(deployments)); mockFetchAllModelDeployments.mockResolvedValue(deployments); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await waitFor(() => { expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(false); @@ -1886,7 +2090,7 @@ describe("AddAutoRouterTab", () => { expect(screen.getByText("Keyword/Semantic Matching")).toBeInTheDocument(); - await user.type(screen.getByPlaceholderText(/smart_router/i), "renamed-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "renamed-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -1908,7 +2112,7 @@ describe("AddAutoRouterTab", () => { mockFetchAvailableModels.mockResolvedValue(groupsFor(anthropicOnly)); mockFetchAllModelDeployments.mockResolvedValue(anthropicOnly); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await waitFor(() => { @@ -1933,7 +2137,7 @@ describe("AddAutoRouterTab", () => { mockFetchAvailableModels.mockResolvedValue(groupsFor(wildcard)); mockFetchAllModelDeployments.mockResolvedValue(wildcard); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await waitFor(() => { @@ -1960,7 +2164,7 @@ describe("AddAutoRouterTab", () => { mockFetchAvailableModels.mockResolvedValue(EXPANDED_HUB_GROUPS); mockFetchAllModelDeployments.mockResolvedValue(WILDCARD_DEPLOYMENTS); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await waitFor(() => { @@ -1974,14 +2178,14 @@ describe("AddAutoRouterTab", () => { mockFetchAvailableModels.mockResolvedValue(EXPANDED_HUB_GROUPS); mockFetchAllModelDeployments.mockResolvedValue(WILDCARD_DEPLOYMENTS); - renderWithProviders(); + renderAdvanced(); openTemplateDropdown(); await waitFor(() => { expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(false); }); await selectTemplate("Anthropic Family"); - await user.type(screen.getByPlaceholderText(/smart_router/i), "wildcard-router"); + fireEvent.change(screen.getByLabelText("Auto Router Name"), { target: { value: "wildcard-router" } }); await waitFor(() => expect(screen.getByRole("button", { name: /add auto router/i })).toBeEnabled()); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -2079,7 +2283,7 @@ describe("preset catalog fetch states", () => { }, ], }); - renderWithProviders(); + renderAdvanced(); await waitForPresetEnabled("Bounded JEV"); await selectTemplate("Bounded JEV"); await userEvent.click(screen.getByRole("combobox", { name: "Decision model" })); @@ -2100,7 +2304,7 @@ describe("preset catalog fetch states", () => { ...LOADED_PRESETS_QUERY, isError: true, } as never); - renderWithProviders(); + renderAdvanced(); expect(screen.queryByText(/Could not load templates/)).not.toBeInTheDocument(); @@ -2114,7 +2318,7 @@ describe("preset catalog fetch states", () => { data: undefined, isPending: true, } as never); - renderWithProviders(); + renderAdvanced(); expect(screen.getByText("Loading templates...")).toBeInTheDocument(); }); @@ -2127,7 +2331,7 @@ describe("preset catalog fetch states", () => { isError: true, refetch, } as never); - renderWithProviders(); + renderAdvanced(); expect(await screen.findByText(/Could not load templates/)).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 82c30395755..8c3f402df1f 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -13,7 +13,7 @@ import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; import { Input } from "@/components/ui/input"; -import { Select, SelectContent, SelectItem, SelectSeparator, SelectTrigger, SelectValue } from "@/components/ui/select"; +import AutoRouterPresetSelect, { type PresetAvailability } from "./AutoRouterPresetSelect"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useZodForm } from "@/lib/forms/useZodForm"; @@ -79,7 +79,16 @@ import { } from "@/lib/autorouter_presets"; import { useAutoRouterPresets } from "@/app/(dashboard)/hooks/autoRouter/useAutoRouterPresets"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; -import { buildAutomaticRouterConfig, buildPreferredTierModels, prefillEmptyTiers } from "./auto_setup"; +import AutoRouterSimpleSetup from "./AutoRouterSimpleSetup"; +import { + deployedRouterModels, + nextRouterName, + simpleSetupError, + recommendSolverPair, + solverPairConfig, + type RouterGoal, + type FamilyPreference, +} from "./auto_setup"; interface AddAutoRouterTabProps { handleOk: () => void; @@ -95,31 +104,6 @@ interface AddAutoRouterTabProps { teams?: Team[] | null; } -type PresetAvailability = - | { kind: "available"; viaDeployments: boolean } - | { kind: "loading" } - | { kind: "unverifiable" } - | { kind: "missing_models"; models: readonly string[] }; - -// Every non-"available" state disables the option. Selection derives from this same function -// (see presetAvailability below), so an option a caller can click is always one that can be applied. -const presetDisabledHint = (availability: PresetAvailability): string | null => { - switch (availability.kind) { - case "available": - return null; - case "loading": - return "Checking model availability..."; - case "unverifiable": - return "Cannot verify these models are available"; - case "missing_models": - return `Missing: ${availability.models.join(", ")}`; - } -}; - -// "loading"/"unverifiable" are transient system states, not a gap specific to this preset; only a -// caller-specific missing-model reason gets the alarming red treatment. -const isPresetHintAlarming = (availability: PresetAvailability): boolean => availability.kind === "missing_models"; - const NO_PRESETS: AutoRouterPreset[] = []; // Why the submit is unavailable, or null when it is available. The button reads this to disable @@ -168,7 +152,7 @@ const autoRouterSchema = (requiresTeamScope: boolean) => type AddAutoRouterFormValues = z.infer>; const EMPTY_FORM_VALUES: AddAutoRouterFormValues = { - auto_router_name: "", + auto_router_name: "auto-router", team_id: null, model_access_group: undefined, }; @@ -210,11 +194,18 @@ const AddAutoRouterTab: React.FC = ({ }); const [modelAccessGroups, setModelAccessGroups] = useState([]); - const [complexityRouterConfig, setComplexityRouterConfig] = useState({ + const [advancedConfig, setComplexityRouterConfig] = useState({ tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, classifier_type: "heuristic", + classification_mode: "user_turn", }); + const [advanced, setAdvanced] = useState(false); + const [prefillAttempted, setPrefillAttempted] = useState(false); + const [goal, setGoal] = useState("balanced"); + const [family, setFamily] = useState("auto"); + const [solverOverrides, setSolverOverrides] = useState<{ everyday?: string; hard?: string }>({}); + const [customTechnicalKeywords, setCustomTechnicalKeywords] = useState([]); const [keywordTierRules, setKeywordTierRules] = useState([]); const [semanticMatchingEnabled, setSemanticMatchingEnabled] = useState(false); @@ -250,17 +241,30 @@ const AddAutoRouterTab: React.FC = ({ isError: modelsError, refetch: refetchModels, } = useQuery({ - queryKey: ["availableModels", "autoRouter", accessToken, ...(isMemberManaged ? [watchedTeamId] : [])], + queryKey: ["availableModels", "autoRouter", accessToken, ...(requiresTeamScope ? [watchedTeamId] : [])], queryFn: () => - isMemberManaged ? fetchAutoRouterModels(accessToken, watchedTeamId) : fetchAvailableModels(accessToken), - enabled: Boolean(accessToken && (!isMemberManaged || watchedTeamId)), + requiresTeamScope ? fetchAutoRouterModels(accessToken, watchedTeamId) : fetchAvailableModels(accessToken), + enabled: Boolean(accessToken && (!requiresTeamScope || watchedTeamId)), }); - const { data: deployments, isLoading: deploymentsLoading } = useQuery({ + const { + data: deployments, + isLoading: deploymentsLoading, + isError: deploymentsError, + refetch: refetchDeployments, + } = useQuery({ queryKey: autoRouterListKey(userId ?? "", userRole), queryFn: () => fetchAllModelDeployments(accessToken, userId ?? "", userRole), enabled: Boolean(accessToken), }); const modelsLoading = groupsLoading || deploymentsLoading; + useEffect(() => { + if (!deployments || form.getFieldState("auto_router_name").isDirty) return; + form.setValue( + "auto_router_name", + nextRouterName(deployments.flatMap((row) => (row.model_name ? [row.model_name] : []))), + ); + }, [deployments, form]); + const modelInfo = React.useMemo(() => data ?? [], [data]); const { data: presetsData, @@ -269,12 +273,12 @@ const AddAutoRouterTab: React.FC = ({ refetch: refetchPresets, } = useAutoRouterPresets(); const presets = presetsData ?? NO_PRESETS; - const automaticSetupLoading = modelsLoading || presetsPending; const presetsUnavailable = presetsError && presetsData === undefined; // react-query keeps the last successful list around when a later refetch fails, so isError alone // can't tell "never loaded" apart from "loaded, then a background refetch errored" - only the // former leaves us with nothing trustworthy to verify a preset's models against. const modelsUnverifiable = modelsError && data === undefined; + const inventoryUnavailable = modelsUnverifiable || (deploymentsError && deployments === undefined); const isAdmin = all_admin_roles.includes(userRole); @@ -294,14 +298,33 @@ const AddAutoRouterTab: React.FC = ({ ), [modelInfo], ); - const preferredTierModels = React.useMemo( - () => buildPreferredTierModels(presets, availability), - [presets, availability], + const solverModels = React.useMemo( + () => deployedRouterModels(modelInfo, deployments ?? []), + [modelInfo, deployments], ); - const automaticRouterConfig = React.useMemo( - () => buildAutomaticRouterConfig(modelInfo, deployments ?? [], preferredTierModels), - [modelInfo, deployments, preferredTierModels], + const recommended = React.useMemo( + () => recommendSolverPair(solverModels, goal, family), + [solverModels, goal, family], ); + const everyday = solverOverrides.everyday + ? solverModels.find((model) => model.name === solverOverrides.everyday) + : recommended?.everyday; + const hard = solverOverrides.hard + ? solverModels.find((model) => model.name === solverOverrides.hard) + : recommended?.hard; + const simpleConfig = everyday && hard ? solverPairConfig(everyday.name, hard.name) : null; + const complexityRouterConfig = advanced + ? advancedConfig + : simpleConfig ?? { + classifier_type: "heuristic" as const, + classification_mode: "user_turn" as const, + tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, + }; + const openAdvanced = () => { + setComplexityRouterConfig(complexityRouterConfig); + setPrefillAttempted(simpleConfig !== null); + setAdvanced(true); + }; // A preset's models can only be trusted against a successfully loaded list. Selection and the // greyed-out state derive from this one function, so a preset that cannot be selected can never @@ -322,22 +345,6 @@ const AddAutoRouterTab: React.FC = ({ [modelsLoading, modelsUnverifiable, availability, groupsOnlyAvailability], ); - const sortedPresetOptions = React.useMemo( - () => - presets - .map((preset) => ({ preset, availability: presetAvailability(preset) })) - .sort((a, b) => Number(b.availability.kind === "available") - Number(a.availability.kind === "available")), - [presets, presetAvailability], - ); - - const templateItems = React.useMemo( - () => [ - ...sortedPresetOptions.map(({ preset }) => ({ value: preset.key, label: preset.label })), - { value: "custom", label: "Start from scratch" }, - ], - [sortedPresetOptions], - ); - const applyPrefill = (prefill: PresetPrefill) => { setEditingTiers(false); setComplexityRouterConfig(prefill.complexityRouterConfig); @@ -349,16 +356,20 @@ const AddAutoRouterTab: React.FC = ({ setEscalationKeywords(prefill.escalationKeywords); }; - const [prefillAttempted, setPrefillAttempted] = useState(false); - const automaticPrefill = automaticSetupLoading || modelsUnverifiable ? null : automaticRouterConfig; - if (automaticPrefill && !prefillAttempted) { + if (advanced && !prefillAttempted && !modelsLoading && simpleConfig) { setPrefillAttempted(true); - setComplexityRouterConfig((current) => prefillEmptyTiers(current, automaticPrefill)); + setComplexityRouterConfig((current) => + !current.custom_tier_set && + !isForecastClassifier(current.classifier_type) && + activeTierRows(current).every((row) => row.models.length === 0) + ? { ...current, tiers: simpleConfig.tiers } + : current, + ); } const handlePresetChange = (presetKey: string | undefined) => { + setPrefillAttempted(true); if (!presetKey || presetKey === "custom") { - setPrefillAttempted(true); setSelectedPreset(presetKey); applyPrefill(buildEmptyPrefill()); @@ -372,7 +383,6 @@ const AddAutoRouterTab: React.FC = ({ const presetState = presetAvailability(preset); if (presetState.kind !== "available") return; - setPrefillAttempted(true); setSelectedPreset(presetKey); applyPrefill(buildPresetPrefill(preset.complexity_router_config, availability)); }; @@ -413,7 +423,11 @@ const AddAutoRouterTab: React.FC = ({ }, !requiresTeamScope || Boolean(watchedTeamId), ); - const saveBlockedReason = submitBlockedReason ?? routerAvailability.saveBlockedReason; + const simpleBlockedReason = advanced ? null : simpleSetupError(modelsLoading, inventoryUnavailable, everyday, hard); + const saveBlockedReason = + !advanced && requiresTeamScope && !watchedTeamId + ? "Select a team to see available models." + : simpleBlockedReason ?? submitBlockedReason ?? routerAvailability.saveBlockedReason; const handleClassifierChange = (config: ComplexityRouterConfigValue) => { setSelectedPreset(undefined); setComplexityRouterConfig(config); @@ -540,71 +554,20 @@ const AddAutoRouterTab: React.FC = ({ setIsTestModalVisible(true); }; - const tierSetup = ( -
-
- -

Replaces classifier and model choices

- {presetsPending &&
Loading templates...
} - {presetsUnavailable && ( -
- Could not load templates, so only Start from scratch is shown.{" "} - -
- )} -
-
- ); - const configurationForm = ( void refetchPresets()} + /> + } editingTiers={editingTiers} onEditingTiersChange={setEditingTiers} modelInfo={modelInfo} @@ -646,8 +609,8 @@ const AddAutoRouterTab: React.FC = ({ )}
-
- {isAdmin && ( +
+ {isAdmin && advanced && (
= ({ )} - - } - judgeModel={ - model.mode !== "embedding") - .map((model) => ({ value: model.model_group, label: model.model_group }))} - showValidationErrors={showValidationErrors} - /> - } - > - - {modelsUnverifiable && ( -

- Could not load available models.{" "} - -

- )} - {configurationForm} - -
- - - Need Help? - - } - /> - Get help on our github - -
- - - - +

+ )} + {configurationForm} + + + )} + {!advanced && routerAvailability.isError && ( +

+ Could not check availability.{" "} + +

+ )} + {!advanced && routerAvailability.data?.error && ( +

+ {routerAvailability.data.error} +

+ )} +
+ + - {isTestingConnection && } - Test Connection - - - - -
-
- - + Need Help? + + } + /> + Get help on our github + +
+ + + + + + + +
+
diff --git a/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts b/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts index 7b12a0d54fd..4097afc17fe 100644 --- a/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/auto_setup.test.ts @@ -1,274 +1,325 @@ -import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { readFileSync } from "node:fs"; import { describe, expect, it } from "vitest"; import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; -import { buildModelAvailability, deploymentRefsFromModelInfo } from "@/lib/autorouter_presets"; import { - buildAutomaticRouterConfig, - buildPreferredTierModels, - prefillEmptyTiers, - type PreferredTierModels, + deployedRouterModels, + recommendSolverPair, + solverPairConfig, + nextRouterName, + simpleSetupError, } from "./auto_setup"; -const models = (...names: string[]) => names.map((model_group) => ({ model_group, mode: "chat" })); -const reasoningModel = (model_group: string, supported_reasoning_efforts: string[]) => ({ - model_group, - mode: "chat", - supports_reasoning: true, - supported_reasoning_efforts, -}); -const deployment = (model_name: string, model = model_name): AutoRouterDeployment => ({ - model_name, +const deployment = (name: string, model: string, extra: Partial = {}): AutoRouterDeployment => ({ + model_name: name, litellm_params: { model }, + ...extra, }); -const tierModels = (config: ReturnType) => - config && Object.values(config.tiers).map((tier) => (typeof tier === "string" ? tier : tier[0])); +const resolve = (...rows: AutoRouterDeployment[]) => + deployedRouterModels( + rows.map((row) => ({ model_group: row.model_name!, mode: "chat" })), + rows, + ); +const names = (pair: ReturnType) => pair && [pair.everyday.name, pair.hard.name]; +const ladder = () => + resolve( + deployment("small", "anthropic/claude-haiku-5-5"), + deployment("medium", "anthropic/claude-sonnet-5-5"), + deployment("large", "anthropic/claude-opus-5-5"), + ); -describe("buildPreferredTierModels", () => { - it("recognizes curated models that are not in a preset", () => { - const available = ["gpt-5.6-luna", "claude-sonnet-5", "grok-4.6", "deepseek-v4-pro"]; - const availability = buildModelAvailability(available, []); - const preferred = buildPreferredTierModels([], availability); - const expected: PreferredTierModels = { - SIMPLE: ["gpt-5.6-luna"], - MEDIUM: ["claude-sonnet-5"], - COMPLEX: ["deepseek-v4-pro", "grok-4.6"], - REASONING: ["deepseek-v4-pro", "grok-4.6"], - }; +const catalog: Record = JSON.parse( + readFileSync(new URL("../../../../../model_prices_and_context_window.json", import.meta.url), "utf8"), +); +const catalogOpusModels = Object.entries(catalog).filter( + ([id, metadata]) => metadata.mode === "chat" && /claude-opus-\d/.test(id) && !/-(?:fast|think)$/.test(id), +); - expect(preferred).toEqual(expected); - }); - - it("prefers the current model ladder over older preset entries", () => { - const availability = buildModelAvailability(["gpt-5.6-luna", "gpt-4o-mini"], []); - const preferred = buildPreferredTierModels( - [ - { - key: "old", - label: "Old", - description: "Old model", - complexity_router_config: { - tiers: { SIMPLE: ["gpt-4o-mini"], MEDIUM: [], COMPLEX: [], REASONING: [] }, - classifier_type: "heuristic_v2", - }, - }, - ], - availability, - ); - - expect(preferred.SIMPLE).toEqual(["gpt-5.6-luna", "gpt-4o-mini"]); - }); -}); - -describe("buildAutomaticRouterConfig", () => { - it("selects native Terra and Sol groups with their reasoning settings, retaining cloud fallback", () => { - const modelNames = ["gpt-5.6-terra", "gpt-5.6-sol"]; - const deployments = modelNames.flatMap((model) => [ - deployment(model, `azure/${model}`), - deployment(`z-native-${model}`, `openai/${model}`), - ]); - const available = deployments.map(({ model_name }) => reasoningModel(model_name!, ["none", "high"])); - const availability = buildModelAvailability( - available.map(({ model_group }) => model_group), - deploymentRefsFromModelInfo(deployments), - ); - const preferred = buildPreferredTierModels([], availability); - const config = buildAutomaticRouterConfig(available, deployments, preferred); - - expect(tierModels(config)).toEqual([ - "z-native-gpt-5.6-terra", - "z-native-gpt-5.6-terra", - "z-native-gpt-5.6-sol", - "z-native-gpt-5.6-sol", - ]); - expect(config?.tier_model_params).toEqual({ - REASONING: { "z-native-gpt-5.6-sol": { reasoning_effort: "high" } }, - }); - - const cloudOnly = models(...modelNames); - const cloudAvailability = buildModelAvailability(modelNames, deploymentRefsFromModelInfo(deployments)); - const cloudPreferred = buildPreferredTierModels([], cloudAvailability); - - expect(tierModels(buildAutomaticRouterConfig(cloudOnly, deployments, cloudPreferred))).toEqual([ - "gpt-5.6-terra", - "gpt-5.6-terra", - "gpt-5.6-sol", - "gpt-5.6-sol", - ]); - }); - - it("selects one preferred model for each tier", () => { - const preferred: PreferredTierModels = { - SIMPLE: ["simple"], - MEDIUM: ["medium"], - COMPLEX: ["complex"], - REASONING: ["reasoning"], - }; - - expect( - tierModels(buildAutomaticRouterConfig(models("simple", "medium", "complex", "reasoning"), [], preferred)), - ).toEqual(["simple", "medium", "complex", "reasoning"]); - }); - - it.each([ - { - provider: "OpenAI", - available: ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-astra"], - expected: ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-6-astra", "gpt-6-astra"], - supportedEfforts: ["low", "medium", "high", "xhigh", "max"], - effort: "max", - }, - { - provider: "Anthropic", - available: ["claude-haiku-4-5", "claude-sonnet-5", "claude-opus-5"], - expected: ["claude-haiku-4-5", "claude-sonnet-5", "claude-opus-5", "claude-opus-5"], - supportedEfforts: ["low", "medium", "high", "max"], - effort: "max", - }, - { - provider: "Google", - available: ["gemini-3.5-flash-lite", "gemini-3.8-flash", "gemini-3.1-pro-preview"], - expected: ["gemini-3.5-flash-lite", "gemini-3.8-flash", "gemini-3.1-pro-preview", "gemini-3.1-pro-preview"], - supportedEfforts: ["low", "medium", "high"], - effort: "high", - }, - { - provider: "DeepSeek", - available: ["deepseek-v4-flash", "deepseek-v4-pro"], - expected: ["deepseek-v4-flash", "deepseek-v4-flash", "deepseek-v4-pro", "deepseek-v4-pro"], - supportedEfforts: ["none", "high"], - effort: "high", - }, - { - provider: "xAI", - available: ["grok-4.6"], - expected: ["grok-4.6", "grok-4.6", "grok-4.6", "grok-4.6"], - supportedEfforts: ["low", "medium", "high", "xhigh"], - effort: "xhigh", - }, - ])( - "uses the current $provider ladder and strongest advertised reasoning effort", - ({ available, expected, supportedEfforts, effort }) => { - const availability = buildModelAvailability(available, []); - const preferred = buildPreferredTierModels([], availability); - const modelInfo = models(...available).map((model) => - model.model_group === expected[3] ? reasoningModel(model.model_group, supportedEfforts) : model, - ); - - const config = buildAutomaticRouterConfig(modelInfo, [], preferred); - - expect(tierModels(config)).toEqual(expected); - expect(config?.tier_model_params).toEqual({ - REASONING: { [expected[3]]: { reasoning_effort: effort } }, - }); +describe("deployedRouterModels", () => { + it.each(catalogOpusModels)( + "recognizes the catalog's standard Opus deployment %s behind a custom alias", + (id, metadata) => { + const row = { + model_name: "custom-deployment", + litellm_params: { model: id, custom_llm_provider: metadata.litellm_provider }, + }; + const models = deployedRouterModels([{ model_group: "custom-deployment", mode: "chat" }], [row]); + expect(models).toEqual([expect.objectContaining({ name: "custom-deployment", family: "claude", line: "opus" })]); }, ); - it("never exceeds the selected model group's advertised reasoning efforts", () => { - const available = ["gpt-5.6-luna", "gpt-5.6-terra", "gpt-5.6-sol"]; - const availability = buildModelAvailability(available, []); - const preferred = buildPreferredTierModels([], availability); - const modelInfo = [ - ...models("gpt-5.6-luna", "gpt-5.6-terra"), - reasoningModel("gpt-5.6-sol", ["none", "low", "medium", "high", "xhigh"]), - ]; - - const config = buildAutomaticRouterConfig(modelInfo, [], preferred); - - expect(config?.tier_model_params).toEqual({ - REASONING: { "gpt-5.6-sol": { reasoning_effort: "xhigh" } }, - }); + it("resolves underlying identities independently of deployment aliases and serving providers", () => { + const models = resolve( + deployment("everyday", "bedrock/us.openai.gpt-6.1-sol"), + deployment("hard", "bedrock_converse/global.openai.gpt-6-astra"), + ); + const expectedEveryday = { name: "everyday", family: "openai", line: "sol", version: [6, 1], provider: "Bedrock" }; + const expectedHard = { name: "hard", family: "openai", line: "astra", version: [6], provider: "Bedrock" }; + expect(models).toEqual( + expect.arrayContaining([expect.objectContaining(expectedEveryday), expect.objectContaining(expectedHard)]), + ); + expect(names(recommendSolverPair(models, "quality", "openai"))).toEqual(["everyday", "hard"]); }); - it("leaves reasoning effort unset when the proxy does not report supported values", () => { - const available = ["grok-4.6"]; - const availability = buildModelAvailability(available, []); - const preferred = buildPreferredTierModels([], availability); - - const config = buildAutomaticRouterConfig(models(...available), [], preferred); - - expect(config?.tier_model_params).toBeUndefined(); - }); - - it("reuses the closest available tier when a tier has no match", () => { - const preferred: PreferredTierModels = { - SIMPLE: ["simple"], - MEDIUM: [], - COMPLEX: ["complex"], - REASONING: [], - }; - - expect(tierModels(buildAutomaticRouterConfig(models("simple", "complex"), [], preferred))).toEqual([ - "simple", - "simple", - "complex", - "complex", + it("recognizes provider-specific Claude names and a declared base model for opaque deployments", () => { + const models = resolve( + deployment("small", "vertex_ai/claude-haiku-4-5@20251001"), + deployment("large", "bedrock/arn:example:inference-profile/us.anthropic.claude-opus-5-5-v1:0"), + deployment("middle", "azure/private-alias", { + litellm_params: { model: "azure/private-alias" }, + model_info: { base_model: "gpt-6.1-sol" }, + }), + ); + expect(models.map(({ name, family, provider }) => ({ name, family, provider }))).toEqual([ + { name: "large", family: "claude", provider: "Bedrock" }, + { name: "middle", family: "openai", provider: "Azure" }, + { name: "small", family: "claude", provider: "Vertex AI" }, ]); }); - it("returns null when none of the available models are recommended", () => { - const preferred: PreferredTierModels = { - SIMPLE: ["missing-simple"], - MEDIUM: ["missing-medium"], - COMPLEX: ["missing-complex"], - REASONING: ["missing-reasoning"], - }; - - expect(buildAutomaticRouterConfig(models("unknown-model"), [], preferred)).toBeNull(); + it("uses only caller-visible chat models and excludes blocked and recursive targets", () => { + const rows = [ + deployment("hidden", "anthropic/claude-opus-6"), + deployment("blocked", "anthropic/claude-opus-5", { model_info: { blocked: true } }), + deployment("embedding", "anthropic/claude-opus-5"), + deployment("recursive", "auto_router/complexity_router"), + deployment("ready", "anthropic/claude-sonnet-5"), + ]; + expect( + deployedRouterModels( + [ + { model_group: "blocked", mode: "chat" }, + { model_group: "embedding", mode: "embedding" }, + { model_group: "recursive" }, + { model_group: "ready", mode: "chat" }, + ], + rows, + ).map((model) => model.name), + ).toEqual(["ready"]); }); - it("ignores non-chat models and existing auto routers", () => { - const preferred: PreferredTierModels = { - SIMPLE: ["gpt-4o-mini", "smart-router"], - MEDIUM: [], - COMPLEX: [], - REASONING: [], - }; - const available = [ - { model_group: "gpt-4o-mini", mode: "chat" }, - { model_group: "image-model", mode: "image_generation" }, - { model_group: "smart-router", mode: "chat" }, + it("does not claim a mixed-version or unknown model group is a known version", () => { + const rows = [ + deployment("mixed", "anthropic/claude-opus-5"), + deployment("mixed", "anthropic/claude-opus-5-5"), + deployment("claude-sonnet-5", "openai/custom-model"), + deployment("wildcard", "openai/*"), ]; + const models = deployedRouterModels( + [ + { model_group: "mixed", mode: "chat" }, + { model_group: "claude-sonnet-5", mode: "chat" }, + { model_group: "wildcard", mode: "chat" }, + ], + rows, + ); + expect(models.map((model) => model.name)).toEqual(["claude-sonnet-5", "mixed"]); + expect(models.every((model) => model.family === undefined)).toBe(true); + expect(recommendSolverPair(models, "balanced", "auto")).toBeNull(); + }); + it("keeps homogeneous replicas and ignores blocked replicas when resolving identity", () => { + const rows = [ + deployment("group", "anthropic/claude-opus-5-5"), + deployment("group", "bedrock/anthropic.claude-opus-5-5-v1:0"), + deployment("group", "anthropic/claude-opus-5", { model_info: { blocked: true } }), + ]; + const models = deployedRouterModels([{ model_group: "group" }], rows); + expect(models).toHaveLength(1); + expect(models[0]).toMatchObject({ name: "group", version: [5, 5], provider: "Anthropic, Bedrock" }); + }); + + it("keeps arbitrary canonical names available without pretending an unversioned alias is latest", () => { expect( - tierModels( - buildAutomaticRouterConfig(available, [deployment("smart-router", "auto_router/complexity_router")], preferred), + deployedRouterModels([{ model_group: "claude-opus-8-12" }, { model_group: "claude-opus-latest" }], []).map( + (model) => model.name, ), - ).toEqual(["gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini", "gpt-4o-mini"]); + ).toEqual(["claude-opus-8-12", "claude-opus-latest"]); }); }); -describe("prefillEmptyTiers", () => { - const empty: ComplexityRouterConfigValue = { - classifier_type: "jev", - jev_classifier_config: { model: "jev-custom", timeout_ms: 4200 }, - classification_mode: "user_turn", - tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, - }; - const automatic: ComplexityRouterConfigValue = { - classifier_type: "heuristic_v2", - tiers: { SIMPLE: ["fast"], MEDIUM: ["mid"], COMPLEX: ["strong"], REASONING: ["strong"] }, - tier_model_params: { REASONING: { strong: { reasoning_effort: "high" } } }, - }; - it("fills tiers while preserving a classifier configured before models load", () => { - expect(prefillEmptyTiers(empty, automatic)).toEqual({ - ...empty, - tiers: automatic.tiers, - tier_model_params: automatic.tier_model_params, +describe("recommendSolverPair", () => { + it.each([ + ["cheapest", ["small", "medium"]], + ["balanced", ["small", "large"]], + ["quality", ["medium", "large"]], + ] as const)("uses the requested %s pair", (goal, expected) => { + expect(names(recommendSolverPair(ladder(), goal, "claude"))).toEqual(expected); + }); + + it("selects future deployed releases numerically within each line, without a version allowlist", () => { + const models = resolve( + deployment("old-small", "anthropic/claude-haiku-8-9"), + deployment("new-small", "bedrock/anthropic.claude-haiku-8-10-v1:0"), + deployment("old-large", "anthropic/claude-opus-8-9"), + deployment("new-large", "anthropic/claude-opus-9"), + ); + expect(names(recommendSolverPair(models, "balanced", "claude"))).toEqual(["new-small", "new-large"]); + }); + + it("orders snapshots within a version without confusing a snapshot with the next version", () => { + const models = resolve( + deployment("old", "anthropic/claude-haiku-4-5-20251001"), + deployment("snapshot", "anthropic/claude-haiku-4-5-20251101"), + deployment("next", "anthropic/claude-haiku-4-6"), + deployment("large", "anthropic/claude-opus-5"), + ); + expect(names(recommendSolverPair(models, "balanced", "claude"))).toEqual(["next", "large"]); + expect( + names( + recommendSolverPair( + models.filter((model) => model.name !== "next"), + "balanced", + "claude", + ), + ), + ).toEqual(["snapshot", "large"]); + }); + + it("prefers rolling aliases over snapshots of the same release", () => { + const models = resolve( + deployment("pin", "anthropic/claude-haiku-4-5-20251101"), + deployment("alias", "anthropic/claude-haiku-4-5"), + deployment("hard", "anthropic/claude-opus-5"), + deployment("luna-pin", "openai/gpt-6-luna-2026-01-01"), + deployment("luna", "openai/gpt-6-luna"), + deployment("astra", "openai/gpt-6-astra"), + ); + expect(names(recommendSolverPair(models, "balanced", "claude"))).toEqual(["alias", "hard"]); + expect(names(recommendSolverPair(models, "balanced", "openai"))).toEqual(["luna", "astra"]); + }); + + it.each([ + { models: ["grok-4.20", "grok-4.20-0309"], expected: "grok-4.20" }, + { models: ["grok-4.20-0309", "grok-4.20-0410"], expected: "grok-4.20-0410" }, + { models: ["grok-4.20-0309", "grok-4.21"], expected: "grok-4.21" }, + ])("orders Grok releases and short-date snapshots: $expected", ({ models, expected }) => { + const available = resolve( + deployment("everyday", "openai/gpt-6-luna"), + ...models.map((model) => deployment(model, `xai/${model}`)), + ); + expect(names(recommendSolverPair(available, "balanced", "grok"))).toEqual(["everyday", expected]); + }); + + it("does not rank a major-release snapshot above a newer minor release", () => { + const models = resolve( + deployment("old", "anthropic/claude-5-haiku-20260101"), + deployment("next", "anthropic/claude-haiku-5-1"), + deployment("hard", "anthropic/claude-opus-5"), + ); + expect(names(recommendSolverPair(models, "balanced", "claude"))).toEqual(["next", "hard"]); + }); + + it("prefers a same-family substitute over a matching rank from another family", () => { + const models = [ + ...ladder().filter((model) => model.line !== "haiku"), + ...resolve(deployment("luna", "openai/gpt-6-luna")), + ]; + const pair = recommendSolverPair(models, "balanced", "claude"); + expect(names(pair)).toEqual(["medium", "large"]); + expect(pair?.notice).toContain("closest available pair"); + }); + + it("prefers a complete family for Auto and explains a cross-family substitute when explicitly requested", () => { + const models = [ + ...ladder().filter((model) => model.line === "sonnet"), + ...resolve( + deployment("luna", "bedrock/us.openai.gpt-6-luna"), + deployment("astra", "bedrock/us.openai.gpt-6-astra"), + ), + ]; + expect(names(recommendSolverPair(models, "balanced", "auto"))).toEqual(["luna", "astra"]); + const selected = recommendSolverPair(models, "quality", "claude"); + expect(names(selected)).toEqual(["medium", "astra"]); + expect(selected?.notice).toContain("two model families"); + }); + + it.each([ + ["glm", "zai/glm-7.3-flash", "zai/glm-7.3"], + ["deepseek", "deepseek/deepseek-v9-flash", "deepseek/deepseek-v9-pro"], + ] as const)("handles the %s two-line family for every goal", (family, small, large) => { + const models = resolve(deployment("small", small), deployment("large", large)); + for (const goal of ["cheapest", "balanced", "quality"] as const) + expect(names(recommendSolverPair(models, goal, family))).toEqual(["small", "large"]); + }); + + it("does not fabricate a second solver or substitute a different family when none of the requested family exists", () => { + expect(recommendSolverPair(ladder().slice(0, 1), "balanced", "auto")).toBeNull(); + expect(recommendSolverPair(ladder(), "balanced", "openai")).toBeNull(); + expect(recommendSolverPair([], "balanced", "auto")).toBeNull(); + }); + + it("retains Gemini family suggestions and recognizes the existing Grok fallback", () => { + const models = resolve( + deployment("lite", "gemini/gemini-3.5-flash-lite"), + deployment("flash-old", "vertex_ai/gemini-3.5-flash"), + deployment("flash", "vertex_ai/gemini-3.8-flash"), + deployment("pro", "gemini/gemini-3.1-pro-preview"), + deployment("grok", "xai/grok-4.6"), + ); + expect(names(recommendSolverPair(models, "cheapest", "gemini"))).toEqual(["lite", "flash"]); + expect(names(recommendSolverPair(models, "balanced", "gemini"))).toEqual(["lite", "pro"]); + expect(names(recommendSolverPair(models, "quality", "gemini"))).toEqual(["flash", "pro"]); + expect(models.find((m) => m.name === "grok")).toMatchObject({ family: "grok", provider: "xAI" }); + expect(names(recommendSolverPair(models, "balanced", "grok"))).toEqual(["lite", "grok"]); + }); + + it("uses the latest deployed Terra as a fallback without replacing the preferred OpenAI pairs", () => { + const models = resolve( + deployment("terra-old", "openai/gpt-5.6-terra"), + deployment("terra", "openai/gpt-6.1-terra"), + deployment("sol", "openai/gpt-6.1-sol"), + ); + expect(names(recommendSolverPair(models, "balanced", "openai"))).toEqual(["terra", "sol"]); + const complete = [ + ...models, + ...resolve(deployment("luna", "openai/gpt-6-luna"), deployment("astra", "openai/gpt-6-astra")), + ]; + expect(names(recommendSolverPair(complete, "cheapest", "openai"))).toEqual(["luna", "sol"]); + expect(names(recommendSolverPair(complete, "balanced", "openai"))).toEqual(["luna", "astra"]); + expect(names(recommendSolverPair(complete, "quality", "openai"))).toEqual(["sol", "astra"]); + }); + + it("suggests a deterministic equal-rank cross-family fallback only when no ordered pair exists", () => { + const models = resolve(deployment("opus", "anthropic/claude-opus-5-5"), deployment("astra", "openai/gpt-6-astra")); + const pair = recommendSolverPair(models, "balanced", "auto"); + expect(names(pair)).toEqual(["astra", "opus"]); + expect(pair?.notice).toContain("two model families"); + expect(names(recommendSolverPair([...models].reverse(), "balanced", "auto"))).toEqual(names(pair)); + const withLuna = [...models, ...resolve(deployment("luna", "openai/gpt-6-luna"))]; + expect(names(recommendSolverPair(withLuna, "balanced", "auto"))).toEqual(["luna", "astra"]); + }); + + it("does not prefer native-provider deployments over a newer cloud deployment", () => { + const models = resolve( + deployment("native", "openai/gpt-6-luna"), + deployment("cloud", "bedrock/global.openai.gpt-6.2-luna"), + deployment("strong", "bedrock/global.openai.gpt-6-astra"), + ); + expect(names(recommendSolverPair(models, "balanced", "openai"))).toEqual(["cloud", "strong"]); + }); +}); + +describe("solverPairConfig", () => { + it("builds two solver references using built-in rules without paid classifiers, custom tiers or output overrides", () => { + expect(solverPairConfig("everyday", "hard")).toEqual({ + classifier_type: "heuristic", + classification_mode: "user_turn", + tiers: { SIMPLE: ["everyday"], MEDIUM: ["everyday"], COMPLEX: ["hard"], REASONING: ["hard"] }, }); }); - it("leaves the draft alone when no recommended models are available", () => { - expect(prefillEmptyTiers(empty, null)).toBe(empty); +}); + +describe("simple setup validation", () => { + it("avoids existing names without requiring a naming decision", () => { + expect(nextRouterName([])).toBe("auto-router"); + expect(nextRouterName(["auto-router", "auto-router-2", "auto-router-4"])).toBe("auto-router-3"); }); - it("preserves a partial manual assignment", () => { - const partial = { ...empty, tiers: { ...empty.tiers, SIMPLE: ["manual"] } }; - expect(prefillEmptyTiers(partial, automatic)).toBe(partial); - }); - it("preserves a custom tier set created before models load", () => { - const custom = { ...empty, custom_tier_set: { tiers: [], fallback_tier_id: "" } }; - expect(prefillEmptyTiers(custom, automatic)).toBe(custom); - }); - it.each(["capability", "llm_v2"] as const)("does not prefill %s solver pools", (classifier_type) => { - const solver = { ...empty, classifier_type }; - expect(prefillEmptyTiers(solver, automatic)).toBe(solver); + it("blocks stale or incomplete inventory and same-line selections", () => { + const [small, large] = [ladder()[0], ladder()[2]]; + expect(simpleSetupError(true, false, small, large)).not.toBeNull(); + expect(simpleSetupError(false, true, small, large)).not.toBeNull(); + expect(simpleSetupError(false, false, small)).not.toBeNull(); + expect(simpleSetupError(false, false, small, small)).not.toBeNull(); + expect(simpleSetupError(false, false, small, large)).toBeNull(); }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/auto_setup.ts b/ui/litellm-dashboard/src/components/add_model/auto_setup.ts index 8b60336cd27..39476470928 100644 --- a/ui/litellm-dashboard/src/components/add_model/auto_setup.ts +++ b/ui/litellm-dashboard/src/components/add_model/auto_setup.ts @@ -1,110 +1,263 @@ -import { activeTierRows } from "./tier_rows"; -import { isForecastClassifier } from "./forecast_classifier_config"; -import { isAutoRouterDeployment, type AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; +import type { AutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; +import { isAutoRouterDeployment } from "@/app/(dashboard)/hooks/models/useModels"; import type { ModelGroup } from "@/components/llm_calls/fetch_models"; -import { resolveAvailableModel, type AutoRouterPreset, type ModelAvailability } from "@/lib/autorouter_presets"; import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; -const TIER_NAMES = ["SIMPLE", "MEDIUM", "COMPLEX", "REASONING"] as const; -type TierName = (typeof TIER_NAMES)[number]; -export type PreferredTierModels = Record; +export const ROUTER_FAMILIES = { + claude: "Claude", + openai: "OpenAI", + gemini: "Gemini", + grok: "Grok", + glm: "GLM", + deepseek: "DeepSeek", +} as const; +export const ROUTER_GOALS = { cheapest: "Cheapest", balanced: "Balanced", quality: "Best quality" } as const; +export type RouterFamily = keyof typeof ROUTER_FAMILIES; +export type RouterGoal = keyof typeof ROUTER_GOALS; +export type FamilyPreference = RouterFamily | "auto"; -const REASONING_EFFORT_STRENGTH = ["max", "xhigh", "high", "medium", "low", "minimal", "none"] as const; +type ModelIdentity = { + family: RouterFamily; + line: string; + rank: number; + version: readonly number[]; + snapshot: number; + model: string; +}; +export type SolverModel = Partial & { name: string; provider: string; model: string }; +export type RouterModel = ModelIdentity & SolverModel; +export type SolverPair = { everyday: RouterModel; hard: RouterModel; notice: string | null }; -const CURRENT_TIER_MODELS: PreferredTierModels = { - SIMPLE: ["gpt-5.6-luna", "claude-haiku-4-5", "gemini-3.5-flash-lite", "deepseek-v4-flash"], - MEDIUM: ["gpt-5.6-terra", "claude-sonnet-5", "gemini-3.8-flash", "deepseek-v4-flash"], - COMPLEX: ["gpt-6-astra", "gpt-5.6-sol", "claude-opus-5", "gemini-3.1-pro-preview", "deepseek-v4-pro", "grok-4.6"], - REASONING: ["gpt-6-astra", "gpt-5.6-sol", "claude-opus-5", "gemini-3.1-pro-preview", "deepseek-v4-pro", "grok-4.6"], +const releaseVersion = (raw: string, date?: string) => { + const snapshot = raw.match(/[-.](\d{8})$/); + const release = snapshot ? raw.slice(0, -9) : raw; + return { + version: release.split(/[.-]/).map(Number), + // An undated alias follows the current snapshot of its release. + snapshot: date || snapshot ? Number((date ?? snapshot![1]).replaceAll("-", "")) : Number.MAX_SAFE_INTEGER, + }; }; -export const buildPreferredTierModels = ( - presets: AutoRouterPreset[], - availability: ModelAvailability, -): PreferredTierModels => - Object.fromEntries( - TIER_NAMES.map((tier) => [ - tier, - Array.from( - new Set( - [ - ...CURRENT_TIER_MODELS[tier], - ...presets.flatMap((preset) => preset.complexity_router_config.tiers[tier]), - ].flatMap((model) => { - const resolved = resolveAvailableModel(model, availability); - return resolved ? [resolved] : []; - }), - ), +const parseModel = (raw: string): ModelIdentity | null => { + const model = raw + .toLowerCase() + .split("/") + .at(-1)! + .replace(/^(?:us|eu|apac|global|au|jp|us-gov)\./, "") + .replace(/^(?:anthropic|openai|zai|deepseek)\./, "") + .replace(/^databricks-/, "") + .replace(/@default$/, "") + .replace(/-v\d+(?::\d+)?$/, ""); + const versionFirst = model.match(/^claude-(\d+(?:[.-]\d+)*)-(haiku|sonnet|opus)(?:[-@]([0-9]{8}))?$/); + const claude = + model.match(/^claude-(haiku|sonnet|opus)-(\d+(?:[.-]\d+)*)(?:@([0-9]+))?$/) ?? + (versionFirst ? [versionFirst[0], versionFirst[2], versionFirst[1], versionFirst[3]] : null); + if (claude) + return { + family: "claude", + line: claude[1], + rank: ["haiku", "sonnet", "opus"].indexOf(claude[1]), + ...releaseVersion(claude[2], claude[3]), + model, + }; + const openai = model.match(/^gpt-(\d+(?:[.-]\d+)*)-(luna|terra|sol|astra)(?:-(\d{4}-\d{2}-\d{2}))?$/); + if (openai) + return { + family: "openai", + line: openai[2], + rank: openai[2] === "terra" ? 0.5 : ["luna", "sol", "astra"].indexOf(openai[2]), + ...releaseVersion(openai[1], openai[3]), + model, + }; + const gemini = model.match(/^gemini-(\d+(?:[.-]\d+)*)-(flash-lite|flash|pro)(?:-preview)?$/); + if (gemini) + return { + family: "gemini", + line: gemini[2], + rank: ["flash-lite", "flash", "pro"].indexOf(gemini[2]), + ...releaseVersion(gemini[1]), + model, + }; + const grok = model.match(/^grok-(\d+(?:[.-]\d+)*?)(?:-(\d{4}))?$/); + if (grok) return { family: "grok", line: "grok", rank: 2, ...releaseVersion(grok[1], grok[2]), model }; + const glm = model.match(/^glm-(\d+(?:[.-]\d+)*)(-flash)?$/); + if (glm) + return { + family: "glm", + line: glm[2] ? "flash" : "glm", + rank: glm[2] ? 0 : 2, + ...releaseVersion(glm[1]), + model, + }; + const deepseek = model.match(/^deepseek-v(\d+(?:[.-]\d+)*)-(flash|pro)$/); + if (deepseek) + return { + family: "deepseek", + line: deepseek[2], + rank: deepseek[2] === "flash" ? 0 : 2, + ...releaseVersion(deepseek[1]), + model, + }; + return null; +}; + +const identityOf = (deployment: AutoRouterDeployment): ModelIdentity | null => + parseModel(deployment.litellm_params?.model ?? "") ?? + parseModel(deployment.model_info?.base_model ?? deployment.litellm_params?.base_model ?? ""); + +const providerOf = (deployment: AutoRouterDeployment, group: ModelGroup, family?: RouterFamily): string => { + const raw = deployment.litellm_params?.model ?? ""; + const provider = + deployment.litellm_params?.custom_llm_provider || + (raw.includes("/") ? raw.split("/")[0] : group.providers?.join(", ")) || + (family + ? { claude: "anthropic", glm: "zai", openai: "openai", gemini: "gemini", grok: "xai", deepseek: "deepseek" }[ + family + ] + : ""); + if (provider.startsWith("bedrock")) return "Bedrock"; + return ( + ( + { + anthropic: "Anthropic", + openai: "OpenAI", + gemini: "Google AI", + xai: "xAI", + azure: "Azure", + azure_ai: "Azure AI", + databricks: "Databricks", + github_copilot: "GitHub Copilot", + "vertex_ai-anthropic_models": "Vertex AI", + vertex_ai: "Vertex AI", + zai: "Z.ai", + deepseek: "DeepSeek", + } as Record + )[provider] ?? provider + ); +}; + +const compareVersion = (left: readonly number[], right: readonly number[]): number => { + const index = Array.from({ length: Math.max(left.length, right.length) }, (_, i) => i).find( + (i) => (left[i] ?? 0) !== (right[i] ?? 0), + ); + return index === undefined ? 0 : (left[index] ?? 0) - (right[index] ?? 0); +}; + +const compareRelease = (left: ModelIdentity, right: ModelIdentity): number => + compareVersion(left.version, right.version) || left.snapshot - right.snapshot; + +export const deployedRouterModels = ( + groups: readonly ModelGroup[], + deployments: readonly AutoRouterDeployment[], +): SolverModel[] => + groups + .flatMap((group) => { + if (group.mode && group.mode !== "chat") return []; + const rows = deployments.filter((row) => row.model_name === group.model_group); + if ( + rows.some(isAutoRouterDeployment) || + group.model_group.includes("*") || + group.model_group.startsWith("auto_router/") + ) + return []; + const active = rows.filter((row) => !row.model_info?.blocked); + if (rows.length > 0 && active.length === 0) return []; + if (active.some((row) => row.model_info?.mode && row.model_info.mode !== "chat")) return []; + if (active.some((row) => row.litellm_params?.model?.includes("*"))) return []; + const identities = active.length > 0 ? active.map(identityOf) : [parseModel(group.model_group)]; + const first = identities[0]; + const homogeneous = + first && + identities.every((identity) => { + if (!identity) return false; + const sameLine = identity.family === first.family && identity.line === first.line; + return sameLine && compareRelease(identity, first) === 0; + }); + const identity = homogeneous ? first : undefined; + const providers = + active.length > 0 + ? active.map((row) => providerOf(row, group, identity?.family)) + : [providerOf({}, group, identity?.family)]; + return [ + { + ...identity, + name: group.model_group, + model: identity?.model ?? group.model_group, + provider: [...new Set(providers)].toSorted().join(", "), + }, + ]; + }) + .toSorted((a, b) => a.name.localeCompare(b.name)); + +const latestLines = (models: readonly RouterModel[]): RouterModel[] => + models.filter( + (model) => + !models.some( + (other) => other.family === model.family && other.line === model.line && compareRelease(other, model) > 0, ), - ]), - ) as PreferredTierModels; - -const selectPreferredTierModels = ( - preferredByTier: PreferredTierModels, - usableNames: ReadonlySet, -): [string, string, string, string] | null => { - const preferred = TIER_NAMES.map((tier) => preferredByTier[tier].find((name) => usableNames.has(name))); - const candidates = preferred.flatMap((model, tier) => (model ? [{ model, tier }] : [])); - if (candidates.length === 0) return null; - - const nearest = (tier: number): string => - [...candidates].sort( - (left, right) => Math.abs(left.tier - tier) - Math.abs(right.tier - tier) || left.tier - right.tier, - )[0].model; - return preferred.map((model, tier) => model ?? nearest(tier)) as [string, string, string, string]; -}; - -export const buildAutomaticRouterConfig = ( - models: ModelGroup[], - deployments: AutoRouterDeployment[], - preferredByTier: PreferredTierModels, -): ComplexityRouterConfigValue | null => { - const autoRouterNames: ReadonlySet = new Set( - deployments - .filter(isAutoRouterDeployment) - .flatMap((deployment) => (deployment.model_name ? [deployment.model_name] : [])), ); - const names = Array.from( - new Set( - models - .filter((model) => model.mode === undefined || model.mode === "chat") - .map((model) => model.model_group) - .filter((name) => name && !name.startsWith("auto_router/") && !autoRouterNames.has(name)), - ), + +const requestedRanks = (goal: RouterGoal): readonly [number, number] => + (({ cheapest: [0, 1], quality: [1, 2], balanced: [0, 2] }) as const)[goal]; + +export const recommendSolverPair = ( + models: readonly SolverModel[], + goal: RouterGoal, + family: FamilyPreference, +): SolverPair | null => { + const recognized = models.filter((model): model is RouterModel => model.family !== undefined); + const latest = latestLines(recognized); + const [low, high] = requestedRanks(goal); + const pairs = latest.flatMap((everyday) => + latest + .filter((hard) => { + const matchesFamily = family === "auto" || everyday.family === family || hard.family === family; + const ordered = hard.rank > everyday.rank; + const equalRankFallback = hard.rank === everyday.rank && hard.family !== everyday.family; + return (ordered || equalRankFallback) && matchesFamily; + }) + .map((hard) => ({ everyday, hard })), ); - if (names.length === 0) return null; - const usableNames: ReadonlySet = new Set(names); - const selected = selectPreferredTierModels(preferredByTier, usableNames); - if (selected === null) return null; - - const supportedReasoningEfforts = models.find( - (model) => model.model_group === selected[3], - )?.supported_reasoning_efforts; - const reasoningEffort = REASONING_EFFORT_STRENGTH.find((effort) => supportedReasoningEfforts?.includes(effort)); - - return { - tiers: { - SIMPLE: [selected[0]], - MEDIUM: [selected[1]], - COMPLEX: [selected[2]], - REASONING: [selected[3]], - }, - classifier_type: "heuristic_v2", - ...(reasoningEffort && { - tier_model_params: { REASONING: { [selected[3]]: { reasoning_effort: reasoningEffort } } }, - }), - }; + const score = ({ everyday, hard }: { everyday: RouterModel; hard: RouterModel }): number => + (everyday.family === hard.family ? 0 : 100) + + (everyday.rank === hard.rank ? 10 : 0) + + Math.abs(everyday.rank - low) + + Math.abs(hard.rank - high); + const chosen = pairs.toSorted( + (a, b) => + score(a) - score(b) || a.everyday.name.localeCompare(b.everyday.name) || a.hard.name.localeCompare(b.hard.name), + )[0]; + if (!chosen) return null; + const mixed = chosen.everyday.family !== chosen.hard.family; + const substituted = chosen.everyday.rank !== low || chosen.hard.rank !== high; + const mixedNotice = mixed ? "A complete pair is unavailable in one family. Using two model families." : null; + const substituteNotice = substituted ? "Using the closest available pair in this family." : null; + return { ...chosen, notice: mixedNotice ?? substituteNotice }; }; -export const prefillEmptyTiers = ( - current: ComplexityRouterConfigValue, - automatic: ComplexityRouterConfigValue | null, -): ComplexityRouterConfigValue => { - if (!automatic || current.custom_tier_set || isForecastClassifier(current.classifier_type)) return current; - if (activeTierRows(current).some((row) => row.models.length > 0)) return current; - return { - ...current, - tiers: { ...current.tiers, ...automatic.tiers }, - tier_model_params: { ...current.tier_model_params, ...automatic.tier_model_params }, - }; +export const solverPairConfig = (everyday: string, hard: string): ComplexityRouterConfigValue => ({ + classifier_type: "heuristic", + classification_mode: "user_turn", + tiers: { SIMPLE: [everyday], MEDIUM: [everyday], COMPLEX: [hard], REASONING: [hard] }, +}); + +export const simpleSetupError = ( + loading: boolean, + failed: boolean, + everyday?: SolverModel, + hard?: SolverModel, +): string | null => { + if (loading) return "Loading deployed models"; + if (failed) return "Could not verify deployed models"; + if (!everyday || !hard) return "Select two different deployed models"; + if (everyday.name === hard.name) return "Select two different deployed models"; + return null; +}; + +export const nextRouterName = (names: readonly string[]): string => { + const candidates = [ + "auto-router", + ...Array.from({ length: names.length + 1 }, (_, index) => `auto-router-${index + 2}`), + ]; + return candidates.find((name) => !names.includes(name))!; }; From 7045139c3a6fe81f80895cb3887882056b5755fe Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 26 Oct 2026 15:40:19 -0700 Subject: [PATCH 15/19] fix(router): stop concurrent deployment deletes from crashing deployment lookups (#47680) * fix(router): stop concurrent deployment deletes from crashing deployment lookups delete_deployment and upsert_deployment shrink model_list before they repair the index maps, so a lookup on another thread in that window read a stale index: IndexError, or the neighbouring deployment and its credentials. An add racing a delete could corrupt the index for good. One RLock now covers model_list and the three index maps. Delete, upsert, add and the set_model_list reset change them under it, and every lookup that turns an index into a deployment reads under it. Fixes #45242 * test(router): release the held delete when the read reaches the index lock The race tests released the held write after a fixed 0.2 s, so a reader the runner stalled past it read the repaired list and passed on unfixed code. The read now runs on its own thread and the write stays held until the read returns or reaches the router's index lock; the 10 s bound only fails a broken synchronization. The integration test gains the audit cells: chat, streaming, Responses and Anthropic messages calls (sync and async, by model name and team public name) during a held delete, the same for upsert, add and set_model_list, a request for the deployment being deleted, a provider error, a double delete, and thread and async bursts during upsert churn. * test(router): cover the early return of a delete and the team filter of get_model_ids delete_deployment now returns None from inside the index lock, and get_model_ids by name reads its deployments through _deployments_named; pin both: deleting an unknown id returns None and leaves the list alone, and exclude_team_models drops only the team deployment. * fix(router): read deployments by name under the index lock outside the Router The response polling provider check, the complexity router's adaptive pool and the savings baseline turned the name index into list positions themselves, so a concurrent delete could still hand them a stale slot. They now read through Router.deployments_named, which takes the deployment index lock. * fix(router): type the locked deployment lookups as DeploymentTypedDict deployments_named and _team_deployments_named returned tuple[Any, ...]. They now return DeploymentTypedDict rows, which is what model_list holds. The adaptive pool setup in the Router and the complexity router drop their branches for pydantic Deployment rows, which these lookups never return, and should_include_deployment accepts the typed row. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../proxy/response_polling/polling_handler.py | 5 +- litellm/router.py | 243 ++++----- .../complexity_router/complexity_router.py | 15 +- litellm/router_strategy/savings_baseline.py | 11 +- .../router_code_coverage.py | 1 + ...router_concurrent_deployment_delete_sdk.py | 474 ++++++++++++++++++ .../test_response_polling_handler.py | 37 +- .../router_strategy/test_complexity_router.py | 21 +- .../router_strategy/test_savings_baseline.py | 3 +- tests/unit/test_router/test_router.py | 321 +++++++++++- 10 files changed, 964 insertions(+), 167 deletions(-) create mode 100644 tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index 787f27f56c8..e50587b1a0e 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -308,10 +308,7 @@ def should_use_polling_for_request( # Otherwise, check ALL deployments for this model_name in router elif llm_router is not None: try: - # Get all deployment indices for this model name - indices: Final = llm_router.model_name_to_deployment_indices.get(model, []) - for idx in indices: - deployment_dict = llm_router.model_list[idx] + for deployment_dict in llm_router.deployments_named(model) or (): litellm_params = deployment_dict.get("litellm_params", {}) # Check custom_llm_provider first diff --git a/litellm/router.py b/litellm/router.py index 7993b5b921c..3af39736e25 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1440,6 +1440,9 @@ class Router: # Maps (team_id, team_public_model_name) -> list of indices in model_list self.team_model_to_deployment_indices: dict[tuple[str, str], list[int]] = {} self.team_public_model_names: frozenset[str] = frozenset() + # Guards model_list and the three index maps above: writers change them, and readers turn an + # index into a deployment, under it, so a lookup never lands between a removal and its re-index + self._deployment_index_lock = threading.RLock() # Initialize cache attributes that ``_invalidate_model_group_info_cache`` # and ``_invalidate_access_groups_cache`` touch *before* the first @@ -2180,8 +2183,8 @@ class Router: public_model_name: Final = deployment.model_info.team_public_model_name if team_id is None or public_model_name is None: return False - sibling_indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) or () - routable_siblings: Final = self._filter_blocked_deployments([self.model_list[idx] for idx in sibling_indices]) + siblings: Final = self._team_deployments_named(team_id, public_model_name) or () + routable_siblings: Final = self._filter_blocked_deployments([dict(sibling) for sibling in siblings]) return len(routable_siblings) > 1 _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) @@ -10827,20 +10830,17 @@ class Router: # O(k) via the name→indices map: only touch deployments whose name # is listed in `available_models`, instead of scanning model_list. for name in config.available_models: - indices = self.model_name_to_deployment_indices.get(name, []) - if not indices: + deployments = self.deployments_named(name) + if not deployments: continue - d = (self.model_list or [])[indices[0]] - mi = d.get("model_info") if isinstance(d, dict) else d.model_info - mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {}) + d = deployments[0] + mi_dict: dict[str, Any] = d.get("model_info") or {} prefs_raw = mi_dict.get("adaptive_router_preferences") if prefs_raw is not None: model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set. - lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params - lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) - cost = lp_dict.get("input_cost_per_token") + cost = d.get("litellm_params", {}).get("input_cost_per_token") if cost is None: cost = mi_dict.get("input_cost_per_token") if cost is not None: @@ -10932,10 +10932,11 @@ class Router: def set_model_list(self, model_list: list): original_model_list: Final = copy.deepcopy(model_list) self._discovered_model_info_cache.flush_cache() - self.model_list = [] - self.model_id_to_deployment_index_map = {} # Reset the index - self.model_name_to_deployment_indices = {} # Reset the model_name index - self.team_model_to_deployment_indices = {} # Reset the team_model index + with self._deployment_index_lock: + self.model_list = [] + self.model_id_to_deployment_index_map = {} # Reset the index + self.model_name_to_deployment_indices = {} # Reset the model_name index + self.team_model_to_deployment_indices = {} # Reset the team_model index self.team_pattern_routers = {} self.team_public_model_names = frozenset() # Reset per-strategy router registries so hot-reload doesn't leave @@ -11326,26 +11327,27 @@ class Router: - model: dict - the model to add to the list - model_id: Optional[str] - the model ID to use for indexing. If None, will try to get from model["model_info"]["id"] """ - idx: Final = len(self.model_list) - self.model_list.append(model) - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() + with self._deployment_index_lock: + idx: Final = len(self.model_list) + self.model_list.append(model) + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() - # Update model_id index for O(1) lookup - if model_id is not None: - self.model_id_to_deployment_index_map[model_id] = idx - elif model.get("model_info", {}).get("id") is not None: - self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx + # Update model_id index for O(1) lookup + if model_id is not None: + self.model_id_to_deployment_index_map[model_id] = idx + elif model.get("model_info", {}).get("id") is not None: + self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx - # Update model_name index for O(1) lookup - model_name: Final = model.get("model_name") - if model_name: - if model_name not in self.model_name_to_deployment_indices: - self.model_name_to_deployment_indices[model_name] = [] - self.model_name_to_deployment_indices[model_name].append(idx) + # Update model_name index for O(1) lookup + model_name: Final = model.get("model_name") + if model_name: + if model_name not in self.model_name_to_deployment_indices: + self.model_name_to_deployment_indices[model_name] = [] + self.model_name_to_deployment_indices[model_name].append(idx) - # Update team_model index for O(1) team-scoped lookup - self._update_team_model_index(model, idx) + # Update team_model index for O(1) team-scoped lookup + self._update_team_model_index(model, idx) def upsert_deployment(self, deployment: Deployment) -> Deployment | None: """ @@ -11372,18 +11374,15 @@ class Router: # if there is a new litellm param -> then update the deployment # remove the previous deployment - removal_idx: int | None = None - deployment_id: Final = deployment.model_info.id - deployment_fast_mapping: Final = self.model_id_to_deployment_index_map - - if deployment_id in deployment_fast_mapping: - removal_idx = deployment_fast_mapping[deployment_id] - + with self._deployment_index_lock: + removal_idx: Final = self.model_id_to_deployment_index_map.get(_deployment_model_id) if removal_idx is not None: self.model_list.pop(removal_idx) self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() - self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx) + self._update_deployment_indices_after_removal( + model_id=_deployment_model_id, removal_idx=removal_idx + ) # Free the outgoing deployment's pre-routing strategy slot (keyed by the # OLD model_name/tags) before the re-add below re-registers it. @@ -11617,35 +11616,32 @@ class Router: - The deleted deployment - OR None (if deleted deployment not found) """ - deployment_idx = None - if id in self.model_id_to_deployment_index_map: - deployment_idx = self.model_id_to_deployment_index_map[id] - try: - if deployment_idx is not None: + with self._deployment_index_lock: + deployment_idx: Final = self.model_id_to_deployment_index_map.get(id) + if deployment_idx is None: + return None # Pop the item from the list first item: Final = self.model_list.pop(deployment_idx) self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx) - _budget_limiter: Final = self.get_router_deployment_budget_limiter() - if _budget_limiter is not None: - _budget_limiter.unregister_deployment_budget(model_id=id) - try: - self._unregister_pre_routing_strategy_for_deployment( - deployment=item if isinstance(item, Deployment) else Deployment(**item) - ) - except Exception: - verbose_router_logger.exception( - "delete_deployment: could not release pre-routing strategies for model_id=%s; " - "the deployment is out of the model_list and its indices are repaired", - id, - ) - return item - else: - return None + _budget_limiter: Final = self.get_router_deployment_budget_limiter() + if _budget_limiter is not None: + _budget_limiter.unregister_deployment_budget(model_id=id) + try: + self._unregister_pre_routing_strategy_for_deployment( + deployment=item if isinstance(item, Deployment) else Deployment(**item) + ) + except Exception: + verbose_router_logger.exception( + "delete_deployment: could not release pre-routing strategies for model_id=%s; " + "the deployment is out of the model_list and its indices are repaired", + id, + ) except Exception: return None + return item def get_router_deployment_budget_limiter( self, @@ -11694,6 +11690,18 @@ class Router: if _budget_limiter is not None: _budget_limiter.register_deployment_budget(deployment=deployment.to_json(exclude_none=True)) + def deployments_named(self, model_name: str) -> tuple[DeploymentTypedDict, ...] | None: + """The deployments indexed under `model_name`, read under the index lock, or None when the name is not indexed.""" + with self._deployment_index_lock: + indices: Final = self.model_name_to_deployment_indices.get(model_name) + return None if indices is None else tuple(self.model_list[idx] for idx in indices) + + def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None: + """The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed.""" + with self._deployment_index_lock: + indices: Final = self.team_model_to_deployment_indices.get((team_id, public_model_name)) + return None if indices is None else tuple(self.model_list[idx] for idx in indices) + def get_deployment(self, model_id: str) -> Deployment | None: """ Returns -> Deployment or None @@ -11701,9 +11709,10 @@ class Router: Raise Exception -> if model found in invalid format """ # Use O(1) lookup via model_id_to_deployment_index_map only - if model_id in self.model_id_to_deployment_index_map: - idx: Final = self.model_id_to_deployment_index_map[model_id] - model: Final = self.model_list[idx] + with self._deployment_index_lock: + idx: Final = self.model_id_to_deployment_index_map.get(model_id) + model: Final = None if idx is None else self.model_list[idx] + if model is not None: if isinstance(model, dict): return Deployment(**model) elif isinstance(model, Deployment): @@ -11737,17 +11746,17 @@ class Router: Optimized with O(1) index lookup instead of O(n) linear scan. """ # O(1) lookup in model_name index - if model_group_name in self.model_name_to_deployment_indices: - indices: Final = self.model_name_to_deployment_indices[model_group_name] - if indices: - # Return first deployment for this model_name - model: Final = self.model_list[indices[0]] - if isinstance(model, dict): - return Deployment(**model) - elif isinstance(model, Deployment): - return model - else: - raise Exception(f"Model Name invalid - {type(model)}") + with self._deployment_index_lock: + indices: Final = self.model_name_to_deployment_indices.get(model_group_name) + # Return first deployment for this model_name + model: Final = self.model_list[indices[0]] if indices else None + if model is not None: + if isinstance(model, dict): + return Deployment(**model) + elif isinstance(model, Deployment): + return model + else: + raise Exception(f"Model Name invalid - {type(model)}") return None @staticmethod @@ -11770,11 +11779,12 @@ class Router: by other teams so a shared model name never resolves another team's credentials. """ - indices: Final = self.model_name_to_deployment_indices.get(model_group_name) or () - usable: Final = ( - self.model_list[idx] for idx in indices if self.deployment_usable_by_team(self.model_list[idx], team_id) - ) - first_usable: Final = next(usable, None) + with self._deployment_index_lock: + indices: Final = self.model_name_to_deployment_indices.get(model_group_name) or () + usable: Final = ( + self.model_list[idx] for idx in indices if self.deployment_usable_by_team(self.model_list[idx], team_id) + ) + first_usable: Final = next(usable, None) if first_usable is None: return None return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable @@ -11877,11 +11887,10 @@ class Router: this never triggers pattern matching or deep copies, so it is safe to call per listed model on the /v1/models hot path. """ - indices: Final = self.model_name_to_deployment_indices.get(model_name) - if not indices: + deployments: Final = self.deployments_named(model_name) + if not deployments: return None - deployments: Final = tuple(self.model_list[index] for index in indices) model_infos: Final = tuple( MappingProxyType( { @@ -12054,10 +12063,11 @@ class Router: def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: if team_id is None: return None - team_indices: Final = self.team_model_to_deployment_indices.get((team_id, model_id)) - if not team_indices: + with self._deployment_index_lock: + team_indices: Final = self.team_model_to_deployment_indices.get((team_id, model_id)) + team_model: Final = self.model_list[team_indices[0]] if team_indices else None + if team_model is None: return None - team_model: Final = self.model_list[team_indices[0]] return Deployment(**team_model) if isinstance(team_model, dict) else team_model def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: @@ -12296,10 +12306,9 @@ class Router: Optimized with O(1) index lookup instead of O(n) linear scan. """ # O(1) lookup via model_id_to_deployment_index_map - if id in self.model_id_to_deployment_index_map: - idx: Final = self.model_id_to_deployment_index_map[id] - return self.model_list[idx] - return None + with self._deployment_index_lock: + idx: Final = self.model_id_to_deployment_index_map.get(id) + return None if idx is None else self.model_list[idx] def get_model_group(self, id: str) -> list | None: """ @@ -12969,20 +12978,19 @@ class Router: if model_name is not None: # O(1) lookup in model_name index, then O(k) iteration where k = deployments for this model_name - if model_name in self.model_name_to_deployment_indices: - indices: Final = self.model_name_to_deployment_indices[model_name] - for idx in indices: - model = self.model_list[idx] - if "model_info" in model and "id" in model["model_info"]: - if exclude_team_models and model["model_info"].get("team_id"): - continue - ids.append(model["model_info"]["id"]) + for model in self.deployments_named(model_name) or (): + if "model_info" in model and "id" in model["model_info"]: + if exclude_team_models and model["model_info"].get("team_id"): + continue + ids.append(model["model_info"]["id"]) else: # When model_name is None, return all model IDs # Use the index map keys for O(n) where n = total deployments - for model_id in self.model_id_to_deployment_index_map: - idx = self.model_id_to_deployment_index_map[model_id] - model = self.model_list[idx] + with self._deployment_index_lock: + id_indices: Final = tuple(self.model_id_to_deployment_index_map.items()) + models: Final = tuple(self.model_list) + for model_id, idx in id_indices: + model = models[idx] if "model_info" in model and "id" in model["model_info"]: if exclude_team_models and model["model_info"].get("team_id"): continue @@ -13260,7 +13268,9 @@ class Router: # handled downstream by the pattern_router in _common_checks_available_deployment. return None - def should_include_deployment(self, model_name: str, model: dict, team_id: str | None = None) -> bool: + def should_include_deployment( + self, model_name: str, model: DeploymentTypedDict | dict, team_id: str | None = None + ) -> bool: """ Get the team-specific model name if team_id matches the deployment. """ @@ -13308,12 +13318,10 @@ class Router: # O(1) lookup in team_model index when team_id is provided if team_id is not None: - key: Final = (team_id, model_name) - if key in self.team_model_to_deployment_indices: - indices = self.team_model_to_deployment_indices[key] + team_models: Final = self._team_deployments_named(team_id, model_name) + if team_models is not None: # O(k) where k = team deployments for this model_name (typically 1-10) - for idx in indices: - model = self.model_list[idx] + for model in team_models: if not self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): continue if model_alias is not None: @@ -13326,12 +13334,10 @@ class Router: return returned_models # O(1) lookup in model_name index - if model_name in self.model_name_to_deployment_indices: - indices = self.model_name_to_deployment_indices[model_name] - + named_models: Final = self.deployments_named(model_name) + if named_models is not None: # O(k) where k = deployments for this model_name (typically 1-10) - for idx in indices: - model = self.model_list[idx] + for model in named_models: if self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): if model_alias is not None: # Optimized: Use shallow copy since we only modify top-level model_name @@ -13345,7 +13351,7 @@ class Router: # Fallback: if team_id is provided and model_name not in index, # check if model_name matches any team_public_model_name # O(n) scan but only when team_id lookup fails - for idx, model in enumerate(self.model_list): + for _idx, model in enumerate(self.model_list): if self.should_include_deployment(model_name=model_name, model=model, team_id=team_id): if model_alias is not None: # Optimized: Use shallow copy since we only modify top-level model_name @@ -14610,12 +14616,13 @@ class Router: def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]: """Every team's deployments under public name `model`, for a proxy admin calling without a team.""" - team_deployments: Final = [ - self.model_list[index] - for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() - if public_model_name == model - for index in indices - ] + with self._deployment_index_lock: + team_deployments: Final = [ + self.model_list[index] + for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() + if public_model_name == model + for index in indices + ] team_ids: Final = { team_id for deployment in team_deployments diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a475358e1f0..cc592bbcbc5 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -3126,17 +3126,14 @@ class ComplexityRouter(CustomLogger): model_to_prefs: Final[dict[str, AdaptiveRouterPreferences]] = {} model_to_cost: Final[dict[str, float]] = {} - model_list: Final = getattr(self.litellm_router_instance, "model_list", None) or [] - name_to_indices: Final = getattr(self.litellm_router_instance, "model_name_to_deployment_indices", {}) or {} for name in available_models: - indices = name_to_indices.get(name, []) - if not indices: + deployments = self.litellm_router_instance.deployments_named(name) + if not deployments: model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) model_to_cost[name] = 0.0 continue - deployment = model_list[indices[0]] - mi = deployment.get("model_info") if isinstance(deployment, dict) else deployment.model_info - mi_dict: dict[str, Any] = mi if isinstance(mi, dict) else (mi.model_dump() if mi else {}) + deployment = deployments[0] + mi_dict: dict[str, Any] = deployment.get("model_info") or {} prefs_raw = mi_dict.get("adaptive_router_preferences") if prefs_raw is not None: model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) @@ -3144,9 +3141,7 @@ class ComplexityRouter(CustomLogger): model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[]) # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set. - lp = deployment.get("litellm_params") if isinstance(deployment, dict) else deployment.litellm_params - lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {}) - cost = lp_dict.get("input_cost_per_token") + cost = deployment.get("litellm_params", {}).get("input_cost_per_token") if cost is None: cost = mi_dict.get("input_cost_per_token") model_to_cost[name] = float(cost) if cost is not None else 0.0 diff --git a/litellm/router_strategy/savings_baseline.py b/litellm/router_strategy/savings_baseline.py index b5b80ddae6d..add0f78fd3a 100644 --- a/litellm/router_strategy/savings_baseline.py +++ b/litellm/router_strategy/savings_baseline.py @@ -14,7 +14,7 @@ bare string with no provider beside them; an operator who writes ``deepseek-r1`` Azure would otherwise be priced against whoever else owns that name. """ -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Final, NamedTuple from litellm._logging import verbose_router_logger @@ -80,12 +80,11 @@ def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]: straight at a provider model rather than at a configured group, and prices under its own name because there is no deployment to override it. """ - indices: Final = router.model_name_to_deployment_indices.get(group_name) - if not indices: + deployments: Final = router.deployments_named(group_name) + if not deployments: return (Baseline(qualified),) if (qualified := canonical_model(group_name)) else () - def candidate(index: int) -> Baseline | None: - deployment: Final = router.model_list[index] + def candidate(deployment: Mapping[str, object]) -> Baseline | None: params: Final = deployment.get("litellm_params") if not isinstance(params, dict): return None @@ -98,7 +97,7 @@ def _models_in(router: "Router", group_name: str) -> tuple[Baseline, ...]: deployment_id: Final = info.get("id") if isinstance(info, dict) else None return Baseline(qualified, str(deployment_id) if deployment_id else None) - return tuple(c for index in indices if (c := candidate(index)) is not None) + return tuple(c for deployment in deployments if (c := candidate(deployment)) is not None) def _priced(router: "Router", candidate: Baseline) -> tuple[float, Baseline] | None: diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index b5e36055354..853b10ad97b 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -171,6 +171,7 @@ ignored_function_names = [ "_deployment_field", # Read by is_batch_capable, tested through the mixed-group tests in test_router.py "_deployment_provider", # Read by without_deployments_that_cannot_batch, tested through the mixed-group tests in test_router.py "_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py + "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py ] diff --git a/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py b/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py new file mode 100644 index 00000000000..dc94c96da05 --- /dev/null +++ b/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py @@ -0,0 +1,474 @@ +from __future__ import annotations + +import asyncio +import json +import linecache +import sys +import threading +import uuid +from collections.abc import AsyncIterable, Awaitable, Callable, Iterable +from concurrent.futures import ThreadPoolExecutor +from types import FrameType +from typing import Final, TypeVar + +import pytest +from integration._support.client import object_value +from integration._support.openai_wire import chat_reply, openai_error, responses_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +import litellm +from litellm import Router + +_DEPLOYMENTS: Final = 4 +_TARGET: Final = _DEPLOYMENTS - 1 +_SYNC_SECONDS: Final = 10 +_RATE_LIMITED_KEY: Final = "key-rate-limited" +_BURST_THREADS: Final = 8 +_BURST_CALLS: Final = 24 +_CHURN_UPSERTS: Final = 120 +_UNREPORTED_MODEL_ID: Final = frozenset({"aresponses-stream", "aanthropic-messages"}) +_ROUTER_FILE: Final = Router.__init__.__code__.co_filename +_INDEX_LOCK_LINES: Final = frozenset( + number + for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1) + if "with self._deployment_index_lock" in line +) + +T = TypeVar("T") + + +def _peer(request: Request) -> Reply: + if request.headers["authorization"] == f"Bearer {_RATE_LIMITED_KEY}": + return openai_error(429) + body: Final = object_value(json.loads(request.body)) + stream: Final = body.get("stream") is True + if request.target == "/v1/responses": + return responses_reply(f"resp_{uuid.uuid4().hex}", "gpt-4o-mini", "ok", stream=stream) + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "gpt-4o-mini", "ok", stream=stream) + + +def _deployment( + index: int, + api_base: str, + team: bool, + *, + model: int | None = None, + public_name: str | None = None, + timeout: float = 10, +) -> dict[str, JsonValue]: + team_info: Final[dict[str, JsonValue]] = ( + {"team_id": "team-1", "team_public_model_name": public_name or f"public-{index}"} if team else {} + ) + return { + "model_name": f"model-{index if model is None else model}", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": api_base, + "api_key": f"key-{index}", + "timeout": timeout, + }, + "model_info": {"id": f"deployment-{index}", **team_info}, + } + + +def _router(wire: Wire, team: bool = False) -> Router: + return Router(model_list=[_deployment(index, f"{wire.url}/v1", team) for index in range(_DEPLOYMENTS)]) + + +def _read_during_held_write(read: Callable[[], T], write: Callable[[], object]) -> T: + """Run `read` on its own thread while `write` waits where the list has shrunk but the index maps have not caught + up. The write resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets + the half-updated list and a read that takes the lock waits for the write""" + reindex: Final = Router._update_deployment_indices_after_removal.__code__ + held: Final = threading.Event() + release: Final = threading.Event() + settled: Final = threading.Event() + released_in_time: Final[list[bool]] = [] + outcome: Final[list[T]] = [] + failure: Final[list[BaseException]] = [] + + def hold_at_reindex(frame: FrameType, event: str, _arg: object) -> None: + if event == "call" and frame.f_code is reindex and not held.is_set(): + held.set() + released_in_time.append(release.wait(timeout=_SYNC_SECONDS)) + + def run_write() -> None: + previous: Final = sys.getprofile() + sys.setprofile(hold_at_reindex) + try: + write() + finally: + sys.setprofile(previous) + + def watch_read(frame: FrameType, event: str, _arg: object) -> Callable[..., object] | None: + if frame.f_code.co_filename != _ROUTER_FILE: + return None + if event == "line" and frame.f_lineno in _INDEX_LOCK_LINES: + settled.set() + return watch_read + + def run_read() -> None: + previous: Final = sys.gettrace() + sys.settrace(watch_read) + try: + outcome.append(read()) + except BaseException as error: + failure.append(error) + finally: + sys.settrace(previous) + settled.set() + + writer: Final = threading.Thread(target=run_write) + reader: Final = threading.Thread(target=run_read) + writer.start() + try: + assert held.wait(timeout=_SYNC_SECONDS), "the write never reached re-indexing" + reader.start() + assert settled.wait(timeout=_SYNC_SECONDS), "the read neither returned nor reached the index lock" + finally: + release.set() + if reader.ident is not None: + reader.join() + writer.join() + assert released_in_time == [True], "the write was not held for the whole read" + if failure: + raise failure[0] + return outcome[0] + + +def _read_during_held_delete(router: Router, deployment_id: str, read: Callable[[], T]) -> T: + return _read_during_held_write(read, lambda: router.delete_deployment(deployment_id)) + + +def _keys(wire: Wire) -> tuple[str, ...]: + return tuple(request.headers["authorization"] for request in wire.drain()) + + +def _team_metadata(team: bool) -> dict[str, JsonValue]: + return {"user_api_key_team_id": "team-1"} if team else {} + + +def _completion(router: Router, model: str, team: bool) -> str | None: + response: Final = router.completion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team) + ) + return response._hidden_params["model_id"] + + +def _completion_stream(router: Router, model: str, team: bool) -> str | None: + stream: Final = router.completion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team), stream=True + ) + assert isinstance(stream, Iterable), stream + chunks: Final = list(stream) + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok", chunks + return stream._hidden_params["model_id"] + + +def _acompletion(router: Router, model: str, team: bool) -> str | None: + async def call() -> str | None: + response: Final = await router.acompletion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team) + ) + return response._hidden_params["model_id"] + + return asyncio.run(call()) + + +def _acompletion_stream(router: Router, model: str, team: bool) -> str | None: + async def call() -> str | None: + stream: Final = await router.acompletion( + model=model, messages=[{"role": "user", "content": "hi"}], metadata=_team_metadata(team), stream=True + ) + assert isinstance(stream, AsyncIterable), stream + chunks: Final = [chunk async for chunk in stream] + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == "ok", chunks + return stream._hidden_params["model_id"] + + return asyncio.run(call()) + + +def _responses(router: Router, model: str, team: bool) -> str | None: + response: Final = router.responses(model=model, input="hi", metadata=_team_metadata(team)) + assert response.output_text == "ok", response + return response._hidden_params["model_id"] + + +def _aresponses(router: Router, model: str, team: bool) -> str | None: + async def call() -> str | None: + response: Final = await router.aresponses(model=model, input="hi", metadata=_team_metadata(team)) + assert response.output_text == "ok", response + return response._hidden_params["model_id"] + + return asyncio.run(call()) + + +def _aresponses_stream(router: Router, model: str, team: bool) -> str | None: + async def call() -> None: + stream: Final = await router.aresponses(model=model, input="hi", metadata=_team_metadata(team), stream=True) + assert isinstance(stream, AsyncIterable), stream + events: Final = [event.type async for event in stream] + assert events[-1] == "response.completed", events + + asyncio.run(call()) + return None + + +def _aanthropic_messages(router: Router, model: str, team: bool) -> str | None: + async def call() -> None: + response: Final = await router.aanthropic_messages( + model=model, messages=[{"role": "user", "content": "hi"}], max_tokens=16, metadata=_team_metadata(team) + ) + assert response["content"][0]["text"] == "ok", response + + asyncio.run(call()) + return None + + +_CALLS: Final[dict[str, Callable[[Router, str, bool], str | None]]] = { + "completion": _completion, + "completion-stream": _completion_stream, + "acompletion": _acompletion, + "acompletion-stream": _acompletion_stream, + "responses": _responses, + "aresponses": _aresponses, + "aresponses-stream": _aresponses_stream, + "aanthropic-messages": _aanthropic_messages, +} + + +@pytest.mark.parametrize("call", list(_CALLS)) +@pytest.mark.parametrize("team", [False, True], ids=["model-name", "team-public-name"]) +def test_call_reaches_its_deployment_while_another_thread_deletes_one(call: str, team: bool) -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire, team) + served_by: Final = _read_during_held_delete( + router, + "deployment-0", + lambda: _CALLS[call](router, f"public-{_TARGET}" if team else f"model-{_TARGET}", team), + ) + keys: Final = _keys(wire) + + assert served_by == (None if call in _UNREPORTED_MODEL_ID else f"deployment-{_TARGET}") + assert keys == (f"Bearer key-{_TARGET}",) + assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)] + + +@pytest.mark.parametrize("deleted", [0, _TARGET - 1], ids=["first", "next-to-target"]) +def test_completion_reaches_its_deployment_whichever_earlier_deployment_is_deleted(deleted: int) -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + served_by: Final = _read_during_held_delete( + router, f"deployment-{deleted}", lambda: _completion(router, f"model-{_TARGET}", team=False) + ) + keys: Final = _keys(wire) + + assert served_by == f"deployment-{_TARGET}" + assert keys == (f"Bearer key-{_TARGET}",) + + +def test_team_completion_reaches_the_remaining_deployment_of_a_shared_public_name_while_its_sibling_is_deleted() -> ( + None +): + with wire_server(_peer) as wire: + api_base: Final = f"{wire.url}/v1" + router: Final = Router( + model_list=[ + _deployment(0, api_base, team=False), + _deployment(1, api_base, team=False), + _deployment(2, api_base, team=True, public_name="public-shared"), + _deployment(3, api_base, team=True, public_name="public-shared"), + ] + ) + served_by: Final = _read_during_held_delete( + router, "deployment-2", lambda: _completion(router, "public-shared", team=True) + ) + keys: Final = _keys(wire) + + assert served_by == "deployment-3" + assert keys == ("Bearer key-3",) + + +def test_completion_for_the_deployment_being_deleted_fails_cleanly_and_sends_nothing() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + with pytest.raises(litellm.BadRequestError): + _read_during_held_delete( + router, + "deployment-0", + lambda: router.completion(model="model-0", messages=[{"role": "user", "content": "hi"}]), + ) + keys: Final = _keys(wire) + + assert keys == () + assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)] + + +def test_provider_error_during_a_concurrent_delete_comes_from_the_requested_deployment() -> None: + with wire_server(_peer) as wire: + api_base: Final = f"{wire.url}/v1" + rate_limited: Final = _deployment(_TARGET, api_base, team=False) + rate_limited["litellm_params"] = {**object_value(rate_limited["litellm_params"]), "api_key": _RATE_LIMITED_KEY} + router: Final = Router( + model_list=[*(_deployment(index, api_base, team=False) for index in range(_TARGET)), rate_limited], + num_retries=0, + ) + with pytest.raises(litellm.RateLimitError): + _read_during_held_delete( + router, + "deployment-0", + lambda: router.completion(model=f"model-{_TARGET}", messages=[{"role": "user", "content": "hi"}]), + ) + served_by: Final = _completion(router, "model-1", team=False) + keys: Final = _keys(wire) + + assert served_by == "deployment-1" + assert keys == (f"Bearer {_RATE_LIMITED_KEY}", "Bearer key-1") + + +def test_two_threads_deleting_the_same_deployment_remove_it_once() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + second: Final = _read_during_held_delete( + router, "deployment-0", lambda: router.delete_deployment("deployment-0") + ) + served_by: Final = _completion(router, f"model-{_TARGET}", team=False) + keys: Final = _keys(wire) + + assert second is None + assert served_by == f"deployment-{_TARGET}" + assert keys == (f"Bearer key-{_TARGET}",) + assert router.get_model_ids() == [f"deployment-{index}" for index in range(1, _DEPLOYMENTS)] + + +def test_completion_reaches_its_deployment_while_another_thread_upserts_an_earlier_one() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + changed: Final = litellm.Deployment(**_deployment(0, f"{wire.url}/v1", team=False, timeout=11)) + served_by: Final = _read_during_held_write( + lambda: _completion(router, f"model-{_TARGET}", team=False), lambda: router.upsert_deployment(changed) + ) + keys: Final = _keys(wire) + + assert served_by == f"deployment-{_TARGET}" + assert keys == (f"Bearer key-{_TARGET}",) + assert sorted(router.get_model_ids()) == [f"deployment-{index}" for index in range(_DEPLOYMENTS)] + + +def test_deployment_added_during_a_concurrent_delete_serves_its_own_requests() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + added: Final = litellm.Deployment(**_deployment(9, f"{wire.url}/v1", team=False)) + _read_during_held_delete(router, "deployment-0", lambda: router.add_deployment(added)) + served: Final = tuple(_completion(router, f"model-{index}", team=False) for index in (9, _TARGET)) + keys: Final = _keys(wire) + + assert served == ("deployment-9", f"deployment-{_TARGET}") + assert keys == ("Bearer key-9", f"Bearer key-{_TARGET}") + + +def test_model_list_set_during_a_concurrent_delete_serves_its_own_requests() -> None: + with wire_server(_peer) as wire: + router: Final = _router(wire) + replacement: Final = [_deployment(index, f"{wire.url}/v1", team=False) for index in range(10, 14)] + _read_during_held_delete(router, "deployment-0", lambda: router.set_model_list(replacement)) + served: Final = tuple(_completion(router, f"model-{index}", team=False) for index in range(10, 14)) + keys: Final = _keys(wire) + + assert served == tuple(f"deployment-{index}" for index in range(10, 14)) + assert keys == tuple(f"Bearer key-{index}" for index in range(10, 14)) + assert router.get_model_ids() == [f"deployment-{index}" for index in range(10, 14)] + + +def _burst_router(api_base: str) -> Router: + """Two deployments per model name, so an upsert of one leaves its model group served.""" + return Router( + model_list=[ + _deployment(index, api_base, team=False, model=index % _DEPLOYMENTS) for index in range(2 * _DEPLOYMENTS) + ] + ) + + +def _churn(router: Router, api_base: str, stop: threading.Event) -> int: + """Upsert the deployments in turn with a changed timeout; each upsert moves one to the end of the list.""" + upserts = 0 + while upserts < _CHURN_UPSERTS and not stop.is_set(): + index, round_number = upserts % (2 * _DEPLOYMENTS), upserts // (2 * _DEPLOYMENTS) + router.upsert_deployment( + litellm.Deployment( + **_deployment(index, api_base, team=False, model=index % _DEPLOYMENTS, timeout=11 + round_number % 2) + ) + ) + upserts += 1 + return upserts + + +def _served_keys(wire: Wire) -> dict[str, str]: + """Which key reached the upstream for each request marker.""" + return { + object_value(json.loads(request.body))["messages"][0]["content"]: request.headers["authorization"] + for request in wire.drain() + } + + +def _assert_each_request_used_its_models_deployment(served: list[tuple[str, str | None]], keys: dict[str, str]) -> None: + for marker, model_id in served: + model: Final = int(marker.split(":")[0]) + assert model_id in (f"deployment-{model}", f"deployment-{model + _DEPLOYMENTS}"), (marker, model_id) + assert keys[marker] == f"Bearer key-{str(model_id).removeprefix('deployment-')}", (marker, keys[marker]) + assert len(keys) == len(served) + + +def test_thread_burst_during_upsert_churn_sends_every_request_with_its_deployments_key() -> None: + with wire_server(_peer) as wire: + api_base: Final = f"{wire.url}/v1" + router: Final = _burst_router(api_base) + stop: Final = threading.Event() + + def call(number: int) -> tuple[str, str | None]: + marker: Final = f"{number % _DEPLOYMENTS}:{number}" + response: Final = router.completion( + model=f"model-{number % _DEPLOYMENTS}", messages=[{"role": "user", "content": marker}] + ) + return marker, response._hidden_params["model_id"] + + with ThreadPoolExecutor(max_workers=_BURST_THREADS + 1) as pool: + churn: Final = pool.submit(_churn, router, api_base, stop) + served: Final = list(pool.map(call, range(_BURST_THREADS * _BURST_CALLS))) + stop.set() + upserts: Final = churn.result() + keys: Final = _served_keys(wire) + + assert upserts > 0 + _assert_each_request_used_its_models_deployment(served, keys) + assert sorted(router.get_model_ids()) == sorted(f"deployment-{index}" for index in range(2 * _DEPLOYMENTS)) + + +async def _gathered(calls: Iterable[Awaitable[tuple[str, str | None]]]) -> list[tuple[str, str | None]]: + return list(await asyncio.gather(*calls)) + + +def test_async_burst_during_upsert_churn_sends_every_request_with_its_deployments_key() -> None: + with wire_server(_peer) as wire: + api_base: Final = f"{wire.url}/v1" + router: Final = _burst_router(api_base) + stop: Final = threading.Event() + + async def call(number: int) -> tuple[str, str | None]: + marker: Final = f"{number % _DEPLOYMENTS}:{number}" + response: Final = await router.acompletion( + model=f"model-{number % _DEPLOYMENTS}", messages=[{"role": "user", "content": marker}] + ) + return marker, response._hidden_params["model_id"] + + with ThreadPoolExecutor(max_workers=1) as pool: + churn: Final = pool.submit(_churn, router, api_base, stop) + served: Final = asyncio.run(_gathered(call(number) for number in range(_BURST_THREADS * _BURST_CALLS))) + stop.set() + upserts: Final = churn.result() + keys: Final = _served_keys(wire) + + assert upserts > 0 + _assert_each_request_used_its_models_deployment(served, keys) + assert sorted(router.get_model_ids()) == sorted(f"deployment-{index}" for index in range(2 * _DEPLOYMENTS)) diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 979ab499f3b..63e38ce533a 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -1011,10 +1011,9 @@ class TestPollingConditionChecks: # Create mock router mock_router = Mock() - mock_router.model_name_to_deployment_indices = {"gpt-5": [0]} - mock_router.model_list = [ - {"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}} - ] + mock_router.deployments_named.side_effect = { + "gpt-5": ({"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5"}},) + }.get result = should_use_polling_for_request( background_mode=True, @@ -1033,13 +1032,14 @@ class TestPollingConditionChecks: ) mock_router = Mock() - mock_router.model_name_to_deployment_indices = {"claude-3": [0]} - mock_router.model_list = [ - { - "model_name": "claude-3", - "litellm_params": {"model": "anthropic/claude-3-sonnet"}, - } - ] + mock_router.deployments_named.side_effect = { + "claude-3": ( + { + "model_name": "claude-3", + "litellm_params": {"model": "anthropic/claude-3-sonnet"}, + }, + ) + }.get result = should_use_polling_for_request( background_mode=True, @@ -1183,13 +1183,14 @@ class TestPollingConditionChecks: ) mock_router = Mock() - mock_router.model_name_to_deployment_indices = {"deep-research": [0]} - mock_router.model_list = [ - { - "model_name": "deep-research", - "litellm_params": {"model": "openai/o4-mini-deep-research"}, - } - ] + mock_router.deployments_named.side_effect = { + "deep-research": ( + { + "model_name": "deep-research", + "litellm_params": {"model": "openai/o4-mini-deep-research"}, + }, + ) + }.get # Model alias "deep-research" is in native_background_mode result = should_use_polling_for_request( diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 9b4019bead6..a46f415e10a 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -112,6 +112,13 @@ def _heuristic_v2_artifact() -> TrainedTierArtifact: ) +def _answer_deployments_named(router: MagicMock) -> None: + """Make a mock router answer `deployments_named` from its own `model_list`, the way Router answers from its index.""" + router.deployments_named.side_effect = lambda name: ( + tuple(deployment for deployment in router.model_list if deployment["model_name"] == name) or None + ) + + @pytest.fixture def mock_router_instance(): """Create a mock LiteLLM Router instance.""" @@ -3181,7 +3188,7 @@ class TestCapabilityClassifier: {"model_name": "efficient-model", "litellm_params": {"input_cost_per_token": 0.000001}}, {"model_name": "capable-model", "litellm_params": {"input_cost_per_token": 0.00001}}, ] - mock_router_instance.model_name_to_deployment_indices = {"efficient-model": [0], "capable-model": [1]} + _answer_deployments_named(mock_router_instance) router: Final = self._router( mock_router_instance, adaptive=True, @@ -5670,7 +5677,7 @@ class TestAdaptiveSoftFloors: "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, }, ] - router.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]} + _answer_deployments_named(router) return router @pytest.fixture @@ -6829,7 +6836,7 @@ class TestTierModelAffinity: ) -> ComplexityRouter: mock_router_instance.cache = DualCache() mock_router_instance.model_list = [] - mock_router_instance.model_name_to_deployment_indices = {} + _answer_deployments_named(mock_router_instance) return ComplexityRouter( model_name="affinity-router", litellm_router_instance=mock_router_instance, @@ -7547,7 +7554,7 @@ class TestSessionAffinity: "model_info": {}, }, ] - mock_router_instance.model_name_to_deployment_indices = {"cheap": [0]} + _answer_deployments_named(mock_router_instance) router = ComplexityRouter( model_name="hybrid", litellm_router_instance=mock_router_instance, @@ -11351,7 +11358,7 @@ class TestClientHousekeepingCalls: "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, }, ] - adaptive_instance.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]} + _answer_deployments_named(adaptive_instance) router = ComplexityRouter( model_name="hybrid", litellm_router_instance=adaptive_instance, @@ -13442,7 +13449,7 @@ class TestPlanModeTierFloor: "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": []}}, }, ] - adaptive_instance.model_name_to_deployment_indices = {"cheap": [0], "premium": [1]} + _answer_deployments_named(adaptive_instance) router = ComplexityRouter( model_name="hybrid", litellm_router_instance=adaptive_instance, @@ -14941,7 +14948,7 @@ class TestModalityRouting: elif path == "adaptive_pick_rewritten": config["adaptive"] = True mock_router_instance.model_list = [] - mock_router_instance.model_name_to_deployment_indices = {} + _answer_deployments_named(mock_router_instance) router = self._router(mock_router_instance, config, vision) result = await router.async_pre_routing_hook(model="m", request_kwargs=request_kwargs, messages=messages) assert result.model == expected_model diff --git a/tests/unit/router_strategy/test_savings_baseline.py b/tests/unit/router_strategy/test_savings_baseline.py index 5efc73d2dcb..4241a541202 100644 --- a/tests/unit/router_strategy/test_savings_baseline.py +++ b/tests/unit/router_strategy/test_savings_baseline.py @@ -107,8 +107,7 @@ class TestResolveBaseline: served; a dashboard counterfactual must not be able to take routing down.""" class Exploding: - @property - def model_name_to_deployment_indices(self): + def deployments_named(self, model_name): raise RuntimeError("router is mid-reload") assert resolve_baseline(Exploding(), ["anything"]) is None diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 568a05e8ff4..4814d36ee39 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -5,6 +5,7 @@ import functools import gc import itertools import json +import linecache import logging import os import sys @@ -13,8 +14,8 @@ import warnings from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping from datetime import datetime, timedelta, timezone from pathlib import Path -from types import MappingProxyType, SimpleNamespace -from typing import TYPE_CHECKING, Final, Literal, cast +from types import CodeType, FrameType, MappingProxyType, SimpleNamespace +from typing import TYPE_CHECKING, Final, Literal, TypeVar, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -60,7 +61,10 @@ from litellm.router import ( _without_line_breaks, Span, ) +from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router_strategy import simple_shuffle +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.router_strategy.savings_baseline import Baseline, _models_in from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.router_utils.cooldown_handlers import ( _async_get_cooldown_deployments, @@ -30004,6 +30008,319 @@ async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_s assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + fallback_tokens +_SYNC_SECONDS: Final = 10 +_ROUTER_FILE: Final = Router.__init__.__code__.co_filename +_INDEX_LOCK_LINES: Final = frozenset( + number + for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1) + if "with self._deployment_index_lock" in line +) + +_Read = TypeVar("_Read") + + +def _concurrent_delete_deployment(index: int) -> dict[str, object]: + return { + "model_name": f"model-{index}", + "litellm_params": {"model": f"openai/gpt-{index}", "api_key": "k"}, + "model_info": { + "id": f"deployment-{index}", + "team_id": "team-1", + "team_public_model_name": "public-shared" if index >= 2 else f"public-{index}", + }, + } + + +def _concurrent_delete_router() -> Router: + return Router(model_list=[_concurrent_delete_deployment(index) for index in range(4)]) + + +def _read_during_held_write( + read: Callable[[], _Read], + write: Callable[[], object], + event: Literal["call", "c_call"] = "call", + write_after: CodeType | None = None, +) -> _Read: + """Run `read` on its own thread while `write` waits at its first `event` in the re-index, after the list shrank and + before the maps caught up. With `write_after`, the write starts when `read` returns from that function. The write + resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets the half-updated + list and a read that takes the lock waits for the write""" + reindex: Final = Router._update_deployment_indices_after_removal.__code__ + held: Final = threading.Event() + release: Final = threading.Event() + settled: Final = threading.Event() + released_in_time: Final[list[bool]] = [] + outcome: Final[list[_Read]] = [] + failure: Final[list[BaseException]] = [] + + def hold_in_reindex(frame: FrameType, observed: str, _arg: object) -> None: + if observed == event and frame.f_code is reindex and not held.is_set(): + held.set() + released_in_time.append(release.wait(timeout=_SYNC_SECONDS)) + + def run_write() -> None: + previous: Final = sys.getprofile() + sys.setprofile(hold_in_reindex) + try: + write() + finally: + sys.setprofile(previous) + + writer: Final = threading.Thread(target=run_write) + + def start_write() -> None: + writer.start() + assert held.wait(timeout=_SYNC_SECONDS), "the write never reached re-indexing" + + def watch_read(frame: FrameType, observed: str, _arg: object) -> Callable[..., object] | None: + if frame.f_code.co_filename != _ROUTER_FILE: + return None + if observed == "return" and frame.f_code is write_after and writer.ident is None: + start_write() + elif observed == "line" and held.is_set() and frame.f_lineno in _INDEX_LOCK_LINES: + settled.set() + return watch_read + + def run_read() -> None: + previous: Final = sys.gettrace() + sys.settrace(watch_read) + try: + outcome.append(read()) + except BaseException as error: + failure.append(error) + finally: + sys.settrace(previous) + settled.set() + + reader: Final = threading.Thread(target=run_read) + try: + if write_after is None: + start_write() + reader.start() + assert settled.wait(timeout=_SYNC_SECONDS), "the read neither returned nor reached the index lock" + finally: + release.set() + if reader.ident is not None: + reader.join() + if writer.ident is not None: + writer.join() + assert released_in_time == [True], "the write was not held for the whole read" + if failure: + raise failure[0] + return outcome[0] + + +def _deployment_ids(deployments: list[DeploymentTypedDict]) -> tuple[str, ...]: + return tuple(deployment["model_info"]["id"] for deployment in deployments) + + +@pytest.mark.parametrize( + ("read", "expected"), + [ + pytest.param( + lambda router: router.get_deployment("deployment-3").model_info.id, "deployment-3", id="get_deployment" + ), + pytest.param( + lambda router: router.get_model_info("deployment-3")["model_info"]["id"], + "deployment-3", + id="get_model_info", + ), + pytest.param( + lambda router: _deployment_ids(router.get_all_deployments("model-3")), + ("deployment-3",), + id="get_all_deployments", + ), + pytest.param( + lambda router: _deployment_ids(router.get_all_deployments("public-shared", team_id="team-1")), + ("deployment-2", "deployment-3"), + id="get_all_deployments-team", + ), + pytest.param( + lambda router: router.map_team_model("public-shared", "team-1"), "public-shared", id="map_team_model" + ), + pytest.param(lambda router: router.get_model_ids("model-3"), ["deployment-3"], id="get_model_ids-by-name"), + pytest.param( + lambda router: router.get_model_ids(), ["deployment-1", "deployment-2", "deployment-3"], id="get_model_ids" + ), + pytest.param( + lambda router: router.get_deployment_by_model_group_name("model-3").model_info.id, + "deployment-3", + id="get_deployment_by_model_group_name", + ), + pytest.param( + lambda router: router.get_model_listing_info("model-3").cost_map_keys, + ("openai/gpt-3",), + id="get_model_listing_info", + ), + pytest.param( + lambda router: router.team_model_has_alternatives("deployment-3"), True, id="team_model_has_alternatives" + ), + pytest.param( + lambda router: ( + router._get_model_group_deployment_usable_by_team( + model_group_name="model-3", team_id="team-1" + ).model_info.id + ), + "deployment-3", + id="model_group_deployment_usable_by_team", + ), + pytest.param( + lambda router: ( + router._get_team_public_name_deployment(model_id="public-shared", team_id="team-1").model_info.id + ), + "deployment-2", + id="team_public_name_deployment", + ), + pytest.param( + lambda router: _deployment_ids(router._team_deployments_across_teams("public-shared")), + ("deployment-2", "deployment-3"), + id="team_deployments_across_teams", + ), + pytest.param( + lambda router: _deployment_ids(list(router.deployments_named("model-3") or ())), + ("deployment-3",), + id="deployments_named", + ), + pytest.param( + lambda router: should_use_polling_for_request( + background_mode=True, + polling_via_cache_enabled=["openai"], + redis_cache=MagicMock(), + model="model-3", + llm_router=router, + ), + True, + id="response_polling_provider_lookup", + ), + pytest.param( + lambda router: _models_in(router, "model-3"), + (Baseline("openai/gpt-3", "deployment-3"),), + id="savings_baseline_candidates", + ), + ], +) +def test_lookup_during_a_concurrent_delete_returns_the_requested_deployment( + read: Callable[[Router], object], expected: object +) -> None: + router: Final = _concurrent_delete_router() + + observed: Final = _read_during_held_write(lambda: read(router), lambda: router.delete_deployment("deployment-0")) + + assert observed == expected + + +def test_lookup_during_an_upsert_returns_the_requested_deployment() -> None: + router: Final = _concurrent_delete_router() + rotated: Final = Deployment( + **{**_concurrent_delete_deployment(0), "litellm_params": {"model": "openai/gpt-0", "api_key": "rotated"}} + ) + + observed: Final = _read_during_held_write( + lambda: router.get_deployment("deployment-3").model_info.id, lambda: router.upsert_deployment(rotated) + ) + + assert observed == "deployment-3" + + +def test_add_during_a_concurrent_delete_indexes_the_added_deployment() -> None: + router: Final = _concurrent_delete_router() + + _read_during_held_write( + lambda: router.add_deployment(Deployment(**_concurrent_delete_deployment(4))), + lambda: router.delete_deployment("deployment-0"), + ) + + assert [router.get_deployment(f"deployment-{index}").model_info.id for index in range(1, 5)] == [ + f"deployment-{index}" for index in range(1, 5) + ] + + +def test_set_model_list_during_a_concurrent_delete_indexes_the_new_list() -> None: + router: Final = _concurrent_delete_router() + + _read_during_held_write( + lambda: router.set_model_list([_concurrent_delete_deployment(index) for index in range(10, 14)]), + lambda: router.delete_deployment("deployment-0"), + event="c_call", + ) + + assert router.get_model_ids() == [f"deployment-{index}" for index in range(10, 14)] + + +def test_team_model_has_alternatives_when_a_delete_starts_after_its_deployment_lookup() -> None: + router: Final = _concurrent_delete_router() + + has_alternatives: Final = _read_during_held_write( + lambda: router.team_model_has_alternatives("deployment-3"), + lambda: router.delete_deployment("deployment-0"), + write_after=Router.get_deployment.__code__, + ) + + assert has_alternatives is True + + +def test_init_adaptive_router_deployment_during_a_concurrent_delete_prices_the_listed_model() -> None: + priced: Final = { + **_concurrent_delete_deployment(3), + "litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06}, + } + router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) + adaptive: Final = Deployment( + model_name="adaptive", + litellm_params=LiteLLM_Params( + model="auto_router/adaptive_router", adaptive_router_config={"available_models": ["model-3"]} + ), + model_info=ModelInfo(id="adaptive"), + ) + + _read_during_held_write( + lambda: router.init_adaptive_router_deployment(deployment=adaptive), + lambda: router.delete_deployment("deployment-0"), + ) + + assert router.adaptive_routers["adaptive"][0].strategy.model_to_cost == {"model-3": 3e-06} + + +def test_complexity_router_adaptive_pool_during_a_concurrent_delete_prices_the_listed_model() -> None: + priced: Final = { + **_concurrent_delete_deployment(3), + "litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06}, + } + router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) + complexity: Final = ComplexityRouter( + model_name="hybrid", + litellm_router_instance=router, + complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["model-3"]}}, + ) + + adaptive: Final = _read_during_held_write( + complexity._ensure_adaptive_router, lambda: router.delete_deployment("deployment-0") + ) + + assert adaptive is not None + assert adaptive.model_to_cost == {"model-3": 3e-06} + + +def test_delete_of_an_unknown_deployment_returns_none_and_keeps_the_list() -> None: + router: Final = _concurrent_delete_router() + + assert router.delete_deployment("deployment-9") is None + assert router.get_model_ids() == [f"deployment-{index}" for index in range(4)] + + +def test_get_model_ids_by_name_skips_team_deployments_only_when_asked() -> None: + shared: Final = {"model_name": "shared", "litellm_params": {"model": "openai/gpt-shared", "api_key": "k"}} + router: Final = Router( + model_list=[ + {**shared, "model_info": {"id": "global"}}, + {**shared, "model_info": {"id": "team", "team_id": "team-1", "team_public_model_name": "public"}}, + ] + ) + + assert router.get_model_ids("shared", exclude_team_models=True) == ["global"] + assert router.get_model_ids("shared") == ["global", "team"] + + _IMAGE_PART = {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}} _IMAGE_CHAT = [{"role": "user", "content": [{"type": "text", "text": "What is in this image?"}, _IMAGE_PART]}] _IMAGE_RESPONSES_INPUT = [ From 4da24fa01e2248a44c3a62e2022b43f5a38794b8 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Mon, 26 Oct 2026 15:42:17 -0700 Subject: [PATCH 16/19] fix(ui): show search tool errors instead of empty results (#48478) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../_components/SearchToolTester.test.tsx | 32 +++++++++++++++++-- .../_components/SearchToolTester.tsx | 2 +- .../src/components/networking.test.ts | 26 +++++++++++++++ .../src/components/networking.tsx | 2 +- 4 files changed, 57 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.test.tsx index 54be24d2401..08a0ab7d02b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.test.tsx @@ -1,4 +1,4 @@ -import { render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { SearchToolTester } from "./SearchToolTester"; @@ -271,7 +271,15 @@ describe("SearchToolTester", () => { }); it("should handle search errors and show notification", async () => { - const error = new Error("Search failed"); + const error = new Error( + JSON.stringify({ + error: { + message: "User not allowed to access search tool: exa-search-2", + type: "user_search_tool_access_denied", + code: "403", + }, + }), + ); vi.mocked(networking.searchToolQueryCall).mockRejectedValue(error); const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); const user = userEvent.setup(); @@ -281,11 +289,29 @@ describe("SearchToolTester", () => { const searchButton = screen.getByRole("button", { name: /search/i }); await user.click(searchButton); await waitFor(() => { - expect(toast.fromError).toHaveBeenCalledWith("Failed to query search tool"); + expect(toast.fromError).toHaveBeenCalledWith(error); }); consoleSpy.mockRestore(); }); + it("should not show empty results or add history when search fails", async () => { + vi.mocked(networking.searchToolQueryCall).mockRejectedValue(new Error("Search failed")); + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {}); + const user = userEvent.setup(); + render(); + const input = screen.getByPlaceholderText("Enter your search query..."); + fireEvent.change(input, { target: { value: "test query" } }); + await user.click(screen.getByRole("button", { name: /search/i })); + + await waitFor(() => { + expect(toast.fromError).toHaveBeenCalled(); + }); + expect(screen.queryByText("No results found")).not.toBeInTheDocument(); + expect(screen.queryByText("test query")).not.toBeInTheDocument(); + expect(screen.getByText("Test your search tool")).toBeInTheDocument(); + consoleSpy.mockRestore(); + }); + it("should maintain search history", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx index 7dcef3fc6ba..646120443c7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTester.tsx @@ -60,7 +60,7 @@ export const SearchToolTester: React.FC = ({ searchToolNa setSearchHistory((prev) => [historyEntry, ...prev]); } catch (error) { console.error("Error querying search tool:", error); - toast.fromError("Failed to query search tool"); + toast.fromError(error); } finally { setIsLoading(false); } diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index e7a55870345..8b0f0389772 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -1178,3 +1178,29 @@ describe("getUiConfig runtime backend discovery", () => { }); }); }); + +describe("searchToolQueryCall", () => { + const originalFetch = global.fetch; + + afterEach(() => { + global.fetch = originalFetch; + }); + + it("rejects with the backend error message for a denied search", async () => { + const errorMessage = "User not allowed to access search tool: exa-search-2"; + global.fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + error: { + message: errorMessage, + type: "user_search_tool_access_denied", + code: "403", + }, + }), + { status: 403 }, + ), + ); + + await expect(Networking.searchToolQueryCall("token", "exa-search-2", "litellm")).rejects.toThrow(errorMessage); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index fd25242ec34..9aade717566 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6860,7 +6860,7 @@ export const searchToolQueryCall = async ( if (!response.ok) { const errorData = await response.text(); await handleError(errorData); - return null; + throw new Error(errorData); } const data = await response.json(); From 2e27a02a9fce9bdc6ece13ba3cb8222dc5b6ecfa Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:44:05 -0700 Subject: [PATCH 17/19] feat(chatgpt): forward client subscriptions through the gateway (#41522) * feat(chatgpt): forward each Codex user's ChatGPT OAuth login through the proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): cover _model_group_llm_provider metadata resolution Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(chatgpt): keep the gateway ChatGPT login bound to the authenticator api_base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(chatgpt): move client credential header selection into the chatgpt provider Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(chatgpt): cover credential recognition and stored account id fallbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(chatgpt): support per-model chatgpt_auth_file for multiple ChatGPT accounts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(chatgpt): resolve bare chatgpt_auth_file paths and drop redundant test docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(chatgpt): pass per-model auth file to early provider detection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(chatgpt): normalize chatgpt_auth_file before caching the authenticator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(chatgpt): type account selection parameters * fix(chatgpt): preserve account selection through saved prompts * fix(chatgpt): allow administrator account configuration through the proxy * test(chatgpt): construct immutable request parameters * fix(router): preserve typed deployment credential resolution * fix(chatgpt): allow admins to probe unsaved account configurations * fix(chatgpt): retain SDK login headers during provider resolution * fix(chatgpt): allow admins to probe unsaved account configurations * fix(chatgpt): normalize subscription health probe input * fix(chatgpt): normalize health probe input and cover credential destinations * test(chatgpt): keep destination rejection assertions focused * feat(proxy): forward client ChatGPT subscriptions on native Codex routes * docs: lead ChatGPT forwarding with standard model setup * fix(chatgpt): preserve selected credentials in health provider lookup * fix(chatgpt): preserve gateway credentials and workspace selection * refactor(chatgpt): avoid rebinding streaming metadata inputs * fix(chatgpt): keep account lookup token argument immutable * fix(chatgpt): scope forwarded credential routing to ChatGPT * refactor: keep forwarded credential isolation opt-in by provider --------- Co-authored-by: Moe Khalil Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Urbanonymous Co-authored-by: mateo --- .../chatgpt_subscription_passthrough.md | 76 +++ .../litellm_core_utils/get_litellm_params.py | 1 + .../get_llm_provider_logic.py | 12 +- .../get_provider_specific_headers.py | 13 + litellm/llms/chatgpt/authenticator.py | 85 +-- litellm/llms/chatgpt/chat/transformation.py | 83 ++- litellm/llms/chatgpt/common_utils.py | 169 ++++- .../llms/chatgpt/responses/transformation.py | 38 +- litellm/llms/custom_httpx/llm_http_handler.py | 48 +- litellm/main.py | 13 +- litellm/proxy/_lazy_features.py | 5 + litellm/proxy/_lazy_openapi_snapshot.json | 178 +++++ litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_utils.py | 4 +- .../health_endpoints/_health_endpoints.py | 10 +- litellm/proxy/litellm_pre_call_utils.py | 15 +- .../codex_passthrough_endpoints.py | 87 +++ .../openai_passthrough_logging_handler.py | 9 +- .../pass_through_endpoints.py | 15 +- litellm/proxy/spend_tracking/savings.py | 5 +- litellm/responses/main.py | 53 +- litellm/responses/utils.py | 20 + litellm/router.py | 62 +- litellm/types/litellm_params.py | 1 + litellm/types/router.py | 4 + litellm/types/utils.py | 1 + litellm/utils.py | 6 +- .../test_get_llm_provider_logic.py | 49 ++ .../chat/test_chatgpt_chat_transformation.py | 638 ++++++++++++++++++ .../test_chatgpt_responses_transformation.py | 165 +++++ .../chatgpt/test_chatgpt_authenticator.py | 148 ++++ .../llms/chatgpt/test_chatgpt_common_utils.py | 90 +++ .../custom_httpx/test_llm_http_handler.py | 41 ++ tests/unit/proxy/auth/test_auth_utils.py | 36 + .../health_endpoints/test_health_endpoints.py | 61 ++ .../test_model_management_endpoints.py | 24 +- .../test_codex_passthrough_endpoints.py | 185 +++++ .../unit/proxy/test_litellm_pre_call_utils.py | 92 +++ tests/unit/responses/test_responses_utils.py | 36 + tests/unit/test_router/test_router.py | 9 + tests/unit/types/test_litellm_params.py | 1 + tests/unit/types/test_router.py | 3 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 105 +++ 43 files changed, 2547 insertions(+), 150 deletions(-) create mode 100644 cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md create mode 100644 litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py create mode 100644 tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py create mode 100644 tests/unit/llms/chatgpt/test_chatgpt_common_utils.py create mode 100644 tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py diff --git a/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md new file mode 100644 index 00000000000..c79632e9195 --- /dev/null +++ b/cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md @@ -0,0 +1,76 @@ +# Use your ChatGPT subscription through LiteLLM + +Sign into ChatGPT in Codex or Pi and send model requests through your existing LiteLLM proxy. The client keeps your login and refreshes it. Each request uses that client's ChatGPT account, while a separate LiteLLM virtual key controls gateway access and usage attribution. + +An administrator adds a `chatgpt/` model using the usual model configuration. No ChatGPT login needs to be added to the proxy. This uses the same separate-key convention as provider passthrough, including Claude Code subscription forwarding. + +## Add a model to the proxy + +Reuse your existing proxy. If you need one, follow the [proxy quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start). Merge this entry into its model list and restart the proxy: + +```yaml +model_list: + - model_name: gpt-6-luna + litellm_params: + model: chatgpt/gpt-6-luna + model_info: + mode: responses +``` + +Use a model available to the caller's subscription. The virtual key must allow the exposed model name. Leave saved ChatGPT logins and API credentials unset for a client-only deployment. Explicit cross-provider fallbacks can incur separate API charges; do not configure them for subscription-only traffic. + +## Configure Codex + +Use your existing LiteLLM URL and a virtual key assigned to your gateway user. Export that key as `LITELLM_API_KEY` in the environment that launches Codex. Keep the value out of committed configuration. + +Add this provider to `~/.codex/config.toml`, preserving your other settings: + +```toml +model = "gpt-6-luna" +model_provider = "litellm_subscription" + +[model_providers.litellm_subscription] +name = "LiteLLM subscription" +base_url = "https://litellm.example.com/v1" +wire_api = "responses" +supports_websockets = false +requires_openai_auth = true +env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" } +``` + +Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway. + +For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway). + +## Native compatibility and Pi + +For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name. + +### Configure Pi + +Sign into the built-in `openai-codex` provider in Pi. Merge this override into `~/.pi/agent/models.json` and export your LiteLLM virtual key as `LITELLM_API_KEY` before launching Pi: + +```json +{ + "providers": { + "openai-codex": { + "baseUrl": "https://litellm.example.com/backend-api/codex", + "headers": { + "x-litellm-api-key": "LITELLM_API_KEY" + } + } + } +} +``` + +Keep Pi's subscription login as the provider credential. The header value names an environment variable, not a key to paste into the file. + +## Verify account usage + +Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person. + +The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions. + +The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate. + +This flow does not connect accounts through the dashboard or pool subscription allowances. Use a Codex build that supports ChatGPT authentication with a custom model provider. The integration forwards HTTP and SSE requests; it does not add native WebSocket or other Codex backend routes. diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 85557090ae9..ca8e2af18e6 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -76,6 +76,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "itpm", "otpm", "use_xai_oauth", + "chatgpt_auth_file", "github_copilot_auth_type", "github_copilot_host", "github_copilot_user_session", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index de51ecd0b21..19c1d5c414c 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -158,6 +158,14 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None +def resolve_model_and_provider_for_metadata(model: str, custom_llm_provider: str | None = None) -> tuple[str, str]: + declared: Final = declared_authenticating_provider(model, custom_llm_provider) + if declared is not None: + return model.removeprefix(f"{declared}/"), declared + resolved_model, resolved_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + return resolved_model, resolved_provider + + def inferred_provider(model: str | None) -> str | None: if not model: return None @@ -842,7 +850,9 @@ def _get_openai_compatible_provider_info( api_base, dynamic_api_key, custom_llm_provider, - ) = litellm.ChatGPTConfig().get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider) + ) = litellm.ChatGPTConfig().get_openai_compatible_provider_info( + model, api_base, api_key, custom_llm_provider, litellm_params + ) elif custom_llm_provider == "novita": api_base = api_base or get_secret("NOVITA_API_BASE") or "https://api.novita.ai/v3/openai" dynamic_api_key = api_key or get_secret_str("NOVITA_API_KEY") diff --git a/litellm/litellm_core_utils/get_provider_specific_headers.py b/litellm/litellm_core_utils/get_provider_specific_headers.py index bc325c3c702..62dfcb073e5 100644 --- a/litellm/litellm_core_utils/get_provider_specific_headers.py +++ b/litellm/litellm_core_utils/get_provider_specific_headers.py @@ -1,10 +1,23 @@ from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter + from litellm.types.utils import ProviderSpecificHeader +_PROVIDER_HEADERS: Final[TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None]] = TypeAdapter( + ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None +) + class ProviderSpecificHeaderUtils: + @staticmethod + def forwarded_authorization(provider_specific_header: object, custom_llm_provider: str | None) -> str | None: + headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( + _PROVIDER_HEADERS.validate_python(provider_specific_header), custom_llm_provider + ) + return next((value for key, value in headers.items() if key.lower() == "authorization"), None) + @staticmethod def get_provider_specific_headers( provider_specific_header: ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None, diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index ff9c6476cba..a5fb1c9c6b9 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -1,9 +1,9 @@ -import base64 import json import os import time from collections.abc import Mapping -from typing import Final, TypeAlias +from functools import lru_cache +from typing import TYPE_CHECKING, Final import httpx from pydantic import JsonValue, TypeAdapter, ValidationError @@ -24,19 +24,21 @@ from .common_utils import ( CHATGPT_OAUTH_TOKEN_URL, GetAccessTokenError, GetDeviceCodeError, + JsonObject, RefreshAccessTokenError, + decode_jwt_claims, + extract_chatgpt_account_id, ) +if TYPE_CHECKING: + from litellm.types.router import GenericLiteLLMParams + TOKEN_EXPIRY_SKEW_SECONDS: Final = 60 TOKEN_REFRESH_TIMEOUT_SECONDS: Final = 30 DEVICE_CODE_TIMEOUT_SECONDS: Final = 15 * 60 DEVICE_CODE_COOLDOWN_SECONDS: Final = 5 * 60 DEVICE_CODE_POLL_SLEEP_SECONDS: Final = 5 -OPENAI_AUTH_CLAIM_KEY: Final = "https://api.openai.com/auth" - -JsonObject: TypeAlias = Mapping[str, JsonValue] - _JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject) @@ -44,6 +46,26 @@ def _optional_str(value: JsonValue | None) -> str | None: return value if isinstance(value, str) else None +def get_chatgpt_auth_file( + litellm_params: "Mapping[str, object] | GenericLiteLLMParams | None", +) -> str | None: + if litellm_params is None: + return None + value: Final = ( + litellm_params.get("chatgpt_auth_file") + if isinstance(litellm_params, Mapping) + else litellm_params.chatgpt_auth_file + ) + if not isinstance(value, str) or not value: + return None + return os.path.abspath(os.path.expanduser(value)) + + +@lru_cache(maxsize=128) +def get_cached_authenticator(auth_file: str) -> "Authenticator": + return Authenticator(auth_file=auth_file) + + def _token_refresh_timeout() -> httpx.Timeout: configured: Final = get_configured_request_timeout() seconds: Final = ( @@ -53,13 +75,13 @@ def _token_refresh_timeout() -> httpx.Timeout: class Authenticator: - def __init__(self) -> None: - self.token_dir = os.getenv( - "CHATGPT_TOKEN_DIR", - os.path.expanduser("~/.config/litellm/chatgpt"), + def __init__(self, auth_file: str | None = None) -> None: + default_auth_file: Final = os.path.join( + os.getenv("CHATGPT_TOKEN_DIR", os.path.expanduser("~/.config/litellm/chatgpt")), + os.getenv("CHATGPT_AUTH_FILE", "auth.json"), ) - self.auth_file = os.path.join(self.token_dir, os.getenv("CHATGPT_AUTH_FILE", "auth.json")) - self._ensure_token_dir() + self.auth_file = os.path.abspath(os.path.expanduser(auth_file or default_auth_file)) + self.token_dir = os.path.dirname(self.auth_file) def get_api_base(self) -> str: return os.getenv("CHATGPT_API_BASE") or os.getenv("OPENAI_CHATGPT_API_BASE") or CHATGPT_API_BASE @@ -99,16 +121,18 @@ class Authenticator: tokens: Final = self._login_device_code() return tokens["access_token"] - def get_account_id(self) -> str | None: + def get_account_id(self, access_token: str | None = None) -> str | None: auth_data: Final = self._read_auth_file() + if access_token and (not auth_data or auth_data.get("access_token") != access_token): + return extract_chatgpt_account_id(access_token) if not auth_data: return None account_id: Final = _optional_str(auth_data.get("account_id")) if account_id: return account_id id_token: Final = auth_data.get("id_token") - access_token: Final = auth_data.get("access_token") - derived: Final = self._extract_account_id(_optional_str(id_token or access_token)) + saved_access_token: Final = auth_data.get("access_token") + derived: Final = extract_chatgpt_account_id(_optional_str(id_token or saved_access_token)) if derived: self._write_auth_file({**auth_data, "account_id": derived}) return derived @@ -129,6 +153,7 @@ class Authenticator: def _write_auth_file(self, data: JsonObject) -> None: try: + self._ensure_token_dir() with open(self.auth_file, "w") as f: json.dump(data, f) except OSError as exc: @@ -146,34 +171,11 @@ class Authenticator: return time.time() >= float(derived_expires_at) - TOKEN_EXPIRY_SKEW_SECONDS def _get_expires_at(self, token: str) -> int | None: - claims: Final = self._decode_jwt_claims(token) - exp: Final = claims.get("exp") + exp: Final = decode_jwt_claims(token).get("exp") if isinstance(exp, (int, float)): return int(exp) return None - def _decode_jwt_claims(self, token: str) -> JsonObject: - try: - parts: Final = token.split(".") - if len(parts) < 2: - return {} - payload_b64: Final = parts[1] + "=" * (-len(parts[1]) % 4) - payload_bytes: Final = base64.urlsafe_b64decode(payload_b64) - return _JSON_OBJECT_ADAPTER.validate_python(json.loads(payload_bytes.decode("utf-8"))) - except Exception: - return {} - - def _extract_account_id(self, token: str | None) -> str | None: - if not token: - return None - claims: Final = self._decode_jwt_claims(token) - auth_claims: Final = claims.get(OPENAI_AUTH_CLAIM_KEY) - if isinstance(auth_claims, dict): - account_id: Final = auth_claims.get("chatgpt_account_id") - if isinstance(account_id, str) and account_id: - return account_id - return None - def _login_device_code(self) -> dict[str, str]: cooldown_remaining: Final = self._get_device_code_cooldown_remaining(self._read_auth_file()) if cooldown_remaining > 0: @@ -361,15 +363,16 @@ class Authenticator: "refresh_token": _optional_str(data.get("refresh_token")) or refresh_token, "id_token": id_token, } + account_id: Final = self.get_account_id() auth_data: Final = self._build_auth_record(refreshed) - self._write_auth_file(auth_data) + self._write_auth_file({**auth_data, "account_id": account_id} if account_id else auth_data) return refreshed def _build_auth_record(self, tokens: dict[str, str]) -> JsonObject: access_token: Final = tokens.get("access_token") id_token: Final = tokens.get("id_token") expires_at: Final = self._get_expires_at(access_token) if access_token else None - account_id: Final = self._extract_account_id(id_token or access_token) + account_id: Final = extract_chatgpt_account_id(id_token or access_token) return { "access_token": access_token, "refresh_token": tokens.get("refresh_token"), diff --git a/litellm/llms/chatgpt/chat/transformation.py b/litellm/llms/chatgpt/chat/transformation.py index a1aac7f315b..9005185931f 100644 --- a/litellm/llms/chatgpt/chat/transformation.py +++ b/litellm/llms/chatgpt/chat/transformation.py @@ -1,14 +1,20 @@ +from collections.abc import Mapping from typing import Any, Final from litellm.exceptions import AuthenticationError from litellm.llms.openai.openai import OpenAIConfig from litellm.types.llms.openai import AllMessageValues +from litellm.types.router import GenericLiteLLMParams -from ..authenticator import Authenticator +from ..authenticator import Authenticator, get_cached_authenticator, get_chatgpt_auth_file from ..common_utils import ( GetAccessTokenError, + chatgpt_client_api_base, ensure_chatgpt_session_id, + get_chatgpt_client_credential, get_chatgpt_default_headers, + get_chatgpt_request_credential, + strip_chatgpt_client_credential_headers, ) from .streaming_utils import ChatGPTToolCallNormalizer @@ -19,9 +25,27 @@ class ChatGPTConfig(OpenAIConfig): api_key: str | None = None, api_base: str | None = None, custom_llm_provider: str = "openai", + authenticator: Authenticator | None = None, ) -> None: super().__init__() - self.authenticator = Authenticator() + self.authenticator = authenticator if authenticator is not None else Authenticator() + + def _resolve_authenticator( + self, litellm_params: Mapping[str, object] | GenericLiteLLMParams | None + ) -> Authenticator: + auth_file: Final = get_chatgpt_auth_file(litellm_params) + return get_cached_authenticator(auth_file) if auth_file else self.authenticator + + @staticmethod + def _get_access_token_or_raise(authenticator: Authenticator, model: str, llm_provider: str) -> str: + try: + return authenticator.get_access_token() + except GetAccessTokenError as e: + raise AuthenticationError( + model=model, + llm_provider=llm_provider, + message=str(e), + ) def api_base_without_login(self) -> str: return self.authenticator.get_api_base() @@ -32,16 +56,16 @@ class ChatGPTConfig(OpenAIConfig): api_base: str | None, api_key: str | None, custom_llm_provider: str, + litellm_params: Mapping[str, object] | GenericLiteLLMParams | None = None, ) -> tuple[str | None, str | None, str]: - dynamic_api_base: Final = self.api_base_without_login() - try: - dynamic_api_key: Final = self.authenticator.get_access_token() - except GetAccessTokenError as e: - raise AuthenticationError( - model=model, - llm_provider=custom_llm_provider, - message=str(e), - ) + client_credential: Final = get_chatgpt_request_credential(litellm_params) + if client_credential is not None: + return chatgpt_client_api_base(api_base), client_credential.access_token, custom_llm_provider + if api_key: + return api_base or self.api_base_without_login(), api_key, custom_llm_provider + authenticator: Final = self._resolve_authenticator(litellm_params) + dynamic_api_base: Final = authenticator.get_api_base() + dynamic_api_key: Final = self._get_access_token_or_raise(authenticator, model, custom_llm_provider) return dynamic_api_base, dynamic_api_key, custom_llm_provider def get_openai_compatible_provider_info( @@ -50,26 +74,49 @@ class ChatGPTConfig(OpenAIConfig): api_base: str | None, api_key: str | None, custom_llm_provider: str, + litellm_params: Mapping[str, object] | GenericLiteLLMParams | None = None, ) -> tuple[str | None, str | None, str]: - return self._get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider) + return self._get_openai_compatible_provider_info(model, api_base, api_key, custom_llm_provider, litellm_params) def validate_environment( self, - headers: dict, + headers: dict[str, str], model: str, messages: list[AllMessageValues], optional_params: dict, - litellm_params: dict, + litellm_params: dict[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: + client_credential: Final = get_chatgpt_client_credential(headers) + if client_credential is not None: + return dict( + get_chatgpt_default_headers( + client_credential.access_token, + client_credential.account_id, + ensure_chatgpt_session_id(litellm_params), + ), + **strip_chatgpt_client_credential_headers(headers), + ) + authenticator: Final = self._resolve_authenticator(litellm_params) + resolved_api_key: Final = ( + self._get_access_token_or_raise(authenticator, model, "chatgpt") + if get_chatgpt_auth_file(litellm_params) + else api_key + ) validated_headers: Final = super().validate_environment( - headers, model, messages, optional_params, litellm_params, api_key, api_base + dict(strip_chatgpt_client_credential_headers(headers)), + model, + messages, + optional_params, + litellm_params, + resolved_api_key, + api_base, ) - account_id: Final = self.authenticator.get_account_id() + account_id: Final = authenticator.get_account_id(resolved_api_key) session_id: Final = ensure_chatgpt_session_id(litellm_params) - default_headers: Final = get_chatgpt_default_headers(api_key or "", account_id, session_id) + default_headers: Final = get_chatgpt_default_headers(resolved_api_key or "", account_id, session_id) return {**default_headers, **validated_headers} def post_stream_processing(self, stream: Any) -> ChatGPTToolCallNormalizer: diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py index fe33219f110..ec7fbfa92da 100644 --- a/litellm/llms/chatgpt/common_utils.py +++ b/litellm/llms/chatgpt/common_utils.py @@ -2,14 +2,32 @@ Constants and helpers for ChatGPT subscription OAuth. """ +import base64 +import json import os import platform -from typing import Any, Final +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias +from urllib.parse import urlsplit from uuid import uuid4 import httpx +from pydantic import JsonValue, TypeAdapter +from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.utils import ProviderSpecificHeader + +if TYPE_CHECKING: + from litellm.types.router import GenericLiteLLMParams + +_SCOPED_HEADERS: Final[TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None]] = TypeAdapter( + ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None +) +_STRING_HEADERS: Final = TypeAdapter(Mapping[str, str]) +_PARAMS: Final = TypeAdapter(Mapping[str, object]) # OAuth + API constants (derived from openai/codex) CHATGPT_AUTH_BASE: Final = "https://auth.openai.com" @@ -19,6 +37,12 @@ CHATGPT_OAUTH_TOKEN_URL: Final = f"{CHATGPT_AUTH_BASE}/oauth/token" CHATGPT_DEVICE_VERIFY_URL: Final = f"{CHATGPT_AUTH_BASE}/codex/device" CHATGPT_API_BASE: Final = "https://chatgpt.com/backend-api/codex" CHATGPT_CLIENT_ID: Final = "app_EMoamEEZ73f0CkXaXp7hrann" +CHATGPT_ACCOUNT_ID_HEADER: Final = "ChatGPT-Account-Id" +OPENAI_AUTH_CLAIM_KEY: Final = "https://api.openai.com/auth" + +JsonObject: TypeAlias = Mapping[str, JsonValue] +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject) +_EMPTY_CLAIMS: Final[JsonObject] = MappingProxyType({}) DEFAULT_ORIGINATOR: Final = "codex_cli_rs" DEFAULT_USER_AGENT: Final = "codex_cli_rs/0.0.0 (Unknown 0; unknown) unknown" @@ -229,7 +253,7 @@ def get_chatgpt_default_headers( access_token: str, account_id: str | None, session_id: str | None = None, -) -> dict: +) -> dict[str, str]: originator: Final = get_chatgpt_originator() user_agent: Final = get_chatgpt_user_agent(originator) headers: Final = { @@ -242,30 +266,133 @@ def get_chatgpt_default_headers( if session_id: headers["session_id"] = session_id if account_id: - headers["ChatGPT-Account-Id"] = account_id + headers[CHATGPT_ACCOUNT_ID_HEADER] = account_id return headers +def decode_jwt_claims(token: str) -> JsonObject: + try: + parts: Final = token.split(".") + if len(parts) < 2: + return _EMPTY_CLAIMS + payload_b64: Final = parts[1] + "=" * (-len(parts[1]) % 4) + payload_bytes: Final = base64.urlsafe_b64decode(payload_b64) + return _JSON_OBJECT_ADAPTER.validate_python(json.loads(payload_bytes.decode("utf-8"))) + except Exception: + return _EMPTY_CLAIMS + + +def extract_chatgpt_account_id(token: str | None) -> str | None: + if not token: + return None + auth_claims: Final = decode_jwt_claims(token).get(OPENAI_AUTH_CLAIM_KEY) + if not isinstance(auth_claims, dict): + return None + account_id: Final = auth_claims.get("chatgpt_account_id") + return account_id if isinstance(account_id, str) and account_id else None + + +def _bearer_token(value: str) -> str: + scheme, separator, token = value.strip().partition(" ") + if not separator: + return scheme + return token.strip() if scheme.lower() == "bearer" else "" + + +def is_chatgpt_oauth_key(value: str | None) -> bool: + if value is None: + return False + return OPENAI_AUTH_CLAIM_KEY in decode_jwt_claims(_bearer_token(value)) + + +@dataclass(frozen=True, slots=True) +class ChatGPTClientCredential: + access_token: str + account_id: str | None + + +def _header_value(headers: Mapping[str, str], name: str) -> str | None: + return next((value for key, value in headers.items() if key.lower() == name), None) + + +def get_chatgpt_client_credential(headers: Mapping[str, str]) -> ChatGPTClientCredential | None: + auth_header: Final = _header_value(headers, "authorization") + if auth_header is None or not is_chatgpt_oauth_key(auth_header): + return None + access_token: Final = _bearer_token(auth_header) + account_id: Final = _header_value(headers, CHATGPT_ACCOUNT_ID_HEADER.lower()) or extract_chatgpt_account_id( + access_token + ) + return ChatGPTClientCredential(access_token=access_token, account_id=account_id) + + +def select_chatgpt_client_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + if get_chatgpt_client_credential(headers) is None: + return MappingProxyType({}) + return MappingProxyType( + { + key: value + for key, value in headers.items() + if (key.lower() == "authorization" and is_chatgpt_oauth_key(value)) + or key.lower() == CHATGPT_ACCOUNT_ID_HEADER.lower() + } + ) + + +def strip_chatgpt_client_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + return MappingProxyType( + { + key: value + for key, value in headers.items() + if key.lower() not in ("authorization", CHATGPT_ACCOUNT_ID_HEADER.lower()) + } + ) + + +def get_chatgpt_request_credential( + params: "Mapping[str, object] | GenericLiteLLMParams | None", +) -> ChatGPTClientCredential | None: + normalized: Final = _normalize_litellm_params(params) + scoped: Final = normalized.get("provider_specific_header") + extra: Final = normalized.get("extra_headers") + headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=_SCOPED_HEADERS.validate_python(scoped), custom_llm_provider="chatgpt" + ) + return get_chatgpt_client_credential(headers) or get_chatgpt_client_credential( + _STRING_HEADERS.validate_python(extra) if extra is not None else MappingProxyType({}) + ) + + +def chatgpt_client_api_base(api_base: str | None) -> str: + base: Final = api_base or CHATGPT_API_BASE + parsed: Final = urlsplit(base) + if ( + parsed.scheme != "https" + or parsed.hostname != "chatgpt.com" + or parsed.port not in (None, 443) + or parsed.username + or parsed.password + ): + raise ChatGPTAuthError( + status_code=400, message="Client ChatGPT credentials require the official HTTPS ChatGPT endpoint" + ) + return base + + def get_chatgpt_default_instructions() -> str: return os.getenv("CHATGPT_DEFAULT_INSTRUCTIONS") or CHATGPT_DEFAULT_INSTRUCTIONS -def _normalize_litellm_params(litellm_params: Any | None) -> dict: - if litellm_params is None: - return {} - if isinstance(litellm_params, dict): - return litellm_params - if hasattr(litellm_params, "model_dump"): - try: - return litellm_params.model_dump() - except Exception: - return {} - if hasattr(litellm_params, "dict"): - try: - return litellm_params.dict() - except Exception: - return {} - return {} +def _normalize_litellm_params(litellm_params: object) -> Mapping[str, object]: + if isinstance(litellm_params, Mapping): + return _PARAMS.validate_python(litellm_params) + serializer: Final = getattr(litellm_params, "model_dump", None) or getattr(litellm_params, "dict", None) + if not callable(serializer): + return MappingProxyType({}) + try: + return _PARAMS.validate_python(serializer()) + except Exception: + return MappingProxyType({}) def get_chatgpt_session_id(litellm_params: object) -> str | None: @@ -274,8 +401,8 @@ def get_chatgpt_session_id(litellm_params: object) -> str | None: value = params.get(key) if value: return str(value) - metadata: Final = params.get("metadata") - if isinstance(metadata, dict): + metadata: Final = _normalize_litellm_params(params.get("metadata")) + if metadata: value = metadata.get("session_id") if value: return str(value) diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 34953a79c60..2352de35b5b 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -25,13 +25,15 @@ from litellm.types.llms.openai import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders -from ..authenticator import Authenticator +from ..authenticator import Authenticator, get_cached_authenticator, get_chatgpt_auth_file from ..common_utils import ( CHATGPT_API_BASE, GetAccessTokenError, ensure_chatgpt_session_id, + get_chatgpt_client_credential, get_chatgpt_default_headers, get_chatgpt_default_instructions, + strip_chatgpt_client_credential_headers, ) if TYPE_CHECKING: @@ -49,14 +51,35 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.CHATGPT + def _resolve_authenticator(self, litellm_params: GenericLiteLLMParams | None) -> Authenticator: + auth_file: Final = get_chatgpt_auth_file(litellm_params) + return get_cached_authenticator(auth_file) if auth_file else self.authenticator + def validate_environment( self, - headers: dict, + headers: dict[str, str], model: str, litellm_params: GenericLiteLLMParams | None, - ) -> dict: + ) -> dict[str, str]: + session_id: Final = ensure_chatgpt_session_id(litellm_params) + client_credential: Final = get_chatgpt_client_credential(headers) + if client_credential is not None: + return { + **get_chatgpt_default_headers(client_credential.access_token, client_credential.account_id, session_id), + **strip_chatgpt_client_credential_headers(headers), + } + authenticator: Final = self._resolve_authenticator(litellm_params) + safe_headers: Final = strip_chatgpt_client_credential_headers(headers) + configured_api_key: Final = litellm_params.api_key if litellm_params is not None else None + if configured_api_key and not get_chatgpt_auth_file(litellm_params): + return { + **get_chatgpt_default_headers( + configured_api_key, authenticator.get_account_id(configured_api_key), session_id + ), + **safe_headers, + } try: - access_token: Final = self.authenticator.get_access_token() + access_token: Final = authenticator.get_access_token() except GetAccessTokenError as e: raise AuthenticationError( model=model, @@ -64,10 +87,9 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): message=str(e), ) - account_id: Final = self.authenticator.get_account_id() - session_id: Final = ensure_chatgpt_session_id(litellm_params) + account_id: Final = authenticator.get_account_id() default_headers: Final = get_chatgpt_default_headers(access_token, account_id, session_id) - return {**default_headers, **headers} + return {**default_headers, **safe_headers} def transform_responses_api_request( self, @@ -79,7 +101,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): ) -> dict: request: Final = super().transform_responses_api_request( model, - input, + [{"role": "user", "content": input}] if isinstance(input, str) else input, response_api_optional_request_params, litellm_params, headers, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 1b491f1614d..7b808cf2be2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -198,6 +198,26 @@ from litellm.utils import ( ) +def _merge_headers_case_insensitive( + headers: Mapping[str, object], extra_headers: Mapping[str, object] | None +) -> dict[str, str]: # mutable-ok: legacy HTTP handler methods require a dict of headers + overridden: Final = frozenset(name.lower() for name in extra_headers or {}) + return { + **{name: str(value) for name, value in headers.items() if name.lower() not in overridden}, + **{name: str(value) for name, value in (extra_headers or {}).items()}, + } + + +def _client_request_headers( + response_api_optional_request_params: Mapping[str, object], + extra_headers: Mapping[str, object] | None, +) -> dict[str, str]: # mutable-ok: provider validate_environment methods require a dict + from litellm.litellm_core_utils.core_helpers import as_str_mapping + + optional_headers: Final = as_str_mapping(response_api_optional_request_params.get("extra_headers")) + return _merge_headers_case_insensitive(optional_headers or {}, extra_headers) + + def _rust_responses_websocket_enabled( custom_llm_provider: str | None, ) -> bool: @@ -2906,14 +2926,18 @@ class BaseLLMHTTPHandler: else: sync_httpx_client = client - headers = responses_api_provider_config.validate_environment( - headers=response_api_optional_request_params.get("extra_headers", {}) or {}, - model=model, - litellm_params=litellm_params, - ) - - if extra_headers: - headers.update(extra_headers) + headers = { # rebind-ok: validation replaces the method's incoming headers + **responses_api_provider_config.validate_environment( + headers=( + _client_request_headers(response_api_optional_request_params, extra_headers) + if custom_llm_provider == "chatgpt" + else response_api_optional_request_params.get("extra_headers", {}) or {} + ), + model=model, + litellm_params=litellm_params, + ), + **((extra_headers or {}) if custom_llm_provider != "chatgpt" else {}), + } # Check if streaming is requested stream = response_api_optional_request_params.get("stream", False) @@ -3104,11 +3128,15 @@ class BaseLLMHTTPHandler: validated_headers: Final = { **await asyncio.to_thread( responses_api_provider_config.validate_environment, - headers=response_api_optional_request_params.get("extra_headers", {}) or {}, + headers=( + _client_request_headers(response_api_optional_request_params, extra_headers) + if custom_llm_provider == "chatgpt" + else response_api_optional_request_params.get("extra_headers", {}) or {} + ), model=model, litellm_params=litellm_params, ), - **(extra_headers or {}), + **((extra_headers or {}) if custom_llm_provider != "chatgpt" else {}), } # Check if streaming is requested diff --git a/litellm/main.py b/litellm/main.py index f30d93361e5..b3d62daa4d8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -677,7 +677,11 @@ async def acompletion( } if custom_llm_provider is None: _supplemental: Final = cast( # cast-ok: kwargs values are request-scoped objects - "dict[str, object]", {k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS if k in kwargs} + "dict[str, object]", + { + **{k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS | {"provider_specific_header"} if k in kwargs}, + "extra_headers": extra_headers, + }, ) _, custom_llm_provider, _, _ = get_llm_provider( model=model, @@ -5678,7 +5682,10 @@ def completion( if deployment_id is not None: # azure llms model = deployment_id custom_llm_provider = "azure" - _supplemental_provider_params: Final = {k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS if k in kwargs} + _supplemental_provider_params: Final = { + **{k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS | {"provider_specific_header"} if k in kwargs}, + "extra_headers": extra_headers, + } model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, @@ -5998,6 +6005,7 @@ def completion( tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), + chatgpt_auth_file=routing_kwargs.get("chatgpt_auth_file"), gigachat_scope=kwargs.get("gigachat_scope"), gigachat_auth_url=kwargs.get("gigachat_auth_url"), gigachat_access_token=kwargs.get("gigachat_access_token"), @@ -9140,6 +9148,7 @@ async def ahealth_check( custom_llm_provider=custom_llm_provider_from_params, api_base=api_base_from_params, api_key=api_key_from_params, + litellm_params=GenericLiteLLMParams(**model_params), ) model_params["cache"] = {"no-cache": True} # don't used cached responses for making health check calls diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index aa67d36af4f..4e54d6ac733 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -207,6 +207,11 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.management_endpoints.config_override_endpoints", path_prefixes=("/config_overrides",), ), + LazyFeature( + name="codex_passthrough", + module_path="litellm.proxy.pass_through_endpoints.codex_passthrough_endpoints", + path_prefixes=("/backend-api/codex/",), + ), LazyFeature( name="llm_passthrough", module_path="litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 4e183bf9d5c..166072e9613 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -8673,6 +8673,184 @@ } } }, + "codex_passthrough": { + "components": { + "schemas": { + "HTTPValidationError": { + "properties": { + "detail": { + "items": { + "$ref": "#/components/schemas/ValidationError" + }, + "title": "Detail", + "type": "array" + } + }, + "title": "HTTPValidationError", + "type": "object" + }, + "ValidationError": { + "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, + "loc": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "title": "Location", + "type": "array" + }, + "msg": { + "title": "Message", + "type": "string" + }, + "type": { + "title": "Error Type", + "type": "string" + } + }, + "required": [ + "loc", + "msg", + "type" + ], + "title": "ValidationError", + "type": "object" + } + } + }, + "paths": { + "/backend-api/codex/{endpoint}": { + "get": { + "description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication.", + "operationId": "codex_passthrough_route_backend_api_codex__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "enum": [ + "models", + "responses" + ], + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Codex Passthrough Route", + "tags": [ + "codex_passthrough" + ] + }, + "post": { + "description": "Forward a client's ChatGPT subscription with separate LiteLLM authentication.", + "operationId": "codex_passthrough_route_backend_api_codex__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "enum": [ + "models", + "responses" + ], + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "text/event-stream": { + "schema": { + "type": "string" + } + } + }, + "description": "Unmodified Codex response or server-sent events" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Codex Passthrough Route", + "tags": [ + "codex_passthrough" + ] + } + } + } + }, "compliance": { "components": { "schemas": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index db433f07f3b..6280e2caf87 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -391,6 +391,7 @@ class LiteLLMRoutes(enum.Enum): ] mapped_pass_through_routes = [ + "/backend-api/codex", "/bedrock", "/comprehendmedical", "/azure_speech", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 25f900a5f9f..30ede935cfa 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -736,7 +736,9 @@ def is_request_body_safe( def _reject_request_credential_params(litellm_params: Mapping[str, object], *, manages_deployments: bool) -> None: if manages_deployments: - reject_server_owned_wif_fields(litellm_params) + reject_server_owned_wif_fields( + MappingProxyType({key: value for key, value in litellm_params.items() if key != "chatgpt_auth_file"}) + ) else: reject_server_owned_wif_params(litellm_params) reject_federated_credential_reference(litellm_params) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 684f1619512..53c54fe15bf 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -2304,7 +2304,15 @@ async def test_model_connection( ) selected_by_id: Final = deployment_by_id is not None - reject_server_owned_wif_fields(request_litellm_params) + reject_server_owned_wif_fields( + MappingProxyType( + { + key: value + for key, value in _OBJECT_MAPPING.validate_python(request_litellm_params).items() + if key != "chatgpt_auth_file" + } + ) + ) # Merge: config params (from proxy config) as base, request params override litellm_params = { **_config_base_for_health_check( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8dc7501bc3b..b51eea7a876 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1417,13 +1417,14 @@ def clean_headers( Cleaned headers dict """ from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + from litellm.llms.chatgpt.common_utils import is_chatgpt_oauth_key clean_headers: Final = {} litellm_key_lower: Final = litellm_key_header_name.lower() if litellm_key_header_name is not None else None for header, value in headers.items(): header_lower = header.lower() - if header_lower == "authorization" and is_anthropic_oauth_key(value): + if header_lower == "authorization" and (is_anthropic_oauth_key(value) or is_chatgpt_oauth_key(value)): if authenticated_with_header is None or authenticated_with_header.lower() != "authorization": clean_headers[header] = value continue @@ -1452,7 +1453,7 @@ def clean_headers( def _is_credential_header(header: str) -> bool: """Whether `header` carries a caller credential rather than request context.""" - return header.lower() in _CREDENTIAL_HEADER_NAMES + return header.lower().removeprefix("x-pass-") in _CREDENTIAL_HEADER_NAMES def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: @@ -4282,6 +4283,7 @@ _NATIVE_MESSAGES_ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join( (LlmProviders.AZURE_AI.value, LlmProviders.DATABRICKS.value) ) _ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = ",".join(sorted(ANTHROPIC_OAUTH_FORWARD_PROVIDERS)) +_CHATGPT_OAUTH_CREDENTIAL_PROVIDERS: Final = LlmProviders.CHATGPT.value def add_provider_specific_headers_to_request( @@ -4289,6 +4291,7 @@ def add_provider_specific_headers_to_request( headers: Mapping[str, str], ) -> bool: from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key + from litellm.llms.chatgpt.common_utils import select_chatgpt_client_credential_headers anthropic_api_headers: Final = {header: headers[header] for header in ANTHROPIC_API_HEADERS if header in headers} anthropic_oauth_credential_headers: Final = { @@ -4296,6 +4299,8 @@ def add_provider_specific_headers_to_request( for header, value in headers.items() if header.lower() == "authorization" and is_anthropic_oauth_key(value) } + chatgpt_oauth_credential_headers: Final = dict(select_chatgpt_client_credential_headers(headers)) + typesafe_danger_zone_headers: Final = { header: value for header, value in headers.items() @@ -4320,13 +4325,17 @@ def add_provider_specific_headers_to_request( ProviderSpecificHeader( custom_llm_provider=LlmProviders.TYPESAFE.value, extra_headers=typesafe_danger_zone_headers ), + ProviderSpecificHeader( + custom_llm_provider=_CHATGPT_OAUTH_CREDENTIAL_PROVIDERS, + extra_headers=chatgpt_oauth_credential_headers, + ), ) if scoped_header["extra_headers"] ] if scoped_headers: data["provider_specific_header"] = scoped_headers[0] if len(scoped_headers) == 1 else scoped_headers - return bool(anthropic_oauth_credential_headers) + return bool(anthropic_oauth_credential_headers or chatgpt_oauth_credential_headers) def _add_otel_traceparent_to_data(data: dict, request: Request): diff --git a/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py new file mode 100644 index 00000000000..3087faa9f67 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/codex_passthrough_endpoints.py @@ -0,0 +1,87 @@ +import hmac +from types import MappingProxyType +from typing import Annotated, Final, Literal + +from fastapi import APIRouter, Depends, HTTPException, Request, Response + +from litellm.llms.chatgpt.common_utils import CHATGPT_API_BASE +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import pass_through_request + +router: Final = APIRouter() +_HEADERS: Final = frozenset( + { + "authorization", + "chatgpt-account-id", + "content-type", + "accept", + "user-agent", + "originator", + "openai-beta", + "session-id", + "thread-id", + "session_id", + "conversation_id", + "x-codex-turn-metadata", + "x-codex-turn-state", + } +) + + +@router.get("/backend-api/codex/{endpoint}", tags=["ChatGPT Pass-through", "pass-through"]) +@router.post( + "/backend-api/codex/{endpoint}", + tags=["ChatGPT Pass-through", "pass-through"], + openapi_extra={ + "requestBody": { + "required": True, + "content": {"application/json": {"schema": {"type": "object", "additionalProperties": True}}}, + } + }, + responses={ + 200: { + "description": "Unmodified Codex response or server-sent events", + "content": { + "application/json": {"schema": {"type": "object", "additionalProperties": True}}, + "text/event-stream": {"schema": {"type": "string"}}, + }, + } + }, +) +async def codex_passthrough_route( + endpoint: Literal["models", "responses"], + request: Request, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> Response: + """Forward a client's ChatGPT subscription with separate LiteLLM authentication.""" + expected_method: Final = "GET" if endpoint == "models" else "POST" + if request.method != expected_method: + raise HTTPException(status_code=405, detail="Method not allowed", headers={"Allow": expected_method}) + + authorization: Final = request.headers.get("authorization", "") + scheme, _, token = authorization.partition(" ") + gateway_key: Final = request.headers.get("x-litellm-api-key", "") + gateway_parts: Final = gateway_key.split() + gateway_token: Final = ( + gateway_parts[-1] + if len(gateway_parts) == 1 or (len(gateway_parts) == 2 and gateway_parts[0].lower() == "bearer") + else "" + ) + if scheme.lower() != "bearer" or not token.strip() or not gateway_token: + raise HTTPException( + status_code=401, + detail="Send the ChatGPT bearer in Authorization and the LiteLLM key in x-litellm-api-key", + ) + if hmac.compare_digest(token.strip().encode(), gateway_token.encode()): + raise HTTPException(status_code=401, detail="ChatGPT and LiteLLM credentials must be separate") + + headers: Final = MappingProxyType({name: value for name, value in request.headers.items() if name in _HEADERS}) + return await pass_through_request( + request=request, + target=f"{CHATGPT_API_BASE}/{endpoint}", + custom_headers=dict(headers), + user_api_key_dict=user_api_key_dict, + custom_llm_provider="chatgpt", + upstream_headers_filter=lambda _: headers, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index c763b5dc26e..c410150df70 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -10,6 +10,7 @@ from typing import Final, Literal, TypeAlias from urllib.parse import urlparse import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -95,6 +96,8 @@ def is_openai_compatible_url(url_route: str | None) -> bool: hostname: Final = parsed_url.hostname if not hostname: return False + if hostname == "chatgpt.com" and parsed_url.path.startswith("/backend-api/codex/"): + return True if _hostname_matches(hostname, _OPENAI_HOSTNAMES): return True if _hostname_matches(hostname, _AZURE_OPENAI_HOSTNAMES): @@ -198,7 +201,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url: Final = urlparse(url_route) - return _is_openai_compatible_host(parsed_url.hostname) and ( + return is_openai_compatible_url(url_route) and ( "/v1/responses" in parsed_url.path or "/responses" in parsed_url.path ) @@ -424,7 +427,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) = None handler_instance: Final = OpenAIPassthroughLoggingHandler() - custom_llm_provider: Final = kwargs.get("custom_llm_provider", "openai") + custom_llm_provider: Final = TypeAdapter(str).validate_python( + kwargs.get("custom_llm_provider", logging_obj.model_call_details.get("custom_llm_provider", "openai")) + ) if is_chat_completions: # Handle chat completions with existing logic diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e1e8864678e..fc0e5be3f4c 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -118,6 +118,7 @@ from litellm.proxy.litellm_pre_call_utils import ( # noqa: F401 # legacy modul _key_or_team_allows_client_pricing_override, # pyright: ignore[reportPrivateUsage] # reuse the proxy's pricing trust policy _strip_client_pricing_overrides, # pyright: ignore[reportPrivateUsage] # sanitize before trusted hooks add guardrail costs get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above + redact_credential_headers, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.route_priority import ( @@ -409,6 +410,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): @staticmethod def get_endpoint_type(url: str, custom_llm_provider: str | None = None) -> EndpointType: parsed_url: Final = urlparse(url) + if custom_llm_provider == "chatgpt": + return EndpointType.OPENAI if custom_llm_provider in ("typesafe", "laya", "bespoke") and parsed_url.path.removesuffix("/").endswith( "/v1/systemone" ): @@ -723,7 +726,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "url": str(request.url), "method": request.method, "body": copy.copy(_parsed_body), # use copy instead of deepcopy - "headers": request.headers, + "headers": redact_credential_headers(request.headers), }, }, "call_type": "pass_through_endpoint", @@ -1490,7 +1493,7 @@ async def pass_through_request( additional_args={ "complete_input_dict": _parsed_body, "api_base": str(logging_url), - "headers": upstream_headers, + "headers": redact_credential_headers(upstream_headers), }, ) stream = HttpPassThroughEndpointHelpers.update_stream_param_based_on_request_body( @@ -1561,7 +1564,7 @@ async def pass_through_request( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, response=relay_response, - request_headers=dict(request.headers), + request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: _response_headers.update(callback_headers) @@ -1653,7 +1656,7 @@ async def pass_through_request( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, response=detected_relay_response, - request_headers=dict(request.headers), + request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: _response_headers.update(callback_headers) @@ -1703,7 +1706,7 @@ async def pass_through_request( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, response=response, - request_headers=dict(request.headers), + request_headers=dict(redact_credential_headers(request.headers)), ) if relay_callback_headers: relay_custom_headers.update(relay_callback_headers) @@ -1878,7 +1881,7 @@ async def pass_through_request( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, response=response, - request_headers=dict(request.headers), + request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: custom_headers.update(callback_headers) diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 2aab8deb034..604a7a389e6 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -19,6 +19,7 @@ from pydantic import ConfigDict, Field import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_prompt_caching_savings, generic_cost_per_token, @@ -69,7 +70,9 @@ def _resolve_model(model: str | None, custom_llm_provider: str | None) -> _Model if not model: return None try: - resolved_model, provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + resolved_model, provider = resolve_model_and_provider_for_metadata( + model=model, custom_llm_provider=custom_llm_provider + ) except Exception as e: # noqa: BLE001 # get_llm_provider raises for unroutable names; degrade to an unavailable estimate verbose_proxy_logger.debug( "savings: cannot resolve provider for model=%s custom_llm_provider=%s (%s)", model, custom_llm_provider, e diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 1ff80b0d71d..d993b05f5b7 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -705,10 +705,12 @@ def _api_base_kwarg(kwargs: Mapping[str, object]) -> str | None: return api_base if isinstance(api_base, str) else None -def _dispatched_model_name(model: str, custom_llm_provider: str, api_base: str | None) -> str: - provider_model, _, _, _ = litellm.get_llm_provider( - model=model, custom_llm_provider=custom_llm_provider, api_base=api_base - ) +def _dispatched_model_name( + model: str, custom_llm_provider: str, api_base: str | None, litellm_params: GenericLiteLLMParams | None = None +) -> str: + from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata + + provider_model, _ = resolve_model_and_provider_for_metadata(model=model, custom_llm_provider=custom_llm_provider) return _strip_responses_routing_prefix(provider_model) @@ -718,6 +720,7 @@ def _will_bridge_to_chat_completions( use_chat_completions_api: bool, model_info: object, api_base: str | None, + litellm_params: GenericLiteLLMParams | None = None, ) -> bool: """``_bridges_to_chat_completions`` for callers running before the provider config is resolved. @@ -734,7 +737,7 @@ def _will_bridge_to_chat_completions( return True return _bridges_to_chat_completions( _resolve_responses_api_provider_config( - _dispatched_model_name(normalized_model[0], custom_llm_provider, api_base), + _dispatched_model_name(normalized_model[0], custom_llm_provider, api_base, litellm_params), custom_llm_provider, model_info, api_base, @@ -750,12 +753,15 @@ def provider_stores_the_conversation(model: str, custom_llm_provider: str | None first MCP round runs as an internal call that writes no spend log at all. """ api_base: Final = _api_base_kwarg(kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) return not _will_bridge_to_chat_completions( model, - custom_llm_provider or litellm.get_llm_provider(model=model, api_base=api_base)[1], + custom_llm_provider + or litellm.get_llm_provider(model=model, api_base=api_base, litellm_params=litellm_params)[1], bool(kwargs.get("use_chat_completions_api")), kwargs.get("model_info"), api_base, + litellm_params, ) @@ -905,7 +911,8 @@ async def aresponses( model=model, api_base=local_vars.get("base_url", None), litellm_params=GenericLiteLLMParams( - **cast("dict[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict + extra_headers=extra_headers, + **cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict ), ) # Update local_vars with detected provider (fixes #19782) @@ -936,6 +943,7 @@ async def aresponses( bool(kwargs.get("use_chat_completions_api")), kwargs.get("model_info"), _api_base_kwarg(kwargs), + GenericLiteLLMParams(**kwargs), ), ): _stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools) @@ -967,7 +975,7 @@ async def aresponses( original_model=original_model, swapped_model=model, custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs={**kwargs, "extra_headers": extra_headers}, prompt_id=prompt_id, ) kwargs.pop("prompt_id", None) @@ -1051,11 +1059,14 @@ def _resolve_prompt_swapped_provider( kwargs: Mapping[str, object], prompt_id: str | None, ) -> str: - swapped_provider: Final = litellm.get_llm_provider(model=swapped_model)[1] + litellm_params: Final = GenericLiteLLMParams(**kwargs) + swapped_provider: Final = litellm.get_llm_provider(model=swapped_model, litellm_params=litellm_params)[1] if kwargs.get("api_key") is None and kwargs.get("api_base") is None: return swapped_provider try: - original_provider: Final = custom_llm_provider or litellm.get_llm_provider(model=original_model)[1] + original_provider: Final = ( + custom_llm_provider or litellm.get_llm_provider(model=original_model, litellm_params=litellm_params)[1] + ) except litellm.BadRequestError: return swapped_provider if swapped_provider == original_provider: @@ -1108,6 +1119,7 @@ def _apply_prompt_management_to_responses_call( use_chat_completions_api, kwargs.get("model_info"), _api_base_kwarg(kwargs), + GenericLiteLLMParams(**kwargs), ), ): _stamp_injection_points_for_the_hook(kwargs, model, custom_llm_provider, input, tools) @@ -1142,7 +1154,7 @@ def _apply_prompt_management_to_responses_call( original_model=original_model, swapped_model=model, custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs={**kwargs, "extra_headers": local_vars.get("extra_headers")}, prompt_id=prompt_id, ) ) @@ -1498,7 +1510,8 @@ def responses( model=model, api_base=local_vars.get("base_url", None), litellm_params=GenericLiteLLMParams( - **cast("dict[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict + extra_headers=extra_headers, + **cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict ), ) local_vars["custom_llm_provider"] = custom_llm_provider @@ -1514,8 +1527,16 @@ def responses( use_chat_completions_api=use_chat_completions_api, ) - # get llm provider logic - litellm_params: Final = GenericLiteLLMParams(**kwargs) + from litellm.litellm_core_utils.core_helpers import attribute_or_none + + litellm_params: Final = GenericLiteLLMParams(extra_headers=extra_headers, **kwargs) + provider_headers: Final = ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers=extra_headers, + provider_specific_header=attribute_or_none(litellm_params, "provider_specific_header"), + custom_llm_provider=custom_llm_provider, + ) + local_vars["extra_headers"] = provider_headers + affinity_api_base: Final = ( get_openai_credentials(api_base=litellm_params.api_base).api_base if custom_llm_provider == "openai" @@ -1525,7 +1546,7 @@ def responses( try: effective_extra_headers: Final = ( add_provider_affinity_header( - headers=extra_headers or MappingProxyType({}), + headers=provider_headers or MappingProxyType({}), litellm_params=MappingProxyType( { "provider_affinity_header": litellm_params.provider_affinity_header, @@ -1538,7 +1559,7 @@ def responses( default_header=default_affinity_header, ) if litellm_params.provider_affinity_header is not None or default_affinity_header is not None - else extra_headers + else provider_headers ) except ValueError as affinity_error: raise litellm.BadRequestError( diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index dc75e77d685..2d0873a218a 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -12,6 +12,7 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, iter_items from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path +from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel from litellm.types.llms.openai import ( AllMessageValues, @@ -29,11 +30,16 @@ from litellm.types.utils import ( CacheCreationTokenDetails, CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, + ProviderSpecificHeader, SpecialEnums, Usage, text_tokens_without_nested_reasoning, ) +_PROVIDER_SPECIFIC_HEADER_ADAPTER: Final[ + TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None] +] = TypeAdapter(ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None) + class _UsageExtraFields(ExtraItemsTypedDict, total=False, extra_items=object): pass @@ -233,6 +239,20 @@ class ResponsesAPIRequestUtils: **extra_headers, } + @staticmethod + def merge_provider_specific_headers( + extra_headers: Mapping[str, object] | None, + provider_specific_header: object, + custom_llm_provider: str | None, + ) -> dict[str, object] | None: # mutable-ok: legacy Responses request fields require a dict + scoped_headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=_PROVIDER_SPECIFIC_HEADER_ADAPTER.validate_python(provider_specific_header), + custom_llm_provider=custom_llm_provider, + ) + if extra_headers is None and not scoped_headers: + return None + return {**(extra_headers or {}), **scoped_headers} + @staticmethod def _check_valid_arg( supported_params: list[str] | None, diff --git a/litellm/router.py b/litellm/router.py index 3af39736e25..c76e4ec843f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -98,7 +98,9 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, r from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import ( declared_authenticating_provider, + inferred_provider, is_registered_custom_provider, + resolve_model_and_provider_for_metadata, ) from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR as _HIDDEN_PARAMS_ATTR, @@ -561,6 +563,7 @@ _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( {"model", "api_base", "api_key", "api_version", "responses_api_version"} ) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +_FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS: Final = frozenset({"chatgpt"}) _CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") _CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" @@ -4801,9 +4804,39 @@ class Router: deployment_litellm_model_name = deployment["litellm_params"]["model"] deployment_api_base = deployment["litellm_params"].get("api_base") deployment_model_name: Final = deployment["model_name"] - if is_clientside_credential(request_kwargs=kwargs): + from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils + + request_context: Final = cast( # cast-ok: the router's legacy request dict has string keys + Mapping[str, object], kwargs + ) + raw_deployment_params: Final[object] = deployment["litellm_params"] + deployment_context: Final = _LITELLM_PARAMS_ADAPTER.validate_python( + raw_deployment_params.model_dump(exclude_none=True) + if isinstance(raw_deployment_params, LiteLLM_Params) + else raw_deployment_params + ) + declared_provider: Final = deployment_context.get("custom_llm_provider") + declared_model: Final = deployment_context.get("model") + deployment_provider: Final = ( + declared_provider + if isinstance(declared_provider, str) + else inferred_provider(declared_model if isinstance(declared_model, str) else None) + ) + forwarded_authorization: Final = ( + ProviderSpecificHeaderUtils.forwarded_authorization( + request_context.get("provider_specific_header"), deployment_provider + ) + if deployment_provider in _FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS + else None + ) + if is_clientside_credential(request_kwargs=kwargs) or forwarded_authorization is not None: deployment_pydantic_obj: Final = self._handle_clientside_credential( - deployment=deployment, kwargs=kwargs, function_name=function_name + deployment=deployment, + kwargs={ + **kwargs, + **({"api_key": forwarded_authorization} if forwarded_authorization is not None else {}), + }, + function_name=function_name, ) model_info = deployment_pydantic_obj.model_info.model_dump() deployment_litellm_model_name = deployment_pydantic_obj.litellm_params.model @@ -11048,6 +11081,10 @@ class Router: custom_llm_provider = None dynamic_api_key = None api_base = None + elif declared_authenticating_provider(litellm_model, deployment.litellm_params.custom_llm_provider) is not None: + _model, custom_llm_provider = ( # rebind-ok: each branch resolves the deployment provider + self._model_group_llm_provider(deployment.litellm_params) + ) else: # check if model provider in supported providers ( @@ -12401,6 +12438,17 @@ class Router: ) return model_info + @staticmethod + def _model_group_llm_provider(litellm_params: LiteLLM_Params) -> tuple[str, str]: + try: + return resolve_model_and_provider_for_metadata( + model=litellm_params.model, + custom_llm_provider=litellm_params.custom_llm_provider, + ) + except litellm.exceptions.BadRequestError as e: + verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e) + return "", "" + def _set_model_group_info( self, model_group: str, user_facing_model_group_name: str, team_id: str | None = None ) -> ModelGroupInfo | None: @@ -12489,15 +12537,7 @@ class Router: deployment_is_mapped = deployment_is_catalog_mapped(model_info, model_info_dict) # get llm provider - litellm_model, llm_provider = "", "" - try: - litellm_model, llm_provider, _, _ = litellm.get_llm_provider( - model=litellm_params.model, - custom_llm_provider=litellm_params.custom_llm_provider, - litellm_params=litellm_params, - ) - except litellm.exceptions.BadRequestError as e: - verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e) + litellm_model, llm_provider = self._model_group_llm_provider(litellm_params) if model_info is None: supported_openai_params = litellm.get_supported_openai_params( diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 089f21a33c1..717b74e5741 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -91,6 +91,7 @@ class ProviderConnection: litellm_credential_name: str | None = None configurable_clientside_auth_params: "Sequence[str | ConfigurableClientsideParamsCustomAuth] | None" = None use_xai_oauth: bool | None = None + chatgpt_auth_file: str | None = None github_copilot_auth_type: str | None = None github_copilot_host: str | None = None github_copilot_user_session: object | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index cd1389aea9d..ba9b82e2a57 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -572,6 +572,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): default=False, description="Use stored xAI OAuth credentials when no xAI API key is configured.", ) + chatgpt_auth_file: str | None = Field( + default=None, + description="Path to a ChatGPT OAuth auth.json for this deployment.", + ) fireworks_forward_user_id: bool | None = Field( default=None, description="Send the LiteLLM user id of the calling key as the `user` field on Fireworks AI chat, responses and messages requests.", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cc7ef851111..4f39741a079 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4360,6 +4360,7 @@ server_owned_wif_litellm_params: Final = ( + openai_wif_litellm_params + oauth_token_exchange_litellm_params + github_copilot_oauth_litellm_params + + ("chatgpt_auth_file",) ) secret_bearing_wif_litellm_params: Final = tuple(sorted(WIF_SECRET_BEARING_KEYS)) diff --git a/litellm/utils.py b/litellm/utils.py index 4315ae426cf..1cc7c63e6a5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2964,12 +2964,14 @@ def supports_native_streaming(model: str, custom_llm_provider: str | None) -> bo Raises: Exception: If the given model is not found in model_prices_and_context_window.json. """ + from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata + try: - model, custom_llm_provider, _, _ = litellm.get_llm_provider( + resolved_model, resolved_provider = resolve_model_and_provider_for_metadata( model=model, custom_llm_provider=custom_llm_provider ) - model_info: Final = get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + model_info: Final = get_model_info_helper(model=resolved_model, custom_llm_provider=resolved_provider) supports_native_streaming = model_info.get("supports_native_streaming", True) if supports_native_streaming is None: supports_native_streaming = True diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index 57e66805678..e9e3f7b8174 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -1,5 +1,9 @@ +from collections.abc import Iterator +from pathlib import Path from typing import Final +import respx + import asyncio, httpx, importlib, os import pytest @@ -8,6 +12,7 @@ from litellm import CustomLLM from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.get_llm_provider_logic import ( get_llm_provider, + resolve_model_and_provider_for_metadata, get_provider_affinity_header, inferred_provider, is_registered_custom_provider, @@ -715,3 +720,47 @@ def test_opencode_affinity_resolution_preserves_existing_transport(prefix: str) model, provider, _, destination = get_llm_provider(model=f"{prefix}/test-model", api_base=api_base) assert (model, provider, destination) == ("test-model", prefix, api_base) assert get_provider_affinity_header(destination) == "x-opencode-session" + + +@pytest.fixture +def chatgpt_login_tripwire(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[None]: + blocked: Final = tmp_path / "blocked" + blocked.write_text("ordinary file") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(blocked / "unused")) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + with respx.mock as transport: + yield + assert len(transport.calls) == 0 + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected"), + [ + ("chatgpt/gpt-5.4", None, ("gpt-5.4", "chatgpt")), + ("gpt-5.4", "chatgpt", ("gpt-5.4", "chatgpt")), + ("github_copilot/gpt-4o", None, ("gpt-4o", "github_copilot")), + ("openai/gpt-5.4", None, ("gpt-5.4", "openai")), + ], +) +def test_resolve_model_and_provider_for_metadata_never_logs_in( + chatgpt_login_tripwire: None, model: str, custom_llm_provider: str | None, expected: tuple[str, str] +) -> None: + assert resolve_model_and_provider_for_metadata(model, custom_llm_provider) == expected + + +def test_metadata_callers_never_log_in_for_a_declared_chatgpt_deployment( + chatgpt_login_tripwire: None, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base + from litellm.proxy.spend_tracking.savings import _resolve_model + from litellm.router import Router + from litellm.types.router import LiteLLM_Params + from litellm.utils import supports_native_streaming + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + + assert supports_native_streaming("chatgpt/gpt-5.4", None) is True + assert get_api_base("chatgpt/gpt-5.4", {"model": "chatgpt/gpt-5.4"}) == "https://chatgpt.com/backend-api/codex" + identity = _resolve_model("chatgpt/gpt-5.4", None) + assert identity is not None and identity.provider == "chatgpt" + assert Router._model_group_llm_provider(LiteLLM_Params(model="chatgpt/gpt-5.4")) == ("gpt-5.4", "chatgpt") diff --git a/tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py b/tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py new file mode 100644 index 00000000000..f9d1dc28d12 --- /dev/null +++ b/tests/unit/llms/chatgpt/chat/test_chatgpt_chat_transformation.py @@ -0,0 +1,638 @@ +from unittest.mock import MagicMock + +import pytest + +from litellm.exceptions import AuthenticationError +from litellm.llms.chatgpt.chat.transformation import ChatGPTConfig +from litellm.llms.chatgpt.common_utils import GetAccessTokenError + +GATEWAY_BASE = "https://chatgpt.example/backend-api/codex" + + +def _gateway_authenticator() -> MagicMock: + authenticator = MagicMock() + authenticator.get_api_base.return_value = GATEWAY_BASE + authenticator.get_access_token.return_value = "gateway-access-token" + return authenticator + + +def _provider_info( + config: ChatGPTConfig, api_base: str | None, api_key: str | None +) -> tuple[str | None, str | None, str]: + return config._get_openai_compatible_provider_info( + model="gpt-5.4", api_base=api_base, api_key=api_key, custom_llm_provider="chatgpt" + ) + + +class TestChatGPTChatProviderInfo: + def test_configured_key_uses_configured_base_without_gateway_login(self): + authenticator = _gateway_authenticator() + config = ChatGPTConfig(authenticator=authenticator) + + assert _provider_info(config, "http://127.0.0.1:9999", "configured-key") == ( + "http://127.0.0.1:9999", + "configured-key", + "chatgpt", + ) + authenticator.get_access_token.assert_not_called() + + def test_configured_key_without_base_uses_gateway_base(self): + config = ChatGPTConfig(authenticator=_gateway_authenticator()) + + assert _provider_info(config, None, "configured-key") == (GATEWAY_BASE, "configured-key", "chatgpt") + + def test_gateway_login_is_never_sent_to_a_custom_base(self): + config = ChatGPTConfig(authenticator=_gateway_authenticator()) + + assert _provider_info(config, "http://attacker.example", None) == ( + GATEWAY_BASE, + "gateway-access-token", + "chatgpt", + ) + + def test_missing_gateway_login_raises_authentication_error(self): + authenticator = _gateway_authenticator() + authenticator.get_access_token.side_effect = GetAccessTokenError(status_code=401, message="not logged in") + config = ChatGPTConfig(authenticator=authenticator) + + with pytest.raises(AuthenticationError, match="not logged in"): + _provider_info(config, None, None) + + +import asyncio +import json +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import respx + +import litellm +from litellm.integrations.dotprompt.dotprompt_manager import DotpromptManager +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.chatgpt.chat.transformation import ChatGPTConfig +from litellm.types.router import GenericLiteLLMParams + + +def _write_auth_record(path: Path, token: str, account_id: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "access_token": token, + "account_id": account_id, + "expires_at": 4102444800, + } + ) + ) + + +@pytest.fixture +def auth_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, Path]: + monkeypatch.setattr("time.time", lambda: 1700000000) + default_dir: Final = tmp_path / "default" + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + _write_auth_record(default_dir / "auth.json", "token-default", "acct-default") + file_a: Final = tmp_path / "account-a" / "auth.json" + file_b: Final = tmp_path / "account-b" / "auth.json" + _write_auth_record(file_a, "token-a", "acct-a") + _write_auth_record(file_b, "token-b", "acct-b") + return file_a, file_b + + +class TestChatGPTMultiAccountChat: + def test_validate_environment_uses_per_model_auth_file(self, auth_files: tuple[Path, Path]) -> None: + file_a, file_b = auth_files + config: Final = ChatGPTConfig() + + headers_a: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + messages=[], + optional_params={}, + litellm_params={"chatgpt_auth_file": str(file_a)}, + api_key="token-default", + ) + headers_b: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + messages=[], + optional_params={}, + litellm_params={"chatgpt_auth_file": str(file_b)}, + api_key="token-default", + ) + + assert headers_a["Authorization"] == "Bearer token-a" + assert headers_a["ChatGPT-Account-Id"] == "acct-a" + assert headers_b["Authorization"] == "Bearer token-b" + assert headers_b["ChatGPT-Account-Id"] == "acct-b" + + def test_validate_environment_defaults_to_global_auth_file(self, auth_files: tuple[Path, Path]) -> None: + config: Final = ChatGPTConfig() + + headers: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + messages=[], + optional_params={}, + litellm_params={}, + api_key="token-default", + ) + + assert headers["Authorization"] == "Bearer token-default" + assert headers["ChatGPT-Account-Id"] == "acct-default" + + def test_provider_info_uses_per_model_auth_file(self, auth_files: tuple[Path, Path]) -> None: + file_a, _ = auth_files + config: Final = ChatGPTConfig() + + _, per_model_key, _ = config._get_openai_compatible_provider_info( + "gpt-6.1-sol", None, None, "chatgpt", {"chatgpt_auth_file": str(file_a)} + ) + _, default_key, _ = config._get_openai_compatible_provider_info("gpt-6.1-sol", None, None, "chatgpt") + + assert per_model_key == "token-a" + assert default_key == "token-default" + + def test_get_llm_provider_threads_auth_file(self, auth_files: tuple[Path, Path]) -> None: + file_a, _ = auth_files + + model, provider, dynamic_api_key, _ = get_llm_provider( + model="chatgpt/gpt-6.1-sol", + litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_a)), + ) + + assert model == "gpt-6.1-sol" + assert provider == "chatgpt" + assert dynamic_api_key == "token-a" + + +def _subscription_response(request: httpx.Request) -> httpx.Response: + account: Final = request.headers["ChatGPT-Account-Id"] + assert request.headers["Authorization"] in ( + f"Bearer token-{account.removeprefix('acct-')}", + f"Bearer {_client_token(account)}", + ) + assert "chatgpt_auth_file" not in request.content.decode() + message: Final = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": account, "annotations": []}], + } + response: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [message], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + events: Final = ( + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + { + "type": "response.output_text.delta", + "item_id": "msg_test", + "output_index": 0, + "content_index": 0, + "delta": account, + }, + {"type": "response.output_item.done", "output_index": 0, "item": message}, + {"type": "response.completed", "response": response}, + ) + return httpx.Response( + 200, + text="".join(f"data: {json.dumps(event)}\n\n" for event in events), + headers={"content-type": "text/event-stream"}, + ) + + +@pytest.mark.parametrize("blocked_default", [False, True]) +@pytest.mark.parametrize("with_prompt", [False, True]) +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("through_router", [False, True]) +@pytest.mark.asyncio +async def test_requests_use_selected_subscription( + auth_files: tuple[Path, Path], + endpoint: str, + through_router: bool, + with_prompt: bool, + blocked_default: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + default_dir: Final = auth_files[0].parent / "missing-default" + if blocked_default: + default_dir.write_text("ordinary file") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir / "unused")) + if with_prompt: + monkeypatch.setattr( + litellm, + "callbacks", + [ + DotpromptManager( + prompt_id="account-check", + prompt_data={"content": "Saved account prompt", "metadata": {"model": "chatgpt/gpt-6.1-sol"}}, + ) + ], + ) + router: Final = ( + litellm.Router( + model_list=[ + { + "model_name": f"account-{index}", + "litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(path)}, + } + for index, path in enumerate(auth_files) + ] + ) + if through_router + else None + ) + + async def request(index: int, path: Path) -> str: + params: Final = MappingProxyType( + { + "model": f"account-{index}" if router else "chatgpt/gpt-6.1-sol", + **({"chatgpt_auth_file": str(path)} if router is None else {}), + **({"prompt_id": "account-check"} if with_prompt else {}), + } + ) + if endpoint == "chat": + result: Final = await (router.acompletion if router else litellm.acompletion)( + **params, messages=[{"role": "user", "content": "Hello"}] + ) + return result.choices[0].message.content + response: Final = await (router.aresponses if router else litellm.aresponses)( + **params, input=[{"role": "user", "content": "Hello"}] + ) + return response.output[0].content[0].text + + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=_subscription_response) + results: Final = await asyncio.gather(*(request(index, path) for index, path in enumerate(auth_files))) + assert results == ["acct-a", "acct-b"] + if with_prompt: + assert all("Saved account prompt" in call.request.content.decode() for call in transport.calls) + + +def _client_token(account: str) -> str: + import base64 + + payload: Final = ( + base64.urlsafe_b64encode(json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account}}).encode()) + .rstrip(b"=") + .decode() + ) + return f"header.{payload}.signature" + + +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("configured_login", [False, True]) +@pytest.mark.parametrize("with_prompt", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_client_login_keeps_its_account_across_gateway_deployments( + auth_files: tuple[Path, Path], + endpoint: str, + configured_login: bool, + with_prompt: bool, + stream: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(auth_files[0].parent / "no-gateway-login")) + if with_prompt: + monkeypatch.setattr( + litellm, + "callbacks", + [ + DotpromptManager( + prompt_id="client-account-check", + prompt_data={"content": "Saved client prompt", "metadata": {"model": "chatgpt/gpt-6.1-sol"}}, + ) + ], + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "subscriptions", + "litellm_params": { + "model": "chatgpt/gpt-6.1-sol", + **({"chatgpt_auth_file": str(path)} if configured_login else {}), + }, + } + for path in auth_files + ], + num_retries=0, + ) + + def reply(request: httpx.Request) -> httpx.Response: + account: Final = request.headers["ChatGPT-Account-Id"] + assert request.headers["Authorization"] == f"Bearer {_client_token(account)}" + return _subscription_response( + httpx.Request( + "POST", + request.url, + content=request.content, + headers={ + **{key: value for key, value in request.headers.items() if key.lower() != "authorization"}, + "Authorization": f"Bearer token-{account.removeprefix('acct-')}", + }, + ) + ) + + async def call(account: str) -> str: + params: Final = MappingProxyType( + { + "model": "subscriptions", + "stream": stream, + **({"prompt_id": "client-account-check"} if with_prompt else {}), + "provider_specific_header": { + "custom_llm_provider": "chatgpt", + "extra_headers": {"Authorization": f"Bearer {_client_token(account)}"}, + }, + } + ) + response: Final = await ( + router.acompletion(**params, messages=[{"role": "user", "content": "Hello"}]) + if endpoint == "chat" + else router.aresponses(**params, input=[{"role": "user", "content": "Hello"}]) + ) + if stream: + chunks: Final = [part async for part in response] + return ( + "".join(part.choices[0].delta.content or "" for part in chunks) + if endpoint == "chat" + else "".join( + getattr(part, "delta", "") + for part in chunks + if getattr(part, "type", "") == "response.output_text.delta" + ) + ) + return response.choices[0].message.content if endpoint == "chat" else response.output[0].content[0].text + + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=reply) + results: Final = await asyncio.gather(call("acct-client-a"), call("acct-client-b")) + assert results == ["acct-client-a", "acct-client-b"] + assert len(transport.calls) == 2 + + +@pytest.mark.asyncio +async def test_one_client_limit_does_not_cool_down_other_subscriptions( + auth_files: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + {"model_name": "pool", "litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(path)}} + for path in auth_files + ], + num_retries=0, + allowed_fails=0, + cooldown_time=60, + ) + headers: Final = { + "custom_llm_provider": "chatgpt", + "extra_headers": {"Authorization": f"Bearer {_client_token('acct-limited')}"}, + } + with respx.mock as transport: + route: Final = transport.post("https://chatgpt.com/backend-api/codex/responses") + route.respond( + 429, + headers={"retry-after": "60"}, + json={"error": {"message": "Usage limit reached", "type": "rate_limit_error"}}, + ) + for _ in range(2): + with pytest.raises(litellm.RateLimitError): + await router.aresponses(model="pool", input="Hello", provider_specific_header=headers) + await asyncio.sleep(0) + route.mock(side_effect=_subscription_response) + response: Final = await router.aresponses(model="pool", input="Hello") + assert response.output[0].content[0].text in ("acct-a", "acct-b") + + +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("async_call", [False, True]) +@pytest.mark.parametrize("direct_headers", [False, True]) +@pytest.mark.parametrize("configured_key", [False, True]) +@pytest.mark.parametrize( + "target", ["https://example.com/v1", "http://chatgpt.com/backend-api/codex", "https://chatgpt.com.example.com/v1"] +) +@pytest.mark.asyncio +async def test_client_login_cannot_be_redirected_to_another_host( + endpoint: str, target: str, async_call: bool, direct_headers: bool, configured_key: bool +) -> None: + headers: Final = {"Authorization": f"Bearer {_client_token('acct-client')}"} + credential: Final = ( + {"extra_headers": headers} + if direct_headers + else {"provider_specific_header": {"custom_llm_provider": "chatgpt", "extra_headers": headers}} + ) + params: Final = MappingProxyType( + { + "model": "chatgpt/gpt-6.1-sol", + "api_base": target, + **({"api_key": "configured-test-key"} if configured_key else {}), + **credential, + } + ) + + async def request() -> None: + if endpoint == "chat": + if async_call: + await litellm.acompletion(**params, messages=[{"role": "user", "content": "Hello"}]) + else: + litellm.completion(**params, messages=[{"role": "user", "content": "Hello"}]) + elif async_call: + await litellm.aresponses(**params, input="Hello") + else: + litellm.responses(**params, input="Hello") + + with respx.mock as transport: + with pytest.raises(litellm.BadRequestError, match="official HTTPS ChatGPT endpoint"): + await request() + assert len(transport.calls) == 0 + + +@pytest.mark.asyncio +async def test_cross_provider_fallback_does_not_receive_the_chatgpt_login( + auth_files: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "subscription", + "litellm_params": {"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(auth_files[0])}, + }, + { + "model_name": "api-fallback", + "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": "sk-example-fallback"}, + }, + ], + fallbacks=[{"subscription": ["api-fallback"]}], + num_retries=0, + ) + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").respond( + 429, json={"error": {"message": "Usage limit reached", "type": "rate_limit_error"}} + ) + fallback: Final = transport.post("https://api.openai.com/v1/responses").respond( + 200, + json={ + "id": "resp_fallback", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [ + { + "id": "msg_fallback", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "fallback", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + ) + response: Final = await router.aresponses( + model="subscription", + input="Hello", + provider_specific_header={ + "custom_llm_provider": "chatgpt", + "extra_headers": { + "Authorization": f"Bearer {_client_token('acct-client')}", + "ChatGPT-Account-Id": "acct-client", + }, + }, + ) + assert response.output[0].content[0].text == "fallback" + assert fallback.calls.last.request.headers["Authorization"] == "Bearer sk-example-fallback" + assert "ChatGPT-Account-Id" not in fallback.calls.last.request.headers + + +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("async_call", [False, True]) +@pytest.mark.parametrize("with_prompt", [False, True]) +@pytest.mark.asyncio +async def test_sdk_extra_headers_use_client_login_without_gateway_auth( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, endpoint: str, async_call: bool, with_prompt: bool +) -> None: + blocked: Final = tmp_path / "blocked" + blocked.write_text("ordinary file") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(blocked / "unused")) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + headers: Final = {"Authorization": f"Bearer {_client_token('acct-client-a')}"} + prompt_params: Final = {"prompt_id": "sdk-client"} if with_prompt else {} + if with_prompt: + monkeypatch.setattr( + litellm, + "callbacks", + [ + DotpromptManager( + prompt_id="sdk-client", + prompt_data={"content": "Hello", "metadata": {"model": "chatgpt/gpt-6-luna"}}, + ) + ], + ) + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=_subscription_response) + if endpoint == "chat": + params: Final = { + "model": "chatgpt/gpt-6.1-sol", + "messages": [{"role": "user", "content": "Hello"}], + "extra_headers": headers, + **prompt_params, + } + result: Final = await litellm.acompletion(**params) if async_call else litellm.completion(**params) + assert result.choices[0].message.content == "acct-client-a" + else: + response: Final = ( + await litellm.aresponses( + model="chatgpt/gpt-6.1-sol", input="Hello", extra_headers=headers, **prompt_params + ) + if async_call + else litellm.responses( + model="chatgpt/gpt-6.1-sol", input="Hello", extra_headers=headers, **prompt_params + ) + ) + assert response.output[0].content[0].text == "acct-client-a" + assert transport.calls.last.request.headers["Authorization"] == headers["Authorization"] + assert len(transport.calls) == 1 + + +@pytest.mark.parametrize("endpoint", ["chat", "responses"]) +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize("credential_source", ["default", "file", "api_key"]) +@pytest.mark.parametrize("authorization_header", [None, "Authorization", "authorization"]) +@pytest.mark.asyncio +async def test_requests_preserve_gateway_token_and_selected_workspace( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + endpoint: str, + sync: bool, + credential_source: str, + authorization_header: str | None, +) -> None: + monkeypatch.setattr("time.time", lambda: 1700000000) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("CHATGPT_AUTH_FILE", str(tmp_path / "default.json")) + token: Final = _client_token("acct-token-default") + auth_file: Final = tmp_path / ("selected.json" if credential_source == "file" else "default.json") + _write_auth_record(auth_file, token, "acct-saved-workspace") + configured_token: Final = _client_token("acct-configured") + expected_token: Final = configured_token if credential_source == "api_key" else token + expected_account: Final = "acct-configured" if credential_source == "api_key" else "acct-saved-workspace" + + def reply(request: httpx.Request) -> httpx.Response: + assert request.headers.get_list("authorization") == [f"Bearer {expected_token}"] + assert request.headers.get_list("chatgpt-account-id") == [expected_account] + assert request.headers["originator"] == "codex_cli_rs" + return _subscription_response( + httpx.Request( + "POST", + request.url, + content=request.content, + headers={ + "Authorization": f"Bearer token-{expected_account.removeprefix('acct-')}", + "ChatGPT-Account-Id": expected_account, + }, + ) + ) + + params: Final = MappingProxyType( + { + "model": "chatgpt/gpt-6.1-sol", + **({"chatgpt_auth_file": str(auth_file)} if credential_source == "file" else {}), + **({"api_key": configured_token} if credential_source == "api_key" else {}), + "extra_headers": { + "originator": "codex_cli_rs", + **({authorization_header: "Bearer sk-unrelated-example"} if authorization_header else {}), + }, + } + ) + with respx.mock as transport: + transport.post("https://chatgpt.com/backend-api/codex/responses").mock(side_effect=reply) + if endpoint == "chat": + result: Final = ( + await asyncio.to_thread(litellm.completion, **params, messages=[{"role": "user", "content": "Hello"}]) + if sync + else await litellm.acompletion(**params, messages=[{"role": "user", "content": "Hello"}]) + ) + assert result.choices[0].message.content == expected_account + else: + response: Final = ( + await asyncio.to_thread(litellm.responses, **params, input=[{"role": "user", "content": "Hello"}]) + if sync + else await litellm.aresponses(**params, input=[{"role": "user", "content": "Hello"}]) + ) + assert response.output[0].content[0].text == expected_account diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index e47d7fc702c..d88bdfc0b41 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -4,8 +4,10 @@ Tests for ChatGPT subscription Responses API transformation Source: litellm/llms/chatgpt/responses/transformation.py """ +import base64 import json from collections.abc import Generator +from pathlib import Path from typing import Final from unittest.mock import MagicMock, patch @@ -13,6 +15,7 @@ import httpx import pytest import litellm +from litellm.llms.chatgpt.authenticator import Authenticator from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig from litellm.llms.openai.common_utils import OpenAIError from litellm.main import responses_api_bridge_check @@ -395,3 +398,165 @@ class TestChatGPTResponsesAPITransformation: assert exc_info.value.message == "ChatGPT upstream failed" assert exc_info.value.status_code == 502 + + +def _fake_chatgpt_oauth_token(account_id: str) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode( + json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}}).encode() + ) + .rstrip(b"=") + .decode() + ) + return f"{header}.{payload}.fake-signature" + + +def _gateway_authenticator() -> MagicMock: + authenticator = MagicMock() + authenticator.get_access_token.return_value = "gateway-access-token" + authenticator.get_account_id.return_value = "acct-gateway" + return authenticator + + +def _config_with(authenticator: MagicMock) -> ChatGPTResponsesAPIConfig: + return ChatGPTResponsesAPIConfig(authenticator=authenticator) + + +class TestChatGPTClientCredentialForwarding: + def test_client_chatgpt_oauth_wins_over_configured_key_and_gateway_login(self): + authenticator = _gateway_authenticator() + token = _fake_chatgpt_oauth_token("acct-alice") + + headers = _config_with(authenticator).validate_environment( + headers={"Authorization": f"Bearer {token}", "originator": "codex_cli_rs"}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key", litellm_session_id="session-1"), + ) + + assert headers["Authorization"] == f"Bearer {token}" + assert headers["ChatGPT-Account-Id"] == "acct-alice" + assert headers["originator"] == "codex_cli_rs" + assert headers["accept"] == "text/event-stream" + assert [name for name in headers if name.lower() == "authorization"] == ["Authorization"] + authenticator.get_access_token.assert_not_called() + + def test_explicit_account_header_wins_over_jwt_claim(self): + token = _fake_chatgpt_oauth_token("acct-from-jwt") + + headers = _config_with(_gateway_authenticator()).validate_environment( + headers={"authorization": f"Bearer {token}", "chatgpt-account-id": "acct-from-header"}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers["ChatGPT-Account-Id"] == "acct-from-header" + assert [name for name in headers if name.lower() == "chatgpt-account-id"] == ["ChatGPT-Account-Id"] + + @pytest.mark.parametrize( + "bearer", ["Bearer sk-proj-fake-platform-key", "Bearer not-a-jwt", "Bearer a.bm90LWpzb24.c"] + ) + def test_non_chatgpt_bearer_falls_back_to_configured_key(self, bearer, tmp_path): + authenticator = Authenticator(auth_file=str(tmp_path / "missing-default.json")) + + headers = ChatGPTResponsesAPIConfig(authenticator=authenticator).validate_environment( + headers={"Authorization": bearer}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key"), + ) + + assert "ChatGPT-Account-Id" not in headers + assert headers["Authorization"] == "Bearer configured-gateway-key" + + def test_configured_key_is_used_without_client_credential(self): + authenticator = _gateway_authenticator() + + headers = _config_with(authenticator).validate_environment( + headers={}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(api_key="configured-gateway-key"), + ) + + assert headers["Authorization"] == "Bearer configured-gateway-key" + authenticator.get_access_token.assert_not_called() + + def test_gateway_login_is_used_without_client_credential_or_configured_key(self): + headers = _config_with(_gateway_authenticator()).validate_environment( + headers={}, + model="gpt-5.4", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers["Authorization"] == "Bearer gateway-access-token" + assert headers["ChatGPT-Account-Id"] == "acct-gateway" + + def test_two_clients_never_share_a_credential(self): + config = _config_with(_gateway_authenticator()) + alice = _fake_chatgpt_oauth_token("acct-alice") + bob = _fake_chatgpt_oauth_token("acct-bob") + + forwarded = [ + config.validate_environment( + headers={"Authorization": f"Bearer {token}"}, model="gpt-5.4", litellm_params=GenericLiteLLMParams() + ) + for token in (alice, bob, alice) + ] + + assert [h["Authorization"] for h in forwarded] == [f"Bearer {alice}", f"Bearer {bob}", f"Bearer {alice}"] + assert [h["ChatGPT-Account-Id"] for h in forwarded] == ["acct-alice", "acct-bob", "acct-alice"] + + +class TestChatGPTResponsesMultiAccount: + @pytest.fixture(autouse=True) + def fixed_clock(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("time.time", lambda: 1700000000) + + @staticmethod + def _write_auth_record(path: Path, token: str, account_id: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "access_token": token, + "account_id": account_id, + "expires_at": 4102444800, + } + ) + ) + + def test_validate_environment_uses_per_model_auth_file( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + default_dir: Final = tmp_path / "default" + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(default_dir)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + self._write_auth_record(default_dir / "auth.json", "token-default", "acct-default") + file_a: Final = tmp_path / "account-a" / "auth.json" + file_b: Final = tmp_path / "account-b" / "auth.json" + self._write_auth_record(file_a, "token-a", "acct-a") + self._write_auth_record(file_b, "token-b", "acct-b") + + config: Final = ChatGPTResponsesAPIConfig() + + headers_a: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_a)), + ) + headers_b: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + litellm_params=GenericLiteLLMParams(chatgpt_auth_file=str(file_b)), + ) + headers_default: Final = config.validate_environment( + headers={}, + model="gpt-6.1-sol", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers_a["Authorization"] == "Bearer token-a" + assert headers_a["ChatGPT-Account-Id"] == "acct-a" + assert headers_b["Authorization"] == "Bearer token-b" + assert headers_b["ChatGPT-Account-Id"] == "acct-b" + assert headers_default["Authorization"] == "Bearer token-default" + assert headers_default["ChatGPT-Account-Id"] == "acct-default" diff --git a/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py index 43a4c26b6d2..2e9e3462bee 100644 --- a/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py +++ b/tests/unit/llms/chatgpt/test_chatgpt_authenticator.py @@ -2,17 +2,24 @@ import base64 import json import time from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final from unittest.mock import MagicMock, mock_open, patch import pytest +import httpx +import respx import litellm from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS, HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS from litellm.llms.chatgpt.authenticator import ( TOKEN_REFRESH_TIMEOUT_SECONDS, Authenticator, + get_cached_authenticator, + get_chatgpt_auth_file, ) from litellm.llms.chatgpt.common_utils import GetAccessTokenError +from litellm.types.router import GenericLiteLLMParams def _make_jwt(payload: dict) -> str: @@ -161,3 +168,144 @@ class TestChatGPTAuthenticator: assert account_id == "acct-123" mock_write.assert_called_once() assert mock_write.call_args[0][0]["account_id"] == "acct-123" + + def test_get_account_id_prefers_stored_value_without_rewriting_file(self, authenticator): + id_token = _make_jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-from-token"}}) + auth_data = json.dumps({"account_id": "acct-stored", "id_token": id_token}) + + with ( + patch("builtins.open", mock_open(read_data=auth_data)), + patch.object(authenticator, "_write_auth_file") as mock_write, + ): + assert authenticator.get_account_id() == "acct-stored" + mock_write.assert_not_called() + + def test_get_account_id_without_login_returns_none(self, authenticator): + with ( + patch("builtins.open", mock_open(read_data="{}")), + patch.object(authenticator, "_write_auth_file") as mock_write, + ): + assert authenticator.get_account_id() is None + mock_write.assert_not_called() + + +def _write_auth_record(path: Path, token: str, account_id: str) -> None: + path.write_text( + json.dumps( + { + "access_token": token, + "account_id": account_id, + "expires_at": 4102444800, + } + ) + ) + + +class TestChatGPTMultiAccountAuthenticator: + @pytest.fixture(autouse=True) + def fixed_clock(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("time.time", lambda: 1700000000) + + def test_explicit_auth_file_overrides_env(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "env-dir")) + monkeypatch.setenv("CHATGPT_AUTH_FILE", "env.json") + custom_file: Final = tmp_path / "account-a" / "auth.json" + + authenticator: Final = Authenticator(auth_file=str(custom_file)) + + assert authenticator.auth_file == str(custom_file) + assert authenticator.token_dir == str(custom_file.parent) + assert not custom_file.parent.exists() + authenticator._write_auth_file({"access_token": "token-a"}) + assert json.loads(custom_file.read_text())["access_token"] == "token-a" + + def test_default_auth_file_from_env(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "env-dir")) + monkeypatch.setenv("CHATGPT_AUTH_FILE", "env.json") + + authenticator: Final = Authenticator() + + assert authenticator.auth_file == str(tmp_path / "env-dir" / "env.json") + + def test_bare_filename_auth_file_resolves_against_cwd( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.chdir(tmp_path) + + authenticator: Final = Authenticator(auth_file="auth.json") + + assert authenticator.auth_file == str(tmp_path / "auth.json") + assert authenticator.token_dir == str(tmp_path) + + def test_two_auth_files_are_isolated(self, tmp_path: Path) -> None: + file_a: Final = tmp_path / "account-a.json" + file_b: Final = tmp_path / "account-b.json" + _write_auth_record(file_a, "token-a", "acct-a") + _write_auth_record(file_b, "token-b", "acct-b") + + authenticator_a: Final = Authenticator(auth_file=str(file_a)) + authenticator_b: Final = Authenticator(auth_file=str(file_b)) + + assert authenticator_a.get_access_token() == "token-a" + assert authenticator_a.get_account_id() == "acct-a" + assert authenticator_b.get_access_token() == "token-b" + assert authenticator_b.get_account_id() == "acct-b" + + def test_get_cached_authenticator_reuses_instance_per_path(self, tmp_path: Path) -> None: + file_a: Final = tmp_path / "account-a.json" + file_b: Final = tmp_path / "account-b.json" + + authenticator_a: Final = get_cached_authenticator(str(file_a)) + + assert get_cached_authenticator(str(file_a)) is authenticator_a + assert get_cached_authenticator(str(file_b)) is not authenticator_a + + def test_get_chatgpt_auth_file(self) -> None: + assert get_chatgpt_auth_file(None) is None + assert get_chatgpt_auth_file({}) is None + assert get_chatgpt_auth_file({"chatgpt_auth_file": ""}) is None + assert get_chatgpt_auth_file({"chatgpt_auth_file": "/a/auth.json"}) == "/a/auth.json" + assert get_chatgpt_auth_file(GenericLiteLLMParams()) is None + assert get_chatgpt_auth_file(GenericLiteLLMParams(chatgpt_auth_file="/b/auth.json")) == "/b/auth.json" + + def test_relative_and_absolute_auth_file_share_cached_authenticator( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.chdir(tmp_path) + get_cached_authenticator.cache_clear() + auth_file: Final = tmp_path / "x" / "auth.json" + auth_file.parent.mkdir(parents=True) + auth_file.write_text("{}") + + relative_key: Final = get_chatgpt_auth_file({"chatgpt_auth_file": "x/auth.json"}) + absolute_key: Final = get_chatgpt_auth_file({"chatgpt_auth_file": str(auth_file)}) + assert relative_key is not None and absolute_key is not None + + assert get_cached_authenticator(relative_key) is get_cached_authenticator(absolute_key) + + +def test_refresh_keeps_saved_workspace(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("time.time", lambda: 1700000000) + auth_file: Final = tmp_path / "selected.json" + auth_file.write_text( + json.dumps( + { + "access_token": "expired-example-token", + "refresh_token": "example-refresh-token", + "expires_at": 1, + "account_id": "acct-selected-workspace", + } + ) + ) + token: Final = _make_jwt({"exp": 4102444800, "https://api.openai.com/auth": {"chatgpt_account_id": "acct-default"}}) + authenticator: Final = Authenticator(auth_file=str(auth_file)) + with respx.mock as transport: + transport.post("https://auth.openai.com/oauth/token").mock( + return_value=httpx.Response( + 200, + json={"access_token": token, "id_token": token, "refresh_token": "example-rotated-refresh"}, + ) + ) + assert authenticator.get_access_token() == token + assert authenticator.get_account_id(token) == "acct-selected-workspace" + assert json.loads(auth_file.read_text())["refresh_token"] == "example-rotated-refresh" diff --git a/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py new file mode 100644 index 00000000000..6b280d97089 --- /dev/null +++ b/tests/unit/llms/chatgpt/test_chatgpt_common_utils.py @@ -0,0 +1,90 @@ +import base64 +import json + +import pytest + +from litellm.llms.chatgpt.common_utils import ( + extract_chatgpt_account_id, + get_chatgpt_client_credential, + is_chatgpt_oauth_key, + select_chatgpt_client_credential_headers, +) + + +def _make_jwt(payload: dict) -> str: + def _b64(obj: dict) -> str: + raw = json.dumps(obj, separators=(",", ":")).encode("utf-8") + return base64.urlsafe_b64encode(raw).decode("utf-8").rstrip("=") + + return f"{_b64({'alg': 'none', 'typ': 'JWT'})}.{_b64(payload)}." + + +CHATGPT_JWT = _make_jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct-alice"}}) +CHATGPT_JWT_WITHOUT_ACCOUNT = _make_jwt({"https://api.openai.com/auth": {"user_id": "user-1"}}) +NON_CHATGPT_JWT = _make_jwt({"sub": "someone", "https://api.openai.com/auth": "not-an-object"}) + + +class TestExtractChatGPTAccountId: + def test_reads_account_id_from_auth_claim(self): + assert extract_chatgpt_account_id(CHATGPT_JWT) == "acct-alice" + + @pytest.mark.parametrize( + "token", + [None, "", "sk-proj-not-a-jwt", "header.only", NON_CHATGPT_JWT, CHATGPT_JWT_WITHOUT_ACCOUNT], + ) + def test_returns_none_without_a_usable_account_claim(self, token): + assert extract_chatgpt_account_id(token) is None + + +class TestIsChatGPTOAuthKey: + def test_recognizes_bearer_and_raw_chatgpt_jwt(self): + assert is_chatgpt_oauth_key(CHATGPT_JWT) + assert is_chatgpt_oauth_key(f"Bearer {CHATGPT_JWT}") + assert is_chatgpt_oauth_key(f"Bearer {CHATGPT_JWT_WITHOUT_ACCOUNT}") + + @pytest.mark.parametrize( + "value", [None, "", "Bearer sk-proj-platform-key", "Bearer not.a.jwt", _make_jwt({"sub": "x"})] + ) + def test_rejects_non_chatgpt_credentials(self, value): + assert not is_chatgpt_oauth_key(value) + + +class TestSelectChatGPTClientCredentialHeaders: + def test_account_header_alone_is_not_forwarded(self): + selected = select_chatgpt_client_credential_headers( + {"chatgpt-account-id": "acct-alice", "authorization": "Bearer sk-proj-platform-key"} + ) + assert dict(selected) == {} + assert get_chatgpt_client_credential({"chatgpt-account-id": "acct-alice"}) is None + + def test_forwards_only_bearer_and_account_context(self): + selected = select_chatgpt_client_credential_headers( + { + "Authorization": f"Bearer {CHATGPT_JWT}", + "ChatGPT-Account-Id": "acct-explicit", + "x-litellm-api-key": "sk-gateway", + "user-agent": "codex", + } + ) + assert dict(selected) == {"Authorization": f"Bearer {CHATGPT_JWT}", "ChatGPT-Account-Id": "acct-explicit"} + credential = get_chatgpt_client_credential(selected) + assert credential is not None + assert credential.access_token == CHATGPT_JWT + assert credential.account_id == "acct-explicit" + + def test_account_id_falls_back_to_jwt_claim(self): + credential = get_chatgpt_client_credential({"authorization": f"Bearer {CHATGPT_JWT}"}) + assert credential is not None + assert credential.account_id == "acct-alice" + + +@pytest.mark.parametrize("scheme", ["Bearer", "bearer", "BEARER"]) +def test_bearer_scheme_is_case_insensitive(scheme): + credential = get_chatgpt_client_credential({"authorization": f"{scheme} {CHATGPT_JWT}"}) + assert credential is not None + assert credential.access_token == CHATGPT_JWT + assert credential.account_id == "acct-alice" + + +def test_other_authentication_schemes_are_not_chatgpt_logins(): + assert not is_chatgpt_oauth_key(f"Basic {CHATGPT_JWT}") diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index bbcd7f70e63..f7850aaf9a1 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5654,3 +5654,44 @@ def test_vector_store_retrieve_handler_sends_id_before_azure_api_version_query() sent: Final = received[0] assert str(sent.url) == "https://x.openai.azure.com/openai/vector_stores/vs_abc?api-version=2025-04-01-preview" assert len(received) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", (False, True)) +@pytest.mark.parametrize("provider", ("openai", "azure")) +async def test_non_chatgpt_responses_preserves_caller_headers_after_validation( + is_async: bool, provider: str, respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=r"https://responses\.example\.com/.*").respond( + 200, + json={ + "id": "resp_headers", "object": "response", "created_at": 1, + "model": "gpt-6-luna", "status": "completed", "output": [], + "parallel_tool_calls": False, "tool_choice": "auto", "tools": [], + }, + ) + client: Final = AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport()) if is_async else HTTPHandler(client=httpx.Client()) + kwargs: Final = { + "model": f"{provider}/gpt-6-luna", + "input": "hello", + "api_key": "fake-configured-key", + "api_base": "https://responses.example.com/v1", + "api_version": "2025-04-01-preview", + "extra_headers": {"Authorization": "Bearer fake-caller-key", "content-type": "application/vnd.example+json"}, + "client": client, + } + try: + response: Final = await litellm.aresponses(**kwargs) if is_async else litellm.responses(**kwargs) + assert response.status == "completed" + assert route.calls[0].request.headers["authorization"] == "Bearer fake-caller-key" + expected_content_types: Final = ( + ["application/json", "application/vnd.example+json"] + if provider == "openai" + else ["application/vnd.example+json"] + ) + assert route.calls[0].request.headers.get_list("content-type") == expected_content_types + finally: + if isinstance(client, AsyncHTTPHandler): + await client.close() + else: + client.client.close() diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index ee1ae43fd3a..10a0f974a42 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3626,6 +3626,18 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: ) +@pytest.mark.parametrize("allow_client_side_credentials", [False, True]) +@pytest.mark.parametrize("nested", [False, True]) +def test_request_cannot_select_chatgpt_auth_file(allow_client_side_credentials: bool, nested: bool) -> None: + selector: Final = {"chatgpt_auth_file": "/tokens/account-b/auth.json"} + with pytest.raises(ValueError, match="chatgpt_auth_file"): + is_request_body_safe( + request_body={"model": "chatgpt-account-a", **({"extra_body": selector} if nested else selector)}, + general_settings={"allow_client_side_credentials": allow_client_side_credentials}, + llm_router=None, + model="chatgpt-account-a", + ) + class TestIsRequestBodySafeBlocksBedrockProjectOverride: """``aws_bedrock_project_id`` pins a deployment to a Bedrock project so that project's data-retention policy applies to its requests. A @@ -5304,3 +5316,27 @@ async def test_a_non_numeric_max_request_size_mb_is_ignored_without_premium(): patch("litellm.proxy.proxy_server.premium_user", False), ): assert await check_if_request_size_is_safe(request=_request_with_content_length(1024)) is True + + +@pytest.mark.parametrize("allow_client_side_credentials", [False, True]) +@pytest.mark.parametrize("route", ["/model/new", "/model/update", "/model/account-a/update", "/health/test_connection"]) +def test_chatgpt_model_configuration_reaches_role_authorization( + allow_client_side_credentials: bool, route: str +) -> None: + body: Final = {"litellm_params": {"model": "chatgpt/gpt-6-luna", "chatgpt_auth_file": "/tokens/account-a/auth.json"}} + assert is_request_body_safe( + request_body=body, + general_settings={"allow_client_side_credentials": allow_client_side_credentials}, + llm_router=None, + model="", + route=route, + ) + for inference_route in ("/v1/chat/completions", "/v1/responses"): + with pytest.raises(ValueError, match="chatgpt_auth_file"): + is_request_body_safe( + request_body=body, + general_settings={"allow_client_side_credentials": allow_client_side_credentials}, + llm_router=None, + model="", + route=inference_route, + ) diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 950d1742d49..67b280eb379 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -5437,3 +5437,64 @@ async def test_health_services_endpoint_langfuse_missing_keys_errors(monkeypatch with pytest.raises(ProxyException, match="auth_check failed") as raised: await health_services_endpoint(service="langfuse") assert "connection refused by lf.internal.example" in str(raised.value.message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("admin", [False, True]) +async def test_unsaved_chatgpt_account_probe_requires_proxy_admin( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, admin: bool +) -> None: + from starlette.requests import Request + from litellm.proxy import proxy_server + + monkeypatch.setenv("CHATGPT_AUTH_FILE", str(tmp_path / "missing-default-auth.json")) + auth_file: Final = tmp_path / "auth.json" + auth_file.write_text( + json.dumps({"access_token": "token-probe", "account_id": "acct-probe", "expires_at": 4102444800}) + ) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + output: Final = { + "id": "msg_probe", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + response: Final = { + "id": "resp_probe", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [output], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + with respx.mock(assert_all_called=False) as upstream: + upstream.post("https://chatgpt.com/backend-api/codex/responses").respond( + 200, + text="data: " + json.dumps({"type": "response.completed", "response": response}) + "\n\n", + headers={"content-type": "text/event-stream"}, + ) + call: Final = health_test_model_connection( + request=Request({"type": "http", "headers": []}), + mode="responses", + litellm_params={"model": "chatgpt/gpt-6.1-sol", "chatgpt_auth_file": str(auth_file)}, + model_info={}, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN if admin else LitellmUserRoles.INTERNAL_USER, token="test-key" + ), + ) + if admin: + result: Final = await call + assert result["status"] == "success" + assert upstream.calls.last.request.headers["Authorization"] == "Bearer token-probe" + assert upstream.calls.last.request.headers["ChatGPT-Account-Id"] == "acct-probe" + assert isinstance(json.loads(upstream.calls.last.request.content)["input"], list) + else: + with pytest.raises(ProxyException) as error: + await call + assert error.value.code == "403" + assert len(upstream.calls) == 0 diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 303b8e060ed..a718426eea6 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -8933,7 +8933,14 @@ class TestNonAdminCannotPersistWifFieldsOnModel: mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() @pytest.mark.asyncio - async def test_patch_model_non_admin_cannot_set_wif_field(self): + @pytest.mark.parametrize( + "field,value", + [ + ("anthropic_keycloak_token_url", "https://identity.example.com/token"), + ("chatgpt_auth_file", "/tokens/account-b/auth.json"), + ], + ) + async def test_patch_model_non_admin_cannot_set_wif_field(self, field: str, value: str) -> None: from litellm.proxy.management_endpoints.model_management_endpoints import ( patch_model, ) @@ -8979,13 +8986,13 @@ class TestNonAdminCannotPersistWifFieldsOnModel: model_id="m1", patch_data=updateDeployment( litellm_params=updateLiteLLMParams( - anthropic_keycloak_token_url="https://attacker.example/token", + **{field: value}, ) ), user_api_key_dict=non_admin, ) err = exc_info.value - assert getattr(err, "param", "") == "anthropic_keycloak_token_url" + assert getattr(err, "param", "") == field mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() @pytest.mark.asyncio @@ -9044,7 +9051,14 @@ class TestNonAdminCannotPersistWifFieldsOnModel: mock_prisma.db.litellm_proxymodeltable.update.assert_not_called() @pytest.mark.asyncio - async def test_patch_model_admin_can_set_wif_field(self): + @pytest.mark.parametrize( + "field,value", + [ + ("anthropic_keycloak_token_url", "https://identity.example.com/token"), + ("chatgpt_auth_file", "/tokens/account-b/auth.json"), + ], + ) + async def test_patch_model_admin_can_set_wif_field(self, field: str, value: str) -> None: from litellm.proxy.management_endpoints.model_management_endpoints import ( patch_model, ) @@ -9095,7 +9109,7 @@ class TestNonAdminCannotPersistWifFieldsOnModel: model_id="m1", patch_data=updateDeployment( litellm_params=updateLiteLLMParams( - anthropic_keycloak_token_url="https://keycloak.internal/token", + **{field: value}, ) ), user_api_key_dict=admin, diff --git a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py new file mode 100644 index 00000000000..354c65776dc --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py @@ -0,0 +1,185 @@ +import asyncio +import json +from collections.abc import Iterator +from datetime import datetime +from typing import Final + +import httpx +import pytest +import respx +from fastapi.testclient import TestClient + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy import proxy_server +from tests._master_key import MASTER_KEY + +_BASE: Final = "https://chatgpt.com/backend-api/codex" +_TOKEN: Final = "test-subscription-token" +_HEADERS: Final = {"authorization": f"Bearer {_TOKEN}", "x-litellm-api-key": MASTER_KEY} + + +@pytest.fixture +def client(monkeypatch: pytest.MonkeyPatch, httpx_transport: None) -> Iterator[TestClient]: + monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY) + yield TestClient(proxy_server.app) + + +def test_native_catalog_preserves_account_and_filters_gateway_credentials( + client: TestClient, respx_mock: respx.MockRouter +) -> None: + catalog: Final = {"models": [{"slug": "subscription-model", "supported_reasoning_levels": ["high"]}]} + upstream: Final = respx_mock.get(f"{_BASE}/models", params={"client_version": "1.0.0"}).respond(json=catalog) + response: Final = client.get( + "/backend-api/codex/models?client_version=1.0.0", + headers={ + **_HEADERS, + "chatgpt-account-id": "account-a", + "cookie": "session=test-cookie", + "x-pass-authorization": "Bearer test-injected-token", + "x-pass-x-litellm-api-key": "test-injected-key", + }, + ) + assert response.status_code == 200 + assert response.json() == catalog + assert upstream.calls.last.request.headers["authorization"] == f"Bearer {_TOKEN}" + assert upstream.calls.last.request.headers["chatgpt-account-id"] == "account-a" + assert "x-litellm-api-key" not in upstream.calls.last.request.headers + assert "cookie" not in upstream.calls.last.request.headers + + +@pytest.mark.parametrize("account", ["account-a", "account-b"]) +def test_native_responses_preserves_body_and_stream( + client: TestClient, respx_mock: respx.MockRouter, account: str +) -> None: + body: Final = { + "model": "subscription-model", + "input": [{"role": "user", "content": "Hello"}], + "reasoning": {"effort": "high"}, + "store": False, + "stream": True, + } + stream: Final = ( + b'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"Hello"}\n\n' + ) + upstream: Final = respx_mock.post(f"{_BASE}/responses").respond( + content=stream, headers={"content-type": "text/event-stream"} + ) + response: Final = client.post( + "/backend-api/codex/responses", + json=body, + headers={**_HEADERS, "authorization": f"Bearer test-{account}", "chatgpt-account-id": account}, + ) + assert response.status_code == 200 + assert response.content == stream + assert response.headers["content-type"].startswith("text/event-stream") + assert json.loads(upstream.calls.last.request.content) == body + assert upstream.calls.last.request.headers["authorization"] == f"Bearer test-{account}" + assert upstream.calls.last.request.headers["chatgpt-account-id"] == account + + +@pytest.mark.parametrize("gateway_key", [None, "sk-invalid-test-key"]) +def test_gateway_auth_rejects_before_upstream( + client: TestClient, respx_mock: respx.MockRouter, gateway_key: str | None +) -> None: + headers: Final = { + "authorization": f"Bearer {_TOKEN}", + **({"x-litellm-api-key": gateway_key} if gateway_key else {}), + } + response: Final = client.get("/backend-api/codex/models", headers=headers) + assert response.status_code in (400, 401, 403) + assert not respx_mock.calls + + +@pytest.mark.parametrize("authorization", ["", "Basic test", "Bearer ", f"Bearer {MASTER_KEY}"]) +def test_subscription_credential_required(client: TestClient, respx_mock: respx.MockRouter, authorization: str) -> None: + response: Final = client.get("/backend-api/codex/models", headers={**_HEADERS, "authorization": authorization}) + assert response.status_code == 401 + assert not respx_mock.calls + + +@pytest.mark.parametrize("status", [401, 403, 429]) +def test_provider_errors_do_not_fall_back(client: TestClient, respx_mock: respx.MockRouter, status: int) -> None: + upstream: Final = respx_mock.post(f"{_BASE}/responses").respond( + status, json={"error": {"message": "subscription unavailable"}}, headers={"retry-after": "60"} + ) + response: Final = client.post( + "/backend-api/codex/responses", headers=_HEADERS, json={"model": "subscription-model", "stream": True} + ) + assert response.status_code == status + assert upstream.call_count == 1 + + +@pytest.mark.parametrize("method,path", [("POST", "models"), ("GET", "responses"), ("GET", "deviceauth")]) +def test_only_native_inference_routes_are_forwarded( + client: TestClient, respx_mock: respx.MockRouter, method: str, path: str +) -> None: + response: Final = client.request(method, f"/backend-api/codex/{path}", headers=_HEADERS) + assert response.status_code in (405, 422) + assert not respx_mock.calls + + +class UsageRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + await self.events.put(kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_usage_logged_without_credentials( + monkeypatch: pytest.MonkeyPatch, httpx_transport: None, respx_mock: respx.MockRouter, stream: bool +) -> None: + monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY) + recorder: Final = UsageRecorder() + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + litellm.register_model( + { + "chatgpt/subscription-model": { + "litellm_provider": "chatgpt", + "mode": "responses", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + } + ) + result: Final = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "subscription-model", + "output": [], + "usage": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15}, + } + event: Final = {"type": "response.completed", "response": result} + response_body: Final = ( + httpx.Response( + 200, + content=f"event: response.completed\ndata: {json.dumps(event)}\n\n", + headers={"content-type": "text/event-stream"}, + ) + if stream + else httpx.Response(200, json=result) + ) + respx_mock.post(f"{_BASE}/responses").mock(return_value=response_body) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test") as caller: + response: Final = await caller.post( + "/backend-api/codex/responses", + json={"model": "subscription-model", "stream": stream, "input": "Hello"}, + headers={**_HEADERS, "x-pass-authorization": "Bearer test-hidden-token", "cookie": "test-hidden-cookie"}, + ) + assert response.status_code == 200 + logged: Final = await asyncio.wait_for(recorder.events.get(), timeout=5) + payload: Final = logged["standard_logging_object"] + assert isinstance(payload, dict) + assert payload["prompt_tokens"] == 12 + assert payload["completion_tokens"] == 3 + serialized: Final = json.dumps(logged, default=str) + for secret in (_TOKEN, MASTER_KEY, "test-hidden-token", "test-hidden-cookie"): + assert secret not in serialized diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 64ff8392c02..ed7817bd660 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1,4 +1,5 @@ import asyncio +import base64 import copy import json import os @@ -9208,6 +9209,97 @@ def test_default_team_settings_bool_turn_off_message_logging_redacts(): ) +def _fake_chatgpt_oauth_token(account_id: str) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"RS256"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode( + json.dumps({"https://api.openai.com/auth": {"chatgpt_account_id": account_id}, "exp": 4102444800}).encode() + ) + .rstrip(b"=") + .decode() + ) + return f"{header}.{payload}.fake-signature" + + +CHATGPT_OAUTH_TOKEN = f"Bearer {_fake_chatgpt_oauth_token('acct-alice')}" +CHATGPT_ACCOUNT_ID_HEADER_NAME = "ChatGPT-Account-Id" +PLATFORM_API_KEY_BEARER = "Bearer sk-proj-fake-platform-key-0123456789" +CHATGPT_LEAK_TARGET_PROVIDERS = ["openai", "azure", "anthropic", "bedrock", "vertex_ai"] + + +def _codex_client_headers( + authorization_header_name: str = "authorization", account_id: str | None = "acct-alice" +) -> dict: + return { + "content-type": "application/json", + "originator": "codex_cli_rs", + authorization_header_name: CHATGPT_OAUTH_TOKEN, + **({CHATGPT_ACCOUNT_ID_HEADER_NAME: account_id} if account_id is not None else {}), + } + + +def test_clean_headers_keeps_chatgpt_oauth_bearer_when_gateway_key_is_separate(): + result = clean_headers( + headers=Headers({**_codex_client_headers(), "x-litellm-api-key": "sk-proxy-virtual-key"}), + litellm_key_header_name="x-litellm-api-key", + forward_llm_provider_auth_headers=False, + authenticated_with_header="x-litellm-api-key", + ) + + assert result["authorization"] == CHATGPT_OAUTH_TOKEN + assert "x-litellm-api-key" not in result + + +def test_clean_headers_drops_bearer_that_authenticated_to_the_gateway(): + result = clean_headers( + headers=Headers({"authorization": CHATGPT_OAUTH_TOKEN, "content-type": "application/json"}), + litellm_key_header_name=None, + forward_llm_provider_auth_headers=False, + authenticated_with_header="authorization", + ) + + assert _authorization_values(result) == [] + + +@pytest.mark.parametrize("bearer", [PLATFORM_API_KEY_BEARER, "Bearer not.a-jwt", "Bearer a.bm90LWpzb24.c"]) +def test_clean_headers_drops_non_chatgpt_bearers(bearer): + result = clean_headers( + headers=Headers({"authorization": bearer, "x-litellm-api-key": "sk-proxy-virtual-key"}), + litellm_key_header_name="x-litellm-api-key", + forward_llm_provider_auth_headers=False, + authenticated_with_header="x-litellm-api-key", + ) + + assert _authorization_values(result) == [] + + +@pytest.mark.parametrize("authorization_header_name", AUTHORIZATION_HEADER_CASINGS) +def test_chatgpt_oauth_credential_and_account_reach_chatgpt_provider(authorization_header_name): + forwarded = _headers_forwarded_to(_codex_client_headers(authorization_header_name), "chatgpt") + + assert _authorization_values(forwarded) == [CHATGPT_OAUTH_TOKEN] + assert forwarded[CHATGPT_ACCOUNT_ID_HEADER_NAME] == "acct-alice" + + +@pytest.mark.parametrize("custom_llm_provider", CHATGPT_LEAK_TARGET_PROVIDERS) +def test_chatgpt_oauth_credential_is_never_forwarded_to_other_providers(custom_llm_provider): + forwarded = _headers_forwarded_to(_codex_client_headers(), custom_llm_provider) + + assert _authorization_values(forwarded) == [] + assert CHATGPT_OAUTH_TOKEN not in forwarded.values() + assert CHATGPT_ACCOUNT_ID_HEADER_NAME not in forwarded + + +def test_chatgpt_account_header_is_not_forwarded_without_chatgpt_oauth_bearer(): + headers = {**_codex_client_headers(), "authorization": PLATFORM_API_KEY_BEARER} + + assert _headers_forwarded_to(headers, "chatgpt") == {} + + +def test_anthropic_oauth_bearer_does_not_reach_chatgpt_provider(): + assert _authorization_values(_headers_forwarded_to(_client_headers(), "chatgpt")) == [] + + def test_add_user_api_key_auth_to_request_metadata_attributes_a_cli_session_to_its_alias(): data = {"model": "gpt-5.4-nano", "messages": [{"role": "user", "content": "hi"}], "litellm_metadata": {}} session = UserAPIKeyAuth( diff --git a/tests/unit/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py index 97a39f64001..0863abd1f1e 100644 --- a/tests/unit/responses/test_responses_utils.py +++ b/tests/unit/responses/test_responses_utils.py @@ -933,6 +933,42 @@ class TestResponsesInputToChatMessages: ) == [user_message] +class TestMergeProviderSpecificHeaders: + def test_scoped_headers_merge_only_for_matching_provider(self): + scoped = {"custom_llm_provider": "chatgpt", "extra_headers": {"Authorization": "Bearer client-token"}} + + merged = ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers={"originator": "codex_cli_rs"}, + provider_specific_header=scoped, + custom_llm_provider="chatgpt", + ) + untouched = ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers={"originator": "codex_cli_rs"}, + provider_specific_header=scoped, + custom_llm_provider="openai", + ) + + assert merged == {"originator": "codex_cli_rs", "Authorization": "Bearer client-token"} + assert untouched == {"originator": "codex_cli_rs"} + + def test_sequence_of_scoped_entries_each_keep_their_own_scope(self): + entries = [ + {"custom_llm_provider": "anthropic", "extra_headers": {"Authorization": "Bearer anthropic-token"}}, + {"custom_llm_provider": "chatgpt", "extra_headers": {"Authorization": "Bearer chatgpt-token"}}, + ] + + merged = ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers=None, provider_specific_header=entries, custom_llm_provider="chatgpt" + ) + + assert merged == {"Authorization": "Bearer chatgpt-token"} + + def test_no_scoped_headers_returns_extra_headers_unchanged(self): + assert ResponsesAPIRequestUtils.merge_provider_specific_headers( + extra_headers=None, provider_specific_header=None, custom_llm_provider="chatgpt" + ) is None + + def test_responses_object_bills_an_explicit_cache_hit_at_the_explicit_rate(): model = "dashscope/qwen-responses-explicit-cache-test" previous_row = litellm.model_cost.get(model) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 4814d36ee39..6aaa3379f47 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -283,6 +283,15 @@ def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): assert kwargs["litellm_metadata"] == {"baz": 123} +def test_router_model_group_llm_provider_resolves_chatgpt_without_login(tmp_path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + params: Final = LiteLLM_Params(model="chatgpt/gpt-5.4") + + assert Router._model_group_llm_provider(params) == ("gpt-5.4", "chatgpt") + assert Router._model_group_llm_provider(LiteLLM_Params(model="no-such-provider-model-xyz")) == ("", "") + assert not any(tmp_path.iterdir()) + + def test_router_with_model_info_and_model_group(): """ Test edge case where user specifies model_group in model_info diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index b0086662f1c..740342f54d9 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -94,6 +94,7 @@ CONNECTION_NAMES: Final = ( "github_copilot_host", "configurable_clientside_auth_params", "use_xai_oauth", + "chatgpt_auth_file", "fireworks_forward_user_id", "aws_batch_role_arn", "aws_bedrock_control_plane_endpoint", diff --git a/tests/unit/types/test_router.py b/tests/unit/types/test_router.py index 2892a829eab..7855eba0042 100644 --- a/tests/unit/types/test_router.py +++ b/tests/unit/types/test_router.py @@ -16,8 +16,8 @@ from litellm.types.router import ( holds_secret_pointer, reject_server_owned_wif_params, server_owned_wif_fields_named, - server_owned_wif_fields_written, server_owned_wif_fields_present, + server_owned_wif_fields_written, ) from litellm.types.utils import ( CustomPricingLiteLLMParams, @@ -367,6 +367,7 @@ def test_server_owned_registry_includes_anthropic_openai_and_oauth_token_exchang + openai_wif_litellm_params + oauth_token_exchange_litellm_params + github_copilot_oauth_litellm_params + + ("chatgpt_auth_file",) ) assert set(openai_wif_litellm_params) == { "openai_identity_provider_id", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 262e8b6deae..d3c1e0422df 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1792,6 +1792,30 @@ export interface paths { patch: operations["azure_speech_proxy_route_azure_speech__endpoint__patch"]; trace?: never; }; + "/backend-api/codex/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Codex Passthrough Route + * @description Forward a client's ChatGPT subscription with separate LiteLLM authentication. + */ + get: operations["codex_passthrough_route_backend_api_codex__endpoint__get"]; + put?: never; + /** + * Codex Passthrough Route + * @description Forward a client's ChatGPT subscription with separate LiteLLM authentication. + */ + post: operations["codex_passthrough_route_backend_api_codex__endpoint__post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/batches": { parameters: { query?: never; @@ -35309,6 +35333,11 @@ export interface components { cache_read_input_token_cost_priority?: number | null; /** Cache Read Input Token Cost Ultrafast */ cache_read_input_token_cost_ultrafast?: number | null; + /** + * Chatgpt Auth File + * @description Path to a ChatGPT OAuth auth.json for this deployment. + */ + chatgpt_auth_file?: string | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; /** Client Id */ @@ -50660,6 +50689,11 @@ export interface components { cache_read_input_token_cost_priority?: number | null; /** Cache Read Input Token Cost Ultrafast */ cache_read_input_token_cost_ultrafast?: number | null; + /** + * Chatgpt Auth File + * @description Path to a ChatGPT OAuth auth.json for this deployment. + */ + chatgpt_auth_file?: string | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; /** Client Id */ @@ -53737,6 +53771,77 @@ export interface operations { }; }; }; + codex_passthrough_route_backend_api_codex__endpoint__get: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: "models" | "responses"; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + codex_passthrough_route_backend_api_codex__endpoint__post: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: "models" | "responses"; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + responses: { + /** @description Unmodified Codex response or server-sent events */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + "text/event-stream": string; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; list_batches_batches_get: { parameters: { query?: { From ffb79b10ee706ebd5a7f4a5784deb0ad127fc2a3 Mon Sep 17 00:00:00 2001 From: Zap777 <124182405+devanshc777@users.noreply.github.com> Date: Tue, 27 Oct 2026 04:30:02 +0530 Subject: [PATCH 18/19] fix(mcp): send tool-call stream headers once the call is authorized (#47897) OAuthScopeResponse holds a tools/call response's http.response.start so a managed OAuth scope challenge can replace it with a 403. Once the call was allowed, the held start still waited for the next body message, which for a slow tool is the 15s SSE keep-alive ping. Send it as soon as the call is allowed instead. Fixes #46659 --- .../mcp_server/scope_challenge.py | 51 +++++---- .../test_mcp_oauth_passthrough_tools.py | 108 ++++++++++++++++-- .../mcp_server/test_mcp_server_manager.py | 4 +- .../test_mcp_server_tool_calls_and_headers.py | 3 +- 4 files changed, 127 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/scope_challenge.py b/litellm/proxy/_experimental/mcp_server/scope_challenge.py index c326418bee7..121234e305d 100644 --- a/litellm/proxy/_experimental/mcp_server/scope_challenge.py +++ b/litellm/proxy/_experimental/mcp_server/scope_challenge.py @@ -33,9 +33,8 @@ class OAuthScopeResponse: self._request_path = request_path self._ready = asyncio.Event() self._lock = asyncio.Lock() - self._start: Message | None = None self._challenge: MCPUpstreamAuthError | None = None - self._refused = False + self._closed = False @property def pending(self) -> bool: @@ -50,31 +49,33 @@ class OAuthScopeResponse: self._ready.set() async def __call__(self, message: Message) -> None: - if message["type"] == "http.response.start": - self._start = message - if message["status"] != 200: - self.allow() - return - await self._ready.wait() async with self._lock: - if self._refused: + if self._closed: return - if self._start is not None: - start: Final = self._start - self._start = None - if self._challenge is not None: - error: Final = self._challenge.to_http_exception(self._base_url, self._request_path) - response: Final = JSONResponse( - status_code=error.status_code, content={"detail": error.detail}, headers=error.headers - ) - self._refused = True - await self._send( - {"type": "http.response.start", "status": response.status_code, "headers": response.raw_headers} - ) - await self._send({"type": "http.response.body", "body": bytes(response.body)}) - return - await self._send(start) - await self._send(message) + try: + if message["type"] == "http.response.start": + if message["status"] != 200: + self.allow() + await self._ready.wait() + if self._challenge is not None: + error: Final = self._challenge.to_http_exception(self._base_url, self._request_path) + response: Final = JSONResponse( + status_code=error.status_code, content={"detail": error.detail}, headers=error.headers + ) + self._closed = True + await self._send( + { + "type": "http.response.start", + "status": response.status_code, + "headers": response.raw_headers, + } + ) + await self._send({"type": "http.response.body", "body": bytes(response.body)}) + return + await self._send(message) + except BaseException: + self._closed = True + raise def scope_response_from_request(request: object) -> OAuthScopeResponse | None: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py index b0d53f9c93b..34b021df2b9 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py @@ -5,6 +5,7 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations import asyncio import logging import sys +from typing import Final from unittest.mock import AsyncMock, MagicMock import httpx @@ -531,12 +532,12 @@ async def test_scope_challenge_replaces_pending_stream_headers_once(): send = AsyncMock() response = OAuthScopeResponse(send, "https://gateway.example", "/docs/mcp") - await response({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(response({"type": "http.response.start", "status": 200, "headers": []})) body = asyncio.create_task(response({"type": "http.response.body", "body": b"event: message\n"})) await asyncio.sleep(0) send.assert_not_awaited() response.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) - await asyncio.wait_for(body, 1) + await asyncio.wait_for(asyncio.gather(headers, body), 1) await response({"type": "http.response.body", "body": b"private tool error"}) assert send.await_count == 2 start, result = [call.args[0] for call in send.await_args_list] @@ -567,7 +568,7 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=request) token = active_mcp_request_ctx_var.set(context) try: - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) await asyncio.sleep(0) for method, status in [("initialize", 200), ("tools/list", 200), ("tools/call", 403)]: @@ -582,7 +583,7 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re upstream = MCPClient(server_url="https://upstream.example/mcp") async with upstream._create_httpx_client_factory(transport=httpx2.MockTransport(lambda request: httpx2.Response(200, content=b"tool result")))() as client: async with client.stream("POST", "https://upstream.example/mcp", json={"method": "tools/call"}): - await asyncio.wait_for(progress, 1) + await asyncio.wait_for(asyncio.gather(headers, progress), 1) assert send.await_args_list[-1].args[0]["body"] == b"progress" gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) await gate({"type": "http.response.body", "body": b"result"}) @@ -594,21 +595,105 @@ async def test_scope_authorization_releases_progress_without_waiting_for_tool_re @pytest.mark.asyncio @pytest.mark.parametrize("failure", [ValueError, asyncio.CancelledError]) -async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation(failure): +@pytest.mark.parametrize("fail_send", [False, True]) +async def test_scope_middleware_releases_waiting_response_on_validation_or_cancellation( + failure: type[BaseException], fail_send: bool, +) -> None: from mcp.server.context import ServerRequestContext from starlette.requests import Request from litellm.proxy._experimental.mcp_server.scope_challenge import SCOPE_RESPONSE_KEY, OAuthScopeResponse, finish_scope_response - send = AsyncMock() - gate = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") - context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + send: Final = AsyncMock(side_effect=BrokenPipeError("peer disconnected") if fail_send else None) + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + context: Final = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + await asyncio.sleep(0) with pytest.raises(failure): await finish_scope_response(context, AsyncMock(side_effect=failure)) - await gate({"type": "http.response.start", "status": 200, "headers": []}) + if fail_send: + with pytest.raises(BrokenPipeError, match="peer disconnected"): + await headers + else: + await headers await gate({"type": "http.response.body", "body": b"error"}) assert not gate.pending - assert send.await_args_list[-1].args[0]["body"] == b"error" + assert send.await_args_list[-1].args[0]["type"] == ("http.response.start" if fail_send else "http.response.body") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_send", [False, True]) +async def test_scope_transport_owns_header_send_cancellation_and_errors(fail_send: bool) -> None: + import httpx2 + from starlette.types import Message + + from litellm.proxy._experimental.mcp_server.legacy_callbacks import record_upstream_tool_authorization + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send_started: Final = asyncio.Event() + send_stopped: Final = asyncio.Event() + writable: Final = asyncio.Event() + sent: Final[list[Message]] = [] + + async def send(message: Message) -> None: + sent.append(message) + assert message["type"] == "http.response.start" + send_started.set() + try: + if fail_send: + raise BrokenPipeError("peer disconnected") + await writable.wait() + finally: + send_stopped.set() + + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) + body: Final = asyncio.create_task(gate({"type": "http.response.body", "body": b"result"})) + await asyncio.sleep(0) + assert not send_started.is_set() + response: Final = httpx2.Response(200, request=httpx2.Request( + "POST", "https://upstream.example/mcp", json={"method": "tools/call"}, + )) + try: + await record_upstream_tool_authorization(response, pending=gate) + await send_started.wait() + if fail_send: + with pytest.raises(BrokenPipeError, match="peer disconnected"): + await headers + else: + headers.cancel() + with pytest.raises(asyncio.CancelledError): + await headers + await body + assert len(sent) == 1 + assert send_stopped.is_set() + finally: + writable.set() + await send_stopped.wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("allowed_before_start", [False, True]) +async def test_scope_authorization_sends_stream_headers_before_the_first_body_chunk(allowed_before_start: bool) -> None: + from starlette.types import Message + + from litellm.proxy._experimental.mcp_server.scope_challenge import OAuthScopeResponse + + send: Final = AsyncMock() + gate: Final = OAuthScopeResponse(send, "https://gateway.example", "/mcp/docs") + start: Final[Message] = {"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")]} + if allowed_before_start: + gate.allow() + headers: Final = asyncio.create_task(gate(start)) + await asyncio.sleep(0) + if not allowed_before_start: + send.assert_not_awaited() + gate.allow() + await headers + send.assert_awaited_once_with(start) + await gate({"type": "http.response.body", "body": b"event: message\n", "more_body": True}) + assert send.await_args_list[-1].args[0]["body"] == b"event: message\n" + assert send.await_count == 2 @pytest.mark.asyncio @@ -697,13 +782,14 @@ async def test_scope_challenges_are_isolated_between_concurrent_requests(): context = ServerRequestContext(session=MagicMock(), lifespan_context={}, protocol_version="2025-11-25", method="tools/call", request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate})) token = active_mcp_request_ctx_var.set(context) try: - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) await asyncio.sleep(0) assert get_scope_response() is gate if denied: gate.deny(MCPUpstreamAuthError(403, None, "docs", required_scope="tools.write")) else: await record_upstream_tool_authorization(httpx2.Response(200, request=httpx2.Request("POST", "https://upstream.example/mcp", json={"method": "tools/call"})), pending=get_scope_response()) + await headers await gate({"type": "http.response.body", "body": b"done"}) return send.await_args_list[0].args[0]["status"] finally: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 7315315a619..ecd08e17cda 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -7681,7 +7681,7 @@ class TestMCPServerManager: request=Request({"type": "http", SCOPE_RESPONSE_KEY: gate}), ) token = active_mcp_request_ctx_var.set(context) - await gate({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(gate({"type": "http.response.start", "status": 200, "headers": []})) progress = asyncio.create_task(gate({"type": "http.response.body", "body": b"progress", "more_body": True})) async def upstream_call(*args, **kwargs): @@ -7701,7 +7701,7 @@ class TestMCPServerManager: assert not result.is_error finally: gate.allow() - await asyncio.wait_for(progress, 1) + await asyncio.wait_for(asyncio.gather(headers, progress), 1) active_mcp_request_ctx_var.reset(token) @pytest.mark.asyncio diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index b8e28c3f9b8..a873e2ab1ae 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py @@ -280,7 +280,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror( send = AsyncMock() response = OAuthScopeResponse(send, "https://gateway.example", "/mcp/pt") request = Request({"type": "http", "headers": [], SCOPE_RESPONSE_KEY: response}) - await response({"type": "http.response.start", "status": 200, "headers": []}) + headers: Final = asyncio.create_task(response({"type": "http.response.start", "status": 200, "headers": []})) try: from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.server import ( @@ -328,6 +328,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror( info_calls = [str(c.args[0]) for c in mock_logger.info.call_args_list if c.args] assert any("Upstream auth failure" in m for m in info_calls) response.allow() + await headers await response({"type": "http.response.body", "body": b"tool error"}) start = send.await_args_list[0].args[0] assert start["status"] == (403 if required_scope is not None else 200) From 593c640971728851d5dbf56941a26ef73e8095c6 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 19:00:38 -0400 Subject: [PATCH 19/19] fix(models): declare Azure audio snapshot capabilities (#48475) --- .../model_prices_and_context_window_backup.json | 4 ++++ model_prices_and_context_window.json | 4 ++++ .../test_model_prices_and_context_window.py | 17 +++++++++++++++++ 3 files changed, 25 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b7e0c881455..d9df9d97d7b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6206,6 +6206,8 @@ "text", "audio" ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6270,6 +6272,8 @@ "text", "audio" ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b7e0c881455..d9df9d97d7b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6206,6 +6206,8 @@ "text", "audio" ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -6270,6 +6272,8 @@ "text", "audio" ], + "supports_audio_input": true, + "supports_audio_output": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index c0abb18ca8c..45222b2c9ab 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -27,6 +27,23 @@ def prices() -> dict: return json.loads(PRICES_PATH.read_text()) +@pytest.mark.parametrize("model", ("azure/gpt-audio-1.5-2026-02-23", "azure/gpt-audio-mini-2025-10-06")) +def test_azure_audio_snapshot_capabilities_match_declared_modalities(model: str, local_model_cost_map: None) -> None: + catalog: Final = TypeAdapter(dict[str, JsonValue]).validate_json(PRICES_PATH.read_bytes()) + backup: Final = TypeAdapter(dict[str, JsonValue]).validate_json(BACKUP_PRICES_PATH.read_bytes()) + entry: Final = catalog[model] + assert isinstance(entry, dict) + assert entry == backup[model] + input_modalities: Final = entry["supported_modalities"] + output_modalities: Final = entry["supported_output_modalities"] + assert isinstance(input_modalities, list) + assert isinstance(output_modalities, list) + expected: Final = ("audio" in input_modalities, "audio" in output_modalities) + info: Final = litellm.get_model_info(model) + assert (info["supports_audio_input"], info["supports_audio_output"]) == expected + assert (litellm.supports_audio_input(model), litellm.supports_audio_output(model)) == expected + + class _CopilotEndpointRow(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore")