= ({
paramsByModel,
onEffortChange,
}) => {
- const rows = models
- .map((model) => {
- const effort = storedEffort(paramsByModel?.[model]);
- const supported = effortOptionsByModel[model] ?? [];
- // A stored effort outside the supported set (hand-authored, or capabilities changed since it
- // was saved) stays listed so it renders and can be cleared.
- const options = effort !== undefined && !supported.includes(effort) ? [...supported, effort] : supported;
- return { model, effort, options };
- })
- .filter(({ model, options }) => options.length > 0 || Object.keys(paramsByModel?.[model] ?? {}).length > 0);
+ const rows = tierEffortRows({ models, effortOptionsByModel, paramsByModel });
if (rows.length === 0) return null;
return (
diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx
index bd691c7f629..998924ce36d 100644
--- a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx
+++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.tsx
@@ -1,6 +1,6 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
-import { modelAvailableCall } from "@/components/networking";
-import { fetchAvailableModelsForTeam } from "./fetch_models";
+import { modelAvailableCall, modelHubCall } from "@/components/networking";
+import { fetchAvailableModels, fetchAvailableModelsForTeam } from "./fetch_models";
vi.mock("@/components/networking", () => ({
modelAvailableCall: vi.fn(),
@@ -8,6 +8,7 @@ vi.mock("@/components/networking", () => ({
}));
const modelAvailableCallMock = vi.mocked(modelAvailableCall);
+const modelHubCallMock = vi.mocked(modelHubCall);
describe("fetchAvailableModelsForTeam", () => {
beforeEach(() => {
@@ -31,3 +32,33 @@ describe("fetchAvailableModelsForTeam", () => {
expect(await fetchAvailableModelsForTeam("token", "team-123")).toEqual([]);
});
});
+
+describe("fetchAvailableModels", () => {
+ beforeEach(() => {
+ vi.clearAllMocks();
+ });
+
+ it("carries the reasoning capabilities the model hub reports for each group", async () => {
+ modelHubCallMock.mockResolvedValue({
+ data: [
+ { model_group: "smart", mode: "chat", supports_reasoning: true, supported_reasoning_efforts: ["low", "high"] },
+ { model_group: "plain", mode: "chat", supports_reasoning: false },
+ ],
+ });
+
+ expect(await fetchAvailableModels("token")).toEqual([
+ { model_group: "plain", mode: "chat" },
+ { model_group: "smart", mode: "chat", supports_reasoning: true, supported_reasoning_efforts: ["low", "high"] },
+ ]);
+ });
+
+ it.each([
+ ["an error payload in place of the list", { data: { error: "no access" } }],
+ ["a missing data key", {}],
+ ["no body at all", undefined],
+ ])("returns an empty list on %s rather than throwing", async (_label, response) => {
+ modelHubCallMock.mockResolvedValue(response);
+
+ expect(await fetchAvailableModels("token")).toEqual([]);
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx
index 96b87887d43..1f812cf0377 100644
--- a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx
+++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx
@@ -44,7 +44,8 @@ export const fetchAvailableModelsForTeam = async (accessToken: string, teamId: s
export const fetchAvailableModels = async (accessToken: string): Promise => {
try {
const fetchedModels = await modelHubCall(accessToken);
- const models: ModelGroup[] = (fetchedModels?.data ?? [])
+ const fetchedData: unknown = fetchedModels?.data;
+ const models: ModelGroup[] = (Array.isArray(fetchedData) ? fetchedData : [])
.map(toModelGroup)
.filter((model: ModelGroup) => model.model_group !== "")
.sort((a: ModelGroup, b: ModelGroup) => a.model_group.localeCompare(b.model_group));
From b71b574af410f436954f9e6fcb28b44eff6c1d34 Mon Sep 17 00:00:00 2001
From: Tin Chi Lo
Date: Mon, 24 Aug 2026 15:12:33 -0400
Subject: [PATCH 20/70] ci: raise three unit shard job timeouts to satisfy the
startup safety gate
---
.github/workflows/test-unit.yml | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index a7c67f2b35d..2dfca3d308f 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -211,7 +211,7 @@ jobs:
workers: 2
reruns: 2
timeout-minutes: 20
- job-timeout-minutes: 55
+ job-timeout-minutes: 60
- shard: proxy-extras
artifact-name: proxy-extras
@@ -219,7 +219,7 @@ jobs:
workers: 2
reruns: 2
timeout-minutes: 20
- job-timeout-minutes: 55
+ job-timeout-minutes: 60
- shard: enterprise-package
artifact-name: enterprise-package
@@ -227,7 +227,7 @@ jobs:
workers: 4
reruns: 2
timeout-minutes: 20
- job-timeout-minutes: 55
+ job-timeout-minutes: 60
- shard: responses-caching-types
artifact-name: responses-caching-types
From 31a67561ab23794bd9bdb72489b05f06e10b1c03 Mon Sep 17 00:00:00 2001
From: tin-berri
Date: Tue, 25 Aug 2026 01:56:00 -0400
Subject: [PATCH 21/70] feat(complexity_router): bound the classifier context
block, not each turn in it (#38145)
The LLM classifier capped every prior turn at 200 characters independently, so a
785 character turn was cut even when the whole block it belonged to was 353
characters. A character budget now bounds the block: turns are taken newest first
and quoted whole while they fit, older turns are dropped whole once it runs out,
and only the turn straddling the boundary is cut. The per-turn cap stays as an
optional clamp for operators who set it deliberately, defaulting to unset.
---
.../complexity_router/complexity_router.py | 61 +++++--
.../complexity_router/config.py | 34 +++-
.../router_strategy/test_complexity_router.py | 169 +++++++++++++++++-
.../add_model/ClassificationMethodConfig.tsx | 33 ++--
.../add_model/ComplexityRouterConfig.test.tsx | 65 ++++---
.../add_model/ComplexityRouterConfig.tsx | 4 +-
.../add_model/add_auto_router_tab.tsx | 2 +-
.../build_complexity_router_config.test.ts | 18 +-
.../build_complexity_router_config.ts | 9 +-
...d_updated_complexity_router_config.test.ts | 23 ++-
.../edit_auto_router_modal.test.tsx | 2 +-
.../edit_auto_router_modal.tsx | 12 +-
.../src/lib/autorouter_presets.ts | 1 +
ui/litellm-dashboard/src/lib/http/schema.d.ts | 13 +-
14 files changed, 356 insertions(+), 90 deletions(-)
diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py
index 4c07679c5ca..087c1f7278d 100644
--- a/litellm/router_strategy/complexity_router/complexity_router.py
+++ b/litellm/router_strategy/complexity_router/complexity_router.py
@@ -19,7 +19,7 @@ import asyncio
import random
import re
from collections.abc import Iterator, Mapping, Sequence
-from itertools import accumulate, islice
+from itertools import accumulate, islice, takewhile
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
@@ -275,6 +275,7 @@ _DEFAULT_REMINDER_MARKERS: Final = ((_REMINDER_OPEN, _REMINDER_CLOSE),)
_TRUNCATION_MARKER: Final = "..."
_TRUNCATION_HEAD_FRACTION: Final = 0.3
+_MIN_QUOTED_TURN_CHARS: Final = 120
_CJK_CHARACTER: Final = re.compile("[-ヿㇰ-ㇿ㐀-䶿一-鿿豈-ｦ-ﾝ\U00020000-\U0003ffff]")
@@ -593,11 +594,40 @@ def _iter_context_turns_newest_first(
)
+def _turns_within_budget(
+ turns: Sequence[tuple[str, str]],
+ budget_chars: int,
+) -> tuple[tuple[str, str], ...]:
+ """The newest-first turns that fit budget_chars, quoted whole wherever they fit.
+
+ Bounding the block rather than every turn in it is what lets an ordinary conversation reach the
+ classifier intact: a per-turn cap cuts a 785 character turn even when the whole block would have
+ been 353 characters, which is three orders of magnitude below anything the classifier call is
+ near. Once the budget does run out the older turns are dropped entire rather than shortened, so
+ at most one turn is ever cut and the rest read as themselves. A remainder too small to carry a
+ sentence buys less signal than the ellipses it would arrive wrapped in, so that turn is dropped.
+
+ The boundary turn is cut to leave room for the marker rather than to the remainder itself, so the
+ quoted block never exceeds budget_chars; the marker is part of what the budget buys, not an extra
+ charged on top of it.
+ """
+ spent: Final = accumulate(len(text) for _, text in turns)
+ fitting: Final = tuple(takewhile(lambda pair: pair[1] <= budget_chars, zip(turns, spent)))
+ remaining: Final = budget_chars - (fitting[-1][1] if fitting else 0)
+ whole: Final = tuple(turn for turn, _ in fitting)
+ cut_to: Final = remaining - len(_TRUNCATION_MARKER)
+ if len(whole) == len(turns) or cut_to < _MIN_QUOTED_TURN_CHARS:
+ return whole
+ boundary_role, boundary_text = turns[len(whole)]
+ return (*whole, (boundary_role, _truncate(boundary_text, cut_to)))
+
+
def _extract_prior_turns(
messages: Sequence[Mapping[str, object]],
current_ask: str | None,
window_size: int,
- per_turn_chars: int,
+ budget_chars: int,
+ per_turn_chars: int | None,
include_assistant: bool,
marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS,
) -> tuple[tuple[str, str], ...]:
@@ -612,19 +642,29 @@ def _extract_prior_turns(
window_size counts turns of every eligible role, so with assistant turns included it is the last N
of the conversation rather than the last N asks. A turn carrying only tool calls or thinking
blocks flattens to empty text and is skipped, so it never spends a slot.
+
+ Three bounds apply and the tightest wins: window_size caps how many turns, budget_chars caps the
+ block they form, and per_turn_chars optionally caps any single one of them before the block is
+ measured. They are separate because they answer separate questions, and only the block bound
+ tracks what the classifier call actually costs.
"""
if window_size <= 0 or not messages:
return ()
- prior: Final = islice(
- (
- turn
- for turn in _iter_context_turns_newest_first(messages, include_assistant, marker_pairs)
- if turn[1] != current_ask
- ),
- window_size,
+ prior: Final = tuple(
+ islice(
+ (
+ turn
+ for turn in _iter_context_turns_newest_first(messages, include_assistant, marker_pairs)
+ if turn[1] != current_ask
+ ),
+ window_size,
+ )
)
- return tuple((role, _truncate(text, per_turn_chars)) for role, text in reversed(tuple(prior)))
+ clamped: Final = (
+ prior if per_turn_chars is None else tuple((role, _truncate(text, per_turn_chars)) for role, text in prior)
+ )
+ return tuple(reversed(_turns_within_budget(clamped, budget_chars)))
def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bool:
@@ -1363,6 +1403,7 @@ class ComplexityRouter(CustomLogger):
messages,
current_ask=prompt,
window_size=self.config.classifier_context_window_size,
+ budget_chars=self.config.classifier_context_budget_chars,
per_turn_chars=self.config.classifier_context_per_turn_chars,
include_assistant=include_assistant,
marker_pairs=self._reminder_markers,
diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py
index d3c4bd7938b..9907407d84d 100644
--- a/litellm/router_strategy/complexity_router/config.py
+++ b/litellm/router_strategy/complexity_router/config.py
@@ -49,7 +49,7 @@ TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = (
DEFAULT_TIER_DISTANCE_PENALTY: Final[float] = 0.5
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE: Final[int] = 3
-DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS: Final[int] = 200
+DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS: Final[int] = 8000
class KeywordTierRule(BaseModel):
@@ -645,12 +645,30 @@ class ComplexityRouterConfig(BaseModel):
"classifier_type is 'llm'."
),
)
- classifier_context_per_turn_chars: int = Field(
- default=DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS,
+ classifier_context_budget_chars: int = Field(
+ default=DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS,
+ ge=0,
+ description=(
+ "Maximum characters of prior-turn text quoted to the LLM classifier, across the whole "
+ "context window, per classification call. Turns are taken newest first and quoted whole "
+ "while they fit, so a conversation small enough to quote entirely is never cut; once the "
+ "budget runs out the older turns are dropped whole and only the turn straddling the "
+ "boundary is truncated, into whatever space is left. The current ask and the caller's "
+ "system prompt sit outside this budget and are always sent in full, as does the numbering "
+ "each quoted turn carries. A budget under 120 leaves no room to quote a turn and "
+ "suppresses the block; set classifier_context_window_size to 0 to turn context off "
+ "deliberately. Only applies when classifier_type is 'llm'."
+ ),
+ )
+ classifier_context_per_turn_chars: int | None = Field(
+ default=None,
gt=0,
description=(
- "Maximum character length for each prior turn's text in the classifier context window. "
- "Turns exceeding this are truncated. Only applies when classifier_type is 'llm'."
+ "Optional cap on each individual prior turn's text, applied before "
+ "classifier_context_budget_chars bounds the block. Unset by default, so one long turn may "
+ "spend the whole budget, which is usually what a follow-up needs; set it when no single "
+ "turn should dominate the context the classifier sees. A capped turn keeps its opening "
+ "and its ending with the middle elided. Only applies when classifier_type is 'llm'."
),
)
classifier_context_include_assistant_turns: bool = Field(
@@ -662,9 +680,9 @@ class ComplexityRouterConfig(BaseModel):
"word 'yes'. When enabled, classifier_context_window_size counts the last N turns of the "
"conversation across both roles rather than the last N user turns, and assistant text is "
"sent to the classifier model, which may be a different deployment or provider than the "
- "routed completion model. Assistant replies share classifier_context_per_turn_chars with "
- "user turns, so raise it if replies are truncated before the part that carries the "
- "difficulty. Off by default because enabling it shifts tier decisions, and therefore "
+ "routed completion model. Assistant replies spend classifier_context_budget_chars "
+ "alongside user turns, so raise it if the oldest turns stop being quoted once replies "
+ "join the window. Off by default because enabling it shifts tier decisions, and therefore "
"spend, for an already-deployed router. Only applies when classifier_type is 'llm'."
),
)
diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py
index d4333fde163..a29b4d03bc5 100644
--- a/tests/test_litellm/router_strategy/test_complexity_router.py
+++ b/tests/test_litellm/router_strategy/test_complexity_router.py
@@ -6052,8 +6052,9 @@ class TestContextAwareClassifier:
[{"role": "user", "content": turn}, {"role": "user", "content": "go ahead"}],
"go ahead",
3,
- 200,
- False,
+ budget_chars=10_000,
+ per_turn_chars=200,
+ include_assistant=False,
)
assert "multi-region gateway" in quoted[0][1]
@@ -6216,7 +6217,167 @@ class TestContextAwareClassifier:
"""
from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns
- assert _extract_prior_turns(messages, current_ask, window, per_turn_chars, include_assistant) == expected
+ assert (
+ _extract_prior_turns(
+ messages,
+ current_ask,
+ window,
+ budget_chars=10_000,
+ per_turn_chars=per_turn_chars,
+ include_assistant=include_assistant,
+ )
+ == expected
+ )
+
+ @pytest.mark.parametrize(
+ "turn_lengths,budget_chars,expected_lengths",
+ [
+ pytest.param((50, 50, 50), 10_000, (50, 50, 50), id="a-block-that-fits-is-quoted-whole"),
+ pytest.param((100, 100, 100), 250, (100, 100), id="oldest-turn-is-dropped-whole"),
+ pytest.param((500, 100), 400, (300, 100), id="only-the-boundary-turn-is-cut"),
+ pytest.param((900,), 300, (300,), id="a-turn-larger-than-the-budget-is-still-quoted"),
+ pytest.param((500, 100), 180, (100,), id="a-remainder-too-small-to-carry-a-sentence-is-dropped"),
+ pytest.param((50,), 0, (), id="a-zero-budget-quotes-nothing"),
+ ],
+ )
+ def test_budget_bounds_the_block_not_each_turn(self, turn_lengths, budget_chars, expected_lengths):
+ """Turns are taken newest first and quoted whole while they fit.
+
+ The defect this replaces capped every turn independently, so a 785 character turn was cut even
+ though the whole block it belonged to was 353 characters. Bounding the block instead means an
+ ordinary conversation arrives intact, and when the budget really does run out the older turns
+ are dropped entire rather than each arriving mangled. At most one turn is ever cut, and a
+ remainder too small to carry a sentence is dropped rather than quoted as two ellipses around a
+ fragment. A single turn bigger than the whole budget is still quoted, cut to the budget, since
+ dropping it would leave the classifier with no context at all.
+ """
+ from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns
+
+ messages = [{"role": "user", "content": f"{i}" * length} for i, length in enumerate(turn_lengths)]
+
+ quoted = _extract_prior_turns(
+ [*messages, {"role": "user", "content": "go ahead"}],
+ "go ahead",
+ len(turn_lengths),
+ budget_chars=budget_chars,
+ per_turn_chars=None,
+ include_assistant=False,
+ )
+
+ assert tuple(len(text) for _, text in quoted) == expected_lengths
+
+ @pytest.mark.parametrize("budget_chars", [130, 200, 351, 400, 999, 8000])
+ @pytest.mark.parametrize("turn_lengths", [(900,), (500, 100), (100, 100, 100), (50, 50, 50)])
+ def test_the_quoted_block_never_exceeds_the_budget(self, turn_lengths, budget_chars):
+ """The budget is a ceiling on what is quoted, marker included.
+
+ Cutting the boundary turn to the remainder and then appending the marker put the block three
+ characters over the number an operator configured, which is the kind of drift that makes a
+ documented ceiling untrue. Asserted across shapes rather than at the one boundary that happened
+ to be wrong, so any future off-by-marker anywhere in the fill is caught here.
+ """
+ from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns
+
+ messages = [{"role": "user", "content": f"{i}" * length} for i, length in enumerate(turn_lengths)]
+
+ quoted = _extract_prior_turns(
+ [*messages, {"role": "user", "content": "go ahead"}],
+ "go ahead",
+ len(turn_lengths),
+ budget_chars=budget_chars,
+ per_turn_chars=None,
+ include_assistant=False,
+ )
+
+ assert sum(len(text) for _, text in quoted) <= budget_chars
+
+ def test_per_turn_cap_still_clamps_when_an_operator_sets_it(self):
+ """An operator who set the per-turn cap keeps exactly what they configured.
+
+ The cap stopped being the default, so it has to keep working for the deployments that named it
+ deliberately; it applies before the block budget rather than instead of it.
+ """
+ from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns
+
+ quoted = _extract_prior_turns(
+ [{"role": "user", "content": "z" * 900}, {"role": "user", "content": "go ahead"}],
+ "go ahead",
+ 3,
+ budget_chars=10_000,
+ per_turn_chars=200,
+ include_assistant=False,
+ )
+
+ assert len(quoted[0][1]) == 203
+
+ @pytest.mark.asyncio
+ async def test_a_long_turn_reaches_the_classifier_whole_by_default(
+ self, mock_router_instance, llm_classifier_config
+ ):
+ """The shipped defaults quote an ordinary long turn without cutting it anywhere.
+
+ This is the whole point of the change, asserted where a deployment actually meets it: no knob
+ set, one turn well past the retired 200 character cap, and no truncation marker in the payload.
+ """
+ from litellm.router_strategy.complexity_router.complexity_router import _TRUNCATION_MARKER
+
+ router = ComplexityRouter(
+ model_name="test-complexity-router",
+ litellm_router_instance=mock_router_instance,
+ complexity_router_config=llm_classifier_config,
+ )
+ mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}'))
+ turn = "The incident ran from 02:10 to 02:40 and only streaming was affected. " * 10 + "Now rewrite it"
+
+ await router.aclassify(
+ "go ahead",
+ messages=[{"role": "user", "content": turn}, {"role": "user", "content": "go ahead"}],
+ )
+
+ user_payload = mock_router_instance.acompletion.call_args.kwargs["messages"][1]["content"]
+ assert turn in user_payload
+ assert _TRUNCATION_MARKER not in user_payload
+
+ @pytest.mark.asyncio
+ async def test_a_turn_dropped_for_budget_still_counts_as_prior_conversation(
+ self, mock_router_instance, llm_classifier_config
+ ):
+ """Dropping turns to fit the budget must not make a long conversation look single-turn.
+
+ The depth line gates on whether prior conversation exists, not on whether any of it was worth
+ quoting, exactly so a continuation is never reported as a context-free first request. A budget
+ tight enough to drop every turn is the newest way to reach that mismatch.
+ """
+ router = ComplexityRouter(
+ model_name="test-complexity-router",
+ litellm_router_instance=mock_router_instance,
+ complexity_router_config={**llm_classifier_config, "classifier_context_budget_chars": 1},
+ )
+ mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}'))
+
+ await router.aclassify(
+ "go ahead",
+ messages=[
+ {"role": "user", "content": "a long earlier request that cannot fit a one character budget"},
+ {"role": "user", "content": "go ahead"},
+ ],
+ )
+
+ user_payload = mock_router_instance.acompletion.call_args.kwargs["messages"][1]["content"]
+ assert "Recent conversation" not in user_payload
+ assert "Conversation so far" in user_payload
+
+ def test_context_defaults_bound_the_block_and_leave_turns_uncapped(self):
+ """The shipped defaults: a block budget, and no per-turn cap unless one is named."""
+ from litellm.router_strategy.complexity_router.config import (
+ DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS,
+ ComplexityRouterConfig,
+ )
+
+ config = ComplexityRouterConfig()
+
+ assert config.classifier_context_budget_chars == DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS
+ assert config.classifier_context_per_turn_chars is None
def test_prior_turn_context_strips_every_configured_pair(self):
"""The classifier's context window is stripped with the same pairs as the ask.
@@ -6235,7 +6396,7 @@ class TestContextAwareClassifier:
{"role": "user", "content": "current ask"},
]
- assert _extract_prior_turns(messages, "current ask", 5, 200, False, pairs) == (
+ assert _extract_prior_turns(messages, "current ask", 5, 10_000, 200, False, pairs) == (
("user", "what about b-trees?"),
("user", "and heaps?"),
)
diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx
index ef6e521de42..c3fc23034f1 100644
--- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx
@@ -16,7 +16,8 @@ import {
ClassifierFallback,
ClassifierType,
ComplexityRouterConfigValue,
- DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS,
+ DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS,
+ MIN_QUOTED_CONTEXT_TURN_CHARS,
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
DEFAULT_CLASSIFIER_FALLBACK,
DEFAULT_CLASSIFIER_TIMEOUT_MS,
@@ -149,6 +150,8 @@ const ClassificationMethodConfig: React.FC = ({
const classifierModelMissing =
showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model;
const usesCustomPrompt = Boolean(value.classifier_llm_config?.system_prompt?.trim());
+ const contextBudget = value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS;
+ const contextBudgetQuotesNothing = contextBudget > 0 && contextBudget < MIN_QUOTED_CONTEXT_TURN_CHARS;
const classificationRubric = value.classifier_llm_config?.classification_rubric ?? DEFAULT_CLASSIFICATION_RUBRIC;
const handleClassifierTypeChange = (classifierType: ClassifierType) => {
@@ -167,9 +170,9 @@ const ClassificationMethodConfig: React.FC = ({
classifierType === "llm"
? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE
: undefined,
- classifier_context_per_turn_chars:
+ classifier_context_budget_chars:
classifierType === "llm"
- ? value.classifier_context_per_turn_chars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS
+ ? value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS
: undefined,
classifier_context_include_assistant_turns:
classifierType === "llm" ? value.classifier_context_include_assistant_turns : undefined,
@@ -235,10 +238,10 @@ const ClassificationMethodConfig: React.FC = ({
});
};
- const handleClassifierContextPerTurnCharsChange = (perTurnChars: number | null) => {
+ const handleClassifierContextBudgetCharsChange = (budgetChars: number | null) => {
onChange({
...value,
- classifier_context_per_turn_chars: perTurnChars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS,
+ classifier_context_budget_chars: budgetChars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS,
});
};
@@ -411,17 +414,27 @@ const ClassificationMethodConfig: React.FC = ({
- Context Per-Turn Character Limit
+ Context Character Budget
- handleClassifierContextPerTurnCharsChange(event.target.value === "" ? null : event.target.valueAsNumber)
+ handleClassifierContextBudgetCharsChange(event.target.value === "" ? null : event.target.valueAsNumber)
}
- min={1}
+ min={0}
className="w-full"
/>
- Prior turns longer than this are truncated.
+
+ Total characters of prior conversation sent to the classifier. Turns are taken newest first and quoted
+ whole while they fit, so a short conversation is never cut.
+
+ {contextBudgetQuotesNothing && (
+
+ Under {MIN_QUOTED_CONTEXT_TURN_CHARS} characters there is no room to quote a turn that does not already
+ fit, so a long conversation reaches the classifier with no context at all. Set Context Window Size to 0
+ to turn context off deliberately.
+
+ )}
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
index 0a848ed7deb..303ef3cdddc 100644
--- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
@@ -108,7 +108,7 @@ describe("ComplexityRouterConfig", () => {
classifier_type: "llm",
classifier_llm_config: { model: "", timeout_ms: 3000, classification_rubric: "agentic" },
classifier_context_window_size: 3,
- classifier_context_per_turn_chars: 200,
+ classifier_context_budget_chars: 8000,
};
expect(onChange).toHaveBeenCalledWith(expectedValue);
});
@@ -130,11 +130,10 @@ describe("ComplexityRouterConfig", () => {
expect(screen.getByDisplayValue("750")).toBeInTheDocument();
expect(screen.getByText("Context Window Size")).toBeInTheDocument();
expect(screen.getByDisplayValue("5")).toBeInTheDocument();
- expect(screen.getByText("Context Per-Turn Character Limit")).toBeInTheDocument();
- expect(screen.getByDisplayValue("400")).toBeInTheDocument();
+ expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument();
});
- it("should default classifier context fields to 3 and 200 when llm is selected without explicit values", () => {
+ it("should default the context window and budget when llm is selected", () => {
const llmValue: ComplexityRouterConfigValue = {
...defaultValue,
classifier_type: "llm",
@@ -147,8 +146,42 @@ describe("ComplexityRouterConfig", () => {
const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement;
expect(within(windowSizeSection).getByDisplayValue("3")).toBeInTheDocument();
- const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement;
- expect(within(perTurnCharsSection).getByDisplayValue("200")).toBeInTheDocument();
+ const budgetSection = screen.getByText("Context Character Budget").closest("div") as HTMLElement;
+ expect(within(budgetSection).getByDisplayValue("8000")).toBeInTheDocument();
+ });
+
+ it("should warn when the budget is too small to quote any turn that does not already fit", () => {
+ const llmValue: ComplexityRouterConfigValue = {
+ ...defaultValue,
+ classifier_type: "llm",
+ classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 },
+ classifier_context_budget_chars: 50,
+ };
+ renderWithProviders(
);
+
+ fireEvent.click(screen.getByText("Advanced: Classification Method"));
+
+ expect(screen.getByText(/no room to quote a turn/i)).toBeInTheDocument();
+ });
+
+ it("should not warn on a budget large enough to quote a turn, nor on a deliberate zero", () => {
+ for (const budget of [120, 8000, 0]) {
+ const { unmount } = renderWithProviders(
+
,
+ );
+ fireEvent.click(screen.getByText("Advanced: Classification Method"));
+ expect(screen.queryByText(/no room to quote a turn/i)).not.toBeInTheDocument();
+ unmount();
+ }
});
it("should show the assistant-turns switch with its configured value when classifier_type is llm", () => {
@@ -229,26 +262,6 @@ describe("ComplexityRouterConfig", () => {
});
});
- it("should call onChange with the updated classifier_context_per_turn_chars when edited", () => {
- const onChange = vi.fn();
- const llmValue: ComplexityRouterConfigValue = {
- ...defaultValue,
- classifier_type: "llm",
- classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 },
- };
- renderWithProviders(
);
- fireEvent.click(screen.getByText("Advanced: Classification Method"));
-
- const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement;
- const input = within(perTurnCharsSection).getByRole("spinbutton");
- fireEvent.change(input, { target: { value: "500" } });
-
- expect(onChange).toHaveBeenCalledWith({
- ...llmValue,
- classifier_context_per_turn_chars: 500,
- });
- });
-
it("should render the custom technical keywords field", () => {
renderWithProviders(
);
fireEvent.click(screen.getByText("Advanced: Classification Method"));
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
index fc731e2c77f..c98dc20d3bc 100644
--- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
@@ -31,7 +31,8 @@ export type { DimensionWeights, TierBoundaries, TokenThresholds };
export const DEFAULT_CLASSIFIER_TIMEOUT_MS = 3000;
export const DEFAULT_TIER_DISTANCE_PENALTY = 0.5;
export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3;
-export const DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS = 200;
+export const DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS = 8000;
+export const MIN_QUOTED_CONTEXT_TURN_CHARS = 120;
export const DEFAULT_SESSION_AFFINITY = false;
export const DEFAULT_DEPLOYMENT_AFFINITY = true;
@@ -137,6 +138,7 @@ export interface ComplexityRouterConfigValue {
classifier_type: ClassifierType;
classifier_llm_config?: ClassifierLLMConfig;
classifier_context_window_size?: number;
+ classifier_context_budget_chars?: number;
classifier_context_per_turn_chars?: number;
classifier_context_include_assistant_turns?: boolean;
classifier_fallback?: ClassifierFallback;
diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
index 9be5391040a..98ee2b7ae7c 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
@@ -352,7 +352,7 @@ const AddAutoRouterTab: React.FC
= ({
classifierType: complexityRouterConfig.classifier_type,
classifierLlmConfig: complexityRouterConfig.classifier_llm_config,
classifierContextWindowSize: complexityRouterConfig.classifier_context_window_size,
- classifierContextPerTurnChars: complexityRouterConfig.classifier_context_per_turn_chars,
+ classifierContextBudgetChars: complexityRouterConfig.classifier_context_budget_chars,
classifierContextIncludeAssistantTurns: complexityRouterConfig.classifier_context_include_assistant_turns,
classifierFallback: complexityRouterConfig.classifier_fallback,
sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY,
diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts
index cb55362c6a7..33f0fb8a539 100644
--- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts
+++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts
@@ -23,7 +23,7 @@ const baseParams: BuildComplexityRouterConfigParams = {
classifierType: "heuristic",
classifierLlmConfig: undefined,
classifierContextWindowSize: undefined,
- classifierContextPerTurnChars: undefined,
+ classifierContextBudgetChars: undefined,
classifierContextIncludeAssistantTurns: undefined,
classifierFallback: undefined,
sessionAffinity: false,
@@ -93,39 +93,39 @@ describe("buildComplexityRouterConfig", () => {
expect(config.classifier_llm_config).toBeUndefined();
});
- it("includes classifier_context_window_size and classifier_context_per_turn_chars only when classifier_type is llm", () => {
+ it("includes classifier_context_window_size and classifier_context_budget_chars only when classifier_type is llm", () => {
const params: BuildComplexityRouterConfigParams = {
...baseParams,
classifierType: "llm",
classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 },
classifierContextWindowSize: 5,
- classifierContextPerTurnChars: 300,
+ classifierContextBudgetChars: 4000,
};
const config = buildComplexityRouterConfig(params);
expect(config.classifier_context_window_size).toBe(5);
- expect(config.classifier_context_per_turn_chars).toBe(300);
+ expect(config.classifier_context_budget_chars).toBe(4000);
});
- it("omits classifier_context_window_size and classifier_context_per_turn_chars when classifier_type is heuristic even if values linger in state", () => {
+ it("omits classifier_context_window_size and classifier_context_budget_chars when classifier_type is heuristic even if values linger in state", () => {
const params: BuildComplexityRouterConfigParams = {
...baseParams,
classifierType: "heuristic",
classifierContextWindowSize: 5,
- classifierContextPerTurnChars: 300,
+ classifierContextBudgetChars: 4000,
};
const config = buildComplexityRouterConfig(params);
expect(config.classifier_context_window_size).toBeUndefined();
- expect(config.classifier_context_per_turn_chars).toBeUndefined();
+ expect(config.classifier_context_budget_chars).toBeUndefined();
});
- it("omits classifier_context_window_size and classifier_context_per_turn_chars when classifier_type is llm but neither was set, leaving the backend default", () => {
+ it("omits classifier_context_window_size and classifier_context_budget_chars when classifier_type is llm but neither was set, leaving the backend default", () => {
const config = buildComplexityRouterConfig({
...baseParams,
classifierType: "llm",
classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 },
});
expect(config.classifier_context_window_size).toBeUndefined();
- expect(config.classifier_context_per_turn_chars).toBeUndefined();
+ expect(config.classifier_context_budget_chars).toBeUndefined();
});
it("allows classifier_context_window_size of 0, distinct from unset, to send no prior-turn context", () => {
diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts
index 677cbe7063f..bd95bea226a 100644
--- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts
+++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts
@@ -80,7 +80,7 @@ export interface BuildComplexityRouterConfigParams {
classifierType: ClassifierType;
classifierLlmConfig: ClassifierLLMConfig | undefined;
classifierContextWindowSize: number | undefined;
- classifierContextPerTurnChars: number | undefined;
+ classifierContextBudgetChars: number | undefined;
classifierContextIncludeAssistantTurns: boolean | undefined;
classifierFallback: ClassifierFallback | undefined;
sessionAffinity: boolean;
@@ -111,6 +111,7 @@ export interface ComplexityRouterConfigPayload {
classifier_type: ClassifierType;
classifier_llm_config?: ClassifierLLMConfig;
classifier_context_window_size?: number;
+ classifier_context_budget_chars?: number;
classifier_context_per_turn_chars?: number;
classifier_context_include_assistant_turns?: boolean;
classifier_fallback?: ClassifierFallback;
@@ -219,7 +220,7 @@ export const buildComplexityRouterConfig = ({
classifierType,
classifierLlmConfig,
classifierContextWindowSize,
- classifierContextPerTurnChars,
+ classifierContextBudgetChars,
classifierContextIncludeAssistantTurns,
classifierFallback,
sessionAffinity,
@@ -270,8 +271,8 @@ export const buildComplexityRouterConfig = ({
classifier_context_window_size: classifierContextWindowSize,
}),
...(classifierType === "llm" &&
- classifierContextPerTurnChars !== undefined && {
- classifier_context_per_turn_chars: classifierContextPerTurnChars,
+ classifierContextBudgetChars !== undefined && {
+ classifier_context_budget_chars: classifierContextBudgetChars,
}),
...(classifierType === "llm" &&
classifierContextIncludeAssistantTurns !== undefined && {
diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts
index 59254dbbe6f..1a8f5f34909 100644
--- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts
+++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts
@@ -113,18 +113,30 @@ describe("buildUpdatedComplexityRouterConfig classifier context window", () => {
expect(result.classifier_context_per_turn_chars).toBe(300);
});
- it("persists an edited classifier context window size and per-turn char limit", () => {
+ it("persists an edited classifier context window size", () => {
const formValue = {
tiers: STORED_LLM.tiers,
classifier_type: "llm" as const,
classifier_llm_config: STORED_LLM.classifier_llm_config,
classifier_context_window_size: 10,
- classifier_context_per_turn_chars: 500,
};
const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue);
expect(result.classifier_context_window_size).toBe(10);
- expect(result.classifier_context_per_turn_chars).toBe(500);
+ });
+
+ it("carries a stored per-turn cap through untouched now that no control sets it", () => {
+ // The modal stopped rendering a per-turn control, so the key left MANAGED_COMPLEXITY_ROUTER_KEYS.
+ // Had it stayed managed, every open-and-save would have silently dropped an operator's cap.
+ const formValue = {
+ tiers: STORED_LLM.tiers,
+ classifier_type: "llm" as const,
+ classifier_llm_config: STORED_LLM.classifier_llm_config,
+ classifier_context_window_size: 10,
+ };
+ const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue);
+
+ expect(result.classifier_context_per_turn_chars).toBe(300);
});
it("omits classifier context fields when classifier_type is heuristic even if values linger in state", () => {
@@ -132,12 +144,12 @@ describe("buildUpdatedComplexityRouterConfig classifier context window", () => {
tiers: STORED_LLM.tiers,
classifier_type: "heuristic" as const,
classifier_context_window_size: 5,
- classifier_context_per_turn_chars: 300,
+ classifier_context_budget_chars: 4000,
};
const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue);
expect(result.classifier_context_window_size).toBeUndefined();
- expect(result.classifier_context_per_turn_chars).toBeUndefined();
+ expect(result.classifier_context_budget_chars).toBeUndefined();
});
it("does not resurrect a stale stored classifier_context_window_size once the form's own value is unset", () => {
@@ -151,7 +163,6 @@ describe("buildUpdatedComplexityRouterConfig classifier context window", () => {
const result = buildUpdatedComplexityRouterConfig(STORED_LLM, formValue);
expect(result.classifier_context_window_size).toBeUndefined();
- expect(result.classifier_context_per_turn_chars).toBeUndefined();
});
});
diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx
index a124bd08476..6b49ebe740f 100644
--- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx
+++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx
@@ -245,7 +245,7 @@ describe("EditAutoRouterModal classifier context window", () => {
await user.click(await screen.findByText("Advanced: Classification Method"));
await screen.findByText("Context Window Size");
expect(screen.getByDisplayValue("5")).toBeInTheDocument();
- expect(screen.getByDisplayValue("300")).toBeInTheDocument();
+ expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /save changes/i }));
diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx
index 8e5317d9c1f..bce96ec76f5 100644
--- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx
+++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx
@@ -78,7 +78,7 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
"classifier_type",
"classifier_llm_config",
"classifier_context_window_size",
- "classifier_context_per_turn_chars",
+ "classifier_context_budget_chars",
"classifier_context_include_assistant_turns",
"classifier_fallback",
"session_affinity",
@@ -175,8 +175,8 @@ export const buildUpdatedComplexityRouterConfig = (
classifier_context_window_size: value.classifier_context_window_size,
}),
...(value.classifier_type === "llm" &&
- value.classifier_context_per_turn_chars !== undefined && {
- classifier_context_per_turn_chars: value.classifier_context_per_turn_chars,
+ value.classifier_context_budget_chars !== undefined && {
+ classifier_context_budget_chars: value.classifier_context_budget_chars,
}),
...(value.classifier_type === "llm" &&
value.classifier_context_include_assistant_turns !== undefined && {
@@ -368,9 +368,9 @@ const EditAutoRouterModal: React.FC = ({
typeof parsedConfig.classifier_context_window_size === "number"
? parsedConfig.classifier_context_window_size
: undefined,
- classifier_context_per_turn_chars:
- typeof parsedConfig.classifier_context_per_turn_chars === "number"
- ? parsedConfig.classifier_context_per_turn_chars
+ classifier_context_budget_chars:
+ typeof parsedConfig.classifier_context_budget_chars === "number"
+ ? parsedConfig.classifier_context_budget_chars
: undefined,
classifier_context_include_assistant_turns:
typeof parsedConfig.classifier_context_include_assistant_turns === "boolean"
diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts
index 7ef0e06e2e6..f20bd3adb8a 100644
--- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts
+++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts
@@ -268,6 +268,7 @@ export const buildPresetPrefill = (
model: resolve(config.classifier_llm_config.model),
},
classifier_context_window_size: config.classifier_context_window_size,
+ classifier_context_budget_chars: config.classifier_context_budget_chars,
classifier_context_per_turn_chars: config.classifier_context_per_turn_chars,
classifier_context_include_assistant_turns: config.classifier_context_include_assistant_turns,
session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY,
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 87f050e417f..78329a1e53c 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -32487,18 +32487,23 @@ export interface components {
* @description Replaces the opening instructions of the LLM classifier rubric (the judging-criteria prose) for a custom tier set. The per-tier bullets and the trust-boundary paragraph telling the classifier to ignore tier requests embedded in quoted caller text are always appended after it and cannot be overridden. Requires tier_definitions; a built-in-tier router customizes its prompt via classifier_llm_config.system_prompt or classification_rubric instead.
*/
classification_prompt?: string | null;
+ /**
+ * Classifier Context Budget Chars
+ * @description Maximum characters of prior-turn text quoted to the LLM classifier, across the whole context window, per classification call. Turns are taken newest first and quoted whole while they fit, so a conversation small enough to quote entirely is never cut; once the budget runs out the older turns are dropped whole and only the turn straddling the boundary is truncated, into whatever space is left. The current ask and the caller's system prompt sit outside this budget and are always sent in full, as does the numbering each quoted turn carries. A budget under 120 leaves no room to quote a turn and suppresses the block; set classifier_context_window_size to 0 to turn context off deliberately. Only applies when classifier_type is 'llm'.
+ * @default 8000
+ */
+ classifier_context_budget_chars: number;
/**
* Classifier Context Include Assistant Turns
- * @description Include assistant turns in the classifier context window, so difficulty stated by the model rather than by the user stays visible: a plan the assistant calls complex, which the user approves with 'yes', is classified on the work being approved instead of on the word 'yes'. When enabled, classifier_context_window_size counts the last N turns of the conversation across both roles rather than the last N user turns, and assistant text is sent to the classifier model, which may be a different deployment or provider than the routed completion model. Assistant replies share classifier_context_per_turn_chars with user turns, so raise it if replies are truncated before the part that carries the difficulty. Off by default because enabling it shifts tier decisions, and therefore spend, for an already-deployed router. Only applies when classifier_type is 'llm'.
+ * @description Include assistant turns in the classifier context window, so difficulty stated by the model rather than by the user stays visible: a plan the assistant calls complex, which the user approves with 'yes', is classified on the work being approved instead of on the word 'yes'. When enabled, classifier_context_window_size counts the last N turns of the conversation across both roles rather than the last N user turns, and assistant text is sent to the classifier model, which may be a different deployment or provider than the routed completion model. Assistant replies spend classifier_context_budget_chars alongside user turns, so raise it if the oldest turns stop being quoted once replies join the window. Off by default because enabling it shifts tier decisions, and therefore spend, for an already-deployed router. Only applies when classifier_type is 'llm'.
* @default false
*/
classifier_context_include_assistant_turns: boolean;
/**
* Classifier Context Per Turn Chars
- * @description Maximum character length for each prior turn's text in the classifier context window. Turns exceeding this are truncated. Only applies when classifier_type is 'llm'.
- * @default 200
+ * @description Optional cap on each individual prior turn's text, applied before classifier_context_budget_chars bounds the block. Unset by default, so one long turn may spend the whole budget, which is usually what a follow-up needs; set it when no single turn should dominate the context the classifier sees. A capped turn keeps its opening and its ending with the middle elided. Only applies when classifier_type is 'llm'.
*/
- classifier_context_per_turn_chars: number;
+ classifier_context_per_turn_chars?: number | null;
/**
* Classifier Context Window Size
* @description Number of prior user turns (tool output and harness reminders excluded) to include as context in the LLM classifier prompt, so a follow-up like 'now do the same for the streaming path' is classified against what it refers to. Counts turns of both roles when classifier_context_include_assistant_turns is enabled. These turns are sent to the classifier model, which may be a different deployment or provider than the routed completion model; that call already carries the current user ask and the caller's system prompt in full. Set to 0 to send neither prior turns nor any conversation context beyond the current ask. Only applies when classifier_type is 'llm'.
From bb27bfd9a7457e69b79ff2901f165f3a0e4c8ef0 Mon Sep 17 00:00:00 2001
From: Anmol Jaiswal <68013660+anmolg1997@users.noreply.github.com>
Date: Tue, 25 Aug 2026 20:42:10 +0530
Subject: [PATCH 22/70] fix(http_handler): dispose aiohttp session when
AsyncHTTPHandler is finalized without a running loop (#36670)
* fix(http_handler): dispose aiohttp session when finalized without a running loop
AsyncHTTPHandler.__del__ can only schedule an async close when a running
event loop exists at finalization time; in any other context (worker
threads whose loop has closed, sync contexts, interpreter shutdown) the
RuntimeError from get_running_loop() is swallowed and the underlying
aiohttp ClientSession is abandoned to GC, emitting 'Unclosed client
session' / 'Unclosed connector' warnings.
This is the disposal gap left after the recycle-time fix: clients created
for short-lived event loops (the loop-id-keyed LLM client cache mints one
handler per loop) are never recycled - they live and die with their loop,
and their finalization is precisely the loop-less case.
Fix:
- no running loop: fall back to the connector's synchronous teardown via
LiteLLMAiohttpTransport._mark_connector_closed - the same finalizer-safe
path used for dead-loop recycles - honoring _owns_session so a shared
session is never closed.
- running loop: keep the async close, but hold a strong reference to the
scheduled task until it completes (a bare create_task() result may be
collected before running), mirroring _background_close_tasks.
Tests: loop-less finalization closes a dead-loop session; running-loop
finalization registers and drains the close task; the sync fallback
respects session ownership. All three fail without the fix.
* lint: conform new finalizer code to the type-discipline budget
Final on the five never-rebound locals (LIT010); the class-level task
registry keeps its mutable set with the sanctioned mutable-ok reason,
mirroring the aiohttp transport's registry (LIT001).
* lint: reasoned pyright ignore on the cross-class teardown call
The handler deliberately reuses the transport's finalizer-safe connector
teardown; no public seam exists and an async close can never run at
loop-less finalization. Clears the net-new reportPrivateUsage the
basedpyright budget gate flagged once the LIT stage passed.
* fix(http_handler): retrieve exceptions from finalizer close tasks
A bare discard done-callback dropped the task without consuming its
exception, so a failing aclose() emitted "Task exception was never
retrieved" at GC, the same noise class this path exists to remove.
Mirror the transport's _on_close_task_done: discard, early-return on
cancellation, retrieve and debug-log the exception.
* fix(http_handler): dispose foreign-loop sessions instead of scheduling aclose on the live loop
GC on a live loop (e.g. the app's) of a handler whose session belongs to
another, possibly dead, loop scheduled aclose() on the current loop, the
cross-loop path the transport refuses. Route both that case and the
loop-less case through the transport's lifecycle-aware
_close_recycled_session, which picks async close on the session's own
loop, threadsafe handoff, or the synchronous connector teardown.
Regression test: a dead-loop session collected while another loop runs
is disposed without scheduling anything on that loop.
* chore: retrigger CI (test_mcp_logging payload-order flake, also failed on litellm_spendlogs_fallback_metadata minutes earlier)
* test(mcp): select the MCP tool-call payload instead of the last-delivered one
TestMCPLogger kept a single last-writer slot; an async success event from
another call (a mocked acompletion whose log task lands late) races the
MCP event for it, so the cost assertions intermittently read the wrong
payload. This PR's finalizer change shifts task interleaving on the loop
and tips that latent race over (also seen on an unrelated PR minutes
earlier). Collect call_type=call_mcp_tool payloads in their own list and
assert on those.
* test(mcp): MCPLoggerHook inherits the order-independent payload capture
It duplicated TestMCPLogger's init and success handler verbatim; the
hook test reads the same MCP payload selection, so subclass instead.
---
litellm/llms/custom_httpx/http_handler.py | 76 ++++++++-
tests/mcp_tests/test_mcp_logging.py | 48 ++++--
.../llms/custom_httpx/test_http_handler.py | 161 +++++++++++++++---
3 files changed, 246 insertions(+), 39 deletions(-)
diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py
index 52f30e31641..777ab576de2 100644
--- a/litellm/llms/custom_httpx/http_handler.py
+++ b/litellm/llms/custom_httpx/http_handler.py
@@ -9,7 +9,7 @@ import threading
import time
from collections.abc import AsyncIterable, Callable, Iterable, Mapping
from http.cookiejar import CookieJar, DefaultCookiePolicy
-from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, TypedDict
+from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional, TypeAlias, TypedDict
import certifi
import httpx
@@ -933,11 +933,83 @@ class AsyncHTTPHandler:
response.raise_for_status()
return response
+ # Strong references to finalizer-scheduled client-close tasks. A bare
+ # create_task() result may be garbage-collected before it runs, leaving
+ # the underlying aiohttp session unclosed ("Unclosed client session").
+ # Mirrors LiteLLMAiohttpTransport._background_close_tasks.
+ _finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes
+
+ @classmethod
+ def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None:
+ cls._finalizer_close_tasks.discard(task)
+ if task.cancelled():
+ return
+ exc: Final = task.exception()
+ if exc is not None:
+ verbose_logger.debug("Error closing client at finalization: %s", exc)
+
+ def _aiohttp_session_bound_elsewhere(self, loop: asyncio.AbstractEventLoop) -> bool:
+ """True when the wrapped aiohttp session is bound to a loop other than
+ ``loop`` — awaiting ``aclose()`` here would touch that loop's internals."""
+ from litellm.llms.custom_httpx.aiohttp_transport import (
+ LiteLLMAiohttpTransport,
+ )
+
+ transport: Final = getattr(self._client, "_transport", None)
+ if not isinstance(transport, LiteLLMAiohttpTransport):
+ return False
+ session: Final = transport.client
+ if not isinstance(session, ClientSession) or session.closed:
+ return False
+ return getattr(session, "_loop", None) is not loop
+
+ def _dispose_wrapped_aiohttp_session(self) -> None:
+ """Dispose the wrapped aiohttp session when ``aclose()`` cannot run here.
+
+ Finalization either has no running loop, or a loop the session is not
+ bound to. Delegating to the transport's lifecycle-aware disposal picks
+ the safe path per session state (async close on its own loop, threadsafe
+ handoff to a loop running elsewhere, or the synchronous connector
+ teardown that flips the flags ``ClientSession.__del__`` checks), so no
+ "Unclosed client session" / "Unclosed connector" warnings fire at
+ garbage collection.
+ """
+ from litellm.llms.custom_httpx.aiohttp_transport import (
+ LiteLLMAiohttpTransport,
+ )
+
+ transport: Final = getattr(self._client, "_transport", None)
+ if not isinstance(transport, LiteLLMAiohttpTransport):
+ return
+ # A shared session (e.g. the proxy's) is never this handler's to close.
+ if not getattr(transport, "_owns_session", False):
+ return
+ session: Final = transport.client
+ if isinstance(session, ClientSession) and not session.closed:
+ transport._close_recycled_session(session) # pyright: ignore[reportPrivateUsage] # deliberate reuse of the transport's lifecycle-aware disposal; an async close can never run in this context
+
def __del__(self) -> None:
try:
if not _handler_may_close_client(sys.getrefcount(self._client), self._owns_client):
return
- asyncio.get_running_loop().create_task(self._client.aclose())
+ try:
+ loop: Final = asyncio.get_running_loop()
+ except RuntimeError:
+ # No running loop at finalization time (worker threads after
+ # their loop closed, interpreter/worker shutdown, GC in a
+ # sync context). An async close can never run here.
+ self._dispose_wrapped_aiohttp_session()
+ return
+ if self._aiohttp_session_bound_elsewhere(loop):
+ # GC ran on a live loop (e.g. the app's) but the session
+ # belongs to another, possibly dead, loop — awaiting aclose()
+ # here is the cross-loop path the transport refuses.
+ self._dispose_wrapped_aiohttp_session()
+ return
+ task: Final = loop.create_task(self._client.aclose())
+ cls: Final = type(self)
+ cls._finalizer_close_tasks.add(task)
+ task.add_done_callback(cls._on_finalizer_close_done)
except Exception:
pass
diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py
index 7ee745b311e..1903f29001f 100644
--- a/tests/mcp_tests/test_mcp_logging.py
+++ b/tests/mcp_tests/test_mcp_logging.py
@@ -24,12 +24,20 @@ from mcp.types import Tool as MCPTool, CallToolResult, TextContent
class TestMCPLogger(CustomLogger):
def __init__(self):
self.standard_logging_payload = None
+ self.mcp_tool_call_payloads = []
super().__init__()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
print("success event")
- self.standard_logging_payload = kwargs.get("standard_logging_object", None)
- print(f"Captured standard_logging_payload: {self.standard_logging_payload}")
+ payload = kwargs.get("standard_logging_object", None)
+ self.standard_logging_payload = payload
+ # Async success events from other calls (e.g. a mocked acompletion whose
+ # log task is delivered late) race with the MCP event for the single
+ # last-writer slot; keep MCP tool calls in their own list so assertions
+ # are order-independent.
+ if payload is not None and payload.get("call_type") == "call_mcp_tool":
+ self.mcp_tool_call_payloads.append(payload)
+ print(f"Captured standard_logging_payload: {payload}")
def _set_authorized_user(server_ids):
@@ -138,7 +146,11 @@ async def test_mcp_cost_tracking():
# wait 1-2 seconds for logging to be processed
await asyncio.sleep(2)
- logged_standard_logging_payload = test_logger.standard_logging_payload
+ logged_standard_logging_payload = (
+ test_logger.mcp_tool_call_payloads[-1]
+ if test_logger.mcp_tool_call_payloads
+ else None
+ )
print("logged_standard_logging_payload", logged_standard_logging_payload)
# Add assertions
@@ -277,7 +289,11 @@ async def test_mcp_cost_tracking_per_tool():
# wait for logging to be processed
await asyncio.sleep(2)
- logged_standard_logging_payload_1 = test_logger.standard_logging_payload
+ logged_standard_logging_payload_1 = (
+ test_logger.mcp_tool_call_payloads[-1]
+ if test_logger.mcp_tool_call_payloads
+ else None
+ )
print(
"logged_standard_logging_payload_1", logged_standard_logging_payload_1
)
@@ -290,6 +306,7 @@ async def test_mcp_cost_tracking_per_tool():
# Reset logger for second test
test_logger.standard_logging_payload = None
+ test_logger.mcp_tool_call_payloads.clear()
# Test 2: Call cheap_tool - should cost 0.1
response2 = await mcp_server_tool_call(
@@ -300,7 +317,11 @@ async def test_mcp_cost_tracking_per_tool():
# wait for logging to be processed
await asyncio.sleep(2)
- logged_standard_logging_payload_2 = test_logger.standard_logging_payload
+ logged_standard_logging_payload_2 = (
+ test_logger.mcp_tool_call_payloads[-1]
+ if test_logger.mcp_tool_call_payloads
+ else None
+ )
print(
"logged_standard_logging_payload_2", logged_standard_logging_payload_2
)
@@ -329,16 +350,7 @@ async def test_mcp_cost_tracking_per_tool():
assert mock_client.call_tool.call_count == 2
-class MCPLoggerHook(CustomLogger):
- def __init__(self):
- self.standard_logging_payload = None
- super().__init__()
-
- async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
- print("success event")
- self.standard_logging_payload = kwargs.get("standard_logging_object", None)
- print(f"Captured standard_logging_payload: {self.standard_logging_payload}")
-
+class MCPLoggerHook(TestMCPLogger):
async def async_post_mcp_tool_call_hook(
self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time
) -> Optional[MCPPostCallResponseObject]:
@@ -436,7 +448,11 @@ async def test_mcp_tool_call_hook():
await asyncio.sleep(2)
# check logged standard logging payload
- logged_standard_logging_payload = test_logger.standard_logging_payload
+ logged_standard_logging_payload = (
+ test_logger.mcp_tool_call_payloads[-1]
+ if test_logger.mcp_tool_call_payloads
+ else None
+ )
print("logged_standard_logging_payload", logged_standard_logging_payload)
assert (
logged_standard_logging_payload is not None
diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py
index f7f89cd1d8d..16d57437043 100644
--- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py
+++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py
@@ -56,9 +56,7 @@ async def test_async_post_streaming_status_error_should_not_wait_forever_for_bod
litellm_handler = AsyncHTTPHandler()
await litellm_handler.client.aclose()
- litellm_handler.client = httpx.AsyncClient(
- transport=httpx.MockTransport(mock_handler)
- )
+ litellm_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler))
try:
with pytest.raises(MaskedHTTPStatusError) as exc_info:
await asyncio.wait_for(
@@ -202,9 +200,7 @@ async def test_ssl_verification_with_aiohttp_transport(monkeypatch: pytest.Monke
transport_connector = transport._get_valid_client_session().connector
assert isinstance(transport_connector, TCPConnector)
- aiohttp_session = aiohttp.ClientSession(
- connector=aiohttp.TCPConnector(ssl=False)
- )
+ aiohttp_session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False))
try:
aiohttp_connector = aiohttp_session.connector
assert isinstance(aiohttp_connector, aiohttp.TCPConnector)
@@ -378,7 +374,8 @@ async def test_get_async_httpx_client_with_shared_session():
# Test with shared session
client = get_async_httpx_client(
- llm_provider=LlmProviders.ANTHROPIC, shared_session=mock_session # type: ignore
+ llm_provider=LlmProviders.ANTHROPIC,
+ shared_session=mock_session, # type: ignore
)
# Verify the client was created successfully
@@ -397,9 +394,7 @@ async def test_get_async_httpx_client_without_shared_session():
from litellm.types.utils import LlmProviders
# Test without shared session
- client = get_async_httpx_client(
- llm_provider=LlmProviders.ANTHROPIC, shared_session=None
- )
+ client = get_async_httpx_client(llm_provider=LlmProviders.ANTHROPIC, shared_session=None)
# Verify the client was created successfully
assert client is not None
@@ -476,11 +471,13 @@ async def test_session_reuse_integration():
# Create two clients with the same session
client1 = get_async_httpx_client(
- llm_provider=LlmProviders.ANTHROPIC, shared_session=mock_session # type: ignore
+ llm_provider=LlmProviders.ANTHROPIC,
+ shared_session=mock_session, # type: ignore
)
client2 = get_async_httpx_client(
- llm_provider=LlmProviders.OPENAI, shared_session=mock_session # type: ignore
+ llm_provider=LlmProviders.OPENAI,
+ shared_session=mock_session, # type: ignore
)
# Both clients should be created successfully
@@ -512,9 +509,7 @@ async def test_session_reuse_integration():
(None, None, None, False), # None value - skip configuration
],
)
-def test_ssl_ecdh_curve(
- env_curve, litellm_curve, expected_curve, should_call, monkeypatch
-):
+def test_ssl_ecdh_curve(env_curve, litellm_curve, expected_curve, should_call, monkeypatch):
"""Test SSL ECDH curve configuration with valid curves and precedence"""
from litellm.llms.custom_httpx.http_handler import _ssl_context_cache
@@ -717,9 +712,7 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout:
_default_cached_client_timeout,
)
- monkeypatch.setattr(
- litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS
- )
+ monkeypatch.setattr(litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS)
monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False)
assert _default_cached_client_timeout() is _DEFAULT_TIMEOUT
@@ -734,9 +727,7 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout:
assert resolved.read == 300.0
assert resolved.connect == 5.0
- def test_cached_async_client_built_with_explicit_request_timeout(
- self, monkeypatch: pytest.MonkeyPatch
- ):
+ def test_cached_async_client_built_with_explicit_request_timeout(self, monkeypatch: pytest.MonkeyPatch):
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.utils import LlmProviders
@@ -1195,3 +1186,131 @@ async def test_aiohttp_session_never_replays_one_upstreams_cookie_to_another():
assert len(jar) == 0
assert dict(jar.filter_cookies(URL("https://upstream-a.example.com"))) == {}
await session.close()
+
+
+def _mint_session_on_dead_loop(handler: AsyncHTTPHandler) -> ClientSession:
+ """Create the transport's real ClientSession on a loop that then closes.
+
+ This is the lifecycle of every client minted for a short-lived event loop
+ (the loop-id-keyed LLM client cache creates one handler per loop): the
+ session outlives its loop and can only ever be disposed loop-lessly.
+ """
+ transport = handler.client._transport
+ assert isinstance(transport, LiteLLMAiohttpTransport)
+ loop = asyncio.new_event_loop()
+
+ async def _create() -> ClientSession:
+ return transport._get_valid_client_session()
+
+ session = loop.run_until_complete(_create())
+ loop.close()
+ return session
+
+
+def test_finalizer_without_running_loop_closes_dead_loop_session():
+ """A handler finalized with no running event loop must still dispose its
+ aiohttp session.
+
+ The async close can never run in that context; without the synchronous
+ fallback the session and its connector are abandoned to GC and emit
+ "Unclosed client session" / "Unclosed connector" warnings."""
+ handler = AsyncHTTPHandler(timeout=61.0)
+ session = _mint_session_on_dead_loop(handler)
+ assert not session.closed
+
+ del handler
+ gc.collect()
+
+ assert session.closed
+
+
+@pytest.mark.asyncio
+async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref():
+ """With a running loop, finalization schedules an async close and must keep
+ a strong reference to the task until it completes — a bare create_task()
+ result may be collected before it runs, leaving the session unclosed."""
+ handler = AsyncHTTPHandler(timeout=61.0)
+ transport = handler.client._transport
+ assert isinstance(transport, LiteLLMAiohttpTransport)
+ session = transport._get_valid_client_session()
+ assert not session.closed
+ del transport
+
+ baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks)
+ del handler
+ gc.collect()
+
+ scheduled = AsyncHTTPHandler._finalizer_close_tasks - baseline_tasks
+ assert len(scheduled) == 1
+
+ await asyncio.gather(*scheduled)
+ assert session.closed
+ assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled)
+
+
+@pytest.mark.asyncio
+async def test_sync_close_helper_respects_session_ownership():
+ """The loop-less fallback closes only sessions the transport owns; a
+ shared session (e.g. the proxy's) must never be closed by a handler."""
+ owned_handler = AsyncHTTPHandler(timeout=61.0)
+ owned_transport = owned_handler.client._transport
+ assert isinstance(owned_transport, LiteLLMAiohttpTransport)
+ owned_session = owned_transport._get_valid_client_session()
+
+ baseline = set(LiteLLMAiohttpTransport._background_close_tasks)
+ owned_handler._dispose_wrapped_aiohttp_session()
+ scheduled = LiteLLMAiohttpTransport._background_close_tasks - baseline
+ await asyncio.gather(*scheduled)
+ assert owned_session.closed
+
+ shared_session = ClientSession()
+ shared_handler = AsyncHTTPHandler(timeout=61.0, shared_session=shared_session)
+ shared_transport = shared_handler.client._transport
+ assert isinstance(shared_transport, LiteLLMAiohttpTransport)
+ assert shared_transport._owns_session is False
+
+ shared_handler._dispose_wrapped_aiohttp_session()
+ assert not shared_session.closed
+
+ await shared_session.close()
+ await shared_handler.close()
+ await owned_handler.close()
+
+
+@pytest.mark.asyncio
+async def test_finalizer_close_done_consumes_exception():
+ """A failing finalizer close must have its exception retrieved by the done
+ callback, or asyncio emits "Task exception was never retrieved" at GC —
+ the same log noise the finalizer path exists to eliminate."""
+
+ async def failing_close() -> None:
+ raise RuntimeError("close failed")
+
+ task = asyncio.get_running_loop().create_task(failing_close())
+ AsyncHTTPHandler._finalizer_close_tasks.add(task)
+ await asyncio.sleep(0)
+
+ AsyncHTTPHandler._on_finalizer_close_done(task)
+ assert task not in AsyncHTTPHandler._finalizer_close_tasks
+
+ cancelled = asyncio.get_running_loop().create_task(asyncio.sleep(30))
+ cancelled.cancel()
+ await asyncio.sleep(0)
+ AsyncHTTPHandler._on_finalizer_close_done(cancelled)
+
+
+@pytest.mark.asyncio
+async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_scheduling():
+ """GC on a live loop (e.g. the app's) of a handler whose session belongs to
+ another, dead loop must not schedule aclose() here — that is the cross-loop
+ path the transport refuses — and must still dispose the session."""
+ handler = AsyncHTTPHandler(timeout=61.0)
+ session = await asyncio.to_thread(_mint_session_on_dead_loop, handler)
+ assert not session.closed
+
+ baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks)
+ del handler
+ gc.collect()
+
+ assert AsyncHTTPHandler._finalizer_close_tasks == baseline_tasks
+ assert session.closed
From fc810484f790f283e2f0a897fa972783e707431c Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 09:28:55 -0700
Subject: [PATCH 23/70] Revert "ci: raise three unit shard job timeouts to
satisfy the startup safety gate"
This reverts commit b71b574af410f436954f9e6fcb28b44eff6c1d34.
---
.github/workflows/test-unit.yml | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index 2dfca3d308f..a7c67f2b35d 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -211,7 +211,7 @@ jobs:
workers: 2
reruns: 2
timeout-minutes: 20
- job-timeout-minutes: 60
+ job-timeout-minutes: 55
- shard: proxy-extras
artifact-name: proxy-extras
@@ -219,7 +219,7 @@ jobs:
workers: 2
reruns: 2
timeout-minutes: 20
- job-timeout-minutes: 60
+ job-timeout-minutes: 55
- shard: enterprise-package
artifact-name: enterprise-package
@@ -227,7 +227,7 @@ jobs:
workers: 4
reruns: 2
timeout-minutes: 20
- job-timeout-minutes: 60
+ job-timeout-minutes: 55
- shard: responses-caching-types
artifact-name: responses-caching-types
From a73f11ae9c736059299c2078ee602bc05e43559d Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 09:40:23 -0700
Subject: [PATCH 24/70] fix(completion_extras): forward reasoning_effort=max
through the Responses API bridge
---
.../transformation.py | 22 +++-------
litellm/types/llms/openai.py | 2 +-
...responses_transformation_transformation.py | 42 ++++++++++++++++---
.../response_api_endpoints/test_endpoints.py | 2 +
4 files changed, 46 insertions(+), 22 deletions(-)
diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py
index b94e91b3034..17815976b4a 100644
--- a/litellm/completion_extras/litellm_responses_transformation/transformation.py
+++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py
@@ -5,7 +5,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req
import json
import os
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args
from openai.types.responses.custom_tool_param import CustomToolParam
from openai.types.responses.response_input_param import (
@@ -35,6 +35,7 @@ from litellm.responses.sse_output_recovery import (
)
from litellm.responses.utils import normalize_responses_api_stream_options
from litellm.types.llms.openai import (
+ REASONING_EFFORT,
ChatCompletionAnnotation,
ChatCompletionReasoningItem,
ChatCompletionToolCallChunk,
@@ -1113,22 +1114,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"
)
- # If string is passed, map with optional summary based on flag/env var
- if reasoning_effort == "none":
- return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none")
- elif reasoning_effort == "high":
- return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high")
- elif reasoning_effort == "xhigh":
- return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh")
- elif reasoning_effort == "medium":
+ if reasoning_effort in get_args(REASONING_EFFORT):
return (
- Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium")
- )
- elif reasoning_effort == "low":
- return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low")
- elif reasoning_effort == "minimal":
- return (
- Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal")
+ Reasoning(effort=reasoning_effort, summary="detailed")
+ if auto_summary_enabled
+ else Reasoning(effort=reasoning_effort)
)
return None
diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py
index e7a3f825455..4a6c4a5bbb5 100644
--- a/litellm/types/llms/openai.py
+++ b/litellm/types/llms/openai.py
@@ -1840,7 +1840,7 @@ ResponsesAPIStreamingResponse = Annotated[
]
-REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
+REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh", "max"]
class OpenAIRealtimeStreamSession(TypedDict, total=False):
diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
index 1fb74b2b7bf..6ca48ce63b8 100644
--- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
+++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
@@ -2,7 +2,7 @@ import datetime
import json
import os
import unittest
-from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
+from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple
from unittest.mock import ANY, MagicMock, Mock, patch
import httpx
@@ -1585,10 +1585,16 @@ def test_map_reasoning_effort_adds_summary_detailed(monkeypatch):
assert result_dict["summary"] == "custom_summary"
print("✓ Dict input is passed through without modification")
- # Test 5: None/unknown values return None
- result_unknown = handler._map_reasoning_effort("unknown_value")
- assert result_unknown is None
- print("✓ Unknown reasoning_effort values return None")
+ # Test 5: every REASONING_EFFORT level reaches the provider, and anything else (a typo, an
+ # unshipped level, "default") is dropped so the request still succeeds at the provider default
+ from litellm.types.llms.openai import Reasoning
+
+ for effort in ("max", "xhigh", "none"):
+ result_passthrough = handler._map_reasoning_effort(effort)
+ assert result_passthrough == Reasoning(effort=effort)
+ for dropped in ("ultra", "hgih", "unknown_value", "", "default"):
+ assert handler._map_reasoning_effort(dropped) is None
+ print("✓ Enumerated levels pass through and unknown ones are dropped")
print(
"✓ All reasoning_effort behaviors work correctly with flag/env var control"
@@ -2438,6 +2444,32 @@ def test_map_optional_params_preserves_reasoning_summary():
assert responses_api_request["reasoning"]["summary"] == "detailed"
+@pytest.mark.parametrize("reasoning_effort", ["max", "high"])
+def test_transform_request_bedrock_mantle_tools_keeps_reasoning_effort(monkeypatch, reasoning_effort):
+ """Regression for reasoning_effort=max being dropped on the chat -> Responses bridge (issue #38084)."""
+ from litellm.completion_extras.litellm_responses_transformation.transformation import (
+ LiteLLMResponsesTransformationHandler,
+ )
+
+ monkeypatch.setattr(litellm, "reasoning_auto_summary", False)
+ monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False)
+ handler: Final = LiteLLMResponsesTransformationHandler()
+
+ result: Final = handler.transform_request(
+ model="openai.gpt-5.6-sol",
+ messages=[{"role": "user", "content": "Say pong"}],
+ optional_params={
+ "reasoning_effort": reasoning_effort,
+ "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}],
+ },
+ litellm_params={"custom_llm_provider": "bedrock_mantle"},
+ headers={},
+ litellm_logging_obj=Mock(),
+ )
+
+ assert result["reasoning"] == {"effort": reasoning_effort}
+
+
def test_map_optional_params_tool_choice_chat_nested_to_responses_api():
"""Chat tool_choice must become Responses ToolChoiceFunction (top-level name)."""
from litellm.completion_extras.litellm_responses_transformation.transformation import (
diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
index 9177944df2d..791d64c6428 100644
--- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
+++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py
@@ -1353,6 +1353,8 @@ class TestParseCursorModelVariant:
("claude-opus-5-fast", "claude-opus-5", None),
("gpt-5.6-sol", "gpt-5.6-sol", None),
("foo-thinking-ultra-fast", "foo-thinking-ultra", None),
+ ("gpt-5.6-thinking-max", "gpt-5.6", "max"),
+ ("foo-thinking-mega-fast", "foo-thinking-mega", None),
("-thinking-high", "-thinking-high", None),
],
)
From 1d695a714b41d2f4ebc0cb87ea560be2d620b0f4 Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Tue, 25 Aug 2026 09:50:09 -0700
Subject: [PATCH 25/70] fix(proxy): reset a stuck team member's budget (#37971)
* fix(proxy): reset a stuck team member's budget
A per-team-member budget check reads a cross-pod spend counter that
nothing ever invalidates. Once a member exceeds their per-member
budget, resetting the key's spend, raising the user's or the team's
own budget, or issuing a new key all leave the member stuck, because
none of them touch this counter or its cached membership object.
Add POST /team/{team_id}/member/{user_id}/reset_spend to reset a
member's tracked spend, and invalidate the same cached state from
/team/member_update when it raises a member's own budget, so that
path also takes effect immediately instead of waiting on the
membership cache's TTL. Name the entity in the check's error message
so a stuck member is diagnosable from the 429 body alone.
* fix(proxy): close reset-vs-floor-read race and surface double Redis write failure on member spend reset
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): broadcast spend reset as a SET so the handler's self-delivered message cannot erase the reset guard
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(proxy): omit null fields from the invalidation message so plain evictions keep the old wire format
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/proxy/_types.py | 11 +
litellm/proxy/auth/auth_checks.py | 114 +++++-
litellm/proxy/auth/user_api_key_auth.py | 50 ++-
.../auth_cache_invalidation_pubsub.py | 57 ++-
.../proxy/common_utils/user_api_key_cache.py | 15 +
.../management_endpoints/team_endpoints.py | 131 +++++-
litellm/proxy/proxy_server.py | 7 +
.../spend_tracking/budget_reservation.py | 10 +-
.../test_team_member_reset_spend.py | 152 +++++++
.../proxy/auth/test_auth_checks.py | 305 ++++++++++++++
.../test_auth_cache_invalidation_pubsub.py | 32 ++
.../test_team_endpoints.py | 380 ++++++++++++++++++
tests/test_litellm/proxy/test_proxy_server.py | 35 ++
ui/litellm-dashboard/src/lib/http/schema.d.ts | 63 +++
14 files changed, 1324 insertions(+), 38 deletions(-)
create mode 100644 tests/proxy_behavior/management/test_team_member_reset_spend.py
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 0840d37ffa1..628e569e1b8 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -815,6 +815,7 @@ class LiteLLMRoutes(enum.Enum):
"/team/member_add",
"/team/member_delete",
"/team/member_update",
+ "/team/{team_id}/member/{user_id}/reset_spend",
"/team/permissions_list",
"/team/permissions_update",
"/team/daily/activity",
@@ -1287,6 +1288,16 @@ class RegenerateKeyRequest(GenerateKeyRequest):
class ResetSpendRequest(LiteLLMPydanticObjectBase):
reset_to: float
+ @field_validator("reset_to", mode="before")
+ @classmethod
+ def reject_bool_reset_to(cls, v):
+ # bool is a subclass of int, so pydantic silently coerces True/False into
+ # 1.0/0.0 for a `float` field: a caller who accidentally sends a boolean
+ # would otherwise get an unintended spend reset instead of a 422.
+ if isinstance(v, bool):
+ raise ValueError("reset_to must be a number, not a boolean") # noqa: TRY004 # pydantic needs ValueError
+ return v
+
class KeyRequest(LiteLLMPydanticObjectBase):
keys: list[str] | None = None
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index e7b98b3cc7f..4af942a357e 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -87,6 +87,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
object_permission_cache_key,
tag_cache_key,
tag_registry_cache_key,
+ team_membership_auth_cache_key,
+ team_membership_reservation_cache_key,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.guardrails.tool_name_extraction import (
@@ -1967,7 +1969,7 @@ async def get_team_membership(
if user_id is None or team_id is None:
return None
- _key: Final = f"team_membership:{user_id}:{team_id}"
+ _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id)
# check if in cache
cached_membership_obj: Final = await user_api_key_cache.async_get_cache(
@@ -2402,6 +2404,116 @@ async def _cache_team_object(
)
+async def invalidate_team_member_spend_state(
+ user_id: str,
+ team_id: str,
+ user_api_key_cache: UserApiKeyCache,
+ new_spend: float | None = None,
+) -> None:
+ """
+ Clear every cached read path for one team member's budget so a spend
+ reset or a raised cap takes effect on the next request instead of
+ waiting on the membership cache's TTL.
+
+ Two independently-keyed cache entries hold the same LiteLLM_TeamMembership
+ row: user_api_key_auth.py's admission check writes ``{team_id}_{user_id}``,
+ while budget_reservation.py's pre-call reservation and auth_checks.py's own
+ get_team_membership() (used by _check_team_member_budget) both write
+ ``team_membership:{user_id}:{team_id}``. Both formats must be invalidated
+ explicitly; writing one does not refresh the other. All keys are also
+ broadcast (LIT-3803): each worker's own in-memory copy (membership object,
+ spend counter, or the counter's own short-TTL DB-floor marker) survives
+ eviction elsewhere until its TTL, so the handling worker alone clearing its
+ copy leaves every other worker still enforcing the pre-reset budget.
+
+ ``new_spend`` is only passed by reset_team_member_spend_fn, which knows the
+ exact post-reset value: it is SET everywhere (matching /key/{key}/reset_spend's
+ own precedent) rather than deleted, so a worker's next read reflects it
+ directly instead of re-deriving it through a DB reseed. team_member_update
+ only changes the budget cap, not the tracked spend, so it passes no
+ new_spend; the live spend counter is untouched in that case (deleting it
+ would force a reseed from the DB's own spend column, which lags the live
+ counter via periodic batch writes, briefly under-enforcing the raised cap
+ against a spend value lower than what was actually tracked) and only the
+ membership caches carrying the new cap are invalidated.
+
+ The floor marker (``spend_db_floor:``, proxy_server.py's
+ _authoritative_floor_spend) caches the pre-reset DB spend for
+ SPEND_DB_FLOOR_CACHE_TTL_SECONDS; left stale after a real reset, a request
+ landing on the pod that cached it can read that higher floor and raise the
+ counter right back above the just-reset spend. It is overwritten here with
+ the post-reset floor (not merely deleted) and _authoritative_floor_spend
+ re-checks the marker after its DB read, so a floor read already in flight
+ on this pod when the reset commits cannot clobber it with the pre-reset
+ value. Both keys are broadcast as SETs carrying new_spend, not deletes:
+ every subscriber (remote pods AND this pod's own, which receives its own
+ message) writes the post-reset value, so the self-delivered message cannot
+ erase the guard just written here.
+
+ Raises HTTPException(503) if Redis still holds the stale pre-reset counter
+ after both the SET and the fallback DELETE fail: budget checks read Redis
+ first, so returning success would leave the old value authoritative for
+ every worker despite the DB write having committed.
+ """
+ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
+ evict_and_broadcast,
+ publish_auth_cache_invalidation,
+ )
+
+ if new_spend is not None:
+ from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
+
+ spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
+ spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
+
+ spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
+ if spend_counter_cache.redis_cache is not None:
+ try:
+ await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
+ except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
+ verbose_proxy_logger.warning(
+ "Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
+ "read reseeds from the DB rather than keeping the stale pre-reset value authoritative",
+ spend_counter_key,
+ e,
+ )
+ try:
+ await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key)
+ except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success
+ verbose_proxy_logger.warning(
+ "Failed to delete stale spend counter %s in Redis after a failed reset write",
+ spend_counter_key,
+ exc_info=True,
+ )
+ raise HTTPException(
+ status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+ detail={ # mutable-ok: HTTPException.detail takes a dict
+ "error": "Spend was reset in the database, but Redis is unreachable and still "
+ "holds the pre-reset counter. Retry once Redis is reachable."
+ },
+ ) from e
+
+ spend_counter_cache.in_memory_cache.set_cache(
+ key=spend_db_floor_key,
+ value=new_spend,
+ ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
+ )
+ await publish_auth_cache_invalidation(cache_key=spend_counter_key, new_value=new_spend, ttl=60)
+ await publish_auth_cache_invalidation(
+ cache_key=spend_db_floor_key,
+ new_value=new_spend,
+ ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
+ )
+
+ await evict_and_broadcast(
+ cache_keys=(
+ team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
+ team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
+ ),
+ user_api_key_cache=user_api_key_cache,
+ )
+
+
async def delete_cache_team_object(
team_id: str,
team_alias: str | None,
diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py
index 658d176f6a7..28d76e6799c 100644
--- a/litellm/proxy/auth/user_api_key_auth.py
+++ b/litellm/proxy/auth/user_api_key_auth.py
@@ -87,7 +87,10 @@ from litellm.proxy.common_utils.http_parsing_utils import (
populate_request_with_path_params,
)
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
-from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+from litellm.proxy.common_utils.user_api_key_cache import (
+ UserApiKeyCache,
+ team_membership_auth_cache_key,
+)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.utils import (
@@ -1970,8 +1973,10 @@ async def _user_api_key_auth_builder(
# Check 3. Check if user is in their team budget
if not skip_budget_checks and valid_token.team_member_spend is not None:
- if prisma_client is not None:
- _cache_key: Final = f"{valid_token.team_id}_{valid_token.user_id}"
+ _user_id: Final = valid_token.user_id
+ _team_id: Final = valid_token.team_id
+ if prisma_client is not None and _user_id is not None and _team_id is not None:
+ _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id)
team_member_info = await user_api_key_cache.async_get_cache(
key=_cache_key,
@@ -1979,25 +1984,21 @@ async def _user_api_key_auth_builder(
)
if team_member_info is None:
# read from DB
- _user_id: Final = valid_token.user_id
- _team_id: Final = valid_token.team_id
-
- if _user_id is not None and _team_id is not None:
- _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first(
- where={
- "user_id": _user_id,
- "team_id": _team_id,
- },
- include={"litellm_budget_table": True},
+ _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first(
+ where={
+ "user_id": _user_id,
+ "team_id": _team_id,
+ },
+ include={"litellm_budget_table": True},
+ )
+ if _db_member is not None:
+ team_member_info = LiteLLM_TeamMembership(**_db_member.dict())
+ await user_api_key_cache.async_set_cache(
+ key=_cache_key,
+ value=team_member_info,
+ model_type=LiteLLM_TeamMembership,
+ ttl=5,
)
- if _db_member is not None:
- team_member_info = LiteLLM_TeamMembership(**_db_member.dict())
- await user_api_key_cache.async_set_cache(
- key=_cache_key,
- value=team_member_info,
- model_type=LiteLLM_TeamMembership,
- ttl=5,
- )
if team_member_info is not None and team_member_info.litellm_budget_table is not None:
team_member_budget: Final = team_member_info.litellm_budget_table.max_budget
@@ -2013,11 +2014,16 @@ async def _user_api_key_auth_builder(
max_budget=team_member_budget,
)
if team_member_spend > team_member_budget:
+ _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}"
raise litellm.BudgetExceededError(
current_cost=team_member_spend,
max_budget=team_member_budget,
+ message=(
+ f"Budget has been exceeded! TeamMember={_entity_id} "
+ f"Current cost: {team_member_spend}, Max budget: {team_member_budget}"
+ ),
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
- entity_id=f"{valid_token.user_id}:{valid_token.team_id}",
+ entity_id=_entity_id,
)
# Check 3. If token is expired
diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
index acdc9728390..fb2ca6372c0 100644
--- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
+++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
@@ -12,6 +12,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
)
if TYPE_CHECKING:
+ from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@@ -30,15 +31,24 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str:
@dataclass(frozen=True, slots=True)
class _CacheInvalidationMessage:
cache_key: str
+ new_value: float | None = None
+ ttl: float | None = None
-def _cache_invalidation_message_json(cache_key: str) -> str:
- return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key)))
+def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str:
+ message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl))
+ return json.dumps({field: value for field, value in message.items() if value is not None})
-def _cache_key_from_message_data(data: object) -> str | None:
+def _finite_number_or_none(value: object) -> float | None:
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ return None
+ return float(value)
+
+
+def _message_from_data(data: object) -> _CacheInvalidationMessage | None:
if isinstance(data, bytes):
- data = data.decode("utf-8", errors="replace")
+ data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str
if not isinstance(data, str):
return None
try:
@@ -48,14 +58,28 @@ def _cache_key_from_message_data(data: object) -> str | None:
if not isinstance(parsed, dict):
return None
cache_key: Final = parsed.get("cache_key")
- return cache_key if isinstance(cache_key, str) else None
+ if not isinstance(cache_key, str):
+ return None
+ return _CacheInvalidationMessage(
+ cache_key=cache_key,
+ new_value=_finite_number_or_none(parsed.get("new_value")),
+ ttl=_finite_number_or_none(parsed.get("ttl")),
+ )
-async def publish_auth_cache_invalidation(cache_key: str) -> None:
+async def publish_auth_cache_invalidation(
+ cache_key: str, new_value: float | None = None, ttl: float | None = None
+) -> None:
"""
Best-effort broadcast so every worker drops its local in-memory copy of a
mutated management object; without this, only the handling worker and Redis
are evicted and other workers keep serving the stale object until its TTL.
+
+ Passing ``new_value`` broadcasts a SET instead of a delete: every subscriber
+ (including the publishing worker's own, which receives its own message)
+ writes the value into its additional in-memory caches rather than deleting
+ the key. A spend reset uses this so the handler's self-delivered message
+ cannot erase the freshly-written post-reset counter or floor marker.
"""
redis_cache: Final = coordination_redis_cache()
if redis_cache is None:
@@ -68,7 +92,10 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None:
cache_key,
)
return
- await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key))
+ await client.publish(
+ auth_cache_invalidation_channel(redis_cache),
+ _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl),
+ )
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
@@ -95,15 +122,17 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
class AuthCacheInvalidationSubscriber:
- __slots__ = ("_redis_cache", "_task", "_user_api_key_cache")
+ __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache")
def __init__(
self,
redis_cache: "RedisCache",
user_api_key_cache: "UserApiKeyCache",
+ additional_in_memory_caches: Sequence["InMemoryCache"] = (),
) -> None:
self._redis_cache = redis_cache
self._user_api_key_cache = user_api_key_cache
+ self._additional_in_memory_caches = tuple(additional_in_memory_caches)
self._task: asyncio.Task[None] | None = None
def start(self) -> None:
@@ -160,12 +189,18 @@ class AuthCacheInvalidationSubscriber:
def _apply_message(self, message: object) -> None:
data: Final = message.get("data") if isinstance(message, dict) else None
- cache_key: Final = _cache_key_from_message_data(data)
- if cache_key is None:
+ parsed: Final = _message_from_data(data)
+ if parsed is None:
+ return
+ if parsed.new_value is not None:
+ for additional_cache in self._additional_in_memory_caches:
+ additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl)
return
in_memory_cache: Final = self._user_api_key_cache.in_memory_cache
if in_memory_cache is not None:
- in_memory_cache.delete_cache(cache_key)
+ in_memory_cache.delete_cache(parsed.cache_key)
+ for additional_cache in self._additional_in_memory_caches:
+ additional_cache.delete_cache(parsed.cache_key)
@staticmethod
async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None:
diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py
index 93d51bdd461..b8df0105b7b 100644
--- a/litellm/proxy/common_utils/user_api_key_cache.py
+++ b/litellm/proxy/common_utils/user_api_key_cache.py
@@ -200,6 +200,21 @@ def end_user_restricted_registry_cache_key() -> str:
return "end_user_restricted_registry"
+def team_membership_auth_cache_key(team_id: str, user_id: str) -> str:
+ """Cache key one team member's ``LiteLLM_TeamMembership`` row is stored under for the admission check."""
+ return f"{team_id}_{user_id}"
+
+
+def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str:
+ """Cache key the pre-call budget reservation stores the same ``LiteLLM_TeamMembership`` row under.
+
+ Deliberately not unified with ``team_membership_auth_cache_key``: the two readers wrote independent
+ keys before this file existed, so a fix that invalidates one must invalidate both explicitly rather
+ than assume a single write is visible to both.
+ """
+ return f"team_membership:{user_id}:{team_id}"
+
+
def get_management_object_ttl(cache: DualCache) -> float:
"""
In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...).
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 01254d5c064..49461d7841d 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -16,7 +16,7 @@ import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
-from typing import Annotated, Final, NamedTuple, Protocol, TypedDict, TypeVar, cast
+from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypedDict, TypeVar, cast
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@@ -56,6 +56,7 @@ from litellm.proxy._types import (
PatchTeamRequest,
ProxyErrorTypes,
ProxyException,
+ ResetSpendRequest,
SpecialManagementEndpointEnums,
SpecialModelNames,
SpecialProxyStrings,
@@ -84,6 +85,7 @@ from litellm.proxy.auth.auth_checks import (
get_team_membership,
get_team_object,
get_user_object,
+ invalidate_team_member_spend_state,
)
from litellm.proxy.auth.auth_utils import (
enforce_batch_enqueued_token_limit_is_admin_only,
@@ -3392,7 +3394,7 @@ async def team_member_update(
Update team member budgets and team member role
"""
- from litellm.proxy.proxy_server import premium_user, prisma_client
+ from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})
@@ -3491,6 +3493,12 @@ async def team_member_update(
budget_patch=budget_patch,
team_default_budget_id=team_default_budget_id,
)
+ if budget_patch:
+ await invalidate_team_member_spend_state(
+ user_id=received_user_id,
+ team_id=data.team_id,
+ user_api_key_cache=user_api_key_cache,
+ )
### update team member role
if data.role is not None:
@@ -3527,6 +3535,125 @@ async def team_member_update(
)
+def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAuth) -> None:
+ """
+ _verify_team_access authorizes a team admin (or org admin) over their own
+ team, with no check that the target user_id differs from the caller. Left
+ unchecked, that admin could target their own LiteLLM_TeamMembership row and
+ repeatedly reset it to 0 right before it crosses their per-member cap,
+ consuming the shared team budget without the configured limit ever binding.
+ Only a proxy admin may reset an admin's own spend.
+ """
+ if user_id == user_api_key_dict.user_id and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
+ _raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.")
+
+
+def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn:
+ detail: Final = {"error": message} # mutable-ok: HTTPException.detail takes a dict
+ raise HTTPException(status_code=status_code, detail=detail)
+
+
+def _validate_team_member_reset_spend_value(
+ reset_to: object,
+ membership: LiteLLM_TeamMembership,
+) -> float:
+ if not isinstance(reset_to, (int, float)):
+ _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a float")
+
+ reset_to_float: Final = float(reset_to)
+ if not math.isfinite(reset_to_float) or reset_to_float < 0:
+ _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a finite number >= 0")
+
+ current_spend: Final = membership.spend or 0.0
+ if reset_to_float > current_spend:
+ _raise_reset_spend_error(
+ status.HTTP_400_BAD_REQUEST,
+ f"reset_to ({reset_to_float}) must be <= current spend ({current_spend})",
+ )
+
+ max_budget: Final = membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None
+ if max_budget is not None and reset_to_float > max_budget:
+ _raise_reset_spend_error(
+ status.HTTP_400_BAD_REQUEST,
+ f"reset_to ({reset_to_float}) must be <= budget ({max_budget})",
+ )
+
+ return reset_to_float
+
+
+@router.post(
+ "/team/{team_id}/member/{user_id}/reset_spend",
+ tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence
+ dependencies=(Depends(user_api_key_auth),),
+)
+@management_endpoint_wrapper
+async def reset_team_member_spend_fn(
+ team_id: str,
+ user_id: str,
+ data: ResetSpendRequest,
+ user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+):
+ """
+ Reset a team member's tracked spend against their per-member budget.
+
+ A member's spend is tracked separately from both their own personal
+ budget and the team's own budget (LiteLLM_TeamMembership.spend), so
+ neither /user/update nor /team/update can clear it: this is the only
+ endpoint that does. The cross-pod spend counter and cached membership
+ reads are invalidated so the reset takes effect on the member's next
+ request rather than waiting on the membership cache's TTL.
+ """
+ from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
+
+ if prisma_client is None:
+ _raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None")
+
+ team_obj: Final = await get_team_object(
+ team_id=team_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=None,
+ proxy_logging_obj=proxy_logging_obj,
+ check_db_only=True,
+ )
+ await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict)
+ _check_not_resetting_own_spend(user_id=user_id, user_api_key_dict=user_api_key_dict)
+
+ membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument
+ "user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument
+ }
+ _membership_row: Final = await _team_membership_db(prisma_client).find_unique(
+ where=membership_where,
+ include={"litellm_budget_table": True}, # mutable-ok: prisma client requires a plain dict include= argument
+ )
+ if _membership_row is None:
+ _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.")
+ membership: Final = LiteLLM_TeamMembership.model_validate(_membership_row.model_dump())
+
+ current_spend: Final = membership.spend or 0.0
+ reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership)
+
+ await _team_membership_db(prisma_client).update(
+ where=membership_where,
+ data={"spend": reset_to}, # mutable-ok: prisma client requires a plain dict data= argument
+ )
+
+ await invalidate_team_member_spend_state(
+ user_id=user_id,
+ team_id=team_id,
+ user_api_key_cache=user_api_key_cache,
+ new_spend=reset_to,
+ )
+
+ return { # mutable-ok: matches this router's established untyped-response-dict convention
+ "team_id": team_id,
+ "user_id": user_id,
+ "spend": reset_to,
+ "previous_spend": current_spend,
+ "max_budget": membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None,
+ }
+
+
def _create_results_from_response(
members: list[Member],
response: TeamAddMemberResponse,
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 7dced4e26b6..0abcdeaf3f6 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -2555,6 +2555,12 @@ async def _authoritative_floor_spend(
if db_spend is None:
return None
+ # a spend reset that committed during the DB read above wrote the post-reset
+ # floor to the marker; keep it over this read's now-stale pre-commit value
+ rechecked: Final = spend_counter_cache.in_memory_cache.get_cache(key=marker_key)
+ if rechecked is not None:
+ return float(rechecked)
+
spend_counter_cache.in_memory_cache.set_cache(
key=marker_key,
value=db_spend,
@@ -6798,6 +6804,7 @@ class ProxyConfig:
subscriber: Final = AuthCacheInvalidationSubscriber(
redis_cache=redis_cache,
user_api_key_cache=user_api_key_cache,
+ additional_in_memory_caches=(spend_counter_cache.in_memory_cache,),
)
self.auth_cache_invalidation_subscriber = subscriber
subscriber.start()
diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py
index ce6c9330620..149f9b960a1 100644
--- a/litellm/proxy/spend_tracking/budget_reservation.py
+++ b/litellm/proxy/spend_tracking/budget_reservation.py
@@ -25,7 +25,11 @@ from litellm.proxy._types import (
from litellm.proxy.auth.auth_utils import get_model_from_request
from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded
from litellm.proxy.auth.route_checks import RouteChecks
-from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key
+from litellm.proxy.common_utils.user_api_key_cache import (
+ end_user_cache_key,
+ tag_cache_key,
+ team_membership_reservation_cache_key,
+)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@@ -546,7 +550,9 @@ async def _get_team_member_budget_counter(
if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None:
return None
- membership_cache_key: Final = f"team_membership:{valid_token.user_id}:{team_object.team_id}"
+ membership_cache_key: Final = team_membership_reservation_cache_key(
+ user_id=valid_token.user_id, team_id=team_object.team_id
+ )
cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key)
team_membership: LiteLLM_TeamMembership | None = None
if isinstance(cached_team_membership, LiteLLM_TeamMembership):
diff --git a/tests/proxy_behavior/management/test_team_member_reset_spend.py b/tests/proxy_behavior/management/test_team_member_reset_spend.py
new file mode 100644
index 00000000000..ec2c78139fe
--- /dev/null
+++ b/tests/proxy_behavior/management/test_team_member_reset_spend.py
@@ -0,0 +1,152 @@
+import uuid
+
+import pytest
+
+from .actors import Actor
+from .conftest import create_scratch_team
+
+pytestmark = pytest.mark.asyncio(loop_scope="session")
+
+_SEED_SPEND = 5.0
+_RESET_TO = 2.0
+
+
+# POST /team/{team_id}/member/{user_id}/reset_spend. The handler gate is
+# _verify_team_access (proxy admin / team admin of this team / org admin of
+# the team's org) — the same gate /team/member_update uses, so this mirrors
+# that file's matrix exactly.
+_MATRIX = [
+ ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
+ ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
+ ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200),
+ ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
+ ("alpha/owner", Actor.OWNER, "alpha", 403),
+ ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
+ ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
+ ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
+ ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
+ ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
+ ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
+ ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
+ ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
+]
+
+
+async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None:
+ if shape == "alpha":
+ await create_scratch_team(
+ prisma,
+ team_id,
+ organization_id=world.org_a_id,
+ admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
+ )
+ elif shape == "beta":
+ await create_scratch_team(prisma, team_id, organization_id=world.org_b_id)
+ else: # pragma: no cover - guard
+ pytest.fail(f"unknown shape={shape}")
+ await prisma.db.litellm_teammembership.create(
+ data={"user_id": member_id, "team_id": team_id, "spend": _SEED_SPEND}
+ )
+
+
+@pytest.mark.parametrize(
+ "actor,shape,expected_status",
+ [(a, sh, s) for (_id, a, sh, s) in _MATRIX],
+ ids=[s[0] for s in _MATRIX],
+)
+async def test_team_member_reset_spend_authz_matrix(
+ actor: Actor,
+ shape: str,
+ expected_status: int,
+ proxy_client,
+ prisma,
+ scratch,
+ world,
+):
+ member_id = scratch.tag("member")
+ await _seed_target(prisma, world, shape, scratch.prefix, member_id)
+ caller = world.keys[actor]
+
+ resp = await proxy_client.post(
+ f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
+ headers={"Authorization": f"Bearer {caller.cleartext}"},
+ json={"reset_to": _RESET_TO},
+ )
+ assert (
+ resp.status_code == expected_status
+ ), f"{actor.value} {shape}: {resp.status_code} {resp.text}"
+
+ row = await prisma.db.litellm_teammembership.find_unique(
+ where={"user_id_team_id": {"user_id": member_id, "team_id": scratch.prefix}}
+ )
+ assert row is not None
+ if expected_status == 200:
+ assert row.spend == _RESET_TO
+ else:
+ assert row.spend == _SEED_SPEND, "denied but spend reset"
+
+
+async def test_team_member_reset_spend_missing_team_is_404(proxy_client, world):
+ resp = await proxy_client.post(
+ f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_spend",
+ headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
+ json={"reset_to": 0.0},
+ )
+ assert resp.status_code == 404, resp.text
+
+
+async def test_team_member_reset_spend_missing_membership_is_404(
+ proxy_client, prisma, scratch, world
+):
+ """A well-formed team but a user_id with no LiteLLM_TeamMembership row is 404."""
+ await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
+ resp = await proxy_client.post(
+ f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_spend",
+ headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
+ json={"reset_to": 0.0},
+ )
+ assert resp.status_code == 404, resp.text
+
+
+async def test_team_member_reset_spend_above_current_spend_is_400(
+ proxy_client, prisma, scratch, world
+):
+ member_id = scratch.tag("member")
+ await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
+ await prisma.db.litellm_teammembership.create(
+ data={"user_id": member_id, "team_id": scratch.prefix, "spend": 1.0}
+ )
+ resp = await proxy_client.post(
+ f"/team/{scratch.prefix}/member/{member_id}/reset_spend",
+ headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
+ json={"reset_to": 5.0},
+ )
+ assert resp.status_code == 400, resp.text
+
+
+async def test_team_member_reset_spend_team_admin_cannot_reset_own_spend(
+ proxy_client, prisma, scratch, world
+):
+ """A team admin targeting their own LiteLLM_TeamMembership row is 403: unchecked, an
+ admin could repeatedly zero their own spend right before it crosses their per-member
+ cap, consuming the shared team budget without the configured limit ever binding."""
+ team_admin = world.keys[Actor.TEAM_ADMIN]
+ await create_scratch_team(
+ prisma,
+ scratch.prefix,
+ organization_id=world.org_a_id,
+ admin_user_ids=[team_admin.user_id],
+ )
+ await prisma.db.litellm_teammembership.create(
+ data={"user_id": team_admin.user_id, "team_id": scratch.prefix, "spend": _SEED_SPEND}
+ )
+ resp = await proxy_client.post(
+ f"/team/{scratch.prefix}/member/{team_admin.user_id}/reset_spend",
+ headers={"Authorization": f"Bearer {team_admin.cleartext}"},
+ json={"reset_to": 0.0},
+ )
+ assert resp.status_code == 403, resp.text
+ row = await prisma.db.litellm_teammembership.find_unique(
+ where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": scratch.prefix}}
+ )
+ assert row is not None and row.spend == _SEED_SPEND, "denied but spend reset"
diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py
index 04f38b5e2ed..abe73f7d05c 100644
--- a/tests/test_litellm/proxy/auth/test_auth_checks.py
+++ b/tests/test_litellm/proxy/auth/test_auth_checks.py
@@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import (
_virtual_key_soft_budget_check,
get_key_object,
get_user_object,
+ invalidate_team_member_spend_state,
vector_store_access_check,
)
from litellm.caching.in_memory_cache import InMemoryCache
@@ -6939,3 +6940,307 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is
router = _router_with_a_group_priced_through_model_info()
assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_clears_both_membership_cache_keys():
+ """A team-member budget reset (new_spend passed) must SET the spend counter to the reset
+ value, clear its DB-floor marker, AND invalidate both independently-keyed membership caches
+ (user_api_key_auth.py's admission check writes one key format, budget_reservation.py and
+ auth_checks.py's own get_team_membership() write the other) or a stale read keeps 429ing
+ after the reset. Asserted against real cache reads, not mock call args, so a change that
+ keeps the call but drops its effect still fails."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership")
+ await real_cache.async_set_cache(key="team_membership:user-1:team-1", value="stale-membership")
+
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0)
+ real_spend_counter_cache.in_memory_cache.set_cache(
+ key="spend_db_floor:spend:team_member:user-1:team-1", value=999.0
+ )
+
+ with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ new_spend=0.0,
+ )
+
+ assert await real_cache.async_get_cache(key="team-1_user-1") is None
+ assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0
+ assert (
+ real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1")
+ == 0.0
+ ), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up"
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_leaves_the_live_spend_counter_alone_without_new_spend():
+ """team_member_update only changes the budget cap, not the tracked spend, so it calls
+ invalidate_team_member_spend_state with no new_spend. Deleting the live spend counter in that
+ case would force the next read to reseed from the DB's own spend column, which lags the live
+ counter via periodic batch writes, briefly UNDER-enforcing the raised cap against a spend
+ value lower than what was actually tracked (regression: PR #37971 Bugbot finding)."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership")
+
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0)
+
+ with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ )
+
+ assert await real_cache.async_get_cache(key="team-1_user-1") is None
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 999.0
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_deleting():
+ """/key/{key}/reset_spend SETs its counter to the reset value rather than deleting it, so a
+ worker's next read reflects it directly instead of falling back through a DB reseed. A reset
+ caller passing new_spend must match that precedent, not merely delete the counter."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ real_spend_counter_cache = DualCache()
+ fake_redis_cache = MagicMock()
+ fake_redis_cache.async_set_cache = AsyncMock()
+ real_spend_counter_cache.redis_cache = fake_redis_cache
+
+ with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ new_spend=2.5,
+ )
+
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5
+ fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60)
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_deletes_redis_counter_when_set_fails(): # test-quality-ok: only observable effect is the fallback call on the same fake client
+ """Redis reads take priority over the local in-memory copy (get_current_spend reads Redis
+ first), so a failed Redis SET would otherwise leave the OLD pre-reset value authoritative
+ for every worker even though the reset reported success. On a failed SET, the stale Redis
+ entry must be deleted instead, so the next read clean-misses and reseeds from the DB."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ real_spend_counter_cache = DualCache()
+ fake_redis_cache = MagicMock()
+ fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down"))
+ fake_redis_cache.async_delete_cache = AsyncMock()
+ real_spend_counter_cache.redis_cache = fake_redis_cache
+
+ with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ new_spend=2.5,
+ )
+
+ fake_redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1")
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_raises_503_when_both_redis_writes_fail():
+ """If the Redis SET fails AND the fallback DELETE fails, the stale pre-reset counter is still
+ authoritative in Redis for every worker. Reporting success would silently keep 429ing the
+ member, so the reset must surface a 503 instead (regression: PR #37971 Greptile finding)."""
+ from fastapi import HTTPException
+
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ real_cache = UserApiKeyCache()
+ real_spend_counter_cache = DualCache()
+ fake_redis_cache = MagicMock()
+ fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down"))
+ fake_redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis still down"))
+ real_spend_counter_cache.redis_cache = fake_redis_cache
+
+ with (
+ patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
+ ),
+ pytest.raises(HTTPException) as exc_info,
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=real_cache,
+ new_spend=2.5,
+ )
+
+ assert exc_info.value.status_code == 503
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_to_remote_workers():
+ """The test above only proves the handling worker's own spend counter is
+ cleared. A remote worker's spend counter is a separate DualCache instance;
+ if the reset never reaches it, that worker keeps enforcing the pre-reset
+ spend the moment its own Redis read for the counter fails and it falls
+ back to its own (now-stale) in-memory copy. Drives the actual message
+ published onto the invalidation channel through a second, independent
+ AuthCacheInvalidationSubscriber standing in for that remote worker, rather
+ than asserting on the publish call args."""
+ from redis.asyncio import Redis
+
+ from litellm.caching.dual_cache import DualCache
+ from litellm.caching.in_memory_cache import InMemoryCache
+ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ published: list[tuple[str, str]] = []
+
+ class _RecordingRedisClient(Redis):
+ def __init__(self) -> None:
+ pass
+
+ async def publish(self, channel: str, message: str) -> int:
+ published.append((channel, message))
+ return 1
+
+ class _FakeRedisCache:
+ def __init__(self) -> None:
+ self.namespace = None
+
+ def init_async_client(self) -> object:
+ return _RecordingRedisClient()
+
+ local_spend_counter_cache = DualCache()
+
+ remote_user_api_key_cache = UserApiKeyCache()
+ remote_spend_counter_in_memory_cache = InMemoryCache()
+ remote_spend_counter_in_memory_cache.set_cache("spend:team_member:user-1:team-1", 999.0)
+ remote_spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:team_member:user-1:team-1", 999.0)
+
+ with (
+ patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache
+ ),
+ patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test
+ "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
+ return_value=_FakeRedisCache(),
+ ),
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=UserApiKeyCache(),
+ new_spend=0.0,
+ )
+
+ def _published_message_for(cache_key: str) -> str:
+ matches = [message for _, message in published if json.loads(message)["cache_key"] == cache_key]
+ assert matches, f"{cache_key} never reached the cross-worker invalidation channel"
+ return matches[-1]
+
+ remote_subscriber = AuthCacheInvalidationSubscriber(
+ redis_cache=_FakeRedisCache(),
+ user_api_key_cache=remote_user_api_key_cache,
+ additional_in_memory_caches=(remote_spend_counter_in_memory_cache,),
+ )
+ for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"):
+ remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
+ {"type": "message", "data": _published_message_for(cache_key)}
+ )
+
+ assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0
+ assert (
+ remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0
+ ), "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor"
+
+
+@pytest.mark.asyncio
+async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_not_erase_the_reset():
+ """The handling worker subscribes to the same invalidation channel it publishes on, so it
+ receives its own reset message. A delete-style broadcast would erase the post-reset counter
+ and floor marker the handler just wrote, reopening the stale-floor race the reset closed
+ (regression: PR #37971 Greptile finding). The broadcast carries the reset value as a SET, so
+ applying the self-delivered message must leave both keys at the post-reset value."""
+ from redis.asyncio import Redis
+
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ published: list[tuple[str, str]] = []
+
+ class _RecordingRedisClient(Redis):
+ def __init__(self) -> None:
+ pass
+
+ async def publish(self, channel: str, message: str) -> int:
+ published.append((channel, message))
+ return 1
+
+ class _FakeRedisCache:
+ def __init__(self) -> None:
+ self.namespace = None
+
+ def init_async_client(self) -> object:
+ return _RecordingRedisClient()
+
+ local_spend_counter_cache = DualCache()
+ local_user_api_key_cache = UserApiKeyCache()
+
+ with (
+ patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache
+ ),
+ patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test
+ "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
+ return_value=_FakeRedisCache(),
+ ),
+ ):
+ await invalidate_team_member_spend_state(
+ user_id="user-1",
+ team_id="team-1",
+ user_api_key_cache=local_user_api_key_cache,
+ new_spend=0.0,
+ )
+
+ own_subscriber = AuthCacheInvalidationSubscriber(
+ redis_cache=_FakeRedisCache(),
+ user_api_key_cache=local_user_api_key_cache,
+ additional_in_memory_caches=(local_spend_counter_cache.in_memory_cache,),
+ )
+ for _, message in published:
+ own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
+ {"type": "message", "data": message}
+ )
+
+ assert local_spend_counter_cache.in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0, (
+ "the handler's self-delivered broadcast erased the post-reset spend counter"
+ )
+ assert (
+ local_spend_counter_cache.in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0
+ ), "the handler's self-delivered broadcast erased the post-reset floor marker, reopening the stale-floor race"
diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
index 468e8aabae8..7d5fc1a3544 100644
--- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
+++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
@@ -6,6 +6,7 @@ from unittest.mock import patch
import pytest
from redis.asyncio import Redis
+from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
AUTH_CACHE_INVALIDATION_CHANNEL,
AuthCacheInvalidationSubscriber,
@@ -144,6 +145,37 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None:
assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL]
+@pytest.mark.asyncio
+async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None:
+ """
+ The spend-counter half of the same cross-worker gap: a remote worker's own
+ spend counter can hold a stale value (its fallback path when that worker's
+ own Redis read for the counter fails), and only clearing user_api_key_cache
+ on message would leave that separate DualCache's in-memory copy untouched.
+ """
+ cache = UserApiKeyCache()
+ spend_counter_in_memory_cache = InMemoryCache()
+ spend_counter_in_memory_cache.set_cache("spend:team_member:u-1:t-1", 999.0)
+ assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is not None
+
+ pubsub = _QueuePubSub(initial_messages=[_invalidation_message("spend:team_member:u-1:t-1")])
+ subscriber = AuthCacheInvalidationSubscriber(
+ redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])),
+ user_api_key_cache=cache,
+ additional_in_memory_caches=(spend_counter_in_memory_cache,),
+ )
+ subscriber.start()
+ try:
+ for _ in range(200):
+ if spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None:
+ break
+ await asyncio.sleep(0.01)
+ finally:
+ await subscriber.stop()
+
+ assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None
+
+
@pytest.mark.asyncio
async def test_subscriber_ignores_malformed_messages() -> None:
cache = UserApiKeyCache()
diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
index f6d74a189bc..7f5d3eb0a14 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
@@ -9,11 +9,13 @@ from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
+from pydantic import ValidationError
from litellm._uuid import uuid
from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth
from litellm.proxy._types import (
+ LiteLLM_BudgetTable,
LiteLLM_BudgetTableFull,
LiteLLM_ModelTable,
LiteLLM_OrganizationMembershipTable,
@@ -27,7 +29,9 @@ from litellm.proxy._types import (
Member,
ProxyErrorTypes,
ProxyException,
+ ResetSpendRequest,
TeamMemberAddRequest,
+ TeamMemberUpdateRequest,
UpdateTeamRequest,
)
from litellm.proxy.management_endpoints.team_endpoints import (
@@ -42,12 +46,15 @@ from litellm.proxy.management_endpoints.team_endpoints import (
_transform_teams_to_deleted_records,
_update_model_table,
_validate_and_populate_member_user_info,
+ _validate_team_member_reset_spend_value,
_verify_team_access,
delete_team,
list_available_teams,
+ reset_team_member_spend_fn,
router,
team_member_add_duplication_check,
team_member_delete,
+ team_member_update,
update_team,
validate_team_org_change,
)
@@ -12603,3 +12610,376 @@ async def test_invalidate_access_group_cache_deletes_the_cached_object():
"user_api_key_cache": cache,
"proxy_logging_obj": logging_obj,
}
+
+
+def test_validate_team_member_reset_spend_value_rejects_non_numeric():
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to="not-a-number",
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert exc.value.status_code == 400
+
+
+def test_validate_team_member_reset_spend_value_rejects_negative():
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to=-1.0,
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert exc.value.status_code == 400
+
+
+@pytest.mark.parametrize("reset_to", [float("nan"), float("inf"), float("-inf")])
+def test_validate_team_member_reset_spend_value_rejects_non_finite(reset_to):
+ """NaN and +/-inf are instances of float and compare False against every bound
+ below (`nan < 0`, `nan > current_spend` are both False), so an isinstance-and-range
+ check alone lets them through to persist as the member's spend and silently
+ disable every later budget comparison against it."""
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to=reset_to,
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert exc.value.status_code == 400
+
+
+@pytest.mark.parametrize("reset_to", [True, False])
+def test_reset_spend_request_rejects_bool_reset_to(reset_to):
+ """bool is a subclass of int, so pydantic silently coerces True/False into 1.0/0.0 for a
+ ``float`` field: {"reset_to": true} would otherwise reach _validate_team_member_reset_spend_value
+ as an indistinguishable 1.0 and reset the member's spend instead of failing the request."""
+ with pytest.raises(ValidationError):
+ ResetSpendRequest(reset_to=reset_to)
+
+
+def test_validate_team_member_reset_spend_value_rejects_above_current_spend():
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to=20.0,
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert exc.value.status_code == 400
+
+
+def test_validate_team_member_reset_spend_value_rejects_above_max_budget():
+ with pytest.raises(HTTPException) as exc:
+ _validate_team_member_reset_spend_value(
+ reset_to=10.0,
+ membership=LiteLLM_TeamMembership(
+ user_id="u1",
+ team_id="t1",
+ spend=10.0,
+ litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=5.0),
+ ),
+ )
+ assert exc.value.status_code == 400
+
+
+def test_validate_team_member_reset_spend_value_accepts_valid_reset():
+ result = _validate_team_member_reset_spend_value(
+ reset_to=0.0,
+ membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0),
+ )
+ assert result == 0.0
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_success(monkeypatch):
+ """A proxy admin resetting a stuck team member's spend must write the DB
+ row to reset_to AND invalidate the cached spend/membership state, or the
+ 429 the endpoint exists to clear keeps firing off the stale cache.
+ Asserted against real cache reads, not mock call args, so a change that
+ keeps the call but drops its effect still fails."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ mock_prisma_client = MagicMock()
+ mock_proxy_logging_obj = MagicMock()
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
+ await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
+
+ membership_row = LiteLLM_TeamMembership(
+ user_id="member-1",
+ team_id="team-1",
+ spend=10.0,
+ litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0),
+ )
+ mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
+ mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj)
+ monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
+ ):
+ response = await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="member-1",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+
+ assert response["spend"] == 0.0
+ assert response["previous_spend"] == 10.0
+ assert response["max_budget"] == 50.0
+ mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with(
+ where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}},
+ data={"spend": 0.0},
+ )
+ assert await real_cache.async_get_cache(key="team-1_member-1") is None
+ assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch):
+ mock_prisma_client = MagicMock()
+ mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
+ ):
+ with pytest.raises(HTTPException) as exc:
+ await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="ghost-user",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+ assert exc.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_team_not_found(monkeypatch):
+ mock_prisma_client = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})),
+ ):
+ with pytest.raises(HTTPException) as exc:
+ await reset_team_member_spend_fn(
+ team_id="ghost-team",
+ user_id="member-1",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+ assert exc.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_forbidden_for_non_admin(monkeypatch):
+ """A caller who is neither proxy admin, org admin, nor this team's admin must be refused,
+ matching every other team-mutating endpoint's authorization."""
+ mock_prisma_client = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])),
+ ):
+ with pytest.raises(HTTPException) as exc:
+ await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="member-1",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user"
+ ),
+ )
+ assert exc.value.status_code == 403
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_team_admin_cannot_reset_own_spend(monkeypatch):
+ """_verify_team_access authorizes a team admin over their own team with no check that the
+ target differs from the caller. Unchecked, that admin could target their own membership row
+ and repeatedly zero it right before it crosses their per-member cap, consuming the shared
+ team budget without the configured limit ever binding (Veria finding on PR #37971)."""
+ mock_prisma_client = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-admin", user_id="team-admin-1")
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(
+ return_value=LiteLLM_TeamTable(
+ team_id="team-1",
+ members_with_roles=[Member(user_id="team-admin-1", role="admin")],
+ )
+ ),
+ ):
+ with pytest.raises(HTTPException) as exc:
+ await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="team-admin-1",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=team_admin,
+ )
+ assert exc.value.status_code == 403
+ mock_prisma_client.db.litellm_teammembership.update.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkeypatch):
+ """The self-reset guard is scoped to non-proxy-admin roles: a proxy admin resetting their
+ own membership spend is the platform-wide trust boundary, not a team-scoped one."""
+ mock_prisma_client = MagicMock()
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+
+ membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0)
+ mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
+ mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
+
+ with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+ AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
+ ):
+ response = await reset_team_member_spend_fn(
+ team_id="team-1",
+ user_id="admin-user",
+ data=ResetSpendRequest(reset_to=0.0),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+ assert response["spend"] == 0.0
+
+
+@pytest.mark.asyncio
+async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch):
+ """Raising a stuck member's max_budget_in_team via the documented /team/member_update
+ endpoint must invalidate the cached membership state, or the raised cap never reaches the
+ admission check and the member stays 429ing. The live spend counter itself must be left
+ untouched: only the cap changed, and deleting the counter would force a reseed from the
+ DB's own spend column, which lags the live counter via periodic batch writes, briefly
+ UNDER-enforcing the raised cap against a spend value lower than what was actually tracked.
+ Asserted against real cache reads, not mock call args, so a change that keeps the call but
+ drops its effect still fails."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ mock_prisma_client = MagicMock()
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
+ await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
+
+ team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[])
+ team_info_response = {
+ "team_info": team_row,
+ "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)],
+ }
+
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
+ monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
+
+ mock_tx = AsyncMock()
+ mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
+ mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None)
+
+ with (
+ patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.team_info",
+ AsyncMock(return_value=team_info_response),
+ ),
+ patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership",
+ AsyncMock(),
+ ),
+ ):
+ await team_member_update(
+ data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", max_budget_in_team=999999.0),
+ http_request=MagicMock(),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+
+ assert await real_cache.async_get_cache(key="team-1_member-1") is None
+ assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0
+
+
+@pytest.mark.asyncio
+async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(monkeypatch):
+ """A role-only update carries an empty budget_patch and touches no budget state,
+ so the member's cached spend/membership state must be left untouched."""
+ from litellm.caching.dual_cache import DualCache
+ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
+
+ mock_prisma_client = MagicMock()
+ real_cache = UserApiKeyCache()
+ await real_cache.async_set_cache(key="team-1_member-1", value="still-fresh-membership")
+ real_spend_counter_cache = DualCache()
+ real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=1.5)
+
+ team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[])
+ team_info_response = {
+ "team_info": team_row,
+ "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)],
+ }
+
+ mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
+ monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
+
+ mock_tx = AsyncMock()
+ mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
+ mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None)
+
+ with (
+ patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints.team_info",
+ AsyncMock(return_value=team_info_response),
+ ),
+ patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
+ "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership",
+ AsyncMock(),
+ ),
+ ):
+ await team_member_update(
+ data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1"),
+ http_request=MagicMock(),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
+ ),
+ )
+
+ assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership"
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 31d2a6cef98..3383527e932 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -11347,3 +11347,38 @@ class TestRouterModelNameOnStreamingChunks:
assert len(frames) >= 3
assert '"router_model_name":"deep-model"' in frames[0]
assert all('"router_model_name":"backup-tier"' in frame for frame in frames[1:])
+
+
+@pytest.mark.asyncio
+async def test_authoritative_floor_spend_keeps_a_reset_marker_written_during_the_db_read():
+ """A team-member spend reset writes the post-reset floor to the spend_db_floor marker
+ (auth_checks.invalidate_team_member_spend_state). A floor read already in flight when the
+ reset commits would otherwise cache its stale pre-reset DB value over the fresh marker,
+ letting a budget check raise the counter right back above the just-reset spend
+ (regression: PR #37971 Greptile finding)."""
+ from litellm.proxy.proxy_server import _authoritative_floor_spend
+
+ real_spend_counter_cache = DualCache()
+ counter_key = "spend:team_member:user-1:team-1"
+ marker_key = f"spend_db_floor:{counter_key}"
+
+ async def db_read_racing_with_a_reset(prisma_client, counter_key):
+ real_spend_counter_cache.in_memory_cache.set_cache(key=marker_key, value=0.0)
+ return 999.0
+
+ with (
+ patch.object( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
+ proxy_server_module, "spend_counter_cache", real_spend_counter_cache
+ ),
+ patch.object( # test-quality-ok: the DB read must race the reset; no injectable seam for module-global prisma reads
+ proxy_server_module.SpendCounterReseed,
+ "from_db",
+ AsyncMock(side_effect=db_read_racing_with_a_reset),
+ ),
+ ):
+ result = await _authoritative_floor_spend(counter_key=counter_key)
+
+ assert result == 0.0
+ assert real_spend_counter_cache.in_memory_cache.get_cache(key=marker_key) == 0.0, (
+ "the in-flight DB read clobbered the post-reset floor marker with the stale pre-reset value"
+ )
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 78329a1e53c..d51bff27784 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -14936,6 +14936,33 @@ export interface paths {
patch?: never;
trace?: never;
};
+ "/team/{team_id}/member/{user_id}/reset_spend": {
+ parameters: {
+ query?: never;
+ header?: never;
+ path?: never;
+ cookie?: never;
+ };
+ get?: never;
+ put?: never;
+ /**
+ * Reset Team Member Spend Fn
+ * @description Reset a team member's tracked spend against their per-member budget.
+ *
+ * A member's spend is tracked separately from both their own personal
+ * budget and the team's own budget (LiteLLM_TeamMembership.spend), so
+ * neither /user/update nor /team/update can clear it: this is the only
+ * endpoint that does. The cross-pod spend counter and cached membership
+ * reads are invalidated so the reset takes effect on the member's next
+ * request rather than waiting on the membership cache's TTL.
+ */
+ post: operations["reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post"];
+ delete?: never;
+ options?: never;
+ head?: never;
+ patch?: never;
+ trace?: never;
+ };
"/team/{team_id}/members/me": {
parameters: {
query?: never;
@@ -55083,6 +55110,42 @@ export interface operations {
};
};
};
+ reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post: {
+ parameters: {
+ query?: never;
+ header?: never;
+ path: {
+ team_id: string;
+ user_id: string;
+ };
+ cookie?: never;
+ };
+ requestBody: {
+ content: {
+ "application/json": components["schemas"]["ResetSpendRequest"];
+ };
+ };
+ responses: {
+ /** @description Successful Response */
+ 200: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": unknown;
+ };
+ };
+ /** @description Validation Error */
+ 422: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": components["schemas"]["HTTPValidationError"];
+ };
+ };
+ };
+ };
team_member_me_team__team_id__members_me_get: {
parameters: {
query?: never;
From 9224b2ce5d3e0327bcc0b02e6f111794574716d6 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 09:50:23 -0700
Subject: [PATCH 26/70] fix(router): freeze reasoning effort flag mappings
---
.../router_utils/reasoning_effort_capability.py | 16 ++++++++++------
1 file changed, 10 insertions(+), 6 deletions(-)
diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py
index d20b1151803..3e4478e5e21 100644
--- a/litellm/router_utils/reasoning_effort_capability.py
+++ b/litellm/router_utils/reasoning_effort_capability.py
@@ -25,12 +25,14 @@ it above.
"""
from collections.abc import Mapping, Sequence
+from types import MappingProxyType
from typing import Final, get_args
import litellm
from litellm.types.llms.openai import REASONING_EFFORT
REASONING_EFFORT_ADVERTISEMENT_ORDER: Final = get_args(REASONING_EFFORT)
+_EMPTY_ENTRY: Final[Mapping[str, object]] = MappingProxyType({})
_EFFORT_FLAGS: Final = (
("none", "supports_none_reasoning_effort"),
@@ -52,17 +54,19 @@ def _bare_model_entry(model_info: Mapping[str, object]) -> Mapping[str, object]:
key: Final = model_info.get("key")
provider: Final = model_info.get("litellm_provider")
if not isinstance(key, str) or not isinstance(provider, str) or not key.startswith(f"{provider}/"):
- return {}
+ return _EMPTY_ENTRY
entry: Final[Mapping[str, object] | None] = litellm.model_cost.get(key.removeprefix(f"{provider}/"))
- return entry if entry is not None else {}
+ return entry if entry is not None else _EMPTY_ENTRY
def _declared_effort_flags(model_info: Mapping[str, object]) -> Mapping[str, object]:
bare: Final = _bare_model_entry(model_info)
- return {
- effort: model_info.get(flag) if model_info.get(flag) is not None else bare.get(flag)
- for effort, flag in _EFFORT_FLAGS
- }
+ return MappingProxyType(
+ {
+ effort: model_info.get(flag) if model_info.get(flag) is not None else bare.get(flag)
+ for effort, flag in _EFFORT_FLAGS
+ }
+ )
def _supports_none_reasoning_effort(model_info: Mapping[str, object], flag: object) -> bool:
From f583151a5b8928361237e715abe75b305fc4b3a5 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 09:53:19 -0700
Subject: [PATCH 27/70] fix(model_prices): raise bedrock_mantle gpt-5.6
max_input_tokens to Mantle's enforced 1050000
---
...odel_prices_and_context_window_backup.json | 9 ++--
model_prices_and_context_window.json | 9 ++--
.../llm_cost_calc/test_llm_cost_calc_utils.py | 4 +-
...bedrock_mantle_responses_transformation.py | 44 ++++++++++++++++++-
4 files changed, 56 insertions(+), 10 deletions(-)
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 9f953e11df1..1aa4c7cd060 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -49016,12 +49016,13 @@
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@@ -49048,12 +49049,13 @@
"output_cost_per_token": 1.32e-05,
"output_cost_per_token_above_272k_tokens": 1.98e-05,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@@ -49080,12 +49082,13 @@
"output_cost_per_token": 1.32e-06,
"output_cost_per_token_above_272k_tokens": 1.98e-06,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 9f953e11df1..1aa4c7cd060 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -49016,12 +49016,13 @@
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@@ -49048,12 +49049,13 @@
"output_cost_per_token": 1.32e-05,
"output_cost_per_token_above_272k_tokens": 1.98e-05,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
@@ -49080,12 +49082,13 @@
"output_cost_per_token": 1.32e-06,
"output_cost_per_token_above_272k_tokens": 1.98e-06,
"litellm_provider": "bedrock_mantle",
- "max_input_tokens": 1000000,
+ "max_input_tokens": 1050000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "responses",
"use_openai_responses_path": true,
"supported_endpoints": [
+ "/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
index c8c36032793..6f513ce1bd4 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
@@ -478,10 +478,10 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_m
],
)
def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model):
- """Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K."""
+ """Bedrock GPT-5.6 enforces a 1,050,000-token context window, billed at the long-context rates above 272K."""
model_cost_map = litellm.model_cost[model]
- assert model_cost_map["max_input_tokens"] == 1000000
+ assert model_cost_map["max_input_tokens"] == 1050000
cached_tokens = 100000
completion_tokens = 1000
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
index 9e05d48a18f..fd279a2bc1f 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
@@ -8,7 +8,8 @@ gate, the URL construction for both paths, and the shared Bearer auth.
"""
import copy
-
+import json
+from pathlib import Path
import pytest
from botocore.exceptions import (
@@ -1523,7 +1524,7 @@ class TestBedrockMantleResponsesPricing:
assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost)
assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost)
assert info["output_cost_per_token"] == pytest.approx(output_cost)
- assert info["max_input_tokens"] == 1000000
+ assert info["max_input_tokens"] == 1050000
assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2)
assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2)
assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2)
@@ -1565,3 +1566,42 @@ class TestBedrockMantleResponsesPricing:
def test_models_registered(self, local_cost_map):
assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models
+
+
+def _repo_cost_map(map_name: str) -> dict:
+ repo_root = Path(__file__).resolve().parents[4]
+ paths = {
+ "root": repo_root / "model_prices_and_context_window.json",
+ "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json",
+ }
+ return json.loads(paths[map_name].read_text())
+
+
+class TestGpt56MantleRegistryEntries:
+ """Locks the gpt-5.6 frontier entries to Bedrock Mantle's live behavior.
+
+ Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna
+ (oversize requests 400 with "prompt tokens (N) exceed model maximum
+ (1050000)", and a 1,030,590-token request completes), matching the OpenAI
+ Bedrock guide. mode must stay "responses": Mantle's native
+ /v1/chat/completions rejects function tools unless reasoning_effort is
+ "none", so chat traffic has to keep bridging to the Responses API
+ (see the responses_api_bridge tests above).
+ """
+
+ @pytest.mark.parametrize("map_name", ("root", "bundled_backup"))
+ @pytest.mark.parametrize(
+ "key",
+ (
+ "bedrock_mantle/openai.gpt-5.6-sol",
+ "bedrock_mantle/openai.gpt-5.6-terra",
+ "bedrock_mantle/openai.gpt-5.6-luna",
+ ),
+ )
+ def test_entry_matches_mantle_enforced_limits(self, map_name, key):
+ entry = _repo_cost_map(map_name)[key]
+ assert entry["max_input_tokens"] == 1050000
+ assert entry["max_output_tokens"] == 128000
+ assert entry["mode"] == "responses"
+ assert entry["use_openai_responses_path"] is True
+ assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"]
From 530dab32b9f5d308fa628586b35bc97eae95a670 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 09:55:13 -0700
Subject: [PATCH 28/70] feat(vertex_ai): add native Vertex AI Interactions API
support
---
litellm/__init__.py | 3 +
litellm/_lazy_imports_registry.py | 5 +
litellm/interactions/utils.py | 7 +
.../llms/vertex_ai/interactions/__init__.py | 0
.../vertex_ai/interactions/transformation.py | 149 +++++++++++
...t_vertex_ai_interactions_transformation.py | 231 ++++++++++++++++++
6 files changed, 395 insertions(+)
create mode 100644 litellm/llms/vertex_ai/interactions/__init__.py
create mode 100644 litellm/llms/vertex_ai/interactions/transformation.py
create mode 100644 tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
diff --git a/litellm/__init__.py b/litellm/__init__.py
index e95b553c5d4..ee2c551481c 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -1801,6 +1801,9 @@ if TYPE_CHECKING:
from .llms.gemini.interactions.transformation import (
GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
)
+ from .llms.vertex_ai.interactions.transformation import (
+ VertexAIInteractionsConfig as VertexAIInteractionsConfig,
+ )
from .llms.openai.chat.o_series_transformation import (
OpenAIOSeriesConfig as OpenAIOSeriesConfig,
OpenAIOSeriesConfig as OpenAIO1Config,
diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py
index 89c72acc06d..c34c9eefe85 100644
--- a/litellm/_lazy_imports_registry.py
+++ b/litellm/_lazy_imports_registry.py
@@ -242,6 +242,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenRouterResponsesAPIConfig",
"BedrockMantleResponsesAPIConfig",
"GoogleAIStudioInteractionsConfig",
+ "VertexAIInteractionsConfig",
"OpenAIOSeriesConfig",
"AnthropicSkillsConfig",
"BaseSkillsAPIConfig",
@@ -977,6 +978,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.gemini.interactions.transformation",
"GoogleAIStudioInteractionsConfig",
),
+ "VertexAIInteractionsConfig": (
+ ".llms.vertex_ai.interactions.transformation",
+ "VertexAIInteractionsConfig",
+ ),
"OpenAIOSeriesConfig": (
".llms.openai.chat.o_series_transformation",
"OpenAIOSeriesConfig",
diff --git a/litellm/interactions/utils.py b/litellm/interactions/utils.py
index 8a1e8836894..3895a85061d 100644
--- a/litellm/interactions/utils.py
+++ b/litellm/interactions/utils.py
@@ -47,6 +47,13 @@ def get_provider_interactions_api_config(
return GoogleAIStudioInteractionsConfig()
+ if provider in (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value):
+ from litellm.llms.vertex_ai.interactions.transformation import (
+ VertexAIInteractionsConfig,
+ )
+
+ return VertexAIInteractionsConfig()
+
return None
diff --git a/litellm/llms/vertex_ai/interactions/__init__.py b/litellm/llms/vertex_ai/interactions/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py
new file mode 100644
index 00000000000..0764a8bea62
--- /dev/null
+++ b/litellm/llms/vertex_ai/interactions/transformation.py
@@ -0,0 +1,149 @@
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+from litellm.litellm_core_utils.url_utils import encode_url_path_segment
+from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig
+from litellm.llms.vertex_ai.common_utils import validate_vertex_location
+from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+VERTEX_INTERACTIONS_API_VERSION: Final = "v1beta1"
+VERTEX_INTERACTIONS_DEFAULT_LOCATION: Final = "global"
+
+
+@dataclass(frozen=True, slots=True)
+class VertexInteractionsTarget:
+ base_url: str
+ project_id: str
+ location: str
+
+ @property
+ def collection_url(self) -> str:
+ return (
+ f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}"
+ f"/projects/{self.project_id}/locations/{self.location}/interactions"
+ )
+
+ def interaction_url(self, interaction_id: str) -> str:
+ encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id")
+ return f"{self.collection_url}/{encoded_interaction_id}"
+
+
+class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig):
+ def __init__(
+ self,
+ mint_access_token: Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]] | None = None,
+ ) -> None:
+ super().__init__()
+ self._mint_access_token: Final[Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]]] = (
+ mint_access_token or self._mint_access_token_with_vertex_base
+ )
+
+ def _mint_access_token_with_vertex_base(
+ self,
+ credentials: VERTEX_CREDENTIALS_TYPES | None,
+ project_id: str | None,
+ ) -> tuple[str, str]:
+ return self._ensure_access_token(
+ credentials=credentials, project_id=project_id, custom_llm_provider="vertex_ai"
+ )
+
+ @property
+ def custom_llm_provider(self) -> LlmProviders:
+ return LlmProviders.VERTEX_AI
+
+ @property
+ def api_version(self) -> str:
+ return VERTEX_INTERACTIONS_API_VERSION
+
+ def get_default_vertex_location(self) -> str:
+ return VERTEX_INTERACTIONS_DEFAULT_LOCATION
+
+ def _mint(self, litellm_params: GenericLiteLLMParams) -> tuple[str, str]:
+ raw_params: Final = litellm_params.model_dump()
+ return self._mint_access_token(
+ self.safe_get_vertex_ai_credentials(raw_params),
+ self.safe_get_vertex_ai_project(raw_params),
+ )
+
+ def _target(self, api_base: str | None, litellm_params: GenericLiteLLMParams) -> VertexInteractionsTarget:
+ _, project_id = self._mint(litellm_params)
+ if not project_id:
+ raise ValueError(
+ "Vertex AI project is required. Set vertex_project, litellm.vertex_project, or VERTEXAI_PROJECT"
+ )
+ location: Final = validate_vertex_location(
+ self.explicit_vertex_ai_location(litellm_params.model_dump()) or VERTEX_INTERACTIONS_DEFAULT_LOCATION
+ )
+ return VertexInteractionsTarget(
+ base_url=self.get_api_base(api_base or None, location),
+ project_id=project_id,
+ location=location,
+ )
+
+ def validate_environment(
+ self,
+ headers: Mapping[str, str],
+ model: str,
+ litellm_params: GenericLiteLLMParams | None,
+ ) -> dict: # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers
+ access_token, _ = self._mint(litellm_params or GenericLiteLLMParams())
+ return { # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers
+ "Content-Type": "application/json",
+ "Authorization": f"Bearer {access_token}",
+ **headers,
+ }
+
+ def get_complete_url(
+ self,
+ api_base: str | None,
+ model: str | None,
+ agent: str | None = None,
+ litellm_params: Mapping[str, object] | None = None,
+ stream: bool | None = None,
+ ) -> str:
+ params: Final = (
+ GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams()
+ )
+ collection_url: Final = self._target(api_base, params).collection_url
+ return f"{collection_url}?alt=sse" if stream else collection_url
+
+ def _interaction_by_id_request(
+ self,
+ interaction_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ url_suffix: str = "",
+ ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+ target: Final = self._target(api_base or None, litellm_params)
+ return f"{target.interaction_url(interaction_id)}{url_suffix}", {} # mutable-ok: same base contract
+
+ def transform_get_interaction_request(
+ self,
+ interaction_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: Mapping[str, str],
+ ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+ return self._interaction_by_id_request(interaction_id, api_base, litellm_params)
+
+ def transform_delete_interaction_request(
+ self,
+ interaction_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: Mapping[str, str],
+ ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+ return self._interaction_by_id_request(interaction_id, api_base, litellm_params)
+
+ def transform_cancel_interaction_request(
+ self,
+ interaction_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: Mapping[str, str],
+ ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
+ return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix=":cancel")
diff --git a/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
new file mode 100644
index 00000000000..3364b1b3872
--- /dev/null
+++ b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py
@@ -0,0 +1,231 @@
+import pytest
+
+import litellm
+from litellm.interactions.utils import get_provider_interactions_api_config
+from litellm.llms.gemini.interactions.transformation import (
+ GoogleAIStudioInteractionsConfig,
+)
+from litellm.llms.vertex_ai.interactions.transformation import (
+ VertexAIInteractionsConfig,
+)
+from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import LlmProviders
+
+GLOBAL_BASE = "https://aiplatform.googleapis.com/v1beta1/projects/test-proj/locations/global/interactions"
+
+
+class MinterRecorder:
+ def __init__(self, resolved_project: str = "creds-proj") -> None:
+ self.calls: list[tuple[VERTEX_CREDENTIALS_TYPES | None, str | None]] = []
+ self.resolved_project = resolved_project
+
+ def __call__(
+ self,
+ credentials: VERTEX_CREDENTIALS_TYPES | None,
+ project_id: str | None,
+ ) -> tuple[str, str]:
+ self.calls.append((credentials, project_id))
+ return "test-token", project_id or self.resolved_project
+
+
+@pytest.fixture
+def minter():
+ return MinterRecorder()
+
+
+@pytest.fixture
+def config(minter):
+ return VertexAIInteractionsConfig(mint_access_token=minter)
+
+
+@pytest.fixture
+def litellm_params():
+ return GenericLiteLLMParams(vertex_project="test-proj", vertex_credentials="creds.json")
+
+
+class TestRegistration:
+ def test_vertex_ai_returns_vertex_config(self):
+ assert isinstance(get_provider_interactions_api_config("vertex_ai"), VertexAIInteractionsConfig)
+
+ def test_vertex_ai_beta_returns_vertex_config(self):
+ assert isinstance(get_provider_interactions_api_config("vertex_ai_beta"), VertexAIInteractionsConfig)
+
+ def test_gemini_still_returns_google_ai_studio_config(self):
+ gemini_config = get_provider_interactions_api_config("gemini")
+ assert isinstance(gemini_config, GoogleAIStudioInteractionsConfig)
+ assert not isinstance(gemini_config, VertexAIInteractionsConfig)
+
+ def test_lazy_import_resolves(self):
+ assert litellm.VertexAIInteractionsConfig is VertexAIInteractionsConfig
+
+ def test_custom_llm_provider_is_vertex_ai(self, config):
+ assert config.custom_llm_provider == LlmProviders.VERTEX_AI
+
+
+class TestValidateEnvironment:
+ def test_sets_bearer_auth_without_gemini_headers(self, config, minter, litellm_params):
+ headers = config.validate_environment(
+ headers={},
+ model="gemini-omni-flash-preview",
+ litellm_params=litellm_params,
+ )
+
+ assert headers["Authorization"] == "Bearer test-token"
+ assert headers["Content-Type"] == "application/json"
+ assert "x-goog-api-key" not in headers
+ assert "Api-Revision" not in headers
+ assert minter.calls == [("creds.json", "test-proj")]
+
+ def test_caller_authorization_wins(self, config, litellm_params):
+ headers = config.validate_environment(
+ headers={"Authorization": "Bearer caller-token"},
+ model="gemini-omni-flash-preview",
+ litellm_params=litellm_params,
+ )
+
+ assert headers["Authorization"] == "Bearer caller-token"
+
+
+class TestGetCompleteUrl:
+ def test_defaults_to_global_v1beta1(self, config, litellm_params):
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params=dict(litellm_params),
+ )
+
+ assert url == GLOBAL_BASE
+
+ def test_stream_appends_alt_sse(self, config, litellm_params):
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params=dict(litellm_params),
+ stream=True,
+ )
+
+ assert url == f"{GLOBAL_BASE}?alt=sse"
+
+ def test_multi_region_location_uses_rep_host(self, config):
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_project": "test-proj", "vertex_location": "us"},
+ )
+
+ assert url == "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/test-proj/locations/us/interactions"
+
+ def test_regional_location_uses_regional_host(self, config):
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_project": "test-proj", "vertex_location": "us-central1"},
+ )
+
+ assert url == (
+ "https://us-central1-aiplatform.googleapis.com"
+ "/v1beta1/projects/test-proj/locations/us-central1/interactions"
+ )
+
+ def test_location_env_fallback_is_ignored(self, config, monkeypatch):
+ monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5")
+
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_project": "test-proj"},
+ )
+
+ assert url == GLOBAL_BASE
+
+ def test_api_base_override(self, config, litellm_params):
+ url = config.get_complete_url(
+ api_base="https://proxy.example.test",
+ model="gemini-omni-flash-preview",
+ litellm_params=dict(litellm_params),
+ )
+
+ assert url == "https://proxy.example.test/v1beta1/projects/test-proj/locations/global/interactions"
+
+ def test_project_resolved_from_credentials_when_not_passed(self, config, monkeypatch):
+ monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
+ monkeypatch.setattr(litellm, "vertex_project", None)
+
+ url = config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_credentials": "creds.json"},
+ )
+
+ assert url == "https://aiplatform.googleapis.com/v1beta1/projects/creds-proj/locations/global/interactions"
+
+ def test_invalid_location_rejected(self, config):
+ with pytest.raises(ValueError, match="Invalid vertex_location"):
+ config.get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={"vertex_project": "test-proj", "vertex_location": "evil.com#"},
+ )
+
+ def test_missing_project_rejected(self, monkeypatch):
+ monkeypatch.delenv("VERTEXAI_PROJECT", raising=False)
+ monkeypatch.setattr(litellm, "vertex_project", None)
+
+ def unresolved_minter(
+ credentials: VERTEX_CREDENTIALS_TYPES | None,
+ project_id: str | None,
+ ) -> tuple[str, str]:
+ return "test-token", ""
+
+ with pytest.raises(ValueError, match="Vertex AI project is required"):
+ VertexAIInteractionsConfig(mint_access_token=unresolved_minter).get_complete_url(
+ api_base=None,
+ model="gemini-omni-flash-preview",
+ litellm_params={},
+ )
+
+
+class TestInteractionByIdRequests:
+ def test_get_url(self, config, litellm_params):
+ url, request_body = config.transform_get_interaction_request(
+ interaction_id="abc123",
+ api_base="",
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert url == f"{GLOBAL_BASE}/abc123"
+ assert request_body == {}
+
+ def test_get_url_encodes_interaction_id(self, config, litellm_params):
+ url, _ = config.transform_get_interaction_request(
+ interaction_id="id/with space",
+ api_base="",
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert url == f"{GLOBAL_BASE}/id%2Fwith%20space"
+
+ def test_delete_url(self, config, litellm_params):
+ url, request_body = config.transform_delete_interaction_request(
+ interaction_id="abc123",
+ api_base="",
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert url == f"{GLOBAL_BASE}/abc123"
+ assert request_body == {}
+
+ def test_cancel_url(self, config, litellm_params):
+ url, request_body = config.transform_cancel_interaction_request(
+ interaction_id="abc123",
+ api_base="",
+ litellm_params=litellm_params,
+ headers={},
+ )
+
+ assert url == f"{GLOBAL_BASE}/abc123:cancel"
+ assert request_body == {}
From ed28581d791e289c2a5b20e0f91678232d5d17ee Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 09:57:25 -0700
Subject: [PATCH 29/70] fix(bedrock_mantle): normalize Codex input item types
Mantle rejects
Mantle 400s ("Invalid 'input': value did not match any expected variant")
on the Codex history item types agent_message, context_compaction, and
local_shell_call, killing every Codex multi-agent session on the first
sub-agent turn. Rewrite agent_message into an assistant output_text message
(preserving encrypted_content slot payloads, which carry the plaintext task
through Mantle), context_compaction into Mantle's supported compaction
spelling, and local_shell_call into the function_call its recorded
function_call_output already pairs with.
---
.../responses/transformation.py | 118 +++++++++++-
...bedrock_mantle_responses_transformation.py | 176 ++++++++++++++++++
2 files changed, 293 insertions(+), 1 deletion(-)
diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py
index 92da5835b2d..9068a641940 100644
--- a/litellm/llms/bedrock_mantle/responses/transformation.py
+++ b/litellm/llms/bedrock_mantle/responses/transformation.py
@@ -15,8 +15,12 @@ role / access key / profile / web identity), signed via the shared
BaseAWSLLM._sign_request after the request body is finalized.
"""
+import json
+from collections.abc import Mapping
from typing import Any, Final
+from typing_extensions import ReadOnly, TypedDict
+
import litellm
from litellm._logging import verbose_logger
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
@@ -50,6 +54,33 @@ _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})
_CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
+_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message"
+_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction"
+_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call"
+
+
+class _RewrittenOutputTextBlock(TypedDict):
+ type: ReadOnly[str]
+ text: ReadOnly[str]
+
+
+class _RewrittenAssistantMessageItem(TypedDict):
+ type: ReadOnly[str]
+ role: ReadOnly[str]
+ content: ReadOnly[tuple[_RewrittenOutputTextBlock, ...]]
+
+
+class _RewrittenCompactionItem(TypedDict):
+ type: ReadOnly[str]
+ encrypted_content: ReadOnly[str]
+
+
+class _RewrittenFunctionCallItem(TypedDict):
+ type: ReadOnly[str]
+ call_id: ReadOnly[str]
+ name: ReadOnly[str]
+ arguments: ReadOnly[str]
+
class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig):
def __init__(
@@ -155,6 +186,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
headers: dict,
) -> dict:
remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input)
+ normalized_input: Final = self._normalize_codex_input_items(remaining_input)
request_params: Final = (
{
**response_api_optional_request_params,
@@ -168,7 +200,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
)
return super().transform_responses_api_request(
model=model,
- input=remaining_input,
+ input=normalized_input,
response_api_optional_request_params=request_params,
litellm_params=litellm_params,
headers=headers,
@@ -210,6 +242,90 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
)
return remaining_input, cls._filter_unsupported_tools(hoisted_tools)
+ @staticmethod
+ def _agent_message_text(item: "Mapping[str, Any]") -> str:
+ content: Final = item.get("content")
+ if not isinstance(content, list):
+ return ""
+ return "".join(
+ str(block.get("text") or block.get("encrypted_content") or "")
+ for block in content
+ if isinstance(block, dict)
+ )
+
+ @classmethod
+ def _normalize_agent_message_item(cls, item: "Mapping[str, Any]") -> "_RewrittenAssistantMessageItem | None":
+ text: Final = cls._agent_message_text(item)
+ if not text:
+ return None
+ rewritten: Final[_RewrittenAssistantMessageItem] = {
+ "type": "message",
+ "role": "assistant",
+ "content": ({"type": "output_text", "text": text},),
+ }
+ return rewritten
+
+ @staticmethod
+ def _normalize_context_compaction_item(item: "Mapping[str, Any]") -> "_RewrittenCompactionItem | None":
+ encrypted_content: Final = item.get("encrypted_content")
+ if not isinstance(encrypted_content, str) or not encrypted_content:
+ return None
+ rewritten: Final[_RewrittenCompactionItem] = {"type": "compaction", "encrypted_content": encrypted_content}
+ return rewritten
+
+ @staticmethod
+ def _normalize_local_shell_call_item(item: "Mapping[str, Any]") -> "_RewrittenFunctionCallItem | None":
+ call_id: Final = item.get("call_id")
+ if not isinstance(call_id, str) or not call_id:
+ return None
+ action: Final = item.get("action")
+ rewritten: Final[_RewrittenFunctionCallItem] = {
+ "type": "function_call",
+ "call_id": call_id,
+ "name": "local_shell",
+ "arguments": json.dumps(action) if isinstance(action, dict) else "{}",
+ }
+ return rewritten
+
+ @classmethod
+ def _normalize_codex_input_item(cls, item: object) -> "tuple[Any, str | None]":
+ """Returns (normalized item or None to drop it, original type when rewritten)."""
+ if not isinstance(item, dict):
+ return item, None
+ item_type: Final = item.get("type")
+ if item_type == _CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE:
+ return cls._normalize_agent_message_item(item), item_type
+ if item_type == _CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE:
+ return cls._normalize_context_compaction_item(item), item_type
+ if item_type == _CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE:
+ return cls._normalize_local_shell_call_item(item), item_type
+ return item, None
+
+ @classmethod
+ def _normalize_codex_input_items(
+ cls,
+ input: "str | ResponseInputParam",
+ ) -> "str | ResponseInputParam":
+ """Rewrite Codex history item types Mantle rejects with 400 "Invalid
+ 'input': value did not match any expected variant" into supported
+ equivalents. `agent_message` (Codex multi-agent traffic; its
+ encrypted_content slot carries the plaintext payload when the model
+ never issued encrypted args) becomes an assistant message,
+ `context_compaction` becomes the `compaction` spelling Mantle accepts,
+ and `local_shell_call` becomes the function_call its recorded
+ function_call_output already pairs with.
+ """
+ if not isinstance(input, list):
+ return input
+ normalized: Final = tuple(cls._normalize_codex_input_item(item) for item in input)
+ rewritten_types: Final = sorted(frozenset(item_type for _, item_type in normalized if item_type is not None))
+ if rewritten_types:
+ verbose_logger.warning(
+ "Bedrock Mantle Responses API: rewrote Codex input item type(s) %s that Mantle rejects.",
+ rewritten_types,
+ )
+ return [item for item, _ in normalized if item is not None] # mutable-ok: ResponseInputParam is a list
+
def map_openai_params(
self,
response_api_optional_params: ResponsesAPIOptionalRequestParams,
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
index 9e05d48a18f..984ff997292 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
@@ -8,6 +8,7 @@ gate, the URL construction for both paths, and the shared Bearer auth.
"""
import copy
+import logging
import pytest
@@ -623,6 +624,181 @@ class TestBedrockMantleCodexAdditionalTools:
assert "additional_tools" in str(mock_debug.call_args)
+class TestBedrockMantleCodexInputItemNormalization:
+ """Mantle 400s ("Invalid 'input': value did not match any expected variant")
+ on the Codex history item types agent_message, context_compaction, and
+ local_shell_call (verified against bedrock-mantle.us-east-1.api.aws with
+ openai.gpt-5.6-sol), so the config must rewrite them into supported
+ equivalents. agent_message is what every Codex multi-agent v2 session sends,
+ and its encrypted_content slot carries the verbatim plaintext payload when
+ the upstream model never issued encrypted args, so that slot must be
+ preserved, not dropped. Mantle also rejects assistant messages with
+ input_text content, so the rewrite must use output_text."""
+
+ _USER_MESSAGE = {
+ "type": "message",
+ "role": "user",
+ "content": [{"type": "input_text", "text": "Continue."}],
+ }
+
+ def _transform(self, input):
+ cfg = BedrockMantleResponsesAPIConfig()
+ return cfg.transform_responses_api_request(
+ model="openai.gpt-5.6-sol",
+ input=input,
+ response_api_optional_request_params={},
+ litellm_params=GenericLiteLLMParams(),
+ headers={},
+ )
+
+ def test_plaintext_agent_message_becomes_assistant_output_text_message(self):
+ body = self._transform(
+ input=[
+ self._USER_MESSAGE,
+ {
+ "type": "agent_message",
+ "id": "amsg_1",
+ "author": "/root/arithmetic",
+ "recipient": "/root",
+ "content": [{"type": "input_text", "text": "Message Type: FINAL_ANSWER\nPayload:\n2+2 is 4."}],
+ },
+ ]
+ )
+ assert body["input"] == [
+ self._USER_MESSAGE,
+ {
+ "type": "message",
+ "role": "assistant",
+ "content": ({"type": "output_text", "text": "Message Type: FINAL_ANSWER\nPayload:\n2+2 is 4."},),
+ },
+ ]
+
+ def test_agent_message_encrypted_content_payload_is_preserved(self):
+ body = self._transform(
+ input=[
+ {
+ "type": "agent_message",
+ "author": "/root",
+ "recipient": "/root/arithmetic",
+ "content": [
+ {"type": "input_text", "text": "Message Type: NEW_TASK\nPayload:\n"},
+ {"type": "encrypted_content", "encrypted_content": "Answer the question 'what is 2+2'."},
+ ],
+ },
+ self._USER_MESSAGE,
+ ]
+ )
+ assert body["input"][0] == {
+ "type": "message",
+ "role": "assistant",
+ "content": (
+ {
+ "type": "output_text",
+ "text": "Message Type: NEW_TASK\nPayload:\nAnswer the question 'what is 2+2'.",
+ },
+ ),
+ }
+
+ def test_agent_message_without_any_text_is_dropped(self):
+ body = self._transform(
+ input=[
+ {"type": "agent_message", "author": "/root", "recipient": "/root/a", "content": []},
+ self._USER_MESSAGE,
+ ]
+ )
+ assert body["input"] == [self._USER_MESSAGE]
+
+ def test_context_compaction_becomes_compaction_with_same_ciphertext(self):
+ body = self._transform(
+ input=[
+ {"type": "context_compaction", "id": "cc_1", "encrypted_content": "smry_abc123"},
+ self._USER_MESSAGE,
+ ]
+ )
+ assert body["input"] == [
+ {"type": "compaction", "encrypted_content": "smry_abc123"},
+ self._USER_MESSAGE,
+ ]
+
+ def test_context_compaction_without_ciphertext_is_dropped(self):
+ body = self._transform(
+ input=[
+ {"type": "context_compaction", "id": "cc_1"},
+ self._USER_MESSAGE,
+ ]
+ )
+ assert body["input"] == [self._USER_MESSAGE]
+
+ def test_local_shell_call_becomes_function_call_keeping_call_id_pairing(self):
+ body = self._transform(
+ input=[
+ {
+ "type": "local_shell_call",
+ "id": "lsh_1",
+ "call_id": "call_1",
+ "status": "completed",
+ "action": {"type": "exec", "command": ["echo", "hi"]},
+ },
+ {"type": "function_call_output", "call_id": "call_1", "output": "hi\n"},
+ self._USER_MESSAGE,
+ ]
+ )
+ assert body["input"] == [
+ {
+ "type": "function_call",
+ "call_id": "call_1",
+ "name": "local_shell",
+ "arguments": '{"type": "exec", "command": ["echo", "hi"]}',
+ },
+ {"type": "function_call_output", "call_id": "call_1", "output": "hi\n"},
+ self._USER_MESSAGE,
+ ]
+
+ def test_local_shell_call_without_call_id_is_dropped(self):
+ body = self._transform(
+ input=[
+ {"type": "local_shell_call", "status": "completed", "action": {"type": "exec", "command": ["ls"]}},
+ self._USER_MESSAGE,
+ ]
+ )
+ assert body["input"] == [self._USER_MESSAGE]
+
+ def test_mantle_supported_item_types_pass_through_untouched(self):
+ supported_items = [
+ self._USER_MESSAGE,
+ {"type": "compaction", "encrypted_content": "smry_abc123"},
+ {"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "call_2"},
+ {"type": "function_call_output", "call_id": "call_2", "output": "ok"},
+ {"type": "tool_search_call", "call_id": "call_3", "execution": "server", "arguments": {"query": "x"}},
+ {"type": "tool_search_output", "call_id": "call_3", "status": "completed", "execution": "server", "tools": []},
+ {"type": "compaction_trigger"},
+ ]
+ body = self._transform(input=copy.deepcopy(supported_items))
+ assert body["input"] == supported_items
+
+ def test_string_input_passes_through(self):
+ body = self._transform(input="Say hi.")
+ assert body["input"] == "Say hi."
+
+ def test_rewrite_is_logged_as_warning_naming_the_types(self, caplog):
+ with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+ body = self._transform(
+ input=[
+ {"type": "agent_message", "author": "a", "recipient": "b", "content": [{"type": "input_text", "text": "hi"}]},
+ self._USER_MESSAGE,
+ ]
+ )
+ assert body["input"][0]["role"] == "assistant"
+ rewrite_warnings = [
+ record.getMessage()
+ for record in caplog.records
+ if record.levelno == logging.WARNING and "rewrote Codex input item type" in record.getMessage()
+ ]
+ assert rewrite_warnings == [
+ "Bedrock Mantle Responses API: rewrote Codex input item type(s) ['agent_message'] that Mantle rejects."
+ ]
+
+
class TestBedrockMantleResponsesRegistry:
def test_registry_returns_config_for_gpt_5_5(self, local_cost_map):
# gpt-5.x advertises /v1/responses in supported_endpoints (capability)
From 44c7cb20aee5832734f626ad522c867b851fde40 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:10:07 -0700
Subject: [PATCH 30/70] feat(models): add missing Together AI serverless models
to the cost map
Backfill 21 serverless chat models, the multilingual-e5 embedding model, and
Llama-Guard-4-12B from the live Together catalog with per-token pricing and
capability flags. Mark 25 delisted together_ai entries with their documented
deprecation_date and point superseded models at a live successor via metadata.
Reprice Llama-3.3-70B-Instruct-Turbo to Together's current rate.
---
...odel_prices_and_context_window_backup.json | 349 +++++++++++++++++-
model_prices_and_context_window.json | 349 +++++++++++++++++-
.../test_together_ai_model_metadata.py | 149 ++++++++
3 files changed, 843 insertions(+), 4 deletions(-)
create mode 100644 tests/test_litellm/test_together_ai_model_metadata.py
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 9f953e11df1..9b2b910e007 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -37886,6 +37886,7 @@
"output_cost_per_token": 1e-07
},
"together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": {
+ "deprecation_date": "2026-02-06",
"litellm_provider": "together_ai",
"mode": "chat",
"supports_function_calling": true,
@@ -37902,6 +37903,7 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": {
+ "deprecation_date": "2026-07-10",
"input_cost_per_token": 2e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 262000,
@@ -37914,6 +37916,7 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": {
+ "deprecation_date": "2026-04-16",
"input_cost_per_token": 6.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
@@ -37926,6 +37929,7 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": {
+ "deprecation_date": "2026-02-06",
"input_cost_per_token": 2e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 40000,
@@ -37937,6 +37941,7 @@
"supports_tool_choice": false
},
"together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": {
+ "deprecation_date": "2026-06-04",
"input_cost_per_token": 2e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
@@ -37949,11 +37954,15 @@
"supports_tool_choice": true
},
"together_ai/deepseek-ai/DeepSeek-R1": {
+ "deprecation_date": "2026-05-14",
"input_cost_per_token": 3e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 128000,
"max_output_tokens": 20480,
"max_tokens": 20480,
+ "metadata": {
+ "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
+ },
"mode": "chat",
"output_cost_per_token": 7e-06,
"supports_function_calling": true,
@@ -37962,6 +37971,7 @@
"supports_tool_choice": true
},
"together_ai/deepseek-ai/DeepSeek-R1-0528-tput": {
+ "deprecation_date": "2026-02-03",
"input_cost_per_token": 5.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 128000,
@@ -37979,6 +37989,9 @@
"max_input_tokens": 65536,
"max_output_tokens": 8192,
"max_tokens": 8192,
+ "metadata": {
+ "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
+ },
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"supports_function_calling": true,
@@ -37987,9 +38000,13 @@
"supports_tool_choice": true
},
"together_ai/deepseek-ai/DeepSeek-V3.1": {
+ "deprecation_date": "2026-05-14",
"input_cost_per_token": 6e-07,
"litellm_provider": "together_ai",
"max_tokens": 16384,
+ "metadata": {
+ "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
+ },
"mode": "chat",
"output_cost_per_token": 1.7e-06,
"source": "https://www.together.ai/models/deepseek-v3-1",
@@ -38001,6 +38018,7 @@
"max_output_tokens": 16384
},
"together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": {
+ "deprecation_date": "2026-03-06",
"litellm_provider": "together_ai",
"mode": "chat",
"supports_function_calling": true,
@@ -38009,16 +38027,21 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": {
- "input_cost_per_token": 8.8e-07,
+ "input_cost_per_token": 1.04e-06,
"litellm_provider": "together_ai",
+ "max_input_tokens": 131072,
+ "max_output_tokens": 131072,
+ "max_tokens": 131072,
"mode": "chat",
- "output_cost_per_token": 8.8e-07,
+ "output_cost_per_token": 1.04e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true
},
"together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": {
+ "deprecation_date": "2025-11-13",
"input_cost_per_token": 0,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38029,6 +38052,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": {
+ "deprecation_date": "2026-03-31",
"input_cost_per_token": 2.7e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38039,6 +38063,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": {
+ "deprecation_date": "2026-02-06",
"input_cost_per_token": 1.8e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38049,6 +38074,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": {
+ "deprecation_date": "2026-02-06",
"input_cost_per_token": 3.5e-06,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38059,6 +38085,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": {
+ "deprecation_date": "2026-02-25",
"input_cost_per_token": 8.8e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38069,6 +38096,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": {
+ "deprecation_date": "2026-03-06",
"input_cost_per_token": 1.8e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38079,6 +38107,7 @@
"supports_tool_choice": true
},
"together_ai/mistralai/Mistral-7B-Instruct-v0.1": {
+ "deprecation_date": "2025-11-13",
"litellm_provider": "together_ai",
"mode": "chat",
"supports_function_calling": true,
@@ -38087,6 +38116,7 @@
"supports_tool_choice": true
},
"together_ai/mistralai/Mistral-Small-24B-Instruct-2501": {
+ "deprecation_date": "2026-04-02",
"litellm_provider": "together_ai",
"mode": "chat",
"supports_function_calling": true,
@@ -38094,6 +38124,7 @@
"supports_tool_choice": true
},
"together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": {
+ "deprecation_date": "2026-04-16",
"input_cost_per_token": 6e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38106,6 +38137,9 @@
"together_ai/moonshotai/Kimi-K2-Instruct": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
+ "metadata": {
+ "successor": "together_ai/moonshotai/Kimi-K3"
+ },
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://www.together.ai/models/kimi-k2-instruct",
@@ -38149,6 +38183,7 @@
"supports_tool_choice": true
},
"together_ai/zai-org/GLM-4.5-Air-FP8": {
+ "deprecation_date": "2026-04-02",
"input_cost_per_token": 2e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 128000,
@@ -38166,6 +38201,9 @@
"max_input_tokens": 200000,
"max_output_tokens": 200000,
"max_tokens": 200000,
+ "metadata": {
+ "successor": "together_ai/zai-org/GLM-5.2"
+ },
"mode": "chat",
"output_cost_per_token": 2.2e-06,
"source": "https://www.together.ai/models/glm-4-6",
@@ -38175,11 +38213,15 @@
"supports_tool_choice": true
},
"together_ai/zai-org/GLM-4.7": {
+ "deprecation_date": "2026-04-02",
"input_cost_per_token": 4.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 200000,
"max_output_tokens": 200000,
"max_tokens": 200000,
+ "metadata": {
+ "successor": "together_ai/zai-org/GLM-5.2"
+ },
"mode": "chat",
"output_cost_per_token": 2e-06,
"source": "https://www.together.ai/models/glm-4-7",
@@ -38189,11 +38231,15 @@
"supports_tool_choice": true
},
"together_ai/moonshotai/Kimi-K2.5": {
+ "deprecation_date": "2026-05-21",
"input_cost_per_token": 5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
+ "metadata": {
+ "successor": "together_ai/moonshotai/Kimi-K3"
+ },
"mode": "chat",
"output_cost_per_token": 2.8e-06,
"source": "https://www.together.ai/models/kimi-k2-5",
@@ -38203,9 +38249,13 @@
"supports_reasoning": true
},
"together_ai/moonshotai/Kimi-K2-Instruct-0905": {
+ "deprecation_date": "2026-03-06",
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 262144,
+ "metadata": {
+ "successor": "together_ai/moonshotai/Kimi-K3"
+ },
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://www.together.ai/models/kimi-k2-0905",
@@ -38214,9 +38264,13 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": {
+ "deprecation_date": "2026-04-02",
"input_cost_per_token": 1.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 262144,
+ "metadata": {
+ "successor": "together_ai/Qwen/Qwen3.7-Plus"
+ },
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct",
@@ -38226,9 +38280,13 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": {
+ "deprecation_date": "2026-02-25",
"input_cost_per_token": 1.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 262144,
+ "metadata": {
+ "successor": "together_ai/Qwen/Qwen3.6-Plus"
+ },
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking",
@@ -38238,6 +38296,7 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3.5-397B-A17B": {
+ "deprecation_date": "2026-06-29",
"input_cost_per_token": 6e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 262144,
@@ -38249,6 +38308,292 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
+ "together_ai/MiniMaxAI/MiniMax-M3": {
+ "input_cost_per_token": 3e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 524288,
+ "max_output_tokens": 524288,
+ "max_tokens": 524288,
+ "mode": "chat",
+ "output_cost_per_token": 1.2e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/Prism-ML/Ternary-Bonsai-27B": {
+ "input_cost_per_token": 0.0,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 0.0,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/Qwen/Qwen3.5-9B": {
+ "input_cost_per_token": 1.7e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 2.5e-07,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/Qwen/Qwen3.6-Plus": {
+ "input_cost_per_token": 5e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1000000,
+ "max_output_tokens": 1000000,
+ "max_tokens": 1000000,
+ "mode": "chat",
+ "output_cost_per_token": 3e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_reasoning": true
+ },
+ "together_ai/Qwen/Qwen3.7-Max": {
+ "input_cost_per_token": 1.25e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1000000,
+ "max_output_tokens": 1000000,
+ "max_tokens": 1000000,
+ "mode": "chat",
+ "output_cost_per_token": 3.75e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/Qwen/Qwen3.7-Plus": {
+ "input_cost_per_token": 3.2e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1000000,
+ "max_output_tokens": 1000000,
+ "max_tokens": 1000000,
+ "mode": "chat",
+ "output_cost_per_token": 1.28e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/Qwen/Qwen3.8-2.4T-A95B": {
+ "input_cost_per_token": 2.5e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1010000,
+ "max_output_tokens": 1010000,
+ "max_tokens": 1010000,
+ "mode": "chat",
+ "output_cost_per_token": 6.25e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/arize-ai/qwen-2-1.5b-instruct": {
+ "input_cost_per_token": 1e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 32768,
+ "max_output_tokens": 32768,
+ "max_tokens": 32768,
+ "mode": "chat",
+ "output_cost_per_token": 1e-07,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": {
+ "input_cost_per_token": 1.4e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048576,
+ "max_output_tokens": 1048576,
+ "max_tokens": 1048576,
+ "mode": "chat",
+ "output_cost_per_token": 2.8e-07,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/deepseek-ai/DeepSeek-V4-Pro": {
+ "input_cost_per_token": 1.74e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 512000,
+ "max_output_tokens": 512000,
+ "max_tokens": 512000,
+ "mode": "chat",
+ "output_cost_per_token": 3.48e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813": {
+ "input_cost_per_token": 1.32e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048576,
+ "max_output_tokens": 1048576,
+ "max_tokens": 1048576,
+ "mode": "chat",
+ "output_cost_per_token": 3.96e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/google/gemma-3n-E4B-it": {
+ "input_cost_per_token": 6e-08,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 32768,
+ "max_output_tokens": 32768,
+ "max_tokens": 32768,
+ "mode": "chat",
+ "output_cost_per_token": 1.2e-07,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/google/gemma-4-31B-it": {
+ "input_cost_per_token": 3.9e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 9.7e-07,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/intfloat/multilingual-e5-large-instruct": {
+ "input_cost_per_token": 2e-08,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 514,
+ "max_tokens": 514,
+ "mode": "embedding",
+ "output_cost_per_token": 2e-08,
+ "output_vector_size": 1024,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/meta-llama/Llama-Guard-4-12B": {
+ "input_cost_per_token": 2e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048576,
+ "max_output_tokens": 1048576,
+ "max_tokens": 1048576,
+ "mode": "chat",
+ "output_cost_per_token": 2e-07,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/meta-models/Muse-Glimmer-30B": {
+ "input_cost_per_token": 3.5e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 131072,
+ "max_output_tokens": 131072,
+ "max_tokens": 131072,
+ "mode": "chat",
+ "output_cost_per_token": 1.5e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/moonshotai/Kimi-K2.7-Code": {
+ "input_cost_per_token": 9.5e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 4e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/moonshotai/Kimi-K3": {
+ "input_cost_per_token": 3e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048576,
+ "max_output_tokens": 1048576,
+ "max_tokens": 1048576,
+ "mode": "chat",
+ "output_cost_per_token": 1.5e-05,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/nvidia/nemotron-3-ultra-550b-a55b": {
+ "input_cost_per_token": 6e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 512288,
+ "max_output_tokens": 512288,
+ "max_tokens": 512288,
+ "mode": "chat",
+ "output_cost_per_token": 3.6e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/pearl-ai/gemma-4-31b-it": {
+ "input_cost_per_token": 2.8e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 8.6e-07,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/thinkingmachines/Inkling": {
+ "input_cost_per_token": 1e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 524288,
+ "max_output_tokens": 524288,
+ "max_tokens": 524288,
+ "mode": "chat",
+ "output_cost_per_token": 4.05e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/thinkingmachines/Inkling-Small": {
+ "input_cost_per_token": 5e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 524288,
+ "max_output_tokens": 524288,
+ "max_tokens": 524288,
+ "mode": "chat",
+ "output_cost_per_token": 1.2e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/zai-org/GLM-5.2": {
+ "input_cost_per_token": 1.4e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048575,
+ "max_output_tokens": 1048575,
+ "max_tokens": 1048575,
+ "mode": "chat",
+ "output_cost_per_token": 4.4e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
"tts-1": {
"input_cost_per_character": 1.5e-05,
"litellm_provider": "openai",
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 9f953e11df1..9b2b910e007 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -37886,6 +37886,7 @@
"output_cost_per_token": 1e-07
},
"together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": {
+ "deprecation_date": "2026-02-06",
"litellm_provider": "together_ai",
"mode": "chat",
"supports_function_calling": true,
@@ -37902,6 +37903,7 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": {
+ "deprecation_date": "2026-07-10",
"input_cost_per_token": 2e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 262000,
@@ -37914,6 +37916,7 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": {
+ "deprecation_date": "2026-04-16",
"input_cost_per_token": 6.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
@@ -37926,6 +37929,7 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": {
+ "deprecation_date": "2026-02-06",
"input_cost_per_token": 2e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 40000,
@@ -37937,6 +37941,7 @@
"supports_tool_choice": false
},
"together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": {
+ "deprecation_date": "2026-06-04",
"input_cost_per_token": 2e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
@@ -37949,11 +37954,15 @@
"supports_tool_choice": true
},
"together_ai/deepseek-ai/DeepSeek-R1": {
+ "deprecation_date": "2026-05-14",
"input_cost_per_token": 3e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 128000,
"max_output_tokens": 20480,
"max_tokens": 20480,
+ "metadata": {
+ "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
+ },
"mode": "chat",
"output_cost_per_token": 7e-06,
"supports_function_calling": true,
@@ -37962,6 +37971,7 @@
"supports_tool_choice": true
},
"together_ai/deepseek-ai/DeepSeek-R1-0528-tput": {
+ "deprecation_date": "2026-02-03",
"input_cost_per_token": 5.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 128000,
@@ -37979,6 +37989,9 @@
"max_input_tokens": 65536,
"max_output_tokens": 8192,
"max_tokens": 8192,
+ "metadata": {
+ "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
+ },
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"supports_function_calling": true,
@@ -37987,9 +38000,13 @@
"supports_tool_choice": true
},
"together_ai/deepseek-ai/DeepSeek-V3.1": {
+ "deprecation_date": "2026-05-14",
"input_cost_per_token": 6e-07,
"litellm_provider": "together_ai",
"max_tokens": 16384,
+ "metadata": {
+ "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
+ },
"mode": "chat",
"output_cost_per_token": 1.7e-06,
"source": "https://www.together.ai/models/deepseek-v3-1",
@@ -38001,6 +38018,7 @@
"max_output_tokens": 16384
},
"together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": {
+ "deprecation_date": "2026-03-06",
"litellm_provider": "together_ai",
"mode": "chat",
"supports_function_calling": true,
@@ -38009,16 +38027,21 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": {
- "input_cost_per_token": 8.8e-07,
+ "input_cost_per_token": 1.04e-06,
"litellm_provider": "together_ai",
+ "max_input_tokens": 131072,
+ "max_output_tokens": 131072,
+ "max_tokens": 131072,
"mode": "chat",
- "output_cost_per_token": 8.8e-07,
+ "output_cost_per_token": 1.04e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true
},
"together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": {
+ "deprecation_date": "2025-11-13",
"input_cost_per_token": 0,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38029,6 +38052,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": {
+ "deprecation_date": "2026-03-31",
"input_cost_per_token": 2.7e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38039,6 +38063,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": {
+ "deprecation_date": "2026-02-06",
"input_cost_per_token": 1.8e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38049,6 +38074,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": {
+ "deprecation_date": "2026-02-06",
"input_cost_per_token": 3.5e-06,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38059,6 +38085,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": {
+ "deprecation_date": "2026-02-25",
"input_cost_per_token": 8.8e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38069,6 +38096,7 @@
"supports_tool_choice": true
},
"together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": {
+ "deprecation_date": "2026-03-06",
"input_cost_per_token": 1.8e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38079,6 +38107,7 @@
"supports_tool_choice": true
},
"together_ai/mistralai/Mistral-7B-Instruct-v0.1": {
+ "deprecation_date": "2025-11-13",
"litellm_provider": "together_ai",
"mode": "chat",
"supports_function_calling": true,
@@ -38087,6 +38116,7 @@
"supports_tool_choice": true
},
"together_ai/mistralai/Mistral-Small-24B-Instruct-2501": {
+ "deprecation_date": "2026-04-02",
"litellm_provider": "together_ai",
"mode": "chat",
"supports_function_calling": true,
@@ -38094,6 +38124,7 @@
"supports_tool_choice": true
},
"together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": {
+ "deprecation_date": "2026-04-16",
"input_cost_per_token": 6e-07,
"litellm_provider": "together_ai",
"mode": "chat",
@@ -38106,6 +38137,9 @@
"together_ai/moonshotai/Kimi-K2-Instruct": {
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
+ "metadata": {
+ "successor": "together_ai/moonshotai/Kimi-K3"
+ },
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://www.together.ai/models/kimi-k2-instruct",
@@ -38149,6 +38183,7 @@
"supports_tool_choice": true
},
"together_ai/zai-org/GLM-4.5-Air-FP8": {
+ "deprecation_date": "2026-04-02",
"input_cost_per_token": 2e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 128000,
@@ -38166,6 +38201,9 @@
"max_input_tokens": 200000,
"max_output_tokens": 200000,
"max_tokens": 200000,
+ "metadata": {
+ "successor": "together_ai/zai-org/GLM-5.2"
+ },
"mode": "chat",
"output_cost_per_token": 2.2e-06,
"source": "https://www.together.ai/models/glm-4-6",
@@ -38175,11 +38213,15 @@
"supports_tool_choice": true
},
"together_ai/zai-org/GLM-4.7": {
+ "deprecation_date": "2026-04-02",
"input_cost_per_token": 4.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 200000,
"max_output_tokens": 200000,
"max_tokens": 200000,
+ "metadata": {
+ "successor": "together_ai/zai-org/GLM-5.2"
+ },
"mode": "chat",
"output_cost_per_token": 2e-06,
"source": "https://www.together.ai/models/glm-4-7",
@@ -38189,11 +38231,15 @@
"supports_tool_choice": true
},
"together_ai/moonshotai/Kimi-K2.5": {
+ "deprecation_date": "2026-05-21",
"input_cost_per_token": 5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"max_tokens": 256000,
+ "metadata": {
+ "successor": "together_ai/moonshotai/Kimi-K3"
+ },
"mode": "chat",
"output_cost_per_token": 2.8e-06,
"source": "https://www.together.ai/models/kimi-k2-5",
@@ -38203,9 +38249,13 @@
"supports_reasoning": true
},
"together_ai/moonshotai/Kimi-K2-Instruct-0905": {
+ "deprecation_date": "2026-03-06",
"input_cost_per_token": 1e-06,
"litellm_provider": "together_ai",
"max_input_tokens": 262144,
+ "metadata": {
+ "successor": "together_ai/moonshotai/Kimi-K3"
+ },
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://www.together.ai/models/kimi-k2-0905",
@@ -38214,9 +38264,13 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": {
+ "deprecation_date": "2026-04-02",
"input_cost_per_token": 1.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 262144,
+ "metadata": {
+ "successor": "together_ai/Qwen/Qwen3.7-Plus"
+ },
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct",
@@ -38226,9 +38280,13 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": {
+ "deprecation_date": "2026-02-25",
"input_cost_per_token": 1.5e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 262144,
+ "metadata": {
+ "successor": "together_ai/Qwen/Qwen3.6-Plus"
+ },
"mode": "chat",
"output_cost_per_token": 1.5e-06,
"source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking",
@@ -38238,6 +38296,7 @@
"supports_tool_choice": true
},
"together_ai/Qwen/Qwen3.5-397B-A17B": {
+ "deprecation_date": "2026-06-29",
"input_cost_per_token": 6e-07,
"litellm_provider": "together_ai",
"max_input_tokens": 262144,
@@ -38249,6 +38308,292 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
+ "together_ai/MiniMaxAI/MiniMax-M3": {
+ "input_cost_per_token": 3e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 524288,
+ "max_output_tokens": 524288,
+ "max_tokens": 524288,
+ "mode": "chat",
+ "output_cost_per_token": 1.2e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/Prism-ML/Ternary-Bonsai-27B": {
+ "input_cost_per_token": 0.0,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 0.0,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/Qwen/Qwen3.5-9B": {
+ "input_cost_per_token": 1.7e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 2.5e-07,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/Qwen/Qwen3.6-Plus": {
+ "input_cost_per_token": 5e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1000000,
+ "max_output_tokens": 1000000,
+ "max_tokens": 1000000,
+ "mode": "chat",
+ "output_cost_per_token": 3e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_reasoning": true
+ },
+ "together_ai/Qwen/Qwen3.7-Max": {
+ "input_cost_per_token": 1.25e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1000000,
+ "max_output_tokens": 1000000,
+ "max_tokens": 1000000,
+ "mode": "chat",
+ "output_cost_per_token": 3.75e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/Qwen/Qwen3.7-Plus": {
+ "input_cost_per_token": 3.2e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1000000,
+ "max_output_tokens": 1000000,
+ "max_tokens": 1000000,
+ "mode": "chat",
+ "output_cost_per_token": 1.28e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/Qwen/Qwen3.8-2.4T-A95B": {
+ "input_cost_per_token": 2.5e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1010000,
+ "max_output_tokens": 1010000,
+ "max_tokens": 1010000,
+ "mode": "chat",
+ "output_cost_per_token": 6.25e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/arize-ai/qwen-2-1.5b-instruct": {
+ "input_cost_per_token": 1e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 32768,
+ "max_output_tokens": 32768,
+ "max_tokens": 32768,
+ "mode": "chat",
+ "output_cost_per_token": 1e-07,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": {
+ "input_cost_per_token": 1.4e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048576,
+ "max_output_tokens": 1048576,
+ "max_tokens": 1048576,
+ "mode": "chat",
+ "output_cost_per_token": 2.8e-07,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/deepseek-ai/DeepSeek-V4-Pro": {
+ "input_cost_per_token": 1.74e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 512000,
+ "max_output_tokens": 512000,
+ "max_tokens": 512000,
+ "mode": "chat",
+ "output_cost_per_token": 3.48e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813": {
+ "input_cost_per_token": 1.32e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048576,
+ "max_output_tokens": 1048576,
+ "max_tokens": 1048576,
+ "mode": "chat",
+ "output_cost_per_token": 3.96e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/google/gemma-3n-E4B-it": {
+ "input_cost_per_token": 6e-08,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 32768,
+ "max_output_tokens": 32768,
+ "max_tokens": 32768,
+ "mode": "chat",
+ "output_cost_per_token": 1.2e-07,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/google/gemma-4-31B-it": {
+ "input_cost_per_token": 3.9e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 9.7e-07,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/intfloat/multilingual-e5-large-instruct": {
+ "input_cost_per_token": 2e-08,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 514,
+ "max_tokens": 514,
+ "mode": "embedding",
+ "output_cost_per_token": 2e-08,
+ "output_vector_size": 1024,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/meta-llama/Llama-Guard-4-12B": {
+ "input_cost_per_token": 2e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048576,
+ "max_output_tokens": 1048576,
+ "max_tokens": 1048576,
+ "mode": "chat",
+ "output_cost_per_token": 2e-07,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/meta-models/Muse-Glimmer-30B": {
+ "input_cost_per_token": 3.5e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 131072,
+ "max_output_tokens": 131072,
+ "max_tokens": 131072,
+ "mode": "chat",
+ "output_cost_per_token": 1.5e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/moonshotai/Kimi-K2.7-Code": {
+ "input_cost_per_token": 9.5e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 4e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/moonshotai/Kimi-K3": {
+ "input_cost_per_token": 3e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048576,
+ "max_output_tokens": 1048576,
+ "max_tokens": 1048576,
+ "mode": "chat",
+ "output_cost_per_token": 1.5e-05,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true,
+ "supports_vision": true
+ },
+ "together_ai/nvidia/nemotron-3-ultra-550b-a55b": {
+ "input_cost_per_token": 6e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 512288,
+ "max_output_tokens": 512288,
+ "max_tokens": 512288,
+ "mode": "chat",
+ "output_cost_per_token": 3.6e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/pearl-ai/gemma-4-31b-it": {
+ "input_cost_per_token": 2.8e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 262144,
+ "max_output_tokens": 262144,
+ "max_tokens": 262144,
+ "mode": "chat",
+ "output_cost_per_token": 8.6e-07,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/thinkingmachines/Inkling": {
+ "input_cost_per_token": 1e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 524288,
+ "max_output_tokens": 524288,
+ "max_tokens": 524288,
+ "mode": "chat",
+ "output_cost_per_token": 4.05e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
+ "together_ai/thinkingmachines/Inkling-Small": {
+ "input_cost_per_token": 5e-07,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 524288,
+ "max_output_tokens": 524288,
+ "max_tokens": 524288,
+ "mode": "chat",
+ "output_cost_per_token": 1.2e-06,
+ "source": "https://docs.together.ai/docs/serverless-models"
+ },
+ "together_ai/zai-org/GLM-5.2": {
+ "input_cost_per_token": 1.4e-06,
+ "litellm_provider": "together_ai",
+ "max_input_tokens": 1048575,
+ "max_output_tokens": 1048575,
+ "max_tokens": 1048575,
+ "mode": "chat",
+ "output_cost_per_token": 4.4e-06,
+ "source": "https://docs.together.ai/docs/serverless-models",
+ "supports_function_calling": true,
+ "supports_parallel_function_calling": true,
+ "supports_reasoning": true,
+ "supports_response_schema": true,
+ "supports_tool_choice": true
+ },
"tts-1": {
"input_cost_per_character": 1.5e-05,
"litellm_provider": "openai",
diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py
new file mode 100644
index 00000000000..7d7712f3c56
--- /dev/null
+++ b/tests/test_litellm/test_together_ai_model_metadata.py
@@ -0,0 +1,149 @@
+import json
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
+
+REPO_ROOT: Final = Path(__file__).parents[2]
+
+SERVERLESS_CHAT_MODELS: Final = (
+ "together_ai/moonshotai/Kimi-K3",
+ "together_ai/zai-org/GLM-5.2",
+ "together_ai/deepseek-ai/DeepSeek-V4-Pro",
+ "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813",
+ "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731",
+ "together_ai/moonshotai/Kimi-K2.7-Code",
+ "together_ai/MiniMaxAI/MiniMax-M3",
+ "together_ai/thinkingmachines/Inkling",
+ "together_ai/thinkingmachines/Inkling-Small",
+ "together_ai/Qwen/Qwen3.8-2.4T-A95B",
+ "together_ai/Qwen/Qwen3.7-Max",
+ "together_ai/Qwen/Qwen3.7-Plus",
+ "together_ai/Qwen/Qwen3.6-Plus",
+ "together_ai/Qwen/Qwen3.5-9B",
+ "together_ai/nvidia/nemotron-3-ultra-550b-a55b",
+ "together_ai/meta-models/Muse-Glimmer-30B",
+ "together_ai/google/gemma-4-31B-it",
+ "together_ai/pearl-ai/gemma-4-31b-it",
+ "together_ai/google/gemma-3n-E4B-it",
+ "together_ai/arize-ai/qwen-2-1.5b-instruct",
+ "together_ai/Prism-ML/Ternary-Bonsai-27B",
+ "together_ai/meta-llama/Llama-Guard-4-12B",
+ "together_ai/openai/gpt-oss-120b",
+ "together_ai/openai/gpt-oss-20b",
+ "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
+)
+
+DEPRECATED_MODELS: Final = {
+ "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": "2026-07-10",
+ "together_ai/Qwen/Qwen3.5-397B-A17B": "2026-06-29",
+ "together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": "2026-06-04",
+ "together_ai/moonshotai/Kimi-K2.5": "2026-05-21",
+ "together_ai/deepseek-ai/DeepSeek-R1": "2026-05-14",
+ "together_ai/deepseek-ai/DeepSeek-V3.1": "2026-05-14",
+ "together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": "2026-04-16",
+ "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": "2026-04-16",
+ "together_ai/zai-org/GLM-4.5-Air-FP8": "2026-04-02",
+ "together_ai/zai-org/GLM-4.7": "2026-04-02",
+ "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": "2026-04-02",
+ "together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": "2026-04-02",
+ "together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": "2026-03-31",
+ "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": "2026-03-06",
+ "together_ai/moonshotai/Kimi-K2-Instruct-0905": "2026-03-06",
+ "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": "2026-03-06",
+ "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": "2026-02-25",
+ "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": "2026-02-25",
+ "together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": "2026-02-06",
+ "together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": "2026-02-06",
+ "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": "2026-02-06",
+ "together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": "2026-02-06",
+ "together_ai/deepseek-ai/DeepSeek-R1-0528-tput": "2026-02-03",
+ "together_ai/mistralai/Mistral-7B-Instruct-v0.1": "2025-11-13",
+ "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": "2025-11-13",
+}
+
+
+@pytest.fixture(scope="module")
+def cost_map() -> dict:
+ with open(REPO_ROOT / "model_prices_and_context_window.json") as f:
+ return json.load(f)
+
+
+@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS)
+def test_together_serverless_chat_model_is_mapped(cost_map: dict, model: str):
+ info = cost_map.get(model)
+ assert info is not None, f"{model} missing from model_prices_and_context_window.json"
+ assert info["litellm_provider"] == "together_ai"
+ assert info["mode"] == "chat"
+ assert info["input_cost_per_token"] >= 0
+ assert info["output_cost_per_token"] >= info["input_cost_per_token"]
+ assert "deprecation_date" not in info
+
+ routed_model, provider, _, _ = get_llm_provider(model=model)
+ assert routed_model == model.removeprefix("together_ai/")
+ assert provider == "together_ai"
+
+
+def test_together_kimi_k3_pricing_and_capabilities(cost_map: dict):
+ info = cost_map["together_ai/moonshotai/Kimi-K3"]
+ assert info["input_cost_per_token"] == 3e-06
+ assert info["output_cost_per_token"] == 1.5e-05
+ assert info["max_input_tokens"] == 1048576
+ assert info["supports_function_calling"] is True
+ assert info["supports_tool_choice"] is True
+ assert info["supports_response_schema"] is True
+ assert info["supports_vision"] is True
+ assert info["supports_reasoning"] is True
+
+
+def test_together_glm_52_pricing(cost_map: dict):
+ info = cost_map["together_ai/zai-org/GLM-5.2"]
+ assert info["input_cost_per_token"] == 1.4e-06
+ assert info["output_cost_per_token"] == 4.4e-06
+ assert info["supports_function_calling"] is True
+ assert info["supports_reasoning"] is True
+
+
+def test_together_multilingual_e5_embedding_entry(cost_map: dict):
+ info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"]
+ assert info["mode"] == "embedding"
+ assert info["input_cost_per_token"] == 2e-08
+ assert info["max_input_tokens"] == 514
+ assert info["output_vector_size"] == 1024
+
+
+def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: dict):
+ info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"]
+ assert info["input_cost_per_token"] == 1.04e-06
+ assert info["output_cost_per_token"] == 1.04e-06
+ assert info["max_input_tokens"] == 131072
+
+
+@pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS))
+def test_together_deprecated_model_carries_deprecation_date(cost_map: dict, model: str):
+ info = cost_map.get(model)
+ assert info is not None, f"{model} missing from model_prices_and_context_window.json"
+ assert info.get("deprecation_date") == DEPRECATED_MODELS[model]
+
+
+def test_together_successor_metadata_points_at_live_models(cost_map: dict):
+ successors = {
+ model: info["metadata"]["successor"]
+ for model, info in cost_map.items()
+ if model.startswith("together_ai/") and "successor" in info.get("metadata", {})
+ }
+ assert len(successors) >= 10
+ for model, successor in successors.items():
+ target = cost_map.get(successor)
+ assert target is not None, f"{model} names successor {successor} that is not in the map"
+ assert "deprecation_date" not in target, f"{model} names deprecated successor {successor}"
+
+
+def test_together_backup_cost_map_in_sync(cost_map: dict):
+ with open(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json") as f:
+ backup = json.load(f)
+ together_main = {k: v for k, v in cost_map.items() if k.startswith("together_ai/")}
+ together_backup = {k: v for k, v in backup.items() if k.startswith("together_ai/")}
+ assert together_backup == together_main
From 68ad575fc21077af8d0e038c92826110ccc0d7c7 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:10:38 -0700
Subject: [PATCH 31/70] fix(bedrock_mantle): register a Bedrock runtime
passthrough config so /bedrock/model//invoke works
---
.../bedrock/passthrough/transformation.py | 5 +
litellm/llms/bedrock_mantle/common_utils.py | 42 +++--
.../passthrough/transformation.py | 44 +++++
litellm/passthrough/main.py | 2 +-
litellm/utils.py | 6 +
...drock_mantle_passthrough_transformation.py | 152 ++++++++++++++++++
6 files changed, 234 insertions(+), 17 deletions(-)
create mode 100644 litellm/llms/bedrock_mantle/passthrough/transformation.py
create mode 100644 tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py
index 0ce2e6f60d3..d0a3c37ffb3 100644
--- a/litellm/llms/bedrock/passthrough/transformation.py
+++ b/litellm/llms/bedrock/passthrough/transformation.py
@@ -1,4 +1,5 @@
import json
+from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Optional, cast
from httpx import Response
@@ -93,6 +94,9 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
endpoint_url,
)
+ def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None:
+ return None
+
def sign_request(
self,
headers: dict,
@@ -109,6 +113,7 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
request_data=request_data or {},
api_base=api_base,
model=model,
+ api_key=self.get_bedrock_bearer_token(optional_params),
)
def logging_non_streaming_response(
diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py
index 889361cd808..d877fbb4e09 100644
--- a/litellm/llms/bedrock_mantle/common_utils.py
+++ b/litellm/llms/bedrock_mantle/common_utils.py
@@ -13,6 +13,7 @@ global state.
"""
import re
+from collections.abc import Mapping
from typing import Final
from botocore.exceptions import (
@@ -31,30 +32,39 @@ BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1"
MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE)
+def resolve_mantle_bearer_token(api_key: str | None) -> str | None:
+ return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
+
+
+def resolve_mantle_region(params: Mapping[str, object]) -> str:
+ region: Final = params.get("aws_region_name")
+ if isinstance(region, str) and region:
+ BaseAWSLLM._validate_aws_region_name(region)
+ return region
+ api_base: Final = params.get("api_base")
+ base: Final = (api_base if isinstance(api_base, str) else None) or get_secret_str("BEDROCK_MANTLE_API_BASE")
+ if base:
+ match: Final = MANTLE_HOST_RE.match(base.rstrip("/"))
+ if match:
+ return match.group(1)
+ return (
+ get_secret_str("BEDROCK_MANTLE_REGION")
+ or get_secret_str("AWS_REGION_NAME")
+ or get_secret_str("AWS_REGION")
+ or BEDROCK_MANTLE_DEFAULT_REGION
+ )
+
+
class BedrockMantleAuthMixin:
_aws_signer: BaseAWSLLM
@staticmethod
def _resolve_bearer_token(api_key: str | None) -> str | None:
- return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
+ return resolve_mantle_bearer_token(api_key)
@staticmethod
def _resolve_region(params: dict) -> str:
- region: Final = params.get("aws_region_name")
- if region:
- BaseAWSLLM._validate_aws_region_name(region)
- return region
- base: Final = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE")
- if base:
- match: Final = MANTLE_HOST_RE.match(base.rstrip("/"))
- if match:
- return match.group(1)
- return (
- get_secret_str("BEDROCK_MANTLE_REGION")
- or get_secret_str("AWS_REGION_NAME")
- or get_secret_str("AWS_REGION")
- or BEDROCK_MANTLE_DEFAULT_REGION
- )
+ return resolve_mantle_region(params)
def sign_request(
self,
diff --git a/litellm/llms/bedrock_mantle/passthrough/transformation.py b/litellm/llms/bedrock_mantle/passthrough/transformation.py
new file mode 100644
index 00000000000..1393ac7c6e7
--- /dev/null
+++ b/litellm/llms/bedrock_mantle/passthrough/transformation.py
@@ -0,0 +1,44 @@
+from collections.abc import Mapping
+from typing import Final, Literal
+
+from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig
+from litellm.llms.bedrock_mantle.common_utils import (
+ MANTLE_HOST_RE,
+ resolve_mantle_bearer_token,
+ resolve_mantle_region,
+)
+
+
+class BedrockMantlePassthroughConfig(BedrockPassthroughConfig):
+ """Native Bedrock runtime passthrough (InvokeModel, Converse) for deployments declared as bedrock_mantle.
+
+ The Mantle host only serves the OpenAI-compatible surface, so a Mantle api_base lends its region and the
+ request itself goes to bedrock-runtime, signed with the deployment's Bearer token or SigV4 credentials.
+ """
+
+ def _get_aws_region_name(
+ self,
+ optional_params: Mapping[str, object],
+ model: str | None = None,
+ model_id: str | None = None,
+ ) -> str:
+ return resolve_mantle_region(optional_params)
+
+ def get_runtime_endpoint(
+ self,
+ api_base: str | None,
+ aws_bedrock_runtime_endpoint: str | None,
+ aws_region_name: str,
+ endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime",
+ ) -> tuple[str, str]:
+ is_mantle_host: Final = api_base is not None and MANTLE_HOST_RE.match(api_base.rstrip("/")) is not None
+ return super().get_runtime_endpoint(
+ api_base=None if is_mantle_host else api_base,
+ aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
+ aws_region_name=aws_region_name,
+ endpoint_type=endpoint_type,
+ )
+
+ def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None:
+ api_key: Final = litellm_params.get("api_key")
+ return resolve_mantle_bearer_token(api_key if isinstance(api_key, str) else None)
diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py
index 8a2ee2a3af8..4b30afb2f98 100644
--- a/litellm/passthrough/main.py
+++ b/litellm/passthrough/main.py
@@ -199,7 +199,7 @@ def llm_passthrough_route(
api_key=api_key,
)
- litellm_params_dict: Final = get_litellm_params(**kwargs)
+ litellm_params_dict: Final = get_litellm_params(api_key=api_key, api_base=api_base, **kwargs)
if client is None:
from litellm.llms.custom_httpx.http_handler import (
diff --git a/litellm/utils.py b/litellm/utils.py
index 012e8785321..5cbd0519032 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -8610,6 +8610,12 @@ class ProviderConfigManager:
)
return BedrockPassthroughConfig()
+ elif LlmProviders.BEDROCK_MANTLE == provider:
+ from litellm.llms.bedrock_mantle.passthrough.transformation import (
+ BedrockMantlePassthroughConfig,
+ )
+
+ return BedrockMantlePassthroughConfig()
elif LlmProviders.VLLM == provider or LlmProviders.HOSTED_VLLM == provider:
from litellm.llms.vllm.passthrough.transformation import (
VLLMPassthroughConfig,
diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
new file mode 100644
index 00000000000..b7f9e492e14
--- /dev/null
+++ b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
@@ -0,0 +1,152 @@
+import json
+from unittest.mock import MagicMock, patch
+
+import httpx
+import pytest
+from botocore.credentials import Credentials
+
+from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig
+from litellm.llms.bedrock_mantle.passthrough.transformation import BedrockMantlePassthroughConfig
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.passthrough.main import llm_passthrough_route
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager
+
+MANTLE_API_BASE = "https://bedrock-mantle.us-east-2.api.aws"
+INVOKE_ENDPOINT = "model/us.openai.gpt-5.6-sol/invoke"
+REQUEST_BODY = {"messages": [{"role": "user", "content": "say pong"}], "max_completion_tokens": 64}
+
+
+@pytest.fixture
+def no_ambient_aws(monkeypatch):
+ for name in (
+ "AWS_BEARER_TOKEN_BEDROCK",
+ "BEDROCK_MANTLE_API_KEY",
+ "BEDROCK_MANTLE_API_BASE",
+ "BEDROCK_MANTLE_REGION",
+ "AWS_BEDROCK_RUNTIME_ENDPOINT",
+ "AWS_REGION_NAME",
+ "AWS_REGION",
+ "AWS_DEFAULT_REGION",
+ ):
+ monkeypatch.delenv(name, raising=False)
+
+
+def test_bedrock_mantle_registers_its_own_bedrock_passthrough_config():
+ config = ProviderConfigManager.get_provider_passthrough_config(
+ model="us.openai.gpt-5.6-sol", provider=LlmProviders.BEDROCK_MANTLE
+ )
+ assert isinstance(config, BedrockMantlePassthroughConfig)
+ assert isinstance(config, BedrockPassthroughConfig)
+
+
+def test_mantle_api_base_only_lends_its_region_to_the_runtime_url(no_ambient_aws):
+ url, base_url = BedrockMantlePassthroughConfig().get_complete_url(
+ api_base=MANTLE_API_BASE,
+ api_key=None,
+ model="us.openai.gpt-5.6-sol",
+ endpoint=INVOKE_ENDPOINT,
+ request_query_params=None,
+ litellm_params={"api_base": MANTLE_API_BASE},
+ )
+ assert str(url) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}"
+ assert base_url == "https://bedrock-runtime.us-east-2.amazonaws.com"
+
+
+def test_explicit_region_and_non_mantle_api_base_are_kept(no_ambient_aws):
+ vpc_endpoint = "https://vpce-0123.bedrock-runtime.us-east-1.vpce.amazonaws.com"
+ url, base_url = BedrockMantlePassthroughConfig().get_complete_url(
+ api_base=vpc_endpoint,
+ api_key=None,
+ model="us.openai.gpt-5.6-sol",
+ endpoint=INVOKE_ENDPOINT,
+ request_query_params=None,
+ litellm_params={"api_base": vpc_endpoint, "aws_region_name": "us-east-1"},
+ )
+ assert str(url) == f"{vpc_endpoint}/{INVOKE_ENDPOINT}"
+ assert base_url == vpc_endpoint
+
+
+def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws):
+ url, _ = BedrockMantlePassthroughConfig().get_complete_url(
+ api_base=None,
+ api_key=None,
+ model="us.openai.gpt-5.6-sol",
+ endpoint=INVOKE_ENDPOINT,
+ request_query_params=None,
+ litellm_params={},
+ )
+ assert str(url) == f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}"
+
+
+@pytest.mark.parametrize(
+ ("litellm_params", "env", "expected_bearer"),
+ [
+ ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"),
+ ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"),
+ ({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"),
+ ],
+)
+def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer):
+ for name, value in env.items():
+ monkeypatch.setenv(name, value)
+ headers, body = BedrockMantlePassthroughConfig().sign_request(
+ headers={},
+ litellm_params=litellm_params,
+ request_data=REQUEST_BODY,
+ api_base=f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}",
+ model="us.openai.gpt-5.6-sol",
+ )
+ assert headers["Authorization"] == f"Bearer {expected_bearer}"
+ assert body is not None
+ assert json.loads(body) == REQUEST_BODY
+
+
+def test_sign_request_falls_back_to_sigv4_scoped_to_the_mantle_region(no_ambient_aws):
+ config = BedrockMantlePassthroughConfig()
+ with patch.object(config, "get_credentials", return_value=Credentials("AKIA", "secret")):
+ headers, body = config.sign_request(
+ headers={},
+ litellm_params={"api_base": MANTLE_API_BASE},
+ request_data=REQUEST_BODY,
+ api_base=f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}",
+ model="us.openai.gpt-5.6-sol",
+ )
+ assert headers["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIA/")
+ assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"]
+ assert body is not None
+ assert json.loads(body) == REQUEST_BODY
+
+
+@pytest.mark.parametrize(
+ ("route_kwargs", "env", "expected_bearer"),
+ [
+ ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"),
+ ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"),
+ ],
+)
+def test_invoke_passthrough_route_reaches_bedrock_runtime_for_a_mantle_deployment(
+ no_ambient_aws, monkeypatch, route_kwargs, env, expected_bearer
+):
+ for name, value in env.items():
+ monkeypatch.setenv(name, value)
+ client = HTTPHandler()
+ with (
+ patch.object(client.client, "send", return_value=MagicMock(status_code=200)),
+ patch.object(client.client, "build_request", wraps=client.client.build_request) as build_request,
+ ):
+ response = llm_passthrough_route(
+ model="bedrock_mantle/us.openai.gpt-5.6-sol",
+ endpoint=INVOKE_ENDPOINT,
+ method="POST",
+ api_base=MANTLE_API_BASE,
+ json=dict(REQUEST_BODY),
+ client=client,
+ litellm_logging_obj=MagicMock(),
+ **route_kwargs,
+ )
+ assert response.status_code == 200
+ sent = build_request.call_args.kwargs
+ assert str(sent["url"]) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}"
+ assert sent["headers"]["Authorization"] == f"Bearer {expected_bearer}"
+ assert json.loads(sent["content"]) == REQUEST_BODY
From b46f17faf5d56ce853fff94d0daa5ffa0d2fb428 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:18:55 -0700
Subject: [PATCH 32/70] fix(together_ai): default endpoints to api.together.ai
instead of api.together.xyz
Together AI moved its canonical API host from api.together.xyz to
api.together.ai. Default the provider api_base and the rerank handler to
the new host, make rerank honor api_base and TOGETHER_AI_API_BASE like
chat already does, map both hosts to together_ai when passed as
api_base, and delete the dead models/info fetch in factory.py.
---
basedpyright-code-budget.json | 8 +--
litellm/constants.py | 1 +
.../get_llm_provider_logic.py | 10 ++-
.../prompt_templates/factory.py | 43 ------------
litellm/llms/together_ai/rerank/handler.py | 12 +++-
litellm/rerank_api/main.py | 3 +
ruff-strict-budget.json | 6 +-
.../test_get_llm_provider_endpoint_match.py | 39 +++++++++++
tests/test_litellm/rerank_api/test_main.py | 67 +++++++++++++++++++
type-discipline-budget.json | 2 +-
10 files changed, 136 insertions(+), 55 deletions(-)
diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json
index 664e1669834..f4d4e25859a 100644
--- a/basedpyright-code-budget.json
+++ b/basedpyright-code-budget.json
@@ -1,6 +1,6 @@
{
"reportAny": {
- "limit": 19955
+ "limit": 19949
},
"reportArgumentType": {
"limit": 2566
@@ -54,7 +54,7 @@
"limit": 0
},
"reportMissingParameterType": {
- "limit": 5663
+ "limit": 5661
},
"reportMissingTypeArgument": {
"limit": 15555
@@ -105,10 +105,10 @@
"limit": 109
},
"reportUnknownMemberType": {
- "limit": 39011
+ "limit": 39009
},
"reportUnknownParameterType": {
- "limit": 19885
+ "limit": 19883
},
"reportUnknownVariableType": {
"limit": 30569
diff --git a/litellm/constants.py b/litellm/constants.py
index 0a1ada3bab2..78aba30f9c0 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -750,6 +750,7 @@ openai_compatible_endpoints: Final[list] = [
"api.groq.com/openai/v1",
"https://integrate.api.nvidia.com/v1",
"api.deepseek.com/v1",
+ "api.together.ai/v1",
"api.together.xyz/v1",
"app.empower.dev/api/v1",
"https://api.friendli.ai/serverless/v1",
diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py
index e674fc37673..d2d82064c47 100644
--- a/litellm/litellm_core_utils/get_llm_provider_logic.py
+++ b/litellm/litellm_core_utils/get_llm_provider_logic.py
@@ -272,6 +272,14 @@ def get_llm_provider(
elif endpoint == "api.deepseek.com/v1":
custom_llm_provider = "deepseek"
dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY")
+ elif endpoint == "api.together.ai/v1" or endpoint == "api.together.xyz/v1":
+ custom_llm_provider = "together_ai"
+ dynamic_api_key = (
+ get_secret_str("TOGETHER_API_KEY")
+ or get_secret_str("TOGETHER_AI_API_KEY")
+ or get_secret_str("TOGETHERAI_API_KEY")
+ or get_secret_str("TOGETHER_AI_TOKEN")
+ )
elif endpoint == "ollama.com":
custom_llm_provider = "ollama"
dynamic_api_key = get_secret_str("OLLAMA_API_KEY")
@@ -707,7 +715,7 @@ def _get_openai_compatible_provider_info(
dynamic_api_key,
) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key)
elif custom_llm_provider == "together_ai":
- api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.xyz/v1"
+ api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.ai/v1"
dynamic_api_key = api_key or (
get_secret_str("TOGETHER_API_KEY")
or get_secret_str("TOGETHER_AI_API_KEY")
diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py
index 826a890eca9..86cfbf70255 100644
--- a/litellm/litellm_core_utils/prompt_templates/factory.py
+++ b/litellm/litellm_core_utils/prompt_templates/factory.py
@@ -643,49 +643,6 @@ def claude_2_1_pt(
return prompt
-### TOGETHER AI
-
-
-def get_model_info(token, model):
- try:
- headers: Final = {"Authorization": f"Bearer {token}"}
- client: Final = HTTPHandler(concurrent_limit=1)
- response: Final = client.get("https://api.together.xyz/models/info", headers=headers)
- if response.status_code == 200:
- model_info: Final = response.json()
- for m in model_info:
- if m["name"].lower().strip() == model.strip():
- return m["config"].get("prompt_format", None), m["config"].get("chat_template", None)
- return None, None
- else:
- return None, None
- except Exception: # safely fail a prompt template request
- return None, None
-
-
-## OLD TOGETHER AI FLOW
-# def format_prompt_togetherai(messages, prompt_format, chat_template):
-# if prompt_format is None:
-# return default_pt(messages)
-
-# human_prompt, assistant_prompt = prompt_format.split("{prompt}")
-
-# if chat_template is not None:
-# prompt = hf_chat_template(
-# model=None, messages=messages, chat_template=chat_template
-# )
-# elif prompt_format is not None:
-# prompt = custom_prompt(
-# role_dict={},
-# messages=messages,
-# initial_prompt_value=human_prompt,
-# final_prompt_value=assistant_prompt,
-# )
-# else:
-# prompt = default_pt(messages)
-# return prompt
-
-
### IBM Granite
diff --git a/litellm/llms/together_ai/rerank/handler.py b/litellm/llms/together_ai/rerank/handler.py
index 10246451a9d..8407018b898 100644
--- a/litellm/llms/together_ai/rerank/handler.py
+++ b/litellm/llms/together_ai/rerank/handler.py
@@ -16,11 +16,16 @@ from litellm.llms.together_ai.rerank.transformation import TogetherAIRerankConfi
from litellm.types.rerank import RerankRequest, RerankResponse
+def _rerank_url(api_base: str) -> str:
+ return f"{api_base.rstrip('/')}/rerank"
+
+
class TogetherAIRerank(BaseLLM):
def rerank(
self,
model: str,
api_key: str,
+ api_base: str,
query: str,
documents: list[str | dict[str, Any]],
top_n: int | None = None,
@@ -46,10 +51,10 @@ class TogetherAIRerank(BaseLLM):
raise ValueError("TogetherAI does not support max_chunks_per_doc")
if _is_async:
- return self.async_rerank(request_data_dict, api_key) # Call async method
+ return self.async_rerank(request_data_dict, api_key, api_base) # Call async method
response: Final = client.post(
- "https://api.together.xyz/v1/rerank",
+ _rerank_url(api_base),
headers={
"accept": "application/json",
"content-type": "application/json",
@@ -69,11 +74,12 @@ class TogetherAIRerank(BaseLLM):
self,
request_data_dict: dict[str, Any],
api_key: str,
+ api_base: str,
) -> RerankResponse:
client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.TOGETHER_AI) # Use async client
response: Final = await client.post(
- "https://api.together.xyz/v1/rerank",
+ _rerank_url(api_base),
headers={
"accept": "application/json",
"content-type": "application/json",
diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py
index 15a6f18a6bb..c8f7842aebf 100644
--- a/litellm/rerank_api/main.py
+++ b/litellm/rerank_api/main.py
@@ -277,6 +277,8 @@ def rerank(
if api_key is None:
raise ValueError("TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment")
+ api_base = dynamic_api_base or optional_params.api_base or litellm.api_base or "https://api.together.ai/v1"
+
response = together_rerank.rerank(
model=model,
query=query,
@@ -286,6 +288,7 @@ def rerank(
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
api_key=api_key,
+ api_base=api_base,
_is_async=_is_async,
)
elif _custom_llm_provider == litellm.LlmProviders.JINA_AI:
diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json
index 03318718fb5..1ca152985f9 100644
--- a/ruff-strict-budget.json
+++ b/ruff-strict-budget.json
@@ -1,6 +1,6 @@
{
"ANN001": {
- "limit": 3018
+ "limit": 3016
},
"ANN002": {
"limit": 71
@@ -9,7 +9,7 @@
"limit": 827
},
"ANN201": {
- "limit": 2016
+ "limit": 2015
},
"ANN202": {
"limit": 852
@@ -57,7 +57,7 @@
"limit": 3
},
"BLE001": {
- "limit": 2919
+ "limit": 2918
},
"C401": {
"limit": 8
diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
index bda7ab4afc6..5c20284282a 100644
--- a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
+++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
@@ -133,3 +133,42 @@ class TestGetLlmProviderRejectsAttackerSmuggledApiBase:
assert provider == "groq"
assert dynamic_api_key == "server-real-groq-key"
+
+
+class TestTogetherApiBaseResolvesProvider:
+ """
+ Regression for the Together host migration: both the current
+ ``api.together.ai`` host and the legacy ``api.together.xyz`` host must
+ resolve to ``together_ai`` when passed as ``api_base``. Before the fix
+ the endpoint list carried the legacy host but the provider-mapping
+ chain had no branch for it, so the match fell through with a None
+ provider and the deployment failed with "LLM Provider NOT provided".
+ """
+
+ @pytest.mark.parametrize(
+ "api_base",
+ [
+ "https://api.together.ai/v1",
+ "https://api.together.xyz/v1",
+ ],
+ )
+ def test_together_api_base_resolves_to_together_ai(self, api_base, monkeypatch):
+ monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env")
+
+ model, provider, dynamic_api_key, returned_api_base = get_llm_provider(
+ model="some-model",
+ api_base=api_base,
+ )
+
+ assert provider == "together_ai"
+ assert dynamic_api_key == "together-key-from-env"
+ assert returned_api_base == api_base
+ assert model == "some-model"
+
+ def test_together_default_api_base_is_together_ai(self, monkeypatch):
+ monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False)
+
+ _, provider, _, api_base = get_llm_provider(model="together_ai/some-model")
+
+ assert provider == "together_ai"
+ assert api_base == "https://api.together.ai/v1"
diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py
index 85777afe81c..587be59c550 100644
--- a/tests/test_litellm/rerank_api/test_main.py
+++ b/tests/test_litellm/rerank_api/test_main.py
@@ -1,6 +1,10 @@
import logging
from unittest.mock import MagicMock, patch
+import httpx
+import pytest
+import respx
+
import litellm
@@ -62,3 +66,66 @@ def test_rerank_does_not_log_request_content_at_info(caplog):
assert all(
r.levelno == logging.DEBUG for r in optional_params_logs
), "optional_rerank_params must be logged at DEBUG, not INFO"
+
+
+TOGETHER_RERANK_BODY = {
+ "id": "rerank-mock-id",
+ "results": [{"index": 0, "relevance_score": 0.95}],
+ "usage": {"prompt_tokens": 10, "total_tokens": 10},
+}
+
+
+def test_together_rerank_defaults_to_together_ai_host(respx_mock: respx.MockRouter, monkeypatch):
+ """Regression for the Together host migration: rerank used to hardcode
+ https://api.together.xyz/v1/rerank. The default must now be api.together.ai."""
+ monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False)
+
+ mock_route = respx_mock.post("https://api.together.ai/v1/rerank")
+ mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
+
+ response = litellm.rerank(
+ model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
+ query=MARKER_QUERY,
+ documents=[MARKER_DOC],
+ api_key="fake-together-key",
+ )
+
+ assert mock_route.called
+ assert response.results[0]["relevance_score"] == 0.95
+
+
+def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter):
+ """Regression: a custom api_base was silently ignored by the Together rerank handler."""
+ mock_route = respx_mock.post("https://custom-together.example/v1/rerank")
+ mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
+
+ litellm.rerank(
+ model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
+ query=MARKER_QUERY,
+ documents=[MARKER_DOC],
+ api_key="fake-together-key",
+ api_base="https://custom-together.example/v1",
+ )
+
+ assert mock_route.called
+ assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key"
+
+
+@pytest.mark.asyncio
+async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch):
+ """Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank."""
+ monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://env-together.example/v1")
+ monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
+
+ mock_route = respx_mock.post("https://env-together.example/v1/rerank")
+ mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
+
+ response = await litellm.arerank(
+ model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
+ query=MARKER_QUERY,
+ documents=[MARKER_DOC],
+ api_key="fake-together-key",
+ )
+
+ assert mock_route.called
+ assert response.results[0]["relevance_score"] == 0.95
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 627811a7f1d..f9fe3042f1e 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -3,7 +3,7 @@
"limit": 22805
},
"LIT002": {
- "limit": 26873
+ "limit": 26872
},
"LIT003": {
"limit": 269
From fd1dca05deb0fd4d50153b42412f716e341236c3 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:19:56 -0700
Subject: [PATCH 33/70] fix(bedrock_mantle): type the codex item dispatcher
without Any
---
litellm/llms/bedrock_mantle/responses/transformation.py | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py
index 9068a641940..3e5dd4ff87d 100644
--- a/litellm/llms/bedrock_mantle/responses/transformation.py
+++ b/litellm/llms/bedrock_mantle/responses/transformation.py
@@ -288,7 +288,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
return rewritten
@classmethod
- def _normalize_codex_input_item(cls, item: object) -> "tuple[Any, str | None]":
+ def _normalize_codex_input_item(cls, item: object) -> "tuple[object, str | None]":
"""Returns (normalized item or None to drop it, original type when rewritten)."""
if not isinstance(item, dict):
return item, None
@@ -324,7 +324,8 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
"Bedrock Mantle Responses API: rewrote Codex input item type(s) %s that Mantle rejects.",
rewritten_types,
)
- return [item for item, _ in normalized if item is not None] # mutable-ok: ResponseInputParam is a list
+ kept: Final = [item for item, _ in normalized if item is not None] # mutable-ok: ResponseInputParam is a list
+ return kept # pyright: ignore[reportReturnType] # Codex passthrough items sit outside the OpenAI input union
def map_openai_params(
self,
From 6be000f1f35091cbbdfedb7df4e0dd8d494c0eaf Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:23:12 -0700
Subject: [PATCH 34/70] test(bedrock_mantle): type _repo_cost_map return
instead of bare dict
---
.../test_bedrock_mantle_responses_transformation.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
index fd279a2bc1f..28c6060e5cc 100644
--- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py
@@ -1568,7 +1568,7 @@ class TestBedrockMantleResponsesPricing:
assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models
-def _repo_cost_map(map_name: str) -> dict:
+def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]:
repo_root = Path(__file__).resolve().parents[4]
paths = {
"root": repo_root / "model_prices_and_context_window.json",
From db8c49305ead5151c2c9b7d9bdfb5cf388b1c140 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:24:28 -0700
Subject: [PATCH 35/70] test: type the cost map fixture instead of bare dict
---
.../test_together_ai_model_metadata.py | 38 ++++++++++++-------
1 file changed, 25 insertions(+), 13 deletions(-)
diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py
index 7d7712f3c56..5a0aadf4737 100644
--- a/tests/test_litellm/test_together_ai_model_metadata.py
+++ b/tests/test_litellm/test_together_ai_model_metadata.py
@@ -3,11 +3,15 @@ from pathlib import Path
from typing import Final
import pytest
+from pydantic import TypeAdapter
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
REPO_ROOT: Final = Path(__file__).parents[2]
+CostMap = dict[str, dict[str, object]]
+COST_MAP_ADAPTER: Final = TypeAdapter(CostMap)
+
SERVERLESS_CHAT_MODELS: Final = (
"together_ai/moonshotai/Kimi-K3",
"together_ai/zai-org/GLM-5.2",
@@ -66,13 +70,13 @@ DEPRECATED_MODELS: Final = {
@pytest.fixture(scope="module")
-def cost_map() -> dict:
+def cost_map() -> CostMap:
with open(REPO_ROOT / "model_prices_and_context_window.json") as f:
- return json.load(f)
+ return COST_MAP_ADAPTER.validate_python(json.load(f))
@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS)
-def test_together_serverless_chat_model_is_mapped(cost_map: dict, model: str):
+def test_together_serverless_chat_model_is_mapped(cost_map: CostMap, model: str):
info = cost_map.get(model)
assert info is not None, f"{model} missing from model_prices_and_context_window.json"
assert info["litellm_provider"] == "together_ai"
@@ -86,7 +90,7 @@ def test_together_serverless_chat_model_is_mapped(cost_map: dict, model: str):
assert provider == "together_ai"
-def test_together_kimi_k3_pricing_and_capabilities(cost_map: dict):
+def test_together_kimi_k3_pricing_and_capabilities(cost_map: CostMap):
info = cost_map["together_ai/moonshotai/Kimi-K3"]
assert info["input_cost_per_token"] == 3e-06
assert info["output_cost_per_token"] == 1.5e-05
@@ -98,7 +102,7 @@ def test_together_kimi_k3_pricing_and_capabilities(cost_map: dict):
assert info["supports_reasoning"] is True
-def test_together_glm_52_pricing(cost_map: dict):
+def test_together_glm_52_pricing(cost_map: CostMap):
info = cost_map["together_ai/zai-org/GLM-5.2"]
assert info["input_cost_per_token"] == 1.4e-06
assert info["output_cost_per_token"] == 4.4e-06
@@ -106,7 +110,7 @@ def test_together_glm_52_pricing(cost_map: dict):
assert info["supports_reasoning"] is True
-def test_together_multilingual_e5_embedding_entry(cost_map: dict):
+def test_together_multilingual_e5_embedding_entry(cost_map: CostMap):
info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"]
assert info["mode"] == "embedding"
assert info["input_cost_per_token"] == 2e-08
@@ -114,7 +118,7 @@ def test_together_multilingual_e5_embedding_entry(cost_map: dict):
assert info["output_vector_size"] == 1024
-def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: dict):
+def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: CostMap):
info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"]
assert info["input_cost_per_token"] == 1.04e-06
assert info["output_cost_per_token"] == 1.04e-06
@@ -122,17 +126,25 @@ def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: dict)
@pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS))
-def test_together_deprecated_model_carries_deprecation_date(cost_map: dict, model: str):
+def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str):
info = cost_map.get(model)
assert info is not None, f"{model} missing from model_prices_and_context_window.json"
assert info.get("deprecation_date") == DEPRECATED_MODELS[model]
-def test_together_successor_metadata_points_at_live_models(cost_map: dict):
+def _successor(info: dict[str, object]) -> str | None:
+ metadata = info.get("metadata")
+ if not isinstance(metadata, dict):
+ return None
+ successor = metadata.get("successor")
+ return successor if isinstance(successor, str) else None
+
+
+def test_together_successor_metadata_points_at_live_models(cost_map: CostMap):
successors = {
- model: info["metadata"]["successor"]
+ model: successor
for model, info in cost_map.items()
- if model.startswith("together_ai/") and "successor" in info.get("metadata", {})
+ if model.startswith("together_ai/") and (successor := _successor(info)) is not None
}
assert len(successors) >= 10
for model, successor in successors.items():
@@ -141,9 +153,9 @@ def test_together_successor_metadata_points_at_live_models(cost_map: dict):
assert "deprecation_date" not in target, f"{model} names deprecated successor {successor}"
-def test_together_backup_cost_map_in_sync(cost_map: dict):
+def test_together_backup_cost_map_in_sync(cost_map: CostMap):
with open(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json") as f:
- backup = json.load(f)
+ backup = COST_MAP_ADAPTER.validate_python(json.load(f))
together_main = {k: v for k, v in cost_map.items() if k.startswith("together_ai/")}
together_backup = {k: v for k, v in backup.items() if k.startswith("together_ai/")}
assert together_backup == together_main
From 0bd4d323da3a5969a7a3940faef03ecd239d2a98 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:33:40 -0700
Subject: [PATCH 36/70] fix(router): resolve provider from api_base in
deployment validation and acompletion
Router._add_deployment called get_llm_provider without the deployment's api_base, so a config entry with a bare model plus a known OpenAI-compatible endpoint failed startup validation with LLM Provider NOT provided and the proxy returned 400 no healthy deployments for that model group. acompletion had the same gap at request time: it forwarded only base_url into its get_llm_provider call, dropping the api_base kwarg the router passes. Both now forward api_base so endpoint matching resolves the provider the same way sync completion already does
---
litellm/main.py | 2 +-
litellm/router.py | 1 +
tests/test_litellm/test_main.py | 13 +++++++
tests/test_litellm/test_router.py | 62 +++++++++++++++++++++++++++++++
4 files changed, 77 insertions(+), 1 deletion(-)
diff --git a/litellm/main.py b/litellm/main.py
index d3967473f99..6dfd8c2d675 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -602,7 +602,7 @@ async def acompletion(
_, custom_llm_provider, _, _ = get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
- api_base=base_url,
+ api_base=kwargs.get("api_base") or base_url,
)
fallbacks = fallbacks or litellm.model_fallbacks
diff --git a/litellm/router.py b/litellm/router.py
index 6ee474730c9..33da39e2677 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -8341,6 +8341,7 @@ class Router:
) = litellm.get_llm_provider(
model=deployment.litellm_params.model,
custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None),
+ api_base=deployment.litellm_params.api_base,
)
# done reading model["litellm_params"]
# Check if provider is supported: either in enum or JSON-configured
diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py
index 99b1cc826aa..8f2b06be4b3 100644
--- a/tests/test_litellm/test_main.py
+++ b/tests/test_litellm/test_main.py
@@ -2944,3 +2944,16 @@ def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map):
assert cost == pytest.approx(
_priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens)
)
+
+
+@pytest.mark.asyncio
+async def test_acompletion_resolves_provider_from_api_base():
+ response = await litellm.acompletion(
+ model="deepseek-chat",
+ api_base="https://api.deepseek.com/v1",
+ api_key="fake-key",
+ messages=[{"role": "user", "content": "hi"}],
+ mock_response="resolved",
+ )
+
+ assert response.choices[0].message.content == "resolved"
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 910b874c2ac..df6754cd7ca 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -8921,3 +8921,65 @@ class TestAzureBaseModelFallbackLogging:
deployment=None, received_model_name="my-group", id="azure-base-model-test-id"
)
assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"]
+
+
+class TestAddDeploymentApiBaseProviderResolution:
+ def test_bare_model_with_known_api_base_initializes(self):
+ router = litellm.Router(
+ model_list=[
+ {
+ "model_name": "groq-pinned",
+ "litellm_params": {
+ "model": "llama-3.3-70b-versatile",
+ "api_base": "https://api.groq.com/openai/v1",
+ "api_key": "fake-key",
+ },
+ },
+ {
+ "model_name": "deepseek-pinned",
+ "litellm_params": {
+ "model": "deepseek-chat",
+ "api_base": "https://api.deepseek.com/v1",
+ "api_key": "fake-key",
+ },
+ },
+ ]
+ )
+
+ model_list = router.get_model_list()
+ assert model_list is not None
+ assert {m["model_name"] for m in model_list} == {"groq-pinned", "deepseek-pinned"}
+
+ def test_bare_model_with_unknown_api_base_still_raises(self):
+ with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"):
+ litellm.Router(
+ model_list=[
+ {
+ "model_name": "mystery",
+ "litellm_params": {
+ "model": "some-unknown-model",
+ "api_base": "https://llm.internal.example.com/v1",
+ "api_key": "fake-key",
+ },
+ }
+ ]
+ )
+
+ def test_explicit_custom_llm_provider_beats_api_base_endpoint_match(self):
+ router = litellm.Router(
+ model_list=[
+ {
+ "model_name": "openai-via-gateway",
+ "litellm_params": {
+ "model": "gpt-3.5-turbo",
+ "custom_llm_provider": "openai",
+ "api_base": "https://api.groq.com/openai/v1",
+ "api_key": "fake-key",
+ },
+ }
+ ]
+ )
+
+ deployment = router.get_deployment_by_model_group_name("openai-via-gateway")
+ assert deployment is not None
+ assert deployment.litellm_params.custom_llm_provider == "openai"
From 367a6e5dc5fc1e23d31b7395d96fa2a5332fd6d2 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:37:48 -0700
Subject: [PATCH 37/70] test(router): pin the guard that keeps a junk-typed
operator effort value out of model group info
---
tests/test_litellm/test_router.py | 34 +++++++++++++++++++++++++++++++
1 file changed, 34 insertions(+)
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 5f1941574eb..e8830d05065 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -9092,6 +9092,40 @@ def test_model_group_info_reasoning_efforts_ignore_a_value_declared_in_model_inf
assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high")
+def test_model_group_info_survives_a_junk_typed_operator_effort_value():
+ """A deployment's registered model_info reads back with whatever the operator wrote under any
+ key, so a wrong-typed supported_reasoning_efforts must not fail the group's info. Only the
+ constructor's trailing override keeps the junk away from ModelGroupInfo validation."""
+ router = litellm.Router(
+ model_list=[
+ {
+ "model_name": "junk-declared-group",
+ "litellm_params": {"model": "openai/lone-reasoner"},
+ "model_info": {"id": "junk-deployment"},
+ },
+ ]
+ )
+
+ def _model_info(model_id: str, model_name: str):
+ return {
+ "key": model_name,
+ "litellm_provider": "openai",
+ "mode": "chat",
+ "supports_reasoning": True,
+ "supports_none_reasoning_effort": True,
+ "supported_reasoning_efforts": "high",
+ }
+
+ with patch.object(router, "get_deployment_model_info", side_effect=_model_info):
+ result = router._set_model_group_info(
+ model_group="junk-declared-group",
+ user_facing_model_group_name="junk-declared-group",
+ )
+
+ assert result is not None
+ assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high")
+
+
def test_model_group_info_reasoning_efforts_ignore_a_mode_the_operator_declared():
"""A deployment is registered in the cost map under its own id with whatever model_info the
operator wrote, so a mode they set themselves reads back exactly like one the map supplied. Only
From 5e6b6c6281f8d26bda113cdd54be0fe71afa4d77 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:39:10 -0700
Subject: [PATCH 38/70] fix(together_ai): let an explicit api_key beat the
Together env key on api_base match
---
litellm/litellm_core_utils/get_llm_provider_logic.py | 2 +-
litellm/llms/together_ai/rerank/handler.py | 2 +-
.../test_get_llm_provider_endpoint_match.py | 12 ++++++++++++
3 files changed, 14 insertions(+), 2 deletions(-)
diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py
index d2d82064c47..005e94ebe82 100644
--- a/litellm/litellm_core_utils/get_llm_provider_logic.py
+++ b/litellm/litellm_core_utils/get_llm_provider_logic.py
@@ -274,7 +274,7 @@ def get_llm_provider(
dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY")
elif endpoint == "api.together.ai/v1" or endpoint == "api.together.xyz/v1":
custom_llm_provider = "together_ai"
- dynamic_api_key = (
+ dynamic_api_key = api_key or (
get_secret_str("TOGETHER_API_KEY")
or get_secret_str("TOGETHER_AI_API_KEY")
or get_secret_str("TOGETHERAI_API_KEY")
diff --git a/litellm/llms/together_ai/rerank/handler.py b/litellm/llms/together_ai/rerank/handler.py
index 8407018b898..b8079e52c97 100644
--- a/litellm/llms/together_ai/rerank/handler.py
+++ b/litellm/llms/together_ai/rerank/handler.py
@@ -51,7 +51,7 @@ class TogetherAIRerank(BaseLLM):
raise ValueError("TogetherAI does not support max_chunks_per_doc")
if _is_async:
- return self.async_rerank(request_data_dict, api_key, api_base) # Call async method
+ return self.async_rerank(request_data_dict, api_key, api_base)
response: Final = client.post(
_rerank_url(api_base),
diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
index 5c20284282a..6cacd119030 100644
--- a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
+++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py
@@ -165,6 +165,18 @@ class TestTogetherApiBaseResolvesProvider:
assert returned_api_base == api_base
assert model == "some-model"
+ def test_explicit_api_key_beats_together_env_key(self, monkeypatch):
+ monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env")
+
+ _, provider, dynamic_api_key, _ = get_llm_provider(
+ model="some-model",
+ api_base="https://api.together.ai/v1",
+ api_key="explicit-caller-key",
+ )
+
+ assert provider == "together_ai"
+ assert dynamic_api_key == "explicit-caller-key"
+
def test_together_default_api_base_is_together_ai(self, monkeypatch):
monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False)
From 0ea6f5e159350aa75e9118ba027b7286074cb0fe Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Date: Tue, 25 Aug 2026 10:50:11 -0700
Subject: [PATCH 39/70] fix(azure_ai): stamp the model router's selected model
instead of matching on the model name
The model Azure Model Router served was recovered by checking whether the text
"model_router" or "model-router" appeared in a model string. Spend logs applied that
check to the litellm model path, where the route prefix guarantees a match, but the
proxy applied it to the client's model group alias, which carries no prefix. A model
group named anything else therefore lost the selected model in both the response and
the spend row.
AzureModelRouterConfig now stamps the served model onto _hidden_params, and the spend
log payload and the proxy's response restamping read that stamp. The name heuristic
survives as a fallback for callers with no response in hand, routed through
get_azure_ai_route so it lives in one place.
---
litellm/litellm_core_utils/litellm_logging.py | 12 +-
.../azure_model_router/transformation.py | 23 +-
litellm/llms/azure_ai/common_utils.py | 36 ++
litellm/proxy/common_request_processing.py | 21 +-
.../test_litellm_logging.py | 472 +++++++-----------
.../chat/test_azure_ai_transformation.py | 97 +++-
.../proxy/test_common_request_processing.py | 443 +++++++---------
7 files changed, 538 insertions(+), 566 deletions(-)
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 9b7707eabe1..01802c3bbb3 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -5762,11 +5762,15 @@ def get_standard_logging_object_payload(
response_model_name = final_response_obj.get("model")
# For Azure Model Router, preserve the actual model in the top-level standard
- # logging payload only when the user has opted in.
+ # logging payload.
+ from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
requested_model: Final = kwargs.get("model")
- if (
- isinstance(requested_model, str)
- and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower())
+ stamped_selected_model: Final = AzureFoundryModelInfo.get_model_router_selected_model(hidden_params)
+ if stamped_selected_model is not None:
+ model_name = stamped_selected_model
+ elif (
+ AzureFoundryModelInfo.is_model_router_call(model=requested_model, hidden_params=hidden_params)
and isinstance(response_model_name, str)
and response_model_name
):
diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py
index 1a924088390..d33564c0f9a 100644
--- a/litellm/llms/azure_ai/azure_model_router/transformation.py
+++ b/litellm/llms/azure_ai/azure_model_router/transformation.py
@@ -65,15 +65,24 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07)
and returns it with the azure_ai/ prefix for proper display and cost tracking.
+
+ Also stamps that model onto ``_hidden_params`` so downstream consumers (spend logs,
+ response restamping) can read it instead of guessing the route from the model string.
"""
- from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+ from litellm.llms.azure_ai.common_utils import (
+ AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
+ AzureFoundryModelInfo,
+ )
+ from litellm.router_utils.add_retry_fallback_headers import (
+ get_hidden_params_dict,
+ )
# Get base model for the parent call (strips routing prefixes for API compatibility)
base_model: Final[str] = AzureFoundryModelInfo.get_base_model(model)
# Call parent transform_response first - this will extract the actual model
# from the raw response (e.g., "gpt-5-nano-2025-08-07")
- model_response = super().transform_response(
+ transformed_response: Final = super().transform_response(
model=base_model,
raw_response=raw_response,
model_response=model_response,
@@ -86,7 +95,15 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
api_key=api_key,
json_mode=json_mode,
)
- return model_response
+ selected_model: Final = transformed_response.model
+ if selected_model:
+ # Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a
+ # class-level dict, so an in-place write can bleed into unrelated responses.
+ transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter
+ **get_hidden_params_dict(transformed_response),
+ AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model,
+ }
+ return transformed_response
def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> dict | None:
"""
diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py
index d25a8fd6561..9d37d8d1185 100644
--- a/litellm/llms/azure_ai/common_utils.py
+++ b/litellm/llms/azure_ai/common_utils.py
@@ -1,3 +1,4 @@
+from collections.abc import Mapping
from typing import Final, Literal
import litellm
@@ -5,6 +6,8 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
+AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
+
class AzureFoundryModelInfo(BaseLLMModelInfo):
"""Model info for Azure AI / Azure Foundry models."""
@@ -37,6 +40,39 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
return "model_router"
return "default"
+ @staticmethod
+ def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None:
+ """The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``.
+
+ Reading this beats re-deriving the route from a model string: the stamp is set on the
+ code path that was actually taken, so it holds no matter what the caller named the model.
+ """
+ if not hidden_params:
+ return None
+ selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY)
+ if isinstance(selected, str) and selected:
+ return selected
+ return None
+
+ @staticmethod
+ def is_model_router_call(
+ model: str | None = None,
+ hidden_params: Mapping[str, object] | None = None,
+ ) -> bool:
+ """Whether a request went down the Azure Model Router route.
+
+ Prefers the response stamp, then the deployment's litellm model path, and only then the
+ caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router
+ name heuristic lives in exactly one place.
+ """
+ if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None:
+ return True
+ deployment_model: Final = (hidden_params or {}).get("litellm_model_name") or (hidden_params or {}).get("model")
+ return any(
+ isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router"
+ for candidate in (deployment_model, model)
+ )
+
@staticmethod
def get_api_base(api_base: str | None = None) -> str | None:
return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE")
diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
index e5714ef66fb..37a947f7f19 100644
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -1136,24 +1136,25 @@ async def open_sse_before_first_byte(
)
-def _is_azure_model_router_request(model: str) -> bool:
+def _is_azure_model_router_request(model: str, hidden_params: Mapping[str, object] | None = None) -> bool:
"""
- Check if the requested model is an Azure Model Router.
+ Check if a request went down the Azure Model Router route.
- Azure Model Router models follow the pattern:
- - azure_ai/model_router/
- - azure_ai/model-router
- - model_router/
- - model-router
+ ``model`` here is what the *client* sent, a model group alias with no ``model_router/``
+ prefix, so matching on it alone only works when the operator happened to put "model-router"
+ in the alias. Where the response is in hand its stamp answers this outright, so callers
+ should pass ``hidden_params``.
Args:
model: The requested model name
+ hidden_params: ``_hidden_params`` from the response, when the caller has it
Returns:
bool: True if this is an Azure Model Router request
"""
- model_lower: Final = model.lower()
- return "model-router" in model_lower or "model_router" in model_lower
+ from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
+ return AzureFoundryModelInfo.is_model_router_call(model=model, hidden_params=hidden_params)
def _override_openai_response_model(
@@ -1221,7 +1222,7 @@ def _override_openai_response_model(
return
# Check if this is an Azure Model Router request - if so, preserve the actual model used
- if _is_azure_model_router_request(requested_model):
+ if _is_azure_model_router_request(requested_model, hidden_params):
verbose_proxy_logger.debug(
"%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.",
log_context,
diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
index 82de634b488..a0f7de6b320 100644
--- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
@@ -6,9 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
-sys.path.insert(
- 0, os.path.abspath("../../..")
-) # Adds the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path
import time
@@ -277,9 +275,7 @@ def test_response_cost_calculator_uses_router_model_id_from_litellm_metadata():
assert cost is not None, "Cost should not be None"
expected_cost = (10 * custom_input_cost) + (5 * custom_output_cost)
- assert cost == pytest.approx(
- expected_cost
- ), f"Expected {expected_cost}, got {cost}"
+ assert cost == pytest.approx(expected_cost), f"Expected {expected_cost}, got {cost}"
finally:
litellm.model_cost.pop(custom_model_id, None)
@@ -876,13 +872,8 @@ async def test_datadog_logger_not_shadowed_by_llm_obs(monkeypatch):
# Regression check: we expect a distinct DataDogLogger, not the LLM Obs logger
assert type(datadog_logger) is DataDogLogger
- assert any(
- isinstance(cb, DataDogLLMObsLogger)
- for cb in logging_module._in_memory_loggers
- )
- assert any(
- type(cb) is DataDogLogger for cb in logging_module._in_memory_loggers
- )
+ assert any(isinstance(cb, DataDogLLMObsLogger) for cb in logging_module._in_memory_loggers)
+ assert any(type(cb) is DataDogLogger for cb in logging_module._in_memory_loggers)
finally:
logging_module._in_memory_loggers.clear()
@@ -893,9 +884,7 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch):
# Required env vars for Logfire integration
monkeypatch.setenv("LOGFIRE_TOKEN", "test-token")
- monkeypatch.setenv(
- "LOGFIRE_BASE_URL", "https://logfire-api-custom.pydantic.dev"
- ) # no trailing slash on purpose
+ monkeypatch.setenv("LOGFIRE_BASE_URL", "https://logfire-api-custom.pydantic.dev") # no trailing slash on purpose
# Import after env vars are set (important if module-level caching exists)
from litellm.integrations.opentelemetry import OpenTelemetry # logger class
@@ -914,9 +903,7 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch):
# Sanity: we got the right logger type and it is cached
assert type(logger) is OpenTelemetry
- assert any(
- type(cb) is OpenTelemetry for cb in logging_module._in_memory_loggers
- )
+ assert any(type(cb) is OpenTelemetry for cb in logging_module._in_memory_loggers)
# Core regression check: base URL env var should influence the exporter endpoint.
#
@@ -927,9 +914,7 @@ async def test_logfire_logger_accepts_env_vars_for_base_url(monkeypatch):
or getattr(logger, "config", None)
or getattr(logger, "_otel_config", None)
)
- assert (
- cfg is not None
- ), "Expected OpenTelemetry logger to keep an otel config on the instance"
+ assert cfg is not None, "Expected OpenTelemetry logger to keep an otel config on the instance"
endpoint = getattr(cfg, "endpoint", None) or getattr(cfg, "otlp_endpoint", None)
assert endpoint is not None, "Expected otel config to expose the OTLP endpoint"
@@ -1087,9 +1072,7 @@ async def test_logging_non_streaming_request():
# Use the filtered call for assertions
call_args = calls_with_expected_input[0]
- standard_logging_object = call_args.kwargs["kwargs"][
- "standard_logging_object"
- ]
+ standard_logging_object = call_args.kwargs["kwargs"]["standard_logging_object"]
assert standard_logging_object["stream"] is not True
finally:
# Restore original callbacks to ensure test isolation
@@ -1107,18 +1090,14 @@ async def test_logging_non_streaming_request():
"agenerate_content_stream",
],
)
-def test_success_handler_skips_sync_callbacks_for_async_requests(
- logging_obj, async_flag
-):
+def test_success_handler_skips_sync_callbacks_for_async_requests(logging_obj, async_flag):
"""Ensure sync success callbacks are skipped when async call type flags are set."""
from litellm.integrations.custom_logger import CustomLogger
class DummyLogger(CustomLogger):
pass
- logging_obj.stream = (
- False # simulate non-streaming request where sync callbacks would normally run
- )
+ logging_obj.stream = False # simulate non-streaming request where sync callbacks would normally run
logging_obj.model_call_details["litellm_params"] = {async_flag: True}
logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"]
@@ -1194,21 +1173,11 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call
def test_is_sync_litellm_request():
assert LitellmLogging._is_sync_litellm_request({}) is True
assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False
- assert (
- LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True})
- is False
- )
- assert (
- LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False
- )
+ assert LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) is False
+ assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False
assert LitellmLogging._is_sync_litellm_request({"agenerate_content": True}) is False
- assert (
- LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True})
- is False
- )
- assert (
- LitellmLogging._is_sync_litellm_request({"aanthropic_messages": False}) is True
- )
+ assert LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) is False
+ assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": False}) is True
def test_get_litellm_params_propagates_allm_passthrough_route():
@@ -1255,9 +1224,7 @@ async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream
logging_obj.model_call_details["litellm_params"] = {"acompletion": True}
with (
- patch.object(
- mock_callback, "async_log_success_event", new_callable=AsyncMock
- ) as mock_async_log,
+ patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log,
patch.object(mock_callback, "log_success_event") as mock_sync_log,
patch.object(
logging_obj,
@@ -1318,9 +1285,7 @@ async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_fin
with (
patch.object(mock_callback, "log_success_event") as mock_sync_log,
- patch.object(
- mock_callback, "async_log_success_event", new_callable=AsyncMock
- ) as mock_async_log,
+ patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log,
patch.object(
logging_obj,
"_success_handler_helper_fn",
@@ -1362,20 +1327,14 @@ async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks(
logging_obj.model_call_details["litellm_params"] = {}
with (
- patch.object(
- logging_obj, "async_success_handler", new_callable=AsyncMock
- ) as mock_async,
- patch.object(
- logging_obj, "success_handler", new_callable=MagicMock
- ) as mock_sync,
+ patch.object(logging_obj, "async_success_handler", new_callable=AsyncMock) as mock_async,
+ patch.object(logging_obj, "success_handler", new_callable=MagicMock) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=True,
),
- patch(
- "litellm.litellm_core_utils.litellm_logging.executor.submit"
- ) as mock_submit,
+ patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit,
):
await logging_obj.dispatch_success_handlers(
result=result,
@@ -1409,9 +1368,7 @@ async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through
try:
with (
- patch.object(
- mock_callback, "async_log_success_event", new_callable=AsyncMock
- ) as mock_async_log,
+ patch.object(mock_callback, "async_log_success_event", new_callable=AsyncMock) as mock_async_log,
patch.object(mock_callback, "log_success_event") as mock_sync_log,
):
await logging_obj.dispatch_success_handlers(result={"id": "pt-1"})
@@ -1438,20 +1395,14 @@ async def test_dispatch_failure_handlers_prefer_async_does_not_submit_sync_handl
logging_obj.model_call_details["litellm_params"] = {}
with (
- patch.object(
- logging_obj, "async_failure_handler", new_callable=AsyncMock
- ) as mock_async,
- patch.object(
- logging_obj, "failure_handler", new_callable=MagicMock
- ) as mock_sync,
+ patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock) as mock_async,
+ patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_failure_callbacks_for_async_calls",
return_value=False,
),
- patch(
- "litellm.litellm_core_utils.litellm_logging.executor.submit"
- ) as mock_submit,
+ patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit,
):
await logging_obj.dispatch_failure_handlers(
exception,
@@ -1534,12 +1485,8 @@ async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_c
patch.object(litellm, "success_callback", []),
patch.object(litellm, "failure_callback", [_sync_failure_callback]),
patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock),
- patch.object(
- logging_obj, "failure_handler", new_callable=MagicMock
- ) as mock_sync,
- patch(
- "litellm.litellm_core_utils.litellm_logging.executor.submit"
- ) as mock_submit,
+ patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync,
+ patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit,
):
await logging_obj.dispatch_failure_handlers(
exception,
@@ -1566,15 +1513,9 @@ async def test_dispatch_failure_handlers_sync_sdk_shortcut_runs_sync_handler_inl
logging_obj.model_call_details["litellm_params"] = {}
with (
- patch.object(
- logging_obj, "async_failure_handler", new_callable=AsyncMock
- ) as mock_async,
- patch.object(
- logging_obj, "failure_handler", new_callable=MagicMock
- ) as mock_sync,
- patch(
- "litellm.litellm_core_utils.litellm_logging.executor.submit"
- ) as mock_submit,
+ patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock) as mock_async,
+ patch.object(logging_obj, "failure_handler", new_callable=MagicMock) as mock_sync,
+ patch("litellm.litellm_core_utils.litellm_logging.executor.submit") as mock_submit,
):
await logging_obj.dispatch_failure_handlers(
exception,
@@ -1621,14 +1562,10 @@ def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj)
event_hook=GuardrailEventHooks.logging_only,
)
guardrail.should_run_guardrail = MagicMock(return_value=False)
- guardrail.logging_hook = MagicMock(
- return_value=(logging_obj.model_call_details, model_response)
- )
+ guardrail.logging_hook = MagicMock(return_value=(logging_obj.model_call_details, model_response))
dummy_logger = DummyLogger()
- dummy_logger.logging_hook = MagicMock(
- return_value=(logging_obj.model_call_details, model_response)
- )
+ dummy_logger.logging_hook = MagicMock(return_value=(logging_obj.model_call_details, model_response))
with patch.object(
logging_obj,
@@ -1762,11 +1699,7 @@ def test_get_request_tags_from_metadata_and_litellm_metadata():
# Test case 2: Tags in litellm_metadata only
tags = StandardLoggingPayloadSetup._get_request_tags(
- litellm_params={
- "litellm_metadata": {
- "tags": ["litellm-metadata-tag-1", "litellm-metadata-tag-2"]
- }
- },
+ litellm_params={"litellm_metadata": {"tags": ["litellm-metadata-tag-1", "litellm-metadata-tag-2"]}},
proxy_server_request={},
)
assert "litellm-metadata-tag-1" in tags
@@ -1871,15 +1804,9 @@ def test_get_request_tags_does_not_mutate_original_tags():
user_agent_count_2 = len([t for t in tags2 if t.startswith("User-Agent:")])
user_agent_count_3 = len([t for t in tags3 if t.startswith("User-Agent:")])
- assert (
- user_agent_count_1 == 2
- ), f"Expected 2 User-Agent tags, got {user_agent_count_1}"
- assert (
- user_agent_count_2 == 2
- ), f"Expected 2 User-Agent tags, got {user_agent_count_2}"
- assert (
- user_agent_count_3 == 2
- ), f"Expected 2 User-Agent tags, got {user_agent_count_3}"
+ assert user_agent_count_1 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_1}"
+ assert user_agent_count_2 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_2}"
+ assert user_agent_count_3 == 2, f"Expected 2 User-Agent tags, got {user_agent_count_3}"
# Verify all returned lists are independent (different objects)
assert tags1 is not tags2
@@ -1912,9 +1839,7 @@ def test_get_extra_header_tags():
# Test case 3: Extra headers configured but request has no headers dict
litellm.extra_spend_tag_headers = ["x-custom", "x-tenant"]
- result = StandardLoggingPayloadSetup._get_extra_header_tags(
- proxy_server_request={"headers": "not-a-dict"}
- )
+ result = StandardLoggingPayloadSetup._get_extra_header_tags(proxy_server_request={"headers": "not-a-dict"})
assert result is None
# Test case 4: Extra headers configured but none match request headers
@@ -2215,9 +2140,7 @@ def test_get_masked_values():
"presidio_anonymizer_api_base": None,
"vertex_credentials": "{sensitive_api_key}",
}
- masked_values = _get_masked_values(
- sensitive_object, unmasked_length=4, number_of_asterisks=4
- )
+ masked_values = _get_masked_values(sensitive_object, unmasked_length=4, number_of_asterisks=4)
assert masked_values["presidio_anonymizer_api_base"] is None
assert masked_values["vertex_credentials"] == "{s****y}"
@@ -2242,9 +2165,7 @@ async def test_e2e_generate_cold_storage_object_key_successful():
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key,
):
# Mock the S3 object key generation to return a predictable result
- mock_get_s3_key.return_value = (
- "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
- )
+ mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
@@ -2285,16 +2206,12 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
with (
patch("litellm.cold_storage_custom_logger", "s3_v2"),
- patch(
- "litellm.logging_callback_manager.get_active_custom_logger_for_callback_name"
- ) as mock_get_logger,
+ patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger,
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key,
):
# Setup mocks
mock_get_logger.return_value = mock_custom_logger
- mock_get_s3_key.return_value = (
- "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
- )
+ mock_get_s3_key.return_value = "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
@@ -2313,9 +2230,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
)
# Verify the result
- assert (
- result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
- )
+ assert result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
@pytest.mark.asyncio
@@ -2338,16 +2253,12 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
with (
patch("litellm.cold_storage_custom_logger", "s3_v2"),
- patch(
- "litellm.logging_callback_manager.get_active_custom_logger_for_callback_name"
- ) as mock_get_logger,
+ patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger,
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key,
):
# Setup mocks
mock_get_logger.return_value = mock_custom_logger
- mock_get_s3_key.return_value = (
- "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
- )
+ mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
@@ -2463,9 +2374,7 @@ def test_get_usage_as_dict():
assert result == {"prompt_tokens": 20, "completion_tokens": 30}
# Test case 5: response_obj with no usage key returns empty
- result = StandardLoggingPayloadSetup.get_usage_as_dict(
- response_obj={"id": "resp-1", "choices": []}
- )
+ result = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj={"id": "resp-1", "choices": []})
assert result == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
@@ -2478,26 +2387,20 @@ def test_append_system_prompt_messages():
# Test case 1: system in kwargs with existing messages
kwargs = {"system": "You are a helpful assistant"}
messages = [{"role": "user", "content": "Hello"}]
- result = StandardLoggingPayloadSetup.append_system_prompt_messages(
- kwargs=kwargs, messages=messages
- )
+ result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages)
assert len(result) == 2
assert result[0] == {"role": "system", "content": "You are a helpful assistant"}
assert result[1] == {"role": "user", "content": "Hello"}
# Test case 2: system in kwargs with None messages
kwargs = {"system": "You are a helpful assistant"}
- result = StandardLoggingPayloadSetup.append_system_prompt_messages(
- kwargs=kwargs, messages=None
- )
+ result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=None)
assert len(result) == 1
assert result[0] == {"role": "system", "content": "You are a helpful assistant"}
# Test case 3: system in kwargs with empty messages list
kwargs = {"system": "You are a helpful assistant"}
- result = StandardLoggingPayloadSetup.append_system_prompt_messages(
- kwargs=kwargs, messages=[]
- )
+ result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=[])
assert len(result) == 1
assert result[0] == {"role": "system", "content": "You are a helpful assistant"}
@@ -2507,24 +2410,18 @@ def test_append_system_prompt_messages():
{"role": "system", "content": "You are a helpful assistant"},
{"role": "user", "content": "Hello"},
]
- result = StandardLoggingPayloadSetup.append_system_prompt_messages(
- kwargs=kwargs, messages=messages
- )
+ result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages)
assert len(result) == 2
assert result[0] == {"role": "system", "content": "You are a helpful assistant"}
# Test case 5: no system in kwargs returns messages unchanged
kwargs = {}
messages = [{"role": "user", "content": "Hello"}]
- result = StandardLoggingPayloadSetup.append_system_prompt_messages(
- kwargs=kwargs, messages=messages
- )
+ result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=kwargs, messages=messages)
assert result == messages
# Test case 6: None kwargs returns messages unchanged
- result = StandardLoggingPayloadSetup.append_system_prompt_messages(
- kwargs=None, messages=messages
- )
+ result = StandardLoggingPayloadSetup.append_system_prompt_messages(kwargs=None, messages=messages)
assert result == messages
@@ -2585,12 +2482,11 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu
# Verify that standard_logging_object was set
assert "standard_logging_object" in logging_obj.model_call_details, (
- "standard_logging_object should be set for pass-through endpoints "
- "even when complete_streaming_response is None"
+ "standard_logging_object should be set for pass-through endpoints even when complete_streaming_response is None"
+ )
+ assert logging_obj.model_call_details["standard_logging_object"] is not None, (
+ "standard_logging_object should not be None for pass-through endpoints"
)
- assert (
- logging_obj.model_call_details["standard_logging_object"] is not None
- ), "standard_logging_object should not be None for pass-through endpoints"
# Verify that async_complete_streaming_response was set to prevent re-processing
# This is consistent with the existing code pattern for regular streaming
@@ -2598,15 +2494,13 @@ async def test_async_success_handler_sets_standard_logging_object_for_pass_throu
"async_complete_streaming_response should be set to prevent re-processing, "
"consistent with the existing code pattern"
)
- assert (
- logging_obj.model_call_details["async_complete_streaming_response"] is result
- ), "async_complete_streaming_response should be set to the result"
+ assert logging_obj.model_call_details["async_complete_streaming_response"] is result, (
+ "async_complete_streaming_response should be set to the result"
+ )
# Verify that response_cost is set to None (cost calculation not possible for pass-through)
# This is consistent with the error handling in the non-pass-through code path
- assert (
- "response_cost" in logging_obj.model_call_details
- ), "response_cost should be set for pass-through endpoints"
+ assert "response_cost" in logging_obj.model_call_details, "response_cost should be set for pass-through endpoints"
assert logging_obj.model_call_details["response_cost"] is None, (
"response_cost should be None for pass-through endpoints since "
"StandardPassThroughResponseObject doesn't have standard usage info"
@@ -2665,14 +2559,10 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp
# Verify first call set the values
assert "standard_logging_object" in logging_obj.model_call_details
assert "async_complete_streaming_response" in logging_obj.model_call_details
- first_standard_logging_object = logging_obj.model_call_details[
- "standard_logging_object"
- ]
+ first_standard_logging_object = logging_obj.model_call_details["standard_logging_object"]
# Second call - should return early due to async_complete_streaming_response guard
- with patch.object(
- logging_obj, "get_combined_callback_list", return_value=[]
- ) as mock_callbacks:
+ with patch.object(logging_obj, "get_combined_callback_list", return_value=[]) as mock_callbacks:
await logging_obj.async_success_handler(
result=result,
start_time=start_time,
@@ -2683,10 +2573,9 @@ async def test_async_success_handler_prevents_reprocessing_for_pass_through_endp
mock_callbacks.assert_not_called()
# Verify standard_logging_object wasn't modified by second call
- assert (
- logging_obj.model_call_details["standard_logging_object"]
- is first_standard_logging_object
- ), "standard_logging_object should not be modified on re-processing"
+ assert logging_obj.model_call_details["standard_logging_object"] is first_standard_logging_object, (
+ "standard_logging_object should not be modified on re-processing"
+ )
@pytest.mark.asyncio
@@ -2725,9 +2614,7 @@ async def test_async_success_handler_sets_standard_logging_object_for_streaming_
}
# Create a pass-through response object (simulating unparseable streaming response)
- result = StandardPassThroughResponseObject(
- response='data: {"chunk": 1}\ndata: {"chunk": 2}\ndata: [DONE]'
- )
+ result = StandardPassThroughResponseObject(response='data: {"chunk": 1}\ndata: {"chunk": 2}\ndata: [DONE]')
start_time = datetime.now()
end_time = datetime.now()
@@ -2747,9 +2634,9 @@ async def test_async_success_handler_sets_standard_logging_object_for_streaming_
"standard_logging_object should be set for streaming pass-through endpoints "
"even when the response cannot be parsed into a ModelResponse"
)
- assert (
- logging_obj.model_call_details["standard_logging_object"] is not None
- ), "standard_logging_object should not be None for streaming pass-through endpoints"
+ assert logging_obj.model_call_details["standard_logging_object"] is not None, (
+ "standard_logging_object should not be None for streaming pass-through endpoints"
+ )
def test_get_error_information_error_code_priority():
@@ -2791,30 +2678,22 @@ def test_get_error_information_error_code_priority():
self.message = message
super().__init__(message)
- both_exception = BothAttributesException(
- code="400", status_code=500, message="Bad Request"
- )
+ both_exception = BothAttributesException(code="400", status_code=500, message="Bad Request")
result = StandardLoggingPayloadSetup.get_error_information(both_exception)
assert result["error_code"] == "400" # Should prefer 'code' over 'status_code'
# Test case 4: Exception with 'code' as empty string - should fall back to 'status_code'
- empty_code_exception = BothAttributesException(
- code="", status_code=404, message="Not Found"
- )
+ empty_code_exception = BothAttributesException(code="", status_code=404, message="Not Found")
result = StandardLoggingPayloadSetup.get_error_information(empty_code_exception)
assert result["error_code"] == "404" # Should fall back to status_code
# Test case 5: Exception with 'code' as "None" string - should fall back to 'status_code'
- none_string_exception = BothAttributesException(
- code="None", status_code=503, message="Service Unavailable"
- )
+ none_string_exception = BothAttributesException(code="None", status_code=503, message="Service Unavailable")
result = StandardLoggingPayloadSetup.get_error_information(none_string_exception)
assert result["error_code"] == "503" # Should fall back to status_code
# Test case 6: Exception with 'code' as None - should fall back to 'status_code'
- none_code_exception = BothAttributesException(
- code=None, status_code=401, message="Unauthorized"
- )
+ none_code_exception = BothAttributesException(code=None, status_code=401, message="Unauthorized")
result = StandardLoggingPayloadSetup.get_error_information(none_code_exception)
assert result["error_code"] == "401" # Should fall back to status_code
@@ -2863,9 +2742,7 @@ def test_get_error_information_prefers_message_attribute_over_str():
)
result = StandardLoggingPayloadSetup.get_error_information(exc)
- assert (
- result["error_message"] == msg
- ), f"expected message from .message attribute, got {result['error_message']!r}"
+ assert result["error_message"] == msg, f"expected message from .message attribute, got {result['error_message']!r}"
assert result["error_code"] == "401"
assert result["error_class"] == "ProxyExceptionLike"
@@ -2940,8 +2817,7 @@ def test_get_error_information_preserves_explicit_empty_message():
exc = ProxyExceptionLike(message="", code=500)
result = StandardLoggingPayloadSetup.get_error_information(exc)
assert result["error_message"] == "", (
- "explicit empty .message must survive verbatim; got "
- f"{result['error_message']!r}"
+ f"explicit empty .message must survive verbatim; got {result['error_message']!r}"
)
@@ -3204,9 +3080,7 @@ def test_process_hidden_params_recalculates_cost_after_failure_handler_zero():
choices=[{"message": {"role": "assistant", "content": "ok"}}],
usage=Usage(prompt_tokens=9698, completion_tokens=30, total_tokens=9728),
)
- logging_obj._process_hidden_params_and_response_cost(
- result, datetime.now(), datetime.now()
- )
+ logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now())
cost = logging_obj.model_call_details.get("response_cost")
assert cost is not None and cost > 0
@@ -3230,9 +3104,7 @@ def test_process_hidden_params_preserves_zero_cost_in_hidden_params():
litellm_call_id="test-hidden-zero-cost",
function_id="test-hidden-zero-cost",
)
- logging_obj.model_call_details["litellm_params"] = {
- "model": "gemini-2.5-flash-lite"
- }
+ logging_obj.model_call_details["litellm_params"] = {"model": "gemini-2.5-flash-lite"}
logging_obj.optional_params = {}
result = ModelResponse(
@@ -3242,9 +3114,7 @@ def test_process_hidden_params_preserves_zero_cost_in_hidden_params():
)
result._hidden_params = {"response_cost": 0.0}
- logging_obj._process_hidden_params_and_response_cost(
- result, datetime.now(), datetime.now()
- )
+ logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now())
assert logging_obj.model_call_details.get("response_cost") == 0.0
slo = logging_obj.model_call_details.get("standard_logging_object") or {}
@@ -3293,9 +3163,7 @@ def test_process_hidden_params_uses_hidden_params_cost_after_failure_handler_zer
)
result._hidden_params = {"response_cost": passthrough_cost}
- logging_obj._process_hidden_params_and_response_cost(
- result, datetime.now(), datetime.now()
- )
+ logging_obj._process_hidden_params_and_response_cost(result, datetime.now(), datetime.now())
assert logging_obj.model_call_details.get("response_cost") == passthrough_cost
slo = logging_obj.model_call_details.get("standard_logging_object") or {}
@@ -3352,9 +3220,7 @@ def test_function_setup_litellm_metadata_populates_metadata():
assert litellm_metadata.get("user_api_key_hash") == test_api_key_hash
# metadata should be a COPY, not an alias — mutating one must not affect the other
- assert (
- metadata is not litellm_metadata
- ), "litellm_params['metadata'] should be a copy, not the same object"
+ assert metadata is not litellm_metadata, "litellm_params['metadata'] should be a copy, not the same object"
def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup():
@@ -3399,9 +3265,9 @@ def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup():
litellm_params = logging_obj.model_call_details.get("litellm_params", {})
litellm_metadata = litellm_params.get("litellm_metadata")
assert litellm_metadata is not None
- assert litellm_metadata.get("standard_logging_guardrail_information") == [
- guardrail_entry
- ], "guardrail writes after function_setup must be visible to the logging object"
+ assert litellm_metadata.get("standard_logging_guardrail_information") == [guardrail_entry], (
+ "guardrail writes after function_setup must be visible to the logging object"
+ )
assert litellm_metadata.get("applied_guardrails") == ["pam-ethical-request"]
merged = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params)
@@ -3570,9 +3436,7 @@ def test_failure_handler_skips_sync_callbacks_for_pass_through_requests(logging_
@pytest.mark.parametrize("call_type", ["completion", "acompletion"])
-def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests(
- logging_obj, call_type
-):
+def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests(logging_obj, call_type):
"""Ensure sync failure callbacks still fire for normal (non-pass-through) requests."""
from litellm.integrations.custom_logger import CustomLogger
@@ -3733,9 +3597,7 @@ def test_standard_logging_hidden_params_backfills_response_cost_without_mutating
)
response._hidden_params = {"response_cost": None, "model_id": "mid-test"}
- payload = logging_obj._build_standard_logging_payload(
- response, datetime.now(), datetime.now()
- )
+ payload = logging_obj._build_standard_logging_payload(response, datetime.now(), datetime.now())
assert payload is not None
assert payload["hidden_params"]["response_cost"] == 0.002
@@ -3789,10 +3651,7 @@ def test_merge_hidden_params_from_response_into_metadata_no_op_when_empty():
_hidden_params = {}
logging_obj._merge_hidden_params_from_response_into_metadata(_NoHp())
- assert (
- "hidden_params"
- not in logging_obj.model_call_details["litellm_params"]["metadata"]
- )
+ assert "hidden_params" not in logging_obj.model_call_details["litellm_params"]["metadata"]
# ── StandardLoggingPayloadSetup.get_additional_headers ───────────────────────
@@ -3870,6 +3729,82 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
assert payload["litellm_call_id"] == call_id
+# ── Azure Model Router selected-model attribution ────────────────────────────
+
+
+def _model_router_response(selected_model: str, stamp: bool):
+ """A ModelResponse as AzureModelRouterConfig hands it back, with or without the stamp."""
+ from litellm.llms.azure_ai.common_utils import (
+ AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
+ )
+ from litellm.types.utils import ModelResponse
+
+ response = ModelResponse(model=selected_model)
+ response._hidden_params = {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {}
+ return response
+
+
+def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj):
+ """
+ The selected model must win off the stamp, not off "model-router" appearing in the
+ requested model. An operator whose model group is named anything else was invisible
+ to the name check, so their logs and spend rows named the router instead.
+ """
+ import datetime
+
+ from litellm.litellm_core_utils.litellm_logging import (
+ get_standard_logging_object_payload,
+ )
+
+ now = datetime.datetime.now()
+ payload = get_standard_logging_object_payload(
+ kwargs={
+ "model": "azure_ai/smart-pick",
+ "custom_llm_provider": "azure_ai",
+ "messages": [],
+ "litellm_params": {"metadata": {}},
+ },
+ init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=True),
+ start_time=now,
+ end_time=now,
+ logging_obj=logging_obj,
+ status="success",
+ )
+
+ assert payload is not None
+ assert payload["model"] == "azure_ai/grok-4-1-fast-reasoning"
+
+
+def test_standard_logging_payload_keeps_requested_model_without_router_stamp(logging_obj):
+ """
+ Control for the test above: an ordinary azure_ai deployment is unaffected, so the stamp
+ is what redirects attribution rather than the response model winning unconditionally.
+ """
+ import datetime
+
+ from litellm.litellm_core_utils.litellm_logging import (
+ get_standard_logging_object_payload,
+ )
+
+ now = datetime.datetime.now()
+ payload = get_standard_logging_object_payload(
+ kwargs={
+ "model": "azure_ai/smart-pick",
+ "custom_llm_provider": "azure_ai",
+ "messages": [],
+ "litellm_params": {"metadata": {}},
+ },
+ init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=False),
+ start_time=now,
+ end_time=now,
+ logging_obj=logging_obj,
+ status="success",
+ )
+
+ assert payload is not None
+ assert payload["model"] == "azure_ai/smart-pick"
+
+
def _make_dict_logging_obj():
"""Build a Logging instance configured for a non-streaming dict result."""
obj = LitellmLogging(
@@ -3905,9 +3840,7 @@ def test_success_handler_computes_cost_for_dict_response():
"_build_standard_logging_payload",
return_value={"response_cost": expected_cost},
),
- patch(
- "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
- ),
+ patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
@@ -3944,9 +3877,7 @@ def test_success_handler_preserves_precomputed_cost_for_dict_response():
"_build_standard_logging_payload",
return_value={"response_cost": precomputed_cost},
),
- patch(
- "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
- ),
+ patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
@@ -3985,9 +3916,7 @@ def test_success_handler_unified_helper_runs_for_typed_results():
"_build_standard_logging_payload",
return_value={"response_cost": expected_cost},
),
- patch(
- "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"
- ),
+ patch("litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload"),
patch.object(
logging_obj,
"_is_recognized_call_type_for_logging",
@@ -4042,9 +3971,7 @@ class TestFirstApiCallStartTimeSetOnce:
assert first == obj.model_call_details["api_call_start_time"]
# Set on the logging object only — user metadata untouched.
assert user_meta == {}
- assert (
- "first_api_call_start_time" not in obj.model_call_details["litellm_params"]
- )
+ assert "first_api_call_start_time" not in obj.model_call_details["litellm_params"]
time.sleep(0.002) # ensure a distinct retry timestamp
obj.pre_call(input="hi", api_key="sk-test")
@@ -4061,18 +3988,16 @@ def test_get_error_information_for_logging_payload_ignores_spoofed_disconnect_wi
baseline = StandardLoggingPayloadSetup.get_error_information(
original_exception=ValueError("provider failure"),
)
- error_information, error_str = (
- StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
- metadata={
- "error_information": {
- "error_code": "499",
- "error_message": "Client disconnected the request",
- "error_class": "ClientDisconnected",
- }
- },
- original_exception=ValueError("provider failure"),
- error_str="provider failure",
- )
+ error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
+ metadata={
+ "error_information": {
+ "error_code": "499",
+ "error_message": "Client disconnected the request",
+ "error_class": "ClientDisconnected",
+ }
+ },
+ original_exception=ValueError("provider failure"),
+ error_str="provider failure",
)
assert error_information == baseline
assert error_str == "provider failure"
@@ -4086,22 +4011,18 @@ def test_get_error_information_for_logging_payload_client_disconnect():
"error_message": "Client disconnected the request",
"error_class": "ClientDisconnected",
}
- error_information, error_str = (
- StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
- metadata={"client_disconnected": True, "error_information": custom_error},
- original_exception=None,
- error_str=None,
- )
+ error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
+ metadata={"client_disconnected": True, "error_information": custom_error},
+ original_exception=None,
+ error_str=None,
)
assert error_information == custom_error
assert error_str == "Client disconnected the request"
- error_information, error_str = (
- StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
- metadata={"client_disconnected": True},
- original_exception=None,
- error_str="existing error",
- )
+ error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
+ metadata={"client_disconnected": True},
+ original_exception=None,
+ error_str="existing error",
)
assert error_information["error_code"] == "499"
assert error_str == "existing error"
@@ -4109,12 +4030,10 @@ def test_get_error_information_for_logging_payload_client_disconnect():
baseline = StandardLoggingPayloadSetup.get_error_information(
original_exception=None,
)
- error_information, error_str = (
- StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
- metadata={},
- original_exception=None,
- error_str=None,
- )
+ error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
+ metadata={},
+ original_exception=None,
+ error_str=None,
)
assert error_information == baseline
assert error_str is None
@@ -4149,9 +4068,7 @@ def test_get_error_information_prefers_message_attribute_over_empty_str():
def __str__(self):
return ""
- info = StandardLoggingPayloadSetup.get_error_information(
- original_exception=_SilentExc()
- )
+ info = StandardLoggingPayloadSetup.get_error_information(original_exception=_SilentExc())
assert info["error_message"] == "real failure detail"
assert info["error_code"] == "401"
@@ -4182,9 +4099,7 @@ def _responses_api_response_with_text(text="hello world"):
type="message",
role="assistant",
status="completed",
- content=[
- ResponseOutputText(annotations=[], text=text, type="output_text")
- ],
+ content=[ResponseOutputText(annotations=[], text=text, type="output_text")],
)
],
usage=ResponseAPIUsage(input_tokens=11, output_tokens=7, total_tokens=18),
@@ -4199,9 +4114,7 @@ def _responses_api_response_with_text(text="hello world"):
("ResponseFailedEvent", "response.failed"),
],
)
-def test_handle_anthropic_messages_response_logging_translates_terminal_responses_api_event(
- event_cls, event_type
-):
+def test_handle_anthropic_messages_response_logging_translates_terminal_responses_api_event(event_cls, event_type):
"""Regression for #28595 / #28943. When anthropic_messages routes to the OpenAI
Responses backend and stream=True, success_handler receives a terminal Responses
API event. The handler must translate it to a ModelResponse whose choices carry
@@ -4240,10 +4153,7 @@ def test_handle_anthropic_messages_response_logging_passes_model_response_throug
"""Anthropic-native path already yields a ModelResponse; it must be returned unchanged."""
logging_obj = _anthropic_messages_logging_obj()
model_response = ModelResponse()
- assert (
- logging_obj._handle_anthropic_messages_response_logging(result=model_response)
- is model_response
- )
+ assert logging_obj._handle_anthropic_messages_response_logging(result=model_response) is model_response
def test_handle_anthropic_messages_response_logging_degrades_on_unparseable_responses_payload():
@@ -4539,9 +4449,7 @@ def test_non_image_response_has_no_output_image_count(logging_obj):
def test_zero_token_video_usage_preserves_duration_seconds(logging_obj):
"""Video usage bills by duration; the payload must keep duration_seconds even with zero tokens."""
- payload = _build_payload_for_media_response(
- logging_obj, {"id": "video-1", "usage": {"duration_seconds": 4.0}}
- )
+ payload = _build_payload_for_media_response(logging_obj, {"id": "video-1", "usage": {"duration_seconds": 4.0}})
assert payload is not None
assert payload["metadata"]["usage_object"]["duration_seconds"] == 4.0
diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py
index 900372f3e54..1f684ecc3b5 100644
--- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py
+++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py
@@ -5,9 +5,7 @@ from unittest.mock import MagicMock, patch
import pytest
-sys.path.insert(
- 0, os.path.abspath("../../../../..")
-) # Adds the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path
from litellm.llms.azure_ai.azure_model_router.transformation import (
AzureModelRouterConfig,
)
@@ -120,9 +118,7 @@ def test_azure_ai_grok_stop_parameter_handling():
# Test supported parameters for Grok models
for model in ("grok-4-fast", "grok-4.3"):
grok_params = config.get_supported_openai_params(model)
- assert (
- "stop" not in grok_params
- ), "Grok models should not support stop parameter"
+ assert "stop" not in grok_params, "Grok models should not support stop parameter"
# Test supported parameters for non-Grok models
gpt_params = config.get_supported_openai_params("gpt-4")
@@ -201,11 +197,84 @@ def test_azure_model_router_response_shows_actual_model():
# Verify that the response contains the actual model used, not the router model
assert result.model == "azure_ai/gpt-5-nano-2025-08-07", (
- f"Expected model to be 'azure_ai/gpt-5-nano-2025-08-07' (actual model used), "
- f"but got '{result.model}'"
+ f"Expected model to be 'azure_ai/gpt-5-nano-2025-08-07' (actual model used), but got '{result.model}'"
)
+def test_azure_model_router_stamps_selected_model_on_hidden_params():
+ """
+ The selected model must be stamped on _hidden_params, not left for downstream code to
+ re-derive by looking for "model-router" in the model string. Deployments whose alias
+ does not contain that text are invisible to the string check.
+ """
+ from httpx import Response
+
+ from litellm.llms.azure_ai.common_utils import (
+ AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
+ AzureFoundryModelInfo,
+ )
+ from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
+ from litellm.types.utils import ModelResponse
+
+ raw_response_json = {
+ "id": "chatcmpl-test456",
+ "object": "chat.completion",
+ "created": 1234567890,
+ "model": "grok-4-1-fast-reasoning",
+ "choices": [
+ {
+ "index": 0,
+ "message": {"role": "assistant", "content": "pong"},
+ "finish_reason": "stop",
+ }
+ ],
+ "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+ }
+
+ mock_response = MagicMock(spec=Response)
+ mock_response.json.return_value = raw_response_json
+ mock_response.text = json.dumps(raw_response_json)
+ mock_response.headers = {}
+
+ logging_obj = MagicMock(spec=LiteLLMLoggingObj)
+ logging_obj.post_call = MagicMock()
+ logging_obj.model_call_details = {}
+
+ result = AzureModelRouterConfig().transform_response(
+ model="smart-pick",
+ raw_response=mock_response,
+ model_response=ModelResponse(),
+ logging_obj=logging_obj,
+ request_data={},
+ messages=[{"role": "user", "content": "Reply with just pong"}],
+ optional_params={},
+ litellm_params={"model": "azure_ai/model_router/smart-pick"},
+ encoding=None,
+ api_key="test-key",
+ json_mode=False,
+ )
+
+ assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == result.model
+ assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == "azure_ai/grok-4-1-fast-reasoning"
+ assert AzureFoundryModelInfo.get_model_router_selected_model(result._hidden_params) == (
+ "azure_ai/grok-4-1-fast-reasoning"
+ )
+ assert AzureFoundryModelInfo.is_model_router_call(model="smart-pick", hidden_params=result._hidden_params) is True
+
+
+def test_azure_model_router_stamp_does_not_leak_across_responses():
+ """
+ ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written
+ as a fresh dict. Mutating in place would bleed the selected model into unrelated responses.
+ """
+ from litellm.llms.azure_ai.common_utils import AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY
+ from litellm.types.utils import ModelResponse
+
+ untouched = ModelResponse()
+
+ assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {})
+
+
def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name():
"""
Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name.
@@ -226,14 +295,10 @@ def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name():
mock_response.text = error_text
mock_response.json.return_value = json.loads(error_text)
mock_response.status_code = 400
- e = httpx.HTTPStatusError(
- message="400", request=MagicMock(), response=mock_response
- )
+ e = httpx.HTTPStatusError(message="400", request=MagicMock(), response=mock_response)
assert config._error_has_tool_level_extra_fields(error_text) is True
- assert (
- config.should_retry_llm_api_inside_llm_translation_on_http_error(e, {}) is True
- )
+ assert config.should_retry_llm_api_inside_llm_translation_on_http_error(e, {}) is True
request_data = {
"model": "FW-Kimi-K2.6",
@@ -354,9 +419,7 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages():
{
"role": "assistant",
"content": "I can help.",
- "thinking_blocks": [
- {"type": "thinking", "thinking": "Reading the file.", "signature": "sig"}
- ],
+ "thinking_blocks": [{"type": "thinking", "thinking": "Reading the file.", "signature": "sig"}],
"provider_specific_fields": {"thought_signature": "sig-top"},
"tool_calls": [
{
diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py
index 716fba370df..562046b2f81 100644
--- a/tests/test_litellm/proxy/test_common_request_processing.py
+++ b/tests/test_litellm/proxy/test_common_request_processing.py
@@ -126,16 +126,12 @@ class TestProxyBaseLLMRequestProcessing:
assert json.loads(result.body) == guardrailed_body
@pytest.mark.asyncio
- async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(
- self, monkeypatch
- ):
+ async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch):
"""The guardrail JSON path must forward upstream response headers (e.g.
x-amzn-requestid) alongside the x-litellm-* headers, matching the
non-guardrail passthrough path, while dropping length headers that no
longer match the rewritten body."""
- processing_obj = ProxyBaseLLMRequestProcessing(
- data={"custom_llm_provider": "bedrock"}
- )
+ processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
monkeypatch.setattr(
processing_obj,
"_has_post_call_guardrails_for_passthrough",
@@ -175,14 +171,10 @@ class TestProxyBaseLLMRequestProcessing:
assert result.headers["content-length"] == str(len(result.body))
@pytest.mark.asyncio
- async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(
- self, monkeypatch
- ):
+ async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch):
"""The guardrail event-stream branch must also forward upstream response
headers alongside the x-litellm-* headers."""
- processing_obj = ProxyBaseLLMRequestProcessing(
- data={"custom_llm_provider": "bedrock"}
- )
+ processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
monkeypatch.setattr(
processing_obj,
"_has_post_call_guardrails_for_passthrough",
@@ -224,15 +216,11 @@ class TestProxyBaseLLMRequestProcessing:
assert result.headers["x-litellm-call-id"] == "test-call-id"
@pytest.mark.asyncio
- async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(
- self, monkeypatch
- ):
+ async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(self, monkeypatch):
"""Guardrailed non-streaming passthrough responses must include headers
injected by post_call_response_headers_hook, matching the headers a
non-guardrailed passthrough response would carry."""
- processing_obj = ProxyBaseLLMRequestProcessing(
- data={"custom_llm_provider": "bedrock"}
- )
+ processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"})
monkeypatch.setattr(
processing_obj,
"_has_post_call_guardrails_for_passthrough",
@@ -251,9 +239,7 @@ class TestProxyBaseLLMRequestProcessing:
return kwargs["response"]
proxy_logging_obj.post_call_success_hook = fake_post_call_success_hook
- proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
- return_value={"x-litellm-custom": "from-hook"}
- )
+ proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-litellm-custom": "from-hook"})
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
response=upstream,
@@ -2221,6 +2207,52 @@ class TestOverrideOpenAIResponseModel:
assert response_obj.model == actual_model_used
assert response_obj.model != requested_model
+ def test_override_model_preserves_model_router_model_for_alias_without_router_in_name(self):
+ """
+ The client sends a model group alias, which carries no model_router/ prefix, so the
+ name check alone only fires when the operator happened to put "model-router" in the
+ alias. With the stamp on the response the actual model survives whatever it is named.
+ """
+ from litellm.llms.azure_ai.common_utils import (
+ AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
+ )
+
+ requested_model = "smart-pick"
+ actual_model_used = "azure_ai/grok-4-1-fast-reasoning"
+
+ response_obj = MagicMock()
+ response_obj.model = actual_model_used
+ response_obj._hidden_params = {
+ "additional_headers": {},
+ AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: actual_model_used,
+ }
+
+ _override_openai_response_model(
+ response_obj=response_obj,
+ requested_model=requested_model,
+ log_context="test_context",
+ )
+ assert response_obj.model == actual_model_used
+
+ def test_override_model_still_restamps_non_router_alias_without_stamp(self):
+ """
+ Control for the test above: absent the stamp, an ordinary deployment keeps being
+ restamped to the requested model, so the stamp is doing the work rather than the
+ preserve branch having gone unconditional.
+ """
+ requested_model = "smart-pick"
+
+ response_obj = MagicMock()
+ response_obj.model = "azure_ai/grok-4-1-fast-reasoning"
+ response_obj._hidden_params = {"additional_headers": {}}
+
+ _override_openai_response_model(
+ response_obj=response_obj,
+ requested_model=requested_model,
+ log_context="test_context",
+ )
+ assert response_obj.model == requested_model
+
def test_override_model_uses_winning_model_for_fastest_response(self):
"""
Test that when fastest_response batch completion is used with a
@@ -2793,9 +2825,7 @@ class TestStreamCloseOnDisconnect:
finally:
closed.set()
- response = _UpstreamClosingStreamingResponse(
- body(), media_type="text/event-stream"
- )
+ response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream")
async def receive():
await asyncio.Event().wait()
@@ -2826,9 +2856,7 @@ class TestStreamCloseOnDisconnect:
finally:
closed.set()
- response = _UpstreamClosingStreamingResponse(
- body(), media_type="text/event-stream"
- )
+ response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream")
async def receive():
await disconnected.wait()
@@ -2899,9 +2927,7 @@ class TestStreamCloseOnDisconnect:
finally:
inner_closed.set()
- response = await create_response(
- generator=wrapped(), media_type="text/event-stream", headers={}
- )
+ response = await create_response(generator=wrapped(), media_type="text/event-stream", headers={})
async def receive():
await asyncio.Event().wait()
@@ -3097,9 +3123,7 @@ class TestStreamCloseOnDisconnect:
with pytest.raises(_ClientDisconnectedBeforeFirstChunk):
await asyncio.wait_for(
- _buffer_first_chunk_honoring_disconnect(
- AcloseRaises(), request=self._request_that_disconnects()
- ),
+ _buffer_first_chunk_honoring_disconnect(AcloseRaises(), request=self._request_that_disconnects()),
timeout=5,
)
@@ -3115,9 +3139,7 @@ class TestStreamCloseOnDisconnect:
with pytest.raises(_ClientDisconnectedBeforeFirstChunk):
await asyncio.wait_for(
- _buffer_first_chunk_honoring_disconnect(
- blocking_gen(), request=self._request_that_disconnects()
- ),
+ _buffer_first_chunk_honoring_disconnect(blocking_gen(), request=self._request_that_disconnects()),
timeout=5,
)
assert closed.is_set()
@@ -3133,9 +3155,7 @@ class TestHandleLLMApiExceptionRetryAfter:
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test")
proxy_logging_obj = MagicMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
- proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
- return_value=callback_headers or {}
- )
+ proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {})
try:
await processor._handle_llm_api_exception(
@@ -3187,9 +3207,7 @@ class TestHandleLLMApiExceptionRetryAfter:
enable_pre_call_checks=False,
cooldown_list=[],
)
- proxy_exc = await self._invoke(
- exc, callback_headers={"retry-after": "", "x-custom": "1"}
- )
+ proxy_exc = await self._invoke(exc, callback_headers={"retry-after": "", "x-custom": "1"})
assert proxy_exc.headers["retry-after"] == "43"
assert proxy_exc.headers["x-custom"] == "1"
@@ -3385,9 +3403,7 @@ class TestDisconnectGatherCleanup:
return Request(scope={"type": "http", "headers": []}, receive=receive)
@pytest.mark.asyncio
- async def test_base_process_llm_request_raises_499_on_client_disconnect(
- self, monkeypatch
- ):
+ async def test_base_process_llm_request_raises_499_on_client_disconnect(self, monkeypatch):
"""With cancel_on_disconnect enabled, base_process_llm_request returns 499."""
import asyncio
@@ -3416,9 +3432,7 @@ class TestDisconnectGatherCleanup:
"common_processing_pre_call_logic",
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
)
- monkeypatch.setattr(
- processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
- )
+ monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
with pytest.raises(HTTPException) as exc_info:
await processing_obj.base_process_llm_request(
@@ -3436,9 +3450,7 @@ class TestDisconnectGatherCleanup:
assert "disconnected" in exc_info.value.detail.lower()
@pytest.mark.asyncio
- async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(
- self, monkeypatch
- ):
+ async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(self, monkeypatch):
import asyncio
import litellm.proxy.common_request_processing as cpr
@@ -3463,9 +3475,7 @@ class TestDisconnectGatherCleanup:
"common_processing_pre_call_logic",
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
)
- monkeypatch.setattr(
- processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
- )
+ monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
monkeypatch.setattr(
cpr,
"route_request",
@@ -3526,9 +3536,7 @@ class TestDisconnectGatherCleanup:
"common_processing_pre_call_logic",
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
)
- monkeypatch.setattr(
- processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
- )
+ monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
with pytest.raises(HTTPException):
await processing_obj.base_process_llm_request(
@@ -3579,9 +3587,7 @@ class TestDisconnectGatherCleanup:
assert task.done()
@pytest.mark.asyncio
- async def test_base_process_llm_request_preserves_llm_error_after_gather(
- self, monkeypatch
- ):
+ async def test_base_process_llm_request_preserves_llm_error_after_gather(self, monkeypatch):
import litellm.proxy.common_request_processing as cpr
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@@ -3610,9 +3616,7 @@ class TestDisconnectGatherCleanup:
"common_processing_pre_call_logic",
AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)),
)
- monkeypatch.setattr(
- processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)
- )
+ monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False))
mock_request = MagicMock(spec=Request)
mock_request.is_disconnected = AsyncMock(return_value=False)
@@ -3649,19 +3653,13 @@ class TestStreamingClientDisconnectLogging:
"litellm_params": {"metadata": {}},
}
- recorded = await _record_streaming_client_disconnect_if_needed(
- mock_request, request_data
- )
+ recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
assert recorded is True
assert request_data["metadata"]["client_disconnected"] is True
+ assert request_data["metadata"]["error_information"]["error_code"] == "499"
assert (
- request_data["metadata"]["error_information"]["error_code"] == "499"
- )
- assert (
- mock_logging_obj.model_call_details["litellm_params"]["metadata"][
- "error_information"
- ]["error_code"]
+ mock_logging_obj.model_call_details["litellm_params"]["metadata"]["error_information"]["error_code"]
== "499"
)
@@ -3675,9 +3673,7 @@ class TestStreamingClientDisconnectLogging:
mock_request.is_disconnected = AsyncMock(return_value=False)
request_data = {"metadata": {}}
- recorded = await _record_streaming_client_disconnect_if_needed(
- mock_request, request_data
- )
+ recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
assert recorded is False
assert "client_disconnected" not in request_data["metadata"]
@@ -3702,22 +3698,12 @@ class TestStreamingClientDisconnectLogging:
"litellm_params": {"metadata": {}},
}
- recorded = await _record_streaming_client_disconnect_if_needed(
- mock_request, request_data
- )
+ recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
assert recorded is True
assert request_data["metadata"]["client_disconnected"] is True
- assert (
- mock_logging_obj.model_call_details["litellm_params"]["metadata"][
- "client_disconnected"
- ]
- is True
- )
- assert (
- mock_logging_obj.model_call_details["metadata"]["client_disconnected"]
- is True
- )
+ assert mock_logging_obj.model_call_details["litellm_params"]["metadata"]["client_disconnected"] is True
+ assert mock_logging_obj.model_call_details["metadata"]["client_disconnected"] is True
@pytest.mark.asyncio
async def test_record_streaming_client_disconnect_handles_none_request_data_metadata(self):
@@ -3733,15 +3719,11 @@ class TestStreamingClientDisconnectLogging:
"litellm_params": {"metadata": None},
}
- recorded = await _record_streaming_client_disconnect_if_needed(
- mock_request, request_data
- )
+ recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data)
assert recorded is True
assert request_data["metadata"]["client_disconnected"] is True
- assert (
- request_data["litellm_params"]["metadata"]["client_disconnected"] is True
- )
+ assert request_data["litellm_params"]["metadata"]["client_disconnected"] is True
@pytest.mark.asyncio
async def test_apply_client_disconnect_metadata_none_returns_early(self):
@@ -3752,9 +3734,7 @@ class TestStreamingClientDisconnectLogging:
_apply_client_disconnect_metadata(None)
@pytest.mark.asyncio
- async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(
- self, monkeypatch
- ):
+ async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(self, monkeypatch):
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
@@ -3786,9 +3766,7 @@ class TestStreamingClientDisconnectLogging:
assert request_data["metadata"]["error_information"]["error_code"] == "499"
@pytest.mark.asyncio
- async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(
- self, monkeypatch
- ):
+ async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(self, monkeypatch):
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
@@ -3818,9 +3796,7 @@ class TestStreamingClientDisconnectLogging:
assert "client_disconnected" not in request_data["metadata"]
@pytest.mark.asyncio
- async def test_async_streaming_data_generator_records_499_on_early_aclose(
- self, monkeypatch
- ):
+ async def test_async_streaming_data_generator_records_499_on_early_aclose(self, monkeypatch):
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
@@ -3835,9 +3811,7 @@ class TestStreamingClientDisconnectLogging:
yield {"choices": [{"delta": {"content": " there"}}]}
mock_proxy_logging = MagicMock(spec=ProxyLogging)
- mock_proxy_logging.async_post_call_streaming_iterator_hook = (
- mock_streaming_iterator
- )
+ mock_proxy_logging.async_post_call_streaming_iterator_hook = mock_streaming_iterator
ProxyLogging._callback_capabilities_cache.clear()
mock_request = MagicMock(spec=Request)
@@ -3848,9 +3822,7 @@ class TestStreamingClientDisconnectLogging:
"model": "gemini-2.0-flash",
"metadata": {},
"litellm_params": {"metadata": {}},
- "litellm_logging_obj": MagicMock(
- model_call_details={"metadata": {}, "litellm_params": {}}
- ),
+ "litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}),
}
gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
@@ -3869,6 +3841,8 @@ class TestStreamingClientDisconnectLogging:
assert request_data["metadata"]["error_information"]["error_code"] == "499"
ProxyLogging._callback_capabilities_cache.clear()
+
+
class TestCancelOnDisconnect:
"""
Coverage for the opt-in `general_settings.cancel_on_disconnect` flag:
@@ -3895,23 +3869,17 @@ class TestCancelOnDisconnect:
llm_call = asyncio.get_running_loop().create_future()
disconnect_event = asyncio.Event()
- await _cancel_llm_call_on_client_disconnect(
- request, llm_call, disconnect_event
- )
+ await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
assert llm_call.cancelled()
assert disconnect_event.is_set()
async def test_monitor_is_noop_while_client_stays_connected(self):
- request = self._request(
- [{"type": "http.request", "body": b"", "more_body": False}]
- )
+ request = self._request([{"type": "http.request", "body": b"", "more_body": False}])
llm_call = asyncio.get_running_loop().create_future()
disconnect_event = asyncio.Event()
- monitor = asyncio.create_task(
- _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
- )
+ monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event))
await asyncio.sleep(0.01)
assert not monitor.done()
@@ -3930,9 +3898,7 @@ class TestCancelOnDisconnect:
llm_call = asyncio.get_running_loop().create_future()
disconnect_event = asyncio.Event()
- await _cancel_llm_call_on_client_disconnect(
- request, llm_call, disconnect_event
- )
+ await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)
assert not llm_call.cancelled()
assert not disconnect_event.is_set()
@@ -3947,9 +3913,7 @@ class TestCancelOnDisconnect:
with pytest.raises(asyncio.CancelledError):
await _await_llm_call_cancelling_on_disconnect(request, llm_call)
- async def _drive_base_process_llm_request(
- self, monkeypatch, general_settings: dict, llm_call, request: Request
- ):
+ async def _drive_base_process_llm_request(self, monkeypatch, general_settings: dict, llm_call, request: Request):
from litellm.proxy._types import UserAPIKeyAuth
logging_obj = MagicMock()
@@ -3958,9 +3922,7 @@ class TestCancelOnDisconnect:
logging_obj._on_deferred_stream_complete = None
logging_obj.cost_breakdown = None
- processor = ProxyBaseLLMRequestProcessing(
- data={"model": "fake-model", "litellm_logging_obj": logging_obj}
- )
+ processor = ProxyBaseLLMRequestProcessing(data={"model": "fake-model", "litellm_logging_obj": logging_obj})
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
@@ -3968,9 +3930,7 @@ class TestCancelOnDisconnect:
proxy_logging_obj.post_call_success_hook = AsyncMock(
side_effect=lambda data, user_api_key_dict, response: response
)
- proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
- return_value=None
- )
+ proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None)
async def fake_route_request(**kwargs):
return llm_call()
@@ -4049,9 +4009,7 @@ class TestCancelOnDisconnect:
with pytest.raises(ProxyException) as exc_info:
await processor._handle_llm_api_exception(
- e=HTTPException(
- status_code=499, detail="Client disconnected the request"
- ),
+ e=HTTPException(status_code=499, detail="Client disconnected the request"),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=proxy_logging_obj,
)
@@ -4117,7 +4075,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", capture_hook)
- with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
+ with patch.object(
+ ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
+ ):
processing_obj = ProxyBaseLLMRequestProcessing(data={})
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
response=httpx_response,
@@ -4167,7 +4127,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock())
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", non_dict_hook)
- with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
+ with patch.object(
+ ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
+ ):
processing_obj = ProxyBaseLLMRequestProcessing(data={})
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
response=httpx_response,
@@ -4205,7 +4167,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
hook_spy = AsyncMock()
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy)
- with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
+ with patch.object(
+ ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
+ ):
processing_obj = ProxyBaseLLMRequestProcessing(data={})
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
response=httpx_response,
@@ -4246,7 +4210,9 @@ class TestAllmPassthroughRoutePostCallGuardrails:
hook_spy = AsyncMock()
monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy)
- with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False):
+ with patch.object(
+ ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False
+ ):
processing_obj = ProxyBaseLLMRequestProcessing(data={})
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
response=httpx_response,
@@ -4358,7 +4324,9 @@ class TestEventStreamAllmPassthroughRoute:
"content-length": "99",
}
- with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True):
+ with patch.object(
+ ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True
+ ):
processing_obj = ProxyBaseLLMRequestProcessing(data={})
result = await processing_obj._handle_non_streaming_allm_passthrough_route(
response=mock_response,
@@ -4389,9 +4357,7 @@ class TestAllmPassthroughStreamingProviderGate:
de-anonymized.
"""
- def _build_processing_obj(
- self, custom_llm_provider: str, endpoint: str = ""
- ) -> ProxyBaseLLMRequestProcessing:
+ def _build_processing_obj(self, custom_llm_provider: str, endpoint: str = "") -> ProxyBaseLLMRequestProcessing:
logging_obj = MagicMock()
logging_obj.litellm_call_id = "call-123"
logging_obj.cost_breakdown = None
@@ -4442,14 +4408,17 @@ class TestAllmPassthroughStreamingProviderGate:
processing_obj = self._build_processing_obj("anthropic")
chunks = [b"chunk-1", b"chunk-2"]
- with patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails",
- return_value=False,
- ), patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails_for_passthrough",
- return_value=True,
+ with (
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails",
+ return_value=False,
+ ),
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails_for_passthrough",
+ return_value=True,
+ ),
):
result = await self._run(processing_obj, monkeypatch, chunks)
@@ -4458,27 +4427,27 @@ class TestAllmPassthroughStreamingProviderGate:
assert streamed == chunks
@pytest.mark.asyncio
- async def test_bedrock_converse_stream_is_buffered_through_handler(
- self, monkeypatch
- ):
- processing_obj = self._build_processing_obj(
- "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream"
- )
+ async def test_bedrock_converse_stream_is_buffered_through_handler(self, monkeypatch):
+ processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream")
chunks = [b"raw-1", b"raw-2"]
- with patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails",
- return_value=False,
- ), patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails_for_passthrough",
- return_value=True,
- ), patch(
- "litellm.llms.bedrock.passthrough.guardrail_translation.handler."
- "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
- new=AsyncMock(return_value=b"modified-body"),
- ) as mock_handler:
+ with (
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails",
+ return_value=False,
+ ),
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails_for_passthrough",
+ return_value=True,
+ ),
+ patch(
+ "litellm.llms.bedrock.passthrough.guardrail_translation.handler."
+ "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
+ new=AsyncMock(return_value=b"modified-body"),
+ ) as mock_handler,
+ ):
result = await self._run(processing_obj, monkeypatch, chunks)
assert isinstance(result, Response)
@@ -4494,19 +4463,23 @@ class TestAllmPassthroughStreamingProviderGate:
)
chunks = [b"raw-1", b"raw-2"]
- with patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails",
- return_value=False,
- ), patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails_for_passthrough",
- return_value=True,
- ), patch(
- "litellm.llms.bedrock.passthrough.guardrail_translation.handler."
- "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
- new=AsyncMock(return_value=b"modified-body"),
- ) as mock_handler:
+ with (
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails",
+ return_value=False,
+ ),
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails_for_passthrough",
+ return_value=True,
+ ),
+ patch(
+ "litellm.llms.bedrock.passthrough.guardrail_translation.handler."
+ "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream",
+ new=AsyncMock(return_value=b"modified-body"),
+ ) as mock_handler,
+ ):
result = await self._run(processing_obj, monkeypatch, chunks)
assert isinstance(result, StreamingResponse)
@@ -4528,14 +4501,17 @@ class TestAllmPassthroughStreamingProviderGate:
)
chunks = [b"raw-1", b"raw-2"]
- with patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails",
- return_value=False,
- ), patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails_for_passthrough",
- return_value=False,
+ with (
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails",
+ return_value=False,
+ ),
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails_for_passthrough",
+ return_value=False,
+ ),
):
result = await self._run(processing_obj, monkeypatch, chunks)
@@ -4554,14 +4530,17 @@ class TestAllmPassthroughStreamingProviderGate:
processing_obj = self._build_processing_obj("anthropic")
chunks = [b"chunk-1", b"chunk-2"]
- with patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails",
- return_value=False,
- ), patch.object(
- ProxyBaseLLMRequestProcessing,
- "_has_post_call_guardrails_for_passthrough",
- return_value=False,
+ with (
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails",
+ return_value=False,
+ ),
+ patch.object(
+ ProxyBaseLLMRequestProcessing,
+ "_has_post_call_guardrails_for_passthrough",
+ return_value=False,
+ ),
):
result = await self._run(processing_obj, monkeypatch, chunks)
@@ -4902,7 +4881,6 @@ class TestResponseCostHeaderForTypedDictResponses:
class TestPreCallWithFallbacksOnLocalRateLimit:
-
@pytest.mark.asyncio
async def test_fallback_triggered_on_local_rate_limit(self):
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
@@ -5054,9 +5032,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}]
user_api_key_dict = MagicMock()
- user_api_key_dict.router_settings = {
- "fallbacks": [{"gpt-4": ["claude-3-haiku"]}]
- }
+ user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]}
with patch.object(
processor,
@@ -5087,9 +5063,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
- processor = ProxyBaseLLMRequestProcessing(
- data={"model": "gpt-4", "disable_fallbacks": True}
- )
+ processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True})
async def mock_pre_call_logic(**kwargs):
raise ProxyRateLimitError(
@@ -5215,9 +5189,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
# Real per-key per-model TPM limiter + a key carrying the customer's
# `model_tpm_limit` metadata (only the primary is capped).
- limiter = _PROXY_MaxParallelRequestsHandler(
- internal_usage_cache=InternalUsageCache(DualCache())
- )
+ limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
user_api_key_dict = UserAPIKeyAuth(
api_key="sk-lit3890",
metadata={"model_tpm_limit": {primary_model: 100}},
@@ -5225,10 +5197,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
# Pre-seed the primary's per-model token counter at the cap so the very
# next request trips it. The counter key uses the *hashed* api_key.
- counter_key = (
- f"{user_api_key_dict.api_key}::{primary_model}"
- f"::{precise_minute}::request_count"
- )
+ counter_key = f"{user_api_key_dict.api_key}::{primary_model}::{precise_minute}::request_count"
await limiter.internal_usage_cache.async_set_cache(
key=counter_key,
value={"current_requests": 0, "current_tpm": 100, "current_rpm": 0},
@@ -5259,9 +5228,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
mock_router = MagicMock()
mock_router.fallbacks = [{primary_model: [fallback_model]}]
- with patch(
- "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock
- ):
+ with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock):
with patch.object(
processor,
"common_processing_pre_call_logic",
@@ -5291,9 +5258,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
# Sanity-check the premise: the limiter genuinely raises a
# ProxyRateLimitError for the capped primary under the frozen clock.
- with patch(
- "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock
- ):
+ with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock):
with pytest.raises(ProxyRateLimitError):
await limiter.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
@@ -5654,16 +5619,12 @@ class TestStreamingClientDisconnectBilling:
prompt_tokens=1000,
completion_tokens=10,
total_tokens=1010,
- prompt_tokens_details=PromptTokensDetailsWrapper(
- cached_tokens=500
- ),
+ prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500),
),
)
)
- event = await self._bill_and_collect_success_event(
- append_openai_style_cached_usage_chunk
- )
+ event = await self._bill_and_collect_success_event(append_openai_style_cached_usage_chunk)
usage = event["response_obj"].usage
assert getattr(usage, "cache_read_input_tokens", None) == 500
@@ -6433,9 +6394,7 @@ class TestInjectCostIntoUsageDict:
logging_obj.model_call_details["custom_llm_provider"] = "anthropic"
assert logging_obj.cost_breakdown is None
- model_response = ModelResponse(
- usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)
- )
+ model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224))
cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj)
assert cost is not None and cost > 0
@@ -6464,9 +6423,7 @@ class TestInjectCostIntoUsageDict:
)
existing = logging_obj.cost_breakdown
- model_response = ModelResponse(
- usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)
- )
+ model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224))
ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj)
assert logging_obj.cost_breakdown is existing
@@ -6761,9 +6718,7 @@ def test_ttft_keepalive_interval_only_arms_for_a_streaming_request(request_data,
@pytest.mark.asyncio
@pytest.mark.parametrize("stream_requested, expect_ping", [(True, True), (False, False)])
-async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(
- stream_requested, expect_ping
-):
+async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(stream_requested, expect_ping):
"""The wiring, not the helper: every route funnels through this method, and the
whole time-to-first-token is spent inside the call it wraps."""
@@ -6909,9 +6864,7 @@ async def test_a_late_failure_is_reported_to_the_failure_hook():
async def record(exc):
audited.append(exc)
- response = await open_sse_before_first_byte(
- slow_failure(), ping_interval_seconds=0.05, on_late_failure=record
- )
+ response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=record)
collected = await _drain(response)
assert [type(exc).__name__ for exc in audited] == ["HTTPException"]
@@ -6928,9 +6881,7 @@ async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame():
async def broken_hook(exc):
raise RuntimeError("the audit backend is down")
- response = await open_sse_before_first_byte(
- slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook
- )
+ response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook)
collected = await _drain(response)
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
@@ -6984,9 +6935,7 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke
[(0, False), (None, True)],
ids=["operator-hard-disabled-this-deployment", "deployment-says-nothing"],
)
-async def test_base_process_llm_request_honours_a_deployment_hard_disable(
- deployment_keepalive, expect_ping
-):
+async def test_base_process_llm_request_honours_a_deployment_hard_disable(deployment_keepalive, expect_ping):
"""`keepalive_seconds: 0` is documented as a disable a request cannot lift. The
funnel has to hand its router to the gate for that to hold before the upstream
has answered, since no deployment has served the request yet."""
@@ -7032,9 +6981,7 @@ async def test_a_hook_returning_a_replacement_decides_what_the_client_sees():
async def sanitize(exc):
return HTTPException(status_code=502, detail="upstream unavailable")
- response = await open_sse_before_first_byte(
- slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize
- )
+ response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize)
collected = await _drain(response)
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
@@ -7073,9 +7020,7 @@ async def test_a_hook_that_returns_nothing_leaves_the_real_error_intact():
async def audit_only(exc):
return None
- response = await open_sse_before_first_byte(
- slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only
- )
+ response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only)
collected = await _drain(response)
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
@@ -7092,9 +7037,7 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug():
async def broken_hook(exc):
raise RuntimeError("the audit backend is down")
- response = await open_sse_before_first_byte(
- slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook
- )
+ response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook)
collected = await _drain(response)
error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip())
From c3bcb6f64f787e256d106d98d3ef17dea525c78b Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Date: Tue, 25 Aug 2026 10:50:35 -0700
Subject: [PATCH 40/70] test(mcp): drain the logging worker after each test so
queued callbacks cannot leak into the next test (#38228)
LoggingWorker now carries still-queued coroutines onto the next event loop (12a34a10d8). Under xdist,
a success-logging coroutine queued by test_acompletion_mcp_respects_manual_approval ran nine seconds
later inside test_mcp_tool_call_hook on the same worker, resolved litellm.callbacks at run time and
overwrote that test's captured payload with a gpt-4o-mini completion (assert 1.35e-05 == 1.42).
Run clear_queue() in the suite's autouse teardown so every coroutine a test enqueues finishes before the
next test registers its callbacks, and add a subprocess regression test that runs the real conftest
against a stopped worker with work still queued.
---
tests/mcp_tests/conftest.py | 3 ++
tests/mcp_tests/test_mcp_logging.py | 45 +++++++++++++++++++++++++++++
2 files changed, 48 insertions(+)
diff --git a/tests/mcp_tests/conftest.py b/tests/mcp_tests/conftest.py
index d1dc3ec7216..5823893afc0 100644
--- a/tests/mcp_tests/conftest.py
+++ b/tests/mcp_tests/conftest.py
@@ -7,6 +7,7 @@ import pytest
import litellm
import asyncio
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@pytest.fixture(scope="session")
@@ -38,6 +39,8 @@ def setup_and_teardown():
yield
# Teardown code (executes after the yield point)
+ # LoggingWorker carries still-queued coroutines onto the next test's loop, where they'd log into that test's callbacks
+ asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
loop.close() # Close the loop created earlier
asyncio.set_event_loop(None) # Remove the reference to the loop
diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py
index 1903f29001f..fc9f675f837 100644
--- a/tests/mcp_tests/test_mcp_logging.py
+++ b/tests/mcp_tests/test_mcp_logging.py
@@ -1,6 +1,9 @@
import os
import pytest
import asyncio
+import subprocess
+import sys
+from pathlib import Path
from typing import Optional
from unittest.mock import AsyncMock, patch
@@ -458,3 +461,45 @@ async def test_mcp_tool_call_hook():
logged_standard_logging_payload is not None
), "Standard logging payload should not be None"
assert logged_standard_logging_payload["response_cost"] == 1.42
+
+
+_QUEUED_LOGGING_OUTLIVES_TEST = '''
+import time
+
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+ran_at = []
+
+
+async def _record_run():
+ ran_at.append(time.monotonic())
+
+
+async def test_1_leaves_logging_queued_behind_a_stopped_worker():
+ GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run())
+ await GLOBAL_LOGGING_WORKER.stop()
+ assert ran_at == []
+
+
+async def test_2_starts_after_the_previous_tests_logging_ran():
+ started_at = time.monotonic()
+ GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run())
+ await GLOBAL_LOGGING_WORKER.flush()
+ assert [t < started_at for t in ran_at] == [True, False]
+'''
+
+
+def test_logging_queued_by_one_test_is_drained_before_the_next(tmp_path: Path):
+ """Regression: a logging coroutine queued by one test must not run inside a later test (it would log into that
+ test's callbacks, which is how test_mcp_tool_call_hook captured a gpt-4o-mini payload under xdist)."""
+ (tmp_path / "conftest.py").write_text((Path(__file__).parent / "conftest.py").read_text())
+ (tmp_path / "pyproject.toml").write_text('[tool.pytest.ini_options]\nasyncio_mode = "auto"\n')
+ (tmp_path / "test_queued_logging.py").write_text(_QUEUED_LOGGING_OUTLIVES_TEST)
+ result = subprocess.run(
+ [sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider", "test_queued_logging.py"],
+ cwd=tmp_path,
+ capture_output=True,
+ text=True,
+ timeout=120,
+ )
+ assert result.returncode == 0, result.stdout + result.stderr
From e8bdbcd1cf176914a2b110f95e8fc9d87c574436 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 10:56:55 -0700
Subject: [PATCH 41/70] fix(bedrock_mantle): parse converse passthrough bodies
with the converse shape config for logging
---
.../passthrough/transformation.py | 29 +++++++++++-
...drock_mantle_passthrough_transformation.py | 45 +++++++++++++++++++
2 files changed, 73 insertions(+), 1 deletion(-)
diff --git a/litellm/llms/bedrock_mantle/passthrough/transformation.py b/litellm/llms/bedrock_mantle/passthrough/transformation.py
index 1393ac7c6e7..e6b831efa57 100644
--- a/litellm/llms/bedrock_mantle/passthrough/transformation.py
+++ b/litellm/llms/bedrock_mantle/passthrough/transformation.py
@@ -1,12 +1,19 @@
from collections.abc import Mapping
-from typing import Final, Literal
+from typing import TYPE_CHECKING, Final, Literal, Optional
+from httpx import Response
+
+from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig
from litellm.llms.bedrock_mantle.common_utils import (
MANTLE_HOST_RE,
resolve_mantle_bearer_token,
resolve_mantle_region,
)
+from litellm.types.utils import LlmProviders
+
+if TYPE_CHECKING:
+ from litellm.types.utils import CostResponseTypes
class BedrockMantlePassthroughConfig(BedrockPassthroughConfig):
@@ -42,3 +49,23 @@ class BedrockMantlePassthroughConfig(BedrockPassthroughConfig):
def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None:
api_key: Final = litellm_params.get("api_key")
return resolve_mantle_bearer_token(api_key if isinstance(api_key, str) else None)
+
+ def logging_non_streaming_response(
+ self,
+ model: str,
+ custom_llm_provider: str,
+ httpx_response: Response,
+ request_data: dict, # mutable-ok: mirrors the inherited BedrockPassthroughConfig signature
+ logging_obj: Logging,
+ endpoint: str,
+ ) -> Optional["CostResponseTypes"]:
+ is_converse: Final = "invoke" not in endpoint and "converse" in endpoint
+ shape_provider: Final = LlmProviders.BEDROCK.value if is_converse else custom_llm_provider
+ return super().logging_non_streaming_response(
+ model=model,
+ custom_llm_provider=shape_provider,
+ httpx_response=httpx_response,
+ request_data=request_data,
+ logging_obj=logging_obj,
+ endpoint=endpoint,
+ )
diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
index b7f9e492e14..8c6eda605ca 100644
--- a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
+++ b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py
@@ -14,6 +14,7 @@ from litellm.utils import ProviderConfigManager
MANTLE_API_BASE = "https://bedrock-mantle.us-east-2.api.aws"
INVOKE_ENDPOINT = "model/us.openai.gpt-5.6-sol/invoke"
+CONVERSE_ENDPOINT = "model/us.openai.gpt-5.6-sol/converse"
REQUEST_BODY = {"messages": [{"role": "user", "content": "say pong"}], "max_completion_tokens": 64}
@@ -150,3 +151,47 @@ def test_invoke_passthrough_route_reaches_bedrock_runtime_for_a_mantle_deploymen
assert str(sent["url"]) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}"
assert sent["headers"]["Authorization"] == f"Bearer {expected_bearer}"
assert json.loads(sent["content"]) == REQUEST_BODY
+
+
+def _logged_model_response(endpoint, body):
+ request = httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com/{endpoint}")
+ return BedrockMantlePassthroughConfig().logging_non_streaming_response(
+ model="us.openai.gpt-5.6-sol",
+ custom_llm_provider="bedrock_mantle",
+ httpx_response=httpx.Response(200, json=body, request=request),
+ request_data={"messages": [{"role": "user", "content": [{"text": "say pong"}]}]},
+ logging_obj=MagicMock(),
+ endpoint=endpoint,
+ )
+
+
+def test_converse_logging_parses_the_converse_response_shape():
+ result = _logged_model_response(
+ CONVERSE_ENDPOINT,
+ {
+ "metrics": {"latencyMs": 800.0},
+ "output": {"message": {"content": [{"text": "pong"}], "role": "assistant"}},
+ "stopReason": "end_turn",
+ "usage": {"inputTokens": 8, "outputTokens": 5, "totalTokens": 13},
+ },
+ )
+ assert result.choices[0].message.content == "pong"
+ assert result.usage.prompt_tokens == 8
+ assert result.usage.completion_tokens == 5
+
+
+def test_invoke_logging_parses_the_openai_chat_response_shape():
+ result = _logged_model_response(
+ INVOKE_ENDPOINT,
+ {
+ "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "pong", "role": "assistant"}}],
+ "created": 1787677792,
+ "id": "chatcmpl-regression",
+ "model": "us.openai.gpt-5.6-sol",
+ "object": "chat.completion",
+ "usage": {"completion_tokens": 5, "prompt_tokens": 8, "total_tokens": 13},
+ },
+ )
+ assert result.choices[0].message.content == "pong"
+ assert result.usage.prompt_tokens == 8
+ assert result.usage.completion_tokens == 5
From 5470c1bccbaa31aa1fccc5a5801c402588b43e80 Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Tue, 25 Aug 2026 10:59:26 -0700
Subject: [PATCH 42/70] fix(ui): forward OAuth
issuer/authorization/token/registration URLs from the MCP server edit form
(#38154)
The edit form's Authorize & Fetch Token button built its temporary OAuth
session payload without issuer, authorization_url, token_url, or
registration_url, unlike the create form's equivalent payload builder. The
backend's temporary-session endpoint builds its ephemeral server purely from
that payload, so any admin-configured OAuth endpoints on an existing server
were silently dropped, endpoint discovery fell back to (and failed against)
the plain server url, and Authorize & Fetch Token 400'd with "authorization
url is not configured" even though the saved server had those fields filled
in. Add the four missing fields to the edit form's temporary payload builder,
mirroring the create form.
---
.../_components/mcp_server_edit.test.tsx | 38 +++++++++++++++++++
.../_components/mcp_server_edit.tsx | 4 ++
2 files changed, 42 insertions(+)
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
index 438caa2f5e6..5aec78ba926 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx
@@ -381,6 +381,44 @@ describe("MCPServerEdit (true passthrough warning)", () => {
});
});
+describe("MCPServerEdit (OAuth authorize temp payload)", () => {
+ beforeEach(() => {
+ vi.clearAllMocks();
+ });
+
+ it("forwards issuer/authorization_url/token_url/registration_url to the temp OAuth session payload", async () => {
+ // Without these fields the ephemeral server the temp OAuth session endpoint builds has no
+ // admin-configured OAuth endpoints on it, discovery falls back to (and fails against) the
+ // plain server url, and Authorize & Fetch Token 400s with "authorization url is not
+ // configured" even though the saved server (and the visible form) has all four fields filled in.
+ render(
+ ,
+ );
+
+ await waitFor(() => {
+ expect(mockOauth.getTemporaryPayload).toBeTruthy();
+ });
+ const payload = mockOauth.getTemporaryPayload!();
+ expect(payload).toBeTruthy();
+ expect(payload?.issuer).toBe("https://github.com/login/oauth");
+ expect(payload?.authorization_url).toBe("https://github.com/login/oauth/authorize");
+ expect(payload?.token_url).toBe("https://github.com/login/oauth/access_token");
+ expect(payload?.registration_url).toBe("https://github.com/login/oauth/register");
+ });
+});
+
describe("MCPServerEdit (auth type switch)", () => {
beforeEach(() => {
vi.clearAllMocks();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
index 5e79b20825c..8793c45371a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx
@@ -282,6 +282,10 @@ const MCPServerEdit: React.FC = ({
credentials: isClientForwardedTokenMode(values.auth_type)
? preservedAdminCredentials(values.credentials)
: values.credentials,
+ issuer: values.issuer,
+ authorization_url: values.authorization_url,
+ token_url: values.token_url,
+ registration_url: values.registration_url,
mcp_access_groups: values.mcp_access_groups || mcpServer.mcp_access_groups,
static_headers: staticHeaders,
command: values.command,
From 559588a4732e94971781983eaf433cb11da9fbb9 Mon Sep 17 00:00:00 2001
From: milan
Date: Tue, 25 Aug 2026 18:23:09 +0000
Subject: [PATCH 43/70] fix(azure_ai): remove new LIT002 violations to satisfy
type-discipline budget
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/llms/azure_ai/azure_model_router/transformation.py | 2 +-
litellm/llms/azure_ai/common_utils.py | 6 +++++-
2 files changed, 6 insertions(+), 2 deletions(-)
diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py
index d33564c0f9a..61cbc213b11 100644
--- a/litellm/llms/azure_ai/azure_model_router/transformation.py
+++ b/litellm/llms/azure_ai/azure_model_router/transformation.py
@@ -99,7 +99,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
if selected_model:
# Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a
# class-level dict, so an in-place write can bleed into unrelated responses.
- transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter
+ transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter # mutable-ok: ModelResponse requires _hidden_params to be a plain dict
**get_hidden_params_dict(transformed_response),
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model,
}
diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py
index e8431f29f07..55a51176666 100644
--- a/litellm/llms/azure_ai/common_utils.py
+++ b/litellm/llms/azure_ai/common_utils.py
@@ -112,7 +112,11 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
"""
if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None:
return True
- deployment_model: Final = (hidden_params or {}).get("litellm_model_name") or (hidden_params or {}).get("model")
+ deployment_model: Final = (
+ hidden_params.get("litellm_model_name") or hidden_params.get("model")
+ if hidden_params is not None
+ else None
+ )
return any(
isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router"
for candidate in (deployment_model, model)
From 77f22be2075858dae62de1f44ecec85e6131ab86 Mon Sep 17 00:00:00 2001
From: milan
Date: Tue, 25 Aug 2026 18:27:19 +0000
Subject: [PATCH 44/70] style: apply ruff format to common_utils
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/llms/azure_ai/common_utils.py | 4 +---
1 file changed, 1 insertion(+), 3 deletions(-)
diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py
index 55a51176666..26a90157455 100644
--- a/litellm/llms/azure_ai/common_utils.py
+++ b/litellm/llms/azure_ai/common_utils.py
@@ -113,9 +113,7 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None:
return True
deployment_model: Final = (
- hidden_params.get("litellm_model_name") or hidden_params.get("model")
- if hidden_params is not None
- else None
+ hidden_params.get("litellm_model_name") or hidden_params.get("model") if hidden_params is not None else None
)
return any(
isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router"
From 6cd1fcdcf05734cff34d0dfed7e098da73a0a913 Mon Sep 17 00:00:00 2001
From: milan
Date: Tue, 25 Aug 2026 18:37:26 +0000
Subject: [PATCH 45/70] test: drop unneeded proxy_server patches and ratchet
lint budgets
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
ruff-strict-budget.json | 2 +-
.../proxy/spend_tracking/test_spend_tracking_utils.py | 11 ++---------
type-discipline-budget.json | 2 +-
3 files changed, 4 insertions(+), 11 deletions(-)
diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json
index 03318718fb5..9815ccfa23e 100644
--- a/ruff-strict-budget.json
+++ b/ruff-strict-budget.json
@@ -168,7 +168,7 @@
"limit": 3
},
"RET504": {
- "limit": 176
+ "limit": 175
},
"RUF012": {
"limit": 240
diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
index dab9415de77..843e1d1296f 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
@@ -3,13 +3,10 @@ import datetime
import json
from datetime import timezone
from typing import Any, Final, cast
-
-from typing_extensions import ReadOnly, TypedDict
+from unittest.mock import AsyncMock, MagicMock, patch
import pytest
-
-
-from unittest.mock import AsyncMock, MagicMock, patch
+from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.constants import (
@@ -3526,8 +3523,6 @@ def _model_router_spend_log_kwargs(slp_model: str | None) -> _ModelRouterSpendLo
}
-@patch("litellm.proxy.proxy_server.master_key", None)
-@patch("litellm.proxy.proxy_server.general_settings", {})
def test_get_logging_payload_uses_standard_logging_payload_model():
payload = get_logging_payload(
kwargs=_model_router_spend_log_kwargs(slp_model="azure_ai/gpt-5-mini"),
@@ -3538,8 +3533,6 @@ def test_get_logging_payload_uses_standard_logging_payload_model():
assert payload["model"] == "azure_ai/gpt-5-mini"
-@patch("litellm.proxy.proxy_server.master_key", None)
-@patch("litellm.proxy.proxy_server.general_settings", {})
def test_get_logging_payload_falls_back_to_kwargs_model_when_slp_model_missing():
payload = get_logging_payload(
kwargs=_model_router_spend_log_kwargs(slp_model=None),
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 05098546325..542762f1cea 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -30,7 +30,7 @@
"limit": 16673
},
"LIT011": {
- "limit": 5588
+ "limit": 5587
},
"LIT012": {
"limit": 4510
From 62ec3b61167d780b93839a7b03be30f2305f7e86 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 11:44:06 -0700
Subject: [PATCH 46/70] fix(together_ai): route chat completions through a
dedicated TogetherAIChatConfig
---
litellm/__init__.py | 3 +
litellm/_lazy_imports_registry.py | 5 +
.../get_supported_openai_params.py | 2 +-
litellm/llms/together_ai/chat.py | 58 -----
litellm/llms/together_ai/chat/__init__.py | 3 +
.../llms/together_ai/chat/transformation.py | 49 ++++
litellm/main.py | 62 ++++-
litellm/utils.py | 4 +-
.../test_together_ai_chat_transformation.py | 232 ++++++++++++++++++
9 files changed, 348 insertions(+), 70 deletions(-)
delete mode 100644 litellm/llms/together_ai/chat.py
create mode 100644 litellm/llms/together_ai/chat/__init__.py
create mode 100644 litellm/llms/together_ai/chat/transformation.py
create mode 100644 tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py
diff --git a/litellm/__init__.py b/litellm/__init__.py
index ee2c551481c..39556d1f04d 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -1628,6 +1628,9 @@ if TYPE_CHECKING:
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
)
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
+ from .llms.together_ai.chat.transformation import (
+ TogetherAIChatConfig as TogetherAIChatConfig,
+ )
from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig
from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig as VertexGeminiConfig,
diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py
index c34c9eefe85..1c833256598 100644
--- a/litellm/_lazy_imports_registry.py
+++ b/litellm/_lazy_imports_registry.py
@@ -177,6 +177,7 @@ LLM_CONFIG_NAMES: Final = (
"AmazonAnthropicClaudeMessagesConfig",
"AmazonMantleMessagesConfig",
"TogetherAIConfig",
+ "TogetherAIChatConfig",
"NLPCloudConfig",
"VertexGeminiConfig",
"GoogleAIStudioGeminiConfig",
@@ -741,6 +742,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
"AmazonMantleMessagesConfig",
),
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
+ "TogetherAIChatConfig": (
+ ".llms.together_ai.chat.transformation",
+ "TogetherAIChatConfig",
+ ),
"NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"),
"VertexGeminiConfig": (
".llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini",
diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py
index 72f36661f4c..7a16ffe4d85 100644
--- a/litellm/litellm_core_utils/get_supported_openai_params.py
+++ b/litellm/litellm_core_utils/get_supported_openai_params.py
@@ -172,7 +172,7 @@ def get_supported_openai_params(
if request_type == "embeddings":
return litellm.JinaAIEmbeddingConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "together_ai":
- return litellm.TogetherAIConfig().get_supported_openai_params(model=model)
+ return litellm.TogetherAIChatConfig().get_supported_openai_params(model=model)
elif custom_llm_provider == "databricks":
if request_type == "chat_completion":
return litellm.DatabricksConfig().get_supported_openai_params(model=model)
diff --git a/litellm/llms/together_ai/chat.py b/litellm/llms/together_ai/chat.py
deleted file mode 100644
index 58d47e45faa..00000000000
--- a/litellm/llms/together_ai/chat.py
+++ /dev/null
@@ -1,58 +0,0 @@
-"""
-Support for OpenAI's `/v1/chat/completions` endpoint.
-
-Calls done in OpenAI/openai.py as TogetherAI is openai-compatible.
-
-Docs: https://docs.together.ai/reference/completions-1
-"""
-
-from typing import Final
-
-from litellm._logging import verbose_logger
-from litellm.utils import supports_function_calling
-
-from ..openai.chat.gpt_transformation import OpenAIGPTConfig
-
-
-class TogetherAIConfig(OpenAIGPTConfig):
- def get_supported_openai_params(self, model: str) -> list:
- """
- Only some together models support response_format / tool calling
-
- Docs: https://docs.together.ai/docs/json-mode
- """
- # Use supports_function_calling() — which reads _get_model_info_helper
- # directly — instead of get_model_info(). get_model_info() calls
- # get_supported_openai_params() as its first step, which routes back
- # into this method for together_ai models, creating a recursion that
- # only terminates when Python's recursion limit or the "not mapped"
- # exception in _get_model_info_helper is hit (~332 deep calls).
- supports_fc: bool | None = None
- try:
- supports_fc = supports_function_calling(model, custom_llm_provider="together_ai")
- except Exception as e:
- verbose_logger.debug("Error getting supported openai params: %s", e)
-
- optional_params: Final = super().get_supported_openai_params(model)
- if supports_fc is not True:
- verbose_logger.debug(
- "Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling"
- )
- optional_params.remove("tools")
- optional_params.remove("tool_choice")
- optional_params.remove("function_call")
- optional_params.remove("response_format")
- return optional_params
-
- def map_openai_params(
- self,
- non_default_params: dict,
- optional_params: dict,
- model: str,
- drop_params: bool,
- ) -> dict:
- mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params)
-
- if "response_format" in mapped_openai_params and mapped_openai_params["response_format"] == {"type": "text"}:
- mapped_openai_params.pop("response_format")
- return mapped_openai_params
diff --git a/litellm/llms/together_ai/chat/__init__.py b/litellm/llms/together_ai/chat/__init__.py
new file mode 100644
index 00000000000..f260d9126d7
--- /dev/null
+++ b/litellm/llms/together_ai/chat/__init__.py
@@ -0,0 +1,3 @@
+from .transformation import TogetherAIChatConfig as TogetherAIChatConfig
+
+TogetherAIConfig = TogetherAIChatConfig
diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py
new file mode 100644
index 00000000000..eb0954bceef
--- /dev/null
+++ b/litellm/llms/together_ai/chat/transformation.py
@@ -0,0 +1,49 @@
+"""
+Translates from OpenAI's `/v1/chat/completions` to Together AI's `/v1/chat/completions`.
+
+Docs: https://docs.together.ai/docs/chat-overview
+"""
+
+from types import MappingProxyType
+from typing import Final
+
+from litellm._logging import verbose_logger
+from litellm.utils import supports_function_calling
+
+from ...openai.chat.gpt_transformation import OpenAIGPTConfig
+
+FUNCTION_CALLING_ONLY_PARAMS: Final = ("tools", "tool_choice", "function_call", "response_format")
+PLAIN_TEXT_RESPONSE_FORMAT: Final = MappingProxyType({"type": "text"})
+
+
+class TogetherAIChatConfig(OpenAIGPTConfig):
+ def get_supported_openai_params(self, model: str) -> list:
+ supports_fc: bool | None = None
+ try:
+ supports_fc = supports_function_calling(model, custom_llm_provider="together_ai")
+ except Exception as e:
+ verbose_logger.debug("Error getting supported openai params: %s", e)
+
+ supported_params: Final = super().get_supported_openai_params(model)
+ if supports_fc is True:
+ return supported_params
+ verbose_logger.debug(
+ "Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling"
+ )
+ for param in FUNCTION_CALLING_ONLY_PARAMS:
+ if param in supported_params:
+ supported_params.remove(param)
+ return supported_params
+
+ def map_openai_params(
+ self,
+ non_default_params: dict,
+ optional_params: dict,
+ model: str,
+ drop_params: bool,
+ ) -> dict:
+ mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params)
+
+ if mapped_openai_params.get("response_format") == PLAIN_TEXT_RESPONSE_FORMAT:
+ mapped_openai_params.pop("response_format")
+ return mapped_openai_params
diff --git a/litellm/main.py b/litellm/main.py
index d3967473f99..8ddcef2b9fa 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -24,6 +24,7 @@ from concurrent import futures
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from copy import deepcopy
from functools import partial
+from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
from litellm._logging import _redact_string
@@ -1811,6 +1812,56 @@ def _complete_fireworks_ai(
return response
+def _complete_together_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult:
+ acompletion: Final = ctx.acompletion
+ api_base: Final = ctx.api_base
+ api_key: Final = ctx.api_key
+ client: Final = _dispatch_client_http(ctx)
+ custom_llm_provider: Final = ctx.custom_llm_provider
+ headers: Final = ctx.headers
+ litellm_params: Final = ctx.litellm_params
+ logging: Final = ctx.logging
+ messages: Final = ctx.messages
+ model: Final = ctx.model
+ model_response: Final = ctx.model_response
+ optional_params: Final = ctx.optional_params
+ provider_config: Final = ctx.provider_config
+ shared_session: Final = ctx.shared_session
+ stream: Final = ctx.stream
+ timeout: Final = ctx.timeout
+
+ try:
+ response: Final = base_llm_http_handler.completion(
+ model=model,
+ messages=messages,
+ headers=headers,
+ model_response=model_response,
+ api_key=api_key,
+ api_base=api_base,
+ acompletion=acompletion,
+ logging_obj=logging,
+ optional_params=optional_params,
+ litellm_params=litellm_params,
+ shared_session=shared_session,
+ timeout=timeout,
+ client=client,
+ custom_llm_provider=custom_llm_provider,
+ encoding=_get_encoding(),
+ stream=stream,
+ provider_config=provider_config,
+ )
+ except Exception as e:
+ logging.post_call(
+ input=messages,
+ api_key=api_key,
+ original_response=str(e),
+ additional_args=MappingProxyType({"headers": headers}),
+ )
+ raise
+
+ return response
+
+
def _complete_heroku(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult:
acompletion: Final = ctx.acompletion
api_base: Final = ctx.api_base
@@ -5600,6 +5651,8 @@ def completion(
elif custom_llm_provider == "fireworks_ai":
## COMPLETION CALL
response = _complete_fireworks_ai(_dispatch_ctx)
+ elif custom_llm_provider == "together_ai":
+ response = _complete_together_ai(_dispatch_ctx)
elif custom_llm_provider == "heroku":
response = _complete_heroku(_dispatch_ctx)
@@ -5649,7 +5702,6 @@ def completion(
or custom_llm_provider == "volcengine"
or custom_llm_provider == "anyscale"
or custom_llm_provider == "openai"
- or custom_llm_provider == "together_ai"
or custom_llm_provider == "nebius"
or custom_llm_provider == "wandb"
or custom_llm_provider == "clarifai"
@@ -5699,14 +5751,6 @@ def completion(
response = _complete_openrouter(_dispatch_ctx)
elif custom_llm_provider == "vercel_ai_gateway":
response = _complete_vercel_ai_gateway(_dispatch_ctx)
- elif (
- custom_llm_provider == "together_ai"
- or ("togethercomputer" in model)
- or (model in litellm.together_ai_models)
- ):
- """
- Deprecated. We now do together ai calls via the openai client - https://docs.together.ai/docs/openai-api-compatibility
- """
elif custom_llm_provider == "palm":
raise ValueError(
"Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en"
diff --git a/litellm/utils.py b/litellm/utils.py
index 012e8785321..43f146d52d5 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -4130,7 +4130,7 @@ def get_optional_params(
drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False),
)
elif custom_llm_provider == "together_ai":
- optional_params = litellm.TogetherAIConfig().map_openai_params(
+ optional_params = litellm.TogetherAIChatConfig().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
@@ -7898,7 +7898,7 @@ class ProviderConfigManager:
LlmProviders.GALADRIEL: (lambda: litellm.GaladrielChatConfig(), False),
LlmProviders.REPLICATE: (lambda: litellm.ReplicateConfig(), False),
LlmProviders.HUGGINGFACE: (lambda: litellm.HuggingFaceChatConfig(), False),
- LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIConfig(), False),
+ LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIChatConfig(), False),
LlmProviders.OPENROUTER: (lambda: litellm.OpenrouterConfig(), False),
LlmProviders.VERCEL_AI_GATEWAY: (
lambda: litellm.VercelAIGatewayConfig(),
diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py
new file mode 100644
index 00000000000..6216d3bf225
--- /dev/null
+++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py
@@ -0,0 +1,232 @@
+import json
+from unittest.mock import MagicMock
+
+import httpx
+import pytest
+
+import litellm
+from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
+from litellm.llms.openai.chat.gpt_transformation import (
+ OpenAIChatCompletionStreamingHandler,
+)
+from litellm.llms.together_ai.chat.transformation import TogetherAIChatConfig
+from litellm.types.utils import LlmProviders, ModelResponse
+
+TOOL_CALLING_MODEL = "openai/gpt-oss-20b"
+REASONING_MODEL = "deepseek-ai/DeepSeek-V3.1"
+PLAIN_MODEL = "Qwen/Qwen3-235B-A22B-fp8-tput"
+UNMAPPED_MODEL = "MiniMaxAI/MiniMax-M3"
+
+FUNCTION_CALLING_PARAMS = ("tools", "tool_choice", "function_call", "response_format")
+
+
+@pytest.fixture(autouse=True)
+def force_local_model_cost(monkeypatch):
+ monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
+
+ monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url))
+
+
+def test_supported_params_tool_calling_model():
+ supported = TogetherAIChatConfig().get_supported_openai_params(model=TOOL_CALLING_MODEL)
+
+ for param in FUNCTION_CALLING_PARAMS:
+ assert param in supported
+
+
+def test_supported_params_plain_model():
+ supported = TogetherAIChatConfig().get_supported_openai_params(model=PLAIN_MODEL)
+
+ for param in FUNCTION_CALLING_PARAMS:
+ assert param not in supported
+ assert "temperature" in supported
+ assert "max_tokens" in supported
+
+
+def test_supported_params_unmapped_model_treated_as_plain():
+ supported = TogetherAIChatConfig().get_supported_openai_params(model=UNMAPPED_MODEL)
+
+ for param in FUNCTION_CALLING_PARAMS:
+ assert param not in supported
+ assert "stream" in supported
+
+
+def test_map_openai_params_tool_calling_model_passes_tools():
+ tools = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}]
+
+ mapped = TogetherAIChatConfig().map_openai_params(
+ non_default_params={"tools": tools, "tool_choice": "auto"},
+ optional_params={},
+ model=TOOL_CALLING_MODEL,
+ drop_params=False,
+ )
+
+ assert mapped["tools"] == tools
+ assert mapped["tool_choice"] == "auto"
+
+
+def test_map_openai_params_reasoning_model_passes_sampling_params():
+ mapped = TogetherAIChatConfig().map_openai_params(
+ non_default_params={"temperature": 0.2, "max_tokens": 512},
+ optional_params={},
+ model=REASONING_MODEL,
+ drop_params=False,
+ )
+
+ assert mapped["temperature"] == 0.2
+ assert mapped["max_tokens"] == 512
+
+
+def test_map_openai_params_drops_text_response_format():
+ mapped = TogetherAIChatConfig().map_openai_params(
+ non_default_params={"response_format": {"type": "text"}, "temperature": 0.5},
+ optional_params={},
+ model=REASONING_MODEL,
+ drop_params=False,
+ )
+
+ assert "response_format" not in mapped
+ assert mapped["temperature"] == 0.5
+
+
+def test_map_openai_params_keeps_json_response_format():
+ response_format = {"type": "json_object"}
+
+ mapped = TogetherAIChatConfig().map_openai_params(
+ non_default_params={"response_format": response_format},
+ optional_params={},
+ model=TOOL_CALLING_MODEL,
+ drop_params=False,
+ )
+
+ assert mapped["response_format"] == response_format
+
+
+def _transform_response(message: dict) -> ModelResponse:
+ raw_response_json = {
+ "id": "chatcmpl-test",
+ "object": "chat.completion",
+ "created": 1234567890,
+ "model": REASONING_MODEL,
+ "choices": [{"index": 0, "message": message, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+ }
+ mock_response = MagicMock(spec=httpx.Response)
+ mock_response.json.return_value = raw_response_json
+ mock_response.text = json.dumps(raw_response_json)
+ mock_response.headers = {}
+ logging_obj = MagicMock(spec=LiteLLMLoggingObj)
+ logging_obj.post_call = MagicMock()
+ logging_obj.model_call_details = {}
+
+ return TogetherAIChatConfig().transform_response(
+ model=REASONING_MODEL,
+ raw_response=mock_response,
+ model_response=ModelResponse(),
+ logging_obj=logging_obj,
+ request_data={},
+ messages=[{"role": "user", "content": "What is 2+2?"}],
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ api_key="test-key",
+ json_mode=False,
+ )
+
+
+def test_transform_response_maps_reasoning_to_reasoning_content():
+ result = _transform_response(
+ {"role": "assistant", "content": "4", "reasoning": "2+2 equals 4"}
+ )
+
+ assert result.choices[0].message.content == "4"
+ assert result.choices[0].message.reasoning_content == "2+2 equals 4"
+
+
+def test_transform_response_preserves_reasoning_content_field():
+ result = _transform_response(
+ {"role": "assistant", "content": "4", "reasoning_content": "adding 2 and 2"}
+ )
+
+ assert result.choices[0].message.reasoning_content == "adding 2 and 2"
+
+
+def test_streaming_chunk_maps_delta_reasoning_to_reasoning_content():
+ iterator = TogetherAIChatConfig().get_model_response_iterator(
+ streaming_response=iter(()), sync_stream=True
+ )
+ assert isinstance(iterator, OpenAIChatCompletionStreamingHandler)
+
+ parsed = iterator.chunk_parser(
+ {
+ "id": "chunk-1",
+ "created": 1234567890,
+ "model": REASONING_MODEL,
+ "choices": [{"index": 0, "delta": {"reasoning": "thinking about 2+2"}}],
+ }
+ )
+
+ assert parsed.choices[0]["delta"]["reasoning_content"] == "thinking about 2+2"
+
+
+def test_together_ai_config_alias_points_at_chat_config():
+ assert litellm.TogetherAIConfig is litellm.TogetherAIChatConfig
+ config = litellm.TogetherAIConfig(max_tokens=10)
+ assert isinstance(config, TogetherAIChatConfig)
+
+
+def test_provider_config_manager_returns_together_chat_config():
+ from litellm.utils import ProviderConfigManager
+
+ config = ProviderConfigManager.get_provider_chat_config(
+ model=REASONING_MODEL, provider=LlmProviders.TOGETHER_AI
+ )
+
+ assert isinstance(config, TogetherAIChatConfig)
+
+
+def test_completion_routes_through_together_chat_config():
+ from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+ captured_requests = []
+
+ def respond(request: httpx.Request) -> httpx.Response:
+ captured_requests.append(request)
+ return httpx.Response(
+ 200,
+ json={
+ "id": "chatcmpl-together",
+ "object": "chat.completion",
+ "created": 1234567890,
+ "model": REASONING_MODEL,
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": "4",
+ "reasoning": "2+2 equals 4",
+ },
+ "finish_reason": "stop",
+ }
+ ],
+ "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+ },
+ )
+
+ client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
+
+ response = litellm.completion(
+ model=f"together_ai/{REASONING_MODEL}",
+ messages=[{"role": "user", "content": "What is 2+2?"}],
+ api_key="fake-key",
+ client=client,
+ )
+
+ request = captured_requests[0]
+ assert str(request.url) == "https://api.together.ai/v1/chat/completions"
+ assert request.headers["authorization"] == "Bearer fake-key"
+ assert json.loads(request.content)["model"] == REASONING_MODEL
+ assert response.choices[0].message.content == "4"
+ assert response.choices[0].message.reasoning_content == "2+2 equals 4"
From 32ebfba5ed7810ead375c613ee2419e167ba831c Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 12:03:51 -0700
Subject: [PATCH 47/70] refactor(together_ai): build the trimmed
supported-params list without mutating the inherited list
---
litellm/llms/together_ai/chat/transformation.py | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py
index eb0954bceef..88fd79f2366 100644
--- a/litellm/llms/together_ai/chat/transformation.py
+++ b/litellm/llms/together_ai/chat/transformation.py
@@ -30,10 +30,9 @@ class TogetherAIChatConfig(OpenAIGPTConfig):
verbose_logger.debug(
"Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling"
)
- for param in FUNCTION_CALLING_ONLY_PARAMS:
- if param in supported_params:
- supported_params.remove(param)
- return supported_params
+ return [ # mutable-ok: the inherited contract returns a plain list; building fresh avoids mutating the base class's value
+ param for param in supported_params if param not in FUNCTION_CALLING_ONLY_PARAMS
+ ]
def map_openai_params(
self,
From 17845b4fb01b8ec3d0c90254bece1b802807f77c Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 12:20:51 -0700
Subject: [PATCH 48/70] fix(anthropic): translate tool_result document blocks
in the /v1/messages bridge
---
.../adapters/transformation.py | 6 +-
...al_pass_through_adapters_transformation.py | 71 +++++++++++++++++++
2 files changed, 74 insertions(+), 3 deletions(-)
diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
index 7c89da81fe6..109017bda27 100644
--- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
+++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
@@ -434,7 +434,7 @@ class LiteLLMAnthropicMessagesAdapter:
content_items = list(content.get("content", []))
# Single-item text keeps the backward-compatible string format; a single
- # image becomes a structured image_url part
+ # image or document becomes a structured image_url part
if len(content_items) == 1:
c = content_items[0]
if isinstance(c, str):
@@ -454,7 +454,7 @@ class LiteLLMAnthropicMessagesAdapter:
)
self._add_cache_control_if_applicable(content, tool_result, model)
tool_message_list.append(tool_result)
- elif c.get("type") == "image":
+ elif c.get("type") in ("image", "document"):
image_part = self._tool_result_image_part(c.get("source"))
tool_result = ChatCompletionToolMessage(
role="tool",
@@ -482,7 +482,7 @@ class LiteLLMAnthropicMessagesAdapter:
text=c.get("text", ""),
)
)
- elif c.get("type") == "image":
+ elif c.get("type") in ("image", "document"):
image_part = self._tool_result_image_part(c.get("source"))
if image_part:
combined_content_parts.append(image_part)
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py
index d0169963962..a7fbd069e61 100644
--- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py
@@ -1,3 +1,4 @@
+import base64
from typing import Any, cast
import pytest
@@ -11,6 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from litellm.litellm_core_utils.prompt_templates.factory import (
THOUGHT_SIGNATURE_SEPARATOR,
+ _bedrock_converse_messages_pt,
)
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import (
OPENAI_MAX_TOOL_NAME_LENGTH,
@@ -3872,6 +3874,75 @@ def test_tool_result_plain_text_unchanged_by_openai_transform():
assert _image_urls_in_user_messages(result) == []
+TOOL_RESULT_PDF_B64 = base64.b64encode(b"%PDF-1.4 minimal regression fixture").decode()
+
+
+def _base64_pdf_block():
+ return {
+ "type": "document",
+ "source": {"type": "base64", "media_type": "application/pdf", "data": TOOL_RESULT_PDF_B64},
+ }
+
+
+def test_tool_result_single_document_kept_as_pdf_data_url():
+ adapter = LiteLLMAnthropicMessagesAdapter()
+ translated = adapter.translate_anthropic_messages_to_openai(
+ messages=[
+ _anthropic_tool_use_turn("toolu_01"),
+ _anthropic_tool_result_turn({"toolu_01": [_base64_pdf_block()]}),
+ ]
+ )
+
+ tool_messages = [m for m in translated if m.get("role") == "tool"]
+ assert len(tool_messages) == 1
+ assert tool_messages[0]["content"] == [
+ {
+ "type": "image_url",
+ "image_url": {"url": f"data:application/pdf;base64,{TOOL_RESULT_PDF_B64}"},
+ }
+ ]
+
+
+def test_tool_result_text_and_document_reach_bedrock_converse_tool_result():
+ """Claude Code >= 2.1.245 sends Read-tool PDF output as a document block inside
+ tool_result; dropping it left bedrock converse models blind to the PDF content."""
+ adapter = LiteLLMAnthropicMessagesAdapter()
+ translated = adapter.translate_anthropic_messages_to_openai(
+ messages=[
+ AnthropicMessagesUserMessageParam(role="user", content="Read pong.pdf"),
+ _anthropic_tool_use_turn("toolu_01"),
+ _anthropic_tool_result_turn(
+ {
+ "toolu_01": [
+ {"type": "text", "text": "PDF file read: pong.pdf (579 bytes)"},
+ _base64_pdf_block(),
+ ]
+ }
+ ),
+ ]
+ )
+
+ converse_messages = _bedrock_converse_messages_pt(
+ messages=translated,
+ model="anthropic.claude-haiku-4-5-20251001-v1:0",
+ llm_provider="bedrock_converse",
+ )
+
+ tool_results = [
+ block["toolResult"]
+ for message in converse_messages
+ for block in message["content"]
+ if "toolResult" in block
+ ]
+ assert len(tool_results) == 1
+ documents = [part["document"] for part in tool_results[0]["content"] if "document" in part]
+ assert len(documents) == 1
+ assert documents[0]["format"] == "pdf"
+ assert documents[0]["source"]["bytes"] == TOOL_RESULT_PDF_B64
+ texts = [part["text"] for part in tool_results[0]["content"] if "text" in part]
+ assert texts == ["PDF file read: pong.pdf (579 bytes)"]
+
+
def test_translate_anthropic_to_openai_carries_prompt_cache_breakpoint_on_system_and_user_blocks():
explicit = {"mode": "explicit"}
openai_request, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
From 2bd2c1393cb231f6572b4e99f5cff427bcb66562 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 12:31:43 -0700
Subject: [PATCH 49/70] docs(pr-template): split Caveats bullets into severity
tiers and call for plain engineering language
---
.github/pull_request_template.md | 16 ++++++++++++++--
CLAUDE.md | 1 +
2 files changed, 15 insertions(+), 2 deletions(-)
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
index 4e428d8cebf..bcdb228746a 100644
--- a/.github/pull_request_template.md
+++ b/.github/pull_request_template.md
@@ -1,7 +1,10 @@
+
+
## TLDR
-
+
Problem this solves:
@@ -112,6 +115,15 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
## QA runbook
diff --git a/CLAUDE.md b/CLAUDE.md
index b3383b4a895..03053b8392c 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -44,6 +44,7 @@ If you ever make public-facing PR descriptions, comments, issues, commit message
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
+- do use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When structure genuinely helps the reader, prefer nested bullets (any depth is fine) over one dense line. This applies to all human-facing text: discussion posts, release notes, and docs included
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
From d749b186de18b1861a996044eca01156e1ae2a4e Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 12:34:31 -0700
Subject: [PATCH 50/70] docs(pr-template): make intent the severe-vs-high
discriminator
---
.github/pull_request_template.md | 10 ++++++----
1 file changed, 6 insertions(+), 4 deletions(-)
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
index bcdb228746a..8a19547cb34 100644
--- a/.github/pull_request_template.md
+++ b/.github/pull_request_template.md
@@ -116,10 +116,12 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
-### Final Attestation
+## Final Attestation
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
From 27ca05a70759d8c2e78b2e4c0bc08aa526640ed2 Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Tue, 25 Aug 2026 13:01:51 -0700
Subject: [PATCH 55/70] fix(ui): read reasoning tokens from Responses API
output_tokens_details (#37952)
---
.../src/components/llm_calls/responses_api.test.tsx | 12 ++++++++++++
.../src/components/llm_calls/responses_api.tsx | 6 ++++--
2 files changed, 16 insertions(+), 2 deletions(-)
diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx
index 065eaaf3632..033813397fc 100644
--- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx
+++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx
@@ -400,4 +400,16 @@ describe("responses_api prompt cache usage", () => {
expect(usageData).not.toHaveProperty("cacheCreationTokens");
expect(usageData.promptTokens).toBe(5000);
});
+
+ it("surfaces reasoning tokens from Responses-shape output_tokens_details", async () => {
+ await expect(captureUsage({ output_tokens_details: { reasoning_tokens: 42 } })).resolves.toMatchObject({
+ reasoningTokens: 42,
+ });
+ });
+
+ it("falls back to completion_tokens_details reasoning tokens when output_tokens_details is absent", async () => {
+ await expect(captureUsage({ completion_tokens_details: { reasoning_tokens: 17 } })).resolves.toMatchObject({
+ reasoningTokens: 17,
+ });
+ });
});
diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx
index 94e8cb46765..8d71a4e29a8 100644
--- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx
+++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx
@@ -295,8 +295,10 @@ export async function makeOpenAIResponsesRequest(
};
// Add reasoning tokens if available
- if (usage.completion_tokens_details?.reasoning_tokens) {
- usageData.reasoningTokens = usage.completion_tokens_details.reasoning_tokens;
+ const reasoningTokens =
+ usage.output_tokens_details?.reasoning_tokens ?? usage.completion_tokens_details?.reasoning_tokens;
+ if (reasoningTokens) {
+ usageData.reasoningTokens = reasoningTokens;
}
if (usage.cost !== undefined && usage.cost !== null) {
From 104fe73113cafa167a11edfa7c268b1d17b5ca29 Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Tue, 25 Aug 2026 13:07:09 -0700
Subject: [PATCH 56/70] fix(dashboard): don't show a stale provider
prompt-cache chip on a response-cache hit (#37951)
* fix(dashboard): don't show a stale provider prompt-cache chip on a response-cache hit
The playground's non-streaming chat completion and responses paths replayed a cache hit's original usage payload verbatim, so ResponseMetrics kept rendering the provider's prompt-cache-write/read chips using token counts from the original request. Detect the hit via the x-litellm-cache-key response header and render a Response Cache indicator instead.
* fix(dashboard): expose x-litellm-cache-key through CORS for the playground cache-hit indicator
---
litellm/constants.py | 1 +
tests/test_litellm/proxy/test_proxy_server.py | 10 +
.../chat_ui/ResponseMetrics.test.tsx | 18 ++
.../components/chat_ui/ResponseMetrics.tsx | 20 ++
.../llm_calls/chat_completion.test.tsx | 210 +++++++++++++++---
.../components/llm_calls/chat_completion.tsx | 10 +-
.../llm_calls/responses_api.test.tsx | 188 +++++++++++++---
.../components/llm_calls/responses_api.tsx | 12 +-
8 files changed, 408 insertions(+), 61 deletions(-)
diff --git a/litellm/constants.py b/litellm/constants.py
index 78aba30f9c0..765bbfe1e54 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -147,6 +147,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [
"x-litellm-adaptive-router-model",
"x-litellm-applied-guardrails",
"x-litellm-guardrail-scan-id",
+ "x-litellm-cache-key",
]
# Gemini model-specific minimal thinking budget constants
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 3383527e932..b9a31acca96 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -78,6 +78,16 @@ def client_no_auth():
return TestClient(app)
+def test_cors_exposes_cache_key_header_to_browser_js():
+ from fastapi.middleware.cors import CORSMiddleware
+
+ from litellm.constants import LITELLM_UI_ALLOW_HEADERS
+
+ cors_middleware = next(m for m in app.user_middleware if m.cls is CORSMiddleware)
+ assert cors_middleware.kwargs["expose_headers"] is LITELLM_UI_ALLOW_HEADERS
+ assert "x-litellm-cache-key" in cors_middleware.kwargs["expose_headers"]
+
+
def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch):
mock_login_result = {"user_id": "test-user"}
mock_prisma_client = MagicMock()
diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx
index 5afc94eb043..f31e1839739 100644
--- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx
+++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.test.tsx
@@ -33,4 +33,22 @@ describe("ResponseMetrics prompt cache chips", () => {
expect(screen.queryByText(/Cache Read/)).not.toBeInTheDocument();
expect(screen.queryByText(/Cache Write/)).not.toBeInTheDocument();
});
+
+ it("shows the response cache indicator instead of the provider cache chips on a response-cache hit", () => {
+ render(
+ ,
+ );
+
+ expect(screen.getByText("Response Cache: Hit")).toBeInTheDocument();
+ expect(screen.queryByText(/Cache Read/)).not.toBeInTheDocument();
+ expect(screen.queryByText(/Cache Write/)).not.toBeInTheDocument();
+ });
+
+ it("does not show the response cache indicator when the flag is absent", () => {
+ render();
+
+ expect(screen.queryByText(/Response Cache/)).not.toBeInTheDocument();
+ });
});
diff --git a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx
index 3e7f2884b23..ec62d0618d7 100644
--- a/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx
+++ b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx
@@ -7,12 +7,16 @@ import {
DatabaseBackup,
DollarSign,
Hash,
+ History,
Lightbulb,
Wrench,
} from "lucide-react";
import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
import { PROMPT_CACHE_CREATION_TOOLTIP, PROMPT_CACHE_READ_TOOLTIP } from "@/utils/promptCacheUsage";
+const RESPONSE_CACHE_TOOLTIP =
+ "This response was replayed from LiteLLM's response cache. The request never reached the provider, so it did not read from or write to the provider's own prompt cache.";
+
export interface TokenUsage {
completionTokens?: number;
promptTokens?: number;
@@ -21,6 +25,7 @@ export interface TokenUsage {
cacheReadTokens?: number;
cacheCreationTokens?: number;
cost?: number;
+ servedFromResponseCache?: boolean;
}
interface ResponseMetricsProps {
@@ -51,7 +56,22 @@ function MetricItem({ label, tooltip, icon, value }: MetricItemProps) {
);
}
+function ResponseCacheIndicator() {
+ return (
+ }
+ value="Hit"
+ />
+ );
+}
+
function PromptCacheChips({ usage }: { usage?: TokenUsage }) {
+ if (usage?.servedFromResponseCache) {
+ return ;
+ }
+
const readTokens = usage?.cacheReadTokens ?? 0;
const creationTokens = usage?.cacheCreationTokens ?? 0;
diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx
index c91de3f6b20..bc4b4e3a351 100644
--- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx
+++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx
@@ -24,6 +24,10 @@ vi.mock("openai", () => ({
},
}));
+const nonStreamingResponse = (data: unknown, headers: Record = {}) => ({
+ withResponse: async () => ({ data, response: { headers: new Headers(headers) } }),
+});
+
describe("chat_completion", () => {
const mockUpdateUI = vi.fn();
const mockChatHistory = [{ role: "user", content: "Hello" }];
@@ -226,25 +230,27 @@ describe("chat_completion", () => {
});
it("should send a non-streaming request and render the whole message at once when streaming is disabled", async () => {
- mockCreate.mockResolvedValueOnce({
- id: "chatcmpl-1",
- object: "chat.completion",
- created: 1,
- model: "gpt-4",
- choices: [
- {
- index: 0,
- finish_reason: "stop",
- message: { role: "assistant", content: "Hello there" },
+ mockCreate.mockReturnValueOnce(
+ nonStreamingResponse({
+ id: "chatcmpl-1",
+ object: "chat.completion",
+ created: 1,
+ model: "gpt-4",
+ choices: [
+ {
+ index: 0,
+ finish_reason: "stop",
+ message: { role: "assistant", content: "Hello there" },
+ },
+ ],
+ usage: {
+ completion_tokens: 2,
+ prompt_tokens: 5,
+ total_tokens: 7,
+ cost: 0.25,
},
- ],
- usage: {
- completion_tokens: 2,
- prompt_tokens: 5,
- total_tokens: 7,
- cost: 0.25,
- },
- });
+ }),
+ );
const onTimingData = vi.fn();
const onUsageData = vi.fn();
@@ -298,24 +304,26 @@ describe("chat_completion", () => {
});
it("should surface reasoning content and MCP metadata from a non-streaming response", async () => {
- mockCreate.mockResolvedValueOnce({
- model: "gpt-4",
- choices: [
- {
- index: 0,
- finish_reason: "stop",
- message: {
- role: "assistant",
- content: "done",
- reasoning_content: "thinking",
- provider_specific_fields: {
- mcp_tool_calls: [{ id: "call_1", function: { name: "search_docs", arguments: "{}" } }],
- mcp_call_results: [{ tool_call_id: "call_1", result: "found it" }],
+ mockCreate.mockReturnValueOnce(
+ nonStreamingResponse({
+ model: "gpt-4",
+ choices: [
+ {
+ index: 0,
+ finish_reason: "stop",
+ message: {
+ role: "assistant",
+ content: "done",
+ reasoning_content: "thinking",
+ provider_specific_fields: {
+ mcp_tool_calls: [{ id: "call_1", function: { name: "search_docs", arguments: "{}" } }],
+ mcp_call_results: [{ tool_call_id: "call_1", result: "found it" }],
+ },
},
},
- },
- ],
- });
+ ],
+ }),
+ );
const onReasoningContent = vi.fn();
const onMCPEvent = vi.fn();
@@ -459,3 +467,137 @@ describe("chat_completion prompt cache usage", () => {
expect(usageData).not.toHaveProperty("cacheCreationTokens");
});
});
+
+describe("chat_completion response cache", () => {
+ const mockUpdateUI = vi.fn();
+ const mockChatHistory = [{ role: "user", content: "Hello" }];
+
+ afterEach(() => {
+ vi.clearAllMocks();
+ });
+
+ it("flags a non-streaming response-cache hit even though it replays provider prompt-cache usage", async () => {
+ mockCreate.mockReturnValueOnce(
+ nonStreamingResponse(
+ {
+ id: "chatcmpl-replayed",
+ model: "gpt-4",
+ choices: [{ index: 0, finish_reason: "stop", message: { role: "assistant", content: "Hello there" } }],
+ usage: {
+ completion_tokens: 2,
+ prompt_tokens: 5000,
+ total_tokens: 5002,
+ prompt_tokens_details: { cached_tokens: 4695 },
+ },
+ },
+ { "x-litellm-cache-key": "cache-key-abc" },
+ ),
+ );
+
+ const onUsageData = vi.fn();
+
+ await makeOpenAIChatCompletionRequest(
+ mockChatHistory,
+ mockUpdateUI,
+ "gpt-4",
+ "test-token",
+ undefined, // tags
+ undefined, // signal
+ undefined, // onReasoningContent
+ undefined, // onTimingData
+ onUsageData,
+ undefined, // traceId
+ undefined, // vector_store_ids
+ undefined, // guardrails
+ undefined, // policies
+ undefined, // selectedMCPServers
+ undefined, // onImageGenerated
+ undefined, // onSearchResults
+ undefined, // temperature
+ undefined, // max_tokens
+ undefined, // onTotalLatency
+ undefined, // customBaseUrl
+ undefined, // mcpServers
+ undefined, // mcpServerToolRestrictions
+ undefined, // onMCPEvent
+ undefined, // mockTestFallbacks
+ undefined, // mcpToolsets
+ false, // streamingEnabled
+ );
+
+ expect(onUsageData).toHaveBeenCalledWith(
+ expect.objectContaining({ cacheReadTokens: 4695, servedFromResponseCache: true }),
+ );
+ });
+
+ it("does not flag a non-streaming response that missed the response cache", async () => {
+ mockCreate.mockReturnValueOnce(
+ nonStreamingResponse({
+ id: "chatcmpl-fresh",
+ model: "gpt-4",
+ choices: [{ index: 0, finish_reason: "stop", message: { role: "assistant", content: "Hello there" } }],
+ usage: { completion_tokens: 2, prompt_tokens: 5, total_tokens: 7 },
+ }),
+ );
+
+ const onUsageData = vi.fn();
+
+ await makeOpenAIChatCompletionRequest(
+ mockChatHistory,
+ mockUpdateUI,
+ "gpt-4",
+ "test-token",
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ onUsageData,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ false, // streamingEnabled
+ );
+
+ expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }));
+ });
+
+ it("never flags a streaming response, even when the proxy reports a cache key", async () => {
+ async function* mockStream() {
+ yield {
+ choices: [{ delta: {}, index: 0 }],
+ model: "gpt-4",
+ usage: { completion_tokens: 2, prompt_tokens: 5, total_tokens: 7 },
+ };
+ }
+ mockCreate.mockResolvedValueOnce(mockStream());
+
+ const onUsageData = vi.fn();
+
+ await makeOpenAIChatCompletionRequest(
+ mockChatHistory,
+ mockUpdateUI,
+ "gpt-4",
+ "test-token",
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ onUsageData,
+ );
+
+ expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }));
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx
index ea66006349b..cd0852bc06e 100644
--- a/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx
+++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx
@@ -73,6 +73,7 @@ export async function makeOpenAIChatCompletionRequest(
const startTime = Date.now();
let firstTokenReceived = false;
let timeToFirstToken: number | undefined = undefined;
+ let servedFromResponseCache = false;
// Track MCP metadata cumulatively across chunks
let mcpMetadata: {
@@ -143,7 +144,13 @@ export async function makeOpenAIChatCompletionRequest(
{ ...requestBody, stream: true, stream_options: { include_usage: true } },
{ signal },
)
- : [completionAsSingleChunk(await client.chat.completions.create({ ...requestBody, stream: false }, { signal }))];
+ : await (async () => {
+ const nonStreamingResponse = await client.chat.completions
+ .create({ ...requestBody, stream: false }, { signal })
+ .withResponse();
+ servedFromResponseCache = nonStreamingResponse.response.headers.get("x-litellm-cache-key") !== null;
+ return [completionAsSingleChunk(nonStreamingResponse.data)];
+ })();
for await (const chunk of response) {
// Process content and measure time to first token
@@ -228,6 +235,7 @@ export async function makeOpenAIChatCompletionRequest(
promptTokens: chunkWithUsage.usage.prompt_tokens,
totalTokens: chunkWithUsage.usage.total_tokens,
...extractPromptCacheTokens(chunkWithUsage.usage),
+ ...(servedFromResponseCache ? { servedFromResponseCache: true } : {}),
};
// Check for reasoning tokens
diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx
index 033813397fc..0b94c093acd 100644
--- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx
+++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx
@@ -20,6 +20,10 @@ vi.mock("openai", () => ({
},
}));
+const nonStreamingResponse = (data: unknown, headers: Record = {}) => ({
+ withResponse: async () => ({ data, response: { headers: new Headers(headers) } }),
+});
+
describe("responses_api", () => {
const mockUpdateTextUI = vi.fn();
const messages: MessageType[] = [{ role: "user", content: "Hello" }];
@@ -71,19 +75,21 @@ describe("responses_api", () => {
});
it("should send a non-streaming request and render the whole output at once when streaming is disabled", async () => {
- mockResponsesCreate.mockResolvedValueOnce({
- id: "resp_456",
- output: [
- {
- type: "message",
- content: [
- { type: "output_text", text: "Full " },
- { type: "output_text", text: "answer" },
- ],
- },
- ],
- usage: { output_tokens: 3, input_tokens: 4, total_tokens: 7 },
- });
+ mockResponsesCreate.mockReturnValueOnce(
+ nonStreamingResponse({
+ id: "resp_456",
+ output: [
+ {
+ type: "message",
+ content: [
+ { type: "output_text", text: "Full " },
+ { type: "output_text", text: "answer" },
+ ],
+ },
+ ],
+ usage: { output_tokens: 3, input_tokens: 4, total_tokens: 7 },
+ }),
+ );
const onTimingData = vi.fn();
const onUsageData = vi.fn();
@@ -162,10 +168,12 @@ describe("responses_api", () => {
expect(onTotalLatency).toHaveBeenCalledTimes(1);
expect(onTotalLatency).toHaveBeenLastCalledWith(expect.any(Number));
- mockResponsesCreate.mockResolvedValueOnce({
- id: "resp_latency",
- output: [{ type: "message", content: [{ type: "output_text", text: "Answer" }] }],
- });
+ mockResponsesCreate.mockReturnValueOnce(
+ nonStreamingResponse({
+ id: "resp_latency",
+ output: [{ type: "message", content: [{ type: "output_text", text: "Answer" }] }],
+ }),
+ );
await callWithStreaming(false);
expect(onTotalLatency).toHaveBeenCalledTimes(2);
@@ -224,14 +232,16 @@ describe("responses_api", () => {
});
it("should replay MCP output items as events for a non-streaming response", async () => {
- mockResponsesCreate.mockResolvedValueOnce({
- id: "resp_789",
- output: [
- { type: "mcp_call", id: "mcp_1", name: "search_docs", arguments: "{}", output: "found it" },
- { type: "message", content: [{ type: "output_text", text: "Answer" }] },
- ],
- usage: { output_tokens: 1, input_tokens: 1, total_tokens: 2 },
- });
+ mockResponsesCreate.mockReturnValueOnce(
+ nonStreamingResponse({
+ id: "resp_789",
+ output: [
+ { type: "mcp_call", id: "mcp_1", name: "search_docs", arguments: "{}", output: "found it" },
+ { type: "message", content: [{ type: "output_text", text: "Answer" }] },
+ ],
+ usage: { output_tokens: 1, input_tokens: 1, total_tokens: 2 },
+ }),
+ );
const onMCPEvent = vi.fn();
const onUsageData = vi.fn();
@@ -413,3 +423,131 @@ describe("responses_api prompt cache usage", () => {
});
});
});
+
+describe("responses_api response cache", () => {
+ const mockUpdateTextUI = vi.fn();
+ const messages: MessageType[] = [{ role: "user", content: "Hello" }];
+
+ afterEach(() => {
+ vi.clearAllMocks();
+ });
+
+ it("flags a non-streaming response-cache hit even though it replays provider prompt-cache usage", async () => {
+ mockResponsesCreate.mockReturnValueOnce(
+ nonStreamingResponse(
+ {
+ id: "resp_replayed",
+ output: [{ type: "message", content: [{ type: "output_text", text: "Full answer" }] }],
+ usage: {
+ output_tokens: 2,
+ input_tokens: 5000,
+ total_tokens: 5002,
+ input_tokens_details: { cached_tokens: 4695 },
+ },
+ },
+ { "x-litellm-cache-key": "cache-key-abc" },
+ ),
+ );
+
+ const onUsageData = vi.fn();
+
+ await makeOpenAIResponsesRequest(
+ messages,
+ mockUpdateTextUI,
+ "gpt-4",
+ "test-token",
+ undefined, // tags
+ undefined, // signal
+ undefined, // onReasoningContent
+ undefined, // onTimingData
+ onUsageData,
+ undefined, // traceId
+ undefined, // vector_store_ids
+ undefined, // guardrails
+ undefined, // policies
+ undefined, // selectedMCPServers
+ undefined, // previousResponseId
+ undefined, // onResponseId
+ undefined, // onMCPEvent
+ undefined, // codeInterpreterEnabled
+ undefined, // onCodeInterpreterResult
+ undefined, // customBaseUrl
+ undefined, // mcpServers
+ undefined, // mcpServerToolRestrictions
+ undefined, // mcpToolsets
+ false, // streamingEnabled
+ );
+
+ expect(onUsageData).toHaveBeenCalledWith(
+ expect.objectContaining({ cacheReadTokens: 4695, servedFromResponseCache: true }),
+ "",
+ );
+ });
+
+ it("does not flag a non-streaming response that missed the response cache", async () => {
+ mockResponsesCreate.mockReturnValueOnce(
+ nonStreamingResponse({
+ id: "resp_fresh",
+ output: [{ type: "message", content: [{ type: "output_text", text: "Full answer" }] }],
+ usage: { output_tokens: 2, input_tokens: 5, total_tokens: 7 },
+ }),
+ );
+
+ const onUsageData = vi.fn();
+
+ await makeOpenAIResponsesRequest(
+ messages,
+ mockUpdateTextUI,
+ "gpt-4",
+ "test-token",
+ undefined, // tags
+ undefined, // signal
+ undefined, // onReasoningContent
+ undefined, // onTimingData
+ onUsageData,
+ undefined, // traceId
+ undefined, // vector_store_ids
+ undefined, // guardrails
+ undefined, // policies
+ undefined, // selectedMCPServers
+ undefined, // previousResponseId
+ undefined, // onResponseId
+ undefined, // onMCPEvent
+ undefined, // codeInterpreterEnabled
+ undefined, // onCodeInterpreterResult
+ undefined, // customBaseUrl
+ undefined, // mcpServers
+ undefined, // mcpServerToolRestrictions
+ undefined, // mcpToolsets
+ false, // streamingEnabled
+ );
+
+ expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }), "");
+ });
+
+ it("never flags a streaming response, even when the proxy reports a cache key", async () => {
+ async function* mockStream() {
+ yield {
+ type: "response.completed",
+ response: { id: "resp_stream", usage: { output_tokens: 2, input_tokens: 5, total_tokens: 7 } },
+ };
+ }
+ mockResponsesCreate.mockResolvedValueOnce(mockStream());
+
+ const onUsageData = vi.fn();
+
+ await makeOpenAIResponsesRequest(
+ messages,
+ mockUpdateTextUI,
+ "gpt-4",
+ "test-token",
+ undefined,
+ undefined,
+ undefined,
+ undefined,
+ onUsageData,
+ );
+
+ expect(onUsageData).toHaveBeenCalledWith(expect.not.objectContaining({ servedFromResponseCache: true }), "");
+ });
+});
diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx
index 8d71a4e29a8..7ab76488504 100644
--- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx
+++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx
@@ -116,6 +116,7 @@ export async function makeOpenAIResponsesRequest(
try {
const startTime = Date.now();
let firstTokenReceived = false;
+ let servedFromResponseCache = false;
// Format messages for the API
const formattedInput = messages.map((message) => {
@@ -202,7 +203,15 @@ export async function makeOpenAIResponsesRequest(
// Create request to OpenAI responses API
// Use 'any' type to avoid TypeScript issues with the experimental API
- const response = await (client as any).responses.create({ ...requestBody, stream: streamingEnabled }, { signal });
+ const response = streamingEnabled
+ ? await (client as any).responses.create({ ...requestBody, stream: true }, { signal })
+ : await (async () => {
+ const nonStreamingResponse = await (client as any).responses
+ .create({ ...requestBody, stream: false }, { signal })
+ .withResponse();
+ servedFromResponseCache = nonStreamingResponse.response.headers.get("x-litellm-cache-key") !== null;
+ return nonStreamingResponse.data;
+ })();
const events = streamingEnabled ? response : responseAsEvents(response);
let mcpToolUsed = "";
@@ -292,6 +301,7 @@ export async function makeOpenAIResponsesRequest(
promptTokens: usage.input_tokens,
totalTokens: usage.total_tokens,
...extractPromptCacheTokens(usage),
+ ...(servedFromResponseCache ? { servedFromResponseCache: true } : {}),
};
// Add reasoning tokens if available
From 5998c0d1d2ecba4b1405420b6dcbcccbfce4845c Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 13:32:26 -0700
Subject: [PATCH 57/70] fix(anthropic): carry tool_result document blocks
through the /v1/messages responses bridge
---
.../responses_adapters/transformation.py | 54 +++++++-
.../test_responses_adapters_transformation.py | 118 +++++++++++++++++-
2 files changed, 170 insertions(+), 2 deletions(-)
diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
index 6d47d0de19f..43c3e504964 100644
--- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
+++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
@@ -87,6 +87,51 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
return source.get("url")
return None
+ @staticmethod
+ def _translate_anthropic_document_block_to_file_part(
+ block: Mapping[str, object],
+ ) -> dict[str, str] | None: # mutable-ok: API message payload
+ """Convert an Anthropic document block to a Responses input_file part."""
+ raw_source: Final = block.get("source")
+ if not isinstance(raw_source, Mapping):
+ return None
+ source: Final = cast(Mapping[str, object], raw_source) # cast-ok: untrusted client payload
+ source_type: Final = source.get("type")
+ if source_type == "base64":
+ data: Final = source.get("data")
+ if not isinstance(data, str) or not data:
+ return None
+ raw_media_type: Final = source.get("media_type")
+ media_type: Final = (
+ raw_media_type if isinstance(raw_media_type, str) and raw_media_type else "application/pdf"
+ )
+ raw_title: Final = block.get("title")
+ filename: Final = raw_title if isinstance(raw_title, str) and raw_title else "document.pdf"
+ return { # mutable-ok: API message payload
+ "type": "input_file",
+ "filename": filename,
+ "file_data": f"data:{media_type};base64,{data}",
+ }
+ if source_type == "url":
+ url: Final = source.get("url")
+ if not isinstance(url, str) or not url:
+ return None
+ return {"type": "input_file", "file_url": url} # mutable-ok: API message payload
+ return None
+
+ @staticmethod
+ def _tool_result_output_value(
+ output_text: str,
+ file_parts: tuple[dict[str, str], ...], # mutable-ok: json content parts
+ ) -> str | list[dict[str, str]]: # mutable-ok: API message payload
+ """Plain string output, or a part list when document file parts are present."""
+ if not file_parts:
+ return output_text
+ text_parts: Final = (
+ [{"type": "input_text", "text": output_text}] if output_text else [] # mutable-ok: API message payload
+ )
+ return [*text_parts, *file_parts] # mutable-ok: API message payload
+
@staticmethod
def _translate_midturn_system_content_to_responses(
content: str | Iterable[AnthropicSystemMessageContent],
@@ -226,6 +271,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
elif btype == "tool_result":
tool_use_id = block.get("tool_use_id", "")
inner = block.get("content")
+ tool_file_parts: tuple[dict[str, str], ...] = () # mutable-ok: json content parts
if inner is None:
output_text = ""
elif isinstance(inner, str):
@@ -251,6 +297,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
{"type": "input_image", "image_url": url} # mutable-ok: json content part
for url in image_urls
)
+ document_candidates = tuple(
+ self._translate_anthropic_document_block_to_file_part(c)
+ for c in inner
+ if isinstance(c, dict) and c.get("type") == "document"
+ )
+ tool_file_parts = tuple(part for part in document_candidates if part is not None)
else:
output_text = str(inner)
# tool_result is a top-level item, not inside the message
@@ -258,7 +310,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
{
"type": "function_call_output",
"call_id": tool_use_id,
- "output": output_text,
+ "output": self._tool_result_output_value(output_text, tool_file_parts),
}
)
if tool_image_parts:
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
index 8225e7cff39..f0cdaa4e8e3 100644
--- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
@@ -16,7 +16,10 @@ from litellm.constants import (
DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET,
DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
)
-from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+ TOOL_RESULT_IMAGE_BOUNDARY,
+ TOOL_RESULT_IMAGE_PLACEHOLDER,
+)
from litellm.llms.anthropic.experimental_pass_through.responses_adapters.transformation import (
LiteLLMAnthropicToResponsesAPIAdapter,
)
@@ -1555,6 +1558,119 @@ class TestToolResultImages:
assert self._input_images(items) == []
+class TestToolResultDocuments:
+ """Documents inside tool_result blocks must survive translation (LIT-6135):
+ the function_call_output output becomes a list of parts carrying the joined
+ text as input_text and each document as an input_file. Without documents the
+ output stays the plain string it always was."""
+
+ PDF_B64 = "JVBERi0xLjQKJSBQT05H"
+ PDF_DATA_URI = "data:application/pdf;base64,JVBERi0xLjQKJSBQT05H"
+ PDF_URL = "https://example.com/report.pdf"
+ PNG_B64 = "iVBORw0KGgoAAAANSUhEUg=="
+
+ def _messages(self, tool_result_content):
+ return [
+ {"role": "user", "content": "read the pdf"},
+ {
+ "role": "assistant",
+ "content": [{"type": "tool_use", "id": "toolu_01", "name": "read", "input": {}}],
+ },
+ {
+ "role": "user",
+ "content": [
+ {"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}
+ ],
+ },
+ ]
+
+ def _translate(self, tool_result_content):
+ return _ADAPTER.translate_messages_to_responses_input(self._messages(tool_result_content))
+
+ @staticmethod
+ def _tool_output(items):
+ return next(item for item in items if item.get("type") == "function_call_output")["output"]
+
+ def _base64_document(self, **extra):
+ return {
+ "type": "document",
+ "source": {"type": "base64", "media_type": "application/pdf", "data": self.PDF_B64},
+ **extra,
+ }
+
+ def test_text_and_base64_document_produce_part_list(self):
+ output = self._tool_output(
+ self._translate([{"type": "text", "text": "PDF file read: mystery.pdf"}, self._base64_document()])
+ )
+ assert output == [
+ {"type": "input_text", "text": "PDF file read: mystery.pdf"},
+ {"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI},
+ ]
+
+ def test_document_only_produces_single_file_part(self):
+ output = self._tool_output(self._translate([self._base64_document()]))
+ assert output == [{"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}]
+
+ def test_document_title_becomes_filename(self):
+ output = self._tool_output(self._translate([self._base64_document(title="quarterly-report.pdf")]))
+ assert output == [
+ {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI}
+ ]
+
+ def test_url_document_becomes_file_url_part(self):
+ output = self._tool_output(
+ self._translate([{"type": "document", "source": {"type": "url", "url": self.PDF_URL}}])
+ )
+ assert output == [{"type": "input_file", "file_url": self.PDF_URL}]
+
+ def test_document_with_empty_data_falls_back_to_string_output(self):
+ output = self._tool_output(
+ self._translate(
+ [
+ {"type": "text", "text": "PDF file read"},
+ {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": ""}},
+ ]
+ )
+ )
+ assert output == "PDF file read"
+
+ def test_document_without_source_dict_keeps_string_output(self):
+ output = self._tool_output(
+ self._translate([{"type": "text", "text": "stub"}, {"type": "document", "source": self.PDF_URL}])
+ )
+ assert output == "stub"
+
+ def test_text_only_tool_result_keeps_plain_string_output(self):
+ output = self._tool_output(self._translate([{"type": "text", "text": "plain result"}]))
+ assert output == "plain result"
+
+ def test_text_image_and_document_mix(self):
+ items = self._translate(
+ [
+ {"type": "text", "text": "captured"},
+ {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": self.PNG_B64}},
+ self._base64_document(),
+ ]
+ )
+
+ output = self._tool_output(items)
+ assert output == [
+ {"type": "input_text", "text": f"captured\n{TOOL_RESULT_IMAGE_PLACEHOLDER}"},
+ {"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI},
+ ]
+
+ image_message = next(
+ item
+ for item in items
+ if item.get("type") == "message"
+ and any(part.get("type") == "input_image" for part in item.get("content", []))
+ )
+ assert image_message["content"] == [
+ {"type": "input_text", "text": TOOL_RESULT_IMAGE_BOUNDARY},
+ {"type": "input_image", "image_url": f"data:image/png;base64,{self.PNG_B64}"},
+ ]
+
+
def _contains_key(value, key) -> bool:
if isinstance(value, dict):
return key in value or any(_contains_key(v, key) for v in value.values())
From 6d0cc1423ecb0e5ecc66da4553ab0e1c2a1c02f6 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
<158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 25 Aug 2026 13:37:45 -0700
Subject: [PATCH 58/70] fix(scim): return user_id as Group members[].value on
transformed group responses (#38161)
* fix(scim): return user_id as Group members[].value on transformed group responses
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style: format scim transformation tests per ruff
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test: drop redundant assertion comment
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* chore: retrigger ci
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: yassin
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../scim/scim_transformations.py | 11 +-
.../scim/test_scim_transformations.py | 117 +++++-------------
2 files changed, 32 insertions(+), 96 deletions(-)
diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py
index 2d95d0bea29..496be05b4b0 100644
--- a/litellm/proxy/management_endpoints/scim/scim_transformations.py
+++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py
@@ -198,15 +198,8 @@ class ScimTransformations:
@staticmethod
def _get_scim_member_value(member: Member) -> str:
- """
- Get the SCIM member value. Use user_email if available, otherwise use user_id.
- SCIM member value should be the unique identifier for the user.
- """
- if hasattr(member, "user_email") and member.user_email:
- return member.user_email
- elif hasattr(member, "user_id"):
- return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
- return ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
+ """The member's SCIM resource id, which LiteLLM serves as user_id (RFC 7643 §8.7.1)."""
+ return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
@staticmethod
def _get_scim_member_display(member: Member) -> str:
diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py
index 9853ce7e1cf..135175dd29d 100644
--- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py
+++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py
@@ -2,7 +2,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
-
from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, Member
from litellm.proxy.management_endpoints.scim.scim_transformations import (
ScimTransformations,
@@ -95,25 +94,17 @@ def mock_prisma_client():
class TestScimTransformations:
@pytest.mark.asyncio
- async def test_transform_litellm_user_to_scim_user(
- self, mock_user, mock_prisma_client
- ):
+ async def test_transform_litellm_user_to_scim_user(self, mock_user, mock_prisma_client):
mock_client, mock_find_unique = mock_prisma_client
# Mock the team lookup
- team1 = LiteLLM_TeamTable(
- team_id="team-1", team_alias="Team One", members_with_roles=[]
- )
- team2 = LiteLLM_TeamTable(
- team_id="team-2", team_alias="Team Two", members_with_roles=[]
- )
+ team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[])
+ team2 = LiteLLM_TeamTable(team_id="team-2", team_alias="Team Two", members_with_roles=[])
mock_find_unique.side_effect = [team1, team2]
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
- mock_user
- )
+ scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user)
assert scim_user.id == mock_user.user_id
assert scim_user.userName == mock_user.user_email
@@ -129,21 +120,15 @@ class TestScimTransformations:
assert scim_user.groups[1].display == "Team Two"
@pytest.mark.asyncio
- async def test_transform_user_with_scim_metadata(
- self, mock_user_with_scim_metadata, mock_prisma_client
- ):
+ async def test_transform_user_with_scim_metadata(self, mock_user_with_scim_metadata, mock_prisma_client):
mock_client, mock_find_unique = mock_prisma_client
# Mock the team lookup
- team1 = LiteLLM_TeamTable(
- team_id="team-1", team_alias="Team One", members_with_roles=[]
- )
+ team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[])
mock_find_unique.return_value = team1
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
- mock_user_with_scim_metadata
- )
+ scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user_with_scim_metadata)
assert scim_user.name.givenName == "Test"
assert scim_user.name.familyName == "User"
@@ -160,15 +145,11 @@ class TestScimTransformations:
teams=[],
created_at=None,
updated_at=None,
- metadata={
- "scim_enterprise": {"costCenter": "CC-42", "department": "Platform"}
- },
+ metadata={"scim_enterprise": {"costCenter": "CC-42", "department": "Platform"}},
)
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
- user
- )
+ scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user)
assert scim_user.enterprise_user is not None
assert scim_user.enterprise_user.costCenter == "CC-42"
@@ -176,9 +157,7 @@ class TestScimTransformations:
assert SCIM_ENTERPRISE_USER_SCHEMA in scim_user.schemas
@pytest.mark.asyncio
- async def test_transform_user_with_entitlements_and_roles_metadata(
- self, mock_prisma_client
- ):
+ async def test_transform_user_with_entitlements_and_roles_metadata(self, mock_prisma_client):
mock_client, mock_find_unique = mock_prisma_client
mock_find_unique.return_value = None
@@ -190,17 +169,13 @@ class TestScimTransformations:
created_at=None,
updated_at=None,
metadata={
- "scim_entitlements": [
- {"value": "jira-software", "display": "Jira Software"}
- ],
+ "scim_entitlements": [{"value": "jira-software", "display": "Jira Software"}],
"scim_roles": [{"value": "engineering-admin", "primary": True}],
},
)
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
- user
- )
+ scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user)
assert scim_user.entitlements is not None
assert scim_user.entitlements[0].value == "jira-software"
@@ -210,9 +185,7 @@ class TestScimTransformations:
assert scim_user.roles[0].primary is True
@pytest.mark.asyncio
- async def test_transform_user_with_malformed_directory_metadata_fails_soft(
- self, mock_prisma_client
- ):
+ async def test_transform_user_with_malformed_directory_metadata_fails_soft(self, mock_prisma_client):
"""Metadata is writable outside the SCIM surface; a corrupted value on one
user must omit the attribute, not fail the whole directory response"""
mock_client, mock_find_unique = mock_prisma_client
@@ -233,9 +206,7 @@ class TestScimTransformations:
)
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
- user
- )
+ scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user)
assert scim_user.id == "user-corrupt"
assert scim_user.entitlements is None
@@ -244,22 +215,14 @@ class TestScimTransformations:
assert SCIM_ENTERPRISE_USER_SCHEMA not in scim_user.schemas
@pytest.mark.asyncio
- async def test_transform_user_without_enterprise_metadata_omits_schema(
- self, mock_user, mock_prisma_client
- ):
+ async def test_transform_user_without_enterprise_metadata_omits_schema(self, mock_user, mock_prisma_client):
mock_client, mock_find_unique = mock_prisma_client
- team1 = LiteLLM_TeamTable(
- team_id="team-1", team_alias="Team One", members_with_roles=[]
- )
- team2 = LiteLLM_TeamTable(
- team_id="team-2", team_alias="Team Two", members_with_roles=[]
- )
+ team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[])
+ team2 = LiteLLM_TeamTable(team_id="team-2", team_alias="Team Two", members_with_roles=[])
mock_find_unique.side_effect = [team1, team2]
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
- mock_user
- )
+ scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user)
assert scim_user.enterprise_user is None
assert SCIM_ENTERPRISE_USER_SCHEMA not in scim_user.schemas
@@ -309,36 +272,28 @@ class TestScimTransformations:
assert dumped_attrs["roles"][0]["value"] == "engineering-admin"
@pytest.mark.asyncio
- async def test_transform_litellm_team_to_scim_group(
- self, mock_team, mock_prisma_client
- ):
+ async def test_transform_litellm_team_to_scim_group(self, mock_team, mock_prisma_client):
mock_client, _ = mock_prisma_client
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(
- mock_team
- )
+ scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(mock_team)
assert scim_group.id == mock_team.team_id
assert scim_group.displayName == mock_team.team_alias
assert len(scim_group.members) == 2
- assert scim_group.members[0].value == "test@example.com"
+ assert scim_group.members[0].value == "user-123"
assert scim_group.members[0].display == "test@example.com"
- assert scim_group.members[1].value == "test2@example.com"
+ assert scim_group.members[1].value == "user-456"
assert scim_group.members[1].display == "test2@example.com"
@pytest.mark.asyncio
- async def test_transform_team_marks_members_as_users(
- self, mock_team, mock_prisma_client
- ):
+ async def test_transform_team_marks_members_as_users(self, mock_team, mock_prisma_client):
"""A LiteLLM team only holds users, and stating the member type keeps the
response from emitting a null ``type`` now that SCIMMember carries one."""
mock_client, _ = mock_prisma_client
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(
- mock_team
- )
+ scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(mock_team)
assert [member.type for member in scim_group.members] == ["User", "User"]
@@ -351,9 +306,7 @@ class TestScimTransformations:
result = ScimTransformations._get_scim_user_name(mock_user_minimal)
assert result == ScimTransformations.DEFAULT_SCIM_DISPLAY_NAME
- def test_get_scim_family_name(
- self, mock_user, mock_user_with_scim_metadata, mock_user_minimal
- ):
+ def test_get_scim_family_name(self, mock_user, mock_user_with_scim_metadata, mock_user_minimal):
# User with alias
result = ScimTransformations._get_scim_family_name(mock_user)
assert result == mock_user.user_alias
@@ -366,9 +319,7 @@ class TestScimTransformations:
result = ScimTransformations._get_scim_family_name(mock_user_minimal)
assert result == ScimTransformations.DEFAULT_SCIM_FAMILY_NAME
- def test_get_scim_given_name(
- self, mock_user, mock_user_with_scim_metadata, mock_user_minimal
- ):
+ def test_get_scim_given_name(self, mock_user, mock_user_with_scim_metadata, mock_user_minimal):
# User with alias
result = ScimTransformations._get_scim_given_name(mock_user)
assert result == mock_user.user_alias
@@ -382,14 +333,10 @@ class TestScimTransformations:
assert result == ScimTransformations.DEFAULT_SCIM_NAME
def test_get_scim_member_value(self):
- # Member with email
- member_with_email = Member(
- user_id="user-123", user_email="test@example.com", role="admin"
- )
+ member_with_email = Member(user_id="user-123", user_email="test@example.com", role="admin")
result = ScimTransformations._get_scim_member_value(member_with_email)
- assert result == member_with_email.user_email
+ assert result == member_with_email.user_id
- # Member without email should fall back to user_id
member_without_email = Member(user_id="user-456", user_email=None, role="user")
result = ScimTransformations._get_scim_member_value(member_without_email)
assert result == member_without_email.user_id
@@ -415,9 +362,7 @@ class TestScimTransformations:
mock_find_unique.return_value = None
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
- user_with_uuid_email
- )
+ scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user_with_uuid_email)
assert scim_user.id == user_with_uuid_email.user_id
assert scim_user.emails is None or len(scim_user.emails) == 0
@@ -443,9 +388,7 @@ class TestScimTransformations:
mock_find_unique.return_value = None
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
- scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
- user_with_none_email
- )
+ scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user_with_none_email)
assert scim_user.id == user_with_none_email.user_id
assert scim_user.emails is None or len(scim_user.emails) == 0
From c96245d7bfc6c9446fa04e1ab31e824843f4259b Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
<158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 25 Aug 2026 13:39:34 -0700
Subject: [PATCH 59/70] fix(scim): preserve existing team memberships when POST
/Users adoption carries no groups (#38166)
Co-authored-by: yassin
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../management_endpoints/scim/scim_v2.py | 7 +-
.../scim/test_scim_v2_endpoints.py | 73 ++++++++++++++++++-
2 files changed, 76 insertions(+), 4 deletions(-)
diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py
index 7183e6cb402..6658963d024 100644
--- a/litellm/proxy/management_endpoints/scim/scim_v2.py
+++ b/litellm/proxy/management_endpoints/scim/scim_v2.py
@@ -176,6 +176,10 @@ class UserProvisionerHelpers:
is persisted too, so re-upserting an existing email demotes a user who is no
longer in the admin group instead of leaving the stale role.
+ IdPs like Entra manage membership exclusively through /Groups and never send
+ ``groups`` on POST /Users, so a request without teams means "unspecified",
+ not "remove from every team": existing memberships are preserved then.
+
Args:
prisma_client: Database client
new_user_request: New user request data
@@ -194,7 +198,8 @@ class UserProvisionerHelpers:
if not existing_user:
return None
- new_teams: Final = list(dict.fromkeys(new_user_request.teams or []))
+ requested_teams: Final = list(dict.fromkeys(new_user_request.teams or []))
+ new_teams: Final = requested_teams if requested_teams else list(existing_user.teams or [])
if new_user_request.user_id != existing_user.user_id:
verbose_proxy_logger.info(
diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
index 5f6c1a2375b..d51e3abde83 100644
--- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
@@ -752,6 +752,63 @@ async def test_handle_existing_user_by_email_syncs_roster_and_dedups_teams(mocke
assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"]
+@pytest.mark.asyncio
+async def test_handle_existing_user_by_email_without_teams_preserves_memberships(mocker):
+ """Adoption via POST /Users without ``groups`` must keep the user's existing teams.
+
+ Regression: Entra manages membership exclusively through /Groups and never sends
+ ``groups`` on POST /Users, so the empty team list was treated as the desired
+ state and the adopted user was removed from every team roster and had ``teams``
+ overwritten with [].
+ """
+ existing_user = mocker.MagicMock()
+ existing_user.user_id = "adopted-id"
+ existing_user.user_email = "member@example.com"
+ existing_user.user_alias = "Member"
+ existing_user.teams = ["team-a", "team-b"]
+ existing_user.metadata = {}
+
+ mock_prisma_client = mocker.MagicMock()
+ mock_prisma_client.db = mocker.MagicMock()
+ mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
+ mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user)
+ mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={})
+
+ mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
+ "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
+ AsyncMock(return_value=None),
+ )
+ mock_team_member_add = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
+ "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add",
+ AsyncMock(),
+ )
+ mock_team_member_delete = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
+ "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete",
+ AsyncMock(),
+ )
+
+ new_user_request = NewUserRequest(
+ user_id="entra-object-id",
+ user_email="member@example.com",
+ user_alias="Member",
+ teams=[],
+ metadata={},
+ auto_create_key=False,
+ )
+
+ await UserProvisionerHelpers.handle_existing_user_by_email(
+ prisma_client=mock_prisma_client, new_user_request=new_user_request
+ )
+
+ mock_team_member_add.assert_not_awaited()
+ mock_team_member_delete.assert_not_awaited()
+
+ update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list
+ assert len(update_calls) == 1
+ assert update_calls[0].kwargs["where"] == {"user_id": "adopted-id"}
+ assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"]
+
+
@pytest.mark.asyncio
async def test_handle_existing_user_by_email_roster_add_failure_blocks_teams_write(mocker):
"""A genuine roster add failure must propagate and must not persist the teams array.
@@ -872,11 +929,16 @@ async def test_handle_existing_user_by_email_roster_remove_failure_blocks_teams_
AsyncMock(side_effect=HTTPException(status_code=500, detail={"error": "No db connected"})),
)
+ mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
+ "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add",
+ AsyncMock(),
+ )
+
new_user_request = NewUserRequest(
user_id="uid",
user_email="member@example.com",
user_alias="Member",
- teams=[],
+ teams=["replacement-team"],
metadata={},
auto_create_key=False,
)
@@ -917,11 +979,16 @@ async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noo
AsyncMock(return_value=None),
)
+ mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper
+ "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add",
+ AsyncMock(),
+ )
+
new_user_request = NewUserRequest(
user_id="uid",
user_email="member@example.com",
user_alias="Member",
- teams=[],
+ teams=["replacement-team"],
metadata={},
auto_create_key=False,
)
@@ -933,7 +1000,7 @@ async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noo
mock_team_member_delete.assert_awaited_once()
update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list
assert len(update_calls) == 1
- assert update_calls[0].kwargs["data"]["teams"] == []
+ assert update_calls[0].kwargs["data"]["teams"] == ["replacement-team"]
@pytest.mark.asyncio
From c70b911122fd8d50dde4a6af525449afeff87aac Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Tue, 25 Aug 2026 13:40:46 -0700
Subject: [PATCH 60/70] fix(router): support mid-stream fallback for
anthropic_messages route type (#38153)
anthropic_messages goes through _ageneric_api_call_with_fallbacks rather
than _acompletion, so its returned streaming iterator was never wrapped
by the chat-completions fallback handler. A retriable SSE event: error
frame (overloaded_error, internal_server_error) from a native
Anthropic/Bedrock passthrough passed through to the client unchanged,
and a MidStreamFallbackError raised by the completion-bridge path's
CustomStreamWrapper propagated unhandled.
Add _aanthropic_messages_streaming_iterator, mirroring
_acompletion_streaming_iterator: it detects a retriable SSE error event
via the new parse_anthropic_error_event helper, raises
MidStreamFallbackError once real generated content (a content_block_delta
frame) has not yet reached the caller, and re-enters the Router's
fallback chain. A MidStreamFallbackError raised directly by the source
iterator (the completion-bridge path) is gated the same way via its own
is_pre_first_chunk flag. The raised MidStreamFallbackError carries a
status-coded original_exception built from the parsed error type, so
status_code/cooldown logic sees the real 429/500/503/etc. instead of a
hardcoded 503.
Lifecycle/bookkeeping frames (message_start, content_block_start, ping,
...) never disqualify a fallback attempt by themselves, since Anthropic
routinely sends message_start before an overload error - but they are
buffered rather than forwarded immediately, since forwarding one and
then appending a fallback attempt's own message_start would produce two
overlapping message lifecycles on one SSE stream. Buffered frames flush,
in order, once real content arrives or the stream ends without error.
Once real content has streamed, or the error is a non-retriable 4xx, the
chunk (or exception) is forwarded as-is rather than starting a second
lifecycle. Content and error coalesced into a single physical read are
handled the same way: once the client has genuinely received the content
(bundled in that same forwarded chunk), no fallback is attempted. A
`ping` keepalive is dropped outright before any real content arrives
(it recurs indefinitely on a slow-starting connection and carries
nothing worth buffering), and the pre-content lifecycle buffer is capped
at MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, forcing an early commit to
the primary stream so a hostile or pathological upstream can't grow it
without bound. is_anthropic_ping_chunk only matches a chunk whose every
event: line is event: ping, so a ping coalesced with real content or a
retriable error into one physical transport chunk is never dropped.
The fallback request kwargs also deep-copy nested litellm_metadata/metadata
(matching the Responses API path) so the primary attempt's
deployment-specific fields never leak into the fallback request, and the
fallback deployment's own provider headers are merged onto the wrapper's
_hidden_params so they still reach the client/logging pipeline. A
fallback that resolves to a non-streaming response (e.g. an agentic
tool-use interception loop) is synthesized into a real Anthropic SSE
event sequence via the new anthropic_messages_response_as_sse_events
helper, instead of yielding a raw dict into the byte stream - including
a trailing signature_delta for a thinking block, and a message_start
whose stop_reason/stop_sequence/output_tokens stay null/zero the way a
real stream's does instead of leaking the completed response's final
state.
Resolves #24004
---
litellm/llms/anthropic/common_utils.py | 15 +
.../messages/streaming_iterator.py | 235 +++-
litellm/llms/anthropic/files/handler.py | 14 +-
litellm/router.py | 363 +++++-
.../messages/test_streaming_iterator.py | 208 ++++
tests/test_litellm/test_router.py | 1101 +++++++++++++++++
6 files changed, 1911 insertions(+), 25 deletions(-)
diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py
index 9461e40cf2e..53cc3464761 100644
--- a/litellm/llms/anthropic/common_utils.py
+++ b/litellm/llms/anthropic/common_utils.py
@@ -38,6 +38,21 @@ DROP_DISABLED_THINKING_WARNING: Final = (
"thinking blocks, and those thinking tokens are billed as output tokens."
)
+# Anthropic error `type` (both the JSON error body and SSE `event: error`
+# payloads use this field) mapped to the HTTP status code it corresponds to.
+ANTHROPIC_ERROR_STATUS_CODE_MAP: Final = MappingProxyType(
+ {
+ "invalid_request_error": 400,
+ "authentication_error": 401,
+ "permission_error": 403,
+ "not_found_error": 404,
+ "rate_limit_error": 429,
+ "api_error": 500,
+ "overloaded_error": 503,
+ "timeout_error": 504,
+ }
+)
+
_BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$")
_INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
_DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
index 922769dbbfd..0a12bc3135f 100644
--- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
+++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
@@ -1,6 +1,6 @@
import asyncio
import json
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Mapping
from datetime import datetime
from typing import Any, Final, Protocol, runtime_checkable
@@ -11,9 +11,11 @@ from typing_extensions import TypedDict
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
)
+from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
@@ -33,26 +35,239 @@ def _is_message_stop_chunk(chunk: object) -> bool:
return False
-def _is_provider_error_chunk(chunk: object) -> bool:
+def is_anthropic_ping_chunk(chunk: object) -> bool:
+ """
+ Whether a chunk is a pure ``ping`` keepalive frame. It carries no content
+ and can recur indefinitely on a slow-starting or idle connection, so a
+ mid-stream fallback wrapper drops it outright while still deciding
+ whether to commit to the primary stream, rather than buffering it.
+
+ A physical transport chunk that coalesces a ping with any other SSE
+ event (``message_start``, ``content_block_delta``, ``event: error``, ...)
+ is NOT a pure ping - dropping it whole would discard those events - so
+ only a chunk whose every ``event:`` line is ``event: ping`` qualifies.
+ """
if isinstance(chunk, dict):
- return chunk.get("type") == "error"
+ return chunk.get("type") == "ping"
if isinstance(chunk, (bytes, bytearray)):
- return any(line == b"event: error" for line in chunk.splitlines())
+ event_lines: Final = tuple(line for line in chunk.splitlines() if line.startswith(b"event:"))
+ return bool(event_lines) and all(line == b"event: ping" for line in event_lines)
return False
+def is_anthropic_content_delta_chunk(chunk: object) -> bool:
+ """
+ Whether a chunk carries actual assistant-generated output (a
+ ``content_block_delta`` frame), as opposed to a lifecycle/bookkeeping
+ frame (``message_start``, ``content_block_start``/``stop``,
+ ``message_delta``, ``message_stop``, ``ping``) that carries nothing
+ worth preserving before an invisible mid-stream fallback retry.
+ """
+ if isinstance(chunk, dict):
+ return chunk.get("type") == "content_block_delta"
+ if isinstance(chunk, (bytes, bytearray)):
+ return any(line == b"event: content_block_delta" for line in chunk.splitlines())
+ return False
+
+
+def _decoded_sse_data_line(line: bytes) -> object | None:
+ if not line.startswith(b"data:"):
+ return None
+ try:
+ return json.loads(line[len(b"data:") :].strip())
+ except (ValueError, TypeError):
+ return None
+
+
+def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None:
+ if isinstance(chunk, dict):
+ return chunk if chunk.get("type") == "error" else None
+ if isinstance(chunk, (bytes, bytearray)):
+ decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines())
+ return next(
+ (
+ candidate
+ for candidate in decoded_lines
+ if isinstance(candidate, dict) and candidate.get("type") == "error"
+ ),
+ None,
+ )
+ return None
+
+
+def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None:
+ """Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None."""
+ payload: Final = _anthropic_error_event_payload(chunk)
+ error_body: Final = payload.get("error") if payload is not None else None
+ return error_body if isinstance(error_body, dict) else None
+
+
+def _is_provider_error_chunk(chunk: object) -> bool:
+ return _anthropic_error_body(chunk) is not None
+
+
+def parse_anthropic_error_event(chunk: object) -> tuple[str, str, int] | None:
+ """
+ Extract ``(error_type, message, http_status_code)`` from an Anthropic SSE
+ ``event: error`` chunk (raw bytes or an already-decoded dict), or None if
+ ``chunk`` is not an error event.
+
+ The status code is looked up via ANTHROPIC_ERROR_STATUS_CODE_MAP,
+ defaulting to 500 for an error ``type`` Anthropic hasn't documented yet.
+ """
+ error_body: Final = _anthropic_error_body(chunk)
+ if error_body is None:
+ return None
+ error_type: Final = error_body.get("type")
+ if not isinstance(error_type, str):
+ return None
+ message: Final = error_body.get("message")
+ return (
+ error_type,
+ message if isinstance(message, str) else error_type,
+ ANTHROPIC_ERROR_STATUS_CODE_MAP.get(error_type, 500),
+ )
+
+
def _is_terminal_stream_chunk(chunk: object) -> bool:
return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk)
+def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes:
+ return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
+
+
def _incomplete_stream_error_sse_event() -> bytes:
- payload: Final = json.dumps(
- {
- "type": "error",
- "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE},
- }
+ return _sse_event( # mutable-ok: one-shot JSON payload, never mutated after construction
+ "error",
+ {"type": "error", "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}},
+ )
+
+
+def _anthropic_content_block_start_and_deltas(
+ block: Mapping[str, object],
+) -> tuple[Mapping[str, object], tuple[Mapping[str, object], ...]]:
+ """
+ ``(content_block_start.content_block, content_block_delta.delta events)``
+ for one Anthropic response content block. A thinking block emits both a
+ thinking_delta and a trailing signature_delta - a real Anthropic stream
+ does the same, and dropping the signature makes any replay of that
+ assistant message (a follow-up turn, a tool-use continuation) fail
+ Anthropic's thinking-signature verification. redacted_thinking has no
+ delta at all - it is sent complete in content_block_start.
+ """
+ match block.get("type"):
+ case "tool_use":
+ return (
+ { # mutable-ok: one-shot payload
+ "id": block.get("id"),
+ "name": block.get("name"),
+ "input": {}, # mutable-ok: one-shot payload
+ "type": "tool_use",
+ },
+ (
+ { # mutable-ok: one-shot payload
+ "partial_json": json.dumps(block.get("input") or {}), # mutable-ok: one-shot payload
+ "type": "input_json_delta",
+ },
+ ),
+ )
+ case "thinking":
+ signature: Final = block.get("signature")
+ signature_deltas: Final = (
+ ({"signature": signature, "type": "signature_delta"},) # mutable-ok: one-shot payload
+ if isinstance(signature, str) and signature
+ else ()
+ )
+ return (
+ {"thinking": "", "signature": "", "type": "thinking"}, # mutable-ok: one-shot payload
+ (
+ {"thinking": block.get("thinking") or "", "type": "thinking_delta"}, # mutable-ok: one-shot payload
+ *signature_deltas,
+ ),
+ )
+ case "redacted_thinking":
+ return ({"type": "redacted_thinking", "data": block.get("data")}, ()) # mutable-ok: one-shot JSON payload
+ case _:
+ return (
+ {"type": "text", "text": ""}, # mutable-ok: one-shot JSON payload
+ ({"type": "text_delta", "text": block.get("text") or ""},), # mutable-ok: one-shot JSON payload
+ )
+
+
+def anthropic_messages_response_as_sse_events(response: AnthropicMessagesResponse) -> tuple[bytes, ...]:
+ """
+ Render a complete (non-streaming) AnthropicMessagesResponse as the SSE
+ event sequence a real streaming request would have produced.
+
+ A mid-stream fallback can resolve to a non-streaming response even
+ though the client asked to stream (e.g. an agentic tool-use loop that
+ intercepts and returns a complete message) - yielding that dict directly
+ into a `/v1/messages` SSE byte stream would produce a malformed
+ response, so it's synthesized into the message_start/content_block_*/
+ message_delta/message_stop lifecycle a real stream would have sent.
+ """
+ content_blocks: Final = response.get("content") or ()
+ content_events: Final = (
+ event for index, block in enumerate(content_blocks) for event in _anthropic_content_block_events(index, block)
+ )
+ # A real message_start always carries a null stop_reason/stop_sequence and
+ # a zero output_tokens - those are only known once generation finishes, so
+ # copying the completed response's final values here would let a client
+ # treat the message as already finished, or double-count output tokens.
+ message_start_usage: Final = { # mutable-ok: one-shot JSON payload
+ **(response.get("usage") or {}),
+ "output_tokens": 0,
+ }
+ message_start_payload: Final = { # mutable-ok: one-shot JSON payload, never mutated after construction
+ "type": "message_start",
+ "message": { # mutable-ok: one-shot JSON payload
+ **response,
+ "content": [], # mutable-ok: one-shot JSON payload
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": message_start_usage,
+ },
+ }
+ message_delta_payload: Final = { # mutable-ok: one-shot JSON payload, never mutated after construction
+ "type": "message_delta",
+ "delta": { # mutable-ok: one-shot JSON payload
+ "stop_reason": response.get("stop_reason"),
+ "stop_sequence": response.get("stop_sequence"),
+ },
+ "usage": response.get("usage") or {}, # mutable-ok: one-shot JSON payload
+ }
+ return (
+ _sse_event("message_start", message_start_payload),
+ *content_events,
+ _sse_event("message_delta", message_delta_payload),
+ _sse_event("message_stop", {"type": "message_stop"}), # mutable-ok: one-shot JSON payload
+ )
+
+
+def _anthropic_content_block_events(index: int, block: Mapping[str, object]) -> tuple[bytes, ...]:
+ start_block, deltas = _anthropic_content_block_start_and_deltas(block)
+ start_payload: Final = { # mutable-ok: one-shot payload
+ "type": "content_block_start",
+ "index": index,
+ "content_block": start_block,
+ }
+ stop_payload: Final = { # mutable-ok: one-shot payload
+ "type": "content_block_stop",
+ "index": index,
+ }
+ delta_events: Final = tuple(
+ _sse_event(
+ "content_block_delta",
+ {"type": "content_block_delta", "index": index, "delta": delta}, # mutable-ok: one-shot payload
+ )
+ for delta in deltas
+ )
+ return (
+ _sse_event("content_block_start", start_payload),
+ *delta_events,
+ _sse_event("content_block_stop", stop_payload),
)
- return f"event: error\ndata: {payload}\n\n".encode()
class AnthropicMessagesStreamHiddenParams(TypedDict):
diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py
index 0c62418708f..5fdf2ceff7f 100644
--- a/litellm/llms/anthropic/files/handler.py
+++ b/litellm/llms/anthropic/files/handler.py
@@ -22,19 +22,7 @@ from litellm.types.llms.openai import (
from litellm.types.utils import CallTypes, LlmProviders, ModelResponse
from ..chat.transformation import AnthropicConfig
-from ..common_utils import AnthropicModelInfo
-
-# Map Anthropic error types to HTTP status codes
-ANTHROPIC_ERROR_STATUS_CODE_MAP: Final = {
- "invalid_request_error": 400,
- "authentication_error": 401,
- "permission_error": 403,
- "not_found_error": 404,
- "rate_limit_error": 429,
- "api_error": 500,
- "overloaded_error": 503,
- "timeout_error": 504,
-}
+from ..common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP, AnthropicModelInfo
class AnthropicFilesHandler:
diff --git a/litellm/router.py b/litellm/router.py
index c658bc441ba..d07effd0d90 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -8,6 +8,7 @@
# Thank you ! We ❤️ you! - Krrish & Ishaan
import asyncio
+import contextlib
import copy
import enum
import hashlib
@@ -20,7 +21,7 @@ import time
import traceback
import weakref
from collections import defaultdict
-from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence
+from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
@@ -248,6 +249,7 @@ from .router_utils.pattern_match_deployments import PatternMatchRouter
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
+ from litellm.exceptions import MidStreamFallbackError
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
)
@@ -264,6 +266,9 @@ if TYPE_CHECKING:
from litellm.router_strategy.quality_router.quality_router import (
QualityRouter,
)
+ from litellm.types.llms.anthropic_messages.anthropic_response import (
+ AnthropicMessagesResponse,
+ )
from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
from litellm.types.llms.openai import (
ResponseAPIUsage,
@@ -361,6 +366,101 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
return False
+# Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks
+# until real content commits the primary stream; a hostile or slow-starting
+# upstream that never emits content or an error could otherwise grow that
+# buffer without bound, so hitting this cap forces an early commit instead.
+MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200
+
+
+def _anthropic_stream_should_drop_pre_content_ping(chunk: object, has_generated_content: bool) -> bool:
+ """A `ping` keepalive seen before any real content is dropped outright - it recurs indefinitely on a
+ slow-starting connection and carries nothing worth buffering toward a possible fallback."""
+ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import is_anthropic_ping_chunk
+
+ if has_generated_content:
+ return False
+ return is_anthropic_ping_chunk(chunk)
+
+
+def _is_retriable_anthropic_status(status_code: int) -> bool:
+ return status_code == 429 or status_code >= 500
+
+
+def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error: "MidStreamFallbackError") -> bool:
+ """
+ A MidStreamFallbackError raised directly by the source iterator (the
+ completion-bridge path's CustomStreamWrapper, e.g. on a transport drop)
+ carries its own pre_first_chunk bookkeeping - gated the same way a
+ detected SSE error event is, so a fallback is never appended after real
+ content already reached the client on either path.
+ """
+ return has_generated_content or not error.is_pre_first_chunk
+
+
+def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool:
+ """
+ Whether `chunk` should make Router._aanthropic_messages_streaming_iterator
+ commit to the primary Anthropic stream (real content arrived, or the
+ pre-content buffer cap was hit) rather than keep buffering lifecycle
+ frames toward a possible fallback.
+ """
+ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+ is_anthropic_content_delta_chunk,
+ )
+
+ if has_generated_content:
+ return False
+ return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS
+
+
+class FallbackAwareAnthropicMessagesStream:
+ """
+ Bare async generators can't carry the `_hidden_params` attribute the
+ proxy reads response headers off of (see
+ router_utils.add_retry_fallback_headers.get_hidden_params_dict), so this
+ thin wrapper carries it through from the source iterator - mirrors
+ AnthropicMessagesStreamingResponse. Used by
+ Router._aanthropic_messages_streaming_iterator.
+ """
+
+ def __init__(self, async_generator: AsyncGenerator[bytes, None], source_iterator: object) -> None:
+ self._async_generator = async_generator
+ self._hidden_params = dict( # mutable-ok: mutated in place by merge_fallback_hidden_params
+ getattr(source_iterator, "_hidden_params", None) or {}
+ )
+
+ def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream":
+ return self
+
+ async def __anext__(self) -> bytes:
+ return await self._async_generator.__anext__()
+
+ async def aclose(self) -> None:
+ await self._async_generator.aclose()
+
+ def merge_fallback_hidden_params(
+ self,
+ fallback_hidden_params: Mapping[str, object],
+ fallback_headers: Mapping[str, object],
+ ) -> None:
+ """
+ Raw bytes can't carry their own _hidden_params the way a
+ ModelResponseStream/ResponsesAPI event can, so a mid-stream
+ fallback's provider headers (e.g. Bedrock's x-amzn-requestid) are
+ merged onto the wrapper itself instead - mirrors
+ Router._apply_fallback_hidden_params_to_item's merge shape.
+ """
+ existing_headers: Final = cast( # cast-ok: additional_headers is always a dict[str, object] when present
+ "dict[str, object]", self._hidden_params.get("additional_headers") or {}
+ )
+ self._hidden_params = { # mutable-ok: matches _hidden_params' existing dict[str, object] shape
+ **self._hidden_params,
+ **fallback_hidden_params,
+ "additional_headers": {**existing_headers, **fallback_headers}, # mutable-ok: same shape
+ }
+
+
class RoutingArgs(enum.Enum):
ttl = 60 # 1min (RPM/TPM expire key)
@@ -4806,6 +4906,264 @@ class Router:
)
return response
+ async def _aanthropic_messages_streaming_iterator(
+ self,
+ response: AsyncIterator[bytes],
+ initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain
+ ) -> AsyncIterator[bytes]:
+ """
+ Wrap an anthropic_messages (/v1/messages) streaming response so a
+ mid-stream provider error triggers the Router's fallback chain
+ (parity with _acompletion_streaming_iterator for the
+ chat-completions path). See #24004.
+
+ anthropic_messages goes through _ageneric_api_call_with_fallbacks
+ rather than _acompletion, so the returned byte iterator is never
+ wrapped by the chat-completions fallback handler. Two failure
+ shapes land here:
+ - the completion-bridge path (deployments with no native
+ /v1/messages endpoint, via
+ LiteLLMMessagesToCompletionTransformationHandler) already
+ raises MidStreamFallbackError out of its underlying
+ CustomStreamWrapper; this wrapper only needs to catch it.
+ - a native Anthropic/Bedrock passthrough never raises anything
+ for a provider SSE `event: error` frame (e.g. `overloaded_error`,
+ `internal_server_error`) - it is forwarded to the client as-is -
+ so this wrapper detects it via parse_anthropic_error_event and
+ raises MidStreamFallbackError itself.
+
+ Only an error before any real content (a content_block_delta frame)
+ has reached the caller triggers a fallback attempt, mirroring the
+ restriction _acompletion_streaming_iterator applies: once generated
+ output has already reached the caller, retrying would start a
+ second, overlapping Anthropic message lifecycle on the same SSE
+ stream, so the error is left to propagate instead of being retried
+ invisibly. A non-retriable client error (4xx other than 429) is
+ never worth a fallback attempt either, so it is also left to
+ propagate.
+
+ Lifecycle/bookkeeping frames (message_start, content_block_start,
+ ping, ...) do not by themselves disqualify a fallback attempt -
+ Anthropic routinely sends message_start before an overload error -
+ but they are BUFFERED rather than forwarded immediately, since
+ forwarding one and then appending a fallback attempt's own
+ message_start would produce two overlapping message lifecycles on
+ one SSE stream. Buffered frames are flushed, in order, the moment
+ real content arrives (the primary attempt has committed by then
+ anyway) or once the stream ends without ever producing content or
+ an error.
+ """
+ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+ aclose_if_supported,
+ parse_anthropic_error_event,
+ )
+
+ source_iterator: Final = response
+
+ async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]:
+ from litellm.exceptions import MidStreamFallbackError
+
+ # Lifecycle/bookkeeping frames (message_start, content_block_start,
+ # ping, ...) are held back rather than forwarded immediately:
+ # Anthropic routinely sends message_start before an overload
+ # error, and once a byte reaches the client a fallback attempt
+ # can only append its OWN message_start, producing two
+ # overlapping message lifecycles on one SSE stream. Buffered
+ # frames are flushed the moment real content (content_block_delta)
+ # arrives - at that point the primary attempt has committed and a
+ # clean retry is no longer possible anyway - or once the primary
+ # stream ends without ever producing content. A `ping` keepalive
+ # is dropped outright rather than buffered, since it can recur
+ # indefinitely on a slow-starting connection and carries nothing
+ # worth preserving; hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS
+ # forces the same early commit as real content arriving, so a
+ # hostile or pathological upstream can't grow the buffer forever.
+ has_generated_content = False # rebind-ok: set once real content is seen, or the buffer cap is hit
+ buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline
+ model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group
+ try:
+ async for chunk in source_iterator:
+ if _anthropic_stream_should_drop_pre_content_ping(chunk, has_generated_content):
+ continue
+ if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)):
+ has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit
+ error_event = parse_anthropic_error_event(chunk)
+ retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over
+ not has_generated_content
+ and error_event is not None
+ and _is_retriable_anthropic_status(error_event[2])
+ )
+ if not has_generated_content and not retriable_pending_error and error_event is None:
+ buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk)
+ continue
+ if retriable_pending_error:
+ assert error_event is not None # guard-ok: retriable_pending_error implies this
+ _error_type, message, status_code = error_event
+ raise MidStreamFallbackError(
+ message=message,
+ model=model,
+ llm_provider="anthropic",
+ original_exception=litellm.exceptions.APIError(
+ status_code=status_code,
+ message=message,
+ llm_provider="anthropic",
+ model=model,
+ ),
+ is_pre_first_chunk=True,
+ )
+ for buffered_chunk in buffered_lifecycle_chunks:
+ yield buffered_chunk
+ buffered_lifecycle_chunks = ()
+ yield chunk
+ for buffered_chunk in buffered_lifecycle_chunks:
+ yield buffered_chunk
+ except MidStreamFallbackError as e:
+ if _anthropic_stream_should_decline_fallback(has_generated_content, e):
+ for buffered_chunk in buffered_lifecycle_chunks:
+ yield buffered_chunk
+ if e.original_exception is not None:
+ raise e.original_exception from e
+ raise
+ async for item in self._aanthropic_messages_fallback_attempt(e, initial_kwargs, wrapper):
+ yield item
+ finally:
+ with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
+ await aclose_if_supported(source_iterator)
+
+ # Referenced by stream_with_fallbacks via closure - assigned here, before
+ # the generator body ever runs, so the reference resolves fine despite
+ # being defined textually after the function that captures it.
+ wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator)
+ return wrapper
+
+ async def _aanthropic_messages_fallback_attempt(
+ self,
+ e: "MidStreamFallbackError",
+ initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain
+ wrapper: "FallbackAwareAnthropicMessagesStream",
+ ) -> AsyncGenerator[bytes, None]:
+ """
+ Re-enters the Router's fallback chain for a mid-stream
+ anthropic_messages error and yields whatever the fallback attempt
+ produces. Split out of _aanthropic_messages_streaming_iterator to
+ keep each function's cyclomatic complexity within the repo's C901
+ budget.
+ """
+ from litellm.exceptions import MidStreamFallbackError
+ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
+ aclose_if_supported,
+ anthropic_messages_response_as_sse_events,
+ )
+
+ fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted
+ try:
+ model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: model group
+ fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the common_utils list|None param
+ "fallbacks", self.fallbacks
+ )
+ context_window_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below
+ "context_window_fallbacks", self.context_window_fallbacks
+ )
+ content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below
+ "content_policy_fallbacks", self.content_policy_fallbacks
+ )
+ initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper
+ self._update_kwargs_before_fallbacks(
+ model=model_group,
+ kwargs=initial_kwargs,
+ metadata_variable_name="litellm_metadata",
+ )
+ fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success
+ e=e,
+ disable_fallbacks=False,
+ fallbacks=fallbacks,
+ context_window_fallbacks=context_window_fallbacks,
+ content_policy_fallbacks=content_policy_fallbacks,
+ model_group=model_group,
+ args=(),
+ kwargs=initial_kwargs,
+ include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True,
+ )
+ fallback_hidden_params, fallback_headers = Router._prepare_fallback_hidden_params(fallback_response)
+ wrapper.merge_fallback_hidden_params(fallback_hidden_params, fallback_headers)
+ if hasattr(fallback_response, "__aiter__"):
+ async for fallback_item in fallback_response:
+ yield fallback_item
+ else:
+ # A fallback can resolve to a complete AnthropicMessagesResponse
+ # dict even for a streaming request (e.g. an agentic tool-use
+ # interception loop) - yielding it as-is would put a raw dict
+ # into a byte stream, so it's synthesized into the SSE
+ # lifecycle a real stream would have sent instead.
+ for event in anthropic_messages_response_as_sse_events(
+ cast("AnthropicMessagesResponse", fallback_response) # cast-ok: non-streaming shape by elimination
+ ):
+ yield event
+ except Exception as fallback_error:
+ verbose_router_logger.error("Anthropic messages streaming fallback also failed: %s", fallback_error)
+ if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None:
+ raise fallback_error.original_exception from fallback_error
+ raise
+ finally:
+ if fallback_response is not None:
+ with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
+ await aclose_if_supported(fallback_response)
+
+ async def _aanthropic_messages_with_streaming_fallbacks(
+ self,
+ original_function: Callable,
+ **kwargs: object, # kwargs-ok: forwarded verbatim to original_function, shape varies per call site
+ ) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]:
+ """
+ _ageneric_api_call_with_fallbacks for anthropic_messages, with the
+ addition of mid-stream fallback handling (see
+ _aanthropic_messages_streaming_iterator). Parity with
+ _aresponses_with_streaming_fallbacks for the Responses API.
+ """
+ from litellm.litellm_core_utils.core_helpers import safe_deep_copy
+
+ # Snapshot the request kwargs before the primary attempt mutates them
+ # in place: _update_kwargs_with_deployment writes deployment-specific
+ # fields (deployment, model_info, api_base, tags, ...) into the
+ # SAME litellm_metadata/metadata dicts a shallow .copy() would still
+ # share, leaking primary-deployment metadata into the mid-stream
+ # fallback request. safe_deep_copy avoids deep-copying the full
+ # kwargs (which can hold non-deepcopyable logging handles/clients).
+ fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry
+ if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
+ fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
+ if isinstance(fallback_kwargs.get("metadata"), dict):
+ fallback_kwargs["metadata"] = safe_deep_copy(fallback_kwargs["metadata"])
+ fallback_kwargs["original_generic_function"] = original_function
+
+ response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
+
+ if kwargs.get("stream") and hasattr(response, "__aiter__"):
+ return await self._aanthropic_messages_streaming_iterator(
+ response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator
+ initial_kwargs=fallback_kwargs,
+ )
+ return response
+
+ async def _dispatch_generic_call_type(
+ self,
+ call_type: str,
+ original_function: Callable,
+ **kwargs: object, # kwargs-ok: forwarded verbatim to the per-call-type helper, shape varies per call site
+ ):
+ """
+ factory_function's shared dispatch for call types with no
+ call-specific handling, except anthropic_messages: kept out of
+ factory_function's own async_wrapper (already at the repo's C901
+ complexity ceiling) so routing its mid-stream fallback handling
+ (#24004) doesn't add another branch there.
+ """
+ if call_type == "anthropic_messages":
+ return await self._aanthropic_messages_with_streaming_fallbacks(
+ original_function=original_function, **kwargs
+ )
+ return await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
+
def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs):
"""
Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router
@@ -5992,7 +6350,8 @@ class Router:
"aget_skill",
"adelete_skill",
):
- return await self._ageneric_api_call_with_fallbacks(
+ return await self._dispatch_generic_call_type(
+ call_type=call_type,
original_function=original_function,
**kwargs,
)
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py
index f33bb3dda8b..652c1f077a9 100644
--- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py
@@ -11,6 +11,10 @@ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterato
BaseAnthropicMessagesStreamingIterator,
_incomplete_stream_error_sse_event,
_is_message_stop_chunk,
+ _is_provider_error_chunk,
+ anthropic_messages_response_as_sse_events,
+ is_anthropic_content_delta_chunk,
+ parse_anthropic_error_event,
)
@@ -157,6 +161,96 @@ def test_is_message_stop_chunk_ignores_substring_in_payload():
assert _is_message_stop_chunk(delta_frame_with_substring) is False
+def test_parse_anthropic_error_event_from_dict_chunk():
+ """Regression for #24004: dict-shaped error chunks parse to
+ (type, message, status) so the Router can decide whether to fall back."""
+ chunk = {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}
+ assert parse_anthropic_error_event(chunk) == ("overloaded_error", "Overloaded", 503)
+ assert _is_provider_error_chunk(chunk) is True
+
+
+def test_parse_anthropic_error_event_from_sse_bytes():
+ """Regression for #24004: a raw `event: error` SSE frame (what a native
+ Anthropic/Bedrock passthrough forwards verbatim today) must parse
+ identically to the dict shape so the Router can raise a fallback."""
+ sse_chunk = (
+ b"event: error\n"
+ b'data: {"type": "error", "error": {"type": "internal_server_error", "message": "boom"}}\n\n'
+ )
+ assert parse_anthropic_error_event(sse_chunk) == ("internal_server_error", "boom", 500)
+ assert _is_provider_error_chunk(sse_chunk) is True
+
+
+def test_parse_anthropic_error_event_defaults_status_for_unknown_type():
+ chunk = {"type": "error", "error": {"type": "some_future_error_type", "message": "?"}}
+ assert parse_anthropic_error_event(chunk) == ("some_future_error_type", "?", 500)
+
+
+def test_parse_anthropic_error_event_missing_message_falls_back_to_type():
+ chunk = {"type": "error", "error": {"type": "overloaded_error"}}
+ assert parse_anthropic_error_event(chunk) == ("overloaded_error", "overloaded_error", 503)
+
+
+def test_parse_anthropic_error_event_non_string_error_type_returns_none():
+ """A malformed error body whose `type` field isn't a string (e.g. an
+ upstream bug sends null or a number) must not be treated as an error
+ event rather than crashing or forwarding a garbage error_type."""
+ chunk = {"type": "error", "error": {"type": None, "message": "boom"}}
+ assert parse_anthropic_error_event(chunk) is None
+
+
+def test_decoded_sse_data_line_swallows_invalid_json():
+ """A `data:` line that isn't valid JSON (a malformed/truncated frame)
+ must not be treated as an error event or raise, just be ignored."""
+ malformed_frame = b"event: error\ndata: {not valid json\n\n"
+ assert parse_anthropic_error_event(malformed_frame) is None
+ assert _is_provider_error_chunk(malformed_frame) is False
+
+
+class TestIsAnthropicContentDeltaChunk:
+ def test_dict_content_block_delta(self):
+ assert is_anthropic_content_delta_chunk({"type": "content_block_delta"}) is True
+
+ def test_dict_other_type(self):
+ assert is_anthropic_content_delta_chunk({"type": "message_start"}) is False
+
+ def test_bytes_content_block_delta(self):
+ assert is_anthropic_content_delta_chunk(b"event: content_block_delta\ndata: {}\n\n") is True
+
+ def test_bytes_other_event(self):
+ assert is_anthropic_content_delta_chunk(b"event: message_start\ndata: {}\n\n") is False
+
+ def test_neither_dict_nor_bytes(self):
+ assert is_anthropic_content_delta_chunk("content_block_delta") is False
+ assert is_anthropic_content_delta_chunk(None) is False
+
+
+@pytest.mark.parametrize(
+ "chunk",
+ [
+ {"type": "content_block_delta", "delta": {"type": "text_delta", "text": "hi"}},
+ b'event: content_block_delta\ndata: {"type": "content_block_delta"}\n\n',
+ b"raw-bytes",
+ "error",
+ None,
+ ],
+)
+def test_parse_anthropic_error_event_non_error_chunks_return_none(chunk):
+ assert parse_anthropic_error_event(chunk) is None
+ assert _is_provider_error_chunk(chunk) is False
+
+
+def test_parse_anthropic_error_event_ignores_substring_in_payload():
+ """A content_block_delta whose partial_json happens to contain the
+ literal string `"type": "error"` must not be misread as an error event."""
+ delta_frame_with_substring = (
+ b"event: content_block_delta\n"
+ b'data: {"type": "content_block_delta", "delta": '
+ b'{"type": "input_json_delta", "partial_json": "\\"type\\": \\"error\\""}}\n\n'
+ )
+ assert parse_anthropic_error_event(delta_frame_with_substring) is None
+
+
@pytest.mark.asyncio
async def test_async_sse_wrapper_emits_error_when_bytes_stream_only_mentions_message_stop_in_payload():
"""
@@ -307,3 +401,117 @@ def test_incomplete_stream_error_sse_event_is_valid_anthropic_error():
"error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE},
}
assert event.endswith("\n\n")
+
+
+def _decode_sse_events(events: tuple[bytes, ...]) -> list[tuple[str, dict]]:
+ decoded = []
+ for event in events:
+ assert isinstance(event, bytes)
+ lines = event.decode().split("\n")
+ assert lines[0].startswith("event: ")
+ decoded.append((lines[0].removeprefix("event: "), json.loads(lines[1].removeprefix("data: "))))
+ return decoded
+
+
+def test_anthropic_messages_response_as_sse_events_text_block():
+ response = {
+ "id": "msg_1",
+ "model": "claude-haiku",
+ "role": "assistant",
+ "type": "message",
+ "stop_reason": "end_turn",
+ "stop_sequence": None,
+ "content": [{"type": "text", "text": "hello"}],
+ "usage": {"input_tokens": 3, "output_tokens": 2},
+ }
+ decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
+ types = [event_type for event_type, _ in decoded]
+ assert types == [
+ "message_start",
+ "content_block_start",
+ "content_block_delta",
+ "content_block_stop",
+ "message_delta",
+ "message_stop",
+ ]
+ # message_start must not carry generated content itself, matching a real
+ # streaming response - it arrives via the content_block_delta that follows.
+ assert decoded[0][1]["message"]["content"] == []
+ assert decoded[0][1]["message"]["id"] == "msg_1"
+ # Bugbot regression: message_start must not carry the completed response's
+ # final stop_reason/stop_sequence/output_tokens - a real stream keeps those
+ # null/zero until message_delta, so a client could otherwise treat the
+ # message as already finished, or double-count output tokens.
+ assert decoded[0][1]["message"]["stop_reason"] is None
+ assert decoded[0][1]["message"]["stop_sequence"] is None
+ assert decoded[0][1]["message"]["usage"] == {"input_tokens": 3, "output_tokens": 0}
+ assert decoded[1][1]["content_block"] == {"type": "text", "text": ""}
+ assert decoded[2][1]["delta"] == {"type": "text_delta", "text": "hello"}
+ assert decoded[4][1]["delta"]["stop_reason"] == "end_turn"
+ assert decoded[4][1]["usage"] == {"input_tokens": 3, "output_tokens": 2}
+
+
+def test_anthropic_messages_response_as_sse_events_tool_use_block():
+ response = {
+ "id": "msg_2",
+ "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "NYC"}}],
+ "stop_reason": "tool_use",
+ }
+ decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
+ content_block_start = dict(decoded)["content_block_start"]
+ assert content_block_start["content_block"] == {
+ "type": "tool_use",
+ "id": "toolu_1",
+ "name": "get_weather",
+ "input": {},
+ }
+ content_block_delta = dict(decoded)["content_block_delta"]
+ assert json.loads(content_block_delta["delta"]["partial_json"]) == {"city": "NYC"}
+ assert content_block_delta["delta"]["type"] == "input_json_delta"
+
+
+def test_anthropic_messages_response_as_sse_events_thinking_block_emits_signature_delta():
+ """Bugbot regression: a thinking block's real `signature` must reach the
+ client via a trailing signature_delta, not be silently dropped - Anthropic
+ rejects a replayed assistant message (a follow-up turn, a tool-use
+ continuation) whose thinking block lacks its original signature."""
+ response = {
+ "id": "msg_5",
+ "content": [{"type": "thinking", "thinking": "let me think", "signature": "sig-abc123"}],
+ "stop_reason": "end_turn",
+ }
+ decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
+ deltas = [payload["delta"] for event_type, payload in decoded if event_type == "content_block_delta"]
+ assert deltas == [
+ {"type": "thinking_delta", "thinking": "let me think"},
+ {"type": "signature_delta", "signature": "sig-abc123"},
+ ]
+
+
+def test_anthropic_messages_response_as_sse_events_thinking_block_without_signature_omits_delta():
+ response = {
+ "id": "msg_6",
+ "content": [{"type": "thinking", "thinking": "let me think", "signature": None}],
+ "stop_reason": "end_turn",
+ }
+ decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
+ deltas = [payload["delta"] for event_type, payload in decoded if event_type == "content_block_delta"]
+ assert deltas == [{"type": "thinking_delta", "thinking": "let me think"}]
+
+
+def test_anthropic_messages_response_as_sse_events_multiple_blocks_are_indexed():
+ response = {
+ "id": "msg_3",
+ "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
+ }
+ decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
+ starts = [payload for event_type, payload in decoded if event_type == "content_block_start"]
+ assert [s["index"] for s in starts] == [0, 1]
+ deltas = [payload for event_type, payload in decoded if event_type == "content_block_delta"]
+ assert [d["delta"]["text"] for d in deltas] == ["a", "b"]
+
+
+def test_anthropic_messages_response_as_sse_events_no_content_blocks():
+ response = {"id": "msg_4", "content": [], "stop_reason": "end_turn"}
+ decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response))
+ assert [event_type for event_type, _ in decoded] == ["message_start", "message_delta", "message_stop"]
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 0b82ab971bf..56fd6df446f 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -12,8 +12,17 @@ import pytest
import litellm
+from litellm import Router
from litellm.exceptions import MidStreamFallbackError
from litellm.integrations.custom_logger import CustomLogger
+from litellm.router import (
+ MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS,
+ FallbackAwareAnthropicMessagesStream,
+ _anthropic_stream_commits_now,
+ _anthropic_stream_should_decline_fallback,
+ _anthropic_stream_should_drop_pre_content_ping,
+ _is_retriable_anthropic_status,
+)
def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata():
@@ -9224,3 +9233,1095 @@ class TestAddDeploymentApiBaseProviderResolution:
deployment = router.get_deployment_by_model_group_name("openai-via-gateway")
assert deployment is not None
assert deployment.litellm_params.custom_llm_provider == "openai"
+
+# =====================================================================
+# anthropic_messages mid-stream-fallback helpers, added for #24004
+# (mid-stream fallback not supported for anthropic_messages route type).
+#
+# anthropic_messages goes through _ageneric_api_call_with_fallbacks rather
+# than _acompletion, so its returned iterator was never wrapped by the chat
+# completions fallback handler: an SSE `event: error` frame from a native
+# Anthropic/Bedrock passthrough passed through to the client silently, and a
+# MidStreamFallbackError raised by the completion-bridge path's
+# CustomStreamWrapper (e.g. a Vertex AI transport drop) propagated
+# unhandled.
+#
+# Targets the helpers introduced on Router:
+# - _aanthropic_messages_streaming_iterator
+# - _aanthropic_messages_fallback_attempt
+# - _aanthropic_messages_with_streaming_fallbacks
+# - _dispatch_generic_call_type
+# =====================================================================
+
+
+async def _anthropic_messages_empty_generator():
+ return
+ yield # pragma: no cover - makes this an async generator
+
+
+def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream:
+ """A minimal wrapper for tests that call _aanthropic_messages_fallback_attempt
+ directly, bypassing _aanthropic_messages_streaming_iterator."""
+ return FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object())
+
+
+def _anthropic_messages_make_router() -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": "primary",
+ "litellm_params": {
+ "model": "anthropic/claude-sonnet-4-5",
+ "api_key": "sk-test",
+ },
+ },
+ {
+ "model_name": "fallback",
+ "litellm_params": {
+ "model": "bedrock/anthropic.claude-sonnet-4-5",
+ },
+ },
+ ]
+ )
+
+
+class _AnthropicMessagesFakeByteStream:
+ """Minimal AsyncIterator[bytes], carrying _hidden_params like
+ AnthropicMessagesStreamingResponse does."""
+
+ def __init__(self, chunks: list) -> None:
+ self._chunks = list(chunks)
+ self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}}
+ self.closed = False
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self) -> bytes:
+ if not self._chunks:
+ raise StopAsyncIteration
+ return self._chunks.pop(0)
+
+ async def aclose(self) -> None:
+ self.closed = True
+
+
+class _AnthropicMessagesRaisingByteStream:
+ """Simulates the completion-bridge path: no error SSE chunk is ever
+ yielded, the underlying CustomStreamWrapper raises MidStreamFallbackError
+ directly out of the iterator instead (a Vertex AI transport drop)."""
+
+ def __init__(self, chunks: list, error: Exception) -> None:
+ self._chunks = list(chunks)
+ self._error = error
+ self._hidden_params: dict = {}
+ self.closed = False
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self) -> bytes:
+ if self._chunks:
+ return self._chunks.pop(0)
+ raise self._error
+
+ async def aclose(self) -> None:
+ self.closed = True
+
+
+class _AnthropicMessagesFallbackByteStream:
+ def __init__(self, chunks: list, hidden_params: dict | None = None) -> None:
+ self._chunks = list(chunks)
+ self._hidden_params = hidden_params if hidden_params is not None else {}
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self) -> bytes:
+ if not self._chunks:
+ raise StopAsyncIteration
+ return self._chunks.pop(0)
+
+
+def _anthropic_messages_overloaded_error_chunk() -> bytes:
+ return (
+ b"event: error\n"
+ b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n'
+ )
+
+
+def _anthropic_messages_invalid_request_error_chunk() -> bytes:
+ return (
+ b"event: error\n"
+ b'data: {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}\n\n'
+ )
+
+
+def _anthropic_messages_rate_limit_error_chunk() -> bytes:
+ return (
+ b"event: error\n"
+ b'data: {"type": "error", "error": {"type": "rate_limit_error", "message": "Too many requests"}}\n\n'
+ )
+
+
+def _anthropic_messages_content_chunk(text: str = "hi") -> bytes:
+ payload = f'{{"type": "content_block_delta", "delta": {{"type": "text_delta", "text": "{text}"}}}}'
+ return f"event: content_block_delta\ndata: {payload}\n\n".encode()
+
+
+def _anthropic_messages_message_start_chunk() -> bytes:
+ """A lifecycle/bookkeeping frame Anthropic sends before any real content -
+ routinely the very first event before an overload error."""
+ return b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_1"}}\n\n'
+
+
+def _anthropic_messages_ping_chunk() -> bytes:
+ return b'event: ping\ndata: {"type": "ping"}\n\n'
+
+
+# -------- _aanthropic_messages_streaming_iterator (passthrough) --------
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_streaming_iterator_passthrough():
+ """Without any error chunk, the wrapper forwards every chunk unchanged
+ and carries the source iterator's _hidden_params through (so response
+ headers like Bedrock's request-id keep flowing to the client)."""
+ router = _anthropic_messages_make_router()
+ source = _AnthropicMessagesFakeByteStream(
+ [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")]
+ )
+
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source, initial_kwargs={"model": "primary"}
+ )
+
+ collected = [chunk async for chunk in wrapped]
+ assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")]
+ assert wrapped._hidden_params["additional_headers"]["x-amzn-requestid"] == "req-1"
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_frames_in_order():
+ """Regression: lifecycle frames held back to guard against a mid-stream
+ fallback must still reach the client, in order, once real content
+ arrives - buffering them for the fallback-safety check must not silently
+ drop them on the happy path."""
+ router = _anthropic_messages_make_router()
+ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n'
+ source = _AnthropicMessagesFakeByteStream(
+ [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop]
+ )
+
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source, initial_kwargs={"model": "primary"}
+ )
+
+ collected = [chunk async for chunk in wrapped]
+ assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop]
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_stream_end():
+ """Regression: if the primary stream ends with only lifecycle frames and
+ no content and no error, the buffered frames must still reach the
+ client rather than being silently swallowed."""
+ router = _anthropic_messages_make_router()
+ message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n'
+ source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop])
+
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source, initial_kwargs={"model": "primary"}
+ )
+
+ collected = [chunk async for chunk in wrapped]
+ assert collected == [_anthropic_messages_message_start_chunk(), message_stop]
+
+ with pytest.raises(StopAsyncIteration):
+ await wrapped.__anext__()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_chunk_skips_fallback():
+ """Greptile review round: transport-level buffering can coalesce a real
+ content_block_delta and a following retriable error into ONE physical
+ read from the source iterator. Since the whole chunk (content and error
+ together) is forwarded to the client atomically, the client genuinely
+ receives the content - so no fallback must be attempted, exactly as if
+ the two events had arrived as separate reads."""
+ router = _anthropic_messages_make_router()
+ coalesced_chunk = _anthropic_messages_content_chunk("partial") + _anthropic_messages_overloaded_error_chunk()
+ source = _AnthropicMessagesFakeByteStream([coalesced_chunk])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ assert collected == [coalesced_chunk]
+ mock_fallback.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_ping_keepalive_never_buffered_or_forwarded():
+ """Bugbot regression: a `ping` keepalive carries no content and must be
+ dropped outright before any real content arrives, rather than buffered -
+ otherwise a slow-starting connection sending many pings could grow the
+ pre-content buffer without bound."""
+ router = _anthropic_messages_make_router()
+ source = _AnthropicMessagesFakeByteStream(
+ [_anthropic_messages_ping_chunk(), _anthropic_messages_content_chunk("hi")]
+ )
+
+ wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"})
+ collected = [chunk async for chunk in wrapped]
+
+ assert _anthropic_messages_ping_chunk() not in collected
+ assert collected == [_anthropic_messages_content_chunk("hi")]
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_pre_content_buffer_cap_forces_commit():
+ """Bugbot regression: a hostile or pathological upstream that never emits
+ real content or an error must not grow the pre-content lifecycle buffer
+ without bound - hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS commits
+ to the primary stream early, exactly as real content arriving would."""
+ router = _anthropic_messages_make_router()
+ lifecycle_chunk = _anthropic_messages_message_start_chunk()
+ error_chunk = _anthropic_messages_overloaded_error_chunk()
+ chunks = [lifecycle_chunk] * (MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5) + [error_chunk]
+ source = _AnthropicMessagesFakeByteStream(chunks)
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source, initial_kwargs={"model": "primary"}
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ mock_fallback.assert_not_awaited()
+ assert collected.count(lifecycle_chunk) == MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5
+ assert collected[-1] == error_chunk
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_ping_coalesced_with_content_in_one_physical_chunk_is_forwarded():
+ """Greptile/Bugbot regression: transport-level buffering can coalesce a
+ `ping` keepalive and a real content_block_delta into ONE physical read.
+ The pre-content ping-drop must only discard PURE ping frames - dropping
+ the whole coalesced chunk would silently lose generated content."""
+ router = _anthropic_messages_make_router()
+ coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_content_chunk("hi")
+ source = _AnthropicMessagesFakeByteStream([coalesced_chunk])
+
+ wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"})
+ collected = [chunk async for chunk in wrapped]
+
+ assert collected == [coalesced_chunk]
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_ping_coalesced_with_retriable_error_still_falls_back():
+ """Greptile/Bugbot regression: a physical chunk coalescing a `ping` with a
+ retriable `event: error` must not be discarded as a keepalive - the error
+ inside it must still trigger the mid-stream fallback."""
+ router = _anthropic_messages_make_router()
+ coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_overloaded_error_chunk()
+ source = _AnthropicMessagesFakeByteStream([coalesced_chunk])
+ fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=fallback_stream),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source, initial_kwargs={"model": "primary"}
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ mock_fallback.assert_awaited_once()
+ assert collected == [_anthropic_messages_content_chunk("fallback answer")]
+
+
+# -------- _aanthropic_messages_fallback_attempt --------
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_fallback_attempt_yields_fallback_stream():
+ """Direct-call regression: the fallback-attempt helper re-enters the
+ Router's fallback chain and forwards whatever the fallback produces."""
+ router = _anthropic_messages_make_router()
+ fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")])
+ error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic")
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=fallback_stream),
+ ) as mock_fallback:
+ collected = [
+ chunk
+ async for chunk in router._aanthropic_messages_fallback_attempt(
+ error,
+ {"model": "primary", "messages": [{"role": "user", "content": "hi"}]},
+ _anthropic_messages_make_wrapper(),
+ )
+ ]
+
+ assert collected == [_anthropic_messages_content_chunk("fallback answer")]
+ mock_fallback.assert_awaited_once()
+ assert mock_fallback.await_args.kwargs["e"] is error
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_fallback_attempt_raises_original_exception_on_double_failure():
+ """Direct-call regression: when the fallback attempt itself fails with a
+ MidStreamFallbackError wrapping a real provider exception, that real
+ exception must surface rather than the internal wrapper exception."""
+ router = _anthropic_messages_make_router()
+ error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic")
+ original_exception = litellm.APIError(
+ status_code=503, message="fallback also overloaded", llm_provider="bedrock", model="fallback"
+ )
+ fallback_failure = MidStreamFallbackError(
+ message="fallback failed", model="fallback", llm_provider="bedrock", original_exception=original_exception
+ )
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(side_effect=fallback_failure),
+ ):
+ with pytest.raises(litellm.APIError) as exc_info:
+ async for _ in router._aanthropic_messages_fallback_attempt(
+ error, {"model": "primary"}, _anthropic_messages_make_wrapper()
+ ):
+ pass
+
+ assert exc_info.value is original_exception
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_fallback_attempt_yields_non_streaming_fallback_response():
+ """Bugbot regression: a fallback that resolves to a non-streaming
+ response (no __aiter__, e.g. an agentic tool-use interception loop) must
+ be synthesized into a valid SSE byte sequence, not yielded as a raw dict
+ into a byte stream - the generator is typed AsyncGenerator[bytes, None]
+ and every item reaching the client must be a real SSE frame."""
+ router = _anthropic_messages_make_router()
+ error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic")
+ non_streaming_response = {"id": "msg_1", "type": "message", "content": [{"type": "text", "text": "hi"}]}
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=non_streaming_response),
+ ):
+ collected = [
+ item
+ async for item in router._aanthropic_messages_fallback_attempt(
+ error, {"model": "primary"}, _anthropic_messages_make_wrapper()
+ )
+ ]
+
+ assert all(isinstance(item, bytes) for item in collected)
+ event_types = [item.split(b"\n")[0].removeprefix(b"event: ") for item in collected]
+ assert event_types == [
+ b"message_start",
+ b"content_block_start",
+ b"content_block_delta",
+ b"content_block_stop",
+ b"message_delta",
+ b"message_stop",
+ ]
+ assert b'"text": "hi"' in collected[2]
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_fallback_attempt_reraises_plain_exception_on_double_failure():
+ """Direct-call regression: when the fallback attempt fails with a plain
+ exception (not a MidStreamFallbackError), that exception itself must
+ propagate unchanged."""
+ router = _anthropic_messages_make_router()
+ error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic")
+ fallback_failure = ValueError("no healthy deployments")
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(side_effect=fallback_failure),
+ ):
+ with pytest.raises(ValueError, match="no healthy deployments") as exc_info:
+ async for _ in router._aanthropic_messages_fallback_attempt(
+ error, {"model": "primary"}, _anthropic_messages_make_wrapper()
+ ):
+ pass
+
+ assert exc_info.value is fallback_failure
+
+
+# -------- _aanthropic_messages_with_streaming_fallbacks --------
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_with_streaming_fallbacks_non_streaming_passthrough():
+ """A non-streaming response (plain dict) is returned unchanged, never wrapped."""
+ router = _anthropic_messages_make_router()
+ plain_response = {"id": "msg_1", "type": "message"}
+
+ async def fake_original(**_kwargs):
+ return plain_response
+
+ with patch.object(
+ router,
+ "_ageneric_api_call_with_fallbacks",
+ new=AsyncMock(return_value=plain_response),
+ ):
+ out = await router._aanthropic_messages_with_streaming_fallbacks(
+ original_function=fake_original,
+ model="primary",
+ stream=False,
+ )
+ assert out is plain_response
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_with_streaming_fallbacks_wraps_streaming_iterator():
+ """A streaming response is wrapped via _aanthropic_messages_streaming_iterator."""
+ router = _anthropic_messages_make_router()
+ streaming_iter = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk()])
+ wrapped_marker = object()
+
+ async def fake_original(**_kwargs):
+ return streaming_iter
+
+ with (
+ patch.object(
+ router,
+ "_ageneric_api_call_with_fallbacks",
+ new=AsyncMock(return_value=streaming_iter),
+ ),
+ patch.object(
+ router,
+ "_aanthropic_messages_streaming_iterator",
+ new=AsyncMock(return_value=wrapped_marker),
+ ) as mock_wrap,
+ ):
+ out = await router._aanthropic_messages_with_streaming_fallbacks(
+ original_function=fake_original,
+ model="primary",
+ stream=True,
+ )
+ assert out is wrapped_marker
+ mock_wrap.assert_awaited_once()
+
+
+# -------- mid-stream error handling --------
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_fallback_on_pre_first_chunk_error_event():
+ """Regression for #24004: a retriable SSE `event: error` frame
+ (overloaded_error/internal_server_error) that arrives before any real
+ content must trigger the router's fallback chain instead of passing
+ through to the client silently."""
+ router = _anthropic_messages_make_router()
+ source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()])
+ fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=fallback_stream),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]},
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ assert collected == [_anthropic_messages_content_chunk("fallback answer")]
+ mock_fallback.assert_awaited_once()
+ raised = mock_fallback.await_args.kwargs["e"]
+ assert isinstance(raised, MidStreamFallbackError)
+ assert raised.status_code == 503
+ assert raised.is_pre_first_chunk is True
+ assert source.closed is True
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_mid_stream_error_preserves_real_status_code():
+ """Bugbot regression: the MidStreamFallbackError raised for a detected SSE
+ `event: error` frame must carry the error's REAL parsed status code
+ (via original_exception), not silently default to 503 for every error
+ type - a rate_limit_error (429) must surface as 429, not 503."""
+ router = _anthropic_messages_make_router()
+ source = _AnthropicMessagesFakeByteStream([_anthropic_messages_rate_limit_error_chunk()])
+ fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=fallback_stream),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]},
+ )
+ [chunk async for chunk in wrapped]
+
+ raised = mock_fallback.await_args.kwargs["e"]
+ assert isinstance(raised, MidStreamFallbackError)
+ assert raised.status_code == 429
+ assert raised.original_exception is not None
+ assert raised.original_exception.status_code == 429
+ assert raised.original_exception.llm_provider == "anthropic"
+
+
+def test_merge_fallback_hidden_params_direct_call():
+ """Direct-call regression: merge_fallback_hidden_params combines the
+ fallback's hidden params/headers with whatever was already present,
+ with the fallback's values winning on key collisions."""
+ wrapper = FallbackAwareAnthropicMessagesStream(
+ _anthropic_messages_empty_generator(),
+ _AnthropicMessagesFakeByteStream([]), # carries {"additional_headers": {"x-amzn-requestid": "req-1"}}
+ )
+ wrapper.merge_fallback_hidden_params(
+ {"model_id": "fallback-deployment"},
+ {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"},
+ )
+ assert wrapper._hidden_params["model_id"] == "fallback-deployment"
+ assert wrapper._hidden_params["additional_headers"] == {
+ "x-amzn-requestid": "req-2",
+ "x-fallback-only": "yes",
+ }
+
+
+def test_anthropic_stream_should_drop_pre_content_ping_direct_call():
+ ping = _anthropic_messages_ping_chunk()
+ content = _anthropic_messages_content_chunk("hi")
+ assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=False) is True
+ assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=True) is False
+ assert _anthropic_stream_should_drop_pre_content_ping(content, has_generated_content=False) is False
+
+
+def test_is_retriable_anthropic_status_direct_call():
+ assert _is_retriable_anthropic_status(429) is True
+ assert _is_retriable_anthropic_status(503) is True
+ assert _is_retriable_anthropic_status(500) is True
+ assert _is_retriable_anthropic_status(400) is False
+ assert _is_retriable_anthropic_status(404) is False
+
+
+def test_anthropic_stream_should_decline_fallback_direct_call():
+ pre_first_chunk_error = MidStreamFallbackError(
+ message="overloaded", model="primary", llm_provider="anthropic", is_pre_first_chunk=True
+ )
+ post_first_chunk_error = MidStreamFallbackError(
+ message="overloaded", model="primary", llm_provider="anthropic", is_pre_first_chunk=False
+ )
+ assert _anthropic_stream_should_decline_fallback(False, pre_first_chunk_error) is False
+ assert _anthropic_stream_should_decline_fallback(True, pre_first_chunk_error) is True
+ assert _anthropic_stream_should_decline_fallback(False, post_first_chunk_error) is True
+
+
+def test_anthropic_stream_commits_now_direct_call():
+ content = _anthropic_messages_content_chunk("hi")
+ lifecycle_chunk = _anthropic_messages_message_start_chunk()
+ assert _anthropic_stream_commits_now(content, has_generated_content=False, buffered_chunk_count=0) is True
+ assert _anthropic_stream_commits_now(content, has_generated_content=True, buffered_chunk_count=0) is False
+ assert (
+ _anthropic_stream_commits_now(
+ lifecycle_chunk,
+ has_generated_content=False,
+ buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS,
+ )
+ is True
+ )
+ assert (
+ _anthropic_stream_commits_now(
+ lifecycle_chunk,
+ has_generated_content=False,
+ buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS - 1,
+ )
+ is False
+ )
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_fallback_merges_fallback_hidden_params():
+ """Bugbot regression: after a successful mid-stream fallback, the
+ wrapper's _hidden_params must reflect the FALLBACK deployment's own
+ provider headers (e.g. a different Bedrock request-id), not stay
+ frozen on the primary's - raw bytes can't carry per-item _hidden_params
+ the way a ModelResponseStream/ResponsesAPI event can, so the wrapper
+ itself is the only place left to expose them."""
+ router = _anthropic_messages_make_router()
+ source = _AnthropicMessagesFakeByteStream(
+ [_anthropic_messages_overloaded_error_chunk()]
+ ) # carries x-amzn-requestid: req-1
+ fallback_stream = _AnthropicMessagesFallbackByteStream(
+ [_anthropic_messages_content_chunk("fallback answer")],
+ hidden_params={"additional_headers": {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}},
+ )
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=fallback_stream),
+ ):
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ _ = [chunk async for chunk in wrapped]
+
+ headers = wrapped._hidden_params["additional_headers"]
+ assert headers["x-amzn-requestid"] == "req-2"
+ assert headers["x-fallback-only"] == "yes"
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_nested_metadata():
+ """Bugbot regression: a shallow .copy() of kwargs still shares the
+ nested litellm_metadata/metadata dict objects with the primary attempt.
+ _update_kwargs_with_deployment mutates that dict in place with
+ deployment-specific fields, which must not leak into the fallback
+ request's metadata."""
+ router = _anthropic_messages_make_router()
+ primary_metadata = {"model_group": "primary"}
+ streaming_iter_kwargs = {}
+
+ async def fake_original(**_kwargs):
+ # Simulate _update_kwargs_with_deployment mutating the primary's
+ # litellm_metadata in place, as the real helper does.
+ primary_metadata["deployment"] = "primary-deployment-object"
+ return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")])
+
+ with patch.object(
+ router,
+ "_aanthropic_messages_streaming_iterator",
+ new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"),
+ ):
+ with patch.object(
+ router,
+ "_ageneric_api_call_with_fallbacks",
+ new=AsyncMock(side_effect=fake_original),
+ ):
+ await router._aanthropic_messages_with_streaming_fallbacks(
+ original_function=fake_original,
+ model="primary",
+ stream=True,
+ litellm_metadata=primary_metadata,
+ )
+
+ fallback_kwargs = streaming_iter_kwargs["initial_kwargs"]
+ assert fallback_kwargs["litellm_metadata"] is not primary_metadata
+ assert "deployment" not in fallback_kwargs["litellm_metadata"]
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_metadata_field():
+ """Same regression as above for the (separate) `metadata` kwarg some
+ call sites use instead of `litellm_metadata`."""
+ router = _anthropic_messages_make_router()
+ primary_metadata = {"tag": "primary"}
+ streaming_iter_kwargs = {}
+
+ async def fake_original(**_kwargs):
+ primary_metadata["deployment"] = "primary-deployment-object"
+ return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")])
+
+ with patch.object(
+ router,
+ "_aanthropic_messages_streaming_iterator",
+ new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"),
+ ):
+ with patch.object(
+ router,
+ "_ageneric_api_call_with_fallbacks",
+ new=AsyncMock(side_effect=fake_original),
+ ):
+ await router._aanthropic_messages_with_streaming_fallbacks(
+ original_function=fake_original,
+ model="primary",
+ stream=True,
+ metadata=primary_metadata,
+ )
+
+ fallback_kwargs = streaming_iter_kwargs["initial_kwargs"]
+ assert fallback_kwargs["metadata"] is not primary_metadata
+ assert "deployment" not in fallback_kwargs["metadata"]
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_fallback_triggers_after_lifecycle_only_frame():
+ """Regression: Anthropic routinely sends a message_start lifecycle frame
+ before an overload error even fires. A lifecycle-only frame (no real
+ content) must not disqualify the fallback attempt, and must not reach
+ the client either - forwarding it and then appending the fallback's own
+ message_start would produce two overlapping message lifecycles on one
+ SSE stream. The primary's buffered lifecycle frame is discarded and the
+ client sees only the fallback's own, single, clean lifecycle."""
+ router = _anthropic_messages_make_router()
+ source = _AnthropicMessagesFakeByteStream(
+ [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()]
+ )
+ fallback_message_start = b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_2"}}\n\n'
+ fallback_stream = _AnthropicMessagesFallbackByteStream(
+ [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")]
+ )
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=fallback_stream),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ assert collected == [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")]
+ assert collected.count(_anthropic_messages_message_start_chunk()) == 0, (
+ "the primary's message_start must never reach the client"
+ )
+ assert sum(1 for c in collected if c.startswith(b"event: message_start")) == 1, (
+ "exactly one message_start must reach the client"
+ )
+ mock_fallback.assert_awaited_once()
+ raised = mock_fallback.await_args.kwargs["e"]
+ assert raised.is_pre_first_chunk is True
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_raised_error_after_real_content_does_not_restart_stream():
+ """Regression: a MidStreamFallbackError raised directly by the source
+ iterator (the completion-bridge path's CustomStreamWrapper, e.g. a
+ transport drop) must not trigger a fallback once real content already
+ reached the client - that would append a second, overlapping message
+ lifecycle onto the same SSE stream. The original exception must
+ propagate to the caller instead."""
+ router = _anthropic_messages_make_router()
+ content = _anthropic_messages_content_chunk("partial answer")
+ original_exception = litellm.APIError(
+ status_code=503,
+ message="stream reset",
+ llm_provider="vertex_ai",
+ model="primary",
+ )
+ raised_error = MidStreamFallbackError(
+ message="stream reset",
+ model="primary",
+ llm_provider="vertex_ai",
+ original_exception=original_exception,
+ is_pre_first_chunk=False,
+ )
+ source = _AnthropicMessagesRaisingByteStream([content], raised_error)
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = []
+
+ async def _consume():
+ async for chunk in wrapped:
+ collected.append(chunk)
+
+ with pytest.raises(litellm.APIError) as exc_info:
+ await _consume()
+
+ assert collected == [content]
+ assert exc_info.value is original_exception
+ mock_fallback.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_fallback_also_catches_raised_midstream_error():
+ """Regression for the completion-bridge path (deployments with no native
+ /v1/messages endpoint): its CustomStreamWrapper raises
+ MidStreamFallbackError directly (e.g. on a Vertex AI transport drop)
+ instead of yielding an SSE error chunk - the wrapper must catch that too."""
+ router = _anthropic_messages_make_router()
+ raised_error = MidStreamFallbackError(
+ message="stream reset",
+ model="primary",
+ llm_provider="vertex_ai",
+ is_pre_first_chunk=True,
+ )
+ source = _AnthropicMessagesRaisingByteStream([], raised_error)
+ fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=fallback_stream),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ assert collected == [_anthropic_messages_content_chunk("fallback answer")]
+ mock_fallback.assert_awaited_once()
+ assert mock_fallback.await_args.kwargs["e"] is raised_error
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_non_retriable_client_error_skips_fallback():
+ """A 4xx (non-429) error type (e.g. invalid_request_error) is a client
+ error a fallback attempt cannot fix, so it must be forwarded to the
+ client as-is rather than burning a fallback attempt."""
+ router = _anthropic_messages_make_router()
+ error_chunk = _anthropic_messages_invalid_request_error_chunk()
+ source = _AnthropicMessagesFakeByteStream([error_chunk])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ assert collected == [error_chunk]
+ mock_fallback.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_post_first_chunk_error_skips_fallback():
+ """Once content has already reached the caller, retrying would start a
+ second, overlapping Anthropic message lifecycle on the same SSE stream -
+ the error must be forwarded instead of triggering an invisible retry."""
+ router = _anthropic_messages_make_router()
+ content = _anthropic_messages_content_chunk("partial answer")
+ error_chunk = _anthropic_messages_overloaded_error_chunk()
+ source = _AnthropicMessagesFakeByteStream([content, error_chunk])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ assert collected == [content, error_chunk]
+ mock_fallback.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_non_retriable_error_flushes_buffered_lifecycle_frames():
+ """A non-retriable error arriving while lifecycle frames are still
+ buffered (no content seen yet) must flush those buffered frames before
+ forwarding the error, so the client still sees the whole primary
+ attempt rather than losing the buffered message_start silently."""
+ router = _anthropic_messages_make_router()
+ error_chunk = _anthropic_messages_invalid_request_error_chunk()
+ source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), error_chunk])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = [chunk async for chunk in wrapped]
+
+ assert collected == [_anthropic_messages_message_start_chunk(), error_chunk]
+ mock_fallback.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_raised_error_declined_flushes_buffered_lifecycle_frames():
+ """When a raised MidStreamFallbackError is declined (source says content
+ was not pre-first-chunk) while lifecycle frames are still buffered, they
+ must be flushed to the client before the exception propagates."""
+ router = _anthropic_messages_make_router()
+ raised_error = MidStreamFallbackError(
+ message="stream reset",
+ model="primary",
+ llm_provider="vertex_ai",
+ is_pre_first_chunk=False,
+ )
+ source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error)
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(),
+ ) as mock_fallback:
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = []
+
+ async def _consume():
+ async for chunk in wrapped:
+ collected.append(chunk)
+
+ with pytest.raises(MidStreamFallbackError) as exc_info:
+ await _consume()
+
+ assert collected == [_anthropic_messages_message_start_chunk()]
+ assert exc_info.value is raised_error
+ mock_fallback.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_raised_error_without_original_exception_reraises_itself():
+ """When a declined MidStreamFallbackError carries no original_exception,
+ the bare exception itself must propagate rather than being swallowed."""
+ router = _anthropic_messages_make_router()
+ content = _anthropic_messages_content_chunk("partial answer")
+ raised_error = MidStreamFallbackError(
+ message="stream reset",
+ model="primary",
+ llm_provider="vertex_ai",
+ is_pre_first_chunk=False,
+ )
+ source = _AnthropicMessagesRaisingByteStream([content], raised_error)
+
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ collected = []
+
+ async def _consume():
+ async for chunk in wrapped:
+ collected.append(chunk)
+
+ with pytest.raises(MidStreamFallbackError) as exc_info:
+ await _consume()
+
+ assert collected == [content]
+ assert exc_info.value is raised_error
+
+
+@pytest.mark.asyncio
+async def test_anthropic_messages_fallback_also_failing_raises_original_exception():
+ """If the fallback attempt itself fails with a MidStreamFallbackError
+ wrapping a real provider exception, the client must see that real
+ exception, not the internal MidStreamFallbackError."""
+ router = _anthropic_messages_make_router()
+ source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()])
+ original_exception = litellm.APIError(
+ status_code=503,
+ message="fallback also overloaded",
+ llm_provider="bedrock",
+ model="fallback",
+ )
+ fallback_failure = MidStreamFallbackError(
+ message="fallback failed",
+ model="fallback",
+ llm_provider="bedrock",
+ original_exception=original_exception,
+ )
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(side_effect=fallback_failure),
+ ):
+ wrapped = await router._aanthropic_messages_streaming_iterator(
+ response=source,
+ initial_kwargs={"model": "primary"},
+ )
+ with pytest.raises(litellm.APIError) as exc_info:
+ async for _ in wrapped:
+ pass
+
+ assert exc_info.value is original_exception
+
+
+# -------- _dispatch_generic_call_type --------
+
+
+@pytest.mark.asyncio
+async def test_dispatch_generic_call_type_routes_anthropic_messages_through_streaming_fallbacks():
+ router = _anthropic_messages_make_router()
+
+ async def fake_original(**_kwargs):
+ return {"id": "msg_1"}
+
+ with patch.object(
+ router,
+ "_aanthropic_messages_with_streaming_fallbacks",
+ new=AsyncMock(return_value="anthropic-result"),
+ ) as mock_anthropic:
+ out = await router._dispatch_generic_call_type(
+ call_type="anthropic_messages",
+ original_function=fake_original,
+ model="primary",
+ )
+ assert out == "anthropic-result"
+ mock_anthropic.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_dispatch_generic_call_type_other_call_types_use_generic_fallback():
+ router = _anthropic_messages_make_router()
+
+ async def fake_original(**_kwargs):
+ return {"id": "file_1"}
+
+ with patch.object(
+ router,
+ "_ageneric_api_call_with_fallbacks",
+ new=AsyncMock(return_value="generic-result"),
+ ) as mock_generic:
+ out = await router._dispatch_generic_call_type(
+ call_type="afile_delete",
+ original_function=fake_original,
+ model="primary",
+ )
+ assert out == "generic-result"
+ mock_generic.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_factory_function_anthropic_messages_uses_streaming_fallback_dispatch():
+ """anthropic_messages must be wired through the mid-stream-fallback-aware
+ path rather than the bare generic dispatch every other call type without
+ special handling uses."""
+ router = _anthropic_messages_make_router()
+ wrapped = router.factory_function(litellm.anthropic_messages, call_type="anthropic_messages")
+ assert callable(wrapped)
+
+ with patch.object(
+ router,
+ "_aanthropic_messages_with_streaming_fallbacks",
+ new=AsyncMock(return_value="ok"),
+ ) as mock_anthropic:
+ result = await wrapped(model="primary")
+ assert result == "ok"
+ mock_anthropic.assert_awaited_once()
From e1a092044599ffeda5122fafd9adce879bb941e5 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 13:54:34 -0700
Subject: [PATCH 61/70] refactor(anthropic): bind tool_result file parts once
instead of seeding an empty tuple
---
.../responses_adapters/transformation.py | 17 ++++++++++-------
1 file changed, 10 insertions(+), 7 deletions(-)
diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
index 43c3e504964..d9d62e0719f 100644
--- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
+++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
@@ -271,7 +271,16 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
elif btype == "tool_result":
tool_use_id = block.get("tool_use_id", "")
inner = block.get("content")
- tool_file_parts: tuple[dict[str, str], ...] = () # mutable-ok: json content parts
+ document_candidates = (
+ tuple(
+ self._translate_anthropic_document_block_to_file_part(c)
+ for c in inner
+ if isinstance(c, dict) and c.get("type") == "document"
+ )
+ if isinstance(inner, list)
+ else ()
+ )
+ tool_file_parts = tuple(part for part in document_candidates if part is not None)
if inner is None:
output_text = ""
elif isinstance(inner, str):
@@ -297,12 +306,6 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
{"type": "input_image", "image_url": url} # mutable-ok: json content part
for url in image_urls
)
- document_candidates = tuple(
- self._translate_anthropic_document_block_to_file_part(c)
- for c in inner
- if isinstance(c, dict) and c.get("type") == "document"
- )
- tool_file_parts = tuple(part for part in document_candidates if part is not None)
else:
output_text = str(inner)
# tool_result is a top-level item, not inside the message
From c3c9e3a91968eec197805dd4218cfec4f987b738 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 13:58:37 -0700
Subject: [PATCH 62/70] test(anthropic): pin dropped file-id and empty-url
document sources to string output
---
.../test_responses_adapters_transformation.py | 17 +++++++++++++++++
1 file changed, 17 insertions(+)
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
index f0cdaa4e8e3..44c956c8ee8 100644
--- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
@@ -1644,6 +1644,23 @@ class TestToolResultDocuments:
output = self._tool_output(self._translate([{"type": "text", "text": "plain result"}]))
assert output == "plain result"
+ def test_file_id_source_document_keeps_string_output(self):
+ output = self._tool_output(
+ self._translate(
+ [
+ {"type": "text", "text": "stub"},
+ {"type": "document", "source": {"type": "file", "file_id": "file_abc123"}},
+ ]
+ )
+ )
+ assert output == "stub"
+
+ def test_url_source_without_url_keeps_string_output(self):
+ output = self._tool_output(
+ self._translate([{"type": "text", "text": "stub"}, {"type": "document", "source": {"type": "url"}}])
+ )
+ assert output == "stub"
+
def test_text_image_and_document_mix(self):
items = self._translate(
[
From 75c4565dde1ed63ecd8a3ef67ed1829da4888759 Mon Sep 17 00:00:00 2001
From: Deepanshu Lulla
Date: Tue, 25 Aug 2026 17:10:56 -0400
Subject: [PATCH 63/70] fix(cerebras): add max_retries and extra_headers to
get_supported_openai_params (#36601)
Co-authored-by: Deepanshu
---
litellm/llms/cerebras/chat.py | 2 +
tests/test_litellm/llms/cerebras/__init__.py | 0
.../test_cerebras_chat_transformation.py | 61 +++++++++++++++++++
3 files changed, 63 insertions(+)
create mode 100644 tests/test_litellm/llms/cerebras/__init__.py
create mode 100644 tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py
diff --git a/litellm/llms/cerebras/chat.py b/litellm/llms/cerebras/chat.py
index 8827b0afd87..c3aa26ade35 100644
--- a/litellm/llms/cerebras/chat.py
+++ b/litellm/llms/cerebras/chat.py
@@ -68,6 +68,8 @@ class CerebrasConfig(OpenAIGPTConfig):
"tool_choice",
"tools",
"user",
+ "max_retries",
+ "extra_headers",
]
# Only add reasoning_effort for models that support it
diff --git a/tests/test_litellm/llms/cerebras/__init__.py b/tests/test_litellm/llms/cerebras/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py b/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py
new file mode 100644
index 00000000000..09718b1e6e0
--- /dev/null
+++ b/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py
@@ -0,0 +1,61 @@
+from litellm.llms.cerebras.chat import CerebrasConfig
+
+
+def test_max_retries_in_supported_params() -> None:
+ config = CerebrasConfig()
+ params = config.get_supported_openai_params(model="llama-3.3-70b")
+ assert "max_retries" in params, (
+ f"max_retries must be in CerebrasConfig.get_supported_openai_params(); got: {params!r}"
+ )
+
+
+def test_extra_headers_in_supported_params() -> None:
+ config = CerebrasConfig()
+ params = config.get_supported_openai_params(model="llama-3.3-70b")
+ assert "extra_headers" in params, (
+ f"extra_headers must be in CerebrasConfig.get_supported_openai_params(); got: {params!r}"
+ )
+
+
+def test_core_openai_params_still_supported() -> None:
+ config = CerebrasConfig()
+ params = config.get_supported_openai_params(model="llama-3.3-70b")
+ for expected in (
+ "max_tokens",
+ "max_completion_tokens",
+ "response_format",
+ "seed",
+ "stop",
+ "stream",
+ "temperature",
+ "top_p",
+ "tool_choice",
+ "tools",
+ "user",
+ ):
+ assert expected in params, f"{expected!r} unexpectedly missing from Cerebras supported params: {params!r}"
+
+
+def test_map_openai_params_preserves_max_retries() -> None:
+ config = CerebrasConfig()
+ result = config.map_openai_params(
+ non_default_params={"max_retries": 0, "temperature": 0.7},
+ optional_params={},
+ model="llama-3.3-70b",
+ drop_params=False,
+ )
+ assert result.get("max_retries") == 0, f"map_openai_params must preserve max_retries=0; got: {result!r}"
+ assert result.get("temperature") == 0.7
+
+
+def test_map_openai_params_preserves_max_retries_zero_falsy() -> None:
+ config = CerebrasConfig()
+ result = config.map_openai_params(
+ non_default_params={"max_retries": 0},
+ optional_params={},
+ model="llama-3.3-70b",
+ drop_params=False,
+ )
+ assert "max_retries" in result and result["max_retries"] == 0, (
+ f"max_retries=0 (falsy) must not be silently omitted; got: {result!r}"
+ )
From ec365a3c12991b924bfa2f85cd62e69e39569c42 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 14:38:34 -0700
Subject: [PATCH 64/70] fix(together_ai): pass tools through for models missing
from the registry
Tools passed for Together models with no supports_function_calling entry in
model_prices_and_context_window.json were rejected with UnsupportedParamsError
by default and silently dropped under drop_params, which made models emit tool
calls as plain text. Fail open instead: pass tool params through with a warning
and let Together validate. Models the registry explicitly marks as not
supporting function calling keep the loud contract: raise by default, drop with
a warning under drop_params.
---
.../llms/together_ai/chat/transformation.py | 61 +++++-
.../test_together_ai_chat_transformation.py | 195 ++++++++++++++++--
2 files changed, 226 insertions(+), 30 deletions(-)
diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py
index 88fd79f2366..3162a34f1b9 100644
--- a/litellm/llms/together_ai/chat/transformation.py
+++ b/litellm/llms/together_ai/chat/transformation.py
@@ -4,34 +4,74 @@ Translates from OpenAI's `/v1/chat/completions` to Together AI's `/v1/chat/compl
Docs: https://docs.together.ai/docs/chat-overview
"""
+from collections.abc import Container
from types import MappingProxyType
from typing import Final
+import litellm
from litellm._logging import verbose_logger
+from litellm.exceptions import UnsupportedParamsError
from litellm.utils import supports_function_calling
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
-FUNCTION_CALLING_ONLY_PARAMS: Final = ("tools", "tool_choice", "function_call", "response_format")
+TOOL_CALLING_PARAMS: Final = ("tools", "tool_choice", "function_call")
PLAIN_TEXT_RESPONSE_FORMAT: Final = MappingProxyType({"type": "text"})
+FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling"
+
+
+def _function_calling_verdict(model: str) -> bool | None:
+ try:
+ if supports_function_calling(model, custom_llm_provider="together_ai"):
+ return True
+ except Exception as e:
+ verbose_logger.debug("Error checking together_ai function calling support for %s: %s", model, e)
+ registry_entry: Final = litellm.model_cost.get(f"together_ai/{model}")
+ if isinstance(registry_entry, dict) and registry_entry.get("supports_function_calling") is False:
+ return False
+ return None
+
+
+def _tool_params_to_drop(passed_params: Container[str], model: str, drop_params: bool) -> tuple[str, ...]:
+ passed_tool_params: Final = tuple(param for param in TOOL_CALLING_PARAMS if param in passed_params)
+ if not passed_tool_params:
+ return ()
+ verdict: Final = _function_calling_verdict(model)
+ if verdict is True:
+ return ()
+ if verdict is None:
+ verbose_logger.warning(
+ "together_ai model %s has no function calling entry in the model registry; passing %s through for Together to validate. Docs - %s",
+ model,
+ ", ".join(passed_tool_params),
+ FUNCTION_CALLING_DOCS_URL,
+ )
+ return ()
+ if drop_params or litellm.drop_params:
+ verbose_logger.warning(
+ "together_ai model %s does not support function calling per the model registry; dropping %s. Docs - %s",
+ model,
+ ", ".join(passed_tool_params),
+ FUNCTION_CALLING_DOCS_URL,
+ )
+ return passed_tool_params
+ raise UnsupportedParamsError(
+ status_code=500,
+ message=f"together_ai does not support parameters: {', '.join(passed_tool_params)}, for model={model}. To drop it from the call, set `litellm.drop_params = True`.",
+ )
class TogetherAIChatConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list:
- supports_fc: bool | None = None
- try:
- supports_fc = supports_function_calling(model, custom_llm_provider="together_ai")
- except Exception as e:
- verbose_logger.debug("Error getting supported openai params: %s", e)
-
+ supports_fc: Final = _function_calling_verdict(model)
supported_params: Final = super().get_supported_openai_params(model)
if supports_fc is True:
return supported_params
verbose_logger.debug(
- "Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling"
+ "Only some together models support response_format. Docs - https://docs.together.ai/docs/function-calling"
)
return [ # mutable-ok: the inherited contract returns a plain list; building fresh avoids mutating the base class's value
- param for param in supported_params if param not in FUNCTION_CALLING_ONLY_PARAMS
+ param for param in supported_params if param != "response_format"
]
def map_openai_params(
@@ -42,7 +82,8 @@ class TogetherAIChatConfig(OpenAIGPTConfig):
drop_params: bool,
) -> dict:
mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params)
-
+ for param in _tool_params_to_drop(mapped_openai_params, model, drop_params):
+ mapped_openai_params.pop(param)
if mapped_openai_params.get("response_format") == PLAIN_TEXT_RESPONSE_FORMAT:
mapped_openai_params.pop("response_format")
return mapped_openai_params
diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py
index 6216d3bf225..0b9fd5364f9 100644
--- a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py
+++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py
@@ -1,10 +1,12 @@
import json
+import logging
from unittest.mock import MagicMock
import httpx
import pytest
import litellm
+from litellm.exceptions import UnsupportedParamsError
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.llms.openai.chat.gpt_transformation import (
OpenAIChatCompletionStreamingHandler,
@@ -14,10 +16,12 @@ from litellm.types.utils import LlmProviders, ModelResponse
TOOL_CALLING_MODEL = "openai/gpt-oss-20b"
REASONING_MODEL = "deepseek-ai/DeepSeek-V3.1"
-PLAIN_MODEL = "Qwen/Qwen3-235B-A22B-fp8-tput"
-UNMAPPED_MODEL = "MiniMaxAI/MiniMax-M3"
+UNMAPPED_MODEL = "example-org/brand-new-model"
+NO_TOOLS_MODEL = "example-org/no-tools-model"
-FUNCTION_CALLING_PARAMS = ("tools", "tool_choice", "function_call", "response_format")
+TOOL_PARAMS = ("tools", "tool_choice", "function_call")
+
+WEATHER_TOOLS = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}]
@pytest.fixture(autouse=True)
@@ -28,44 +32,103 @@ def force_local_model_cost(monkeypatch):
monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url))
+@pytest.fixture
+def registry_disables_function_calling(monkeypatch):
+ monkeypatch.setitem(
+ litellm.model_cost,
+ f"together_ai/{NO_TOOLS_MODEL}",
+ {"litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": False},
+ )
+
+
+@pytest.fixture
+def together_warning_log(caplog):
+ from litellm._logging import verbose_logger
+
+ verbose_logger.addHandler(caplog.handler)
+ with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+ yield caplog
+ verbose_logger.removeHandler(caplog.handler)
+
+
def test_supported_params_tool_calling_model():
supported = TogetherAIChatConfig().get_supported_openai_params(model=TOOL_CALLING_MODEL)
- for param in FUNCTION_CALLING_PARAMS:
+ for param in (*TOOL_PARAMS, "response_format"):
assert param in supported
-def test_supported_params_plain_model():
- supported = TogetherAIChatConfig().get_supported_openai_params(model=PLAIN_MODEL)
-
- for param in FUNCTION_CALLING_PARAMS:
- assert param not in supported
- assert "temperature" in supported
- assert "max_tokens" in supported
-
-
-def test_supported_params_unmapped_model_treated_as_plain():
+def test_supported_params_unmapped_model_keeps_tool_params():
supported = TogetherAIChatConfig().get_supported_openai_params(model=UNMAPPED_MODEL)
- for param in FUNCTION_CALLING_PARAMS:
- assert param not in supported
+ for param in TOOL_PARAMS:
+ assert param in supported
+ assert "response_format" not in supported
assert "stream" in supported
+ assert "temperature" in supported
+
+
+def test_supported_params_no_tools_model_keeps_tool_params(registry_disables_function_calling):
+ supported = TogetherAIChatConfig().get_supported_openai_params(model=NO_TOOLS_MODEL)
+
+ for param in TOOL_PARAMS:
+ assert param in supported
+ assert "response_format" not in supported
def test_map_openai_params_tool_calling_model_passes_tools():
- tools = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}]
-
mapped = TogetherAIChatConfig().map_openai_params(
- non_default_params={"tools": tools, "tool_choice": "auto"},
+ non_default_params={"tools": WEATHER_TOOLS, "tool_choice": "auto"},
optional_params={},
model=TOOL_CALLING_MODEL,
drop_params=False,
)
- assert mapped["tools"] == tools
+ assert mapped["tools"] == WEATHER_TOOLS
assert mapped["tool_choice"] == "auto"
+@pytest.mark.parametrize("drop_params", [False, True])
+def test_map_openai_params_unmapped_model_passes_tools_through(drop_params, together_warning_log):
+ mapped = TogetherAIChatConfig().map_openai_params(
+ non_default_params={"tools": WEATHER_TOOLS, "tool_choice": "required"},
+ optional_params={},
+ model=UNMAPPED_MODEL,
+ drop_params=drop_params,
+ )
+
+ assert mapped["tools"] == WEATHER_TOOLS
+ assert mapped["tool_choice"] == "required"
+ assert UNMAPPED_MODEL in together_warning_log.text
+ assert "passing tools, tool_choice through" in together_warning_log.text
+
+
+def test_map_openai_params_no_tools_model_drops_tools_with_warning(
+ registry_disables_function_calling, together_warning_log
+):
+ mapped = TogetherAIChatConfig().map_openai_params(
+ non_default_params={"tools": WEATHER_TOOLS, "temperature": 0.5},
+ optional_params={},
+ model=NO_TOOLS_MODEL,
+ drop_params=True,
+ )
+
+ assert "tools" not in mapped
+ assert mapped["temperature"] == 0.5
+ assert NO_TOOLS_MODEL in together_warning_log.text
+ assert "dropping tools" in together_warning_log.text
+
+
+def test_map_openai_params_no_tools_model_raises_without_drop_params(registry_disables_function_calling):
+ with pytest.raises(UnsupportedParamsError, match="does not support parameters"):
+ TogetherAIChatConfig().map_openai_params(
+ non_default_params={"tools": WEATHER_TOOLS},
+ optional_params={},
+ model=NO_TOOLS_MODEL,
+ drop_params=False,
+ )
+
+
def test_map_openai_params_reasoning_model_passes_sampling_params():
mapped = TogetherAIChatConfig().map_openai_params(
non_default_params={"temperature": 0.2, "max_tokens": 512},
@@ -170,6 +233,41 @@ def test_streaming_chunk_maps_delta_reasoning_to_reasoning_content():
assert parsed.choices[0]["delta"]["reasoning_content"] == "thinking about 2+2"
+def test_streaming_chunk_preserves_tool_call_index_and_id():
+ iterator = TogetherAIChatConfig().get_model_response_iterator(
+ streaming_response=iter(()), sync_stream=True
+ )
+
+ def parse_tool_call_chunk(tool_call: dict):
+ parsed = iterator.chunk_parser(
+ {
+ "id": "chunk-1",
+ "created": 1234567890,
+ "model": TOOL_CALLING_MODEL,
+ "choices": [{"index": 0, "delta": {"role": "assistant", "content": "", "tool_calls": [tool_call]}}],
+ }
+ )
+ return parsed.choices[0]["delta"]["tool_calls"][0]
+
+ opener = parse_tool_call_chunk(
+ {
+ "index": 1,
+ "id": "call_abc123",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": ""},
+ }
+ )
+ continuation = parse_tool_call_chunk(
+ {"index": 1, "id": "", "type": "function", "function": {"arguments": '{"city": "San'}}
+ )
+
+ assert opener["index"] == 1
+ assert opener["id"] == "call_abc123"
+ assert opener["function"]["name"] == "get_weather"
+ assert continuation["index"] == 1
+ assert continuation["function"]["arguments"] == '{"city": "San'
+
+
def test_together_ai_config_alias_points_at_chat_config():
assert litellm.TogetherAIConfig is litellm.TogetherAIChatConfig
config = litellm.TogetherAIConfig(max_tokens=10)
@@ -230,3 +328,60 @@ def test_completion_routes_through_together_chat_config():
assert json.loads(request.content)["model"] == REASONING_MODEL
assert response.choices[0].message.content == "4"
assert response.choices[0].message.reasoning_content == "2+2 equals 4"
+
+
+def test_completion_unmapped_model_sends_tools_to_together():
+ from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+ captured_requests = []
+
+ def respond(request: httpx.Request) -> httpx.Response:
+ captured_requests.append(request)
+ return httpx.Response(
+ 200,
+ json={
+ "id": "chatcmpl-together-tools",
+ "object": "chat.completion",
+ "created": 1234567890,
+ "model": UNMAPPED_MODEL,
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {
+ "id": "call_abc123",
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "arguments": '{"city": "San Francisco"}',
+ },
+ }
+ ],
+ },
+ "finish_reason": "tool_calls",
+ }
+ ],
+ "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+ },
+ )
+
+ client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
+
+ response = litellm.completion(
+ model=f"together_ai/{UNMAPPED_MODEL}",
+ messages=[{"role": "user", "content": "What is the weather in San Francisco?"}],
+ tools=WEATHER_TOOLS,
+ tool_choice="auto",
+ api_key="fake-key",
+ client=client,
+ )
+
+ request_body = json.loads(captured_requests[0].content)
+ assert request_body["tools"] == WEATHER_TOOLS
+ assert request_body["tool_choice"] == "auto"
+ tool_call = response.choices[0].message.tool_calls[0]
+ assert tool_call.function.name == "get_weather"
+ assert json.loads(tool_call.function.arguments) == {"city": "San Francisco"}
From 99c1b33fd32a21bf764bf3fa32f28d3e0062d90c Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 14:50:14 -0700
Subject: [PATCH 65/70] chore: humanize the CLAUDE.md
---
CLAUDE.md | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/CLAUDE.md b/CLAUDE.md
index 03053b8392c..d3aa552a120 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -41,10 +41,10 @@ If you ever make public-facing PR descriptions, comments, issues, commit message
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
-- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
+- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
-- do use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When structure genuinely helps the reader, prefer nested bullets (any depth is fine) over one dense line. This applies to all human-facing text: discussion posts, release notes, and docs included
+- do use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure genuinely helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure. This applies to all human-facing text: discussion posts, release notes, and docs included
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
From d61ccd0f583044a8cbc5fdd9b7284088e246d535 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 14:51:53 -0700
Subject: [PATCH 66/70] chore: make it more clear
---
CLAUDE.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/CLAUDE.md b/CLAUDE.md
index d3aa552a120..f9d0bcdc7ba 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -44,7 +44,7 @@ If you ever make public-facing PR descriptions, comments, issues, commit message
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
-- do use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure genuinely helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure. This applies to all human-facing text: discussion posts, release notes, and docs included
+- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure. This applies to all human-facing text: pull requests, issues, commit messages, discussion posts, release notes, docs, etc.
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
From a40157c55def33d31aaa6dfc199b15cb89833b3a Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 14:53:24 -0700
Subject: [PATCH 67/70] chore: make it more concise
---
CLAUDE.md | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/CLAUDE.md b/CLAUDE.md
index f9d0bcdc7ba..6f36c2f9ecb 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -37,14 +37,14 @@ If you're resolving a linear ticket, in the "## Linear ticket" section of the PR
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
-If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
+If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
-- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure. This applies to all human-facing text: pull requests, issues, commit messages, discussion posts, release notes, docs, etc.
+- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
From 6c0c91c5adec0f26b126321cfba44f66d318082a Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Tue, 25 Aug 2026 14:55:11 -0700
Subject: [PATCH 68/70] fix(team): serialize member_add, member_delete, and
delete under the team's advisory lock (#37969)
* fix(proxy): make /team/member_delete's four cleanups atomic
The team roster update, the user.teams update, the team membership
delete, and the team-scoped verification token delete ran as four
sequential writes with no transaction around them, so a failure
between any two left the removal half applied. Thread a single
prisma transaction through all four writes, following the same
tx. pattern /team/member_add and /team/member_update already
use, so either all four land or none do.
* fix(team): serialize member_add, member_delete, and delete under the team's advisory lock
/team/member_add validated a team exists and then wrote the user's teams array and
a membership row without holding anything across that gap, so a /team/delete could
commit its reference sweeps in between and leave a member pointing at a team id that
no longer exists. The write path already re-read members_with_roles under a row lock
before this change, but SELECT ... FOR UPDATE can deadlock with the access-group
endpoints, which lock an access group and then a team.
member_add now takes pg_advisory_xact_lock(hashtext(team_id)) before re-reading the
team and only writes if it is still there, so a delete that already committed is
visible before any write happens. delete_team takes the same lock around its own
row delete and reference sweep, so the two requests can never interleave: whichever
acquires the lock first runs to completion before the other's read can proceed.
Dropping the row lock from member_add's read also dropped the incidental protection
it gave against a concurrent member_delete, which still wrote from the snapshot it
validated against, unlocked, and could silently overwrite whatever member_add had
just committed. member_delete now takes the same advisory lock and re-reads the
roster under it before computing its own write, so it can never resurrect a member
by overwriting from stale data.
Resolves LIT-5544
* fix(team): run member writes on the advisory lock's transaction
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(team): keep member writes on the lock holder's connection after merge
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(team): keep the transactional member create an upsert on user_id
The transaction path was creating the email-identified user row outright, where the
regular client path upserts on user_id. Share one upsert helper between both member
paths so the create stays idempotent on the lock holder's connection.
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(team): read member_delete's user and key rows on the lock-holding transaction
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../management_endpoints/team_endpoints.py | 258 +++++++++------
.../access_group_team_sync.py | 10 +-
litellm/proxy/management_helpers/utils.py | 128 +++++--
litellm/repositories/team_repository.py | 21 +-
.../test_team_delete_member_add_race.py | 307 +++++++++++++++++
tests/proxy_unit_tests/test_proxy_server.py | 34 +-
.../test_team_endpoints.py | 311 +++++++++++++++---
.../test_management_helpers_utils.py | 78 +++++
.../repositories/test_repositories.py | 12 +-
9 files changed, 949 insertions(+), 210 deletions(-)
create mode 100644 tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 49461d7841d..ba4fb81323f 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -16,11 +16,12 @@ import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
-from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypedDict, TypeVar, cast
+from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, JsonValue
+from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
@@ -116,6 +117,7 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
+ TEAM_ADVISORY_LOCK_SQL,
AccessGroupSyncTx,
invalidate_access_group_caches,
reconcile_team_access_group_membership,
@@ -134,6 +136,7 @@ from litellm.proxy.management_helpers.team_metadata_validation import (
validate_team_metadata_if_configured,
)
from litellm.proxy.management_helpers.utils import (
+ MemberWriteTx,
add_new_member,
management_endpoint_wrapper,
)
@@ -330,11 +333,44 @@ class _TeamIdInFilter(TypedDict, total=False):
team_id: Mapping[str, Sequence[str]]
+class _DeletedTeamsResult(TypedDict):
+ deleted_teams: ReadOnly[Sequence[str]]
+
+
+class _ErrorDetail(TypedDict):
+ error: ReadOnly[str]
+
+
class _TeamCreateTx(AccessGroupSyncTx, Protocol):
@property
def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
+class _MemberDeleteTx(Protocol):
+ """The tables `/team/member_delete` reads while it holds the team's advisory lock.
+
+ Reading them off the transaction keeps the whole endpoint on the one pooled connection
+ it already checked out: a request that has the lock but still needs another connection
+ can be starved by the lock waiters, which is a deadlock rather than a wait when enough
+ of them hold the rest of the pool."""
+
+ @property
+ def litellm_usertable(self) -> "_PrismaTableActions[LiteLLM_UserTable]": ...
+
+ @property
+ def litellm_verificationtoken(self) -> "_PrismaTableActions[LiteLLM_VerificationToken]": ...
+
+
+class _TeamDeleteTx(AccessGroupSyncTx, Protocol):
+ async def execute_raw(self, query: str, *args: object) -> int: ...
+
+ @property
+ def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
+
+ @property
+ def litellm_teammembership(self) -> "_PrismaTableActions[LiteLLM_TeamMembership]": ...
+
+
_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """
UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams)
"""
@@ -2580,8 +2616,13 @@ async def _process_team_members(
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
+ tx: MemberWriteTx | None = None,
) -> tuple[list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
- """Process and add new team members."""
+ """Process and add new team members.
+
+ ``tx`` is the caller's open transaction, when it has one, so the member writes run on the
+ connection it already holds instead of checking out a second one.
+ """
updated_users: Final[list[LiteLLM_UserTable]] = []
updated_team_memberships: Final[list[LiteLLM_TeamMembership]] = []
@@ -2607,6 +2648,7 @@ async def _process_team_members(
default_team_budget_id=default_team_budget_id,
allowed_models=member_allowed_models,
budget_duration=data.budget_duration,
+ tx=tx,
)
except Exception as e:
raise HTTPException(
@@ -2629,6 +2671,7 @@ async def _process_team_members(
default_team_budget_id=default_team_budget_id,
allowed_models=member_allowed_models,
budget_duration=data.budget_duration,
+ tx=tx,
)
except Exception as e:
raise HTTPException(
@@ -2708,65 +2751,40 @@ async def _add_team_members_to_team(
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
) -> tuple[LiteLLM_TeamTable, list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
- """Add team members to the team.
+ """Add team members to the team, under the team's advisory lock.
- The members_with_roles reconciliation runs inside a transaction that locks
- the team row with ``SELECT ... FOR UPDATE`` before reading the current
- membership. Concurrent /team/member_add calls for the same team therefore
- serialize on the row lock and each appends onto the other's committed
- result, instead of both rewriting the whole JSON array from a stale
- snapshot (which silently drops one member on the losing write).
+ The lock (``TEAM_ADVISORY_LOCK_SQL``, keyed on the team id) is taken first, and the
+ team is re-read under it before any write, so a delete that already committed is
+ visible here before this call writes anything: the user and membership writes only
+ happen once the re-read proves the team is still live. /team/delete takes the same
+ lock around its own sweep-and-delete, so the two can never interleave; whichever
+ acquires the lock first runs to completion before the other's re-read can proceed.
- The same lock serializes this against /team/delete: the delete cannot remove
- the row while the reconcile holds it, and a reconcile that finds the row
- already gone cleans up after itself rather than leaving the member pointing
- at a deleted team id.
- """
- # Process and add new members
- updated_users, updated_team_memberships = await _process_team_members(
- data=data,
- complete_team_data=complete_team_data,
- prisma_client=prisma_client,
- user_api_key_dict=user_api_key_dict,
- litellm_proxy_admin_name=litellm_proxy_admin_name,
- )
-
- updated_team: Final = await _write_members_with_roles_locked(
- data=data,
- complete_team_data=complete_team_data,
- prisma_client=prisma_client,
- updated_users=updated_users,
- )
- if updated_team is None:
- await _sweep_deleted_team_references(team_ids=(data.team_id,), prisma_client=prisma_client)
- raise HTTPException(
- status_code=404,
- detail={"error": f"Team={data.team_id} was deleted while this member add was running"},
- )
-
- return updated_team, updated_users, updated_team_memberships
-
-
-async def _write_members_with_roles_locked(
- data: TeamMemberAddRequest,
- complete_team_data: LiteLLM_TeamTable,
- prisma_client: PrismaClient,
- updated_users: list[LiteLLM_UserTable],
-) -> LiteLLM_TeamTable | None:
- """Reconcile members_with_roles under the team row lock. None when the team row is gone.
-
- That read is at least as recent as the user and membership writes the caller
- already made, so a missing row means /team/delete committed after them. Its
- post-delete sweep can have run before those writes landed, which is why the
- caller sweeps this team id again rather than only reporting the 404.
+ The user and membership writes run on this transaction too, not on a second
+ connection from the pool: a lock waiter that needs a connection it hasn't got yet is
+ a waiter that can deadlock the pool, since enough concurrent adds for one team would
+ hold every connection waiting on the lock while the holder waits for a free one.
"""
async with prisma_client.tx() as tx:
+ await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, data.team_id)
+
locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id)
if locked_members is None:
- return None
-
+ gone_detail: Final[_ErrorDetail] = {
+ "error": f"Team={data.team_id} was deleted while this member add was running"
+ }
+ raise HTTPException(status_code=404, detail=gone_detail)
complete_team_data.members_with_roles = locked_members
+ updated_users, updated_team_memberships = await _process_team_members(
+ data=data,
+ complete_team_data=complete_team_data,
+ prisma_client=prisma_client,
+ user_api_key_dict=user_api_key_dict,
+ litellm_proxy_admin_name=litellm_proxy_admin_name,
+ tx=tx,
+ )
+
await _update_team_members_list(
data=data,
complete_team_data=complete_team_data,
@@ -2774,11 +2792,13 @@ async def _write_members_with_roles_locked(
)
_db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles]
- return await tx.litellm_teamtable.update(
+ updated_team: Final = await tx.litellm_teamtable.update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_team_members)},
)
+ return updated_team, updated_users, updated_team_memberships
+
def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None:
"""Update the Prometheus team members gauge after a membership change.
@@ -3159,10 +3179,6 @@ async def team_member_add(
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
- # Check if updated_team is None
- if updated_team is None:
- raise HTTPException(status_code=404, detail={"error": f"Team with id {data.team_id} not found"})
-
_emit_team_members_metric(complete_team_data)
await _create_team_member_add_audit_logs(
@@ -3276,45 +3292,63 @@ async def team_member_delete(
)
## DELETE MEMBER FROM TEAM
- removed_team_members, new_team_members = _cleanup_members_with_roles(
- existing_team_row=existing_team_row,
- data=data,
- )
-
- if not removed_team_members:
- raise HTTPException(status_code=400, detail={"error": "User not found in team"})
-
- existing_team_row.members_with_roles = new_team_members
-
- _db_new_team_members: Final[list[dict]] = [m.model_dump() for m in new_team_members]
-
- ## DELETE TEAM ID from USER ROW, IF EXISTS ##
- # get user row
- removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None)
- key_val: Final[Mapping[str, object]] = (
- {"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
- )
- existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(where=key_val)
-
- # Also clean up any existing team membership rows for this user and team
- user_ids_to_delete: Final = removed_user_ids.union(
- (data.user_id,) if data.user_id is not None else (),
- (user.user_id for user in existing_user_rows if user.user_id),
- )
-
- ## DELETE KEYS CREATED BY USER FOR THIS TEAM
- # Fetch keys before deletion so their audit records can be persisted alongside the delete.
- # An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows.
- keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many(
- where={
- "user_id": {"in": sorted(user_ids_to_delete)},
- "team_id": data.team_id,
- }
- )
-
- # All four cleanups run on one connection so a failure between them leaves
- # no partial removal: either every write below lands, or none of them do.
+ # Everything from here on runs under the team's advisory lock, the same one
+ # /team/member_add and /team/delete take: without it, this endpoint's own row-level
+ # update lock used to be the only thing serializing it against a concurrent member_add,
+ # and only by accident (their SELECT ... FOR UPDATE contended for the same row lock this
+ # UPDATE takes). Now that member_add reads under the advisory lock instead, this has to
+ # take it too, and re-read the roster under it rather than off the snapshot validated
+ # above, or a member_add that commits in between can have its addition silently
+ # overwritten by this delete computing from stale data.
async with prisma_client.tx() as tx:
+ await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, data.team_id)
+
+ fresh_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id)
+ if fresh_members is None:
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"Team id={data.team_id} does not exist in db"},
+ )
+
+ removed_team_members, new_team_members = _cleanup_members_with_roles(
+ existing_team_row=LiteLLM_TeamTable(team_id=data.team_id, members_with_roles=fresh_members),
+ data=data,
+ )
+
+ if not removed_team_members:
+ raise HTTPException(status_code=400, detail={"error": "User not found in team"})
+
+ existing_team_row.members_with_roles = new_team_members
+
+ _db_new_team_members: Final[list[dict]] = [m.model_dump() for m in new_team_members]
+
+ ## DELETE TEAM ID from USER ROW, IF EXISTS ##
+ # get user row
+ removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None)
+ key_val: Final[Mapping[str, object]] = (
+ {"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
+ )
+ member_tx: Final[_MemberDeleteTx] = tx
+ existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await member_tx.litellm_usertable.find_many(
+ where=key_val
+ )
+
+ # Also clean up any existing team membership rows for this user and team
+ user_ids_to_delete: Final = removed_user_ids.union(
+ (data.user_id,) if data.user_id is not None else (),
+ (user.user_id for user in existing_user_rows if user.user_id),
+ )
+
+ ## DELETE KEYS CREATED BY USER FOR THIS TEAM
+ # Fetch keys before deletion so their audit records can be persisted alongside the delete.
+ # An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows.
+ keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await member_tx.litellm_verificationtoken.find_many(
+ where={
+ "user_id": {"in": sorted(user_ids_to_delete)},
+ "team_id": data.team_id,
+ }
+ )
+
await tx.litellm_teamtable.update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_new_team_members)},
@@ -4009,7 +4043,21 @@ async def delete_team(
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
## DELETE TEAMS
- deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team")
+ # Both the delete and the reconcile sweep run under every team's advisory lock
+ # (TEAM_ADVISORY_LOCK_SQL, the same one /team/member_add takes before its own writes),
+ # sorted so two overlapping batch deletes always request their locks in the same order.
+ # A member_add mid-flight for one of these teams either finishes its write and releases
+ # the lock before this transaction starts, in which case this sweep reaches what it wrote,
+ # or is still waiting on the lock, in which case its own re-read happens after this commits
+ # and sees the row gone before it writes anything.
+ delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}}
+ async with prisma_client.tx() as tx:
+ for team_id in sorted(data.team_ids):
+ await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
+ await tx.litellm_teamtable.delete_many(where=delete_filter)
+ await _sweep_deleted_team_references_tx(team_ids=data.team_ids, tx=tx)
+
+ deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids}
# Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and
# `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a
@@ -4022,12 +4070,6 @@ async def delete_team(
proxy_logging_obj=proxy_logging_obj,
)
- # Sweep again now the team is gone. A `/team/member_add` that landed between the first sweep
- # and the delete would have re-appended the reference; an add still in flight sees the row
- # missing under its own row lock and sweeps what it wrote. Both passes are idempotent, and
- # keeping the first one means a failure here still leaves a team the admin can retry deleting.
- await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
-
for deleted_team in team_rows:
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
@@ -4056,6 +4098,16 @@ async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client:
_ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)}))
+async def _sweep_deleted_team_references_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> None:
+ """Same sweep as `_sweep_deleted_team_references`, run on the transaction that holds
+ every id's advisory lock and deletes the team rows, so it commits or rolls back with them."""
+ for team_id in team_ids:
+ _ = await tx.execute_raw(_STRIP_DELETED_TEAM_FROM_USERS_SQL, team_id)
+
+ membership_filter: Final[_TeamIdInFilter] = {"team_id": {"in": tuple(team_ids)}}
+ _ = await tx.litellm_teammembership.delete_many(where=membership_filter)
+
+
async def _invalidate_deleted_key_cache(
keys: Sequence[LiteLLM_VerificationToken],
user_api_key_cache: UserApiKeyCache,
diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py
index 55c0346e375..664e36c9f10 100644
--- a/litellm/proxy/management_helpers/access_group_team_sync.py
+++ b/litellm/proxy/management_helpers/access_group_team_sync.py
@@ -23,9 +23,11 @@ from pydantic import BaseModel, TypeAdapter
from litellm.proxy.auth.auth_checks import _delete_cache_access_object
# hashtext collisions only cost two unrelated teams a little serialization, and the
-# lock is never taken by the access-group endpoints, so it cannot join their
-# access-group-then-team lock order to form a cycle.
-_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
+# lock is never taken by the access-group endpoints as a SELECT ... FOR UPDATE row lock,
+# so it cannot join their access-group-then-team lock order to form a cycle. team_endpoints
+# reuses this exact statement to serialize /team/member_add and /team/delete against each
+# other and against this mirror, rather than defining a second, divergent lock on the same key.
+TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1'
@@ -138,7 +140,7 @@ async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id:
concurrent write for a different team cannot be lost the way a read-modify-write of
the whole array can, and the pair commits together or not at all.
"""
- await tx.query_raw(_LOCK_TEAM_SQL, team_id)
+ await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id))
desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else ()
affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired))
diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py
index cb30ce90c7f..e2d7262fb69 100644
--- a/litellm/proxy/management_helpers/utils.py
+++ b/litellm/proxy/management_helpers/utils.py
@@ -34,7 +34,7 @@ from litellm.proxy._types import ( # key request types; user request types; tea
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
-from litellm.proxy.utils import PrismaClient
+from litellm.proxy.utils import PrismaClient, jsonify_object
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.table_repositories import TeamMembershipRepository
from litellm.repositories.user_repository import UserRepository
@@ -79,6 +79,8 @@ class _PrismaUserTable(Protocol):
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
) -> _PrismaUserRecord | None: ...
+ async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_PrismaUserRecord]: ...
+
class _PrismaTeamMembershipTable(Protocol):
"""Team membership table actions the management helpers issue."""
@@ -86,6 +88,73 @@ class _PrismaTeamMembershipTable(Protocol):
async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ...
+class MemberWriteTx(Protocol):
+ """Transaction surface `add_new_member` writes through when the caller owns one.
+
+ A caller already holding a transaction, and with it a pooled connection plus that
+ transaction's locks, passes it here so these writes reuse that connection rather than
+ checking out another one that lock waiters may already have drained from the pool.
+ """
+
+ @property
+ def litellm_usertable(self) -> _PrismaUserTable: ...
+
+ @property
+ def litellm_budgettable(self) -> _PrismaBudgetTable: ...
+
+ @property
+ def litellm_teammembership(self) -> _PrismaTeamMembershipTable: ...
+
+
+def _user_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaUserTable:
+ return tx.litellm_usertable if tx is not None else UserRepository(prisma_client).table
+
+
+def _budget_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaBudgetTable:
+ return tx.litellm_budgettable if tx is not None else BudgetRepository(prisma_client).table
+
+
+def _team_membership_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaTeamMembershipTable:
+ return tx.litellm_teammembership if tx is not None else TeamMembershipRepository(prisma_client).table
+
+
+async def _find_users_by_email(
+ prisma_client: PrismaClient, tx: MemberWriteTx | None, user_email: str
+) -> Sequence[_PrismaUserRecord]:
+ if tx is not None:
+ return await tx.litellm_usertable.find_many(where={"user_email": user_email})
+ rows: Final[Sequence[_PrismaUserRecord] | None] = await prisma_client.get_data(
+ key_val={"user_email": user_email},
+ table_name="user",
+ query_type="find_all",
+ )
+ return rows if rows is not None else ()
+
+
+async def _upsert_user_row(
+ user_table: _PrismaUserTable, user_id: str, create_data: Mapping[str, object]
+) -> _PrismaUserRecord | None:
+ """Insert the user row if it is absent, leaving an existing row as it is.
+
+ Upserting keeps concurrent provisioning of the same new user from racing on create.
+ The update branch re-states user_id rather than being empty because Prisma only
+ compiles an upsert down to INSERT ... ON CONFLICT when the update is non-empty, and
+ otherwise falls back to a racy SELECT-then-INSERT.
+ """
+ return await user_table.upsert(
+ where={"user_id": user_id},
+ data={"create": create_data, "update": {"user_id": user_id}},
+ )
+
+
+async def _create_user_row(
+ prisma_client: PrismaClient, tx: MemberWriteTx | None, user_data: dict[str, object]
+) -> _PrismaUserRecord | None:
+ if tx is not None:
+ return await _upsert_user_row(tx.litellm_usertable, str(user_data["user_id"]), jsonify_object(user_data))
+ return await prisma_client.insert_data(data=user_data, table_name="user")
+
+
def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]:
user_info: Final = litellm.default_internal_user_params or {}
@@ -206,6 +275,7 @@ async def _clone_team_default_budget_for_member(
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
budget_duration_override: str | None = None,
+ tx: MemberWriteTx | None = None,
) -> str | None:
"""
Create a new budget row that copies the values from the team's default
@@ -220,7 +290,7 @@ async def _clone_team_default_budget_for_member(
member while keeping the default's other limits, so an admin can set a
member's reset cadence without discarding the team default's max_budget.
"""
- budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
+ budget_table: Final[_PrismaBudgetTable] = _budget_table(prisma_client, tx)
default_budget: Final = await budget_table.find_unique(where={"budget_id": default_team_budget_id})
if default_budget is None:
return None
@@ -248,7 +318,7 @@ async def _clone_team_default_budget_for_member(
if cloned_data.get("budget_duration"):
cloned_data["budget_reset_at"] = get_budget_reset_time(cloned_data["budget_duration"])
- new_budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(data=cloned_data)
+ new_budget: Final[_PrismaBudgetRecord] = await budget_table.create(data=cloned_data)
return new_budget.budget_id
@@ -260,6 +330,7 @@ async def _resolve_member_budget_id(
allowed_models: list[str] | None,
budget_duration: str | None,
default_team_budget_id: str | None,
+ tx: MemberWriteTx | None = None,
) -> str | None:
"""
Resolve the budget a new team member should be linked to.
@@ -279,6 +350,7 @@ async def _resolve_member_budget_id(
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
budget_duration_override=budget_duration,
+ tx=tx,
)
if not has_explicit_limit and budget_duration is None:
@@ -295,12 +367,14 @@ async def _resolve_member_budget_id(
if budget_duration is not None:
budget_data["budget_duration"] = budget_duration
budget_data["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration)
- budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
+ budget_table: Final[_PrismaBudgetTable] = _budget_table(prisma_client, tx)
response: Final = await budget_table.create(data=budget_data)
return response.budget_id
-async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, team_id: str) -> None:
+async def _append_team_id_if_absent(
+ prisma_client: PrismaClient, user_id: str, team_id: str, tx: MemberWriteTx | None = None
+) -> None:
"""Append team_id to a user's teams array, only if it is not already present.
The row-level filter makes the append a no-op once the team is present, so
@@ -309,7 +383,7 @@ async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, t
number of teams a user belongs to). Teams added concurrently for a different
team id are unaffected, since each update filters on its own team id.
"""
- user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
+ user_table: Final[_PrismaUserTable] = _user_table(prisma_client, tx)
await user_table.update_many(
where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}},
data={"teams": {"push": [team_id]}},
@@ -326,6 +400,7 @@ async def add_new_member(
default_team_budget_id: str | None = None,
allowed_models: list[str] | None = None,
budget_duration: str | None = None,
+ tx: MemberWriteTx | None = None,
) -> tuple[LiteLLM_UserTable, LiteLLM_TeamMembership | None]:
"""
Add a new member to a team
@@ -334,49 +409,41 @@ async def add_new_member(
- add team member w/ budget to team member table
Returns created/existing user + team membership w/ budget id
+
+ Callers already inside a transaction pass it as ``tx`` so every write here runs on that
+ connection instead of borrowing more from the pool while the caller's locks are held.
"""
returned_user: LiteLLM_UserTable | None = None
returned_team_membership: LiteLLM_TeamMembership | None = None
## ADD TEAM ID, to USER TABLE IF NEW ##
if new_member.user_id is not None:
new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id)
- # Upsert ensures the user row exists atomically (no create race when the
- # same new user is provisioned concurrently), seeding teams on create.
- # The teams append lives in the filtered update below rather than the
- # upsert's update branch so an already-existing user does not get a
- # duplicate team id. The update branch still has to write something:
- # Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it
- # is non-empty, and falls back to a racy SELECT-then-INSERT when it is
- # not, so this re-states user_id as a no-op rather than being empty.
- user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
- _returned_user: _PrismaUserRecord | None = await user_table.upsert(
- where={"user_id": new_member.user_id},
- data={
- "create": {"teams": [team_id], **new_user_defaults},
- "update": {"user_id": new_member.user_id},
- },
+ # The teams append lives in the filtered update below rather than the upsert's
+ # update branch so an already-existing user does not get a duplicate team id.
+ _returned_user: _PrismaUserRecord | None = await _upsert_user_row(
+ _user_table(prisma_client, tx),
+ new_member.user_id,
+ {"teams": [team_id], **new_user_defaults},
)
- await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id)
+ await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id, tx)
if _returned_user is not None:
returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif new_member.user_email is not None:
new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email)
## user email is not unique acc. to prisma schema -> future improvement
### for now: check if it exists in db, if not - insert it
- existing_user_row: Final[list[_PrismaUserRecord] | None] = await prisma_client.get_data(
- key_val={"user_email": new_member.user_email},
- table_name="user",
- query_type="find_all",
+ existing_user_row: Final[Sequence[_PrismaUserRecord]] = await _find_users_by_email(
+ prisma_client, tx, new_member.user_email
)
- if existing_user_row is None or (isinstance(existing_user_row, list) and len(existing_user_row) == 0):
+ if len(existing_user_row) == 0:
new_user_defaults["teams"] = [team_id]
- _returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user")
+ _returned_user = await _create_user_row(prisma_client, tx, new_user_defaults)
if _returned_user is not None:
returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif len(existing_user_row) == 1:
user_info: Final = existing_user_row[0]
- await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id)
+ await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id, tx)
returned_user = LiteLLM_UserTable.model_validate(user_info.model_dump())
elif len(existing_user_row) > 1:
raise HTTPException(
@@ -392,10 +459,11 @@ async def add_new_member(
allowed_models=allowed_models,
budget_duration=budget_duration,
default_team_budget_id=default_team_budget_id,
+ tx=tx,
)
if _budget_id and returned_user is not None and returned_user.user_id is not None:
- membership_table: Final[_PrismaTeamMembershipTable] = TeamMembershipRepository(prisma_client).table
+ membership_table: Final[_PrismaTeamMembershipTable] = _team_membership_table(prisma_client, tx)
_returned_team_membership: Final = await membership_table.create(
data={
"team_id": team_id,
diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py
index 7efd32288e4..d636592f925 100644
--- a/litellm/repositories/team_repository.py
+++ b/litellm/repositories/team_repository.py
@@ -58,19 +58,22 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
return LiteLLM_TeamTable.model_validate(data)
async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member] | None:
- """Return the team's members_with_roles, locking the row FOR UPDATE.
+ """Return the team's members_with_roles. The caller must already hold
+ ``TEAM_ADVISORY_LOCK_SQL`` for this team_id on ``tx`` before calling this.
- ``None`` when the team row is gone, which a caller holding the lock can
- only see if a delete committed under it, as opposed to ``[]`` for a team
- that simply has no members.
+ ``None`` when the team row is gone, which is only possible under that lock if
+ a delete committed before this read, as opposed to ``[]`` for a team that
+ simply has no members.
- Must be called inside a transaction so the row lock is held until
- commit. This serializes concurrent membership writers on the team row
- so the losing writer appends onto the winner's committed result instead
- of overwriting it from a stale snapshot.
+ A plain read is enough here because the advisory lock, not a row lock, is what
+ serializes this against a concurrent writer: ``SELECT ... FOR UPDATE`` would
+ additionally take a row lock on ``LiteLLM_TeamTable``, and the access-group
+ endpoints lock an access group and then a team row, so a team-row-first lock
+ here can deadlock with them. The advisory lock cannot, since those endpoints
+ never take it.
"""
rows: Final = await tx.query_raw(
- 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1 FOR UPDATE',
+ 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1',
team_id,
)
if not rows:
diff --git a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py
new file mode 100644
index 00000000000..30544a8bb81
--- /dev/null
+++ b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py
@@ -0,0 +1,307 @@
+"""
+Real-Postgres coverage for the /team/member_add vs /team/delete race (LIT-5544), and for
+/team/member_delete's participation in the same lock.
+
+A member_add that validated the team before a delete began could previously still commit
+its writes after the delete's reference sweeps had already run, leaving a user record and
+a membership row pointing at a team id that no longer exists. Neither side of that race can
+be forced by a sequential script: it needs one request to be genuinely mid-flight while the
+other commits. A mocked prisma cannot arbitrate that either, since the property under test
+is whether Postgres's own advisory lock actually serializes the two requests.
+
+These tests pin the interleaving the same way test_access_group_team_sync.py does: a second
+real connection holds the team's advisory lock in its own transaction, so the function under
+test is provably blocked on it rather than hoping a sleep lands in the right gap.
+"""
+
+import asyncio
+import json
+import os
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from unittest.mock import MagicMock
+
+import pytest
+from fastapi import HTTPException
+
+from litellm.proxy._types import (
+ DeleteTeamRequest,
+ LitellmUserRoles,
+ Member,
+ TeamMemberAddRequest,
+ UserAPIKeyAuth,
+)
+from litellm.caching.caching import DualCache
+from litellm.proxy.utils import PrismaClient, ProxyLogging
+
+TEAM = "lit5544-race-team"
+USER = "lit5544-race-user"
+_DELETE_SEEDED = 'DELETE FROM "LiteLLM_TeamMembership" WHERE team_id = $1'
+_DELETE_USER = 'DELETE FROM "LiteLLM_UserTable" WHERE user_id = $1'
+_DELETE_TEAM = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = $1'
+_LOCK_SQL = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
+
+
+@asynccontextmanager
+async def _clean_db():
+ """Connects inside the running test's loop: an async fixture would be torn up on a
+ different loop than the test body, which prisma's engine lock refuses outright."""
+ from prisma import Prisma
+
+ if not os.getenv("DATABASE_URL"):
+ pytest.fail("DATABASE_URL is required; these tests must not silently skip")
+
+ db = Prisma()
+ await db.connect()
+ try:
+ await db.execute_raw(_DELETE_SEEDED, TEAM)
+ await db.execute_raw(_DELETE_USER, USER)
+ await db.execute_raw(_DELETE_TEAM, TEAM)
+ yield db
+ finally:
+ await db.execute_raw(_DELETE_SEEDED, TEAM)
+ await db.execute_raw(_DELETE_USER, USER)
+ await db.execute_raw(_DELETE_TEAM, TEAM)
+ await db.disconnect()
+
+
+@asynccontextmanager
+async def _real_prisma_client():
+ """The full app-level PrismaClient, not the raw generated client: add_new_member reads
+ and writes through PrismaClient.get_data/insert_data, which the raw client doesn't have."""
+ proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
+ client = PrismaClient(database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj)
+ await client.connect()
+ try:
+ yield client
+ finally:
+ await client.db.disconnect()
+
+
+def _admin_auth():
+ return UserAPIKeyAuth(user_id="lit5544-admin", api_key="sk-lit5544", user_role=LitellmUserRoles.PROXY_ADMIN.value)
+
+
+@pytest.mark.asyncio
+async def test_member_add_blocked_by_delete_writes_no_dangling_reference():
+ """
+ member_add re-reads the team under the advisory lock before writing anything. When a
+ delete already holds that lock and then removes the row, member_add's re-read must see
+ the row gone and raise, without ever calling the write that appends the user/membership
+ references, which is the only way this leaves zero trace after the delete wins.
+ """
+ from litellm.proxy._types import LiteLLM_TeamTable
+ from litellm.proxy.management_endpoints.team_endpoints import (
+ _add_team_members_to_team,
+ )
+
+ async with _clean_db() as db:
+ await db.litellm_teamtable.create(data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": "[]"})
+
+ async with _real_prisma_client() as prisma_client:
+ from prisma import Prisma
+
+ blocker = Prisma()
+ await blocker.connect()
+ lock_acquired = asyncio.Event()
+
+ async def add_member():
+ lock_acquired.set()
+ await _add_team_members_to_team(
+ data=TeamMemberAddRequest(
+ team_id=TEAM,
+ member=Member(user_id=USER, role="user"),
+ max_budget_in_team=5.0,
+ ),
+ complete_team_data=LiteLLM_TeamTable(team_id=TEAM, members_with_roles=[]),
+ prisma_client=prisma_client,
+ user_api_key_dict=_admin_auth(),
+ litellm_proxy_admin_name="lit5544-admin",
+ )
+
+ try:
+ async with blocker.tx(timeout=timedelta(seconds=30)) as held:
+ await held.query_raw(_LOCK_SQL, TEAM)
+ task = asyncio.create_task(add_member())
+ await lock_acquired.wait()
+ await asyncio.sleep(0.2)
+ assert not task.done(), "member_add did not wait on the team's advisory lock"
+
+ # the delete wins the race: strip the team row while the lock is held
+ await held.execute_raw(_DELETE_TEAM, TEAM)
+
+ with pytest.raises(HTTPException) as exc_info:
+ await asyncio.wait_for(task, timeout=30)
+ assert exc_info.value.status_code == 404
+ finally:
+ await blocker.disconnect()
+
+ user_row = await db.litellm_usertable.find_unique(where={"user_id": USER})
+ assert user_row is None, "member_add must not have written a user row for a team that was gone under its lock"
+
+ membership_row = await db.litellm_teammembership.find_first(where={"team_id": TEAM, "user_id": USER})
+ assert membership_row is None
+
+
+@pytest.mark.asyncio
+async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster():
+ """
+ team_member_delete takes the same advisory lock and re-reads the roster under it, so a
+ member_add that committed while member_delete was waiting on the lock is not silently
+ undone. Without the re-read, member_delete would compute its new roster from the stale
+ snapshot it validated against before the lock, and its write would overwrite the
+ member_add's addition right back out even though member_add's request already succeeded.
+ """
+ import litellm.proxy.proxy_server as proxy_server_module
+ from litellm.proxy._types import TeamMemberDeleteRequest
+ from litellm.proxy.management_endpoints.team_endpoints import team_member_delete
+
+ other_user = f"{USER}-other"
+ seeded_roster = '[{"user_id": "%s", "user_email": null, "role": "user"}]' % USER
+ winning_add_roster = (
+ '[{"user_id": "%s", "user_email": null, "role": "user"}, '
+ '{"user_id": "%s", "user_email": null, "role": "user"}]' % (USER, other_user)
+ )
+
+ async with _clean_db() as db:
+ await db.litellm_teamtable.create(
+ data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": seeded_roster}
+ )
+
+ async with _real_prisma_client() as prisma_client:
+ original_prisma_client = proxy_server_module.prisma_client
+ proxy_server_module.prisma_client = prisma_client
+
+ try:
+ from prisma import Prisma
+
+ blocker = Prisma()
+ await blocker.connect()
+ lock_acquired = asyncio.Event()
+
+ async def run_delete():
+ lock_acquired.set()
+ return await team_member_delete(
+ data=TeamMemberDeleteRequest(team_id=TEAM, user_id=USER),
+ user_api_key_dict=_admin_auth(),
+ )
+
+ try:
+ async with blocker.tx(timeout=timedelta(seconds=30)) as held:
+ await held.query_raw(_LOCK_SQL, TEAM)
+ task = asyncio.create_task(run_delete())
+ await lock_acquired.wait()
+ await asyncio.sleep(0.2)
+ assert not task.done(), "member_delete did not wait on the team's advisory lock"
+
+ # member_add wins the race: it adds `other_user` while holding the lock
+ await held.litellm_teamtable.update(
+ where={"team_id": TEAM},
+ data={"members_with_roles": winning_add_roster},
+ )
+
+ await asyncio.wait_for(task, timeout=30)
+ finally:
+ await blocker.disconnect()
+ finally:
+ proxy_server_module.prisma_client = original_prisma_client
+
+ team_row = await db.litellm_teamtable.find_unique(where={"team_id": TEAM})
+ raw_roster = team_row.members_with_roles
+ parsed_roster = json.loads(raw_roster) if isinstance(raw_roster, str) else raw_roster
+ remaining_ids = {m["user_id"] for m in parsed_roster}
+ assert remaining_ids == {other_user}, (
+ "member_delete must remove only the user it targeted from the roster it actually "
+ "committed to, not silently drop the member the winning add just committed"
+ )
+
+
+@pytest.mark.asyncio
+async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference():
+ """
+ A member_add that wins the lock race writes its reference and releases the lock; the
+ delete that was waiting on it must then run its locked sweep against the row as it
+ actually is, not a stale snapshot, and reap that reference rather than leaving it
+ stranded on a team id the delete is about to remove.
+ """
+ import litellm.proxy.proxy_server as proxy_server_module
+ from litellm.proxy._types import LiteLLM_TeamTable
+ from litellm.proxy.management_endpoints.team_endpoints import delete_team
+
+ async with _clean_db() as db:
+ await db.litellm_teamtable.create(data={"team_id": TEAM, "team_alias": TEAM, "members_with_roles": "[]"})
+
+ async with _real_prisma_client() as prisma_client:
+ proxy_logging_obj = prisma_client.proxy_logging_obj
+ original_prisma_client = proxy_server_module.prisma_client
+ original_admin_name = proxy_server_module.litellm_proxy_admin_name
+ original_proxy_logging_obj = proxy_server_module.proxy_logging_obj
+ original_cache = proxy_server_module.user_api_key_cache
+ original_router = proxy_server_module.llm_router
+ proxy_server_module.prisma_client = prisma_client
+ proxy_server_module.litellm_proxy_admin_name = "lit5544-admin"
+ proxy_server_module.proxy_logging_obj = proxy_logging_obj
+ proxy_server_module.user_api_key_cache = original_cache or proxy_logging_obj.internal_usage_cache
+ proxy_server_module.llm_router = None
+
+ async def restore():
+ proxy_server_module.prisma_client = original_prisma_client
+ proxy_server_module.litellm_proxy_admin_name = original_admin_name
+ proxy_server_module.proxy_logging_obj = original_proxy_logging_obj
+ proxy_server_module.user_api_key_cache = original_cache
+ proxy_server_module.llm_router = original_router
+
+ try:
+ from prisma import Prisma
+
+ blocker = Prisma()
+ await blocker.connect()
+ lock_acquired = asyncio.Event()
+
+ async def run_delete():
+ lock_acquired.set()
+ return await delete_team(
+ data=DeleteTeamRequest(team_ids=[TEAM]),
+ http_request=MagicMock(),
+ user_api_key_dict=_admin_auth(),
+ litellm_changed_by="lit5544-admin",
+ )
+
+ try:
+ async with blocker.tx(timeout=timedelta(seconds=30)) as held:
+ await held.query_raw(_LOCK_SQL, TEAM)
+ task = asyncio.create_task(run_delete())
+ await lock_acquired.wait()
+ await asyncio.sleep(0.3)
+ assert not task.done(), "delete_team did not wait on the team's advisory lock"
+
+ # member_add wins the race: write the reference while holding the lock
+ await held.litellm_usertable.upsert(
+ where={"user_id": USER},
+ data={
+ "create": {"user_id": USER, "teams": [TEAM]},
+ "update": {"teams": {"push": [TEAM]}},
+ },
+ )
+ await held.litellm_teammembership.create(data={"team_id": TEAM, "user_id": USER})
+ await held.litellm_teamtable.update(
+ where={"team_id": TEAM},
+ data={"members_with_roles": '[{"user_id": "%s", "role": "user"}]' % USER},
+ )
+
+ await asyncio.wait_for(task, timeout=30)
+ finally:
+ await blocker.disconnect()
+ finally:
+ await restore()
+
+ team_row = await db.litellm_teamtable.find_unique(where={"team_id": TEAM})
+ assert team_row is None
+
+ user_row = await db.litellm_usertable.find_unique(where={"user_id": USER})
+ assert user_row is not None and TEAM not in user_row.teams, (
+ "delete_team's locked sweep must reap the reference member_add wrote just before losing the lock"
+ )
+
+ membership_row = await db.litellm_teammembership.find_first(where={"team_id": TEAM, "user_id": USER})
+ assert membership_row is None
diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py
index 21dbf3e090f..ceaabf0a70f 100644
--- a/tests/proxy_unit_tests/test_proxy_server.py
+++ b/tests/proxy_unit_tests/test_proxy_server.py
@@ -1169,6 +1169,22 @@ async def test_create_user_default_budget(prisma_client, user_role): # noqa: F8
assert mock_client.call_args.kwargs["data"]["budget_duration"] is None
+def _member_add_tx_cm(team_table):
+ """Transaction whose member writes land on whatever tables are mocked on `prisma_client.db`"""
+
+ class _Tx:
+ query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
+ litellm_teamtable = team_table
+
+ def __getattr__(self, table_name):
+ return getattr(litellm.proxy.proxy_server.prisma_client.db, table_name)
+
+ tx_cm = MagicMock()
+ tx_cm.__aenter__ = AsyncMock(return_value=_Tx())
+ tx_cm.__aexit__ = AsyncMock(return_value=None)
+ return tx_cm
+
+
@pytest.mark.parametrize("new_member_method", ["user_id", "user_email"])
@pytest.mark.asyncio
@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).")
@@ -1230,7 +1246,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): # noqa
)
)
mock_litellm_usertable.upsert = mock_client
- mock_litellm_usertable.find_many = AsyncMock(return_value=None)
+ mock_litellm_usertable.find_many = AsyncMock(return_value=[])
# Mock find_first for user_email validation (returns None for new users)
mock_litellm_usertable.find_first = AsyncMock(return_value=None)
# Mock find_unique for user_id validation (returns None for new users)
@@ -1245,12 +1261,7 @@ async def test_create_team_member_add(prisma_client, new_member_method): # noqa
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
)
- tx_mock = AsyncMock()
- tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
- tx_mock.litellm_teamtable = team_mock_client
- tx_cm = MagicMock()
- tx_cm.__aenter__ = AsyncMock(return_value=tx_mock)
- tx_cm.__aexit__ = AsyncMock(return_value=None)
+ tx_cm = _member_add_tx_cm(team_mock_client)
original_tx = litellm.proxy.proxy_server.prisma_client.tx
litellm.proxy.proxy_server.prisma_client.tx = MagicMock(
return_value=tx_cm
@@ -1432,7 +1443,7 @@ async def test_create_team_member_add_team_admin(
)
)
mock_litellm_usertable.upsert = mock_client
- mock_litellm_usertable.find_many = AsyncMock(return_value=None)
+ mock_litellm_usertable.find_many = AsyncMock(return_value=[])
# Mock find_first for user_email validation (returns None for new users)
mock_litellm_usertable.find_first = AsyncMock(return_value=None)
# Mock find_unique for user_id validation (returns None for new users)
@@ -1443,12 +1454,7 @@ async def test_create_team_member_add_team_admin(
return_value=LiteLLM_TeamTableCachedObj(team_id="1234")
)
- tx_mock = AsyncMock()
- tx_mock.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
- tx_mock.litellm_teamtable = team_mock_client
- tx_cm = MagicMock()
- tx_cm.__aenter__ = AsyncMock(return_value=tx_mock)
- tx_cm.__aexit__ = AsyncMock(return_value=None)
+ tx_cm = _member_add_tx_cm(team_mock_client)
with (
patch.object(
diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
index 7f5d3eb0a14..f33854eeb81 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
@@ -4,7 +4,7 @@ from contextlib import asynccontextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Optional, cast
-from unittest.mock import AsyncMock, MagicMock, call, patch
+from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch
import pytest
from fastapi import HTTPException
@@ -58,6 +58,9 @@ from litellm.proxy.management_endpoints.team_endpoints import (
update_team,
validate_team_org_change,
)
+from litellm.proxy.management_helpers.access_group_team_sync import (
+ TEAM_ADVISORY_LOCK_SQL,
+)
from litellm.proxy.management_helpers.team_member_permission_checks import (
TeamMemberPermissionChecks,
)
@@ -75,7 +78,11 @@ client = TestClient(app)
def _wire_team_create_tx(prisma_client):
"""`/team/new` inserts the team and mirrors it onto the access groups in one transaction,
- so a mocked client has to hand its team table back out of `db.tx()`."""
+ so a mocked client has to hand its team table back out of `db.tx()`.
+
+ A `/team/new` carrying members then adds them under the team's advisory lock, and those
+ writes run on that lock's transaction, so `tx()` has to hand back the mocked tables too
+ for the per-table assertions on `prisma_client.db.*` to keep seeing them."""
@asynccontextmanager
async def _tx():
@@ -85,18 +92,67 @@ def _wire_team_create_tx(prisma_client):
)
prisma_client.db.tx = lambda *_args, **_kwargs: _tx()
+ _wire_member_add_tx(prisma_client)
+
+
+def _wire_member_add_tx(prisma_client):
+ """/team/member_add takes the team's advisory lock, re-reads the roster under it, and runs
+ the user, budget, and membership writes on that same transaction, so a mocked client has
+ to hand its own table mocks back out of `tx()`.
+
+ Tables resolve on access, not here, since tests routinely replace `db.` after
+ wiring the transaction."""
+
+ class _Tx:
+ query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
+
+ def __getattr__(self, table_name):
+ return getattr(prisma_client.db, table_name)
+
+ tx = _Tx()
+ tx_cm = MagicMock()
+ tx_cm.__aenter__ = AsyncMock(return_value=tx)
+ tx_cm.__aexit__ = AsyncMock(return_value=None)
+ prisma_client.tx = MagicMock(return_value=tx_cm)
def _wire_member_delete_tx(prisma_client):
- """/team/member_delete's four cleanups run inside one transaction, so a mocked
- client has to hand back its own table mocks out of `tx()` for the existing
- per-table assertions to keep seeing the calls."""
+ """/team/member_delete's four cleanups, plus the advisory-lock re-read that now guards
+ them, run inside one transaction, so a mocked client has to hand back its own table
+ mocks (and a `query_raw` that answers the locked re-read from the same team row the
+ test already configured on `find_unique`) out of `tx()` for the existing per-table
+ assertions to keep seeing the calls."""
+
+ async def _query_raw(sql, team_id):
+ if sql != TEAM_ADVISORY_LOCK_SQL:
+ team_row = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id})
+ if team_row is not None:
+ return [{"members_with_roles": team_row.model_dump()["members_with_roles"]}]
+ return []
+
+ class _Tx:
+ query_raw = staticmethod(_query_raw)
+
+ def __getattr__(self, table_name):
+ return getattr(prisma_client.db, table_name)
+
+ tx = _Tx()
+ tx_cm = MagicMock()
+ tx_cm.__aenter__ = AsyncMock(return_value=tx)
+ tx_cm.__aexit__ = AsyncMock(return_value=None)
+ prisma_client.tx = MagicMock(return_value=tx_cm)
+
+
+def _wire_team_delete_tx(prisma_client):
+ """`/team/delete` deletes the team rows and runs its post-delete reference sweep under
+ every team's advisory lock in one transaction, so a mocked client has to hand its own
+ table mocks (and db-level execute_raw) back out of `tx()` for existing per-table
+ assertions on `prisma_client.db.*` to keep seeing those calls."""
tx = SimpleNamespace(
litellm_teamtable=prisma_client.db.litellm_teamtable,
- litellm_usertable=prisma_client.db.litellm_usertable,
litellm_teammembership=prisma_client.db.litellm_teammembership,
- litellm_verificationtoken=prisma_client.db.litellm_verificationtoken,
- litellm_deletedverificationtoken=prisma_client.db.litellm_deletedverificationtoken,
+ query_raw=AsyncMock(return_value=[]),
+ execute_raw=prisma_client.db.execute_raw,
)
tx_cm = MagicMock()
tx_cm.__aenter__ = AsyncMock(return_value=tx)
@@ -1669,6 +1725,7 @@ async def test_process_team_members_single_member():
default_team_budget_id="budget-123",
allowed_models=None,
budget_duration=None,
+ tx=None,
)
@@ -1809,8 +1866,8 @@ async def test_update_team_members_list_duplicate_prevention():
async def test_add_team_members_reconciles_against_freshly_locked_row():
"""
Regression: _add_team_members_to_team must build the new members_with_roles
- from the row it re-reads under a lock inside the write transaction, not from
- the stale complete_team_data snapshot captured at the start of the request.
+ from the row it re-reads under the team's advisory lock, not from the stale
+ complete_team_data snapshot captured at the start of the request.
Two concurrent /team/member_add calls for the same team read the same
snapshot; without the locked re-read the losing write rewrites the whole
@@ -1871,24 +1928,89 @@ async def test_add_team_members_reconciles_against_freshly_locked_row():
written_ids = sorted(m["user_id"] for m in json.loads(captured["data"]["members_with_roles"]))
assert written_ids == ["alice", "bob", "zed"]
- lock_reads = [call for call in tx.query_raw.call_args_list if "FOR UPDATE" in str(call.args[0])]
- assert lock_reads, "expected a SELECT ... FOR UPDATE row-lock read before the write"
+ assert tx.query_raw.call_args_list[0].args == (TEAM_ADVISORY_LOCK_SQL, "test-team-lock"), (
+ "expected the team's advisory lock to be acquired before the members_with_roles read"
+ )
+ assert not any("FOR UPDATE" in str(call.args[0]) for call in tx.query_raw.call_args_list), (
+ "a row lock here can deadlock with the access-group endpoints; only the advisory lock is safe"
+ )
assert [m.user_id for m in updated_team.members_with_roles] == ["zed", "alice", "bob"]
@pytest.mark.asyncio
-async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request():
+async def test_add_team_members_runs_member_writes_on_the_lock_holding_transaction():
+ """
+ Regression pin against exhausting the connection pool with advisory-lock waiters.
+
+ Every concurrent /team/member_add for one team holds a pooled connection while it waits
+ on the team's advisory lock. If the holder's member writes went to the regular client,
+ it would need a second connection to finish, so enough concurrent adds fill the pool
+ with waiters and the holder can never commit or release the lock. The member writes
+ therefore have to run on the transaction that already owns the connection.
+ """
+ from litellm.proxy.management_endpoints.team_endpoints import (
+ _add_team_members_to_team,
+ )
+
+ added_user = MagicMock()
+ added_user.user_id = "bob"
+ added_user.model_dump.return_value = {"user_id": "bob", "teams": ["team-pool"]}
+ created_budget = MagicMock()
+ created_budget.budget_id = "budget-pool"
+ membership = MagicMock()
+ membership.model_dump.return_value = {
+ "team_id": "team-pool",
+ "user_id": "bob",
+ "budget_id": "budget-pool",
+ "litellm_budget_table": None,
+ }
+
+ tx = MagicMock()
+ tx.query_raw = AsyncMock(return_value=[{"members_with_roles": []}])
+ tx.litellm_teamtable.update = AsyncMock(
+ return_value=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[])
+ )
+ tx.litellm_usertable.upsert = AsyncMock(return_value=added_user)
+ tx.litellm_usertable.update_many = AsyncMock()
+ tx.litellm_budgettable.create = AsyncMock(return_value=created_budget)
+ tx.litellm_teammembership.create = AsyncMock(return_value=membership)
+
+ tx_cm = MagicMock()
+ tx_cm.__aenter__ = AsyncMock(return_value=tx)
+ tx_cm.__aexit__ = AsyncMock(return_value=None)
+
+ prisma_client = MagicMock()
+ prisma_client.tx = MagicMock(return_value=tx_cm)
+ type(prisma_client).db = PropertyMock(
+ side_effect=AssertionError("member writes must not reach for a second pooled connection")
+ )
+
+ _, updated_users, updated_team_memberships = await _add_team_members_to_team(
+ data=TeamMemberAddRequest(
+ team_id="team-pool",
+ member=Member(user_id="bob", role="user"),
+ max_budget_in_team=50.0,
+ ),
+ complete_team_data=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]),
+ prisma_client=cast(object, prisma_client),
+ user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+ litellm_proxy_admin_name="admin",
+ )
+
+ assert [user.user_id for user in updated_users] == ["bob"]
+ assert [tm.budget_id for tm in updated_team_memberships] == ["budget-pool"]
+
+
+@pytest.mark.asyncio
+async def test_add_team_members_writes_nothing_when_the_team_is_deleted_mid_request():
"""
Regression pin for the /team/member_add vs /team/delete race.
- The user row and membership writes land before the reconcile takes the team
- row lock, so a /team/delete that commits in between has already run its own
- reference sweep and cannot see them. The empty locked SELECT is the only
- signal that happened, and leaving it at that would strand the member on a
- deleted team id, which authorization paths that trust `user.teams` would
- treat as membership if the id were ever recreated. So the request must sweep
- the references it just wrote and fail, not report success.
+ The advisory lock is acquired, and the team is gone, before any write is attempted:
+ the empty locked SELECT is proof a /team/delete already committed under the same
+ lock, so this request must fail without writing the user or membership rows in the
+ first place, rather than writing them and then trying to sweep them back out.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
_add_team_members_to_team,
@@ -1907,9 +2029,10 @@ async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request()
prisma_client.db.execute_raw = AsyncMock()
prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
+ process_team_members = AsyncMock(return_value=([], []))
with patch(
"litellm.proxy.management_endpoints.team_endpoints._process_team_members",
- new=AsyncMock(return_value=([], [])),
+ new=process_team_members,
):
with pytest.raises(HTTPException) as exc_info:
await _add_team_members_to_team(
@@ -1924,14 +2047,10 @@ async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request()
)
assert exc_info.value.status_code == 404
+ process_team_members.assert_not_awaited()
tx.litellm_teamtable.update.assert_not_awaited()
-
- assert prisma_client.db.execute_raw.await_args_list == [
- call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-deleted-mid-add")
- ]
- prisma_client.db.litellm_teammembership.delete_many.assert_awaited_once_with(
- where={"team_id": {"in": ("team-deleted-mid-add",)}}
- )
+ prisma_client.db.execute_raw.assert_not_awaited()
+ prisma_client.db.litellm_teammembership.delete_many.assert_not_awaited()
def test_add_new_models_to_team_with_existing_models():
@@ -4246,6 +4365,86 @@ async def test_team_member_delete_cleans_verification_tokens(
)
+@pytest.mark.asyncio
+async def test_team_member_delete_reads_on_the_lock_holding_transaction(
+ mock_db_client, mock_admin_auth
+):
+ """
+ Regression pin against exhausting the connection pool with advisory-lock waiters.
+
+ Every concurrent removal for one team holds a pooled connection while it waits on the
+ team's advisory lock, and /team/delete fans its per-member removals out concurrently.
+ A holder whose reads went to the regular client would need a second connection to
+ finish, so enough waiters fill the pool and the holder can never release the lock.
+ Both reads therefore have to run on the transaction that already owns the connection.
+ """
+ from litellm.proxy._types import TeamMemberDeleteRequest
+ from litellm.proxy.management_endpoints.team_endpoints import team_member_delete
+
+ test_team_id = "team-del-pool-123"
+ test_user_id = "user-del-pool-123"
+ roster_entry = {"user_id": test_user_id, "user_email": None, "role": "user"}
+
+ mock_team_row = MagicMock()
+ mock_team_row.model_dump.return_value = {
+ "team_id": test_team_id,
+ "members_with_roles": [roster_entry],
+ "team_member_permissions": [],
+ "metadata": {},
+ "models": [],
+ "spend": 0.0,
+ }
+ mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(
+ return_value=mock_team_row
+ )
+
+ user_row = MagicMock()
+ user_row.user_id = test_user_id
+ user_row.teams = [test_team_id]
+
+ # Both are wired to answer, so the endpoint completes either way and the awaits below
+ # are what tells which connection it read on.
+ pooled_user_read = AsyncMock(return_value=[user_row])
+ pooled_token_read = AsyncMock(return_value=[])
+ mock_db_client.db.litellm_usertable.find_many = pooled_user_read
+ mock_db_client.db.litellm_verificationtoken.find_many = pooled_token_read
+
+ tx = MagicMock()
+ tx.query_raw = AsyncMock(return_value=[{"members_with_roles": [roster_entry]}])
+ tx.litellm_teamtable.update = AsyncMock(return_value=mock_team_row)
+ tx.litellm_usertable.find_many = AsyncMock(return_value=[user_row])
+ tx.litellm_usertable.update = AsyncMock()
+ tx.litellm_teammembership.delete_many = AsyncMock()
+ tx.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
+ tx.litellm_verificationtoken.delete_many = AsyncMock()
+
+ tx_cm = MagicMock()
+ tx_cm.__aenter__ = AsyncMock(return_value=tx)
+ tx_cm.__aexit__ = AsyncMock(return_value=None)
+ mock_db_client.tx = MagicMock(return_value=tx_cm)
+
+ await team_member_delete(
+ data=TeamMemberDeleteRequest(team_id=test_team_id, user_id=test_user_id),
+ user_api_key_dict=mock_admin_auth,
+ )
+
+ tx.litellm_usertable.find_many.assert_awaited_once_with(
+ where={"user_id": {"in": [test_user_id]}}
+ )
+ tx.litellm_verificationtoken.find_many.assert_awaited_once_with(
+ where={"user_id": {"in": [test_user_id]}, "team_id": test_team_id}
+ )
+ pooled_user_read.assert_not_awaited()
+ pooled_token_read.assert_not_awaited()
+
+ tx.litellm_usertable.update.assert_awaited_once_with(
+ where={"user_id": test_user_id}, data={"teams": {"set": []}}
+ )
+ tx.litellm_teammembership.delete_many.assert_awaited_once_with(
+ where={"team_id": test_team_id, "user_id": test_user_id}
+ )
+
+
@pytest.mark.parametrize(
"roster_email",
["Alice@Example.com", "alice-invited-as@example.com"],
@@ -7411,6 +7610,7 @@ async def test_delete_team_persists_deleted_teams(
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
+ _wire_team_delete_tx(mock_prisma_client)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
@@ -7481,15 +7681,14 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
cache_state_when_rows_deleted = {}
async def record_cache_state_then_delete(*args, **kwargs):
- if kwargs.get("table_name") == "team":
- cache_state_when_rows_deleted["doomed_still_cached"] = (
- fresh_cache.get_cache(key="team_id:team-doomed") is not None
- )
- return {"deleted_teams": ["team-doomed"]}
+ cache_state_when_rows_deleted["doomed_still_cached"] = (
+ fresh_cache.get_cache(key="team_id:team-doomed") is not None
+ )
+ return 1
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=doomed_team)
- mock_prisma_client.delete_data = AsyncMock(side_effect=record_cache_state_then_delete)
+ mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 0})
mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
@@ -7498,6 +7697,7 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
mock_prisma_client.db.execute_raw = mock_execute_raw
mock_membership_delete_many = AsyncMock()
mock_prisma_client.db.litellm_teammembership.delete_many = mock_membership_delete_many
+ mock_prisma_client.db.litellm_teamtable.delete_many = AsyncMock(side_effect=record_cache_state_then_delete)
mock_tx = AsyncMock()
mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
@@ -7506,6 +7706,11 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
+ # The locked delete-and-sweep transaction /team/member_add serializes against, kept
+ # separate from mock_tx above (the BYOK-model-cleanup transaction, unrelated to this lock).
+ _wire_team_delete_tx(mock_prisma_client)
+ mock_lock_tx = mock_prisma_client.tx.return_value.__aenter__.return_value
+
fresh_cache = UserApiKeyCache()
for cached_team_id, cached_alias in (
("team-doomed", "doomed-team"),
@@ -7539,14 +7744,22 @@ async def test_delete_team_sweeps_references_outside_members_with_roles(
assert mock_execute_raw.await_args_list == [
call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"),
call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"),
- ], "the sweep must run once before the team row is deleted and again after, so a member_add racing the delete cannot leave the reference behind"
+ ], (
+ "the unlocked sweep must run once to catch pre-existing drift, and the locked sweep "
+ "(alongside the delete, under the same advisory lock member_add takes) must run again "
+ "so a member_add that wrote its reference just before losing the lock is still reaped"
+ )
- # same two passes: the second one reaps a membership row inserted while the delete was running
+ # same two passes for the membership rows, the second under the lock alongside the delete
assert mock_membership_delete_many.await_args_list == [
call(where={"team_id": {"in": ("team-doomed",)}}),
call(where={"team_id": {"in": ("team-doomed",)}}),
]
+ assert mock_lock_tx.query_raw.await_args_list == [call(TEAM_ADVISORY_LOCK_SQL, "team-doomed")], (
+ "the advisory lock must be acquired before the team row is deleted"
+ )
+
assert fresh_cache.get_cache(key="team_id:team-doomed") is None
assert fresh_cache.get_cache(key="team_alias:doomed-team") is None
assert fresh_cache.get_cache(key="team_id:team-kept") is not None
@@ -7596,6 +7809,7 @@ async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
+ _wire_team_delete_tx(mock_prisma_client)
fresh_cache = UserApiKeyCache()
fresh_cache.set_cache(key="hashed-doomed-key", value=UserAPIKeyAuth(token="hashed-doomed-key", team_id="team-doomed"))
@@ -7623,14 +7837,17 @@ async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(
@pytest.mark.asyncio
-async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cache(
+async def test_delete_team_failing_locked_sweep_rolls_back_the_delete_and_leaves_the_cache_alone(
monkeypatch,
disable_audit_logging_for_mocked_team,
):
"""
- The reconcile sweep runs after the team row is committed deleted. If it ran before cache
- eviction, a sweep failure would return an error with the team gone from the db but still
- served from cache, which is the exact bug this PR exists to fix.
+ The team delete and its post-delete reconcile sweep run inside one transaction, under the
+ team's advisory lock, so a sweep failure rolls the delete back with it rather than leaving
+ the row gone with the sweep half done. Cache eviction only runs after that transaction
+ commits, so a failure here must leave the team exactly as it was: still in the db, and
+ still cached. Evicting a cache entry for a delete that never actually committed would be
+ the same class of bug this PR exists to fix, just on the other side of the transaction.
"""
from litellm.proxy._types import DeleteTeamRequest
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@@ -7650,7 +7867,7 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac
mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock()
- # the first sweep succeeds, the post-delete reconcile sweep blows up
+ # the unlocked pre-delete sweep succeeds, the locked post-delete sweep blows up
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[None, ConnectionError("db went away")])
mock_tx = AsyncMock()
@@ -7659,6 +7876,7 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
+ _wire_team_delete_tx(mock_prisma_client)
fresh_cache = UserApiKeyCache()
cached_obj = LiteLLM_TeamTableCachedObj(team_id="team-doomed", team_alias="doomed-team")
@@ -7682,9 +7900,10 @@ async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cac
litellm_changed_by="admin-user",
)
- # the delete committed, so the cache must not still be serving the team
- assert fresh_cache.get_cache(key="team_id:team-doomed") is None
- assert fresh_cache.get_cache(key="team_alias:doomed-team") is None
+ # the transaction that deletes the row and runs the locked sweep never committed, so
+ # cache eviction (which only runs after that commit) must never have been reached
+ assert fresh_cache.get_cache(key="team_id:team-doomed") is not None
+ assert fresh_cache.get_cache(key="team_alias:doomed-team") is not None
@pytest.mark.asyncio
@@ -7726,6 +7945,7 @@ async def test_delete_team_broadcasts_cache_invalidation_to_other_workers(
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
+ _wire_team_delete_tx(mock_prisma_client)
published = []
@@ -7796,6 +8016,7 @@ async def test_delete_team_survives_a_failing_cache_backend(
mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx)
mock_tx_cm.__aexit__ = AsyncMock(return_value=False)
mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm)
+ _wire_team_delete_tx(mock_prisma_client)
exploding_logging_obj = MagicMock()
exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock(
@@ -7820,7 +8041,7 @@ async def test_delete_team_survives_a_failing_cache_backend(
)
assert result == {"deleted_teams": ["team-doomed"]}
- mock_delete_data.assert_any_await(team_id_list=["team-doomed"], table_name="team")
+ mock_prisma_client.db.litellm_teamtable.delete_many.assert_any_await(where={"team_id": {"in": ["team-doomed"]}})
assert exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache.await_count > 0
diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
index bdc2f9065b9..a6b1fc32eda 100644
--- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
+++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
@@ -1120,3 +1120,81 @@ async def test_add_new_member_creates_missing_user_atomically_via_upsert():
assert upsert_data["create"]["teams"] == ["team-1"]
assert upsert_data["update"], "empty update branch degrades the upsert to a racy SELECT-then-INSERT"
assert "teams" not in upsert_data["update"]
+
+
+def _member_write_tx() -> MagicMock:
+ tx = MagicMock()
+ created_user = MagicMock()
+ created_user.user_id = "pool-user"
+ created_user.model_dump.return_value = {
+ "user_id": "pool-user",
+ "user_email": "pool@example.com",
+ "teams": ["team-pool"],
+ "user_role": "internal_user",
+ }
+ created_budget = MagicMock()
+ created_budget.budget_id = "budget-pool"
+ membership = MagicMock()
+ membership.model_dump.return_value = {
+ "team_id": "team-pool",
+ "user_id": "pool-user",
+ "budget_id": "budget-pool",
+ "litellm_budget_table": None,
+ }
+ tx.litellm_usertable.upsert = AsyncMock(return_value=created_user)
+ tx.litellm_usertable.create = AsyncMock(return_value=created_user)
+ tx.litellm_usertable.update_many = AsyncMock()
+ tx.litellm_usertable.find_many = AsyncMock(return_value=[])
+ tx.litellm_budgettable.find_unique = AsyncMock(return_value=None)
+ tx.litellm_budgettable.create = AsyncMock(return_value=created_budget)
+ tx.litellm_teammembership.create = AsyncMock(return_value=membership)
+ return tx
+
+
+@pytest.mark.parametrize(
+ "new_member",
+ [
+ Member(user_id="pool-user", role="user"),
+ Member(user_email="pool@example.com", role="user"),
+ ],
+ ids=["by_user_id", "by_user_email"],
+)
+@pytest.mark.asyncio
+async def test_add_new_member_runs_every_write_on_the_caller_transaction(new_member):
+ """
+ Regression pin against exhausting the connection pool with advisory-lock waiters.
+
+ /team/member_add calls this while holding the team's advisory lock inside a transaction,
+ so it already owns a pooled connection. Any query issued on the regular client here needs
+ a second one, and enough concurrent adds for one team leave every connection parked on the
+ lock while the holder waits for a free one, so nothing ever commits or releases the lock.
+ Given a transaction, every read and write has to go through it.
+ """
+ from litellm.proxy._types import LitellmUserRoles
+
+ tx = _member_write_tx()
+ prisma_client = AsyncMock()
+
+ result_user, result_membership = await add_new_member(
+ new_member=new_member,
+ max_budget_in_team=50.0,
+ prisma_client=prisma_client,
+ team_id="team-pool",
+ user_api_key_dict=UserAPIKeyAuth(
+ user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN
+ ),
+ litellm_proxy_admin_name="admin",
+ tx=tx,
+ )
+
+ assert result_user.user_id == "pool-user"
+ assert result_membership is not None
+ assert result_membership.budget_id == "budget-pool"
+
+ assert tx.litellm_budgettable.create.await_count == 1
+ assert tx.litellm_teammembership.create.await_count == 1
+ assert tx.litellm_usertable.upsert.await_count + tx.litellm_usertable.create.await_count == 1
+
+ prisma_client.db.assert_not_called()
+ prisma_client.get_data.assert_not_awaited()
+ prisma_client.insert_data.assert_not_awaited()
diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py
index 38af52f165c..758d379f22c 100644
--- a/tests/test_litellm/repositories/test_repositories.py
+++ b/tests/test_litellm/repositories/test_repositories.py
@@ -541,17 +541,19 @@ class TestTeamRepository:
assert [m.user_id for m in members] == expected_ids
sql = tx.query_raw.call_args.args[0]
- assert "FOR UPDATE" in sql
+ assert "FOR UPDATE" not in sql, (
+ "a row lock here can deadlock with the access-group endpoints; the caller must "
+ "already hold the team's advisory lock, so a plain read is all this needs"
+ )
assert tx.query_raw.call_args.args[1] == "team-1"
@pytest.mark.asyncio
async def test_get_members_with_roles_locked_missing_row(self, repo):
"""None, not [], so a caller can tell a deleted team from an empty one.
- /team/member_add reconciles membership under this lock and has to fail,
- and clean up the references it already wrote, when a /team/delete
- committed underneath it. An empty list would look like a live team with
- no members and it would carry on writing.
+ /team/member_add reconciles membership under the team's advisory lock and has to
+ fail, without writing anything, when a /team/delete committed underneath it. An
+ empty list would look like a live team with no members and it would carry on writing.
"""
tx = MagicMock()
tx.query_raw = AsyncMock(return_value=[])
From 61ee2f5ecbb602fb42a8f8b9d2f4316bdc0678a4 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 25 Aug 2026 15:00:57 -0700
Subject: [PATCH 69/70] fix(anthropic): carry user-content document blocks
through the /v1/messages responses bridge
---
.../responses_adapters/transformation.py | 7 ++
.../test_responses_adapters_transformation.py | 81 +++++++++++++++++++
2 files changed, 88 insertions(+)
diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
index d9d62e0719f..492e0050cfe 100644
--- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
+++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py
@@ -214,6 +214,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
system text -> message(role=system, input_text)
user text -> message(role=user, input_text)
user image -> message(role=user, input_image)
+ user document -> message(role=user, input_file)
user tool_result -> function_call_output
assistant text -> message(role=assistant, output_text)
assistant thinking -> reasoning
@@ -268,6 +269,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
{"type": "input_image", "image_url": url}, block.get("prompt_cache_breakpoint")
)
)
+ elif btype == "document":
+ file_part = self._translate_anthropic_document_block_to_file_part(block)
+ if file_part:
+ user_parts.append(
+ with_prompt_cache_breakpoint(file_part, block.get("prompt_cache_breakpoint"))
+ )
elif btype == "tool_result":
tool_use_id = block.get("tool_use_id", "")
inner = block.get("content")
diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
index 44c956c8ee8..56f106e407c 100644
--- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
+++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py
@@ -1688,6 +1688,87 @@ class TestToolResultDocuments:
]
+class TestUserContentDocuments:
+ """Documents in plain user content must survive translation (LIT-6144): each
+ document block becomes an input_file part of the user message, in block order,
+ exactly like image blocks become input_image parts. Untranslatable documents
+ are dropped without disturbing the surrounding parts."""
+
+ PDF_B64 = "JVBERi0xLjQKJSBQT05H"
+ PDF_DATA_URI = "data:application/pdf;base64,JVBERi0xLjQKJSBQT05H"
+ PDF_URL = "https://example.com/report.pdf"
+ EXPLICIT = {"mode": "explicit"}
+
+ def _translate(self, user_content):
+ return _ADAPTER.translate_messages_to_responses_input([{"role": "user", "content": user_content}])
+
+ @staticmethod
+ def _user_content(items):
+ return next(item for item in items if item.get("type") == "message" and item.get("role") == "user")["content"]
+
+ def _base64_document(self, **extra):
+ return {
+ "type": "document",
+ "source": {"type": "base64", "media_type": "application/pdf", "data": self.PDF_B64},
+ **extra,
+ }
+
+ def test_document_then_text_keeps_block_order(self):
+ content = self._user_content(
+ self._translate([self._base64_document(), {"type": "text", "text": "what does the pdf say?"}])
+ )
+ assert content == [
+ {"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI},
+ {"type": "input_text", "text": "what does the pdf say?"},
+ ]
+
+ def test_document_title_becomes_filename(self):
+ content = self._user_content(self._translate([self._base64_document(title="quarterly-report.pdf")]))
+ assert content == [
+ {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI}
+ ]
+
+ def test_url_document_becomes_file_url_part(self):
+ content = self._user_content(
+ self._translate([{"type": "document", "source": {"type": "url", "url": self.PDF_URL}}])
+ )
+ assert content == [{"type": "input_file", "file_url": self.PDF_URL}]
+
+ def test_document_only_content_still_produces_user_message(self):
+ content = self._user_content(self._translate([self._base64_document()]))
+ assert content == [{"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}]
+
+ def test_empty_base64_data_drops_only_the_document_part(self):
+ content = self._user_content(
+ self._translate(
+ [
+ {"type": "text", "text": "still here"},
+ {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": ""}},
+ ]
+ )
+ )
+ assert content == [{"type": "input_text", "text": "still here"}]
+
+ def test_non_dict_source_drops_only_the_document_part(self):
+ content = self._user_content(
+ self._translate([{"type": "text", "text": "still here"}, {"type": "document", "source": self.PDF_URL}])
+ )
+ assert content == [{"type": "input_text", "text": "still here"}]
+
+ def test_document_breakpoint_rides_on_the_file_part(self):
+ content = self._user_content(
+ self._translate([self._base64_document(prompt_cache_breakpoint=self.EXPLICIT)])
+ )
+ assert content == [
+ {
+ "type": "input_file",
+ "filename": "document.pdf",
+ "file_data": self.PDF_DATA_URI,
+ "prompt_cache_breakpoint": self.EXPLICIT,
+ }
+ ]
+
+
def _contains_key(value, key) -> bool:
if isinstance(value, dict):
return key in value or any(_contains_key(v, key) for v in value.values())
From bb22742025065c10baba94b7d587ab000c6174ff Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
<158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 25 Aug 2026 15:54:25 -0700
Subject: [PATCH 70/70] fix(rerank): emit latency and cost headers on /rerank
(#35419)
* fix(rerank): emit latency and cost headers on /rerank
Thread the logging object into the rerank httpx calls and pass hidden_params through to get_custom_headers, so x-litellm-overhead-duration-ms, x-litellm-response-duration-ms, x-litellm-response-cost, x-litellm-call-id and the LITELLM_DETAILED_TIMING x-litellm-timing-* headers show up on rerank like they do on chat completions
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(rerank): keep zero response cost in the /rerank cost header
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* ci: assign the new rerank endpoint tests to the proxy-endpoints shard
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test: suppress TQ008 on the rerank header tests with reasons
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: milan
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yassin
---
.github/workflows/test-unit.yml | 1 +
litellm/llms/bedrock/rerank/handler.py | 3 +
litellm/llms/custom_httpx/llm_http_handler.py | 1 +
litellm/proxy/rerank_endpoints/endpoints.py | 3 +
.../test_bedrock_rerank_header_forwarding.py | 39 +++++-
.../custom_httpx/test_llm_http_handler.py | 31 +++++
.../proxy/rerank_endpoints/__init__.py | 0
.../proxy/rerank_endpoints/test_endpoints.py | 120 ++++++++++++++++++
8 files changed, 193 insertions(+), 5 deletions(-)
create mode 100644 tests/test_litellm/proxy/rerank_endpoints/__init__.py
create mode 100644 tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py
diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index 2dfca3d308f..c23678c51ae 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -164,6 +164,7 @@ jobs:
tests/test_litellm/proxy/public_endpoints
tests/test_litellm/proxy/prompts
tests/test_litellm/proxy/rag_endpoints
+ tests/test_litellm/proxy/rerank_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/config_resolvers
diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py
index 1cc72f265eb..a2a1acec80c 100644
--- a/litellm/llms/bedrock/rerank/handler.py
+++ b/litellm/llms/bedrock/rerank/handler.py
@@ -29,6 +29,7 @@ class BedrockRerankHandler(BaseAWSLLM):
async def arerank(
self,
prepared_request: BedrockPreparedRequest,
+ logging_obj: LitellmLogging,
timeout: float | httpx.Timeout | None = None,
client: AsyncHTTPHandler | None = None,
):
@@ -40,6 +41,7 @@ class BedrockRerankHandler(BaseAWSLLM):
headers=dict(prepared_request["prepped"].headers),
data=prepared_request["body"],
timeout=timeout,
+ logging_obj=logging_obj,
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
@@ -98,6 +100,7 @@ class BedrockRerankHandler(BaseAWSLLM):
if _is_async:
return self.arerank(
prepared_request,
+ logging_obj=logging_obj,
timeout=timeout,
client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None,
)
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index b7ddb55ae89..bb582f677be 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -1203,6 +1203,7 @@ class BaseLLMHTTPHandler:
headers=headers,
data=json.dumps(request_data),
timeout=timeout,
+ logging_obj=logging_obj,
)
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py
index 45b190c1f9d..dd5803796b7 100644
--- a/litellm/proxy/rerank_endpoints/endpoints.py
+++ b/litellm/proxy/rerank_endpoints/endpoints.py
@@ -90,12 +90,15 @@ async def rerank(
fastapi_response.headers.update(
ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
+ call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None),
model_id=model_id,
cache_key=cache_key,
api_base=api_base,
version=version,
+ response_cost=hidden_params.get("response_cost", None),
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
request_data=data,
+ hidden_params=hidden_params,
**additional_headers,
)
)
diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
index b2a2046b131..253edad57e9 100644
--- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
+++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py
@@ -77,7 +77,7 @@ def test_bedrock_rerank_header_forwarding_sync(model):
with (
patch.object(client, "post") as mock_post,
- patch(
+ patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
return_value=mock_credentials_info,
),
@@ -170,7 +170,7 @@ async def test_bedrock_rerank_header_forwarding_async(model):
with (
patch.object(client, "post", new_callable=AsyncMock) as mock_post,
- patch(
+ patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
return_value=mock_credentials_info,
),
@@ -241,7 +241,7 @@ def test_bedrock_rerank_timeout_sync():
with (
patch.object(client, "post") as mock_post,
- patch(
+ patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
return_value=mock_credentials_info,
),
@@ -285,7 +285,7 @@ async def test_bedrock_rerank_timeout_async():
with (
patch.object(client, "post", new_callable=AsyncMock) as mock_post,
- patch(
+ patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
return_value=mock_credentials_info,
),
@@ -340,7 +340,7 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
with (
patch.object(client, "post") as mock_post,
- patch(
+ patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
"litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
return_value=mock_credentials_info,
),
@@ -400,3 +400,32 @@ def test_bedrock_rerank_extra_headers_and_headers_merge():
except Exception as e:
pytest.fail(f"Failed to merge and forward headers: {str(e)}")
+
+
+@pytest.mark.asyncio
+async def test_bedrock_rerank_records_llm_api_duration():
+ """The bedrock rerank handler must feed httpx timing into the logging obj, so the
+ proxy can emit x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank."""
+ import httpx
+
+ def handle(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(200, json=bedrock_rerank_response)
+
+ client = AsyncHTTPHandler()
+ client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+ with patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport
+ "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params",
+ return_value=create_mock_credentials(),
+ ):
+ response = await litellm.arerank(
+ model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0",
+ query=test_query,
+ documents=test_documents,
+ top_n=3,
+ client=client,
+ aws_region_name="us-east-1",
+ )
+
+ assert response._hidden_params["litellm_overhead_time_ms"] is not None
+ assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"]
diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
index a93b14d45f3..694a01cda5f 100644
--- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
+++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
@@ -2528,6 +2528,37 @@ def test_only_callbacks_that_can_charge_a_frame_are_collected_for_ws_quota(monke
assert _collect_ws_project_quota_callbacks() == (quota,)
+@pytest.mark.asyncio
+async def test_async_rerank_records_llm_api_duration():
+ """arerank must feed the httpx timing into the logging obj, so the proxy can emit
+ x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank."""
+
+ def handle(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ 200,
+ json={
+ "id": "rerank-1",
+ "results": [{"index": 0, "relevance_score": 0.9}],
+ "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 1}},
+ },
+ )
+
+ client = AsyncHTTPHandler()
+ client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
+
+ response = await litellm.arerank(
+ model="cohere/rerank-v3.5",
+ query="what is the capital of france",
+ documents=["paris", "berlin"],
+ top_n=1,
+ api_key="fake-key",
+ client=client,
+ )
+
+ assert response._hidden_params["litellm_overhead_time_ms"] is not None
+ assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"]
+
+
class _JSONBodyVideoConfig(OpenAIVideoConfig):
def use_multipart_form_data(self) -> bool:
return False
diff --git a/tests/test_litellm/proxy/rerank_endpoints/__init__.py b/tests/test_litellm/proxy/rerank_endpoints/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py
new file mode 100644
index 00000000000..9f11ff6f20d
--- /dev/null
+++ b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py
@@ -0,0 +1,120 @@
+"""
+Tests for rerank_endpoints/endpoints.py response headers.
+"""
+
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi import Request, Response
+
+import litellm.proxy.common_request_processing as common_request_processing_mod
+import litellm.proxy.proxy_server as proxy_server_mod
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.rerank_endpoints.endpoints import rerank
+from litellm.types.utils import RerankResponse
+
+HIDDEN_PARAMS = {
+ "model_id": "deployment-1",
+ "api_base": "https://bedrock-agent-runtime.us-east-1.amazonaws.com",
+ "response_cost": 0.002,
+ "_response_ms": 1500.5,
+ "litellm_overhead_time_ms": 12.5,
+ "callback_duration_ms": 1.25,
+ "timing_llm_api_ms": 1488.0,
+ "timing_pre_processing_ms": 10.0,
+ "timing_post_processing_ms": 2.5,
+ "timing_message_copy_ms": 0.01,
+}
+
+
+def _build_request() -> Request:
+ body = json.dumps({"model": "rerank-model", "query": "q", "documents": ["a", "b"]}).encode()
+
+ async def receive():
+ return {"type": "http.request", "body": body, "more_body": False}
+
+ return Request(
+ scope={
+ "type": "http",
+ "method": "POST",
+ "path": "/rerank",
+ "headers": [(b"content-type", b"application/json")],
+ "query_string": b"",
+ },
+ receive=receive,
+ )
+
+
+async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response:
+ response = RerankResponse(id="rerank-1", results=[{"index": 0, "relevance_score": 0.9}])
+ response._hidden_params = dict(hidden_params)
+
+ fastapi_response = Response()
+ proxy_logging_obj = MagicMock()
+ proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"])
+ proxy_logging_obj.update_request_status = AsyncMock()
+
+ async def fake_add_litellm_data_to_request(**kwargs):
+ return {**kwargs["data"], "litellm_call_id": "call-123"}
+
+ async def fake_route_request(**kwargs):
+ async def _call():
+ return response
+
+ return _call()
+
+ with (
+ patch.object(proxy_server_mod, "add_litellm_data_to_request", fake_add_litellm_data_to_request), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+ patch.object(proxy_server_mod, "route_request", fake_route_request), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+ patch.object(proxy_server_mod, "proxy_logging_obj", proxy_logging_obj), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+ patch.object(proxy_server_mod, "llm_router", MagicMock()), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+ patch.object(proxy_server_mod, "version", "1.2.3"), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler
+ ):
+ await rerank(
+ request=_build_request(),
+ fastapi_response=fastapi_response,
+ user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+ )
+
+ return fastapi_response
+
+
+@pytest.mark.asyncio
+async def test_rerank_emits_latency_and_cost_headers():
+ """/rerank must surface the same hidden_params-derived headers as /chat/completions."""
+ fastapi_response = await _call_rerank()
+
+ assert fastapi_response.headers["x-litellm-call-id"] == "call-123"
+ assert fastapi_response.headers["x-litellm-response-duration-ms"] == "1500.5"
+ assert fastapi_response.headers["x-litellm-overhead-duration-ms"] == "12.5"
+ assert fastapi_response.headers["x-litellm-callback-duration-ms"] == "1.25"
+ assert fastapi_response.headers["x-litellm-response-cost"] == "0.002"
+
+
+@pytest.mark.asyncio
+async def test_rerank_emits_detailed_timing_headers_when_enabled():
+ """LITELLM_DETAILED_TIMING must also work on /rerank, not just /chat/completions."""
+ with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True): # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test
+ fastapi_response = await _call_rerank()
+
+ assert fastapi_response.headers["x-litellm-timing-llm-api-ms"] == "1488.0"
+ assert fastapi_response.headers["x-litellm-timing-pre-processing-ms"] == "10.0"
+ assert fastapi_response.headers["x-litellm-timing-post-processing-ms"] == "2.5"
+ assert fastapi_response.headers["x-litellm-timing-message-copy-ms"] == "0.01"
+
+
+@pytest.mark.asyncio
+async def test_rerank_emits_zero_response_cost_header():
+ """A free deployment costs 0.0, which is a real cost and must not be dropped."""
+ fastapi_response = await _call_rerank({**HIDDEN_PARAMS, "response_cost": 0.0})
+
+ assert fastapi_response.headers["x-litellm-response-cost"] == "0.0"
+
+
+@pytest.mark.asyncio
+async def test_rerank_omits_detailed_timing_headers_when_disabled():
+ with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False): # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test
+ fastapi_response = await _call_rerank()
+
+ assert "x-litellm-timing-llm-api-ms" not in fastapi_response.headers