From 87d7e86479a934cf4f37b2de0723238b4fbb80a8 Mon Sep 17 00:00:00 2001 From: CHANGE Date: Fri, 10 Apr 2026 09:41:33 -0400 Subject: [PATCH 01/27] feat(helm): add tpl support to extraContainers and extraInitContainers Wrap toYaml with tpl in deployment and migration job templates so users can reference Helm values (e.g. {{ .Values.image.repository }}) inside extraContainers and extraInitContainers definitions. --- .../litellm-helm/templates/deployment.yaml | 4 +- .../templates/migrations-job.yaml | 4 +- .../litellm-helm/tests/deployment_tests.yaml | 58 ++++++++++++++++ .../tests/migrations-job_tests.yaml | 66 +++++++++++++++++++ 4 files changed, 128 insertions(+), 4 deletions(-) diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index 3040fb45d86..97123e5df69 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -47,7 +47,7 @@ spec: {{- toYaml .Values.podSecurityContext | nindent 8 }} {{- with .Values.extraInitContainers }} initContainers: - {{- toYaml . | nindent 8 }} + {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} containers: - name: {{ include "litellm.name" . }} @@ -212,7 +212,7 @@ spec: {{- toYaml . | nindent 12 }} {{- end }} {{- with .Values.extraContainers }} - {{- toYaml . | nindent 8 }} + {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} volumes: {{ if .Values.securityContext.readOnlyRootFilesystem }} diff --git a/deploy/charts/litellm-helm/templates/migrations-job.yaml b/deploy/charts/litellm-helm/templates/migrations-job.yaml index 8b93a60c1a3..c3f32fe32f3 100644 --- a/deploy/charts/litellm-helm/templates/migrations-job.yaml +++ b/deploy/charts/litellm-helm/templates/migrations-job.yaml @@ -37,7 +37,7 @@ spec: serviceAccountName: {{ include "litellm.migrationServiceAccountName" . }} {{- with .Values.migrationJob.extraInitContainers }} initContainers: - {{- toYaml . | nindent 8 }} + {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} containers: - name: prisma-migrations @@ -96,7 +96,7 @@ spec: {{- toYaml . | nindent 12 }} {{- end }} {{- with .Values.migrationJob.extraContainers }} - {{- toYaml . | nindent 8 }} + {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} {{- with .Values.volumes }} volumes: diff --git a/deploy/charts/litellm-helm/tests/deployment_tests.yaml b/deploy/charts/litellm-helm/tests/deployment_tests.yaml index 0d278f25693..b1cbafaf408 100644 --- a/deploy/charts/litellm-helm/tests/deployment_tests.yaml +++ b/deploy/charts/litellm-helm/tests/deployment_tests.yaml @@ -319,3 +319,61 @@ tests: asserts: - notExists: path: spec.minReadySeconds + - it: should work with extraInitContainers + template: deployment.yaml + set: + extraInitContainers: + - name: init-test + image: busybox:latest + command: ["echo", "hello"] + asserts: + - contains: + path: spec.template.spec.initContainers + content: + name: init-test + image: busybox:latest + command: ["echo", "hello"] + - it: should support tpl in extraInitContainers + template: deployment.yaml + set: + image: + repository: ghcr.io/berriai/litellm-database + tag: test + extraInitContainers: + - name: init-tpl + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + command: ["echo", "hello"] + asserts: + - contains: + path: spec.template.spec.initContainers + content: + name: init-tpl + image: "ghcr.io/berriai/litellm-database:test" + command: ["echo", "hello"] + - it: should work with extraContainers + template: deployment.yaml + set: + extraContainers: + - name: sidecar + image: busybox:latest + asserts: + - contains: + path: spec.template.spec.containers + content: + name: sidecar + image: busybox:latest + - it: should support tpl in extraContainers + template: deployment.yaml + set: + image: + repository: ghcr.io/berriai/litellm-database + tag: test + extraContainers: + - name: sidecar-tpl + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + asserts: + - contains: + path: spec.template.spec.containers + content: + name: sidecar-tpl + image: "ghcr.io/berriai/litellm-database:test" diff --git a/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml b/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml index ee684c3c3d7..05dd37b4857 100644 --- a/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml +++ b/deploy/charts/litellm-helm/tests/migrations-job_tests.yaml @@ -188,3 +188,69 @@ tests: - equal: path: spec.template.spec.serviceAccountName value: pre-existing-sa + - it: should work with extraInitContainers + template: migrations-job.yaml + set: + migrationJob: + enabled: true + extraInitContainers: + - name: init-test + image: busybox:latest + command: ["echo", "hello"] + asserts: + - contains: + path: spec.template.spec.initContainers + content: + name: init-test + image: busybox:latest + command: ["echo", "hello"] + - it: should support tpl in extraInitContainers + template: migrations-job.yaml + set: + image: + repository: ghcr.io/berriai/litellm-database + tag: test + migrationJob: + enabled: true + extraInitContainers: + - name: init-tpl + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + command: ["echo", "hello"] + asserts: + - contains: + path: spec.template.spec.initContainers + content: + name: init-tpl + image: "ghcr.io/berriai/litellm-database:test" + command: ["echo", "hello"] + - it: should work with extraContainers + template: migrations-job.yaml + set: + migrationJob: + enabled: true + extraContainers: + - name: sidecar + image: busybox:latest + asserts: + - contains: + path: spec.template.spec.containers + content: + name: sidecar + image: busybox:latest + - it: should support tpl in extraContainers + template: migrations-job.yaml + set: + image: + repository: ghcr.io/berriai/litellm-database + tag: test + migrationJob: + enabled: true + extraContainers: + - name: sidecar-tpl + image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}" + asserts: + - contains: + path: spec.template.spec.containers + content: + name: sidecar-tpl + image: "ghcr.io/berriai/litellm-database:test" From b28b82c1057d700eee2d75c6d7cadb744e3deadf Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 14 Apr 2026 21:46:13 -0700 Subject: [PATCH 02/27] chore(ui): migrate GuardrailTestPlayground off @tremor/react to antd Part of the ongoing migration away from @tremor/react per CLAUDE.md. Swaps tremor Card/Title/Text/TextInput for antd Card/Input and plain h2/h3/span/p tags with Tailwind. Drops the redundant Checkbox avatar in the sidebar list since row clicks already toggle selection and the blue-border + bg-blue-50 styling already signals selected state. --- .../guardrails/GuardrailTestPlayground.tsx | 41 +++++++------------ 1 file changed, 15 insertions(+), 26 deletions(-) diff --git a/ui/litellm-dashboard/src/components/guardrails/GuardrailTestPlayground.tsx b/ui/litellm-dashboard/src/components/guardrails/GuardrailTestPlayground.tsx index 462e6a5f684..751602079ca 100644 --- a/ui/litellm-dashboard/src/components/guardrails/GuardrailTestPlayground.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/GuardrailTestPlayground.tsx @@ -1,6 +1,5 @@ import React, { useState } from "react"; -import { Card, Title, Text, TextInput } from "@tremor/react"; -import { List, Empty, Spin, Checkbox } from "antd"; +import { Card, List, Empty, Spin, Input } from "antd"; import { ExperimentOutlined, SearchOutlined } from "@ant-design/icons"; import GuardrailTestPanel from "./GuardrailTestPanel"; import { applyGuardrail } from "../networking"; @@ -117,18 +116,18 @@ const GuardrailTestPlayground: React.FC = ({ return (
- +
{/* Left Sidebar - Guardrails List */}
- Guardrails - Guardrails + } placeholder="Search guardrails..." value={searchQuery} - onValueChange={setSearchQuery} + onChange={(e) => setSearchQuery(e.target.value)} />
@@ -156,24 +155,14 @@ const GuardrailTestPlayground: React.FC = ({ toggleGuardrailSelection(guardrail.guardrail_name); } }} - className={`cursor-pointer hover:bg-gray-50 transition-colors px-4 ${ + style={{ paddingLeft: 24, paddingRight: 16 }} + className={`cursor-pointer hover:bg-gray-50 transition-colors ${ selectedGuardrails.has(guardrail.guardrail_name || "") ? "bg-blue-50 border-l-4 border-l-blue-500" : "border-l-4 border-l-transparent" }`} > { - e.stopPropagation(); - if (guardrail.guardrail_name) { - toggleGuardrailSelection(guardrail.guardrail_name); - } - }} - /> - } title={
@@ -206,29 +195,29 @@ const GuardrailTestPlayground: React.FC = ({
- + {selectedGuardrails.size} of {filteredGuardrails.length} selected - +
{/* Right Panel - Test Area */}
- Guardrail Testing Playground +

Guardrail Testing Playground

{selectedGuardrails.size === 0 ? (
- +

Select Guardrails to Test - - +

+

Choose one or more guardrails from the left sidebar to start testing and comparing results. - +

) : (
From 17568e81f28befd3215b4f2e98123a1b7b16bd5b Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Wed, 15 Apr 2026 08:45:11 -0700 Subject: [PATCH 03/27] chore(ui): use antd Typography in GuardrailTestPlayground - Replace plain text/heading tags with Typography.Text, Title, Paragraph - Document Typography preference in CLAUDE.md UI guidelines --- CLAUDE.md | 2 +- .../guardrails/GuardrailTestPlayground.tsx | 18 ++++++++++-------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 043055408c2..a2716876b12 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -110,7 +110,7 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components: - When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select ### UI Component Library -- **Always use `antd` for new UI components** — we are migrating off of `@tremor/react`. Do not introduce new `Badge`, `Text`, `Card`, `Grid`, `Title`, or other imports from `@tremor/react` in any new or modified file. Use `antd` equivalents: `Tag` for labels, plain ``/`
` with Tailwind classes (or `Typography.Text`) for text, `Card` from `antd`, etc. Note that `antd` has no `"yellow"` Tag color — use `"gold"` for amber/yellow. +- **Always use `antd` for new UI components** — we are migrating off of `@tremor/react`. Do not introduce new `Badge`, `Text`, `Card`, `Grid`, `Title`, or other imports from `@tremor/react` in any new or modified file. Use `antd` equivalents: `Tag` for labels, `Typography.Text` / `Typography.Title` / `Typography.Paragraph` for textual content (avoid plain text-only ``, `

`, `` when Typography fits), and `Card` from `antd`. Note that `antd` has no `"yellow"` Tag color — use `"gold"` for amber/yellow. ### MCP OAuth / OpenAPI Transport Mapping - `TRANSPORT.OPENAPI` is a UI-only concept. The backend only accepts `"http"`, `"sse"`, or `"stdio"`. Always map it to `"http"` before any API call (including pre-OAuth temp-session calls). diff --git a/ui/litellm-dashboard/src/components/guardrails/GuardrailTestPlayground.tsx b/ui/litellm-dashboard/src/components/guardrails/GuardrailTestPlayground.tsx index 751602079ca..7ede8f4c8d7 100644 --- a/ui/litellm-dashboard/src/components/guardrails/GuardrailTestPlayground.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/GuardrailTestPlayground.tsx @@ -1,5 +1,5 @@ import React, { useState } from "react"; -import { Card, List, Empty, Spin, Input } from "antd"; +import { Card, List, Empty, Spin, Input, Typography } from "antd"; import { ExperimentOutlined, SearchOutlined } from "@ant-design/icons"; import GuardrailTestPanel from "./GuardrailTestPanel"; import { applyGuardrail } from "../networking"; @@ -195,29 +195,31 @@ const GuardrailTestPlayground: React.FC = ({

- + {selectedGuardrails.size} of {filteredGuardrails.length} selected - +
{/* Right Panel - Test Area */}
-

Guardrail Testing Playground

+ + Guardrail Testing Playground +
{selectedGuardrails.size === 0 ? (
-

+ Select Guardrails to Test -

-

+ + Choose one or more guardrails from the left sidebar to start testing and comparing results. -

+
) : (
From 3c9bf32611095f9f1e4e92a0cbc724391baae6fb Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 16 Apr 2026 13:08:31 +0530 Subject: [PATCH 04/27] fix(bedrock): prevent negative streaming costs with start-only cache usage Bedrock /v1/messages streams can report cache tokens only on message_start while message_delta carries only uncached input tokens. Merge cache fields onto the final delta usage and clamp negative text-token remainders in cost calc to keep usage/cost consistent. Made-with: Cursor --- .../litellm_core_utils/llm_cost_calc/utils.py | 3 + .../anthropic_claude3_transformation.py | 72 ++++++--- .../test_anthropic_claude3_transformation.py | 141 ++++++++++++++++++ 3 files changed, 199 insertions(+), 17 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 191231f3e66..3fd913958da 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -684,6 +684,9 @@ def generic_cost_per_token( # noqa: PLR0915 - cache_creation - image_tokens ) + # Clamp to zero: inconsistent streaming usage + if text_tokens < 0: + text_tokens = 0 prompt_tokens_details["text_tokens"] = text_tokens ( diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index d00a18fe7e3..542c43b6594 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -507,11 +507,14 @@ class AmazonAnthropicClaudeMessagesConfig( """ Bedrock invoke does not return SSE formatted data. This function is a wrapper to ensure litellm chunks are SSE formatted. - Bedrock's Anthropic-compatible streaming puts cache usage fields - (cache_creation_input_tokens, cache_read_input_tokens) only on - message_stop, not on message_start or message_delta. Claude Code's - SDK only merges usage from message_delta, so we promote those fields - from message_stop onto message_delta before yielding. + Bedrock's Anthropic-compatible streaming usually puts cache usage fields + (cache_creation_input_tokens, cache_read_input_tokens) on message_stop. + Some deployments (including GovCloud) emit the cache breakdown only on + ``message_start.message.usage``; ``message_delta`` / ``message_stop`` then + repeat uncached ``input_tokens`` only. We promote cache fields from + ``message_stop`` onto ``message_delta``, and when those are absent we + merge them from ``message_start`` so logging/cost sees a consistent usage + object (fixes negative input costs: LIT-2411). """ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, @@ -527,6 +530,27 @@ class AmazonAnthropicClaudeMessagesConfig( async for chunk in handler.async_sse_wrapper(patched_stream): yield chunk + @staticmethod + def _merge_message_start_cache_into_delta_usage( + delta_usage: Dict[str, Any], + start_usage: Optional[Dict[str, Any]], + ) -> None: + """ + Copy cache breakdown from message_start onto message_delta usage when + those keys are missing on the delta (GovCloud / some Bedrock streams). + """ + if not start_usage: + return + for field in ("cache_creation_input_tokens", "cache_read_input_tokens"): + if field not in delta_usage: + val = start_usage.get(field) + if val is not None: + delta_usage[field] = val + if "cache_creation" not in delta_usage: + cc = start_usage.get("cache_creation") + if cc is not None: + delta_usage["cache_creation"] = cc + @staticmethod async def _promote_message_stop_usage( completion_stream: AsyncIterator[ @@ -534,20 +558,13 @@ class AmazonAnthropicClaudeMessagesConfig( ], ) -> AsyncIterator[Union[bytes, GenericStreamingChunk, ModelResponseStream, dict]]: """ - Promote cache usage fields from message_stop onto message_delta. - - Bedrock reports input_tokens (uncached only) on message_start, and - the full breakdown (input_tokens, cache_creation_input_tokens, - cache_read_input_tokens) only on message_stop. Claude Code's SDK - merges usage from message_start and message_delta but ignores - message_stop. This method buffers message_delta and, when - message_stop arrives with cache usage, merges those fields into the - message_delta usage. input_tokens is kept as the uncached-only - count; downstream calculate_usage adds cache tokens to - prompt_tokens. + Promote cache usage fields onto message_delta from message_stop (and, + when stop lacks them, from message_start). Ensures the final usage + chunk that logging/cost sees is always self-consistent. """ _CACHE_FIELDS = ("cache_creation_input_tokens", "cache_read_input_tokens") - pending_delta = None + pending_delta: Optional[Dict[str, Any]] = None + start_usage_snapshot: Optional[Dict[str, Any]] = None async for chunk in completion_stream: if not isinstance(chunk, dict): @@ -559,6 +576,17 @@ class AmazonAnthropicClaudeMessagesConfig( chunk_type = chunk.get("type") + if chunk_type == "message_start": + msg = chunk.get("message") or {} + u = msg.get("usage") + if isinstance(u, dict): + start_usage_snapshot = dict(u) + if pending_delta is not None: + yield pending_delta + pending_delta = None + yield chunk + continue + if chunk_type == "message_delta": pending_delta = chunk continue @@ -575,6 +603,10 @@ class AmazonAnthropicClaudeMessagesConfig( if raw_input is not None: delta_usage["input_tokens"] = raw_input if isinstance(raw_input, int) else 0 + AmazonAnthropicClaudeMessagesConfig._merge_message_start_cache_into_delta_usage( + delta_usage, start_usage_snapshot + ) + if delta_usage: pending_delta["usage"] = delta_usage # type: ignore[arg-type] @@ -590,6 +622,12 @@ class AmazonAnthropicClaudeMessagesConfig( yield chunk if pending_delta is not None: + delta_usage = dict(pending_delta.get("usage") or {}) + AmazonAnthropicClaudeMessagesConfig._merge_message_start_cache_into_delta_usage( + delta_usage, start_usage_snapshot + ) + if delta_usage: + pending_delta["usage"] = delta_usage # type: ignore[arg-type] yield pending_delta diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index f0186f7891f..6b6beca4b66 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -587,6 +587,147 @@ async def test_promote_message_stop_usage_preserves_message_delta_output_tokens( assert delta_out["usage"]["input_tokens"] == 3 +@pytest.mark.asyncio +async def test_promote_message_start_cache_when_message_stop_omits_cache_fields(): + """ + GovCloud / some Bedrock streams put cache_read only on message_start; delta and + stop repeat uncached input_tokens only. Merging start cache onto message_delta + avoids inconsistent usage and negative input costs (LIT-2411). + """ + cfg = AmazonAnthropicClaudeMessagesConfig() + + async def _stream(): # type: ignore[return-type] + yield { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "model": "claude-sonnet-4-5-20250929", + "usage": { + "input_tokens": 10, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 22167, + "cache_creation": { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 0, + }, + "output_tokens": 4, + }, + }, + } + yield { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"input_tokens": 10, "output_tokens": 181}, + } + yield {"type": "message_stop", "usage": {"input_tokens": 10, "output_tokens": 181}} + + merged: list[dict] = [] + async for chunk in cfg._promote_message_stop_usage(_stream()): + if isinstance(chunk, dict): + merged.append(chunk) + + delta_chunks = [c for c in merged if c.get("type") == "message_delta"] + assert len(delta_chunks) == 1 + u = delta_chunks[0]["usage"] + assert u["input_tokens"] == 10 + assert u["output_tokens"] == 181 + assert u["cache_read_input_tokens"] == 22167 + assert u["cache_creation_input_tokens"] == 0 + + +@pytest.mark.asyncio +async def test_unified_bedrock_messages_cache_on_start_only_never_negative_cost(): + """ + Regression guard for LIT-2411: + If cache usage is present only on message_start (and omitted from + message_delta/message_stop), final reconstructed usage + cost must still + be consistent and non-negative. + """ + from litellm import completion_cost + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( + AnthropicPassthroughLoggingHandler, + ) + + cfg = AmazonAnthropicClaudeMessagesConfig() + + async def _stream(): # type: ignore[return-type] + yield { + "type": "message_start", + "message": { + "id": "msg_bdrk_01WuFzkDbE9KWgiWakMRNKcA", + "type": "message", + "role": "assistant", + "content": [], + "model": "claude-sonnet-4-5-20250929", + "stop_reason": None, + "stop_sequence": None, + "usage": { + "input_tokens": 10, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 22167, + "cache_creation": { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 0, + }, + "output_tokens": 4, + }, + }, + } + yield {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + yield { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello from regression test"}, + } + yield {"type": "content_block_stop", "index": 0} + yield { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 181, "input_tokens": 10}, + } + yield {"type": "message_stop", "usage": {"input_tokens": 10, "output_tokens": 181}} + + logging_obj = LiteLLMLoggingObj( + model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + call_type="chat", + start_time=datetime.now(), + litellm_call_id="test_cache_on_start_only_never_negative_cost", + function_id="test_cache_on_start_only_never_negative_cost", + ) + + collected: list[bytes] = [] + async for sse in cfg.bedrock_sse_wrapper( + completion_stream=_stream(), + litellm_logging_obj=logging_obj, + request_body={"model": "anthropic.claude-3-5-sonnet-20240620-v1:0"}, + ): + collected.append(sse) + + built = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( + all_chunks=collected, + model="anthropic.claude-3-5-sonnet-20240620-v1:0", + litellm_logging_obj=Mock(), + ) + assert built.usage is not None + assert built.usage.prompt_tokens == 22177 + assert built.usage.completion_tokens == 181 + assert built.usage.cache_creation_input_tokens == 0 + assert built.usage.cache_read_input_tokens == 22167 + + cost = completion_cost( + completion_response=built, + model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + custom_llm_provider="bedrock", + ) + assert cost > 0 + assert cost == pytest.approx(0.0093951, rel=0, abs=1e-9) + + @pytest.mark.asyncio async def test_unified_bedrock_messages_sse_usage_and_cost_claude_sonnet_46(): """ From 72be35f9b81da33c34a2c82c37ff69f9d998a3e0 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 16 Apr 2026 11:01:30 -0700 Subject: [PATCH 05/27] chore(ui): migrate router_settings page from Tremor to antd MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace @tremor/react components with antd equivalents: - Button → antd Button - TextInput → antd Input - Switch → antd Switch --- .../router_settings/LatencyBasedConfiguration.tsx | 4 ++-- .../router_settings/ReliabilityRetriesSection.tsx | 4 ++-- .../src/components/router_settings/TagFilteringToggle.tsx | 2 +- .../src/components/router_settings/index.tsx | 6 +++--- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx index 9c776c94a0b..1c2d52cf0c5 100644 --- a/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/LatencyBasedConfiguration.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { TextInput } from "@tremor/react"; +import { Input } from "antd"; interface routingStrategyArgs { ttl?: number; @@ -42,7 +42,7 @@ const LatencyBasedConfiguration: React.FC = ({

{paramExplanation[param] || ""}

- = ({

{routerFieldsMetadata[param]?.field_description || ""}

- = ({ accessToken, userRole, {/* Actions - Sticky at bottom */}
- -
From 260679679f7881d96cbbf53bbf23b4ffaf7091dc Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 16 Apr 2026 14:04:33 -0700 Subject: [PATCH 06/27] fix(ui): repair router_settings tests broken by full antd mock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The antd mocks in RouterSettingsForm.test.tsx and index.test.tsx replaced the entire antd module with only Select, so the Switch and Button used by nested components failed to render. Use importOriginal to preserve the rest of antd and override only Select. Also fix the TagFilteringToggle click assertion — antd's Switch fires onChange with (checked, event), so toHaveBeenCalledWith(true) was always going to miss. Assert the checked arg directly instead of coupling to antd's call signature. --- .../RouterSettingsForm.test.tsx | 41 +++++++++++-------- .../TagFilteringToggle.test.tsx | 3 +- .../components/router_settings/index.test.tsx | 38 +++++++++-------- 3 files changed, 46 insertions(+), 36 deletions(-) diff --git a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx index 3988f012814..a8ff485cf9b 100644 --- a/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/RouterSettingsForm.test.tsx @@ -4,25 +4,30 @@ import userEvent from "@testing-library/user-event"; import RouterSettingsForm from "./RouterSettingsForm"; import type { RouterSettingsFormValue } from "./RouterSettingsForm"; -// Use the same antd mock as RoutingStrategySelector to keep things consistent -vi.mock("antd", () => ({ - Select: Object.assign( - ({ value, onChange, children }: any) => ( - - ), - { - Option: ({ value, children }: any) => ( - +// Override antd Select (complex to drive in JSDOM) while preserving the rest +// of antd (Switch, Button, etc.) so nested components render normally. +vi.mock("antd", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + Select: Object.assign( + ({ value, onChange, children }: any) => ( + ), - } - ), -})); + { + Option: ({ value, children }: any) => ( + + ), + } + ), + }; +}); const defaultValue: RouterSettingsFormValue = { routerSettings: {}, diff --git a/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx b/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx index b469751ac2d..cc071117b82 100644 --- a/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/TagFilteringToggle.test.tsx @@ -90,6 +90,7 @@ describe("TagFilteringToggle", () => { await user.click(screen.getByRole("switch")); - expect(onToggle).toHaveBeenCalledWith(true); + expect(onToggle).toHaveBeenCalledTimes(1); + expect(onToggle.mock.calls[0][0]).toBe(true); }); }); diff --git a/ui/litellm-dashboard/src/components/router_settings/index.test.tsx b/ui/litellm-dashboard/src/components/router_settings/index.test.tsx index eea2717298a..80f0ed98c81 100644 --- a/ui/litellm-dashboard/src/components/router_settings/index.test.tsx +++ b/ui/litellm-dashboard/src/components/router_settings/index.test.tsx @@ -3,24 +3,28 @@ import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils" import userEvent from "@testing-library/user-event"; import RouterSettings from "./index"; -vi.mock("antd", () => ({ - Select: Object.assign( - ({ value, onChange, children }: any) => ( - - ), - { - Option: ({ value, children }: any) => ( - +vi.mock("antd", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + Select: Object.assign( + ({ value, onChange, children }: any) => ( + ), - } - ), -})); + { + Option: ({ value, children }: any) => ( + + ), + } + ), + }; +}); vi.mock("@/components/networking", () => ({ getCallbacksCall: vi.fn(), From b6d5728134c2d7f9410139628bae7708cfe62f82 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 16 Apr 2026 19:07:02 -0700 Subject: [PATCH 07/27] add support for pages param --- .../document_intelligence/transformation.py | 103 +++++++++++++- .../test_ocr_azure_document_intelligence.py | 126 ++++++++++++++++++ 2 files changed, 223 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 6ef309ca679..f2888b1cd2e 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -54,10 +54,87 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ Get supported OCR parameters for Azure Document Intelligence. - Azure DI has minimal optional parameters compared to Mistral OCR. - Most Mistral-specific params are ignored during transformation. + Azure DI exposes a `pages` query parameter on the analyze endpoint + (1-based, e.g. "1-3,5,7-9"). To keep the public request shape + aligned with Mistral OCR, callers pass `pages` using Mistral + semantics — a list of 0-based integers — or a pre-formatted + Azure-style string. Other Mistral-specific params (e.g. + `include_image_base64`) are not supported by Azure DI and are + ignored during transformation. """ - return [] + return ["pages"] + + def map_ocr_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + ) -> dict: + """ + Map OCR params to Azure DI format. + + Translates Mistral-style `pages` (list[int], 0-based) into Azure's + `pages` query string (1-based, e.g. "1,2,3" or "1-3,5"). A raw + string that already matches Azure's format is passed through + unchanged. + """ + pages = non_default_params.get("pages") + if pages is None: + return optional_params + + normalized = self._normalize_pages_param(pages) + if normalized: + optional_params["pages"] = normalized + return optional_params + + @staticmethod + def _normalize_pages_param(pages: Any) -> str: + """ + Convert a caller-provided `pages` value to Azure DI's query-string + form. Azure expects 1-based page numbers, grammar: `^(\\d+(-\\d+)?)(,\\s*(\\d+(-\\d+)?))*$`. + + Accepted inputs: + - list[int]: Mistral-style 0-based indices. Converted to 1-based + and joined (e.g. [0,1,2] -> "1,2,3"). + - list[str]: tokens like "1" or "3-5". Validated, joined as-is + (treated as Azure-native, i.e. 1-based). + - str: already in Azure format. Validated and whitespace-stripped. + """ + pages_pattern = re.compile(r"^\s*\d+(-\d+)?(\s*,\s*\d+(-\d+)?)*\s*$") + + if isinstance(pages, str): + if not pages_pattern.match(pages): + raise ValueError( + f"Invalid `pages` string for Azure Document Intelligence: " + f"{pages!r}. Expected format like '1-3,5,7-9'." + ) + return pages.replace(" ", "") + + if isinstance(pages, list): + if len(pages) == 0: + return "" + if all(isinstance(p, bool) for p in pages): + raise ValueError("`pages` must be integers, not booleans") + if all(isinstance(p, int) for p in pages): + if any(p < 0 for p in pages): + raise ValueError( + "`pages` integers must be >= 0 (Mistral 0-based indices)" + ) + # Mistral 0-based -> Azure 1-based. + return ",".join(str(p + 1) for p in sorted(set(pages))) + if all(isinstance(p, str) for p in pages): + joined = ",".join(p.strip() for p in pages) + if not pages_pattern.match(joined): + raise ValueError( + f"Invalid `pages` list for Azure Document Intelligence: " + f"{pages!r}. Expected tokens like '1' or '3-5'." + ) + return joined + + raise ValueError( + "`pages` must be a list[int] (0-based, Mistral-style) or a " + "string like '1-3,5,7-9'." + ) def validate_environment( self, @@ -141,7 +218,20 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): # Azure Document Intelligence analyze endpoint # Note: API version 2024-11-30+ uses /documentintelligence/ (not /formrecognizer/) - return f"{api_base}/documentintelligence/documentModels/{model_id}:analyze?api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" + url = ( + f"{api_base}/documentintelligence/documentModels/{model_id}:analyze" + f"?api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" + ) + + # Azure DI accepts `pages` as a query param (1-based, e.g. "1-3,5"). + # `optional_params` has already been normalized in `map_ocr_params`. + pages = optional_params.get("pages") if optional_params else None + if pages: + from urllib.parse import quote + + url += f"&pages={quote(str(pages), safe=',-')}" + + return url def _extract_base64_from_data_uri(self, data_uri: str) -> str: """ @@ -233,8 +323,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): data["urlSource"] = document_url verbose_logger.debug("Using urlSource for Azure Document Intelligence") - # Azure DI doesn't support most Mistral-specific params - # Ignore pages, include_image_base64, etc. + # Azure DI: `pages` is a query param (wired in get_complete_url), + # not a body field. Other Mistral-specific params (e.g. + # include_image_base64, image_limit) are unsupported and ignored. return OCRRequestData(data=data, files=None) diff --git a/tests/ocr_tests/test_ocr_azure_document_intelligence.py b/tests/ocr_tests/test_ocr_azure_document_intelligence.py index 9c1c9e134db..c67cfcbc98a 100644 --- a/tests/ocr_tests/test_ocr_azure_document_intelligence.py +++ b/tests/ocr_tests/test_ocr_azure_document_intelligence.py @@ -9,6 +9,10 @@ import os import pytest from base_ocr_unit_tests import BaseOCRTest +from litellm.constants import AZURE_DOCUMENT_INTELLIGENCE_API_VERSION +from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, +) class TestAzureDocumentIntelligenceOCR(BaseOCRTest): @@ -42,3 +46,125 @@ class TestAzureDocumentIntelligenceOCR(BaseOCRTest): "api_base": endpoint, } + +class TestAzureDocumentIntelligencePagesParam: + """ + Unit tests for the Mistral-compatible `pages` parameter translation to + Azure Document Intelligence's `pages` query string. + + These tests exercise the transformation layer directly and do not + require Azure credentials or a network call. + """ + + @pytest.fixture + def cfg(self) -> AzureDocumentIntelligenceOCRConfig: + return AzureDocumentIntelligenceOCRConfig() + + def test_get_supported_ocr_params_includes_pages(self, cfg): + assert cfg.get_supported_ocr_params("prebuilt-layout") == ["pages"] + + def test_map_ocr_params_mistral_zero_based_int_list(self, cfg): + mapped = cfg.map_ocr_params({"pages": [0, 1, 2]}, {}, "prebuilt-layout") + assert mapped == {"pages": "1,2,3"} + + def test_map_ocr_params_dedupes_and_sorts(self, cfg): + mapped = cfg.map_ocr_params({"pages": [2, 0, 0, 1]}, {}, "prebuilt-layout") + assert mapped == {"pages": "1,2,3"} + + def test_map_ocr_params_empty_list_omits_pages(self, cfg): + mapped = cfg.map_ocr_params({"pages": []}, {}, "prebuilt-layout") + assert mapped == {} + + def test_map_ocr_params_azure_native_string_range(self, cfg): + mapped = cfg.map_ocr_params({"pages": "3-9"}, {}, "prebuilt-layout") + assert mapped == {"pages": "3-9"} + + def test_map_ocr_params_azure_native_string_with_spaces_stripped(self, cfg): + mapped = cfg.map_ocr_params({"pages": "1-3, 5"}, {}, "prebuilt-layout") + assert mapped == {"pages": "1-3,5"} + + def test_map_ocr_params_list_of_string_tokens(self, cfg): + mapped = cfg.map_ocr_params({"pages": ["1", "3-5"]}, {}, "prebuilt-layout") + assert mapped == {"pages": "1,3-5"} + + def test_map_ocr_params_invalid_string_raises(self, cfg): + with pytest.raises(ValueError, match="Invalid `pages` string"): + cfg.map_ocr_params({"pages": "a,b"}, {}, "prebuilt-layout") + + def test_map_ocr_params_negative_index_raises(self, cfg): + with pytest.raises(ValueError, match="must be >= 0"): + cfg.map_ocr_params({"pages": [-1]}, {}, "prebuilt-layout") + + def test_map_ocr_params_bool_list_raises(self, cfg): + with pytest.raises(ValueError, match="must be integers, not booleans"): + cfg.map_ocr_params({"pages": [True, False]}, {}, "prebuilt-layout") + + def test_map_ocr_params_unsupported_type_raises(self, cfg): + with pytest.raises(ValueError): + cfg.map_ocr_params({"pages": 5}, {}, "prebuilt-layout") + + def test_get_complete_url_appends_pages_query(self, cfg): + url = cfg.get_complete_url( + api_base="https://example.cognitiveservices.azure.com/", + model="azure_ai/doc-intelligence/prebuilt-layout", + optional_params={"pages": "1-3,5"}, + ) + assert ( + f"api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" in url + ), url + assert "pages=1-3,5" in url, url + assert "/documentintelligence/documentModels/prebuilt-layout:analyze" in url + + def test_get_complete_url_no_pages_when_optional_params_empty(self, cfg): + url = cfg.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="prebuilt-layout", + optional_params={}, + ) + assert "pages=" not in url + + def test_transform_ocr_request_does_not_put_pages_in_body(self, cfg): + req = cfg.transform_ocr_request( + model="prebuilt-layout", + document={ + "type": "document_url", + "document_url": "https://example.com/x.pdf", + }, + optional_params={"pages": "1,2,3"}, + headers={}, + ) + assert req.data is not None + assert "pages" not in req.data + assert req.data.get("urlSource") == "https://example.com/x.pdf" + + def test_end_to_end_mistral_shape_to_azure_query(self, cfg): + """ + Caller sends Mistral-style `pages: [2,3,4,5,6,7,8]` (0-based, + meaning human pages 3-9). LiteLLM should turn that into Azure's + `&pages=3,4,5,6,7,8,9` on the analyze URL, and the body should + still only contain urlSource. + """ + non_default_params = {"pages": [2, 3, 4, 5, 6, 7, 8]} + optional_params = cfg.map_ocr_params( + non_default_params=non_default_params, + optional_params={}, + model="prebuilt-layout", + ) + url = cfg.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="prebuilt-layout", + optional_params=optional_params, + ) + req = cfg.transform_ocr_request( + model="prebuilt-layout", + document={ + "type": "document_url", + "document_url": "https://example.com/x.pdf", + }, + optional_params=optional_params, + headers={}, + ) + + assert "pages=3,4,5,6,7,8,9" in url + assert req.data == {"urlSource": "https://example.com/x.pdf"} + From 089ca5fc0f00f87a9ecb0420b83f54c0f9b160e1 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 16 Apr 2026 19:27:14 -0700 Subject: [PATCH 08/27] fix(azure-ocr): reject mixed bool+int in pages list validation Use any() instead of all() for bool check so lists like [True, 1, 2] raise ValueError; bool is a subclass of int so all(int) alone was insufficient. Made-with: Cursor --- .../llms/azure_ai/ocr/document_intelligence/transformation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index f2888b1cd2e..7eb47b03b48 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -113,7 +113,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): if isinstance(pages, list): if len(pages) == 0: return "" - if all(isinstance(p, bool) for p in pages): + if any(isinstance(p, bool) for p in pages): raise ValueError("`pages` must be integers, not booleans") if all(isinstance(p, int) for p in pages): if any(p < 0 for p in pages): From a84c27674fe3dcfad9832fe93b0020a53836ec1e Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 16 Apr 2026 19:28:32 -0700 Subject: [PATCH 09/27] refactor(azure-ocr): move urllib.parse.quote to module imports Remove inline import in get_complete_url; quote is stdlib with no circular import risk per project style. Made-with: Cursor --- .../llms/azure_ai/ocr/document_intelligence/transformation.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 7eb47b03b48..f80252ba6d8 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -11,6 +11,7 @@ import asyncio import re import time from typing import Any, Dict, Optional +from urllib.parse import quote import httpx @@ -227,8 +228,6 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): # `optional_params` has already been normalized in `map_ocr_params`. pages = optional_params.get("pages") if optional_params else None if pages: - from urllib.parse import quote - url += f"&pages={quote(str(pages), safe=',-')}" return url From d86c6a5b2fe357f8d03ac1de6292c3de25d0f115 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 17 Apr 2026 12:36:58 +0530 Subject: [PATCH 10/27] fix(proxy): prioritize reasoning health check token defaults Apply reasoning-first precedence for background health-check max tokens, parse reasoning env as optional, and raise non-wildcard fallback max_tokens from 1 to 5 for better reliability. Made-with: Cursor --- docs/my-website/docs/proxy/config_settings.md | 3 +- docs/my-website/docs/proxy/health.md | 26 ++++- litellm/constants.py | 19 ++++ litellm/proxy/health_check.py | 82 ++++++++++++-- .../proxy/test_health_check_max_tokens.py | 105 +++++++++++++++++- 5 files changed, 221 insertions(+), 14 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index be819a04ca1..a886a754f5e 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -487,7 +487,8 @@ router_settings: | AZURE_STORAGE_CLIENT_ID | The Application Client ID to use for Authentication to Azure Blob Storage logging | AZURE_STORAGE_CLIENT_SECRET | The Application Client Secret to use for Authentication to Azure Blob Storage logging | AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY | Cost per GB per day for Azure Vector Store service -| BACKGROUND_HEALTH_CHECK_MAX_TOKENS | Optional global default for `max_tokens` on proxy background health checks when a model has no `health_check_max_tokens`. If unset, non-wildcard models default to 1. Applies to wildcard routes when set. Default is unset +| BACKGROUND_HEALTH_CHECK_MAX_TOKENS | Optional global default for `max_tokens` on proxy background health checks when a model has no `health_check_max_tokens`. If unset, non-wildcard models default to 5. Applies to wildcard routes when set. Default is unset +| BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING | For **non-wildcard** reasoning models (`supports_reasoning(model)=true`), this takes precedence over `BACKGROUND_HEALTH_CHECK_MAX_TOKENS` when set. If unset, reasoning models fall back to `BACKGROUND_HEALTH_CHECK_MAX_TOKENS` (if set) or default behavior. Wildcard routes ignore this. Default is unset | BATCH_STATUS_POLL_INTERVAL_SECONDS | Interval in seconds for polling batch status. Default is 3600 (1 hour) | BATCH_STATUS_POLL_MAX_ATTEMPTS | Maximum number of attempts for polling batch status. Default is 24 (for 24 hours) | BEDROCK_MAX_POLICY_SIZE | Maximum size for Bedrock policy. Default is 75 diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 1d893961b62..535c90154bd 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -338,7 +338,7 @@ model_list: ## Health Check Max Tokens -By default, health checks use `max_tokens=1` to minimize cost and latency. For wildcard models, the default is `max_tokens=10`. +By default, health checks use `max_tokens=5` to balance reliability with low cost and latency. For wildcard models, the default is `max_tokens=10`. You can override this per-model by setting `health_check_max_tokens` in the `model_info` section of your config.yaml. @@ -352,6 +352,30 @@ model_list: health_check_max_tokens: 5 # 👈 OVERRIDE HEALTH CHECK MAX TOKENS ``` +### Reasoning vs non-reasoning defaults + +Reasoning models (per `supports_reasoning` in the model map) often need a higher health-check `max_tokens` because providers count reasoning tokens toward the completion budget. You can set **separate** limits without listing every model: + +**Per deployment (`model_info`)** — used when `health_check_max_tokens` is not set. Ignored for wildcard routes (`*` in `litellm_params.model`, i.e. the deployment model string; not `health_check_model`). + +```yaml +model_list: + - model_name: openai-stack + litellm_params: + model: openai/gpt-5-nano + api_key: os.environ/OPENAI_API_KEY + model_info: + health_check_max_tokens_reasoning: 128 + health_check_max_tokens_non_reasoning: 1 +``` + +**Global (environment)**: + +- `BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING` — for non-wildcard reasoning models, this value takes precedence when set +- `BACKGROUND_HEALTH_CHECK_MAX_TOKENS` — global fallback for all models (including wildcard routes) + +If neither is set, non-wildcard models default to `5` and wildcard routes omit `max_tokens`. + ## `/health/readiness` Unprotected endpoint for checking if proxy is ready to accept requests diff --git a/litellm/constants.py b/litellm/constants.py index e5f637e9f15..6c89cf5946d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1360,6 +1360,25 @@ try: ) except (ValueError, TypeError): BACKGROUND_HEALTH_CHECK_MAX_TOKENS = None + + +_background_health_check_max_tokens_reasoning_env = os.getenv( + "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING" +) +try: + _raw_background_health_check_max_tokens_reasoning = ( + _background_health_check_max_tokens_reasoning_env.strip() + if _background_health_check_max_tokens_reasoning_env is not None + else "" + ) + BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING: Optional[int] = ( + int(_raw_background_health_check_max_tokens_reasoning) + if _raw_background_health_check_max_tokens_reasoning + else None + ) +except (ValueError, TypeError): + BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING = None + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check" LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli" LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs" diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 1518ed66ab1..fbe056d41cb 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -13,6 +13,7 @@ import litellm logger = logging.getLogger(__name__) from litellm.constants import ( BACKGROUND_HEALTH_CHECK_MAX_TOKENS, + BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING, DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS, ) @@ -292,6 +293,75 @@ def build_deployment_health_states( return states +def _deployment_model_string_for_health_check(litellm_params: dict) -> str: + """Deployment model from litellm_params (before Bedrock rewrite). + + Used for reasoning vs non-reasoning max_tokens and wildcard detection only. + Does not use ``health_check_model``; that override applies later to the request. + """ + return litellm_params.get("model") or "" + + +def _health_check_deployment_is_wildcard(litellm_params: dict) -> bool: + return "*" in _deployment_model_string_for_health_check(litellm_params) + + +def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> Optional[int]: + """ + Pick max_tokens for the health check request. + + Priority: + 1. model_info.health_check_max_tokens (explicit override) + 2. For non-wildcard routes: health_check_max_tokens_reasoning / _non_reasoning + from model_info based on litellm.supports_reasoning(litellm_params["model"]) + 3. For non-wildcard reasoning routes: BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING + from env (if set) + 4. BACKGROUND_HEALTH_CHECK_MAX_TOKENS (global, any route including wildcards) + 5. Non-wildcard default: 5 + 6. Wildcard and nothing from (1)(4): leave unset (caller omits max_tokens) + """ + explicit = model_info.get("health_check_max_tokens", None) + if explicit is not None: + return int(explicit) + + is_wildcard = _health_check_deployment_is_wildcard(litellm_params) + deployment_model = _deployment_model_string_for_health_check(litellm_params) + + if not is_wildcard: + tokens_reasoning = model_info.get("health_check_max_tokens_reasoning", None) + tokens_non_reasoning = model_info.get( + "health_check_max_tokens_non_reasoning", None + ) + if tokens_reasoning is not None or tokens_non_reasoning is not None: + try: + is_reasoning = litellm.supports_reasoning(deployment_model) + except Exception: + is_reasoning = False + if is_reasoning and tokens_reasoning is not None: + return int(tokens_reasoning) + if not is_reasoning and tokens_non_reasoning is not None: + return int(tokens_non_reasoning) + + if not is_wildcard: + try: + is_reasoning = litellm.supports_reasoning(deployment_model) + except Exception: + is_reasoning = False + if ( + is_reasoning + and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None + ): + return int(BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING) + + if BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None: + return int(BACKGROUND_HEALTH_CHECK_MAX_TOKENS) + + if not is_wildcard: + return 5 + + return None + + def _update_litellm_params_for_health_check( model_info: dict, litellm_params: dict ) -> dict: @@ -304,15 +374,9 @@ def _update_litellm_params_for_health_check( - for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID """ litellm_params["messages"] = _get_random_llm_message() - _health_check_max_tokens = model_info.get("health_check_max_tokens", None) - if _health_check_max_tokens is not None: - litellm_params["max_tokens"] = _health_check_max_tokens - elif BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None: - litellm_params["max_tokens"] = BACKGROUND_HEALTH_CHECK_MAX_TOKENS - elif "*" not in ( - model_info.get("health_check_model") or litellm_params.get("model") or "" - ): - litellm_params["max_tokens"] = 1 + _resolved_max_tokens = _resolve_health_check_max_tokens(model_info, litellm_params) + if _resolved_max_tokens is not None: + litellm_params["max_tokens"] = _resolved_max_tokens _health_check_model = model_info.get("health_check_model", None) if _health_check_model is not None: diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py index 4d417b40b59..d5011e51fef 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -4,20 +4,25 @@ import pytest from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers from litellm.proxy import health_check as hc_module -from litellm.proxy.health_check import _update_litellm_params_for_health_check +from litellm.proxy.health_check import ( + _resolve_health_check_max_tokens, + _update_litellm_params_for_health_check, +) @pytest.mark.asyncio -async def test_update_litellm_params_max_tokens_default(): +async def test_update_litellm_params_max_tokens_default(monkeypatch): """ Test that max_tokens defaults to 1 for non-wildcard models. """ + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None) + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", None) model_info = {} litellm_params = {"model": "gpt-4"} updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) - assert updated_params["max_tokens"] == 1 + assert updated_params["max_tokens"] == 5 @pytest.mark.asyncio @@ -126,3 +131,97 @@ async def test_global_env_var_applies_to_wildcard_models(monkeypatch): updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) assert updated_params["max_tokens"] == 15 + + +def test_resolve_health_check_max_tokens_reasoning_specific_model_info(): + model_info = { + "health_check_max_tokens_reasoning": 64, + "health_check_max_tokens_non_reasoning": 2, + } + litellm_params = {"model": "openai/gpt-4o"} + + with patch.object(hc_module.litellm, "supports_reasoning", return_value=False): + assert _resolve_health_check_max_tokens(model_info, litellm_params) == 2 + + with patch.object(hc_module.litellm, "supports_reasoning", return_value=True): + assert _resolve_health_check_max_tokens(model_info, litellm_params) == 64 + + +def test_explicit_health_check_max_tokens_beats_reasoning_specific(): + model_info = { + "health_check_max_tokens": 9, + "health_check_max_tokens_reasoning": 64, + "health_check_max_tokens_non_reasoning": 2, + } + litellm_params = {"model": "openai/gpt-4o"} + + with patch.object(hc_module.litellm, "supports_reasoning", return_value=True): + assert _resolve_health_check_max_tokens(model_info, litellm_params) == 9 + + +def test_reasoning_specific_falls_through_when_wrong_branch_only(monkeypatch): + """Only non-reasoning key set but model is reasoning → fall back to default 5.""" + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None) + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", None) + model_info = {"health_check_max_tokens_non_reasoning": 3} + litellm_params = {"model": "openai/o1"} + + with patch.object(hc_module.litellm, "supports_reasoning", return_value=True): + assert _resolve_health_check_max_tokens(model_info, litellm_params) == 5 + + +@pytest.mark.asyncio +async def test_background_split_env_reasoning_vs_non_reasoning(monkeypatch): + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None) + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", 50) + + model_info = {} + litellm_params = {"model": "azure/gpt-4"} + + with patch.object(hc_module.litellm, "supports_reasoning", return_value=False): + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + assert updated["max_tokens"] == 5 + + litellm_params2 = {"model": "openai/o1"} + with patch.object(hc_module.litellm, "supports_reasoning", return_value=True): + updated2 = _update_litellm_params_for_health_check(model_info, litellm_params2) + assert updated2["max_tokens"] == 50 + + +@pytest.mark.asyncio +async def test_reasoning_env_precedence_over_global(monkeypatch): + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 10) + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", 20) + + model_info = {} + litellm_params = {"model": "openai/gpt-5.4"} + + with patch.object(hc_module.litellm, "supports_reasoning", return_value=True): + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + assert updated["max_tokens"] == 20 + + +@pytest.mark.asyncio +async def test_non_reasoning_uses_global_when_reasoning_env_set(monkeypatch): + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 10) + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", 20) + + model_info = {} + litellm_params = {"model": "azure/gpt-4"} + + with patch.object(hc_module.litellm, "supports_reasoning", return_value=False): + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + assert updated["max_tokens"] == 10 + + +def test_wildcard_ignores_reasoning_split_model_info(monkeypatch): + """Wildcard routes do not use reasoning/non-reasoning model_info split.""" + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None) + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", None) + model_info = { + "health_check_max_tokens_reasoning": 99, + "health_check_max_tokens_non_reasoning": 7, + } + litellm_params = {"model": "openai/*"} + + assert _resolve_health_check_max_tokens(model_info, litellm_params) is None From 76c69bb933d5709e4007ea8685bb4c8cdadc53f2 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 17 Apr 2026 12:49:31 +0530 Subject: [PATCH 11/27] fix(proxy): avoid duplicate reasoning capability lookup Compute supports_reasoning once per non-wildcard health-check resolution path and update the stale default-max-tokens test docstring. Made-with: Cursor --- litellm/proxy/health_check.py | 14 ++++---------- .../proxy/test_health_check_max_tokens.py | 2 +- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index fbe056d41cb..e0664703d28 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -328,25 +328,19 @@ def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> deployment_model = _deployment_model_string_for_health_check(litellm_params) if not is_wildcard: + try: + is_reasoning = litellm.supports_reasoning(deployment_model) + except Exception: + is_reasoning = False tokens_reasoning = model_info.get("health_check_max_tokens_reasoning", None) tokens_non_reasoning = model_info.get( "health_check_max_tokens_non_reasoning", None ) if tokens_reasoning is not None or tokens_non_reasoning is not None: - try: - is_reasoning = litellm.supports_reasoning(deployment_model) - except Exception: - is_reasoning = False if is_reasoning and tokens_reasoning is not None: return int(tokens_reasoning) if not is_reasoning and tokens_non_reasoning is not None: return int(tokens_non_reasoning) - - if not is_wildcard: - try: - is_reasoning = litellm.supports_reasoning(deployment_model) - except Exception: - is_reasoning = False if ( is_reasoning and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py index d5011e51fef..09211b72c3e 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -13,7 +13,7 @@ from litellm.proxy.health_check import ( @pytest.mark.asyncio async def test_update_litellm_params_max_tokens_default(monkeypatch): """ - Test that max_tokens defaults to 1 for non-wildcard models. + Test that max_tokens defaults to 5 for non-wildcard models. """ monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", None) monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING", None) From 66167f532e24e0fa3322d63c9eb4bfa4b433b5ca Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 16 Apr 2026 23:48:14 -0700 Subject: [PATCH 12/27] [Feature] UI - Settings: add forward_llm_provider_auth_headers toggle --- .../proxy_setting_endpoints.py | 11 +++ .../test_proxy_setting_endpoints.py | 77 +++++++++++++++++++ 2 files changed, 88 insertions(+) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 0349f289b4e..dd2b5aea7c8 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -100,6 +100,15 @@ class UISettings(BaseModel): description="If enabled, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription.", ) + forward_llm_provider_auth_headers: bool = Field( + default=False, + description=( + "If enabled, forwards LLM provider auth headers (x-api-key, " + "x-goog-api-key, api-key, ocp-apim-subscription-key) to the upstream " + "provider. Required for Claude Code BYOK (bring-your-own-key) flows." + ), + ) + enable_projects_ui: bool = Field( default=False, description="If enabled, shows the Projects feature in the UI sidebar and the project field in key management.", @@ -149,6 +158,7 @@ ALLOWED_UI_SETTINGS_FIELDS = { "enabled_ui_pages_internal_users", "require_auth_for_public_ai_hub", "forward_client_headers_to_llm_api", + "forward_llm_provider_auth_headers", "enable_projects_ui", "disable_agents_for_internal_users", "allow_agents_for_team_admins", @@ -162,6 +172,7 @@ ALLOWED_UI_SETTINGS_FIELDS = { # general_settings at runtime (on both read and write). _RUNTIME_GENERAL_SETTINGS_FLAGS = [ "forward_client_headers_to_llm_api", + "forward_llm_provider_auth_headers", "disable_agents_for_internal_users", "allow_agents_for_team_admins", "disable_vector_stores_for_internal_users", diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index bd9968ae936..f82017b14ae 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1016,6 +1016,83 @@ class TestProxySettingEndpoints: assert "unsupported_flag" not in stored_settings assert stored_settings["disable_model_add_for_internal_users"] is False + def test_update_ui_settings_persists_forward_llm_provider_auth_headers( + self, mock_auth, monkeypatch + ): + """BYOK flag must be allowlisted and persisted to litellm_uisettings.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + mock_prisma = MagicMock() + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + payload = {"forward_llm_provider_auth_headers": True} + + try: + response = client.patch("/update/ui_settings", json=payload) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + assert data["settings"]["forward_llm_provider_auth_headers"] is True + + assert mock_prisma.db.litellm_uisettings.upsert.called + call_args = mock_prisma.db.litellm_uisettings.upsert.call_args + create_data = call_args.kwargs["data"]["create"] + stored_settings = json.loads(create_data["ui_settings"]) + assert stored_settings["forward_llm_provider_auth_headers"] is True + + def test_update_ui_settings_syncs_forward_llm_provider_auth_headers_to_general_settings( + self, mock_auth, monkeypatch + ): + """BYOK flag must be synced into general_settings dict so the request path sees it.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Reset general_settings so the test is hermetic + general_settings: dict = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", general_settings + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + payload = {"forward_llm_provider_auth_headers": True} + + try: + response = client.patch("/update/ui_settings", json=payload) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + assert general_settings.get("forward_llm_provider_auth_headers") is True + def test_get_sso_settings_from_database(self, mock_proxy_config, mock_auth, monkeypatch): """Test getting SSO settings from the dedicated database table""" import json From d0f17064797970819e48891ec5c8fd7992e57903 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 16 Apr 2026 23:54:00 -0700 Subject: [PATCH 13/27] [Test] pre-call utils: regression coverage for BYOK x-api-key forwarding --- .../proxy/test_litellm_pre_call_utils.py | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index cf7e71b14d4..4c14aad077a 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import Request +from starlette.datastructures import Headers import litellm from litellm.proxy._types import TeamCallbackMetadata, UserAPIKeyAuth @@ -23,6 +24,7 @@ from litellm.proxy.litellm_pre_call_utils import ( add_guardrails_from_policy_engine, add_litellm_data_to_request, check_if_token_is_service_account, + clean_headers, ) from litellm.types.utils import CredentialItem @@ -2455,3 +2457,79 @@ def test_resolve_provider_hint_from_model_name(): "azure/gpt-4", config, None, pre_alias_model_name="gpt-4", provider="azure" ) assert result == "azure-cred" + + +def test_clean_headers_preserves_x_api_key_when_byok_enabled(): + """ + Regression test: when forward_llm_provider_auth_headers=True, + clean_headers() must preserve the client-supplied x-api-key header + so it can be forwarded to the upstream Anthropic API (BYOK flow). + """ + headers = Headers( + { + "x-api-key": "sk-ant-api03-client-key", + "x-litellm-api-key": "sk-proxy-virtual-key", + "content-type": "application/json", + } + ) + + result = clean_headers( + headers=headers, + litellm_key_header_name="x-litellm-api-key", + forward_llm_provider_auth_headers=True, + authenticated_with_header="x-litellm-api-key", + ) + + # x-api-key must be preserved for BYOK + assert result.get("x-api-key") == "sk-ant-api03-client-key" + # x-litellm-api-key must NOT leak to the upstream + assert "x-litellm-api-key" not in result + + +def test_clean_headers_strips_x_api_key_when_byok_disabled(): + """ + Regression test: with forward_llm_provider_auth_headers=False (default), + x-api-key must be stripped so proxy-configured keys are not overridden + by a client-supplied one. + """ + headers = Headers( + { + "x-api-key": "sk-ant-api03-client-key", + "x-litellm-api-key": "sk-proxy-virtual-key", + } + ) + + result = clean_headers( + headers=headers, + litellm_key_header_name="x-litellm-api-key", + forward_llm_provider_auth_headers=False, + authenticated_with_header="x-litellm-api-key", + ) + + assert "x-api-key" not in result + + +def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth_header(): + """ + Anti-replay regression: even when forward_llm_provider_auth_headers=True, + if the client authenticated TO the proxy using x-api-key (i.e., the proxy + key arrived as x-api-key), clean_headers() must NOT forward that header + upstream. Otherwise a proxy-auth key would leak to the LLM provider. + """ + headers = Headers( + { + "x-api-key": "sk-proxy-auth-key-masquerading-as-anthropic-key", + "content-type": "application/json", + } + ) + + result = clean_headers( + headers=headers, + litellm_key_header_name="x-litellm-api-key", + forward_llm_provider_auth_headers=True, + authenticated_with_header="x-api-key", + ) + + # Even with BYOK enabled, x-api-key must be stripped when it was used + # as the LiteLLM auth header (anti-replay guard). + assert "x-api-key" not in result From ad48af69aafabcbd655b7aaab4f8a63f73e868e2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 17 Apr 2026 00:00:15 -0700 Subject: [PATCH 14/27] [Feature] UI - Models: allow empty api_key for Anthropic provider (BYOK) --- .../provider_create_fields.json | 4 +-- .../public_endpoints/test_public_endpoints.py | 28 +++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 860593a6eab..b24f3df6272 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -218,8 +218,8 @@ "key": "api_key", "label": "API Key", "placeholder": "sk-", - "tooltip": null, - "required": true, + "tooltip": "Leave empty for BYOK (bring-your-own-key) flows, where clients forward their own Anthropic key via the x-api-key header. Requires the 'Forward LLM provider auth headers' UI setting to be enabled.", + "required": false, "field_type": "password", "options": null, "default_value": null diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index d62f88bf169..dc6def48fcb 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -138,6 +138,34 @@ def test_azure_provider_fields_include_entra_id(): assert fields_by_key["client_secret"]["required"] is False +def test_anthropic_provider_fields_support_byok(): + """ + The Anthropic provider form must allow BYOK: + - api_key is optional (not required) so admins can create models without a key + - api_key has a non-null tooltip explaining the BYOK use case + """ + app_instance = FastAPI() + app_instance.include_router(router) + test_client = TestClient(app_instance) + + response = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers = response.json() + + anthropic = next((p for p in providers if p["provider"] == "Anthropic"), None) + assert anthropic is not None, "Anthropic provider entry not found" + + fields_by_key = {f["key"]: f for f in anthropic["credential_fields"]} + assert "api_key" in fields_by_key + assert fields_by_key["api_key"]["required"] is False, ( + "Anthropic api_key must be optional so admins can configure BYOK models " + "without entering a key. See BYOK tutorial." + ) + assert fields_by_key["api_key"].get("tooltip"), ( + "Anthropic api_key must have a tooltip explaining the BYOK use case." + ) + + def test_public_model_hub_with_healthy_model(): """Test that health information is populated for a healthy model""" app = FastAPI() From 9a8aa4fae4270e83063f688d422ea088e9e4b2f7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 17 Apr 2026 00:04:38 -0700 Subject: [PATCH 15/27] [Feature] UI - Models: add api_base field for Anthropic provider form --- .../provider_create_fields.json | 10 ++++++++++ .../public_endpoints/test_public_endpoints.py | 17 +++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index b24f3df6272..98ec977b072 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -214,6 +214,16 @@ "provider_display_name": "Anthropic", "litellm_provider": "anthropic", "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": null, + "tooltip": "Optional. Override the upstream Anthropic API URL. Defaults to https://api.anthropic.com.", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, { "key": "api_key", "label": "API Key", diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index dc6def48fcb..ced5e2f43b0 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -164,6 +164,23 @@ def test_anthropic_provider_fields_support_byok(): assert fields_by_key["api_key"].get("tooltip"), ( "Anthropic api_key must have a tooltip explaining the BYOK use case." ) + assert "api_base" in fields_by_key, ( + "Anthropic provider form must expose api_base so cloud customers " + "can override the upstream URL without env var access." + ) + api_base_field = fields_by_key["api_base"] + assert api_base_field["required"] is False + assert api_base_field["field_type"] == "text" + assert api_base_field.get("tooltip"), ( + "api_base should have a tooltip explaining it is optional." + ) + + # UI forms render fields in credential_fields order; api_base should come first + # so an admin sees the URL override before the key field. + field_order = [f["key"] for f in anthropic["credential_fields"]] + assert field_order.index("api_base") < field_order.index("api_key"), ( + "api_base must appear before api_key in credential_fields (matches AI21 and ANTHROPIC_TEXT convention)." + ) def test_public_model_hub_with_healthy_model(): From 7eae18d1587959c4f2589bfca599059623356fc6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 17 Apr 2026 00:11:32 -0700 Subject: [PATCH 16/27] [Feature] UI - Settings: toggle row for forward_llm_provider_auth_headers --- .../AdminSettings/UISettings/UISettings.tsx | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index fb7c38449b0..9751ad435d5 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -17,6 +17,8 @@ export default function UISettings() { const disableTeamAdminDeleteProperty = schema?.properties?.disable_team_admin_delete_team_user; const requireAuthForPublicAIHubProperty = schema?.properties?.require_auth_for_public_ai_hub; const forwardClientHeadersProperty = schema?.properties?.forward_client_headers_to_llm_api; + const forwardLLMProviderAuthHeadersProperty = + schema?.properties?.forward_llm_provider_auth_headers; const enableProjectsUIProperty = schema?.properties?.enable_projects_ui; const enabledPagesProperty = schema?.properties?.enabled_ui_pages_internal_users; const disableAgentsProperty = schema?.properties?.disable_agents_for_internal_users; @@ -84,6 +86,20 @@ export default function UISettings() { ); }; + const handleToggleForwardLLMProviderAuthHeaders = (checked: boolean) => { + updateSettings( + { forward_llm_provider_auth_headers: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + const handleToggleEnableProjectsUI = (checked: boolean) => { updateSettings( { enable_projects_ui: checked }, @@ -284,6 +300,26 @@ export default function UISettings() { + + + + Forward LLM provider auth headers + + {forwardLLMProviderAuthHeadersProperty?.description ?? + "If enabled, forwards provider auth headers (x-api-key, x-goog-api-key, api-key) to the LLM provider. Required for Claude Code BYOK."} + + + + Date: Fri, 17 Apr 2026 00:13:56 -0700 Subject: [PATCH 17/27] [Docs] BYOK tutorial: document the UI-only configuration path --- docs/my-website/docs/tutorials/claude_code_byok.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/my-website/docs/tutorials/claude_code_byok.md b/docs/my-website/docs/tutorials/claude_code_byok.md index e1deac623bb..cbb937a59e5 100644 --- a/docs/my-website/docs/tutorials/claude_code_byok.md +++ b/docs/my-website/docs/tutorials/claude_code_byok.md @@ -35,6 +35,17 @@ By default, LiteLLM strips `x-api-key` from client requests for security. Settin ::: +:::tip Configure via UI instead of config.yaml + +You can also complete this setup from the LiteLLM admin UI: + +- Add the model via **Models → Add Model**, leaving the **API Key** field blank. +- Enable the toggle at **Settings → UI Settings → "Forward LLM provider auth headers"**. + +Both UI actions write to the database and override `config.yaml` at runtime. + +::: + ## Step 2: Create a LiteLLM Virtual Key Create a virtual key in the LiteLLM UI or via API. From 89365628c905852e73eafddf7f3542752c7a5511 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 4 Apr 2026 11:49:12 -0700 Subject: [PATCH 18/27] [Fix] Persist default end-user budget_id to DB so budget reset job picks up implicitly created users Previously, _apply_default_budget_to_end_user() only set the budget in-memory, leaving budget_id NULL in the database. This caused the budget reset job to skip these users since it filters by budget_id. Now the function also persists budget_id via a Prisma update call (non-fatal on failure). Fixes #22019 Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/auth/auth_checks.py | 16 ++- .../test_default_end_user_budget_simple.py | 127 +++++++++++++++++- 2 files changed, 136 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56958a88f6d..dd81afe78cb 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -857,8 +857,22 @@ async def _apply_default_budget_to_end_user( ) if default_budget is not None: - # Apply default budget to end user object + # Apply default budget to end user object (in-memory for this request) end_user_obj.litellm_budget_table = default_budget + + # Persist budget_id to DB so the budget reset job can find this user + try: + await prisma_client.db.litellm_endusertable.update( + where={"user_id": end_user_obj.user_id}, + data={"budget_id": litellm.max_end_user_budget_id}, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to persist default budget_id for end user %s: %s", + end_user_obj.user_id, + e, + ) + verbose_proxy_logger.debug( f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}" ) diff --git a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py b/tests/proxy_unit_tests/test_default_end_user_budget_simple.py index 92ca1f71703..8c85a79cbf8 100644 --- a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py +++ b/tests/proxy_unit_tests/test_default_end_user_budget_simple.py @@ -52,21 +52,22 @@ async def test_default_budget_applied_to_end_user_without_budget(): mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: mock_end_user_data) ) + mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None) mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: default_budget.dict()) ) - + mock_cache = AsyncMock(spec=DualCache) mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() - + result = await get_end_user_object( end_user_id=end_user_id, prisma_client=mock_prisma_client, user_api_key_cache=mock_cache, route="/chat/completions", ) - + # Verify default budget was applied assert result is not None assert result.litellm_budget_table is not None @@ -74,7 +75,13 @@ async def test_default_budget_applied_to_end_user_without_budget(): assert result.litellm_budget_table.max_budget == 10.0 assert result.litellm_budget_table.rpm_limit == 2 assert result.litellm_budget_table.tpm_limit == 10 - + + # Verify budget_id was persisted to DB + mock_prisma_client.db.litellm_endusertable.update.assert_called_once_with( + where={"user_id": end_user_id}, + data={"budget_id": default_budget_id}, + ) + litellm.max_end_user_budget_id = None @@ -162,14 +169,15 @@ async def test_budget_enforcement_blocks_over_budget_users(): mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: mock_end_user_data) ) + mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None) mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: default_budget.dict()) ) - + mock_cache = AsyncMock(spec=DualCache) mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() - + # Should raise BudgetExceededError with pytest.raises(litellm.BudgetExceededError) as exc_info: await get_end_user_object( @@ -185,6 +193,113 @@ async def test_budget_enforcement_blocks_over_budget_users(): litellm.max_end_user_budget_id = None +@pytest.mark.asyncio +async def test_default_budget_db_persist_failure_is_nonfatal(): + """ + If the DB update to persist budget_id fails, the budget should still be + applied in-memory for the current request (non-fatal warning). + """ + end_user_id = f"test_user_{uuid.uuid4().hex}" + default_budget_id = str(uuid.uuid4()) + litellm.max_end_user_budget_id = default_budget_id + + default_budget = LiteLLM_BudgetTable( + budget_id=default_budget_id, + max_budget=10.0, + ) + + mock_end_user_data = { + "user_id": end_user_id, + "spend": 1.0, + "litellm_budget_table": None, + "alias": None, + "allowed_model_region": None, + "default_model": None, + "blocked": False, + } + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( + return_value=MagicMock(dict=lambda: mock_end_user_data) + ) + # Simulate DB update failure + mock_prisma_client.db.litellm_endusertable.update = AsyncMock( + side_effect=Exception("DB connection lost") + ) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=MagicMock(dict=lambda: default_budget.model_dump()) + ) + + mock_cache = AsyncMock(spec=DualCache) + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + result = await get_end_user_object( + end_user_id=end_user_id, + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + route="/chat/completions", + ) + + # Budget should still be applied in-memory despite DB failure + assert result is not None + assert result.litellm_budget_table is not None + assert result.litellm_budget_table.budget_id == default_budget_id + + litellm.max_end_user_budget_id = None + + +@pytest.mark.asyncio +async def test_explicit_budget_skips_db_update(): + """ + End users with an explicit budget should NOT trigger a DB update + for the default budget_id. + """ + end_user_id = f"test_user_{uuid.uuid4().hex}" + explicit_budget_id = str(uuid.uuid4()) + default_budget_id = str(uuid.uuid4()) + litellm.max_end_user_budget_id = default_budget_id + + explicit_budget = LiteLLM_BudgetTable( + budget_id=explicit_budget_id, + max_budget=100.0, + ) + + mock_end_user_data = { + "user_id": end_user_id, + "spend": 10.0, + "litellm_budget_table": explicit_budget.model_dump(), + "alias": None, + "allowed_model_region": None, + "default_model": None, + "blocked": False, + } + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( + return_value=MagicMock(dict=lambda: mock_end_user_data) + ) + mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None) + + mock_cache = AsyncMock(spec=DualCache) + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + + result = await get_end_user_object( + end_user_id=end_user_id, + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + route="/chat/completions", + ) + + # Should keep explicit budget, NOT call update + assert result is not None + assert result.litellm_budget_table.budget_id == explicit_budget_id + mock_prisma_client.db.litellm_endusertable.update.assert_not_called() + + litellm.max_end_user_budget_id = None + + @pytest.mark.asyncio async def test_system_works_without_default_budget_configured(): """ From bb9955beca7baee74d0339f075b237cbe812ad70 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 4 Apr 2026 18:10:05 -0700 Subject: [PATCH 19/27] [Fix] Budget reset job now resets implicitly-created end users with NULL budget_id When litellm.max_end_user_budget_id is configured, implicitly-created end users (via /chat/completions) have budget_id=NULL in the DB since the default budget is only applied in-memory. The budget reset job filtered by budget_id, so these users were never reset and eventually permanently blocked. Fix: when the default budget is in the reset list, also query for and reset end users with budget_id=NULL and spend > 0. This keeps the hot auth path unchanged (no DB writes on every request). Fixes #22019 Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/auth/auth_checks.py | 16 +- .../proxy/common_utils/reset_budget_job.py | 47 ++++- .../test_default_end_user_budget_simple.py | 154 ++------------- .../common_utils/test_reset_budget_job.py | 185 ++++++++++++++++++ 4 files changed, 247 insertions(+), 155 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dd81afe78cb..56958a88f6d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -857,22 +857,8 @@ async def _apply_default_budget_to_end_user( ) if default_budget is not None: - # Apply default budget to end user object (in-memory for this request) + # Apply default budget to end user object end_user_obj.litellm_budget_table = default_budget - - # Persist budget_id to DB so the budget reset job can find this user - try: - await prisma_client.db.litellm_endusertable.update( - where={"user_id": end_user_obj.user_id}, - data={"budget_id": litellm.max_end_user_budget_id}, - ) - except Exception as e: - verbose_proxy_logger.warning( - "Failed to persist default budget_id for end user %s: %s", - end_user_obj.user_id, - e, - ) - verbose_proxy_logger.debug( f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}" ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 16243038b78..d935c50ffd3 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -4,6 +4,7 @@ import time from datetime import datetime, timezone from typing import List, Literal, Optional, Union +import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( LiteLLM_BudgetTableFull, @@ -162,16 +163,35 @@ class ResetBudgetJob: table_name="budget", ) + budget_ids_to_reset = [ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ] + endusers_to_reset = await self.prisma_client.get_data( table_name="enduser", query_type="find_all", - budget_id_list=[ - budget.budget_id - for budget in budgets_to_reset - if budget.budget_id is not None - ], + budget_id_list=budget_ids_to_reset, ) + # Also reset end users with no budget_id (NULL) who use the + # default budget via litellm.max_end_user_budget_id. These + # users are enforced in-memory but never had budget_id + # persisted, so the query above misses them. + if ( + litellm.max_end_user_budget_id is not None + and litellm.max_end_user_budget_id in budget_ids_to_reset + ): + default_budget_endusers = ( + await self._get_endusers_with_no_budget_id() + ) + if default_budget_endusers: + if endusers_to_reset is None: + endusers_to_reset = default_budget_endusers + else: + endusers_to_reset.extend(default_budget_endusers) + await self.reset_budget_for_litellm_team_members( budgets_to_reset=budgets_to_reset ) @@ -279,6 +299,23 @@ class ResetBudgetJob: ) verbose_proxy_logger.exception("Failed to reset budget for endusers: %s", e) + async def _get_endusers_with_no_budget_id( + self, + ) -> List[LiteLLM_EndUserTable]: + """ + Fetch end users that have no explicit budget_id set (NULL) and have + accumulated spend > 0. These are implicitly-created end users that + rely on the default budget (litellm.max_end_user_budget_id) applied + in-memory during auth checks. + """ + rows = await self.prisma_client.db.litellm_endusertable.find_many( + where={ + "budget_id": None, + "spend": {"gt": 0}, + }, + ) + return [LiteLLM_EndUserTable(**row.dict()) for row in rows] + async def reset_budget_for_litellm_keys(self): """ Resets the budget for all the litellm keys diff --git a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py b/tests/proxy_unit_tests/test_default_end_user_budget_simple.py index 8c85a79cbf8..970a7ab4718 100644 --- a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py +++ b/tests/proxy_unit_tests/test_default_end_user_budget_simple.py @@ -29,14 +29,14 @@ async def test_default_budget_applied_to_end_user_without_budget(): end_user_id = f"test_user_{uuid.uuid4().hex}" default_budget_id = str(uuid.uuid4()) litellm.max_end_user_budget_id = default_budget_id - + default_budget = LiteLLM_BudgetTable( budget_id=default_budget_id, max_budget=10.0, rpm_limit=2, tpm_limit=10, ) - + # Mock end user in DB without budget mock_end_user_data = { "user_id": end_user_id, @@ -47,12 +47,11 @@ async def test_default_budget_applied_to_end_user_without_budget(): "default_model": None, "blocked": False, } - + mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: mock_end_user_data) ) - mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None) mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: default_budget.dict()) ) @@ -76,12 +75,6 @@ async def test_default_budget_applied_to_end_user_without_budget(): assert result.litellm_budget_table.rpm_limit == 2 assert result.litellm_budget_table.tpm_limit == 10 - # Verify budget_id was persisted to DB - mock_prisma_client.db.litellm_endusertable.update.assert_called_once_with( - where={"user_id": end_user_id}, - data={"budget_id": default_budget_id}, - ) - litellm.max_end_user_budget_id = None @@ -95,13 +88,13 @@ async def test_explicit_budget_not_overridden_by_default(): explicit_budget_id = str(uuid.uuid4()) default_budget_id = str(uuid.uuid4()) litellm.max_end_user_budget_id = default_budget_id - + explicit_budget = LiteLLM_BudgetTable( budget_id=explicit_budget_id, max_budget=100.0, rpm_limit=50, ) - + # Mock end user with explicit budget mock_end_user_data = { "user_id": end_user_id, @@ -112,29 +105,29 @@ async def test_explicit_budget_not_overridden_by_default(): "default_model": None, "blocked": False, } - + mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: mock_end_user_data) ) - + mock_cache = AsyncMock(spec=DualCache) mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() - + result = await get_end_user_object( end_user_id=end_user_id, prisma_client=mock_prisma_client, user_api_key_cache=mock_cache, route="/chat/completions", ) - + # Verify explicit budget is kept (not replaced with default) assert result is not None assert result.litellm_budget_table.budget_id == explicit_budget_id assert result.litellm_budget_table.max_budget == 100.0 assert result.litellm_budget_table.rpm_limit == 50 - + litellm.max_end_user_budget_id = None @@ -147,13 +140,13 @@ async def test_budget_enforcement_blocks_over_budget_users(): end_user_id = f"test_user_{uuid.uuid4().hex}" default_budget_id = str(uuid.uuid4()) litellm.max_end_user_budget_id = default_budget_id - + default_budget = LiteLLM_BudgetTable( budget_id=default_budget_id, max_budget=10.0, rpm_limit=2, ) - + # Mock end user who has already spent more than budget mock_end_user_data = { "user_id": end_user_id, @@ -164,12 +157,11 @@ async def test_budget_enforcement_blocks_over_budget_users(): "default_model": None, "blocked": False, } - + mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: mock_end_user_data) ) - mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None) mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: default_budget.dict()) ) @@ -186,116 +178,9 @@ async def test_budget_enforcement_blocks_over_budget_users(): user_api_key_cache=mock_cache, route="/chat/completions", ) - + assert "ExceededBudget" in str(exc_info.value) assert end_user_id in str(exc_info.value) - - litellm.max_end_user_budget_id = None - - -@pytest.mark.asyncio -async def test_default_budget_db_persist_failure_is_nonfatal(): - """ - If the DB update to persist budget_id fails, the budget should still be - applied in-memory for the current request (non-fatal warning). - """ - end_user_id = f"test_user_{uuid.uuid4().hex}" - default_budget_id = str(uuid.uuid4()) - litellm.max_end_user_budget_id = default_budget_id - - default_budget = LiteLLM_BudgetTable( - budget_id=default_budget_id, - max_budget=10.0, - ) - - mock_end_user_data = { - "user_id": end_user_id, - "spend": 1.0, - "litellm_budget_table": None, - "alias": None, - "allowed_model_region": None, - "default_model": None, - "blocked": False, - } - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( - return_value=MagicMock(dict=lambda: mock_end_user_data) - ) - # Simulate DB update failure - mock_prisma_client.db.litellm_endusertable.update = AsyncMock( - side_effect=Exception("DB connection lost") - ) - mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( - return_value=MagicMock(dict=lambda: default_budget.model_dump()) - ) - - mock_cache = AsyncMock(spec=DualCache) - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_set_cache = AsyncMock() - - result = await get_end_user_object( - end_user_id=end_user_id, - prisma_client=mock_prisma_client, - user_api_key_cache=mock_cache, - route="/chat/completions", - ) - - # Budget should still be applied in-memory despite DB failure - assert result is not None - assert result.litellm_budget_table is not None - assert result.litellm_budget_table.budget_id == default_budget_id - - litellm.max_end_user_budget_id = None - - -@pytest.mark.asyncio -async def test_explicit_budget_skips_db_update(): - """ - End users with an explicit budget should NOT trigger a DB update - for the default budget_id. - """ - end_user_id = f"test_user_{uuid.uuid4().hex}" - explicit_budget_id = str(uuid.uuid4()) - default_budget_id = str(uuid.uuid4()) - litellm.max_end_user_budget_id = default_budget_id - - explicit_budget = LiteLLM_BudgetTable( - budget_id=explicit_budget_id, - max_budget=100.0, - ) - - mock_end_user_data = { - "user_id": end_user_id, - "spend": 10.0, - "litellm_budget_table": explicit_budget.model_dump(), - "alias": None, - "allowed_model_region": None, - "default_model": None, - "blocked": False, - } - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( - return_value=MagicMock(dict=lambda: mock_end_user_data) - ) - mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=None) - - mock_cache = AsyncMock(spec=DualCache) - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_set_cache = AsyncMock() - - result = await get_end_user_object( - end_user_id=end_user_id, - prisma_client=mock_prisma_client, - user_api_key_cache=mock_cache, - route="/chat/completions", - ) - - # Should keep explicit budget, NOT call update - assert result is not None - assert result.litellm_budget_table.budget_id == explicit_budget_id - mock_prisma_client.db.litellm_endusertable.update.assert_not_called() litellm.max_end_user_budget_id = None @@ -308,7 +193,7 @@ async def test_system_works_without_default_budget_configured(): """ end_user_id = f"test_user_{uuid.uuid4().hex}" litellm.max_end_user_budget_id = None # Not configured - + # Mock end user without budget mock_end_user_data = { "user_id": end_user_id, @@ -319,25 +204,24 @@ async def test_system_works_without_default_budget_configured(): "default_model": None, "blocked": False, } - + mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_endusertable.find_unique = AsyncMock( return_value=MagicMock(dict=lambda: mock_end_user_data) ) - + mock_cache = AsyncMock(spec=DualCache) mock_cache.async_get_cache = AsyncMock(return_value=None) mock_cache.async_set_cache = AsyncMock() - + result = await get_end_user_object( end_user_id=end_user_id, prisma_client=mock_prisma_client, user_api_key_cache=mock_cache, route="/chat/completions", ) - + # Should work fine, just without budget limits assert result is not None assert result.user_id == end_user_id assert result.litellm_budget_table is None # No budget applied - diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 1f2d4f4905f..db531bb3bb3 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -36,10 +36,24 @@ class MockLiteLLMVerificationToken: return {"count": 1} +class MockLiteLLMEndUserTable: + def __init__(self): + self.find_many_calls: List[Dict[str, Any]] = [] + self._find_many_results: List[Any] = [] + + def set_find_many_results(self, results: List[Any]): + self._find_many_results = results + + async def find_many(self, where: Dict[str, Any]) -> List[Any]: + self.find_many_calls.append({"where": where}) + return self._find_many_results + + class MockDB: def __init__(self): self.litellm_teammembership = MockLiteLLMTeamMembership() self.litellm_verificationtoken = MockLiteLLMVerificationToken() + self.litellm_endusertable = MockLiteLLMEndUserTable() class MockPrismaClient: @@ -613,3 +627,174 @@ def test_budget_table_reset_also_resets_linked_keys( ) assert calls[0]["where"]["budget_id"] == {"in": ["7d-budget-tier"]} assert calls[0]["data"]["spend"] == 0 + + +def test_reset_budget_resets_endusers_with_null_budget_id( + reset_budget_job, mock_prisma_client +): + """ + When litellm.max_end_user_budget_id is configured and that budget is + being reset, end users with budget_id=NULL should also have their spend + reset. These users were implicitly created and have no budget_id persisted, + but are enforced against the default budget in-memory. + """ + import litellm + + now = datetime.now(timezone.utc) + default_budget_id = "default-enduser-budget" + litellm.max_end_user_budget_id = default_budget_id + + # Budget that is due for reset — matches the default end user budget + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 50.0, + "budget_duration": "1d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": default_budget_id, + "created_at": now - timedelta(days=1), + }, + ) + + # End user WITH explicit budget_id (found by the normal budget_id_list query) + enduser_with_budget = type( + "LiteLLM_EndUserTable", + (), + { + "spend": 30.0, + "litellm_budget_table": test_budget, + "user_id": "enduser-explicit", + }, + ) + + # End user WITHOUT budget_id (NULL) — should also be reset + enduser_no_budget_row = type( + "EndUserRow", + (), + { + "spend": 25.0, + "user_id": "enduser-implicit", + "budget_id": None, + "alias": None, + "allowed_model_region": None, + "default_model": None, + "blocked": False, + "object_permission_id": None, + "object_permission": None, + "litellm_budget_table": None, + "dict": lambda self=None: { + "spend": 25.0, + "user_id": "enduser-implicit", + "blocked": False, + "alias": None, + "allowed_model_region": None, + "default_model": None, + "litellm_budget_table": None, + "object_permission_id": None, + "object_permission": None, + }, + }, + ) + + mock_prisma_client.data["budget"] = [test_budget] + mock_prisma_client.data["enduser"] = [enduser_with_budget] + + # Set up the DB mock for NULL-budget-id end users + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [enduser_no_budget_row] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + # Both end users should have been reset + updated = mock_prisma_client.updated_data["enduser"] + assert len(updated) == 2, ( + f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}" + ) + + user_ids = {u.user_id for u in updated} + assert "enduser-explicit" in user_ids + assert "enduser-implicit" in user_ids + + for u in updated: + assert u.spend == 0.0, f"Expected spend=0 for {u.user_id}, got {u.spend}" + + # Verify find_many was called to fetch NULL-budget-id end users + find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls + assert len(find_many_calls) == 1 + assert find_many_calls[0]["where"] == {"budget_id": None, "spend": {"gt": 0}} + + litellm.max_end_user_budget_id = None + + +def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured( + reset_budget_job, mock_prisma_client +): + """ + When litellm.max_end_user_budget_id is NOT configured, end users with + budget_id=NULL should NOT be fetched or reset. + """ + import litellm + + now = datetime.now(timezone.utc) + litellm.max_end_user_budget_id = None + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 50.0, + "budget_duration": "1d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": "some-budget", + "created_at": now - timedelta(days=1), + }, + ) + + mock_prisma_client.data["budget"] = [test_budget] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + # Should NOT have queried for NULL-budget-id end users + find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls + assert len(find_many_calls) == 0 + + litellm.max_end_user_budget_id = None + + +def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_list( + reset_budget_job, mock_prisma_client +): + """ + When litellm.max_end_user_budget_id IS configured but the corresponding + budget is NOT in the budgets-to-reset list (not yet expired), end users + with budget_id=NULL should NOT be reset. + """ + import litellm + + now = datetime.now(timezone.utc) + litellm.max_end_user_budget_id = "default-budget-not-expired" + + # A different budget that IS expiring (not the default one) + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 50.0, + "budget_duration": "1d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": "other-budget", + "created_at": now - timedelta(days=1), + }, + ) + + mock_prisma_client.data["budget"] = [test_budget] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + # Should NOT have queried for NULL-budget-id end users + find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls + assert len(find_many_calls) == 0 + + litellm.max_end_user_budget_id = None From acb4f8925400d4dc41c7f1de56f18f4db2977cbe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 17 Apr 2026 21:18:35 -0700 Subject: [PATCH 20/27] [Fix] UI - Models: clarify Anthropic api_base label and tooltip to prevent recursive-loop misconfiguration --- litellm/proxy/public_endpoints/provider_create_fields.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 98ec977b072..163c9648de7 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -216,9 +216,9 @@ "credential_fields": [ { "key": "api_base", - "label": "API Base", - "placeholder": null, - "tooltip": "Optional. Override the upstream Anthropic API URL. Defaults to https://api.anthropic.com.", + "label": "Upstream API Base", + "placeholder": "https://api.anthropic.com", + "tooltip": "Optional. Where the proxy forwards requests upstream. Leave blank to use Anthropic's public API. Set this only for private Anthropic deployments or reverse proxies. Do NOT set this to your LiteLLM proxy URL — that causes a recursive loop.", "required": false, "field_type": "text", "options": null, From 6fe79035d9cddca10f17e5e6bb1a8c4274b126f1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 17 Apr 2026 21:21:43 -0700 Subject: [PATCH 21/27] [Fix] UI - Settings: clarify Max-subscription vs BYOK toggle descriptions (independent of each other) --- .../ui_crud_endpoints/proxy_setting_endpoints.py | 16 ++++++++++++---- .../AdminSettings/UISettings/UISettings.tsx | 4 ++-- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index dd2b5aea7c8..1523e1cb9be 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -97,15 +97,23 @@ class UISettings(BaseModel): forward_client_headers_to_llm_api: bool = Field( default=False, - description="If enabled, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription.", + description=( + "Forwards client headers (Authorization, anthropic-beta, and x-* " + "custom headers) to the upstream LLM. Enable for Claude Code with a " + "Max subscription (forwards the OAuth token) or to pass custom/tracing " + "headers through to the provider. Independent of the BYOK toggle — " + "enable only the one(s) you need." + ), ) forward_llm_provider_auth_headers: bool = Field( default=False, description=( - "If enabled, forwards LLM provider auth headers (x-api-key, " - "x-goog-api-key, api-key, ocp-apim-subscription-key) to the upstream " - "provider. Required for Claude Code BYOK (bring-your-own-key) flows." + "Forwards provider auth headers (x-api-key, x-goog-api-key, api-key, " + "ocp-apim-subscription-key) to the upstream LLM, overriding any " + "deployment-configured key for that request. Enable for Claude Code " + "BYOK (clients bring their own API key). Independent of the " + "client-headers toggle — enable only the one(s) you need." ), ) diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 9751ad435d5..c1042e3cd35 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -295,7 +295,7 @@ export default function UISettings() { Forward client headers to LLM API {forwardClientHeadersProperty?.description ?? - "If enabled, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription."} + "Forwards client headers (Authorization, anthropic-beta, and x-* custom headers) to the upstream LLM. Enable for Claude Code with a Max subscription (forwards the OAuth token) or to pass custom/tracing headers through to the provider. Independent of the BYOK toggle — enable only the one(s) you need."} @@ -315,7 +315,7 @@ export default function UISettings() { Forward LLM provider auth headers {forwardLLMProviderAuthHeadersProperty?.description ?? - "If enabled, forwards provider auth headers (x-api-key, x-goog-api-key, api-key) to the LLM provider. Required for Claude Code BYOK."} + "Forwards provider auth headers (x-api-key, x-goog-api-key, api-key, ocp-apim-subscription-key) to the upstream LLM, overriding any deployment-configured key for that request. Enable for Claude Code BYOK (clients bring their own API key). Independent of the client-headers toggle — enable only the one(s) you need."} From c28877757cf9aa9b83034b439ed0caeeb8a1fc7f Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 17 Apr 2026 22:48:39 -0700 Subject: [PATCH 22/27] fix(ui): extra_headers not persisting on MCP server edit Set extra_headers explicitly in initialValues instead of relying on a useEffect setFieldValue call that races with Antd form initialization. Also avoid sending empty array on submit so the backend's exclude_none doesn't overwrite stored values. --- .../src/components/mcp_tools/MCPPermissionManagement.tsx | 4 ---- .../src/components/mcp_tools/mcp_server_edit.tsx | 3 ++- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.tsx b/ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.tsx index 89aa756cb27..dcb7298c830 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/MCPPermissionManagement.tsx @@ -27,10 +27,6 @@ const MCPPermissionManagement: React.FC = ({ // Set initial values when mcpServer changes useEffect(() => { if (mcpServer) { - // Set extra_headers if they exist - if (mcpServer.extra_headers) { - form.setFieldValue("extra_headers", mcpServer.extra_headers); - } if (mcpServer.static_headers) { const staticHeaders = Object.entries(mcpServer.static_headers).map(([header, value]) => ({ header, diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index 574e7871759..e97222acbd0 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -189,6 +189,7 @@ const MCPServerEdit: React.FC = ({ ...mcpServer, transport: effectiveTransport, static_headers: initialStaticHeaders, + extra_headers: mcpServer.extra_headers || [], oauth_flow_type: mcpServer.token_url ? OAUTH_FLOW.M2M : OAUTH_FLOW.INTERACTIVE, token_validation_json: mcpServer.token_validation ? JSON.stringify(mcpServer.token_validation, null, 2) @@ -563,7 +564,7 @@ const MCPServerEdit: React.FC = ({ mcp_access_groups: accessGroups, alias: restValues.alias, // Include permission management fields - extra_headers: restValues.extra_headers || [], + extra_headers: restValues.extra_headers?.length ? restValues.extra_headers : undefined, allowed_tools: allowedTools.length > 0 ? allowedTools : null, tool_name_to_display_name: Object.keys(toolNameToDisplayName).length > 0 ? toolNameToDisplayName : null, tool_name_to_description: Object.keys(toolNameToDescription).length > 0 ? toolNameToDescription : null, From e2e5b63b413ec4116c90649cd80aef0f0f454c6a Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Fri, 17 Apr 2026 23:24:46 -0700 Subject: [PATCH 23/27] fix(ui): revert submit guard to allow intentional extra_headers clearing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ?.length guard prevented intentional clears from persisting — send the form value as-is since initialValues now populates it correctly. --- .../src/components/mcp_tools/mcp_server_edit.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index e97222acbd0..c64214e080e 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -564,7 +564,7 @@ const MCPServerEdit: React.FC = ({ mcp_access_groups: accessGroups, alias: restValues.alias, // Include permission management fields - extra_headers: restValues.extra_headers?.length ? restValues.extra_headers : undefined, + extra_headers: restValues.extra_headers || [], allowed_tools: allowedTools.length > 0 ? allowedTools : null, tool_name_to_display_name: Object.keys(toolNameToDisplayName).length > 0 ? toolNameToDisplayName : null, tool_name_to_description: Object.keys(toolNameToDescription).length > 0 ? toolNameToDescription : null, From 6bf90fb10a332629c45898a81413f175261e95c0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 18 Apr 2026 10:37:30 -0700 Subject: [PATCH 24/27] fix(mypy): cast msg and pending_delta in _promote_message_stop_usage to resolve union-attr and assignment errors --- .../anthropic_claude3_transformation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 542c43b6594..67b4ad1b23a 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -577,7 +577,7 @@ class AmazonAnthropicClaudeMessagesConfig( chunk_type = chunk.get("type") if chunk_type == "message_start": - msg = chunk.get("message") or {} + msg: Dict[str, Any] = cast(Dict[str, Any], chunk.get("message") or {}) u = msg.get("usage") if isinstance(u, dict): start_usage_snapshot = dict(u) @@ -588,7 +588,7 @@ class AmazonAnthropicClaudeMessagesConfig( continue if chunk_type == "message_delta": - pending_delta = chunk + pending_delta = cast(Dict[str, Any], chunk) continue if chunk_type == "message_stop" and pending_delta is not None: From 2c41f3c291208c4f78f6c56a2f4b4587f335bd1f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 18 Apr 2026 12:17:45 -0700 Subject: [PATCH 25/27] [Fix] UI - Keys: strip empty premium fields from key update payload The /key/update response echoes top-level defaults like policies:[] into client state. On a subsequent edit, the form resends policies:[], which the backend treats as "user is setting policies" and blocks with a 403 enterprise check regardless of value. Drop premium metadata fields from the update payload when the current form value and the previously persisted value are both empty. Genuine clears (non-empty -> empty) still pass through so premium users can clear policies as intended. --- .../templates/key_info_view.test.tsx | 78 +++++++++++++++++++ .../components/templates/key_info_view.tsx | 26 +++++++ 2 files changed, 104 insertions(+) diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index f269ad96a27..6b5732436a4 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -7,8 +7,20 @@ import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { useResetKeySpend } from "@/app/(dashboard)/hooks/keys/useResetKeySpend"; import { KeyResponse, Team } from "../key_team_helpers/key_list"; +import { keyUpdateCall } from "../networking"; import KeyInfoView from "./key_info_view"; +const editViewMocks = vi.hoisted(() => ({ + onSubmit: undefined as ((v: Record) => Promise) | undefined, +})); + +vi.mock("./key_edit_view", () => ({ + KeyEditView: ({ onSubmit }: { onSubmit: (v: Record) => Promise }) => { + editViewMocks.onSubmit = onSubmit; + return
; + }, +})); + vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: vi.fn(), })); @@ -680,4 +692,70 @@ describe("KeyInfoView", () => { }); }); }); + + describe("premium metadata payload normalization", () => { + const enterEditMode = async (keyData: KeyResponse) => { + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "proxy-admin-user", + userRole: "proxy_admin", + }); + renderWithProviders( + {}} + keyId="test-key-id" + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + await userEvent.click(screen.getByRole("tab", { name: /settings/i })); + await userEvent.click(screen.getByRole("button", { name: /edit settings/i })); + await waitFor(() => expect(editViewMocks.onSubmit).toBeDefined()); + }; + + beforeEach(() => { + editViewMocks.onSubmit = undefined; + vi.mocked(keyUpdateCall).mockClear(); + vi.mocked(keyUpdateCall).mockResolvedValue({}); + }); + + it("should drop an empty policies field when the key previously had no policies", async () => { + // Reproduces the real bug: after a successful /key/update, the response echoes + // top-level `policies: []` into client state. Without stripping, the next save + // resends `[]` and trips the premium gate in prepare_metadata_fields. + const keyData: KeyResponse = { + ...MOCK_KEY_DATA, + user_id: "proxy-admin-user", + metadata: {}, + policies: [], + } as KeyResponse; + + await enterEditMode(keyData); + await editViewMocks.onSubmit!({ key: keyData.token, token: keyData.token, policies: [] }); + + expect(keyUpdateCall).toHaveBeenCalledWith( + expect.anything(), + expect.not.objectContaining({ policies: expect.anything() }), + ); + }); + + it("should keep an empty policies field when the key previously had policies set", async () => { + // Premium users must still be able to clear existing policies by sending `[]`. + const keyData: KeyResponse = { + ...MOCK_KEY_DATA, + user_id: "proxy-admin-user", + metadata: { policies: ["existing-policy"] }, + policies: ["existing-policy"], + } as KeyResponse; + + await enterEditMode(keyData); + await editViewMocks.onSubmit!({ key: keyData.token, token: keyData.token, policies: [] }); + + expect(keyUpdateCall).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ policies: [] }), + ); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 5b5e7722c09..bdb8254c726 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -34,6 +34,21 @@ interface KeyInfoViewProps { backButtonText?: string; } +// Must stay in sync with LiteLLM_ManagementEndpoint_MetadataFields_Premium +// in litellm/proxy/_types.py — limited to fields the key-edit form submits. +const PREMIUM_METADATA_FIELDS = [ + "policies", + "guardrails", + "prompts", + "tags", + "allowed_passthrough_routes", +] as const; + +const isEmptyValue = (v: unknown): boolean => + v == null || + (Array.isArray(v) && v.length === 0) || + (typeof v === "string" && v.trim() === ""); + /** * ───────────────────────────────────────────────────────────────────────── * @deprecated @@ -146,6 +161,17 @@ export default function KeyInfoView({ delete formValues.prompts; } + // Drop premium metadata fields that are empty AND were empty before. + // The /key/update response echoes defaults like `policies: []` back into + // state; without this, the next save resends `[]` and trips the premium + // gate in prepare_metadata_fields for non-premium users. + for (const field of PREMIUM_METADATA_FIELDS) { + const previousValue = (currentKeyData.metadata as Record | undefined)?.[field]; + if (isEmptyValue(formValues[field]) && isEmptyValue(previousValue)) { + delete formValues[field]; + } + } + // Handle max budget empty string formValues.max_budget = mapEmptyStringToNull(formValues.max_budget); From cae8b74b0b13dd714209fb296ca5258d1244512b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 18 Apr 2026 12:23:58 -0700 Subject: [PATCH 26/27] Fall back to top-level keyData when resolving previous premium value Premium fields like policies are echoed at the top level of the /key/update response, not necessarily mirrored into metadata. Read metadata first then fall back to the top-level property so an intentional clear is preserved in either shape. --- .../templates/key_info_view.test.tsx | 19 +++++++++++++++++++ .../components/templates/key_info_view.tsx | 4 +++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 6b5732436a4..724c961bbdc 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -757,5 +757,24 @@ describe("KeyInfoView", () => { expect.objectContaining({ policies: [] }), ); }); + + it("should keep an empty policies field when the previous value lives only at the top level of keyData", async () => { + // Defensive: some premium fields may be present at the top level but not + // mirrored into metadata. A genuine clear must still be forwarded. + const keyData: KeyResponse = { + ...MOCK_KEY_DATA, + user_id: "proxy-admin-user", + metadata: {}, + policies: ["existing-policy"], + } as KeyResponse; + + await enterEditMode(keyData); + await editViewMocks.onSubmit!({ key: keyData.token, token: keyData.token, policies: [] }); + + expect(keyUpdateCall).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ policies: [] }), + ); + }); }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index bdb8254c726..492e43cbc81 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -166,7 +166,9 @@ export default function KeyInfoView({ // state; without this, the next save resends `[]` and trips the premium // gate in prepare_metadata_fields for non-premium users. for (const field of PREMIUM_METADATA_FIELDS) { - const previousValue = (currentKeyData.metadata as Record | undefined)?.[field]; + const previousValue = + (currentKeyData.metadata as Record | undefined)?.[field] ?? + (currentKeyData as unknown as Record)[field]; if (isEmptyValue(formValues[field]) && isEmptyValue(previousValue)) { delete formValues[field]; } From cfdf893226e26264bdd99d2605c94ea460445979 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 18 Apr 2026 13:04:32 -0700 Subject: [PATCH 27/27] [Fix] Merge prometheus_helpers.py into prometheus_helpers/__init__.py to resolve file/package collision A previous refactor added `litellm/integrations/prometheus_helpers.py` as a sibling to the existing `litellm/integrations/prometheus_helpers/` directory (which contains `prometheus_api.py` and has no `__init__.py`). The file shadowed the namespace-package directory, so any deferred `from litellm.integrations.prometheus_helpers.prometheus_api import ...` raised `ModuleNotFoundError: 'litellm.integrations.prometheus_helpers' is not a package` at request time. Two runtime call sites hit that path: - /global/spend/logs (spend_management_endpoints.py) returned plain-text 500 "Internal Server Error" for every call, breaking the Admin UI Usage tab and programmatic consumers. - SlackAlerting.send_fallback_stats_from_prometheus silently failed inside its own try/except. Fix: move prometheus_helpers.py content into prometheus_helpers/__init__.py and delete the stray .py. The directory becomes a regular package, so both the package-root import (from ...prometheus_helpers import X) and the submodule import (from ...prometheus_helpers.prometheus_api import X) resolve correctly. No call sites change. --- .../{prometheus_helpers.py => prometheus_helpers/__init__.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename litellm/integrations/{prometheus_helpers.py => prometheus_helpers/__init__.py} (100%) diff --git a/litellm/integrations/prometheus_helpers.py b/litellm/integrations/prometheus_helpers/__init__.py similarity index 100% rename from litellm/integrations/prometheus_helpers.py rename to litellm/integrations/prometheus_helpers/__init__.py