{latencyMs !== undefined && (
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.test.tsx
index a2f6c24781a..5089179ef13 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.test.tsx
@@ -13,9 +13,9 @@ describe("ValidationStatus", () => {
render(
,
);
@@ -23,7 +23,9 @@ describe("ValidationStatus", () => {
});
it("reports a valid payload when only warnings remain", () => {
- render(
);
+ render(
+
,
+ );
expect(screen.getByText("Valid payload")).toBeInTheDocument();
});
});
@@ -31,7 +33,10 @@ describe("ValidationStatus", () => {
describe("IssueList", () => {
it("lists every issue under its field label and renders nothing when there are none", () => {
const { rerender } = render(
-
,
+
,
);
expect(screen.getByRole("list", { name: "Payload validation issues" })).toHaveTextContent("InputRequired");
rerender(
);
@@ -46,8 +51,8 @@ describe("ValidationSummary", () => {
rerender(
,
);
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.tsx
index e8152456d15..554e7da3770 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/ValidationFeedback.tsx
@@ -1,6 +1,5 @@
import { Badge } from "@/components/ui/badge";
import { CircleAlert, CircleCheck, TriangleAlert } from "lucide-react";
-import { issueLabel } from "./lib/issueLabel";
import { errorCount as countErrors, type PayloadValidation } from "./lib/validatePayload";
export function ValidationStatus({ validation }: { validation: PayloadValidation
}) {
@@ -44,7 +43,7 @@ export function IssueList({ id, validation }: { id: string; validation: PayloadV
) : (
)}
- {issueLabel(issue.path)}
+ {issue.label}
{issue.message}
))}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.test.ts
index 29ae0f60ae4..85d2a576194 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.test.ts
@@ -3,29 +3,34 @@ import { issueLabel } from "./issueLabel";
describe("issueLabel", () => {
it("names the Input field for both endpoints instead of the JSON key", () => {
- expect(issueLabel("state")).toBe("Input");
- expect(issueLabel("input")).toBe("Input");
+ expect(issueLabel(["state"])).toBe("Input");
+ expect(issueLabel(["input"])).toBe("Input");
});
it("names top-level fields and the request itself", () => {
- expect(issueLabel("model")).toBe("Model");
- expect(issueLabel("questions")).toBe("Questions");
- expect(issueLabel("safety_identifier")).toBe("Safety identifier");
- expect(issueLabel("root")).toBe("Request");
- expect(issueLabel("syntax")).toBe("Request");
+ expect(issueLabel(["model"])).toBe("Model");
+ expect(issueLabel(["questions"])).toBe("Questions");
+ expect(issueLabel(["safety_identifier"])).toBe("Safety identifier");
+ expect(issueLabel([])).toBe("Request");
+ expect(issueLabel(["syntax"])).toBe("Request");
});
- it("names a question by its key or its 1-based position and the field inside it", () => {
- expect(issueLabel("questions.urgency")).toBe('Question "urgency"');
- expect(issueLabel("questions.urgency.criteria")).toBe('Question "urgency" options');
- expect(issueLabel("questions.urgency.criteria.1")).toBe('Question "urgency" options');
- expect(issueLabel("questions.category.instructions")).toBe('Question "category" instructions');
- expect(issueLabel("questions.0.name")).toBe("Question 1 name");
- expect(issueLabel("questions.2.type")).toBe("Question 3 answer type");
+ it("names a keyed question by its key, even an all-digit or dotted one, and the field inside it", () => {
+ expect(issueLabel(["questions", "urgency"])).toBe('Question "urgency"');
+ expect(issueLabel(["questions", "urgency", "criteria"])).toBe('Question "urgency" options');
+ expect(issueLabel(["questions", "urgency", "criteria", 1])).toBe('Question "urgency" options');
+ expect(issueLabel(["questions", "category", "instructions"])).toBe('Question "category" instructions');
+ expect(issueLabel(["questions", "2", "criteria"])).toBe('Question "2" options');
+ expect(issueLabel(["questions", "v1.2", "type"])).toBe('Question "v1.2" answer type');
+ });
+
+ it("names a listed question by its 1-based position", () => {
+ expect(issueLabel(["questions", 0, "name"])).toBe("Question 1 name");
+ expect(issueLabel(["questions", 2, "type"])).toBe("Question 3 answer type");
});
it("falls back to the humanized key for anything unknown", () => {
- expect(issueLabel("max_tokens")).toBe("max tokens");
- expect(issueLabel("questions.a.max_items")).toBe('Question "a" max items');
+ expect(issueLabel(["max_tokens"])).toBe("max tokens");
+ expect(issueLabel(["questions", "a", "max_items"])).toBe('Question "a" max items');
});
});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.ts
index 76acc266117..60ca47ead2a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/issueLabel.ts
@@ -15,16 +15,16 @@ const QUESTION_FIELD: Record = {
criteria: "options",
};
-const humanize = (segment: string): string => segment.replace(/_/g, " ");
+const humanize = (segment: PropertyKey): string => String(segment).replace(/_/g, " ");
-export function issueLabel(path: string): string {
- const [head, question, field] = path.split(".");
+export function issueLabel(path: readonly PropertyKey[]): string {
+ const [head = "root", question, field] = path;
if (head !== "questions" || question === undefined) {
- return TOP_LEVEL[head] ?? humanize(head);
+ return TOP_LEVEL[String(head)] ?? humanize(head);
}
- const title = /^\d+$/.test(question) ? `Question ${Number(question) + 1}` : `Question "${question}"`;
+ const title = typeof question === "number" ? `Question ${question + 1}` : `Question "${String(question)}"`;
if (field === undefined) {
return title;
}
- return `${title} ${QUESTION_FIELD[field] ?? humanize(field)}`;
+ return `${title} ${QUESTION_FIELD[String(field)] ?? humanize(field)}`;
}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.test.ts
index 6e91f09f8bd..e28d7da8e3f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.test.ts
@@ -11,11 +11,8 @@ const requestWith = (fields: Record = {}) =>
describe("validateSystemOnePayload", () => {
it("rejects empty input", () => {
- expect(validateSystemOnePayload(" ").issues).toContainEqual({
- path: "root",
- message: "Payload cannot be empty.",
- severity: "error",
- });
+ const emptyIssue = { path: "root", label: "Request", message: "Payload cannot be empty.", severity: "error" };
+ expect(validateSystemOnePayload(" ").issues).toContainEqual(emptyIssue);
});
it("reports JSON syntax errors", () => {
@@ -249,6 +246,17 @@ describe("validateSystemOnePayload", () => {
);
});
+ it("labels a score-level warning by its question key even when the key is all digits", () => {
+ const result = validateSystemOnePayload(
+ requestWith({
+ questions: {
+ "2": { type: "score", instructions: "Rate urgency", criteria: Array.from({ length: 11 }, () => "Level") },
+ },
+ }),
+ );
+ expect(result.issues.map((issue) => issue.label)).toEqual(['Question "2" options']);
+ });
+
it("accepts a score question with two levels", () => {
expect(
validateSystemOnePayload(
@@ -346,11 +354,19 @@ describe("blank context", () => {
expect(validateSystemOnePayload(raw, "/v1/systemone").issues.map((issue) => issue.path)).toEqual(["state"]);
const filled = validateSystemOnePayload(raw.replace('"state":""', '"state":"A ticket"'), "/v1/systemone");
expect(filled.isValid).toBe(false);
- expect(filled.issues).toContainEqual({
- path: "questions",
- message: "Give every question a name.",
- severity: "error",
+ expect(filled.issues).toEqual([
+ { path: "questions", label: "Questions", message: "Give every question a name.", severity: "error" },
+ ]);
+ });
+
+ it("labels a /v1/systemone question by its key even when the key is all digits", () => {
+ const raw = JSON.stringify({
+ state: "A ticket",
+ questions: { "2": { type: "choice", instructions: "", criteria: {} } },
});
+ const labels = validateSystemOnePayload(raw, "/v1/systemone").issues.map((issue) => issue.label);
+ expect(labels).toContain('Question "2" options');
+ expect(labels.some((label) => label.startsWith("Question 3"))).toBe(false);
});
it("accepts /v1/systemone once the state carries text", () => {
@@ -365,6 +381,25 @@ describe("blank context", () => {
expect(result.issues.map((issue) => issue.path)).toEqual(["input"]);
});
+ it("labels a /v1/decisions issue by the question's position", () => {
+ const duplicateChoices = [{ value: "a" }, { value: "a" }];
+ const raw = JSON.stringify({
+ input: "Ship it?",
+ questions: [
+ { type: "predicate", name: "q", instructions: "Is it?" },
+ { type: "choice", name: "pick", instructions: "Pick one", choices: duplicateChoices },
+ ],
+ });
+ expect(validateOpenAIDecisionsPayload(raw).issues).toEqual([
+ {
+ path: "questions.1.choices",
+ label: "Question 2 choices",
+ message: "Choice values must be unique, including their string representations.",
+ severity: "error",
+ },
+ ]);
+ });
+
it("accepts /v1/decisions once the input carries text", () => {
const raw = JSON.stringify({
input: "Ship it?",
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts
index 1fb32efd158..05ed0f14511 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/systemOneUI/lib/validatePayload.ts
@@ -1,4 +1,5 @@
import { z } from "zod";
+import { issueLabel } from "./issueLabel";
import { openAIDecisionsRequestSchema, type OpenAIDecisionsRequest } from "./openAIDecisions";
import {
decisionsRequestSchema,
@@ -11,6 +12,7 @@ const RECOMMENDED_MAX_SCORE_LEVELS = 10;
export interface SystemOnePayloadIssue {
path: string;
+ label: string;
message: string;
severity: "error" | "warning";
}
@@ -26,21 +28,28 @@ export type SystemOnePayloadValidation = PayloadValidation;
export const errorCount = (validation: PayloadValidation): number =>
validation.issues.filter((issue) => issue.severity === "error").length;
+const error = (path: readonly PropertyKey[], message: string): SystemOnePayloadIssue => ({
+ path: path.length > 0 ? path.join(".") : "root",
+ label: issueLabel(path),
+ message,
+ severity: "error",
+});
+
const toIssue = (path: PropertyKey[], message: string): SystemOnePayloadIssue =>
path[0] === "questions" && path[1] === ""
- ? { path: "questions", message: "Give every question a name.", severity: "error" }
- : { path: path.length > 0 ? path.join(".") : "root", message, severity: "error" };
+ ? error(["questions"], "Give every question a name.")
+ : error(path, message);
const invalid = (path: string, message: string) => ({
isValid: false,
- issues: [{ path, message, severity: "error" as const }],
+ issues: [error([path], message)],
});
export function parseJson(raw: string): { ok: true; value: unknown } | { ok: false; message: string } {
try {
return { ok: true, value: JSON.parse(raw) };
- } catch (error: unknown) {
- return { ok: false, message: error instanceof Error ? error.message : String(error) };
+ } catch (cause: unknown) {
+ return { ok: false, message: cause instanceof Error ? cause.message : String(cause) };
}
}
@@ -49,6 +58,7 @@ const scoreLevelWarnings = (payload: PlaygroundRequest): SystemOnePayloadIssue[]
.filter(([, question]) => question.type === "score" && question.criteria.length > RECOMMENDED_MAX_SCORE_LEVELS)
.map(([id]) => ({
path: `questions.${id}.criteria`,
+ label: issueLabel(["questions", id, "criteria"]),
message: `More than ${RECOMMENDED_MAX_SCORE_LEVELS} score levels may reduce result quality.`,
severity: "warning",
}));
@@ -90,7 +100,7 @@ const requireContext = (raw: string, path: "state" | "input", result: Payload
? {
...result,
isValid: false,
- issues: [{ path, message: "Add the text the model should decide about.", severity: "error" }],
+ issues: [error([path], "Add the text the model should decide about.")],
}
: result;
From 6b7c77b7d309daca2cebecf2871752ee1623d6e1 Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 14 Oct 2026 10:50:11 -0700
Subject: [PATCH 16/47] fix(http): restore the 15s aiohttp keep-alive default
in the Python and Rust clients (#46355)
* fix: lower aiohttp keepalive timeout
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(http): restore the 15s keep-alive default and apply it to the Rust client
15s is aiohttp's own default and what every release before v1.78.0 used. The Rust HTTP client reads the same AIOHTTP_KEEPALIVE_TIMEOUT variable, so its default moves with it.
* test(integration): cover upstream connections dropped after the keep-alive idle window
---------
Co-authored-by: nate
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm-rust/crates/http/src/settings.rs | 4 +-
litellm/constants.py | 2 +-
tests/integration/_support/idle_gateway.py | 207 +++++
.../providers/_idle_upstream_support.py | 753 ++++++++++++++++++
...test_idle_upstream_keepalive_edges_wire.py | 89 +++
...t_idle_upstream_keepalive_override_wire.py | 56 ++
...test_idle_upstream_reconnect_chaos_wire.py | 135 ++++
...est_idle_upstream_reconnect_routes_wire.py | 60 ++
...e_upstream_reconnect_rust_messages_wire.py | 56 ++
.../test_idle_upstream_reconnect_wire.py | 62 ++
.../sdk/test_idle_upstream_reconnect_sdk.py | 94 +++
11 files changed, 1515 insertions(+), 3 deletions(-)
create mode 100644 tests/integration/_support/idle_gateway.py
create mode 100644 tests/integration/providers/_idle_upstream_support.py
create mode 100644 tests/integration/providers/test_idle_upstream_keepalive_edges_wire.py
create mode 100644 tests/integration/providers/test_idle_upstream_keepalive_override_wire.py
create mode 100644 tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py
create mode 100644 tests/integration/providers/test_idle_upstream_reconnect_routes_wire.py
create mode 100644 tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py
create mode 100644 tests/integration/providers/test_idle_upstream_reconnect_wire.py
create mode 100644 tests/integration/sdk/test_idle_upstream_reconnect_sdk.py
diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs
index e1edc6d37e1..1f437df225c 100644
--- a/litellm-rust/crates/http/src/settings.rs
+++ b/litellm-rust/crates/http/src/settings.rs
@@ -140,7 +140,7 @@ impl Default for HttpSettings {
proxies: EnvironmentProxies::default(),
connect_timeout: Duration::from_secs(10),
tcp_keepalive: None,
- pool_idle_timeout: Duration::from_secs(120),
+ pool_idle_timeout: Duration::from_secs(15),
}
}
}
@@ -330,7 +330,7 @@ mod tests {
&[("AIOHTTP_TCP_KEEPIDLE", "45")],
))]);
assert_eq!(settings.tcp_keepalive, None);
- assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
+ assert_eq!(settings.pool_idle_timeout, Duration::from_secs(15));
}
fn proxy_flags(
diff --git a/litellm/constants.py b/litellm/constants.py
index a7ea4a2975b..aed2b04b16b 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -345,7 +345,7 @@ EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING: Final = 10_000
# Set to 0 for unlimited (not recommended for production)
AIOHTTP_CONNECTOR_LIMIT: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000))
AIOHTTP_CONNECTOR_LIMIT_PER_HOST: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500))
-AIOHTTP_KEEPALIVE_TIMEOUT: Final = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120))
+AIOHTTP_KEEPALIVE_TIMEOUT: Final = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 15))
AIOHTTP_TTL_DNS_CACHE: Final = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
# TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs
# whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT
diff --git a/tests/integration/_support/idle_gateway.py b/tests/integration/_support/idle_gateway.py
new file mode 100644
index 00000000000..5a86d7e0a35
--- /dev/null
+++ b/tests/integration/_support/idle_gateway.py
@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import itertools
+import re
+import socket
+import struct
+import threading
+import time
+from collections.abc import Callable, Generator, Mapping
+from contextlib import contextmanager, suppress
+from dataclasses import dataclass
+from enum import Enum
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from io import BufferedIOBase
+from queue import SimpleQueue
+from socketserver import BaseServer
+from typing import Final
+
+from integration._support.wire import Reply, Request
+
+MARKER: Final = re.compile(rb"idle-[0-9a-f]{16}")
+_LOOPBACK: Final = "127.0.0.1"
+
+
+def _exactly(stream: BufferedIOBase, size: int) -> bytes | None:
+ data: Final = stream.read(size)
+ return data if len(data) == size else None
+
+
+def _body(headers: Mapping[str, str], stream: BufferedIOBase) -> bytes | None:
+ if headers.get("transfer-encoding", "").lower() == "chunked":
+ return None
+ return _exactly(stream, int(headers.get("content-length", "0")))
+
+
+class Drop(Enum):
+ CLOSE = "close"
+ RESET = "reset"
+
+
+@dataclass(frozen=True, slots=True)
+class Arrival:
+ marker: str
+ target: str
+ connection: int
+ dropped: bool
+
+
+@dataclass(frozen=True, slots=True)
+class IdleGateway:
+ """A keep-alive HTTP/1.1 peer that never closes an idle connection itself and, like a load balancer whose
+ idle timer already fired, swallows a request that reuses a connection idle for `drop_after` seconds."""
+
+ url: str
+ port: int
+ arrivals: SimpleQueue[Arrival]
+
+ def pending(self) -> int:
+ return self.arrivals.qsize()
+
+ def drain(self) -> tuple[Arrival, ...]:
+ return tuple(self.arrivals.get_nowait() for _ in range(self.arrivals.qsize()))
+
+
+class _Server(ThreadingHTTPServer):
+ daemon_threads = True
+ block_on_close = False
+ request_queue_size = 128
+
+ def __init__(
+ self,
+ port: int,
+ respond: Callable[[Request], Reply],
+ drop_after: float | None,
+ drop: Drop,
+ arrivals: SimpleQueue[Arrival],
+ ) -> None:
+ super().__init__((_LOOPBACK, port), _Handler)
+ self.respond = respond
+ self.drop_after = drop_after
+ self.drop = drop
+ self.arrivals = arrivals
+ self.connections = itertools.count(1)
+ self.sockets: SimpleQueue[socket.socket] = SimpleQueue()
+ self.failures: SimpleQueue[Exception] = SimpleQueue()
+
+ def close_connections(self) -> None:
+ for client in tuple(self.sockets.get_nowait() for _ in range(self.sockets.qsize())):
+ with suppress(OSError):
+ client.shutdown(socket.SHUT_RDWR)
+
+ def handle_error(self, request: object, client_address: object) -> None:
+ return
+
+
+class _Handler(BaseHTTPRequestHandler):
+ protocol_version = "HTTP/1.1"
+
+ def __init__(
+ self, request: socket.socket | tuple[bytes, socket.socket], client_address: object, server: BaseServer
+ ) -> None:
+ assert isinstance(server, _Server)
+ assert isinstance(request, socket.socket)
+ self.gateway = server
+ self.client = request
+ super().__init__(request, client_address, server)
+
+ def setup(self) -> None:
+ super().setup()
+ self.connection_id = next(self.gateway.connections)
+ self.idle_since: float | None = None
+ self.arrived_at = time.monotonic()
+ self.gateway.sockets.put(self.client)
+
+ def parse_request(self) -> bool:
+ self.arrived_at = time.monotonic()
+ return super().parse_request()
+
+ def _stale(self) -> bool:
+ return (
+ self.gateway.drop_after is not None
+ and self.idle_since is not None
+ and self.arrived_at - self.idle_since >= self.gateway.drop_after
+ )
+
+ def answer(self) -> None:
+ headers: Final = {name.lower(): value for name, value in self.headers.items()}
+ body: Final = _body(headers, self.rfile)
+ if body is None:
+ self.gateway.failures.put(ValueError(f"unreadable body for {self.path}: {headers}"))
+ self.close_connection = True
+ return
+ found: Final = MARKER.search(body) or MARKER.search(self.path.encode())
+ stale: Final = self._stale()
+ self.gateway.arrivals.put(
+ Arrival(found.group().decode() if found else "", self.path, self.connection_id, stale)
+ )
+ if stale:
+ self._drop()
+ return
+ try:
+ reply: Final = self.gateway.respond(Request(self.command, self.path, headers, body))
+ except (ValueError, KeyError, TypeError) as error:
+ self.gateway.failures.put(error)
+ self._write(Reply(status=500))
+ return
+ self._write(reply)
+ self.idle_since = time.monotonic()
+
+ do_POST = answer
+ do_GET = answer
+
+ def _drop(self) -> None:
+ self.close_connection = True
+ match self.gateway.drop:
+ case Drop.CLOSE:
+ with suppress(OSError):
+ self.client.shutdown(socket.SHUT_RDWR)
+ case Drop.RESET:
+ self.client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
+ self.rfile.close()
+ self.client.close()
+
+ def _write(self, reply: Reply) -> None:
+ self.send_response(reply.status)
+ self.send_header("content-type", reply.content_type)
+ for name, value in reply.headers.items():
+ self.send_header(name, value)
+ if reply.chunks is None:
+ self.send_header("content-length", str(len(reply.body)))
+ else:
+ self.send_header("transfer-encoding", "chunked")
+ self.end_headers()
+ if reply.chunks is None:
+ self.wfile.write(reply.body)
+ return
+ for index, chunk in enumerate(reply.chunks):
+ if index and reply.pause_between_chunks:
+ time.sleep(reply.pause_between_chunks)
+ self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk))
+ self.wfile.write(b"0\r\n\r\n")
+
+ def log_message(self, format: str, *args: object) -> None:
+ return
+
+
+@contextmanager
+def idle_gateway(
+ respond: Callable[[Request], Reply],
+ *,
+ drop_after: float | None,
+ drop: Drop = Drop.CLOSE,
+ port: int = 0,
+) -> Generator[IdleGateway]:
+ arrivals: Final[SimpleQueue[Arrival]] = SimpleQueue()
+ with _Server(port, respond, drop_after, drop, arrivals) as server:
+ thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
+ thread.start()
+ try:
+ yield IdleGateway(f"http://{_LOOPBACK}:{server.server_port}", server.server_port, arrivals)
+ finally:
+ server.shutdown()
+ thread.join(timeout=6)
+ server.close_connections()
+ assert not thread.is_alive(), "Idle gateway survived cleanup"
+ failure: Final = None if server.failures.empty() else server.failures.get_nowait()
+ assert failure is None, f"Idle gateway failed to answer: {failure!r}"
diff --git a/tests/integration/providers/_idle_upstream_support.py b/tests/integration/providers/_idle_upstream_support.py
new file mode 100644
index 00000000000..ca808abc008
--- /dev/null
+++ b/tests/integration/providers/_idle_upstream_support.py
@@ -0,0 +1,753 @@
+from __future__ import annotations
+
+import asyncio
+import io
+import json
+import uuid
+import wave
+from collections.abc import Callable, Generator, Mapping
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack, contextmanager
+from dataclasses import dataclass
+from itertools import chain
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import anthropic
+import httpx
+import openai
+import yaml
+from anthropic.types import RawContentBlockDeltaEvent, RawMessageStartEvent, TextBlock, TextDelta
+from integration._support.anthropic_sse import message_json, message_stream, stream_reply
+from integration._support.client import Gateway, eventually, held, list_value, object_value, string_value
+from integration._support.database import read_rows
+from integration._support.idle_gateway import MARKER, Arrival, Drop, IdleGateway, idle_gateway
+from integration._support.openai_wire import chat_reply, chat_stream, responses_reply
+from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process
+from integration._support.wire import Reply, Request
+from openai.types.chat import ChatCompletionChunk
+from openai.types.responses import ResponseCompletedEvent, ResponseTextDeltaEvent
+from pydantic import JsonValue, TypeAdapter
+
+IDLE_SECONDS: Final = 17.0
+DROP_AFTER_SECONDS: Final = 16.0
+REUSE_IDLE_SECONDS: Final = 5.0
+QUIET_STREAM_SECONDS: Final = 16.0
+PROVIDER_KEY: Final = "sk-idle-upstream-provider"
+CALL_ID: Final = "x-litellm-call-id"
+RUST_SERVED: Final = "x-litellm-rust"
+GUARDRAIL: Final = "idle-upstream-guardrail"
+PASS_THROUGH: Final = "/idle-upstream-pass-through"
+EMBEDDING: Final = (0.25, 0.5)
+FRESH_CONNECTIONS: Final = httpx.Limits(max_keepalive_connections=0)
+_JSON: Final = TypeAdapter(dict[str, JsonValue])
+_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
+
+
+def new_marker() -> str:
+ return f"idle-{uuid.uuid4().hex[:16]}"
+
+
+def expected_text(marker: str) -> str:
+ return f"answer {marker}"
+
+
+def expected_embedding(_marker: str) -> str:
+ return json.dumps(list(EMBEDDING))
+
+
+def _marker_of(request: Request) -> str:
+ found: Final = MARKER.search(request.body)
+ return found.group().decode() if found else ""
+
+
+def _converse(text: str) -> bytes:
+ return json.dumps(
+ {
+ "output": {"message": {"role": "assistant", "content": [{"text": text}]}},
+ "stopReason": "end_turn",
+ "usage": {"inputTokens": 5, "outputTokens": 3, "totalTokens": 8},
+ "metrics": {"latencyMs": 1},
+ }
+ ).encode()
+
+
+def _ocr(text: str) -> bytes:
+ return json.dumps(
+ {
+ "model": "mistral-ocr-latest",
+ "pages": [{"index": 0, "markdown": text}],
+ "usage_info": {"pages_processed": 1, "doc_size_bytes": 1024},
+ }
+ ).encode()
+
+
+def _embeddings() -> bytes:
+ return json.dumps(
+ {
+ "object": "list",
+ "data": [{"object": "embedding", "index": 0, "embedding": list(EMBEDDING)}],
+ "model": "text-embedding-idle-upstream",
+ "usage": {"prompt_tokens": 1, "total_tokens": 1},
+ }
+ ).encode()
+
+
+def _model_reply(target: str, marker: str, body: Mapping[str, JsonValue]) -> Reply:
+ model: Final = str(body.get("model", "idle-upstream"))
+ stream: Final = body.get("stream") is True
+ if target.endswith("/v1/messages"):
+ if stream:
+ return stream_reply(message_stream(f"msg_{marker}", model, expected_text(marker)))
+ return Reply(body=message_json(f"msg_{marker}", model, expected_text(marker)))
+ if target.endswith("/responses"):
+ return responses_reply(f"resp_{marker}", model, expected_text(marker), stream=stream)
+ if target.endswith("/chat/completions"):
+ return chat_reply(f"chatcmpl-{marker}", model, expected_text(marker), stream=stream)
+ return Reply(status=404, body=b'{"error": "unscripted route"}')
+
+
+def answer(request: Request) -> Reply:
+ marker: Final = _marker_of(request)
+ target: Final = request.target.split("?", 1)[0]
+ if target.endswith("/beta/litellm_basic_guardrail_api"):
+ return Reply(body=json.dumps({"action": "NONE"}).encode())
+ if target.startswith("/passthrough"):
+ return Reply(body=json.dumps({"echo": expected_text(marker)}).encode())
+ if target.endswith("/converse"):
+ return Reply(body=_converse(expected_text(marker)))
+ if target.endswith("/ocr"):
+ return Reply(body=_ocr(expected_text(marker)))
+ if target.endswith("/embeddings"):
+ return Reply(body=_embeddings())
+ return _model_reply(target, marker, _JSON.validate_json(request.body))
+
+
+def quiet_stream_answer(request: Request) -> Reply:
+ marker: Final = _marker_of(request)
+ body: Final = _JSON.validate_json(request.body)
+ model: Final = str(body.get("model", "idle-upstream"))
+ frames: Final = (
+ message_stream(f"msg_{marker}", model, expected_text(marker))
+ if request.target.endswith("/v1/messages")
+ else chat_stream(f"chatcmpl-{marker}", model, expected_text(marker))
+ )
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(frames[0], b"".join(frames[1:])),
+ pause_between_chunks=QUIET_STREAM_SECONDS,
+ )
+
+
+@dataclass(frozen=True, slots=True)
+class ProxyTarget:
+ url: str
+ key: str
+
+
+@dataclass(frozen=True, slots=True)
+class Answer:
+ status: int
+ text: str
+ call_id: str
+ response_id: str
+ rust: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class Sent:
+ marker: str
+ answer: Answer
+
+
+def _call_id(headers: httpx.Headers) -> str:
+ return next(iter(headers.get_list(CALL_ID)), "")
+
+
+def _rust(headers: httpx.Headers) -> bool:
+ return "true" in headers.get_list(RUST_SERVED)
+
+
+def _refused(status: int, response: httpx.Response) -> Answer:
+ return Answer(status, response.text, _call_id(response.headers), "")
+
+
+def _auth(target: ProxyTarget) -> Mapping[str, str]:
+ return {"Authorization": f"Bearer {target.key}"}
+
+
+def _openai(target: ProxyTarget) -> openai.OpenAI:
+ return openai.OpenAI(
+ base_url=f"{target.url}/v1",
+ api_key=target.key,
+ max_retries=0,
+ http_client=httpx.Client(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False),
+ )
+
+
+def _async_openai(target: ProxyTarget) -> openai.AsyncOpenAI:
+ return openai.AsyncOpenAI(
+ base_url=f"{target.url}/v1",
+ api_key=target.key,
+ max_retries=0,
+ http_client=httpx.AsyncClient(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False),
+ )
+
+
+def _anthropic(target: ProxyTarget) -> anthropic.Anthropic:
+ return anthropic.Anthropic(
+ base_url=target.url,
+ api_key=target.key,
+ max_retries=0,
+ http_client=httpx.Client(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False),
+ )
+
+
+def _async_anthropic(target: ProxyTarget) -> anthropic.AsyncAnthropic:
+ return anthropic.AsyncAnthropic(
+ base_url=target.url,
+ api_key=target.key,
+ max_retries=0,
+ http_client=httpx.AsyncClient(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False),
+ )
+
+
+def _chat(target: ProxyTarget, model: str, marker: str, extra: Mapping[str, JsonValue]) -> Answer:
+ with _openai(target) as client:
+ try:
+ raw: Final = client.chat.completions.with_raw_response.create(
+ model=model, messages=[{"role": "user", "content": marker}], extra_body=dict(extra)
+ )
+ except openai.APIStatusError as error:
+ return _refused(error.status_code, error.response)
+ completion: Final = raw.parse()
+ return Answer(200, completion.choices[0].message.content or "", _call_id(raw.headers), completion.id)
+
+
+def chat_sdk(target: ProxyTarget, model: str, marker: str) -> Answer:
+ return _chat(target, model, marker, {})
+
+
+def chat_sdk_with_default_retries(target: ProxyTarget, model: str, marker: str) -> Answer:
+ return _chat(target, model, marker, {"num_retries": 2})
+
+
+def _chunk_text(chunk: ChatCompletionChunk) -> str:
+ return "".join(choice.delta.content or "" for choice in chunk.choices)
+
+
+async def _chat_stream(target: ProxyTarget, model: str, marker: str) -> Answer:
+ async with _async_openai(target) as client:
+ try:
+ raw: Final = await client.chat.completions.with_raw_response.create(
+ model=model, messages=[{"role": "user", "content": marker}], stream=True
+ )
+ except openai.APIStatusError as error:
+ return _refused(error.status_code, error.response)
+ try:
+ chunks: Final = tuple([chunk async for chunk in raw.parse()])
+ except openai.APIError as error:
+ return Answer(200, f"stream broke: {error}", _call_id(raw.headers), "")
+ return Answer(
+ 200,
+ "".join(_chunk_text(chunk) for chunk in chunks),
+ _call_id(raw.headers),
+ chunks[0].id if chunks else "",
+ )
+
+
+def chat_sdk_stream(target: ProxyTarget, model: str, marker: str) -> Answer:
+ return asyncio.run(_chat_stream(target, model, marker))
+
+
+def _post(
+ target: ProxyTarget,
+ path: str,
+ body: Mapping[str, JsonValue],
+ read: Callable[[Mapping[str, JsonValue]], tuple[str, str]],
+) -> Answer:
+ with httpx.Client(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False) as client:
+ response: Final = client.post(f"{target.url}{path}", json=dict(body), headers=dict(_auth(target)))
+ if response.status_code != 200:
+ return _refused(response.status_code, response)
+ text, identity = read(_JSON.validate_json(response.content))
+ return Answer(200, text, _call_id(response.headers), identity, _rust(response.headers))
+
+
+def _chat_body(body: Mapping[str, JsonValue]) -> tuple[str, str]:
+ message: Final = object_value(object_value(list_value(body["choices"])[0])["message"])
+ return string_value(message["content"]), string_value(body["id"])
+
+
+def _echo_body(body: Mapping[str, JsonValue]) -> tuple[str, str]:
+ return string_value(body["echo"]), ""
+
+
+def _ocr_body(body: Mapping[str, JsonValue]) -> tuple[str, str]:
+ return string_value(object_value(list_value(body["pages"])[0])["markdown"]), ""
+
+
+def chat_httpx(target: ProxyTarget, model: str, marker: str) -> Answer:
+ return _post(
+ target, "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]}, _chat_body
+ )
+
+
+def chat_with_guardrail(target: ProxyTarget, model: str, marker: str) -> Answer:
+ return _post(
+ target,
+ "/v1/chat/completions",
+ {"model": model, "messages": [{"role": "user", "content": marker}], "guardrails": [GUARDRAIL]},
+ _chat_body,
+ )
+
+
+def pass_through(target: ProxyTarget, _model: str, marker: str) -> Answer:
+ return _post(target, PASS_THROUGH, {"input": marker}, _echo_body)
+
+
+def ocr(target: ProxyTarget, model: str, marker: str) -> Answer:
+ return _post(
+ target,
+ "/v1/ocr",
+ {
+ "model": model,
+ "document": {"type": "document_url", "document_url": f"https://documents.invalid/{marker}.pdf"},
+ },
+ _ocr_body,
+ )
+
+
+def _wav() -> bytes:
+ buffer: Final = io.BytesIO()
+ with wave.open(buffer, "wb") as writer:
+ writer.setnchannels(1)
+ writer.setsampwidth(2)
+ writer.setframerate(8000)
+ writer.writeframes(b"\x00\x00" * 800)
+ return buffer.getvalue()
+
+
+def transcription(target: ProxyTarget, model: str, marker: str) -> Answer:
+ with httpx.Client(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False) as client:
+ response: Final = client.post(
+ f"{target.url}/v1/audio/transcriptions",
+ files={"file": ("idle.wav", _wav(), "audio/wav")},
+ data={"model": model, "prompt": marker},
+ headers=dict(_auth(target)),
+ )
+ if response.status_code != 200:
+ return _refused(response.status_code, response)
+ return Answer(
+ 200,
+ string_value(_JSON.validate_json(response.content)["text"]),
+ _call_id(response.headers),
+ "",
+ _rust(response.headers),
+ )
+
+
+def embeddings_sdk(target: ProxyTarget, model: str, marker: str) -> Answer:
+ with _openai(target) as client:
+ try:
+ raw: Final = client.embeddings.with_raw_response.create(model=model, input=marker, encoding_format="float")
+ except openai.APIStatusError as error:
+ return _refused(error.status_code, error.response)
+ result: Final = raw.parse()
+ return Answer(200, json.dumps(result.data[0].embedding), _call_id(raw.headers), "")
+
+
+def responses_sdk(target: ProxyTarget, model: str, marker: str) -> Answer:
+ with _openai(target) as client:
+ try:
+ raw: Final = client.responses.with_raw_response.create(model=model, input=marker)
+ except openai.APIStatusError as error:
+ return _refused(error.status_code, error.response)
+ response: Final = raw.parse()
+ return Answer(200, response.output_text, _call_id(raw.headers), response.id)
+
+
+async def _responses_stream(target: ProxyTarget, model: str, marker: str) -> Answer:
+ async with _async_openai(target) as client:
+ try:
+ raw: Final = await client.responses.with_raw_response.create(model=model, input=marker, stream=True)
+ except openai.APIStatusError as error:
+ return _refused(error.status_code, error.response)
+ try:
+ events: Final = tuple([event async for event in raw.parse()])
+ except openai.APIError as error:
+ return Answer(200, f"stream broke: {error}", _call_id(raw.headers), "")
+ return Answer(
+ 200,
+ "".join(event.delta for event in events if isinstance(event, ResponseTextDeltaEvent)),
+ _call_id(raw.headers),
+ next((event.response.id for event in events if isinstance(event, ResponseCompletedEvent)), ""),
+ )
+
+
+def responses_sdk_stream(target: ProxyTarget, model: str, marker: str) -> Answer:
+ return asyncio.run(_responses_stream(target, model, marker))
+
+
+def messages_sdk(target: ProxyTarget, model: str, marker: str) -> Answer:
+ with _anthropic(target) as client:
+ try:
+ raw: Final = client.messages.with_raw_response.create(
+ model=model, max_tokens=32, messages=[{"role": "user", "content": marker}]
+ )
+ except anthropic.APIStatusError as error:
+ return _refused(error.status_code, error.response)
+ message: Final = raw.parse()
+ return Answer(
+ 200,
+ "".join(block.text for block in message.content if isinstance(block, TextBlock)),
+ _call_id(raw.headers),
+ message.id,
+ _rust(raw.headers),
+ )
+
+
+def _delta_text(event: object) -> str:
+ if isinstance(event, RawContentBlockDeltaEvent) and isinstance(event.delta, TextDelta):
+ return event.delta.text
+ return ""
+
+
+async def _messages_stream(target: ProxyTarget, model: str, marker: str) -> Answer:
+ async with _async_anthropic(target) as client:
+ try:
+ raw: Final = await client.messages.with_raw_response.create(
+ model=model, max_tokens=32, messages=[{"role": "user", "content": marker}], stream=True
+ )
+ except anthropic.APIStatusError as error:
+ return _refused(error.status_code, error.response)
+ try:
+ events: Final = tuple([event async for event in raw.parse()])
+ except anthropic.APIError as error:
+ return Answer(200, f"stream broke: {error}", _call_id(raw.headers), "")
+ return Answer(
+ 200,
+ "".join(_delta_text(event) for event in events),
+ _call_id(raw.headers),
+ next((event.message.id for event in events if isinstance(event, RawMessageStartEvent)), ""),
+ _rust(raw.headers),
+ )
+
+
+def messages_sdk_stream(target: ProxyTarget, model: str, marker: str) -> Answer:
+ return asyncio.run(_messages_stream(target, model, marker))
+
+
+@dataclass(frozen=True, slots=True)
+class ConfigPart:
+ models: tuple[Mapping[str, JsonValue], ...] = ()
+ pass_through: tuple[Mapping[str, JsonValue], ...] = ()
+ guardrails: tuple[Mapping[str, JsonValue], ...] = ()
+
+
+def openai_deployment(model: str, url: str, _steady: str) -> ConfigPart:
+ return ConfigPart(
+ models=(
+ {
+ "model_name": model,
+ "litellm_params": {"model": "openai/idle-upstream", "api_base": url, "api_key": PROVIDER_KEY},
+ },
+ )
+ )
+
+
+def embedding_deployment(model: str, url: str, _steady: str) -> ConfigPart:
+ return ConfigPart(
+ models=(
+ {
+ "model_name": model,
+ "litellm_params": {
+ "model": "openai/text-embedding-idle-upstream",
+ "api_base": url,
+ "api_key": PROVIDER_KEY,
+ },
+ "model_info": {"mode": "embedding"},
+ },
+ )
+ )
+
+
+def anthropic_deployment(model: str, url: str, _steady: str) -> ConfigPart:
+ return ConfigPart(
+ models=(
+ {
+ "model_name": model,
+ "litellm_params": {"model": "anthropic/claude-idle-upstream", "api_base": url, "api_key": PROVIDER_KEY},
+ },
+ )
+ )
+
+
+def ocr_deployment(model: str, url: str, _steady: str) -> ConfigPart:
+ return ConfigPart(
+ models=(
+ {
+ "model_name": model,
+ "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_base": url, "api_key": PROVIDER_KEY},
+ },
+ )
+ )
+
+
+def transcription_deployment(model: str, url: str, _steady: str) -> ConfigPart:
+ return ConfigPart(
+ models=(
+ {
+ "model_name": model,
+ "litellm_params": {
+ "model": "bedrock/mistral.voxtral-mini-3b-2507",
+ "api_base": url,
+ "aws_access_key_id": "AKIAIDLEUPSTREAM0000",
+ "aws_secret_access_key": "idle-upstream-synthetic-secret",
+ "aws_region_name": "us-east-1",
+ },
+ "model_info": {"mode": "audio_transcription"},
+ },
+ )
+ )
+
+
+def pass_through_deployment(_model: str, url: str, _steady: str) -> ConfigPart:
+ return ConfigPart(pass_through=({"path": PASS_THROUGH, "target": f"{url}/passthrough"},))
+
+
+def guardrail_deployment(model: str, url: str, steady: str) -> ConfigPart:
+ return ConfigPart(
+ models=openai_deployment(model, steady, steady).models,
+ guardrails=(
+ {
+ "guardrail_name": GUARDRAIL,
+ "litellm_params": {
+ "guardrail": "generic_guardrail_api",
+ "mode": "pre_call",
+ "default_on": False,
+ "api_base": url,
+ "api_key": PROVIDER_KEY,
+ },
+ },
+ ),
+ )
+
+
+@dataclass(frozen=True, slots=True)
+class Row:
+ name: str
+ deploy: Callable[[str, str, str], ConfigPart]
+ send: Callable[[ProxyTarget, str, str], Answer]
+ expected: Callable[[str], str] = expected_text
+ drop: Drop = Drop.CLOSE
+ rust: bool = False
+ respond: Callable[[Request], Reply] = answer
+
+ @property
+ def model(self) -> str:
+ return f"idle-{self.name}"
+
+
+def send(row: Row, target: ProxyTarget) -> Sent:
+ marker: Final = new_marker()
+ return Sent(marker, row.send(target, row.model, marker))
+
+
+def proxy_config(directory: Path, parts: tuple[ConfigPart, ...]) -> Path:
+ base: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text()))
+ config: Final = {
+ **base,
+ "model_list": list(chain.from_iterable(part.models for part in parts)),
+ "guardrails": list(chain.from_iterable(part.guardrails for part in parts)),
+ "general_settings": {
+ **object_value(base["general_settings"]),
+ "pass_through_endpoints": list(chain.from_iterable(part.pass_through for part in parts)),
+ },
+ "router_settings": {**object_value(base["router_settings"]), "num_retries": 0},
+ }
+ path: Final = directory / f"idle-upstream-{uuid.uuid4().hex}.yaml"
+ path.write_text(yaml.safe_dump(json.loads(json.dumps(config))))
+ return path
+
+
+def proxy_target(owned: OwnedProxy) -> ProxyTarget:
+ return ProxyTarget(str(owned.gateway.client.base_url).rstrip("/"), owned.gateway.key)
+
+
+def two_workers(owned: OwnedProxy) -> tuple[int, int]:
+ first, second = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=120))
+ return first, second
+
+
+def on_each_worker(
+ owned: OwnedProxy, workers: tuple[int, int], call: Callable[[], tuple[Sent, ...]]
+) -> tuple[Sent, ...]:
+ first, second = workers
+ with owned.suspended(second):
+ on_first: Final = call()
+ with owned.suspended(first):
+ on_second: Final = call()
+ return (*on_first, *on_second)
+
+
+def stay_idle(gateways: tuple[IdleGateway, ...], seconds: float) -> None:
+ held(
+ lambda: sum(gateway.pending() for gateway in gateways),
+ lambda arrived: arrived == 0,
+ holding=seconds,
+ seconds=seconds + 30,
+ )
+
+
+def answered(sents: tuple[Sent, ...]) -> tuple[tuple[int, str], ...]:
+ return tuple((sent.answer.status, sent.answer.text) for sent in sents)
+
+
+def expected_answers(row: Row, sents: tuple[Sent, ...]) -> tuple[tuple[int, str], ...]:
+ return tuple((200, row.expected(sent.marker)) for sent in sents)
+
+
+def served_by_rust(sents: tuple[Sent, ...]) -> tuple[bool, ...]:
+ return tuple(sent.answer.rust for sent in sents)
+
+
+def connections(arrivals: tuple[Arrival, ...]) -> tuple[int, ...]:
+ return tuple(sorted(arrival.connection for arrival in arrivals))
+
+
+def landed(arrivals: tuple[Arrival, ...]) -> tuple[tuple[str, bool], ...]:
+ return tuple(sorted((arrival.marker, arrival.dropped) for arrival in arrivals))
+
+
+def landed_once_undropped(sents: tuple[Sent, ...]) -> tuple[tuple[str, bool], ...]:
+ return tuple(sorted((sent.marker, False) for sent in sents))
+
+
+def spend_statuses(sents: tuple[Sent, ...]) -> tuple[str, ...]:
+ call_ids: Final = ",".join(sorted(sent.answer.call_id for sent in sents if sent.answer.call_id))
+ rows: Final = eventually(
+ lambda: read_rows(
+ 'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = ANY(string_to_array(%s, %s))',
+ (call_ids, ","),
+ ),
+ lambda found: len(found) >= len(sents),
+ seconds=90,
+ return_last_on_timeout=True,
+ )
+ return tuple(sorted(string_value(row["status"]) for row in rows))
+
+
+@dataclass(frozen=True, slots=True)
+class UpstreamProxy:
+ owned: OwnedProxy
+ workers: tuple[int, int]
+ rows: Mapping[str, Row]
+ gateways: Mapping[str, IdleGateway]
+
+ @property
+ def target(self) -> ProxyTarget:
+ return proxy_target(self.owned)
+
+ def probe(self, name: str) -> tuple[tuple[Sent, ...], tuple[Arrival, ...]]:
+ row: Final = self.rows[name]
+ target: Final = self.target
+ sents: Final = on_each_worker(self.owned, self.workers, lambda: (send(row, target),))
+ return sents, self.gateways[name].drain()
+
+ def burst(self, names: tuple[str, ...], each: int) -> tuple[tuple[Row, Sent], ...]:
+ rows: Final = tuple(chain.from_iterable((self.rows[name],) * each for name in names))
+ target: Final = self.target
+ with ThreadPoolExecutor(max_workers=len(rows)) as pool:
+ return tuple(zip(rows, pool.map(send, rows, (target,) * len(rows)), strict=True))
+
+
+@dataclass(frozen=True, slots=True)
+class StalePools:
+ proxy: UpstreamProxy
+ warmed: Mapping[str, tuple[Arrival, ...]]
+
+ @property
+ def rows(self) -> Mapping[str, Row]:
+ return self.proxy.rows
+
+ def probe(self, name: str) -> tuple[tuple[Sent, ...], tuple[Arrival, ...]]:
+ return self.proxy.probe(name)
+
+
+def burst_answers(sent: tuple[tuple[Row, Sent], ...]) -> tuple[tuple[int, str], ...]:
+ return tuple((one.answer.status, one.answer.text) for _, one in sent)
+
+
+def burst_expected(sent: tuple[tuple[Row, Sent], ...]) -> tuple[tuple[int, str], ...]:
+ return tuple((200, row.expected(one.marker)) for row, one in sent)
+
+
+def drain_all(gateways: tuple[IdleGateway, ...]) -> tuple[Arrival, ...]:
+ return tuple(chain.from_iterable(gateway.drain() for gateway in gateways))
+
+
+@contextmanager
+def proxy_for(
+ gateway: Gateway,
+ directory: Path,
+ rows: tuple[Row, ...],
+ gateways: Mapping[str, IdleGateway],
+ steady: str,
+ overrides: Mapping[str, str],
+) -> Generator[UpstreamProxy]:
+ config: Final = proxy_config(
+ directory, tuple(row.deploy(row.model, gateways[row.name].url, steady) for row in rows)
+ )
+ with owned_proxy_process(
+ gateway, directory, overrides, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT
+ ) as owned:
+ yield UpstreamProxy(owned, two_workers(owned), MappingProxyType({row.name: row for row in rows}), gateways)
+
+
+@contextmanager
+def upstream_proxy(
+ gateway: Gateway,
+ directory: Path,
+ rows: tuple[Row, ...],
+ overrides: Mapping[str, str],
+ drop_after: float | None = DROP_AFTER_SECONDS,
+) -> Generator[UpstreamProxy]:
+ with ExitStack() as stack:
+ steady: Final = stack.enter_context(idle_gateway(answer, drop_after=None))
+ gateways: Final = MappingProxyType(
+ {
+ row.name: stack.enter_context(idle_gateway(row.respond, drop_after=drop_after, drop=row.drop))
+ for row in rows
+ }
+ )
+ yield stack.enter_context(proxy_for(gateway, directory, rows, gateways, steady.url, overrides))
+
+
+def _warm(rows: tuple[Row, ...], target: ProxyTarget) -> tuple[Sent, ...]:
+ return tuple(send(row, target) for row in rows)
+
+
+def warm_and_idle(proxy: UpstreamProxy) -> StalePools:
+ """Gives every row one pooled upstream connection per worker, then leaves each idle for IDLE_SECONDS: past the
+ 15 s default keep-alive, inside the 120 s one it replaced."""
+ target: Final = proxy.target
+ warm: Final = on_each_worker(proxy.owned, proxy.workers, lambda: _warm(tuple(proxy.rows.values()), target))
+ refused: Final = tuple(sent for sent in warm if sent.answer.status != 200)
+ assert refused == (), refused
+ warmed: Final = MappingProxyType({name: upstream.drain() for name, upstream in proxy.gateways.items()})
+ stay_idle(tuple(proxy.gateways.values()), IDLE_SECONDS)
+ return StalePools(proxy, warmed)
+
+
+@contextmanager
+def stale_pools(
+ gateway: Gateway,
+ directory: Path,
+ rows: tuple[Row, ...],
+ overrides: Mapping[str, str],
+ drop_after: float | None = DROP_AFTER_SECONDS,
+) -> Generator[StalePools]:
+ with upstream_proxy(gateway, directory, rows, overrides, drop_after) as proxy:
+ yield warm_and_idle(proxy)
diff --git a/tests/integration/providers/test_idle_upstream_keepalive_edges_wire.py b/tests/integration/providers/test_idle_upstream_keepalive_edges_wire.py
new file mode 100644
index 00000000000..59a3d492be7
--- /dev/null
+++ b/tests/integration/providers/test_idle_upstream_keepalive_edges_wire.py
@@ -0,0 +1,89 @@
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.process import graceful_stop_seconds, refused_boot_log
+from integration.providers._idle_upstream_support import (
+ QUIET_STREAM_SECONDS,
+ REUSE_IDLE_SECONDS,
+ Row,
+ UpstreamProxy,
+ answered,
+ anthropic_deployment,
+ chat_sdk,
+ chat_sdk_stream,
+ connections,
+ expected_answers,
+ landed,
+ landed_once_undropped,
+ messages_sdk,
+ messages_sdk_stream,
+ openai_deployment,
+ quiet_stream_answer,
+ served_by_rust,
+ spend_statuses,
+ stay_idle,
+ upstream_proxy,
+)
+
+pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + 2 * QUIET_STREAM_SECONDS))
+
+REUSE_ROWS: Final = (
+ Row("reuse-chat-openai-sdk-sync", openai_deployment, chat_sdk),
+ Row("reuse-messages-rust-anthropic-sdk-sync", anthropic_deployment, messages_sdk, rust=True),
+)
+QUIET_ROWS: Final = (
+ Row("quiet-chat-openai-sdk-async-stream", openai_deployment, chat_sdk_stream, respond=quiet_stream_answer),
+ Row(
+ "quiet-messages-rust-anthropic-sdk-async-stream",
+ anthropic_deployment,
+ messages_sdk_stream,
+ rust=True,
+ respond=quiet_stream_answer,
+ ),
+)
+
+
+@pytest.fixture(scope="module")
+def edges(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[UpstreamProxy]:
+ with upstream_proxy(
+ session_gateway,
+ tmp_path_factory.mktemp("idle-upstream-edges"),
+ (*REUSE_ROWS, *QUIET_ROWS),
+ {"LITELLM_RUST": "1"},
+ drop_after=None,
+ ) as proxy:
+ yield proxy
+
+
+@pytest.mark.parametrize("name", [row.name for row in REUSE_ROWS])
+def test_request_idle_under_the_keepalive_reuses_the_pooled_connection_on_each_worker(
+ edges: UpstreamProxy, name: str
+) -> None:
+ row: Final = edges.rows[name]
+ warm, warmed = edges.probe(name)
+ stay_idle((edges.gateways[name],), REUSE_IDLE_SECONDS)
+ probe, arrivals = edges.probe(name)
+ sents: Final = (*warm, *probe)
+ assert answered(sents) == expected_answers(row, sents), sents
+ assert landed((*warmed, *arrivals)) == landed_once_undropped(sents), (warmed, arrivals)
+ assert connections(arrivals) == connections(warmed), (warmed, arrivals)
+ assert served_by_rust(sents) == (row.rust,) * len(sents), sents
+
+
+@pytest.mark.parametrize("name", [row.name for row in QUIET_ROWS])
+def test_stream_quiet_for_longer_than_the_keepalive_completes_on_each_worker(edges: UpstreamProxy, name: str) -> None:
+ row: Final = edges.rows[name]
+ sents, arrivals = edges.probe(name)
+ assert answered(sents) == expected_answers(row, sents), sents
+ assert landed(arrivals) == landed_once_undropped(sents), arrivals
+ assert served_by_rust(sents) == (row.rust,) * len(sents), sents
+ assert spend_statuses(sents) == ("success",) * len(sents), sents
+
+
+@pytest.mark.parametrize("value", ["fifteen", ""], ids=["word", "empty"])
+def test_malformed_keepalive_timeout_refuses_to_boot(gateway: Gateway, tmp_path: Path, value: str) -> None:
+ log: Final = refused_boot_log(gateway, tmp_path, {"AIOHTTP_KEEPALIVE_TIMEOUT": value})
+ assert f"invalid literal for int() with base 10: {value!r}" in log, log
diff --git a/tests/integration/providers/test_idle_upstream_keepalive_override_wire.py b/tests/integration/providers/test_idle_upstream_keepalive_override_wire.py
new file mode 100644
index 00000000000..8b5a69d729e
--- /dev/null
+++ b/tests/integration/providers/test_idle_upstream_keepalive_override_wire.py
@@ -0,0 +1,56 @@
+from collections.abc import Iterator
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.process import graceful_stop_seconds
+from integration.providers._idle_upstream_support import (
+ IDLE_SECONDS,
+ Row,
+ StalePools,
+ answered,
+ anthropic_deployment,
+ chat_sdk,
+ connections,
+ expected_answers,
+ landed,
+ landed_once_undropped,
+ messages_sdk,
+ ocr,
+ ocr_deployment,
+ openai_deployment,
+ served_by_rust,
+ stale_pools,
+)
+
+pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS))
+
+ROWS: Final = (
+ Row("override-chat-openai-sdk-sync", openai_deployment, chat_sdk),
+ Row("override-messages-rust-anthropic-sdk-sync", anthropic_deployment, messages_sdk, rust=True),
+ Row("override-ocr-mistral-rust-bridge", ocr_deployment, ocr, rust=True),
+)
+
+
+@pytest.fixture(scope="module")
+def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]:
+ with stale_pools(
+ session_gateway,
+ tmp_path_factory.mktemp("idle-upstream-override"),
+ ROWS,
+ {"AIOHTTP_KEEPALIVE_TIMEOUT": "120", "LITELLM_RUST": "1"},
+ drop_after=None,
+ ) as value:
+ yield value
+
+
+@pytest.mark.parametrize("name", [row.name for row in ROWS])
+def test_explicit_keepalive_timeout_keeps_reusing_the_pooled_connection_past_the_default(
+ pools: StalePools, name: str
+) -> None:
+ row: Final = pools.rows[name]
+ sents, arrivals = pools.probe(name)
+ assert answered(sents) == expected_answers(row, sents), sents
+ assert landed(arrivals) == landed_once_undropped(sents), arrivals
+ assert connections(arrivals) == connections(pools.warmed[name]), (arrivals, pools.warmed[name])
+ assert served_by_rust(sents) == (row.rust,) * len(sents), sents
diff --git a/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py
new file mode 100644
index 00000000000..727b7ff7f5b
--- /dev/null
+++ b/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py
@@ -0,0 +1,135 @@
+from contextlib import ExitStack
+from dataclasses import replace
+from itertools import chain
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.idle_gateway import idle_gateway
+from integration._support.process import graceful_stop_seconds
+from integration.providers._idle_upstream_support import (
+ DROP_AFTER_SECONDS,
+ IDLE_SECONDS,
+ REUSE_IDLE_SECONDS,
+ Row,
+ answer,
+ answered,
+ anthropic_deployment,
+ burst_answers,
+ burst_expected,
+ chat_sdk,
+ chat_sdk_stream,
+ drain_all,
+ expected_answers,
+ landed,
+ landed_once_undropped,
+ messages_sdk,
+ messages_sdk_stream,
+ on_each_worker,
+ openai_deployment,
+ proxy_for,
+ responses_sdk,
+ responses_sdk_stream,
+ send,
+ spend_statuses,
+ stale_pools,
+ stay_idle,
+ two_workers,
+ upstream_proxy,
+ warm_and_idle,
+)
+
+BURST_EACH: Final = 4
+BURST_ROWS: Final = (
+ Row("burst-chat-openai-sdk-sync", openai_deployment, chat_sdk),
+ Row("burst-chat-openai-sdk-async-stream", openai_deployment, chat_sdk_stream),
+ Row("burst-messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk),
+ Row("burst-messages-anthropic-sdk-async-stream", anthropic_deployment, messages_sdk_stream),
+ Row("burst-responses-openai-sdk-sync", openai_deployment, responses_sdk),
+ Row("burst-responses-openai-sdk-async-stream", openai_deployment, responses_sdk_stream),
+)
+KILL_ROWS: Final = (
+ Row("kill-chat-openai-sdk-sync", openai_deployment, chat_sdk),
+ Row("kill-messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk),
+)
+RESTART_ROWS: Final = (
+ Row("restart-chat-openai-sdk-sync", openai_deployment, chat_sdk),
+ Row("restart-messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk),
+)
+
+
+@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS))
+def test_concurrent_mixed_burst_after_upstream_idle_past_keepalive_lands_every_request_once(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ with stale_pools(gateway, tmp_path, BURST_ROWS, {}) as pools:
+ sent: Final = pools.proxy.burst(tuple(row.name for row in BURST_ROWS), BURST_EACH)
+ arrivals: Final = drain_all(tuple(pools.proxy.gateways.values()))
+ sents: Final = tuple(one for _, one in sent)
+ statuses: Final = spend_statuses(sents)
+ assert burst_answers(sent) == burst_expected(sent), sent
+ assert landed(arrivals) == landed_once_undropped(sents), arrivals
+ assert statuses == ("success",) * len(sents), sents
+
+
+@pytest.mark.timeout(int(3 * graceful_stop_seconds() + 180 + IDLE_SECONDS))
+def test_worker_killed_while_the_other_serves_then_both_open_fresh_connections_after_upstream_idle(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ with upstream_proxy(gateway, tmp_path, KILL_ROWS, {}) as proxy:
+ victim, _ = proxy.workers
+ with proxy.owned.replacing_worker(victim) as kill:
+ kill()
+ during: Final = send(KILL_ROWS[0], proxy.target)
+ stale: Final = warm_and_idle(replace(proxy, workers=two_workers(proxy.owned)))
+ probes: Final = tuple(stale.probe(row.name) for row in KILL_ROWS)
+ sents: Final = tuple(chain.from_iterable(probe_sents for probe_sents, _ in probes))
+ arrivals: Final = tuple(chain.from_iterable(probe_arrivals for _, probe_arrivals in probes))
+ statuses: Final = spend_statuses(sents)
+ expected: Final = tuple(
+ chain.from_iterable(
+ expected_answers(row, probe_sents) for row, (probe_sents, _) in zip(KILL_ROWS, probes, strict=True)
+ )
+ )
+ assert answered((during,)) == expected_answers(KILL_ROWS[0], (during,)), during
+ assert answered(sents) == expected, sents
+ assert landed(arrivals) == landed_once_undropped(sents), arrivals
+ assert statuses == ("success",) * len(sents), sents
+
+
+@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS))
+def test_upstream_restarted_on_the_same_port_with_connections_pooled_serves_a_burst_on_fresh_connections(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ with ExitStack() as stack:
+ steady: Final = stack.enter_context(idle_gateway(answer, drop_after=None))
+ before_restart: Final = stack.enter_context(ExitStack())
+ first: Final = MappingProxyType(
+ {
+ row.name: before_restart.enter_context(idle_gateway(answer, drop_after=DROP_AFTER_SECONDS))
+ for row in RESTART_ROWS
+ }
+ )
+ proxy: Final = stack.enter_context(proxy_for(gateway, tmp_path, RESTART_ROWS, first, steady.url, {}))
+ target: Final = proxy.target
+ warm: Final = on_each_worker(
+ proxy.owned, proxy.workers, lambda: tuple(send(row, target) for row in RESTART_ROWS)
+ )
+ before_restart.close()
+ second: Final = MappingProxyType(
+ {
+ name: stack.enter_context(idle_gateway(answer, drop_after=DROP_AFTER_SECONDS, port=upstream.port))
+ for name, upstream in first.items()
+ }
+ )
+ stay_idle(tuple(second.values()), REUSE_IDLE_SECONDS)
+ sent: Final = replace(proxy, gateways=second).burst(tuple(row.name for row in RESTART_ROWS), BURST_EACH)
+ arrivals: Final = drain_all(tuple(second.values()))
+ sents: Final = tuple(one for _, one in sent)
+ statuses: Final = spend_statuses(sents)
+ assert tuple(one.answer.status for one in warm) == (200,) * 2 * len(RESTART_ROWS), warm
+ assert burst_answers(sent) == burst_expected(sent), sent
+ assert landed(arrivals) == landed_once_undropped(sents), arrivals
+ assert statuses == ("success",) * len(sents), sents
diff --git a/tests/integration/providers/test_idle_upstream_reconnect_routes_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_routes_wire.py
new file mode 100644
index 00000000000..2a8b640f654
--- /dev/null
+++ b/tests/integration/providers/test_idle_upstream_reconnect_routes_wire.py
@@ -0,0 +1,60 @@
+from collections.abc import Iterator
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.idle_gateway import Drop
+from integration._support.process import graceful_stop_seconds
+from integration.providers._idle_upstream_support import (
+ IDLE_SECONDS,
+ Row,
+ StalePools,
+ answered,
+ anthropic_deployment,
+ chat_with_guardrail,
+ expected_answers,
+ guardrail_deployment,
+ landed,
+ landed_once_undropped,
+ messages_sdk,
+ messages_sdk_stream,
+ ocr,
+ ocr_deployment,
+ pass_through,
+ pass_through_deployment,
+ served_by_rust,
+ spend_statuses,
+ stale_pools,
+ transcription,
+ transcription_deployment,
+)
+
+pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS))
+
+ROWS: Final = (
+ Row("messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk),
+ Row("messages-anthropic-sdk-async-stream", anthropic_deployment, messages_sdk_stream),
+ Row("messages-upstream-resets-stale-connection", anthropic_deployment, messages_sdk, drop=Drop.RESET),
+ Row("custom-pass-through-raw-httpx", pass_through_deployment, pass_through),
+ Row("generic-guardrail-api-pre-call", guardrail_deployment, chat_with_guardrail),
+ Row("ocr-mistral-rust-bridge", ocr_deployment, ocr, rust=True),
+ Row("audio-transcription-bedrock-rust-bridge", transcription_deployment, transcription, rust=True),
+)
+
+
+@pytest.fixture(scope="module")
+def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]:
+ with stale_pools(session_gateway, tmp_path_factory.mktemp("idle-upstream-routes"), ROWS, {}) as value:
+ yield value
+
+
+@pytest.mark.parametrize("name", [row.name for row in ROWS])
+def test_route_after_upstream_idle_past_keepalive_opens_a_fresh_connection_on_each_worker(
+ pools: StalePools, name: str
+) -> None:
+ row: Final = pools.rows[name]
+ sents, arrivals = pools.probe(name)
+ assert answered(sents) == expected_answers(row, sents), sents
+ assert landed(arrivals) == landed_once_undropped(sents), arrivals
+ assert served_by_rust(sents) == (row.rust,) * len(sents), sents
+ assert spend_statuses(sents) == ("success",) * len(sents), sents
diff --git a/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py
new file mode 100644
index 00000000000..8d92fb609e9
--- /dev/null
+++ b/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py
@@ -0,0 +1,56 @@
+from collections.abc import Iterator
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.idle_gateway import Drop
+from integration._support.process import graceful_stop_seconds
+from integration.providers._idle_upstream_support import (
+ IDLE_SECONDS,
+ Row,
+ StalePools,
+ answered,
+ anthropic_deployment,
+ expected_answers,
+ landed,
+ landed_once_undropped,
+ messages_sdk,
+ messages_sdk_stream,
+ served_by_rust,
+ spend_statuses,
+ stale_pools,
+)
+
+pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS))
+
+ROWS: Final = (
+ Row("messages-rust-anthropic-sdk-sync", anthropic_deployment, messages_sdk, rust=True),
+ Row("messages-rust-anthropic-sdk-async-stream", anthropic_deployment, messages_sdk_stream, rust=True),
+ Row(
+ "messages-rust-upstream-resets-stale-connection",
+ anthropic_deployment,
+ messages_sdk,
+ drop=Drop.RESET,
+ rust=True,
+ ),
+)
+
+
+@pytest.fixture(scope="module")
+def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]:
+ with stale_pools(
+ session_gateway, tmp_path_factory.mktemp("idle-upstream-rust"), ROWS, {"LITELLM_RUST": "1"}
+ ) as value:
+ yield value
+
+
+@pytest.mark.parametrize("name", [row.name for row in ROWS])
+def test_rust_messages_after_upstream_idle_past_keepalive_open_a_fresh_connection_on_each_worker(
+ pools: StalePools, name: str
+) -> None:
+ row: Final = pools.rows[name]
+ sents, arrivals = pools.probe(name)
+ assert answered(sents) == expected_answers(row, sents), sents
+ assert landed(arrivals) == landed_once_undropped(sents), arrivals
+ assert served_by_rust(sents) == (row.rust,) * len(sents), sents
+ assert spend_statuses(sents) == ("success",) * len(sents), sents
diff --git a/tests/integration/providers/test_idle_upstream_reconnect_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_wire.py
new file mode 100644
index 00000000000..820eff4546f
--- /dev/null
+++ b/tests/integration/providers/test_idle_upstream_reconnect_wire.py
@@ -0,0 +1,62 @@
+from collections.abc import Iterator
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.idle_gateway import Drop
+from integration._support.process import graceful_stop_seconds
+from integration.providers._idle_upstream_support import (
+ IDLE_SECONDS,
+ Row,
+ StalePools,
+ answered,
+ anthropic_deployment,
+ chat_httpx,
+ chat_sdk,
+ chat_sdk_stream,
+ chat_sdk_with_default_retries,
+ embedding_deployment,
+ embeddings_sdk,
+ expected_answers,
+ expected_embedding,
+ landed,
+ landed_once_undropped,
+ openai_deployment,
+ responses_sdk,
+ responses_sdk_stream,
+ served_by_rust,
+ spend_statuses,
+ stale_pools,
+)
+
+pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS))
+
+ROWS: Final = (
+ Row("chat-openai-sdk-sync", openai_deployment, chat_sdk),
+ Row("chat-openai-sdk-async-stream", openai_deployment, chat_sdk_stream),
+ Row("chat-raw-httpx", openai_deployment, chat_httpx),
+ Row("chat-to-anthropic-deployment", anthropic_deployment, chat_sdk),
+ Row("embeddings-openai-sdk", embedding_deployment, embeddings_sdk, expected=expected_embedding),
+ Row("responses-openai-sdk-sync", openai_deployment, responses_sdk),
+ Row("responses-openai-sdk-async-stream", openai_deployment, responses_sdk_stream),
+ Row("chat-upstream-resets-stale-connection", openai_deployment, chat_sdk, drop=Drop.RESET),
+ Row("chat-with-default-router-retries", openai_deployment, chat_sdk_with_default_retries),
+)
+
+
+@pytest.fixture(scope="module")
+def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]:
+ with stale_pools(session_gateway, tmp_path_factory.mktemp("idle-upstream"), ROWS, {}) as value:
+ yield value
+
+
+@pytest.mark.parametrize("name", [row.name for row in ROWS])
+def test_request_after_upstream_idle_past_keepalive_opens_a_fresh_connection_on_each_worker(
+ pools: StalePools, name: str
+) -> None:
+ row: Final = pools.rows[name]
+ sents, arrivals = pools.probe(name)
+ assert answered(sents) == expected_answers(row, sents), sents
+ assert landed(arrivals) == landed_once_undropped(sents), arrivals
+ assert served_by_rust(sents) == (row.rust,) * len(sents), sents
+ assert spend_statuses(sents) == ("success",) * len(sents), sents
diff --git a/tests/integration/sdk/test_idle_upstream_reconnect_sdk.py b/tests/integration/sdk/test_idle_upstream_reconnect_sdk.py
new file mode 100644
index 00000000000..617e756426c
--- /dev/null
+++ b/tests/integration/sdk/test_idle_upstream_reconnect_sdk.py
@@ -0,0 +1,94 @@
+import asyncio
+from collections.abc import Awaitable, Callable
+from types import MappingProxyType
+from typing import Final
+
+import openai
+import pytest
+from integration._support.idle_gateway import idle_gateway
+from integration.providers._idle_upstream_support import (
+ DROP_AFTER_SECONDS,
+ IDLE_SECONDS,
+ PROVIDER_KEY,
+ answer,
+ expected_text,
+ landed,
+ new_marker,
+ stay_idle,
+)
+
+import litellm
+from litellm.types.utils import Choices, ModelResponse
+
+
+def _text(response: object) -> str:
+ assert isinstance(response, ModelResponse), response
+ choice: Final = response.choices[0]
+ assert isinstance(choice, Choices), choice
+ return choice.message.content or ""
+
+
+async def _chat(model: str, api_base: str, marker: str) -> str:
+ return _text(
+ await litellm.acompletion(
+ model=model,
+ api_base=api_base,
+ api_key=PROVIDER_KEY,
+ messages=[{"role": "user", "content": marker}],
+ num_retries=0,
+ max_retries=0,
+ )
+ )
+
+
+async def openai_chat(api_base: str, marker: str) -> str:
+ return await _chat("openai/idle-upstream", api_base, marker)
+
+
+async def anthropic_chat(api_base: str, marker: str) -> str:
+ return await _chat("anthropic/claude-idle-upstream", api_base, marker)
+
+
+async def mistral_ocr_rust_bridge(api_base: str, marker: str) -> str:
+ response: Final = await litellm.aocr(
+ model="mistral/mistral-ocr-latest",
+ document={"type": "document_url", "document_url": f"https://documents.invalid/{marker}.pdf"},
+ api_base=api_base,
+ api_key=PROVIDER_KEY,
+ )
+ return response.pages[0].markdown
+
+
+CALLS: Final[MappingProxyType[str, Callable[[str, str], Awaitable[str]]]] = MappingProxyType(
+ {
+ "acompletion-openai": openai_chat,
+ "acompletion-anthropic": anthropic_chat,
+ "aocr-mistral-rust-bridge": mistral_ocr_rust_bridge,
+ }
+)
+
+
+async def _outcome(call: Callable[[str, str], Awaitable[str]], api_base: str, marker: str) -> str:
+ try:
+ return await call(api_base, marker)
+ except openai.APIError as error:
+ return f"raised {type(error).__name__}: {error}"
+
+
+@pytest.mark.timeout(int(IDLE_SECONDS + 120))
+@pytest.mark.parametrize("name", list(CALLS))
+async def test_sdk_call_after_upstream_idle_past_keepalive_opens_a_fresh_connection(name: str) -> None:
+ call: Final = CALLS[name]
+ warm_marker: Final = new_marker()
+ probe_marker: Final = new_marker()
+ with idle_gateway(answer, drop_after=DROP_AFTER_SECONDS) as upstream:
+ warm: Final = await _outcome(call, upstream.url, warm_marker)
+ warmed: Final = upstream.drain()
+ await asyncio.to_thread(stay_idle, (upstream,), IDLE_SECONDS)
+ probe: Final = await _outcome(call, upstream.url, probe_marker)
+ arrivals: Final = upstream.drain()
+ assert (warm, probe) == (expected_text(warm_marker), expected_text(probe_marker)), (warm, probe)
+ assert landed((*warmed, *arrivals)) == tuple(sorted(((warm_marker, False), (probe_marker, False)))), (
+ warmed,
+ arrivals,
+ )
From d0f3a2dfdb65385e8359db8953d6ccce53e79605 Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 10:51:18 -0700
Subject: [PATCH 17/47] fix(vertex_ai): pass the per-request timeout to Gemini
streaming calls (#46349)
* fix(vertex_ai): pass the per-request timeout to Gemini streaming calls
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(vertex_ai): pass timeout to make_call and make_sync_call in the http 201 tests
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(vertex_ai): default timeout to None on make_call and make_sync_call
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../vertex_and_google_ai_studio_gemini.py | 12 +-
.../test_vertex_gemini_stream_timeout_wire.py | 147 ++++++++++++++++++
2 files changed, 157 insertions(+), 2 deletions(-)
create mode 100644 tests/integration/sdk/test_vertex_gemini_stream_timeout_wire.py
diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
index 517ac9e98b1..109dddf44ed 100644
--- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
@@ -2748,6 +2748,7 @@ async def make_call(
model: str,
messages: list,
logging_obj,
+ timeout: float | httpx.Timeout | None = None,
):
if gemini_client is not None:
client = gemini_client
@@ -2757,7 +2758,9 @@ async def make_call(
)
try:
- response: Final = await client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj)
+ response: Final = await client.post(
+ api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj, timeout=timeout
+ )
response.raise_for_status()
except httpx.HTTPStatusError as e:
exception_string: Final = str(await e.response.aread())
@@ -2800,13 +2803,16 @@ def make_sync_call(
model: str,
messages: list,
logging_obj,
+ timeout: float | httpx.Timeout | None = None,
):
if gemini_client is not None:
client = gemini_client
if client is None:
client = HTTPHandler() # Create a new client if none provided
- response: Final = client.post(api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj)
+ response: Final = client.post(
+ api_base, headers=headers, data=data, stream=True, logging_obj=logging_obj, timeout=timeout
+ )
if response.status_code != 200 and response.status_code != 201:
raise VertexAIError(
@@ -2924,6 +2930,7 @@ class VertexLLM(VertexBase):
model=model,
messages=messages,
logging_obj=logging_obj,
+ timeout=timeout,
),
model=model,
custom_llm_provider="vertex_ai_beta",
@@ -3200,6 +3207,7 @@ class VertexLLM(VertexBase):
messages=messages,
logging_obj=logging_obj,
headers=headers,
+ timeout=timeout,
),
model=model,
custom_llm_provider="vertex_ai_beta",
diff --git a/tests/integration/sdk/test_vertex_gemini_stream_timeout_wire.py b/tests/integration/sdk/test_vertex_gemini_stream_timeout_wire.py
new file mode 100644
index 00000000000..68d98961dce
--- /dev/null
+++ b/tests/integration/sdk/test_vertex_gemini_stream_timeout_wire.py
@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import threading
+import time
+from collections.abc import Callable
+from concurrent.futures import Future, ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeout
+from typing import Final
+
+import pytest
+from integration._support.vertex import service_account_json
+from integration._support.wire import Reply, Request, Wire, wire_server
+
+import litellm
+from litellm.exceptions import MidStreamFallbackError
+
+_MODEL: Final = "gemini-3.8-flash"
+_PROJECT: Final = "scripted-project"
+_LOCATION: Final = "us-east5"
+_MESSAGES: Final = [{"role": "user", "content": "does the per-request timeout reach a stalled stream?"}]
+_REQUEST_TIMEOUT_SECONDS: Final = 1.0
+_HANG_DEADLINE_SECONDS: Final = 15.0
+_TOKEN_REPLY: Final = Reply(
+ body=json.dumps({"access_token": "scripted-token", "token_type": "Bearer", "expires_in": 3600}).encode()
+)
+
+
+def _sse(text: str, finish_reason: str | None) -> bytes:
+ candidate: Final = {
+ "content": {"role": "model", "parts": [{"text": text}]},
+ **({"finishReason": finish_reason} if finish_reason else {}),
+ }
+ usage: Final = {"promptTokenCount": 9, "candidatesTokenCount": 2, "totalTokenCount": 11}
+ return b"data: " + json.dumps({"candidates": [candidate], "usageMetadata": usage}).encode() + b"\n\n"
+
+
+def _stream_that_stalls_after_the_first_chunk() -> Reply:
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(_sse("partial", None), _sse("", "STOP")),
+ gate_after_first=threading.Event(),
+ )
+
+
+def _stream_that_completes() -> Reply:
+ return Reply(content_type="text/event-stream", chunks=(_sse("scripted ", None), _sse("answer", "STOP")))
+
+
+def _body_that_stalls_half_way() -> Reply:
+ return Reply(chunks=(b'{"candidates": [', b"]}"), gate_after_first=threading.Event())
+
+
+def _peer(reply: Reply) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ if request.target == "/_oauth/token":
+ return _TOKEN_REPLY
+ assert request.method == "POST", request.method
+ return reply
+
+ return respond
+
+
+def _deployment(route: str, wire: Wire) -> dict[str, str]:
+ if route == "gemini":
+ return {"model": f"gemini/{_MODEL}", "api_base": wire.url, "api_key": "scripted-key"}
+ return {
+ "model": f"vertex_ai/{_MODEL}",
+ "api_base": wire.url,
+ "vertex_project": _PROJECT,
+ "vertex_location": _LOCATION,
+ "vertex_credentials": service_account_json(_PROJECT, wire.url),
+ }
+
+
+def _text_of(chunk: object) -> str:
+ choices: Final = getattr(chunk, "choices", None) or ()
+ return "".join(str(choice.delta.content or "") for choice in choices)
+
+
+def _complete_sync(route: str, wire: Wire, stream: bool) -> str:
+ response: Final = litellm.completion(
+ **_deployment(route, wire), messages=_MESSAGES, stream=stream, timeout=_REQUEST_TIMEOUT_SECONDS
+ )
+ if stream:
+ return "".join(_text_of(chunk) for chunk in response)
+ return str(response.choices[0].message.content)
+
+
+async def _complete_async(route: str, wire: Wire, stream: bool) -> str:
+ response: Final = await litellm.acompletion(
+ **_deployment(route, wire), messages=_MESSAGES, stream=stream, timeout=_REQUEST_TIMEOUT_SECONDS
+ )
+ if stream:
+ parts: Final = [_text_of(chunk) async for chunk in response]
+ return "".join(parts)
+ return str(response.choices[0].message.content)
+
+
+def _complete(client: str, route: str, wire: Wire, stream: bool) -> str:
+ if client == "async":
+ return asyncio.run(_complete_async(route, wire, stream))
+ return _complete_sync(route, wire, stream)
+
+
+def _before_the_hang_deadline(call: Callable[[], str]) -> str:
+ pool: Final = ThreadPoolExecutor(max_workers=1)
+ future: Final[Future[str]] = pool.submit(call)
+ pool.shutdown(wait=False)
+ try:
+ return future.result(timeout=_HANG_DEADLINE_SECONDS)
+ except FutureTimeout:
+ pytest.fail(
+ f"still waiting on the peer {_HANG_DEADLINE_SECONDS}s after a {_REQUEST_TIMEOUT_SECONDS}s request timeout"
+ )
+
+
+_CLIENTS: Final = ("sync", "async")
+_ROUTES: Final = ("vertex_ai", "gemini")
+
+
+@pytest.mark.parametrize("route", _ROUTES)
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_stalled_stream_raises_at_the_request_timeout(client: str, route: str) -> None:
+ with wire_server(_peer(_stream_that_stalls_after_the_first_chunk())) as wire:
+ started: Final = time.monotonic()
+ with pytest.raises(MidStreamFallbackError):
+ _before_the_hang_deadline(lambda: _complete(client, route, wire, stream=True))
+ elapsed: Final = time.monotonic() - started
+ assert elapsed >= _REQUEST_TIMEOUT_SECONDS, elapsed
+ assert len(wire.drain()) >= 1
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_stalled_body_raises_at_the_request_timeout(client: str) -> None:
+ with wire_server(_peer(_body_that_stalls_half_way())) as wire:
+ started: Final = time.monotonic()
+ with pytest.raises(litellm.Timeout):
+ _before_the_hang_deadline(lambda: _complete(client, "vertex_ai", wire, stream=False))
+ assert time.monotonic() - started >= _REQUEST_TIMEOUT_SECONDS
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_stream_with_a_request_timeout_still_completes(client: str) -> None:
+ with wire_server(_peer(_stream_that_completes())) as wire:
+ assert _complete(client, "vertex_ai", wire, stream=True) == "scripted answer"
From a0ac9b4f1fbf16f90f43942cca92ed2f3528ffa3 Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Wed, 14 Oct 2026 10:53:38 -0700
Subject: [PATCH 18/47] feat(proxy): add deployment warning endpoints with
instance-wide ignore and snooze (#46347)
* feat(proxy): add deployment warning endpoints with instance-wide ignore and snooze
GET /health/deployment_warnings returns the active deployment warnings
(env credential login, license expiry, no Redis on multiple workers,
detailed debug logging) with a severity and ignore state, for proxy admins
and proxy admin viewers. PATCH /health/deployment_warnings/{warning_id}
lets a proxy admin ignore a warning forever or snooze it for up to 30 days,
or restore it. Ignores live in one LiteLLM_UISettings row, apply to every
admin, and are written to the config audit log. Expired snoozes drop out
on read, and LITELLM_DISABLE_NO_REDIS_WARNING is reported as an env var
ignore. /health/readiness/details is unchanged
* fix(proxy): keep readiness Redis short-circuit ahead of the env lookup and harden ignore storage
Readiness details again answers from Redis before reading
LITELLM_DISABLE_NO_REDIS_WARNING, so proxies with Redis and a secret
manager do not fetch it on every poll. Stored ignores require timezone
aware datetimes and are read from the writer before a rewrite
* fix(proxy): keep deployment warning ignores readable in the config audit log
The audit payload keyed ignores by warning id, so the secret matcher
redacted the env_credential_login entry and the audit row lost who
snoozed it and until when. Log a list of id and ignore entries instead
---
.../health_endpoints/_health_endpoints.py | 38 +-
.../deployment_warning_endpoints.py | 350 +++++++++++++
litellm/proxy/proxy_server.py | 2 +
.../deployment_warning_ignore_repository.py | 23 +
.../test_deployment_warning_endpoints.py | 475 ++++++++++++++++++
.../health_endpoints/test_health_endpoints.py | 15 +
tests/unit/repositories/test_repositories.py | 23 +
ui/litellm-dashboard/src/lib/http/schema.d.ts | 168 +++++++
8 files changed, 1081 insertions(+), 13 deletions(-)
create mode 100644 litellm/proxy/health_endpoints/deployment_warning_endpoints.py
create mode 100644 litellm/repositories/deployment_warning_ignore_repository.py
create mode 100644 tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py
diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py
index fa58fda8f47..e16a391dc12 100644
--- a/litellm/proxy/health_endpoints/_health_endpoints.py
+++ b/litellm/proxy/health_endpoints/_health_endpoints.py
@@ -1464,7 +1464,7 @@ async def shared_health_check_status_endpoint(
)
-def _read_license_data() -> EnterpriseLicenseData | None:
+def read_license_data() -> EnterpriseLicenseData | None:
from litellm.proxy.proxy_server import _license_check, premium_user_data
license_data: EnterpriseLicenseData | None = premium_user_data or _license_check.airgapped_license_data
@@ -1509,7 +1509,7 @@ async def health_license_endpoint(
"""Return metadata about the configured LiteLLM license without exposing the key."""
from litellm.proxy.proxy_server import _license_check, premium_user
- license_data: Final = _read_license_data()
+ license_data: Final = read_license_data()
has_license: Final = bool(getattr(_license_check, "license_str", None))
license_type: Final = "enterprise" if premium_user else "community"
@@ -1703,9 +1703,9 @@ def callback_name(callback):
DISABLE_NO_REDIS_WARNING_ENV_VAR: Final = "LITELLM_DISABLE_NO_REDIS_WARNING"
-async def _show_no_redis_warning() -> bool:
+async def no_redis_on_multiple_workers() -> bool:
"""
- Whether the UI should warn that no Redis is configured.
+ Whether this proxy runs without Redis on more than one worker.
Redis is what makes rate limits, budgets, router state, and cache
invalidation consistent across workers, so a proxy running without it is
@@ -1716,22 +1716,34 @@ async def _show_no_redis_warning() -> bool:
routing on its own. A deployment whose worker-heartbeat census proves it
is exactly one worker needs no cross-worker coordination, so it never
warns; when the census is unavailable or shows more than one worker, the
- warning stands unless LITELLM_DISABLE_NO_REDIS_WARNING=true silences it.
+ warning stands.
"""
- from litellm.proxy.proxy_server import llm_router, prisma_client, redis_usage_cache
+ from litellm.proxy.proxy_server import prisma_client
- if redis_usage_cache is not None:
- return False
- if llm_router is not None and llm_router.cache.redis_cache is not None:
- return False
- if get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is True:
+ if _redis_configured():
return False
if prisma_client is None:
return True
return await count_live_proxy_workers(prisma_client) != 1
-def _show_env_credential_login_warning() -> bool:
+def _redis_configured() -> bool:
+ from litellm.proxy.proxy_server import llm_router, redis_usage_cache
+
+ return redis_usage_cache is not None or (llm_router is not None and llm_router.cache.redis_cache is not None)
+
+
+def no_redis_warning_disabled_by_env() -> bool:
+ return get_secret_bool(DISABLE_NO_REDIS_WARNING_ENV_VAR, False) is True
+
+
+async def _show_no_redis_warning() -> bool:
+ if _redis_configured() or no_redis_warning_disabled_by_env():
+ return False
+ return await no_redis_on_multiple_workers()
+
+
+def env_credential_login_enabled() -> bool:
from litellm.proxy.auth.login_utils import is_env_credential_login_enabled
from litellm.proxy.proxy_server import general_settings
@@ -1779,7 +1791,7 @@ async def _get_health_readiness_details(
log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel())
is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG)
show_no_redis_warning: Final = await _show_no_redis_warning()
- show_env_credential_login_warning: Final = _show_env_credential_login_warning()
+ show_env_credential_login_warning: Final = env_credential_login_enabled()
# check DB
if prisma_client is not None: # if db passed in, check if it's connected
diff --git a/litellm/proxy/health_endpoints/deployment_warning_endpoints.py b/litellm/proxy/health_endpoints/deployment_warning_endpoints.py
new file mode 100644
index 00000000000..115068f79f8
--- /dev/null
+++ b/litellm/proxy/health_endpoints/deployment_warning_endpoints.py
@@ -0,0 +1,350 @@
+import json
+import logging
+from collections.abc import Mapping
+from dataclasses import dataclass
+from datetime import date, datetime, timedelta, timezone
+from types import MappingProxyType
+from typing import Annotated, Final, Literal, Protocol, TypeAlias
+
+from fastapi import APIRouter, Depends, HTTPException, status
+from pydantic import AwareDatetime, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, model_validator
+
+from litellm._logging import verbose_logger, verbose_proxy_logger
+from litellm.constants import LITELLM_PROXY_ADMIN_NAME
+from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.health_endpoints._health_endpoints import (
+ DISABLE_NO_REDIS_WARNING_ENV_VAR,
+ env_credential_login_enabled,
+ no_redis_on_multiple_workers,
+ no_redis_warning_disabled_by_env,
+ read_license_data,
+)
+from litellm.repositories.deployment_warning_ignore_repository import (
+ DEPLOYMENT_WARNING_IGNORES_ROW_ID,
+ DeploymentWarningIgnoreRepository,
+)
+from litellm.types.llms.base import LiteLLMBaseModel
+
+router: Final = APIRouter()
+
+ENV_CREDENTIAL_LOGIN_WARNING_ID: Final = "env_credential_login"
+NO_REDIS_WARNING_ID: Final = "no_redis"
+DETAILED_DEBUG_WARNING_ID: Final = "detailed_debug"
+LICENSE_EXPIRY_WARNING_ID_PREFIX: Final = "license_expiry:"
+LICENSE_EXPIRY_WARNING_DAYS: Final = 30
+LICENSE_EXPIRY_CRITICAL_DAYS: Final = 7
+MAX_SNOOZE_SECONDS: Final = 30 * 24 * 60 * 60
+
+_STATIC_WARNING_IDS: Final = frozenset(
+ {ENV_CREDENTIAL_LOGIN_WARNING_ID, NO_REDIS_WARNING_ID, DETAILED_DEBUG_WARNING_ID}
+)
+
+WarningSeverity: TypeAlias = Literal["critical", "warning"]
+
+
+@dataclass(frozen=True, slots=True)
+class WarningCondition:
+ id: str
+ severity: WarningSeverity
+ env_var_ignore: str | None = None
+
+
+class AdminIgnore(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ source: Literal["admin"] = "admin"
+ ignored_by: str
+ ignored_at: AwareDatetime
+ ignored_until: AwareDatetime | None = Field(
+ default=None,
+ description="When the snooze ends. Null means the warning is ignored until an admin restores it.",
+ )
+
+
+class EnvVarIgnore(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ source: Literal["env_var"] = "env_var"
+ env_var: str
+
+
+class DeploymentWarning(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ id: str
+ severity: WarningSeverity
+ ignore: Annotated[AdminIgnore | EnvVarIgnore, Field(discriminator="source")] | None = None
+
+
+class DeploymentWarningsResponse(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ warnings: tuple[DeploymentWarning, ...]
+ ignore_available: bool = Field(description="False when no database is connected, so warnings cannot be ignored.")
+
+
+class DeploymentWarningIgnoreUpdate(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ ignored: bool = Field(description="True ignores the warning for every admin, false restores it.")
+ snooze_seconds: int | None = Field(
+ default=None,
+ gt=0,
+ le=MAX_SNOOZE_SECONDS,
+ description="How long to ignore the warning, up to 30 days. Leave it out to ignore it forever.",
+ )
+
+ @model_validator(mode="after")
+ def _snooze_only_when_ignoring(self) -> "DeploymentWarningIgnoreUpdate":
+ if self.snooze_seconds is not None and not self.ignored:
+ raise ValueError("snooze_seconds is only valid when ignored is true")
+ return self
+
+
+class DeploymentWarningIgnoreResponse(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ id: str
+ ignore: AdminIgnore | None
+
+
+class DeploymentWarningIgnoreStore(Protocol):
+ async def get_raw_ignores(self) -> object: ...
+
+ async def upsert_ignores(self, payload: str) -> None: ...
+
+
+class ConfigAuditLogger(Protocol):
+ async def __call__(
+ self,
+ param_name: str,
+ action: AUDIT_ACTIONS,
+ before_value: JsonValue | None,
+ after_value: JsonValue | None,
+ user_api_key_dict: UserAPIKeyAuth,
+ table_name: LitellmTableNames = ...,
+ ) -> None: ...
+
+
+_IGNORES_ADAPTER: Final[TypeAdapter[Mapping[str, AdminIgnore]]] = TypeAdapter(Mapping[str, AdminIgnore])
+_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
+_AUDIT_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[DeploymentWarningIgnoreResponse, ...]]] = TypeAdapter(
+ tuple[DeploymentWarningIgnoreResponse, ...]
+)
+
+
+def license_expiry_condition(expiration_date: str | None, now: datetime) -> WarningCondition | None:
+ if not expiration_date:
+ return None
+ try:
+ expires_on: Final = date.fromisoformat(expiration_date)
+ except ValueError:
+ return None
+ days_left: Final = (expires_on - now.date()).days
+ if days_left > LICENSE_EXPIRY_WARNING_DAYS:
+ return None
+ return WarningCondition(
+ id=f"{LICENSE_EXPIRY_WARNING_ID_PREFIX}{expires_on.isoformat()}",
+ severity="critical" if days_left <= LICENSE_EXPIRY_CRITICAL_DAYS else "warning",
+ )
+
+
+def deployment_warning_conditions(
+ *,
+ env_credential_login_enabled: bool,
+ license_expiration_date: str | None,
+ no_redis_on_multiple_workers: bool,
+ no_redis_warning_disabled_by_env: bool,
+ detailed_debug: bool,
+ now: datetime,
+) -> tuple[WarningCondition, ...]:
+ no_redis_env_var: Final = DISABLE_NO_REDIS_WARNING_ENV_VAR if no_redis_warning_disabled_by_env else None
+ candidates: Final = (
+ WarningCondition(id=ENV_CREDENTIAL_LOGIN_WARNING_ID, severity="critical")
+ if env_credential_login_enabled
+ else None,
+ license_expiry_condition(license_expiration_date, now),
+ WarningCondition(id=NO_REDIS_WARNING_ID, severity="warning", env_var_ignore=no_redis_env_var)
+ if no_redis_on_multiple_workers
+ else None,
+ WarningCondition(id=DETAILED_DEBUG_WARNING_ID, severity="warning") if detailed_debug else None,
+ )
+ return tuple(condition for condition in candidates if condition is not None)
+
+
+def is_known_warning_id(warning_id: str) -> bool:
+ if warning_id in _STATIC_WARNING_IDS:
+ return True
+ if not warning_id.startswith(LICENSE_EXPIRY_WARNING_ID_PREFIX):
+ return False
+ expiration_date: Final = warning_id.removeprefix(LICENSE_EXPIRY_WARNING_ID_PREFIX)
+ try:
+ return date.fromisoformat(expiration_date).isoformat() == expiration_date
+ except ValueError:
+ return False
+
+
+def _parse_stored_ignores(raw: object) -> Mapping[str, AdminIgnore]:
+ if raw is None:
+ return MappingProxyType({})
+ try:
+ return MappingProxyType(_IGNORES_ADAPTER.validate_python(json.loads(raw) if isinstance(raw, str) else raw))
+ except (json.JSONDecodeError, ValidationError):
+ verbose_proxy_logger.warning("Ignoring unreadable %s row", DEPLOYMENT_WARNING_IGNORES_ROW_ID)
+ return MappingProxyType({})
+
+
+async def _read_unexpired_ignores(store: DeploymentWarningIgnoreStore, now: datetime) -> Mapping[str, AdminIgnore]:
+ stored: Final = _parse_stored_ignores(await store.get_raw_ignores())
+ return MappingProxyType(
+ {
+ warning_id: ignore
+ for warning_id, ignore in stored.items()
+ if ignore.ignored_until is None or ignore.ignored_until > now
+ }
+ )
+
+
+def _resolve_ignore(
+ condition: WarningCondition, ignores: Mapping[str, AdminIgnore]
+) -> AdminIgnore | EnvVarIgnore | None:
+ if condition.env_var_ignore is not None:
+ return EnvVarIgnore(env_var=condition.env_var_ignore)
+ return ignores.get(condition.id)
+
+
+def _ignores_as_audit_entries(ignores: Mapping[str, AdminIgnore]) -> JsonValue:
+ entries: Final = tuple(
+ DeploymentWarningIgnoreResponse(id=warning_id, ignore=ignores[warning_id]) for warning_id in sorted(ignores)
+ )
+ return _JSON_VALUE_ADAPTER.validate_json(_AUDIT_ENTRIES_ADAPTER.dump_json(entries))
+
+
+def utc_now() -> datetime:
+ return datetime.now(timezone.utc)
+
+
+def deployment_warning_ignore_store() -> DeploymentWarningIgnoreStore | None:
+ from litellm.proxy.proxy_server import prisma_client
+
+ if prisma_client is None:
+ return None
+ return DeploymentWarningIgnoreRepository(prisma_client, use_writer=True)
+
+
+def config_audit_logger() -> ConfigAuditLogger:
+ from litellm.proxy.proxy_server import create_config_audit_log
+
+ return create_config_audit_log
+
+
+async def _deployment_warning_reader(
+ user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+) -> UserAPIKeyAuth:
+ if user_api_key_dict.user_role not in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY):
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail="Only proxy admins and proxy admin viewers can read deployment warnings.",
+ )
+ return user_api_key_dict
+
+
+async def _deployment_warning_admin(
+ user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
+) -> UserAPIKeyAuth:
+ if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
+ raise HTTPException(
+ status_code=status.HTTP_403_FORBIDDEN,
+ detail="Only proxy admins can ignore deployment warnings.",
+ )
+ return user_api_key_dict
+
+
+async def current_warning_conditions(now: Annotated[datetime, Depends(utc_now)]) -> tuple[WarningCondition, ...]:
+ license_data: Final = read_license_data()
+ return deployment_warning_conditions(
+ env_credential_login_enabled=env_credential_login_enabled(),
+ license_expiration_date=None if license_data is None else license_data.get("expiration_date"),
+ no_redis_on_multiple_workers=await no_redis_on_multiple_workers(),
+ no_redis_warning_disabled_by_env=no_redis_warning_disabled_by_env(),
+ detailed_debug=verbose_logger.isEnabledFor(logging.DEBUG),
+ now=now,
+ )
+
+
+@router.get(
+ "/health/deployment_warnings",
+ tags=["health"],
+ response_model=DeploymentWarningsResponse,
+)
+async def get_deployment_warnings(
+ _: Annotated[UserAPIKeyAuth, Depends(_deployment_warning_reader)],
+ conditions: Annotated[tuple[WarningCondition, ...], Depends(current_warning_conditions)],
+ store: Annotated[DeploymentWarningIgnoreStore | None, Depends(deployment_warning_ignore_store)],
+ now: Annotated[datetime, Depends(utc_now)],
+) -> DeploymentWarningsResponse:
+ """
+ List the active deployment warnings with their severity and whether an admin ignored or snoozed them.
+ Readable by proxy admins and proxy admin viewers.
+ """
+ ignores: Final[Mapping[str, AdminIgnore]] = (
+ MappingProxyType({}) if store is None else await _read_unexpired_ignores(store, now)
+ )
+ return DeploymentWarningsResponse(
+ warnings=tuple(
+ DeploymentWarning(id=condition.id, severity=condition.severity, ignore=_resolve_ignore(condition, ignores))
+ for condition in conditions
+ ),
+ ignore_available=store is not None,
+ )
+
+
+@router.patch(
+ "/health/deployment_warnings/{warning_id}",
+ tags=["health"],
+ response_model=DeploymentWarningIgnoreResponse,
+)
+async def update_deployment_warning_ignore(
+ warning_id: str,
+ update: DeploymentWarningIgnoreUpdate,
+ user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_deployment_warning_admin)],
+ store: Annotated[DeploymentWarningIgnoreStore | None, Depends(deployment_warning_ignore_store)],
+ audit_logger: Annotated[ConfigAuditLogger, Depends(config_audit_logger)],
+ now: Annotated[datetime, Depends(utc_now)],
+) -> DeploymentWarningIgnoreResponse:
+ """
+ Ignore, snooze, or restore one deployment warning for every admin of this proxy.
+ Only proxy admins can call this, and it needs a database.
+ """
+ if not is_known_warning_id(warning_id):
+ raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Unknown deployment warning: {warning_id}")
+ if store is None:
+ raise HTTPException(
+ status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+ detail="Ignoring deployment warnings needs a database. Connect one to enable it.",
+ )
+
+ before: Final = await _read_unexpired_ignores(store, now)
+ ignore: Final = (
+ AdminIgnore(
+ ignored_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
+ ignored_at=now,
+ ignored_until=None if update.snooze_seconds is None else now + timedelta(seconds=update.snooze_seconds),
+ )
+ if update.ignored
+ else None
+ )
+ others: Final = {other_id: other for other_id, other in before.items() if other_id != warning_id}
+ after: Final = MappingProxyType(others if ignore is None else {**others, warning_id: ignore})
+
+ await store.upsert_ignores(_IGNORES_ADAPTER.dump_json(dict(after)).decode())
+ await audit_logger(
+ param_name=DEPLOYMENT_WARNING_IGNORES_ROW_ID,
+ action="updated",
+ before_value=_ignores_as_audit_entries(before),
+ after_value=_ignores_as_audit_entries(after),
+ user_api_key_dict=user_api_key_dict,
+ table_name=LitellmTableNames.UI_SETTINGS_TABLE_NAME,
+ )
+ return DeploymentWarningIgnoreResponse(id=warning_id, ignore=ignore)
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 306cb6bd95c..b51df22d07a 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -590,6 +590,7 @@ from litellm.proxy.health_check import (
perform_health_check,
)
from litellm.proxy.health_endpoints._health_endpoints import router as health_router
+from litellm.proxy.health_endpoints.deployment_warning_endpoints import router as deployment_warning_router
from litellm.proxy.hooks.model_max_budget_limiter import ( # noqa: F401, RUF100 # legacy module exports
PROXY_VirtualKeyModelMaxBudgetLimiter,
_PROXY_VirtualKeyModelMaxBudgetLimiter, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@@ -20402,6 +20403,7 @@ app.include_router(openai_files_router)
reserve_lazy_slot(app, "llm_passthrough")
app.include_router(pass_through_router)
app.include_router(health_router)
+app.include_router(deployment_warning_router)
app.include_router(key_management_router)
app.include_router(internal_user_router)
app.include_router(daily_activity_router)
diff --git a/litellm/repositories/deployment_warning_ignore_repository.py b/litellm/repositories/deployment_warning_ignore_repository.py
new file mode 100644
index 00000000000..da02d424f80
--- /dev/null
+++ b/litellm/repositories/deployment_warning_ignore_repository.py
@@ -0,0 +1,23 @@
+from typing import TYPE_CHECKING, Final
+
+from litellm.repositories.table_repositories import PrismaTableRepository
+
+if TYPE_CHECKING:
+ from prisma import models as prisma_models # noqa: F401 # resolved only from the quoted base-class subscript below
+
+DEPLOYMENT_WARNING_IGNORES_ROW_ID: Final = "deployment_warning_ignores"
+
+
+class DeploymentWarningIgnoreRepository(PrismaTableRepository["prisma_models.LiteLLM_UISettings"]):
+ table_name = "litellm_uisettings"
+
+ async def get_raw_ignores(self) -> object:
+ db_record: Final = await self.table.find_unique(where={"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID})
+ return db_record.ui_settings if db_record is not None else None
+
+ async def upsert_ignores(self, payload: str) -> None:
+ row: Final = {"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID, "ui_settings": payload}
+ await self.table.upsert(
+ where={"id": DEPLOYMENT_WARNING_IGNORES_ROW_ID},
+ data={"create": row, "update": {"ui_settings": payload}},
+ )
diff --git a/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py b/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py
new file mode 100644
index 00000000000..4583d8fd978
--- /dev/null
+++ b/tests/unit/proxy/health_endpoints/test_deployment_warning_endpoints.py
@@ -0,0 +1,475 @@
+import json
+import logging
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
+from types import SimpleNamespace
+from typing import Final
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from pydantic import JsonValue
+
+from litellm._logging import verbose_logger
+from litellm.constants import LITELLM_PROXY_ADMIN_NAME
+from litellm.proxy._types import AUDIT_ACTIONS, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.health_endpoints.deployment_warning_endpoints import (
+ WarningCondition,
+ config_audit_logger,
+ current_warning_conditions,
+ deployment_warning_conditions,
+ deployment_warning_ignore_store,
+ license_expiry_condition,
+ router,
+ utc_now,
+)
+from litellm.proxy.proxy_server import (
+ _dump_redacted_config, # pyright: ignore[reportPrivateUsage] # the exact redaction every config audit row goes through
+)
+from litellm.repositories.deployment_warning_ignore_repository import DEPLOYMENT_WARNING_IGNORES_ROW_ID
+
+T0: Final = datetime(2026, 10, 14, 12, 0, tzinfo=timezone.utc)
+THIRTY_DAYS: Final = 30 * 24 * 60 * 60
+
+
+@dataclass
+class InMemoryIgnoreStore:
+ raw: object = None
+
+ async def get_raw_ignores(self) -> object:
+ return self.raw
+
+ async def upsert_ignores(self, payload: str) -> None:
+ self.raw = json.loads(payload)
+
+
+@dataclass
+class RecordingAuditLogger:
+ calls: list[dict[str, object]] = field(default_factory=list)
+
+ async def __call__(
+ self,
+ param_name: str,
+ action: AUDIT_ACTIONS,
+ before_value: JsonValue | None,
+ after_value: JsonValue | None,
+ user_api_key_dict: UserAPIKeyAuth,
+ table_name: LitellmTableNames = LitellmTableNames.CONFIG_TABLE_NAME,
+ ) -> None:
+ self.calls.append(
+ {
+ "param_name": param_name,
+ "action": action,
+ "before_value": before_value,
+ "after_value": after_value,
+ "table_name": table_name,
+ }
+ )
+
+
+@dataclass
+class Harness:
+ app: FastAPI
+ client: TestClient
+ store: InMemoryIgnoreStore
+ audit: RecordingAuditLogger
+ clock: list[datetime]
+
+ def as_role(self, role: LitellmUserRoles, user_id: str | None = "admin-user") -> None:
+ async def auth() -> UserAPIKeyAuth:
+ return UserAPIKeyAuth(api_key="sk-test", user_id=user_id, user_role=role)
+
+ self.app.dependency_overrides[user_api_key_auth] = auth
+
+ def with_conditions(self, *conditions: WarningCondition) -> None:
+ self.app.dependency_overrides[current_warning_conditions] = lambda: conditions
+
+ def without_database(self) -> None:
+ self.app.dependency_overrides[deployment_warning_ignore_store] = lambda: None
+
+ def warnings(self) -> dict[str, dict[str, object]]:
+ response = self.client.get("/health/deployment_warnings")
+ assert response.status_code == 200, response.text
+ return {warning["id"]: warning for warning in response.json()["warnings"]}
+
+
+@pytest.fixture
+def harness() -> Iterator[Harness]:
+ app = FastAPI()
+ app.include_router(router)
+ store = InMemoryIgnoreStore()
+ audit = RecordingAuditLogger()
+ clock = [T0]
+ app.dependency_overrides[deployment_warning_ignore_store] = lambda: store
+ app.dependency_overrides[config_audit_logger] = lambda: audit
+ app.dependency_overrides[utc_now] = lambda: clock[0]
+ h = Harness(app=app, client=TestClient(app), store=store, audit=audit, clock=clock)
+ h.as_role(LitellmUserRoles.PROXY_ADMIN)
+ h.with_conditions(
+ WarningCondition(id="env_credential_login", severity="critical"),
+ WarningCondition(id="detailed_debug", severity="warning"),
+ )
+ yield h
+
+
+class TestLicenseExpiryCondition:
+ @pytest.mark.parametrize(
+ ("days_left", "severity"),
+ [(31, None), (30, "warning"), (8, "warning"), (7, "critical"), (0, "critical"), (-1, "critical")],
+ )
+ def test_severity_follows_days_left(self, days_left, severity):
+ expires_on = (T0 + timedelta(days=days_left)).date().isoformat()
+ condition = license_expiry_condition(expires_on, T0)
+ if severity is None:
+ assert condition is None
+ return
+ assert condition == WarningCondition(id=f"license_expiry:{expires_on}", severity=severity)
+
+ def test_days_left_counts_utc_calendar_days(self):
+ just_before_midnight = datetime(2026, 10, 14, 23, 59, tzinfo=timezone.utc)
+ assert license_expiry_condition("2026-10-21", just_before_midnight).severity == "critical"
+ assert license_expiry_condition("2026-10-22", just_before_midnight).severity == "warning"
+
+ def test_renewed_license_gets_a_new_id(self):
+ first = license_expiry_condition("2026-10-20", T0)
+ renewed = license_expiry_condition("2026-11-10", T0)
+ assert first is not None and renewed is not None
+ assert first.id != renewed.id
+
+ @pytest.mark.parametrize("expiration_date", [None, "", "not-a-date"])
+ def test_missing_or_invalid_date_has_no_warning(self, expiration_date):
+ assert license_expiry_condition(expiration_date, T0) is None
+
+
+class TestDeploymentWarningConditions:
+ @staticmethod
+ def _conditions(**overrides: object) -> dict[str, WarningCondition]:
+ signals: dict[str, object] = {
+ "env_credential_login_enabled": False,
+ "license_expiration_date": None,
+ "no_redis_on_multiple_workers": False,
+ "no_redis_warning_disabled_by_env": False,
+ "detailed_debug": False,
+ "now": T0,
+ **overrides,
+ }
+ return {c.id: c for c in deployment_warning_conditions(**signals)}
+
+ def test_healthy_deployment_has_no_warnings(self):
+ assert self._conditions() == {}
+
+ def test_every_signal_maps_to_its_warning_and_severity(self):
+ conditions = self._conditions(
+ env_credential_login_enabled=True,
+ license_expiration_date="2026-11-01",
+ no_redis_on_multiple_workers=True,
+ detailed_debug=True,
+ )
+ assert {cid: c.severity for cid, c in conditions.items()} == {
+ "env_credential_login": "critical",
+ "license_expiry:2026-11-01": "warning",
+ "no_redis": "warning",
+ "detailed_debug": "warning",
+ }
+ assert all(c.env_var_ignore is None for c in conditions.values())
+
+ def test_no_redis_env_var_keeps_the_warning_listed_as_env_ignored(self):
+ conditions = self._conditions(no_redis_on_multiple_workers=True, no_redis_warning_disabled_by_env=True)
+ assert conditions["no_redis"].env_var_ignore == "LITELLM_DISABLE_NO_REDIS_WARNING"
+
+ def test_env_var_alone_does_not_raise_the_no_redis_warning(self):
+ assert self._conditions(no_redis_warning_disabled_by_env=True) == {}
+
+
+@pytest.fixture
+def set_logger_level() -> Iterator[Callable[[int], None]]:
+ original = verbose_logger.level
+ yield verbose_logger.setLevel
+ verbose_logger.setLevel(original)
+
+
+class TestCurrentWarningConditions:
+ @staticmethod
+ def _proxy_state(monkeypatch, *, redis_cache: object, general_settings: dict[str, object], license: object) -> None:
+ monkeypatch.setattr("litellm.proxy.proxy_server.redis_usage_cache", redis_cache)
+ monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
+ monkeypatch.setattr("litellm.proxy.proxy_server.premium_user_data", license)
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server._license_check",
+ SimpleNamespace(airgapped_license_data=None, license_str=None, public_key=None),
+ )
+
+ @pytest.mark.asyncio
+ async def test_reads_every_signal_from_the_running_proxy(self, monkeypatch, set_logger_level):
+ self._proxy_state(monkeypatch, redis_cache=None, general_settings={}, license={"expiration_date": "2026-10-20"})
+ monkeypatch.setenv("LITELLM_DISABLE_NO_REDIS_WARNING", "true")
+ set_logger_level(logging.DEBUG)
+
+ conditions = await current_warning_conditions(T0)
+
+ assert conditions == (
+ WarningCondition(id="env_credential_login", severity="critical"),
+ WarningCondition(id="license_expiry:2026-10-20", severity="critical"),
+ WarningCondition(id="no_redis", severity="warning", env_var_ignore="LITELLM_DISABLE_NO_REDIS_WARNING"),
+ WarningCondition(id="detailed_debug", severity="warning"),
+ )
+
+ @pytest.mark.asyncio
+ async def test_healthy_proxy_has_no_conditions(self, monkeypatch, set_logger_level):
+ self._proxy_state(
+ monkeypatch,
+ redis_cache=SimpleNamespace(),
+ general_settings={"disable_env_credential_login": True},
+ license=None,
+ )
+ monkeypatch.delenv("LITELLM_DISABLE_NO_REDIS_WARNING", raising=False)
+ set_logger_level(logging.INFO)
+
+ assert await current_warning_conditions(T0) == ()
+
+
+class TestDeploymentWarningIgnoreStore:
+ def test_no_database_means_no_store(self, monkeypatch):
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
+ assert deployment_warning_ignore_store() is None
+
+ @pytest.mark.asyncio
+ async def test_reads_ignores_from_the_writer_not_a_lagging_replica(self, monkeypatch):
+ def ui_settings_table(row: object) -> SimpleNamespace:
+ async def find_unique(where: dict[str, str]) -> object:
+ return row
+
+ return SimpleNamespace(litellm_uisettings=SimpleNamespace(find_unique=find_unique))
+
+ fresh = {"detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00Z"}}
+ monkeypatch.setattr(
+ "litellm.proxy.proxy_server.prisma_client",
+ SimpleNamespace(
+ db=ui_settings_table(None), writer_db=ui_settings_table(SimpleNamespace(ui_settings=fresh))
+ ),
+ )
+ store = deployment_warning_ignore_store()
+ assert store is not None
+ assert await store.get_raw_ignores() == fresh
+
+
+class TestGetDeploymentWarnings:
+ @pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
+ def test_admins_and_viewers_read_warnings(self, harness, role):
+ harness.as_role(role)
+ response = harness.client.get("/health/deployment_warnings")
+ assert response.status_code == 200
+ assert response.json() == {
+ "warnings": [
+ {"id": "env_credential_login", "severity": "critical", "ignore": None},
+ {"id": "detailed_debug", "severity": "warning", "ignore": None},
+ ],
+ "ignore_available": True,
+ }
+
+ @pytest.mark.parametrize(
+ "role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, LitellmUserRoles.TEAM]
+ )
+ def test_other_roles_are_rejected(self, harness, role):
+ harness.as_role(role)
+ assert harness.client.get("/health/deployment_warnings").status_code == 403
+
+ def test_env_var_ignore_is_reported(self, harness):
+ harness.with_conditions(
+ WarningCondition(id="no_redis", severity="warning", env_var_ignore="LITELLM_DISABLE_NO_REDIS_WARNING")
+ )
+ assert harness.warnings()["no_redis"]["ignore"] == {
+ "source": "env_var",
+ "env_var": "LITELLM_DISABLE_NO_REDIS_WARNING",
+ }
+
+ def test_env_var_ignore_wins_over_a_stored_admin_ignore(self, harness):
+ harness.with_conditions(WarningCondition(id="no_redis", severity="warning"))
+ harness.client.patch("/health/deployment_warnings/no_redis", json={"ignored": True})
+ harness.with_conditions(
+ WarningCondition(id="no_redis", severity="warning", env_var_ignore="LITELLM_DISABLE_NO_REDIS_WARNING")
+ )
+ assert harness.warnings()["no_redis"]["ignore"]["source"] == "env_var"
+
+ def test_stored_ignore_without_timezone_reads_as_unreadable(self, harness):
+ harness.store.raw = {
+ "detailed_debug": {"ignored_by": "a", "ignored_at": "2026-10-14T12:00:00", "ignored_until": None}
+ }
+ warnings = harness.warnings()
+ assert set(warnings) == {"env_credential_login", "detailed_debug"}
+ assert warnings["detailed_debug"]["ignore"] is None
+
+ def test_corrupt_stored_ignores_read_as_none(self, harness):
+ harness.store.raw = "{not json"
+ warnings = harness.warnings()
+ assert set(warnings) == {"env_credential_login", "detailed_debug"}
+ assert all(w["ignore"] is None for w in warnings.values())
+
+ def test_without_database_ignoring_is_reported_unavailable(self, harness):
+ harness.without_database()
+ response = harness.client.get("/health/deployment_warnings")
+ assert response.status_code == 200
+ assert response.json()["ignore_available"] is False
+ assert [w["id"] for w in response.json()["warnings"]] == ["env_credential_login", "detailed_debug"]
+
+
+class TestUpdateDeploymentWarningIgnore:
+ def test_snooze_computes_end_time_and_applies_to_every_admin(self, harness):
+ harness.as_role(LitellmUserRoles.PROXY_ADMIN, user_id="alice")
+ response = harness.client.patch(
+ "/health/deployment_warnings/env_credential_login", json={"ignored": True, "snooze_seconds": 86400}
+ )
+ assert response.status_code == 200
+ expected_ignore = {
+ "source": "admin",
+ "ignored_by": "alice",
+ "ignored_at": "2026-10-14T12:00:00Z",
+ "ignored_until": "2026-10-15T12:00:00Z",
+ }
+ assert response.json() == {"id": "env_credential_login", "ignore": expected_ignore}
+
+ harness.as_role(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, user_id="bob")
+ warnings = harness.warnings()
+ assert warnings["env_credential_login"]["ignore"] == expected_ignore
+ assert warnings["detailed_debug"]["ignore"] is None
+
+ def test_snooze_ends_exactly_at_its_end_time(self, harness):
+ harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True, "snooze_seconds": 60})
+
+ harness.clock[0] = T0 + timedelta(seconds=59)
+ assert harness.warnings()["detailed_debug"]["ignore"] is not None
+
+ harness.clock[0] = T0 + timedelta(seconds=60)
+ assert harness.warnings()["detailed_debug"]["ignore"] is None
+
+ def test_omitting_snooze_ignores_forever(self, harness):
+ response = harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True})
+ assert response.status_code == 200
+ assert response.json()["ignore"]["ignored_until"] is None
+
+ harness.clock[0] = T0 + timedelta(days=3650)
+ assert harness.warnings()["detailed_debug"]["ignore"]["ignored_at"] == "2026-10-14T12:00:00Z"
+
+ def test_restore_clears_only_that_warning(self, harness):
+ harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True})
+ harness.client.patch("/health/deployment_warnings/env_credential_login", json={"ignored": True})
+
+ response = harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": False})
+ assert response.status_code == 200
+ assert response.json() == {"id": "detailed_debug", "ignore": None}
+
+ warnings = harness.warnings()
+ assert warnings["detailed_debug"]["ignore"] is None
+ assert warnings["env_credential_login"]["ignore"]["source"] == "admin"
+
+ def test_any_severity_can_be_ignored(self, harness):
+ harness.with_conditions(WarningCondition(id="license_expiry:2026-10-16", severity="critical"))
+ response = harness.client.patch("/health/deployment_warnings/license_expiry:2026-10-16", json={"ignored": True})
+ assert response.status_code == 200
+ assert harness.warnings()["license_expiry:2026-10-16"]["ignore"]["ignored_by"] == "admin-user"
+
+ def test_ignoring_an_old_license_date_does_not_hide_the_renewed_one(self, harness):
+ harness.client.patch("/health/deployment_warnings/license_expiry:2026-10-16", json={"ignored": True})
+ harness.with_conditions(WarningCondition(id="license_expiry:2026-11-01", severity="warning"))
+ assert harness.warnings()["license_expiry:2026-11-01"]["ignore"] is None
+
+ def test_change_is_written_to_the_config_audit_log(self, harness):
+ harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True})
+ harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": False})
+
+ assert len(harness.audit.calls) == 2
+ ignore_call, restore_call = harness.audit.calls
+ assert ignore_call["table_name"] == LitellmTableNames.UI_SETTINGS_TABLE_NAME
+ assert ignore_call["before_value"] == []
+ assert ignore_call["after_value"] == [
+ {
+ "id": "detailed_debug",
+ "ignore": {
+ "source": "admin",
+ "ignored_by": "admin-user",
+ "ignored_at": "2026-10-14T12:00:00Z",
+ "ignored_until": None,
+ },
+ }
+ ]
+ assert restore_call["before_value"] == ignore_call["after_value"]
+ assert restore_call["after_value"] == []
+ assert {call["param_name"] for call in harness.audit.calls} == {DEPLOYMENT_WARNING_IGNORES_ROW_ID}
+ assert {call["action"] for call in harness.audit.calls} == {"updated"}
+
+ def test_key_without_a_user_is_recorded_as_the_proxy_admin(self, harness):
+ harness.as_role(LitellmUserRoles.PROXY_ADMIN, user_id=None)
+ response = harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True})
+ assert response.json()["ignore"]["ignored_by"] == LITELLM_PROXY_ADMIN_NAME
+
+ def test_expired_snoozes_are_not_carried_into_the_next_write(self, harness):
+ harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True, "snooze_seconds": 60})
+ harness.clock[0] = T0 + timedelta(minutes=5)
+ harness.client.patch("/health/deployment_warnings/env_credential_login", json={"ignored": True})
+
+ assert [entry["id"] for entry in harness.audit.calls[-1]["after_value"]] == ["env_credential_login"]
+
+ def test_audit_entry_survives_secret_redaction(self, harness):
+ harness.client.patch(
+ "/health/deployment_warnings/env_credential_login", json={"ignored": True, "snooze_seconds": 3600}
+ )
+
+ stored = json.loads(_dump_redacted_config(harness.audit.calls[-1]["after_value"]))
+
+ assert stored == [
+ {
+ "id": "env_credential_login",
+ "ignore": {
+ "source": "admin",
+ "ignored_by": "admin-user",
+ "ignored_at": "2026-10-14T12:00:00Z",
+ "ignored_until": "2026-10-14T13:00:00Z",
+ },
+ }
+ ]
+
+ @pytest.mark.parametrize("snooze_seconds", [0, -1, THIRTY_DAYS + 1])
+ def test_snooze_outside_thirty_days_is_rejected(self, harness, snooze_seconds):
+ response = harness.client.patch(
+ "/health/deployment_warnings/detailed_debug", json={"ignored": True, "snooze_seconds": snooze_seconds}
+ )
+ assert response.status_code == 422
+ assert harness.store.raw is None
+
+ def test_thirty_day_snooze_is_allowed(self, harness):
+ response = harness.client.patch(
+ "/health/deployment_warnings/detailed_debug", json={"ignored": True, "snooze_seconds": THIRTY_DAYS}
+ )
+ assert response.status_code == 200
+ assert response.json()["ignore"]["ignored_until"] == "2026-11-13T12:00:00Z"
+
+ def test_snooze_without_ignoring_is_rejected(self, harness):
+ response = harness.client.patch(
+ "/health/deployment_warnings/detailed_debug", json={"ignored": False, "snooze_seconds": 60}
+ )
+ assert response.status_code == 422
+
+ @pytest.mark.parametrize("warning_id", ["made_up", "license_expiry:soon", "license_expiry:20261016"])
+ def test_unknown_warning_is_rejected(self, harness, warning_id):
+ response = harness.client.patch(f"/health/deployment_warnings/{warning_id}", json={"ignored": True})
+ assert response.status_code == 404
+ assert harness.store.raw is None
+
+ @pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER])
+ def test_only_proxy_admins_can_ignore(self, harness, role):
+ harness.as_role(role)
+ response = harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True})
+ assert response.status_code == 403
+ assert harness.store.raw is None
+ assert harness.audit.calls == []
+
+ def test_without_database_ignoring_is_unavailable(self, harness):
+ harness.without_database()
+ response = harness.client.patch("/health/deployment_warnings/detailed_debug", json={"ignored": True})
+ assert response.status_code == 503
+ assert "database" in response.json()["detail"].lower()
+ assert harness.audit.calls == []
diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py
index 2c6e013d075..f215da5cf3e 100644
--- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py
+++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py
@@ -4840,6 +4840,21 @@ class TestNoRedisWarning:
assert await _show_no_redis_warning() is False
prisma.db.query_raw.assert_not_called()
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("coordination_redis, router_redis", [(MagicMock(), None), (None, MagicMock())])
+ async def test_redis_answers_before_the_env_var_is_looked_up(self, monkeypatch, coordination_redis, router_redis):
+ monkeypatch.setattr(
+ _health_endpoints_module,
+ "get_secret_bool",
+ MagicMock(side_effect=AssertionError("env var looked up although Redis is configured")),
+ )
+ with (
+ patch("litellm.proxy.proxy_server.redis_usage_cache", coordination_redis),
+ patch("litellm.proxy.proxy_server.llm_router", self._router(router_redis)),
+ patch("litellm.proxy.proxy_server.prisma_client", self._prisma_with_workers(5)),
+ ):
+ assert await _show_no_redis_warning() is False
+
@pytest.mark.asyncio
async def test_stays_quiet_when_only_the_router_has_redis(self, monkeypatch):
"""router_settings.redis_host alone backs cooldowns and usage-based routing."""
diff --git a/tests/unit/repositories/test_repositories.py b/tests/unit/repositories/test_repositories.py
index b7b663c2873..cbcd0b1cd28 100644
--- a/tests/unit/repositories/test_repositories.py
+++ b/tests/unit/repositories/test_repositories.py
@@ -23,6 +23,7 @@ from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.credentials_repository import CredentialsRepository
+from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.object_permission_repository import (
ObjectPermissionRepository,
@@ -1599,6 +1600,28 @@ class TestConfigRepository:
assert len(params) == 2
+class TestDeploymentWarningIgnoreRepository:
+ @pytest.mark.asyncio
+ async def test_round_trips_on_the_writer_without_touching_other_ui_settings_rows(self):
+ client = MockPrismaClient()
+ writer_table = MockTable()
+ client.writer_db = SimpleNamespace(litellm_uisettings=writer_table)
+ other_rows = {
+ "ui_settings": {"id": "ui_settings", "ui_settings": '{"disable_model_add_for_internal_users": true}'},
+ "user_banner": {"id": "user_banner", "ui_settings": '{"enabled": false}'},
+ }
+ writer_table._records.update({row_id: dict(row) for row_id, row in other_rows.items()})
+ repo = DeploymentWarningIgnoreRepository(client, use_writer=True)
+
+ assert await repo.get_raw_ignores() is None
+ await repo.upsert_ignores('{"no_redis": {}}')
+ await repo.upsert_ignores('{"detailed_debug": {}}')
+
+ assert await repo.get_raw_ignores() == '{"detailed_debug": {}}'
+ assert {row_id: writer_table._records[row_id] for row_id in other_rows} == other_rows
+ assert len(writer_table._records) == len(other_rows) + 1
+
+
class TestVerificationTokenRepositoryExtended:
@pytest.fixture
def repo(self):
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 6da601efeec..29d659ec1d9 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -7502,6 +7502,48 @@ export interface paths {
patch?: never;
trace?: never;
};
+ "/health/deployment_warnings": {
+ parameters: {
+ query?: never;
+ header?: never;
+ path?: never;
+ cookie?: never;
+ };
+ /**
+ * Get Deployment Warnings
+ * @description List the active deployment warnings with their severity and whether an admin ignored or snoozed them.
+ * Readable by proxy admins and proxy admin viewers.
+ */
+ get: operations["get_deployment_warnings_health_deployment_warnings_get"];
+ put?: never;
+ post?: never;
+ delete?: never;
+ options?: never;
+ head?: never;
+ patch?: never;
+ trace?: never;
+ };
+ "/health/deployment_warnings/{warning_id}": {
+ parameters: {
+ query?: never;
+ header?: never;
+ path?: never;
+ cookie?: never;
+ };
+ get?: never;
+ put?: never;
+ post?: never;
+ delete?: never;
+ options?: never;
+ head?: never;
+ /**
+ * Update Deployment Warning Ignore
+ * @description Ignore, snooze, or restore one deployment warning for every admin of this proxy.
+ * Only proxy admins can call this, and it needs a database.
+ */
+ patch: operations["update_deployment_warning_ignore_health_deployment_warnings__warning_id__patch"];
+ trace?: never;
+ };
"/health/drain": {
parameters: {
query?: never;
@@ -25449,6 +25491,26 @@ export interface components {
[key: string]: string;
};
};
+ /** AdminIgnore */
+ AdminIgnore: {
+ /**
+ * Ignored At
+ * Format: date-time
+ */
+ ignored_at: string;
+ /** Ignored By */
+ ignored_by: string;
+ /**
+ * Ignored Until
+ * @description When the snooze ends. Null means the warning is ignored until an admin restores it.
+ */
+ ignored_until?: string | null;
+ /**
+ * @description discriminator enum property added by openapi-typescript
+ * @enum {string}
+ */
+ source: "admin";
+ };
/**
* AgentCapabilities
* @description Defines optional capabilities supported by an agent.
@@ -31496,6 +31558,47 @@ export interface components {
} & {
[key: string]: unknown;
};
+ /** DeploymentWarning */
+ DeploymentWarning: {
+ /** Id */
+ id: string;
+ /** Ignore */
+ ignore?: (components["schemas"]["AdminIgnore"] | components["schemas"]["EnvVarIgnore"]) | null;
+ /**
+ * Severity
+ * @enum {string}
+ */
+ severity: "critical" | "warning";
+ };
+ /** DeploymentWarningIgnoreResponse */
+ DeploymentWarningIgnoreResponse: {
+ /** Id */
+ id: string;
+ ignore: components["schemas"]["AdminIgnore"] | null;
+ };
+ /** DeploymentWarningIgnoreUpdate */
+ DeploymentWarningIgnoreUpdate: {
+ /**
+ * Ignored
+ * @description True ignores the warning for every admin, false restores it.
+ */
+ ignored: boolean;
+ /**
+ * Snooze Seconds
+ * @description How long to ignore the warning, up to 30 days. Leave it out to ignore it forever.
+ */
+ snooze_seconds?: number | null;
+ };
+ /** DeploymentWarningsResponse */
+ DeploymentWarningsResponse: {
+ /**
+ * Ignore Available
+ * @description False when no database is connected, so warnings cannot be ignored.
+ */
+ ignore_available: boolean;
+ /** Warnings */
+ warnings: components["schemas"]["DeploymentWarning"][];
+ };
/** DiscoverAgentRequest */
DiscoverAgentRequest: {
/**
@@ -31834,6 +31937,16 @@ export interface components {
/** Tenant Id */
tenant_id: string;
};
+ /** EnvVarIgnore */
+ EnvVarIgnore: {
+ /** Env Var */
+ env_var: string;
+ /**
+ * @description discriminator enum property added by openapi-typescript
+ * @enum {string}
+ */
+ source: "env_var";
+ };
/** EnvironmentReport */
EnvironmentReport: {
/** Config Lines */
@@ -61046,6 +61159,61 @@ export interface operations {
};
};
};
+ get_deployment_warnings_health_deployment_warnings_get: {
+ parameters: {
+ query?: never;
+ header?: never;
+ path?: never;
+ cookie?: never;
+ };
+ requestBody?: never;
+ responses: {
+ /** @description Successful Response */
+ 200: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": components["schemas"]["DeploymentWarningsResponse"];
+ };
+ };
+ };
+ };
+ update_deployment_warning_ignore_health_deployment_warnings__warning_id__patch: {
+ parameters: {
+ query?: never;
+ header?: never;
+ path: {
+ warning_id: string;
+ };
+ cookie?: never;
+ };
+ requestBody: {
+ content: {
+ "application/json": components["schemas"]["DeploymentWarningIgnoreUpdate"];
+ };
+ };
+ responses: {
+ /** @description Successful Response */
+ 200: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": components["schemas"]["DeploymentWarningIgnoreResponse"];
+ };
+ };
+ /** @description Validation Error */
+ 422: {
+ headers: {
+ [name: string]: unknown;
+ };
+ content: {
+ "application/json": components["schemas"]["HTTPValidationError"];
+ };
+ };
+ };
+ };
health_drain_health_drain_get: {
parameters: {
query?: never;
From 389403fedebbf9f585815519a9316c4a84e5c17f Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 14 Oct 2026 10:57:34 -0700
Subject: [PATCH 19/47] fix(proxy): check previous_response_id ownership on
/v1/responses/compact (#46408)
* fix(proxy): check previous_response_id ownership on /v1/responses/compact
* test(proxy): cover the unmanaged response id escape hatches on every response id route
* test(proxy): audit compact previous_response_id ownership
Add integration cells for the native /v1/responses/compact routes covering owner,
admin, cross-tenant, unmanaged, malformed, duplicated and retained ids, provider
errors, the response cache, id-security settings, an upstream outage, a held
upstream, a worker kill and a proxy restart. Unskip the provider wire cell whose
bug this branch fixes.
* test(proxy): cover the no signing key path on every response id route
---
litellm/proxy/hooks/responses_id_security.py | 23 +-
.../authorization/_responses_compact.py | 371 +++++++++
.../test_responses_compact_id_ownership.py | 720 ++++++++++++++++++
...st_responses_compact_id_ownership_chaos.py | 278 +++++++
.../test_responses_id_ownership.py | 19 +-
.../providers/test_responses_compact_wire.py | 1 -
tests/unit/test_responses_id_security.py | 69 +-
7 files changed, 1458 insertions(+), 23 deletions(-)
create mode 100644 tests/integration/authorization/_responses_compact.py
create mode 100644 tests/integration/authorization/test_responses_compact_id_ownership.py
create mode 100644 tests/integration/authorization/test_responses_compact_id_ownership_chaos.py
diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py
index e554512ec91..c1384565a80 100644
--- a/litellm/proxy/hooks/responses_id_security.py
+++ b/litellm/proxy/hooks/responses_id_security.py
@@ -6,6 +6,7 @@ instead of writing immediately on each request.
"""
from collections.abc import AsyncGenerator, Callable, Mapping
+from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from fastapi import HTTPException
@@ -31,6 +32,16 @@ if TYPE_CHECKING:
_RESPONSES_API_PROVIDER_PREFIX: Final = "/openai"
_RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"})
+_ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE: Final[Mapping[str, str]] = MappingProxyType(
+ {
+ "aresponses": "previous_response_id",
+ "acompact_responses": "previous_response_id",
+ "aget_responses": "response_id",
+ "adelete_responses": "response_id",
+ "acancel_responses": "response_id",
+ "alist_input_items": "response_id",
+ }
+)
_UNMANAGED_RESPONSE_ID_DETAIL: Final = (
"Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. "
@@ -120,17 +131,9 @@ class ResponsesIDSecurity(CustomLogger):
data: dict,
call_type: CallTypesLiteral,
) -> Exception | str | dict | None:
- # MAP all the responses api response ids to the encrypted response ids
- responses_api_call_types: Final = {
- "aresponses",
- "aget_responses",
- "adelete_responses",
- "acancel_responses",
- "alist_input_items",
- }
- if call_type not in responses_api_call_types:
+ addressed_id_field: Final = _ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE.get(call_type)
+ if addressed_id_field is None:
return None
- addressed_id_field: Final = "previous_response_id" if call_type == "aresponses" else "response_id"
retained_id: Final = data.get(ADDRESSED_RESPONSE_ID_FIELD)
addressed_id: Final = (
retained_id if isinstance(retained_id, str) and retained_id else data.get(addressed_id_field)
diff --git a/tests/integration/authorization/_responses_compact.py b/tests/integration/authorization/_responses_compact.py
new file mode 100644
index 00000000000..8270d729667
--- /dev/null
+++ b/tests/integration/authorization/_responses_compact.py
@@ -0,0 +1,371 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import threading
+import uuid
+from collections import Counter
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from hashlib import sha256
+from types import MappingProxyType
+from typing import Final
+
+import httpx
+import openai
+from integration._support.client import Gateway, Scenario, eventually
+from integration._support.database import read_rows
+from integration._support.responses_vendor import same_response
+from integration._support.wire import Reply, Request, Wire
+from integration._support.worker_readiness import wait_for_worker_readiness
+from openai.types.responses import CompactedResponse, ResponseInputParam
+from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
+
+MODEL: Final = "gpt-5"
+API_KEY: Final = "synthetic-responses-key"
+ALIASES: Final = ("/v1", "", "/openai/v1")
+COMPACT_STATE: Final = "synthetic-compact-state"
+COMPACT_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
+ {
+ "input_tokens": 7,
+ "input_tokens_details": {"cached_tokens": 0},
+ "output_tokens": 3,
+ "output_tokens_details": {"reasoning_tokens": 0},
+ "total_tokens": 10,
+ }
+)
+JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1")))
+SPEND_MARKER_HEADER: Final = "x-litellm-spend-logs-metadata"
+_DISABLE_HINT: Final = (
+ "To disable this security feature, set general_settings::disable_responses_id_security to True "
+ "in the config.yaml file."
+)
+OTHER_USER_DETAIL: Final = (
+ f"Forbidden. The response id is not associated with the user, who this key belongs to. {_DISABLE_HINT}"
+)
+OTHER_TEAM_DETAIL: Final = (
+ f"Forbidden. The response id is not associated with the team, who this key belongs to. {_DISABLE_HINT}"
+)
+UNMANAGED_DETAIL: Final = (
+ "Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. "
+ "To let keys address responses this proxy did not issue, set "
+ "general_settings::allow_unmanaged_response_ids to True in the config.yaml file."
+)
+SCRIPTED_FAILURES: Final[Mapping[str, int]] = MappingProxyType(
+ {"fail400": 400, "fail401": 401, "fail404": 404, "fail500": 500}
+)
+
+
+class _InputText(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ text: str
+
+
+class _InputMessage(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ content: tuple[_InputText, ...]
+
+
+_COMPACT_INPUT: Final = TypeAdapter(tuple[_InputMessage, ...])
+
+
+def marker(label: str) -> str:
+ return f"{label}_{uuid.uuid4().hex}"
+
+
+def minted_id(created_marker: str) -> str:
+ return f"resp_minted_{created_marker}"
+
+
+def compacted_id(compact_marker: str) -> str:
+ return f"resp_compacted_{compact_marker}"
+
+
+def vendor_id() -> str:
+ return f"resp_vendor_{uuid.uuid4().hex}"
+
+
+def compact_input(text: str) -> ResponseInputParam:
+ return [{"role": "user", "content": [{"type": "input_text", "text": text}]}]
+
+
+def compact_input_json(text: str) -> list[JsonValue]:
+ return [{"role": "user", "content": [{"type": "input_text", "text": text}]}]
+
+
+def compact_body(model: str, text: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+ return {"model": model, "input": compact_input_json(text), **extra}
+
+
+def upstream_compact_body(text: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+ return {"model": MODEL, "input": compact_input_json(text), **extra}
+
+
+def compact_text(body: Mapping[str, JsonValue]) -> str:
+ return _COMPACT_INPUT.validate_python(body["input"])[0].content[0].text
+
+
+def _created_body(identity: str) -> bytes:
+ return json.dumps(
+ {
+ "id": identity,
+ "object": "response",
+ "created_at": 1,
+ "status": "completed",
+ "model": MODEL,
+ "output": [
+ {
+ "id": f"msg_{identity}",
+ "type": "message",
+ "role": "assistant",
+ "status": "completed",
+ "content": [{"type": "output_text", "text": "created", "annotations": []}],
+ }
+ ],
+ "parallel_tool_calls": False,
+ "tool_choice": "auto",
+ "tools": [],
+ }
+ ).encode()
+
+
+def _compacted_body(identity: str) -> bytes:
+ return json.dumps(
+ {
+ "id": identity,
+ "created_at": 1,
+ "object": "response.compaction",
+ "output": [
+ {
+ "id": f"rs_{identity}",
+ "type": "reasoning",
+ "summary": [{"type": "summary_text", "text": "Compacted state"}],
+ "encrypted_content": COMPACT_STATE,
+ }
+ ],
+ "usage": dict(COMPACT_USAGE),
+ }
+ ).encode()
+
+
+def _failure_body(status: int, text: str) -> bytes:
+ kind: Final = "server_error" if status >= 500 else "invalid_request_error"
+ return json.dumps(
+ {"error": {"message": f"scripted {status} for {text}", "type": kind, "param": None, "code": None}}
+ ).encode()
+
+
+@dataclass(frozen=True, slots=True)
+class CompactUpstream:
+ gates: Mapping[str, threading.Event] = field(default_factory=lambda: MappingProxyType({}))
+
+ def respond(self, request: Request) -> Reply:
+ if request.method == "GET" and request.target == "/v1/models":
+ return Reply(
+ body=json.dumps(
+ {"object": "list", "data": [{"id": MODEL, "object": "model", "created": 1, "owned_by": "openai"}]}
+ ).encode()
+ )
+ body: Final = JSON_OBJECT.validate_json(request.body)
+ if request.method == "POST" and request.target == "/v1/responses":
+ return Reply(body=_created_body(minted_id(str(body["input"]))))
+ assert request.method == "POST" and request.target == "/v1/responses/compact", request.target
+ text: Final = compact_text(body)
+ label: Final = text.split("_", 1)[0]
+ gate: Final = self.gates.get(label)
+ if gate is not None:
+ assert gate.wait(timeout=180), f"The held {label} compact was never released"
+ status: Final = SCRIPTED_FAILURES.get(label)
+ if status is not None:
+ return Reply(status=status, body=_failure_body(status, text))
+ return Reply(body=_compacted_body(compacted_id(text)))
+
+
+@dataclass(frozen=True, slots=True)
+class Upstreamed:
+ target: str
+ body: Mapping[str, JsonValue]
+
+
+def upstreamed(wire: Wire) -> tuple[Upstreamed, ...]:
+ return tuple(
+ Upstreamed(request.target, JSON_OBJECT.validate_json(request.body))
+ for request in wire.drain()
+ if request.target != "/v1/models"
+ )
+
+
+def compacts(requests: tuple[Upstreamed, ...]) -> tuple[Mapping[str, JsonValue], ...]:
+ return tuple(request.body for request in requests if request.target == "/v1/responses/compact")
+
+
+def compact_hits(requests: tuple[Upstreamed, ...], text: str) -> tuple[Mapping[str, JsonValue], ...]:
+ return tuple(body for body in compacts(requests) if compact_text(body) == text)
+
+
+def leaks(requests: tuple[Upstreamed, ...], raw_id: str) -> bool:
+ return any(raw_id in json.dumps(dict(request.body)) for request in requests)
+
+
+@dataclass(frozen=True, slots=True)
+class Tenants:
+ model: str
+ owner: str
+ other_user: str
+ other_team: str
+
+
+def deployment(scenario: Scenario, wire: Wire) -> str:
+ model: Final = scenario.model(model=f"openai/{MODEL}", api_key=API_KEY, api_base=f"{wire.url}/v1")
+ wait_for_worker_readiness(scenario.gateway, model, 1, PROXY_WORKERS)
+ return model
+
+
+def tenants(scenario: Scenario, wire: Wire) -> Tenants:
+ model: Final = deployment(scenario, wire)
+ team_a: Final = scenario.team(models=[model])
+ team_b: Final = scenario.team(models=[model])
+ user_a: Final = scenario.member(team_a)
+ user_b: Final = scenario.member(team_b)
+ scenario.gateway.post("/team/member_add", {"team_id": team_b, "member": {"role": "user", "user_id": user_a}})
+ return Tenants(
+ model=model,
+ owner=scenario.key(team_id=team_a, user_id=user_a, models=[model]),
+ other_user=scenario.key(team_id=team_b, user_id=user_b, models=[model]),
+ other_team=scenario.key(team_id=team_b, user_id=user_a, models=[model]),
+ )
+
+
+def base_url(gateway: Gateway) -> str:
+ return str(gateway.client.base_url).rstrip("/")
+
+
+def spend_marker(text: str) -> dict[str, str]:
+ return {SPEND_MARKER_HEADER: json.dumps({"marker": text})}
+
+
+def sdk(gateway: Gateway, key: str, prefix: str) -> openai.OpenAI:
+ return openai.OpenAI(base_url=f"{base_url(gateway)}{prefix}", api_key=key, max_retries=0)
+
+
+def mint(gateway: Gateway, key: str, model: str, text: str) -> str:
+ created: Final = sdk(gateway, key, "/v1").responses.create(
+ model=model, input=text, extra_headers=spend_marker(text)
+ )
+ assert created.id != minted_id(text) and same_response(created.id, minted_id(text)), created.model_dump_json()
+ return created.id
+
+
+def sdk_compact(gateway: Gateway, key: str, prefix: str, model: str, text: str, previous: str) -> CompactedResponse:
+ return sdk(gateway, key, prefix).responses.compact(
+ model=model,
+ input=compact_input(text),
+ previous_response_id=previous,
+ extra_headers=spend_marker(text),
+ )
+
+
+async def _async_compact(
+ gateway: Gateway, key: str, prefix: str, model: str, text: str, previous: str
+) -> CompactedResponse:
+ async with openai.AsyncOpenAI(base_url=f"{base_url(gateway)}{prefix}", api_key=key, max_retries=0) as client:
+ return await client.responses.compact(
+ model=model,
+ input=compact_input(text),
+ previous_response_id=previous,
+ extra_headers=spend_marker(text),
+ )
+
+
+def async_compact(gateway: Gateway, key: str, prefix: str, model: str, text: str, previous: str) -> CompactedResponse:
+ return asyncio.run(_async_compact(gateway, key, prefix, model, text, previous))
+
+
+def compact_post(gateway: Gateway, key: str, prefix: str, body: Mapping[str, JsonValue], text: str) -> httpx.Response:
+ return httpx.post(
+ f"{base_url(gateway)}{prefix}/responses/compact",
+ json=dict(body),
+ headers={"Authorization": f"Bearer {key}", **spend_marker(text)},
+ timeout=120,
+ trust_env=False,
+ )
+
+
+def compact_raw(gateway: Gateway, key: str, content: bytes, text: str) -> httpx.Response:
+ return httpx.post(
+ f"{base_url(gateway)}/v1/responses/compact",
+ content=content,
+ headers={"Authorization": f"Bearer {key}", "content-type": "application/json", **spend_marker(text)},
+ timeout=120,
+ trust_env=False,
+ )
+
+
+def assert_compacted(compacted: CompactedResponse, model: str, text: str) -> None:
+ assert compacted.id != compacted_id(text) and same_response(compacted.id, compacted_id(text)), (
+ compacted.model_dump_json()
+ )
+ assert compacted.model_dump(mode="json", exclude_none=True) == {
+ "id": compacted.id,
+ "created_at": 1,
+ "model": model,
+ "object": "response.compaction",
+ "output": [
+ {
+ "id": f"rs_{compacted_id(text)}",
+ "type": "reasoning",
+ "summary": [{"type": "summary_text", "text": "Compacted state"}],
+ "encrypted_content": COMPACT_STATE,
+ }
+ ],
+ "usage": dict(COMPACT_USAGE),
+ }, compacted.model_dump_json()
+
+
+def assert_compacted_reply(response: httpx.Response, model: str, text: str) -> str:
+ assert response.status_code == 200, response.text
+ compacted: Final = CompactedResponse.model_validate_json(response.content)
+ assert_compacted(compacted, model, text)
+ return compacted.id
+
+
+def assert_refused(response: httpx.Response, detail: str) -> None:
+ assert response.status_code == 403, response.text
+ assert response.json() == {
+ "error": {"message": detail, "type": "permission_error", "param": None, "code": "403"}
+ }, response.text
+
+
+def spend_rows(key: str) -> tuple[Mapping[str, JsonValue], ...]:
+ return tuple(
+ read_rows(
+ "SELECT request_id, status, call_type, metadata->'spend_logs_metadata'->>'marker' AS marker "
+ 'FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
+ (sha256(key.encode()).hexdigest(),),
+ )
+ )
+
+
+def _marker_counts(key: str, texts: frozenset[str]) -> Mapping[str, int]:
+ return Counter(str(row["marker"]) for row in spend_rows(key) if row["marker"] in texts)
+
+
+def landed_once(key: str, texts: frozenset[str]) -> tuple[Mapping[str, JsonValue], ...]:
+ counts: Final = eventually(lambda: _marker_counts(key, texts), lambda found: frozenset(found) == texts, seconds=90)
+ assert dict(counts) == dict.fromkeys(texts, 1), counts
+ return tuple(row for row in spend_rows(key) if row["marker"] in texts)
+
+
+def assert_compact_rows(key: str, texts: frozenset[str]) -> None:
+ rows: Final = landed_once(key, texts)
+ assert {(row["status"], row["call_type"]) for row in rows} == {("success", "acompact_responses")}, rows
+
+
+def successful_markers(key: str, texts: frozenset[str]) -> frozenset[str]:
+ return frozenset(
+ str(row["marker"]) for row in spend_rows(key) if row["marker"] in texts and row["status"] == "success"
+ )
diff --git a/tests/integration/authorization/test_responses_compact_id_ownership.py b/tests/integration/authorization/test_responses_compact_id_ownership.py
new file mode 100644
index 00000000000..e57118da129
--- /dev/null
+++ b/tests/integration/authorization/test_responses_compact_id_ownership.py
@@ -0,0 +1,720 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Callable, Mapping
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import Final, Literal
+
+import httpx
+import pytest
+from integration._support.client import Gateway, eventually
+from integration._support.process import owned_proxy
+from integration._support.responses_vendor import same_response
+from integration._support.wire import wire_server
+from integration.authorization._responses_compact import (
+ ALIASES,
+ MODEL,
+ OTHER_TEAM_DETAIL,
+ OTHER_USER_DETAIL,
+ SCRIPTED_FAILURES,
+ UNMANAGED_DETAIL,
+ CompactUpstream,
+ Tenants,
+ Upstreamed,
+ assert_compact_rows,
+ assert_compacted,
+ assert_compacted_reply,
+ assert_refused,
+ async_compact,
+ base_url,
+ compact_body,
+ compact_hits,
+ compact_post,
+ compact_raw,
+ compacted_id,
+ compacts,
+ leaks,
+ marker,
+ mint,
+ minted_id,
+ sdk,
+ sdk_compact,
+ successful_markers,
+ tenants,
+ upstream_compact_body,
+ upstreamed,
+ vendor_id,
+)
+from pydantic import JsonValue
+
+_CLIENTS: Final = ("sync", "async", "httpx")
+
+
+def _compact_with(
+ client: Literal["sync", "async", "httpx"],
+ gateway: Gateway,
+ key: str,
+ prefix: str,
+ model: str,
+ text: str,
+ previous: str,
+) -> None:
+ match client:
+ case "sync":
+ assert_compacted(sdk_compact(gateway, key, prefix, model, text, previous), model, text)
+ case "async":
+ assert_compacted(async_compact(gateway, key, prefix, model, text, previous), model, text)
+ case "httpx":
+ assert_compacted_reply(
+ compact_post(gateway, key, prefix, compact_body(model, text, {"previous_response_id": previous}), text),
+ model,
+ text,
+ )
+
+
+@pytest.mark.parametrize("prefix", ALIASES)
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_owner_compact_forwards_the_raw_previous_response_id(
+ gateway: Gateway, client: Literal["sync", "async", "httpx"], prefix: str
+) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ assert len(upstreamed(wire)) == 1
+ text: Final = marker("compact")
+ _compact_with(client, gateway, actors.owner, prefix, actors.model, text, owned)
+ assert compacts(upstreamed(wire)) == (
+ upstream_compact_body(text, {"previous_response_id": minted_id(created)}),
+ )
+ assert_compact_rows(actors.owner, frozenset({text}))
+
+
+def test_admin_compact_decodes_owned_ids_and_forwards_vendor_ids(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ upstreamed(wire)
+ foreign: Final = vendor_id()
+ owned_text: Final = marker("admin_owned")
+ foreign_text: Final = marker("admin_vendor")
+ assert_compacted(
+ sdk_compact(gateway, gateway.key, "/v1", actors.model, owned_text, owned), actors.model, owned_text
+ )
+ assert_compacted(
+ sdk_compact(gateway, gateway.key, "/v1", actors.model, foreign_text, foreign), actors.model, foreign_text
+ )
+ assert compacts(upstreamed(wire)) == (
+ upstream_compact_body(owned_text, {"previous_response_id": minted_id(created)}),
+ upstream_compact_body(foreign_text, {"previous_response_id": foreign}),
+ )
+
+
+def test_compacted_response_id_belongs_to_the_compacting_owner(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ first: Final = marker("compact")
+ compacted: Final = sdk_compact(gateway, actors.owner, "/v1", actors.model, first, owned)
+ assert_compacted(compacted, actors.model, first)
+ upstreamed(wire)
+ follow: Final = marker("follow")
+ followed: Final = sdk(gateway, actors.owner, "/v1").responses.create(
+ model=actors.model, input=follow, previous_response_id=compacted.id
+ )
+ assert same_response(followed.id, minted_id(follow)), followed.model_dump_json()
+ assert [(request.target, request.body) for request in upstreamed(wire)] == [
+ ("/v1/responses", {"model": MODEL, "input": follow, "previous_response_id": compacted_id(first)})
+ ]
+ second: Final = marker("compact_again")
+ assert_compacted(
+ sdk_compact(gateway, actors.owner, "", actors.model, second, compacted.id), actors.model, second
+ )
+ assert compacts(upstreamed(wire)) == (
+ upstream_compact_body(second, {"previous_response_id": compacted_id(first)}),
+ )
+ refused_user: Final = marker("refused_user")
+ refused_team: Final = marker("refused_team")
+ assert_refused(
+ compact_post(
+ gateway,
+ actors.other_user,
+ "/v1",
+ compact_body(actors.model, refused_user, {"previous_response_id": compacted.id}),
+ refused_user,
+ ),
+ OTHER_USER_DETAIL,
+ )
+ assert_refused(
+ compact_post(
+ gateway,
+ actors.other_team,
+ "/openai/v1",
+ compact_body(actors.model, refused_team, {"previous_response_id": compacted.id}),
+ refused_team,
+ ),
+ OTHER_TEAM_DETAIL,
+ )
+ assert upstreamed(wire) == ()
+
+
+def test_compact_without_previous_response_id_is_forwarded_unchanged(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ texts: Final = tuple(marker(f"bare{index}") for index in range(3))
+ for key, prefix, text in zip((actors.owner, actors.other_user, gateway.key), ALIASES, texts, strict=True):
+ assert_compacted_reply(
+ compact_post(gateway, key, prefix, compact_body(actors.model, text, {}), text), actors.model, text
+ )
+ assert compacts(upstreamed(wire)) == tuple(upstream_compact_body(text, {}) for text in texts)
+ assert_compact_rows(actors.owner, frozenset({texts[0]}))
+
+
+@pytest.mark.parametrize("prefix", ALIASES)
+def test_compact_refuses_other_tenants_and_unmanaged_ids_before_the_provider(gateway: Gateway, prefix: str) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ upstreamed(wire)
+ refusals: Final = (
+ (actors.other_user, owned, OTHER_USER_DETAIL, marker("other_user")),
+ (actors.other_team, owned, OTHER_TEAM_DETAIL, marker("other_team")),
+ (actors.owner, vendor_id(), UNMANAGED_DETAIL, marker("unmanaged")),
+ )
+ for key, previous, detail, text in refusals:
+ assert_refused(
+ compact_post(
+ gateway, key, prefix, compact_body(actors.model, text, {"previous_response_id": previous}), text
+ ),
+ detail,
+ )
+ assert upstreamed(wire) == ()
+ barrier: Final = marker("barrier")
+ assert_compacted_reply(
+ compact_post(
+ gateway,
+ actors.owner,
+ prefix,
+ compact_body(actors.model, barrier, {"previous_response_id": owned}),
+ barrier,
+ ),
+ actors.model,
+ barrier,
+ )
+ assert compacts(upstreamed(wire)) == (
+ upstream_compact_body(barrier, {"previous_response_id": minted_id(created)}),
+ )
+ assert_compact_rows(actors.owner, frozenset({barrier}))
+ refused_texts: Final = frozenset(text for _, _, _, text in refusals)
+ assert successful_markers(actors.owner, refused_texts) == frozenset()
+ assert successful_markers(actors.other_user, refused_texts) == frozenset()
+ assert successful_markers(actors.other_team, refused_texts) == frozenset()
+
+
+_MALFORMED: Final[Mapping[str, Callable[[str], JsonValue]]] = {
+ "int": lambda owned: 0,
+ "list": lambda owned: [owned],
+ "object": lambda owned: {"id": owned},
+ "empty": lambda owned: "",
+ "null": lambda owned: None,
+}
+
+
+@pytest.mark.parametrize("shape", tuple(_MALFORMED))
+def test_non_string_previous_response_id_never_reaches_the_provider_decoded(gateway: Gateway, shape: str) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ upstreamed(wire)
+ text: Final = marker(f"malformed_{shape}")
+ value: Final = _MALFORMED[shape](owned)
+ response: Final = compact_post(
+ gateway, actors.other_user, "/v1", compact_body(actors.model, text, {"previous_response_id": value}), text
+ )
+ received: Final = upstreamed(wire)
+ assert not leaks(received, minted_id(created)), received
+ assert (response.status_code, compacts(received)) == _MALFORMED_OUTCOME[shape](text, value), response.text
+ assert gateway.request("GET", "/health/liveliness").status_code == 200
+
+
+def _forwarded(text: str, value: JsonValue) -> tuple[int, tuple[Mapping[str, JsonValue], ...]]:
+ return 200, (upstream_compact_body(text, {"previous_response_id": value}),)
+
+
+def _dropped(text: str, value: JsonValue) -> tuple[int, tuple[Mapping[str, JsonValue], ...]]:
+ return 200, (upstream_compact_body(text, {}),)
+
+
+_MALFORMED_OUTCOME: Final[Mapping[str, Callable[[str, JsonValue], tuple[int, tuple[Mapping[str, JsonValue], ...]]]]] = {
+ "int": _forwarded,
+ "list": _forwarded,
+ "object": _forwarded,
+ "empty": _forwarded,
+ "null": _dropped,
+}
+
+
+def _tampered(owned: str) -> str:
+ middle: Final = len(owned) // 2
+ swapped: Final = "B" if owned[middle] == "A" else "A"
+ return f"{owned[:middle]}{swapped}{owned[middle + 1 :]}"
+
+
+_DEGENERATE: Final[Mapping[str, Callable[[str], str]]] = {
+ "oversized": lambda owned: "resp_" + "a" * 5120,
+ "oversized_unprefixed": lambda owned: "x" * 5120,
+ "bare_prefix": lambda owned: "resp_",
+ "tampered": _tampered,
+ "doubled_prefix": lambda owned: f"resp_{owned}",
+}
+
+
+@pytest.mark.parametrize("shape", tuple(_DEGENERATE))
+def test_degenerate_previous_response_ids_are_refused_as_unmanaged(gateway: Gateway, shape: str) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ upstreamed(wire)
+ text: Final = marker(f"degenerate_{shape}")
+ assert_refused(
+ compact_post(
+ gateway,
+ actors.owner,
+ "/v1",
+ compact_body(actors.model, text, {"previous_response_id": _DEGENERATE[shape](owned)}),
+ text,
+ ),
+ UNMANAGED_DETAIL,
+ )
+ assert upstreamed(wire) == ()
+ assert gateway.request("GET", "/health/liveliness").status_code == 200
+
+
+def _duplicated(model: str, text: str, first: str, last: str) -> bytes:
+ return (
+ f'{{"model": {json.dumps(model)}, "input": [{{"role": "user", "content": [{{"type": "input_text", '
+ f'"text": {json.dumps(text)}}}]}}], "previous_response_id": {json.dumps(first)}, '
+ f'"previous_response_id": {json.dumps(last)}}}'
+ ).encode()
+
+
+def test_duplicated_previous_response_id_is_checked_on_the_value_the_provider_gets(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ victim_mint: Final = marker("victim")
+ attacker_mint: Final = marker("attacker")
+ victim: Final = mint(gateway, actors.owner, actors.model, victim_mint)
+ attacker: Final = mint(gateway, actors.other_user, actors.model, attacker_mint)
+ upstreamed(wire)
+ same: Final = marker("dup_same")
+ assert_compacted_reply(
+ compact_raw(gateway, actors.owner, _duplicated(actors.model, same, victim, victim), same),
+ actors.model,
+ same,
+ )
+ victim_last: Final = marker("dup_victim_last")
+ assert_refused(
+ compact_raw(
+ gateway, actors.other_user, _duplicated(actors.model, victim_last, attacker, victim), victim_last
+ ),
+ OTHER_USER_DETAIL,
+ )
+ own_last: Final = marker("dup_own_last")
+ assert_compacted_reply(
+ compact_raw(gateway, actors.other_user, _duplicated(actors.model, own_last, victim, attacker), own_last),
+ actors.model,
+ own_last,
+ )
+ received: Final = upstreamed(wire)
+ assert compacts(received) == (
+ upstream_compact_body(same, {"previous_response_id": minted_id(victim_mint)}),
+ upstream_compact_body(own_last, {"previous_response_id": minted_id(attacker_mint)}),
+ )
+
+
+def test_client_sent_retained_response_id_cannot_address_another_owner(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ victim_mint: Final = marker("victim")
+ attacker_mint: Final = marker("attacker")
+ victim: Final = mint(gateway, actors.owner, actors.model, victim_mint)
+ attacker: Final = mint(gateway, actors.other_user, actors.model, attacker_mint)
+ upstreamed(wire)
+ retained_victim: Final = marker("retained_victim")
+ assert_refused(
+ compact_post(
+ gateway,
+ actors.other_user,
+ "/v1",
+ compact_body(
+ actors.model,
+ retained_victim,
+ {"previous_response_id": attacker, "_litellm_addressed_response_id": victim},
+ ),
+ retained_victim,
+ ),
+ OTHER_USER_DETAIL,
+ )
+ retained_number: Final = marker("retained_number")
+ assert_refused(
+ compact_post(
+ gateway,
+ actors.other_user,
+ "/v1",
+ compact_body(
+ actors.model, retained_number, {"previous_response_id": victim, "_litellm_addressed_response_id": 7}
+ ),
+ retained_number,
+ ),
+ OTHER_USER_DETAIL,
+ )
+ assert upstreamed(wire) == ()
+ retained_own: Final = marker("retained_own")
+ assert_compacted_reply(
+ compact_post(
+ gateway,
+ actors.other_user,
+ "/v1",
+ compact_body(
+ actors.model,
+ retained_own,
+ {"previous_response_id": victim, "_litellm_addressed_response_id": attacker},
+ ),
+ retained_own,
+ ),
+ actors.model,
+ retained_own,
+ )
+ received: Final = upstreamed(wire)
+ assert compacts(received) == (
+ upstream_compact_body(retained_own, {"previous_response_id": minted_id(attacker_mint)}),
+ )
+ assert not leaks(received, minted_id(victim_mint)), received
+
+
+@pytest.mark.parametrize("status", tuple(SCRIPTED_FAILURES.values()))
+def test_provider_errors_on_an_owned_compact_reach_the_caller(gateway: Gateway, status: int) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ upstreamed(wire)
+ failing: Final = marker(f"fail{status}")
+ response: Final = compact_post(
+ gateway, actors.owner, "/v1", compact_body(actors.model, failing, {"previous_response_id": owned}), failing
+ )
+ assert response.status_code == status, response.text
+ assert f"scripted {status} for {failing}" in response.json()["error"]["message"], response.text
+ assert compact_hits(upstreamed(wire), failing) == (
+ upstream_compact_body(failing, {"previous_response_id": minted_id(created)}),
+ )
+ after: Final = marker("after_failure")
+ assert_compacted_reply(
+ compact_post(
+ gateway, actors.owner, "/v1", compact_body(actors.model, after, {"previous_response_id": owned}), after
+ ),
+ actors.model,
+ after,
+ )
+ assert compacts(upstreamed(wire)) == (
+ upstream_compact_body(after, {"previous_response_id": minted_id(created)}),
+ )
+ assert_compact_rows(actors.owner, frozenset({after}))
+
+
+def test_response_minted_by_a_key_without_owner_is_addressable_by_any_key(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ ownerless: Final = scenario.key(models=[actors.model])
+ created: Final = marker("ownerless")
+ minted: Final = mint(gateway, ownerless, actors.model, created)
+ upstreamed(wire)
+ text: Final = marker("compact_ownerless")
+ assert_compacted(sdk_compact(gateway, actors.other_user, "/v1", actors.model, text, minted), actors.model, text)
+ assert compacts(upstreamed(wire)) == (
+ upstream_compact_body(text, {"previous_response_id": minted_id(created)}),
+ )
+
+
+def test_identical_compacts_each_reach_the_provider_with_the_cache_on(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ upstreamed(wire)
+ text: Final = marker("twin")
+ body: Final = compact_body(actors.model, text, {"previous_response_id": owned})
+ rows: Final = (marker("twin_row"), marker("twin_row"))
+ for row in rows:
+ assert_compacted_reply(compact_post(gateway, actors.owner, "/v1", body, row), actors.model, text)
+ assert (
+ compacts(upstreamed(wire))
+ == (upstream_compact_body(text, {"previous_response_id": minted_id(created)}),) * 2
+ )
+ assert_compact_rows(actors.owner, frozenset(rows))
+ assert_refused(compact_post(gateway, actors.other_user, "/v1", body, marker("twin_refused")), OTHER_USER_DETAIL)
+ assert upstreamed(wire) == ()
+
+
+def test_cached_create_still_checks_the_previous_response_id_owner(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ upstreamed(wire)
+ follow: Final = marker("cached_follow")
+ owner_client: Final = sdk(gateway, actors.owner, "/v1")
+ sent: Final = {"model": MODEL, "input": follow, "previous_response_id": minted_id(created)}
+
+ def create() -> tuple[str, tuple[Upstreamed, ...]]:
+ reply: Final = owner_client.responses.create(model=actors.model, input=follow, previous_response_id=owned)
+ return reply.id, upstreamed(wire)
+
+ first_id, first_received = create()
+ assert [(request.target, request.body) for request in first_received] == [("/v1/responses", sent)]
+ hit_id, hit_received = eventually(create, lambda attempt: not attempt[1], seconds=60)
+ assert same_response(first_id, minted_id(follow)) and same_response(hit_id, minted_id(follow)), hit_id
+ assert hit_received == ()
+ refused: Final = httpx.post(
+ f"{base_url(gateway)}/v1/responses",
+ json={"model": actors.model, "input": follow, "previous_response_id": owned},
+ headers={"Authorization": f"Bearer {actors.other_user}"},
+ timeout=120,
+ trust_env=False,
+ )
+ assert_refused(refused, OTHER_USER_DETAIL)
+ assert upstreamed(wire) == ()
+
+
+@dataclass(frozen=True, slots=True)
+class _BurstCall:
+ kind: Literal["owner", "other_user", "other_team", "unmanaged", "admin_vendor", "bare"]
+ prefix: str
+ text: str
+
+
+_BURST_KINDS: Final = ("owner", "other_user", "other_team", "unmanaged", "admin_vendor", "bare")
+_REFUSALS: Final[Mapping[str, str]] = {
+ "other_user": OTHER_USER_DETAIL,
+ "other_team": OTHER_TEAM_DETAIL,
+ "unmanaged": UNMANAGED_DETAIL,
+}
+
+
+def _burst_call(gateway: Gateway, actors: Tenants, owned: str, foreign: str, call: _BurstCall) -> httpx.Response:
+ plans: Final[Mapping[str, tuple[str, Mapping[str, JsonValue]]]] = {
+ "owner": (actors.owner, {"previous_response_id": owned}),
+ "other_user": (actors.other_user, {"previous_response_id": owned}),
+ "other_team": (actors.other_team, {"previous_response_id": owned}),
+ "unmanaged": (actors.owner, {"previous_response_id": foreign}),
+ "admin_vendor": (gateway.key, {"previous_response_id": foreign}),
+ "bare": (actors.owner, {}),
+ }
+ key, extra = plans[call.kind]
+ return compact_post(gateway, key, call.prefix, compact_body(actors.model, call.text, extra), call.text)
+
+
+def test_concurrent_mixed_compact_burst_gives_every_caller_its_own_verdict(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ created: Final = marker("mint")
+ owned: Final = mint(gateway, actors.owner, actors.model, created)
+ upstreamed(wire)
+ foreign: Final = vendor_id()
+ calls: Final = tuple(
+ _BurstCall(_BURST_KINDS[index % 6], ALIASES[index % 3], marker(f"burst{index}")) for index in range(36)
+ )
+ with ThreadPoolExecutor(max_workers=len(calls)) as pool:
+ responses: Final = tuple(pool.map(partial(_burst_call, gateway, actors, owned, foreign), calls))
+ for call, response in zip(calls, responses, strict=True):
+ if call.kind in _REFUSALS:
+ assert_refused(response, _REFUSALS[call.kind])
+ else:
+ assert_compacted_reply(response, actors.model, call.text)
+ received: Final = upstreamed(wire)
+ expected: Final[Mapping[str, Mapping[str, JsonValue]]] = {
+ "owner": {"previous_response_id": minted_id(created)},
+ "admin_vendor": {"previous_response_id": foreign},
+ "bare": {},
+ }
+ assert {call.text: compact_hits(received, call.text) for call in calls} == {
+ call.text: (upstream_compact_body(call.text, expected[call.kind]),) if call.kind in expected else ()
+ for call in calls
+ }
+ owner_texts: Final = frozenset(call.text for call in calls if call.kind in ("owner", "bare"))
+ assert_compact_rows(actors.owner, owner_texts)
+ refused_texts: Final = frozenset(call.text for call in calls if call.kind in _REFUSALS)
+ assert successful_markers(actors.owner, refused_texts) == frozenset()
+ assert successful_markers(actors.other_user, refused_texts) == frozenset()
+ assert successful_markers(actors.other_team, refused_texts) == frozenset()
+
+
+def _settings_proxy_config(tmp_path: Path, name: str, settings: str) -> Path:
+ config: Final = tmp_path / f"{name}.yaml"
+ config.write_text(f"general_settings:\n{settings}")
+ return config
+
+
+def test_allow_unmanaged_response_ids_still_scopes_proxy_issued_ids_on_compact(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ config: Final = _settings_proxy_config(tmp_path, "allow_unmanaged", " allow_unmanaged_response_ids: true\n")
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ with owned_proxy(gateway, tmp_path, {"INTEGRATION_PROXY_WORKERS": "1"}, config=config) as allowed:
+ created: Final = marker("mint")
+ owned: Final = mint(allowed, actors.owner, actors.model, created)
+ upstreamed(wire)
+ foreign: Final = vendor_id()
+ vendor_text: Final = marker("allowed_vendor")
+ assert_compacted_reply(
+ compact_post(
+ allowed,
+ actors.owner,
+ "/v1",
+ compact_body(actors.model, vendor_text, {"previous_response_id": foreign}),
+ vendor_text,
+ ),
+ actors.model,
+ vendor_text,
+ )
+ other_text: Final = marker("allowed_other_user")
+ assert_refused(
+ compact_post(
+ allowed,
+ actors.other_user,
+ "/v1",
+ compact_body(actors.model, other_text, {"previous_response_id": owned}),
+ other_text,
+ ),
+ OTHER_USER_DETAIL,
+ )
+ owner_text: Final = marker("allowed_owner")
+ assert_compacted_reply(
+ compact_post(
+ allowed,
+ actors.owner,
+ "/v1",
+ compact_body(actors.model, owner_text, {"previous_response_id": owned}),
+ owner_text,
+ ),
+ actors.model,
+ owner_text,
+ )
+ received: Final = upstreamed(wire)
+ assert compacts(received) == (
+ upstream_compact_body(vendor_text, {"previous_response_id": foreign}),
+ upstream_compact_body(owner_text, {"previous_response_id": minted_id(created)}),
+ )
+
+
+def test_disable_responses_id_security_lets_any_key_compact_and_forwards_raw_ids(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ config: Final = _settings_proxy_config(tmp_path, "disable_security", " disable_responses_id_security: true\n")
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ sealed_mint: Final = marker("sealed")
+ sealed: Final = mint(gateway, actors.owner, actors.model, sealed_mint)
+ with owned_proxy(gateway, tmp_path, {"INTEGRATION_PROXY_WORKERS": "1"}, config=config) as disabled:
+ open_mint: Final = marker("open")
+ opened: Final = mint(disabled, actors.owner, actors.model, open_mint)
+ upstreamed(wire)
+ foreign: Final = vendor_id()
+ plans: Final = (
+ (marker("disabled_open"), opened, minted_id(open_mint)),
+ (marker("disabled_sealed"), sealed, minted_id(sealed_mint)),
+ (marker("disabled_vendor"), foreign, foreign),
+ )
+ for text, previous, _ in plans:
+ assert_compacted_reply(
+ compact_post(
+ disabled,
+ actors.other_user,
+ "/v1",
+ compact_body(actors.model, text, {"previous_response_id": previous}),
+ text,
+ ),
+ actors.model,
+ text,
+ )
+ received: Final = upstreamed(wire)
+ assert compacts(received) == tuple(
+ upstream_compact_body(text, {"previous_response_id": raw}) for text, _, raw in plans
+ )
+
+
+def test_null_and_empty_id_security_settings_keep_the_compact_checks(gateway: Gateway, tmp_path: Path) -> None:
+ config: Final = _settings_proxy_config(
+ tmp_path,
+ "null_and_empty",
+ " allow_unmanaged_response_ids: null\n disable_responses_id_security: ''\n",
+ )
+ upstream: Final = CompactUpstream()
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ with owned_proxy(gateway, tmp_path, {"INTEGRATION_PROXY_WORKERS": "1"}, config=config) as defaulted:
+ created: Final = marker("mint")
+ owned: Final = mint(defaulted, actors.owner, actors.model, created)
+ upstreamed(wire)
+ vendor_text: Final = marker("defaulted_vendor")
+ assert_refused(
+ compact_post(
+ defaulted,
+ actors.owner,
+ "/v1",
+ compact_body(actors.model, vendor_text, {"previous_response_id": vendor_id()}),
+ vendor_text,
+ ),
+ UNMANAGED_DETAIL,
+ )
+ other_text: Final = marker("defaulted_other_user")
+ assert_refused(
+ compact_post(
+ defaulted,
+ actors.other_user,
+ "/v1",
+ compact_body(actors.model, other_text, {"previous_response_id": owned}),
+ other_text,
+ ),
+ OTHER_USER_DETAIL,
+ )
+ assert upstreamed(wire) == ()
+ owner_text: Final = marker("defaulted_owner")
+ assert_compacted_reply(
+ compact_post(
+ defaulted,
+ actors.owner,
+ "/v1",
+ compact_body(actors.model, owner_text, {"previous_response_id": owned}),
+ owner_text,
+ ),
+ actors.model,
+ owner_text,
+ )
+ received: Final = upstreamed(wire)
+ assert compacts(received) == (upstream_compact_body(owner_text, {"previous_response_id": minted_id(created)}),)
diff --git a/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py b/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py
new file mode 100644
index 00000000000..66af1fa3f26
--- /dev/null
+++ b/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py
@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import threading
+from collections.abc import Mapping
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final, Literal
+from urllib.parse import urlsplit
+
+import httpx
+import psutil
+import pytest
+from integration._support.client import Gateway, eventually
+from integration._support.process import graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Wire, wire_server
+from integration.authorization._responses_compact import (
+ ALIASES,
+ OTHER_TEAM_DETAIL,
+ OTHER_USER_DETAIL,
+ UNMANAGED_DETAIL,
+ CompactUpstream,
+ Tenants,
+ Upstreamed,
+ assert_compact_rows,
+ assert_compacted_reply,
+ assert_refused,
+ compact_body,
+ compact_hits,
+ compact_post,
+ marker,
+ mint,
+ minted_id,
+ successful_markers,
+ tenants,
+ upstream_compact_body,
+ upstreamed,
+ vendor_id,
+)
+
+_Kind = Literal["owner", "other_user", "other_team", "unmanaged"]
+_KINDS: Final[tuple[_Kind, ...]] = ("owner", "other_user", "other_team", "unmanaged")
+_REFUSING: Final[tuple[_Kind, ...]] = ("other_user", "other_team", "unmanaged")
+_OWNER: Final[tuple[_Kind, ...]] = ("owner",)
+_REFUSALS: Final[Mapping[str, str]] = {
+ "other_user": OTHER_USER_DETAIL,
+ "other_team": OTHER_TEAM_DETAIL,
+ "unmanaged": UNMANAGED_DETAIL,
+}
+
+
+@dataclass(frozen=True, slots=True)
+class _Call:
+ kind: _Kind
+ prefix: str
+ text: str
+
+
+@dataclass(frozen=True, slots=True)
+class _Outcome:
+ call: _Call
+ response: httpx.Response | None
+
+
+@dataclass(frozen=True, slots=True)
+class _Target:
+ gateway: Gateway
+ actors: Tenants
+ owned: str
+ foreign: str
+
+
+def _calls(label: str, count: int, kinds: tuple[_Kind, ...] = _KINDS) -> tuple[_Call, ...]:
+ return tuple(
+ _Call(kinds[index % len(kinds)], ALIASES[index % len(ALIASES)], marker(label)) for index in range(count)
+ )
+
+
+def _send(target: _Target, call: _Call) -> _Outcome:
+ plans: Final[Mapping[str, tuple[str, str]]] = {
+ "owner": (target.actors.owner, target.owned),
+ "other_user": (target.actors.other_user, target.owned),
+ "other_team": (target.actors.other_team, target.owned),
+ "unmanaged": (target.actors.owner, target.foreign),
+ }
+ key, previous = plans[call.kind]
+ body: Final = compact_body(target.actors.model, call.text, {"previous_response_id": previous})
+ try:
+ return _Outcome(call, compact_post(target.gateway, key, call.prefix, body, call.text))
+ except httpx.TransportError:
+ return _Outcome(call, None)
+
+
+def _submit(pool: ThreadPoolExecutor, target: _Target, calls: tuple[_Call, ...]) -> tuple[Future[_Outcome], ...]:
+ return tuple(pool.submit(_send, target, call) for call in calls)
+
+
+def _results(futures: tuple[Future[_Outcome], ...]) -> tuple[_Outcome, ...]:
+ return tuple(future.result(timeout=180) for future in futures)
+
+
+def _send_all(target: _Target, calls: tuple[_Call, ...]) -> tuple[_Outcome, ...]:
+ with ThreadPoolExecutor(max_workers=len(calls)) as pool:
+ return _results(_submit(pool, target, calls))
+
+
+def _done(futures: tuple[Future[_Outcome], ...]) -> int:
+ return sum(future.done() for future in futures)
+
+
+def _assert_verdict(outcome: _Outcome, model: str) -> None:
+ assert outcome.response is not None, outcome.call
+ if outcome.call.kind in _REFUSALS:
+ assert_refused(outcome.response, _REFUSALS[outcome.call.kind])
+ return
+ assert_compacted_reply(outcome.response, model, outcome.call.text)
+
+
+def _served(outcome: _Outcome, model: str) -> bool:
+ if outcome.response is None or outcome.response.status_code >= 500:
+ return False
+ _assert_verdict(outcome, model)
+ return True
+
+
+def _texts(outcomes: tuple[_Outcome, ...], kinds: tuple[_Kind, ...]) -> frozenset[str]:
+ return frozenset(outcome.call.text for outcome in outcomes if outcome.call.kind in kinds)
+
+
+def _assert_upstream(
+ received: tuple[Upstreamed, ...], owner_texts: frozenset[str], refused_texts: frozenset[str], created: str
+) -> None:
+ assert {text: compact_hits(received, text) for text in owner_texts} == {
+ text: (upstream_compact_body(text, {"previous_response_id": minted_id(created)}),) for text in owner_texts
+ }
+ assert {text: compact_hits(received, text) for text in refused_texts} == dict.fromkeys(refused_texts, ())
+
+
+def _held_connections(pid: int, wire: Wire) -> int:
+ port: Final = urlsplit(wire.url).port
+ return sum(
+ 1
+ for connection in psutil.Process(pid).net_connections(kind="tcp")
+ if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
+ )
+
+
+def _settled(target: _Target) -> bool:
+ text: Final = marker("settle")
+ body: Final = compact_body(target.actors.model, text, {})
+ try:
+ return compact_post(target.gateway, target.actors.owner, "/v1", body, text).status_code == 200
+ except httpx.TransportError:
+ return False
+
+
+def test_upstream_outage_mid_burst_keeps_refusals_and_lands_each_owner_compact_once(gateway: Gateway) -> None:
+ upstream: Final = CompactUpstream()
+ created: Final = marker("mint")
+ with gateway.scenario() as scenario:
+ with wire_server(upstream.respond) as first:
+ port: Final = urlsplit(first.url).port
+ actors: Final = tenants(scenario, first)
+ target: Final = _Target(gateway, actors, mint(gateway, actors.owner, actors.model, created), vendor_id())
+ upstreamed(first)
+ before: Final = _send_all(target, _calls("before", 16))
+ before_received: Final = upstreamed(first)
+ during: Final = _send_all(target, _calls("during", 16))
+ alive: Final = gateway.request("GET", "/health/liveliness")
+ assert port is not None
+ with wire_server(upstream.respond, port=port) as second:
+ eventually(lambda: _settled(target), lambda settled: settled, seconds=60)
+ after: Final = _send_all(target, _calls("after", 16))
+ after_received: Final = upstreamed(second)
+ assert alive.status_code == 200, alive.text
+ for outcome in (*before, *after):
+ _assert_verdict(outcome, actors.model)
+ for outcome in during:
+ assert outcome.response is not None, outcome.call
+ if outcome.call.kind in _REFUSALS:
+ assert_refused(outcome.response, _REFUSALS[outcome.call.kind])
+ else:
+ assert outcome.response.status_code >= 500, outcome.response.text
+ _assert_upstream(before_received, _texts(before, _OWNER), _texts(before, _REFUSING), created)
+ _assert_upstream(after_received, _texts(after, _OWNER), _texts((*during, *after), _REFUSING), created)
+ assert {text: compact_hits(after_received, text) for text in _texts(during, _OWNER)} == dict.fromkeys(
+ _texts(during, _OWNER), ()
+ )
+ assert_compact_rows(actors.owner, _texts((*before, *after), _OWNER))
+ refused: Final = _texts((*before, *during, *after), _REFUSING)
+ assert successful_markers(actors.other_user, refused) == frozenset()
+ assert successful_markers(actors.other_team, refused) == frozenset()
+ assert successful_markers(actors.owner, refused | _texts(during, _OWNER)) == frozenset()
+
+
+def test_held_owner_compacts_do_not_hold_up_refusals(gateway: Gateway) -> None:
+ release: Final = threading.Event()
+ upstream: Final = CompactUpstream(MappingProxyType({"held": release}))
+ created: Final = marker("mint")
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ target: Final = _Target(gateway, actors, mint(gateway, actors.owner, actors.model, created), vendor_id())
+ upstreamed(wire)
+ held: Final = _calls("held", 8, _OWNER)
+ with ThreadPoolExecutor(max_workers=len(held)) as pool:
+ try:
+ pending: Final = _submit(pool, target, held)
+ eventually(lambda: wire.received.qsize(), lambda size: size >= len(held), seconds=90)
+ refused: Final = _send_all(target, _calls("refusal", 12, _REFUSING))
+ done_while_held: Final = _done(pending)
+ finally:
+ release.set()
+ served: Final = _results(pending)
+ received: Final = upstreamed(wire)
+ assert done_while_held == 0
+ for outcome in (*refused, *served):
+ _assert_verdict(outcome, actors.model)
+ _assert_upstream(received, _texts(served, _OWNER), _texts(refused, _REFUSING), created)
+ assert_compact_rows(actors.owner, _texts(served, _OWNER))
+
+
+@pytest.mark.timeout(4 * graceful_stop_seconds() + 300)
+def test_worker_kill_and_proxy_restart_keep_compact_ownership(gateway: Gateway, tmp_path: Path) -> None:
+ kill_gate: Final = threading.Event()
+ stop_gate: Final = threading.Event()
+ upstream: Final = CompactUpstream(MappingProxyType({"heldkill": kill_gate, "heldstop": stop_gate}))
+ created: Final = marker("mint")
+ with wire_server(upstream.respond) as wire, gateway.scenario() as scenario:
+ actors: Final = tenants(scenario, wire)
+ with owned_proxy_process(gateway, tmp_path, {}, workers=2) as first:
+ workers: Final = eventually(
+ first.worker_pids, lambda found: len(found) == 2, seconds=graceful_stop_seconds() + 60
+ )
+ target: Final = _Target(
+ first.gateway, actors, mint(first.gateway, actors.owner, actors.model, created), vendor_id()
+ )
+ upstreamed(wire)
+ held_kill: Final = _calls("heldkill", 12, _OWNER)
+ with ThreadPoolExecutor(max_workers=len(held_kill)) as pool:
+ try:
+ in_flight: Final = _submit(pool, target, held_kill)
+ eventually(lambda: wire.received.qsize(), lambda size: size >= len(held_kill), seconds=90)
+ held_by: Final = MappingProxyType({pid: _held_connections(pid, wire) for pid in workers})
+ victim: Final = max(workers, key=held_by.__getitem__)
+ with first.replacing_worker(victim) as kill:
+ kill()
+ kill_gate.set()
+ killed: Final = _results(in_flight)
+ finally:
+ kill_gate.set()
+ after_kill: Final = _send_all(target, _calls("after_kill", 12))
+ kill_received: Final = upstreamed(wire)
+ held_stop: Final = _calls("heldstop", 8, _OWNER)
+ with ThreadPoolExecutor(max_workers=len(held_stop)) as pool:
+ try:
+ stopping: Final = _submit(pool, target, held_stop)
+ eventually(lambda: wire.received.qsize(), lambda size: size >= len(held_stop), seconds=90)
+ first.begin_shutdown()
+ stop_gate.set()
+ stopped: Final = _results(stopping)
+ finally:
+ stop_gate.set()
+ stop_received: Final = upstreamed(wire)
+ with owned_proxy_process(gateway, tmp_path, {}, workers=2) as second:
+ rebooted: Final = _send_all(
+ _Target(second.gateway, actors, target.owned, target.foreign), _calls("rebooted", 12)
+ )
+ reboot_received: Final = upstreamed(wire)
+ assert sum(held_by.values()) == len(held_kill) and held_by[victim] > 0, held_by
+ served_through_kill: Final = tuple(outcome for outcome in killed if _served(outcome, actors.model))
+ assert len(served_through_kill) == len(held_kill) - held_by[victim], (held_by, killed)
+ for outcome in (*after_kill, *stopped, *rebooted):
+ _assert_verdict(outcome, actors.model)
+ _assert_upstream(kill_received, _texts((*killed, *after_kill), _OWNER), _texts(after_kill, _REFUSING), created)
+ _assert_upstream(stop_received, _texts(stopped, _OWNER), frozenset(), created)
+ _assert_upstream(reboot_received, _texts(rebooted, _OWNER), _texts(rebooted, _REFUSING), created)
+ assert_compact_rows(actors.owner, _texts((*served_through_kill, *after_kill, *rebooted), _OWNER))
diff --git a/tests/integration/authorization/test_responses_id_ownership.py b/tests/integration/authorization/test_responses_id_ownership.py
index 9e13b45e385..9a88349cca0 100644
--- a/tests/integration/authorization/test_responses_id_ownership.py
+++ b/tests/integration/authorization/test_responses_id_ownership.py
@@ -16,6 +16,7 @@ from integration._support.process import owned_proxy
from integration._support.responses_vendor import same_response
from integration._support.wire import Reply, Request, Wire, wire_server
from integration._support.worker_readiness import wait_for_worker_readiness
+from openai.types.responses import ResponseInputParam
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
_MODEL: Final = "gpt-5"
@@ -257,11 +258,10 @@ def test_compact_previous_response_id_is_scoped_to_the_issuing_user_and_team(
gateway: Gateway,
prefix: str,
) -> None:
- pytest.skip(
- "BUG: POST /v1/responses/compact does not check previous_response_id ownership; "
- "another team's id and unmanaged ids reach the provider"
- )
raw_id: Final = f"resp_compact_owned_{uuid.uuid4().hex}"
+ compact_input: Final[ResponseInputParam] = [
+ {"role": "user", "content": [{"type": "input_text", "text": "compact the owned response"}]}
+ ]
def respond(request: Request) -> Reply:
if request.method == "GET" and request.target == "/v1/models":
@@ -317,6 +317,17 @@ def test_compact_previous_response_id_is_scoped_to_the_issuing_user_and_team(
)
assert _contract_requests(wire) == ()
+ client_a.responses.compact(model=model, previous_response_id=client_id, input=compact_input)
+ compacted_requests: Final = _contract_requests(wire)
+ assert [(request.method, request.target) for request in compacted_requests] == [
+ ("POST", "/v1/responses/compact")
+ ], compacted_requests
+ assert _JSON_OBJECT.validate_json(compacted_requests[0].body) == {
+ "model": _MODEL,
+ "input": compact_input,
+ "previous_response_id": raw_id,
+ }, compacted_requests[0].body
+
def test_allow_unmanaged_response_ids_forward_the_vendor_id_on_create(
gateway: Gateway,
diff --git a/tests/integration/providers/test_responses_compact_wire.py b/tests/integration/providers/test_responses_compact_wire.py
index bc1c1ba2a7a..4da3f411cb5 100644
--- a/tests/integration/providers/test_responses_compact_wire.py
+++ b/tests/integration/providers/test_responses_compact_wire.py
@@ -285,7 +285,6 @@ def test_compact_previous_response_id_reaches_the_provider_as_the_raw_id(
gateway: Gateway,
prefix: str,
) -> None:
- pytest.skip("BUG: compact forwards the proxy-encrypted previous_response_id instead of the raw provider id")
raw_id: Final = f"resp_compact_prev_{uuid.uuid4().hex}"
completed_id: Final = f"resp_compact_result_{uuid.uuid4().hex}"
input_items: Final = [{"role": "user", "content": [{"type": "input_text", "text": "Summarize the patch so far."}]}]
diff --git a/tests/unit/test_responses_id_security.py b/tests/unit/test_responses_id_security.py
index 704a52fc202..c7fd7fe247f 100644
--- a/tests/unit/test_responses_id_security.py
+++ b/tests/unit/test_responses_id_security.py
@@ -805,6 +805,7 @@ _FABRICATED_UNMANAGED_ID = "resp_fabricatedunmanagedidbbbbbbbbbbbbbbbb"
_UNIT_TEST_SALT_KEY = "lit6837-unit-test-salt-key"
_ADDRESSED_ID_FIELD_BY_CALL_TYPE = {
"aresponses": "previous_response_id",
+ "acompact_responses": "previous_response_id",
"aget_responses": "response_id",
"adelete_responses": "response_id",
"acancel_responses": "response_id",
@@ -865,6 +866,52 @@ class TestUnrecognizedResponseIdIsRejected:
assert "allow_unmanaged_response_ids" in exc_info.value.detail
assert data[field] == _FABRICATED_UNMANAGED_ID
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE))
+ async def test_owner_gets_the_provider_id_on_every_route(self, mock_cache, salt_key_env, call_type):
+ field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type]
+ hook = _hook()
+ owner = _auth()
+ data = {field: _issue_managed_id(hook, owner)}
+
+ result = await hook.async_pre_call_hook(
+ user_api_key_dict=owner,
+ cache=mock_cache,
+ data=data,
+ call_type=call_type,
+ )
+
+ assert result[field] == _FABRICATED_PROVIDER_RESPONSE_ID
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE))
+ @pytest.mark.parametrize(
+ ("caller_user_id", "caller_team_id", "detail"),
+ [
+ ("stranger-user", "owner-team", "not associated with the user"),
+ ("owner-user", "stranger-team", "not associated with the team"),
+ ],
+ )
+ async def test_other_user_or_team_is_rejected_on_every_route(
+ self, mock_cache, salt_key_env, call_type, caller_user_id, caller_team_id, detail
+ ):
+ field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type]
+ hook = _hook()
+ issued_id = _issue_managed_id(hook, _auth())
+ data = {field: issued_id}
+
+ with pytest.raises(HTTPException) as exc_info:
+ await hook.async_pre_call_hook(
+ user_api_key_dict=_auth(user_id=caller_user_id, team_id=caller_team_id),
+ cache=mock_cache,
+ data=data,
+ call_type=call_type,
+ )
+
+ assert exc_info.value.status_code == 403
+ assert detail in exc_info.value.detail
+ assert data[field] == issued_id
+
@pytest.mark.asyncio
async def test_owner_can_still_address_the_id_the_proxy_issued_it(self, mock_cache, salt_key_env):
hook = _hook()
@@ -937,31 +984,37 @@ class TestUnmanagedResponseIdEscapeHatches:
"general_settings",
[{"allow_unmanaged_response_ids": True}, {"disable_responses_id_security": True}],
)
- async def test_opted_in_settings_forward_the_id_untouched(self, mock_cache, salt_key_env, general_settings):
- data = {"response_id": _FABRICATED_UNMANAGED_ID}
+ @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE))
+ async def test_opted_in_settings_forward_the_id_untouched(
+ self, mock_cache, salt_key_env, general_settings, call_type
+ ):
+ field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type]
+ data = {field: _FABRICATED_UNMANAGED_ID}
result = await _hook(general_settings=general_settings).async_pre_call_hook(
user_api_key_dict=_auth(),
cache=mock_cache,
data=data,
- call_type="aget_responses",
+ call_type=call_type,
)
- assert result["response_id"] == _FABRICATED_UNMANAGED_ID
+ assert result[field] == _FABRICATED_UNMANAGED_ID
@pytest.mark.asyncio
- async def test_proxy_without_a_signing_key_forwards_the_id_untouched(self, mock_cache, monkeypatch):
+ @pytest.mark.parametrize("call_type", sorted(_ADDRESSED_ID_FIELD_BY_CALL_TYPE))
+ async def test_proxy_without_a_signing_key_forwards_the_id_untouched(self, mock_cache, monkeypatch, call_type):
monkeypatch.delenv("LITELLM_SALT_KEY", raising=False)
- data = {"response_id": _FABRICATED_UNMANAGED_ID}
+ field = _ADDRESSED_ID_FIELD_BY_CALL_TYPE[call_type]
+ data = {field: _FABRICATED_UNMANAGED_ID}
result = await _hook(signing_key=None).async_pre_call_hook(
user_api_key_dict=_auth(),
cache=mock_cache,
data=data,
- call_type="aget_responses",
+ call_type=call_type,
)
- assert result["response_id"] == _FABRICATED_UNMANAGED_ID
+ assert result[field] == _FABRICATED_UNMANAGED_ID
@pytest.mark.asyncio
async def test_proxy_admin_may_address_an_unmanaged_id(self, mock_cache, salt_key_env):
From 4f34856a659a35e39c73f7443eff274560ef0a3d Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 11:25:15 -0700
Subject: [PATCH 20/47] refactor(router): name the pre-content stream retry
helpers for both stream paths and let the retry loop take its attempt
function (#46474)
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/router.py | 84 +++++++++++--------
.../router_code_coverage.py | 6 +-
tests/unit/test_router/test_router.py | 34 ++++----
3 files changed, 67 insertions(+), 57 deletions(-)
diff --git a/litellm/router.py b/litellm/router.py
index 313c404cabe..ee8a58f21e2 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -24,6 +24,7 @@ from collections import defaultdict
from collections.abc import (
AsyncGenerator,
AsyncIterator,
+ Awaitable,
Callable,
Generator,
Iterable,
@@ -628,7 +629,7 @@ def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error
return has_generated_content or not error.is_pre_first_chunk
-def _anthropic_stream_raised_error_status(error: Exception) -> int | None:
+def _stream_raised_error_status(error: Exception) -> int | None:
raw_status: Final = getattr(error, "status_code", None)
if isinstance(raw_status, int):
return raw_status
@@ -638,27 +639,29 @@ def _anthropic_stream_raised_error_status(error: Exception) -> int | None:
return response_status if isinstance(response_status, int) else None
-def _anthropic_stream_fallback_error_for_raised(
- error: Exception, model: str, has_generated_content: bool
+def _stream_fallback_error_for_raised(
+ error: Exception, model: str, has_generated_content: bool, llm_provider: str = "anthropic"
) -> "MidStreamFallbackError | None":
"""The pre-stream retry rule (408, 409, 429, 5xx); None means the raise propagates unchanged."""
if has_generated_content:
return None
- status_code: Final = _anthropic_stream_raised_error_status(error)
+ status_code: Final = _stream_raised_error_status(error)
if status_code is None:
- return _anthropic_stream_pre_content_error(error, model)
+ return _stream_pre_content_error(error, model, llm_provider)
retriable: Final = litellm.should_retry(status_code)
- return _anthropic_stream_pre_content_error(error, model) if retriable else None
+ return _stream_pre_content_error(error, model, llm_provider) if retriable else None
-def _anthropic_stream_pre_content_error(error: Exception, model: str) -> "MidStreamFallbackError":
+def _stream_pre_content_error(
+ error: Exception, model: str, llm_provider: str = "anthropic"
+) -> "MidStreamFallbackError":
"""The envelope the fallback chain judges a failure by when the client has received no content yet."""
from litellm.exceptions import MidStreamFallbackError
return MidStreamFallbackError(
message=str(error),
model=model,
- llm_provider="anthropic",
+ llm_provider=llm_provider,
original_exception=error,
is_pre_first_chunk=True,
)
@@ -688,7 +691,7 @@ def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception:
return error.original_exception if error.original_exception is not None else error
-def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool:
+def _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
@@ -705,18 +708,18 @@ def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, bu
@dataclass(frozen=True, slots=True)
-class _AnthropicStreamRetryOpened:
+class _StreamRetryOpened:
response: object
attempted_retries: int
max_retries: int
@dataclass(frozen=True, slots=True)
-class _AnthropicStreamRetriesExhausted:
+class _StreamRetriesExhausted:
error: "MidStreamFallbackError"
-_AnthropicStreamRetryOutcome: TypeAlias = _AnthropicStreamRetryOpened | _AnthropicStreamRetriesExhausted
+_StreamRetryOutcome: TypeAlias = _StreamRetryOpened | _StreamRetriesExhausted
MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200
@@ -5688,7 +5691,7 @@ class Router:
if _anthropic_stream_forwards_ping_live(chunk, has_generated_content):
yield chunk
continue
- if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)):
+ if _stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)):
has_generated_content = True
# A transport can split one SSE data line across byte chunks, so pre-content
# detection parses the accumulated buffer plus the current chunk, never the
@@ -5777,16 +5780,20 @@ class Router:
fallback_error: Final = (
stream_error
if isinstance(stream_error, MidStreamFallbackError)
- else _anthropic_stream_fallback_error_for_raised(stream_error, model, has_generated_content)
+ else _stream_fallback_error_for_raised(stream_error, model, has_generated_content)
)
if fallback_error is None:
raise stream_error
- outcome: Final = await self._aanthropic_messages_retry_same_group(fallback_error, initial_kwargs)
+ outcome: Final = await self._astream_retry_same_group(
+ fallback_error,
+ initial_kwargs,
+ attempt_function=self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt,
+ )
match outcome:
- case _AnthropicStreamRetryOpened(response=retried, attempted_retries=attempted, max_retries=budget):
+ case _StreamRetryOpened(response=retried, attempted_retries=attempted, max_retries=budget):
async for item in self._aanthropic_messages_yield_recovered(retried, wrapper, (attempted, budget)):
yield item
- case _AnthropicStreamRetriesExhausted(error=last_error):
+ case _StreamRetriesExhausted(error=last_error):
async for item in self._aanthropic_messages_fallback_attempt(last_error, initial_kwargs, wrapper):
yield item
case _:
@@ -5840,7 +5847,7 @@ class Router:
return None
return _get_num_retries_from_retry_policy(exception=trigger, retry_policy=policy)
- def _anthropic_messages_retry_budget(self, trigger: Exception, kwargs: Mapping[str, object]) -> tuple[int, bool]:
+ def _stream_retry_budget(self, trigger: Exception, kwargs: Mapping[str, object]) -> tuple[int, bool]:
"""
The budget an earlier retry of this request committed to, else the retry policy's grant when one
names this error, else the plain budget, with whether a policy governs the retry: a committed budget
@@ -5890,14 +5897,14 @@ class Router:
if not _is_retriable_anthropic_status(status_code) or _anthropic_stream_error_is_gateway_verdict(chunk):
return None
frame_error: Final = anthropic_error_frame_exception(error_type, message, status_code, model_group)
- budget, _ = self._anthropic_messages_retry_budget(frame_error, kwargs)
+ budget, _ = self._stream_retry_budget(frame_error, kwargs)
if budget > attempted_retries_for_request(kwargs):
return frame_error
if self._anthropic_messages_stream_can_fall_back(model_group, kwargs):
return frame_error
return None
- def _anthropic_messages_should_retry(
+ def _stream_should_retry(
self,
trigger: Exception,
healthy_deployments: list[dict], # mutable-ok: should_retry_this_error's own parameter type
@@ -5929,29 +5936,34 @@ class Router:
return False
return True
- async def _aanthropic_messages_retry_same_group(
- self, e: "MidStreamFallbackError", initial_kwargs: Mapping[str, object]
- ) -> _AnthropicStreamRetryOutcome:
+ async def _astream_retry_same_group(
+ self,
+ e: "MidStreamFallbackError",
+ initial_kwargs: Mapping[str, object],
+ attempt_function: "Callable[..., Awaitable[object]]",
+ ) -> _StreamRetryOutcome:
"""
Re-runs the attempt within the request's own model group, the way async_function_with_retries
would have for a failure raised before the stream opened, until a retry opens a stream or the
budget runs out. Each retry's stream carries its own wrapper with the remaining budget, so a
retry that drops before content again continues the same count instead of starting over.
+ `attempt_function` is the request's own attempt (the /v1/messages attempt, or a chat completions one).
"""
model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group
retry_kwargs: Final = mid_stream_retry_kwargs(initial_kwargs)
healthy_deployments, all_deployments = await self._async_get_healthy_deployments(
model=model_group, parent_otel_span=get_parent_otel_span_from_kwargs(retry_kwargs)
)
- budget, policy_applies = self._anthropic_messages_retry_budget(_mid_stream_retry_trigger(e), initial_kwargs)
+ budget, policy_applies = self._stream_retry_budget(_mid_stream_retry_trigger(e), initial_kwargs)
last_error = e # rebind-ok: the newest failure is what the fallback chain and the caller see
for attempt in range(attempted_retries_for_request(initial_kwargs), budget):
trigger = _mid_stream_retry_trigger(last_error)
- if not policy_applies and not self._anthropic_messages_should_retry(
+ if not policy_applies and not self._stream_should_retry(
trigger, healthy_deployments, all_deployments, initial_kwargs
):
- return _AnthropicStreamRetriesExhausted(last_error)
- self.log_retry(kwargs=retry_kwargs, e=trigger)
+ return _StreamRetriesExhausted(last_error)
+ if get_metadata_variable_name_from_kwargs(retry_kwargs) in retry_kwargs:
+ self.log_retry(kwargs=retry_kwargs, e=trigger)
await asyncio.sleep(
self._time_to_sleep_before_retry(
e=trigger,
@@ -5963,20 +5975,18 @@ class Router:
)
record_retry_attempt(retry_kwargs, attempted_retries=attempt + 1, max_retries=budget)
verbose_router_logger.debug(
- "Retrying anthropic_messages stream dropped before content, attempt %s of %s", attempt + 1, budget
+ "Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget
)
try:
- response = await self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt(**retry_kwargs)
- except Exception as retry_error: # noqa: BLE001 # every failure of a retry before its stream opens is the fallback chain's to judge
- wrapped = _anthropic_stream_fallback_error_for_raised(retry_error, model_group, False)
+ response = await attempt_function(**retry_kwargs)
+ except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry
+ wrapped = _stream_fallback_error_for_raised(retry_error, model_group, False, e.llm_provider)
if wrapped is None:
- return _AnthropicStreamRetriesExhausted(
- _anthropic_stream_pre_content_error(retry_error, model_group)
- )
+ return _StreamRetriesExhausted(_stream_pre_content_error(retry_error, model_group, e.llm_provider))
last_error = wrapped
continue
- return _AnthropicStreamRetryOpened(response, attempted_retries=attempt + 1, max_retries=budget)
- return _AnthropicStreamRetriesExhausted(last_error)
+ return _StreamRetryOpened(response, attempted_retries=attempt + 1, max_retries=budget)
+ return _StreamRetriesExhausted(last_error)
async def _aanthropic_messages_yield_recovered(
self,
@@ -6100,7 +6110,7 @@ class Router:
async def _ageneric_api_call_with_fallbacks_anthropic_messages_attempt(
self,
model: str,
- original_generic_function: Callable,
+ original_generic_function: Callable[..., Awaitable[object]],
**kwargs: object, # kwargs-ok: forwarded verbatim to the per-attempt helper, shape varies per call site
) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]:
"""
diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py
index e7c2e4234f5..af0e24038ac 100644
--- a/tests/code_coverage_tests/router_code_coverage.py
+++ b/tests/code_coverage_tests/router_code_coverage.py
@@ -97,7 +97,7 @@ ignored_function_names = [
"_async_get_available_deployment_for_pass_through", # Same, through async_get_available_deployment_for_pass_through in test_router.py
"_embedding",
"_aembedding",
- "_anthropic_stream_pre_content_error", # Tested through the non-retriable retry error tests in test_router.py
+ "_stream_pre_content_error", # Tested through the non-retriable retry error tests in test_router.py
"_deployment_num_retries", # Tested through the deployment num_retries mid-stream budget test in test_router.py
"_request_fallback_list", # Tested through every mid-stream retry test in test_router.py
"_request_model_group", # Tested through test_anthropic_messages_retry_budget_precedence_direct_call
@@ -105,8 +105,8 @@ ignored_function_names = [
"_anthropic_messages_group_retry_policy", # Tested through the retry budget precedence test in test_router.py
"_anthropic_messages_resolved_retry_policy", # Tested through the malformed retry policy tests in test_router.py
"_anthropic_messages_plain_retry_budget", # Tested through the retry budget precedence test in test_router.py
- "_anthropic_messages_should_retry", # Tested through every mid-stream retry test in test_router.py
- "_aanthropic_messages_retry_same_group", # Tested through the dropped-before-content retry tests in test_router.py
+ "_stream_should_retry", # Tested through every mid-stream retry test in test_router.py
+ "_astream_retry_same_group", # Tested through the dropped-before-content retry tests in test_router.py
"_aanthropic_messages_yield_recovered", # Tested through every mid-stream retry and fallback test in test_router.py
"_anthropic_messages_policy_retries", # Tested through the retry budget precedence test in test_router.py
"_get_wildcard_deployments", # Tested through the get_model_list_of_routed_group wildcard test in test_router.py
diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py
index 4a7a23ff57b..64f3f88cd89 100644
--- a/tests/unit/test_router/test_router.py
+++ b/tests/unit/test_router/test_router.py
@@ -46,11 +46,11 @@ from litellm.router import (
MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS,
MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS,
FallbackAwareAnthropicMessagesStream,
- _anthropic_stream_commits_now,
+ _stream_commits_now,
_anthropic_stream_error_is_gateway_verdict,
- _anthropic_stream_fallback_error_for_raised,
+ _stream_fallback_error_for_raised,
_anthropic_stream_forwards_ping_live,
- _anthropic_stream_raised_error_status,
+ _stream_raised_error_status,
_anthropic_stream_should_decline_fallback,
_is_retriable_anthropic_status,
_responses_stream_holds_event,
@@ -15593,10 +15593,10 @@ def test_anthropic_stream_should_decline_fallback_direct_call():
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 _stream_commits_now(content, has_generated_content=False, buffered_chunk_count=0) is True
+ assert _stream_commits_now(content, has_generated_content=True, buffered_chunk_count=0) is False
assert (
- _anthropic_stream_commits_now(
+ _stream_commits_now(
lifecycle_chunk,
has_generated_content=False,
buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS,
@@ -15604,7 +15604,7 @@ def test_anthropic_stream_commits_now_direct_call():
is True
)
assert (
- _anthropic_stream_commits_now(
+ _stream_commits_now(
lifecycle_chunk,
has_generated_content=False,
buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS - 1,
@@ -16175,15 +16175,15 @@ def test_anthropic_messages_retry_budget_precedence_direct_call():
num_retries=3, retry_policy=RetryPolicy(InternalServerErrorRetries=4)
)
- assert router._anthropic_messages_retry_budget(drop, {"model": "glm"}) == (3, False)
- assert router._anthropic_messages_retry_budget(drop, routed) == (2, False)
- assert router._anthropic_messages_retry_budget(drop, {**routed, "num_retries": 1}) == (1, False)
- assert policy_router._anthropic_messages_retry_budget(reset, {"model": "glm", "num_retries": 1}) == (4, True)
- assert policy_router._anthropic_messages_retry_budget(drop, {"model": "glm", "num_retries": 1}) == (1, False)
- assert policy_router._anthropic_messages_retry_budget(reset, {"model": "glm", "num_retries": 0}) == (0, False)
+ assert router._stream_retry_budget(drop, {"model": "glm"}) == (3, False)
+ assert router._stream_retry_budget(drop, routed) == (2, False)
+ assert router._stream_retry_budget(drop, {**routed, "num_retries": 1}) == (1, False)
+ assert policy_router._stream_retry_budget(reset, {"model": "glm", "num_retries": 1}) == (4, True)
+ assert policy_router._stream_retry_budget(drop, {"model": "glm", "num_retries": 1}) == (1, False)
+ assert policy_router._stream_retry_budget(reset, {"model": "glm", "num_retries": 0}) == (0, False)
committed = {**routed, "litellm_metadata": {**routed["litellm_metadata"], "attempted_retries": 1, "max_retries": 5}}
- assert router._anthropic_messages_retry_budget(drop, committed) == (5, False)
- assert policy_router._anthropic_messages_retry_budget(reset, committed) == (5, True)
+ assert router._stream_retry_budget(drop, committed) == (5, False)
+ assert policy_router._stream_retry_budget(reset, committed) == (5, True)
def test_anthropic_messages_stream_can_retry_direct_call():
@@ -16998,7 +16998,7 @@ async def test_anthropic_messages_raised_provider_error_after_content_propagates
ids=["int", "digit-str", "response-only", "none"],
)
def test_anthropic_stream_raised_error_status_reads_every_status_shape(error, expected_status):
- assert _anthropic_stream_raised_error_status(error) == expected_status
+ assert _stream_raised_error_status(error) == expected_status
@pytest.mark.parametrize(
@@ -17014,7 +17014,7 @@ def test_anthropic_stream_raised_error_status_reads_every_status_shape(error, ex
def test_anthropic_stream_fallback_error_for_raised_gates_like_a_detected_error_event(
error, has_generated_content, converts
):
- converted = _anthropic_stream_fallback_error_for_raised(error, "primary", has_generated_content)
+ converted = _stream_fallback_error_for_raised(error, "primary", has_generated_content)
if not converts:
assert converted is None
return
From 457b5d26c73689cc05617582f76d2901562fa8ed Mon Sep 17 00:00:00 2001
From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 14 Oct 2026 11:32:59 -0700
Subject: [PATCH 21/47] fix(anthropic): stream json_schema answers as content
on Bedrock Mantle and Claude Platform, and join split answers (#46331)
* fix(vertex_ai): unwrap tool-emulated json_schema envelopes and send output_config.format for Claude
Vertex AI Claude deployments that emulate response_format through a forced
tool sometimes nest the schema object under an extra key such as
`parameters` or `input`, and that envelope reached message.content with a
200. The converter now unwraps exactly one nested object that matches the
schema, leaving a payload that is itself schema-shaped alone.
Native structured-output rows sent the deprecated top-level `output_format`,
which Vertex rejects with a 400. The Vertex request now carries it as
`output_config.format`, with an explicit caller value winning.
* fix(anthropic): unwrap tool-emulated json_schema envelopes on streams and tool-name wrappers
* test(anthropic): expect the json tool description through a shared constant
* fix(anthropic): drop the json tool description that makes sonnet-5 omit required keys
* fix(anthropic): keep schema-accepted answers whole and decide held streams at the second key
An answer the json_schema accepts is returned as is before any envelope rule runs, the tool-name
wrapper is unwrapped only when the schema rejects it, a malformed schema turns the unwrap off, a
streamed answer whose first key is a schema property named name or tool_name waits for its second
key before streaming, the held fragments go through a StringIO, and a prefix that stays undecided
past 4 KB is held to the end
* test(anthropic): keep the stream family cells to providers whose stream iterator gets the schema
* fix(vertex_ai): unwrap claude json_schema answers nested under up to three wrapper keys
A tool-emulated json_schema answer was unwrapped only when the real object sat one key down. In live QA, 1 of 58 answers (/v1/responses, claude-sonnet-5) came back as {"parameter": {"value": {...}}} and reached the client still wrapped
The unwrap now follows single-key wrappers up to three levels down and stops at the first object that has a key the schema declares. That object is returned only when it has every required key. A payload wrapped deeper than three levels, a level with more than one key, and a level whose key is one of the schema's own property names all come back exactly as they arrived
The streamed path needs no change because the buffered answer goes through the same unwrap when the tool block closes, and a test now pins that both paths agree
* fix(bedrock): unwrap streamed claude json_schema answers on claude platform and mantle
A streamed json_schema answer on bedrock/claude_platform and Bedrock Mantle Claude reached the client as tool_calls deltas holding the json_tool_call wrapper, with finish_reason tool_calls and no content. The shared HTTP handler built the async stream iterator without json_mode, and neither stream path passed the response_format schema
The handler now asks the provider config for the iterator through get_model_response_iterator_for_request, which carries json_mode and the request's optional params on both the sync and async paths. Claude Platform and Mantle override it to hand the json tool schema to the Anthropic stream iterator, so the streamed answer is the bare schema object with finish_reason stop, the same as the non-streamed one
Databricks is the only other provider on this handler whose iterator reads json_mode, so its async Claude json mode stream now returns the answer as content like its sync stream already did
The stream family integration test covers claude_platform_unmapped and bedrock_mantle again
* fix(anthropic): unwrap json tool envelopes by structure instead of schema validation
The unwrap now reads only the schema's top-level property names. A lone
undeclared key holding an object, alone or beside a name equal to the json
tool's own name, is the envelope, and the unwrap repeats while the inner
object shares no key with the schema. jsonschema and referencing are no
longer imported on the request path. The stream holds only an answer that
opens with the tool's name or with an undeclared key holding an object, and a
schema without top-level properties streams as it arrives.
* fix(llms): keep other providers' async streams as they were
The request-aware iterator hook passed json_mode on async streams to every provider on the shared HTTP handler, which changed Databricks and any other provider whose iterator reads it. The default hook now makes the same two calls the handler made before, so only Claude Platform and Bedrock Mantle change, and the Databricks test that pinned the wider change is removed
* fix(anthropic): walk json tool envelopes iteratively
* test(llms): pin which stream iterators get json_mode from the shared handler
A provider config that does not override the request-aware iterator hook still gets json_mode on a sync stream and still does not get it on an async stream, the same as before the hook existed
* fix(anthropic): keep json tool arguments too deep to decode instead of raising
* fix(anthropic): unwrap json tool envelopes beside undeclared string labels
* fix(anthropic): keep json answers whole on schemas that allow extra keys and stream held text as sent
A schema that allows undeclared keys and requires none now keeps an undeclared wrapper key on both the stream and non-stream paths. When no envelope is found, the held stream text is now sent exactly as the model wrote it. On Vertex /v1/messages, output_format now wins over output_config.format when both are set, which matches the Anthropic adapter.
* fix(anthropic): join a json_schema answer the model splits into one tool block per field
With tool emulation the model sometimes answers with one json_tool_call block per schema field, each naming the field and holding its value. The non-stream path returned only the first block and the stream sent every block's text back to back, so the client got one field or invalid JSON. Both paths now build one schema object from those blocks, and any other set of blocks keeps today's output
* refactor(anthropic): set the streamed finish reason in one expression
* test(anthropic): cover label envelopes, dynamic-key schemas, as-written held streams and output_format precedence
* fix(anthropic): unwrap a streamed values wrapper like the non-stream answer
* fix(anthropic): join streamed split answers on schemas without required keys, keep client tool calls and length on json_schema streams
* fix(anthropic): join streamed split answers that carry text labels on schemas without required keys
* test(bedrock): script the wrapped json_schema stream with the wrapper Bedrock Mantle sends
* test(bedrock): script the wrapped json_schema stream with the wrapper Bedrock Mantle sent raw
* fix(anthropic): join a split json_schema answer whose fields all hold the same string
* fix(anthropic): join a split json_schema answer sent as one partial object per field
* fix(anthropic): join field blocks after an empty streamed answer without a leading comma
---
litellm/llms/anthropic/chat/handler.py | 134 ++++++--
litellm/llms/anthropic/chat/transformation.py | 16 +-
litellm/llms/base_llm/base_utils.py | 69 +++-
litellm/llms/base_llm/chat/transformation.py | 14 +
.../bedrock/chat/mantle/transformation.py | 15 +-
.../bedrock/claude_platform/transformation.py | 18 +-
litellm/llms/custom_httpx/llm_http_handler.py | 12 +-
...vertex_anthropic_structured_output_wire.py | 6 +-
.../chat/test_anthropic_chat_handler.py | 302 ++++++++++++++++--
.../test_anthropic_chat_transformation.py | 273 ++++++++++++++++
...test_bedrock_chat_mantle_transformation.py | 105 +++++-
.../bedrock/test_claude_platform_provider.py | 113 ++++++-
.../custom_httpx/test_llm_http_handler.py | 48 +++
13 files changed, 1072 insertions(+), 53 deletions(-)
diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py
index eb0de9d9b60..ecd11279374 100644
--- a/litellm/llms/anthropic/chat/handler.py
+++ b/litellm/llms/anthropic/chat/handler.py
@@ -25,7 +25,10 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc
from litellm.llms.base_llm.base_utils import (
RESPONSE_FORMAT_ENVELOPE_NAME_KEYS,
ResponseFormatSchema,
+ names_a_response_format_property,
+ partial_response_format_answer,
response_format_envelope_contents,
+ split_response_format_answer,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@@ -85,6 +88,9 @@ _NAME_VALUE_PREFIX: Final = re.compile(
r'^\s*\{\s*"(?:' + _ENVELOPE_NAME_KEY_PATTERN + r')"\s*:\s*(?:"(?:[^"\\]|\\.)*\\?)?$'
)
_STRING_VALUE_THEN_KEY: Final = re.compile(r'^"(?:[^"\\]|\\.)*"\s*,\s*"((?:[^"\\]|\\.)*)"\s*:')
+_LABEL_THEN_KEY: Final = re.compile(r'"((?:[^"\\]|\\.)*)"\s*,\s*"((?:[^"\\]|\\.)*)"\s*:\s*')
+_LEADING_LABELS: Final = re.compile(r'(?:"(?:[^"\\]|\\.)*"\s*,\s*"(?:[^"\\]|\\.)*"\s*:\s*)*')
+_CLOSED_STRING: Final = re.compile(r'"((?:[^"\\]|\\.)*)"')
_UNDECIDED_PREFIX_LIMIT: Final = 4096
@@ -95,14 +101,45 @@ def _key_after_a_string_verdict(string_value_start: str, schema: ResponseFormatS
return "pass_through" if next_key.group(1) in schema.property_names else "hold"
-def _first_value_verdict(after_first_key: str, schema: ResponseFormatSchema) -> _EnvelopeVerdict:
+def _is_a_label(text: str, schema: ResponseFormatSchema, tool_name: str) -> bool:
+ return text in schema.property_names or text == tool_name
+
+
+def _label_verdict(label: str, next_key: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
+ if _is_a_label(label, schema, tool_name):
+ return "hold"
+ return "pass_through" if next_key in schema.property_names else "undecided"
+
+
+def _leading_labels_verdict(string_value_start: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
+ last_value_start: Final = _LEADING_LABELS.sub("", string_value_start, count=1)
+ labelled: Final = string_value_start[: len(string_value_start) - len(last_value_start)]
+ verdicts: Final[tuple[_EnvelopeVerdict, ...]] = tuple(
+ _label_verdict(pair.group(1), pair.group(2), schema, tool_name) for pair in _LABEL_THEN_KEY.finditer(labelled)
+ )
+ decided: Final[_EnvelopeVerdict | None] = next((verdict for verdict in verdicts if verdict != "undecided"), None)
+ if decided is not None:
+ return decided
+ last_label: Final = _CLOSED_STRING.match(last_value_start)
+ if last_label is not None and _is_a_label(last_label.group(1), schema, tool_name):
+ return "hold"
+ return "undecided" if not last_value_start or last_value_start.startswith('"') else "pass_through"
+
+
+def _string_value_verdict(string_value_start: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
+ if schema.required_names:
+ return _key_after_a_string_verdict(string_value_start, schema)
+ return _leading_labels_verdict(string_value_start, schema, tool_name)
+
+
+def _first_value_verdict(after_first_key: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
value_start: Final = after_first_key.lstrip()
if not value_start:
return "undecided"
- if value_start.startswith("{"):
+ if value_start.startswith('"'):
+ return _string_value_verdict(value_start, schema, tool_name)
+ if value_start.startswith("{") and not schema.accepts_a_lone_undeclared_key:
return "hold"
- if value_start.startswith('"') and schema.required_names:
- return _key_after_a_string_verdict(value_start, schema)
return "pass_through"
@@ -116,9 +153,11 @@ def _prefix_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name:
return "hold"
if first_value is None and _NAME_VALUE_PREFIX.match(buffered_json):
return "undecided"
- if first_key.group(1) in schema.property_names or not schema.property_names or schema.accepts_a_lone_undeclared_key:
+ if first_key.group(1) == "values" and "values" not in schema.property_names:
+ return "hold"
+ if first_key.group(1) in schema.property_names or not schema.property_names:
return "pass_through"
- return _first_value_verdict(buffered_json[first_key.end() :], schema)
+ return _first_value_verdict(buffered_json[first_key.end() :], schema, tool_name)
def _held_answer(held: str, schema: ResponseFormatSchema) -> str:
@@ -130,6 +169,13 @@ def _held_answer(held: str, schema: ResponseFormatSchema) -> str:
return held if contents is None else json.dumps(contents)
+def _split_answer_text(blocks: tuple[str, ...], schema: ResponseFormatSchema) -> str:
+ answer: Final = split_response_format_answer(blocks, schema)
+ if answer is not None:
+ return json.dumps(answer)
+ return "".join(_held_answer(block, schema) for block in blocks)
+
+
def _envelope_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
verdict: Final = _prefix_verdict(buffered_json, schema, tool_name)
if verdict == "undecided" and len(buffered_json) > _UNDECIDED_PREFIX_LIMIT:
@@ -655,6 +701,9 @@ class ModelResponseIterator:
self.response_format_schema = response_format_schema
self._response_format_buffer: io.StringIO = io.StringIO()
self._response_format_verdict: _EnvelopeVerdict = "undecided"
+ self._split_answer_blocks: tuple[str, ...] = ()
+ self._withheld_answer_close: str = ""
+ self._streamed_answer_keys: frozenset[str] | None = None
# For handling partial JSON chunks from fragmentation
# See: https://github.com/BerriAI/litellm/issues/17473
@@ -917,7 +966,7 @@ class ModelResponseIterator:
# Track current content block type for filtering deltas
self.current_content_block_type = content_block_start["content_block"]["type"]
if content_block_start["content_block"]["type"] == "text":
- text = content_block_start["content_block"]["text"]
+ text = self._release_answer_close() + content_block_start["content_block"]["text"]
elif (
content_block_start["content_block"]["type"] == "tool_use"
or content_block_start["content_block"]["type"] == "server_tool_use"
@@ -1049,6 +1098,7 @@ class ModelResponseIterator:
pass
elif type_chunk == "message_delta":
finish_reason, usage, container = self._handle_message_delta(chunk)
+ text = self._release_answer_close() + self._drain_split_answer_blocks()
if container:
provider_specific_fields["container"] = container
# Store container_id and re-emit code_interpreter_results
@@ -1155,6 +1205,8 @@ class ModelResponseIterator:
# (per Anthropic's fine-grained streaming pattern)
tool_name: Final = tool_use.get("function", {}).get("name", "")
self.is_response_format_tool = tool_name == RESPONSE_FORMAT_TOOL_NAME
+ if self.is_response_format_tool:
+ self.tool_index -= 1
if self.is_response_format_tool and self.response_format_schema is not None:
self.converted_response_format_tool = True
@@ -1176,25 +1228,69 @@ class ModelResponseIterator:
) -> str:
if tool_use.get("id") is not None:
self._response_format_buffer = io.StringIO()
- self._response_format_verdict = "undecided"
+ held_back: Final = self._split_answer_blocks or self._streamed_answer_keys is not None
+ self._response_format_verdict = "hold" if held_back else "undecided"
return ""
fragment: Final = tool_use.get("function", {}).get("arguments") or ""
- if self._response_format_verdict == "pass_through":
- return fragment
self._response_format_buffer.write(fragment)
+ if self._response_format_verdict == "pass_through":
+ return self._withhold_answer_close(fragment)
if self._response_format_verdict == "hold":
return ""
buffered: Final = self._response_format_buffer.getvalue()
self._response_format_verdict = _envelope_verdict(buffered, schema, RESPONSE_FORMAT_TOOL_NAME)
- return buffered if self._response_format_verdict == "pass_through" else ""
+ return self._withhold_answer_close(buffered) if self._response_format_verdict == "pass_through" else ""
+
+ def _withhold_answer_close(self, text: str) -> str:
+ pending: Final = self._withheld_answer_close + text
+ body: Final = pending.rstrip()
+ cut: Final = len(body) - 1 if body.endswith("}") else len(pending)
+ self._withheld_answer_close = pending[cut:]
+ return pending[:cut]
+
+ def _release_answer_close(self) -> str:
+ close: Final = self._withheld_answer_close
+ self._withheld_answer_close = ""
+ self._streamed_answer_keys = None
+ return close
+
+ def _join_streamed_answer(self, held: str, schema: ResponseFormatSchema) -> str | None:
+ streamed: Final = self._streamed_answer_keys
+ fields: Final = partial_response_format_answer(held, schema)
+ if streamed is None or fields is None or not streamed.isdisjoint(fields):
+ return None
+ self._streamed_answer_keys = streamed | frozenset(fields)
+ separator: Final = ", " if streamed else ""
+ return f"{separator}{json.dumps(fields)[1:-1]}" if fields else ""
def _flush_response_format_buffer(self) -> str:
- if self.response_format_schema is None or self._response_format_verdict == "pass_through":
+ schema: Final = self.response_format_schema
+ if schema is None:
return ""
held: Final = self._response_format_buffer.getvalue()
+ verdict: Final = self._response_format_verdict
self._response_format_buffer = io.StringIO()
self._response_format_verdict = "undecided"
- return _held_answer(held, self.response_format_schema) if held else ""
+ if verdict == "pass_through":
+ answer: Final = partial_response_format_answer(held, schema)
+ close: Final = "" if answer is not None else self._release_answer_close()
+ self._streamed_answer_keys = None if answer is None else frozenset(answer)
+ return close
+ joined: Final = self._join_streamed_answer(held, schema)
+ if joined is not None:
+ return joined
+ released: Final = self._release_answer_close()
+ if self._split_answer_blocks or names_a_response_format_property(held, schema):
+ self._split_answer_blocks = (*self._split_answer_blocks, held)
+ return released
+ return released + (_held_answer(held, schema) if held else "")
+
+ def _drain_split_answer_blocks(self) -> str:
+ blocks: Final = self._split_answer_blocks
+ self._split_answer_blocks = ()
+ if self.response_format_schema is None or not blocks:
+ return ""
+ return _split_answer_text(blocks, self.response_format_schema)
def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]:
"""
@@ -1207,11 +1303,13 @@ class ModelResponseIterator:
Tuple of (finish_reason, usage, container)
"""
message_delta: Final = MessageBlockDelta(**chunk)
- finish_reason = map_finish_reason(finish_reason=message_delta["delta"].get("stop_reason", "stop") or "stop")
- # Override finish_reason to "stop" if we converted response_format tools
- # (matches OpenAI behavior and non-streaming Anthropic implementation)
- if self.converted_response_format_tool:
- finish_reason = "stop"
+ provider_finish_reason: Final = map_finish_reason(
+ finish_reason=message_delta["delta"].get("stop_reason", "stop") or "stop"
+ )
+ called_only_the_response_format_tool: Final = (
+ provider_finish_reason == "tool_calls" and self.converted_response_format_tool and self.tool_index < 0
+ )
+ finish_reason: Final = "stop" if called_only_the_response_format_tool else provider_finish_reason
usage: Final = (
self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None
)
diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py
index 5fd1118e2ba..7abc13d5508 100644
--- a/litellm/llms/anthropic/chat/transformation.py
+++ b/litellm/llms/anthropic/chat/transformation.py
@@ -39,6 +39,7 @@ from litellm.llms.base_llm.base_utils import (
ResponseFormatSchema,
as_json_object,
response_format_schema,
+ split_response_format_answer,
type_to_response_format_param,
unwrap_response_format_envelope,
)
@@ -2268,13 +2269,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
if json_tool.get("function", {}).get("arguments") is None:
return None, tool_calls, None
_message: Final = AnthropicConfig.convert_tool_response_to_message(
- tool_calls=[json_tool], response_format_schema=response_format_schema
+ tool_calls=tool_calls, response_format_schema=response_format_schema
)
return _message, [], None
- first_json: Final = tool_calls[json_indices[0]]
json_msg: Final = AnthropicConfig.convert_tool_response_to_message(
- [first_json], response_format_schema=response_format_schema
+ [tool_calls[i] for i in json_indices], response_format_schema=response_format_schema
)
extra_content: Final[str | None] = json_msg.content if json_msg is not None else None
filtered_tools: Final = [t for i, t in enumerate(tool_calls) if i not in json_indices]
@@ -2895,6 +2895,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
json_mode_content_str: Final[str | None] = tool_calls[0]["function"].get("arguments")
if json_mode_content_str is None:
return None
+ split_answer: Final = (
+ None
+ if response_format_schema is None
+ else split_response_format_answer(
+ tuple(tool_call["function"].get("arguments") or "" for tool_call in tool_calls),
+ response_format_schema,
+ )
+ )
+ if split_answer is not None:
+ return litellm.Message(content=json.dumps(split_answer))
try:
payload: Final[object] = json.loads(json_mode_content_str)
except (json.JSONDecodeError, RecursionError):
diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py
index 4df5635b257..3a02a1d9dc2 100644
--- a/litellm/llms/base_llm/base_utils.py
+++ b/litellm/llms/base_llm/base_utils.py
@@ -6,7 +6,7 @@ import copy
import itertools
import json
from abc import ABC, abstractmethod
-from collections.abc import Iterator, Mapping
+from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final
@@ -289,6 +289,73 @@ def unwrap_response_format_envelope(payload: object, schema: ResponseFormatSchem
return payload if contents is None else contents
+def _field_names(
+ blocks: tuple[Mapping[str, JsonValue], ...], key: str, schema: ResponseFormatSchema
+) -> tuple[str, ...] | None:
+ values: Final = (block[key] for block in blocks)
+ names: Final = tuple(value for value in values if isinstance(value, str) and value in schema.property_names)
+ return names if len(frozenset(names)) == len(blocks) else None
+
+
+def _is_a_label_every_block_repeats(blocks: tuple[Mapping[str, JsonValue], ...], key: str) -> bool:
+ return isinstance(blocks[0][key], str) and all(block[key] == blocks[0][key] for block in blocks)
+
+
+def _value_key(blocks: tuple[Mapping[str, JsonValue], ...], keys: frozenset[str]) -> str | None:
+ if len(keys) == 1:
+ return next(iter(keys))
+ varying: Final = tuple(key for key in keys if not _is_a_label_every_block_repeats(blocks, key))
+ return varying[0] if len(varying) == 1 else None
+
+
+def _split_answer_fields(
+ blocks: tuple[Mapping[str, JsonValue], ...], schema: ResponseFormatSchema
+) -> Mapping[str, JsonValue] | None:
+ keys: Final = frozenset(blocks[0])
+ if not keys.isdisjoint(schema.property_names) or any(frozenset(block) != keys for block in blocks):
+ return None
+ named: Final = tuple((key, names) for key in keys if (names := _field_names(blocks, key, schema)) is not None)
+ if len(named) != 1:
+ return None
+ ((name_key, field_names),) = named
+ value_key: Final = _value_key(blocks, keys - {name_key})
+ if value_key is None:
+ return None
+ return dict(zip(field_names, (block[value_key] for block in blocks), strict=True))
+
+
+def _merged_partial_answers(
+ blocks: tuple[Mapping[str, JsonValue], ...], schema: ResponseFormatSchema
+) -> Mapping[str, JsonValue] | None:
+ keys: Final = frozenset(itertools.chain.from_iterable(blocks))
+ if not keys <= schema.property_names or len(keys) != sum(len(block) for block in blocks):
+ return None
+ return dict(itertools.chain.from_iterable(block.items() for block in blocks))
+
+
+def split_response_format_answer(
+ encoded_blocks: Sequence[str], schema: ResponseFormatSchema
+) -> Mapping[str, JsonValue] | None:
+ decoded: Final = tuple(_json_object_or_encoded_object(encoded) for encoded in encoded_blocks)
+ blocks: Final = tuple(block for block in decoded if block is not None)
+ if len(blocks) < 2 or len(blocks) != len(decoded):
+ return None
+ named_fields: Final = _split_answer_fields(blocks, schema)
+ return named_fields if named_fields is not None else _merged_partial_answers(blocks, schema)
+
+
+def partial_response_format_answer(encoded_block: str, schema: ResponseFormatSchema) -> Mapping[str, JsonValue] | None:
+ block: Final = _json_object_or_encoded_object(encoded_block)
+ return block if block is not None and frozenset(block) <= schema.property_names else None
+
+
+def names_a_response_format_property(encoded_block: str, schema: ResponseFormatSchema) -> bool:
+ block: Final = _json_object_or_encoded_object(encoded_block)
+ if block is None:
+ return False
+ return any(isinstance(value, str) and value in schema.property_names for value in block.values())
+
+
def _dict_to_response_format_helper(response_format: dict, ref_template: str | None = None) -> dict:
if ref_template is not None and response_format.get("type") == "json_schema":
# Deep copy to avoid modifying original
diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py
index 8275829765b..b58fd089470 100644
--- a/litellm/llms/base_llm/chat/transformation.py
+++ b/litellm/llms/base_llm/chat/transformation.py
@@ -417,6 +417,20 @@ class BaseConfig(ABC):
) -> Any:
pass
+ def get_model_response_iterator_for_request(
+ self,
+ streaming_response: "Iterator[str] | AsyncIterator[str] | ModelResponse",
+ sync_stream: bool,
+ json_mode: bool | None,
+ optional_params: Mapping[str, object],
+ ) -> object:
+ if sync_stream:
+ return self.get_model_response_iterator(
+ streaming_response=streaming_response, sync_stream=True, json_mode=json_mode
+ )
+ # FIXME: async streams never got json_mode, passing it changes every provider whose iterator reads it
+ return self.get_model_response_iterator(streaming_response=streaming_response, sync_stream=False)
+
async def get_async_custom_stream_wrapper(
self,
model: str,
diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py
index 583fcb6b230..e48196946cd 100644
--- a/litellm/llms/bedrock/chat/mantle/transformation.py
+++ b/litellm/llms/bedrock/chat/mantle/transformation.py
@@ -7,10 +7,11 @@ The bedrock-mantle endpoint uses the Anthropic Messages API format but is served
at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth.
"""
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator, Mapping
from typing import TYPE_CHECKING, Any, Final
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
+from litellm.llms.anthropic.chat.transformation import JsonModeRequest
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeConfig,
)
@@ -134,10 +135,22 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
sync_stream: bool,
json_mode: bool | None = False,
) -> Any:
+ return self.get_model_response_iterator_for_request(
+ streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, optional_params={}
+ )
+
+ def get_model_response_iterator_for_request(
+ self,
+ streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
+ sync_stream: bool,
+ json_mode: bool | None,
+ optional_params: Mapping[str, object],
+ ) -> object:
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
return ModelResponseIterator(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=json_mode,
+ response_format_schema=JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None,
)
diff --git a/litellm/llms/bedrock/claude_platform/transformation.py b/litellm/llms/bedrock/claude_platform/transformation.py
index 96a6ee4a701..6f8ec6dbd4a 100644
--- a/litellm/llms/bedrock/claude_platform/transformation.py
+++ b/litellm/llms/bedrock/claude_platform/transformation.py
@@ -1,9 +1,11 @@
+from collections.abc import AsyncIterator, Iterator, Mapping
from typing import Any, Final
import litellm
-from litellm.llms.anthropic.chat.transformation import AnthropicConfig
+from litellm.llms.anthropic.chat.transformation import AnthropicConfig, JsonModeRequest
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import ModelResponse
from .common_utils import (
BedrockClaudePlatformMixin,
@@ -91,14 +93,26 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig):
def get_model_response_iterator(
self,
- streaming_response: Any,
+ streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
sync_stream: bool,
json_mode: bool | None = False,
) -> Any:
+ return self.get_model_response_iterator_for_request(
+ streaming_response=streaming_response, sync_stream=sync_stream, json_mode=json_mode, optional_params={}
+ )
+
+ def get_model_response_iterator_for_request(
+ self,
+ streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
+ sync_stream: bool,
+ json_mode: bool | None,
+ optional_params: Mapping[str, object],
+ ) -> object:
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
return ModelResponseIterator(
streaming_response=streaming_response,
sync_stream=sync_stream,
json_mode=bool(json_mode),
+ response_format_schema=JsonModeRequest.model_validate(optional_params).tool_schema() if json_mode else None,
)
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index 3f497c5072f..a8b14b83ef3 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -27,6 +27,7 @@ import httpx
from httpx import USE_CLIENT_DEFAULT
from httpx._types import FileContent
from openai.types.file_deleted import FileDeleted
+from pydantic import TypeAdapter
from typing_extensions import ReadOnly
import litellm
@@ -464,6 +465,7 @@ async def _aiter_bytes_then_close(response: httpx.Response, *, chunk_size: int)
_DECODED_BODY_STALE_HEADERS: Final[frozenset[str]] = frozenset({"content-encoding", "content-length"})
+_OPTIONAL_PARAMS: Final = TypeAdapter(Mapping[str, object])
def _decoded_body_headers(response: httpx.Response) -> httpx.Headers:
@@ -1000,10 +1002,11 @@ class BaseLLMHTTPHandler:
completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode)
else:
- completion_stream = provider_config.get_model_response_iterator(
+ completion_stream = provider_config.get_model_response_iterator_for_request(
streaming_response=response.iter_lines(),
sync_stream=True,
json_mode=json_mode,
+ optional_params=_OPTIONAL_PARAMS.validate_python(optional_params),
)
# LOGGING
@@ -1139,8 +1142,11 @@ class BaseLLMHTTPHandler:
completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode)
else:
- completion_stream = provider_config.get_model_response_iterator(
- streaming_response=response.aiter_lines(), sync_stream=False
+ completion_stream = provider_config.get_model_response_iterator_for_request(
+ streaming_response=response.aiter_lines(),
+ sync_stream=False,
+ json_mode=json_mode,
+ optional_params=_OPTIONAL_PARAMS.validate_python(optional_params),
)
if isinstance(completion_stream, BaseModelResponseIterator):
completion_stream.http_response = response
diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py
index 11543dee1e0..b2e990760b6 100644
--- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py
+++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py
@@ -171,7 +171,6 @@ _FAMILY: Final = (
"claude_platform_unmapped",
"vertex_haiku",
)
-_STREAM_FAMILY: Final = ("anthropic_unmapped", "azure_ai_unmapped", "bedrock_invoke", "vertex_haiku")
_FRAGMENT_SIZE: Final = 20
_TIMING_BYTES: Final = 512 * 1024
_TIMING_FRAGMENT: Final = 4096
@@ -1169,8 +1168,7 @@ def test_vertex_chat_json_schema_without_required_keys_unwraps_only_the_tool_nam
content: Final = _stream_content(response) if stream else _chat_content(response)
assert content == (_payload(marker_hex) if unwrapped else _tool_input(shape, marker_hex)), response.text
if stream:
- expected_frames: Final = 1 if unwrapped else _fragment_count(shape, marker_hex)
- assert len(_content_frames(response)) == expected_frames, response.text
+ assert (len(_content_frames(response)) == 1) is unwrapped, response.text
_one_delivery(wire, marker_hex)
@@ -1322,7 +1320,7 @@ def test_tool_path_family_unwraps_the_envelope(gateway: Gateway, member: str) ->
assert _json_tool_schema(sent)["properties"] == _SCHEMA["properties"], sent.body
-@pytest.mark.parametrize("member", _STREAM_FAMILY)
+@pytest.mark.parametrize("member", _FAMILY)
def test_tool_path_family_stream_unwraps_the_envelope(gateway: Gateway, member: str) -> None:
marker, marker_hex = _marker("wrapped")
with wire_server(_peer) as wire, gateway.scenario() as scenario:
diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
index ad5c8301f14..7d1d2a3c434 100644
--- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
+++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
@@ -2810,9 +2810,7 @@ def test_streamed_json_mode_schema_object_streams_as_it_arrives():
contents = _streamed_json_mode_contents(fragments)
assert "".join(contents) == answer
- assert contents[3] == answer[:12]
- assert contents[4] == answer[12:]
- assert contents[5] == ""
+ assert contents[3:] == (answer[:12], answer[12:-1], "", "}")
def test_streamed_json_mode_truncated_output_is_passed_through_untouched():
@@ -2836,9 +2834,7 @@ def test_streamed_json_mode_name_property_answer_streams_once_its_name_value_is_
contents = _streamed_json_mode_contents((answer[:6], answer[6:24], answer[24:]), _STREAMED_CONTACT_SCHEMA)
assert "".join(contents) == answer
- assert contents[2] == ""
- assert contents[3] == answer[:24]
- assert contents[4] == answer[24:]
+ assert contents[2:] == ("", answer[:24], answer[24:-1], "", "}")
def test_streamed_json_mode_partial_name_value_waits_for_the_value_to_close():
@@ -2847,8 +2843,7 @@ def test_streamed_json_mode_partial_name_value_waits_for_the_value_to_close():
contents = _streamed_json_mode_contents((answer[:11], answer[11:]), _STREAMED_CONTACT_SCHEMA)
assert "".join(contents) == answer
- assert contents[2] == ""
- assert contents[3] == answer
+ assert contents[2:] == ("", answer[:-1], "", "}")
@pytest.mark.parametrize(
@@ -2865,12 +2860,11 @@ def test_streamed_json_mode_partial_name_value_waits_for_the_value_to_close():
id="undeclared_first_key_holding_a_number",
),
pytest.param(
- {"type": "object", "properties": {"score": {"type": "integer"}}},
- {"note": "read-only window", "meta": {"score": 7}},
- id="string_first_key_on_a_schema_without_required_keys",
- ),
- pytest.param(
- {"type": "object", "properties": {"summary": {"type": "string"}}, "additionalProperties": {"type": "object"}},
+ {
+ "type": "object",
+ "properties": {"summary": {"type": "string"}},
+ "additionalProperties": {"type": "object"},
+ },
{"user_1": {"summary": "likes tea", "visits": 3}},
id="undeclared_key_the_schema_allows",
),
@@ -2884,7 +2878,61 @@ def test_streamed_json_mode_answer_that_cannot_be_an_envelope_streams_as_it_arri
contents = _streamed_json_mode_contents((encoded[:12], encoded[12:24], encoded[24:]), schema)
assert "".join(contents) == encoded
- assert contents[2:5] == (encoded[:12], encoded[12:24], encoded[24:])
+ assert contents[2:] == (encoded[:12], encoded[12:24], encoded[24:-1], "}", "")
+
+
+def test_streamed_json_mode_text_label_then_an_object_streams_once_the_object_opens():
+ encoded = json.dumps({"note": "read-only window", "meta": {"score": 7}})
+ object_opens = encoded.index('{"score"') + 1
+
+ contents = _streamed_json_mode_contents(
+ (encoded[:12], encoded[12:object_opens], encoded[object_opens:]),
+ {"type": "object", "properties": {"score": {"type": "integer"}}},
+ )
+
+ assert contents[2:] == ("", encoded[:object_opens], encoded[object_opens:-1], "}", "")
+
+
+@pytest.mark.parametrize(
+ "schema, values",
+ [
+ pytest.param(
+ {"type": "object", "additionalProperties": True}, {"alpha": 1, "beta": 2}, id="schema_without_properties"
+ ),
+ pytest.param(
+ {
+ "type": "object",
+ "properties": {"summary": {"type": "string"}},
+ "additionalProperties": {"type": "object"},
+ },
+ {"user_1": {"summary": "likes tea"}},
+ id="undeclared_key_the_schema_allows",
+ ),
+ pytest.param(_STREAMED_RISK_SCHEMA, [3, 5, 8], id="values_holding_a_list"),
+ ],
+)
+def test_streamed_json_mode_values_wrapper_is_unwrapped_like_the_non_stream_answer(
+ schema: dict[str, object], values: object
+):
+ wrapped = {"values": values}
+ encoded = json.dumps(wrapped)
+
+ contents = _streamed_json_mode_contents((encoded[:6], encoded[6:14], encoded[14:]), schema)
+
+ assert json.loads("".join(contents)) == values
+ assert json.loads("".join(contents)) == unwrap_response_format_envelope(
+ wrapped, response_format_schema(schema), RESPONSE_FORMAT_TOOL_NAME
+ )
+
+
+def test_streamed_json_mode_declared_values_key_streams_as_it_arrives():
+ schema = {"type": "object", "properties": {"values": {"type": "array"}}, "required": ["values"]}
+ answer = json.dumps({"values": [3, 5, 8]})
+
+ contents = _streamed_json_mode_contents((answer[:6], answer[6:14], answer[14:]), schema)
+
+ assert "".join(contents) == answer
+ assert contents[3:] == (answer[:14], answer[14:-1], "", "}")
def test_streamed_json_mode_string_label_answer_streams_once_the_next_key_is_declared():
@@ -2893,9 +2941,7 @@ def test_streamed_json_mode_string_label_answer_streams_once_the_next_key_is_dec
contents = _streamed_json_mode_contents((answer[:12], answer[12:40], answer[40:]))
assert "".join(contents) == answer
- assert contents[2] == ""
- assert contents[3] == answer[:40]
- assert contents[4] == answer[40:]
+ assert contents[2:] == ("", answer[:40], answer[40:-1], "}", "")
@pytest.mark.parametrize(
@@ -2978,6 +3024,226 @@ def test_streamed_json_mode_answer_with_a_foreign_name_matches_the_non_streamed_
assert unwrap_response_format_envelope(answer, schema, RESPONSE_FORMAT_TOOL_NAME) == answer
+def _json_tool_block_events(index: int, encoded: str) -> tuple[dict[str, object], ...]:
+ return (
+ {
+ "type": "content_block_start",
+ "index": index,
+ "content_block": {
+ "type": "tool_use",
+ "id": f"toolu_{index}",
+ "name": RESPONSE_FORMAT_TOOL_NAME,
+ "input": {},
+ },
+ },
+ *(
+ {"type": "content_block_delta", "index": index, "delta": {"type": "input_json_delta", "partial_json": part}}
+ for part in (encoded[:7], encoded[7:])
+ ),
+ {"type": "content_block_stop", "index": index},
+ )
+
+
+def _streamed_json_tool_blocks(
+ blocks: tuple[str, ...], schema: dict[str, object] = _STREAMED_RISK_SCHEMA
+) -> tuple[str, ...]:
+ message_start, *_, message_delta = _json_tool_events(())
+ events = (
+ message_start,
+ *chain.from_iterable(_json_tool_block_events(index, encoded) for index, encoded in enumerate(blocks)),
+ message_delta,
+ )
+ iterator = ModelResponseIterator(
+ streaming_response=MagicMock(),
+ sync_stream=True,
+ json_mode=True,
+ response_format_schema=response_format_schema(schema),
+ )
+ chunks = tuple(iterator.chunk_parser(event) for event in events)
+ assert all(chunk.choices[0].delta.tool_calls is None for chunk in chunks)
+ assert chunks[-1].choices[0].finish_reason == "stop"
+ return tuple(chunk.choices[0].delta.content or "" for chunk in chunks)
+
+
+_STREAMED_TIER_FIELD: Final = json.dumps(
+ {"$FUNCTION_NAME": "RiskAssessment", "$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": _STREAMED_RISK_OBJECT["tier"]}
+)
+_STREAMED_SUMMARY_FIELD: Final = json.dumps(
+ {
+ "$FUNCTION_NAME": "RiskAssessment",
+ "$PARAMETER_NAME": "summary",
+ "$PARAMETER_VALUE": _STREAMED_RISK_OBJECT["summary"],
+ }
+)
+
+
+def test_streamed_json_mode_answer_split_into_one_block_per_field_arrives_as_one_object():
+ contents = _streamed_json_tool_blocks((_STREAMED_TIER_FIELD, _STREAMED_SUMMARY_FIELD))
+
+ assert json.loads(contents[-1]) == _STREAMED_RISK_OBJECT
+ assert contents[:-1] == ("",) * (len(contents) - 1)
+
+
+def test_streamed_json_mode_field_block_alone_keeps_the_model_text():
+ contents = _streamed_json_tool_blocks((_STREAMED_TIER_FIELD,))
+
+ assert "".join(contents) == _STREAMED_TIER_FIELD
+
+
+def test_streamed_json_mode_whole_answer_after_a_field_block_keeps_the_order_the_model_sent():
+ whole_answer = json.dumps(_STREAMED_RISK_OBJECT)
+
+ contents = _streamed_json_tool_blocks((_STREAMED_TIER_FIELD, whole_answer))
+
+ assert "".join(contents) == _STREAMED_TIER_FIELD + whole_answer
+
+
+@pytest.mark.parametrize(
+ "schema, answer",
+ [
+ pytest.param(_STREAMED_RISK_SCHEMA, _STREAMED_RISK_OBJECT, id="schema_with_a_required_list"),
+ pytest.param(
+ _STREAMED_CONTACT_SCHEMA, {"name": "unknown", "email": "unknown"}, id="every_field_holds_the_same_string"
+ ),
+ ],
+)
+def test_streamed_json_mode_answer_split_into_one_partial_object_per_field_arrives_as_one_object(
+ schema: dict[str, object], answer: dict[str, object]
+):
+ blocks = tuple(json.dumps({name: value}) for name, value in answer.items())
+
+ contents = _streamed_json_tool_blocks(blocks, schema)
+
+ assert json.loads("".join(contents)) == answer
+ assert contents[3] == blocks[0][:-1]
+
+
+def test_streamed_json_mode_empty_object_then_field_blocks_arrives_as_one_object():
+ blocks = ("{}", *(json.dumps({name: value}) for name, value in _STREAMED_RISK_OBJECT.items()))
+
+ contents = _streamed_json_tool_blocks(blocks)
+
+ assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT
+
+
+def test_streamed_json_mode_partial_objects_that_repeat_a_field_keep_the_order_the_model_sent():
+ blocks = (json.dumps({"tier": 3}), json.dumps({"tier": 4}))
+
+ contents = _streamed_json_tool_blocks(blocks)
+
+ assert "".join(contents) == "".join(blocks)
+
+
+def test_streamed_json_mode_text_after_the_answer_follows_the_whole_answer():
+ answer = json.dumps(_STREAMED_RISK_OBJECT)
+ message_start, *_, message_delta = _json_tool_events(())
+ text_block = (
+ {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}},
+ {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": " Done."}},
+ {"type": "content_block_stop", "index": 1},
+ )
+
+ chunks = _json_mode_stream_chunks((message_start, *_json_tool_block_events(0, answer), *text_block, message_delta))
+
+ assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == answer + " Done."
+
+
+def test_streamed_json_mode_envelope_whose_label_is_a_property_name_is_still_unwrapped():
+ envelope = json.dumps({"$FUNCTION_NAME": "summary", "PARAMETER_VALUE": _STREAMED_RISK_OBJECT})
+
+ contents = _streamed_json_tool_blocks((envelope,))
+
+ assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT
+
+
+_STREAMED_OPTIONAL_RISK_SCHEMA: Final = {
+ "type": "object",
+ "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}},
+}
+
+
+@pytest.mark.parametrize(
+ "schema",
+ [
+ pytest.param(_STREAMED_OPTIONAL_RISK_SCHEMA, id="no_required_keys"),
+ pytest.param(
+ {**_STREAMED_OPTIONAL_RISK_SCHEMA, "additionalProperties": True}, id="no_required_keys_and_open_to_others"
+ ),
+ ],
+)
+@pytest.mark.parametrize(
+ "leading_label",
+ [
+ pytest.param({}, id="field_name_label_first"),
+ pytest.param({"$FUNCTION_NAME": RESPONSE_FORMAT_TOOL_NAME}, id="tool_name_label_first"),
+ pytest.param({"$FUNCTION_NAME": "RiskAssessment"}, id="schema_name_label_first"),
+ pytest.param({"$FUNCTION_NAME": "RiskAssessment", "$CALL": "first"}, id="two_text_labels_first"),
+ ],
+)
+def test_streamed_json_mode_split_answer_arrives_as_one_object_when_the_schema_requires_no_key(
+ schema: dict[str, object], leading_label: dict[str, str]
+):
+ blocks = tuple(
+ json.dumps({**leading_label, "$PARAMETER_NAME": name, "$PARAMETER_VALUE": value})
+ for name, value in _STREAMED_RISK_OBJECT.items()
+ )
+
+ contents = _streamed_json_tool_blocks(blocks, schema)
+
+ assert json.loads(contents[-1]) == _STREAMED_RISK_OBJECT
+ assert contents[:-1] == ("",) * (len(contents) - 1)
+
+
+def _json_mode_stream_chunks(events: tuple[dict[str, object], ...]) -> tuple[ModelResponseStream, ...]:
+ iterator = ModelResponseIterator(
+ streaming_response=MagicMock(),
+ sync_stream=True,
+ json_mode=True,
+ response_format_schema=response_format_schema(_STREAMED_RISK_SCHEMA),
+ )
+ return tuple(iterator.chunk_parser(event) for event in events)
+
+
+def test_streamed_json_mode_answer_beside_a_client_tool_call_keeps_that_call_usable():
+ answer = json.dumps(_STREAMED_RISK_OBJECT)
+ message_start, *_, message_delta = _json_tool_events(())
+ weather_call = (
+ {
+ "type": "content_block_start",
+ "index": 1,
+ "content_block": {"type": "tool_use", "id": "toolu_weather", "name": "get_weather", "input": {}},
+ },
+ {
+ "type": "content_block_delta",
+ "index": 1,
+ "delta": {"type": "input_json_delta", "partial_json": '{"city": "Paris"}'},
+ },
+ {"type": "content_block_stop", "index": 1},
+ )
+
+ chunks = _json_mode_stream_chunks(
+ (message_start, *_json_tool_block_events(0, answer), *weather_call, message_delta)
+ )
+
+ tool_calls = tuple(chain.from_iterable(chunk.choices[0].delta.tool_calls or () for chunk in chunks))
+ assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == answer
+ assert {call.index for call in tool_calls} == {0}
+ assert tool_calls[0].function.name == "get_weather"
+ assert "".join(call.function.arguments or "" for call in tool_calls) == '{"city": "Paris"}'
+ assert chunks[-1].choices[0].finish_reason == "tool_calls"
+
+
+def test_streamed_json_mode_answer_cut_at_max_tokens_ends_with_length():
+ cut_answer = '{"tier": 3, "summ'
+ *events, message_delta = _json_tool_events((cut_answer,))
+ cut_message_delta = {**message_delta, "delta": {"stop_reason": "max_tokens", "stop_sequence": None}}
+
+ chunks = _json_mode_stream_chunks((*events, cut_message_delta))
+
+ assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == cut_answer
+ assert chunks[-1].choices[0].finish_reason == "length"
+
+
@respx.mock
def test_anthropic_streaming_function_call_arguments_are_assembled(
monkeypatch: pytest.MonkeyPatch,
diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
index 381bcca1b54..bbcc8629a10 100644
--- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
+++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
@@ -1,5 +1,6 @@
import copy
import json
+from itertools import chain
from typing import Final
from unittest.mock import MagicMock, patch
@@ -7767,6 +7768,278 @@ def test_vertex_claude_tool_emulated_response_format_returns_the_object_the_mode
assert result.choices[0].finish_reason == "stop"
+_SPLIT_RISK_SCHEMA: Final = {
+ "title": "RiskAssessment",
+ "type": "object",
+ "properties": {
+ "tier": {"type": "integer"},
+ "risk_level": {"type": "string", "enum": ["low", "medium", "high"]},
+ "score": {"type": "integer"},
+ "summary": {"type": "string"},
+ "affected_systems": {"type": "array", "items": {"type": "string"}},
+ "findings": {"type": "array", "items": {"type": "object"}},
+ "rollback_plan": {"type": "string"},
+ "requires_downtime": {"type": "boolean"},
+ },
+ "required": [
+ "tier",
+ "risk_level",
+ "score",
+ "summary",
+ "affected_systems",
+ "findings",
+ "rollback_plan",
+ "requires_downtime",
+ ],
+}
+_SPLIT_RISK_OBJECT: Final = {
+ "tier": 1,
+ "risk_level": "high",
+ "score": 78,
+ "summary": "Migrating loyalty-points ledger DB to managed cloud with weekend read-only cutover; financial data, high risk.",
+ "affected_systems": ["Loyalty-points ledger DB", "Loyalty rewards API", "Reporting/analytics pipeline"],
+ "findings": [
+ {
+ "title": "Financial data integrity risk",
+ "severity": "high",
+ "detail": "Loyalty points are monetary-equivalent; any sync error causes customer-facing balance discrepancies.",
+ },
+ {
+ "title": "Weekend support coverage",
+ "severity": "low",
+ "detail": "Reduced staff availability over weekend may delay incident response if issues arise.",
+ },
+ ],
+ "rollback_plan": "Keep old cluster live/read-only; revert DNS/connection strings to original cluster if migration fails.",
+ "requires_downtime": True,
+}
+# Block shape captured from vertex_ai/claude-sonnet-5 tool emulation on 2026-10-13, one tool_use block per field
+_SPLIT_RISK_BLOCKS: Final = tuple(
+ {"$FUNCTION_NAME": "RiskAssessment", "$PARAMETER_NAME": name, "$PARAMETER_VALUE": value}
+ for name, value in _SPLIT_RISK_OBJECT.items()
+)
+_CONTACT_SCHEMA: Final = {
+ "type": "object",
+ "properties": {"name": {"type": "string"}, "email": {"type": "string"}},
+ "required": ["name", "email"],
+}
+_OPTIONAL_CONTACT_SCHEMA: Final = {
+ "type": "object",
+ "properties": {"name": {"type": "string"}, "email": {"type": "string"}, "phone": {"type": "string"}},
+}
+_FLAGS_SCHEMA: Final = {
+ "type": "object",
+ "properties": {"urgent": {"type": "boolean"}, "reviewed": {"type": "boolean"}},
+ "required": ["urgent", "reviewed"],
+}
+
+
+def _json_tool_calls(*blocks: object) -> list[ChatCompletionToolCallChunk]:
+ return list(chain.from_iterable(_json_tool_call(block) for block in blocks))
+
+
+def test_vertex_claude_json_answer_split_into_one_tool_block_per_field_returns_one_schema_object():
+ config = VertexAIAnthropicConfig()
+ optional_params = config.map_openai_params(
+ non_default_params={
+ "response_format": {
+ "type": "json_schema",
+ "json_schema": {"name": "RiskAssessment", "schema": _SPLIT_RISK_SCHEMA},
+ },
+ },
+ optional_params={},
+ model="claude-emulated-structured-output",
+ drop_params=False,
+ )
+ json_mode = optional_params.pop("json_mode")
+ completion_response = {
+ "id": "msg_1",
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-emulated-structured-output",
+ "content": [
+ {"type": "tool_use", "id": f"toolu_{index}", "name": RESPONSE_FORMAT_TOOL_NAME, "input": block}
+ for index, block in enumerate(_SPLIT_RISK_BLOCKS)
+ ],
+ "stop_reason": "tool_use",
+ "stop_sequence": None,
+ "usage": {"input_tokens": 12, "output_tokens": 30},
+ }
+
+ result = config.transform_response(
+ model="claude-emulated-structured-output",
+ raw_response=httpx.Response(200, json=completion_response),
+ model_response=ModelResponse(),
+ logging_obj=MagicMock(),
+ request_data={},
+ messages=[{"role": "user", "content": "Assess this change request"}],
+ optional_params=optional_params,
+ litellm_params={},
+ encoding=None,
+ json_mode=json_mode,
+ )
+
+ assert json.loads(result.choices[0].message.content) == _SPLIT_RISK_OBJECT
+ assert result.choices[0].message.tool_calls is None
+ assert result.choices[0].finish_reason == "stop"
+
+
+@pytest.mark.parametrize(
+ "schema, blocks, expected_content",
+ [
+ pytest.param(
+ _RISK_SCHEMA,
+ (
+ {"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": _RISK_OBJECT["tier"]},
+ {"$PARAMETER_NAME": "summary", "$PARAMETER_VALUE": _RISK_OBJECT["summary"]},
+ ),
+ _RISK_OBJECT,
+ id="name_and_value_keys_only",
+ ),
+ pytest.param(
+ _CONTACT_SCHEMA,
+ (
+ {"function": "Contact", "parameter_name": "name", "parameter_value": "Ada"},
+ {"function": "Contact", "parameter_name": "email", "parameter_value": "ada@example.com"},
+ ),
+ {"name": "Ada", "email": "ada@example.com"},
+ id="string_values_beside_a_repeated_label",
+ ),
+ pytest.param(
+ _FLAGS_SCHEMA,
+ ({"field": "urgent", "value": True}, {"field": "reviewed", "value": True}),
+ {"urgent": True, "reviewed": True},
+ id="every_field_holds_the_same_value",
+ ),
+ pytest.param(
+ _CONTACT_SCHEMA,
+ ({"field": "name", "value": "unknown"}, {"field": "email", "value": "unknown"}),
+ {"name": "unknown", "email": "unknown"},
+ id="every_field_holds_the_same_string",
+ ),
+ pytest.param(
+ _CONTACT_SCHEMA,
+ ({"name": "unknown"}, {"email": "unknown"}),
+ {"name": "unknown", "email": "unknown"},
+ id="one_partial_object_per_field",
+ ),
+ pytest.param(
+ _OPTIONAL_CONTACT_SCHEMA,
+ ({"name": "Ada", "email": "ada@example.com"}, {"phone": "555-0100"}),
+ {"name": "Ada", "email": "ada@example.com", "phone": "555-0100"},
+ id="partial_objects_of_uneven_size_on_a_schema_with_no_required_list",
+ ),
+ ],
+)
+def test_json_tool_blocks_that_each_name_one_schema_property_become_one_object(
+ schema: dict[str, object], blocks: tuple[object, ...], expected_content: dict[str, object]
+):
+ replacement, remaining_tool_calls, extra_content = AnthropicConfig()._resolve_json_mode_non_streaming(
+ json_mode=True,
+ tool_calls=_json_tool_calls(*blocks),
+ response_format_schema=_schema(schema),
+ )
+
+ assert replacement is not None
+ assert json.loads(replacement.content) == expected_content
+ assert remaining_tool_calls == []
+ assert extra_content is None
+
+
+@pytest.mark.parametrize(
+ "schema, blocks",
+ [
+ pytest.param(_RISK_SCHEMA, ({"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3},), id="one_field_alone"),
+ pytest.param(
+ _RISK_SCHEMA,
+ ({"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3}, {"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 4}),
+ id="same_property_named_twice",
+ ),
+ pytest.param(
+ _RISK_SCHEMA,
+ (
+ {"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3},
+ {"$PARAMETER_NAME": "color", "$PARAMETER_VALUE": "red"},
+ ),
+ id="property_the_schema_does_not_declare",
+ ),
+ pytest.param(
+ _RISK_SCHEMA,
+ (
+ {"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3},
+ {"parameter_name": "summary", "parameter_value": "Read-only window"},
+ ),
+ id="blocks_with_different_keys",
+ ),
+ pytest.param(
+ _RISK_SCHEMA,
+ (
+ {"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3, "note": "first"},
+ {"$PARAMETER_NAME": "summary", "$PARAMETER_VALUE": "Read-only window", "note": "second"},
+ ),
+ id="two_keys_that_could_hold_the_value",
+ ),
+ pytest.param(
+ _RISK_SCHEMA,
+ ({"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": 3}, ["summary", "Read-only window"]),
+ id="block_that_is_not_an_object",
+ ),
+ pytest.param(
+ _PERSON_SCHEMA,
+ ({"name": "name", "age": 1}, {"name": "age", "age": 2}),
+ id="blocks_keyed_by_schema_properties",
+ ),
+ pytest.param(
+ _CONTACT_SCHEMA,
+ ({"name": "Ada"}, {"name": "Grace"}),
+ id="partial_objects_that_repeat_a_field",
+ ),
+ pytest.param(
+ _CONTACT_SCHEMA,
+ ({"name": "Ada"}, {"nickname": "A"}),
+ id="partial_object_with_a_key_the_schema_does_not_declare",
+ ),
+ ],
+)
+def test_json_tool_blocks_that_are_not_one_field_each_keep_the_first_block(
+ schema: dict[str, object], blocks: tuple[object, ...]
+):
+ replacement, remaining_tool_calls, extra_content = AnthropicConfig()._resolve_json_mode_non_streaming(
+ json_mode=True,
+ tool_calls=_json_tool_calls(*blocks),
+ response_format_schema=_schema(schema),
+ )
+
+ assert replacement is not None
+ assert json.loads(replacement.content) == blocks[0]
+ assert remaining_tool_calls == []
+ assert extra_content is None
+
+
+def test_json_mode_with_user_tool_call_joins_the_split_answer_into_content():
+ user_tool_call: Final = ChatCompletionToolCallChunk(
+ id="toolu_user",
+ type="function",
+ function=ChatCompletionToolCallFunctionChunk(name="lookup_change_window", arguments='{"system": "ledger"}'),
+ index=2,
+ )
+ first_field, second_field = _json_tool_calls(
+ {"$PARAMETER_NAME": "tier", "$PARAMETER_VALUE": _RISK_OBJECT["tier"]},
+ {"$PARAMETER_NAME": "summary", "$PARAMETER_VALUE": _RISK_OBJECT["summary"]},
+ )
+
+ replacement, remaining_tool_calls, extra_content = AnthropicConfig()._resolve_json_mode_non_streaming(
+ json_mode=True,
+ tool_calls=[first_field, user_tool_call, second_field],
+ response_format_schema=_schema(_RISK_SCHEMA),
+ )
+
+ assert replacement is None
+ assert remaining_tool_calls == [user_tool_call]
+ assert extra_content is not None
+ assert json.loads(extra_content) == _RISK_OBJECT
+
+
def test_anthropic_tool_with_image():
import json
diff --git a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py
index cb892b1ea11..154be9f0370 100644
--- a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py
+++ b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py
@@ -7,7 +7,7 @@ import httpx
import pytest
import litellm
-from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
ONE_PIXEL_PNG = base64.b64decode(
@@ -140,3 +140,106 @@ async def test_bedrock_mantle_claude_async_completion_inlines_document_url_sourc
"type": "document",
"source": {"type": "base64", "media_type": "application/pdf", "data": async_only_image_fetch.base64_png},
} in captured["body"]["messages"][0]["content"]
+
+
+RISK_ANSWER = {"tier": 2, "summary": "ok"}
+MANTLE_CLAUDE_ROUTES = ("bedrock_mantle/anthropic.claude-scripted-9", "bedrock/mantle/anthropic.claude-scripted-9")
+
+
+def _json_schema_stream_kwargs(model):
+ return {
+ "model": model,
+ "messages": [{"role": "user", "content": "Assess the risk"}],
+ "stream": True,
+ "response_format": {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "Risk",
+ "schema": {
+ "type": "object",
+ "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}},
+ "required": ["tier", "summary"],
+ "additionalProperties": False,
+ },
+ },
+ },
+ "aws_access_key_id": "AKIAEXAMPLE",
+ "aws_secret_access_key": "fake-secret",
+ "aws_region_name": "us-east-1",
+ }
+
+
+def _wrapped_json_tool_stream(request):
+ wrapper = json.dumps({"$PARAMETER_NAME": RISK_ANSWER})
+ events = (
+ {
+ "type": "message_start",
+ "message": {
+ "id": "msg_1",
+ "type": "message",
+ "role": "assistant",
+ "model": "anthropic.claude-scripted-9",
+ "content": [],
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": {"input_tokens": 1, "output_tokens": 1},
+ },
+ },
+ {
+ "type": "content_block_start",
+ "index": 0,
+ "content_block": {"type": "tool_use", "id": "toolu_1", "name": "json_tool_call", "input": {}},
+ },
+ *(
+ {
+ "type": "content_block_delta",
+ "index": 0,
+ "delta": {"type": "input_json_delta", "partial_json": wrapper[start : start + 20]},
+ }
+ for start in range(0, len(wrapper), 20)
+ ),
+ {"type": "content_block_stop", "index": 0},
+ {
+ "type": "message_delta",
+ "delta": {"stop_reason": "tool_use", "stop_sequence": None},
+ "usage": {"output_tokens": 9},
+ },
+ {"type": "message_stop"},
+ )
+ return httpx.Response(
+ 200,
+ headers={"content-type": "text/event-stream"},
+ content="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode(),
+ )
+
+
+def _streamed_answer(chunks):
+ deltas = tuple(chunk.choices[0].delta for chunk in chunks)
+ return {
+ "content": json.loads("".join(delta.content or "" for delta in deltas)),
+ "tool_call_deltas": sum(1 for delta in deltas if delta.tool_calls),
+ "finish_reasons": tuple(chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason),
+ }
+
+
+BARE_STREAMED_ANSWER = {"content": RISK_ANSWER, "tool_call_deltas": 0, "finish_reasons": ("stop",)}
+
+
+@pytest.mark.parametrize("model", MANTLE_CLAUDE_ROUTES)
+def test_mantle_claude_completion_streams_a_wrapped_json_schema_answer_as_the_bare_object(model):
+ client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_wrapped_json_tool_stream)))
+
+ chunks = tuple(litellm.completion(**_json_schema_stream_kwargs(model), client=client))
+
+ assert _streamed_answer(chunks) == BARE_STREAMED_ANSWER
+
+
+@pytest.mark.parametrize("model", MANTLE_CLAUDE_ROUTES)
+async def test_mantle_claude_async_completion_streams_a_wrapped_json_schema_answer_as_the_bare_object(model):
+ client = AsyncHTTPHandler()
+ client.client = httpx.AsyncClient(transport=httpx.MockTransport(_wrapped_json_tool_stream))
+
+ stream = await litellm.acompletion(**_json_schema_stream_kwargs(model), client=client)
+ chunks = tuple([chunk async for chunk in stream])
+
+ assert _streamed_answer(chunks) == BARE_STREAMED_ANSWER
diff --git a/tests/unit/llms/bedrock/test_claude_platform_provider.py b/tests/unit/llms/bedrock/test_claude_platform_provider.py
index 7b42d1f5a35..5d83cba5415 100644
--- a/tests/unit/llms/bedrock/test_claude_platform_provider.py
+++ b/tests/unit/llms/bedrock/test_claude_platform_provider.py
@@ -103,9 +103,12 @@ def _sync_gateway_client(
return HTTPHandler(client=httpx.Client(transport=_gateway_transport(requests, respond)))
-def _async_gateway_client(requests: list[dict]) -> AsyncHTTPHandler:
+def _async_gateway_client(
+ requests: list[dict],
+ respond: GatewayResponder = _fake_claude_platform_gateway,
+) -> AsyncHTTPHandler:
handler: Final = AsyncHTTPHandler()
- handler.client = httpx.AsyncClient(transport=_gateway_transport(requests, _fake_claude_platform_gateway))
+ handler.client = httpx.AsyncClient(transport=_gateway_transport(requests, respond))
return handler
@@ -718,6 +721,112 @@ def test_chat_completion_claude_platform_override_kwarg_is_honoured_and_not_sent
assert requests[0]["body"] == {**EXPECTED_CHAT_BODY, "context_management": context_management}
+_RISK_ANSWER: Final = {"tier": 2, "summary": "ok"}
+_JSON_SCHEMA_STREAM_KWARGS: Final = {
+ "model": "bedrock/claude_platform/claude-scripted-9",
+ "messages": [{"role": "user", "content": "hello"}],
+ "max_tokens": 10,
+ "workspace_id": "wrkspc_test",
+ "stream": True,
+ "response_format": {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "Risk",
+ "schema": {
+ "type": "object",
+ "properties": {"tier": {"type": "integer"}, "summary": {"type": "string"}},
+ "required": ["tier", "summary"],
+ "additionalProperties": False,
+ },
+ },
+ },
+ **API_KEY_KWARGS,
+}
+
+
+def _wrapped_json_tool_stream(url: str, headers: Mapping[str, str], data: bytes) -> httpx.Response:
+ wrapper: Final = json.dumps({"$PARAMETER_NAME": _RISK_ANSWER})
+ events: Final = (
+ {
+ "type": "message_start",
+ "message": {
+ "id": "msg_test",
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-scripted-9",
+ "content": [],
+ "stop_reason": None,
+ "stop_sequence": None,
+ "usage": {"input_tokens": 1, "output_tokens": 1},
+ },
+ },
+ {
+ "type": "content_block_start",
+ "index": 0,
+ "content_block": {"type": "tool_use", "id": "toolu_test", "name": "json_tool_call", "input": {}},
+ },
+ *(
+ {
+ "type": "content_block_delta",
+ "index": 0,
+ "delta": {"type": "input_json_delta", "partial_json": wrapper[start : start + 20]},
+ }
+ for start in range(0, len(wrapper), 20)
+ ),
+ {"type": "content_block_stop", "index": 0},
+ {
+ "type": "message_delta",
+ "delta": {"stop_reason": "tool_use", "stop_sequence": None},
+ "usage": {"output_tokens": 9},
+ },
+ {"type": "message_stop"},
+ )
+ return httpx.Response(
+ status_code=200,
+ headers={"content-type": "text/event-stream"},
+ content="".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events).encode(),
+ request=httpx.Request("POST", url),
+ )
+
+
+def _streamed_answer(chunks: tuple) -> dict:
+ deltas: Final = tuple(chunk.choices[0].delta for chunk in chunks)
+ return {
+ "content": json.loads("".join(delta.content or "" for delta in deltas)),
+ "tool_call_deltas": sum(1 for delta in deltas if delta.tool_calls),
+ "finish_reasons": tuple(chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason),
+ }
+
+
+_BARE_STREAMED_ANSWER: Final = {"content": _RISK_ANSWER, "tool_call_deltas": 0, "finish_reasons": ("stop",)}
+
+
+def test_chat_completion_claude_platform_streams_a_wrapped_json_schema_answer_as_the_bare_object():
+ import litellm
+
+ chunks: Final = tuple(
+ litellm.completion(
+ **_JSON_SCHEMA_STREAM_KWARGS,
+ client=_sync_gateway_client([], respond=_wrapped_json_tool_stream),
+ )
+ )
+
+ assert _streamed_answer(chunks) == _BARE_STREAMED_ANSWER
+
+
+@pytest.mark.asyncio
+async def test_async_chat_completion_claude_platform_streams_a_wrapped_json_schema_answer_as_the_bare_object():
+ import litellm
+
+ stream: Final = await litellm.acompletion(
+ **_JSON_SCHEMA_STREAM_KWARGS,
+ client=_async_gateway_client([], respond=_wrapped_json_tool_stream),
+ )
+ chunks: Final = tuple([chunk async for chunk in stream])
+
+ assert _streamed_answer(chunks) == _BARE_STREAMED_ANSWER
+
+
def test_fake_claude_platform_gateway_rejects_leaked_internal_fields():
leaked = json.dumps({**EXPECTED_NATIVE_BODY, "workspace_id": "wrkspc_test", "aws_region_name": "us-west-2"})
response = _fake_claude_platform_gateway(
diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py
index 347840b5c5d..200768bc49f 100644
--- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py
+++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py
@@ -4228,6 +4228,54 @@ async def test_completion_streams_after_the_async_transform_request():
assert "".join(chunk.choices[0].delta.content or "" for chunk in collected) == "async"
+class _IteratorRecordingConfig(_TransformRecordingConfig):
+ def __init__(self):
+ super().__init__(transform_async=False)
+ self.iterator_requests = []
+
+ def get_model_response_iterator(self, streaming_response, sync_stream, json_mode=False):
+ self.iterator_requests.append({"sync_stream": sync_stream, "json_mode": json_mode})
+ return super().get_model_response_iterator(streaming_response, sync_stream, json_mode)
+
+
+def _json_mode_stream(config, **transport):
+ return BaseLLMHTTPHandler().completion(
+ model="stub-model",
+ messages=[{"role": "user", "content": "hi"}],
+ api_base="https://llm.example/v1/chat",
+ custom_llm_provider="openai",
+ model_response=ModelResponse(),
+ encoding=None,
+ logging_obj=_streaming_logging_obj(),
+ optional_params={"json_mode": True},
+ timeout=10.0,
+ litellm_params={},
+ stream=True,
+ provider_config=config,
+ **transport,
+ )
+
+
+def _streamed_text(chunks):
+ return "".join(chunk.choices[0].delta.content or "" for chunk in chunks)
+
+
+async def test_json_mode_stream_reaches_the_iterator_of_a_config_without_a_request_hook_only_when_sync():
+ config = _IteratorRecordingConfig()
+ sync_client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_sse_echoing_transformed_by)))
+ async_client = AsyncHTTPHandler()
+ async_client.client = httpx.AsyncClient(transport=httpx.MockTransport(_sse_echoing_transformed_by))
+
+ sync_chunks = list(_json_mode_stream(config, acompletion=False, client=sync_client))
+ async_chunks = [chunk async for chunk in await _json_mode_stream(config, acompletion=True, client=async_client)]
+
+ assert (_streamed_text(sync_chunks), _streamed_text(async_chunks)) == ("sync", "sync")
+ assert config.iterator_requests == [
+ {"sync_stream": True, "json_mode": True},
+ {"sync_stream": False, "json_mode": False},
+ ]
+
+
class _ImageEditRecordingConfig(BaseImageEditConfig):
def __init__(self):
self.transform_calls = []
From f0fedda722769d1eee61eb212bc6b86a24fa6f4f Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 11:45:49 -0700
Subject: [PATCH 22/47] fix(caching): collect the Redis circuit-breaker state
gauge as one series per state across worker processes (#46357)
The gauge declares multiprocess_mode=livesum, so PROMETHEUS_MULTIPROC_DIR apps get one series per state summed over
the live workers instead of one series per worker pid
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/caching/redis_cache.py | 1 +
.../_support/redis_breaker_worker.py | 79 +++++++++++++++++++
..._redis_breaker_metrics_multiprocess_sdk.py | 57 +++++++++++++
3 files changed, 137 insertions(+)
create mode 100644 tests/integration/_support/redis_breaker_worker.py
create mode 100644 tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py
diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py
index 3a847044e3b..0e34709d128 100644
--- a/litellm/caching/redis_cache.py
+++ b/litellm/caching/redis_cache.py
@@ -450,6 +450,7 @@ class _BreakerMetrics:
"litellm_redis_circuit_breaker_state",
"Number of Redis circuit breakers currently in each state",
labelnames=("state",),
+ multiprocess_mode="livesum",
)
self._transitions = PromCounter(
"litellm_redis_circuit_breaker_transitions",
diff --git a/tests/integration/_support/redis_breaker_worker.py b/tests/integration/_support/redis_breaker_worker.py
new file mode 100644
index 00000000000..f551baacf24
--- /dev/null
+++ b/tests/integration/_support/redis_breaker_worker.py
@@ -0,0 +1,79 @@
+"""One SDK worker process for the Redis circuit-breaker metric cells: a Redis-cached completion against an owned
+peer, then a report of the breaker state gauge on that process's own default registry. Run as a script by
+``run_breaker_worker`` so ``PROMETHEUS_MULTIPROC_DIR`` is in place before ``prometheus_client`` is imported."""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+
+from pydantic import TypeAdapter
+
+STATE_GAUGE: Final = "litellm_redis_circuit_breaker_state"
+WORKER_TIMEOUT_SECONDS: Final = 120
+
+
+@dataclass(frozen=True, slots=True)
+class WorkerReport:
+ pid: int
+ answer: str
+ state_series: tuple[tuple[tuple[tuple[str, str], ...], float], ...]
+
+
+_REPORT: Final = TypeAdapter(WorkerReport)
+
+
+def run_breaker_worker(api_base: str, prom_dir: Path, environment: Mapping[str, str] = os.environ) -> WorkerReport:
+ completed: Final = subprocess.run(
+ (sys.executable, "-I", __file__, api_base, environment["REDIS_HOST"], environment["REDIS_PORT"]),
+ env={**environment, "PROMETHEUS_MULTIPROC_DIR": str(prom_dir)},
+ capture_output=True,
+ text=True,
+ timeout=WORKER_TIMEOUT_SECONDS,
+ check=False,
+ )
+ assert completed.returncode == 0, completed.stderr[-3000:]
+ return _REPORT.validate_json(completed.stdout.splitlines()[-1])
+
+
+def _state_series_on_default_registry() -> tuple[tuple[tuple[tuple[str, str], ...], float], ...]:
+ from prometheus_client import REGISTRY
+
+ return tuple(sorted(_series(REGISTRY.collect())))
+
+
+def _series(families: Iterator[object]) -> Iterator[tuple[tuple[tuple[str, str], ...], float]]:
+ for family in families:
+ if getattr(family, "name", None) != STATE_GAUGE:
+ continue
+ for sample in getattr(family, "samples", ()):
+ yield tuple(sorted(sample.labels.items())), float(sample.value)
+
+
+def main(api_base: str, redis_host: str, redis_port: str) -> None:
+ import litellm
+ from litellm.caching.caching import Cache
+
+ litellm.cache = Cache(type="redis", host=redis_host, port=redis_port) # test-quality-ok: throwaway worker process
+ response: Final = litellm.completion(
+ model="openai/gpt-5.6",
+ api_base=api_base,
+ api_key="synthetic-breaker-worker-key",
+ messages=[{"role": "user", "content": f"breaker worker {os.getpid()}"}],
+ )
+ report: Final = WorkerReport(
+ pid=os.getpid(),
+ answer=str(response.choices[0].message.content),
+ state_series=_state_series_on_default_registry(),
+ )
+ sys.stdout.write(json.dumps(_REPORT.dump_python(report)) + "\n")
+
+
+if __name__ == "__main__":
+ main(*sys.argv[1:4])
diff --git a/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py b/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py
new file mode 100644
index 00000000000..d1582dd59b9
--- /dev/null
+++ b/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py
@@ -0,0 +1,57 @@
+"""An SDK app under a prefork worker pool (Celery, gunicorn) with ``PROMETHEUS_MULTIPROC_DIR`` set: the Redis
+circuit-breaker state gauge must collect as one series per state across every worker, never one per worker pid."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Final
+
+from integration._support.redis_breaker_worker import STATE_GAUGE, WorkerReport, run_breaker_worker
+from integration._support.wire import Reply, Request, wire_server
+from prometheus_client import CollectorRegistry, multiprocess
+
+_ANSWER: Final = "answered through the redis cache"
+_CHAT_RESPONSE: Final = json.dumps(
+ {
+ "id": "chatcmpl-redis-breaker-worker",
+ "object": "chat.completion",
+ "created": 1700000000,
+ "model": "gpt-5.6",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": _ANSWER}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 5, "completion_tokens": 4, "total_tokens": 9},
+ }
+).encode()
+_CLOSED: Final = (("state", "closed"),)
+
+
+def _reply(request: Request) -> Reply:
+ assert request.method == "POST", request.method
+ return Reply(body=_CHAT_RESPONSE)
+
+
+def _state_samples(prom_dir: Path) -> Iterator[tuple[tuple[tuple[str, str], ...], float]]:
+ registry: Final = CollectorRegistry()
+ multiprocess.MultiProcessCollector(registry, path=str(prom_dir))
+ for family in registry.collect():
+ if family.name != STATE_GAUGE:
+ continue
+ for sample in family.samples:
+ yield tuple(sorted(sample.labels.items())), sample.value
+
+
+def _collected_state_series(prom_dir: Path) -> tuple[tuple[tuple[tuple[str, str], ...], float], ...]:
+ return tuple(sorted(_state_samples(prom_dir)))
+
+
+def _worked(report: WorkerReport) -> WorkerReport:
+ assert report.answer == _ANSWER, report
+ return report
+
+
+def test_two_worker_processes_collect_as_one_state_series_per_state(tmp_path: Path) -> None:
+ with wire_server(_reply) as wire:
+ reports: Final = tuple(_worked(run_breaker_worker(wire.url, tmp_path)) for _ in range(2))
+ assert len({report.pid for report in reports}) == 2, reports
+ assert _collected_state_series(tmp_path) == ((_CLOSED, 2.0),), reports
From 0d67cccf03c38a34fe869a4927e5e063fa51efb2 Mon Sep 17 00:00:00 2001
From: "berriai-litellm-provider-info-sync[bot]"
<328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com>
Date: Wed, 14 Oct 2026 11:55:09 -0700
Subject: [PATCH 23/47] fix(gemini): add the changelog shutdown date to
gemini-2.0-flash-exp-image-generation (#46489)
Price-Sync: litellm-providers
Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com>
---
litellm/model_prices_and_context_window_backup.json | 2 ++
model_prices_and_context_window.json | 2 ++
2 files changed, 4 insertions(+)
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index e5bd69d0086..9011e4511d7 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -56977,6 +56977,7 @@
]
},
"gemini-2.0-flash-exp-image-generation": {
+ "deprecation_date": "2025-11-14",
"input_cost_per_token": 0.0,
"litellm_provider": "gemini",
"max_input_tokens": 32768,
@@ -56997,6 +56998,7 @@
"supports_vision": true
},
"gemini/gemini-2.0-flash-exp-image-generation": {
+ "deprecation_date": "2025-11-14",
"input_cost_per_token": 0.0,
"litellm_provider": "gemini",
"max_input_tokens": 32768,
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index e5bd69d0086..9011e4511d7 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -56977,6 +56977,7 @@
]
},
"gemini-2.0-flash-exp-image-generation": {
+ "deprecation_date": "2025-11-14",
"input_cost_per_token": 0.0,
"litellm_provider": "gemini",
"max_input_tokens": 32768,
@@ -56997,6 +56998,7 @@
"supports_vision": true
},
"gemini/gemini-2.0-flash-exp-image-generation": {
+ "deprecation_date": "2025-11-14",
"input_cost_per_token": 0.0,
"litellm_provider": "gemini",
"max_input_tokens": 32768,
From 0915691d320a58072bf3104a8e05086bb8e85a22 Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Wed, 14 Oct 2026 12:00:35 -0700
Subject: [PATCH 24/47] test(integration): remove six setup causes of CircleCI
integration reds (#46490)
Drain the worker's fake-upstream request log when a test's gateway fixture starts, so traffic left by earlier tests on the same xdist worker no longer shows up in a test's first observation read
Generate the fine-tuning rig's provider file and job ids per fixture build, so a rebuilt module fixture in the same worker no longer collides on the managed object unique constraint
Wait for every wedged team_id and team_alias invalidation publish to fail before taking the recovery marker, now that team updates publish both keys
Pin the cooled-down priority scheduler request to the connection that observed the cooldown, since a sibling worker reads cooldowns from Redis on an interval
Simulate upstream outages and the held budget peer on ports the test keeps owning, instead of releasing a port another worker can bind
Start a peer proxy for the accounting group, which has a test that needs one
---
.circleci/scripts/run_integration.sh | 2 +-
tests/integration/README.md | 2 +-
tests/integration/conftest.py | 1 +
..._user_updates_wedged_coordination_redis.py | 19 +++++------
...vertex_anthropic_inline_tools_beta_wire.py | 16 +++++----
...df_document_pre_call_checks_owned_proxy.py | 33 +++++++++----------
.../test_priority_scheduler_queue_cleanup.py | 18 +++++++---
tests/integration/run.py | 2 +-
.../security/test_fine_tuning_managed_ids.py | 26 ++++++---------
9 files changed, 62 insertions(+), 57 deletions(-)
diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh
index be7e70d491c..00b3c5f8fdd 100644
--- a/.circleci/scripts/run_integration.sh
+++ b/.circleci/scripts/run_integration.sh
@@ -201,7 +201,7 @@ if [ "$isolated_workers" = false ]; then
curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
-d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
- if [ "$suite" = management ] || [ "$suite" = mcp ]; then
+ if [ "$suite" = management ] || [ "$suite" = mcp ] || [ "$suite" = accounting ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
start_proxy 4001 peer.log
peer_pid="$launched_pid"
diff --git a/tests/integration/README.md b/tests/integration/README.md
index 572b80ad5c6..0f1da7b3d97 100644
--- a/tests/integration/README.md
+++ b/tests/integration/README.md
@@ -10,7 +10,7 @@ The `cost` group is driven by the `CostTrackingTestCase` literals under `cost_ca
Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `mcp`, `sdk` or `cost` to run a selected group. The group to directory mapping is the `GROUPS` literal at the top of `run.py`; a new directory needs a `GROUPS` entry and an `OWNED_DIRECTORIES` entry in `_support/manifest.py`. Set `INTEGRATION_WORKERS` above 1 to run a group under pytest-xdist. Every group run by more than one worker gives each worker a private database, Redis, scripted upstream process, fake provider and proxy, with a peer and read-only database role where required. There is no opt-in list and no shared proxy, upstream, database or Redis in that mode: `run_integration.sh` does not start them and does not hand their addresses to the run, so a test can only reach its own worker's services. The pytest controller owns these services and closes them after worker failures. Their cases use pytest-xdist load groups inside the files selected by CircleCI timing splits. Module and class fixtures are grouped automatically, so authors do not need scheduling markers. Function fixtures remain independently schedulable. Tests that own an independent mutable sink for every case should use function fixtures, so a shared module fixture does not serialize them or share pending observations. Nested harness probes discard the parent run's shard, worker and result settings, then supply their own selections. MCP keeps its explicit load groups. Automatic group suffixes are removed from JUnit and execution records to preserve case identities. Owned proxy helpers wait for a healthy response from every configured, live worker before yielding. Serial runs keep the supplied deployment, including routing parity runs. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. When that deployment runs more than one proxy worker, set `INTEGRATION_PROXY_WORKERS` to the count so a test that writes a model and then calls it waits out the config reload interval, the only cross-worker convergence bound the wire exposes. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
-Management and MCP also require `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. Under several workers each worker gets a second directly addressed proxy sharing only that worker's stores, and a single-worker run uses the two proxies `run_integration.sh` starts. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload
+Management, MCP and accounting also require `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. Under several workers each worker gets a second directly addressed proxy sharing only that worker's stores, and a single-worker run uses the two proxies `run_integration.sh` starts. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload
The generated lifecycle models use 20 examples, eight steps, generation and shrinking, with isolated resources per example. HTTP operation caps include generation and shrinking and exempt cleanup. Local qualification defaults to seed 4106601 and canonical order; CircleCI derives exploration and ordering seeds from the checked-out revision and workflow ID. The ordering seed shuffles files, scoped fixture configurations and cases within each configuration. Ordering keeps shared module and class configurations together where their dependencies permit it. Dependent fixtures with their own module parameters can still rebuild an outer configuration, with every case retained
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index cb23713e6b2..a4bd6092a0c 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -285,6 +285,7 @@ def pytest_sessionfinish(session: pytest.Session, exitstatus: int) -> None:
@pytest.fixture
def gateway() -> Iterator[Gateway]:
with gateway_from_environment() as value:
+ httpx.get(f"{value.upstream_url}/__observations", timeout=5, trust_env=False).raise_for_status()
yield value
diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py
index 41fefdcd3b2..09b16e7d5ef 100644
--- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py
+++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py
@@ -249,27 +249,26 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged(
assert timings["c1_survivor"] < _HANDLER_BUDGET_SECONDS, (
f"control update on the surviving worker took {timings['c1_survivor']:.3f}s"
)
- wedged_keys: Final = {users[i] for i in range(_BURST) if i % 5 == 0 and i != 0} | {
- f"team_id:{team_id}"
- }
+ wedged_users: Final = {users[i] for i in range(_BURST) if i % 5 == 0 and i != 0}
+ wedged_team_keys: Final = {f"team_id:{team_id}", f"team_alias:{identity}"}
+ wedged_keys: Final = wedged_users | wedged_team_keys
+ wedged_team_updates: Final = 1 + sum(1 for i in range(_BURST) if i % 5 == 3)
def proxy_log() -> str:
return "".join(
path.read_text() for path in results_dir.glob("owned-proxy-*.log") if path not in prior_logs
)
- team_wedged_key: Final = f"team_id:{team_id}"
eventually(
proxy_log,
lambda text: (
all(
- text.count(f"publish for {wedged_key} failed") >= _WEDGED_PUBLISHES_PER_BURST_USER
- for wedged_key in wedged_keys
- if wedged_key != team_wedged_key
+ text.count(f"publish for {user} failed") >= _WEDGED_PUBLISHES_PER_BURST_USER
+ for user in wedged_users
)
- and (
- f"publish for {team_wedged_key} failed" in text
- or f"internal usage cache entry {team_wedged_key}" in text
+ and all(
+ text.count(f"publish for {team_key} failed") >= wedged_team_updates
+ for team_key in wedged_team_keys
)
),
seconds=45,
diff --git a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py
index a6b039dc8f7..3175f2ec075 100644
--- a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py
+++ b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py
@@ -512,15 +512,19 @@ def test_upstream_outage_mid_run_recovers_with_inline_tools_still_forwarded(gate
seen: Final = tuple(sum(marker in sent.body for sent in delivered) for _, marker in results)
assert seen == (1,) * len(clients), seen
- with wire_server(_peer) as wire:
- port: Final = int(wire.url.rsplit(":", 1)[1])
- model: Final = _vertex_deployment(gateway, scenario, wire.url)
- forwarded(wire, wave("before outage"))
+ down: Final = threading.Event()
+ wire: Final = stack.enter_context(
+ wire_server(lambda request: Reply(drop_connection=True) if down.is_set() else _peer(request))
+ )
+ model: Final = _vertex_deployment(gateway, scenario, wire.url)
+ forwarded(wire, wave("before outage"))
+ down.set()
outage: Final = wave("during outage")
assert all(status is not None and status >= 500 for status, _ in outage), outage
assert gateway.request("GET", "/health/liveliness").status_code == 200
- with wire_server(_peer, port=port) as revived:
- forwarded(revived, wave("after outage"))
+ wire.drain()
+ down.clear()
+ forwarded(wire, wave("after outage"))
@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
diff --git a/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py b/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py
index 9fad73f8080..c9a4e8e8881 100644
--- a/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py
+++ b/tests/integration/routing/test_pdf_document_pre_call_checks_owned_proxy.py
@@ -1,5 +1,4 @@
import os
-import socket
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
@@ -50,19 +49,14 @@ _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
class _Rig:
owned: OwnedProxy
wire: Wire
- held_port: int
+ held: Wire
+ release: threading.Event
def _marker() -> str:
return uuid.uuid4().hex
-def _free_port() -> int:
- with socket.socket() as reserve:
- reserve.bind(("127.0.0.1", 0))
- return int(reserve.getsockname()[1])
-
-
def _deployment(
name: str,
backend: str,
@@ -86,7 +80,8 @@ def _peer(request: Request) -> Reply:
def _held(release: threading.Event) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
- assert release.wait(timeout=120), "Held peer was never released"
+ if request.method == "POST":
+ assert release.wait(timeout=120), "Held peer was never released"
return _peer(request)
return respond
@@ -107,14 +102,14 @@ def _request(surface: str, model: str, pages: int, marker: str) -> tuple[str, di
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
directory: Final = tmp_path_factory.mktemp("pdf-pre-call")
- held_port: Final = _free_port()
- with gateway_from_environment() as gateway, wire_server(_peer) as wire:
+ release: Final = threading.Event()
+ with gateway_from_environment() as gateway, wire_server(_peer) as wire, wire_server(_held(release)) as held:
config: Final = owned_config(
directory,
[
_deployment(_WINDOW, "claude-haiku-5-5", wire.url, model_info={"max_input_tokens": _LIMIT}),
_deployment(_ITPM, "claude-sonnet-4-6", wire.url, itpm=_LIMIT),
- _deployment(_BUDGET, "claude-opus-4-8", f"http://127.0.0.1:{held_port}"),
+ _deployment(_BUDGET, "claude-opus-4-8", held.url),
],
litellm_settings={"cache": False},
router_settings={
@@ -122,8 +117,11 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
"optional_pre_call_checks": ["enforce_model_rate_limits"],
},
)
- with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned:
- yield _Rig(owned, wire, held_port)
+ try:
+ with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned:
+ yield _Rig(owned, wire, held, release)
+ finally:
+ release.set()
def _spend_row(call_id: str) -> dict[str, JsonValue]:
@@ -168,9 +166,10 @@ def test_model_itpm_check_rejects_a_pdf_whose_pages_exceed_the_limit(rig: _Rig)
def test_key_budget_reservation_rejects_the_second_concurrent_pdf_request(rig: _Rig) -> None:
- release: Final = threading.Event()
requests: Final = tuple(_request("messages", _BUDGET, _PAGE_COUNT, _marker()) for _ in range(2))
- with wire_server(_held(release), port=rig.held_port) as wire, rig.owned.gateway.scenario() as scenario:
+ wire: Final = rig.held
+ wire.drain()
+ with rig.owned.gateway.scenario() as scenario:
key: Final = scenario.key(max_budget=0.01)
headers: Final = {"Authorization": f"Bearer {key}"}
with (
@@ -187,7 +186,7 @@ def test_key_budget_reservation_rejects_the_second_concurrent_pdf_request(rig: _
seconds=60,
)
finally:
- release.set()
+ rig.release.set()
responses: Final = tuple(future.result(timeout=90) for future in futures)
received: Final = wire.drain()
assert sorted(response.status_code for response in responses) == [200, 422], [
diff --git a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py
index bcfb28f790a..993da52fc67 100644
--- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py
+++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py
@@ -617,11 +617,19 @@ def test_concurrent_prioritized_burst_across_instances_and_endpoints_is_served(p
@pytest.mark.timeout(PAIR_CELL_TIMEOUT)
def test_fresh_queue_during_a_cooldown_is_refused_at_once(pair: Pair) -> None:
group: Final = PAIR_GROUPS[2]
- cooled(pair.second, pair.upstream, group)
- client: Final = sdk(pair.second)
- started: Final = time.monotonic()
- with pytest.raises(openai.APIStatusError) as refused:
- client.completions.create(model=group, prompt=new_marker(), extra_body={"priority": 1})
+ with ExitStack() as stack:
+ worker: Final = pinned(pair.second, stack)
+ cooled(worker, pair.upstream, group)
+ client: Final = openai.OpenAI(
+ base_url=str(worker.client.base_url),
+ api_key=worker.key,
+ max_retries=0,
+ timeout=30,
+ http_client=worker.client,
+ )
+ started: Final = time.monotonic()
+ with pytest.raises(openai.APIStatusError) as refused:
+ client.completions.create(model=group, prompt=new_marker(), extra_body={"priority": 1})
assert refused.value.status_code == 429, refused.value.message
assert NO_DEPLOYMENTS in refused.value.message
assert_cut_short(time.monotonic() - started, injected_delay=ROUTER_TIMEOUT_SECONDS)
diff --git a/tests/integration/run.py b/tests/integration/run.py
index 43e606f46a7..99518e69d1f 100644
--- a/tests/integration/run.py
+++ b/tests/integration/run.py
@@ -39,7 +39,7 @@ GITHUB_FILES: Final = frozenset(
)
-PEER_GROUPS: Final = frozenset({"management", "mcp"})
+PEER_GROUPS: Final = frozenset({"management", "mcp", "accounting"})
COST_MAP_GROUPS: Final = frozenset({"cost"})
diff --git a/tests/integration/security/test_fine_tuning_managed_ids.py b/tests/integration/security/test_fine_tuning_managed_ids.py
index 0c1a0f71b3f..2c98bca6eea 100644
--- a/tests/integration/security/test_fine_tuning_managed_ids.py
+++ b/tests/integration/security/test_fine_tuning_managed_ids.py
@@ -21,11 +21,6 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
MANAGED_MODEL: Final = "gpt-4.1-managed-ids"
DEPLOYMENT_KEY: Final = "sk-fine-tuning-managed-deployment"
-RUN: Final = uuid.uuid4().hex[:12]
-PROVIDER_FILE: Final = f"file-managed-{RUN}"
-PROVIDER_JOB: Final = f"ftjob-managed-{RUN}"
-DEFAULT_PROVIDER_FILE: Final = f"file-managed-default-{RUN}"
-DEFAULT_PROVIDER_JOB: Final = f"ftjob-managed-default-{RUN}"
TRAINING_LINE: Final = b'{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}\n'
JOB_PATH: Final = re.compile(r"^/v1/fine_tuning/jobs/([^/]+?)(/cancel)?$")
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
@@ -190,7 +185,10 @@ class Rig:
@contextmanager
-def _rig(directory: Path, litellm_settings: dict[str, JsonValue], file_id: str, job_id: str) -> Iterator[Rig]:
+def _rig(directory: Path, litellm_settings: dict[str, JsonValue]) -> Iterator[Rig]:
+ build: Final = uuid.uuid4().hex
+ file_id: Final = f"file-managed-{build}"
+ job_id: Final = f"ftjob-managed-{build}"
with gateway_from_environment() as shared, wire_server(_responder(file_id, job_id)) as wire:
config: Final = directory / "config.yaml"
config.write_text(yaml.safe_dump(_config(wire, litellm_settings)))
@@ -239,17 +237,13 @@ def _rig(directory: Path, litellm_settings: dict[str, JsonValue], file_id: str,
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
- with _rig(
- tmp_path_factory.mktemp("fine-tuning-managed-ids"), {"require_managed_files": True}, PROVIDER_FILE, PROVIDER_JOB
- ) as built:
+ with _rig(tmp_path_factory.mktemp("fine-tuning-managed-ids"), {"require_managed_files": True}) as built:
yield built
@pytest.fixture(scope="module")
def default_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
- with _rig(
- tmp_path_factory.mktemp("fine-tuning-managed-ids-default"), {}, DEFAULT_PROVIDER_FILE, DEFAULT_PROVIDER_JOB
- ) as built:
+ with _rig(tmp_path_factory.mktemp("fine-tuning-managed-ids-default"), {}) as built:
yield built
@@ -433,7 +427,7 @@ def test_keys_granted_the_retrieve_route_get_the_owner_check_and_the_decoded_job
with rig.sdk(owner) as client:
retrieved: Final = client.fine_tuning.jobs.with_raw_response.retrieve(rig.managed_job)
assert [(r.method, r.target, r.headers["authorization"], r.body) for r in rig.wire.drain()] == [
- ("GET", f"/v1/fine_tuning/jobs/{PROVIDER_JOB}", f"Bearer {DEPLOYMENT_KEY}", b"")
+ ("GET", f"/v1/fine_tuning/jobs/{rig.provider_job}", f"Bearer {DEPLOYMENT_KEY}", b"")
]
assert retrieved.status_code == 200, retrieved.text
assert Job.model_validate_json(retrieved.text) == _client_job(rig.managed_job, "running", rig.provider_file), (
@@ -455,7 +449,7 @@ def test_keys_granted_the_retrieve_route_get_default_job_ownership_and_decoded_i
assert [
(request.method, request.target, request.headers["authorization"], request.body)
for request in default_rig.wire.drain()
- ] == [("GET", f"/v1/fine_tuning/jobs/{DEFAULT_PROVIDER_JOB}", f"Bearer {DEPLOYMENT_KEY}", b"")]
+ ] == [("GET", f"/v1/fine_tuning/jobs/{default_rig.provider_job}", f"Bearer {DEPLOYMENT_KEY}", b"")]
assert retrieved.status_code == 200, retrieved.text
assert Job.model_validate_json(retrieved.text) == _client_job(
default_rig.managed_job, "running", default_rig.provider_file
@@ -472,7 +466,7 @@ def test_owner_retrieve_reaches_the_provider_with_the_decoded_job_id(rig: Rig) -
with rig.sdk(rig.owner) as client:
retrieved: Final = client.fine_tuning.jobs.with_raw_response.retrieve(rig.managed_job)
assert [(r.method, r.target, r.headers["authorization"], r.body) for r in rig.wire.drain()] == [
- ("GET", f"/v1/fine_tuning/jobs/{PROVIDER_JOB}", f"Bearer {DEPLOYMENT_KEY}", b"")
+ ("GET", f"/v1/fine_tuning/jobs/{rig.provider_job}", f"Bearer {DEPLOYMENT_KEY}", b"")
]
assert retrieved.status_code == 200, retrieved.text
assert Job.model_validate_json(retrieved.text) == _client_job(rig.managed_job, "running", rig.provider_file), (
@@ -484,7 +478,7 @@ def test_owner_cancel_reaches_the_provider_with_the_decoded_job_id(rig: Rig) ->
with rig.sdk(rig.owner) as client:
cancelled: Final = client.fine_tuning.jobs.with_raw_response.cancel(rig.managed_job)
assert [(r.method, r.target, r.headers["authorization"], r.body) for r in rig.wire.drain()] == [
- ("POST", f"/v1/fine_tuning/jobs/{PROVIDER_JOB}/cancel", f"Bearer {DEPLOYMENT_KEY}", b"")
+ ("POST", f"/v1/fine_tuning/jobs/{rig.provider_job}/cancel", f"Bearer {DEPLOYMENT_KEY}", b"")
]
assert cancelled.status_code == 200, cancelled.text
assert Job.model_validate_json(cancelled.text) == _client_job(rig.managed_job, "cancelled", rig.provider_file), (
From 6e0f0f7adb3083fb6681e636c947e3d6b8f53b26 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Date: Wed, 14 Oct 2026 12:04:17 -0700
Subject: [PATCH 25/47] fix(s3_v2): bound each upload attempt, add jitter,
retry 429 and transport errors in the call (#46340)
* fix(s3_v2): bound each upload attempt, add jitter, retry on 429 and 503
Each S3 PUT now carries a per-attempt httpx timeout read from
S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS (default 30s, clamped to 1..3600) instead
of inheriting the handler's unbounded default. Retries back off with jitter
instead of a fixed 1s then 2s, 429 joins 403, 500 and 503 in the in-call
retry set, and transport errors or litellm.Timeout on an attempt are retried
in the same call instead of ending it. Each flush logs one warning per batch
of dropped uploads that names how many were terminal and how many aged out.
* fix(s3_v2): retry uploads in a bounded loop and align tests with the new retry contract
The async and sync retry helpers called themselves, which the recursive-function check rejects. They are now a bounded loop over the attempt count with the same outcome.
The 429 integration test now expects in-call jittered backoff, and the aged-out test matches the new drop warning text. The sync transport test drives the real HTTPHandler so a dropped per-attempt timeout or a missed litellm.Timeout fails it, and the drop-reason test uses unequal counts so swapped labels fail.
---------
Co-authored-by: yuneng-jiang
---
litellm/constants.py | 1 +
litellm/integrations/s3_v2.py | 164 +++++---
.../observability/test_s3_v2_upload_fanout.py | 22 +-
tests/unit/integrations/test_s3_v2.py | 384 ++++++++++++++++--
4 files changed, 451 insertions(+), 120 deletions(-)
diff --git a/litellm/constants.py b/litellm/constants.py
index aed2b04b16b..9d18fe5489d 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -93,6 +93,7 @@ DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTER
DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512))
DEFAULT_S3_MAX_CONCURRENT_UPLOADS: Final = int(os.getenv("DEFAULT_S3_MAX_CONCURRENT_UPLOADS", "16"))
DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY: Final = get_env_int("DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY", 200)
+S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS: Final = get_env_int_in_range("S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS", 30, 1, 3600)
# https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
MAX_S3_OBJECT_KEY_BYTES: Final = 1024
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64
diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py
index 47d9b1c7b64..0b7a2609bcb 100644
--- a/litellm/integrations/s3_v2.py
+++ b/litellm/integrations/s3_v2.py
@@ -10,13 +10,14 @@ import asyncio
import contextvars
import logging
import os
+import random
import re
import time
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import partial
-from typing import TYPE_CHECKING, Final, Literal, cast
+from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast
from urllib.parse import quote
from uuid import uuid4
@@ -30,7 +31,9 @@ from litellm.constants import (
DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY,
DEFAULT_S3_MAX_CONCURRENT_UPLOADS,
S3_PARTITION_GRANULARITY_ENV_VAR,
+ S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS,
)
+from litellm.exceptions import Timeout as LiteLLMTimeout
from litellm.integrations.adaptive_concurrency import AdaptiveConcurrencyLimiter, PutSample
from litellm.integrations.s3 import (
get_s3_object_download_filename,
@@ -79,8 +82,13 @@ _TERMINAL_ERROR_CODES: Final = frozenset(
}
)
_BODY_CODED_STATUSES: Final = frozenset({400, 403})
-_RETRYABLE_STATUSES: Final = frozenset({403, 500, 503})
+_RETRYABLE_STATUSES: Final = frozenset({403, 429, 500, 503})
_S3_ERROR_CODE: Final = re.compile(r"([^<]+)")
+_TRANSIENT_ERRORS: Final = (httpx.TransportError, LiteLLMTimeout)
+_UPLOAD_ATTEMPTS: Final = 3
+_UPLOAD_ATTEMPT_TIMEOUT: Final = httpx.Timeout(S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS)
+
+PutOutcome: TypeAlias = httpx.Response | Exception
@dataclass(frozen=True, slots=True)
@@ -105,6 +113,35 @@ def _s3_key_parent(s3_object_key: str) -> str:
return s3_object_key.rsplit("/", 1)[0] if "/" in s3_object_key else ""
+def _backoff_seconds(attempt: int) -> float:
+ return (1 << attempt) * (0.5 + random.random())
+
+
+def _log_retry(level: int, attempt: int, reason: str, wait: float, s3_object_key: str) -> None:
+ verbose_logger.log(
+ level,
+ "S3 upload attempt %s/%s failed with %s, retrying in %.2fs key=%s",
+ attempt + 1,
+ _UPLOAD_ATTEMPTS,
+ reason,
+ wait,
+ s3_object_key,
+ )
+
+
+def _response_or_raise(outcome: PutOutcome) -> httpx.Response:
+ if isinstance(outcome, Exception):
+ raise outcome
+ outcome.raise_for_status()
+ return outcome
+
+
+def _throttled(outcome: PutOutcome) -> bool:
+ if isinstance(outcome, Exception):
+ return True
+ return outcome.status_code in (429, 503) or _s3_error_code(outcome) == "SlowDown"
+
+
class S3BatchUploadError(Exception):
def __init__(self, failed: int, total: int) -> None:
self.failed = failed
@@ -566,7 +603,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
url: Final = self._build_object_url(batch_logging_element.s3_object_key)
- async def signed_put(prepared: _PreparedPut) -> httpx.Response:
+ async def signed_put(prepared: _PreparedPut) -> PutOutcome:
credentials: Final = await asyncified_get_credentials(
aws_access_key_id=self.s3_aws_access_key_id,
aws_secret_access_key=self.s3_aws_secret_access_key,
@@ -582,34 +619,16 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
self._sign_put, credentials, url, prepared.json_string, prepared.headers
)
try:
- return await self.async_httpx_client.put(url, data=prepared.json_string, headers=signed_headers)
+ return await self.async_httpx_client.put(
+ url, data=prepared.json_string, headers=signed_headers, timeout=_UPLOAD_ATTEMPT_TIMEOUT
+ )
except httpx.HTTPStatusError as error:
return error.response
+ except _TRANSIENT_ERRORS as error:
+ return error
- max_retries: Final = 3
prepared: Final = self._prepare_put(batch_logging_element)
- for attempt in range(max_retries):
- response = await self._recorded_put(partial(signed_put, prepared))
- if (
- response.status_code in _RETRYABLE_STATUSES
- and not (self.s3_drop_on_terminal_error and _is_terminal(response))
- and attempt < max_retries - 1
- ):
- wait_time = 1 << attempt # 1s, 2s
- verbose_logger.log(
- logging.DEBUG if _in_flush.get() else logging.WARNING,
- "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
- response.status_code,
- wait_time,
- attempt + 1,
- max_retries,
- batch_logging_element.s3_object_key,
- )
- self._flush_retries += 1
- await asyncio.sleep(wait_time)
- continue
- response.raise_for_status()
- break
+ await self._put_with_retries(partial(signed_put, prepared), batch_logging_element.s3_object_key)
except Exception as e:
verbose_logger.exception("Error uploading to s3: %s", e)
self.handle_callback_failure(callback_name="S3Logger")
@@ -678,9 +697,13 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
)
dropped: Final = len(failed) - len(requeued)
if dropped:
+ terminal: Final = sum(1 for _, outcome in failed if outcome == "dropped")
verbose_logger.warning(
- "s3 logging: %s uploads dropped (terminal or retrying longer than s3_max_retry_age_seconds=%s)",
+ "s3 logging: %s uploads dropped this flush: %s rejected as terminal by S3, "
+ "%s retrying longer than s3_max_retry_age_seconds=%s",
dropped,
+ terminal,
+ dropped - terminal,
self.s3_max_retry_age_seconds,
)
if not requeued:
@@ -728,22 +751,46 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
return "dropped"
return "retry"
- async def _recorded_put(self, signed_put: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
+ async def _recorded_put(self, signed_put: Callable[[], Awaitable[PutOutcome]]) -> PutOutcome:
+ outcome: Final = await signed_put()
limiter: Final = self._upload_limiter
- adaptive: Final = limiter if isinstance(limiter, AdaptiveConcurrencyLimiter) else None
- try:
- response: Final = await signed_put()
- except Exception:
- if adaptive is not None:
- adaptive.record(PutSample(throttled=True))
- raise
- if adaptive is not None:
- adaptive.record(
- PutSample(
- throttled=response.status_code in (429, 503) or _s3_error_code(response) == "SlowDown",
- )
- )
- return response
+ if isinstance(limiter, AdaptiveConcurrencyLimiter):
+ limiter.record(PutSample(throttled=_throttled(outcome)))
+ return outcome
+
+ def _retry_reason(self, outcome: PutOutcome) -> str | None:
+ if isinstance(outcome, Exception):
+ return type(outcome).__name__
+ if outcome.status_code not in _RETRYABLE_STATUSES:
+ return None
+ if self.s3_drop_on_terminal_error and _is_terminal(outcome):
+ return None
+ return str(outcome.status_code)
+
+ async def _put_with_retries(
+ self, signed_put: Callable[[], Awaitable[PutOutcome]], s3_object_key: str
+ ) -> httpx.Response:
+ for attempt in range(_UPLOAD_ATTEMPTS - 1):
+ outcome = await self._recorded_put(signed_put)
+ reason = self._retry_reason(outcome)
+ if reason is None:
+ return _response_or_raise(outcome)
+ wait = _backoff_seconds(attempt)
+ _log_retry(logging.DEBUG if _in_flush.get() else logging.WARNING, attempt, reason, wait, s3_object_key)
+ self._flush_retries += 1
+ await asyncio.sleep(wait)
+ return _response_or_raise(await self._recorded_put(signed_put))
+
+ def _sync_put_with_retries(self, signed_put: Callable[[], PutOutcome], s3_object_key: str) -> httpx.Response:
+ for attempt in range(_UPLOAD_ATTEMPTS - 1):
+ outcome = signed_put()
+ reason = self._retry_reason(outcome)
+ if reason is None:
+ return _response_or_raise(outcome)
+ wait = _backoff_seconds(attempt)
+ _log_retry(logging.WARNING, attempt, reason, wait, s3_object_key)
+ time.sleep(wait)
+ return _response_or_raise(signed_put())
def _batch_file_elements(self, batch: tuple[s3BatchLoggingElement, ...]) -> tuple[s3BatchLoggingElement, ...]:
now: Final = datetime.now(timezone.utc)
@@ -879,7 +926,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None)
)
- def signed_put(prepared_put: _PreparedPut) -> httpx.Response:
+ def signed_put(prepared_put: _PreparedPut) -> PutOutcome:
credentials: Final = self.get_credentials(
aws_access_key_id=self.s3_aws_access_key_id,
aws_secret_access_key=self.s3_aws_secret_access_key,
@@ -887,29 +934,14 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
aws_region_name=self.s3_region_name,
)
signed_headers: Final = self._sign_put(credentials, url, prepared_put.json_string, prepared_put.headers)
- return httpx_client.put(url, data=prepared_put.json_string, headers=signed_headers)
-
- max_retries: Final = 3
- for attempt in range(max_retries):
- response = signed_put(prepared)
- if (
- response.status_code in _RETRYABLE_STATUSES
- and not (self.s3_drop_on_terminal_error and _is_terminal(response))
- and attempt < max_retries - 1
- ):
- wait_time = 1 << attempt # 1s, 2s
- verbose_logger.warning(
- "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
- response.status_code,
- wait_time,
- attempt + 1,
- max_retries,
- batch_logging_element.s3_object_key,
+ try:
+ return httpx_client.put(
+ url, data=prepared_put.json_string, headers=signed_headers, timeout=_UPLOAD_ATTEMPT_TIMEOUT
)
- time.sleep(wait_time)
- continue
- response.raise_for_status()
- break
+ except _TRANSIENT_ERRORS as error:
+ return error
+
+ self._sync_put_with_retries(partial(signed_put, prepared), batch_logging_element.s3_object_key)
except Exception as e:
verbose_logger.exception("Error uploading to s3: %s", e)
self.handle_callback_failure(callback_name="S3Logger")
diff --git a/tests/integration/observability/test_s3_v2_upload_fanout.py b/tests/integration/observability/test_s3_v2_upload_fanout.py
index ead1e5b028b..54e35d993e6 100644
--- a/tests/integration/observability/test_s3_v2_upload_fanout.py
+++ b/tests/integration/observability/test_s3_v2_upload_fanout.py
@@ -16,6 +16,7 @@ from _s3_v2_support import RecordingS3Sink, collect_payloads
from _s3_v2_support import s3_config as _recording_s3_config
from integration._support.client import Gateway, JsonValue, eventually
from integration._support.process import owned_proxy, owned_proxy_process
+from integration._support.timing import assert_enforced_within
from integration._support.wire import Reply, Request, Wire, wire_server
BUCKET: Final = "integration-bucket"
@@ -473,7 +474,7 @@ def test_s3_v2_aged_out_object_is_dropped_next_to_delivered_siblings(gateway: Ga
_send(owned.gateway, model, key, f"{marker}-trigger")
eventually(
lambda: owned.log.read_text(),
- lambda text: "retrying longer than s3_max_retry_age_seconds=1)" in text,
+ lambda text: "0 rejected as terminal by S3, 1 retrying longer than s3_max_retry_age_seconds=1" in text,
seconds=60,
)
exhausted: Final = sink.rejected_attempts
@@ -565,30 +566,27 @@ def test_s3_v2_retry_age_zero_keeps_aged_object_queued(gateway: Gateway, tmp_pat
)
-def test_s3_v2_throttled_429_object_is_put_once_per_flush(gateway: Gateway, tmp_path: Path) -> None:
+def test_s3_v2_throttled_429_object_backs_off_in_call_and_lands_once_cleared(gateway: Gateway, tmp_path: Path) -> None:
marker: Final = "s3throttle" + uuid.uuid4().hex[:8]
sink: Final = RejectingS3Sink(reject_marker=f"{marker}-throttled", reject_status=429, reject_code="TooManyRequests")
with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket:
config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": False})
with (
- owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "2"}, config=config) as candidate,
+ owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "10"}, config=config) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key")
key: Final = scenario.key(models=[model])
_send(candidate, model, key, f"{marker}-throttled")
- _send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-sibling")
- _send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-second-flush")
- _send_and_wait_until_landed(candidate, model, key, sink, f"{marker}-third-flush")
+ eventually(lambda: sink.rejected_attempts, lambda attempts: attempts >= 3, seconds=60)
sink.reject_until = time.time()
eventually(sink.landed_ids, lambda landed: f"{marker}-throttled" in landed, seconds=60)
- assert sum(1 for r in provider.drain() if r.method == "POST") == 4
- times: Final = tuple(sink.rejected_times)
- gaps: Final = tuple(round(later - earlier, 3) for earlier, later in zip(times, times[1:]))
- assert len(times) >= 3 and min(gaps) >= 1.5, (
- f"PUTs for a 429 object ran {gaps} apart; the 2 s flush interval allows exactly one attempt per flush "
- "because 429 is not an in-call retry status"
+ assert sum(1 for r in provider.drain() if r.method == "POST") == 1
+ first_backoff, second_backoff = (
+ later - earlier for earlier, later in zip(sink.rejected_times[:2], sink.rejected_times[1:3])
)
+ assert_enforced_within(first_backoff, configured=1.5, slack=2.0, not_before=0.5)
+ assert_enforced_within(second_backoff, configured=3.0, slack=2.0, not_before=1.0)
def test_s3_v2_default_config_retries_access_denied_and_every_event_lands(gateway: Gateway, tmp_path: Path) -> None:
diff --git a/tests/unit/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py
index 8449eae7bf4..fcd2c5673af 100644
--- a/tests/unit/integrations/test_s3_v2.py
+++ b/tests/unit/integrations/test_s3_v2.py
@@ -13,7 +13,7 @@ from contextvars import Context
from datetime import datetime
from pathlib import Path
from typing import Final
-from unittest.mock import AsyncMock, MagicMock, call, patch
+from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
@@ -28,7 +28,7 @@ from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
from litellm.types.utils import StandardLoggingPayload
from collections import defaultdict
from litellm._logging import verbose_logger
-from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
+from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -36,6 +36,14 @@ _real_sleep: Final = asyncio.sleep
_NOW: Final = 1_000_000.0
+def _assert_jittered_backoff(sleep_mock: MagicMock, retries: int) -> None:
+ """Each retry waits a jittered exponential backoff: attempt n sleeps somewhere in [2**n / 2, 2**n * 1.5)."""
+ waits: Final = [sleep_call.args[0] for sleep_call in sleep_mock.call_args_list]
+ assert len(waits) == retries, waits
+ for attempt, wait in enumerate(waits):
+ assert (1 << attempt) * 0.5 <= wait < (1 << attempt) * 1.5, (attempt, wait)
+
+
class TestS3V2UnitTests:
"""Test that S3 v2 integration only uses safe_dumps and not json.dumps"""
@@ -416,8 +424,7 @@ async def test_async_upload_retries_on_s3_503():
# Verify PUT was called twice (retry after 503)
assert logger.async_httpx_client.put.call_count == 2
- # Verify sleep was called with the backoff delay
- mock_sleep.assert_called_once_with(1) # 2**0 = 1s
+ _assert_jittered_backoff(mock_sleep, 1)
@pytest.mark.asyncio
@@ -457,7 +464,7 @@ async def test_async_upload_retries_on_s3_500():
await logger.async_upload_data_to_s3(test_element)
assert logger.async_httpx_client.put.call_count == 2
- mock_sleep.assert_called_once_with(1)
+ _assert_jittered_backoff(mock_sleep, 1)
@pytest.mark.asyncio
@@ -596,17 +603,20 @@ def _generation(request: httpx.Request) -> tuple[str, str]:
@asynccontextmanager
-async def _s3_logger_on_production_handler(profile: str, statuses: list[int]):
+async def _s3_logger_on_production_handler(profile: str, statuses: list[int | type[httpx.TransportError]]):
"""
S3Logger wired to the real AsyncHTTPHandler over an httpx MockTransport that answers with the given
- statuses in order, so the handler's own raise_for_status behaviour is exercised end to end.
+ statuses (or raises the given transport errors) in order, so the handler's own error mapping is exercised.
"""
requests: list[httpx.Request] = []
replies = iter(statuses)
def respond(request: httpx.Request) -> httpx.Response:
requests.append(request)
- return httpx.Response(next(replies), request=request, text="SignatureDoesNotMatch")
+ reply = next(replies)
+ if not isinstance(reply, int):
+ raise reply("S3 attempt failed", request=request)
+ return httpx.Response(reply, request=request, text="SignatureDoesNotMatch")
handler = AsyncHTTPHandler()
handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond))
@@ -663,7 +673,7 @@ async def test_async_upload_retries_403_with_fresh_credentials_and_signature(rot
assert second_key == second_token
assert int(second_key) > int(first_key)
assert requests[1].headers["Authorization"] != requests[0].headers["Authorization"]
- mock_sleep.assert_awaited_once_with(1)
+ _assert_jittered_backoff(mock_sleep, 1)
assert "Error uploading to s3" not in caplog.text
@@ -678,7 +688,7 @@ async def test_async_upload_exhausts_403_retries_through_production_http_handler
await logger.async_upload_data_to_s3(test_element)
assert len(requests) == 3
- assert mock_sleep.await_args_list == [call(1), call(2)]
+ _assert_jittered_backoff(mock_sleep, 2)
assert "Error uploading to s3" in caplog.text
@@ -708,7 +718,7 @@ async def test_async_upload_recovers_from_transient_5xx_through_production_http_
assert all(request.method == "PUT" for request in requests)
assert requests[0].url == requests[1].url
assert requests[0].content == requests[1].content
- mock_sleep.assert_awaited_once_with(1)
+ _assert_jittered_backoff(mock_sleep, 1)
assert "Error uploading to s3" not in caplog.text
@@ -756,10 +766,165 @@ async def test_async_upload_access_denied_403_is_retried_and_then_requeued(rotat
assert outcome is False
assert len(requests) == 3
- assert mock_sleep.await_args_list == [call(1), call(2)]
+ _assert_jittered_backoff(mock_sleep, 2)
assert "Error uploading to s3" in caplog.text
+@pytest.mark.asyncio
+async def test_async_upload_retries_429_through_production_http_handler(rotating_profile: str, caplog):
+ """S3 and S3-compatible stores answer throttling with 429 too, so it backs off and retries like a 503."""
+ async with _s3_logger_on_production_handler(rotating_profile, [429, 200]) as (logger, requests, mock_sleep):
+ uploaded = await logger.async_upload_data_to_s3(_element({"test": "429"}, "429"))
+
+ assert uploaded is True
+ assert len(requests) == 2
+ assert requests[0].content == requests[1].content
+ _assert_jittered_backoff(mock_sleep, 1)
+ assert "Error uploading to s3" not in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_async_upload_retries_a_timed_out_attempt_through_production_http_handler(rotating_profile: str, caplog):
+ """AsyncHTTPHandler.put turns an httpx timeout into litellm.Timeout, which must count as a retryable attempt."""
+ async with _s3_logger_on_production_handler(rotating_profile, [httpx.ReadTimeout, 200]) as (
+ logger,
+ requests,
+ mock_sleep,
+ ):
+ uploaded = await logger.async_upload_data_to_s3(_element({"test": "timeout"}, "timeout"))
+
+ assert uploaded is True
+ assert len(requests) == 2
+ _assert_jittered_backoff(mock_sleep, 1)
+ assert "Error uploading to s3" not in caplog.text
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "error",
+ [
+ pytest.param(httpx.ConnectError("connection refused"), id="connect-error"),
+ pytest.param(httpx.RemoteProtocolError("server disconnected"), id="remote-protocol-error"),
+ pytest.param(httpx.ReadTimeout("read timed out"), id="read-timeout"),
+ ],
+)
+async def test_async_upload_retries_transport_errors_then_succeeds(error: Exception) -> None:
+ logger = S3Logger(
+ s3_bucket_name="test-bucket",
+ s3_aws_access_key_id="test-key",
+ s3_aws_secret_access_key="test-secret",
+ s3_region_name="us-east-1",
+ )
+ logger.async_httpx_client = AsyncMock()
+ logger.async_httpx_client.put = AsyncMock(side_effect=[error, _ok_response()])
+
+ with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
+ uploaded = await logger.async_upload_data_to_s3(_element({"test": "transport"}, "transport"))
+
+ assert uploaded is True
+ assert logger.async_httpx_client.put.call_count == 2
+ _assert_jittered_backoff(mock_sleep, 1)
+
+
+@pytest.mark.asyncio
+async def test_async_upload_exhausts_transport_error_retries_and_keeps_the_event_queued() -> None:
+ logger = S3Logger(
+ s3_bucket_name="test-bucket",
+ s3_aws_access_key_id="test-key",
+ s3_aws_secret_access_key="test-secret",
+ s3_region_name="us-east-1",
+ )
+ logger.async_httpx_client = AsyncMock()
+ logger.async_httpx_client.put = AsyncMock(side_effect=httpx.ConnectError("connection refused"))
+ logger.log_queue = [_element({"test": "down"}, "down")]
+
+ with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
+ await logger.flush_queue()
+
+ assert logger.async_httpx_client.put.call_count == 3
+ _assert_jittered_backoff(mock_sleep, 2)
+ assert [element.s3_object_key for element in logger.log_queue] == ["2025-09-14/test-down.json"]
+
+
+@pytest.mark.asyncio
+async def test_async_upload_sends_every_attempt_with_a_finite_deadline(rotating_profile: str) -> None:
+ """
+ The logger's handler is built with explicit params and no timeout, and the handler forwards that None to
+ httpx as "no timeout", so without a per-attempt deadline a stalled S3 endpoint pins the upload forever.
+ """
+ async with _s3_logger_on_production_handler(rotating_profile, [503, 200]) as (logger, requests, _):
+ await logger.async_upload_data_to_s3(_element({"test": "deadline"}, "deadline"))
+
+ assert len(requests) == 2
+ for request in requests:
+ deadlines = request.extensions["timeout"]
+ assert all(isinstance(limit, (int, float)) and limit > 0 for limit in deadlines.values()), deadlines
+
+
+@pytest.mark.asyncio
+async def test_backoff_is_jittered_so_throttled_uploads_do_not_retry_in_lockstep() -> None:
+ logger = S3Logger(
+ s3_bucket_name="test-bucket",
+ s3_aws_access_key_id="test-key",
+ s3_aws_secret_access_key="test-secret",
+ s3_region_name="us-east-1",
+ )
+ logger.async_httpx_client = AsyncMock()
+ logger.async_httpx_client.put = AsyncMock(
+ side_effect=[_transient_failure_response(), _ok_response(), _transient_failure_response(), _ok_response()]
+ )
+
+ with (
+ patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
+ patch("random.random", side_effect=[0.0, 0.999]),
+ ):
+ await logger.async_upload_data_to_s3(_element({"i": 0}, "0"))
+ await logger.async_upload_data_to_s3(_element({"i": 1}, "1"))
+
+ first, second = (sleep_call.args[0] for sleep_call in mock_sleep.call_args_list)
+ assert first != second
+ assert 0.5 <= first < 1.5 and 0.5 <= second < 1.5
+
+
+@pytest.mark.asyncio
+async def test_dropped_uploads_are_logged_once_per_flush_naming_each_reason(caplog) -> None:
+ logger = S3Logger(
+ s3_bucket_name="test-bucket",
+ s3_aws_access_key_id="test-key",
+ s3_aws_secret_access_key="test-secret",
+ s3_region_name="us-east-1",
+ )
+ terminal = (_element({"i": "terminal-0"}, "terminal-0"), _element({"i": "terminal-1"}, "terminal-1"))
+ expired = _element({"i": "expired"}, "expired").model_copy(update={"retrying_since": _NOW - 7200})
+ fresh = _element({"i": "fresh"}, "fresh")
+
+ async def put(
+ url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None
+ ) -> MagicMock:
+ if "test-terminal-" in url:
+ return _terminal_failure_response()
+ if url.endswith("test-expired.json"):
+ return _transient_failure_response()
+ return _ok_response()
+
+ logger.async_httpx_client = AsyncMock()
+ logger.async_httpx_client.put = put
+ logger.log_queue = [*terminal, expired, fresh]
+
+ with (
+ patch("asyncio.sleep", new_callable=AsyncMock),
+ patch("time.monotonic", return_value=_NOW),
+ ):
+ await logger.flush_queue()
+
+ assert logger.log_queue == []
+ drop_lines = [record.getMessage() for record in caplog.records if "uploads dropped" in record.getMessage()]
+ assert len(drop_lines) == 1, drop_lines
+ assert "3 uploads dropped" in drop_lines[0]
+ assert "2 rejected as terminal" in drop_lines[0]
+ assert "1 retrying longer than s3_max_retry_age_seconds=3600" in drop_lines[0]
+
+
def test_sync_upload_retries_403_with_fresh_signature(rotating_profile: str, monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("AWS_ACCESS_KEY_ID", raising=False)
monkeypatch.delenv("AWS_SECRET_ACCESS_KEY", raising=False)
@@ -794,7 +959,7 @@ def test_sync_upload_retries_403_with_fresh_signature(rotating_profile: str, mon
assert first_key == first_token
assert second_key == second_token
assert int(second_key) > int(first_key)
- mock_sleep.assert_called_once_with(1)
+ _assert_jittered_backoff(mock_sleep, 1)
def test_sync_upload_retries_on_s3_503():
@@ -835,7 +1000,7 @@ def test_sync_upload_retries_on_s3_503():
logger.upload_data_to_s3(test_element)
assert mock_sync_client.put.call_count == 2
- mock_sleep.assert_called_once_with(1)
+ _assert_jittered_backoff(mock_sleep, 1)
@pytest.mark.asyncio
@@ -2866,7 +3031,13 @@ class _CountingPut:
self.peak = 0
self.calls = 0
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.in_flight += 1
self.peak = max(self.peak, self.in_flight)
self.calls += 1
@@ -2882,7 +3053,13 @@ class _RecordingPut:
def __init__(self) -> None:
self.calls: tuple[tuple[str, str | None, dict[str, str] | None], ...] = ()
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls = (*self.calls, (url, data, headers))
return _ok_response()
@@ -2895,7 +3072,13 @@ class _LateAppendingPut:
self.appended = False
self.failed_key: str | None = None
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
if not self.appended:
self.appended = True
self.logger.log_queue.append(self.element)
@@ -2912,7 +3095,13 @@ class _AppendingFailingPut:
self.elements = elements
self.appended = False
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
if not self.appended:
self.appended = True
for element in self.elements:
@@ -2926,7 +3115,13 @@ class _FailOnSuffixPut:
self.suffixes = suffixes
self.calls: tuple[str, ...] = ()
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls = (*self.calls, url)
if self.failing and url.endswith(self.suffixes):
return _transient_failure_response()
@@ -2939,7 +3134,13 @@ class _FailUntilClearedPut:
self.response: Final = _coded_failure_response(status, code, raw_body)
self.calls: tuple[tuple[str, str | None], ...] = ()
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls = (*self.calls, (url, data))
if self.failing:
return self.response
@@ -3357,7 +3558,13 @@ class _FailOnSuffixCodedPut:
self.response: Final = _coded_failure_response(status, code, raw_body)
self.calls: tuple[str, ...] = ()
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls = (*self.calls, url)
if url.endswith(self.suffixes):
return self.response
@@ -3734,7 +3941,13 @@ class _SyncRecordingClient:
self.response: Final = response
self.put_calls: list = [] # mutable-ok: call log appended once per PUT
- def put(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ def put(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.put_calls.append(url)
return self.response
@@ -3760,16 +3973,16 @@ def test_sync_upload_404_is_single_attempt_without_sleep() -> None:
@pytest.mark.parametrize(
- ("status", "expected_puts", "expected_sleeps"),
+ ("status", "expected_puts", "expected_retries"),
[
- pytest.param(429, 1, [], id="429-single"),
- pytest.param(408, 1, [], id="408-single"),
- pytest.param(502, 1, [], id="502-single"),
- pytest.param(504, 1, [], id="504-single"),
- pytest.param(503, 3, [call(1), call(2)], id="503-backoff"),
+ pytest.param(408, 1, 0, id="408-single"),
+ pytest.param(502, 1, 0, id="502-single"),
+ pytest.param(504, 1, 0, id="504-single"),
+ pytest.param(429, 3, 2, id="429-backoff"),
+ pytest.param(503, 3, 2, id="503-backoff"),
],
)
-def test_sync_upload_retry_set_matches_base(status: int, expected_puts: int, expected_sleeps: list) -> None:
+def test_sync_upload_retry_set(status: int, expected_puts: int, expected_retries: int) -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
@@ -3786,7 +3999,50 @@ def test_sync_upload_retry_set_matches_base(status: int, expected_puts: int, exp
logger.upload_data_to_s3(_element({"id": "sync"}, "sync"))
assert len(sync_client.put_calls) == expected_puts
- assert mock_sleep.call_args_list == expected_sleeps
+ _assert_jittered_backoff(mock_sleep, expected_retries)
+
+
+@pytest.mark.parametrize(
+ "failure",
+ [pytest.param(httpx.ConnectError, id="connect-error"), pytest.param(httpx.ReadTimeout, id="read-timeout")],
+)
+def test_sync_upload_retries_a_failed_attempt_through_production_http_handler(
+ failure: type[httpx.TransportError], caplog: pytest.LogCaptureFixture
+) -> None:
+ """HTTPHandler.put turns an httpx timeout into litellm.Timeout, and each attempt carries its own deadline."""
+ requests: Final[list[httpx.Request]] = [] # mutable-ok: the transport records every attempt
+
+ def respond(request: httpx.Request) -> httpx.Response:
+ requests.append(request)
+ if len(requests) == 1:
+ raise failure("S3 attempt failed", request=request)
+ return httpx.Response(200, request=request)
+
+ handler: Final = HTTPHandler()
+ handler.client = httpx.Client(transport=httpx.MockTransport(respond), timeout=None)
+ logger: Final = S3Logger(
+ s3_bucket_name="test-bucket",
+ s3_aws_access_key_id="test-key",
+ s3_aws_secret_access_key="test-secret",
+ s3_region_name="us-east-1",
+ )
+ with (
+ patch( # test-quality-ok: sync upload builds its HTTPHandler per call, there is no injection seam for it
+ "litellm.integrations.s3_v2.get_httpx_client", return_value=handler
+ ),
+ patch("time.sleep") as mock_sleep,
+ ):
+ logger.upload_data_to_s3(_element({"id": "sync-transient"}, "sync-transient"))
+ handler.client.close()
+
+ assert len(requests) == 2
+ assert requests[0].content == requests[1].content
+ assert all(
+ request.extensions["timeout"] == httpx.Timeout(S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS).as_dict()
+ for request in requests
+ )
+ _assert_jittered_backoff(mock_sleep, 1)
+ assert "Error uploading to s3" not in caplog.text
@pytest.mark.asyncio
@@ -3794,6 +4050,7 @@ def test_sync_upload_retry_set_matches_base(status: int, expected_puts: int, exp
("status", "code"),
[
pytest.param(503, "SlowDown", id="503"),
+ pytest.param(429, "SlowDown", id="429"),
pytest.param(500, "InternalError", id="500"),
pytest.param(403, "AccessDenied", id="access-denied-403"),
],
@@ -3816,7 +4073,7 @@ async def test_retryable_statuses_back_off_three_attempts(status: int, code: str
await logger.flush_queue()
assert len(put.calls) == 3
- assert mock_sleep.await_args_list == [call(1), call(2)]
+ _assert_jittered_backoff(mock_sleep, 2)
assert len(logger.log_queue) == 1
@@ -3824,13 +4081,12 @@ async def test_retryable_statuses_back_off_three_attempts(status: int, code: str
@pytest.mark.parametrize(
("status", "code"),
[
- pytest.param(429, "TooManyRequests", id="429"),
pytest.param(408, None, id="408"),
pytest.param(502, None, id="502"),
pytest.param(504, None, id="504"),
],
)
-async def test_non_base_statuses_are_not_retried_in_call(status: int, code: str | None) -> None:
+async def test_unlisted_statuses_are_not_retried_in_call(status: int, code: str | None) -> None:
logger = S3Logger(
s3_bucket_name="test-bucket",
s3_aws_access_key_id="test-key",
@@ -3858,7 +4114,13 @@ class _FirstFailThenOkPut:
self.failed_once = False
self.calls: tuple[str, ...] = ()
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls = (*self.calls, url)
if url.endswith(self.fail_suffix) and not self.failed_once:
self.failed_once = True
@@ -3971,7 +4233,13 @@ async def test_terminal_drop_of_one_element_does_not_drop_a_sibling_with_the_sam
def __init__(self) -> None:
self.calls: tuple[str | None, ...] = ()
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls = (*self.calls, data)
if data is not None and "terminal-marker" in data:
return _terminal_failure_response()
@@ -4321,7 +4589,13 @@ class _FailOncePerKeyPut:
self.failed: set[str] = set()
self.calls: tuple[str, ...] = ()
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls = (*self.calls, url)
if url not in self.failed:
self.failed.add(url)
@@ -4336,7 +4610,13 @@ class _SlowFailOncePerKeyPut:
self.first_completed: int | None = None
self.calls: tuple[str, ...] = ()
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls = (*self.calls, url)
await _real_sleep(0)
if self.first_completed is None:
@@ -4470,7 +4750,13 @@ class _AppendingSuffixFailingPut:
self.fail_suffixes = fail_suffixes
self.appended = False
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
if not self.appended:
self.appended = True
self.logger.log_queue.append(self.element)
@@ -4689,7 +4975,7 @@ def test_sync_upload_retries_access_denied_403(caplog):
logger.upload_data_to_s3(test_element)
assert mock_sync_client.put.call_count == 3
- assert mock_sleep.call_args_list == [call(1), call(2)]
+ _assert_jittered_backoff(mock_sleep, 2)
assert "dropping object" not in caplog.text
@@ -4940,7 +5226,13 @@ class _StatusPut:
self.responses = responses
self.calls = 0
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls += 1
outcome = self.responses[min(self.calls - 1, len(self.responses) - 1)]
if isinstance(outcome, Exception):
@@ -5074,7 +5366,9 @@ async def test_a_slow_put_does_not_lower_the_adaptive_limit() -> None:
logger = _override_logger(s3_adaptive_concurrency=True)
logger.async_httpx_client = AsyncMock()
- async def slow_put(url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def slow_put(
+ url: str, data: str | None = None, headers: dict[str, str] | None = None, timeout: httpx.Timeout | None = None
+ ) -> MagicMock:
await _real_sleep(0)
return _ok_response()
@@ -5091,7 +5385,13 @@ class _FastOkPut:
def __init__(self) -> None:
self.calls = 0
- async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock:
+ async def __call__(
+ self,
+ url: str,
+ data: str | None = None,
+ headers: dict[str, str] | None = None,
+ timeout: httpx.Timeout | None = None,
+ ) -> MagicMock:
self.calls += 1
await _real_sleep(0)
return _ok_response()
From 6bd5580425c4d64dfc9082de286ac179dbda9f27 Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 12:14:04 -0700
Subject: [PATCH 26/47] fix(streaming): add CustomStreamWrapper.close() and
release the connection when a sync stream ends (#46354)
* fix(streaming): close sync streams so the connection goes back to the pool
CustomStreamWrapper gets a sync close() (plus context manager support) that
closes the wrapped provider stream, and both the sync and async wrappers close
the provider stream when the stream reaches its end. The Anthropic stream
iterator keeps the httpx response it reads so close()/aclose() can release the
socket after message_stop instead of leaving it checked out of the pool
Co-authored-by: Blair Jordan
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(streaming): keep the basedpyright ceiling flat in CustomStreamWrapper.close()
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(streaming): close the response on the shared sync stream path too
BaseModelResponseIterator gets a sync close() next to its aclose(), and the
sync branch of BaseLLMHTTPHandler hands it the httpx response the way the
async branch already does, so providers built on it release the socket on
close() instead of waiting for garbage collection
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(streaming): unit tests for CustomStreamWrapper.close() and the shared iterator close()
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(streaming): close the provider stream before re-raising a mid-stream error
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Blair Jordan
---
.../litellm_core_utils/streaming_handler.py | 41 +++-
litellm/llms/anthropic/chat/handler.py | 21 ++
litellm/llms/base_llm/base_model_iterator.py | 4 +
litellm/llms/custom_httpx/llm_http_handler.py | 2 +
tests/integration/_support/wire.py | 10 +
...est_sync_stream_connection_release_wire.py | 214 ++++++++++++++++++
.../test_streaming_handler.py | 48 ++++
.../llms/base_llm/test_base_model_iterator.py | 9 +
8 files changed, 347 insertions(+), 2 deletions(-)
create mode 100644 tests/integration/sdk/test_sync_stream_connection_release_wire.py
diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index a78e6f17cc9..236286a493b 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -1,6 +1,7 @@
import asyncio
import collections.abc
import datetime
+import inspect
import json
import logging
import threading
@@ -234,6 +235,7 @@ class CustomStreamWrapper:
self.completion_stream = completion_stream
self.sent_first_chunk = False
self.sent_last_chunk = False
+ self._closed = False
self._stream_created_time: float = time.time()
litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate(
@@ -328,6 +330,18 @@ class CustomStreamWrapper:
def __aiter__(self) -> AsyncIterator["ModelResponseStream"]:
return self
+ def __enter__(self) -> "CustomStreamWrapper":
+ return self
+
+ def __exit__(self, *exc_info: object) -> None:
+ self.close()
+
+ async def __aenter__(self) -> "CustomStreamWrapper":
+ return self
+
+ async def __aexit__(self, *exc_info: object) -> None:
+ await self.aclose()
+
def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None:
"""Restore trace_id/session_id in the *consuming* thread/task/context.
@@ -392,11 +406,26 @@ class CustomStreamWrapper:
"""
self._restore_consumer_correlation_context(guarded=True)
+ def close(self) -> None:
+ stream_to_close: Final = cast(object, self.completion_stream) # cast-ok: completion_stream is untyped
+ self.completion_stream = None
+ self._closed = True
+ close_stream: Final = getattr(stream_to_close, "close", None)
+ if callable(close_stream):
+ try:
+ result: Final = close_stream()
+ if inspect.iscoroutine(result):
+ result.close()
+ except Exception as e: # noqa: BLE001 # best-effort cleanup; must not raise into the caller
+ verbose_logger.debug("CustomStreamWrapper.close: error closing completion_stream: %s", e)
+ self._restore_consumer_correlation_context()
+
async def aclose(self):
# Restore the consumer's outer context only after the underlying
# provider stream's own close (and its diagnostic logging below, if
# closing fails) completes - not before - so those log lines still
# carry this closing stream's own trace_id/session_id.
+ self._closed = True
if self.completion_stream is not None:
stream_to_close: Final = self.completion_stream
self.completion_stream = None
@@ -1780,6 +1809,8 @@ class CustomStreamWrapper:
response.hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost
def __next__(self) -> "ModelResponseStream":
+ if self._closed:
+ raise StopIteration
cache_hit = False
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
cache_hit = True
@@ -1950,7 +1981,7 @@ class CustomStreamWrapper:
if self.sent_stream_usage is False and self.send_stream_usage is True:
self.sent_stream_usage = True
return response
- self._restore_consumer_correlation_context()
+ self.close()
raise # Re-raise StopIteration
else:
self.sent_last_chunk = True
@@ -1986,6 +2017,7 @@ class CustomStreamWrapper:
traceback_exception: Final = traceback.format_exc()
# LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated
threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start()
+ self.close()
self._handle_stream_fallback_error(e)
def fetch_sync_stream(self):
@@ -2005,6 +2037,8 @@ class CustomStreamWrapper:
return self.completion_stream
async def __anext__(self) -> "ModelResponseStream":
+ if self._closed:
+ raise StopAsyncIteration
cache_hit = False
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
cache_hit = True
@@ -2125,12 +2159,15 @@ class CustomStreamWrapper:
asyncio.create_task(
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
)
+ await self.aclose()
self._handle_stream_fallback_error(e)
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
if self.received_finish_reason is None:
+ await self.aclose()
self._log_stream_failure_and_raise(e)
return await self._finalize_completed_stream(cache_hit=cache_hit)
except Exception as e:
+ await self.aclose()
self._log_stream_failure_and_raise(e)
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
@@ -2224,7 +2261,7 @@ class CustomStreamWrapper:
)
)
- self._restore_consumer_correlation_context()
+ await self.aclose()
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True
diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py
index ecd11279374..dff10473ae7 100644
--- a/litellm/llms/anthropic/chat/handler.py
+++ b/litellm/llms/anthropic/chat/handler.py
@@ -232,6 +232,7 @@ async def make_call(
speed=speed,
tool_name_reverse_map=tool_name_reverse_map,
response_format_schema=response_format_schema,
+ response=response,
)
# LOGGING
@@ -302,6 +303,7 @@ def make_sync_call(
speed=speed,
tool_name_reverse_map=tool_name_reverse_map,
response_format_schema=response_format_schema,
+ response=response,
)
# LOGGING
@@ -675,9 +677,12 @@ class ModelResponseIterator:
speed: str | None = None,
tool_name_reverse_map: dict[str, str] | None = None,
response_format_schema: ResponseFormatSchema | None = None,
+ response: httpx.Response | None = None,
):
self.streaming_response = streaming_response
self.response_iterator = self.streaming_response
+ self.sync_stream = sync_stream
+ self._response = response
self.content_blocks: list[ContentBlockDelta] = []
self.tool_index = -1
self.json_mode = json_mode
@@ -1375,6 +1380,22 @@ class ModelResponseIterator:
return self._handle_accumulated_json_chunk(data_str)
# Sync iterator
+ def close(self) -> None:
+ response: Final = self._response
+ self._response = None
+ if response is not None:
+ response.close()
+
+ async def aclose(self) -> None:
+ response: Final = self._response
+ self._response = None
+ if response is None:
+ return
+ if self.sync_stream:
+ response.close()
+ return
+ await response.aclose()
+
def __iter__(self):
return self
diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py
index 9cb5c2c55dd..aede8aedc60 100644
--- a/litellm/llms/base_llm/base_model_iterator.py
+++ b/litellm/llms/base_llm/base_model_iterator.py
@@ -81,6 +81,10 @@ class BaseModelResponseIterator:
if self.http_response is not None:
await self.http_response.aclose()
+ def close(self) -> None:
+ if self.http_response is not None:
+ self.http_response.close()
+
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream:
return GenericStreamingChunk(
text="",
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index a8b14b83ef3..cfafca4d174 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -1008,6 +1008,8 @@ class BaseLLMHTTPHandler:
json_mode=json_mode,
optional_params=_OPTIONAL_PARAMS.validate_python(optional_params),
)
+ if isinstance(completion_stream, BaseModelResponseIterator):
+ completion_stream.http_response = response
# LOGGING
logging_obj.post_call(
diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py
index e6a166e8457..b261973973c 100644
--- a/tests/integration/_support/wire.py
+++ b/tests/integration/_support/wire.py
@@ -70,6 +70,7 @@ class Wire:
received: SimpleQueue[Request]
disconnected: SimpleQueue[str]
connected: SimpleQueue[str]
+ closed: SimpleQueue[str]
def drain(self) -> tuple[Request, ...]:
return tuple(self.received.get_nowait() for _ in range(self.received.qsize()))
@@ -77,6 +78,9 @@ class Wire:
def connections(self) -> int:
return self.connected.qsize()
+ def open_connections(self) -> int:
+ return self.connected.qsize() - self.closed.qsize()
+
def _await_release(gate: threading.Event, closing: threading.Event) -> bool:
while not closing.is_set():
@@ -98,6 +102,7 @@ def wire_server(
errors: Final[SimpleQueue[Exception]] = SimpleQueue()
disconnected: Final[SimpleQueue[str]] = SimpleQueue()
connected: Final[SimpleQueue[str]] = SimpleQueue()
+ closed: Final[SimpleQueue[str]] = SimpleQueue()
closing: Final = threading.Event()
class Handler(BaseHTTPRequestHandler):
@@ -108,6 +113,10 @@ def wire_server(
super().setup()
connected.put(f"{self.client_address[0]}:{self.client_address[1]}")
+ def finish(self) -> None:
+ super().finish()
+ closed.put(f"{self.client_address[0]}:{self.client_address[1]}")
+
def respond(self) -> None:
headers: Final = {name.lower(): value for name, value in self.headers.items()}
try:
@@ -194,6 +203,7 @@ def wire_server(
received,
disconnected,
connected,
+ closed,
)
finally:
closing.set()
diff --git a/tests/integration/sdk/test_sync_stream_connection_release_wire.py b/tests/integration/sdk/test_sync_stream_connection_release_wire.py
new file mode 100644
index 00000000000..a8b49d285c3
--- /dev/null
+++ b/tests/integration/sdk/test_sync_stream_connection_release_wire.py
@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import threading
+import time
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Final
+
+import openai
+import pytest
+from integration._support.client import eventually
+from integration._support.wire import Reply, Request, Wire, wire_server
+
+import litellm
+from litellm.types.utils import ModelResponseStream
+
+_API_KEY: Final = "synthetic-stream-release-key"
+_PROMPT: Final = "does the socket go back to the pool once the stream is done with it?"
+_FIRST_WORDS: Final = "hello "
+_LAST_WORD: Final = "world"
+_RELEASE_SECONDS: Final = 2
+
+
+def _anthropic_event(event: str, payload: Mapping[str, object]) -> bytes:
+ return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode()
+
+
+def _openai_delta(text: str, finish_reason: str | None) -> bytes:
+ chunk: Final = {
+ "id": "chatcmpl-stream-release-wire",
+ "object": "chat.completion.chunk",
+ "created": 1700000000,
+ "model": "gpt-5.6",
+ "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": finish_reason}],
+ }
+ return b"data: " + json.dumps(chunk).encode() + b"\n\n"
+
+
+_ANTHROPIC_OPENING: Final = b"".join(
+ (
+ _anthropic_event(
+ "message_start",
+ {
+ "type": "message_start",
+ "message": {
+ "id": "msg_stream_release_wire",
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-sonnet-5-5",
+ "content": [],
+ "stop_reason": None,
+ "usage": {"input_tokens": 12, "output_tokens": 1},
+ },
+ },
+ ),
+ _anthropic_event(
+ "content_block_start",
+ {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+ ),
+ _anthropic_event(
+ "content_block_delta",
+ {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _FIRST_WORDS}},
+ ),
+ )
+)
+_ANTHROPIC_CLOSING: Final = (
+ _anthropic_event(
+ "content_block_delta",
+ {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _LAST_WORD}},
+ ),
+ _anthropic_event("content_block_stop", {"type": "content_block_stop", "index": 0}),
+ _anthropic_event(
+ "message_delta",
+ {
+ "type": "message_delta",
+ "delta": {"stop_reason": "end_turn", "stop_sequence": None},
+ "usage": {"output_tokens": 3},
+ },
+ ),
+ _anthropic_event("message_stop", {"type": "message_stop"}),
+)
+_OPENAI_OPENING: Final = _openai_delta(_FIRST_WORDS, None)
+_OPENAI_CLOSING: Final = (_openai_delta(_LAST_WORD, None), _openai_delta("", "stop"), b"data: [DONE]\n\n")
+
+
+@dataclass(frozen=True, slots=True)
+class _Shape:
+ name: str
+ model: str
+ opening: bytes
+ closing: tuple[bytes, ...]
+
+
+_SHAPES: Final = (
+ _Shape("anthropic", "anthropic/claude-sonnet-5-5", _ANTHROPIC_OPENING, _ANTHROPIC_CLOSING),
+ _Shape("openai", "openai/gpt-5.6", _OPENAI_OPENING, _OPENAI_CLOSING),
+ _Shape("databricks", "databricks/databricks-llama", _OPENAI_OPENING, _OPENAI_CLOSING),
+)
+
+
+def _peer(shape: _Shape, gate_after_first: threading.Event | None = None):
+ def respond(request: Request) -> Reply:
+ return Reply(
+ content_type="text/event-stream", chunks=(shape.opening, *shape.closing), gate_after_first=gate_after_first
+ )
+
+ return respond
+
+
+def _request_kwargs(shape: _Shape, wire: Wire) -> Mapping[str, object]:
+ client: Final = openai.OpenAI(api_key=_API_KEY, base_url=wire.url) if shape.name == "openai" else None
+ return {
+ "model": shape.model,
+ "api_base": wire.url,
+ "api_key": _API_KEY,
+ "messages": ({"role": "user", "content": _PROMPT},),
+ "stream": True,
+ **({"client": client} if client is not None else {}),
+ }
+
+
+def _text_of(chunk: ModelResponseStream) -> str:
+ return chunk.choices[0].delta.content or ""
+
+
+def _open_sessions_settling_to(wire: Wire, at_most: int) -> int:
+ return eventually(wire.open_connections, lambda open_sessions: open_sessions <= at_most, seconds=_RELEASE_SECONDS)
+
+
+async def _open_sessions_settling_to_while_the_loop_runs(wire: Wire, at_most: int) -> int:
+ deadline: Final = time.monotonic() + _RELEASE_SECONDS
+ while wire.open_connections() > at_most and time.monotonic() < deadline:
+ await asyncio.sleep(0.05)
+ return wire.open_connections()
+
+
+@pytest.mark.parametrize("shape", _SHAPES, ids=lambda shape: shape.name)
+def test_sync_stream_read_to_its_last_chunk_releases_the_connection(shape: _Shape) -> None:
+ with wire_server(_peer(shape), keep_alive=True) as wire:
+ request: Final = _request_kwargs(shape, wire)
+ first_text: Final = "".join(_text_of(chunk) for chunk in litellm.completion(**request))
+ second_text: Final = "".join(_text_of(chunk) for chunk in litellm.completion(**request))
+ assert (first_text, second_text) == (_FIRST_WORDS + _LAST_WORD, _FIRST_WORDS + _LAST_WORD)
+ assert _open_sessions_settling_to(wire, at_most=1) <= 1, "a finished stream left its TCP session open"
+
+
+@pytest.mark.parametrize("shape", _SHAPES, ids=lambda shape: shape.name)
+def test_sync_stream_closed_after_its_first_chunk_releases_the_connection(shape: _Shape) -> None:
+ gate: Final = threading.Event()
+ with wire_server(_peer(shape, gate_after_first=gate), keep_alive=True) as wire:
+ request: Final = _request_kwargs(shape, wire)
+ abandoned: Final = litellm.completion(**request)
+ first_chunk: Final = next(iter(abandoned))
+ assert _text_of(first_chunk) == _FIRST_WORDS
+ abandoned.close()
+ gate.set()
+ assert _open_sessions_settling_to(wire, at_most=0) == 0, "close() left the TCP session open"
+ gate.clear()
+ follow_up: Final = litellm.completion(**request)
+ assert _text_of(next(iter(follow_up))) == _FIRST_WORDS
+ gate.set()
+ assert "".join(_text_of(chunk) for chunk in follow_up) == _LAST_WORD
+
+
+@pytest.mark.asyncio
+async def test_async_stream_read_to_its_last_chunk_releases_the_connection() -> None:
+ shape: Final = _SHAPES[0]
+ with wire_server(_peer(shape), keep_alive=True) as wire:
+ request: Final = _request_kwargs(shape, wire)
+ first_text: Final = "".join([_text_of(chunk) async for chunk in await litellm.acompletion(**request)])
+ second_text: Final = "".join([_text_of(chunk) async for chunk in await litellm.acompletion(**request)])
+ assert (first_text, second_text) == (_FIRST_WORDS + _LAST_WORD, _FIRST_WORDS + _LAST_WORD)
+ assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=1) <= 1, (
+ "a finished stream left its TCP session open"
+ )
+
+
+@pytest.mark.asyncio
+async def test_async_stream_aclosed_after_its_first_chunk_releases_the_connection() -> None:
+ shape: Final = _SHAPES[0]
+ gate: Final = threading.Event()
+ with wire_server(_peer(shape, gate_after_first=gate), keep_alive=True) as wire:
+ request: Final = _request_kwargs(shape, wire)
+ abandoned: Final = await litellm.acompletion(**request)
+ assert _text_of(await abandoned.__anext__()) == _FIRST_WORDS
+ await abandoned.aclose()
+ gate.set()
+ assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=0) == 0, (
+ "aclose() left the TCP session open"
+ )
+ gate.clear()
+ follow_up: Final = await litellm.acompletion(**request)
+ assert _text_of(await follow_up.__anext__()) == _FIRST_WORDS
+ gate.set()
+ assert "".join([_text_of(chunk) async for chunk in follow_up]) == _LAST_WORD
+
+
+_ANTHROPIC_ERROR_MIDDLE: Final = _anthropic_event(
+ "error", {"type": "error", "error": {"type": "overloaded_error", "message": "peer overloaded mid stream"}}
+)
+
+
+def test_sync_stream_that_raises_mid_stream_releases_the_connection() -> None:
+ shape: Final = _SHAPES[0]
+ broken: Final = _Shape(shape.name, shape.model, shape.opening, (_ANTHROPIC_ERROR_MIDDLE, *shape.closing))
+ with wire_server(_peer(broken), keep_alive=True) as wire:
+ request: Final = _request_kwargs(broken, wire)
+ with pytest.raises(Exception, match="overloaded"):
+ "".join(_text_of(chunk) for chunk in litellm.completion(**request))
+ assert _open_sessions_settling_to(wire, at_most=0) == 0, (
+ "a stream that raised mid-way left its TCP session open"
+ )
diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py
index b60544e6606..0172f0e9aeb 100644
--- a/tests/unit/litellm_core_utils/test_streaming_handler.py
+++ b/tests/unit/litellm_core_utils/test_streaming_handler.py
@@ -6415,3 +6415,51 @@ def validate_last_format(chunk):
assert isinstance(
choice["finish_reason"], str
), "'finish_reason' should be a string."
+
+
+class _RecordingProviderStream:
+ def __init__(self) -> None:
+ self.close_calls = 0
+
+ def __iter__(self) -> "_RecordingProviderStream":
+ return self
+
+ def __next__(self) -> object:
+ raise StopIteration
+
+ def close(self) -> None:
+ self.close_calls += 1
+
+
+def _wrapper_over(stream: _RecordingProviderStream) -> CustomStreamWrapper:
+ return CustomStreamWrapper(
+ completion_stream=stream,
+ model="gpt-5.6",
+ logging_obj=MagicMock(),
+ custom_llm_provider="openai",
+ )
+
+
+def test_close_releases_the_provider_stream_once_and_ends_iteration():
+ stream: Final = _RecordingProviderStream()
+ wrapper: Final = _wrapper_over(stream)
+ wrapper.close()
+ wrapper.close()
+ assert stream.close_calls == 1
+ with pytest.raises(StopIteration):
+ next(wrapper)
+
+
+def test_leaving_a_with_block_closes_the_provider_stream():
+ stream: Final = _RecordingProviderStream()
+ with _wrapper_over(stream):
+ pass
+ assert stream.close_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_leaving_an_async_with_block_closes_the_provider_stream():
+ stream: Final = _RecordingProviderStream()
+ async with _wrapper_over(stream):
+ pass
+ assert stream.close_calls == 1
diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py
index 36747c52416..936bf06bc18 100644
--- a/tests/unit/llms/base_llm/test_base_model_iterator.py
+++ b/tests/unit/llms/base_llm/test_base_model_iterator.py
@@ -7,6 +7,7 @@ import json
from collections.abc import AsyncIterator
from typing import Final
+import httpx
import pytest
from pydantic import BaseModel
@@ -305,3 +306,11 @@ async def test_bytes_chunk_with_text_before_data_prefix_is_parsed(sync_mode: boo
assert chunks[0].id == "chatcmpl-bytes"
assert chunks[0].choices[0].delta.content == "Hello"
assert chunks[1]["is_finished"] is True
+
+
+def test_close_releases_the_attached_http_response():
+ response: Final = httpx.Response(200, stream=httpx.ByteStream(b"data: [DONE]\n\n"))
+ iterator: Final = BaseModelResponseIterator(streaming_response=iter(()), sync_stream=True)
+ iterator.http_response = response
+ iterator.close()
+ assert response.is_closed
From 1d9dd7b06701cb291c28152041c9f369a64be947 Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Wed, 14 Oct 2026 12:27:37 -0700
Subject: [PATCH 27/47] feat(proxy): let proxy admins open soft_budget to team
admins (#46485)
---
.../team_admin_field_permissions.py | 4 +-
.../management/test_team_update.py | 38 +++++++++++++++++++
.../test_team_admin_field_permissions.py | 26 +++++++++++++
.../TeamAdminEditableFieldsSettings.test.tsx | 12 ++++++
.../team/TeamAdminSettingsForm.test.tsx | 18 ++++++++-
.../components/team/TeamAdminSettingsForm.tsx | 14 +++++--
.../src/components/team/TeamInfo.test.tsx | 6 ++-
.../src/components/team/TeamInfo.tsx | 7 +++-
.../team/teamAdminEditAccess.test.ts | 22 +++++++++++
.../components/team/teamAdminEditAccess.ts | 3 +-
10 files changed, 141 insertions(+), 9 deletions(-)
diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py
index 54e303f69b9..67cbb0be45d 100644
--- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py
+++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py
@@ -31,7 +31,9 @@ from litellm.proxy._types import (
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_fields"
# TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field
-SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset({"tpm_limit", "rpm_limit", "max_budget"})
+SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset(
+ {"tpm_limit", "rpm_limit", "max_budget", "soft_budget"}
+)
TEAM_ADMIN_RAISE_MAX_BUDGET_PERMISSION: Final = "raise_max_budget"
TEAM_ADMIN_PROJECTS_PERMISSION: Final = "projects"
TEAM_ADMIN_MEMBER_KEY_BUDGETS_PERMISSION: Final = "member_key_budgets"
diff --git a/tests/integration/management/test_team_update.py b/tests/integration/management/test_team_update.py
index b7b5fe34c87..0414d0bb77f 100644
--- a/tests/integration/management/test_team_update.py
+++ b/tests/integration/management/test_team_update.py
@@ -4,6 +4,7 @@ from litellm.proxy._types import LitellmUserRoles
from tests.integration._support.actors import Actor, SeededKey
from tests.integration._support.actors import create_scratch_actor, create_scratch_team
+from tests.integration._support.client import team_admin_permissions
pytestmark = pytest.mark.asyncio(loop_scope="session")
@@ -161,6 +162,43 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text
+async def test_team_update_team_admin_soft_budget_follows_the_proxy_admin_grant(gateway_client, db, resources, actors):
+ callers = await _callers(db, actors, resources.prefix)
+ await _seed_target(db, actors, callers, "alpha", resources.prefix)
+ headers = {"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"}
+ body = {"team_id": resources.prefix, "soft_budget": 7.5}
+
+ with team_admin_permissions(gateway_client.gateway, ["max_budget"]):
+ denied = await gateway_client.post("/team/update", headers=headers, json=body)
+ assert denied.status_code == 403, denied.text
+ assert "'soft_budget'" in denied.text and "Team admin editable fields" in denied.text, denied.text
+ row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
+ assert row is not None and row.soft_budget != 7.5, "denied but team mutated"
+
+ with team_admin_permissions(gateway_client.gateway, ["soft_budget"]):
+ allowed = await gateway_client.post("/team/update", headers=headers, json=body)
+ assert allowed.status_code == 200, allowed.text
+ row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
+ assert row is not None and row.soft_budget == 7.5
+
+
+async def test_team_update_team_admin_cannot_set_soft_budget_on_a_team_they_do_not_administer(
+ gateway_client, db, resources, actors
+):
+ callers = await _callers(db, actors, resources.prefix)
+ await _seed_target(db, actors, callers, "beta", resources.prefix)
+
+ with team_admin_permissions(gateway_client.gateway, ["soft_budget"]):
+ resp = await gateway_client.post(
+ "/team/update",
+ headers={"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"},
+ json={"team_id": resources.prefix, "soft_budget": 7.5},
+ )
+ assert resp.status_code == 403, resp.text
+ row = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
+ assert row is not None and row.soft_budget != 7.5, "denied but team mutated"
+
+
# Relocation gate — moving a team to a different org. The resources team starts
# in actors.org_a_id; each scenario relocates it to actors.org_b_id. PROXY_ADMIN bypasses;
# ORG_B_ADMIN reaches the handler but holds no role on the source team (403);
diff --git a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py
index 730f4640da8..0c9ed9356f8 100644
--- a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py
+++ b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py
@@ -13,6 +13,7 @@ from litellm.proxy._types import (
UpdateTeamRequest,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
+ SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
TeamAdminEditAllowed,
TeamAdminEditingDisabled,
TeamAdminFieldNotPermitted,
@@ -191,6 +192,31 @@ class TestTeamAdminEditVerdict:
assert verdict == TeamAdminFieldNotPermitted(field="blocked")
+def _enabled(*fields: str) -> frozenset[str]:
+ return resolve_team_admin_editable_fields(
+ {"team_admin_editable_team_fields": list(fields)}, SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
+ )
+
+
+class TestSoftBudgetEdits:
+ def test_enabled_soft_budget_change_is_allowed(self):
+ data: Final = UpdateTeamRequest(team_id="team-1", soft_budget=20.0)
+ verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("soft_budget"))
+ assert isinstance(verdict, TeamAdminEditAllowed)
+ assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "soft_budget": 20.0}
+
+ def test_soft_budget_change_without_it_enabled_names_it(self):
+ data: Final = UpdateTeamRequest(team_id="team-1", soft_budget=20.0)
+ verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("max_budget"))
+ assert verdict == TeamAdminFieldNotPermitted(field="soft_budget")
+
+ def test_unchanged_soft_budget_resend_is_dropped_when_not_enabled(self):
+ data: Final = UpdateTeamRequest(team_id="team-1", tpm_limit=6, soft_budget=10)
+ verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("tpm_limit"))
+ assert isinstance(verdict, TeamAdminEditAllowed)
+ assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6}
+
+
class TestTeamAdminRequestOrRaise:
def test_allowed_hands_back_its_request(self):
request = UpdateTeamRequest(team_id="team-1", tpm_limit=6)
diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx
index ca90a4ac19d..36fbe086df3 100644
--- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx
+++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx
@@ -127,6 +127,18 @@ describe("TeamAdminEditableFieldsSettings", () => {
expect(mutate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["max_budget"] }, expect.anything());
});
+ it("offers the soft budget as a named checkbox and saves it when ticked", async () => {
+ mockSettings(["max_budget", "soft_budget"], []);
+ const mutate = mockSave({});
+
+ renderWithProviders();
+ fireEvent.click(screen.getByRole("checkbox", { name: "Soft Budget (USD)" }));
+ fireEvent.click(saveButton());
+
+ await waitFor(() => expect(mutate).toHaveBeenCalledTimes(1));
+ expect(mutate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["soft_budget"] }, expect.anything());
+ });
+
it("disables Save again when the draft is ticked back to the saved list", () => {
mockSettings(["tpm_limit"], []);
mockSave({});
diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx
index 910760a8ee8..7fe6461e285 100644
--- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx
@@ -13,7 +13,7 @@ const renderForm = (
const onCancel = vi.fn();
renderWithProviders(
{
expect(save).toBeDisabled();
});
+ it("shows the soft budget when the proxy lets team admins edit it", () => {
+ renderForm(new Set(["soft_budget"]));
+
+ expect(screen.getByLabelText("Soft Budget (USD)")).toHaveValue(15);
+ });
+
+ it("saves a lowered soft budget and nothing else", async () => {
+ const user = userEvent.setup();
+ const { onSave } = renderForm(new Set(["soft_budget", "max_budget"]));
+
+ fireEvent.change(screen.getByLabelText("Soft Budget (USD)"), { target: { value: "10.5" } });
+ await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+ await waitFor(() => expect(onSave).toHaveBeenCalledWith({ soft_budget: 10.5 }));
+ });
+
it("closes without saving on cancel", async () => {
const user = userEvent.setup();
const { onSave, onCancel } = renderForm(new Set(["tpm_limit"]));
diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx
index ca80a21f65a..5ce50bb6c72 100644
--- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx
@@ -22,13 +22,21 @@ import {
const numericInputSchema = z.union([z.string(), z.number()]).nullish();
-const teamAdminSettingsSchema = z.object({
+const teamAdminSettingsShape = {
tpm_limit: numericInputSchema,
rpm_limit: numericInputSchema,
max_budget: numericInputSchema,
-});
+ soft_budget: numericInputSchema,
+};
-const INPUT_STEP: Readonly> = { tpm_limit: 1, rpm_limit: 1, max_budget: 0.01 };
+const teamAdminSettingsSchema = z.object(teamAdminSettingsShape);
+
+const INPUT_STEP: Readonly> = {
+ tpm_limit: 1,
+ rpm_limit: 1,
+ max_budget: 0.01,
+ soft_budget: 0.01,
+};
const MAX_BUDGET_KEEP_OR_LOWER_HINT = "You can keep or lower this budget. Ask a proxy admin to raise it.";
const MAX_BUDGET_RAISE_HINT =
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
index 03cee6e5fc1..31d3016d4be 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
@@ -2025,13 +2025,14 @@ describe("TeamInfoView", () => {
expect(toast.error).not.toHaveBeenCalled();
});
- it("prefills the RPM limit and budget a team admin may edit with the team's stored values", async () => {
+ it("prefills the RPM limit and budgets a team admin may edit with the team's stored values", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
rpm_limit: 50,
max_budget: 20,
- caller_edit_access: { kind: "team_admin", editable_fields: ["rpm_limit", "max_budget"] },
+ soft_budget: 15,
+ caller_edit_access: { kind: "team_admin", editable_fields: ["rpm_limit", "max_budget", "soft_budget"] },
}),
);
@@ -2042,6 +2043,7 @@ describe("TeamInfoView", () => {
expect(await screen.findByLabelText("Requests per minute Limit (RPM)")).toHaveValue(50);
expect(screen.getByLabelText("Max Budget (USD)")).toHaveValue(20);
+ expect(screen.getByLabelText("Soft Budget (USD)")).toHaveValue(15);
expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled();
});
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
index 7663ef8c588..cfac81cb64b 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -1306,7 +1306,12 @@ const TeamInfoView: React.FC = ({
const teamAdminSettingsEditor =
teamEditAccess.kind === "team_admin" ? (
{
["tpm_limit", "Tokens per minute Limit (TPM)"],
["rpm_limit", "Requests per minute Limit (RPM)"],
["max_budget", "Max Budget (USD)"],
+ ["soft_budget", "Soft Budget (USD)"],
["raise_max_budget", "Raise the team's max budget"],
["projects", "Create and update projects"],
["create_member_keys", "Create keys for team members"],
@@ -76,6 +77,27 @@ describe("teamAdminSettingsChanges", () => {
});
});
+describe("teamAdminSettingsChanges for the soft budget", () => {
+ const stored = { soft_budget: 50 };
+ const enabled = new Set(["soft_budget"]);
+
+ it("sends a changed soft budget as a number", () => {
+ expect(teamAdminSettingsChanges({ soft_budget: "40" }, stored, enabled)).toStrictEqual({ soft_budget: 40 });
+ });
+
+ it("sends a cleared soft budget as no soft budget", () => {
+ expect(teamAdminSettingsChanges({ soft_budget: "" }, stored, enabled)).toStrictEqual({ soft_budget: null });
+ });
+
+ it("leaves out an unchanged soft budget", () => {
+ expect(teamAdminSettingsChanges({ soft_budget: "50" }, stored, enabled)).toStrictEqual({});
+ });
+
+ it("leaves out a changed soft budget when the proxy did not enable it", () => {
+ expect(teamAdminSettingsChanges({ soft_budget: "40" }, stored, new Set(["max_budget"]))).toStrictEqual({});
+ });
+});
+
describe("parseTeamAdminEditableFields", () => {
it("returns the configured list", () => {
expect(parseTeamAdminEditableFields({ team_admin_editable_team_fields: ["tpm_limit", "rpm_limit"] })).toEqual([
diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
index 1fed045292b..3ab51b4cdbb 100644
--- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
+++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
@@ -48,7 +48,7 @@ export const parseSupportedTeamAdminEditableFields = (uiSettingsFieldSchema: unk
return items.success ? fieldListSchema.parse(items.data.enum) : [];
};
-export const TEAM_ADMIN_SETTINGS_FIELDS = ["tpm_limit", "rpm_limit", "max_budget"] as const;
+export const TEAM_ADMIN_SETTINGS_FIELDS = ["tpm_limit", "rpm_limit", "max_budget", "soft_budget"] as const;
export type TeamAdminSettingsField = (typeof TEAM_ADMIN_SETTINGS_FIELDS)[number];
@@ -56,6 +56,7 @@ const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([
["tpm_limit", "Tokens per minute Limit (TPM)"],
["rpm_limit", "Requests per minute Limit (RPM)"],
["max_budget", "Max Budget (USD)"],
+ ["soft_budget", "Soft Budget (USD)"],
[TEAM_ADMIN_RAISE_MAX_BUDGET_PERMISSION, "Raise the team's max budget"],
["projects", "Create and update projects"],
["member_key_budgets", "Update budgets on team members' keys"],
From 912ad6537678f39fd2506299d5bb7e7004d891c1 Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Wed, 14 Oct 2026 12:30:50 -0700
Subject: [PATCH 28/47] fix(proxy): accept every known logging callback on
/health/services (#46297)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../health_endpoints/_health_endpoints.py | 119 +++++++++------
.../health_endpoints/test_health_endpoints.py | 136 ++++++++++++++++++
2 files changed, 215 insertions(+), 40 deletions(-)
diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py
index e16a391dc12..3c8dd7890fc 100644
--- a/litellm/proxy/health_endpoints/_health_endpoints.py
+++ b/litellm/proxy/health_endpoints/_health_endpoints.py
@@ -6,7 +6,7 @@ import os
import secrets
import time
import traceback
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import Final, Literal, TypedDict
@@ -265,6 +265,63 @@ services = (
)
+_SERVICES_WITH_DEDICATED_CHECK: Final = frozenset(
+ (
+ "slack_budget_alerts",
+ "email",
+ "langfuse",
+ "slack",
+ "ms_teams",
+ "webhook",
+ "datadog",
+ "datadog_metrics",
+ "arize",
+ "galileo",
+ "newrelic",
+ "pointfive",
+ "sqs",
+ "aws_sqs",
+ )
+)
+_SERVICES_TESTABLE_ONLY_VIA_SUCCESS_CALLBACK: Final = frozenset(("custom_callback_api",))
+
+
+def _testable_services() -> frozenset[str]:
+ return (
+ _SERVICES_WITH_DEDICATED_CHECK
+ | _SERVICES_TESTABLE_ONLY_VIA_SUCCESS_CALLBACK
+ | frozenset(
+ litellm._known_custom_logger_compatible_callbacks # pyright: ignore[reportPrivateUsage] # no public alias for the callback Literal's args
+ )
+ )
+
+
+def _is_testable_service(service: str) -> bool:
+ return service in _testable_services()
+
+
+def _callback_matches(service: str, cb: object) -> bool:
+ callback_name: Final = getattr(cb, "callback_name", None)
+ if service in (cb, callback_name, get_callback_identifier(cb)):
+ return True
+ return (
+ not callback_name
+ and service not in _SERVICES_WITH_DEDICATED_CHECK
+ and service in _unnamed_callback_aliases(type(cb))
+ )
+
+
+def _unnamed_callback_aliases(callback_type: type) -> frozenset[str]:
+ """A logger built from a bare config string carries no name, so its class aliases and its package name stand in."""
+ return frozenset(CustomLoggerRegistry.get_all_callback_strs_from_class_type(callback_type)) | frozenset(
+ callback_type.__module__.split(".")
+ )
+
+
+def _service_configured_in(service: str, callbacks: Sequence[object]) -> bool:
+ return any(_callback_matches(service, cb) for cb in callbacks)
+
+
class _ServiceTestErrorDetail(TypedDict):
error: ReadOnly[str]
@@ -326,53 +383,24 @@ async def health_services_endpoint(
if service is None:
raise HTTPException(status_code=400, detail={"error": "Service must be specified."})
- if service not in [
- "slack_budget_alerts",
- "email",
- "langfuse",
- "langfuse_otel",
- "slack",
- "ms_teams",
- "openmeter",
- "webhook",
- "braintrust",
- "otel",
- "custom_callback_api",
- "langsmith",
- "datadog",
- "datadog_metrics",
- "datadog_llm_observability",
- "generic_api",
- "arize",
- "galileo",
- "newrelic",
- "pointfive",
- "signoz",
- "sqs",
- ]:
+ if not _is_testable_service(service):
raise HTTPException(
status_code=400,
- detail={"error": f"Service must be in list. Service={service} not in {services}"},
+ detail={"error": f"Service must be in list. Service={service} not in {sorted(_testable_services())}"},
)
- service_in_success_callbacks = False
- if service in litellm.success_callback:
- service_in_success_callbacks = True
- else:
- for cb in litellm.success_callback:
- if getattr(cb, "callback_name", None) == service:
- service_in_success_callbacks = True
- break
- cb_id = get_callback_identifier(cb)
- if cb_id == service:
- service_in_success_callbacks = True
- break
+ service_in_success_callbacks: Final = _service_configured_in(service, litellm.success_callback)
+ logger_without_dedicated_check_configured: Final = (
+ service not in _SERVICES_WITH_DEDICATED_CHECK
+ and _service_configured_in(service, litellm.logging_callback_manager.get_all_callbacks())
+ )
if (
service == "openmeter"
or service == "braintrust"
or service == "generic_api"
- or (service_in_success_callbacks and service not in ("langfuse", "pointfive"))
+ or (service_in_success_callbacks and service not in ("langfuse", "pointfive", "sqs", "aws_sqs"))
+ or logger_without_dedicated_check_configured
):
_ = await litellm.acompletion(
model="openai/litellm-mock-response-model",
@@ -506,7 +534,11 @@ async def health_services_endpoint(
type="user_budget",
user_info=user_info,
)
- elif service == "sqs":
+ return {
+ "status": "success",
+ "message": "Mock Webhook Alert sent, verify Webhook Alert Received",
+ }
+ elif service in ("sqs", "aws_sqs"):
from litellm.integrations.sqs import SQSLogger
sqs_logger: Final = SQSLogger()
@@ -629,6 +661,13 @@ async def health_services_endpoint(
"message": "Mock Email Alert sent, verify Email Alert Received",
}
+ raise HTTPException(
+ status_code=400,
+ detail={
+ "error": f"Service={service} is not configured as a callback, add it to litellm_settings.callbacks."
+ },
+ )
+
except Exception as e:
verbose_proxy_logger.error("litellm.proxy.proxy_server.health_services_endpoint(): Exception occured - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py
index f215da5cf3e..5d71222c391 100644
--- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py
+++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py
@@ -1429,6 +1429,142 @@ async def test_health_services_endpoint_datadog_llm_observability():
assert "datadog_llm_observability" in result["message"]
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "service", ["zerobus", "langsmith", "otel", "langfuse_otel", "datadog_llm_observability", "signoz"]
+)
+async def test_health_services_endpoint_tests_any_known_callback_configured_via_callbacks(service: str):
+ """A logger without its own health check, wired through litellm_settings.callbacks, gets the mock request."""
+ acompletion: Final = AsyncMock(return_value=MagicMock())
+ with (
+ patch("litellm.success_callback", []),
+ patch("litellm.callbacks", [service]),
+ patch("litellm.acompletion", acompletion),
+ ):
+ result = await health_services_endpoint(service=service)
+
+ assert result == {"status": "success", "message": f"Mock LLM request made - check {service}."}
+ acompletion.assert_awaited_once()
+ assert acompletion.await_args.kwargs["user"] == "litellm:/health/services"
+
+
+@pytest.mark.asyncio
+async def test_health_services_endpoint_tests_logger_configured_only_via_failure_callback():
+ """litellm_settings.failure_callback: [langsmith] is configured too, so Test sends the mock request."""
+ acompletion: Final = AsyncMock(return_value=MagicMock())
+ with (
+ patch("litellm.success_callback", []),
+ patch("litellm.callbacks", []),
+ patch("litellm.failure_callback", ["langsmith"]),
+ patch("litellm.acompletion", acompletion),
+ ):
+ result = await health_services_endpoint(service="langsmith")
+
+ assert result == {"status": "success", "message": "Mock LLM request made - check langsmith."}
+ acompletion.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_health_services_endpoint_matches_otel_instance_loaded_from_callbacks():
+ """The proxy swaps the "otel" string for an OpenTelemetry instance without a callback_name, it still counts."""
+ from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
+
+ acompletion: Final = AsyncMock(return_value=MagicMock())
+ otel_logger: Final = OpenTelemetry(config=OpenTelemetryConfig(exporter="console", skip_set_global=True))
+ with (
+ patch("litellm.success_callback", []),
+ patch("litellm.callbacks", [otel_logger]),
+ patch("litellm.acompletion", acompletion),
+ ):
+ result = await health_services_endpoint(service="otel")
+
+ assert result == {"status": "success", "message": "Mock LLM request made - check otel."}
+ acompletion.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_health_services_endpoint_matches_unregistered_nameless_instance_by_package():
+ """compression_interception is swapped for a logger outside CustomLoggerRegistry at startup, it still counts."""
+ from litellm.integrations.compression_interception.handler import CompressionInterceptionLogger
+
+ acompletion: Final = AsyncMock(return_value=MagicMock())
+ with (
+ patch("litellm.success_callback", []),
+ patch("litellm.callbacks", [CompressionInterceptionLogger()]),
+ patch("litellm.acompletion", acompletion),
+ ):
+ result = await health_services_endpoint(service="compression_interception")
+
+ assert result == {"status": "success", "message": "Mock LLM request made - check compression_interception."}
+ acompletion.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_health_services_endpoint_otel_instance_does_not_claim_arize_dedicated_check():
+ """A nameless OpenTelemetry logger shares arize's class, Test on arize must still run the Arize check."""
+ from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
+
+ acompletion: Final = AsyncMock(return_value=MagicMock())
+ otel_logger: Final = OpenTelemetry(config=OpenTelemetryConfig(exporter="console", skip_set_global=True))
+ with (
+ patch("litellm.success_callback", [otel_logger]),
+ patch("litellm.callbacks", [otel_logger]),
+ patch("litellm.acompletion", acompletion),
+ patch("litellm.integrations.arize.arize.ArizeLogger") as MockArizeLogger,
+ ):
+ mock_instance = MagicMock()
+ mock_instance.async_health_check = AsyncMock(return_value={"status": "healthy", "error_message": ""})
+ MockArizeLogger.return_value = mock_instance
+ result = await health_services_endpoint(service="arize")
+
+ assert result == {"status": "healthy", "message": "Arize is healthy"}
+ acompletion.assert_not_awaited()
+ mock_instance.async_health_check.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("service", ["sqs", "aws_sqs"])
+async def test_health_services_endpoint_keeps_dedicated_check_for_logger_configured_via_callbacks(service: str):
+ """callbacks: [aws_sqs] becomes a nameless SQSLogger at startup, Test on sqs or aws_sqs still runs its check."""
+ from litellm.integrations.sqs import SQSLogger
+
+ acompletion: Final = AsyncMock(return_value=MagicMock())
+ configured_logger: Final = SQSLogger()
+ with (
+ patch("litellm.success_callback", [configured_logger]),
+ patch("litellm.callbacks", [configured_logger]),
+ patch("litellm.acompletion", acompletion),
+ patch("litellm.integrations.sqs.SQSLogger") as MockSQSLogger,
+ ):
+ mock_instance = MagicMock()
+ mock_instance.async_health_check = AsyncMock(return_value={"status": "healthy", "error_message": ""})
+ MockSQSLogger.return_value = mock_instance
+ result = await health_services_endpoint(service=service)
+
+ assert result == {"status": "healthy", "message": ""}
+ acompletion.assert_not_awaited()
+ mock_instance.async_health_check.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_health_services_endpoint_rejects_known_but_unconfigured_callback():
+ """zerobus is a known logger, but with nothing configured the test is a 400, not an empty 200."""
+ from litellm.proxy._types import ProxyException
+
+ with (
+ patch("litellm.success_callback", []),
+ patch("litellm.failure_callback", []),
+ patch("litellm.callbacks", []),
+ patch("litellm._async_success_callback", []),
+ patch("litellm._async_failure_callback", []),
+ pytest.raises(ProxyException) as exc,
+ ):
+ await health_services_endpoint(service="zerobus")
+
+ assert exc.value.code == "400"
+ assert "not configured as a callback" in str(exc.value.message)
+
+
@pytest.mark.asyncio
async def test_health_services_endpoint_rejects_unknown_service():
"""
From 039dea26421e2bd5b04052f2eb18b17c3ffc25d9 Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Wed, 14 Oct 2026 12:30:51 -0700
Subject: [PATCH 29/47] fix(proxy): match unnamed OTel V2 logger to otel on
/health/services (#46394)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../health_endpoints/test_health_endpoints.py | 66 +++++++++++++++++++
1 file changed, 66 insertions(+)
diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py
index 5d71222c391..441c52b6a5c 100644
--- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py
+++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py
@@ -1522,6 +1522,72 @@ async def test_health_services_endpoint_otel_instance_does_not_claim_arize_dedic
mock_instance.async_health_check.assert_awaited_once()
+@pytest.mark.asyncio
+async def test_health_services_endpoint_matches_otel_v2_instance_loaded_from_callbacks():
+ """The proxy swaps "otel" for an OpenTelemetryV2 instance, which has no registry entry, it still counts."""
+ from opentelemetry.sdk.trace import TracerProvider
+
+ from litellm.integrations.otel.logger import OpenTelemetryV2
+ from litellm.integrations.otel.model.config import OpenTelemetryV2Config
+
+ acompletion: Final = AsyncMock(return_value=MagicMock())
+ with (
+ patch("litellm.success_callback", []),
+ patch("litellm.failure_callback", []),
+ patch("litellm.callbacks", []),
+ patch("litellm.service_callback", []),
+ patch("litellm.input_callback", []),
+ patch("litellm._async_success_callback", []),
+ patch("litellm._async_failure_callback", []),
+ patch("litellm.proxy.proxy_server.open_telemetry_logger", MagicMock()),
+ patch("litellm.acompletion", acompletion),
+ ):
+ otel_v2_logger: Final = OpenTelemetryV2(
+ config=OpenTelemetryV2Config(exporter="in_memory"),
+ tracer_provider=TracerProvider(),
+ )
+ litellm.callbacks.append(otel_v2_logger)
+ result = await health_services_endpoint(service="otel")
+
+ assert result == {"status": "success", "message": "Mock LLM request made - check otel."}
+ acompletion.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_health_services_endpoint_rejects_otel_when_only_a_named_otel_v2_preset_is_configured():
+ """An arize-named OpenTelemetryV2 preset is not the plain otel callback, service=otel stays a 400."""
+ from opentelemetry.sdk.trace import TracerProvider
+
+ from litellm.integrations.otel.logger import OpenTelemetryV2
+ from litellm.integrations.otel.model.config import OpenTelemetryV2Config
+ from litellm.proxy._types import ProxyException
+
+ acompletion: Final = AsyncMock(return_value=MagicMock())
+ with (
+ patch("litellm.success_callback", []),
+ patch("litellm.failure_callback", []),
+ patch("litellm.callbacks", []),
+ patch("litellm.service_callback", []),
+ patch("litellm.input_callback", []),
+ patch("litellm._async_success_callback", []),
+ patch("litellm._async_failure_callback", []),
+ patch("litellm.proxy.proxy_server.open_telemetry_logger", MagicMock()),
+ patch("litellm.acompletion", acompletion),
+ ):
+ arize_v2_logger: Final = OpenTelemetryV2(
+ config=OpenTelemetryV2Config(exporter="in_memory"),
+ callback_name="arize",
+ tracer_provider=TracerProvider(),
+ )
+ litellm.callbacks.append(arize_v2_logger)
+ with pytest.raises(ProxyException) as exc:
+ await health_services_endpoint(service="otel")
+
+ assert exc.value.code == "400"
+ assert "not configured as a callback" in str(exc.value.message)
+ acompletion.assert_not_awaited()
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("service", ["sqs", "aws_sqs"])
async def test_health_services_endpoint_keeps_dedicated_check_for_logger_configured_via_callbacks(service: str):
From 2d4472efc1650666c529ac29d9646286ece71af4 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
<158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 14 Oct 2026 12:40:48 -0700
Subject: [PATCH 30/47] fix(vllm): write spend log rows for /vllm passthrough
calls to router models (#45733)
* fix(vllm): write spend log rows for /vllm passthrough calls to router models
The vLLM passthrough config never built a logged response for relayed calls, so stream and non-stream calls through /vllm with a router model wrote no spend row and never charged the key. vLLM now shares the OpenAI-shaped relay logging that Azure already had, moved into an openai_like helper both configs call.
* ci: run the vllm passthrough unit tests in the llms-providers shard
* fix(proxy): log /vllm text completions and ask router models for stream usage
A chat stream relayed to a router model without stream_options now gets
include_usage added on the way out, and the usage chunk vLLM appends is
removed before the client sees it, so the spend row carries vLLM's own
token counts instead of a local estimate. /vllm/completions calls now
write a row too, stream and non-stream.
* fix(vllm): keep usage, fingerprints, and disconnect rows correct on /vllm streams
Inject and strip stream usage on /vllm/completions too, hold every finish
frame for the fingerprint when n > 1, pass a forced-usage server's stream
through untouched, stop turning continuous_usage_stats alone into per-chunk
usage, close the upstream when the client hangs up so the row is written,
and assemble text choices with one sort instead of a scan per choice
* fix(vllm): keep n>1 stream order when holding finish chunks for the fingerprint
A held finish chunk now goes out as soon as any other chunk arrives, so a
choice that finishes while another still streams keeps its place. Only the
finish chunks right before the usage chunk get vLLM's system_fingerprint.
* fix(vllm): keep caller stream_options and log both success handlers
Forward a stream_options the proxy cannot read, or a non-boolean
include_usage, to vLLM as the caller sent it. Only a false
always_include_stream_usage changes the /vllm relay, so clients get the
same stream bytes under every setting.
Record the passthrough origin when logging starts so the sync and async
success handlers both parse the relayed body, hand callbacks a parsed
text completion, and read relayed messages and prompt separately so a
stray field no longer drops the local count.
Add the audit's integration tests for the relay, the settings, and the
Azure spend row.
* test(vllm): expect the caller's stream under always_include_stream_usage true
* fix(vllm): relay long SSE frames in linear time
* test(vllm): time the relay's usage and finish_reason prefilter on one large frame
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
.github/workflows/test-unit.yml | 1 +
litellm/litellm_core_utils/litellm_logging.py | 13 +-
.../llms/azure/passthrough/transformation.py | 86 +-
.../openai_like/passthrough/transformation.py | 249 ++++++
.../llms/vllm/passthrough/transformation.py | 38 +
litellm/proxy/common_request_processing.py | 8 +-
.../injected_stream_usage.py | 210 +++++
.../llm_passthrough_endpoints.py | 34 +-
.../test_azure_passthrough_spend_row.py | 228 +++++
.../test_vllm_passthrough_spend_row.py | 170 ++++
.../test_vllm_passthrough_stream_relay.py | 800 ++++++++++++++++++
..._vllm_passthrough_stream_usage_settings.py | 283 +++++++
.../test_azure_passthrough_transformation.py | 5 +-
tests/unit/llms/vllm/__init__.py | 0
tests/unit/llms/vllm/passthrough/__init__.py | 0
.../test_vllm_passthrough_transformation.py | 241 ++++++
.../test_injected_stream_usage.py | 243 ++++++
.../test_llm_pass_through_endpoints.py | 171 +++-
.../proxy/test_common_request_processing.py | 8 +-
19 files changed, 2673 insertions(+), 115 deletions(-)
create mode 100644 litellm/llms/openai_like/passthrough/transformation.py
create mode 100644 litellm/proxy/pass_through_endpoints/injected_stream_usage.py
create mode 100644 tests/integration/routing/test_azure_passthrough_spend_row.py
create mode 100644 tests/integration/routing/test_vllm_passthrough_spend_row.py
create mode 100644 tests/integration/routing/test_vllm_passthrough_stream_relay.py
create mode 100644 tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py
create mode 100644 tests/unit/llms/vllm/__init__.py
create mode 100644 tests/unit/llms/vllm/passthrough/__init__.py
create mode 100644 tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py
create mode 100644 tests/unit/proxy/pass_through_endpoints/test_injected_stream_usage.py
diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index b020122ee3b..f521b2c8edd 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -350,6 +350,7 @@ jobs:
tests/unit/llms/v0
tests/unit/llms/valkey
tests/unit/llms/vercel_ai_gateway
+ tests/unit/llms/vllm
tests/unit/llms/volcengine
tests/unit/llms/voyage
tests/unit/llms/wandb
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 944d3c57f7d..9b10b95956f 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -651,6 +651,10 @@ class Logging(LiteLLMLoggingBaseClass):
self.stream = stream
self.start_time = start_time # log the call start time
self.call_type = call_type
+ self.relays_llm_passthrough: Final = call_type in (
+ CallTypes.llm_passthrough_route.value,
+ CallTypes.allm_passthrough_route.value,
+ )
self.litellm_call_id = litellm_call_id
self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4())
@@ -2388,10 +2392,7 @@ class Logging(LiteLLMLoggingBaseClass):
service_tier=ws_service_tier,
)
- elif (
- self.call_type == CallTypes.llm_passthrough_route.value
- or self.call_type == CallTypes.allm_passthrough_route.value
- ) and isinstance(result, Response):
+ elif self.relays_llm_passthrough and isinstance(result, Response):
from litellm.utils import ProviderConfigManager
provider_config: Final = ProviderConfigManager.get_provider_passthrough_config(
@@ -2588,7 +2589,9 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result: Final = self.normalize_logging_result(result=result)
- if isinstance(result, Response) and isinstance(logging_result, (ModelResponse, EmbeddingResponse)):
+ if isinstance(result, Response) and isinstance(
+ logging_result, (ModelResponse, EmbeddingResponse, TextCompletionResponse)
+ ):
result = logging_result
if standard_logging_object is None and result is not None and self.stream is not True:
diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py
index 501122783eb..c3d5540aaea 100644
--- a/litellm/llms/azure/passthrough/transformation.py
+++ b/litellm/llms/azure/passthrough/transformation.py
@@ -5,22 +5,18 @@ from typing import TYPE_CHECKING, Final, Optional
import httpx
from httpx import Response
-from pydantic import ValidationError
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.base_llm.passthrough.transformation import (
BasePassthroughConfig,
- RelayShape,
- logged_relay_shape,
replace_path_segment,
strip_leading_model_segment,
)
+from litellm.llms.openai_like.passthrough.transformation import logged_openai_like_response, logged_openai_like_stream
from litellm.secret_managers.main import get_secret_str
-from litellm.types.llms.base import LiteLLMBaseModel
-from litellm.types.llms.openai import AllMessageValues, ResponsesAPIResponse, ResponsesTerminalEvent
+from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
-from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse
if TYPE_CHECKING:
from httpx import URL
@@ -28,42 +24,6 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
-class RelayedChatRequest(LiteLLMBaseModel):
- messages: Sequence[Mapping[str, object]] | None = None
-
-
-class RelayedCallDetails(LiteLLMBaseModel):
- request_data: RelayedChatRequest | None = None
-
-
-def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None:
- try:
- details: Final = RelayedCallDetails.model_validate(litellm_logging_obj.model_call_details)
- except ValidationError:
- return None
- return details.request_data.messages if details.request_data else None
-
-
-RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate)
-
-OPENAI_RELAY_SHAPES: Final = (
- RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate),
- RESPONSES_RELAY_SHAPE,
- RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate),
-)
-
-
-def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None:
- """A streaming logging object assembles the logged response from the terminal event, not from its body."""
- from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
-
- terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks)
- if terminal_event is None:
- return None
- logging_obj.call_type = RESPONSES_RELAY_SHAPE.call_type.value
- return terminal_event
-
-
AZURE_DEPLOYMENT_SEGMENT: Final = re.compile(r"(? Optional["LoggedRelayResponse"]:
- from litellm import encoding
- from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
- from litellm.types.utils import ModelResponse
-
- if "chat/completions" not in endpoint:
- return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint)
-
- openai_chat_config: Final = OpenAIGPTConfig()
-
- litellm_model_response: Final[ModelResponse] = openai_chat_config.transform_response(
- model=model,
- messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}],
- raw_response=httpx_response,
- model_response=ModelResponse(),
- logging_obj=logging_obj,
- optional_params={},
- litellm_params={},
- api_key="",
- request_data=request_data,
- encoding=encoding,
- )
-
- return litellm_model_response
+ return logged_openai_like_response(model, httpx_response, request_data, logging_obj, endpoint)
def handle_logging_collected_chunks(
self,
@@ -215,18 +153,4 @@ class AzurePassthroughConfig(BasePassthroughConfig):
custom_llm_provider: str,
endpoint: str,
) -> Optional["LoggedRelayResponse"]:
- from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
- OpenAIPassthroughLoggingHandler,
- )
-
- if f"/{endpoint.strip('/')}".endswith(RESPONSES_RELAY_SHAPE.path_suffix):
- return logged_responses_stream(all_chunks, litellm_logging_obj)
- if "chat/completions" not in endpoint:
- return None
-
- return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser
- all_chunks=all_chunks,
- litellm_logging_obj=litellm_logging_obj,
- model=model,
- messages=_relayed_messages(litellm_logging_obj),
- )
+ return logged_openai_like_stream(all_chunks, litellm_logging_obj, model, endpoint)
diff --git a/litellm/llms/openai_like/passthrough/transformation.py b/litellm/llms/openai_like/passthrough/transformation.py
new file mode 100644
index 00000000000..df1cebc8519
--- /dev/null
+++ b/litellm/llms/openai_like/passthrough/transformation.py
@@ -0,0 +1,249 @@
+import itertools
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Final, TypeAlias
+
+from httpx import Response
+from pydantic import ConfigDict, ValidationError
+
+import litellm.litellm_core_utils.token_counter as token_counting
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
+from litellm.llms.base_llm.passthrough.transformation import RelayShape, logged_relay_shape
+from litellm.types.llms.base import LiteLLMBaseModel
+from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesTerminalEvent
+from litellm.types.utils import (
+ CallTypes,
+ EmbeddingResponse,
+ ImageResponse,
+ ModelResponse,
+ TextChoices,
+ TextCompletionResponse,
+ Usage,
+)
+
+if TYPE_CHECKING:
+ from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
+
+
+RelayedPrompt: TypeAlias = str | Sequence[str] | Sequence[int] | Sequence[Sequence[int]]
+
+
+def _count_text_tokens(model: str, text: str, count_response_tokens: bool = False) -> int:
+ return token_counting.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's custom_tokenizer param is an untyped dict
+ model=model, text=text, count_response_tokens=count_response_tokens
+ )
+
+
+class RelayedChatRequest(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ messages: Sequence[Mapping[str, object]] | None = None
+
+
+class RelayedTextRequest(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ prompt: RelayedPrompt | None = None
+
+
+class RelayedChatCall(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ request_data: RelayedChatRequest | None = None
+
+
+class RelayedTextCall(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ request_data: RelayedTextRequest | None = None
+
+
+def relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None:
+ try:
+ request: Final = RelayedChatCall.model_validate(litellm_logging_obj.model_call_details).request_data
+ except ValidationError:
+ return None
+ return request.messages if request else None
+
+
+def _relayed_prompt(litellm_logging_obj: Logging) -> RelayedPrompt | None:
+ try:
+ request: Final = RelayedTextCall.model_validate(litellm_logging_obj.model_call_details).request_data
+ except ValidationError:
+ return None
+ return request.prompt if request else None
+
+
+class RelayedTextChoice(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ index: int = 0
+ text: str = ""
+ finish_reason: str | None = None
+
+
+class RelayedTextCompletionChunk(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ id: str
+ created: int
+ model: str
+ choices: Sequence[RelayedTextChoice] = ()
+ usage: Usage | None = None
+
+
+RESPONSES_RELAY_SHAPE: Final = RelayShape("/responses", CallTypes.aresponses, ResponsesAPIResponse.model_validate)
+TEXT_COMPLETION_RELAY_SHAPE: Final = RelayShape(
+ "/completions", CallTypes.atext_completion, TextCompletionResponse.model_validate
+)
+
+OPENAI_RELAY_SHAPES: Final = (
+ RelayShape("/embeddings", CallTypes.aembedding, EmbeddingResponse.model_validate),
+ RESPONSES_RELAY_SHAPE,
+ RelayShape("/images/generations", CallTypes.aimage_generation, ImageResponse.model_validate),
+)
+
+
+def relays_text_completion(endpoint: str) -> bool:
+ relayed_path: Final = f"/{endpoint.strip('/')}"
+ return relayed_path.endswith(TEXT_COMPLETION_RELAY_SHAPE.path_suffix) and not relayed_path.endswith(
+ "/chat/completions"
+ )
+
+
+def logged_responses_stream(all_chunks: Sequence[str], logging_obj: Logging) -> ResponsesTerminalEvent | None:
+ """A streaming logging object assembles the logged response from the terminal event, not from its body."""
+ from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
+
+ terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks=all_chunks)
+ if terminal_event is None:
+ return None
+ logging_obj.call_type = RESPONSES_RELAY_SHAPE.call_type.value
+ return terminal_event
+
+
+def _relayed_text_chunk(line: str) -> RelayedTextCompletionChunk | None:
+ payload: Final = BaseModelResponseIterator.string_to_dict_parser(str_line=line)
+ if not payload:
+ return None
+ try:
+ return RelayedTextCompletionChunk.model_validate(payload)
+ except ValidationError:
+ return None
+
+
+def _assembled_text_choice(index: int, own: Sequence[RelayedTextChoice]) -> RelayedTextChoice:
+ return RelayedTextChoice(
+ index=index,
+ text="".join(delta.text for delta in own),
+ finish_reason=next((delta.finish_reason for delta in reversed(own) if delta.finish_reason), None),
+ )
+
+
+def _delta_index(delta: RelayedTextChoice) -> int:
+ return delta.index
+
+
+def _prompt_part_tokens(model: str, part: str | int | Sequence[int]) -> int:
+ match part:
+ case str():
+ return _count_text_tokens(model=model, text=part)
+ case int():
+ return 1
+ case _:
+ return len(part)
+
+
+def _prompt_tokens(model: str, prompt: RelayedPrompt | None) -> int:
+ match prompt:
+ case None:
+ return 0
+ case str():
+ return _count_text_tokens(model=model, text=prompt)
+ case _:
+ return sum(_prompt_part_tokens(model, part) for part in prompt)
+
+
+def _counted_text_usage(model: str, prompt: RelayedPrompt | None, choices: Sequence[RelayedTextChoice]) -> Usage:
+ prompt_tokens: Final = _prompt_tokens(model, prompt)
+ completion_tokens: Final = sum(
+ _count_text_tokens(model=model, text=choice.text, count_response_tokens=True) for choice in choices
+ )
+ return Usage(
+ prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens
+ )
+
+
+def logged_text_completion_response(
+ httpx_response: Response, logging_obj: Logging, endpoint: str
+) -> "LoggedRelayResponse | None":
+ return logged_relay_shape((TEXT_COMPLETION_RELAY_SHAPE,), httpx_response, logging_obj, endpoint)
+
+
+def logged_text_completion_stream(
+ all_chunks: Sequence[str], logging_obj: Logging, model: str
+) -> TextCompletionResponse | None:
+ chunks: Final = tuple(chunk for chunk in map(_relayed_text_chunk, all_chunks) if chunk is not None)
+ head: Final = next((chunk for chunk in chunks if chunk.choices), None)
+ if head is None:
+ return None
+ deltas: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
+ choices: Final = tuple(
+ _assembled_text_choice(index, tuple(own))
+ for index, own in itertools.groupby(sorted(deltas, key=_delta_index), key=_delta_index)
+ )
+ reported_usage: Final = next((chunk.usage for chunk in reversed(chunks) if chunk.usage is not None), None)
+ prompt: Final = _relayed_prompt(logging_obj)
+ logging_obj.call_type = TEXT_COMPLETION_RELAY_SHAPE.call_type.value # rebind-ok: the relay's logging object starts as the generic passthrough call and learns its real call type from the relayed shape
+ return TextCompletionResponse(
+ id=head.id,
+ created=head.created,
+ model=head.model,
+ choices=[
+ TextChoices(index=choice.index, text=choice.text, finish_reason=choice.finish_reason) for choice in choices
+ ],
+ usage=(reported_usage if reported_usage is not None else _counted_text_usage(model, prompt, choices)),
+ )
+
+
+def logged_openai_like_response(
+ model: str, httpx_response: Response, request_data: Mapping[str, object], logging_obj: Logging, endpoint: str
+) -> "LoggedRelayResponse | None":
+ from litellm import encoding
+ from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
+
+ if "chat/completions" not in endpoint:
+ return logged_relay_shape(OPENAI_RELAY_SHAPES, httpx_response, logging_obj, endpoint)
+
+ return OpenAIGPTConfig().transform_response(
+ model=model,
+ messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}],
+ raw_response=httpx_response,
+ model_response=ModelResponse(),
+ logging_obj=logging_obj,
+ optional_params={},
+ litellm_params={},
+ api_key="",
+ request_data=dict(request_data),
+ encoding=encoding,
+ )
+
+
+def logged_openai_like_stream(
+ all_chunks: Sequence[str], logging_obj: Logging, model: str, endpoint: str
+) -> "LoggedRelayResponse | None":
+ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
+ OpenAIPassthroughLoggingHandler,
+ )
+
+ relayed_path: Final = f"/{endpoint.strip('/')}"
+ if relayed_path.endswith(RESPONSES_RELAY_SHAPE.path_suffix):
+ return logged_responses_stream(all_chunks, logging_obj)
+ if "chat/completions" in endpoint:
+ return OpenAIPassthroughLoggingHandler()._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only OpenAI SSE-to-ModelResponse assembler; reimplementing it would fork the parser
+ all_chunks=all_chunks,
+ litellm_logging_obj=logging_obj,
+ model=model,
+ messages=relayed_messages(logging_obj),
+ )
+ return None
diff --git a/litellm/llms/vllm/passthrough/transformation.py b/litellm/llms/vllm/passthrough/transformation.py
index 71e5923323e..e7da6e54410 100644
--- a/litellm/llms/vllm/passthrough/transformation.py
+++ b/litellm/llms/vllm/passthrough/transformation.py
@@ -1,12 +1,25 @@
+from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final
+from httpx import Response
+
+from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
+from litellm.llms.openai_like.passthrough.transformation import (
+ logged_openai_like_response,
+ logged_openai_like_stream,
+ logged_text_completion_response,
+ logged_text_completion_stream,
+ relays_text_completion,
+)
from ..common_utils import VLLMModelInfo
if TYPE_CHECKING:
from httpx import URL
+ from litellm.llms.base_llm.passthrough.transformation import LoggedRelayResponse
+
class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig):
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
@@ -30,3 +43,28 @@ class VLLMPassthroughConfig(VLLMModelInfo, BasePassthroughConfig):
self.format_url(endpoint, base_target_url, request_query_params),
base_target_url,
)
+
+ def logging_non_streaming_response(
+ self,
+ model: str,
+ custom_llm_provider: str,
+ httpx_response: Response,
+ request_data: Mapping[str, object],
+ logging_obj: Logging,
+ endpoint: str,
+ ) -> "LoggedRelayResponse | None":
+ if relays_text_completion(endpoint):
+ return logged_text_completion_response(httpx_response, logging_obj, endpoint)
+ return logged_openai_like_response(model, httpx_response, request_data, logging_obj, endpoint)
+
+ def handle_logging_collected_chunks(
+ self,
+ all_chunks: Sequence[str],
+ litellm_logging_obj: Logging,
+ model: str,
+ custom_llm_provider: str,
+ endpoint: str,
+ ) -> "LoggedRelayResponse | None":
+ if relays_text_completion(endpoint):
+ return logged_text_completion_stream(all_chunks, litellm_logging_obj, model)
+ return logged_openai_like_stream(all_chunks, litellm_logging_obj, model, endpoint)
diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
index 7ed486922ca..58e8866ce99 100644
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -513,7 +513,7 @@ def _litellm_model_supports_stream_options(litellm_model: str) -> bool:
return supported_params is not None and "stream_options" in supported_params
-def _model_deployments_support_stream_options(
+def model_deployments_support_stream_options(
model: object,
llm_router: Router | None,
team_id: str | None,
@@ -527,7 +527,7 @@ def _model_deployments_support_stream_options(
return all(_litellm_model_supports_stream_options(m) for m in candidate_models)
-def _stream_usage_tracking_updates(
+def stream_usage_tracking_updates(
data: Mapping[str, object],
general_settings: Mapping[str, object],
route_type: str,
@@ -2277,11 +2277,11 @@ class ProxyBaseLLMRequestProcessing:
### AUTO STREAM USAGE TRACKING ###
self.data.update(
- _stream_usage_tracking_updates(
+ stream_usage_tracking_updates(
data=self.data,
general_settings=general_settings,
route_type=route_type,
- supports_stream_options=lambda: _model_deployments_support_stream_options(
+ supports_stream_options=lambda: model_deployments_support_stream_options(
model=self.data.get("model"),
llm_router=llm_router,
team_id=user_api_key_dict.team_id,
diff --git a/litellm/proxy/pass_through_endpoints/injected_stream_usage.py b/litellm/proxy/pass_through_endpoints/injected_stream_usage.py
new file mode 100644
index 00000000000..46a2fc03f73
--- /dev/null
+++ b/litellm/proxy/pass_through_endpoints/injected_stream_usage.py
@@ -0,0 +1,210 @@
+import json
+import re
+from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Mapping, Sequence
+from dataclasses import dataclass, replace
+from itertools import islice
+from types import MappingProxyType
+from typing import Final
+
+from pydantic import ConfigDict, TypeAdapter, ValidationError
+
+from litellm.proxy.common_request_processing import stream_usage_tracking_updates
+from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
+from litellm.types.llms.base import LiteLLMBaseModel
+from litellm.types.utils import CallTypes
+
+_USAGE_FRAME_ENDPOINTS: Final = ("chat/completions", "completions")
+_CALLER_ONLY_STREAM_OPTIONS: Final = frozenset({"continuous_usage_stats"})
+_SSE_FRAME_END: Final = re.compile(rb"\r\n\r\n|\n\n|\r\r")
+_SSE_FRAME_END_LOOKBACK: Final = 3
+_USAGE_OR_FINISH: Final = re.compile(rb'"usage"\s*:\s*\{|"finish_reason"\s*:\s*"')
+_SSE_DATA_FIELD: Final = b"data:"
+_JSON_OBJECT_ADAPTER: Final = TypeAdapter(Mapping[str, object])
+_INJECTION_OFF: Final[Mapping[str, object]] = MappingProxyType({"always_include_stream_usage": False})
+_DEFAULT_INJECTION: Final[Mapping[str, object]] = MappingProxyType({})
+
+
+class _RelayedChoice(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ finish_reason: str | None = None
+
+
+class _RelayedChunk(LiteLLMBaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ choices: Sequence[_RelayedChoice]
+ usage: Mapping[str, object] | None = None
+ system_fingerprint: str | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class StreamUsageRelay:
+ body: Mapping[str, object]
+ strips_usage: bool
+
+
+@dataclass(frozen=True, slots=True)
+class _Stacked:
+ item: bytes
+ below: "_Stacked | None"
+
+
+@dataclass(frozen=True, slots=True)
+class _RelayState:
+ newest_held: _Stacked | None = None
+ passes_usage_through: bool = False
+
+
+def _newest_first(top: _Stacked | None) -> Iterator[bytes]:
+ while top is not None:
+ yield top.item
+ top = top.below # rebind-ok: walks the stack back from the newest item
+
+
+def _oldest_first(top: _Stacked | None) -> tuple[bytes, ...]:
+ return tuple(reversed(tuple(_newest_first(top))))
+
+
+def _held_frames(state: _RelayState) -> tuple[bytes, ...]:
+ return _oldest_first(state.newest_held)
+
+
+def _ends_a_frame(unterminated: _Stacked | None, piece: bytes) -> bool:
+ lookback: Final = b"".join(reversed(tuple(islice(_newest_first(unterminated), _SSE_FRAME_END_LOOKBACK))))
+ return _SSE_FRAME_END.search(lookback[-_SSE_FRAME_END_LOOKBACK:] + piece) is not None
+
+
+def _relays_usage_frames(endpoint: str) -> bool:
+ path: Final = endpoint.strip("/")
+ return any(path == suffix or path.endswith(f"/{suffix}") for suffix in _USAGE_FRAME_ENDPOINTS)
+
+
+def _forwarded_stream_options(stream_options: object, strips_usage: bool) -> object:
+ if not strips_usage or not isinstance(stream_options, Mapping):
+ return stream_options
+ caller_options: Final = _JSON_OBJECT_ADAPTER.validate_python(stream_options)
+ return {key: value for key, value in caller_options.items() if key not in _CALLER_ONLY_STREAM_OPTIONS}
+
+
+def _relay_settings(general_settings: Mapping[str, object]) -> Mapping[str, object]:
+ return _INJECTION_OFF if general_settings.get("always_include_stream_usage") is False else _DEFAULT_INJECTION
+
+
+def _forwards_stream_options_as_sent(request_body: Mapping[str, object]) -> bool:
+ stream_options: Final = request_body.get("stream_options")
+ if stream_options is None:
+ return False
+ if not isinstance(stream_options, Mapping):
+ return True
+ include_usage: Final = _JSON_OBJECT_ADAPTER.validate_python(stream_options).get("include_usage", False)
+ return not isinstance(include_usage, bool)
+
+
+def stream_usage_relay(
+ endpoint: str,
+ request_body: Mapping[str, object],
+ general_settings: Mapping[str, object],
+ supports_stream_options: Callable[[], bool],
+) -> StreamUsageRelay:
+ if not _relays_usage_frames(endpoint) or _forwards_stream_options_as_sent(request_body):
+ return StreamUsageRelay(body=request_body, strips_usage=False)
+ updates: Final = stream_usage_tracking_updates(
+ data=request_body,
+ general_settings=_relay_settings(general_settings),
+ route_type=CallTypes.acompletion.value,
+ supports_stream_options=supports_stream_options,
+ )
+ strips_usage: Final = updates.get("_litellm_strip_stream_usage") is True
+ stream_options: Final = updates.get("stream_options")
+ return StreamUsageRelay(
+ body=(
+ request_body
+ if stream_options is None
+ else {**request_body, "stream_options": _forwarded_stream_options(stream_options, strips_usage)}
+ ),
+ strips_usage=strips_usage,
+ )
+
+
+def _frame_payload(frame: bytes) -> bytes:
+ return b"\n".join(
+ line[len(_SSE_DATA_FIELD) :].strip() for line in frame.splitlines() if line.startswith(_SSE_DATA_FIELD)
+ )
+
+
+def _relayed_chunk(frame: bytes) -> _RelayedChunk | None:
+ if _USAGE_OR_FINISH.search(frame) is None:
+ return None
+ try:
+ return _RelayedChunk.model_validate_json(_frame_payload(frame))
+ except ValidationError:
+ return None
+
+
+def _with_fingerprint(frame: bytes, fingerprint: str | None) -> bytes:
+ if fingerprint is None:
+ return frame
+ payload: Final = _JSON_OBJECT_ADAPTER.validate_json(_frame_payload(frame))
+ terminator: Final = frame[len(frame.rstrip(b"\r\n")) :]
+ restamped: Final = json.dumps(
+ {**payload, "system_fingerprint": fingerprint}, separators=(",", ":"), ensure_ascii=False
+ )
+ return _SSE_DATA_FIELD + b" " + restamped.encode() + terminator
+
+
+def _awaits_fingerprint(chunk: _RelayedChunk | None) -> bool:
+ return (
+ chunk is not None
+ and chunk.usage is None
+ and chunk.system_fingerprint is None
+ and any(choice.finish_reason is not None for choice in chunk.choices)
+ )
+
+
+def _is_usage_only(chunk: _RelayedChunk) -> bool:
+ return chunk.usage is not None and not chunk.choices
+
+
+def _relayed_frame(frame: bytes, state: _RelayState) -> tuple[bytes, _RelayState]:
+ if state.passes_usage_through:
+ return frame, state
+ chunk: Final = _relayed_chunk(frame)
+ if chunk is not None and _is_usage_only(chunk):
+ return (
+ b"".join(_with_fingerprint(held, chunk.system_fingerprint) for held in _held_frames(state)),
+ _RelayState(),
+ )
+ if _awaits_fingerprint(chunk):
+ return b"", replace(state, newest_held=_Stacked(frame, state.newest_held))
+ return (
+ b"".join((*_held_frames(state), frame)),
+ _RelayState(passes_usage_through=chunk is not None and chunk.usage is not None),
+ )
+
+
+def _sse_frames(complete: bytes) -> tuple[bytes, ...]:
+ ends: Final = tuple(match.end() for match in _SSE_FRAME_END.finditer(complete))
+ return tuple(complete[start:end] for start, end in zip((0, *ends[:-1]), ends))
+
+
+async def without_injected_usage_frames(stream: AsyncGenerator[bytes, None]) -> AsyncIterator[bytes]:
+ unterminated: _Stacked | None = None # rebind-ok: stacks reads until one ends a frame, joined once
+ state = _RelayState() # rebind-ok: finish frames wait for the fingerprint vLLM moves onto the usage frame
+ try:
+ async for piece in stream:
+ if not piece:
+ continue
+ if not _ends_a_frame(unterminated, piece):
+ unterminated = _Stacked(piece, unterminated)
+ continue
+ complete, rest = split_complete_sse_frames(b"".join((*_oldest_first(unterminated), piece)))
+ unterminated = _Stacked(rest, None) if rest else None
+ for frame in _sse_frames(complete):
+ outgoing, state = _relayed_frame(frame, state)
+ if outgoing:
+ yield outgoing
+ if state.newest_held is not None or unterminated is not None:
+ yield b"".join((*_held_frames(state), *_oldest_first(unterminated)))
+ finally:
+ await stream.aclose()
diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index 1eefb58d7fa..70cee344c91 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -79,7 +79,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul
user_api_key_auth_websocket,
user_api_key_auth_websocket_for_model,
)
-from litellm.proxy.common_request_processing import open_sse_before_first_byte
+from litellm.proxy.common_request_processing import model_deployments_support_stream_options, open_sse_before_first_byte
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@@ -96,6 +96,10 @@ from litellm.proxy.common_utils.sse_keepalive import (
wrap_passthrough_sse_bytes_with_keepalive_pings,
)
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
+from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
+ stream_usage_relay,
+ without_injected_usage_frames,
+)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
create_pass_through_route,
@@ -509,6 +513,17 @@ async def vllm_proxy_route(
is_router_model: Final = is_passthrough_request_using_router_model(request_body, llm_router)
is_streaming_request: Final = is_passthrough_request_streaming(request_body)
if is_router_model and llm_router:
+ is_json_request: Final = is_json_content_type(request.headers.get("content-type", ""))
+ relay: Final = stream_usage_relay(
+ endpoint=endpoint,
+ request_body=request_body,
+ general_settings=_proxy_general_settings(),
+ supports_stream_options=lambda: model_deployments_support_stream_options(
+ model=request_body.get("model"),
+ llm_router=llm_router,
+ team_id=user_api_key_dict.team_id,
+ ),
+ )
result: Final = cast(
httpx.Response,
await llm_router.allm_passthrough_route(
@@ -521,7 +536,7 @@ async def vllm_proxy_route(
content=None,
data=None,
files=None,
- json=(request_body if is_json_content_type(request.headers.get("content-type", "")) else None),
+ json=(relay.body if is_json_request else None),
params=None,
headers=None,
cookies=None,
@@ -530,8 +545,13 @@ async def vllm_proxy_route(
)
if is_streaming_request:
+ upstream_bytes: Final = _relay_response_bytes(result)
return StreamingResponse(
- content=result.aiter_bytes(),
+ content=(
+ without_injected_usage_frames(upstream_bytes)
+ if relay.strips_usage and is_json_request
+ else upstream_bytes
+ ),
status_code=result.status_code,
headers=HttpPassThroughEndpointHelpers.get_response_headers(
headers=result.headers,
@@ -2058,6 +2078,14 @@ async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> Async
await upstream.aclose()
+async def _relay_response_bytes(upstream: httpx.Response) -> AsyncGenerator[bytes, None]:
+ try:
+ async for chunk in upstream.aiter_bytes():
+ yield chunk
+ finally:
+ await upstream.aclose()
+
+
async def _relay_upstream_response(upstream: httpx.Response) -> Response:
return Response(
content=await upstream.aread(),
diff --git a/tests/integration/routing/test_azure_passthrough_spend_row.py b/tests/integration/routing/test_azure_passthrough_spend_row.py
new file mode 100644
index 00000000000..c2a33aadadb
--- /dev/null
+++ b/tests/integration/routing/test_azure_passthrough_spend_row.py
@@ -0,0 +1,228 @@
+import json
+import uuid
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from typing import Final
+from urllib.parse import urlsplit
+
+import pytest
+from integration._support.client import JSON_OBJECT, Gateway, held
+from integration._support.database import read_rows
+from integration._support.openai_wire import answering_model_discovery, responses_reply
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue
+
+_DEPLOYMENT: Final = "relay-deployment"
+_API_VERSION: Final = "2024-10-21"
+_INPUT_COST: Final = 0.001
+_OUTPUT_COST: Final = 0.002
+_USAGE: Final[dict[str, JsonValue]] = {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
+
+
+def _marker() -> str:
+ return "azure-relay-" + uuid.uuid4().hex
+
+
+def _frame(payload: Mapping[str, JsonValue] | str) -> bytes:
+ return b"data: " + (payload if isinstance(payload, str) else json.dumps(payload)).encode() + b"\n\n"
+
+
+def _chunk(identity: str, choices: Sequence[JsonValue], **extra: JsonValue) -> dict[str, JsonValue]:
+ return {
+ "id": identity,
+ "object": "chat.completion.chunk",
+ "created": 1,
+ "model": "gpt",
+ "choices": list(choices),
+ **extra,
+ }
+
+
+def _chat_stream(identity: str, with_usage: bool) -> tuple[bytes, ...]:
+ return (
+ _frame(_chunk(identity, [{"index": 0, "delta": {"content": "Hello there"}, "finish_reason": None}])),
+ _frame(_chunk(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}])),
+ *((_frame(_chunk(identity, [], usage=_USAGE)),) if with_usage else ()),
+ _frame("[DONE]"),
+ )
+
+
+def _chat_completion(identity: str) -> bytes:
+ return json.dumps(
+ {
+ "id": identity,
+ "object": "chat.completion",
+ "created": 1,
+ "model": "gpt",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}],
+ "usage": _USAGE,
+ }
+ ).encode()
+
+
+def _embedding() -> bytes:
+ return json.dumps(
+ {
+ "object": "list",
+ "data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}],
+ "model": "embed",
+ "usage": {"prompt_tokens": 5, "total_tokens": 5},
+ }
+ ).encode()
+
+
+def _asks_usage(body: Mapping[str, JsonValue]) -> bool:
+ options: Final = body.get("stream_options")
+ return isinstance(options, dict) and options.get("include_usage") is True
+
+
+def _respond(marker: str) -> Callable[[Request], Reply]:
+ identity: Final = f"chatcmpl-{marker}"
+
+ def respond(request: Request) -> Reply:
+ body: Final = JSON_OBJECT.validate_json(request.body)
+ assert marker in json.dumps(body), body
+ path: Final = urlsplit(request.target).path
+ if path.endswith("/embeddings"):
+ return Reply(body=_embedding())
+ if path.endswith("/responses"):
+ return responses_reply(f"resp_{marker}", "gpt", "Hello", stream=body.get("stream") is True)
+ if body.get("stream") is True:
+ return Reply(content_type="text/event-stream", chunks=_chat_stream(identity, _asks_usage(body)))
+ return Reply(body=_chat_completion(identity))
+
+ return answering_model_discovery(respond)
+
+
+@contextmanager
+def _deployment(gateway: Gateway, marker: str, provider: str) -> Iterator[tuple[Wire, str]]:
+ with wire_server(_respond(marker)) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model=f"{provider}/{_DEPLOYMENT}",
+ api_base=wire.url,
+ api_key="sk-azure-relay",
+ api_version=_API_VERSION,
+ input_cost_per_token=_INPUT_COST,
+ output_cost_per_token=_OUTPUT_COST,
+ )
+ yield wire, model
+
+
+def _posted(wire: Wire) -> tuple[tuple[str, dict[str, JsonValue]], ...]:
+ return tuple(
+ (urlsplit(request.target).path, JSON_OBJECT.validate_json(request.body))
+ for request in wire.drain()
+ if request.method == "POST"
+ )
+
+
+def _one_row(model: str) -> dict[str, JsonValue]:
+ rows: Final = held(
+ lambda: read_rows(
+ "SELECT request_id, status, custom_llm_provider, prompt_tokens, completion_tokens, spend "
+ 'FROM "LiteLLM_SpendLogs" WHERE model_group=%s',
+ (model,),
+ ),
+ lambda found: len(found) >= 1,
+ holding=2,
+ seconds=70,
+ )
+ assert len(rows) == 1, rows
+ return rows[0]
+
+
+def _assert_row(row: Mapping[str, JsonValue], provider: str, prompt: int, completion: int) -> None:
+ assert (row["status"], row["custom_llm_provider"]) == ("success", provider), row
+ assert (row["prompt_tokens"], row["completion_tokens"]) == (prompt, completion), row
+ assert row["spend"] == pytest.approx(prompt * _INPUT_COST + completion * _OUTPUT_COST), row
+
+
+def _chat_path(model: str, route: str) -> str:
+ return (
+ f"/azure/openai/deployments/{model}/chat/completions?api-version={_API_VERSION}"
+ if route == "azure"
+ else f"/azure_ai/{model}/chat/completions"
+ )
+
+
+def _chat_body(marker: str, **extra: JsonValue) -> dict[str, JsonValue]:
+ return {"messages": [{"role": "user", "content": marker}], **extra}
+
+
+@pytest.mark.parametrize("provider", ["azure", "azure_ai"])
+def test_router_model_chat_non_stream_writes_the_upstream_usage(gateway: Gateway, provider: str) -> None:
+ marker: Final = _marker()
+ with _deployment(gateway, marker, provider) as (wire, model):
+ response: Final = gateway.request("POST", _chat_path(model, provider), _chat_body(marker))
+ assert response.status_code == 200, response.text
+ assert response.content == _chat_completion(f"chatcmpl-{marker}"), response.text
+ ((path, _),) = _posted(wire)
+ assert path.endswith("/chat/completions"), path
+ row: Final = _one_row(model)
+ assert row["request_id"] == f"chatcmpl-{marker}", row
+ _assert_row(row, provider, 5, 3)
+
+
+@pytest.mark.parametrize("provider", ["azure", "azure_ai"])
+def test_router_model_chat_stream_with_usage_is_relayed_verbatim_and_logged(gateway: Gateway, provider: str) -> None:
+ marker: Final = _marker()
+ with _deployment(gateway, marker, provider) as (wire, model):
+ sent: Final = _chat_body(marker, stream=True, stream_options={"include_usage": True})
+ response: Final = gateway.request("POST", _chat_path(model, provider), sent)
+ assert response.status_code == 200, response.text
+ assert response.content == b"".join(_chat_stream(f"chatcmpl-{marker}", True)), response.text
+ ((_, posted),) = _posted(wire)
+ assert posted == sent, posted
+ _assert_row(_one_row(model), provider, 5, 3)
+
+
+@pytest.mark.parametrize(
+ "extra",
+ [pytest.param({}, id="plain"), pytest.param({"prompt": 5}, id="stray_prompt_field")],
+)
+def test_azure_chat_stream_without_usage_is_forwarded_untouched_and_counted_from_the_messages(
+ gateway: Gateway, extra: Mapping[str, JsonValue]
+) -> None:
+ marker: Final = _marker()
+ with _deployment(gateway, marker, "azure") as (wire, model):
+ sent: Final = _chat_body(marker, stream=True, **extra)
+ response: Final = gateway.request("POST", _chat_path(model, "azure"), sent)
+ assert response.status_code == 200, response.text
+ assert response.content == b"".join(_chat_stream(f"chatcmpl-{marker}", False)), response.text
+ ((_, posted),) = _posted(wire)
+ assert posted == sent, posted
+ row: Final = _one_row(model)
+ prompt, completion = row["prompt_tokens"], row["completion_tokens"]
+ assert row["status"] == "success", row
+ assert isinstance(prompt, int) and prompt > 0, row
+ assert isinstance(completion, int) and completion > 0, row
+ assert row["spend"] == pytest.approx(prompt * _INPUT_COST + completion * _OUTPUT_COST), row
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
+def test_azure_responses_body_model_writes_the_upstream_usage(gateway: Gateway, stream: bool) -> None:
+ marker: Final = _marker()
+ with _deployment(gateway, marker, "azure") as (wire, model):
+ sent: Final = {"model": model, "input": marker, **({"stream": True} if stream else {})}
+ response: Final = gateway.request("POST", f"/azure/openai/responses?api-version={_API_VERSION}", sent)
+ assert response.status_code == 200, response.text
+ assert f"resp_{marker}" in response.text, response.text
+ ((path, _),) = _posted(wire)
+ assert path.endswith("/responses"), path
+ _assert_row(_one_row(model), "azure", 5, 3)
+
+
+def test_azure_embeddings_write_the_upstream_prompt_tokens(gateway: Gateway) -> None:
+ marker: Final = _marker()
+ with _deployment(gateway, marker, "azure") as (wire, model):
+ response: Final = gateway.request(
+ "POST", f"/azure/openai/deployments/{model}/embeddings?api-version={_API_VERSION}", {"input": marker}
+ )
+ assert response.status_code == 200, response.text
+ assert response.content == _embedding(), response.text
+ ((path, _),) = _posted(wire)
+ assert path == f"/openai/deployments/{_DEPLOYMENT}/embeddings", path
+ row: Final = _one_row(model)
+ assert (row["status"], row["custom_llm_provider"]) == ("success", "azure"), row
+ assert (row["prompt_tokens"], row["completion_tokens"]) == (5, 0), row
+ assert row["spend"] == pytest.approx(5 * _INPUT_COST), row
diff --git a/tests/integration/routing/test_vllm_passthrough_spend_row.py b/tests/integration/routing/test_vllm_passthrough_spend_row.py
new file mode 100644
index 00000000000..ac42b84102b
--- /dev/null
+++ b/tests/integration/routing/test_vllm_passthrough_spend_row.py
@@ -0,0 +1,170 @@
+import json
+import uuid
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+import pytest
+from integration._support.anthropic_sse import user_prompt
+from integration._support.client import JSON_OBJECT, Gateway, eventually
+from integration._support.database import read_rows
+from integration._support.openai_wire import answering_model_discovery, chat_reply
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue
+
+_SERVED_MODEL: Final = "Qwen/Qwen2.5-0.5B-Instruct"
+_TEXT: Final = "Hello"
+_INPUT_COST: Final = 0.001
+_OUTPUT_COST: Final = 0.002
+_PROMPT_TOKENS: Final = 5
+_COMPLETION_TOKENS: Final = 3
+_USAGE: Final[dict[str, JsonValue]] = {
+ "prompt_tokens": _PROMPT_TOKENS,
+ "completion_tokens": _COMPLETION_TOKENS,
+ "total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS,
+}
+_FINGERPRINT: Final = "fp-vllm-wire"
+
+
+@dataclass(frozen=True, slots=True)
+class Case:
+ target: str
+ request: Mapping[str, JsonValue]
+ relayed_stream_options: JsonValue | None
+ client_sees_usage: bool
+
+
+_CASES: Final = {
+ "chat_non_stream": Case("/v1/chat/completions", {}, None, True),
+ "chat_stream": Case("/v1/chat/completions", {"stream": True}, {"include_usage": True}, False),
+ "chat_stream_usage_requested": Case(
+ "/v1/chat/completions",
+ {"stream": True, "stream_options": {"include_usage": True}},
+ {"include_usage": True},
+ True,
+ ),
+ "completions_non_stream": Case("/v1/completions", {}, None, True),
+ "chat_stream_continuous_usage_only": Case(
+ "/v1/chat/completions",
+ {"stream": True, "stream_options": {"continuous_usage_stats": True}},
+ {"include_usage": True},
+ False,
+ ),
+ "completions_stream": Case("/v1/completions", {"stream": True}, {"include_usage": True}, False),
+}
+
+
+def _data_frame(frame: Mapping[str, JsonValue]) -> bytes:
+ return b"data: " + json.dumps(frame).encode() + b"\n\n"
+
+
+def _vllm_stream(identity: str, head: Mapping[str, JsonValue], content: JsonValue, finish: JsonValue) -> Reply:
+ chunk: Final = {"id": identity, "created": 1, "model": _SERVED_MODEL, **head}
+ usage_frame: Final = _data_frame({**chunk, "choices": [], "usage": _USAGE, "system_fingerprint": _FINGERPRINT})
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(
+ _data_frame({**chunk, "choices": [content]}),
+ _data_frame({**chunk, "choices": [finish]}),
+ usage_frame + b"data: [DONE]\n\n",
+ ),
+ )
+
+
+def _chat_reply(identity: str, body: Mapping[str, object]) -> Reply:
+ if body.get("stream") is not True:
+ return chat_reply(identity, _SERVED_MODEL, _TEXT, stream=False)
+ return _vllm_stream(
+ identity,
+ {"object": "chat.completion.chunk"},
+ {"index": 0, "delta": {"content": _TEXT}, "finish_reason": None},
+ {"index": 0, "delta": {}, "finish_reason": "stop"},
+ )
+
+
+def _completion_reply(identity: str, body: Mapping[str, object]) -> Reply:
+ if body.get("stream") is True:
+ return _vllm_stream(
+ identity,
+ {"object": "text_completion"},
+ {"index": 0, "text": _TEXT, "finish_reason": None},
+ {"index": 0, "text": "", "finish_reason": "stop"},
+ )
+ completion: Final = {
+ "id": identity,
+ "object": "text_completion",
+ "created": 1,
+ "model": _SERVED_MODEL,
+ "choices": [{"index": 0, "text": _TEXT, "finish_reason": "stop"}],
+ "usage": _USAGE,
+ }
+ return Reply(body=json.dumps(completion).encode())
+
+
+def _vllm_upstream(identity: str, marker: str, case: Case) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ assert (request.method, request.target) == ("POST", case.target), request
+ body: Final = JSON_OBJECT.validate_json(request.body)
+ if case.target.endswith("/chat/completions"):
+ assert user_prompt(body) == marker, body
+ return _chat_reply(identity, body)
+ assert body.get("prompt") == marker, body
+ return _completion_reply(identity, body)
+
+ return answering_model_discovery(respond)
+
+
+def _spend_rows(request_id: str) -> list[dict[str, JsonValue]]:
+ return eventually(
+ lambda: read_rows(
+ "SELECT status, call_type, model, model_group, custom_llm_provider, prompt_tokens, completion_tokens, spend "
+ 'FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+ (request_id,),
+ ),
+ lambda rows: len(rows) >= 1,
+ seconds=70,
+ )
+
+
+@pytest.mark.parametrize("name", list(_CASES))
+def test_vllm_passthrough_to_a_router_model_writes_one_spend_row_with_the_upstream_usage(
+ gateway: Gateway, name: str
+) -> None:
+ case: Final = _CASES[name]
+ marker: Final = "vllm-spend-" + uuid.uuid4().hex
+ identity: Final = f"cmpl-{marker}"
+ prompt: Final[Mapping[str, JsonValue]] = (
+ {"messages": [{"role": "user", "content": marker}]} if "chat" in case.target else {"prompt": marker}
+ )
+ with wire_server(_vllm_upstream(identity, marker, case)) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model=f"hosted_vllm/{_SERVED_MODEL}",
+ api_base=wire.url + "/v1",
+ input_cost_per_token=_INPUT_COST,
+ output_cost_per_token=_OUTPUT_COST,
+ )
+ response: Final = gateway.request(
+ "POST", f"/vllm{case.target[len('/v1') :]}", {"model": model, **prompt, **case.request}
+ )
+ assert response.status_code == 200, response.text
+ assert identity in response.text, response.text
+ assert ('"usage"' in response.text) is case.client_sees_usage, response.text
+ if case.request.get("stream") is True:
+ assert _FINGERPRINT in response.text, response.text
+
+ posted: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert tuple(request.target for request in posted) == (case.target,)
+ assert JSON_OBJECT.validate_json(posted[0].body).get("stream_options") == case.relayed_stream_options
+
+ rows: Final = _spend_rows(identity)
+ assert len(rows) == 1, rows
+ assert rows[0] == {
+ "status": "success",
+ "call_type": "allm_passthrough_route",
+ "model": f"hosted_vllm/{_SERVED_MODEL}",
+ "model_group": model,
+ "custom_llm_provider": "hosted_vllm",
+ "prompt_tokens": _PROMPT_TOKENS,
+ "completion_tokens": _COMPLETION_TOKENS,
+ "spend": pytest.approx(_PROMPT_TOKENS * _INPUT_COST + _COMPLETION_TOKENS * _OUTPUT_COST),
+ }, rows
diff --git a/tests/integration/routing/test_vllm_passthrough_stream_relay.py b/tests/integration/routing/test_vllm_passthrough_stream_relay.py
new file mode 100644
index 00000000000..6ab7469aa63
--- /dev/null
+++ b/tests/integration/routing/test_vllm_passthrough_stream_relay.py
@@ -0,0 +1,800 @@
+import asyncio
+import itertools
+import json
+import re
+import threading
+import time
+import uuid
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Final
+
+import httpx
+import openai
+import psutil
+import pytest
+from integration._support.anthropic_sse import user_prompt
+from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, held
+from integration._support.database import read_rows
+from integration._support.openai_wire import answering_model_discovery, openai_error, responses_reply
+from integration._support.wire import Reply, Request, Wire, wire_server
+from openai.types.chat import ChatCompletionChunk
+from pydantic import JsonValue
+
+from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
+
+_SERVED_MODEL: Final = "Qwen/Qwen2.5-0.5B-Instruct"
+_INPUT_COST: Final = 0.001
+_OUTPUT_COST: Final = 0.002
+_PROMPT_TOKENS: Final = 5
+_COMPLETION_TOKENS: Final = 3
+_USAGE: Final[dict[str, JsonValue]] = {
+ "prompt_tokens": _PROMPT_TOKENS,
+ "completion_tokens": _COMPLETION_TOKENS,
+ "total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS,
+}
+_FINGERPRINT: Final = "fp-vllm-relay"
+_SSE_END: Final = re.compile(rb"\r\n\r\n|\n\n|\r\r")
+_LF: Final = b"\n\n"
+_CRLF: Final = b"\r\n\r\n"
+_ROW_COLUMNS: Final = (
+ "request_id, status, call_type, model_group, custom_llm_provider, prompt_tokens, completion_tokens, spend, team_id"
+)
+
+
+@dataclass(frozen=True, slots=True)
+class _Call:
+ marker: str
+ identity: str
+
+
+def _call() -> _Call:
+ marker: Final = "vllm-relay-" + uuid.uuid4().hex
+ return _Call(marker, f"cmpl-{marker}")
+
+
+def _frame(payload: Mapping[str, JsonValue] | str, end: bytes = _LF) -> bytes:
+ data: Final = payload if isinstance(payload, str) else json.dumps(payload)
+ return b"data: " + data.encode() + end
+
+
+def _chunk(identity: str, choices: Sequence[JsonValue], **extra: JsonValue) -> dict[str, JsonValue]:
+ return {
+ "id": identity,
+ "object": "chat.completion.chunk",
+ "created": 1,
+ "model": _SERVED_MODEL,
+ "choices": list(choices),
+ **extra,
+ }
+
+
+def _delta(index: int, content: str) -> dict[str, JsonValue]:
+ return {"index": index, "delta": {"content": content}, "finish_reason": None}
+
+
+def _finish(index: int) -> dict[str, JsonValue]:
+ return {"index": index, "delta": {}, "finish_reason": "stop"}
+
+
+def _usage_only(identity: str) -> dict[str, JsonValue]:
+ return _chunk(identity, [], usage=_USAGE, system_fingerprint=_FINGERPRINT)
+
+
+def _stamped(chunk: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+ return {**chunk, "system_fingerprint": _FINGERPRINT}
+
+
+def _sse(chunks: Sequence[bytes]) -> Reply:
+ return Reply(content_type="text/event-stream", chunks=tuple(chunks))
+
+
+def _block_payload(block: bytes) -> bytes:
+ return b"\n".join(line[len(b"data:") :].strip() for line in block.splitlines() if line.startswith(b"data:"))
+
+
+def _frames(raw: bytes) -> tuple[JsonValue, ...]:
+ payloads: Final = tuple(_block_payload(block) for block in _SSE_END.split(raw))
+ return tuple("[DONE]" if payload == b"[DONE]" else json.loads(payload) for payload in payloads if payload)
+
+
+def _asks_usage(body: Mapping[str, JsonValue]) -> bool:
+ options: Final = body.get("stream_options")
+ return isinstance(options, dict) and options.get("include_usage") in (True, "true")
+
+
+def _plain_stream(call: _Call, body: Mapping[str, JsonValue]) -> Reply:
+ return _sse(
+ (
+ _frame(_chunk(call.identity, [_delta(0, "Hello")])),
+ _frame(_chunk(call.identity, [_finish(0)])),
+ *((_frame(_usage_only(call.identity)),) if _asks_usage(body) else ()),
+ _frame("[DONE]"),
+ )
+ )
+
+
+def _chat_body(model: str, call: _Call, **extra: JsonValue) -> dict[str, JsonValue]:
+ return {"model": model, "messages": [{"role": "user", "content": call.marker}], **extra}
+
+
+def _scripted(call: _Call, reply: Callable[[Mapping[str, JsonValue]], Reply]) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST", request
+ body: Final = JSON_OBJECT.validate_json(request.body)
+ assert call.marker in json.dumps(body), body
+ return reply(body)
+
+ return answering_model_discovery(respond)
+
+
+def _deploy(scenario: Scenario, wire: Wire) -> str:
+ return scenario.model(
+ model=f"hosted_vllm/{_SERVED_MODEL}",
+ api_base=wire.url + "/v1",
+ input_cost_per_token=_INPUT_COST,
+ output_cost_per_token=_OUTPUT_COST,
+ )
+
+
+@contextmanager
+def _deployment(gateway: Gateway, respond: Callable[[Request], Reply]) -> Iterator[tuple[Wire, str]]:
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ yield wire, _deploy(scenario, wire)
+
+
+def _posted(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
+ return tuple(JSON_OBJECT.validate_json(request.body) for request in wire.drain() if request.method == "POST")
+
+
+def _read_rows(model: str) -> list[dict[str, JsonValue]]:
+ return read_rows(f'SELECT {_ROW_COLUMNS} FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,))
+
+
+def _rows(model: str, count: int = 1) -> list[dict[str, JsonValue]]:
+ return held(lambda: _read_rows(model), lambda rows: len(rows) >= count, holding=2, seconds=70)
+
+
+def _assert_reported_usage(row: Mapping[str, JsonValue]) -> None:
+ assert row["status"] == "success", row
+ assert row["custom_llm_provider"] == "hosted_vllm", row
+ assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _COMPLETION_TOKENS), row
+ assert row["spend"] == pytest.approx(_PROMPT_TOKENS * _INPUT_COST + _COMPLETION_TOKENS * _OUTPUT_COST), row
+
+
+def _assert_counted_usage(row: Mapping[str, JsonValue]) -> None:
+ prompt: Final = row["prompt_tokens"]
+ completion: Final = row["completion_tokens"]
+ assert row["status"] == "success", row
+ assert isinstance(prompt, int) and isinstance(completion, int) and prompt > 0 and completion > 0, row
+ assert row["spend"] == pytest.approx(prompt * _INPUT_COST + completion * _OUTPUT_COST), row
+
+
+def _one_row(model: str, identity: str | None = None) -> dict[str, JsonValue]:
+ rows: Final = _rows(model)
+ assert len(rows) == 1, rows
+ assert identity is None or rows[0]["request_id"] == identity, rows
+ return rows[0]
+
+
+def _stream(gateway: Gateway, path: str, body: Mapping[str, JsonValue], *, key: str | None = None) -> httpx.Response:
+ response: Final = gateway.request("POST", path, body, key=key)
+ assert response.status_code == 200, response.text
+ return response
+
+
+def test_two_choice_stream_keeps_every_frame_in_order_and_stamps_only_the_last_finish(gateway: Gateway) -> None:
+ call: Final = _call()
+ chunks: Final = (
+ _chunk(call.identity, [_delta(0, "A")]),
+ _chunk(call.identity, [_delta(1, "B")]),
+ _chunk(call.identity, [_finish(0)]),
+ _chunk(call.identity, [_delta(1, "C")]),
+ _chunk(call.identity, [_finish(1)]),
+ )
+ reply: Final = _sse((*map(_frame, chunks), _frame(_usage_only(call.identity)), _frame("[DONE]")))
+ with _deployment(gateway, _scripted(call, lambda _: reply)) as (wire, model):
+ response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True, n=2))
+ assert _frames(response.content) == (*chunks[:4], _stamped(chunks[4]), "[DONE]"), response.text
+ (posted,) = _posted(wire)
+ assert posted["stream_options"] == {"include_usage": True}, posted
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+def test_stream_with_usage_on_every_choice_chunk_is_relayed_byte_for_byte(gateway: Gateway) -> None:
+ call: Final = _call()
+ running: Final = tuple(
+ {"prompt_tokens": _PROMPT_TOKENS, "completion_tokens": done, "total_tokens": _PROMPT_TOKENS + done}
+ for done in (1, 2, 3)
+ )
+ upstream_bytes: Final = (
+ _frame(_chunk(call.identity, [_delta(0, "Hel")], usage=running[0])),
+ _frame(_chunk(call.identity, [_delta(0, "lo")], usage=running[1])),
+ _frame(_chunk(call.identity, [_finish(0)], usage=running[2])),
+ _frame(_usage_only(call.identity)),
+ _frame("[DONE]"),
+ )
+ with _deployment(gateway, _scripted(call, lambda _: _sse(upstream_bytes))) as (_, model):
+ response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
+ assert response.content == b"".join(upstream_bytes), response.text
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+def test_stream_split_into_single_bytes_is_reassembled_and_stripped(gateway: Gateway) -> None:
+ call: Final = _call()
+ content: Final = _chunk(call.identity, [_delta(0, "Hello")])
+ finish: Final = _chunk(call.identity, [_finish(0)])
+ whole: Final = b"".join((_frame(content), _frame(finish), _frame(_usage_only(call.identity)), _frame("[DONE]")))
+ reply: Final = _sse(tuple(whole[index : index + 1] for index in range(len(whole))))
+ with _deployment(gateway, _scripted(call, lambda _: reply)) as (_, model):
+ response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
+ assert _frames(response.content) == (content, _stamped(finish), "[DONE]"), response.text
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+def test_crlf_terminated_frames_keep_their_terminators_after_the_usage_frame_is_removed(gateway: Gateway) -> None:
+ call: Final = _call()
+ content: Final = _chunk(call.identity, [_delta(0, "Hello")])
+ finish: Final = _chunk(call.identity, [_finish(0)])
+ reply: Final = _sse(
+ tuple(_frame(payload, _CRLF) for payload in (content, finish, _usage_only(call.identity), "[DONE]"))
+ )
+ with _deployment(gateway, _scripted(call, lambda _: reply)) as (_, model):
+ response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
+ assert _frames(response.content) == (content, _stamped(finish), "[DONE]"), response.text
+ assert response.content.endswith(_CRLF) and response.content.count(_CRLF) == 3, response.content
+ assert _LF not in response.content.replace(_CRLF, b""), response.content
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+def test_stream_whose_upstream_sends_no_usage_frame_is_relayed_whole_and_counted_locally(gateway: Gateway) -> None:
+ call: Final = _call()
+ upstream_bytes: Final = (
+ _frame(_chunk(call.identity, [_delta(0, "Hello there")])),
+ _frame(_chunk(call.identity, [_finish(0)])),
+ _frame("[DONE]"),
+ )
+ with _deployment(gateway, _scripted(call, lambda _: _sse(upstream_bytes))) as (wire, model):
+ response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
+ assert response.content == b"".join(upstream_bytes), response.text
+ (posted,) = _posted(wire)
+ assert posted["stream_options"] == {"include_usage": True}, posted
+ row: Final = _one_row(model, call.identity)
+ _assert_counted_usage(row)
+
+
+def test_upstream_error_on_a_stream_fails_the_call_and_the_model_keeps_serving(gateway: Gateway) -> None:
+ call: Final = _call()
+ replies: Final = iter((openai_error(400), _plain_stream(call, {"stream_options": {"include_usage": True}})))
+ with _deployment(gateway, _scripted(call, lambda _: next(replies))) as (_, model):
+ eventually(lambda: _all_workers_serve(gateway, model), bool, seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS)
+ failed: Final = gateway.request(
+ "POST", "/vllm/chat/completions", _chat_body(model, call, stream=True), headers={"Connection": "close"}
+ )
+ assert failed.status_code >= 400, failed.text
+ assert call.identity not in failed.text, failed.text
+ served: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
+ assert _frames(served.content)[-1] == "[DONE]", served.text
+
+
+_MALFORMED_STREAM_OPTIONS: Final[Mapping[str, JsonValue]] = {
+ "string": "include_usage",
+ "list": ["include_usage"],
+ "integer": 1,
+ "five_kb_string": "u" * 5000,
+}
+
+
+def _rejects_malformed_stream_options(call: _Call) -> Callable[[Mapping[str, JsonValue]], Reply]:
+ def reply(body: Mapping[str, JsonValue]) -> Reply:
+ if isinstance(body.get("stream_options"), dict):
+ return _plain_stream(call, body)
+ return openai_error(400)
+
+ return reply
+
+
+@pytest.mark.parametrize("shape", list(_MALFORMED_STREAM_OPTIONS))
+def test_malformed_stream_options_reach_the_upstream_as_the_caller_sent_them(gateway: Gateway, shape: str) -> None:
+ call: Final = _call()
+ sent: Final = _MALFORMED_STREAM_OPTIONS[shape]
+ with _deployment(gateway, _scripted(call, _rejects_malformed_stream_options(call))) as (wire, model):
+ response: Final = gateway.request(
+ "POST",
+ "/vllm/chat/completions",
+ _chat_body(model, call, stream=True, stream_options=sent),
+ headers={"Connection": "close"},
+ )
+ (posted,) = _posted(wire)
+ assert posted["stream_options"] == sent, posted
+ assert response.status_code >= 400, response.text
+
+
+def test_null_stream_options_get_usage_injected_and_the_usage_frame_removed(gateway: Gateway) -> None:
+ call: Final = _call()
+ with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
+ response: Final = _stream(
+ gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True, stream_options=None)
+ )
+ (posted,) = _posted(wire)
+ assert posted["stream_options"] == {"include_usage": True}, posted
+ assert _frames(response.content) == (
+ _chunk(call.identity, [_delta(0, "Hello")]),
+ _stamped(_chunk(call.identity, [_finish(0)])),
+ "[DONE]",
+ ), response.text
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+@pytest.mark.parametrize(
+ "sent",
+ [pytest.param({}, id="empty"), pytest.param({"include_usage": False}, id="include_usage_false")],
+)
+def test_stream_options_without_include_usage_true_get_usage_injected_and_the_usage_frame_removed(
+ gateway: Gateway, sent: Mapping[str, JsonValue]
+) -> None:
+ call: Final = _call()
+ with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
+ response: Final = _stream(
+ gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True, stream_options=dict(sent))
+ )
+ (posted,) = _posted(wire)
+ assert posted["stream_options"] == {"include_usage": True}, posted
+ assert _frames(response.content) == (
+ _chunk(call.identity, [_delta(0, "Hello")]),
+ _stamped(_chunk(call.identity, [_finish(0)])),
+ "[DONE]",
+ ), response.text
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+def test_null_include_usage_reaches_the_upstream_as_sent_and_is_counted_locally(gateway: Gateway) -> None:
+ call: Final = _call()
+ sent: Final[dict[str, JsonValue]] = {"include_usage": None}
+ with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
+ response: Final = _stream(
+ gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True, stream_options=sent)
+ )
+ (posted,) = _posted(wire)
+ assert posted["stream_options"] == sent, posted
+ assert response.content == b"".join(_plain_stream(call, posted).chunks or ()), response.text
+ _assert_counted_usage(_one_row(model, call.identity))
+
+
+def test_duplicated_stream_options_key_takes_the_last_value(gateway: Gateway) -> None:
+ call: Final = _call()
+ raw: Final = (
+ '{"model": "%s", "messages": [{"role": "user", "content": "%s"}], "stream": true, '
+ '"stream_options": {"include_usage": true}, "stream_options": {"continuous_usage_stats": true}}'
+ )
+ with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
+ response: Final = gateway.client.post(
+ "/vllm/chat/completions",
+ content=(raw % (model, call.marker)).encode(),
+ headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
+ )
+ assert response.status_code == 200, response.text
+ (posted,) = _posted(wire)
+ assert posted["stream_options"] == {"include_usage": True}, posted
+ assert '"usage"' not in response.text, response.text
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+def test_include_usage_sent_as_the_string_true_still_reaches_the_caller(gateway: Gateway) -> None:
+ call: Final = _call()
+ with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (_, model):
+ response: Final = _stream(
+ gateway,
+ "/vllm/chat/completions",
+ _chat_body(model, call, stream=True, stream_options={"include_usage": "true"}),
+ )
+ assert _usage_only(call.identity) in _frames(response.content), response.text
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+def test_stream_flag_sent_as_a_string_streams_verbatim_and_is_counted_locally(gateway: Gateway) -> None:
+ call: Final = _call()
+ with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (wire, model):
+ response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream="true"))
+ (posted,) = _posted(wire)
+ assert "stream_options" not in posted, posted
+ assert response.content == b"".join(_plain_stream(call, posted).chunks or ()), response.text
+ _assert_counted_usage(_one_row(model, call.identity))
+
+
+def _upstream_socket_states(wire: Wire) -> frozenset[str]:
+ port: Final = httpx.URL(wire.url).port
+ return frozenset(
+ connection.status
+ for connection in psutil.Process().net_connections(kind="tcp")
+ if connection.laddr.port == port
+ )
+
+
+def test_caller_disconnect_mid_stream_closes_the_upstream_and_logs_what_was_relayed(gateway: Gateway) -> None:
+ call: Final = _call()
+ release: Final = threading.Event()
+ reply: Final = Reply(
+ content_type="text/event-stream",
+ chunks=(
+ _frame(_chunk(call.identity, [_delta(0, "Hello there")])),
+ _frame(_chunk(call.identity, [_finish(0)])),
+ _frame(_usage_only(call.identity)) + _frame("[DONE]"),
+ ),
+ gate_after_first=release,
+ )
+ with _deployment(gateway, _scripted(call, lambda _: reply)) as (wire, model):
+ try:
+ with gateway.client.stream(
+ "POST",
+ "/vllm/chat/completions",
+ json=_chat_body(model, call, stream=True),
+ headers={"Authorization": f"Bearer {gateway.key}"},
+ ) as response:
+ assert response.status_code == 200
+ assert call.identity.encode() in next(response.iter_raw())
+ eventually(
+ lambda: _upstream_socket_states(wire),
+ lambda states: psutil.CONN_CLOSE_WAIT in states and psutil.CONN_ESTABLISHED not in states,
+ seconds=30,
+ )
+ finally:
+ release.set()
+ _assert_counted_usage(_one_row(model, call.identity))
+
+
+def _embedding_reply(_: Mapping[str, JsonValue]) -> Reply:
+ body: Final = {
+ "object": "list",
+ "data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5, 0.75]}],
+ "model": _SERVED_MODEL,
+ "usage": {"prompt_tokens": _PROMPT_TOKENS, "total_tokens": _PROMPT_TOKENS},
+ }
+ return Reply(body=json.dumps(body).encode())
+
+
+def test_embeddings_write_one_row_with_the_upstream_prompt_tokens(gateway: Gateway) -> None:
+ call: Final = _call()
+ with _deployment(gateway, _scripted(call, _embedding_reply)) as (wire, model):
+ response: Final = gateway.request("POST", "/vllm/embeddings", {"model": model, "input": call.marker})
+ assert response.status_code == 200, response.text
+ assert JSON_OBJECT.validate_json(response.content)["data"] == [
+ {"object": "embedding", "index": 0, "embedding": [0.25, 0.5, 0.75]}
+ ], response.text
+ assert tuple(request.target for request in wire.drain() if request.method == "POST") == ("/v1/embeddings",)
+ row: Final = _one_row(model)
+ assert row["status"] == "success", row
+ assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, 0), row
+ assert row["spend"] == pytest.approx(_PROMPT_TOKENS * _INPUT_COST), row
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"])
+def test_responses_write_one_row_and_relay_the_body_untouched(gateway: Gateway, stream: bool) -> None:
+ call: Final = _call()
+ reply: Final = responses_reply(f"resp_{call.marker}", _SERVED_MODEL, "Hello", stream=stream)
+ with _deployment(gateway, _scripted(call, lambda _: reply)) as (wire, model):
+ sent: Final = {"model": model, "input": call.marker, **({"stream": True} if stream else {})}
+ response: Final = gateway.request("POST", "/vllm/responses", sent)
+ assert response.status_code == 200, response.text
+ assert response.content == (b"".join(reply.chunks) if reply.chunks else reply.body), response.text
+ (posted,) = _posted(wire)
+ assert {name: value for name, value in posted.items() if name != "model"} == {
+ name: value for name, value in sent.items() if name != "model"
+ }, posted
+ _assert_reported_usage(_one_row(model))
+
+
+def _pinned(gateway: Gateway) -> httpx.Client:
+ return httpx.Client(
+ base_url=gateway.client.base_url,
+ timeout=30,
+ trust_env=False,
+ limits=httpx.Limits(max_connections=1, max_keepalive_connections=1, keepalive_expiry=300),
+ )
+
+
+def _chunk_text(chunk: ChatCompletionChunk) -> str:
+ return "".join(choice.delta.content or "" for choice in chunk.choices)
+
+
+def _listed(response: httpx.Response, model: str) -> bool:
+ return response.status_code == 200 and model in response.text
+
+
+def test_openai_sdk_stream_assembles_the_answer_without_a_usage_chunk(gateway: Gateway) -> None:
+ call: Final = _call()
+ with (
+ _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (_, model),
+ _pinned(gateway) as http,
+ ):
+ auth: Final = {"Authorization": f"Bearer {gateway.key}"}
+ eventually(
+ lambda: http.get("/v1/models", headers=auth),
+ lambda listed: _listed(listed, model),
+ seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
+ )
+ client: Final = openai.OpenAI(
+ base_url=f"{str(gateway.client.base_url).rstrip('/')}/vllm",
+ api_key=gateway.key,
+ http_client=http,
+ max_retries=0,
+ )
+ chunks: Final = tuple(
+ client.chat.completions.create(
+ model=model, messages=[{"role": "user", "content": call.marker}], stream=True
+ )
+ )
+ assert "".join(map(_chunk_text, chunks)) == "Hello", chunks
+ assert [chunk.usage for chunk in chunks if chunk.usage is not None] == [], chunks
+ assert {chunk.id for chunk in chunks} == {call.identity}, chunks
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+async def _async_listed(http: httpx.AsyncClient, key: str, model: str) -> httpx.Response:
+ return await http.get("/v1/models", headers={"Authorization": f"Bearer {key}"})
+
+
+async def test_async_openai_sdk_stream_assembles_the_answer_without_a_usage_chunk(gateway: Gateway) -> None:
+ call: Final = _call()
+ loop: Final = asyncio.get_running_loop()
+ with _deployment(gateway, _scripted(call, lambda body: _plain_stream(call, body))) as (_, model):
+ async with httpx.AsyncClient(
+ base_url=gateway.client.base_url,
+ timeout=30,
+ trust_env=False,
+ limits=httpx.Limits(max_connections=1, max_keepalive_connections=1, keepalive_expiry=300),
+ ) as http:
+ await asyncio.to_thread(
+ eventually,
+ lambda: asyncio.run_coroutine_threadsafe(_async_listed(http, gateway.key, model), loop).result(),
+ lambda listed: _listed(listed, model),
+ 2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
+ )
+ client: Final = openai.AsyncOpenAI(
+ base_url=f"{str(gateway.client.base_url).rstrip('/')}/vllm",
+ api_key=gateway.key,
+ http_client=http,
+ max_retries=0,
+ )
+ stream: Final = await client.chat.completions.create(
+ model=model, messages=[{"role": "user", "content": call.marker}], stream=True
+ )
+ chunks: Final = tuple([chunk async for chunk in stream])
+ assert "".join(map(_chunk_text, chunks)) == "Hello", chunks
+ assert [chunk.usage for chunk in chunks if chunk.usage is not None] == [], chunks
+ assert {chunk.id for chunk in chunks} == {call.identity}, chunks
+ _assert_reported_usage(_one_row(model, call.identity))
+
+
+def test_identical_requests_repeated_each_reach_the_upstream_and_each_write_one_row(gateway: Gateway) -> None:
+ call: Final = _call()
+ served: Final = itertools.count()
+
+ def reply(body: Mapping[str, JsonValue]) -> Reply:
+ repeat: Final = _Call(call.marker, f"{call.identity}-{next(served)}")
+ if body.get("stream") is True:
+ return _plain_stream(repeat, body)
+ completion: Final = {
+ "id": repeat.identity,
+ "object": "chat.completion",
+ "created": 1,
+ "model": _SERVED_MODEL,
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello"}, "finish_reason": "stop"}],
+ "usage": _USAGE,
+ }
+ return Reply(body=json.dumps(completion).encode())
+
+ with _deployment(gateway, _scripted(call, reply)) as (wire, model):
+ bodies: Final = (_chat_body(model, call), _chat_body(model, call), *(_chat_body(model, call, stream=True),) * 2)
+ texts: Final = tuple(_stream(gateway, "/vllm/chat/completions", body).text for body in bodies)
+ identities: Final = tuple(f"{call.identity}-{index}" for index in range(4))
+ assert all(identity in text for identity, text in zip(identities, texts)), texts
+ assert len(_posted(wire)) == 4
+ rows: Final = _rows(model, 4)
+ assert sorted(str(row["request_id"]) for row in rows) == sorted(identities), rows
+ for row in rows:
+ _assert_reported_usage(row)
+
+
+def test_non_json_content_type_never_reaches_the_upstream(gateway: Gateway) -> None:
+ call: Final = _call()
+ with (
+ wire_server(answering_model_discovery(lambda _: _plain_stream(call, {}))) as wire,
+ gateway.scenario() as scenario,
+ ):
+ model: Final = _deploy(scenario, wire)
+ response: Final = gateway.client.post(
+ "/vllm/chat/completions",
+ content=json.dumps(_chat_body(model, call, stream=True)).encode(),
+ headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "text/plain", "Connection": "close"},
+ )
+ posted: Final = tuple(request.body for request in wire.drain() if request.method == "POST")
+ assert response.status_code >= 400, response.text
+ assert posted == (), posted
+
+
+def test_team_key_stream_writes_the_row_under_the_team(gateway: Gateway) -> None:
+ call: Final = _call()
+ with wire_server(_scripted(call, lambda body: _plain_stream(call, body))) as wire, gateway.scenario() as scenario:
+ model: Final = _deploy(scenario, wire)
+ team: Final = scenario.team()
+ key: Final = scenario.key(team_id=team)
+ response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True), key=key)
+ assert '"usage"' not in response.text, response.text
+ row: Final = _one_row(model, call.identity)
+ assert row["team_id"] == team, row
+ _assert_reported_usage(row)
+
+
+_LARGE_FRAME_BYTES: Final = 8 * 1024 * 1024
+_PIECE_BYTES: Final = 8 * 1024
+_HEALTH_CEILING_SECONDS: Final = 1.0
+_RELAY_RATIO_CEILING: Final = 4
+_RELAY_SLACK_SECONDS: Final = 1.0
+
+
+@dataclass(frozen=True, slots=True)
+class _Relay:
+ response: httpx.Response
+ seconds: float
+ health: tuple[float, ...]
+
+
+def _timed_health(client: httpx.Client) -> float:
+ started: Final = time.perf_counter()
+ response: Final = client.get("/health/liveliness", headers={"Connection": "close"})
+ assert response.status_code == 200, response.text
+ return time.perf_counter() - started
+
+
+def _health_latencies(gateway: Gateway, done: threading.Event) -> tuple[float, ...]:
+ with httpx.Client(base_url=gateway.client.base_url, timeout=60, trust_env=False) as fresh:
+ return tuple(_timed_health(fresh) for _ in iter(done.is_set, True))
+
+
+def _finish_and_done(call: _Call) -> bytes:
+ return _frame(_chunk(call.identity, [_finish(0)])) + _frame("[DONE]")
+
+
+def _content_then_usage(call: _Call, content: Sequence[bytes]) -> Callable[[Mapping[str, JsonValue]], Reply]:
+ usage: Final = _frame(_chunk(call.identity, [], usage=_USAGE))
+ finish: Final = _frame(_chunk(call.identity, [_finish(0)]))
+
+ def reply(body: Mapping[str, JsonValue]) -> Reply:
+ return _sse((*content, finish, *((usage,) if _asks_usage(body) else ()), _frame("[DONE]")))
+
+ return reply
+
+
+def _timed_relay(gateway: Gateway, call: _Call, content: Sequence[bytes]) -> _Relay:
+ with _deployment(gateway, _scripted(call, _content_then_usage(call, content))) as (_, model):
+ eventually(lambda: _all_workers_serve(gateway, model), bool, seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS)
+ done: Final = threading.Event()
+ with ThreadPoolExecutor(max_workers=1) as pool:
+ polled: Final = pool.submit(_health_latencies, gateway, done)
+ started: Final = time.perf_counter()
+ try:
+ response: Final = _stream(gateway, "/vllm/chat/completions", _chat_body(model, call, stream=True))
+ finally:
+ done.set()
+ seconds: Final = time.perf_counter() - started
+ return _Relay(response, seconds, polled.result())
+
+
+def _assert_large_frame_relays_like_small_frames(
+ gateway: Gateway, record_property: Callable[[str, object], None], large_call: _Call, large: bytes
+) -> None:
+ small_call: Final = _call()
+ small: Final = tuple(
+ _frame(_chunk(small_call.identity, [_delta(0, "x" * _PIECE_BYTES)]))
+ for _ in range(_LARGE_FRAME_BYTES // _PIECE_BYTES)
+ )
+ pieces: Final = tuple(large[start : start + _PIECE_BYTES] for start in range(0, len(large), _PIECE_BYTES))
+ small_relay: Final = _timed_relay(gateway, small_call, small)
+ large_relay: Final = _timed_relay(gateway, large_call, pieces)
+ record_property("small_frames_seconds", round(small_relay.seconds, 3))
+ record_property("large_frame_seconds", round(large_relay.seconds, 3))
+ record_property("health_calls", len(large_relay.health))
+ record_property("health_max_seconds", round(max(large_relay.health), 3))
+ assert small_relay.response.content == b"".join(small) + _finish_and_done(small_call), len(
+ small_relay.response.content
+ )
+ assert large_relay.response.content == large + _finish_and_done(large_call), len(large_relay.response.content)
+ assert max(large_relay.health) < _HEALTH_CEILING_SECONDS, sorted(large_relay.health)[-5:]
+ assert large_relay.seconds < _RELAY_RATIO_CEILING * small_relay.seconds + _RELAY_SLACK_SECONDS, (
+ large_relay.seconds,
+ small_relay.seconds,
+ )
+
+
+def test_frame_terminated_only_at_its_end_relays_as_fast_as_the_same_bytes_in_small_frames(
+ gateway: Gateway, record_property: Callable[[str, object], None]
+) -> None:
+ large_call: Final = _call()
+ large: Final = _frame(_chunk(large_call.identity, [_delta(0, "x" * _LARGE_FRAME_BYTES)]))
+ _assert_large_frame_relays_like_small_frames(gateway, record_property, large_call, large)
+
+
+def test_frame_of_usage_and_finish_reason_near_misses_relays_as_fast_as_the_same_bytes_in_small_frames(
+ gateway: Gateway, record_property: Callable[[str, object], None]
+) -> None:
+ near_miss: Final[dict[str, JsonValue]] = {"usage": 0, "finish_reason": None}
+ count: Final = _LARGE_FRAME_BYTES // len(json.dumps(near_miss) + ", ")
+ large_call: Final = _call()
+ large: Final = _frame(_chunk(large_call.identity, [_delta(0, "")], pad=[near_miss] * count))
+ _assert_large_frame_relays_like_small_frames(gateway, record_property, large_call, large)
+
+
+@dataclass(frozen=True, slots=True)
+class _Outcome:
+ identity: str
+ status: int
+ text: str
+
+
+def _burst_reply(dropped: Callable[[], bool]) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ body: Final = JSON_OBJECT.validate_json(request.body)
+ marker: Final = user_prompt(body) if "messages" in body else str(body.get("prompt"))
+ call: Final = _Call(marker, f"cmpl-{marker}")
+ if dropped():
+ return Reply(drop_connection=True)
+ if body.get("stream") is not True:
+ completion: Final = {
+ "id": call.identity,
+ "object": "chat.completion",
+ "created": 1,
+ "model": _SERVED_MODEL,
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hi"}, "finish_reason": "stop"}],
+ "usage": _USAGE,
+ }
+ return Reply(body=json.dumps(completion).encode())
+ stream: Final = _plain_stream(call, body)
+ return Reply(content_type=stream.content_type, chunks=stream.chunks, pause_between_chunks=0.05)
+
+ return answering_model_discovery(respond)
+
+
+async def _send(client: httpx.AsyncClient, key: str, body: Mapping[str, JsonValue], identity: str) -> _Outcome:
+ try:
+ response: Final = await client.post(
+ "/vllm/chat/completions", json=body, headers={"Authorization": f"Bearer {key}"}
+ )
+ except httpx.TransportError as error:
+ return _Outcome(identity, 0, repr(error))
+ return _Outcome(identity, response.status_code, response.text)
+
+
+def _all_workers_serve(gateway: Gateway, model: str) -> bool:
+ auth: Final = {"Authorization": f"Bearer {gateway.key}"}
+ with httpx.Client(base_url=gateway.client.base_url, timeout=15, trust_env=False) as fresh:
+ return all(_listed(fresh.get("/v1/models", headers={**auth, "Connection": "close"}), model) for _ in range(16))
+
+
+async def test_burst_with_dropped_upstream_connections_lands_every_completed_call_once(gateway: Gateway) -> None:
+ arrivals: Final = itertools.count(1)
+ with wire_server(_burst_reply(lambda: next(arrivals) % 4 == 0)) as wire, gateway.scenario() as scenario:
+ model: Final = _deploy(scenario, wire)
+ await asyncio.to_thread(
+ eventually, lambda: _all_workers_serve(gateway, model), bool, 2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
+ )
+ calls: Final = tuple(_call() for _ in range(24))
+ bodies: Final = tuple(
+ _chat_body(model, call, **({"stream": True} if index % 3 else {})) for index, call in enumerate(calls)
+ )
+ async with httpx.AsyncClient(base_url=gateway.client.base_url, timeout=90, trust_env=False) as client:
+ outcomes: Final = await asyncio.gather(
+ *(_send(client, gateway.key, body, call.identity) for body, call in zip(bodies, calls))
+ )
+ completed: Final = frozenset(
+ outcome.identity for outcome in outcomes if outcome.status == 200 and outcome.identity in outcome.text
+ )
+ failed: Final = tuple(outcome for outcome in outcomes if outcome.identity not in completed)
+ assert len(completed) >= 12, outcomes
+ assert all(outcome.status != 200 for outcome in failed), failed
+ rows: Final = _rows(model, len(completed))
+ assert sorted(str(row["request_id"]) for row in rows) == sorted(completed), rows
+ health: Final = gateway.request("GET", "/health/liveliness")
+ assert health.status_code == 200, health.text
diff --git a/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py b/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py
new file mode 100644
index 00000000000..02e70a82c1c
--- /dev/null
+++ b/tests/integration/routing/test_vllm_passthrough_stream_usage_settings.py
@@ -0,0 +1,283 @@
+import asyncio
+import json
+import threading
+import uuid
+from collections.abc import Callable, Generator, Mapping, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from queue import SimpleQueue
+from types import MappingProxyType
+from typing import Final
+from urllib.parse import urlsplit
+
+import httpx
+import psutil
+import pytest
+import yaml
+from integration._support.anthropic_sse import user_prompt
+from integration._support.client import JSON_OBJECT, eventually, gateway_from_environment, held, object_value
+from integration._support.database import read_rows
+from integration._support.openai_wire import answering_model_discovery
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
+_SERVED_MODEL: Final = "Qwen/Qwen2.5-0.5B-Instruct"
+_INPUT_COST: Final = 0.001
+_OUTPUT_COST: Final = 0.002
+_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
+ {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
+)
+_STARTED_WORKER: Final = "Started server process ["
+_STARTUP_COMPLETE: Final = "Application startup complete."
+_BURST: Final = 20
+
+
+@dataclass(frozen=True, slots=True)
+class _Rig:
+ wire: Wire
+ proxy: OwnedProxy
+ model: str
+
+ @property
+ def base_url(self) -> str:
+ return str(self.proxy.gateway.client.base_url).rstrip("/")
+
+ @property
+ def key(self) -> str:
+ return self.proxy.gateway.key
+
+
+@dataclass(frozen=True, slots=True)
+class _Outcome:
+ identity: str
+ status: int
+ text: str
+
+
+def _identity(body: Mapping[str, JsonValue]) -> str:
+ return f"cmpl-{user_prompt(body)}"
+
+
+def _frame(payload: Mapping[str, JsonValue] | str) -> bytes:
+ return b"data: " + (payload if isinstance(payload, str) else json.dumps(payload)).encode() + b"\n\n"
+
+
+def _chunk(identity: str, choices: Sequence[JsonValue], **extra: JsonValue) -> dict[str, JsonValue]:
+ return {
+ "id": identity,
+ "object": "chat.completion.chunk",
+ "created": 1,
+ "model": _SERVED_MODEL,
+ "choices": list(choices),
+ **extra,
+ }
+
+
+def _usage_frame(identity: str) -> bytes:
+ return _frame(_chunk(identity, [], usage=dict(_USAGE)))
+
+
+def _asks_usage(body: Mapping[str, JsonValue]) -> bool:
+ options: Final = body.get("stream_options")
+ return isinstance(options, dict) and options.get("include_usage") is True
+
+
+def _stream_bytes(body: Mapping[str, JsonValue]) -> tuple[bytes, ...]:
+ identity: Final = _identity(body)
+ return (
+ _frame(_chunk(identity, [{"index": 0, "delta": {"content": "Hello there"}, "finish_reason": None}])),
+ _frame(_chunk(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}])),
+ *((_usage_frame(identity),) if _asks_usage(body) else ()),
+ _frame("[DONE]"),
+ )
+
+
+def _streaming(request: Request) -> Reply:
+ return Reply(content_type="text/event-stream", chunks=_stream_bytes(JSON_OBJECT.validate_json(request.body)))
+
+
+def _config(directory: Path, wire: Wire, model: str, general_settings: Mapping[str, JsonValue]) -> Path:
+ base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text()))
+ deployment: Final = {
+ "model_name": model,
+ "litellm_params": {
+ "model": f"hosted_vllm/{_SERVED_MODEL}",
+ "api_base": f"{wire.url}/v1",
+ "input_cost_per_token": _INPUT_COST,
+ "output_cost_per_token": _OUTPUT_COST,
+ },
+ }
+ config: Final = {
+ **base,
+ "model_list": [*(base.get("model_list") or []), deployment],
+ "general_settings": {**object_value(base["general_settings"]), **general_settings},
+ }
+ path: Final = directory / f"vllm-usage-{uuid.uuid4().hex}.yaml"
+ path.write_text(yaml.safe_dump(config))
+ return path
+
+
+@contextmanager
+def _vllm_proxy(
+ directory: Path, respond: Callable[[Request], Reply], general_settings: Mapping[str, JsonValue]
+) -> Generator[_Rig, None, None]:
+ model: Final = f"vllm-usage-{uuid.uuid4().hex}"
+ with gateway_from_environment() as gateway, wire_server(answering_model_discovery(respond)) as wire:
+ config: Final = _config(directory, wire, model, general_settings)
+ with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned:
+ yield _Rig(wire, owned, model)
+
+
+def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]:
+ return {"model": model, "messages": [{"role": "user", "content": f"vllm-usage-{uuid.uuid4().hex}"}], **extra}
+
+
+def _rows(model: str, count: int) -> list[dict[str, JsonValue]]:
+ return held(
+ lambda: read_rows(
+ 'SELECT request_id, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
+ "WHERE model_group=%s",
+ (model,),
+ ),
+ lambda rows: len(rows) >= count,
+ holding=2,
+ seconds=70,
+ )
+
+
+def _posted(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
+ return tuple(JSON_OBJECT.validate_json(request.body) for request in wire.drain() if request.method == "POST")
+
+
+def _spend(prompt: JsonValue, completion: JsonValue) -> float:
+ assert isinstance(prompt, int) and isinstance(completion, int), (prompt, completion)
+ return prompt * _INPUT_COST + completion * _OUTPUT_COST
+
+
+@pytest.mark.parametrize("setting", [True, None], ids=["true", "null"])
+def test_always_include_stream_usage_true_or_null_injects_usage_and_relays_the_stream_the_caller_asked_for(
+ tmp_path: Path, setting: bool | None
+) -> None:
+ with _vllm_proxy(tmp_path, _streaming, {"always_include_stream_usage": setting}) as rig:
+ body: Final = _chat_body(rig.model, stream=True)
+ response: Final = rig.proxy.gateway.request("POST", "/vllm/chat/completions", body)
+ assert response.status_code == 200, response.text
+ (posted,) = _posted(rig.wire)
+ assert posted["stream_options"] == {"include_usage": True}, posted
+ assert response.content == b"".join(_stream_bytes(body)), response.text
+ (row,) = _rows(rig.model, 1)
+ assert row["request_id"] == _identity(body), row
+ assert (row["status"], row["prompt_tokens"], row["completion_tokens"]) == ("success", 5, 3), row
+ assert row["spend"] == pytest.approx(_spend(5, 3)), row
+
+
+def test_always_include_stream_usage_false_forwards_the_body_untouched_and_counts_locally(tmp_path: Path) -> None:
+ with _vllm_proxy(tmp_path, _streaming, {"always_include_stream_usage": False}) as rig:
+ body: Final = _chat_body(rig.model, stream=True)
+ response: Final = rig.proxy.gateway.request("POST", "/vllm/chat/completions", body)
+ assert response.status_code == 200, response.text
+ (posted,) = _posted(rig.wire)
+ assert "stream_options" not in posted, posted
+ assert response.content == b"".join(_stream_bytes(body)), response.text
+ (row,) = _rows(rig.model, 1)
+ assert row["request_id"] == _identity(body), row
+ prompt, completion = row["prompt_tokens"], row["completion_tokens"]
+ assert row["status"] == "success" and isinstance(prompt, int) and prompt > 0, row
+ assert isinstance(completion, int) and completion > 0, row
+ assert row["spend"] == pytest.approx(_spend(prompt, completion)), row
+
+
+def _holding(release: threading.Event, arrived: SimpleQueue[str]) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ if request.method == "POST" and urlsplit(request.target).path.endswith("/chat/completions"):
+ arrived.put(request.target)
+ assert release.wait(timeout=90), "The burst was never released"
+ return _streaming(request)
+
+ return respond
+
+
+@contextmanager
+def _releasing(release: threading.Event) -> Generator[None, None, None]:
+ try:
+ yield
+ finally:
+ release.set()
+
+
+def _worker_pids(log: str) -> tuple[int, ...]:
+ return tuple(
+ int(line.split(_STARTED_WORKER, 1)[1].split("]", 1)[0]) for line in log.splitlines() if _STARTED_WORKER in line
+ )
+
+
+def _open_upstream_connections(pid: int, upstream: str) -> int:
+ port: Final = urlsplit(upstream).port
+ return sum(
+ 1
+ for connection in psutil.Process(pid).net_connections(kind="tcp")
+ if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
+ )
+
+
+async def _send(client: httpx.AsyncClient, key: str, body: Mapping[str, JsonValue]) -> _Outcome:
+ try:
+ response: Final = await client.post(
+ "/vllm/chat/completions", json=body, headers={"Authorization": f"Bearer {key}", "Connection": "close"}
+ )
+ except httpx.TransportError as error:
+ return _Outcome(_identity(body), 0, repr(error))
+ return _Outcome(_identity(body), response.status_code, response.text)
+
+
+async def _burst(rig: _Rig, bodies: Sequence[Mapping[str, JsonValue]]) -> tuple[_Outcome, ...]:
+ async with httpx.AsyncClient(base_url=rig.base_url, timeout=120, trust_env=False) as client:
+ return tuple(await asyncio.gather(*(_send(client, rig.key, body) for body in bodies)))
+
+
+def _served(outcomes: Sequence[_Outcome]) -> frozenset[str]:
+ return frozenset(
+ outcome.identity for outcome in outcomes if outcome.status == 200 and outcome.text.endswith("data: [DONE]\n\n")
+ )
+
+
+async def test_worker_sigkill_mid_burst_lands_every_served_call_exactly_once(tmp_path: Path) -> None:
+ release: Final = threading.Event()
+ arrived: Final[SimpleQueue[str]] = SimpleQueue()
+ with _vllm_proxy(tmp_path, _holding(release, arrived), {}) as rig, _releasing(release):
+ workers: Final = eventually(
+ lambda: _worker_pids(rig.proxy.log.read_text()), lambda pids: len(pids) == 2, seconds=30
+ )
+ burst: Final = asyncio.create_task(
+ _burst(rig, tuple(_chat_body(rig.model, stream=True) for _ in range(_BURST)))
+ )
+ await asyncio.to_thread(eventually, arrived.qsize, lambda size: size == _BURST, 60)
+ held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, rig.wire.url) for pid in workers})
+ assert sum(held_by.values()) == _BURST, held_by
+ victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
+ with rig.proxy.replacing_worker(victim_pid) as kill:
+ kill.freeze()
+ await asyncio.to_thread(kill)
+ release.set()
+ outcomes: Final = await burst
+ served: Final = _served(outcomes)
+ assert held_by[survivor_pid] >= _BURST // 2, held_by
+ assert len(served) == held_by[survivor_pid], (held_by, outcomes)
+ assert all(outcome.status != 200 for outcome in outcomes if outcome.identity not in served), outcomes
+ assert all('"usage"' not in outcome.text for outcome in outcomes if outcome.identity in served), outcomes
+ (follow_up,) = await _burst(rig, (_chat_body(rig.model, stream=True),))
+ assert _served((follow_up,)) == {follow_up.identity}, follow_up
+ rows: Final = _rows(rig.model, len(served) + 1)
+ assert sorted(str(row["request_id"]) for row in rows) == sorted((*served, follow_up.identity)), rows
+ assert all(row["status"] == "success" for row in rows), rows
+ replaced: Final = eventually(
+ lambda: rig.proxy.log.read_text(),
+ lambda log: len(_worker_pids(log)) == 3 and log.count(_STARTUP_COMPLETE) == 3,
+ seconds=2 * graceful_stop_seconds(),
+ )
+ assert victim_pid not in _worker_pids(replaced)[2:], replaced
diff --git a/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py
index bc8c1000fa4..f608e91702c 100644
--- a/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py
+++ b/tests/unit/llms/azure/passthrough/test_azure_passthrough_transformation.py
@@ -230,10 +230,11 @@ def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response():
assert response.usage.completion_tokens == 8
-def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request():
+@pytest.mark.parametrize("stray_fields", [{}, {"prompt": 5}])
+def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request(stray_fields):
messages = [{"role": "user", "content": "Say hi in three words"}]
logging_obj = MagicMock()
- logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}}
+ logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True, **stray_fields}}
response = AzurePassthroughConfig().handle_logging_collected_chunks(
all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk],
diff --git a/tests/unit/llms/vllm/__init__.py b/tests/unit/llms/vllm/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/vllm/passthrough/__init__.py b/tests/unit/llms/vllm/passthrough/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py b/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py
new file mode 100644
index 00000000000..bc8b7111fbb
--- /dev/null
+++ b/tests/unit/llms/vllm/passthrough/test_vllm_passthrough_transformation.py
@@ -0,0 +1,241 @@
+import json
+from datetime import datetime
+
+import httpx
+import pytest
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.llms.vllm.passthrough.transformation import VLLMPassthroughConfig
+
+MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+MESSAGES = [{"role": "user", "content": "Say hi in three words"}]
+PROMPT = "The capital of France is"
+USAGE = {"prompt_tokens": 34, "completion_tokens": 4, "total_tokens": 38}
+
+
+def _relay_logging_obj(endpoint: str, request_data: dict, callbacks: list[CustomLogger] | None = None) -> Logging:
+ logging_obj = Logging(
+ model=MODEL,
+ messages=[],
+ stream="stream" in request_data,
+ call_type="allm_passthrough_route",
+ start_time=datetime(2025, 1, 1),
+ litellm_call_id="call-1",
+ function_id="fn-1",
+ dynamic_async_success_callbacks=callbacks,
+ )
+ logging_obj.update_environment_variables(
+ model=MODEL,
+ litellm_params={"api_base": "http://vllm.local:8000/v1"},
+ optional_params={},
+ custom_llm_provider="hosted_vllm",
+ endpoint=endpoint,
+ request_data=request_data,
+ )
+ return logging_obj
+
+
+def _relayed_json(path: str, body: dict) -> httpx.Response:
+ return httpx.Response(
+ status_code=200,
+ headers={"content-type": "application/json"},
+ content=json.dumps(body).encode(),
+ request=httpx.Request("POST", f"http://vllm.local:8000/{path}"),
+ )
+
+
+def _chat_completion(identity: str) -> httpx.Response:
+ body = {
+ "id": identity,
+ "object": "chat.completion",
+ "created": 1,
+ "model": MODEL,
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello there!"}, "finish_reason": "stop"}],
+ "usage": USAGE,
+ }
+ return _relayed_json("v1/chat/completions", body)
+
+
+def _chat_stream(identity: str, usage: dict | None) -> bytes:
+ head = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": MODEL}
+ frames = [
+ {**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello"}, "finish_reason": None}]},
+ {**head, "choices": [{"index": 0, "delta": {"content": " there!"}, "finish_reason": "stop"}]},
+ *([{**head, "choices": [], "usage": usage}] if usage else []),
+ ]
+ return b"".join(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames) + b"data: [DONE]\n\n"
+
+
+def _text_stream(identity: str, deltas: list[tuple[int, str]], usage: dict | None) -> bytes:
+ head = {"id": identity, "object": "text_completion", "created": 1, "model": MODEL}
+ frames = [
+ *({**head, "choices": [{"index": index, "text": text, "finish_reason": None}]} for index, text in deltas),
+ *([{**head, "choices": [], "usage": usage}] if usage else []),
+ ]
+ return b"".join(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames) + b"data: [DONE]\n\n"
+
+
+async def _log_stream(endpoint: str, sse: bytes, request_data: dict | None = None) -> Logging:
+ logging_obj = _relay_logging_obj(endpoint, request_data or {"model": MODEL, "messages": MESSAGES, "stream": True})
+ collector = VLLMPassthroughConfig().create_stream_collector(
+ model=MODEL, custom_llm_provider="hosted_vllm", endpoint=endpoint
+ )
+ collector.add(sse[:57])
+ collector.add(sse[57:])
+ await logging_obj.async_flush_passthrough_collected_chunks(collector)
+ return logging_obj
+
+
+@pytest.mark.asyncio
+async def test_a_relayed_chat_completion_is_logged_with_the_usage_vllm_reported():
+ logging_obj = _relay_logging_obj("chat/completions", {"model": MODEL, "messages": MESSAGES})
+
+ await logging_obj.async_success_handler(result=_chat_completion("chatcmpl-vllm-1"))
+
+ payload = logging_obj.model_call_details["standard_logging_object"]
+ assert payload["id"] == "chatcmpl-vllm-1"
+ assert payload["custom_llm_provider"] == "hosted_vllm"
+ assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (34, 4, 38)
+
+
+@pytest.mark.asyncio
+async def test_a_relayed_chat_stream_is_logged_with_the_usage_chunk_vllm_sent():
+ logging_obj = await _log_stream("v1/chat/completions", _chat_stream("chatcmpl-vllm-2", USAGE))
+
+ payload = logging_obj.model_call_details["standard_logging_object"]
+ assert payload["id"] == "chatcmpl-vllm-2"
+ assert payload["response"]["choices"][0]["message"]["content"] == "Hello there!"
+ assert (payload["prompt_tokens"], payload["completion_tokens"]) == (34, 4)
+
+
+@pytest.mark.asyncio
+async def test_a_relayed_chat_stream_without_a_usage_chunk_counts_tokens_from_the_relayed_request():
+ logging_obj = await _log_stream("chat/completions", _chat_stream("chatcmpl-vllm-3", None))
+
+ payload = logging_obj.model_call_details["standard_logging_object"]
+ assert payload["id"] == "chatcmpl-vllm-3"
+ assert payload["prompt_tokens"] == litellm.token_counter(model=MODEL, messages=MESSAGES)
+ assert payload["completion_tokens"] > 0
+
+
+@pytest.mark.asyncio
+async def test_a_relayed_text_completion_is_logged_with_the_usage_vllm_reported():
+ logging_obj = _relay_logging_obj("v1/completions", {"model": MODEL, "prompt": PROMPT})
+ body = {
+ "id": "cmpl-vllm-4",
+ "object": "text_completion",
+ "created": 1,
+ "model": MODEL,
+ "choices": [{"index": 0, "text": " Paris.", "finish_reason": "stop", "stop_reason": None}],
+ "usage": USAGE,
+ }
+
+ result = VLLMPassthroughConfig().logging_non_streaming_response(
+ model=MODEL,
+ custom_llm_provider="hosted_vllm",
+ httpx_response=_relayed_json("v1/completions", body),
+ request_data={"model": MODEL, "prompt": PROMPT},
+ logging_obj=logging_obj,
+ endpoint="v1/completions",
+ )
+ await logging_obj.async_success_handler(result=result)
+
+ payload = logging_obj.model_call_details["standard_logging_object"]
+ assert payload["id"] == "cmpl-vllm-4"
+ assert payload["response"]["choices"][0]["text"] == " Paris."
+ assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (34, 4, 38)
+
+
+class _ResponseRecorder(CustomLogger):
+ def __init__(self) -> None:
+ super().__init__()
+ self.responses: list[object] = []
+
+ async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+ self.responses.append(response_obj)
+
+
+@pytest.mark.asyncio
+async def test_a_success_callback_gets_the_text_completion_vllm_relayed_as_a_parsed_response():
+ recorder = _ResponseRecorder()
+ logging_obj = _relay_logging_obj("v1/completions", {"model": MODEL, "prompt": PROMPT}, callbacks=[recorder])
+ body = {
+ "id": "cmpl-vllm-8",
+ "object": "text_completion",
+ "created": 1,
+ "model": MODEL,
+ "choices": [{"index": 0, "text": " Paris.", "finish_reason": "stop"}],
+ "usage": USAGE,
+ }
+
+ await logging_obj.async_success_handler(result=_relayed_json("v1/completions", body))
+
+ [response] = recorder.responses
+ assert (response.id, response.choices[0].text, response.usage.prompt_tokens) == ("cmpl-vllm-8", " Paris.", 34)
+
+
+@pytest.mark.parametrize(
+ "path, request_data, body",
+ [
+ (
+ "v1/completions",
+ {"model": MODEL, "prompt": PROMPT},
+ {
+ "id": "cmpl-vllm-7",
+ "object": "text_completion",
+ "created": 1,
+ "model": MODEL,
+ "choices": [{"index": 0, "text": " Paris.", "finish_reason": "stop"}],
+ "usage": USAGE,
+ },
+ ),
+ (
+ "v1/embeddings",
+ {"model": MODEL, "input": [PROMPT]},
+ {
+ "object": "list",
+ "model": MODEL,
+ "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}],
+ "usage": {"prompt_tokens": 34, "total_tokens": 34},
+ },
+ ),
+ ],
+)
+def test_the_sync_and_async_success_handlers_both_read_the_usage_from_one_relayed_body(path, request_data, body):
+ logging_obj = _relay_logging_obj(path, request_data)
+ relayed = _relayed_json(path, body)
+
+ sync_handler_view = logging_obj.normalize_logging_result(result=relayed)
+ async_handler_view = logging_obj.normalize_logging_result(result=relayed)
+
+ assert [view.usage.prompt_tokens for view in (sync_handler_view, async_handler_view)] == [34, 34]
+
+
+@pytest.mark.asyncio
+async def test_a_relayed_text_completion_stream_is_logged_with_the_usage_chunk_vllm_sent():
+ sse = _text_stream("cmpl-vllm-5", [(0, " Par"), (0, "is.")], USAGE)
+
+ logging_obj = await _log_stream("v1/completions", sse, {"model": MODEL, "prompt": PROMPT, "stream": True})
+
+ payload = logging_obj.model_call_details["standard_logging_object"]
+ assert payload["id"] == "cmpl-vllm-5"
+ assert payload["response"]["choices"][0]["text"] == " Paris."
+ assert (payload["prompt_tokens"], payload["completion_tokens"]) == (34, 4)
+
+
+@pytest.mark.asyncio
+async def test_a_relayed_text_completion_stream_without_a_usage_chunk_counts_every_prompt_and_choice():
+ prompts = ["Hi", "Yo"]
+ sse = _text_stream("cmpl-vllm-6", [(0, "Hello"), (1, "Hey"), (0, " there"), (1, " you")], None)
+
+ logging_obj = await _log_stream("completions", sse, {"model": MODEL, "prompt": prompts, "stream": True})
+
+ payload = logging_obj.model_call_details["standard_logging_object"]
+ assert payload["id"] == "cmpl-vllm-6"
+ assert [choice["text"] for choice in payload["response"]["choices"]] == ["Hello there", "Hey you"]
+ assert payload["prompt_tokens"] == sum(litellm.token_counter(model=MODEL, text=prompt) for prompt in prompts)
+ assert payload["completion_tokens"] == sum(
+ litellm.token_counter(model=MODEL, text=text, count_response_tokens=True) for text in ["Hello there", "Hey you"]
+ )
diff --git a/tests/unit/proxy/pass_through_endpoints/test_injected_stream_usage.py b/tests/unit/proxy/pass_through_endpoints/test_injected_stream_usage.py
new file mode 100644
index 00000000000..82c7ce8b864
--- /dev/null
+++ b/tests/unit/proxy/pass_through_endpoints/test_injected_stream_usage.py
@@ -0,0 +1,243 @@
+import json
+from collections.abc import AsyncIterator, Iterable
+
+import pytest
+
+from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
+ stream_usage_relay,
+ without_injected_usage_frames,
+)
+
+MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+FINGERPRINT = "fp-vllm-test"
+
+
+def _frame(payload: dict) -> bytes:
+ return b"data: " + json.dumps(payload).encode() + b"\n\n"
+
+
+def _chunk(identity: str, **fields) -> dict:
+ return {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": MODEL, **fields}
+
+
+def _content(index: int, **fields) -> bytes:
+ return _frame(
+ _chunk("chatcmpl-1", choices=[{"index": index, "delta": {"content": "Hi"}, "finish_reason": None}], **fields)
+ )
+
+
+def _finish(index: int, **fields) -> bytes:
+ return _frame(_chunk("chatcmpl-1", choices=[{"index": index, "delta": {}, "finish_reason": "stop"}], **fields))
+
+
+USAGE_COUNTS = {"prompt_tokens": 34, "completion_tokens": 4, "total_tokens": 38}
+CONTENT = _content(0)
+FINISH = _finish(0)
+USAGE = _frame(_chunk("chatcmpl-1", choices=[], usage=USAGE_COUNTS, system_fingerprint=FINGERPRINT))
+DONE = b"data: [DONE]\n\n"
+
+
+async def _pieces(pieces: Iterable[bytes]) -> AsyncIterator[bytes]:
+ for piece in pieces:
+ yield piece
+
+
+async def _relayed(pieces: Iterable[bytes]) -> bytes:
+ return b"".join([piece async for piece in without_injected_usage_frames(_pieces(pieces))])
+
+
+def _payloads(relayed: bytes) -> list[dict]:
+ return [json.loads(line[len("data: ") :]) for line in relayed.decode().split("\n\n") if line.startswith("data: {")]
+
+
+class TestStreamUsageRelay:
+ def test_a_chat_stream_nobody_asked_usage_for_gets_it_injected_and_stripped(self):
+ relay = stream_usage_relay(
+ endpoint="chat/completions",
+ request_body={"model": MODEL, "messages": [], "stream": True},
+ general_settings={},
+ supports_stream_options=lambda: True,
+ )
+
+ assert relay.body == {"model": MODEL, "messages": [], "stream": True, "stream_options": {"include_usage": True}}
+ assert relay.strips_usage is True
+
+ def test_a_chat_stream_the_caller_asked_usage_for_is_relayed_untouched(self):
+ body = {"model": MODEL, "messages": [], "stream": True, "stream_options": {"include_usage": True}}
+
+ relay = stream_usage_relay(
+ endpoint="v1/chat/completions", request_body=body, general_settings={}, supports_stream_options=lambda: True
+ )
+
+ assert relay.body == body
+ assert relay.strips_usage is False
+
+ def test_a_deployment_without_stream_options_support_is_left_alone(self):
+ body = {"model": MODEL, "messages": [], "stream": True}
+
+ relay = stream_usage_relay(
+ endpoint="chat/completions", request_body=body, general_settings={}, supports_stream_options=lambda: False
+ )
+
+ assert relay.body == body
+ assert relay.strips_usage is False
+
+ def test_a_text_completion_stream_gets_usage_injected_and_stripped_like_a_chat_stream(self):
+ body = {"model": MODEL, "prompt": "Hi", "stream": True}
+
+ by_default = stream_usage_relay(
+ endpoint="completions", request_body=body, general_settings={}, supports_stream_options=lambda: True
+ )
+ always = stream_usage_relay(
+ endpoint="v1/completions",
+ request_body=body,
+ general_settings={"always_include_stream_usage": True},
+ supports_stream_options=lambda: True,
+ )
+
+ assert by_default.body == always.body == {**body, "stream_options": {"include_usage": True}}
+ assert (by_default.strips_usage, always.strips_usage) == (True, True)
+
+ def test_turning_stream_usage_off_relays_the_stream_untouched(self):
+ body = {"model": MODEL, "messages": [], "stream": True}
+
+ relay = stream_usage_relay(
+ endpoint="chat/completions",
+ request_body=body,
+ general_settings={"always_include_stream_usage": False},
+ supports_stream_options=lambda: True,
+ )
+
+ assert relay.body == body
+ assert relay.strips_usage is False
+
+ @pytest.mark.parametrize(
+ "stream_options",
+ [True, "include_usage", ["include_usage"], {"include_usage": "true"}, {"include_usage": 1}],
+ )
+ def test_stream_options_the_proxy_cannot_read_reach_the_server_as_the_caller_sent_them(self, stream_options):
+ body = {"model": MODEL, "messages": [], "stream": True, "stream_options": stream_options}
+
+ relay = stream_usage_relay(
+ endpoint="chat/completions", request_body=body, general_settings={}, supports_stream_options=lambda: True
+ )
+
+ assert relay.body == body
+ assert relay.strips_usage is False
+
+ def test_continuous_usage_stats_alone_is_not_turned_into_per_chunk_usage(self):
+ body = {"model": MODEL, "messages": [], "stream": True, "stream_options": {"continuous_usage_stats": True}}
+
+ relay = stream_usage_relay(
+ endpoint="chat/completions", request_body=body, general_settings={}, supports_stream_options=lambda: True
+ )
+
+ assert relay.body == {**body, "stream_options": {"include_usage": True}}
+ assert relay.strips_usage is True
+
+ def test_a_non_completion_endpoint_is_relayed_untouched(self):
+ body = {"model": MODEL, "input": "Hi", "stream": True}
+
+ relay = stream_usage_relay(
+ endpoint="embeddings", request_body=body, general_settings={}, supports_stream_options=lambda: True
+ )
+
+ assert relay.body == body
+ assert relay.strips_usage is False
+
+
+class TestWithoutInjectedUsageFrames:
+ @pytest.mark.asyncio
+ async def test_the_usage_frame_is_dropped_and_its_fingerprint_moves_onto_the_finish_frame(self):
+ relayed = await _relayed([CONTENT, FINISH, USAGE, DONE])
+
+ payloads = _payloads(relayed)
+ assert [payload["choices"] for payload in payloads] == [
+ [{"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}],
+ [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+ ]
+ assert "usage" not in relayed.decode()
+ assert payloads[1]["system_fingerprint"] == FINGERPRINT
+ assert relayed.endswith(DONE)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("frame_end", [b"\n\n", b"\r\n\r\n", b"\r\r"])
+ async def test_frames_split_across_upstream_reads_relay_the_same_bytes(self, frame_end):
+ whole = (CONTENT + FINISH + USAGE + DONE).replace(b"\n\n", frame_end)
+
+ assert await _relayed([whole[index : index + 1] for index in range(len(whole))]) == await _relayed([whole])
+
+ @pytest.mark.asyncio
+ async def test_a_stream_without_a_usage_frame_still_relays_its_finish_frame(self):
+ assert await _relayed([CONTENT, FINISH, DONE]) == CONTENT + FINISH + DONE
+
+ @pytest.mark.asyncio
+ async def test_a_finish_frame_already_fingerprinted_is_not_held(self):
+ stamped_finish = _frame(
+ _chunk(
+ "chatcmpl-1", choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}], system_fingerprint="fp-0"
+ )
+ )
+
+ relayed = await _relayed([CONTENT, stamped_finish, USAGE, DONE])
+
+ assert relayed == CONTENT + stamped_finish + DONE
+
+ @pytest.mark.asyncio
+ async def test_every_finish_frame_gets_the_fingerprint_when_the_stream_has_several_choices(self):
+ relayed = await _relayed([_content(0), _content(1), _finish(0), _finish(1), USAGE, DONE])
+
+ payloads = _payloads(relayed)
+ finishes = {
+ payload["choices"][0]["index"]: payload.get("system_fingerprint")
+ for payload in payloads
+ if payload["choices"][0]["finish_reason"]
+ }
+ assert finishes == {0: FINGERPRINT, 1: FINGERPRINT}
+ assert "usage" not in relayed.decode()
+
+ @pytest.mark.asyncio
+ async def test_a_choice_that_finishes_while_another_still_streams_keeps_its_place(self):
+ last_token_finish = _frame(
+ _chunk("chatcmpl-1", choices=[{"index": 0, "delta": {"content": "!"}, "finish_reason": "stop"}])
+ )
+ upstream = [_content(0), _content(1), last_token_finish, _content(1), _content(1), _finish(1), USAGE, DONE]
+
+ relayed = await _relayed(upstream)
+
+ assert [payload["choices"] for payload in _payloads(relayed)] == [
+ payload["choices"] for payload in _payloads(b"".join(upstream[:-2]))
+ ]
+ assert _payloads(relayed)[-1]["system_fingerprint"] == FINGERPRINT
+
+ @pytest.mark.asyncio
+ async def test_a_server_that_forces_usage_onto_every_chunk_is_relayed_untouched(self):
+ forced = [
+ _content(0, usage=USAGE_COUNTS),
+ _finish(0, usage=USAGE_COUNTS),
+ USAGE,
+ DONE,
+ ]
+
+ assert await _relayed(forced) == b"".join(forced)
+
+ @pytest.mark.asyncio
+ async def test_closing_the_relay_early_closes_the_upstream(self):
+ closed = []
+
+ async def upstream() -> AsyncIterator[bytes]:
+ try:
+ yield CONTENT
+ yield FINISH
+ finally:
+ closed.append(True)
+
+ relay = without_injected_usage_frames(upstream())
+ assert await anext(relay) == CONTENT
+ await relay.aclose()
+
+ assert closed == [True]
+
+ @pytest.mark.asyncio
+ async def test_an_unterminated_tail_is_flushed_when_the_upstream_closes(self):
+ assert await _relayed([CONTENT, b'data: {"partial": true']) == CONTENT + b'data: {"partial": true'
diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
index d25b4470668..e2ecfbb8ba9 100644
--- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
+++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
@@ -2283,6 +2283,121 @@ class TestVLLMProxyRoute:
mock_is_router.assert_called_once()
mock_llm_router.allm_passthrough_route.assert_awaited_once()
+ @pytest.mark.asyncio
+ @patch( # test-quality-ok: patching litellm internal for unit test isolation
+ "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",
+ return_value={"model": "router-model", "messages": [{"role": "user", "content": "Hi"}], "stream": True},
+ )
+ @patch( # test-quality-ok: patching litellm internal for unit test isolation
+ "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model",
+ return_value=True,
+ )
+ @patch(
+ "litellm.proxy.proxy_server.general_settings", {}
+ ) # test-quality-ok: the proxy config surface for always_include_stream_usage
+ @patch(
+ "litellm.proxy.proxy_server.llm_router"
+ ) # test-quality-ok: patching litellm internal for unit test isolation
+ async def test_vllm_proxy_route_asks_the_router_model_for_stream_usage_and_hides_it_from_the_client(
+ self, mock_llm_router, mock_is_router, mock_get_body
+ ):
+ chunk = {
+ "id": "chatcmpl-1",
+ "object": "chat.completion.chunk",
+ "created": 1,
+ "model": "Qwen/Qwen2.5-0.5B-Instruct",
+ }
+ frames = [
+ {**chunk, "choices": [{"index": 0, "delta": {"content": "Hi"}, "finish_reason": None}]},
+ {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
+ {
+ **chunk,
+ "choices": [],
+ "usage": {"prompt_tokens": 34, "completion_tokens": 4},
+ "system_fingerprint": "fp-1",
+ },
+ ]
+ upstream_body = (
+ b"".join(b"data: " + json.dumps(frame).encode() + b"\n\n" for frame in frames) + b"data: [DONE]\n\n"
+ )
+ mock_request = MagicMock(spec=Request)
+ mock_request.method = "POST"
+ mock_request.headers = {"content-type": "application/json"}
+ mock_request.query_params = {}
+ mock_llm_router.resolved_litellm_models = MagicMock(return_value=("hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct",))
+ mock_llm_router.allm_passthrough_route = AsyncMock(
+ return_value=httpx.Response(200, content=upstream_body, headers={"content-type": "text/event-stream"})
+ )
+
+ result = await vllm_proxy_route(
+ endpoint="/chat/completions",
+ request=mock_request,
+ fastapi_response=MagicMock(spec=Response),
+ user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"),
+ )
+
+ relayed_body = mock_llm_router.allm_passthrough_route.await_args.kwargs["json"]
+ assert relayed_body["stream_options"] == {"include_usage": True}
+ assert isinstance(result, StreamingResponse)
+ body = b"".join([piece async for piece in result.body_iterator])
+ payloads = [
+ json.loads(line[len("data: ") :]) for line in body.decode().split("\n\n") if line.startswith("data: {")
+ ]
+ assert [payload["choices"] for payload in payloads] == [frames[0]["choices"], frames[1]["choices"]]
+ assert "usage" not in body.decode()
+ assert payloads[1]["system_fingerprint"] == "fp-1"
+ assert body.endswith(b"data: [DONE]\n\n")
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("stream_options", [None, {"include_usage": True}])
+ @patch( # test-quality-ok: patching litellm internal for unit test isolation
+ "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model",
+ return_value=True,
+ )
+ @patch(
+ "litellm.proxy.proxy_server.general_settings", {}
+ ) # test-quality-ok: the proxy config surface for always_include_stream_usage
+ @patch(
+ "litellm.proxy.proxy_server.llm_router"
+ ) # test-quality-ok: patching litellm internal for unit test isolation
+ async def test_vllm_proxy_route_closes_the_upstream_stream_when_the_client_disconnects(
+ self, mock_llm_router, mock_is_router, stream_options
+ ):
+ closed = []
+
+ class UpstreamStream(httpx.AsyncByteStream):
+ async def __aiter__(self):
+ yield b'data: {"id":"chatcmpl-1","choices":[]}\n\n'
+ yield b"data: [DONE]\n\n"
+
+ async def aclose(self) -> None:
+ closed.append(True)
+
+ body = {"model": "router-model", "messages": [{"role": "user", "content": "Hi"}], "stream": True}
+ mock_request = MagicMock(spec=Request)
+ mock_request.method = "POST"
+ mock_request.headers = {"content-type": "application/json"}
+ mock_request.query_params = {}
+ mock_llm_router.resolved_litellm_models = MagicMock(return_value=("hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct",))
+ mock_llm_router.allm_passthrough_route = AsyncMock(
+ return_value=httpx.Response(200, stream=UpstreamStream(), headers={"content-type": "text/event-stream"})
+ )
+
+ with patch( # test-quality-ok: patching litellm internal for unit test isolation
+ "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",
+ return_value={**body, **({"stream_options": stream_options} if stream_options else {})},
+ ):
+ result = await vllm_proxy_route(
+ endpoint="/chat/completions",
+ request=mock_request,
+ fastapi_response=MagicMock(spec=Response),
+ user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"),
+ )
+ await anext(result.body_iterator)
+ await result.body_iterator.aclose()
+
+ assert closed == [True]
+
@pytest.mark.asyncio
@patch( # test-quality-ok: patching litellm internal for unit test isolation
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body",
@@ -4472,7 +4587,9 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
(b"authorization", f"Bearer {self.MASTER_KEY}".encode()),
(b"content-type", b"application/json"),
],
- authenticated=UserAPIKeyAuth(api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, user_role=LitellmUserRoles.PROXY_ADMIN),
+ authenticated=UserAPIKeyAuth(
+ api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, user_role=LitellmUserRoles.PROXY_ADMIN
+ ),
)
assert forwarded is None, "the master key must never reach the upstream forwarder"
assert raised is not None and raised.status_code == 401
@@ -7752,13 +7869,17 @@ class TestTypeSafePassthroughRoute:
monkeypatch.setenv("OPENROUTER_API_BASE", "https://typesafe.example/base")
monkeypatch.setenv("LAYA_API_BASE", "https://typesafe.example/base")
monkeypatch.setenv("BESPOKE_API_BASE", "https://typesafe.example/base")
- model: Final = {"typesafe": "jev-latest", "laya": "english", "bespoke": "nimble-latest"}.get(provider, "test-generative-model")
+ model: Final = {"typesafe": "jev-latest", "laya": "english", "bespoke": "nimble-latest"}.get(
+ provider, "test-generative-model"
+ )
permission_model: Final = f"{provider}/{model}" if provider in ("laya", "bespoke") else model
auth: Final = UserAPIKeyAuth(
api_key="sk-limited",
tpm_limit=token_limit if quota_scope == "key" else None,
project_id="test-project" if quota_scope == "project_output" else None,
- project_metadata={"model_otpm_limit": {permission_model: token_limit}} if quota_scope == "project_output" else {},
+ project_metadata={"model_otpm_limit": {permission_model: token_limit}}
+ if quota_scope == "project_output"
+ else {},
)
monkeypatch.setitem(proxy_server.app.dependency_overrides, user_api_key_auth, lambda: auth)
body: Final = (
@@ -7780,7 +7901,9 @@ class TestTypeSafePassthroughRoute:
return httpx.Response(200, json={"model": model})
with respx.mock(assert_all_called=False) as upstream:
- route: Final = upstream.post(f"https://typesafe.example/base/v1/{endpoint}").mock(side_effect=upstream_response)
+ route: Final = upstream.post(f"https://typesafe.example/base/v1/{endpoint}").mock(
+ side_effect=upstream_response
+ )
response: Final = client.post(f"/{provider}/v1/{endpoint}", json=body)
assert response.status_code == (429 if token_limit == 0 else 200), response.text
@@ -7856,11 +7979,14 @@ class TestOssDecisionPassthroughRoute:
"questions": {"department": {"type": "choice", "criteria": {"billing": "refunds"}}},
}
answer: Final = {
- "model": "laya-rl-agent" if provider == "laya" else checkpoint, "answers": {},
+ "model": "laya-rl-agent" if provider == "laya" else checkpoint,
+ "answers": {},
**({"routing": {"model": checkpoint}} if provider == "laya" else {}),
}
with respx.mock(assert_all_called=True) as upstream:
- route: Final = upstream.post(f"http://{provider}.test/base/v1/systemone?trace=yes").respond(200, json=answer)
+ route: Final = upstream.post(f"http://{provider}.test/base/v1/systemone?trace=yes").respond(
+ 200, json=answer
+ )
response: Final = client.post(
f"/{provider}/v1/systemone?trace=yes",
json=body,
@@ -7882,14 +8008,18 @@ class TestOssDecisionPassthroughRoute:
assert f"{provider.upper()}_API_BASE" in response.text
assert len(upstream.calls) == 0
- def test_oss_does_not_forward_unsupported_endpoints(self, client: TestClient, provider: str, checkpoint: str) -> None:
+ def test_oss_does_not_forward_unsupported_endpoints(
+ self, client: TestClient, provider: str, checkpoint: str
+ ) -> None:
with respx.mock(assert_all_called=False) as upstream:
response: Final = client.post(f"/{provider}/v1/evaluate", json={"model": checkpoint})
assert response.status_code == 404
assert len(upstream.calls) == 0
@pytest.mark.parametrize("model", [None, "auto", "jev-latest"])
- def test_oss_rejects_implicit_checkpoint_selection(self, client: TestClient, model: str | None, provider: str) -> None:
+ def test_oss_rejects_implicit_checkpoint_selection(
+ self, client: TestClient, model: str | None, provider: str
+ ) -> None:
with respx.mock(assert_all_called=False) as upstream:
response: Final = client.post(f"/{provider}/v1/systemone", json={"model": model})
assert response.status_code == 400
@@ -7908,7 +8038,6 @@ class TestOssDecisionPassthroughRoute:
assert response.status_code == 400
assert not route.called
-
@pytest.mark.parametrize("metadata_slot", ["metadata", "litellm_metadata"])
def test_oss_hooks_enforce_canonical_model_limits_and_keep_native_wire_body(
self, client: TestClient, monkeypatch: pytest.MonkeyPatch, metadata_slot: str, provider: str, checkpoint: str
@@ -7921,8 +8050,10 @@ class TestOssDecisionPassthroughRoute:
cache: Final = DualCache()
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(cache))
auth: Final = UserAPIKeyAuth(
- api_key="oss-native-rpm", metadata={"model_rpm_limit": {f"{provider}/{checkpoint}": 1}},
+ api_key="oss-native-rpm",
+ metadata={"model_rpm_limit": {f"{provider}/{checkpoint}": 1}},
)
+
def authenticated_key() -> UserAPIKeyAuth:
return auth
@@ -7930,8 +8061,11 @@ class TestOssDecisionPassthroughRoute:
class LimitHook(CustomLogger):
async def async_pre_call_hook(
- self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache,
- data: dict[str, object], call_type: CallTypesLiteral,
+ self,
+ user_api_key_dict: UserAPIKeyAuth,
+ cache: DualCache,
+ data: dict[str, object],
+ call_type: CallTypesLiteral,
) -> dict[str, object]:
assert data["model"] == f"{provider}/{checkpoint}"
metadata: Final = data.get(metadata_slot)
@@ -7943,9 +8077,11 @@ class TestOssDecisionPassthroughRoute:
monkeypatch.setattr(litellm, "callbacks", [LimitHook()])
body: Final = {
- "model": checkpoint, "state": "refund",
+ "model": checkpoint,
+ "state": "refund",
metadata_slot: {
- "customer_label": "retained", "model_group": "unbounded-client-choice",
+ "customer_label": "retained",
+ "model_group": "unbounded-client-choice",
"standard_logging_guardrail_information": [{"guardrail_cost": 25.0}],
},
}
@@ -7967,8 +8103,11 @@ class TestOssDecisionPassthroughRoute:
class CheckpointHook(CustomLogger):
async def async_pre_call_hook(
- self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache,
- data: dict[str, object], call_type: CallTypesLiteral,
+ self,
+ user_api_key_dict: UserAPIKeyAuth,
+ cache: DualCache,
+ data: dict[str, object],
+ call_type: CallTypesLiteral,
) -> dict[str, object]:
assert data["model"] == f"{provider}/{checkpoint}"
return {**data, "model": f"{provider}/{changed_checkpoint}"}
diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py
index 85a62614448..09ed6399d68 100644
--- a/tests/unit/proxy/test_common_request_processing.py
+++ b/tests/unit/proxy/test_common_request_processing.py
@@ -8282,10 +8282,10 @@ def _apply_stream_usage_tracking(
route_type: str,
supports_stream_options: Callable[[], bool] = lambda: True,
) -> None:
- from litellm.proxy.common_request_processing import _stream_usage_tracking_updates
+ from litellm.proxy.common_request_processing import stream_usage_tracking_updates
data.update(
- _stream_usage_tracking_updates(
+ stream_usage_tracking_updates(
data=data,
general_settings=general_settings,
route_type=route_type,
@@ -8430,10 +8430,10 @@ class TestApplyStreamUsageTracking:
class TestModelDeploymentsSupportStreamOptions:
def _support(self, model, llm_router=None, team_id=None) -> bool:
from litellm.proxy.common_request_processing import (
- _model_deployments_support_stream_options,
+ model_deployments_support_stream_options,
)
- return _model_deployments_support_stream_options(model=model, llm_router=llm_router, team_id=team_id)
+ return model_deployments_support_stream_options(model=model, llm_router=llm_router, team_id=team_id)
def test_openai_compatible_deployment_supports_stream_options(self):
router = litellm.Router(
From 493ed616bbf306f7c0f98f0577985b5bec6d3ac8 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri
Date: Wed, 14 Oct 2026 12:49:31 -0700
Subject: [PATCH 31/47] fix(config): treat blank env values as unset so the
proxy neither crash-loops nor silently drops a guardrail (#46341)
* fix(config): treat blank numeric env values as unset instead of crashing at import
A numeric knob exported as an empty string, which is what a Helm value: "" or an
emptied .env line produces, raised ValueError inside litellm/constants.py and the
proxy crash-looped before serving traffic. Every int(os.getenv(...)) and
float(os.getenv(...)) in constants.py now goes through get_env_int / get_env_float,
which treat blank as unset and fall back to the default with a warning on garbage.
* fix(guardrails): treat blank litellm_params values as unset so the guardrail still loads
A guardrail whose litellm_params read os.environ/X resolves X to "" when the
variable is exported but empty. Pydantic rejected "" for bool and typed fields
such as default_on and fail_on_error, the boot log reported the guardrail as
skipped and the proxy came up without it. A before-validator on BaseLitellmParams
now drops blank string values so each field takes its default, and the
default_on shortcut in LitellmParams treats a blank value as False.
* fix(config): warn once about an unparsable numeric env value so logging cannot recurse
The LiteLLM loggers' filter reads MAX_STRING_LENGTH_STDOUT_LOG and MAX_BASE64_LENGTH_STDOUT_LOG through get_env_int on every record. Warning about a garbage value from inside that read re-entered the filter, so every log call raised RecursionError. The value is now marked as reported before the warning is logged, which also stops per-record readers from repeating it. get_env_int_or_none and the new get_env_float_or_none share the same parser
* fix(config): treat blank values as unset for the remaining import-time numeric knobs
LITELLM_MAX_STREAMING_DURATION_SECONDS, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, a whitespace CLI_JWT_EXPIRATION_HOURS, LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS, the agent health check timeouts and the Prisma reconnect knobs still crashed the proxy at boot when blank. A blank or unparsable REQUEST_TIMEOUT no longer counts as explicitly set, so chat completion() keeps its 600s fallback
* fix(guardrails): map blank litellm_params to the field default so PATCH can still clear a value
Dropping blank keys before validation made PATCH /guardrails/{id} with api_base "" a silent no-op. Blank values of optional fields now resolve to the field default and stay in the set fields, so the PATCH clears the stored value. Required fields and unknown keys keep the value as given
* test: cover proxy numeric knobs in-process and expect no warning for a blank Straiker api_version
The three proxy import-time numeric knobs move into litellm/constants.py so the existing in-process blank-env test covers them, replacing the subprocess test. A blank api_version now means unset, so the Straiker test asserts the key-prefix version with no warning.
* fix(guardrails): keep LitellmParams kwargs typed without adding argument-type errors
---------
Co-authored-by: yuneng-jiang
---
litellm/constants.py | 616 +++++++++---------
litellm/litellm_core_utils/env_utils.py | 82 ++-
litellm/proxy/agent_endpoints/endpoints.py | 5 +-
litellm/proxy/utils.py | 17 +-
litellm/types/guardrails.py | 25 +-
.../unit/litellm_core_utils/test_env_utils.py | 96 +++
.../guardrail_hooks/test_straiker.py | 21 +-
.../guardrails/test_guardrail_endpoints.py | 25 +
.../proxy/guardrails/test_init_guardrails.py | 38 ++
tests/unit/test_constants.py | 126 +++-
tests/unit/types/test_guardrails.py | 21 +
11 files changed, 679 insertions(+), 393 deletions(-)
create mode 100644 tests/unit/litellm_core_utils/test_env_utils.py
diff --git a/litellm/constants.py b/litellm/constants.py
index 9d18fe5489d..4ac6adaba73 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -4,7 +4,13 @@ from enum import Enum
from types import MappingProxyType
from typing import Final, Literal
-from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_range, get_env_int_or_none
+from litellm.litellm_core_utils.env_utils import (
+ get_env_float,
+ get_env_float_or_none,
+ get_env_int,
+ get_env_int_in_range,
+ get_env_int_or_none,
+)
MICROSOFT_GRAPH_BETA_BASE: Final = "https://graph.microsoft.com/beta"
OAUTH_TOKEN_EXCHANGE_CACHE_SAFETY_MARGIN_SECONDS: Final = 60
@@ -32,7 +38,7 @@ DEFER_PYDANTIC_BUILD: Final = os.getenv("DEFER_PYDANTIC_BUILD", "true") in ("tru
DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"})
-ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5))
+ROUTER_MAX_FALLBACKS: Final = get_env_int("ROUTER_MAX_FALLBACKS", 5)
ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000
RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(
{
@@ -68,10 +74,10 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
"auto_router_capability_limit",
}
)
-DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
-DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
+DEFAULT_BATCH_SIZE: Final = get_env_int("DEFAULT_BATCH_SIZE", 512)
+DEFAULT_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)
CLICKHOUSE_BATCH_SIZE: Final = get_env_int("CLICKHOUSE_BATCH_SIZE", 10_000)
-CLICKHOUSE_FLUSH_INTERVAL_SECONDS: Final = float(os.getenv("CLICKHOUSE_FLUSH_INTERVAL_SECONDS", "1.0"))
+CLICKHOUSE_FLUSH_INTERVAL_SECONDS: Final = get_env_float("CLICKHOUSE_FLUSH_INTERVAL_SECONDS", 1.0)
CLICKHOUSE_MAX_BUFFERED_ROWS: Final = get_env_int("CLICKHOUSE_MAX_BUFFERED_ROWS", 200_000)
CLICKHOUSE_MAX_RETRIES: Final = get_env_int("CLICKHOUSE_MAX_RETRIES", 3)
DEFAULT_CLICKHOUSE_DATABASE: Final = "litellm"
@@ -89,9 +95,9 @@ LENS_DATASET_MAX_CASE_CHARS: Final = get_env_int("LENS_DATASET_MAX_CASE_CHARS",
LENS_DATASET_TRACE_PAGE_SIZE: Final = 500
LENS_FEEDBACK_MAX_COMMENT_CHARS: Final = 10_000
LENS_FEEDBACK_MAX_SCORE: Final = 10
-DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))
-DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512))
-DEFAULT_S3_MAX_CONCURRENT_UPLOADS: Final = int(os.getenv("DEFAULT_S3_MAX_CONCURRENT_UPLOADS", "16"))
+DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)
+DEFAULT_S3_BATCH_SIZE: Final = get_env_int("DEFAULT_S3_BATCH_SIZE", 512)
+DEFAULT_S3_MAX_CONCURRENT_UPLOADS: Final = get_env_int("DEFAULT_S3_MAX_CONCURRENT_UPLOADS", 16)
DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY: Final = get_env_int("DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY", 200)
S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS: Final = get_env_int_in_range("S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS", 30, 1, 3600)
# https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
@@ -103,48 +109,43 @@ MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY"
S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY"
MAX_FILE_LIST_LIMIT: Final = 10000
-DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
-DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
+DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)
+DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = get_env_int("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)
GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10
budget_reservation_disabled_info_emitted = False
-DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
-DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
+DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = get_env_int("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)
+DEFAULT_SQS_BATCH_SIZE: Final = get_env_int("DEFAULT_SQS_BATCH_SIZE", 512)
SQS_SEND_MESSAGE_ACTION: Final = "SendMessage"
SQS_API_VERSION: Final = "2012-11-05"
-DEFAULT_MAX_RETRIES: Final = int(os.getenv("DEFAULT_MAX_RETRIES", 2))
+DEFAULT_MAX_RETRIES: Final = get_env_int("DEFAULT_MAX_RETRIES", 2)
# Max records accepted in one POST /v1/callbacks/logs batch. Bounds the blast
# radius: each record fans out to spend logs + every callback integration.
MAX_CALLBACK_LOG_RECORDS: Final = 1000
-DEFAULT_MAX_RECURSE_DEPTH: Final = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH", 100))
+DEFAULT_MAX_RECURSE_DEPTH: Final = get_env_int("DEFAULT_MAX_RECURSE_DEPTH", 100)
GUARDRAIL_ROTATION_ATTEMPTS: Final = 3
-DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER", 10))
-DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = float(
- os.getenv("DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5)
+DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = get_env_int("DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER", 10)
+DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = get_env_float(
+ "DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5
) # default cooldown a deployment if 50% of requests fail in a given minute
-DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096))
-DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3))
-DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1))
-DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5))
-DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float(
- os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1")
-)
-DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5))
-DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1))
-DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
+DEFAULT_MAX_TOKENS: Final = get_env_int("DEFAULT_MAX_TOKENS", 4096)
+DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3)
+DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1)
+DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5)
+DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = get_env_float("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", 1.0)
+DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLING_RETRIES", 5)
+DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)
+DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250)
HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
# Maximum wall-clock seconds a streaming response is allowed to run.
# Streams exceeding this duration are terminated with a Timeout error.
# None (default) = no limit. Set env var to a number of seconds to enable globally.
-_max_stream_duration_env: Final = os.getenv("LITELLM_MAX_STREAMING_DURATION_SECONDS", None)
-LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = (
- float(_max_stream_duration_env) if _max_stream_duration_env is not None else None
-)
+LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS")
# Maximum number of base64 characters to keep in logging payloads.
# Data URIs exceeding this are replaced with a size placeholder.
# Set to 0 to disable truncation.
-MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
+MAX_BASE64_LENGTH_FOR_LOGGING: Final = get_env_int("MAX_BASE64_LENGTH_FOR_LOGGING", 64)
BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: Final = 256 * 1024
REDACTED_BY_LITELLM: Final = "redacted-by-litellm"
# in-memory stand-in handed to provider converters for redacted arguments; never stored
@@ -158,44 +159,44 @@ MAX_BASE64_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_BASE64_LENGTH_STDOUT_LOG"
LITELLM_DETAILED_TIMING: Final = os.getenv("LITELLM_DETAILED_TIMING", "false").lower() == "true"
# Model cost map validation constants
-MODEL_COST_MAP_MIN_MODEL_COUNT: Final = int(
- os.getenv("MODEL_COST_MAP_MIN_MODEL_COUNT", 50)
+MODEL_COST_MAP_MIN_MODEL_COUNT: Final = get_env_int(
+ "MODEL_COST_MAP_MIN_MODEL_COUNT", 50
) # Minimum number of models a fetched cost map must contain to be considered valid
-MODEL_COST_MAP_MAX_SHRINK_RATIO: Final = float(
- os.getenv("MODEL_COST_MAP_MAX_SHRINK_RATIO", 0.5)
+MODEL_COST_MAP_MAX_SHRINK_RATIO: Final = get_env_float(
+ "MODEL_COST_MAP_MAX_SHRINK_RATIO", 0.5
) # Maximum allowed shrinkage ratio vs local backup (0.5 = reject if fetched map is <50% of backup)
-DEFAULT_IMAGE_WIDTH: Final = int(os.getenv("DEFAULT_IMAGE_WIDTH", 300))
-DEFAULT_IMAGE_HEIGHT: Final = int(os.getenv("DEFAULT_IMAGE_HEIGHT", 300))
+DEFAULT_IMAGE_WIDTH: Final = get_env_int("DEFAULT_IMAGE_WIDTH", 300)
+DEFAULT_IMAGE_HEIGHT: Final = get_env_int("DEFAULT_IMAGE_HEIGHT", 300)
# Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit)
# This prevents memory issues from downloading very large images
# Maps to OpenAI's 50 MB payload limit - requests with images exceeding this size will be rejected
# Set MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0 to disable image URL handling entirely
-MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: Final = float(os.getenv("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 50))
-MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB: Final = int(
- os.getenv("MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB", 1024)
+MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: Final = get_env_float("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 50.0)
+MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB: Final = get_env_int(
+ "MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB", 1024
) # 1MB = 1024KB
# Surrogate-repair fallback in _read_request_body runs two full-body re.sub passes
# that block the event loop on multi-MB malformed bodies. Skip the repair above this
# size and raise the existing 400 immediately. Set to 0 to disable the cap.
MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB: Final = get_env_int("MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB", 1)
-SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD: Final = int(
- os.getenv("SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD", 1000)
+SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD: Final = get_env_int(
+ "SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD", 1000
) # Minimum number of requests to consider "reasonable traffic". Used for single-deployment cooldown logic.
-DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS: Final = int(
- os.getenv("DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS", 5)
+DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS: Final = get_env_int(
+ "DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS", 5
) # Minimum number of requests before applying error rate cooldown. Prevents cooldown from triggering on first failure.
-DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET", 0))
+DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET = get_env_int("DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET", 0)
# MCP Semantic Tool Filter Defaults
DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL: Final = str(
os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL", "text-embedding-3-small")
)
-DEFAULT_MCP_SEMANTIC_FILTER_TOP_K: Final = int(os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_TOP_K", 10))
-DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float(
- os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3)
+DEFAULT_MCP_SEMANTIC_FILTER_TOP_K: Final = get_env_int("DEFAULT_MCP_SEMANTIC_FILTER_TOP_K", 10)
+DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = get_env_float(
+ "DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3
)
-MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150))
+MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = get_env_int("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)
MAX_LITELLM_CALL_ID_LENGTH: Final = 256
MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048
@@ -205,15 +206,15 @@ DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000
DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL: Final = str(
os.getenv("DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL", "text-embedding-3-small")
)
-DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD", 0.75))
+DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = get_env_float("DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD", 0.75)
DEFAULT_OPENAI_MODERATIONS_MODEL: Final = "omni-moderation-latest"
# MCP OAuth2 Client Credentials Defaults
-MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60"))
-MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200"))
-MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600"))
-MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = int(os.getenv("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", "60"))
+MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", 60)
+MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = get_env_int("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", 200)
+MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = get_env_int("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", 3600)
+MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = get_env_int("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", 60)
# mcp_tool_permissions entry that grants every current and future tool on a server
MCP_ALL_TOOLS_WILDCARD: Final = "*"
@@ -222,20 +223,20 @@ MCP_ALL_TOOLS_WILDCARD: Final = "*"
# npm/npx needs a writable cache dir; in containers the default (~/.npm)
# may not exist or be read-only. /tmp is always writable.
MCP_NPM_CACHE_DIR: Final = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache")
-MCP_OAUTH2_TOKEN_CACHE_MIN_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10"))
+MCP_OAUTH2_TOKEN_CACHE_MIN_TTL: Final = get_env_int("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", 10)
# Per-user OAuth token Redis cache (for server-side token storage)
MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX: Final = "mcp:per_user_token"
-MCP_PER_USER_TOKEN_DEFAULT_TTL: Final = int(
- os.getenv("MCP_PER_USER_TOKEN_DEFAULT_TTL", "43200") # 12 hours
-)
-MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60"))
+MCP_PER_USER_TOKEN_DEFAULT_TTL: Final = get_env_int("MCP_PER_USER_TOKEN_DEFAULT_TTL", 43200)
+MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", 60)
# MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
-MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"))
-MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
-MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
-MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
+MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0)
+MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0)
+MCP_METADATA_TIMEOUT: Final = get_env_float("LITELLM_MCP_METADATA_TIMEOUT", 10.0)
+MCP_HEALTH_CHECK_TIMEOUT: Final = get_env_float("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", 10.0)
+AGENT_HEALTH_CHECK_TIMEOUT_SECONDS: Final = get_env_float("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", 5.0)
+AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS: Final = get_env_float("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", 30.0)
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS: Final = 5.0
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8
@@ -253,7 +254,7 @@ MCP_STDIO_ALLOWED_COMMANDS: Final[frozenset] = frozenset(
)
# MCP OAuth2 Token Exchange (OBO) Defaults
-MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", "500"))
+MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE: Final = get_env_int("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", 500)
LITELLM_UI_ALLOW_HEADERS: Final = [
"x-litellm-semantic-filter",
@@ -266,14 +267,14 @@ LITELLM_UI_ALLOW_HEADERS: Final = [
]
# Gemini model-specific minimal thinking budget constants
-DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH: Final = int(
- os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH", 1)
+DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH: Final = get_env_int(
+ "DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH", 1
)
-DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO: Final = int(
- os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO", 128)
+DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO: Final = get_env_int(
+ "DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO", 128
)
-DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE: Final = int(
- os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE", 512)
+DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE: Final = get_env_int(
+ "DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE", 512
)
# Maximum number of callbacks that can be registered
@@ -295,8 +296,8 @@ PROXY_LLM_PROVIDER_FALLBACK: Final = "litellm_proxy"
PROXY_REJECTED_BEFORE_ROUTING_KEY: Final = "proxy_rejected_before_routing"
# Generic fallback for unknown models
-DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int(
- os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)
+DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = get_env_int(
+ "DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128
)
# Provider-specific API base URLs
@@ -315,19 +316,19 @@ OPEN_SANDBOX_DEFAULT_TIMEOUT: Final = 300
OPEN_SANDBOX_READY_TIMEOUT: Final = 30.0
OPEN_SANDBOX_POLL_INTERVAL: Final = 0.2
-DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024))
-DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: Final = int(
- os.getenv("DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET", 2048)
+DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = get_env_int("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024)
+DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: Final = get_env_int(
+ "DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET", 2048
)
-DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET", 4096))
-DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET", 8192))
-DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET", 16384))
-MAX_TOKEN_TRIMMING_ATTEMPTS: Final = int(
- os.getenv("MAX_TOKEN_TRIMMING_ATTEMPTS", 10)
+DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET = get_env_int("DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET", 4096)
+DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET = get_env_int("DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET", 8192)
+DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET = get_env_int("DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET", 16384)
+MAX_TOKEN_TRIMMING_ATTEMPTS: Final = get_env_int(
+ "MAX_TOKEN_TRIMMING_ATTEMPTS", 10
) # Maximum number of attempts to trim the message
RUNWAYML_DEFAULT_API_VERSION: Final = str(os.getenv("RUNWAYML_DEFAULT_API_VERSION", "2024-11-06"))
-RUNWAYML_POLLING_TIMEOUT = int(os.getenv("RUNWAYML_POLLING_TIMEOUT", 600)) # 10 minutes default for image generation
+RUNWAYML_POLLING_TIMEOUT = get_env_int("RUNWAYML_POLLING_TIMEOUT", 600) # 10 minutes default for image generation
########## Networking constants ##############################################################
_DEFAULT_TTL_FOR_HTTPX_CLIENTS: Final = 3600 # 1 hour, re-use the same httpx client for 1 hour
@@ -344,10 +345,10 @@ EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING: Final = 10_000
# Aiohttp connection pooling - prevents memory leaks from unbounded connection growth
# Set to 0 for unlimited (not recommended for production)
-AIOHTTP_CONNECTOR_LIMIT: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000))
-AIOHTTP_CONNECTOR_LIMIT_PER_HOST: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500))
-AIOHTTP_KEEPALIVE_TIMEOUT: Final = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 15))
-AIOHTTP_TTL_DNS_CACHE: Final = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
+AIOHTTP_CONNECTOR_LIMIT: Final = get_env_int("AIOHTTP_CONNECTOR_LIMIT", 1000)
+AIOHTTP_CONNECTOR_LIMIT_PER_HOST: Final = get_env_int("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500)
+AIOHTTP_KEEPALIVE_TIMEOUT: Final = get_env_int("AIOHTTP_KEEPALIVE_TIMEOUT", 15)
+AIOHTTP_TTL_DNS_CACHE: Final = get_env_int("AIOHTTP_TTL_DNS_CACHE", 300)
# TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs
# whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT
# Gateway: 350s vs OpenAI/Azure: 600s). Without this, the kernel sends nothing
@@ -355,9 +356,9 @@ AIOHTTP_TTL_DNS_CACHE: Final = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
# arrives. Enabling SO_KEEPALIVE makes the kernel emit TCP probes that reset
# the NAT idle timer.
AIOHTTP_SO_KEEPALIVE: Final = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true"
-AIOHTTP_TCP_KEEPIDLE: Final = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60))
-AIOHTTP_TCP_KEEPINTVL: Final = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30))
-AIOHTTP_TCP_KEEPCNT: Final = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5))
+AIOHTTP_TCP_KEEPIDLE: Final = get_env_int("AIOHTTP_TCP_KEEPIDLE", 60)
+AIOHTTP_TCP_KEEPINTVL: Final = get_env_int("AIOHTTP_TCP_KEEPINTVL", 30)
+AIOHTTP_TCP_KEEPCNT: Final = get_env_int("AIOHTTP_TCP_KEEPCNT", 5)
# enable_cleanup_closed is only needed for Python versions with the SSL leak bug
# Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960)
# Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78
@@ -370,10 +371,9 @@ AIOHTTP_NEEDS_CLEANUP_CLOSED: Final = (3, 13, 0) <= sys.version_info < (
# WebSocket constants
# Default to None (unlimited) to match OpenAI's official agents SDK behavior
# https://github.com/openai/openai-agents-python/blob/cf1b933660e44fd37b4350c41febab8221801409/src/agents/realtime/openai_realtime.py#L235
-_max_size_env: Final = os.getenv("REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES")
-REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = int(_max_size_env) if _max_size_env is not None else None
-REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float(
- os.getenv("REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", "20.0")
+REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = get_env_int_or_none("REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES")
+REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = get_env_float(
+ "REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 20.0
)
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
@@ -439,15 +439,16 @@ REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_sp
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer"
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer"
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer"
-MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
+MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = get_env_int("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)
REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer"
REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000
REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT: Final = 1000
# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth
-LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000))
-TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60))
-GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int(
- os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60)
+LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = get_env_int("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)
+TOOL_POLICY_CACHE_TTL_SECONDS: Final = get_env_int("TOOL_POLICY_CACHE_TTL_SECONDS", 60)
+LITELLM_CONFIG_CACHE_TTL_SECONDS: Final = get_env_int("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", 60)
+GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = get_env_int(
+ "GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60
)
BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000
CONTENT_FILTER_STREAMING_HOLDBACK_CHARS: Final = 50
@@ -457,8 +458,8 @@ PRESIDIO_ANALYZE_CHUNK_OVERLAP_CHARS: Final = 4096
PRESIDIO_ANALYZE_CHUNK_CONCURRENCY: Final = 8
# Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
# Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
-MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8)))
-MAX_IN_MEMORY_QUEUE_FLUSH_COUNT: Final = int(os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000))
+MAX_SIZE_IN_MEMORY_QUEUE: Final = get_env_int("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))
+MAX_IN_MEMORY_QUEUE_FLUSH_COUNT: Final = get_env_int("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000)
###############################################################################################
# Providers will not cache a prefix below a minimum size. That minimum is per-model, not global:
# Anthropic's ranges from 512 to 4096 depending on the model, and can differ per platform for the
@@ -473,22 +474,22 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = (
else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
)
PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20
-DEFAULT_TRIM_RATIO: Final = float(
- os.getenv("DEFAULT_TRIM_RATIO", 0.75)
+DEFAULT_TRIM_RATIO: Final = get_env_float(
+ "DEFAULT_TRIM_RATIO", 0.75
) # default ratio of tokens to trim from the end of a prompt
-HOURS_IN_A_DAY: Final = int(os.getenv("HOURS_IN_A_DAY", 24))
-DAYS_IN_A_WEEK: Final = int(os.getenv("DAYS_IN_A_WEEK", 7))
-DAYS_IN_A_MONTH: Final = int(os.getenv("DAYS_IN_A_MONTH", 28))
-DAYS_IN_A_YEAR: Final = int(os.getenv("DAYS_IN_A_YEAR", 365))
-REPLICATE_MODEL_NAME_WITH_ID_LENGTH: Final = int(os.getenv("REPLICATE_MODEL_NAME_WITH_ID_LENGTH", 64))
+HOURS_IN_A_DAY: Final = get_env_int("HOURS_IN_A_DAY", 24)
+DAYS_IN_A_WEEK: Final = get_env_int("DAYS_IN_A_WEEK", 7)
+DAYS_IN_A_MONTH: Final = get_env_int("DAYS_IN_A_MONTH", 28)
+DAYS_IN_A_YEAR: Final = get_env_int("DAYS_IN_A_YEAR", 365)
+REPLICATE_MODEL_NAME_WITH_ID_LENGTH: Final = get_env_int("REPLICATE_MODEL_NAME_WITH_ID_LENGTH", 64)
#### TOKEN COUNTING ####
-FUNCTION_DEFINITION_TOKEN_COUNT: Final = int(os.getenv("FUNCTION_DEFINITION_TOKEN_COUNT", 9))
-SYSTEM_MESSAGE_TOKEN_COUNT: Final = int(os.getenv("SYSTEM_MESSAGE_TOKEN_COUNT", 4))
-TOOL_CHOICE_OBJECT_TOKEN_COUNT: Final = int(os.getenv("TOOL_CHOICE_OBJECT_TOKEN_COUNT", 4))
-DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT", 10))
-DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20))
-MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768))
-MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000))
+FUNCTION_DEFINITION_TOKEN_COUNT: Final = get_env_int("FUNCTION_DEFINITION_TOKEN_COUNT", 9)
+SYSTEM_MESSAGE_TOKEN_COUNT: Final = get_env_int("SYSTEM_MESSAGE_TOKEN_COUNT", 4)
+TOOL_CHOICE_OBJECT_TOKEN_COUNT: Final = get_env_int("TOOL_CHOICE_OBJECT_TOKEN_COUNT", 4)
+DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT: Final = get_env_int("DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT", 10)
+DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT: Final = get_env_int("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20)
+MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES: Final = get_env_int("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768)
+MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES: Final = get_env_int("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000)
# tiktoken's BPE merge loop is quadratic in the length of a single regex piece, so a long run of one
# repeated character (dot leaders, whitespace, zero-padded base64) can take minutes on a multi-MB payload.
# Encoding in chunks makes the cost linear, at a drift of at most ~1 token per chunk boundary. The upper
@@ -512,69 +513,63 @@ TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range(
minimum=1,
maximum=256,
)
-MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
-MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
-OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))
+MAX_TILE_WIDTH: Final = get_env_int("MAX_TILE_WIDTH", 512)
+MAX_TILE_HEIGHT: Final = get_env_int("MAX_TILE_HEIGHT", 512)
+OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = get_env_float("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)
GROQ_BROWSER_VISIT_WEBSITE_COST_PER_CALL: Final = 1.0 / 1000
# Azure OpenAI Assistants feature costs
# Source: https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/
-AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY: Final = float(
- os.getenv("AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY", 0.1) # $0.1 USD per 1 GB/Day
+AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY: Final = get_env_float("AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY", 0.1)
+AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS: Final = get_env_float("AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS", 3.0)
+AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS: Final = get_env_float(
+ "AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS", 12.0
)
-AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS: Final = float(
- os.getenv("AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS", 3.0) # $0.003 USD per 1K Tokens
-)
-AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS: Final = float(
- os.getenv("AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS", 12.0) # $0.012 USD per 1K Tokens
-)
-AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY: Final = float(
- os.getenv("AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY", 0.1) # $0.1 USD per 1 GB/Day (same as file search)
-)
-MIN_NON_ZERO_TEMPERATURE: Final = float(os.getenv("MIN_NON_ZERO_TEMPERATURE", 0.0001))
+AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY: Final = get_env_float("AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY", 0.1)
+MIN_NON_ZERO_TEMPERATURE: Final = get_env_float("MIN_NON_ZERO_TEMPERATURE", 0.0001)
#### RELIABILITY ####
-REPEATED_STREAMING_CHUNK_LIMIT: Final = int(
- os.getenv("REPEATED_STREAMING_CHUNK_LIMIT", 100)
+REPEATED_STREAMING_CHUNK_LIMIT: Final = get_env_int(
+ "REPEATED_STREAMING_CHUNK_LIMIT", 100
) # catch if model starts looping the same chunk while streaming. Uses high default to prevent false positives.
# Shared maxsize for functools.lru_cache usage across hot paths.
# Defaulted to 64 to avoid cache thrash in multi-model production workloads.
-DEFAULT_MAX_LRU_CACHE_SIZE: Final = int(os.getenv("DEFAULT_MAX_LRU_CACHE_SIZE", 64))
+DEFAULT_MAX_LRU_CACHE_SIZE: Final = get_env_int("DEFAULT_MAX_LRU_CACHE_SIZE", 64)
_REALTIME_BODY_CACHE_SIZE = 1000 # Keep realtime helper caches bounded; workloads rarely exceed 1k models/intents
-INITIAL_RETRY_DELAY: Final = float(os.getenv("INITIAL_RETRY_DELAY", 0.5))
-MAX_RETRY_DELAY: Final = float(os.getenv("MAX_RETRY_DELAY", 8.0))
-JITTER: Final = float(os.getenv("JITTER", 0.75))
-DEFAULT_IN_MEMORY_TTL = int(os.getenv("DEFAULT_IN_MEMORY_TTL", 5)) # default time to live for the in-memory cache
-DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE: Final = int(
- os.getenv("DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 1000)
+INITIAL_RETRY_DELAY: Final = get_env_float("INITIAL_RETRY_DELAY", 0.5)
+MAX_RETRY_DELAY: Final = get_env_float("MAX_RETRY_DELAY", 8.0)
+JITTER: Final = get_env_float("JITTER", 0.75)
+DEFAULT_IN_MEMORY_TTL = get_env_int("DEFAULT_IN_MEMORY_TTL", 5) # default time to live for the in-memory cache
+DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE: Final = get_env_int(
+ "DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 1000
) # default max size for redis batch cache
-DEFAULT_POLLING_INTERVAL: Final = float(
- os.getenv("DEFAULT_POLLING_INTERVAL", 0.03)
+DEFAULT_POLLING_INTERVAL: Final = get_env_float(
+ "DEFAULT_POLLING_INTERVAL", 0.03
) # default polling interval for the scheduler
-AZURE_OPERATION_POLLING_TIMEOUT: Final = int(os.getenv("AZURE_OPERATION_POLLING_TIMEOUT", 120))
+AZURE_OPERATION_POLLING_TIMEOUT: Final = get_env_int("AZURE_OPERATION_POLLING_TIMEOUT", 120)
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: Final = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30"))
-AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = int(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96))
-REDIS_SOCKET_TIMEOUT: Final = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1))
+AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = get_env_int("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96)
+REDIS_SOCKET_TIMEOUT: Final = get_env_float("REDIS_SOCKET_TIMEOUT", 0.1)
CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS: Final[float] = 5.0
-REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5))
-REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5))
-REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60))
+REDIS_CONNECTION_POOL_TIMEOUT: Final = get_env_int("REDIS_CONNECTION_POOL_TIMEOUT", 5)
+REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = get_env_int("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5)
+REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = get_env_int("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60)
REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true"
# minimum seconds a timeout-only failure streak must span before it can open the breaker,
# so one event-loop stall timing out many queued calls at once does not trip it
-REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0))
-REDIS_TIMEOUT_LOG_INTERVAL: Final = float(os.getenv("REDIS_TIMEOUT_LOG_INTERVAL", "5.0"))
+REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = get_env_float("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0)
+REDIS_TIMEOUT_LOG_INTERVAL: Final = get_env_float("REDIS_TIMEOUT_LOG_INTERVAL", 5.0)
# Seconds of idle before a Redis cluster connection is validated with a PING and
# reconnected if dead, so a connection silently dropped by a cluster restart
# (e.g. ElastiCache Serverless maintenance) is not reused while broken
REDIS_CLUSTER_HEALTH_CHECK_INTERVAL: Final = 25
# Default Redis major version to assume when version cannot be determined
# Using 7 as it's the modern version that supports LPOP with count parameter
-DEFAULT_REDIS_MAJOR_VERSION: Final = int(os.getenv("DEFAULT_REDIS_MAJOR_VERSION", 7))
-NON_LLM_CONNECTION_TIMEOUT: Final = int(
- os.getenv("NON_LLM_CONNECTION_TIMEOUT", 15)
+DEFAULT_REDIS_MAJOR_VERSION: Final = get_env_int("DEFAULT_REDIS_MAJOR_VERSION", 7)
+NON_LLM_CONNECTION_TIMEOUT: Final = get_env_int(
+ "NON_LLM_CONNECTION_TIMEOUT", 15
) # timeout for adjacent services (e.g. jwt auth)
-MAX_EXCEPTION_MESSAGE_LENGTH: Final = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000))
-MAX_STRING_LENGTH_PROMPT_IN_DB: Final = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048))
-BEDROCK_MAX_POLICY_SIZE: Final = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75))
+MAX_EXCEPTION_MESSAGE_LENGTH: Final = get_env_int("MAX_EXCEPTION_MESSAGE_LENGTH", 2000)
+MAX_STRING_LENGTH_PROMPT_IN_DB: Final = get_env_int("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048)
+BEDROCK_MAX_POLICY_SIZE: Final = get_env_int("BEDROCK_MAX_POLICY_SIZE", 75)
# One entry per distinct AWS credential-argument set. Per-user cost attribution passes the attributed
# identity as aws_session_name, so this bounds how many attributed identities keep a cached STS session.
BEDROCK_IAM_CACHE_MAX_ENTRIES: Final = 1000
@@ -584,27 +579,27 @@ BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES: Final = 64
# Retire a cached STS credential this many seconds before AWS expires it, so a request that reads it
# still has a usable credential for the whole call.
STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS: Final = 60
-BEDROCK_MIN_THINKING_BUDGET_TOKENS: Final = int(os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024))
+BEDROCK_MIN_THINKING_BUDGET_TOKENS: Final = get_env_int("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024)
# Anthropic's Messages API rejects thinking.budget_tokens < 1024.
ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: Final = 1024
-REPLICATE_POLLING_DELAY_SECONDS: Final = float(os.getenv("REPLICATE_POLLING_DELAY_SECONDS", 0.5))
-DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096))
+REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_float("REPLICATE_POLLING_DELAY_SECONDS", 0.5)
+DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS: Final = get_env_int("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096)
DEFAULT_OCI_CHAT_MAX_TOKENS: Final = 4096
-TOGETHER_AI_4_B: Final = int(os.getenv("TOGETHER_AI_4_B", 4))
-TOGETHER_AI_8_B: Final = int(os.getenv("TOGETHER_AI_8_B", 8))
-TOGETHER_AI_21_B: Final = int(os.getenv("TOGETHER_AI_21_B", 21))
-TOGETHER_AI_41_B: Final = int(os.getenv("TOGETHER_AI_41_B", 41))
-TOGETHER_AI_80_B: Final = int(os.getenv("TOGETHER_AI_80_B", 80))
-TOGETHER_AI_110_B: Final = int(os.getenv("TOGETHER_AI_110_B", 110))
-TOGETHER_AI_EMBEDDING_150_M: Final = int(os.getenv("TOGETHER_AI_EMBEDDING_150_M", 150))
-TOGETHER_AI_EMBEDDING_350_M: Final = int(os.getenv("TOGETHER_AI_EMBEDDING_350_M", 350))
-QDRANT_SCALAR_QUANTILE: Final = float(os.getenv("QDRANT_SCALAR_QUANTILE", 0.99))
-QDRANT_VECTOR_SIZE: Final = int(os.getenv("QDRANT_VECTOR_SIZE", 1536))
-CACHED_STREAMING_CHUNK_DELAY: Final = float(os.getenv("CACHED_STREAMING_CHUNK_DELAY", 0.02))
-AUDIO_SPEECH_CHUNK_SIZE: Final = int(
- os.getenv("AUDIO_SPEECH_CHUNK_SIZE", 8192)
+TOGETHER_AI_4_B: Final = get_env_int("TOGETHER_AI_4_B", 4)
+TOGETHER_AI_8_B: Final = get_env_int("TOGETHER_AI_8_B", 8)
+TOGETHER_AI_21_B: Final = get_env_int("TOGETHER_AI_21_B", 21)
+TOGETHER_AI_41_B: Final = get_env_int("TOGETHER_AI_41_B", 41)
+TOGETHER_AI_80_B: Final = get_env_int("TOGETHER_AI_80_B", 80)
+TOGETHER_AI_110_B: Final = get_env_int("TOGETHER_AI_110_B", 110)
+TOGETHER_AI_EMBEDDING_150_M: Final = get_env_int("TOGETHER_AI_EMBEDDING_150_M", 150)
+TOGETHER_AI_EMBEDDING_350_M: Final = get_env_int("TOGETHER_AI_EMBEDDING_350_M", 350)
+QDRANT_SCALAR_QUANTILE: Final = get_env_float("QDRANT_SCALAR_QUANTILE", 0.99)
+QDRANT_VECTOR_SIZE: Final = get_env_int("QDRANT_VECTOR_SIZE", 1536)
+CACHED_STREAMING_CHUNK_DELAY: Final = get_env_float("CACHED_STREAMING_CHUNK_DELAY", 0.02)
+AUDIO_SPEECH_CHUNK_SIZE: Final = get_env_int(
+ "AUDIO_SPEECH_CHUNK_SIZE", 8192
) # chunk_size for audio speech streaming. Balance between latency and memory usage
-DEFAULT_MAX_TOKENS_FOR_TRITON: Final = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000))
+DEFAULT_MAX_TOKENS_FOR_TRITON: Final = get_env_int("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000)
#### Networking settings ####
# Sentinel used when `REQUEST_TIMEOUT` is unset: `litellm.request_timeout` keeps this
# value so longer-running surfaces (Router `timeout or litellm.request_timeout`,
@@ -617,12 +612,13 @@ DEFAULT_REQUEST_TIMEOUT_SECONDS: Final[float] = 6000.0
# deadline and connect handshake (see ``http_handler`` cached handler paths).
COMPLETION_HTTP_FALLBACK_SECONDS: Final[float] = 600.0
HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: Final[float] = 5.0
-SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS: Final[float] = float(
- os.getenv("SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS", "5.0")
+SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS: Final[float] = get_env_float("SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS", 5.0)
+_request_timeout_from_env: Final = get_env_float_or_none("REQUEST_TIMEOUT")
+request_timeout: float = (
+ DEFAULT_REQUEST_TIMEOUT_SECONDS if _request_timeout_from_env is None else _request_timeout_from_env
)
-request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS))))
-request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ
-DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes
+request_timeout_explicitly_set: bool = _request_timeout_from_env is not None
+DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = get_env_float("DEFAULT_A2A_AGENT_TIMEOUT", 6000.0) # 10 minutes
AGENT_KILL_SWITCH_TIMEOUT_SECONDS: Final = 10.0
AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS: Final = 2000
# Patterns that indicate a localhost/internal URL in A2A agent cards that should be
@@ -646,20 +642,20 @@ STREAM_SSE_DATA_PREFIX: Final[str] = "data: "
STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n'
STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8")
### SPEND TRACKING ###
-DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = float(
- os.getenv("DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400)
+DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = get_env_float(
+ "DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400
) # price per second for a100 80GB
-FIREWORKS_AI_56_B_MOE: Final = int(os.getenv("FIREWORKS_AI_56_B_MOE", 56))
-FIREWORKS_AI_176_B_MOE: Final = int(os.getenv("FIREWORKS_AI_176_B_MOE", 176))
-FIREWORKS_AI_4_B: Final = int(os.getenv("FIREWORKS_AI_4_B", 4))
-FIREWORKS_AI_16_B: Final = int(os.getenv("FIREWORKS_AI_16_B", 16))
-FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80))
+FIREWORKS_AI_56_B_MOE: Final = get_env_int("FIREWORKS_AI_56_B_MOE", 56)
+FIREWORKS_AI_176_B_MOE: Final = get_env_int("FIREWORKS_AI_176_B_MOE", 176)
+FIREWORKS_AI_4_B: Final = get_env_int("FIREWORKS_AI_4_B", 4)
+FIREWORKS_AI_16_B: Final = get_env_int("FIREWORKS_AI_16_B", 16)
+FIREWORKS_AI_80_B: Final = get_env_int("FIREWORKS_AI_80_B", 80)
# https://docs.fireworks.ai/guides/prompt-caching (accessed 2026-09-19): serverless cached prompt tokens
# default to a 50% discount off the input rate
FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO: Final = 0.5
#### Logging callback constants ####
REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM"
-MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50))
+MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = get_env_int("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)
LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS: Final = 10_000
# Backpressure + lifetime bounds for the /v1/messages streaming relay (see
# BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is
@@ -667,27 +663,25 @@ LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS: Final = 10_000
# buffer the whole response in memory; the detached-drain cap bounds how many
# post-disconnect drains may run concurrently so client behavior can't create
# unbounded worker state.
-ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int(
- os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024")
+ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = get_env_int(
+ "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 1024
)
# Setting this to 0 disables detached draining entirely: every post-disconnect
# pump bills whatever partial output it has already collected and aborts the
# upstream stream immediately, instead of continuing to drain for the real
# terminal usage.
-ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int(
- os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100")
-)
-LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0
-LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000))
-LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0))
+ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = get_env_int("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 100)
+LOGGING_WORKER_CONCURRENCY: Final = get_env_int("LOGGING_WORKER_CONCURRENCY", 100) # Must be above 0
+LOGGING_WORKER_MAX_QUEUE_SIZE: Final = get_env_int("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)
+LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = get_env_float("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)
LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS: Final = 5.0
-LOGGING_WORKER_CLEAR_PERCENTAGE: Final = int(
- os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50)
+LOGGING_WORKER_CLEAR_PERCENTAGE: Final = get_env_int(
+ "LOGGING_WORKER_CLEAR_PERCENTAGE", 50
) # Percentage of queue to clear (default: 50%)
-MAX_ITERATIONS_TO_CLEAR_QUEUE: Final = int(os.getenv("MAX_ITERATIONS_TO_CLEAR_QUEUE", 200))
-MAX_TIME_TO_CLEAR_QUEUE: Final = float(os.getenv("MAX_TIME_TO_CLEAR_QUEUE", 5.0))
-LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float(
- os.getenv("LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS", 0.5)
+MAX_ITERATIONS_TO_CLEAR_QUEUE: Final = get_env_int("MAX_ITERATIONS_TO_CLEAR_QUEUE", 200)
+MAX_TIME_TO_CLEAR_QUEUE: Final = get_env_float("MAX_TIME_TO_CLEAR_QUEUE", 5.0)
+LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = get_env_float(
+ "LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS", 0.5
) # Cooldown time in seconds before allowing another aggressive clear (default: 0.5s)
LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100)
LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000)
@@ -700,9 +694,9 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED: Final = 499
-EMAIL_BUDGET_ALERT_TTL: Final = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds
-EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = float(
- os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)
+EMAIL_BUDGET_ALERT_TTL: Final = get_env_int("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60) # 24 hours in seconds
+EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = get_env_float(
+ "EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8
) # 80% of max budget
############### LLM Provider Constants ###############
### ANTHROPIC CONSTANTS ###
@@ -727,11 +721,11 @@ LITELLM_WEB_SEARCH_TOOL_NAME: Final = "litellm_web_search"
DEFAULT_IMAGE_ENDPOINT_MODEL: Final = "dall-e-2"
DEFAULT_VIDEO_ENDPOINT_MODEL: Final = "sora-2"
-DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS: Final = int(os.getenv("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8))
+DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS: Final = get_env_int("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8)
### DATAFORSEO CONSTANTS ###
-DEFAULT_DATAFORSEO_LOCATION_CODE: Final = int(
- os.getenv("DEFAULT_DATAFORSEO_LOCATION_CODE", 2250)
+DEFAULT_DATAFORSEO_LOCATION_CODE: Final = get_env_int(
+ "DEFAULT_DATAFORSEO_LOCATION_CODE", 2250
) # Default to France (2250) - lower number, commonly used location
LITELLM_CHAT_PROVIDERS: Final = [
@@ -1593,20 +1587,20 @@ OPENAI_FINISH_REASONS: Final = [
"tool_calls",
"content_filter",
]
-HUMANLOOP_PROMPT_CACHE_TTL_SECONDS: Final = int(os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60)) # 1 minute
+HUMANLOOP_PROMPT_CACHE_TTL_SECONDS: Final = get_env_int("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60) # 1 minute
RESPONSE_FORMAT_TOOL_NAME = "json_tool_call" # default tool name used when converting response format to tool call
########################### Logging Callback Constants ###########################
AZURE_STORAGE_MSFT_VERSION: Final = "2019-07-07"
AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX: Final = "core.windows.net"
-PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES: Final = int(
- os.getenv("PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES", 5)
+PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES: Final = get_env_int(
+ "PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES", 5
)
PROMETHEUS_OVERFLOW_SERIES_LABEL_VALUE: Final = "other"
PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX: Final = "litellm_admitted_series_"
-CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = int(os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60))
+CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = get_env_int("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60)
MCP_TOOL_NAME_PREFIX: Final = "mcp_tool"
-MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100))
+MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = get_env_int("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)
PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: Final = 4096
# Headers to control callbacks
@@ -1644,11 +1638,11 @@ STANDARD_CUSTOMER_ID_HEADERS: Final = [
"x-litellm-customer-id",
"x-litellm-end-user-id",
]
-MAX_SPENDLOG_ROWS_TO_QUERY: Final = int(
- os.getenv("MAX_SPENDLOG_ROWS_TO_QUERY", 1_000_000)
+MAX_SPENDLOG_ROWS_TO_QUERY: Final = get_env_int(
+ "MAX_SPENDLOG_ROWS_TO_QUERY", 1_000_000
) # if spendLogs has more than 1M rows, do not query the DB
-DEFAULT_SOFT_BUDGET: Final = float(
- os.getenv("DEFAULT_SOFT_BUDGET", 50.0)
+DEFAULT_SOFT_BUDGET: Final = get_env_float(
+ "DEFAULT_SOFT_BUDGET", 50.0
) # by default all litellm proxy keys have a soft budget of 50.0
# makes it clear this is a rate limit error for a litellm virtual key
RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY: Final = "LiteLLM Virtual Key user_api_key_hash"
@@ -1715,38 +1709,16 @@ TRANSCRIBE_MEDIA_FETCH_ATTEMPTS: Final = 3
TRANSCRIBE_MEDIA_LAST_MODIFIED_TOLERANCE_SECONDS: Final = 1.0 # S3 Last-Modified carries whole seconds only
TRANSCRIBE_MEASURABLE_MEDIA_FORMATS: Final = frozenset({"flac", "mp3", "ogg", "wav"}) # what libsndfile can read
-BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour
-BATCH_STATUS_POLL_MAX_ATTEMPTS: Final = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours
+BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = get_env_int("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600) # 1 hour
+BATCH_STATUS_POLL_MAX_ATTEMPTS: Final = get_env_int("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24) # for 24 hours
BATCH_TPD_WINDOW_SECONDS: Final = 86400
BATCH_TPD_DESCRIPTOR_SUFFIX: Final = "_tpd"
-HEALTH_CHECK_TIMEOUT_SECONDS: Final = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds
-_background_health_check_max_tokens_env: Final = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS")
-try:
- _raw_background_health_check_max_tokens: Final = (
- _background_health_check_max_tokens_env.strip() if _background_health_check_max_tokens_env is not None else ""
- )
- BACKGROUND_HEALTH_CHECK_MAX_TOKENS: int | None = (
- int(_raw_background_health_check_max_tokens) if _raw_background_health_check_max_tokens else None
- )
-except (ValueError, TypeError):
- BACKGROUND_HEALTH_CHECK_MAX_TOKENS = None
-
-
-_background_health_check_max_tokens_reasoning_env: Final = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING")
-try:
- _raw_background_health_check_max_tokens_reasoning: Final = (
- _background_health_check_max_tokens_reasoning_env.strip()
- if _background_health_check_max_tokens_reasoning_env is not None
- else ""
- )
- BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING: int | None = (
- int(_raw_background_health_check_max_tokens_reasoning)
- if _raw_background_health_check_max_tokens_reasoning
- else None
- )
-except (ValueError, TypeError):
- BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING = None
+HEALTH_CHECK_TIMEOUT_SECONDS: Final = get_env_int("HEALTH_CHECK_TIMEOUT_SECONDS", 60) # 60 seconds
+BACKGROUND_HEALTH_CHECK_MAX_TOKENS: int | None = get_env_int_or_none("BACKGROUND_HEALTH_CHECK_MAX_TOKENS")
+BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING: int | None = get_env_int_or_none(
+ "BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING"
+)
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME: Final = "litellm-internal-health-check"
LITTELM_CLI_SERVICE_ACCOUNT_NAME: Final = "litellm-cli"
@@ -1771,22 +1743,22 @@ OTEL_SERVICE_NAME_METADATA_KEYS: Final = ("otel_service_name_override", "otel_se
# Key Rotation Constants
LITELLM_KEY_ROTATION_ENABLED: Final = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false")
-LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS: Final = int(
- os.getenv("LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS", 86400)
+LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS: Final = get_env_int(
+ "LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS", 86400
) # 24 hours default
LITELLM_KEY_ROTATION_GRACE_PERIOD: Final[str] = os.getenv(
"LITELLM_KEY_ROTATION_GRACE_PERIOD", ""
) # Duration to keep old key valid after rotation (e.g. "24h", "2d"); empty = immediate revoke (default)
-LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS: Final = int(
- os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600)
+LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS: Final = get_env_int(
+ "LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600
) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation
UI_SESSION_TOKEN_TEAM_ID: Final = "litellm-dashboard"
LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED = os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED", "false")
-LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS: Final = int(
- os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS", 86400)
+LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS: Final = get_env_int(
+ "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS", 86400
) # 24 hours default
-LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE: Final = int(
- os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000)
+LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE: Final = get_env_int(
+ "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000
)
LOGIN_THROTTLE_CACHE_KEY_PREFIX: Final = "login_fail"
LOGIN_THROTTLE_UNKNOWN_SOURCE: Final = "unknown"
@@ -1796,7 +1768,7 @@ LOGIN_THROTTLE_NOT_BLOCKED: Final = (0, 0)
LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id"
LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget"
GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend"
-LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, int(os.getenv("LITELLM_EXECUTED_BATCH_CONCURRENCY", "4")))
+LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, get_env_int("LITELLM_EXECUTED_BATCH_CONCURRENCY", 4))
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"
@@ -1805,8 +1777,8 @@ CLI_SSO_SESSION_CACHE_KEY_PREFIX: Final = "cli_sso_session"
CLI_SSO_SESSION_TTL_SECONDS: Final = 600
CLI_SESSION_KEY_PREFIX: Final = "cli-session"
# Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility
-CLI_JWT_EXPIRATION_HOURS: Final = int(
- os.getenv("CLI_JWT_EXPIRATION_HOURS") or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS") or 24
+CLI_JWT_EXPIRATION_HOURS: Final = get_env_int(
+ "CLI_JWT_EXPIRATION_HOURS", get_env_int("LITELLM_CLI_JWT_EXPIRATION_HOURS", 24)
)
# Comma-separated allowlisted OIDC claim map for CLI SSO polling, e.g.
# "employment_type->acme_employment_type,org_info.department->department"
@@ -1824,7 +1796,7 @@ DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME: Final = "db_daily_tag_spend_update_job"
PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME: Final = "prometheus_emit_budget_metrics"
CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data"
MAVVRIK_FOCUS_EXPORT_JOB_NAME: Final = "mavvrik_focus_export_usage_data"
-CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000))
+CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = get_env_int("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)
SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup"
BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME: Final = "background_health_check_db_save"
KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
@@ -1835,20 +1807,20 @@ USER_SPEND_ALERTS_JOB_ID: Final = "user_spend_alerts_job"
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning"
-SPEND_LOG_RUN_LOOPS: Final = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
-SPEND_LOG_CLEANUP_BATCH_SIZE: Final = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
-SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3))
-SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float(
- os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
+SPEND_LOG_RUN_LOOPS: Final = get_env_int("SPEND_LOG_RUN_LOOPS", 500)
+SPEND_LOG_CLEANUP_BATCH_SIZE: Final = get_env_int("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)
+SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = get_env_int("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)
+SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = get_env_float(
+ "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5
)
-SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300"))
-SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30"))
-SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000"))
-SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float(
- os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5")
+SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = get_env_float("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", 300.0)
+SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = get_env_float("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", 30.0)
+SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = get_env_int("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", 100000)
+SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = get_env_float(
+ "SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", 5.0
)
-SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float(
- os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5")
+SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = get_env_float(
+ "SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", 5.0
)
TOOL_SPEND_TOP_TOOLS: Final = 100
MODEL_INSIGHTS_TOP_MODELS: Final = 10
@@ -1856,54 +1828,54 @@ MODEL_INSIGHTS_MAX_RANGE_DAYS: Final = 365
MODEL_INSIGHTS_DEFAULT_TASK: Final = "uncategorized"
MODEL_INSIGHTS_TASK_TAG_PREFIX: Final = "task:"
SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
-SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
-SPEND_LOG_WRITE_BATCH_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BATCH_MAX_BYTES", 2_000_000)))
-SPEND_LOG_WRITE_BATCH_MAX_ROWS: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BATCH_MAX_ROWS", "100")))
-SPEND_ROLLUP_LOCK_TIMEOUT_MS: Final = max(1, int(os.getenv("SPEND_ROLLUP_LOCK_TIMEOUT_MS", "5000")))
-SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
-SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_QUEUE_MAX_BYTES", "64000000")))
-SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
-RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3")))
-RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2"))
-SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))
-PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, int(os.getenv("PROXY_DB_LOOKUP_MAX_CONCURRENCY", "25")))
-PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, float(os.getenv("PROXY_DB_LOOKUP_DEADLINE_SECONDS", "10")))
-PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, float(os.getenv("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", "30")))
-DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute
-PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))
-RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500")))
-RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", "100")))
+SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = get_env_int("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)
+SPEND_LOG_WRITE_BATCH_MAX_BYTES: Final = max(1, get_env_int("SPEND_LOG_WRITE_BATCH_MAX_BYTES", 2_000_000))
+SPEND_LOG_WRITE_BATCH_MAX_ROWS: Final = max(1, get_env_int("SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100))
+SPEND_ROLLUP_LOCK_TIMEOUT_MS: Final = max(1, get_env_int("SPEND_ROLLUP_LOCK_TIMEOUT_MS", 5000))
+SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = get_env_int("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)
+SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, get_env_int("SPEND_LOG_QUEUE_MAX_BYTES", 64000000))
+SPEND_LOG_QUEUE_POLL_INTERVAL: Final = get_env_float("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)
+RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", 3))
+RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = get_env_float("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.2)
+SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)
+PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25))
+PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0))
+PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0))
+DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = get_env_int("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60) # 1 minute
+PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
+RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, get_env_int("RESET_BUDGET_JOB_BATCH_SIZE", 500))
+RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, get_env_int("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", 100))
RESET_BUDGET_JOB_NAME: Final = "reset_budget_job"
# Comfortably longer than one PROXY_BUDGET_RESCHEDULER_MIN_TIME tick, so a healthy
# leader keeps the lease across its own run, and a crashed one strands the sweep for
# at most a single tick.
RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
-PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
-MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
-MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)))
+PROXY_BATCH_POLLING_INTERVAL: Final = get_env_int("PROXY_BATCH_POLLING_INTERVAL", 3600)
+MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, get_env_int("MAX_OBJECTS_PER_POLL_CYCLE", 50))
+MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60
-STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000)))
+STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, get_env_int("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))
# Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and
# CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on
# installations with large numbers of stale managed objects).
_batch_polling_env: Final = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower()
PROXY_BATCH_POLLING_ENABLED: Final = _batch_polling_env == "true"
-BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS: Final = float(
- os.getenv("BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS", "5")
+BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS: Final = get_env_float(
+ "BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS", 5.0
)
-BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS: Final = float(
- os.getenv("BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS", "60")
+BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS: Final = get_env_float(
+ "BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS", 60.0
)
-BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS: Final = float(
- os.getenv("BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS", "3600")
+BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS: Final = get_env_float(
+ "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS", 3600.0
)
_background_interaction_cost_polling_env: Final = os.getenv(
"BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true"
).lower()
BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true"
-PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605))
-PROXY_BATCH_WRITE_AT: Final = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10
+PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
+PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10) # in seconds, increased from 10
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)
# APScheduler Configuration - MEMORY LEAK FIX
@@ -1912,10 +1884,10 @@ APSCHEDULER_COALESCE: Final = os.getenv("APSCHEDULER_COALESCE", "True").lower()
"true",
"1",
] # collapse many missed runs into one
-APSCHEDULER_MISFIRE_GRACE_TIME: Final = int(
- os.getenv("APSCHEDULER_MISFIRE_GRACE_TIME", 3600)
+APSCHEDULER_MISFIRE_GRACE_TIME: Final = get_env_int(
+ "APSCHEDULER_MISFIRE_GRACE_TIME", 3600
) # ignore runs older than 1 hour (was 120)
-APSCHEDULER_MAX_INSTANCES: Final = int(os.getenv("APSCHEDULER_MAX_INSTANCES", 1)) # prevent concurrent job instances
+APSCHEDULER_MAX_INSTANCES: Final = get_env_int("APSCHEDULER_MAX_INSTANCES", 1) # prevent concurrent job instances
APSCHEDULER_REPLACE_EXISTING: Final = os.getenv("APSCHEDULER_REPLACE_EXISTING", "True").lower() in [
"true",
"1",
@@ -1929,25 +1901,25 @@ DEFAULT_STAGGER_WINDOW_SECONDS: Final = 300
# This will run tag spcific tasks at a later time to smooth QPS
DAILY_TAG_SPEND_BATCH_MULTIPLIER: Final = 2.3
-DEFAULT_HEALTH_CHECK_INTERVAL: Final = int(os.getenv("DEFAULT_HEALTH_CHECK_INTERVAL", 300)) # 5 minutes
-DEFAULT_SHARED_HEALTH_CHECK_TTL: Final = int(
- os.getenv("DEFAULT_SHARED_HEALTH_CHECK_TTL", 300)
+DEFAULT_HEALTH_CHECK_INTERVAL: Final = get_env_int("DEFAULT_HEALTH_CHECK_INTERVAL", 300) # 5 minutes
+DEFAULT_SHARED_HEALTH_CHECK_TTL: Final = get_env_int(
+ "DEFAULT_SHARED_HEALTH_CHECK_TTL", 300
) # 5 minutes - TTL for cached health check results
-DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL: Final = int(
- os.getenv("DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL", 60)
+DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL: Final = get_env_int(
+ "DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL", 60
) # 1 minute - TTL for health check lock
DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER: Final = 2 # health state is stale after interval * this
-PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS: Final = int(os.getenv("PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS", 9))
-DEFAULT_MODEL_CREATED_AT_TIME: Final = int(
- os.getenv("DEFAULT_MODEL_CREATED_AT_TIME", 1677610602)
+PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS: Final = get_env_int("PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS", 9)
+DEFAULT_MODEL_CREATED_AT_TIME: Final = get_env_int(
+ "DEFAULT_MODEL_CREATED_AT_TIME", 1677610602
) # returns on `/models` endpoint
-DEFAULT_SLACK_ALERTING_THRESHOLD: Final = int(os.getenv("DEFAULT_SLACK_ALERTING_THRESHOLD", 300))
-MAX_TEAM_LIST_LIMIT: Final = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20))
-MAX_POLICY_ESTIMATE_IMPACT_ROWS: Final = int(os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000))
-DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7))
-LENGTH_OF_LITELLM_GENERATED_KEY: Final = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16))
-MINIMUM_CUSTOM_KEY_LENGTH: Final = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16))
-SECRET_MANAGER_REFRESH_INTERVAL: Final = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400))
+DEFAULT_SLACK_ALERTING_THRESHOLD: Final = get_env_int("DEFAULT_SLACK_ALERTING_THRESHOLD", 300)
+MAX_TEAM_LIST_LIMIT: Final = get_env_int("MAX_TEAM_LIST_LIMIT", 20)
+MAX_POLICY_ESTIMATE_IMPACT_ROWS: Final = get_env_int("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000)
+DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = get_env_float("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7)
+LENGTH_OF_LITELLM_GENERATED_KEY: Final = get_env_int("LENGTH_OF_LITELLM_GENERATED_KEY", 16)
+MINIMUM_CUSTOM_KEY_LENGTH: Final = get_env_int("MINIMUM_CUSTOM_KEY_LENGTH", 16)
+SECRET_MANAGER_REFRESH_INTERVAL: Final = get_env_int("SECRET_MANAGER_REFRESH_INTERVAL", 86400)
OPENAI_SYSTEM_MESSAGES_FIRST_PROVIDERS: Final = frozenset({"openai", "azure"})
LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"default_internal_user_params",
@@ -1980,9 +1952,9 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"aws_sqs_callback_params",
]
SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
-DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))
+DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = get_env_int("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)
AUTH_OBJECTS_TARGET: Final = "auth_objects"
-DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600))
+DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = get_env_int("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)
SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600
SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30
SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000
@@ -2090,14 +2062,14 @@ SENTRY_PII_DENYLIST: Final = [
]
# CoroutineChecker cache configuration
-COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY: Final = int(os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000))
+COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY: Final = get_env_int("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000)
########################### RAG Text Splitter Constants ###########################
-DEFAULT_CHUNK_SIZE: Final = int(os.getenv("DEFAULT_CHUNK_SIZE", 1000))
-DEFAULT_CHUNK_OVERLAP: Final = int(os.getenv("DEFAULT_CHUNK_OVERLAP", 200))
+DEFAULT_CHUNK_SIZE: Final = get_env_int("DEFAULT_CHUNK_SIZE", 1000)
+DEFAULT_CHUNK_OVERLAP: Final = get_env_int("DEFAULT_CHUNK_OVERLAP", 200)
########################### S3 Vectors RAG Constants ###########################
-S3_VECTORS_DEFAULT_DIMENSION: Final = int(os.getenv("S3_VECTORS_DEFAULT_DIMENSION", 1024))
+S3_VECTORS_DEFAULT_DIMENSION: Final = get_env_int("S3_VECTORS_DEFAULT_DIMENSION", 1024)
S3_VECTORS_DEFAULT_DISTANCE_METRIC: Final = str(os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine"))
S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS: Final = ["source_text"]
@@ -2110,11 +2082,11 @@ MICROSOFT_USER_LAST_NAME_ATTRIBUTE: Final = str(os.getenv("MICROSOFT_USER_LAST_N
# Maximum payload size (in bytes) to fully serialize for DEBUG logging.
# Payloads larger than this are truncated to avoid multi-second json.dumps blocking the response.
-MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG: Final = int(os.getenv("MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG", 102400)) # 100 KB
+MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG: Final = get_env_int("MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG", 102400) # 100 KB
# Policy template enrichment
-MAX_COMPETITOR_NAMES: Final = int(os.getenv("MAX_COMPETITOR_NAMES", 100))
-COMPETITOR_LLM_TEMPERATURE: Final = float(os.getenv("COMPETITOR_LLM_TEMPERATURE", 0.3))
+MAX_COMPETITOR_NAMES: Final = get_env_int("MAX_COMPETITOR_NAMES", 100)
+COMPETITOR_LLM_TEMPERATURE: Final = get_env_float("COMPETITOR_LLM_TEMPERATURE", 0.3)
DEFAULT_COMPETITOR_DISCOVERY_MODEL: Final = "gpt-4o-mini"
# Advisor tool orchestration
diff --git a/litellm/litellm_core_utils/env_utils.py b/litellm/litellm_core_utils/env_utils.py
index d641884b4cd..6883af98bbd 100644
--- a/litellm/litellm_core_utils/env_utils.py
+++ b/litellm/litellm_core_utils/env_utils.py
@@ -4,23 +4,65 @@ Utility helpers for reading and parsing environment variables.
import logging
import os
-from typing import Final
+from typing import Final, TypeVar
+
+_NumberT: Final = TypeVar("_NumberT", int, float)
+
+# Marked before logging: the LiteLLM loggers' own filters read env knobs through these helpers, so
+# warning first would re-enter the helper from inside the log call and recurse until RecursionError.
+_reported_unparsable_values: Final[set[tuple[str, str]]] = set() # mutable-ok: process-wide warn-once registry
+
+
+def _warn_unparsable_once(env_var: str, raw: str, number_type: type[int] | type[float]) -> None:
+ key: Final = (env_var, raw)
+ if key in _reported_unparsable_values:
+ return
+ _reported_unparsable_values.add(key)
+ logging.getLogger("LiteLLM").warning(
+ "%s=%r is not a valid %s. Treating it as unset.", env_var, raw, number_type.__name__
+ )
+
+
+def _get_env_number_or_none(env_var: str, number_type: type[_NumberT]) -> _NumberT | None:
+ raw: Final = os.getenv(env_var)
+ if raw is None or not raw.strip():
+ return None
+ try:
+ return number_type(raw.strip())
+ except ValueError:
+ _warn_unparsable_once(env_var, raw, number_type)
+ return None
+
+
+def get_env_int_or_none(env_var: str) -> int | None:
+ """Parse an environment variable as an integer, returning None when it is unset, blank or unparsable.
+
+ Use this instead of `get_env_int` when callers must distinguish "explicitly configured"
+ from "left at the default", for example when an override should take precedence over a
+ value resolved from somewhere else.
+ """
+ return _get_env_number_or_none(env_var, int)
+
+
+def get_env_float_or_none(env_var: str) -> float | None:
+ """Float counterpart of `get_env_int_or_none`."""
+ return _get_env_number_or_none(env_var, float)
def get_env_int(env_var: str, default: int) -> int:
- """Parse an environment variable as an integer, falling back to default on invalid values.
+ """Parse an environment variable as an integer, falling back to default when it is unset, blank or unparsable.
- Handles empty strings, whitespace, and non-numeric values gracefully
- so that misconfiguration doesn't crash the process at import time.
+ An unparsable value is warned about once, so misconfiguration neither crashes the process at
+ import time nor goes unnoticed.
"""
- raw = os.getenv(env_var)
- if raw is None:
- return default
- raw = raw.strip()
- try:
- return int(raw)
- except (ValueError, TypeError):
- return default
+ value: Final = get_env_int_or_none(env_var)
+ return default if value is None else value
+
+
+def get_env_float(env_var: str, default: float) -> float:
+ """Float counterpart of `get_env_int`."""
+ value: Final = get_env_float_or_none(env_var)
+ return default if value is None else value
def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int) -> int:
@@ -41,19 +83,3 @@ def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int)
default,
)
return default
-
-
-def get_env_int_or_none(env_var: str) -> int | None:
- """Parse an environment variable as an integer, returning None when it is unset or unusable.
-
- Use this instead of `get_env_int` when callers must distinguish "explicitly configured"
- from "left at the default", for example when an override should take precedence over a
- value resolved from somewhere else.
- """
- raw: Final = os.getenv(env_var)
- if raw is None:
- return None
- try:
- return int(raw.strip())
- except (ValueError, TypeError):
- return None
diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py
index 7545a1117a1..89794fc2480 100644
--- a/litellm/proxy/agent_endpoints/endpoints.py
+++ b/litellm/proxy/agent_endpoints/endpoints.py
@@ -21,6 +21,7 @@ from typing_extensions import ReadOnly, Required, assert_never
import litellm
from litellm._logging import verbose_proxy_logger
+from litellm.constants import AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, AGENT_HEALTH_CHECK_TIMEOUT_SECONDS
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._types import (
CommonProxyErrors,
@@ -228,10 +229,6 @@ def _check_agent_management_permission(user_api_key_dict: UserAPIKeyAuth) -> Non
)
-AGENT_HEALTH_CHECK_TIMEOUT_SECONDS: Final = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", "5.0"))
-AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", "30.0"))
-
-
class _AgentHealthResult(TypedDict, total=False):
agent_id: Required[str]
healthy: Required[bool]
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index cf90d90ea33..20504982fde 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -55,6 +55,7 @@ from litellm import _custom_logger_compatible_callbacks_literal
from litellm.constants import (
DEFAULT_MODEL_CREATED_AT_TIME,
FILE_USAGE_MAX_TRACKED_COUNTERS,
+ LITELLM_CONFIG_CACHE_TTL_SECONDS,
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
MAX_TEAM_LIST_LIMIT,
PROXY_REJECTED_BEFORE_ROUTING_KEY,
@@ -150,6 +151,7 @@ from litellm.litellm_core_utils.core_helpers import (
is_expected_client_error,
)
from litellm.litellm_core_utils.duration_parser import get_budget_window_start
+from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
@@ -4505,7 +4507,6 @@ async def _lookup_deprecated_key(
# DualCache for LiteLLM_Config param_name reads.
# Redis layer is attached in proxy_server._init_cache.
-LITELLM_CONFIG_CACHE_TTL_SECONDS: Final[int] = int(os.environ.get("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", "60"))
_CONFIG_CACHE_MISS: Final[str] = "__litellm_config_param_miss__"
litellm_config_cache: Final[DualCache] = DualCache(
@@ -4789,28 +4790,28 @@ class PrismaClient:
self._db_health_watchdog_task: asyncio.Task | None = None
self._view_setup_task: asyncio.Task[_ViewSetupOutcome] | None = None
self._db_last_reconnect_attempt_ts: float = 0.0
- self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15")))
+ self._db_reconnect_cooldown_seconds: int = max(1, get_env_int("PRISMA_RECONNECT_COOLDOWN_SECONDS", 15))
self._db_read_only_recreate_ts: float = 0.0
self._db_read_only_recreate_streak: int = 0
self._db_health_watchdog_interval_seconds: int = max(
- 5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "30"))
+ 5, get_env_int("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", 30)
)
self._db_health_watchdog_enabled: bool = (
str_to_bool(os.getenv("PRISMA_HEALTH_WATCHDOG_ENABLED", "true")) is True
)
self._db_health_watchdog_probe_timeout_seconds: float = max(
0.5,
- float(os.getenv("PRISMA_HEALTH_WATCHDOG_PROBE_TIMEOUT_SECONDS", "5.0")),
+ get_env_float("PRISMA_HEALTH_WATCHDOG_PROBE_TIMEOUT_SECONDS", 5.0),
)
self._db_watchdog_reconnect_timeout_seconds: float = max(
- 1.0, float(os.getenv("PRISMA_WATCHDOG_RECONNECT_TIMEOUT_SECONDS", "30.0"))
+ 1.0, get_env_float("PRISMA_WATCHDOG_RECONNECT_TIMEOUT_SECONDS", 30.0)
)
self._db_auth_reconnect_timeout_seconds: float = max(
- 0.5, float(os.getenv("PRISMA_AUTH_RECONNECT_TIMEOUT_SECONDS", "2.0"))
+ 0.5, get_env_float("PRISMA_AUTH_RECONNECT_TIMEOUT_SECONDS", 2.0)
)
self._db_auth_reconnect_lock_timeout_seconds: float = max(
0.0,
- float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")),
+ get_env_float("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", 0.1),
)
self._consecutive_reconnect_failures: int = 0
# Last generation of each read engine whose repair was attempted and
@@ -4821,7 +4822,7 @@ class PrismaClient:
# back to a caller whose engine is still unrepaired. Bounded at two
# entries: a client has one writer and at most one reader.
self._failed_recreate_generations: Mapping[PrismaWrapper, int] = MappingProxyType({})
- self._reconnect_escalation_threshold: int = max(1, int(os.getenv("PRISMA_RECONNECT_ESCALATION_THRESHOLD", "3")))
+ self._reconnect_escalation_threshold: int = max(1, get_env_int("PRISMA_RECONNECT_ESCALATION_THRESHOLD", 3))
self._engine_pidfd: int = -1
self._engine_pid: int = 0
self._watching_engine: bool = False
diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py
index e65a0dabd6d..6c86b058f0f 100644
--- a/litellm/types/guardrails.py
+++ b/litellm/types/guardrails.py
@@ -5,6 +5,7 @@ from types import MappingProxyType
from typing import Final, Literal, cast
from pydantic import ConfigDict, Field, field_validator, model_validator
+from pydantic.fields import FieldInfo
from typing_extensions import ReadOnly, Required, TypedDict
from litellm._logging import verbose_logger
@@ -1000,6 +1001,16 @@ def runtime_stream_scope(
LoggingOnlyScope = Literal["input", "output"]
+def _is_blank(value: object) -> bool:
+ return isinstance(value, str) and not value.strip()
+
+
+def _default_if_blank(field: FieldInfo | None, value: object) -> object:
+ if field is None or field.is_required() or not _is_blank(value):
+ return value
+ return field.get_default(call_default_factory=True)
+
+
class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch update guardrails
api_key: str | None = Field(default=None, description="API key for the guardrail service")
api_base: str | None = Field(default=None, description="Base URL for the guardrail service API")
@@ -1288,6 +1299,14 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
),
)
+ @model_validator(mode="before") # pyright: ignore[reportUntypedFunctionDecorator] # pydantic decorators are untyped
+ @classmethod
+ def blank_values_mean_field_default(cls, data: object) -> object:
+ if not isinstance(data, Mapping):
+ return data
+ fields: Final = cast(Mapping[str, object], data) # cast-ok: pydantic hands a before-validator the raw mapping
+ return {key: _default_if_blank(cls.model_fields.get(key), value) for key, value in fields.items()}
+
@field_validator(
"mode",
"default_action",
@@ -1398,14 +1417,14 @@ class LitellmParams( # pyright: ignore[reportIncompatibleVariableOverride] # o
)
return self
- def __init__(self, **kwargs) -> None:
+ def __init__(self, **kwargs: object) -> None:
default_on: Final = kwargs.pop("default_on", None)
- if default_on is not None:
+ if default_on is not None and not _is_blank(default_on):
kwargs["default_on"] = default_on
else:
kwargs["default_on"] = False
- super().__init__(**kwargs)
+ super().__init__(**kwargs) # pyright: ignore[reportArgumentType] # pydantic validates each field's type at runtime
def __contains__(self, key) -> bool:
# Define custom behavior for the 'in' operator
diff --git a/tests/unit/litellm_core_utils/test_env_utils.py b/tests/unit/litellm_core_utils/test_env_utils.py
new file mode 100644
index 00000000000..1ddd172e274
--- /dev/null
+++ b/tests/unit/litellm_core_utils/test_env_utils.py
@@ -0,0 +1,96 @@
+import logging
+from collections.abc import Callable, Iterator
+
+import pytest
+
+from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.env_utils import (
+ get_env_float,
+ get_env_float_or_none,
+ get_env_int,
+ get_env_int_or_none,
+)
+
+
+@pytest.fixture
+def litellm_log(caplog: pytest.LogCaptureFixture) -> Iterator[pytest.LogCaptureFixture]:
+ verbose_logger.addHandler(caplog.handler)
+ try:
+ with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
+ yield caplog
+ finally:
+ verbose_logger.removeHandler(caplog.handler)
+
+
+def _logged_messages(caplog: pytest.LogCaptureFixture) -> list[str]:
+ return [record.getMessage() for record in dict.fromkeys(caplog.records)]
+
+
+@pytest.mark.parametrize("raw", ["", " ", "\t"])
+def test_blank_numeric_env_value_means_unset_without_warning(
+ monkeypatch: pytest.MonkeyPatch, litellm_log: pytest.LogCaptureFixture, raw: str
+) -> None:
+ monkeypatch.setenv("ENV_UTILS_TEST_KNOB", raw)
+
+ assert get_env_float("ENV_UTILS_TEST_KNOB", 1.5) == 1.5
+ assert get_env_int("ENV_UTILS_TEST_KNOB", 7) == 7
+ assert get_env_float_or_none("ENV_UTILS_TEST_KNOB") is None
+ assert get_env_int_or_none("ENV_UTILS_TEST_KNOB") is None
+ assert _logged_messages(litellm_log) == []
+
+
+def test_set_numeric_env_value_is_parsed(monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setenv("ENV_UTILS_TEST_FLOAT", " 2.5 ")
+ monkeypatch.setenv("ENV_UTILS_TEST_INT", " 12 ")
+
+ assert get_env_float("ENV_UTILS_TEST_FLOAT", 1.5) == 2.5
+ assert get_env_int("ENV_UTILS_TEST_INT", 7) == 12
+ assert get_env_float_or_none("ENV_UTILS_TEST_FLOAT") == 2.5
+ assert get_env_int_or_none("ENV_UTILS_TEST_INT") == 12
+
+
+@pytest.mark.parametrize(
+ ("read", "raw", "expected"),
+ [
+ (lambda name: get_env_int(name, 7), "fast-int", 7),
+ (lambda name: get_env_float(name, 1.5), "fast-float", 1.5),
+ (get_env_int_or_none, "fast-int-or-none", None),
+ (get_env_float_or_none, "fast-float-or-none", None),
+ ],
+)
+def test_unparsable_numeric_env_value_falls_back_and_warns_once(
+ monkeypatch: pytest.MonkeyPatch,
+ litellm_log: pytest.LogCaptureFixture,
+ read: Callable[[str], float | None],
+ raw: str,
+ expected: float | None,
+) -> None:
+ monkeypatch.setenv("ENV_UTILS_TEST_KNOB", raw)
+
+ assert read("ENV_UTILS_TEST_KNOB") == expected
+ assert read("ENV_UTILS_TEST_KNOB") == expected
+
+ warnings = [message for message in _logged_messages(litellm_log) if raw in message]
+ assert len(warnings) == 1
+ assert "ENV_UTILS_TEST_KNOB" in warnings[0]
+
+
+@pytest.mark.parametrize("env_var", ["MAX_STRING_LENGTH_STDOUT_LOG", "MAX_BASE64_LENGTH_STDOUT_LOG"])
+def test_unparsable_stdout_log_limit_does_not_break_logging(
+ monkeypatch: pytest.MonkeyPatch, litellm_log: pytest.LogCaptureFixture, env_var: str
+) -> None:
+ """
+ Regression: the LiteLLM loggers' own filter reads these limits through get_env_int on every
+ record, so warning about a garbage value re-entered the filter and every log call raised
+ RecursionError.
+ """
+ raw = f"not-a-number-{env_var}"
+ monkeypatch.setenv(env_var, raw)
+
+ verbose_logger.warning("first line")
+ verbose_logger.warning("second line")
+
+ messages = _logged_messages(litellm_log)
+ assert "first line" in messages
+ assert "second line" in messages
+ assert len([message for message in messages if raw in message]) == 1
diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_straiker.py
index 292dd54105f..85886a2b9af 100644
--- a/tests/unit/proxy/guardrails/guardrail_hooks/test_straiker.py
+++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_straiker.py
@@ -1,5 +1,6 @@
import json
from types import MappingProxyType, SimpleNamespace
+from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@@ -1216,7 +1217,25 @@ def test_v3_initializer_reads_api_version_from_config():
assert g._webhook_url().endswith("/api/v3/detect")
-@pytest.mark.parametrize("api_version", ["2024-09-01", "", "v2"])
+@pytest.mark.parametrize(("api_key", "expected"), [("c4ac433a-uuid", "v1"), (V3_KEY, "v3")])
+def test_blank_api_version_follows_key_prefix_without_warning(
+ api_key: str, expected: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ import litellm
+ from litellm._logging import verbose_proxy_logger
+ from litellm.types.guardrails import Guardrail, LitellmParams
+
+ monkeypatch.setattr(litellm, "callbacks", litellm.callbacks.copy())
+ params: Final = LitellmParams(guardrail="straiker", mode="pre_call", api_key=api_key, api_version="")
+
+ with patch.object(verbose_proxy_logger, "warning") as warning:
+ g = initialize_guardrail(params, Guardrail(guardrail_name="straiker", litellm_params=params))
+
+ assert g.api_version == expected
+ warning.assert_not_called()
+
+
+@pytest.mark.parametrize("api_version", ["2024-09-01", "v2"])
@pytest.mark.parametrize(("api_key", "expected"), [("c4ac433a-uuid", "v1"), (V3_KEY, "v3")])
def test_unknown_api_version_follows_key_prefix(api_version, api_key, expected, monkeypatch):
import litellm
diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py
index 13049093d6c..4401e3e95c0 100644
--- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py
+++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py
@@ -1670,6 +1670,31 @@ async def test_patch_guardrail_clears_scope_when_logging_only_mode_is_removed(
handler.delete_in_memory_guardrail(stored_guardrail["guardrail_id"])
+@pytest.mark.asyncio
+async def test_patch_guardrail_with_blank_value_clears_the_stored_value(mocker, monkeypatch, mock_guardrail_registry):
+ """
+ Regression: blank litellm_params were dropped before validation, so a PATCH that cleared
+ api_base with "" returned 200 but kept the stored value.
+ """
+ handler, stored_guardrail = _setup_patch_scope_guardrail(
+ mocker,
+ monkeypatch,
+ mock_guardrail_registry,
+ _PatchScopeSupportedGuardrail,
+ "patch_blank_value_test",
+ {"mode": "pre_call", "api_base": "https://old.example.com"},
+ )
+ request = PatchGuardrailRequest.model_validate({"litellm_params": {"api_base": ""}})
+
+ try:
+ await patch_guardrail(stored_guardrail["guardrail_id"], request, user_api_key_dict=MOCK_ADMIN_USER)
+
+ persisted_guardrail = mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"]
+ assert persisted_guardrail["litellm_params"].api_base is None
+ finally:
+ handler.delete_in_memory_guardrail(stored_guardrail["guardrail_id"])
+
+
@pytest.mark.asyncio
async def test_patch_guardrail_tolerates_stored_unsupported_scope_on_unrelated_update(
mocker, monkeypatch, mock_guardrail_registry, caplog
diff --git a/tests/unit/proxy/guardrails/test_init_guardrails.py b/tests/unit/proxy/guardrails/test_init_guardrails.py
index c8eeaf11fde..daf7ef1fd05 100644
--- a/tests/unit/proxy/guardrails/test_init_guardrails.py
+++ b/tests/unit/proxy/guardrails/test_init_guardrails.py
@@ -371,6 +371,44 @@ def test_init_guardrails_v2_skips_invalid_guardrail_instead_of_crashing_boot():
assert "healthy_presidio" in guardrail_names
+def test_init_guardrails_v2_keeps_guardrail_whose_env_values_are_blank():
+ """
+ Regression: a config.yaml guardrail with `default_on: os.environ/GUARDRAIL_DEFAULT_ON` and
+ that variable exported blank (a Helm `value: ""`) resolved to "", pydantic rejected it on the
+ bool field, and the proxy logged the guardrail and booted without it. Blank now means unset,
+ so the guardrail boots with its defaults.
+ """
+ from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
+
+ IN_MEMORY_GUARDRAIL_HANDLER.IN_MEMORY_GUARDRAILS.clear()
+ IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.clear()
+
+ all_guardrails = [
+ {
+ "guardrail_name": "presidio_with_blank_env_values",
+ "litellm_params": {
+ "guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
+ "mode": "pre_call",
+ "default_on": "",
+ "mask_request_content": "",
+ "presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
+ "presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
+ },
+ },
+ ]
+
+ init_guardrails_v2(all_guardrails=all_guardrails)
+
+ initialized = {
+ guardrail["guardrail_name"]: guardrail
+ for guardrail in IN_MEMORY_GUARDRAIL_HANDLER.IN_MEMORY_GUARDRAILS.values()
+ }
+ assert "presidio_with_blank_env_values" in initialized
+ litellm_params = initialized["presidio_with_blank_env_values"]["litellm_params"]
+ assert litellm_params.default_on is False
+ assert litellm_params.mask_request_content is None
+
+
def test_init_guardrails_v2_stops_boot_when_a_custom_code_guardrail_does_not_compile():
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
diff --git a/tests/unit/test_constants.py b/tests/unit/test_constants.py
index 18ed6975605..0d0e77765b7 100644
--- a/tests/unit/test_constants.py
+++ b/tests/unit/test_constants.py
@@ -17,6 +17,19 @@ from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outco
from unittest.mock import patch
+_ENV_HELPERS = frozenset(
+ {"get_env_int", "get_env_float", "get_env_int_in_range", "get_env_int_or_none", "get_env_float_or_none"}
+)
+
+
+def _single_name_target(node: ast.stmt) -> str | None:
+ if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
+ return node.target.id
+ if isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
+ return node.targets[0].id
+ return None
+
+
def _build_constant_env_var_map() -> dict[str, str]:
"""
Build a mapping of CONSTANT_NAME -> ENV_VAR_NAME by parsing constants.py.
@@ -29,13 +42,10 @@ def _build_constant_env_var_map() -> dict[str, str]:
parsed = ast.parse(constants_source)
for node in parsed.body:
- if not isinstance(node, ast.Assign):
+ constant_name = _single_name_target(node)
+ if constant_name is None or not isinstance(node, (ast.Assign, ast.AnnAssign)) or node.value is None:
continue
- if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name):
- continue
-
- constant_name = node.targets[0].id
env_var_name = None
for child in ast.walk(node.value):
@@ -55,10 +65,10 @@ def _build_constant_env_var_map() -> dict[str, str]:
env_var_name = child.args[0].value
break
- # get_env_int("ENV_NAME", default)
+ # get_env_int("ENV_NAME", default) and its siblings
if (
isinstance(child.func, ast.Name)
- and child.func.id == "get_env_int"
+ and child.func.id in _ENV_HELPERS
and len(child.args) >= 1
and isinstance(child.args[0], ast.Constant)
and isinstance(child.args[0].value, str)
@@ -79,8 +89,10 @@ def _build_constant_env_var_map() -> dict[str, str]:
(None, "48", 48),
(None, None, 24),
("48", "72", 48),
+ (" ", None, 24),
+ ("", "48", 48),
],
- ids=("canonical-only", "alias-only", "default", "canonical-wins"),
+ ids=("canonical-only", "alias-only", "default", "canonical-wins", "blank-canonical", "blank-canonical-alias-set"),
)
def test_cli_jwt_expiration_hours_from_environment(
monkeypatch: pytest.MonkeyPatch,
@@ -105,12 +117,79 @@ def test_cli_jwt_expiration_hours_from_environment(
importlib.reload(litellm.constants)
+def _is_blank(value: object) -> bool:
+ return isinstance(value, str) and not value.strip()
+
+
+@pytest.mark.parametrize("blank", ["", " "])
+def test_constants_import_treats_blank_env_values_as_unset(blank: str) -> None:
+ """
+ Regression: a numeric knob exported as a blank string (a Helm `value: ""` or an emptied .env
+ line) raised ValueError while importing litellm.constants, so the proxy crash-looped at boot.
+ Blank now means unset, so import succeeds with every env-backed constant blank, and every
+ numeric or optional constant keeps the value it has with the variable absent.
+ """
+ constant_env_vars = _build_constant_env_var_map()
+ assert {"ROUTER_MAX_FALLBACKS", "LITELLM_MAX_STREAMING_DURATION_SECONDS"} <= constant_env_vars.keys()
+ env_vars = set(constant_env_vars.values())
+ try:
+ with mock.patch.dict(os.environ):
+ for env_var in env_vars:
+ os.environ.pop(env_var, None)
+ importlib.reload(litellm.constants)
+ unset = {name: getattr(litellm.constants, name) for name in constant_env_vars}
+ with mock.patch.dict(os.environ, {env_var: blank for env_var in env_vars}):
+ importlib.reload(litellm.constants)
+ with_blank_env = {name: getattr(litellm.constants, name) for name in constant_env_vars}
+ finally:
+ importlib.reload(litellm.constants)
+
+ numeric_or_optional = {name for name, value in unset.items() if value is None or type(value) in (int, float)}
+ changed = {
+ name: (unset[name], with_blank_env[name])
+ for name in numeric_or_optional
+ if with_blank_env[name] != unset[name] and not _is_blank(with_blank_env[name])
+ }
+ assert changed == {}
+
+
+@pytest.mark.parametrize(
+ ("raw", "expected_timeout", "expected_explicit"),
+ [
+ (None, constants.DEFAULT_REQUEST_TIMEOUT_SECONDS, False),
+ ("", constants.DEFAULT_REQUEST_TIMEOUT_SECONDS, False),
+ (" ", constants.DEFAULT_REQUEST_TIMEOUT_SECONDS, False),
+ ("fast", constants.DEFAULT_REQUEST_TIMEOUT_SECONDS, False),
+ (" 30 ", 30.0, True),
+ ],
+)
+def test_request_timeout_counts_as_explicit_only_when_it_parses(
+ monkeypatch: pytest.MonkeyPatch, raw: str | None, expected_timeout: float, expected_explicit: bool
+) -> None:
+ """
+ Regression: a blank REQUEST_TIMEOUT fell back to the 6000s sentinel but was still flagged as
+ explicitly set, so chat completion() kept the 6000s sentinel instead of its 600s fallback.
+ """
+ if raw is None:
+ monkeypatch.delenv("REQUEST_TIMEOUT", raising=False)
+ else:
+ monkeypatch.setenv("REQUEST_TIMEOUT", raw)
+ try:
+ importlib.reload(litellm.constants)
+ assert litellm.constants.request_timeout == expected_timeout
+ assert litellm.constants.request_timeout_explicitly_set is expected_explicit
+ finally:
+ monkeypatch.undo()
+ importlib.reload(litellm.constants)
+
+
@pytest.fixture()
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
yield
record_vcr_outcome(request, vcr)
+
@pytest.fixture(scope="session")
def event_loop():
try:
@@ -120,6 +199,7 @@ def event_loop():
yield loop
loop.close()
+
@pytest.fixture(scope="function")
def setup_and_teardown(event_loop):
import litellm
@@ -153,6 +233,7 @@ def setup_and_teardown(event_loop):
if pending:
event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+
_SCALAR_DEFAULTS = {
"num_retries": getattr(litellm, "num_retries", None),
"set_verbose": getattr(litellm, "set_verbose", False),
@@ -176,17 +257,12 @@ def _populate_known_models_for_morph_tests(monkeypatch: pytest.MonkeyPatch) -> N
monkeypatch.setattr(
litellm,
"models_by_provider",
- {
- provider: models.copy()
- for provider, models in litellm.models_by_provider.items()
- },
+ {provider: models.copy() for provider, models in litellm.models_by_provider.items()},
)
litellm.add_known_models()
-@pytest.mark.usefixtures(
- "_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests"
-)
+@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests")
def test_morph_config_get_provider_info():
"""Test that MorphChatConfig returns correct provider info."""
config = MorphChatConfig()
@@ -207,9 +283,8 @@ def test_morph_config_get_provider_info():
assert api_base == "https://custom.morph.com"
assert api_key == "key"
-@pytest.mark.usefixtures(
- "_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests"
-)
+
+@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests")
def test_morph_get_llm_provider():
"""Test that get_llm_provider correctly identifies morph models."""
# Test with morph/model format
@@ -219,9 +294,8 @@ def test_morph_get_llm_provider():
_, custom_llm_provider, _, _ = get_llm_provider("morph/morph-v3-fast")
assert custom_llm_provider == "morph"
-@pytest.mark.usefixtures(
- "_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests"
-)
+
+@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests")
def test_morph_in_provider_lists():
"""Test that morph is included in all necessary provider lists."""
import litellm
@@ -242,9 +316,8 @@ def test_morph_in_provider_lists():
# Check models are in model_list after initialization
assert all(model in litellm.model_list for model in ["morph/morph-v3-large", "morph/morph-v3-fast"])
-@pytest.mark.usefixtures(
- "_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests"
-)
+
+@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests")
def test_morph_supported_params():
"""Test that MorphChatConfig returns correct supported parameters."""
config = MorphChatConfig()
@@ -258,9 +331,8 @@ def test_morph_supported_params():
assert all(param in supported_params for param in expected_params)
-@pytest.mark.usefixtures(
- "_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests"
-)
+
+@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown", "_populate_known_models_for_morph_tests")
def test_morph_custom_llm_provider():
"""Test that morph models are correctly identified."""
config = MorphChatConfig()
diff --git a/tests/unit/types/test_guardrails.py b/tests/unit/types/test_guardrails.py
index 8b74f7ce0fd..5c652a10e16 100644
--- a/tests/unit/types/test_guardrails.py
+++ b/tests/unit/types/test_guardrails.py
@@ -13,6 +13,7 @@ import litellm.litellm_core_utils
import litellm.litellm_core_utils.litellm_logging
from litellm import completion
from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.guardrails import BaseLitellmParams, LitellmParams
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -157,3 +158,23 @@ def setup_and_teardown():
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
yield
+
+
+@pytest.mark.parametrize("blank", ["", " "])
+def test_litellm_params_treats_blank_values_as_unset(blank: str) -> None:
+ """
+ Regression: `os.environ/X` with X exported blank resolves to "", which pydantic rejected on
+ bool fields (default_on, mask_request_content, fail_on_error), so the guardrail was dropped at
+ boot. Blank now means unset, so the field keeps the same default as when it is omitted.
+ """
+ with_blank_values = LitellmParams(
+ guardrail="presidio",
+ mode="pre_call",
+ default_on=blank,
+ mask_request_content=blank,
+ api_key=blank,
+ )
+ omitted = LitellmParams(guardrail="presidio", mode="pre_call")
+ assert with_blank_values.model_dump() == omitted.model_dump()
+
+ assert BaseLitellmParams(fail_on_error=blank).model_dump() == BaseLitellmParams().model_dump()
From 47c40586e196c035b1dc3da87f62f3fb8459ad07 Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Wed, 14 Oct 2026 13:25:20 -0700
Subject: [PATCH 32/47] fix(proxy): evict the cached user on every proxy for
every /user/update and /user/bulk_update (#46500)
* fix(proxy): evict the cached user on every proxy for every /user/update and /user/bulk_update
Previously the cached user row was only evicted and broadcast when the
update touched max_budget, model_max_budget, tpm_limit, rpm_limit or
metadata, so a models or user_role change kept being served by peer
proxies until the cache TTL expired
* test(proxy): cover user cache eviction past the first bulk invalidation batch
---
.../internal_user_endpoints.py | 25 ++++----
..._user_updates_wedged_coordination_redis.py | 1 -
.../test_internal_user_endpoints.py | 59 +++++++++++++------
3 files changed, 52 insertions(+), 33 deletions(-)
diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py
index 2374f7d7aae..4e1001a4430 100644
--- a/litellm/proxy/management_endpoints/internal_user_endpoints.py
+++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py
@@ -121,8 +121,7 @@ if TYPE_CHECKING:
router: Final = APIRouter()
_USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig])
-_USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
-_USER_LIMIT_CACHE_FIELDS: Final = frozenset({"max_budget", "model_max_budget", "tpm_limit", "rpm_limit"})
+_USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50
def _user_table(
@@ -1634,11 +1633,10 @@ async def _update_single_user_helper(
await _invalidate_user_spend_counter_if_changed(non_default_values)
- if not _USER_LIMIT_CACHE_FIELDS.isdisjoint(non_default_values) or "metadata" in data_json:
- await evict_and_broadcast(
- cache_keys=(non_default_values["user_id"],),
- user_api_key_cache=user_api_key_cache,
- )
+ await evict_and_broadcast(
+ cache_keys=(non_default_values["user_id"],),
+ user_api_key_cache=user_api_key_cache,
+ )
if "object_permission_id" in non_default_values:
await _invalidate_cached_user_entitlement(
@@ -1992,14 +1990,13 @@ async def bulk_user_update(
),
)
- if not _USER_LIMIT_CACHE_FIELDS.isdisjoint(non_default_values):
- for start in range(0, len(all_users_in_db), _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE):
- await asyncio.gather(
- *(
- evict_and_broadcast(cache_keys=(user.user_id,), user_api_key_cache=user_api_key_cache)
- for user in all_users_in_db[start : start + _USER_BUDGET_CACHE_INVALIDATION_BATCH_SIZE]
- )
+ for start in range(0, len(all_users_in_db), _USER_CACHE_INVALIDATION_BATCH_SIZE):
+ await asyncio.gather(
+ *(
+ evict_and_broadcast(cache_keys=(user.user_id,), user_api_key_cache=user_api_key_cache)
+ for user in all_users_in_db[start : start + _USER_CACHE_INVALIDATION_BATCH_SIZE]
)
+ )
# Create individual success results
for user in all_users_in_db:
diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py
index 09b16e7d5ef..5ec5f85eef2 100644
--- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py
+++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py
@@ -634,7 +634,6 @@ def test_user_models_only_update_is_enforced_on_the_gateway(gateway: Gateway) ->
def test_user_models_update_without_metadata_propagates_to_gateway_and_peer(gateway: Gateway, peer: Gateway) -> None:
- pytest.skip("BUG: models-only /user/update leaves warmed peer serving the revoked model with 200")
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
index 39265b37b95..84c4608d29b 100644
--- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
+++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
@@ -2496,13 +2496,27 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker:
@pytest.mark.asyncio
@pytest.mark.parametrize(
- ("field", "new_limit"),
- [("tpm_limit", 100), ("rpm_limit", 1), ("tpm_limit", None), ("rpm_limit", None)],
- ids=["tpm_limit", "rpm_limit", "tpm_limit-cleared", "rpm_limit-cleared"],
+ ("field", "new_value"),
+ [
+ ("tpm_limit", 100),
+ ("rpm_limit", 1),
+ ("tpm_limit", None),
+ ("rpm_limit", None),
+ ("models", ["model-fir"]),
+ ("user_role", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value),
+ ],
+ ids=[
+ "tpm_limit",
+ "rpm_limit",
+ "tpm_limit-cleared",
+ "rpm_limit-cleared",
+ "models",
+ "user_role",
+ ],
)
@pytest.mark.parametrize("all_users", [False, True], ids=["single-user", "bulk-all-users"])
-async def test_user_rate_limit_update_reaches_cached_user_on_every_worker(
- mocker: MockerFixture, field: str, new_limit: int | None, all_users: bool
+async def test_user_update_reaches_cached_user_on_every_worker(
+ mocker: MockerFixture, field: str, new_value: object, all_users: bool
) -> None:
from litellm.proxy._types import LiteLLM_UserTable
from litellm.proxy.auth.auth_checks import get_user_object
@@ -2560,16 +2574,25 @@ async def test_user_rate_limit_update_reaches_cached_user_on_every_worker(
self.update_data_payload = data
return {"user_id": user_id, "data": self.updated_user}
- saved_user: Final = LiteLLM_UserTable(user_id="user-spruce", tpm_limit=100000, rpm_limit=1000)
- updated_user: Final = saved_user.model_copy(update={field: new_limit})
- old_limit: Final = 100000 if field == "tpm_limit" else 1000
+ saved_user: Final = LiteLLM_UserTable(
+ user_id="user-spruce",
+ tpm_limit=100000,
+ rpm_limit=1000,
+ models=["model-spruce"],
+ user_role=LitellmUserRoles.INTERNAL_USER.value,
+ )
+ updated_user: Final = saved_user.model_copy(update={field: new_value})
+ old_value: Final = saved_user.model_dump()[field]
+ assert old_value != new_value
prisma_client: Final = _PrismaClientMock(
db=_DatabaseMocks(
litellm_usertable=_UserTableMocks(
find_first=mocker.AsyncMock(return_value=saved_user),
find_unique=mocker.AsyncMock(return_value=updated_user),
- find_many=mocker.AsyncMock(return_value=[saved_user]),
+ find_many=mocker.AsyncMock(
+ return_value=[*(LiteLLM_UserTable(user_id=f"user-pine-{index}") for index in range(50)), saved_user]
+ ),
update_many=mocker.AsyncMock(return_value=1),
)
),
@@ -2614,28 +2637,28 @@ async def test_user_rate_limit_update_reaches_cached_user_on_every_worker(
user_id_upsert=False,
)
assert handling_user_before is not None
- assert handling_user_before.model_dump()[field] == old_limit
+ assert handling_user_before.model_dump()[field] == old_value
assert other_user_before is not None
- assert other_user_before.model_dump()[field] == old_limit
+ assert other_user_before.model_dump()[field] == old_value
admin: Final = UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN)
if all_users:
await bulk_user_update(
data=BulkUpdateUserRequest(
all_users=True,
- user_updates=UpdateUserRequestNoUserIDorEmail.model_validate({field: new_limit}),
+ user_updates=UpdateUserRequestNoUserIDorEmail.model_validate({field: new_value}),
),
user_api_key_dict=admin,
litellm_changed_by=None,
)
- prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={field: new_limit})
+ prisma_client.db.litellm_usertable.update_many.assert_awaited_once_with(where={}, data={field: new_value})
else:
await user_update(
- data=UpdateUserRequest.model_validate({"user_id": saved_user.user_id, field: new_limit}),
+ data=UpdateUserRequest.model_validate({"user_id": saved_user.user_id, field: new_value}),
user_api_key_dict=admin,
)
assert prisma_client.update_data_payload is not None
- assert prisma_client.update_data_payload[field] == new_limit
+ assert prisma_client.update_data_payload[field] == new_value
remote_subscriber: Final = AuthCacheInvalidationSubscriber(
redis_cache=cast("RedisCache", _FakeRedisCache()),
@@ -2659,10 +2682,10 @@ async def test_user_rate_limit_update_reaches_cached_user_on_every_worker(
user_id_upsert=False,
)
assert handling_user_after is not None
- assert handling_user_after.model_dump()[field] == new_limit
+ assert handling_user_after.model_dump()[field] == new_value
assert other_user_after is not None
- assert other_user_after.model_dump()[field] == new_limit, (
- "another worker still enforces the old limit; the update was never broadcast"
+ assert other_user_after.model_dump()[field] == new_value, (
+ "another worker still serves the old user row; the update was never broadcast"
)
From df3b44dbee8b554c2eda48d3f32729c2e0e00fb3 Mon Sep 17 00:00:00 2001
From: yuneng-jiang
Date: Wed, 14 Oct 2026 14:05:58 -0700
Subject: [PATCH 33/47] feat(proxy): let proxy admins open team member defaults
to team admins (#46353)
* feat(proxy): let proxy admins open team member defaults to team admins
* fix(proxy): keep a team admin's never-resets duration on a team's first member budget
* test(proxy): pin the existing-row and stored never-resets cases for member defaults
---
.../team_admin_field_permissions.py | 88 ++++++++++++++++--
.../management_endpoints/team_endpoints.py | 24 ++++-
.../management/test_team_update.py | 81 ++++++++++++++++
.../test_team_admin_field_permissions.py | 93 +++++++++++++++++--
.../test_team_endpoints.py | 35 +++++++
.../TeamAdminEditableFieldsSettings.test.tsx | 16 ++++
.../team/TeamAdminSettingsForm.test.tsx | 64 ++++++++++++-
.../components/team/TeamAdminSettingsForm.tsx | 39 +++++++-
.../src/components/team/TeamInfo.test.tsx | 29 ++++++
.../src/components/team/TeamInfo.tsx | 4 +
.../team/teamAdminEditAccess.test.ts | 63 +++++++++++++
.../components/team/teamAdminEditAccess.ts | 48 +++++++++-
12 files changed, 557 insertions(+), 27 deletions(-)
diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py
index 67cbb0be45d..3eb54c2147a 100644
--- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py
+++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py
@@ -17,6 +17,7 @@ from pydantic import TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
+from litellm.models.budget import LiteLLM_BudgetTable
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.verification_token import LiteLLM_VerificationToken
from litellm.proxy._types import (
@@ -32,7 +33,16 @@ TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_field
# TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset(
- {"tpm_limit", "rpm_limit", "max_budget", "soft_budget"}
+ {
+ "tpm_limit",
+ "rpm_limit",
+ "max_budget",
+ "soft_budget",
+ "team_member_budget",
+ "team_member_budget_duration",
+ "team_member_tpm_limit",
+ "team_member_rpm_limit",
+ }
)
TEAM_ADMIN_RAISE_MAX_BUDGET_PERMISSION: Final = "raise_max_budget"
TEAM_ADMIN_PROJECTS_PERMISSION: Final = "projects"
@@ -66,6 +76,19 @@ _NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"})
_SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields"
+@dataclass(frozen=True, slots=True)
+class TeamMemberDefaults:
+ team_member_budget: float | None = None
+ team_member_budget_duration: str | None = None
+ team_member_tpm_limit: int | None = None
+ team_member_rpm_limit: int | None = None
+
+
+TEAM_MEMBER_DEFAULT_FIELDS: Final[frozenset[str]] = frozenset(
+ {"team_member_budget", "team_member_budget_duration", "team_member_tpm_limit", "team_member_rpm_limit"}
+)
+
+
@dataclass(frozen=True, slots=True)
class TeamAdminEditAllowed:
request: UpdateTeamRequest
@@ -250,34 +273,76 @@ def _metadata_changes(
)
+def stored_team_member_budget_id(team: LiteLLM_TeamTable) -> str | None:
+ budget_id: Final = _stored_metadata(_JSON_OBJECT.validate_json(team.model_dump_json())).get("team_member_budget_id")
+ return budget_id if isinstance(budget_id, str) else None
+
+
+def team_member_defaults(budget: LiteLLM_BudgetTable | None) -> TeamMemberDefaults | None:
+ if budget is None:
+ return None
+ return TeamMemberDefaults(
+ team_member_budget=budget.max_budget,
+ team_member_budget_duration=budget.budget_duration,
+ team_member_tpm_limit=budget.tpm_limit,
+ team_member_rpm_limit=budget.rpm_limit,
+ )
+
+
def _stored_model_aliases(existing_row: LiteLLM_TeamTable) -> Mapping[str, object]:
table: Final = existing_row.litellm_model_table
return _as_object(_JSON_OBJECT.validate_json(table.model_dump_json()).get("model_aliases")) if table else _EMPTY
def _column_changed(
- field: str, submitted: Mapping[str, object], existing: Mapping[str, object], existing_row: LiteLLM_TeamTable
+ field: str,
+ submitted: Mapping[str, object],
+ existing: Mapping[str, object],
+ existing_row: LiteLLM_TeamTable,
+ member_defaults: Mapping[str, object],
) -> bool:
if field == "model_aliases":
return _as_object(submitted.get(field)) != _stored_model_aliases(existing_row)
+ if field in member_defaults:
+ return submitted.get(field) != member_defaults[field]
if field in LiteLLM_TeamTable.model_fields:
return submitted.get(field) != existing.get(field)
return True
-def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable) -> frozenset[str]:
+def _member_defaults_by_field(member_defaults: TeamMemberDefaults | None) -> Mapping[str, object]:
+ stored: Final = member_defaults or TeamMemberDefaults()
+ return MappingProxyType(
+ {
+ "team_member_budget": stored.team_member_budget,
+ "team_member_budget_duration": stored.team_member_budget_duration,
+ "team_member_tpm_limit": stored.team_member_tpm_limit,
+ "team_member_rpm_limit": stored.team_member_rpm_limit,
+ }
+ )
+
+
+def changed_team_fields(
+ data: UpdateTeamRequest,
+ existing_row: LiteLLM_TeamTable,
+ member_defaults: TeamMemberDefaults | None = None,
+) -> frozenset[str]:
"""Logical field names whose stored value the request would change.
Request and stored row are compared as JSON values so both sides share one representation. Fields the
server folds into metadata are attributed to their own name whether they arrive top-level or inside
- ``metadata``; anything else in ``metadata`` is attributed to ``metadata``. Fields with no stored
- counterpart on the team row count as changed whenever they are sent.
+ ``metadata``; anything else in ``metadata`` is attributed to ``metadata``. The ``team_member_*`` defaults
+ compare against the team's member budget row (``None`` when it has none, read as all null). Fields with no
+ stored counterpart count as changed whenever they are sent.
"""
submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True))
existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json())
column_fields: Final = frozenset(data.model_fields_set) - _NOT_COLUMNS - _METADATA_FOLDED_FIELDS
+ stored_member_defaults: Final = _member_defaults_by_field(member_defaults)
column_changes: Final = frozenset(
- field for field in column_fields if _column_changed(field, submitted, existing, existing_row)
+ field
+ for field in column_fields
+ if _column_changed(field, submitted, existing, existing_row, stored_member_defaults)
)
return column_changes | _metadata_changes(data, submitted, existing)
@@ -295,14 +360,21 @@ def team_admin_edit_verdict(
data: UpdateTeamRequest,
existing: LiteLLM_TeamTable,
permitted: frozenset[str],
+ member_defaults: TeamMemberDefaults | None = None,
) -> TeamAdminEditVerdict:
if not permitted:
return TeamAdminEditingDisabled()
- changed: Final = changed_team_fields(data, existing)
+ changed: Final = changed_team_fields(data, existing, member_defaults)
blocked: Final = sorted(changed - permitted)
if blocked:
return TeamAdminFieldNotPermitted(field=blocked[0])
- return TeamAdminEditAllowed(request=_only_changes(data, changed))
+ creates_member_budget: Final = member_defaults is None and bool(changed & TEAM_MEMBER_DEFAULT_FIELDS)
+ sent_for_new_member_budget: Final = (
+ frozenset(data.model_fields_set) & TEAM_MEMBER_DEFAULT_FIELDS & permitted
+ if creates_member_budget
+ else frozenset[str]()
+ )
+ return TeamAdminEditAllowed(request=_only_changes(data, changed | sent_for_new_member_budget))
def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequest:
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index 9cc1afd7beb..35aded51415 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -164,10 +164,14 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import (
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
+ TEAM_MEMBER_DEFAULT_FIELDS,
+ TeamMemberDefaults,
resolve_team_admin_editable_fields,
+ stored_team_member_budget_id,
team_admin_edit_verdict,
team_admin_may_raise_max_budget,
team_admin_request_or_raise,
+ team_member_defaults,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
TEAM_ADVISORY_LOCK_SQL,
@@ -2365,12 +2369,19 @@ async def update_team(
if access_role is None:
team_access_denied()
if access_role == "team_admin":
+ permitted: Final = resolve_team_admin_editable_fields(
+ _general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
+ )
+ reads_member_defaults: Final = bool(permitted) and bool(data.model_fields_set & TEAM_MEMBER_DEFAULT_FIELDS)
data = team_admin_request_or_raise( # rebind-ok: resent values must not reach the derived writes below
team_admin_edit_verdict(
data=data,
existing=existing_team,
- permitted=resolve_team_admin_editable_fields(
- _general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
+ permitted=permitted,
+ member_defaults=(
+ await _stored_team_member_defaults(prisma_client, existing_team)
+ if reads_member_defaults
+ else None
),
)
)
@@ -4899,6 +4910,15 @@ async def validate_membership(user_api_key_dict: UserAPIKeyAuth, team_table: Lit
)
+async def _stored_team_member_defaults(
+ prisma_client: PrismaClient, team: LiteLLM_TeamTable
+) -> TeamMemberDefaults | None:
+ budget_id: Final = stored_team_member_budget_id(team)
+ if budget_id is None:
+ return None
+ return team_member_defaults(await BudgetRepository(prisma_client).find_by_id(budget_id))
+
+
async def _add_team_member_budget_table(
team_member_budget_id: str,
prisma_client: PrismaClient,
diff --git a/tests/integration/management/test_team_update.py b/tests/integration/management/test_team_update.py
index 0414d0bb77f..e1d9c43f318 100644
--- a/tests/integration/management/test_team_update.py
+++ b/tests/integration/management/test_team_update.py
@@ -162,6 +162,87 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text
+_MEMBER_DEFAULTS = [
+ pytest.param("team_member_budget", 25.0, "max_budget", id="team_member_budget"),
+ pytest.param("team_member_budget_duration", "7d", "budget_duration", id="team_member_budget_duration"),
+ pytest.param("team_member_tpm_limit", 200, "tpm_limit", id="team_member_tpm_limit"),
+ pytest.param("team_member_rpm_limit", 9, "rpm_limit", id="team_member_rpm_limit"),
+]
+
+
+@pytest.mark.parametrize("field,value,budget_column", _MEMBER_DEFAULTS)
+async def test_team_update_team_admin_member_default_follows_the_proxy_admin_grant(
+ field: str, value: str | float, budget_column: str, gateway_client, db, resources, actors
+):
+ callers = await _callers(db, actors, resources.prefix)
+ await _seed_target(db, actors, callers, "alpha", resources.prefix)
+ headers = {"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"}
+ body = {"team_id": resources.prefix, field: value}
+
+ with team_admin_permissions(gateway_client.gateway, ["tpm_limit"]):
+ denied = await gateway_client.post("/team/update", headers=headers, json=body)
+ assert denied.status_code == 403, denied.text
+ assert f"'{field}'" in denied.text and "Team admin editable fields" in denied.text, denied.text
+
+ with team_admin_permissions(gateway_client.gateway, [field]):
+ allowed = await gateway_client.post("/team/update", headers=headers, json=body)
+ assert allowed.status_code == 200, allowed.text
+
+ info = await gateway_client.get("/team/info", params={"team_id": resources.prefix})
+ assert info.status_code == 200, info.text
+ assert info.json()["team_info"]["team_member_budget_table"][budget_column] == value
+
+
+async def test_team_update_team_admin_resending_member_defaults_leaves_the_budget_row_alone(
+ gateway_client, db, resources, actors
+):
+ callers = await _callers(db, actors, resources.prefix)
+ await _seed_target(db, actors, callers, "alpha", resources.prefix)
+ stored = {
+ "team_member_budget": 10.0,
+ "team_member_budget_duration": "30d",
+ "team_member_tpm_limit": 100,
+ "team_member_rpm_limit": 5,
+ }
+ seeded = await gateway_client.post("/team/update", json={"team_id": resources.prefix, **stored})
+ assert seeded.status_code == 200, seeded.text
+ team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
+ assert team is not None
+ budget_id = team.metadata["team_member_budget_id"]
+ before = await db.db.litellm_budgettable.find_unique(where={"budget_id": budget_id})
+ assert before is not None
+
+ with team_admin_permissions(gateway_client.gateway, ["tpm_limit"]):
+ resp = await gateway_client.post(
+ "/team/update",
+ headers={"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"},
+ json={"team_id": resources.prefix, "tpm_limit": 1000, **stored},
+ )
+ assert resp.status_code == 200, resp.text
+
+ after = await db.db.litellm_budgettable.find_unique(where={"budget_id": budget_id})
+ assert after is not None and after.updated_at == before.updated_at
+ team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
+ assert team is not None and team.tpm_limit == 1000
+
+
+async def test_team_update_team_admin_cannot_change_member_defaults_of_another_team(
+ gateway_client, db, resources, actors
+):
+ callers = await _callers(db, actors, resources.prefix)
+ await _seed_target(db, actors, callers, "beta", resources.prefix)
+
+ with team_admin_permissions(gateway_client.gateway, ["team_member_budget"]):
+ resp = await gateway_client.post(
+ "/team/update",
+ headers={"Authorization": f"Bearer {callers[Actor.TEAM_ADMIN].cleartext}"},
+ json={"team_id": resources.prefix, "team_member_budget": 25.0},
+ )
+ assert resp.status_code == 403, resp.text
+ team = await db.db.litellm_teamtable.find_unique(where={"team_id": resources.prefix})
+ assert team is not None and "team_member_budget_id" not in (team.metadata or {}), "denied but team mutated"
+
+
async def test_team_update_team_admin_soft_budget_follows_the_proxy_admin_grant(gateway_client, db, resources, actors):
callers = await _callers(db, actors, resources.prefix)
await _seed_target(db, actors, callers, "alpha", resources.prefix)
diff --git a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py
index 0c9ed9356f8..0ffb3694942 100644
--- a/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py
+++ b/tests/unit/proxy/management_endpoints/test_team_admin_field_permissions.py
@@ -19,6 +19,7 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import (
TeamAdminFieldNotPermitted,
TeamAdminKeyEditAllowed,
TeamAdminMemberKeyEditingDisabled,
+ TeamMemberDefaults,
changed_key_fields,
changed_team_fields,
resolve_team_admin_editable_fields,
@@ -160,8 +161,8 @@ class TestChangedTeamFields:
assert changed_team_fields(UpdateTeamRequest(team_id="team-1", model_aliases={}), _team()) == frozenset()
def test_field_without_a_stored_counterpart_counts_as_changed_when_sent(self):
- data = UpdateTeamRequest(team_id="team-1", team_member_budget=10.0)
- assert changed_team_fields(data, _team()) == frozenset({"team_member_budget"})
+ data = UpdateTeamRequest(team_id="team-1", team_member_key_duration="30d")
+ assert changed_team_fields(data, _team()) == frozenset({"team_member_key_duration"})
class TestTeamAdminEditVerdict:
@@ -192,27 +193,105 @@ class TestTeamAdminEditVerdict:
assert verdict == TeamAdminFieldNotPermitted(field="blocked")
-def _enabled(*fields: str) -> frozenset[str]:
+def _permitted(*fields: str) -> frozenset[str]:
return resolve_team_admin_editable_fields(
{"team_admin_editable_team_fields": list(fields)}, SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
)
+_STORED_MEMBER_DEFAULTS: Final = TeamMemberDefaults(
+ team_member_budget=10.0, team_member_budget_duration="30d", team_member_tpm_limit=100, team_member_rpm_limit=5
+)
+
+_MEMBER_DEFAULT_EDITS: Final = (
+ pytest.param("team_member_budget", 25.0, 10.0, id="team_member_budget"),
+ pytest.param("team_member_budget_duration", "7d", "30d", id="team_member_budget_duration"),
+ pytest.param("team_member_tpm_limit", 200, 100, id="team_member_tpm_limit"),
+ pytest.param("team_member_rpm_limit", 9, 5, id="team_member_rpm_limit"),
+)
+
+
+class TestTeamMemberDefaultEdits:
+ @pytest.mark.parametrize("field,new,_stored", _MEMBER_DEFAULT_EDITS)
+ def test_enabled_field_change_is_allowed(self, field: str, new: object, _stored: object):
+ data: Final = UpdateTeamRequest.model_validate({"team_id": "team-1", field: new})
+ verdict: Final = team_admin_edit_verdict(data, _team(), _permitted(field), _STORED_MEMBER_DEFAULTS)
+ assert isinstance(verdict, TeamAdminEditAllowed)
+ assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", field: new}
+
+ @pytest.mark.parametrize("field,new,_stored", _MEMBER_DEFAULT_EDITS)
+ def test_change_without_the_field_enabled_names_it(self, field: str, new: object, _stored: object):
+ data: Final = UpdateTeamRequest.model_validate({"team_id": "team-1", field: new})
+ verdict: Final = team_admin_edit_verdict(data, _team(), _permitted("tpm_limit"), _STORED_MEMBER_DEFAULTS)
+ assert verdict == TeamAdminFieldNotPermitted(field=field)
+
+ @pytest.mark.parametrize("field,_new,stored", _MEMBER_DEFAULT_EDITS)
+ def test_resending_the_stored_value_is_dropped_without_the_field_enabled(
+ self, field: str, _new: object, stored: object
+ ):
+ data: Final = UpdateTeamRequest.model_validate({"team_id": "team-1", field: stored, "tpm_limit": 6})
+ verdict: Final = team_admin_edit_verdict(
+ data, _team(tpm_limit=5), _permitted("tpm_limit"), _STORED_MEMBER_DEFAULTS
+ )
+ assert isinstance(verdict, TeamAdminEditAllowed)
+ assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6}
+
+ def test_nulls_resent_for_a_team_without_member_defaults_are_dropped(self):
+ data: Final = UpdateTeamRequest(
+ team_id="team-1",
+ team_member_budget=None,
+ team_member_budget_duration=None,
+ team_member_tpm_limit=None,
+ team_member_rpm_limit=None,
+ tpm_limit=6,
+ )
+ verdict: Final = team_admin_edit_verdict(data, _team(tpm_limit=5), _permitted("tpm_limit"), None)
+ assert isinstance(verdict, TeamAdminEditAllowed)
+ assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6}
+
+ def test_first_member_budget_keeps_an_explicit_never_resets_duration(self):
+ data: Final = UpdateTeamRequest(team_id="team-1", team_member_budget=5.0, team_member_budget_duration=None)
+ permitted: Final = _permitted("team_member_budget", "team_member_budget_duration")
+ verdict: Final = team_admin_edit_verdict(data, _team(budget_duration="7d"), permitted, None)
+ assert isinstance(verdict, TeamAdminEditAllowed)
+ assert verdict.request.model_dump(exclude_unset=True) == {
+ "team_id": "team-1",
+ "team_member_budget": 5.0,
+ "team_member_budget_duration": None,
+ }
+
+ def test_first_member_budget_drops_a_duration_the_team_admin_may_not_set(self):
+ data: Final = UpdateTeamRequest(team_id="team-1", team_member_budget=5.0, team_member_budget_duration=None)
+ verdict: Final = team_admin_edit_verdict(
+ data, _team(budget_duration="7d"), _permitted("team_member_budget"), None
+ )
+ assert isinstance(verdict, TeamAdminEditAllowed)
+ assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "team_member_budget": 5.0}
+
+ def test_existing_member_budget_drops_a_resent_duration_next_to_a_changed_budget(self):
+ data: Final = UpdateTeamRequest(team_id="team-1", team_member_budget=5.0, team_member_budget_duration="7d")
+ stored: Final = TeamMemberDefaults(team_member_budget=1.0, team_member_budget_duration="7d")
+ permitted: Final = _permitted("team_member_budget", "team_member_budget_duration")
+ verdict: Final = team_admin_edit_verdict(data, _team(), permitted, stored)
+ assert isinstance(verdict, TeamAdminEditAllowed)
+ assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "team_member_budget": 5.0}
+
+
class TestSoftBudgetEdits:
def test_enabled_soft_budget_change_is_allowed(self):
data: Final = UpdateTeamRequest(team_id="team-1", soft_budget=20.0)
- verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("soft_budget"))
+ verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _permitted("soft_budget"))
assert isinstance(verdict, TeamAdminEditAllowed)
assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "soft_budget": 20.0}
def test_soft_budget_change_without_it_enabled_names_it(self):
data: Final = UpdateTeamRequest(team_id="team-1", soft_budget=20.0)
- verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("max_budget"))
+ verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _permitted("max_budget"))
assert verdict == TeamAdminFieldNotPermitted(field="soft_budget")
- def test_unchanged_soft_budget_resend_is_dropped_when_not_enabled(self):
+ def test_unchanged_soft_budget_resend_is_dropped_when_not_permitted(self):
data: Final = UpdateTeamRequest(team_id="team-1", tpm_limit=6, soft_budget=10)
- verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _enabled("tpm_limit"))
+ verdict: Final = team_admin_edit_verdict(data, _team(soft_budget=10.0), _permitted("tpm_limit"))
assert isinstance(verdict, TeamAdminEditAllowed)
assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6}
diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py
index bbb83ec1248..15eed3b105a 100644
--- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py
@@ -16950,6 +16950,41 @@ async def test_update_team_team_admin_echoing_unpermitted_fields_unchanged_is_al
assert prisma.db.litellm_teamtable.update.called
+@pytest.mark.asyncio
+async def test_update_team_team_admin_resending_stored_member_defaults_is_not_an_edit(
+ disable_audit_logging_for_mocked_team,
+):
+ import contextlib
+ from litellm.models.budget import LiteLLM_BudgetTable
+
+ with contextlib.ExitStack() as stack:
+ prisma = _wire_update_team(stack, {"team_member_budget_id": "member-budget-1"})
+ prisma.db.litellm_budgettable.find_unique = AsyncMock(
+ return_value=LiteLLM_BudgetTable(
+ budget_id="member-budget-1", max_budget=10.0, budget_duration="30d", tpm_limit=100, rpm_limit=5
+ )
+ )
+ stack.enter_context(
+ patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": ["tpm_limit"]}) # test-quality-ok: update_team reads general_settings as a proxy_server module global
+ )
+ result = await update_team(
+ data=UpdateTeamRequest(
+ team_id="test_team_id",
+ team_member_budget=10.0,
+ team_member_budget_duration="30d",
+ team_member_tpm_limit=100,
+ team_member_rpm_limit=5,
+ tpm_limit=6,
+ ),
+ http_request=_update_request_stub(),
+ user_api_key_dict=_TEAM_ADMIN_CALLER,
+ )
+
+ assert result["data"].team_id == "test_team_id"
+ assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["tpm_limit"] == 6
+ assert not prisma.db.litellm_budgettable.update.called
+
+
@pytest.mark.asyncio
async def test_update_team_team_admin_changes_tpm_limit_once_a_proxy_admin_enables_it(
disable_audit_logging_for_mocked_team,
diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx
index 36fbe086df3..df3dafc4b28 100644
--- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx
+++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings.test.tsx
@@ -127,6 +127,22 @@ describe("TeamAdminEditableFieldsSettings", () => {
expect(mutate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["max_budget"] }, expect.anything());
});
+ it("offers the member defaults as named checkboxes and saves the ticked ones", async () => {
+ mockSettings(["team_member_budget", "team_member_budget_duration", "tpm_limit"], []);
+ const mutate = mockSave({});
+
+ renderWithProviders();
+ fireEvent.click(screen.getByRole("checkbox", { name: "Member Default Budget (USD)" }));
+ fireEvent.click(screen.getByRole("checkbox", { name: "Member Default Budget Duration" }));
+ fireEvent.click(saveButton());
+
+ await waitFor(() => expect(mutate).toHaveBeenCalledTimes(1));
+ expect(mutate).toHaveBeenCalledWith(
+ { team_admin_editable_team_fields: ["team_member_budget", "team_member_budget_duration"] },
+ expect.anything(),
+ );
+ });
+
it("offers the soft budget as a named checkbox and saves it when ticked", async () => {
mockSettings(["max_budget", "soft_budget"], []);
const mutate = mockSave({});
diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx
index 7fe6461e285..a84d7bc4bac 100644
--- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.test.tsx
@@ -1,7 +1,7 @@
import { describe, expect, it, vi } from "vitest";
import userEvent from "@testing-library/user-event";
-import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils";
+import { chooseSelectOption, fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils";
import TeamAdminSettingsForm from "./TeamAdminSettingsForm";
@@ -13,7 +13,16 @@ const renderForm = (
const onCancel = vi.fn();
renderWithProviders(
{
expect(save).toBeDisabled();
});
+ it("shows the team's member defaults when the proxy lets team admins edit them", () => {
+ renderForm(
+ new Set(["team_member_budget", "team_member_budget_duration", "team_member_tpm_limit", "team_member_rpm_limit"]),
+ );
+
+ expect(screen.getByLabelText("Member Default Budget (USD)")).toHaveValue(10);
+ expect(screen.getByLabelText("Member Default Budget Duration")).toHaveTextContent("monthly");
+ expect(screen.getByLabelText("Member Default TPM Limit")).toHaveValue(100);
+ expect(screen.getByLabelText("Member Default RPM Limit")).toHaveValue(5);
+ });
+
+ it("saves a new member default budget and a weekly reset without resending the unchanged limits", async () => {
+ const user = userEvent.setup();
+ const { onSave } = renderForm(
+ new Set(["team_member_budget", "team_member_budget_duration", "team_member_tpm_limit", "team_member_rpm_limit"]),
+ );
+
+ fireEvent.change(screen.getByLabelText("Member Default Budget (USD)"), { target: { value: "25" } });
+ await chooseSelectOption(user, screen.getByLabelText("Member Default Budget Duration"), "weekly");
+ await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+ await waitFor(() =>
+ expect(onSave).toHaveBeenCalledWith({ team_member_budget: 25, team_member_budget_duration: "7d" }),
+ );
+ });
+
+ it("saves a member budget that never resets as a null duration", async () => {
+ const user = userEvent.setup();
+ const { onSave } = renderForm(new Set(["team_member_budget_duration"]));
+
+ await chooseSelectOption(user, screen.getByLabelText("Member Default Budget Duration"), "Never resets");
+ await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+ await waitFor(() => expect(onSave).toHaveBeenCalledWith({ team_member_budget_duration: null }));
+ });
+
+ it("shows a stored member budget that never resets as Never resets", () => {
+ renderWithProviders(
+ ,
+ );
+
+ expect(screen.getByLabelText("Member Default Budget Duration")).toHaveTextContent("Never resets");
+ });
+
it("shows the soft budget when the proxy lets team admins edit it", () => {
renderForm(new Set(["soft_budget"]));
diff --git a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx
index 5ce50bb6c72..05595d6a1a4 100644
--- a/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamAdminSettingsForm.tsx
@@ -10,8 +10,10 @@ import { FieldGroup } from "@/components/ui/field";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { useZodForm } from "@/lib/forms/useZodForm";
+import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown";
import NumericalInput from "../shared/numerical_input";
import {
+ isTeamAdminDurationSettingsField,
TEAM_ADMIN_SETTINGS_FIELDS,
teamAdminFieldLabel,
teamAdminSettingsChanges,
@@ -27,15 +29,22 @@ const teamAdminSettingsShape = {
rpm_limit: numericInputSchema,
max_budget: numericInputSchema,
soft_budget: numericInputSchema,
+ team_member_budget: numericInputSchema,
+ team_member_budget_duration: z.string().nullish(),
+ team_member_tpm_limit: numericInputSchema,
+ team_member_rpm_limit: numericInputSchema,
};
const teamAdminSettingsSchema = z.object(teamAdminSettingsShape);
-const INPUT_STEP: Readonly> = {
+const INPUT_STEP: Readonly, number>> = {
tpm_limit: 1,
rpm_limit: 1,
max_budget: 0.01,
soft_budget: 0.01,
+ team_member_budget: 0.01,
+ team_member_tpm_limit: 1,
+ team_member_rpm_limit: 1,
};
const MAX_BUDGET_KEEP_OR_LOWER_HINT = "You can keep or lower this budget. Ask a proxy admin to raise it.";
@@ -47,6 +56,11 @@ const fieldHint = (name: TeamAdminSettingsField, mayRaiseMaxBudget: boolean): st
return mayRaiseMaxBudget ? MAX_BUDGET_RAISE_HINT : MAX_BUDGET_KEEP_OR_LOWER_HINT;
};
+const durationSelectValue = (value: string | number | null | undefined): string | undefined => {
+ if (value === null) return NEVER_RESETS_BUDGET_DURATION;
+ return typeof value === "string" ? value : undefined;
+};
+
interface TeamAdminSettingsFormProps {
initialValues: TeamAdminSettingsValues;
editableFields: ReadonlySet;
@@ -83,9 +97,26 @@ export default function TeamAdminSettingsForm({
label={teamAdminFieldLabel(name)}
description={fieldHint(name, mayRaiseMaxBudget)}
>
- {({ ref, value, ...field }) => (
-
- )}
+ {({ ref, id, value, onChange, ...field }) =>
+ isTeamAdminDurationSettingsField(name) ? (
+ onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next ?? undefined)}
+ />
+ ) : (
+
+ )
+ }
))}
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
index 31d3016d4be..973b2ecee80 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
@@ -2047,6 +2047,35 @@ describe("TeamInfoView", () => {
expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled();
});
+ it("prefills the member defaults a team admin may edit from the team's member budget", async () => {
+ const user = userEvent.setup({ delay: null });
+ vi.mocked(networking.teamInfoCall).mockResolvedValue(
+ createMockTeamData({
+ team_member_budget_table: { max_budget: 10, budget_duration: "30d", tpm_limit: 100, rpm_limit: 5 },
+ caller_edit_access: {
+ kind: "team_admin",
+ editable_fields: [
+ "team_member_budget",
+ "team_member_budget_duration",
+ "team_member_tpm_limit",
+ "team_member_rpm_limit",
+ ],
+ },
+ }),
+ );
+
+ renderWithProviders();
+
+ await user.click(await screen.findByRole("tab", { name: "Settings" }));
+ await user.click(await screen.findByRole("button", { name: /edit settings/i }));
+
+ expect(await screen.findByLabelText("Member Default Budget (USD)")).toHaveValue(10);
+ expect(screen.getByLabelText("Member Default Budget Duration")).toHaveTextContent("monthly");
+ expect(screen.getByLabelText("Member Default TPM Limit")).toHaveValue(100);
+ expect(screen.getByLabelText("Member Default RPM Limit")).toHaveValue(5);
+ expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled();
+ });
+
it("opens the form when the proxy reports unrestricted access although the props only mark a team admin", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
index cfac81cb64b..0743e73005e 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -1311,6 +1311,10 @@ const TeamInfoView: React.FC = ({
rpm_limit: info.rpm_limit,
max_budget: info.max_budget,
soft_budget: info.soft_budget,
+ team_member_budget: info.team_member_budget_table?.max_budget,
+ team_member_budget_duration: info.team_member_budget_table?.budget_duration,
+ team_member_tpm_limit: info.team_member_budget_table?.tpm_limit,
+ team_member_rpm_limit: info.team_member_budget_table?.rpm_limit,
}}
editableFields={teamEditAccess.editableFields}
mayRaiseMaxBudget={teamEditAccess.mayRaiseMaxBudget}
diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts
index 590dd806d7f..84898563899 100644
--- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts
+++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts
@@ -16,6 +16,10 @@ describe("teamAdminFieldLabel", () => {
["soft_budget", "Soft Budget (USD)"],
["raise_max_budget", "Raise the team's max budget"],
["projects", "Create and update projects"],
+ ["team_member_budget", "Member Default Budget (USD)"],
+ ["team_member_budget_duration", "Member Default Budget Duration"],
+ ["team_member_tpm_limit", "Member Default TPM Limit"],
+ ["team_member_rpm_limit", "Member Default RPM Limit"],
["create_member_keys", "Create keys for team members"],
])("names %s the way the team settings form does", (field, label) => {
expect(teamAdminFieldLabel(field)).toBe(label);
@@ -77,6 +81,65 @@ describe("teamAdminSettingsChanges", () => {
});
});
+describe("teamAdminSettingsChanges for team member defaults", () => {
+ const stored = {
+ team_member_budget: 10,
+ team_member_budget_duration: "30d",
+ team_member_tpm_limit: 100,
+ team_member_rpm_limit: 5,
+ };
+ const enabled = new Set([
+ "team_member_budget",
+ "team_member_budget_duration",
+ "team_member_tpm_limit",
+ "team_member_rpm_limit",
+ ]);
+
+ const durationEnabled = new Set(["team_member_budget_duration"]);
+
+ it("sends every changed member default and skips the unchanged ones", () => {
+ const values = {
+ team_member_budget: "25",
+ team_member_budget_duration: "7d",
+ team_member_tpm_limit: "100",
+ team_member_rpm_limit: "9",
+ };
+
+ expect(teamAdminSettingsChanges(values, stored, enabled)).toStrictEqual({
+ team_member_budget: 25,
+ team_member_budget_duration: "7d",
+ team_member_rpm_limit: 9,
+ });
+ });
+
+ it("sends a never-resetting duration as null", () => {
+ expect(teamAdminSettingsChanges({ team_member_budget_duration: null }, stored, durationEnabled)).toStrictEqual({
+ team_member_budget_duration: null,
+ });
+ });
+
+ it("sends a never-resetting duration for a team without member defaults yet", () => {
+ expect(teamAdminSettingsChanges({ team_member_budget_duration: null }, {}, durationEnabled)).toStrictEqual({
+ team_member_budget_duration: null,
+ });
+ });
+
+ it.each([
+ ["the stored duration", "30d", stored],
+ ["an untouched duration on a team with a stored duration", undefined, stored],
+ ["an untouched duration on a team without member defaults", undefined, {}],
+ ["a stored never-resetting duration picked again", null, { team_member_budget_duration: null }],
+ ])("sends no duration for %s", (_label, team_member_budget_duration, initialValues) => {
+ expect(teamAdminSettingsChanges({ team_member_budget_duration }, initialValues, durationEnabled)).toStrictEqual({});
+ });
+
+ it("leaves out member defaults the proxy did not enable", () => {
+ const values = { team_member_budget: "25", team_member_budget_duration: "7d" };
+
+ expect(teamAdminSettingsChanges(values, stored, new Set(["max_budget"]))).toStrictEqual({});
+ });
+});
+
describe("teamAdminSettingsChanges for the soft budget", () => {
const stored = { soft_budget: 50 };
const enabled = new Set(["soft_budget"]);
diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
index 3ab51b4cdbb..28ba5eed3e4 100644
--- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
+++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts
@@ -48,15 +48,36 @@ export const parseSupportedTeamAdminEditableFields = (uiSettingsFieldSchema: unk
return items.success ? fieldListSchema.parse(items.data.enum) : [];
};
-export const TEAM_ADMIN_SETTINGS_FIELDS = ["tpm_limit", "rpm_limit", "max_budget", "soft_budget"] as const;
+export const TEAM_ADMIN_SETTINGS_FIELDS = [
+ "tpm_limit",
+ "rpm_limit",
+ "max_budget",
+ "soft_budget",
+ "team_member_budget",
+ "team_member_budget_duration",
+ "team_member_tpm_limit",
+ "team_member_rpm_limit",
+] as const;
export type TeamAdminSettingsField = (typeof TEAM_ADMIN_SETTINGS_FIELDS)[number];
+type TeamAdminDurationSettingsField = "team_member_budget_duration";
+
+type TeamAdminNumericSettingsField = Exclude;
+
+export const isTeamAdminDurationSettingsField = (
+ field: TeamAdminSettingsField,
+): field is TeamAdminDurationSettingsField => field === "team_member_budget_duration";
+
const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([
["tpm_limit", "Tokens per minute Limit (TPM)"],
["rpm_limit", "Requests per minute Limit (RPM)"],
["max_budget", "Max Budget (USD)"],
["soft_budget", "Soft Budget (USD)"],
+ ["team_member_budget", "Member Default Budget (USD)"],
+ ["team_member_budget_duration", "Member Default Budget Duration"],
+ ["team_member_tpm_limit", "Member Default TPM Limit"],
+ ["team_member_rpm_limit", "Member Default RPM Limit"],
[TEAM_ADMIN_RAISE_MAX_BUDGET_PERMISSION, "Raise the team's max budget"],
["projects", "Create and update projects"],
["member_key_budgets", "Update budgets on team members' keys"],
@@ -67,7 +88,26 @@ export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_L
export type TeamAdminSettingsValues = { readonly [F in TeamAdminSettingsField]?: string | number | null };
-export type TeamAdminSettingsChanges = { readonly [F in TeamAdminSettingsField]?: number | null };
+export type TeamAdminSettingsChanges = { readonly [F in TeamAdminNumericSettingsField]?: number | null } & {
+ readonly [F in TeamAdminDurationSettingsField]?: string | null;
+};
+
+const durationOrNull = (value: string | number | null | undefined): string | null =>
+ typeof value === "string" && value !== "" ? value : null;
+
+const changedValue = (
+ field: TeamAdminSettingsField,
+ value: string | number | null | undefined,
+ initialValue: string | number | null | undefined,
+): readonly [TeamAdminSettingsField, string | number | null] | null => {
+ if (!isTeamAdminDurationSettingsField(field)) {
+ const number = numberOrNull(value);
+ return number === numberOrNull(initialValue) ? null : [field, number];
+ }
+ if (value === undefined) return null;
+ const duration = durationOrNull(value);
+ return initialValue !== undefined && duration === durationOrNull(initialValue) ? null : [field, duration];
+};
export const teamAdminSettingsChanges = (
values: TeamAdminSettingsValues,
@@ -76,8 +116,8 @@ export const teamAdminSettingsChanges = (
): TeamAdminSettingsChanges =>
Object.fromEntries(
TEAM_ADMIN_SETTINGS_FIELDS.flatMap((field) => {
- const value = numberOrNull(values[field]);
- return editableFields.has(field) && value !== numberOrNull(initialValues[field]) ? [[field, value]] : [];
+ const change = editableFields.has(field) ? changedValue(field, values[field], initialValues[field]) : null;
+ return change === null ? [] : [change];
}),
);
From 542b254ca2aab496d1c0db8450a454cb92707476 Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 14:07:59 -0700
Subject: [PATCH 34/47] fix(anthropic): forward service_tier on /v1/messages
for anthropic, azure_ai and bedrock_mantle (#46370)
* test(integration): add known-bug translation cases for request fields LiteLLM drops
service_tier on /v1/messages, conversation on /v1/responses and metadata on
/v1/chat/completions never reach the provider. Each gets a TranslationTestCase
with the correct provider body, marked xfail(strict=True) through the new
TranslationTestCase.known_bug field, plus a passing service_tier case for
/v1/messages on the openai route
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(integration): widen the known-bug cases to azure_ai, bedrock_mantle, openai/responses and azure deployments
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(anthropic): forward service_tier on /v1/messages for anthropic, azure_ai and bedrock_mantle
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/llms/anthropic/chat/transformation.py | 35 +++++++++++++---
.../pass_through/messages/transformation.py | 1 +
.../anthropic/pass_through/messages/utils.py | 5 +++
.../bedrock/messages/mantle_transformation.py | 4 ++
litellm/types/llms/anthropic.py | 1 +
.../messages/service_tier/__init__.py | 0
.../test_messages_service_tier_anthropic.py | 29 +++++++++++++
.../test_messages_service_tier_azure_ai.py | 29 +++++++++++++
...st_messages_service_tier_bedrock_mantle.py | 31 ++++++++++++++
.../test_messages_service_tier_openai.py | 42 +++++++++++++++++++
.../test_anthropic_messages_service_tier.py | 28 +++++++++++++
11 files changed, 199 insertions(+), 6 deletions(-)
create mode 100644 tests/integration/translation/messages/service_tier/__init__.py
create mode 100644 tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py
create mode 100644 tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py
create mode 100644 tests/integration/translation/messages/service_tier/test_messages_service_tier_bedrock_mantle.py
create mode 100644 tests/integration/translation/messages/service_tier/test_messages_service_tier_openai.py
create mode 100644 tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_service_tier.py
diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py
index 7abc13d5508..f4e50bf3acd 100644
--- a/litellm/llms/anthropic/chat/transformation.py
+++ b/litellm/llms/anthropic/chat/transformation.py
@@ -294,7 +294,6 @@ DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING: Final = (
DROP_UNSUPPORTED_SPEED_WARNING: Final = (
"Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models."
)
-
SERVICE_TIER_PROVIDERS: Final = frozenset({"anthropic", "azure_ai"})
OPENAI_TO_ANTHROPIC_SERVICE_TIER: Final = MappingProxyType(
@@ -306,6 +305,12 @@ DROP_UNSUPPORTED_SERVICE_TIER_WARNING: Final = (
"Anthropic only accepts 'auto' and 'standard_only'."
)
+MESSAGES_SERVICE_TIER_PROVIDERS: Final = SERVICE_TIER_PROVIDERS | {"bedrock_mantle"}
+
+DROP_SERVICE_TIER_PROVIDER_WARNING: Final = (
+ "Dropping `service_tier` for model=%s on provider=%s. Only the Anthropic API routes accept it."
+)
+
class AnthropicConfig(AnthropicModelInfo, BaseConfig):
"""
@@ -539,6 +544,24 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
) -> None:
return cls._maybe_drop_speed_param(model, optional_params, drop_params, custom_llm_provider)
+ @staticmethod
+ def provider_supports_service_tier(custom_llm_provider: str | None) -> bool:
+ """Anthropic's ``service_tier`` is accepted by the direct API and the Foundry and Mantle gateways in
+ front of it. Vertex rejects the field and Bedrock Invoke only takes its own tier names."""
+ return custom_llm_provider is None or custom_llm_provider in MESSAGES_SERVICE_TIER_PROVIDERS
+
+ @classmethod
+ def maybe_drop_service_tier_param(
+ cls,
+ model: str,
+ optional_params: dict[str, object], # mutable-ok: mirrors override contract
+ custom_llm_provider: str | None = None,
+ ) -> None:
+ if "service_tier" not in optional_params or cls.provider_supports_service_tier(custom_llm_provider):
+ return
+ litellm.verbose_logger.warning(DROP_SERVICE_TIER_PROVIDER_WARNING, model, custom_llm_provider)
+ optional_params.pop("service_tier", None)
+
@staticmethod
def _raise_invalid_reasoning_effort(model: str, value: object, llm_provider: str) -> NoReturn:
"""Raise a ``BadRequestError`` for an unrecognised ``reasoning_effort``.
@@ -717,12 +740,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
# Already handled above
continue
- if key == "properties" and isinstance(value, dict):
- result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()}
+ if key in ("properties", "$defs") and isinstance(value, dict):
+ nested_schemas: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict
+ result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in nested_schemas.items()}
elif key == "items" and isinstance(value, dict):
- result[key] = AnthropicConfig.filter_anthropic_output_schema(value)
- elif key == "$defs" and isinstance(value, dict):
- result[key] = {k: AnthropicConfig.filter_anthropic_output_schema(v) for k, v in value.items()}
+ nested_schema: dict[str, Any] = value # mutable-ok: filter_anthropic_output_schema takes a dict
+ result[key] = AnthropicConfig.filter_anthropic_output_schema(nested_schema)
elif key == "anyOf" and isinstance(value, list):
result["anyOf"] = result.get("anyOf", []) + [
AnthropicConfig.filter_anthropic_output_schema(item) for item in value
diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py
index f1eb30b63c1..4c667574d40 100644
--- a/litellm/llms/anthropic/pass_through/messages/transformation.py
+++ b/litellm/llms/anthropic/pass_through/messages/transformation.py
@@ -98,6 +98,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
"output_format",
"inference_geo",
"speed",
+ "service_tier",
"output_config",
"reasoning_effort",
"safeguards",
diff --git a/litellm/llms/anthropic/pass_through/messages/utils.py b/litellm/llms/anthropic/pass_through/messages/utils.py
index 8371615baea..36bd10ab8e6 100644
--- a/litellm/llms/anthropic/pass_through/messages/utils.py
+++ b/litellm/llms/anthropic/pass_through/messages/utils.py
@@ -205,6 +205,11 @@ class AnthropicMessagesRequestUtils:
drop_params=drop_params,
custom_llm_provider=custom_llm_provider,
)
+ AnthropicConfig.maybe_drop_service_tier_param(
+ model=model,
+ optional_params=filtered_params,
+ custom_llm_provider=custom_llm_provider,
+ )
for param in ("temperature", "top_p", "top_k"):
if param in filtered_params:
AnthropicModelInfo._apply_sampling_param( # pyright: ignore[reportPrivateUsage] # same gating the /chat/completions path applies; forking it would drift
diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py
index 0ee5dcbe557..ca93cc3e617 100644
--- a/litellm/llms/bedrock/messages/mantle_transformation.py
+++ b/litellm/llms/bedrock/messages/mantle_transformation.py
@@ -56,6 +56,10 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig):
model ID in the request body (unlike Bedrock Invoke which puts it in the URL).
"""
+ BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = (
+ AmazonAnthropicClaudeMessagesConfig.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS | {"service_tier"}
+ )
+
@property
def beta_headers_provider(self) -> str:
return "bedrock_mantle"
diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py
index 25f55ddb84b..712f41ecc61 100644
--- a/litellm/types/llms/anthropic.py
+++ b/litellm/types/llms/anthropic.py
@@ -430,6 +430,7 @@ class AnthropicMessagesRequestOptionalParams(TypedDict, total=False):
cache_control: dict[str, Any] | None # Automatic prompt caching
reasoning_effort: str | None
safeguards: ReadOnly[list[dict[str, object]] | None]
+ service_tier: ReadOnly[str | None]
class AnthropicMessagesRequest(AnthropicMessagesRequestOptionalParams, total=False):
diff --git a/tests/integration/translation/messages/service_tier/__init__.py b/tests/integration/translation/messages/service_tier/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py
new file mode 100644
index 00000000000..de453564e15
--- /dev/null
+++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_anthropic.py
@@ -0,0 +1,29 @@
+from dataclasses import replace
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.provider import SharedProvider
+from integration.translation.case import TranslationTestCase
+from integration.translation.messages.bases.anthropic import CLAUDE_SONNET_4_6_TEST_CASE
+from integration.translation.runner import assert_translation
+
+CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace(
+ CLAUDE_SONNET_4_6_TEST_CASE,
+ scenario="service_tier",
+ litellm_request={**CLAUDE_SONNET_4_6_TEST_CASE.litellm_request, "service_tier": "standard_only"},
+ expected_provider_request={
+ **CLAUDE_SONNET_4_6_TEST_CASE.expected_provider_request,
+ "service_tier": "standard_only",
+ },
+ mock_provider_response={**CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, "id": "msg_011Cg1RG36pCcTLE26htFFqC"},
+ expected_litellm_response={
+ **CLAUDE_SONNET_4_6_TEST_CASE.expected_litellm_response,
+ "id": "msg_011Cg1RG36pCcTLE26htFFqC",
+ },
+)
+
+
+@pytest.mark.parametrize("case", [CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE], ids=lambda case: case.id)
+def test_messages_service_tier_anthropic(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None:
+ assert_translation(case, gateway, provider)
diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py
new file mode 100644
index 00000000000..13bf9dd888f
--- /dev/null
+++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_azure_ai.py
@@ -0,0 +1,29 @@
+from dataclasses import replace
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.provider import SharedProvider
+from integration.translation.case import TranslationTestCase
+from integration.translation.messages.bases.azure_ai import CLAUDE_SONNET_4_6_TEST_CASE
+from integration.translation.runner import assert_translation
+
+CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE: Final = replace(
+ CLAUDE_SONNET_4_6_TEST_CASE,
+ scenario="service_tier",
+ litellm_request={**CLAUDE_SONNET_4_6_TEST_CASE.litellm_request, "service_tier": "standard_only"},
+ expected_provider_request={
+ **CLAUDE_SONNET_4_6_TEST_CASE.expected_provider_request,
+ "service_tier": "standard_only",
+ },
+ mock_provider_response={**CLAUDE_SONNET_4_6_TEST_CASE.mock_provider_response, "id": "msg_011Cg1SnrstCdzZPwmTk2qQa"},
+ expected_litellm_response={
+ **CLAUDE_SONNET_4_6_TEST_CASE.expected_litellm_response,
+ "id": "msg_011Cg1SnrstCdzZPwmTk2qQa",
+ },
+)
+
+
+@pytest.mark.parametrize("case", [CLAUDE_SONNET_4_6_SERVICE_TIER_TEST_CASE], ids=lambda case: case.id)
+def test_messages_service_tier_azure_ai(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None:
+ assert_translation(case, gateway, provider)
diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_bedrock_mantle.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_bedrock_mantle.py
new file mode 100644
index 00000000000..261d34271e9
--- /dev/null
+++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_bedrock_mantle.py
@@ -0,0 +1,31 @@
+from dataclasses import replace
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.provider import SharedProvider
+from integration.translation.case import TranslationTestCase
+from integration.translation.messages.bases.bedrock_mantle import CLAUDE_SONNET_5_TEST_CASE
+from integration.translation.runner import assert_translation
+
+CLAUDE_SONNET_5_SERVICE_TIER_TEST_CASE: Final = replace(
+ CLAUDE_SONNET_5_TEST_CASE,
+ scenario="service_tier",
+ litellm_request={**CLAUDE_SONNET_5_TEST_CASE.litellm_request, "service_tier": "standard_only"},
+ expected_provider_request={**CLAUDE_SONNET_5_TEST_CASE.expected_provider_request, "service_tier": "standard_only"},
+ mock_provider_response={
+ **CLAUDE_SONNET_5_TEST_CASE.mock_provider_response,
+ "id": "msg_bdrk_6oqdfuf56lp47vwldyos5f7l4lyyot337z2cv3h3lbaj3a3ei7vq",
+ },
+ expected_litellm_response={
+ **CLAUDE_SONNET_5_TEST_CASE.expected_litellm_response,
+ "id": "msg_bdrk_6oqdfuf56lp47vwldyos5f7l4lyyot337z2cv3h3lbaj3a3ei7vq",
+ },
+)
+
+
+@pytest.mark.parametrize("case", [CLAUDE_SONNET_5_SERVICE_TIER_TEST_CASE], ids=lambda case: case.id)
+def test_messages_service_tier_bedrock_mantle(
+ case: TranslationTestCase, gateway: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(case, gateway, provider)
diff --git a/tests/integration/translation/messages/service_tier/test_messages_service_tier_openai.py b/tests/integration/translation/messages/service_tier/test_messages_service_tier_openai.py
new file mode 100644
index 00000000000..433c5bc4b2f
--- /dev/null
+++ b/tests/integration/translation/messages/service_tier/test_messages_service_tier_openai.py
@@ -0,0 +1,42 @@
+from dataclasses import replace
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.provider import SharedProvider
+from integration.translation.case import TranslationTestCase
+from integration.translation.messages.bases.openai import GPT_5_4_TEST_CASE
+from integration.translation.runner import assert_translation
+
+GPT_5_4_SERVICE_TIER_TEST_CASE: Final = replace(
+ GPT_5_4_TEST_CASE,
+ scenario="service_tier",
+ litellm_request={**GPT_5_4_TEST_CASE.litellm_request, "service_tier": "auto"},
+ expected_provider_request={**GPT_5_4_TEST_CASE.expected_provider_request, "service_tier": "auto"},
+ mock_provider_response={
+ **GPT_5_4_TEST_CASE.mock_provider_response,
+ "id": "resp_0738fe6d3b086cd1006acef6c23a1487d2860e02c53af945a6",
+ "created_at": 1791948482,
+ "access_programs": None,
+ "completed_at": 1791948483,
+ "output": [
+ {
+ "id": "msg_0738fe6d3b086cd1006acef6c310cc87d283ecc2cc1d6ff29c",
+ "type": "message",
+ "status": "completed",
+ "content": [{"type": "output_text", "annotations": [], "logprobs": [], "text": "Hello."}],
+ "phase": "final_answer",
+ "role": "assistant",
+ }
+ ],
+ },
+ expected_litellm_response={
+ **GPT_5_4_TEST_CASE.expected_litellm_response,
+ "id": "resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDo3NGI3ZDQzZjBmNjQ1MjlkYWY4NWQ4MjRmNjE3ZGYzNDRjZTdhNGI3NjI3YTYzMDMzZGUzMDVmNTE2NTRlZjc2O3Jlc3BvbnNlX2lkOnJlc3BfMDczOGZlNmQzYjA4NmNkMTAwNmFjZWY2YzIzYTE0ODdkMjg2MGUwMmM1M2FmOTQ1YTY=",
+ },
+)
+
+
+@pytest.mark.parametrize("case", [GPT_5_4_SERVICE_TIER_TEST_CASE], ids=lambda case: case.id)
+def test_messages_service_tier_openai(case: TranslationTestCase, gateway: Gateway, provider: SharedProvider) -> None:
+ assert_translation(case, gateway, provider)
diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_service_tier.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_service_tier.py
new file mode 100644
index 00000000000..6472ccd1b2e
--- /dev/null
+++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_service_tier.py
@@ -0,0 +1,28 @@
+from typing import Final
+
+import pytest
+
+from litellm.llms.anthropic.pass_through.messages.utils import AnthropicMessagesRequestUtils
+
+REQUEST: Final[dict[str, object]] = {
+ "max_tokens": 64,
+ "service_tier": "standard_only",
+ "messages": [{"role": "user", "content": "Hello"}],
+}
+
+
+@pytest.mark.parametrize("custom_llm_provider", ["anthropic", "azure_ai", "bedrock_mantle", None])
+def test_messages_forwards_service_tier_on_anthropic_api_routes(custom_llm_provider: str | None) -> None:
+ optional_params = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
+ params=dict(REQUEST), model="claude-sonnet-4-6", drop_params=False, custom_llm_provider=custom_llm_provider
+ )
+ assert dict(optional_params).get("service_tier") == "standard_only"
+
+
+@pytest.mark.parametrize("custom_llm_provider", ["vertex_ai", "bedrock"])
+def test_messages_drops_service_tier_where_the_provider_rejects_it(custom_llm_provider: str) -> None:
+ optional_params = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
+ params=dict(REQUEST), model="claude-sonnet-4-6", drop_params=False, custom_llm_provider=custom_llm_provider
+ )
+ assert "service_tier" not in optional_params
+ assert dict(optional_params).get("max_tokens") == 64
From 83980c9d63fe7a205c66ba45639ac2ea6bca7ebf Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 14:19:25 -0700
Subject: [PATCH 35/47] fix(router): read the first stream chunk before
returning so a failure before any token runs through the retry and fallback
budget (#46493)
* fix(router): read the first stream chunk before returning so a failure before any token runs through the retry and fallback budget
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(router): keep only the raising call inside pytest.raises for the pre-content stream tests
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(router): yield the primed first chunk from a helper generator so the stream iterators keep their complexity
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(bedrock): expect the empty invoke stream to be reported from router.completion now that the first chunk is read there
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test: give router stream mocks a first chunk or an ending so the eager first read neither hangs nor retries
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(router): name the first-chunk helpers after what they do (_read_first_chunk, _FirstChunk, _chunks_with_first)
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(router): drop the string annotation on _FirstChunk.chunk, ModelResponseStream is imported at runtime
---------
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/router.py | 75 ++++-
.../test_router_sync_stream_fallback_wire.py | 14 +-
...r_tool_call_stream_fallback_budget_wire.py | 289 ++++++++++++++++++
.../test_base_invoke_transformation.py | 17 +-
tests/unit/llms/chat/test_converse_handler.py | 17 +-
tests/unit/test_router/test_router.py | 119 +++++++-
6 files changed, 496 insertions(+), 35 deletions(-)
create mode 100644 tests/integration/sdk/test_router_tool_call_stream_fallback_budget_wire.py
diff --git a/litellm/router.py b/litellm/router.py
index ee8a58f21e2..16257775d3f 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -23,6 +23,7 @@ import weakref
from collections import defaultdict
from collections.abc import (
AsyncGenerator,
+ AsyncIterable,
AsyncIterator,
Awaitable,
Callable,
@@ -691,6 +692,67 @@ def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception:
return error.original_exception if error.original_exception is not None else error
+@dataclass(frozen=True, slots=True)
+class _FirstChunk:
+ """The chunk Router._acompletion / _completion already read off a stream; None means the provider closed it empty."""
+
+ chunk: ModelResponseStream | None
+
+
+async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk:
+ """
+ Read the first chunk while Router._acompletion still holds the request, so a
+ stream that fails before any token raises there, inside async_function_with_retries
+ and the fallback chain, instead of surfacing to the caller mid-iteration.
+ """
+ from litellm.exceptions import MidStreamFallbackError
+
+ try:
+ return _FirstChunk(chunk=await stream.__anext__())
+ except StopAsyncIteration:
+ return _FirstChunk(chunk=None)
+ except MidStreamFallbackError as e:
+ raise _mid_stream_retry_trigger(e) from e
+
+
+def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk:
+ """Sync twin of _aread_first_chunk for Router._completion."""
+ from litellm.exceptions import MidStreamFallbackError
+
+ try:
+ return _FirstChunk(chunk=next(stream))
+ except StopIteration:
+ return _FirstChunk(chunk=None)
+ except MidStreamFallbackError as e:
+ raise _mid_stream_retry_trigger(e) from e
+
+
+def _chunks_with_first(model_response: Iterable[object], taken: _FirstChunk | None) -> Iterator[object]:
+ """The stream with the chunk _read_first_chunk already read put back in front (None: the stream as is)."""
+ if taken is None:
+ yield from model_response
+ return
+ if taken.chunk is None:
+ return
+ yield taken.chunk
+ yield from model_response
+
+
+async def _achunks_with_first(
+ model_response: AsyncIterable[object], taken: _FirstChunk | None
+) -> AsyncIterator[object]:
+ """Async twin of _chunks_with_first."""
+ if taken is None:
+ async for item in model_response:
+ yield item
+ return
+ if taken.chunk is None:
+ return
+ yield taken.chunk
+ async for item in model_response:
+ yield item
+
+
def _stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool:
"""
Whether `chunk` should make Router._aanthropic_messages_streaming_iterator
@@ -2791,6 +2853,7 @@ class Router:
model_response=response,
messages=messages,
initial_kwargs=input_kwargs_for_streaming_fallback,
+ first_chunk=_read_first_chunk(response),
)
return response
@@ -3050,6 +3113,7 @@ class Router:
messages: list[dict[str, str]],
initial_kwargs: dict,
deployment_slot: contextlib.AsyncExitStack | None = None,
+ first_chunk: _FirstChunk | None = None,
) -> CustomStreamWrapper:
"""
Helper to iterate over a streaming response.
@@ -3058,6 +3122,8 @@ class Router:
`deployment_slot` holds the deployment's max_parallel_requests semaphore; it is
released when the stream is exhausted, closed, or falls back to another deployment
+
+ `first_chunk` carries the chunk _acompletion already read (None: iterate from the start)
"""
from litellm.exceptions import MidStreamFallbackError
@@ -3102,7 +3168,7 @@ class Router:
async def stream_with_fallbacks():
fallback_response = None # Track for cleanup in finally
try:
- async for item in model_response:
+ async for item in _achunks_with_first(model_response, first_chunk):
yield item
except MidStreamFallbackError as e:
with anyio.CancelScope(shield=True):
@@ -3677,6 +3743,7 @@ class Router:
model_response: CustomStreamWrapper,
messages: list[dict[str, str]],
initial_kwargs: dict,
+ first_chunk: _FirstChunk | None = None,
) -> CustomStreamWrapper:
"""
Sync equivalent of _acompletion_streaming_iterator.
@@ -3712,7 +3779,7 @@ class Router:
def stream_with_fallbacks():
fallback_response = None
try:
- for item in model_response:
+ for item in _chunks_with_first(model_response, first_chunk):
yield item
except MidStreamFallbackError as e:
if fallbacks_disabled_for_request(initial_kwargs) or (
@@ -3959,6 +4026,9 @@ class Router:
and response.make_call is not None
):
await response.fetch_stream()
+ first_chunk: Final = (
+ await _aread_first_chunk(response) if isinstance(response, CustomStreamWrapper) else None
+ )
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
@@ -3975,6 +4045,7 @@ class Router:
messages=messages,
initial_kwargs=input_kwargs_for_streaming_fallback,
deployment_slot=deployment_slot.pop_all(),
+ first_chunk=first_chunk,
)
return response
diff --git a/tests/integration/sdk/test_router_sync_stream_fallback_wire.py b/tests/integration/sdk/test_router_sync_stream_fallback_wire.py
index fc5fb028123..36616eda50e 100644
--- a/tests/integration/sdk/test_router_sync_stream_fallback_wire.py
+++ b/tests/integration/sdk/test_router_sync_stream_fallback_wire.py
@@ -154,7 +154,8 @@ def test_walks_every_configured_fallback(client: str) -> None:
router: Final = _router(wire, ("primary", "fb1", "fb2"), fallbacks=[{"primary": ["fb1", "fb2"]}])
streamed: Final = _stream(client, router)
assert streamed.text == "answered by fb2", streamed
- assert streamed.attempted_fallbacks == 2, streamed
+ # the header a non-streaming request with the same chain carries today
+ assert streamed.attempted_fallbacks == 1, streamed
assert _deployments_hit(wire) == ("primary", "fb1", "fb2")
@@ -192,11 +193,12 @@ def test_router_retries_configured(client: str) -> None:
model_list=_router(wire, ("primary", "backup")).model_list,
fallbacks=_PRIMARY_TO_BACKUP,
num_retries=2,
+ retry_after=0,
disable_cooldowns=True,
)
streamed: Final = _stream(client, router)
assert streamed.text == "answered by the backup", streamed
- assert _deployments_hit(wire) == ("primary", "backup")
+ assert _deployments_hit(wire) == ("primary", "primary", "primary", "backup")
@pytest.mark.parametrize("client", _CLIENTS)
@@ -219,12 +221,12 @@ def test_per_request_fallbacks_none_turns_the_router_list_off(client: str) -> No
@pytest.mark.parametrize("client", _CLIENTS)
-def test_max_fallbacks_caps_the_walk(client: str) -> None:
+def test_max_fallbacks_walks_the_list_like_a_non_streaming_request(client: str) -> None:
with wire_server(_peer(_PRIMARY_AND_FB1_DIE)) as wire:
router: Final = _router(wire, ("primary", "fb1", "fb2"), fallbacks=[{"primary": ["fb1", "fb2"]}], max_fallbacks=1)
- with pytest.raises(litellm.APIConnectionError, match="overloaded"):
- _stream(client, router)
- assert _deployments_hit(wire) == ("primary", "fb1")
+ streamed: Final = _stream(client, router)
+ assert streamed.text == "answered by fb2", streamed
+ assert _deployments_hit(wire) == ("primary", "fb1", "fb2")
def test_called_inside_a_running_loop() -> None:
diff --git a/tests/integration/sdk/test_router_tool_call_stream_fallback_budget_wire.py b/tests/integration/sdk/test_router_tool_call_stream_fallback_budget_wire.py
new file mode 100644
index 00000000000..4b9bc8588a7
--- /dev/null
+++ b/tests/integration/sdk/test_router_tool_call_stream_fallback_budget_wire.py
@@ -0,0 +1,289 @@
+from __future__ import annotations
+
+import asyncio
+import itertools
+import json
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+import pytest
+from integration._support.wire import Reply, Request, Wire, wire_server
+
+import litellm
+from litellm import Router
+
+_MODEL: Final = "gpt-5.6"
+_API_KEY: Final = "synthetic-tool-call-fallback-key"
+_PROMPT: Final = "what is the weather in the city the primary deployment never finishes naming?"
+_TOOLS: Final = (
+ {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+ },
+ },
+)
+_ERROR_FRAME: Final = (
+ b"data: " + json.dumps({"error": {"message": "overloaded", "type": "server_error", "code": 500}}).encode() + b"\n\n"
+)
+_DONE: Final = b"data: [DONE]\n\n"
+_PRIMARY_TO_BACKUP: Final = [{"primary": ["backup"]}]
+_BOTH_WAYS: Final = [{"primary": ["backup"]}, {"backup": ["primary"]}]
+
+pytestmark = pytest.mark.timeout(120)
+
+
+def _chunk(delta: Mapping[str, object], finish_reason: str | None) -> bytes:
+ body: Final = {
+ "id": "chatcmpl-tool-call-fallback-wire",
+ "object": "chat.completion.chunk",
+ "created": 1700000000,
+ "model": _MODEL,
+ "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}],
+ }
+ return b"data: " + json.dumps(body).encode() + b"\n\n"
+
+
+def _tool_call_open(call_id: str) -> bytes:
+ return _chunk(
+ {
+ "role": "assistant",
+ "tool_calls": [
+ {"index": 0, "id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": ""}}
+ ],
+ },
+ None,
+ )
+
+
+def _tool_call_arguments(fragment: str) -> bytes:
+ return _chunk({"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]}, None)
+
+
+_FINISH: Final = _chunk({}, "tool_calls")
+
+
+def _serves_tool_call(call_id: str, city: str) -> Reply:
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(
+ _tool_call_open(call_id),
+ _tool_call_arguments('{"city": "'),
+ _tool_call_arguments(f'{city}"}}'),
+ _FINISH,
+ _DONE,
+ ),
+ )
+
+
+def _dies_after_tool_call_opens(call_id: str) -> Reply:
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(_tool_call_open(call_id), _tool_call_arguments('{"city": "'), _ERROR_FRAME, _DONE),
+ )
+
+
+_DIES_BEFORE_CONTENT: Final = Reply(content_type="text/event-stream", chunks=(_ERROR_FRAME, _DONE))
+_DROPS_BEFORE_CONTENT: Final = Reply(content_type="text/event-stream", chunks=(_DONE,), abort_after=0)
+_REJECTS_BEFORE_STREAMING: Final = Reply(
+ status=500, body=json.dumps({"error": {"message": "overloaded", "type": "server_error", "code": 500}}).encode()
+)
+
+
+def _deployment_of(request: Request) -> str:
+ deployment, _, route = request.target.lstrip("/").partition("/")
+ assert request.method == "POST", request.method
+ assert route == "chat/completions", request.target
+ return deployment
+
+
+def _peer(replies: Mapping[str, Reply]) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ body: Final = json.loads(request.body)
+ assert body.get("stream") is True, body
+ assert [tool["function"]["name"] for tool in body["tools"]] == ["get_weather"], body
+ return replies[_deployment_of(request)]
+
+ return respond
+
+
+def _always(reply: Reply) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ _deployment_of(request)
+ return reply
+
+ return respond
+
+
+def _fails_then_serves(failures: int, served: Reply) -> Callable[[Request], Reply]:
+ attempts: Final = itertools.count(1)
+
+ def respond(request: Request) -> Reply:
+ _deployment_of(request)
+ return _DIES_BEFORE_CONTENT if next(attempts) <= failures else served
+
+ return respond
+
+
+def _deployments_hit(wire: Wire) -> tuple[str, ...]:
+ return tuple(_deployment_of(request) for request in wire.drain())
+
+
+def _router(wire: Wire, deployments: tuple[str, ...], **settings: object) -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": name,
+ "litellm_params": {"model": f"openai/{_MODEL}", "api_base": f"{wire.url}/{name}", "api_key": _API_KEY},
+ }
+ for name in deployments
+ ],
+ disable_cooldowns=True,
+ **settings,
+ )
+
+
+@dataclass(frozen=True, slots=True)
+class _ToolCallStream:
+ call_ids: tuple[str, ...]
+ arguments: str
+ finish_reasons: tuple[str, ...]
+
+
+def _tool_call_deltas(chunk: object) -> tuple[object, ...]:
+ choices: Final = getattr(chunk, "choices", None) or ()
+ return tuple(delta for choice in choices for delta in (choice.delta.tool_calls or ()))
+
+
+def _collect(chunks: tuple[object, ...]) -> _ToolCallStream:
+ deltas: Final = tuple(delta for chunk in chunks for delta in _tool_call_deltas(chunk))
+ finish_reasons: Final = tuple(
+ choice.finish_reason
+ for chunk in chunks
+ for choice in (getattr(chunk, "choices", None) or ())
+ if choice.finish_reason
+ )
+ return _ToolCallStream(
+ call_ids=tuple(delta.id for delta in deltas if delta.id),
+ arguments="".join(delta.function.arguments or "" for delta in deltas),
+ finish_reasons=finish_reasons,
+ )
+
+
+def _request(**extra: object) -> dict[str, object]:
+ return {
+ "model": "primary",
+ "messages": [{"role": "user", "content": _PROMPT}],
+ "tools": list(_TOOLS),
+ "stream": True,
+ **extra,
+ }
+
+
+def _stream_sync(router: Router, **extra: object) -> _ToolCallStream:
+ return _collect(tuple(router.completion(**_request(**extra))))
+
+
+async def _stream_async(router: Router, **extra: object) -> _ToolCallStream:
+ stream: Final = await router.acompletion(**_request(**extra))
+ return _collect(tuple([chunk async for chunk in stream]))
+
+
+def _stream(client: str, router: Router, **extra: object) -> _ToolCallStream:
+ if client == "async":
+ return asyncio.run(_stream_async(router, **extra))
+ return _stream_sync(router, **extra)
+
+
+_CLIENTS: Final = ("sync", "async")
+_BACKUP_ANSWERS: Final = _ToolCallStream(
+ call_ids=("call_backup",), arguments='{"city": "Lisbon"}', finish_reasons=("tool_calls",)
+)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_primary_dies_before_the_tool_call_opens(client: str) -> None:
+ replies: Final = {"primary": _DIES_BEFORE_CONTENT, "backup": _serves_tool_call("call_backup", "Lisbon")}
+ with wire_server(_peer(replies)) as wire:
+ router: Final = _router(wire, ("primary", "backup"), num_retries=0, fallbacks=_PRIMARY_TO_BACKUP)
+ assert _stream(client, router) == _BACKUP_ANSWERS
+ assert _deployments_hit(wire) == ("primary", "backup")
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_primary_drops_the_connection_before_the_tool_call_opens(client: str) -> None:
+ replies: Final = {"primary": _DROPS_BEFORE_CONTENT, "backup": _serves_tool_call("call_backup", "Lisbon")}
+ with wire_server(_peer(replies)) as wire:
+ router: Final = _router(wire, ("primary", "backup"), num_retries=0, fallbacks=_PRIMARY_TO_BACKUP)
+ assert _stream(client, router) == _BACKUP_ANSWERS
+ assert _deployments_hit(wire) == ("primary", "backup")
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_primary_dies_after_tool_call_deltas_streamed(client: str) -> None:
+ replies: Final = {
+ "primary": _dies_after_tool_call_opens("call_primary"),
+ "backup": _serves_tool_call("call_backup", "Lisbon"),
+ }
+ with wire_server(_peer(replies)) as wire:
+ router: Final = _router(wire, ("primary", "backup"), num_retries=2, fallbacks=_PRIMARY_TO_BACKUP)
+ with pytest.raises(litellm.APIConnectionError, match="overloaded"):
+ _stream(client, router)
+ assert _deployments_hit(wire) == ("primary",)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_disable_fallbacks_holds_before_the_tool_call_opens(client: str) -> None:
+ replies: Final = {"primary": _DIES_BEFORE_CONTENT, "backup": _serves_tool_call("call_backup", "Lisbon")}
+ with wire_server(_peer(replies)) as wire:
+ router: Final = _router(wire, ("primary", "backup"), num_retries=0, fallbacks=_PRIMARY_TO_BACKUP)
+ with pytest.raises(litellm.APIConnectionError, match="overloaded"):
+ _stream(client, router, disable_fallbacks=True)
+ assert _deployments_hit(wire) == ("primary",)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_primary_dies_before_the_tool_call_opens_is_retried_within_the_group(client: str) -> None:
+ with wire_server(_fails_then_serves(1, _serves_tool_call("call_primary", "Lisbon"))) as wire:
+ router: Final = _router(wire, ("primary",), num_retries=1, retry_after=0)
+ assert _stream(client, router) == _ToolCallStream(
+ call_ids=("call_primary",), arguments='{"city": "Lisbon"}', finish_reasons=("tool_calls",)
+ )
+ assert _deployments_hit(wire) == ("primary", "primary")
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_request_num_retries_overrides_the_router_budget(client: str) -> None:
+ with wire_server(_always(_DIES_BEFORE_CONTENT)) as wire:
+ router: Final = _router(wire, ("primary",), num_retries=3, retry_after=0)
+ with pytest.raises(litellm.APIConnectionError, match="overloaded"):
+ _stream(client, router, num_retries=1)
+ assert _deployments_hit(wire) == ("primary", "primary")
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+@pytest.mark.parametrize("failure", (_DIES_BEFORE_CONTENT, _DROPS_BEFORE_CONTENT), ids=("error-frame", "dropped"))
+def test_every_stream_dies_before_content_spends_exactly_the_budget(client: str, failure: Reply) -> None:
+ num_retries: Final = 2
+ with wire_server(_always(failure)) as wire:
+ router: Final = _router(
+ wire, ("primary", "backup"), num_retries=num_retries, fallbacks=_BOTH_WAYS, retry_after=0
+ )
+ with pytest.raises(litellm.APIConnectionError):
+ _stream(client, router)
+ assert _deployments_hit(wire) == ("primary",) * (num_retries + 1) + ("backup",) * (num_retries + 1)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_every_request_rejected_before_streaming_spends_a_bounded_budget(client: str) -> None:
+ num_retries: Final = 1
+ with wire_server(_always(_REJECTS_BEFORE_STREAMING)) as wire:
+ router: Final = _router(
+ wire, ("primary", "backup"), num_retries=num_retries, fallbacks=_BOTH_WAYS, retry_after=0
+ )
+ with pytest.raises(litellm.InternalServerError, match="overloaded"):
+ _stream(client, router)
+ hit: Final = _deployments_hit(wire)
+ assert hit == ("primary",) * (num_retries + 1) + ("backup",) * (num_retries + 1), hit
diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
index d9b01c4e75a..1890c68b30d 100644
--- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
+++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py
@@ -343,15 +343,18 @@ def test_router_deployment_stream_chunk_size_reaches_iter_bytes(
client: Final = HTTPHandler()
client.post = MagicMock(return_value=mock_response)
router: Final = litellm.Router(
- model_list=[{"model_name": "invoke-chunked", "litellm_params": {**INVOKE_DEPLOYMENT, **deployment_extras}}]
+ model_list=[{"model_name": "invoke-chunked", "litellm_params": {**INVOKE_DEPLOYMENT, **deployment_extras}}],
+ num_retries=0,
)
- router.completion(
- model="invoke-chunked",
- messages=[{"role": "user", "content": "hi"}],
- stream=True,
- client=client,
- )
+ # the empty body is a stream that ends before its first chunk, which the Router now reports from the call
+ with pytest.raises(litellm.BadGatewayError):
+ router.completion(
+ model="invoke-chunked",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=True,
+ client=client,
+ )
mock_response.iter_bytes.assert_called_once_with(chunk_size=expected_chunk_size)
data: Final = client.post.call_args.kwargs["data"]
diff --git a/tests/unit/llms/chat/test_converse_handler.py b/tests/unit/llms/chat/test_converse_handler.py
index 6279b74bc0d..d0360b07957 100644
--- a/tests/unit/llms/chat/test_converse_handler.py
+++ b/tests/unit/llms/chat/test_converse_handler.py
@@ -417,15 +417,18 @@ def test_router_deployment_stream_chunk_size_reaches_iter_bytes(
client: Final = HTTPHandler()
client.post = MagicMock(return_value=mock_response)
router: Final = litellm.Router(
- model_list=[{"model_name": "converse-chunked", "litellm_params": {**CONVERSE_DEPLOYMENT, **deployment_extras}}]
+ model_list=[{"model_name": "converse-chunked", "litellm_params": {**CONVERSE_DEPLOYMENT, **deployment_extras}}],
+ num_retries=0,
)
- router.completion(
- model="converse-chunked",
- messages=[{"role": "user", "content": "hi"}],
- stream=True,
- client=client,
- )
+ # the empty body is a stream that ends before its first chunk, which the Router now reports from the call
+ with pytest.raises(litellm.BadGatewayError):
+ router.completion(
+ model="converse-chunked",
+ messages=[{"role": "user", "content": "hi"}],
+ stream=True,
+ client=client,
+ )
mock_response.iter_bytes.assert_called_once_with(chunk_size=expected_chunk_size)
data: Final = client.post.call_args.kwargs["data"]
diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py
index 64f3f88cd89..5faea8d6c1d 100644
--- a/tests/unit/test_router/test_router.py
+++ b/tests/unit/test_router/test_router.py
@@ -1638,7 +1638,8 @@ def test_arouter_responses_api_bridge():
)
## CONFIRM BRIDGE IS CALLED
- with patch.object(litellm, "responses", return_value=AsyncMock()) as mock_responses:
+ # an empty stream: the Router reads the first chunk before returning, and a stream of mocks never produces one
+ with patch.object(litellm, "responses", return_value=iter([])) as mock_responses:
result = router.completion(
model="[IP-approved] o3-pro",
messages=[{"role": "user", "content": "Hello, world!"}],
@@ -4020,11 +4021,10 @@ async def test_acompletion_mid_stream_fallback_honors_a_per_request_fallbacks_no
return _stream_for(**kwargs)
with patch("litellm.acompletion", side_effect=fake_acompletion) as provider_calls:
- response: Final = await router.acompletion(
- model="primary", messages=[{"role": "user", "content": "hi"}], stream=True, fallbacks=None
- )
with pytest.raises(litellm.InternalServerError, match="provider 500 from openai/primary-model"):
- [chunk async for chunk in response]
+ await router.acompletion(
+ model="primary", messages=[{"role": "user", "content": "hi"}], stream=True, fallbacks=None
+ )
assert _groups_called(provider_calls) == ["primary"]
@@ -16837,11 +16837,10 @@ async def test_acompletion_disable_fallbacks_reaches_the_mid_stream_hop():
return _MidStreamOptOutChatStream(_mid_stream_opt_out_trigger(primary_error), model=kwargs["model"])
with patch("litellm.acompletion", side_effect=primary_stream) as provider_calls:
- response = await router.acompletion(
- model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True, disable_fallbacks=True
- )
with pytest.raises(litellm.InternalServerError) as raised:
- [chunk async for chunk in response]
+ await router.acompletion(
+ model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True, disable_fallbacks=True
+ )
assert raised.value is primary_error
assert [call.kwargs["metadata"]["model_group"] for call in provider_calls.call_args_list] == ["primary"]
@@ -16856,11 +16855,10 @@ def test_completion_disable_fallbacks_reaches_the_mid_stream_hop():
return _MidStreamOptOutChatStream(_mid_stream_opt_out_trigger(primary_error), model=kwargs["model"])
with patch("litellm.completion", side_effect=primary_stream) as provider_calls:
- response = router.completion(
- model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True, disable_fallbacks=True
- )
with pytest.raises(litellm.InternalServerError) as raised:
- list(response)
+ router.completion(
+ model="primary", messages=[{"role": "user", "content": "Hi"}], stream=True, disable_fallbacks=True
+ )
assert raised.value is primary_error
assert [call.kwargs["metadata"]["model_group"] for call in provider_calls.call_args_list] == ["primary"]
@@ -26135,3 +26133,98 @@ async def test_batch_completion_fastest_response_streams_mocked_chunks() -> None
assert text == "stream answer"
assert chunks[-1].choices[0].finish_reason == "stop"
+
+
+class _FirstPullStream:
+ """A stream whose first pull yields `first`, raises it, or ends (first=None)."""
+
+ def __init__(self, first: object) -> None:
+ self._first = first
+
+ def _pull(self, stop: type[Exception]) -> object:
+ if self._first is None:
+ raise stop
+ if isinstance(self._first, Exception):
+ raise self._first
+ return self._first
+
+ def __next__(self) -> object:
+ return self._pull(StopIteration)
+
+ async def __anext__(self) -> object:
+ return self._pull(StopAsyncIteration)
+
+
+def _first_chunk_cases():
+ from litellm.exceptions import MidStreamFallbackError
+ from litellm.router import _FirstChunk
+ from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+ chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="hi"))])
+ provider_error = litellm.InternalServerError(message="primary died at stream start", llm_provider="openai", model="m")
+ wrapped = MidStreamFallbackError(
+ message=str(provider_error),
+ model="m",
+ llm_provider="openai",
+ original_exception=provider_error,
+ is_pre_first_chunk=True,
+ )
+ return [
+ pytest.param(chunk, _FirstChunk(chunk=chunk), None, id="first-chunk"),
+ pytest.param(None, _FirstChunk(chunk=None), None, id="empty-stream"),
+ pytest.param(wrapped, None, provider_error, id="pre-content-failure"),
+ ]
+
+
+@pytest.mark.parametrize("first, expected, raised", _first_chunk_cases())
+def test_read_first_chunk_direct_call(first, expected, raised):
+ from litellm.router import _read_first_chunk
+
+ if raised is None:
+ assert _read_first_chunk(_FirstPullStream(first)) == expected
+ return
+ with pytest.raises(type(raised)) as caught:
+ _read_first_chunk(_FirstPullStream(first))
+ assert caught.value is raised
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("first, expected, raised", _first_chunk_cases())
+async def test_aread_first_chunk_direct_call(first, expected, raised):
+ from litellm.router import _aread_first_chunk
+
+ if raised is None:
+ assert await _aread_first_chunk(_FirstPullStream(first)) == expected
+ return
+ with pytest.raises(type(raised)) as caught:
+ await _aread_first_chunk(_FirstPullStream(first))
+ assert caught.value is raised
+
+
+def _chunks_with_first_cases():
+ from litellm.router import _FirstChunk
+
+ return [
+ pytest.param(None, ["a", "b"], ["a", "b"], id="nothing-read-stream-as-is"),
+ pytest.param(_FirstChunk(chunk=None), ["a", "b"], [], id="empty-stream-not-touched"),
+ pytest.param(_FirstChunk(chunk="first"), ["a", "b"], ["first", "a", "b"], id="first-chunk-put-back"),
+ ]
+
+
+@pytest.mark.parametrize("taken, rest, expected", _chunks_with_first_cases())
+def test_chunks_with_first_direct_call(taken, rest, expected):
+ from litellm.router import _chunks_with_first
+
+ assert list(_chunks_with_first(iter(rest), taken)) == expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("taken, rest, expected", _chunks_with_first_cases())
+async def test_achunks_with_first_direct_call(taken, rest, expected):
+ from litellm.router import _achunks_with_first
+
+ async def stream():
+ for item in rest:
+ yield item
+
+ assert [item async for item in _achunks_with_first(stream(), taken)] == expected
From f10134728deb227277cf12ba59fec8a79b880b19 Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 14:29:54 -0700
Subject: [PATCH 36/47] fix(caching): report Redis circuit-breaker metrics
through the Prometheus services logger (#46484)
* fix(caching): collect the Redis circuit-breaker state gauge as one series per state across worker processes
The gauge declares multiprocess_mode=livesum, so PROMETHEUS_MULTIPROC_DIR apps get one series per state summed over
the live workers instead of one series per worker pid
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(caching): register Redis circuit-breaker metrics only when LiteLLM prometheus logging is on
The breaker collectors join the default registry once the prometheus or prometheus_system callback is configured or
a Prometheus logger starts, with earlier breaker events tallied and seeded in, and
litellm.disable_redis_circuit_breaker_metrics turns them off for apps that export their own metrics
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(caching): drop the disable_redis_circuit_breaker_metrics opt-out
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(caching): keep the Redis breaker metrics gate to a per-state count and a publish hook
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(caching): pin the Redis breaker metrics gate against an injected registry
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(caching): publish the breaker collectors before reading them in the existing metrics test
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(caching): report Redis circuit-breaker metrics through the Prometheus services logger
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(integrations): build the Redis breaker collectors with typed prometheus_client constructors
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/_service_logger.py | 5 +
litellm/caching/redis_cache.py | 96 ++++++-------------
litellm/integrations/prometheus_services.py | 59 +++++++++++-
.../_support/redis_breaker_worker.py | 18 +++-
..._redis_breaker_metrics_multiprocess_sdk.py | 29 +++++-
tests/unit/caching/test_redis_cache.py | 33 +++++--
...test_prometheus_metric_name_consistency.py | 12 +--
7 files changed, 162 insertions(+), 90 deletions(-)
diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py
index 6385aa96272..060b4fb878e 100644
--- a/litellm/_service_logger.py
+++ b/litellm/_service_logger.py
@@ -152,6 +152,11 @@ class ServiceLogging(CustomLogger):
except Exception as e:
verbose_logger.exception("Error dispatching service event - %s", e)
+ def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None:
+ prometheus_logger: Final[PrometheusServicesLogger | None] = getattr(self, "prometheusServicesLogger", None)
+ if prometheus_logger is not None:
+ prometheus_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class)
+
def service_success_hook(
self,
service: ServiceTypes,
diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py
index 0e34709d128..b3715333010 100644
--- a/litellm/caching/redis_cache.py
+++ b/litellm/caching/redis_cache.py
@@ -18,7 +18,9 @@ import json
import logging
import threading
import time
-from collections.abc import Awaitable, Callable, Iterator, Sequence
+import weakref
+from collections import Counter
+from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import timedelta
@@ -51,12 +53,12 @@ from .base_cache import BaseCache
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
- from prometheus_client import Counter as _PromCounter
- from prometheus_client import Gauge as _PromGauge
from redis.asyncio import Redis, RedisCluster
from redis.asyncio.client import Pipeline
from redis.asyncio.cluster import ClusterPipeline
+ from litellm._service_logger import ServiceLogging
+
pipeline = Pipeline
cluster_pipeline = ClusterPipeline
async_redis_client = Redis
@@ -243,7 +245,9 @@ class RedisCircuitBreaker:
recovery_timeout: int,
enabled: bool = True,
timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
+ service_logger: "ServiceLogging | None" = None,
) -> None:
+ self._service_logger = service_logger
self.failure_threshold = failure_threshold
self.recovery_timeout = recovery_timeout
self.enabled = enabled
@@ -255,7 +259,12 @@ class RedisCircuitBreaker:
self._opened_at: float | None = None
self._state = self.CLOSED
self._generation = 0
- _breaker_metrics().record_state_change(None, self._state)
+ _LIVE_BREAKERS.add(self)
+ self._report(transitions_to=None, failure_class=None)
+
+ @property
+ def state(self) -> str:
+ return self._state
@property
def generation(self) -> int:
@@ -301,7 +310,7 @@ class RedisCircuitBreaker:
self._timeout_count = 0
self._timeout_streak_started_at = None
self._opened_at = now
- _breaker_metrics().record_failure("timeout" if is_timeout else "connectivity")
+ self._report(transitions_to=None, failure_class="timeout" if is_timeout else "connectivity")
if self._should_open(now):
if self._state != self.OPEN:
verbose_logger.warning(
@@ -327,10 +336,13 @@ class RedisCircuitBreaker:
def _set_state(self, state: str) -> None:
if state == self._state:
return
- _breaker_metrics().record_transition(state)
- _breaker_metrics().record_state_change(self._state, state)
self._state = state
self._generation += 1
+ self._report(transitions_to=state, failure_class=None)
+
+ def _report(self, transitions_to: str | None, failure_class: str | None) -> None:
+ if self._service_logger is not None:
+ self._service_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class)
_RedisCallResult = TypeVar("_RedisCallResult")
@@ -430,60 +442,6 @@ def is_redis_timeout_failure(exc: BaseException) -> bool:
return any(isinstance(link, timeout_types) for link in _explicit_causes(exc))
-class _BreakerMetrics:
- """Prometheus metrics for the Redis circuit breaker; no-ops when the client is absent.
-
- Registered lazily on the default registry (which /metrics serves) via the module-level
- ``_breaker_metrics`` singleton so repeated RedisCache construction never re-registers.
- """
-
- def __init__(self) -> None:
- self._state_gauge: _PromGauge | None = None
- self._transitions: _PromCounter | None = None
- self._failures: _PromCounter | None = None
- try:
- from prometheus_client import Counter as PromCounter
- from prometheus_client import Gauge
- except ImportError:
- return
- self._state_gauge = Gauge(
- "litellm_redis_circuit_breaker_state",
- "Number of Redis circuit breakers currently in each state",
- labelnames=("state",),
- multiprocess_mode="livesum",
- )
- self._transitions = PromCounter(
- "litellm_redis_circuit_breaker_transitions",
- "Redis circuit breaker state transitions",
- labelnames=("state",),
- )
- self._failures = PromCounter(
- "litellm_redis_circuit_breaker_failures",
- "Redis health failures counted by the circuit breaker",
- labelnames=("failure_class",),
- )
-
- def record_state_change(self, old_state: str | None, new_state: str) -> None:
- if self._state_gauge is None:
- return
- if old_state is not None:
- self._state_gauge.labels(old_state).dec()
- self._state_gauge.labels(new_state).inc()
-
- def record_transition(self, state: str) -> None:
- if self._transitions is not None:
- self._transitions.labels(state).inc()
-
- def record_failure(self, failure_class: str) -> None:
- if self._failures is not None:
- self._failures.labels(failure_class).inc()
-
-
-@functools.lru_cache(maxsize=1)
-def _breaker_metrics() -> _BreakerMetrics:
- return _BreakerMetrics()
-
-
def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None:
"""Record a Redis failure that the calling method is about to swallow.
@@ -499,6 +457,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
+_LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet()
+
+
+def redis_breaker_state_counts() -> Mapping[str, int]:
+ """How many live Redis circuit breakers sit in each state, for the Prometheus services logger."""
+ return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS)))
+
+
class RedisCircuitBreakerOpenError(Exception):
pass
@@ -685,12 +651,9 @@ class RedisCache(BaseCache):
redis_kwargs["socket_timeout"] = socket_timeout
### HEALTH MONITORING OBJECT ###
- if kwargs.get("service_logger_obj", None) is not None and isinstance(
- kwargs["service_logger_obj"], ServiceLogging
- ):
+ self.service_logger_obj: ServiceLogging = ServiceLogging()
+ if isinstance(kwargs.get("service_logger_obj"), ServiceLogging):
self.service_logger_obj = kwargs.pop("service_logger_obj")
- else:
- self.service_logger_obj = ServiceLogging()
redis_kwargs.update(kwargs)
self.redis_client = get_redis_client(**redis_kwargs)
@@ -717,6 +680,7 @@ class RedisCache(BaseCache):
failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
enabled=REDIS_CIRCUIT_BREAKER_ENABLED,
+ service_logger=self.service_logger_obj,
)
self._setup_health_pings()
diff --git a/litellm/integrations/prometheus_services.py b/litellm/integrations/prometheus_services.py
index d0db4888c50..42a3ffac855 100644
--- a/litellm/integrations/prometheus_services.py
+++ b/litellm/integrations/prometheus_services.py
@@ -3,7 +3,9 @@
# On success + failure, log events to Prometheus for litellm / adjacent services (litellm, redis, postgres, llm api providers)
-from typing import Final
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import TYPE_CHECKING, Final
import litellm
from litellm._logging import print_verbose, verbose_logger
@@ -17,6 +19,42 @@ from litellm.types.services import (
FAILED_REQUESTS_LABELS: Final = ["error_class", "function_name"]
+if TYPE_CHECKING:
+ from prometheus_client import Counter as PrometheusCounter
+ from prometheus_client import Gauge
+
+
+@dataclass(frozen=True, slots=True)
+class _RedisBreakerCollectors:
+ state: "Gauge"
+ transitions: "PrometheusCounter"
+ failures: "PrometheusCounter"
+
+
+@lru_cache(maxsize=1)
+def _redis_breaker_collectors() -> _RedisBreakerCollectors:
+ from prometheus_client import Counter as PrometheusCounter
+ from prometheus_client import Gauge
+
+ return _RedisBreakerCollectors(
+ state=Gauge(
+ "litellm_redis_circuit_breaker_state",
+ "Number of Redis circuit breakers currently in each state",
+ labelnames=("state",),
+ multiprocess_mode="livesum",
+ ),
+ transitions=PrometheusCounter(
+ "litellm_redis_circuit_breaker_transitions",
+ "Redis circuit breaker state transitions",
+ labelnames=("state",),
+ ),
+ failures=PrometheusCounter(
+ "litellm_redis_circuit_breaker_failures",
+ "Redis health failures counted by the circuit breaker",
+ labelnames=("failure_class",),
+ ),
+ )
+
class PrometheusServicesLogger:
# Class variables or attributes
@@ -41,6 +79,11 @@ class PrometheusServicesLogger:
self.Counter = Counter
self.Gauge = Gauge
self.REGISTRY = REGISTRY
+ breaker_collectors: Final = _redis_breaker_collectors()
+ self.redis_breaker_state = breaker_collectors.state
+ self.redis_breaker_transitions = breaker_collectors.transitions
+ self.redis_breaker_failures = breaker_collectors.failures
+ self._refresh_redis_breaker_gauge()
verbose_logger.debug("in init prometheus services metrics")
@@ -98,6 +141,20 @@ class PrometheusServicesLogger:
return DEFAULT_METRICS
return metrics
+ def _refresh_redis_breaker_gauge(self) -> None:
+ from litellm.caching.redis_cache import RedisCircuitBreaker, redis_breaker_state_counts
+
+ counts: Final = redis_breaker_state_counts()
+ for state in (RedisCircuitBreaker.CLOSED, RedisCircuitBreaker.OPEN, RedisCircuitBreaker.HALF_OPEN):
+ self.redis_breaker_state.labels(state).set(counts.get(state, 0))
+
+ def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None:
+ self._refresh_redis_breaker_gauge()
+ if transitions_to is not None:
+ self.redis_breaker_transitions.labels(transitions_to).inc()
+ if failure_class is not None:
+ self.redis_breaker_failures.labels(failure_class).inc()
+
def is_metric_registered(self, metric_name) -> bool:
# Use _names_to_collectors (O(1)) instead of REGISTRY.collect() (O(n)) to avoid
# perf regression when a new Router is created per request (e.g. router_settings in DB).
diff --git a/tests/integration/_support/redis_breaker_worker.py b/tests/integration/_support/redis_breaker_worker.py
index f551baacf24..02d9c0ea2d5 100644
--- a/tests/integration/_support/redis_breaker_worker.py
+++ b/tests/integration/_support/redis_breaker_worker.py
@@ -11,10 +11,12 @@ import sys
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from pathlib import Path
-from typing import Final
+from typing import Final, Literal, TypeAlias
from pydantic import TypeAdapter
+Scenario: TypeAlias = Literal["callbacks_before_cache", "callbacks_after_cache", "no_callbacks"]
+
STATE_GAUGE: Final = "litellm_redis_circuit_breaker_state"
WORKER_TIMEOUT_SECONDS: Final = 120
@@ -29,9 +31,11 @@ class WorkerReport:
_REPORT: Final = TypeAdapter(WorkerReport)
-def run_breaker_worker(api_base: str, prom_dir: Path, environment: Mapping[str, str] = os.environ) -> WorkerReport:
+def run_breaker_worker(
+ scenario: Scenario, api_base: str, prom_dir: Path, environment: Mapping[str, str] = os.environ
+) -> WorkerReport:
completed: Final = subprocess.run(
- (sys.executable, "-I", __file__, api_base, environment["REDIS_HOST"], environment["REDIS_PORT"]),
+ (sys.executable, "-I", __file__, scenario, api_base, environment["REDIS_HOST"], environment["REDIS_PORT"]),
env={**environment, "PROMETHEUS_MULTIPROC_DIR": str(prom_dir)},
capture_output=True,
text=True,
@@ -56,11 +60,15 @@ def _series(families: Iterator[object]) -> Iterator[tuple[tuple[tuple[str, str],
yield tuple(sorted(sample.labels.items())), float(sample.value)
-def main(api_base: str, redis_host: str, redis_port: str) -> None:
+def main(scenario: str, api_base: str, redis_host: str, redis_port: str) -> None:
import litellm
from litellm.caching.caching import Cache
+ if scenario == "callbacks_before_cache":
+ litellm.service_callback = ["prometheus_system"] # test-quality-ok: this worker process exits right after the report
litellm.cache = Cache(type="redis", host=redis_host, port=redis_port) # test-quality-ok: throwaway worker process
+ if scenario == "callbacks_after_cache":
+ litellm.service_callback = ["prometheus_system"] # test-quality-ok: throwaway worker process
response: Final = litellm.completion(
model="openai/gpt-5.6",
api_base=api_base,
@@ -76,4 +84,4 @@ def main(api_base: str, redis_host: str, redis_port: str) -> None:
if __name__ == "__main__":
- main(*sys.argv[1:4])
+ main(*sys.argv[1:5])
diff --git a/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py b/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py
index d1582dd59b9..bd02f2a818b 100644
--- a/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py
+++ b/tests/integration/sdk/test_redis_breaker_metrics_multiprocess_sdk.py
@@ -1,5 +1,6 @@
"""An SDK app under a prefork worker pool (Celery, gunicorn) with ``PROMETHEUS_MULTIPROC_DIR`` set: the Redis
-circuit-breaker state gauge must collect as one series per state across every worker, never one per worker pid."""
+circuit-breaker state gauge must collect as one series per state across every worker, and it must stay off
+the registry entirely for an app that never turned on LiteLLM's own Prometheus logging."""
from __future__ import annotations
@@ -24,6 +25,12 @@ _CHAT_RESPONSE: Final = json.dumps(
}
).encode()
_CLOSED: Final = (("state", "closed"),)
+_HALF_OPEN: Final = (("state", "half_open"),)
+_OPEN: Final = (("state", "open"),)
+
+
+def _all_closed(breakers: float) -> tuple[tuple[tuple[tuple[str, str], ...], float], ...]:
+ return ((_CLOSED, breakers), (_HALF_OPEN, 0.0), (_OPEN, 0.0))
def _reply(request: Request) -> Reply:
@@ -52,6 +59,22 @@ def _worked(report: WorkerReport) -> WorkerReport:
def test_two_worker_processes_collect_as_one_state_series_per_state(tmp_path: Path) -> None:
with wire_server(_reply) as wire:
- reports: Final = tuple(_worked(run_breaker_worker(wire.url, tmp_path)) for _ in range(2))
+ reports: Final = tuple(
+ _worked(run_breaker_worker("callbacks_before_cache", wire.url, tmp_path)) for _ in range(2)
+ )
assert len({report.pid for report in reports}) == 2, reports
- assert _collected_state_series(tmp_path) == ((_CLOSED, 2.0),), reports
+ assert _collected_state_series(tmp_path) == _all_closed(2.0), reports
+
+
+def test_prometheus_logging_turned_on_after_the_cache_still_reports_the_closed_breaker(tmp_path: Path) -> None:
+ with wire_server(_reply) as wire:
+ report: Final = _worked(run_breaker_worker("callbacks_after_cache", wire.url, tmp_path))
+ assert report.state_series == _all_closed(1.0), report
+ assert _collected_state_series(tmp_path) == _all_closed(1.0), report
+
+
+def test_app_without_litellm_prometheus_logging_gets_no_breaker_series(tmp_path: Path) -> None:
+ with wire_server(_reply) as wire:
+ report: Final = _worked(run_breaker_worker("no_callbacks", wire.url, tmp_path))
+ assert report.state_series == (), report
+ assert _collected_state_series(tmp_path) == (), report
diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py
index 0a041694d52..7cffa62d423 100644
--- a/tests/unit/caching/test_redis_cache.py
+++ b/tests/unit/caching/test_redis_cache.py
@@ -971,18 +971,17 @@ async def test_hard_failure_resets_timeout_streak_so_a_later_burst_must_earn_its
@pytest.mark.asyncio
-async def test_breaker_metrics_track_state_and_failure_class():
+async def test_breaker_metrics_track_state_and_failure_class(monkeypatch):
"""Breaker accounting must be observable: failure class, transitions, and state."""
from prometheus_client import REGISTRY
from redis.exceptions import ConnectionError as RedisConnectionError
from redis.exceptions import TimeoutError as RedisTimeoutError
- from litellm.caching.redis_cache import RedisCircuitBreaker, _breaker_metrics, is_redis_timeout_failure
+ from litellm._service_logger import ServiceLogging
+ from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure
- metrics: Final = _breaker_metrics()
- for collector in (metrics._state_gauge, metrics._transitions, metrics._failures):
- if collector is not None and collector not in REGISTRY._collector_to_names:
- REGISTRY.register(collector)
+ monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"])
+ service_logger: Final = ServiceLogging()
def sample(name, labels=None):
return REGISTRY.get_sample_value(name, labels) or 0.0
@@ -993,7 +992,9 @@ async def test_breaker_metrics_track_state_and_failure_class():
open_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "open"})
closed_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "closed"})
- breaker = RedisCircuitBreaker(failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0)
+ breaker = RedisCircuitBreaker(
+ failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0, service_logger=service_logger
+ )
breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("t")))
breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused")))
breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused")))
@@ -1011,6 +1012,24 @@ async def test_breaker_metrics_track_state_and_failure_class():
assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + 1
+def test_prometheus_services_logger_starts_from_the_breakers_already_built(monkeypatch):
+ """A breaker built before prometheus_system is on shows up in the state gauge once the services logger starts."""
+ from prometheus_client import REGISTRY
+
+ from litellm._service_logger import ServiceLogging
+ from litellm.caching.redis_cache import RedisCircuitBreaker, redis_breaker_state_counts
+
+ monkeypatch.setattr(litellm, "service_callback", [])
+ early: Final = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, service_logger=ServiceLogging())
+ assert early.is_open() is False
+
+ monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"])
+ ServiceLogging()
+ assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == (
+ redis_breaker_state_counts()["closed"]
+ )
+
+
def test_sync_guard_counts_a_timeout_as_a_timeout():
from redis.exceptions import TimeoutError as RedisTimeoutError
diff --git a/tests/unit/integrations/test_prometheus_metric_name_consistency.py b/tests/unit/integrations/test_prometheus_metric_name_consistency.py
index d648afcd087..4db34274e82 100644
--- a/tests/unit/integrations/test_prometheus_metric_name_consistency.py
+++ b/tests/unit/integrations/test_prometheus_metric_name_consistency.py
@@ -21,7 +21,6 @@ from prometheus_client import REGISTRY, Gauge
from prometheus_client.registry import Collector
import litellm
-from litellm.caching.redis_cache import _breaker_metrics
from litellm.integrations.prometheus import PrometheusLogger
from litellm.integrations.prometheus_services import PrometheusServicesLogger
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import SpendLogCleanupMetrics
@@ -62,10 +61,7 @@ def _lazy_owner_collectors() -> tuple[Collector, ...]:
assert SpendLogCleanupMetrics.runs is not None
in_flight: Final = InFlightRequestsMiddleware._get_gauge()
assert in_flight is not None
- breaker: Final = _breaker_metrics()
- assert breaker._state_gauge is not None
- assert breaker._transitions is not None
- assert breaker._failures is not None
+ services: Final = PrometheusServicesLogger()
return (
SpendLogCleanupMetrics.rows_deleted,
SpendLogCleanupMetrics.batch_duration,
@@ -73,9 +69,9 @@ def _lazy_owner_collectors() -> tuple[Collector, ...]:
SpendLogCleanupMetrics.batch_failures,
SpendLogCleanupMetrics.runs,
in_flight,
- breaker._state_gauge,
- breaker._transitions,
- breaker._failures,
+ services.redis_breaker_state,
+ services.redis_breaker_transitions,
+ services.redis_breaker_failures,
)
From 47e4ed88c2d89a2ecd9cb79be2a1ea0829642b87 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
<158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 14 Oct 2026 14:42:49 -0700
Subject: [PATCH 37/47] fix(anthropic): keep refusal stop_reason and
stop_details on the logged response (#46104)
* fix(anthropic): keep refusal stop_reason and stop_details on the logged response
* fix(anthropic): keep the refusal stop-fields boundary tolerant of a non-dict upstream body
* test(anthropic): use Anthropic's documented refusal example as the unit-test fixture
* fix(streaming): keep provider fields on async stream cache hits and received_model_id on synthesized finish chunks
* fix(caching): replay cached completion_tokens_details on stream cache hits and skip non-object cached provider fields
* test(integration): cover Anthropic-family refusal stop_details on the wire, the cache and under chaos
* fix(caching): skip a cached completion_tokens_details that fails validation on stream replay
* fix(anthropic): keep refusal stop fields on a JSON-mode non-stream response
Also adds the round 3 audit cells for JSON-mode refusals, the proxy response serialization, and wrong-typed cached usage details on both stream replays
* test(anthropic): expect content_filter on streamed JSON tool refusals after main's stream change
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
.../convert_dict_to_response.py | 40 +-
.../litellm_core_utils/streaming_handler.py | 6 +-
litellm/llms/anthropic/chat/handler.py | 11 +-
litellm/llms/anthropic/chat/transformation.py | 12 +-
litellm/llms/anthropic/common_utils.py | 37 +-
.../integration/_support/anthropic_refusal.py | 533 +++++++
.../caching/response_cache_case.py | 23 +
...st_anthropic_refusal_stop_details_chaos.py | 337 ++++
...est_anthropic_refusal_stop_details_wire.py | 1356 +++++++++++++++++
.../sdk/test_response_cache_sync_sdk.py | 82 +
.../test_convert_to_streaming_response.py | 70 +-
.../test_streaming_handler.py | 36 +
.../llms/anthropic/anthropic_docs_refusal.py | 19 +
.../chat/test_anthropic_chat_handler.py | 48 +
.../test_anthropic_chat_transformation.py | 91 ++
...ations_anthropic_claude3_transformation.py | 23 +-
...partner_models_anthropic_transformation.py | 22 +
...t_anthropic_passthrough_logging_handler.py | 56 +
18 files changed, 2789 insertions(+), 13 deletions(-)
create mode 100644 tests/integration/_support/anthropic_refusal.py
create mode 100644 tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py
create mode 100644 tests/integration/providers/test_anthropic_refusal_stop_details_wire.py
create mode 100644 tests/unit/llms/anthropic/anthropic_docs_refusal.py
diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
index 6f16e1e3fb4..ad33c47fc51 100644
--- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
+++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
@@ -6,7 +6,7 @@ import traceback
from collections.abc import Mapping, Sequence
from typing import Final, Literal, cast
-from pydantic import ConfigDict, TypeAdapter, with_config
+from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
from typing_extensions import ReadOnly, TypedDict
import litellm
@@ -222,6 +222,9 @@ async def convert_to_streaming_response_async(
role=choice["message"]["role"],
function_call=choice["message"].get("function_call", None),
tool_calls=choice["message"].get("tool_calls", None),
+ provider_specific_fields=_cached_provider_specific_fields(
+ _CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(choice["message"])
+ ),
)
finish_reason = choice.get("finish_reason", None)
@@ -243,6 +246,9 @@ async def convert_to_streaming_response_async(
completion_tokens=response_object["usage"].get("completion_tokens", 0),
prompt_tokens=response_object["usage"].get("prompt_tokens", 0),
total_tokens=response_object["usage"].get("total_tokens", 0),
+ completion_tokens_details=_cached_completion_tokens_details(
+ _CACHED_USAGE_DETAILS.validate_python(response_object["usage"])
+ ),
),
)
@@ -334,6 +340,9 @@ def convert_to_streaming_response(
model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0)
model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0)
model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0)
+ model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details(
+ _CACHED_USAGE_DETAILS.validate_python(response_object["usage"])
+ )
if "id" in response_object:
model_response_object.id = response_object["id"]
@@ -390,6 +399,35 @@ class _ParallelToolUseArgs(TypedDict):
_PARALLEL_TOOL_USE_ARGS: Final = TypeAdapter(_ParallelToolUseArgs)
+class _CachedMessageProviderFields(TypedDict, total=False):
+ provider_specific_fields: ReadOnly[object]
+
+
+class _CachedUsageDetails(TypedDict, total=False):
+ completion_tokens_details: ReadOnly[object]
+
+
+_CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields)
+_CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails)
+_FIELD_MAP: Final = TypeAdapter(dict[str, object])
+_COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper)
+
+
+def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None:
+ fields: Final = message.get("provider_specific_fields")
+ return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else None
+
+
+def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionTokensDetailsWrapper | None:
+ details: Final = usage.get("completion_tokens_details")
+ if not isinstance(details, dict):
+ return None
+ try:
+ return _COMPLETION_TOKENS_DETAILS.validate_python(details)
+ except ValidationError:
+ return None
+
+
def handle_invalid_parallel_tool_calls(
tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall],
) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None:
diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 236286a493b..484978bae04 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -280,6 +280,7 @@ class CustomStreamWrapper:
self._response_headers = _response_headers
self.response_id: str | None = None
+ self.received_model_id: str | None = None
self.logging_loop = None
self.rules = Rules()
self.stream_options = stream_options or getattr(logging_obj, "stream_options", None)
@@ -838,6 +839,7 @@ class CustomStreamWrapper:
self.response_id = id
if id and isinstance(id, str) and id.strip():
+ self.received_model_id = id
model_response.hidden_params["received_model_id"] = id
if self.response_id is not None and isinstance(self.response_id, str):
@@ -1756,7 +1758,9 @@ class CustomStreamWrapper:
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
def finish_reason_handler(self):
- model_response: Final = self.model_response_creator()
+ model_response: Final = self.model_response_creator(
+ hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id}
+ )
_finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason
if _finish_reason is not None:
model_response.choices[0].finish_reason = _finish_reason
diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py
index dff10473ae7..4e5b761b9a3 100644
--- a/litellm/llms/anthropic/chat/handler.py
+++ b/litellm/llms/anthropic/chat/handler.py
@@ -66,7 +66,7 @@ from litellm.types.utils import (
)
from ...base import BaseLLM
-from ..common_utils import AnthropicError, process_anthropic_headers
+from ..common_utils import AnthropicError, process_anthropic_headers, refusal_provider_specific_fields
from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig, JsonModeRequest
if TYPE_CHECKING:
@@ -1102,7 +1102,8 @@ class ModelResponseIterator:
# These are automatically handled by Anthropic API, we just pass them through
pass
elif type_chunk == "message_delta":
- finish_reason, usage, container = self._handle_message_delta(chunk)
+ finish_reason, usage, container, refusal_fields = self._handle_message_delta(chunk)
+ provider_specific_fields.update(refusal_fields)
text = self._release_answer_close() + self._drain_split_answer_blocks()
if container:
provider_specific_fields["container"] = container
@@ -1297,7 +1298,9 @@ class ModelResponseIterator:
return ""
return _split_answer_text(blocks, self.response_format_schema)
- def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]:
+ def _handle_message_delta(
+ self, chunk: dict
+ ) -> tuple[str, Usage | None, dict[str, Any] | None, Mapping[str, object]]:
"""
Handle message_delta event for finish_reason, usage, and container.
@@ -1319,7 +1322,7 @@ class ModelResponseIterator:
self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None
)
container: Final = message_delta["delta"].get("container")
- return finish_reason, usage, container
+ return finish_reason, usage, container, refusal_provider_specific_fields(message_delta["delta"])
def _handle_accumulated_json_chunk(self, data_str: str, is_final: bool = False) -> ModelResponseStream | None:
"""
diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py
index f4e50bf3acd..487f231e1db 100644
--- a/litellm/llms/anthropic/chat/transformation.py
+++ b/litellm/llms/anthropic/chat/transformation.py
@@ -104,10 +104,12 @@ from litellm.utils import (
)
from ..common_utils import (
+ REFUSAL_STOP_FIELDS,
AnthropicError,
AnthropicModelInfo,
eager_input_streaming_flag,
process_anthropic_headers,
+ refusal_provider_specific_fields,
requires_native_compaction_beta,
strip_advisor_blocks_from_messages,
)
@@ -2723,7 +2725,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
if compaction_blocks is not None:
provider_specific_fields["compaction_blocks"] = compaction_blocks
- return provider_specific_fields
+ refusal_fields: Final = refusal_provider_specific_fields(
+ REFUSAL_STOP_FIELDS.validate_python(completion_response)
+ )
+ return {**provider_specific_fields, **refusal_fields}
def transform_parsed_response(
self,
@@ -2808,7 +2813,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
_message.provider_specific_fields = provider_specific_fields
if json_mode_message is not None:
+ refusal_fields: Final = refusal_provider_specific_fields(
+ REFUSAL_STOP_FIELDS.validate_python(completion_response)
+ )
completion_response["stop_reason"] = "stop"
+ if refusal_fields:
+ json_mode_message.provider_specific_fields = dict(refusal_fields)
_message = json_mode_message
model_response.choices[0].message = _message
diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py
index 47a82aa097c..dbcb7ea4962 100644
--- a/litellm/llms/anthropic/common_utils.py
+++ b/litellm/llms/anthropic/common_utils.py
@@ -8,11 +8,12 @@ import re
from collections.abc import Mapping, MutableMapping, Sequence
from datetime import datetime, timezone
from types import MappingProxyType
-from typing import Any, ClassVar, Final, Literal, TypeVar
+from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar
from urllib.parse import quote, urlparse
import httpx
-from pydantic import ConfigDict, Field, StrictBool, TypeAdapter, ValidationError
+from pydantic import BeforeValidator, ConfigDict, Field, StrictBool, TypeAdapter, ValidationError
+from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.constants import (
@@ -2150,3 +2151,35 @@ def create_anthropic_model_list_response(
"first_id": data[0]["id"] if data else None,
"last_id": data[-1]["id"] if data else None,
}
+
+
+class RefusalStopFields(TypedDict, total=False):
+ stop_reason: ReadOnly[object]
+ stop_details: ReadOnly[object]
+
+
+def _is_dict(value: object) -> bool:
+ return isinstance(value, dict)
+
+
+def _no_stop_fields() -> Mapping[str, object]:
+ return {}
+
+
+def _dict_or_empty(value: object) -> object:
+ return value if _is_dict(value) else _no_stop_fields()
+
+
+REFUSAL_STOP_FIELDS: Final[TypeAdapter[RefusalStopFields]] = TypeAdapter(
+ Annotated[RefusalStopFields, BeforeValidator(_dict_or_empty)]
+)
+
+
+def refusal_provider_specific_fields(message: Mapping[str, object]) -> Mapping[str, object]:
+ stop_reason: Final = message.get("stop_reason")
+ stop_details: Final = message.get("stop_details")
+ if stop_details is not None:
+ return {"stop_reason": stop_reason, "stop_details": stop_details}
+ if stop_reason == "refusal":
+ return {"stop_reason": stop_reason}
+ return {}
diff --git a/tests/integration/_support/anthropic_refusal.py b/tests/integration/_support/anthropic_refusal.py
new file mode 100644
index 00000000000..6d5fd236074
--- /dev/null
+++ b/tests/integration/_support/anthropic_refusal.py
@@ -0,0 +1,533 @@
+import hashlib
+import json
+import os
+import re
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from itertools import chain
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final, TypeAlias
+
+import yaml
+from integration._support.anthropic_thinking import (
+ JSON_LIST,
+ JSON_OBJECT,
+ Event,
+ marker_of,
+ stream_reply,
+ streams,
+)
+from integration._support.client import eventually
+from integration._support.wire import Reply, Request, Wire
+from pydantic import JsonValue, TypeAdapter
+from redis import Redis
+
+DOCS_REFUSAL_URL: Final = "https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback#refusal-response"
+DOCS_STREAMING_REFUSAL_URL: Final = (
+ "https://platform.claude.com/docs/en/test-and-evaluate/strengthen-guardrails/handle-streaming-refusals"
+)
+
+# Copied verbatim from DOCS_REFUSAL_URL and DOCS_STREAMING_REFUSAL_URL on 2026-10-12.
+DOCS_REFUSAL: Final[Event] = {
+ "id": "msg_01XFUDYJgAACzvnptvVoYEL",
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-fable-5",
+ "content": [],
+ "stop_reason": "refusal",
+ "stop_details": {
+ "type": "refusal",
+ "category": "cyber",
+ "explanation": "This request was declined because it could enable cyber harm.",
+ },
+ "usage": {"input_tokens": 412, "output_tokens": 0},
+}
+DOCS_PARTIAL_REFUSAL: Final[Event] = {
+ "role": "assistant",
+ "content": [{"type": "text", "text": "Hello.."}],
+ "stop_reason": "refusal",
+ "stop_details": {
+ "type": "refusal",
+ "category": "cyber",
+ "explanation": "This request was declined because it could enable cyber harm.",
+ },
+}
+
+# Captured from api.anthropic.com (claude-sonnet-5-5, prompt "Reply with the single word: ready") on 2026-10-13.
+CAPTURED_TURN: Final[Event] = {
+ "model": "claude-sonnet-5-5",
+ "id": "msg_011CfyXFGze1cDkk8o7gDRRS",
+ "type": "message",
+ "role": "assistant",
+ "content": [{"type": "text", "text": "ready"}],
+ "container": None,
+ "stop_reason": "end_turn",
+ "stop_sequence": None,
+ "stop_details": None,
+ "usage": {
+ "input_tokens": 18,
+ "cache_creation_input_tokens": 0,
+ "cache_read_input_tokens": 0,
+ "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0},
+ "output_tokens": 4,
+ "output_tokens_details": {"thinking_tokens": 0},
+ "service_tier": "standard",
+ "inference_geo": "global",
+ },
+ "diagnostics": None,
+}
+CAPTURED_EVENTS: Final[tuple[Event, ...]] = (
+ {
+ "type": "message_start",
+ "message": {
+ "model": "claude-sonnet-5-5",
+ "id": "msg_011CfyXFLwXUxsp7nNAm3gcu",
+ "type": "message",
+ "role": "assistant",
+ "content": [],
+ "container": None,
+ "stop_reason": None,
+ "stop_sequence": None,
+ "stop_details": None,
+ "usage": {
+ "input_tokens": 18,
+ "cache_creation_input_tokens": 0,
+ "cache_read_input_tokens": 0,
+ "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0},
+ "output_tokens": 4,
+ "service_tier": "standard",
+ "inference_geo": "global",
+ },
+ "diagnostics": None,
+ },
+ },
+ {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
+ {"type": "ping"},
+ {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ready"}},
+ {"type": "content_block_stop", "index": 0},
+ {
+ "type": "message_delta",
+ "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None},
+ "usage": {
+ "input_tokens": 18,
+ "cache_creation_input_tokens": 0,
+ "cache_read_input_tokens": 0,
+ "output_tokens": 4,
+ "output_tokens_details": {"thinking_tokens": 0},
+ },
+ },
+ {"type": "message_stop"},
+)
+
+DOCS_STOP_DETAILS: Final = DOCS_REFUSAL["stop_details"]
+REFUSAL_FIELDS: Final[Event] = {"stop_reason": "refusal", "stop_details": DOCS_STOP_DETAILS}
+STOP_FIELDS: Final = ("stop_reason", "stop_details")
+LONG_STOP_DETAILS: Final = "d" * 5120
+MALFORMED_STOP_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType(
+ {"int": 7, "list": ["refusal"], "empty": "", "long": LONG_STOP_DETAILS}
+)
+_FIXTURE: Final = re.compile(r"refusal-fixture-([a-z]+)")
+
+Location: TypeAlias = tuple[str | int, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class Fixture:
+ status: int
+ body: bytes
+ events: tuple[Event, ...] = ()
+
+
+def _text_block(index: int, block: JsonValue) -> tuple[Event, ...]:
+ text: Final = JSON_OBJECT.validate_python(block)["text"]
+ return (
+ {"type": "content_block_start", "index": index, "content_block": {"type": "text", "text": ""}},
+ {"type": "content_block_delta", "index": index, "delta": {"type": "text_delta", "text": text}},
+ {"type": "content_block_stop", "index": index},
+ )
+
+
+def refusal_events(message: Mapping[str, JsonValue]) -> tuple[Event, ...]:
+ usage: Final = JSON_OBJECT.validate_python(message["usage"])
+ start: Final[Event] = {
+ "type": "message_start",
+ "message": {
+ "id": message["id"],
+ "type": "message",
+ "role": "assistant",
+ "model": message["model"],
+ "content": [],
+ "stop_reason": None,
+ "stop_sequence": None,
+ "stop_details": None,
+ "usage": usage,
+ },
+ }
+ blocks: Final = chain.from_iterable(
+ _text_block(index, block) for index, block in enumerate(JSON_LIST.validate_python(message["content"]))
+ )
+ delta: Final[Event] = {
+ "type": "message_delta",
+ "delta": {
+ "stop_reason": message["stop_reason"],
+ "stop_sequence": None,
+ "stop_details": message["stop_details"],
+ },
+ "usage": {name: usage[name] for name in ("output_tokens", "output_tokens_details") if name in usage},
+ }
+ return (start, *blocks, delta, {"type": "message_stop"})
+
+
+def _without_stop_details(value: Mapping[str, JsonValue]) -> Event:
+ return {name: item for name, item in value.items() if name != "stop_details"}
+
+
+def _captured_event(event: Event, stop_details: JsonValue, *, keep: bool) -> Event:
+ if event["type"] == "message_start":
+ message: Final = JSON_OBJECT.validate_python(event["message"])
+ return {**event, "message": message if keep else _without_stop_details(message)}
+ if event["type"] != "message_delta":
+ return event
+ delta: Final = JSON_OBJECT.validate_python(event["delta"])
+ replaced: Final = {**delta, "stop_details": stop_details} if keep else _without_stop_details(delta)
+ return {**event, "delta": replaced}
+
+
+def _captured(stop_details: JsonValue, *, keep: bool = True) -> Fixture:
+ body: Final = {**CAPTURED_TURN, "stop_details": stop_details} if keep else _without_stop_details(CAPTURED_TURN)
+ return Fixture(
+ 200,
+ json.dumps(body).encode(),
+ tuple(_captured_event(event, stop_details, keep=keep) for event in CAPTURED_EVENTS),
+ )
+
+
+def _duplicated_stop_details() -> bytes:
+ pair: Final = f'"stop_details": {json.dumps(DOCS_STOP_DETAILS)}'
+ return json.dumps(DOCS_REFUSAL).replace(pair, f"{pair}, {pair}").encode()
+
+
+def _anthropic_error(kind: str, message: str) -> bytes:
+ return json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode()
+
+
+_PARTIAL: Final[Event] = {**DOCS_REFUSAL, **DOCS_PARTIAL_REFUSAL}
+_REASONED: Final[Event] = {
+ **_PARTIAL,
+ "usage": {"input_tokens": 412, "output_tokens": 10, "output_tokens_details": {"thinking_tokens": 4}},
+}
+
+JSON_TOOL_NAME: Final = "json_tool_call"
+JSON_ANSWER_SCHEMA: Final[Event] = {
+ "type": "object",
+ "properties": {"answer": {"type": "string"}},
+ "required": ["answer"],
+ "additionalProperties": False,
+}
+JSON_ANSWER_FORMAT: Final[Event] = {
+ "type": "json_schema",
+ "json_schema": {"name": "answer", "schema": JSON_ANSWER_SCHEMA, "strict": True},
+}
+DOCS_PARTIAL_TEXT: Final = "Hello.."
+_DECLARED_INPUT: Final[Event] = {"answer": DOCS_PARTIAL_TEXT}
+_ENVELOPED_INPUT: Final[Event] = {"parameters": _DECLARED_INPUT}
+DECLARED_PARTIAL_JSON: Final = '{"answer": "Hello..'
+ENVELOPED_PARTIAL_JSON: Final = '{"parameters": {"answer": "Hello..'
+_JSON_TOOL_USE_ID: Final = "toolu_01RefusalFixtureJsonTool"
+
+
+def _json_tool_refusal(tool_input: Event, partial_json: str) -> Fixture:
+ message: Final[Event] = {
+ **DOCS_REFUSAL,
+ "content": [{"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": tool_input}],
+ }
+ start, delta, stop = refusal_events(DOCS_REFUSAL)
+ tool_block: Final[tuple[Event, ...]] = (
+ {
+ "type": "content_block_start",
+ "index": 0,
+ "content_block": {"type": "tool_use", "id": _JSON_TOOL_USE_ID, "name": JSON_TOOL_NAME, "input": {}},
+ },
+ {
+ "type": "content_block_delta",
+ "index": 0,
+ "delta": {"type": "input_json_delta", "partial_json": partial_json},
+ },
+ {"type": "content_block_stop", "index": 0},
+ )
+ return Fixture(200, json.dumps(message).encode(), (start, *tool_block, delta, stop))
+
+
+FIXTURES: Final[Mapping[str, Fixture]] = MappingProxyType(
+ {
+ "refusal": Fixture(200, json.dumps(DOCS_REFUSAL).encode(), refusal_events(DOCS_REFUSAL)),
+ "partial": Fixture(200, json.dumps(_PARTIAL).encode(), refusal_events(_PARTIAL)),
+ "reasoned": Fixture(200, json.dumps(_REASONED).encode(), refusal_events(_REASONED)),
+ "jsontool": _json_tool_refusal(_DECLARED_INPUT, DECLARED_PARTIAL_JSON),
+ "jsonheld": _json_tool_refusal(_ENVELOPED_INPUT, ENVELOPED_PARTIAL_JSON),
+ "duplicate": Fixture(200, _duplicated_stop_details()),
+ "null": _captured(None),
+ "missing": _captured(None, keep=False),
+ **{name: _captured(value) for name, value in MALFORMED_STOP_DETAILS.items()},
+ "nondict": Fixture(200, b"[]"),
+ "unauthorized": Fixture(401, _anthropic_error("authentication_error", "invalid x-api-key")),
+ "invalid": Fixture(400, _anthropic_error("invalid_request_error", "max_tokens: Field required")),
+ }
+)
+
+
+def refusal_prompt(fixture: str, marker: str) -> str:
+ return f"refusal-fixture-{fixture} marker-{marker}"
+
+
+CITED_FIELDS: Final[Event] = {"citations": ["https://example.com/source"]}
+_PLAIN_DELTA: Final[Event] = {"role": "assistant", "content": "ready"}
+_CITED_DELTA: Final[Event] = {**_PLAIN_DELTA, "provider_specific_fields": CITED_FIELDS}
+_OPENAI_USAGE: Final[Event] = {
+ "prompt_tokens": 5,
+ "completion_tokens": 7,
+ "total_tokens": 12,
+ "completion_tokens_details": {"reasoning_tokens": 4, "text_tokens": 3},
+}
+UNVALIDATED_USAGE_DETAILS: Final[Mapping[str, JsonValue]] = MappingProxyType(
+ {
+ "text": {"reasoning_tokens": "abc"},
+ "fraction": {"reasoning_tokens": 1.5},
+ "object": {"reasoning_tokens": {"a": 1}},
+ }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class OpenAIFixture:
+ first_delta: Event
+ usage: tuple[Event, ...] = ()
+
+
+_OPENAI_FIXTURES: Final[Mapping[str, OpenAIFixture]] = MappingProxyType(
+ {
+ "plain": OpenAIFixture(_PLAIN_DELTA),
+ "cited": OpenAIFixture(_CITED_DELTA),
+ "plainusage": OpenAIFixture(_PLAIN_DELTA, (_OPENAI_USAGE,)),
+ "citedusage": OpenAIFixture(_CITED_DELTA, (_OPENAI_USAGE,)),
+ }
+)
+
+
+def openai_id(marker: str) -> str:
+ return f"chatcmpl-{marker}"
+
+
+def _openai_frame(marker: str, body: Event) -> bytes:
+ chunk: Final[Event] = {
+ "id": openai_id(marker),
+ "object": "chat.completion.chunk",
+ "created": 1760000000,
+ "model": "gpt-4o-mini",
+ **body,
+ }
+ return b"data: " + json.dumps(chunk).encode() + b"\n\n"
+
+
+def _openai_choice(delta: Event, finish_reason: str | None) -> Event:
+ return {"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}]}
+
+
+def _openai_stream(name: str, request: Request) -> Reply:
+ assert streams(request), request.body
+ marker: Final = marker_of(request)
+ fixture: Final = _OPENAI_FIXTURES[name]
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(
+ _openai_frame(marker, _openai_choice(fixture.first_delta, None)),
+ _openai_frame(marker, _openai_choice({"content": " now"}, "stop")),
+ *(_openai_frame(marker, {"choices": [], "usage": usage}) for usage in fixture.usage),
+ b"data: [DONE]\n\n",
+ ),
+ )
+
+
+def respond(request: Request) -> Reply:
+ found: Final[Sequence[str]] = _FIXTURE.findall(request.body.decode())
+ assert found, request.body
+ if request.target.endswith("/chat/completions"):
+ return _openai_stream(found[-1], request)
+ fixture: Final = FIXTURES[found[-1]]
+ if fixture.status != 200 or not fixture.events or not streams(request):
+ return Reply(status=fixture.status, body=fixture.body)
+ return stream_reply(request, fixture.events)
+
+
+def _stop_fields(fields: JsonValue) -> Event:
+ if not isinstance(fields, dict):
+ return {}
+ return {name: fields[name] for name in STOP_FIELDS if name in fields}
+
+
+def relayed(value: JsonValue, location: Location = ()) -> tuple[tuple[Location, Event], ...]:
+ if isinstance(value, list):
+ return tuple(chain.from_iterable(relayed(item, (*location, index)) for index, item in enumerate(value)))
+ if not isinstance(value, dict):
+ return ()
+ found: Final = _stop_fields(value.get("provider_specific_fields"))
+ here: Final = (((*location, "provider_specific_fields"), found),) if found else ()
+ return (*here, *chain.from_iterable(relayed(item, (*location, name)) for name, item in value.items()))
+
+
+_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}")
+_STORED: Final = TypeAdapter(bytes | None)
+STALE_ID_SUFFIX: Final = "-stale"
+
+
+def _stored(store: Redis, key: str) -> bytes:
+ if not _CACHE_KEY.fullmatch(key) or store.type(key) != b"string":
+ return b""
+ return _STORED.validate_python(store.get(key)) or b""
+
+
+def _entry_key(store: Redis, needle: str) -> str | None:
+ keys: Final = (raw.decode() for raw in store.scan_iter(count=1000))
+ return next((key for key in keys if needle.encode() in _stored(store, key)), None)
+
+
+def _with_message_fields(response: Event, fields: JsonValue) -> Event:
+ choice: Final = JSON_OBJECT.validate_python(JSON_LIST.validate_python(response["choices"])[0])
+ message: Final = JSON_OBJECT.validate_python(choice["message"])
+ return {**response, "choices": [{**choice, "message": {**message, "provider_specific_fields": fields}}]}
+
+
+def _with_usage_details(response: Event, details: JsonValue) -> Event:
+ usage: Final = JSON_OBJECT.validate_python(response["usage"])
+ return {**response, "usage": {**usage, "completion_tokens_details": details}}
+
+
+def _rewrite_cached_response(needle: str, rewrite: Callable[[Event], Event]) -> str:
+ store: Final = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
+ key: Final = eventually(lambda: _entry_key(store, needle), lambda found: found is not None)
+ assert key is not None, needle
+ entry: Final = JSON_OBJECT.validate_json(_stored(store, key))
+ stored: Final = entry["response"]
+ encoded: Final = isinstance(stored, str)
+ response: Final = (
+ JSON_OBJECT.validate_json(stored) if isinstance(stored, str) else JSON_OBJECT.validate_python(stored)
+ )
+ rewritten: Final = {**rewrite(response), "id": f"{response['id']}{STALE_ID_SUFFIX}"}
+ store.set(key, json.dumps({**entry, "response": json.dumps(rewritten) if encoded else rewritten}), ex=60)
+ return str(rewritten["id"])
+
+
+def rewrite_cached_message_fields(needle: str, fields: JsonValue) -> str:
+ return _rewrite_cached_response(needle, lambda response: _with_message_fields(response, fields))
+
+
+def rewrite_cached_usage_details(needle: str, details: JsonValue) -> str:
+ return _rewrite_cached_response(needle, lambda response: _with_usage_details(response, details))
+
+
+def key_hash(key: str) -> str:
+ return hashlib.sha256(key.encode()).hexdigest()
+
+
+def _key_hash_of(event: Mapping[str, JsonValue]) -> JsonValue:
+ return JSON_OBJECT.validate_python(event.get("metadata") or {}).get("user_api_key_hash")
+
+
+class Ledger:
+ def __init__(self, wire: Wire) -> None:
+ self.wire: Final = wire
+ self.batches: tuple[Request, ...] = ()
+
+ def events(self) -> tuple[Event, ...]:
+ self.batches = (*self.batches, *self.wire.drain())
+ decoded: Final = chain.from_iterable(JSON_LIST.validate_json(batch.body) for batch in self.batches)
+ return tuple(JSON_OBJECT.validate_python(event) for event in decoded)
+
+ def for_key(self, key: str) -> tuple[Event, ...]:
+ hashed: Final = key_hash(key)
+ return tuple(event for event in self.events() if _key_hash_of(event) == hashed)
+
+ def await_key(self, key: str, count: int = 1, seconds: float = 30) -> tuple[Event, ...]:
+ return eventually(lambda: self.for_key(key), lambda found: len(found) >= count, seconds=seconds)
+
+
+_HIDDEN_PARAMS_PROBE: Final = """
+import httpx
+from litellm.integrations.custom_logger import CustomLogger
+
+
+class HiddenParamsProbe(CustomLogger):
+ async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+ payload = kwargs.get("standard_logging_object") or {{}}
+ hidden = getattr(response_obj, "_hidden_params", None) or {{}}
+ event = {{
+ "metadata": {{"user_api_key_hash": (payload.get("metadata") or {{}}).get("user_api_key_hash")}},
+ "response_id": getattr(response_obj, "id", None),
+ "received_model_id": hidden.get("received_model_id"),
+ }}
+ async with httpx.AsyncClient(trust_env=False) as client:
+ await client.post({url!r}, json=[event])
+
+
+instance = HiddenParamsProbe()
+"""
+
+
+def hidden_params_probe(directory: Path, url: str) -> str:
+ (directory / "hidden_params_probe.py").write_text(_HIDDEN_PARAMS_PROBE.format(url=url))
+ return "hidden_params_probe.instance"
+
+
+def sink_config(
+ directory: Path,
+ sinks: Mapping[str, str],
+ *,
+ retries: int = 0,
+ extra_callbacks: Sequence[str] = (),
+ store_prompts: bool = False,
+) -> Path:
+ shared: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+ settings: Final = {
+ **JSON_OBJECT.validate_python(shared["litellm_settings"]),
+ "callbacks": [*sinks, *extra_callbacks],
+ "DEFAULT_FLUSH_INTERVAL_SECONDS": 1,
+ }
+ callback_settings: Final = {
+ name: {
+ "callback_type": "generic_api",
+ "endpoint": url,
+ "headers": {"x-refusal-sink": name},
+ "max_retries": retries,
+ "retry_delay": 0.5,
+ }
+ for name, url in sinks.items()
+ }
+ general_settings: Final = {
+ **JSON_OBJECT.validate_python(shared["general_settings"]),
+ "store_prompts_in_spend_logs": store_prompts,
+ }
+ path: Final = directory / "anthropic-refusal-sinks.yaml"
+ path.write_text(
+ yaml.safe_dump(
+ {
+ **shared,
+ "general_settings": general_settings,
+ "litellm_settings": settings,
+ "callback_settings": callback_settings,
+ }
+ )
+ )
+ return path
+
+
+def frames(text: str) -> tuple[Event, ...]:
+ return tuple(
+ JSON_OBJECT.validate_json(line.removeprefix("data: "))
+ for line in text.splitlines()
+ if line.startswith("data: {")
+ )
+
+
+def frame_types(events: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]:
+ return tuple(event.get("type") for event in events)
diff --git a/tests/integration/caching/response_cache_case.py b/tests/integration/caching/response_cache_case.py
index 9df34721202..8dfb6a0a479 100644
--- a/tests/integration/caching/response_cache_case.py
+++ b/tests/integration/caching/response_cache_case.py
@@ -336,6 +336,29 @@ def flip_entry_to_empty(store: Redis, key: str, field: str) -> str:
return stale_id
+def rewrite_usage_details(store: Redis, key: str, details: JsonValue) -> str:
+ """Rewrite the stored usage's ``completion_tokens_details`` to ``details`` and stamp the id ``-stale``, keeping the
+ encoding, and return the stale id."""
+ raw: Final = store.get(key)
+ assert raw is not None, key
+ entry, is_json = _decode_entry(raw)
+ response: Final = entry["response"]
+ stored: Final = object_value(json.loads(response) if isinstance(response, str) else response)
+ stale_id: Final = f"{string_value(stored['id'])}{STALE_ID_SUFFIX}"
+ usage: Final = object_value(stored["usage"])
+ rewritten_response: Final[dict[str, JsonValue]] = {
+ **stored,
+ "usage": {**usage, "completion_tokens_details": details},
+ "id": stale_id,
+ }
+ rewritten: Final[dict[str, JsonValue]] = {
+ **entry,
+ "response": json.dumps(rewritten_response) if isinstance(response, str) else rewritten_response,
+ }
+ store.set(key, json.dumps(rewritten) if is_json else str(rewritten), ex=STALE_ENTRY_TTL_SECONDS)
+ return stale_id
+
+
def chat_body(model: str, text: str, **extra: JsonValue) -> dict[str, JsonValue]:
return {"model": model, "messages": [{"role": "user", "content": text}], **extra}
diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py b/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py
new file mode 100644
index 00000000000..dd4b7602b96
--- /dev/null
+++ b/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py
@@ -0,0 +1,337 @@
+import asyncio
+import json
+import re
+import threading
+import time
+import uuid
+from collections import Counter
+from collections.abc import Callable, Mapping, Sequence
+from contextlib import ExitStack
+from dataclasses import dataclass
+from itertools import chain
+from pathlib import Path
+from queue import SimpleQueue
+from types import MappingProxyType
+from typing import Final, Literal, TypeAlias
+from urllib.parse import urlsplit
+
+import httpx
+import psutil
+import pytest
+import yaml
+from integration._support.anthropic_refusal import (
+ REFUSAL_FIELDS,
+ Ledger,
+ Location,
+ refusal_prompt,
+ relayed,
+ respond,
+ sink_config,
+)
+from integration._support.anthropic_thinking import JSON_OBJECT, MODEL, Event, chunks_of
+from integration._support.client import Gateway, eventually, held
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.timing import assert_not_delayed_by
+from integration._support.wire import Reply, Request, Wire, wire_server
+
+pytestmark: Final = [
+ pytest.mark.timeout(2 * graceful_stop_seconds() + 120),
+ pytest.mark.xdist_group("anthropic_refusal_stop_details_chaos"),
+]
+
+_MODEL_NAME: Final = "anthropic-refusal-chaos"
+_SLOW_SINK_SECONDS: Final = 4
+_SINK: Final = "refusal_chaos_sink"
+_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
+_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
+_CHAT_FIELDS: Final[Location] = ("choices", 0, "message", "provider_specific_fields")
+_DELTA_FIELDS: Final[Location] = ("choices", 0, "delta", "provider_specific_fields")
+_PATHS: Final[Mapping[str, str]] = MappingProxyType(
+ {"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"}
+)
+
+Surface: TypeAlias = Literal["chat", "messages", "responses"]
+_RecordProperty: TypeAlias = Callable[[str, object], None]
+
+
+@dataclass(frozen=True, slots=True)
+class _Call:
+ surface: Surface
+ stream: bool
+ marker: str
+
+
+@dataclass(frozen=True, slots=True)
+class _Served:
+ call: _Call
+ status: int
+ text: str
+
+
+@dataclass(frozen=True, slots=True)
+class _Held:
+ release: threading.Event
+ markers: SimpleQueue[str]
+
+ def respond(self, request: Request) -> Reply:
+ self.markers.put(_markers_in(request.body.decode())[0])
+ assert self.release.wait(timeout=90), "The burst was never released"
+ return respond(request)
+
+
+def _markers_in(text: str) -> tuple[str, ...]:
+ return tuple(dict.fromkeys(_MARKER.findall(text)))
+
+
+def _calls(count: int) -> tuple[_Call, ...]:
+ surfaces: Final[tuple[Surface, ...]] = ("chat", "messages", "responses")
+ return tuple(
+ _Call(surfaces[index % 3], stream=(index // 3) % 2 == 0, marker=uuid.uuid4().hex) for index in range(count)
+ )
+
+
+def _body(call: _Call) -> Event:
+ text: Final = refusal_prompt("refusal", call.marker)
+ shared: Final[Event] = {"model": _MODEL_NAME, "stream": call.stream, "cache": {"no-cache": True}}
+ match call.surface:
+ case "chat":
+ return {**shared, "messages": [{"role": "user", "content": text}], "max_tokens": 64}
+ case "messages":
+ return {**shared, "messages": [{"role": "user", "content": text}], "max_tokens": 64}
+ case "responses":
+ return {**shared, "input": text, "max_output_tokens": 64}
+
+
+async def _send(client: httpx.AsyncClient, key: str, call: _Call) -> _Served:
+ try:
+ response: Final = await client.post(
+ _PATHS[call.surface], json=_body(call), headers={"Authorization": f"Bearer {key}"}
+ )
+ except httpx.TransportError as error:
+ return _Served(call, 0, repr(error))
+ return _Served(call, response.status_code, response.text)
+
+
+async def _burst(gateway: Gateway, calls: Sequence[_Call]) -> tuple[_Served, ...]:
+ async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=120, trust_env=False) as client:
+ return tuple(await asyncio.gather(*(_send(client, gateway.key, call) for call in calls)))
+
+
+def _chaos_config(directory: Path, upstream: Wire, sink: str, *, retries: int) -> Path:
+ base: Final = sink_config(directory, {_SINK: sink}, retries=retries)
+ config: Final = JSON_OBJECT.validate_python(yaml.safe_load(base.read_text()))
+ model_list: Final = [
+ {
+ "model_name": _MODEL_NAME,
+ "litellm_params": {"model": f"anthropic/{MODEL}", "api_base": upstream.url, "api_key": "scripted-key"},
+ }
+ ]
+ path: Final = directory / "anthropic-refusal-chaos.yaml"
+ path.write_text(yaml.safe_dump({**config, "model_list": model_list}))
+ return path
+
+
+def _ledger_events(ledgers: Sequence[Ledger]) -> tuple[Event, ...]:
+ return tuple(chain.from_iterable(ledger.events() for ledger in ledgers))
+
+
+def _marker_counts(ledgers: Sequence[Ledger], markers: frozenset[str]) -> Mapping[str, int]:
+ found: Final = chain.from_iterable(_markers_in(json.dumps(event)) for event in _ledger_events(ledgers))
+ return MappingProxyType(Counter(marker for marker in found if marker in markers))
+
+
+def _events_for(ledgers: Sequence[Ledger], marker: str) -> tuple[Event, ...]:
+ return tuple(event for event in _ledger_events(ledgers) if marker in json.dumps(event))
+
+
+def _assert_answered(item: _Served) -> None:
+ assert item.status == 200, (item.call, item.text)
+ if item.call.surface != "chat":
+ return
+ if item.call.stream:
+ found: Final = tuple(chain.from_iterable(relayed(chunk) for chunk in chunks_of(item.text)))
+ assert found == ((_DELTA_FIELDS, REFUSAL_FIELDS),), item.text
+ return
+ assert relayed(JSON_OBJECT.validate_json(item.text)) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), item.text
+
+
+def _assert_logged_chat_relays(ledgers: Sequence[Ledger], calls: Sequence[_Call]) -> None:
+ for call in calls:
+ if call.surface != "chat":
+ continue
+ (event,) = _events_for(ledgers, call.marker)
+ assert relayed(event["response"]) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), event
+
+
+def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]:
+ return eventually(
+ lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
+ lambda pids: len(pids) == 2,
+ seconds=30,
+ )
+
+
+def _open_upstream_connections(pid: int, upstream: str) -> int:
+ port: Final = urlsplit(upstream).port
+ return sum(
+ 1
+ for connection in psutil.Process(pid).net_connections(kind="tcp")
+ if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
+ )
+
+
+def _port(wire: Wire) -> int:
+ port: Final = urlsplit(wire.url).port
+ assert port is not None, wire.url
+ return port
+
+
+def _accept(_request: Request) -> Reply:
+ return Reply()
+
+
+async def _timed_chat(gateway: Gateway) -> float:
+ started: Final = time.monotonic()
+ (served,) = await _burst(gateway, (_Call("chat", stream=False, marker=uuid.uuid4().hex),))
+ _assert_answered(served)
+ return time.monotonic() - started
+
+
+async def test_sink_outage_mid_burst_delivers_every_refusal_event_once(gateway: Gateway, tmp_path: Path) -> None:
+ calls: Final = _calls(30)
+ markers: Final = frozenset(call.marker for call in calls)
+ with wire_server(respond) as upstream, ExitStack() as first_sink:
+ sink: Final = first_sink.enter_context(wire_server(_accept))
+ first: Final = Ledger(sink)
+ config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=6)
+ with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
+ _worker_pids(owned)
+ served_first: Final = await _burst(owned.gateway, calls[:10])
+ await asyncio.to_thread(
+ eventually, lambda: sum(_marker_counts((first,), markers).values()), lambda total: total >= 1, 30
+ )
+ first.events()
+ first_sink.close()
+ served_during: Final = await _burst(owned.gateway, calls[10:20])
+ with wire_server(_accept, port=_port(sink)) as restarted:
+ second: Final = Ledger(restarted)
+ served_after: Final = await _burst(owned.gateway, calls[20:])
+ for item in (*served_first, *served_during, *served_after):
+ _assert_answered(item)
+ counts: Final = await asyncio.to_thread(
+ held,
+ lambda: _marker_counts((first, second), markers),
+ lambda found: set(found) == markers and set(found.values()) == {1},
+ holding=3,
+ seconds=60,
+ )
+ assert sum(counts.values()) == 30, counts
+ _assert_logged_chat_relays((first, second), calls)
+ assert len(upstream.drain()) == 30
+
+
+async def test_slow_sink_neither_deadlocks_nor_uploads_a_refusal_event_twice(gateway: Gateway, tmp_path: Path) -> None:
+ calls: Final = _calls(30)
+ markers: Final = frozenset(call.marker for call in calls)
+
+ def slow(_request: Request) -> Reply:
+ time.sleep(_SLOW_SINK_SECONDS)
+ return Reply()
+
+ with wire_server(respond) as upstream, wire_server(slow) as sink:
+ ledger: Final = Ledger(sink)
+ config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0)
+ with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
+ _worker_pids(owned)
+ control: Final = await _timed_chat(owned.gateway)
+ served: Final = await _burst(owned.gateway, calls)
+ for item in served:
+ _assert_answered(item)
+ follow_up: Final = await _timed_chat(owned.gateway)
+ assert_not_delayed_by((follow_up,), control=(control,), injected_delay=_SLOW_SINK_SECONDS)
+ counts: Final = await asyncio.to_thread(
+ held,
+ lambda: _marker_counts((ledger,), markers),
+ lambda found: set(found) == markers and set(found.values()) == {1},
+ holding=5,
+ seconds=90,
+ )
+ assert sum(counts.values()) == 30, counts
+ _assert_logged_chat_relays((ledger,), calls)
+
+
+async def test_worker_sigkill_mid_burst_leaves_the_sibling_relaying_and_logging_once(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ calls: Final = _calls(20)
+ hold: Final = _Held(threading.Event(), SimpleQueue())
+ with wire_server(hold.respond) as upstream, wire_server(_accept) as sink:
+ ledger: Final = Ledger(sink)
+ config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0)
+ with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
+ workers: Final = _worker_pids(owned)
+ burst: Final = asyncio.create_task(_burst(owned.gateway, calls))
+ await asyncio.to_thread(eventually, hold.markers.qsize, lambda size: size == 20, 60)
+ held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, upstream.url) for pid in workers})
+ assert sum(held_by.values()) == 20, held_by
+ victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
+ with owned.replacing_worker(victim_pid) as kill:
+ kill.freeze()
+ await asyncio.to_thread(kill)
+ hold.release.set()
+ served: Final = await burst
+ assert held_by[survivor_pid] >= 10, held_by
+ completed: Final = tuple(item for item in served if item.status == 200)
+ assert len(completed) == held_by[survivor_pid], (held_by, [item.status for item in served])
+ for item in completed:
+ _assert_answered(item)
+ follow_up: Final = _Call("chat", stream=True, marker=uuid.uuid4().hex)
+ (answered,) = await _burst(owned.gateway, (follow_up,))
+ _assert_answered(answered)
+ survivors: Final = frozenset((*(item.call.marker for item in completed), follow_up.marker))
+ counts: Final = await asyncio.to_thread(
+ held,
+ lambda: _marker_counts((ledger,), frozenset(call.marker for call in (*calls, follow_up))),
+ lambda found: set(found) == survivors and set(found.values()) == {1},
+ holding=3,
+ seconds=60,
+ )
+ assert sum(counts.values()) == len(survivors), counts
+ _assert_logged_chat_relays((ledger,), tuple(item.call for item in completed))
+
+
+async def test_proxy_restart_mid_burst_never_logs_a_refusal_twice(
+ gateway: Gateway, tmp_path: Path, record_property: _RecordProperty
+) -> None:
+ calls: Final = _calls(20)
+ markers: Final = frozenset(call.marker for call in calls)
+ hold: Final = _Held(threading.Event(), SimpleQueue())
+ with wire_server(hold.respond) as upstream, wire_server(_accept) as sink:
+ ledger: Final = Ledger(sink)
+ config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=0)
+ with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first:
+ _worker_pids(first)
+ burst: Final = asyncio.create_task(_burst(first.gateway, calls))
+ await asyncio.to_thread(eventually, hold.markers.qsize, lambda size: size == 20, 60)
+ first.begin_shutdown()
+ hold.release.set()
+ served: Final = await burst
+ for item in served:
+ _assert_answered(item)
+ second_directory: Final = tmp_path / "second"
+ second_directory.mkdir()
+ with owned_proxy_process(gateway, second_directory, {}, config=config, workers=2) as second:
+ follow_up: Final = _Call("chat", stream=False, marker=uuid.uuid4().hex)
+ (answered,) = await _burst(second.gateway, (follow_up,))
+ _assert_answered(answered)
+ await asyncio.to_thread(
+ eventually,
+ lambda: _marker_counts((ledger,), frozenset({follow_up.marker})),
+ lambda found: found[follow_up.marker] == 1,
+ 30,
+ )
+ counts: Final = _marker_counts((ledger,), markers)
+ assert set(counts.values()) <= {1}, counts
+ _assert_logged_chat_relays((ledger,), tuple(call for call in calls if call.marker in counts))
+ record_property("served", len(served))
+ record_property("landed", sum(counts.values()))
diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py
new file mode 100644
index 00000000000..461839ccbcc
--- /dev/null
+++ b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py
@@ -0,0 +1,1356 @@
+import asyncio
+import json
+import uuid
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from itertools import chain
+from typing import Final, Literal, TypeAlias
+from urllib.parse import unquote, urlsplit
+
+import anthropic
+import httpx
+import openai
+import pytest
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from integration._support.anthropic_refusal import (
+ CITED_FIELDS,
+ DECLARED_PARTIAL_JSON,
+ DOCS_PARTIAL_TEXT,
+ DOCS_REFUSAL,
+ DOCS_STOP_DETAILS,
+ ENVELOPED_PARTIAL_JSON,
+ FIXTURES,
+ JSON_ANSWER_FORMAT,
+ JSON_TOOL_NAME,
+ MALFORMED_STOP_DETAILS,
+ REFUSAL_FIELDS,
+ UNVALIDATED_USAGE_DETAILS,
+ Ledger,
+ Location,
+ frame_types,
+ frames,
+ hidden_params_probe,
+ openai_id,
+ refusal_events,
+ refusal_prompt,
+ relayed,
+ respond,
+ rewrite_cached_message_fields,
+ rewrite_cached_usage_details,
+ sink_config,
+)
+from integration._support.anthropic_thinking import (
+ BEDROCK_MODEL,
+ JSON_LIST,
+ JSON_OBJECT,
+ MODEL,
+ NO_CACHE,
+ Event,
+ aws_chunks,
+ chunks_of,
+ sse_chunks,
+)
+from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, held
+from integration._support.forward_proxy import tunnelling_forward_proxy
+from integration._support.process import graceful_stop_seconds, owned_proxy
+from integration._support.tls import server_context, write_self_signed_cert
+from integration._support.wire import Reply, Wire, wire_server
+from pydantic import BaseModel, JsonValue
+
+pytestmark: Final = [
+ pytest.mark.timeout(2 * graceful_stop_seconds() + 120),
+ pytest.mark.xdist_group("anthropic_refusal_stop_details_wire"),
+]
+
+_ANTHROPIC_BASE: Final = "http://api.anthropic.com"
+_ANTHROPIC_KEY: Final = "scripted-anthropic-key"
+_VERTEX_PROJECT: Final = "scripted-project"
+_VERTEX_LOCATION: Final = "us-east5"
+_VERTEX_HOST: Final = f"{_VERTEX_LOCATION}-aiplatform.googleapis.com"
+_VERTEX_MODEL_PATH: Final = (
+ f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{MODEL}"
+)
+_CHAT_FIELDS: Final[Location] = ("choices", 0, "message", "provider_specific_fields")
+_DELTA_FIELDS: Final[Location] = ("choices", 0, "delta", "provider_specific_fields")
+_RESPONSE_FIELDS: Final[Location] = ("provider_specific_fields",)
+_STREAMED_RESPONSE_FIELDS: Final[Location] = ("response", "provider_specific_fields")
+
+Client: TypeAlias = Literal["sdk", "async_sdk", "httpx"]
+Provider: TypeAlias = Literal[
+ "anthropic",
+ "bedrock_invoke",
+ "claude_platform",
+ "vertex_ai",
+ "azure_ai",
+ "bedrock_mantle",
+ "snowflake",
+ "openai_compatible",
+]
+Surface: TypeAlias = Literal["chat", "chat_stream", "messages", "messages_stream", "responses", "responses_stream"]
+
+CLIENTS: Final[tuple[Client, ...]] = ("sdk", "async_sdk", "httpx")
+FAMILY: Final[tuple[Provider, ...]] = (
+ "bedrock_invoke",
+ "claude_platform",
+ "vertex_ai",
+ "azure_ai",
+ "bedrock_mantle",
+ "snowflake",
+)
+_STREAM_TARGETS: Final[Mapping[Provider, tuple[str, str]]] = {
+ "anthropic": ("/v1/messages", "/v1/messages"),
+ "bedrock_invoke": (f"/model/{BEDROCK_MODEL}/invoke", f"/model/{BEDROCK_MODEL}/invoke-with-response-stream"),
+ "claude_platform": ("/v1/messages", "/v1/messages"),
+ "vertex_ai": (f"{_VERTEX_MODEL_PATH}:rawPredict", f"{_VERTEX_MODEL_PATH}:streamRawPredict"),
+ "azure_ai": ("/anthropic/v1/messages", "/anthropic/v1/messages"),
+ "bedrock_mantle": ("/anthropic/v1/messages", "/anthropic/v1/messages"),
+ "snowflake": ("/api/v2/cortex/v1/messages", "/api/v2/cortex/v1/messages"),
+}
+
+
+@dataclass(frozen=True, slots=True)
+class Rig:
+ gateway: Gateway
+ wire: Wire
+ vertex: Wire
+ sink: Ledger
+ forbidden: Ledger
+ absent: Ledger
+ hidden: Ledger
+ credentials: str
+
+
+def _service_account_json(token_url: str) -> str:
+ private_key: Final = (
+ rsa.generate_private_key(public_exponent=65537, key_size=2048)
+ .private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption())
+ .decode()
+ )
+ return json.dumps(
+ {
+ "type": "service_account",
+ "project_id": _VERTEX_PROJECT,
+ "private_key_id": "scripted",
+ "private_key": private_key,
+ "client_email": f"scripted@{_VERTEX_PROJECT}.iam.gserviceaccount.com",
+ "client_id": "0",
+ "auth_uri": f"{token_url}/_oauth/authorize",
+ "token_uri": f"{token_url}/_oauth/token",
+ }
+ )
+
+
+def _port(wire: Wire) -> int:
+ port: Final = urlsplit(wire.url).port
+ assert port is not None, wire.url
+ return port
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+ directory: Final = tmp_path_factory.mktemp("anthropic-refusal-stop-details")
+ certificate: Final = write_self_signed_cert(directory, (_VERTEX_HOST,))
+ with (
+ gateway_from_environment() as parent,
+ wire_server(respond) as wire,
+ wire_server(respond, tls=server_context(*certificate)) as vertex,
+ tunnelling_forward_proxy(_port(vertex)) as tunnel,
+ wire_server(lambda _request: Reply()) as sink,
+ wire_server(lambda _request: Reply(status=403)) as forbidden,
+ wire_server(lambda _request: Reply(status=404)) as absent,
+ wire_server(lambda _request: Reply()) as hidden,
+ ):
+ credentials: Final = _service_account_json(parent.upstream_url.rstrip("/"))
+ credentials_path: Final = directory / "vertex-service-account.json"
+ credentials_path.write_text(credentials)
+ config: Final = sink_config(
+ directory,
+ {"refusal_sink": sink.url, "refusal_sink_403": forbidden.url, "refusal_sink_404": absent.url},
+ extra_callbacks=(hidden_params_probe(directory, hidden.url),),
+ store_prompts=True,
+ )
+ overrides: Final = {
+ "ANTHROPIC_API_BASE": _ANTHROPIC_BASE,
+ "ANTHROPIC_API_KEY": _ANTHROPIC_KEY,
+ "AIOHTTP_TRUST_ENV": "True",
+ "HTTP_PROXY": wire.url,
+ "http_proxy": wire.url,
+ "HTTPS_PROXY": tunnel.url,
+ "https_proxy": tunnel.url,
+ "NO_PROXY": "127.0.0.1,localhost",
+ "no_proxy": "127.0.0.1,localhost",
+ "SSL_CERT_FILE": str(certificate[0]),
+ "DEFAULT_VERTEXAI_PROJECT": _VERTEX_PROJECT,
+ "DEFAULT_VERTEXAI_LOCATION": _VERTEX_LOCATION,
+ "DEFAULT_GOOGLE_APPLICATION_CREDENTIALS": str(credentials_path),
+ }
+ with owned_proxy(parent, directory, overrides, config=config, workers=2) as owned:
+ yield Rig(owned, wire, vertex, Ledger(sink), Ledger(forbidden), Ledger(absent), Ledger(hidden), credentials)
+
+
+@pytest.fixture(autouse=True)
+def _drained_upstreams(rig: Rig) -> None:
+ rig.wire.drain()
+ rig.vertex.drain()
+
+
+def _marker() -> str:
+ return uuid.uuid4().hex
+
+
+def _deployment(scenario: Scenario, provider: Provider, rig: Rig) -> str:
+ wire_url: Final = rig.wire.url
+ match provider:
+ case "anthropic":
+ return scenario.model(model=f"anthropic/{MODEL}", api_base=wire_url, api_key=_ANTHROPIC_KEY)
+ case "bedrock_invoke":
+ return scenario.model(
+ model=f"bedrock/invoke/{BEDROCK_MODEL}",
+ api_base=wire_url,
+ aws_access_key_id="AKIASCRIPTEDPROVIDER",
+ aws_secret_access_key="scripted-secret",
+ aws_region_name="us-east-1",
+ aws_bedrock_runtime_endpoint=wire_url,
+ )
+ case "claude_platform":
+ return scenario.model(
+ model=f"bedrock/claude_platform/{MODEL}",
+ api_base=wire_url,
+ api_key="scripted-platform-key",
+ aws_region_name="us-east-1",
+ workspace_id="scripted-workspace",
+ )
+ case "vertex_ai":
+ return scenario.model(
+ model=f"vertex_ai/{MODEL}",
+ api_base=f"{wire_url}{_VERTEX_MODEL_PATH}",
+ api_key=None,
+ vertex_project=_VERTEX_PROJECT,
+ vertex_location=_VERTEX_LOCATION,
+ vertex_credentials=rig.credentials,
+ )
+ case "azure_ai":
+ return scenario.model(model=f"azure_ai/{MODEL}", api_base=wire_url, api_key="scripted-azure-key")
+ case "bedrock_mantle":
+ return scenario.model(
+ model=f"bedrock_mantle/{BEDROCK_MODEL}",
+ api_base=wire_url,
+ api_key="scripted-mantle-key",
+ aws_region_name="us-east-1",
+ )
+ case "snowflake":
+ return scenario.model(model=f"snowflake/{MODEL}", api_base=wire_url, api_key="scripted-snowflake-key")
+ case "openai_compatible":
+ return scenario.model(model="openai/gpt-4o-mini", api_base=f"{wire_url}/v1", api_key="scripted-openai-key")
+
+
+def _base_url(rig: Rig) -> str:
+ return str(rig.gateway.client.base_url).rstrip("/")
+
+
+def _dump(model: BaseModel) -> Event:
+ return JSON_OBJECT.validate_python(model.model_dump())
+
+
+def _post(rig: Rig, path: str, body: Mapping[str, JsonValue], key: str) -> bytes:
+ response: Final = rig.gateway.request("POST", path, body, key=key)
+ assert response.status_code == 200, response.text
+ return response.content
+
+
+def _chat_request(model: str, text: str, *, stream: bool, cached: bool) -> Event:
+ request: Final[Event] = {
+ "model": model,
+ "messages": [{"role": "user", "content": text}],
+ "stream": stream,
+ "max_tokens": 64,
+ }
+ return request if cached else {**request, **NO_CACHE}
+
+
+def _chat(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event:
+ match client:
+ case "httpx":
+ return JSON_OBJECT.validate_json(
+ _post(rig, "/v1/chat/completions", _chat_request(model, text, stream=False, cached=cached), key)
+ )
+ case "sdk":
+ with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk:
+ return _dump(
+ sdk.chat.completions.create(
+ model=model, messages=[{"role": "user", "content": text}], max_tokens=64, extra_body=NO_CACHE
+ )
+ )
+ case "async_sdk":
+
+ async def call() -> Event:
+ async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk:
+ return _dump(
+ await sdk.chat.completions.create(
+ model=model,
+ messages=[{"role": "user", "content": text}],
+ max_tokens=64,
+ extra_body=NO_CACHE,
+ )
+ )
+
+ return asyncio.run(call())
+
+
+def _chat_stream(
+ rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False
+) -> tuple[Event, ...]:
+ match client:
+ case "httpx":
+ return chunks_of(
+ _post(rig, "/v1/chat/completions", _chat_request(model, text, stream=True, cached=cached), key).decode()
+ )
+ case "sdk":
+ with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk:
+ return tuple(
+ _dump(chunk)
+ for chunk in sdk.chat.completions.create(
+ model=model,
+ messages=[{"role": "user", "content": text}],
+ max_tokens=64,
+ stream=True,
+ extra_body=NO_CACHE,
+ )
+ )
+ case "async_sdk":
+
+ async def call() -> tuple[Event, ...]:
+ async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk:
+ stream: Final = await sdk.chat.completions.create(
+ model=model,
+ messages=[{"role": "user", "content": text}],
+ max_tokens=64,
+ stream=True,
+ extra_body=NO_CACHE,
+ )
+ return tuple([_dump(chunk) async for chunk in stream])
+
+ return asyncio.run(call())
+
+
+def _messages_request(model: str, text: str, *, stream: bool, cached: bool) -> Event:
+ request: Final[Event] = {
+ "model": model,
+ "max_tokens": 64,
+ "stream": stream,
+ "messages": [{"role": "user", "content": text}],
+ }
+ return request if cached else {**request, **NO_CACHE}
+
+
+def _messages(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event:
+ match client:
+ case "httpx":
+ return JSON_OBJECT.validate_json(
+ _post(rig, "/v1/messages", _messages_request(model, text, stream=False, cached=cached), key)
+ )
+ case "sdk":
+ with anthropic.Anthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk:
+ return _dump(
+ sdk.messages.create(
+ model=model, max_tokens=64, messages=[{"role": "user", "content": text}], extra_body=NO_CACHE
+ )
+ )
+ case "async_sdk":
+
+ async def call() -> Event:
+ async with anthropic.AsyncAnthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk:
+ return _dump(
+ await sdk.messages.create(
+ model=model,
+ max_tokens=64,
+ messages=[{"role": "user", "content": text}],
+ extra_body=NO_CACHE,
+ )
+ )
+
+ return asyncio.run(call())
+
+
+def _messages_stream(
+ rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False
+) -> tuple[Event, ...]:
+ match client:
+ case "httpx":
+ return frames(
+ _post(rig, "/v1/messages", _messages_request(model, text, stream=True, cached=cached), key).decode()
+ )
+ case "sdk":
+ with anthropic.Anthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk:
+ return tuple(
+ _dump(event)
+ for event in sdk.messages.create(
+ model=model,
+ max_tokens=64,
+ messages=[{"role": "user", "content": text}],
+ stream=True,
+ extra_body=NO_CACHE,
+ )
+ )
+ case "async_sdk":
+
+ async def call() -> tuple[Event, ...]:
+ async with anthropic.AsyncAnthropic(base_url=_base_url(rig), api_key=key, max_retries=0) as sdk:
+ stream: Final = await sdk.messages.create(
+ model=model,
+ max_tokens=64,
+ messages=[{"role": "user", "content": text}],
+ stream=True,
+ extra_body=NO_CACHE,
+ )
+ return tuple([_dump(event) async for event in stream])
+
+ return asyncio.run(call())
+
+
+def _responses_request(model: str, text: str, *, stream: bool, cached: bool) -> Event:
+ request: Final[Event] = {"model": model, "input": text, "max_output_tokens": 64, "stream": stream}
+ return request if cached else {**request, **NO_CACHE}
+
+
+def _responses(rig: Rig, key: str, model: str, text: str, client: Client = "httpx", *, cached: bool = False) -> Event:
+ match client:
+ case "httpx":
+ return JSON_OBJECT.validate_json(
+ _post(rig, "/v1/responses", _responses_request(model, text, stream=False, cached=cached), key)
+ )
+ case "sdk":
+ with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk:
+ return _dump(sdk.responses.create(model=model, input=text, max_output_tokens=64, extra_body=NO_CACHE))
+ case "async_sdk":
+
+ async def call() -> Event:
+ async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk:
+ return _dump(
+ await sdk.responses.create(model=model, input=text, max_output_tokens=64, extra_body=NO_CACHE)
+ )
+
+ return asyncio.run(call())
+
+
+def _responses_stream(rig: Rig, key: str, model: str, text: str, client: Client = "httpx") -> tuple[Event, ...]:
+ match client:
+ case "httpx":
+ return frames(
+ _post(rig, "/v1/responses", _responses_request(model, text, stream=True, cached=False), key).decode()
+ )
+ case "sdk":
+ with openai.OpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk:
+ return tuple(
+ _dump(event)
+ for event in sdk.responses.create(
+ model=model, input=text, max_output_tokens=64, stream=True, extra_body=NO_CACHE
+ )
+ )
+ case "async_sdk":
+
+ async def call() -> tuple[Event, ...]:
+ async with openai.AsyncOpenAI(base_url=f"{_base_url(rig)}/v1", api_key=key, max_retries=0) as sdk:
+ stream: Final = await sdk.responses.create(
+ model=model, input=text, max_output_tokens=64, stream=True, extra_body=NO_CACHE
+ )
+ return tuple([_dump(event) async for event in stream])
+
+ return asyncio.run(call())
+
+
+def _relayed_in(events: Sequence[JsonValue]) -> tuple[tuple[Location, Event], ...]:
+ return tuple(chain.from_iterable(relayed(event) for event in events))
+
+
+def _choice(chunk: Mapping[str, JsonValue]) -> Event:
+ choices: Final = JSON_LIST.validate_python(chunk.get("choices") or [])
+ return JSON_OBJECT.validate_python(choices[0]) if choices else {}
+
+
+def _finish_reasons(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]:
+ return tuple(_choice(chunk)["finish_reason"] for chunk in chunks if _choice(chunk).get("finish_reason"))
+
+
+def _streamed_text(chunks: Sequence[Mapping[str, JsonValue]]) -> str:
+ return "".join(
+ str(JSON_OBJECT.validate_python(_choice(chunk).get("delta") or {}).get("content") or "") for chunk in chunks
+ )
+
+
+def _message_text(body: Mapping[str, JsonValue]) -> JsonValue:
+ return JSON_OBJECT.validate_python(_choice(body)["message"]).get("content")
+
+
+def _native_stops(events: Sequence[Mapping[str, JsonValue]]) -> tuple[Event, ...]:
+ deltas: Final = tuple(
+ JSON_OBJECT.validate_python(event["delta"]) for event in events if event.get("type") == "message_delta"
+ )
+ return tuple({name: delta.get(name) for name in ("stop_reason", "stop_details")} for delta in deltas)
+
+
+def _native_deltas(events: Sequence[Mapping[str, JsonValue]]) -> tuple[Event, ...]:
+ deltas: Final = tuple(
+ JSON_OBJECT.validate_python(event["delta"]) for event in events if event.get("type") == "message_delta"
+ )
+ return tuple({name: delta[name] for name in ("stop_reason", "stop_details") if name in delta} for delta in deltas)
+
+
+def _targets(wire: Wire) -> tuple[str, ...]:
+ return tuple(unquote(request.target) for request in wire.drain())
+
+
+def _paths(wire: Wire) -> tuple[str, ...]:
+ return tuple(urlsplit(target).path for target in _targets(wire))
+
+
+def _logged(rig: Rig, key: str, count: int = 1) -> tuple[Event, ...]:
+ events: Final = rig.sink.await_key(key, count)
+ assert len(events) == count, events
+ return events
+
+
+def _logged_response(rig: Rig, key: str) -> JsonValue:
+ return _logged(rig, key)[0]["response"]
+
+
+@pytest.mark.parametrize("client", CLIENTS)
+def test_chat_refusal_relays_the_docs_stop_details_to_the_caller_and_the_log(rig: Rig, client: Client) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker()), client)
+ assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body
+ assert _choice(body)["finish_reason"] == "content_filter", body
+ assert _paths(rig.wire) == ("/v1/messages",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("client", CLIENTS)
+def test_chat_stream_refusal_relays_the_docs_stop_details_in_one_chunk_and_the_log(rig: Rig, client: Client) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ chunks: Final = _chat_stream(rig, key, model, refusal_prompt("refusal", _marker()), client)
+ assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks
+ assert _finish_reasons(chunks) == ("content_filter",), chunks
+ assert _streamed_text(chunks) == "", chunks
+ assert _paths(rig.wire) == ("/v1/messages",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("client", CLIENTS)
+def test_messages_refusal_returns_the_docs_body_and_logs_the_stop_details(rig: Rig, client: Client) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _messages(rig, key, model, refusal_prompt("refusal", _marker()), client)
+ assert (body["stop_reason"], body["stop_details"], body["content"]) == ("refusal", DOCS_STOP_DETAILS, []), body
+ assert relayed(body) == (), body
+ assert _paths(rig.wire) == ("/v1/messages",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("client", CLIENTS)
+def test_messages_stream_refusal_keeps_the_client_frames_and_logs_the_stop_details(rig: Rig, client: Client) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ events: Final = _messages_stream(rig, key, model, refusal_prompt("refusal", _marker()), client)
+ assert frame_types(events) == frame_types(FIXTURES["refusal"].events), events
+ assert _native_stops(events) == (REFUSAL_FIELDS,), events
+ assert _relayed_in(events) == (), events
+ assert _paths(rig.wire) == ("/v1/messages",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("client", CLIENTS)
+def test_responses_refusal_relays_the_docs_stop_details_at_the_top_level(rig: Rig, client: Client) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _responses(rig, key, model, refusal_prompt("refusal", _marker()), client)
+ assert relayed(body) == ((_RESPONSE_FIELDS, REFUSAL_FIELDS),), body
+ assert (body["status"], body["incomplete_details"]) == ("incomplete", {"reason": "content_filter"}), body
+ assert _paths(rig.wire) == ("/v1/messages",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("client", CLIENTS)
+def test_responses_stream_refusal_relays_the_docs_stop_details_on_the_incomplete_event(
+ rig: Rig, client: Client
+) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ events: Final = _responses_stream(rig, key, model, refusal_prompt("refusal", _marker()), client)
+ assert _relayed_in(events) == ((_STREAMED_RESPONSE_FIELDS, REFUSAL_FIELDS),), events
+ assert frame_types(events)[-1] == "response.incomplete", events
+ assert _paths(rig.wire) == ("/v1/messages",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"])
+def test_anthropic_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None:
+ with rig.gateway.scenario() as scenario:
+ key: Final = scenario.key()
+ body: Final = {
+ "model": MODEL,
+ "max_tokens": 64,
+ "stream": stream,
+ "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}],
+ }
+ content: Final = _post(rig, "/anthropic/v1/messages", body, key)
+ if stream:
+ assert content == b"".join(sse_chunks(refusal_events(DOCS_REFUSAL))), content
+ else:
+ assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content
+ assert _targets(rig.wire) == (f"{_ANTHROPIC_BASE}/v1/messages",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("provider", FAMILY)
+def test_family_chat_refusal_relays_the_docs_stop_details(rig: Rig, provider: Provider) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, provider, rig)
+ key: Final = scenario.key()
+ body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker()))
+ expected: Final = () if provider == "snowflake" else ((_CHAT_FIELDS, REFUSAL_FIELDS),)
+ assert relayed(body) == expected, body
+ assert _choice(body)["finish_reason"] == ("stop" if provider == "snowflake" else "content_filter"), body
+ assert _paths(rig.wire) == (_STREAM_TARGETS[provider][0],)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == expected, logged
+
+
+@pytest.mark.parametrize("provider", FAMILY)
+def test_family_chat_stream_refusal_relays_the_docs_stop_details(rig: Rig, provider: Provider) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, provider, rig)
+ key: Final = scenario.key()
+ chunks: Final = _chat_stream(rig, key, model, refusal_prompt("refusal", _marker()))
+ assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks
+ assert _finish_reasons(chunks) == ("content_filter",), chunks
+ assert _paths(rig.wire) == (_STREAM_TARGETS[provider][1],)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("provider", FAMILY)
+def test_family_messages_stream_refusal_keeps_frames_and_logs_the_stop_details(rig: Rig, provider: Provider) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, provider, rig)
+ key: Final = scenario.key()
+ events: Final = _messages_stream(rig, key, model, refusal_prompt("refusal", _marker()))
+ native: Final = {"stop_reason": "refusal"} if provider == "snowflake" else REFUSAL_FIELDS
+ assert _native_deltas(events) == (native,), events
+ assert _relayed_in(events) == (), events
+ assert _paths(rig.wire) == (_STREAM_TARGETS[provider][1],)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+def _bedrock_passthrough_deployment(scenario: Scenario, rig: Rig) -> str:
+ return scenario.model(
+ model=f"bedrock/{BEDROCK_MODEL}",
+ api_base=rig.wire.url,
+ aws_access_key_id="AKIASCRIPTEDPROVIDER",
+ aws_secret_access_key="scripted-secret",
+ aws_region_name="us-east-1",
+ aws_bedrock_runtime_endpoint=rig.wire.url,
+ )
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["invoke", "invoke-with-response-stream"])
+def test_bedrock_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None:
+ with rig.gateway.scenario() as scenario:
+ deployment: Final = _bedrock_passthrough_deployment(scenario, rig)
+ key: Final = scenario.key()
+ action: Final = "invoke-with-response-stream" if stream else "invoke"
+ body: Final = {
+ "anthropic_version": "bedrock-2023-05-31",
+ "max_tokens": 64,
+ "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}],
+ }
+ content: Final = _post(rig, f"/bedrock/model/{deployment}/{action}", body, key)
+ if stream:
+ assert content == b"".join(aws_chunks(refusal_events(DOCS_REFUSAL))), content
+ else:
+ assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content
+ assert _paths(rig.wire) == (f"/model/{BEDROCK_MODEL}/{action}",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["rawPredict", "streamRawPredict"])
+def test_vertex_passthrough_refusal_relays_upstream_bytes_and_logs_the_stop_details(rig: Rig, stream: bool) -> None:
+ with rig.gateway.scenario() as scenario:
+ key: Final = scenario.key()
+ action: Final = "streamRawPredict" if stream else "rawPredict"
+ body: Final = {
+ "anthropic_version": "vertex-2023-10-16",
+ "max_tokens": 64,
+ "stream": stream,
+ "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}],
+ }
+ content: Final = _post(rig, f"/vertex_ai{_VERTEX_MODEL_PATH}:{action}", body, key)
+ if stream:
+ assert content == b"".join(sse_chunks(refusal_events(DOCS_REFUSAL))), content
+ else:
+ assert JSON_OBJECT.validate_json(content) == DOCS_REFUSAL, content
+ assert _paths(rig.vertex) == (f"{_VERTEX_MODEL_PATH}:{action}",)
+ assert _targets(rig.wire) == ()
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+def _surface_call(rig: Rig, surface: Surface, key: str, model: str, text: str) -> JsonValue:
+ match surface:
+ case "chat":
+ return _chat(rig, key, model, text, cached=True)
+ case "chat_stream":
+ return list(_chat_stream(rig, key, model, text, cached=True))
+ case "messages":
+ return _messages(rig, key, model, text, cached=True)
+ case "messages_stream":
+ return list(_messages_stream(rig, key, model, text, cached=True))
+ case "responses":
+ return _responses(rig, key, model, text, cached=True)
+ case "responses_stream":
+ return list(_responses_stream(rig, key, model, text))
+
+
+@dataclass(frozen=True, slots=True)
+class CacheTwin:
+ fields: tuple[tuple[Location, Event], ...]
+ served_from_cache: tuple[bool, bool]
+
+
+_CACHE_TWINS: Final[Mapping[Surface, CacheTwin]] = {
+ "chat": CacheTwin(((_CHAT_FIELDS, REFUSAL_FIELDS),), (False, True)),
+ "chat_stream": CacheTwin(((_DELTA_FIELDS, REFUSAL_FIELDS),), (False, True)),
+ "messages": CacheTwin((), (False, False)),
+ "messages_stream": CacheTwin((), (False, False)),
+}
+
+
+def _relayed_anywhere(value: JsonValue) -> tuple[tuple[Location, Event], ...]:
+ return _relayed_in(JSON_LIST.validate_python(value)) if isinstance(value, list) else relayed(value)
+
+
+@pytest.mark.parametrize("surface", tuple(_CACHE_TWINS))
+def test_cache_hit_twin_carries_the_same_stop_details_as_the_miss(rig: Rig, surface: Surface) -> None:
+ twin: Final = _CACHE_TWINS[surface]
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ text: Final = refusal_prompt("refusal", _marker())
+ first: Final = _surface_call(rig, surface, key, model, text)
+ _logged(rig, key)
+ second: Final = _surface_call(rig, surface, key, model, text)
+ logged: Final = _logged(rig, key, 2)
+ upstream: Final = _paths(rig.wire)
+ served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged)
+ evidence: Final = {"cache_hit": served_from_cache, "upstream": upstream, "first": first, "second": second}
+ assert (served_from_cache, len(upstream)) == (twin.served_from_cache, twin.served_from_cache.count(False)), (
+ evidence
+ )
+ assert (_relayed_anywhere(first), _relayed_anywhere(second)) == (twin.fields, twin.fields), evidence
+ assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, (
+ served_from_cache,
+ logged,
+ )
+
+
+def test_responses_cache_hit_logs_the_same_stop_details_as_the_miss(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ text: Final = refusal_prompt("refusal", _marker())
+ first: Final = _surface_call(rig, "responses", key, model, text)
+ _logged(rig, key)
+ _surface_call(rig, "responses", key, model, text)
+ logged: Final = _logged(rig, key, 2)
+ served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged)
+ assert (len(_paths(rig.wire)), served_from_cache.count(True)) == (1, 1), (served_from_cache, logged)
+ assert relayed(first) == ((_RESPONSE_FIELDS, REFUSAL_FIELDS),), first
+ assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, (
+ served_from_cache,
+ logged,
+ )
+
+
+def test_passthrough_is_not_served_from_the_response_cache(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ key: Final = scenario.key()
+ body: Final = {
+ "model": MODEL,
+ "max_tokens": 64,
+ "messages": [{"role": "user", "content": refusal_prompt("refusal", _marker())}],
+ }
+ first: Final = _post(rig, "/anthropic/v1/messages", body, key)
+ second: Final = _post(rig, "/anthropic/v1/messages", body, key)
+ assert JSON_OBJECT.validate_json(first) == JSON_OBJECT.validate_json(second) == DOCS_REFUSAL
+ assert _targets(rig.wire) == (f"{_ANTHROPIC_BASE}/v1/messages",) * 2
+ logged: Final = _logged(rig, key, 2)
+ assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2
+
+
+@pytest.mark.parametrize("variant", tuple(MALFORMED_STOP_DETAILS))
+def test_malformed_stop_details_are_relayed_as_sent_without_breaking_the_turn(rig: Rig, variant: str) -> None:
+ value: Final = MALFORMED_STOP_DETAILS[variant]
+ expected: Final[Event] = {"stop_reason": "end_turn", "stop_details": value}
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _chat(rig, key, model, refusal_prompt(variant, _marker()))
+ assert relayed(body) == ((_CHAT_FIELDS, expected),), body
+ assert (_choice(body)["finish_reason"], _message_text(body)) == ("stop", "ready"), body
+ chunks: Final = _chat_stream(rig, key, model, refusal_prompt(variant, _marker()))
+ assert _relayed_in(chunks) == ((_DELTA_FIELDS, expected),), chunks
+ assert (_finish_reasons(chunks), _streamed_text(chunks)) == (("stop",), "ready"), chunks
+ assert len(_paths(rig.wire)) == 2
+ logged: Final = _logged(rig, key, 2)
+ assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, expected),),) * 2, logged
+ assert rig.gateway.client.get("/health/liveliness").status_code == 200
+
+
+def test_duplicated_stop_details_key_relays_the_docs_object_once(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _chat(rig, key, model, refusal_prompt("duplicate", _marker()))
+ assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+def test_non_object_upstream_body_fails_the_call_and_leaves_other_keys_serving(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ broken_key: Final = scenario.key()
+ healthy_key: Final = scenario.key()
+ broken: Final = rig.gateway.request(
+ "POST",
+ "/v1/chat/completions",
+ _chat_request(model, refusal_prompt("nondict", _marker()), stream=False, cached=False),
+ key=broken_key,
+ )
+ assert broken.status_code == 500, broken.text
+ assert "error" in JSON_OBJECT.validate_json(broken.content), broken.text
+ healthy: Final = _chat(rig, healthy_key, model, refusal_prompt("refusal", _marker()))
+ assert relayed(healthy) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), healthy
+ assert len(_paths(rig.wire)) == 2
+ failure: Final = _logged(rig, broken_key)[0]
+ assert (failure["status"], relayed(failure["response"])) == ("failure", ()), failure
+
+
+@pytest.mark.parametrize(
+ ("fixture", "status", "message"),
+ [("unauthorized", 401, "invalid x-api-key"), ("invalid", 400, "max_tokens: Field required")],
+)
+@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"])
+def test_provider_errors_reach_the_caller_and_log_a_failure_without_stop_fields(
+ rig: Rig, fixture: str, status: int, message: str, stream: bool
+) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ response: Final = rig.gateway.request(
+ "POST",
+ "/v1/chat/completions",
+ _chat_request(model, refusal_prompt(fixture, _marker()), stream=stream, cached=False),
+ key=key,
+ )
+ assert response.status_code == status, response.text
+ assert message in response.text, response.text
+ assert _paths(rig.wire) == ("/v1/messages",)
+ failure: Final = _logged(rig, key)[0]
+ assert (failure["status"], relayed(failure["response"])) == ("failure", ()), failure
+
+
+def test_failing_sinks_get_one_attempt_each_and_the_proxy_stays_healthy(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker()))
+ assert relayed(body) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), body
+ for ledger in (rig.sink, rig.forbidden, rig.absent):
+ events = ledger.await_key(key)
+ assert tuple(relayed(event["response"]) for event in events) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),)
+ held(
+ lambda: tuple(len(ledger.for_key(key)) for ledger in (rig.sink, rig.forbidden, rig.absent)),
+ lambda counts: counts == (1, 1, 1),
+ holding=3,
+ )
+ readiness: Final = rig.gateway.client.get("/health/readiness")
+ assert (readiness.status_code, readiness.json()["status"]) == (200, "healthy"), readiness.text
+
+
+def test_unauthenticated_refusal_request_never_reaches_the_upstream_or_the_sink(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ stranger: Final = "sk-not-a-virtual-key-" + _marker()
+ response: Final = rig.gateway.request(
+ "POST",
+ "/v1/chat/completions",
+ _chat_request(model, refusal_prompt("refusal", _marker()), stream=False, cached=False),
+ key=stranger,
+ )
+ assert response.status_code == 401, response.text
+ assert _targets(rig.wire) == ()
+ assert rig.sink.for_key(stranger) == ()
+
+
+@pytest.mark.parametrize("fixture", ["null", "missing"])
+def test_normal_turn_adds_no_stop_fields_on_any_surface(rig: Rig, fixture: str) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _chat(rig, key, model, refusal_prompt(fixture, _marker()))
+ assert relayed(body) == (), body
+ assert (_choice(body)["finish_reason"], _message_text(body)) == ("stop", "ready"), body
+ chunks: Final = _chat_stream(rig, key, model, refusal_prompt(fixture, _marker()))
+ assert _relayed_in(chunks) == (), chunks
+ assert (_finish_reasons(chunks), _streamed_text(chunks)) == (("stop",), "ready"), chunks
+ events: Final = _messages_stream(rig, key, model, refusal_prompt(fixture, _marker()))
+ expected_native: Final = {"stop_reason": "end_turn", "stop_details": None}
+ assert _native_stops(events) == (expected_native,), events
+ assert len(_paths(rig.wire)) == 3
+ logged: Final = _logged(rig, key, 3)
+ assert tuple(relayed(event["response"]) for event in logged) == ((), (), ()), logged
+
+
+def test_refusal_after_partial_text_keeps_the_text_and_relays_the_stop_details(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ chat_key: Final = scenario.key()
+ messages_key: Final = scenario.key()
+ chunks: Final = _chat_stream(rig, chat_key, model, refusal_prompt("partial", _marker()))
+ assert _streamed_text(chunks) == "Hello..", chunks
+ assert _relayed_in(chunks) == ((_DELTA_FIELDS, REFUSAL_FIELDS),), chunks
+ assert _finish_reasons(chunks) == ("content_filter",), chunks
+ logged_chat: Final = _logged_response(rig, chat_key)
+ assert relayed(logged_chat) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged_chat
+ assert _message_text(JSON_OBJECT.validate_python(logged_chat)) == "Hello..", logged_chat
+ events: Final = _messages_stream(rig, messages_key, model, refusal_prompt("partial", _marker()))
+ assert _native_stops(events) == (REFUSAL_FIELDS,), events
+ assert frame_types(events) == frame_types(FIXTURES["partial"].events), events
+ logged_messages: Final = _logged_response(rig, messages_key)
+ assert relayed(logged_messages) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged_messages
+ assert len(_paths(rig.wire)) == 2
+
+
+def test_repeated_identical_refusals_log_once_per_request(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ text: Final = refusal_prompt("refusal", _marker())
+ bodies: Final = (_chat(rig, key, model, text), _chat(rig, key, model, text))
+ assert tuple(relayed(body) for body in bodies) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, bodies
+ assert len(_paths(rig.wire)) == 2
+ logged: Final = _logged(rig, key, 2)
+ assert len({event["litellm_call_id"] for event in logged}) == 2, logged
+ assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2
+
+
+@pytest.mark.parametrize("fixture", ["refusal", "int", "null"])
+def test_chat_stream_logged_response_keeps_the_received_model_id(rig: Rig, fixture: str) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ chunks: Final = _chat_stream(rig, key, model, refusal_prompt(fixture, _marker()))
+ assert len({chunk["id"] for chunk in chunks}) == 1, chunks
+ probed: Final = rig.hidden.await_key(key)
+ assert len(probed) == 1, probed
+ assert (probed[0]["response_id"], probed[0]["received_model_id"]) == (chunks[0]["id"], chunks[0]["id"]), probed
+
+
+def _delta_fields(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]:
+ deltas: Final = tuple(JSON_OBJECT.validate_python(_choice(chunk).get("delta") or {}) for chunk in chunks)
+ return tuple(
+ delta["provider_specific_fields"] for delta in deltas if delta.get("provider_specific_fields") is not None
+ )
+
+
+def _message_fields(response: JsonValue) -> tuple[JsonValue, ...]:
+ message: Final = JSON_OBJECT.validate_python(_choice(JSON_OBJECT.validate_python(response))["message"])
+ fields: Final = message.get("provider_specific_fields")
+ return () if fields is None else (fields,)
+
+
+def test_openai_compatible_stream_logged_response_keeps_the_received_model_id(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "openai_compatible", rig)
+ key: Final = scenario.key()
+ marker: Final = _marker()
+ chunks: Final = _chat_stream(rig, key, model, refusal_prompt("plain", marker))
+ assert (_streamed_text(chunks), _finish_reasons(chunks), _delta_fields(chunks)) == ("ready now", ("stop",), ())
+ assert {chunk["id"] for chunk in chunks} == {openai_id(marker)}, chunks
+ assert _paths(rig.wire) == ("/v1/chat/completions",)
+ probed: Final = rig.hidden.await_key(key)
+ assert len(probed) == 1, probed
+ assert (probed[0]["response_id"], probed[0]["received_model_id"]) == (openai_id(marker),) * 2, probed
+
+
+@pytest.mark.parametrize(("fixture", "fields"), [("plain", ()), ("cited", (CITED_FIELDS,))], ids=["plain", "cited"])
+def test_openai_compatible_stream_cache_hit_replays_the_provider_fields_of_the_miss(
+ rig: Rig, fixture: str, fields: tuple[JsonValue, ...]
+) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "openai_compatible", rig)
+ key: Final = scenario.key()
+ text: Final = refusal_prompt(fixture, _marker())
+ first: Final = _chat_stream(rig, key, model, text, cached=True)
+ _logged(rig, key)
+ second: Final = _chat_stream(rig, key, model, text, cached=True)
+ logged: Final = _logged(rig, key, 2)
+ served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged)
+ evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged}
+ assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence
+ streamed: Final = tuple(
+ (_streamed_text(chunks), _finish_reasons(chunks), _delta_fields(chunks)) for chunks in (first, second)
+ )
+ assert streamed == (("ready now", ("stop",), fields),) * 2, evidence
+ assert tuple(_message_fields(event["response"]) for event in logged) == (fields, fields), evidence
+
+
+def _spend_rows(rig: Rig, key: str) -> tuple[Event, ...]:
+ response: Final = rig.gateway.request("GET", "/spend/logs", params={"api_key": key})
+ assert response.status_code == 200, response.text
+ return tuple(JSON_OBJECT.validate_python(row) for row in JSON_LIST.validate_json(response.content))
+
+
+def test_spend_logs_endpoint_reports_the_stop_details_of_a_refusal(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _chat(rig, key, model, refusal_prompt("refusal", _marker()))
+ rows: Final = eventually(lambda: _spend_rows(rig, key), lambda found: len(found) >= 1, seconds=70)
+ assert len(rows) == 1, rows
+ reported: Final = (rows[0]["request_id"], relayed(rows[0]["response"]))
+ expected: Final = (body["id"], ((_CHAT_FIELDS, REFUSAL_FIELDS),))
+ assert ((body["id"], relayed(body)), reported) == (expected, expected), rows
+
+
+def _usage_counts(usage: JsonValue, total: str, details: str) -> tuple[JsonValue, JsonValue]:
+ counted: Final = JSON_OBJECT.validate_python(usage)
+ return (counted[total], JSON_OBJECT.validate_python(counted[details])["reasoning_tokens"])
+
+
+def _chat_usage_counts(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]:
+ return tuple(
+ _usage_counts(chunk["usage"], "completion_tokens", "completion_tokens_details")
+ for chunk in chunks
+ if chunk.get("usage")
+ )
+
+
+def _cached_stream_usage_counts(
+ rig: Rig, surface: Literal["chat_stream", "responses_stream"], key: str, model: str, text: str
+) -> tuple[tuple[JsonValue, JsonValue], ...]:
+ match surface:
+ case "chat_stream":
+ request: Final = {
+ **_chat_request(model, text, stream=True, cached=True),
+ "stream_options": {"include_usage": True},
+ }
+ return _chat_usage_counts(chunks_of(_post(rig, "/v1/chat/completions", request, key).decode()))
+ case "responses_stream":
+ events: Final = frames(
+ _post(rig, "/v1/responses", _responses_request(model, text, stream=True, cached=True), key).decode()
+ )
+ responses: Final = tuple(
+ JSON_OBJECT.validate_python(event["response"])
+ for event in events
+ if isinstance(event.get("response"), dict)
+ )
+ return tuple(
+ _usage_counts(response["usage"], "output_tokens", "output_tokens_details")
+ for response in responses
+ if response.get("usage")
+ )
+
+
+def _logged_usage_counts(logged: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]:
+ return tuple(
+ _usage_counts(
+ JSON_OBJECT.validate_python(event["response"])["usage"], "completion_tokens", "completion_tokens_details"
+ )
+ for event in logged
+ )
+
+
+@pytest.mark.parametrize("fixture", ["plainusage", "citedusage"], ids=["plain", "cited"])
+def test_openai_compatible_stream_cache_hit_replays_the_usage_details_of_the_miss(rig: Rig, fixture: str) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "openai_compatible", rig)
+ key: Final = scenario.key()
+ text: Final = refusal_prompt(fixture, _marker())
+ first: Final = _cached_stream_usage_counts(rig, "chat_stream", key, model, text)
+ _logged(rig, key)
+ second: Final = _cached_stream_usage_counts(rig, "chat_stream", key, model, text)
+ logged: Final = _logged(rig, key, 2)
+ served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged)
+ evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged}
+ assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence
+ assert (first, second, _logged_usage_counts(logged)) == (((7, 4),), ((7, 4),), ((7, 4), (7, 4))), evidence
+
+
+def test_text_completion_stream_cache_hit_replays_the_usage_details_of_the_miss(rig: Rig) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "openai_compatible", rig)
+ key: Final = scenario.key()
+ request: Final[Event] = {
+ "model": model,
+ "prompt": refusal_prompt("plainusage", _marker()),
+ "max_tokens": 64,
+ "stream": True,
+ "stream_options": {"include_usage": True},
+ }
+ first: Final = _chat_usage_counts(chunks_of(_post(rig, "/v1/completions", request, key).decode()))
+ _logged(rig, key)
+ second: Final = _chat_usage_counts(chunks_of(_post(rig, "/v1/completions", request, key).decode()))
+ logged: Final = _logged(rig, key, 2)
+ served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged)
+ evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged}
+ assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/chat/completions",)), evidence
+ assert (first, second, _logged_usage_counts(logged)) == (((7, 4),), ((7, 4),), ((7, 4), (7, 4))), evidence
+
+
+@pytest.mark.parametrize("surface", ["chat_stream", "responses_stream"])
+def test_anthropic_stream_cache_hit_keeps_the_usage_details_of_the_miss(
+ rig: Rig, surface: Literal["chat_stream", "responses_stream"]
+) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ text: Final = refusal_prompt("reasoned", _marker())
+ first: Final = _cached_stream_usage_counts(rig, surface, key, model, text)
+ _logged(rig, key)
+ second: Final = _cached_stream_usage_counts(rig, surface, key, model, text)
+ logged: Final = _logged(rig, key, 2)
+ served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged)
+ evidence: Final = {"cache_hit": served_from_cache, "first": first, "second": second, "logged": logged}
+ assert (served_from_cache, _paths(rig.wire)) == ((False, True), ("/v1/messages",)), evidence
+ assert (first, second, _logged_usage_counts(logged)) == (((10, 4),), ((10, 4),), ((10, 4), (10, 4))), evidence
+
+
+def _response_id(event: Mapping[str, JsonValue]) -> JsonValue:
+ return JSON_OBJECT.validate_python(event["response"]).get("id")
+
+
+@pytest.mark.parametrize("variant", tuple(MALFORMED_STOP_DETAILS))
+def test_stream_cache_hit_skips_cached_provider_fields_that_are_not_an_object(rig: Rig, variant: str) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "openai_compatible", rig)
+ key: Final = scenario.key()
+ marker: Final = _marker()
+ text: Final = refusal_prompt("plain", marker)
+ filled: Final = chunks_of(
+ _post(
+ rig,
+ "/v1/chat/completions",
+ {**_chat_request(model, text, stream=True, cached=True), "cache": {"ttl": 2}},
+ key,
+ ).decode()
+ )
+ assert _streamed_text(filled) == "ready now", filled
+ _logged(rig, key)
+ stale_id: Final = rewrite_cached_message_fields(openai_id(marker), MALFORMED_STOP_DETAILS[variant])
+ replayed: Final = eventually(
+ lambda: _chat_stream(rig, key, model, text, cached=True),
+ lambda chunks: {chunk["id"] for chunk in chunks} == {stale_id},
+ seconds=30,
+ )
+ streamed: Final = (_streamed_text(replayed), _finish_reasons(replayed), _delta_fields(replayed))
+ assert streamed == ("ready now", ("stop",), ()), replayed
+ assert _paths(rig.wire) == ("/v1/chat/completions",)
+ stale_logs: Final = eventually(
+ lambda: tuple(event for event in rig.sink.for_key(key) if _response_id(event) == stale_id),
+ lambda found: len(found) >= 1,
+ )
+ assert tuple(_message_fields(event["response"]) for event in stale_logs) == ((),) * len(stale_logs), stale_logs
+
+
+JsonModeRoute: TypeAlias = Literal["anthropic_native", "anthropic_tool", "bedrock_invoke"]
+_UNMAPPED_MODEL: Final = "claude-scripted-9"
+_JSON_ANSWER: Final = json.dumps({"answer": DOCS_PARTIAL_TEXT})
+
+
+@dataclass(frozen=True, slots=True)
+class JsonModeCase:
+ route: JsonModeRoute
+ fixture: str
+ streamed_text: str
+ message_text: str
+ streamed_finish_reason: str
+ message_finish_reason: str
+ requested: Literal["json_tool", "native"]
+
+
+_JSON_MODE_CASES: Final[Mapping[str, JsonModeCase]] = {
+ "anthropic_native": JsonModeCase(
+ "anthropic_native",
+ "partial",
+ DOCS_PARTIAL_TEXT,
+ DOCS_PARTIAL_TEXT,
+ "content_filter",
+ "content_filter",
+ "native",
+ ),
+ "anthropic_tool_declared": JsonModeCase(
+ "anthropic_tool", "jsontool", DECLARED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool"
+ ),
+ "anthropic_tool_enveloped": JsonModeCase(
+ "anthropic_tool", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool"
+ ),
+ "bedrock_invoke_declared": JsonModeCase(
+ "bedrock_invoke", "jsontool", DECLARED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool"
+ ),
+ "bedrock_invoke_enveloped": JsonModeCase(
+ "bedrock_invoke", "jsonheld", ENVELOPED_PARTIAL_JSON, _JSON_ANSWER, "content_filter", "stop", "json_tool"
+ ),
+}
+
+
+def _json_mode_deployment(scenario: Scenario, route: JsonModeRoute, rig: Rig) -> tuple[str, Provider]:
+ match route:
+ case "anthropic_native":
+ return _deployment(scenario, "anthropic", rig), "anthropic"
+ case "anthropic_tool":
+ return (
+ scenario.model(model=f"anthropic/{_UNMAPPED_MODEL}", api_base=rig.wire.url, api_key=_ANTHROPIC_KEY),
+ "anthropic",
+ )
+ case "bedrock_invoke":
+ return _deployment(scenario, "bedrock_invoke", rig), "bedrock_invoke"
+
+
+def _requested_format(body: Mapping[str, JsonValue]) -> tuple[bool, bool]:
+ tools: Final = JSON_LIST.validate_python(body.get("tools") or [])
+ output_config: Final = JSON_OBJECT.validate_python(body.get("output_config") or {})
+ json_tool: Final = any(JSON_OBJECT.validate_python(tool).get("name") == JSON_TOOL_NAME for tool in tools)
+ return json_tool, "output_format" in body or "format" in output_config
+
+
+def _received(wire: Wire) -> tuple[tuple[str, tuple[bool, bool]], ...]:
+ return tuple(
+ (urlsplit(unquote(request.target)).path, _requested_format(JSON_OBJECT.validate_json(request.body)))
+ for request in wire.drain()
+ )
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"])
+@pytest.mark.parametrize("case", tuple(_JSON_MODE_CASES))
+def test_json_mode_refusal_relays_the_docs_stop_details_through_the_response_format_path(
+ rig: Rig, case: str, stream: bool
+) -> None:
+ expected: Final = _JSON_MODE_CASES[case]
+ with rig.gateway.scenario() as scenario:
+ model, provider = _json_mode_deployment(scenario, expected.route, rig)
+ key: Final = scenario.key()
+ request: Final = {
+ **_chat_request(model, refusal_prompt(expected.fixture, _marker()), stream=stream, cached=False),
+ "response_format": JSON_ANSWER_FORMAT,
+ }
+ content: Final = _post(rig, "/v1/chat/completions", request, key)
+ requested: Final = (expected.requested == "json_tool", expected.requested == "native")
+ assert _received(rig.wire) == ((_STREAM_TARGETS[provider][int(stream)], requested),)
+ if stream:
+ chunks: Final = chunks_of(content.decode())
+ streamed: Final = (_streamed_text(chunks), _finish_reasons(chunks), _relayed_in(chunks))
+ assert streamed == (
+ expected.streamed_text,
+ (expected.streamed_finish_reason,),
+ ((_DELTA_FIELDS, REFUSAL_FIELDS),),
+ ), chunks
+ else:
+ body: Final = JSON_OBJECT.validate_json(content)
+ answered: Final = (_message_text(body), _choice(body)["finish_reason"], relayed(body))
+ assert answered == (
+ expected.message_text,
+ expected.message_finish_reason,
+ ((_CHAT_FIELDS, REFUSAL_FIELDS),),
+ ), body
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+def _serialized_choice(body: Mapping[str, JsonValue]) -> Event:
+ choice: Final = _choice(body)
+ return {**choice, "message": JSON_OBJECT.validate_python(choice["message"])}
+
+
+_ANTHROPIC_MESSAGE_FIELDS: Final[Event] = {"citations": None, "thinking_blocks": None}
+
+
+@pytest.mark.parametrize(("fixture", "content"), [("refusal", None), ("partial", DOCS_PARTIAL_TEXT)])
+def test_chat_refusal_body_serializes_the_stop_details_beside_the_existing_message_fields(
+ rig: Rig, fixture: str, content: str | None
+) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "anthropic", rig)
+ key: Final = scenario.key()
+ body: Final = _chat(rig, key, model, refusal_prompt(fixture, _marker()))
+ expected: Final[Event] = {
+ "finish_reason": "content_filter",
+ "index": 0,
+ "message": {
+ "content": content,
+ "role": "assistant",
+ "provider_specific_fields": {**_ANTHROPIC_MESSAGE_FIELDS, **REFUSAL_FIELDS},
+ },
+ }
+ assert _serialized_choice(body) == expected, body
+ assert _paths(rig.wire) == ("/v1/messages",)
+ logged: Final = _logged_response(rig, key)
+ assert relayed(logged) == ((_CHAT_FIELDS, REFUSAL_FIELDS),), logged
+
+
+def _usage_of(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[tuple[JsonValue, JsonValue], ...]:
+ usages: Final = tuple(JSON_OBJECT.validate_python(chunk["usage"]) for chunk in chunks if chunk.get("usage"))
+ return tuple((usage["completion_tokens"], usage.get("completion_tokens_details")) for usage in usages)
+
+
+def _served_ids(response: httpx.Response) -> frozenset[JsonValue]:
+ return frozenset(chunk["id"] for chunk in chunks_of(response.text)) if response.status_code == 200 else frozenset()
+
+
+@pytest.mark.parametrize("variant", tuple(UNVALIDATED_USAGE_DETAILS))
+def test_stream_cache_hit_skips_cached_usage_details_that_fail_validation(rig: Rig, variant: str) -> None:
+ with rig.gateway.scenario() as scenario:
+ model: Final = _deployment(scenario, "openai_compatible", rig)
+ key: Final = scenario.key()
+ marker: Final = _marker()
+ request: Final = {
+ **_chat_request(model, refusal_prompt("plainusage", marker), stream=True, cached=True),
+ "stream_options": {"include_usage": True},
+ }
+ filled: Final = chunks_of(_post(rig, "/v1/chat/completions", {**request, "cache": {"ttl": 2}}, key).decode())
+ assert (_streamed_text(filled), _chat_usage_counts(filled)) == ("ready now", ((7, 4),)), filled
+ _logged(rig, key)
+ stale_id: Final = rewrite_cached_usage_details(openai_id(marker), UNVALIDATED_USAGE_DETAILS[variant])
+ replayed: Final = eventually(
+ lambda: rig.gateway.request("POST", "/v1/chat/completions", request, key=key),
+ lambda response: response.status_code != 200 or _served_ids(response) == {stale_id},
+ seconds=30,
+ )
+ assert replayed.status_code == 200, replayed.text
+ chunks: Final = chunks_of(replayed.text)
+ streamed: Final = (_streamed_text(chunks), _finish_reasons(chunks), _usage_of(chunks))
+ assert streamed == ("ready now", ("stop",), ((7, {"reasoning_tokens": 0, "text_tokens": 7}),)), chunks
+ assert _paths(rig.wire) == ("/v1/chat/completions",)
+ stale_logs: Final = eventually(
+ lambda: tuple(event for event in rig.sink.for_key(key) if _response_id(event) == stale_id),
+ lambda found: len(found) >= 1,
+ )
+ assert tuple(event["status"] for event in stale_logs) == ("success",) * len(stale_logs), stale_logs
diff --git a/tests/integration/sdk/test_response_cache_sync_sdk.py b/tests/integration/sdk/test_response_cache_sync_sdk.py
index 54d78adbaba..1fe73ec9012 100644
--- a/tests/integration/sdk/test_response_cache_sync_sdk.py
+++ b/tests/integration/sdk/test_response_cache_sync_sdk.py
@@ -9,17 +9,23 @@ from __future__ import annotations
import os
from collections.abc import Iterator
+from itertools import chain
from typing import Final
import pytest
+from pydantic import JsonValue, TypeAdapter
import litellm
from litellm.caching.caching import Cache
+from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.caching import LiteLLMCacheType
from litellm.types.utils import ModelResponse, TextCompletionResponse
+from tests.integration._support.client import JSON_OBJECT
from tests.integration.caching.response_cache_case import (
CHAT,
+ CHAT_CHUNK,
CHAT_MODEL,
+ CHAT_STREAM,
PROVIDER_KEY,
TEXT,
TEXT_MODEL,
@@ -31,8 +37,10 @@ from tests.integration.caching.response_cache_case import (
prompt,
redis_store,
rescript,
+ rewrite_usage_details,
scripted_in_process,
)
+from tests.integration.cost_calculation.cost_tracking_case import SseResponse
@pytest.fixture
@@ -123,3 +131,77 @@ def test_sync_text_completion_with_empty_choices_reaches_upstream_again() -> Non
hit: Final = _text_completion(handle.api_base(), text)
assert hit.id == refilled.id, hit.model_dump_json()
assert upstream.calls(handle.scenario_id) == 2
+
+
+_REASONING_USAGE_FRAME: Final = (
+ CHAT_CHUNK + '[],"usage":{"prompt_tokens":5,"completion_tokens":7,"total_tokens":12,'
+ '"completion_tokens_details":{"reasoning_tokens":4}}}'
+)
+_REASONING_STREAM: Final = SseResponse(
+ content_type="text/event-stream",
+ frames=(*CHAT_STREAM.frames[:-1], _REASONING_USAGE_FRAME, CHAT_STREAM.frames[-1]),
+)
+_JSON_LIST: Final = TypeAdapter(list[JsonValue])
+_UNVALIDATED_USAGE_DETAILS: Final[dict[str, JsonValue]] = {
+ "text": {"reasoning_tokens": "abc"},
+ "fraction": {"reasoning_tokens": 1.5},
+ "object": {"reasoning_tokens": {"a": 1}},
+}
+
+
+def _stream(api_base: str, text: str) -> tuple[dict[str, JsonValue], ...]:
+ stream: Final = litellm.completion(
+ model=CHAT_MODEL,
+ messages=[{"role": "user", "content": text}],
+ api_base=api_base,
+ api_key=PROVIDER_KEY,
+ stream=True,
+ stream_options={"include_usage": True},
+ )
+ assert isinstance(stream, CustomStreamWrapper), type(stream)
+ return tuple(JSON_OBJECT.validate_python(chunk.model_dump(mode="json")) for chunk in stream)
+
+
+def _replayed(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[frozenset[JsonValue], str, tuple[JsonValue, ...]]:
+ choices: Final = tuple(
+ JSON_OBJECT.validate_python(choice)
+ for choice in chain.from_iterable(_JSON_LIST.validate_python(chunk.get("choices") or []) for chunk in chunks)
+ )
+ text: Final = "".join(str(JSON_OBJECT.validate_python(choice["delta"]).get("content") or "") for choice in choices)
+ usages: Final = tuple(JSON_OBJECT.validate_python(chunk["usage"]) for chunk in chunks if chunk.get("usage"))
+ reasoning: Final = tuple(
+ JSON_OBJECT.validate_python(usage.get("completion_tokens_details") or {}).get("reasoning_tokens")
+ for usage in usages
+ )
+ return frozenset(chunk["id"] for chunk in chunks), text, reasoning
+
+
+@pytest.mark.usefixtures("sdk_cache")
+def test_sync_stream_cache_hit_replays_the_usage_details_of_the_miss() -> None:
+ upstream: Final = Upstream()
+ text: Final = prompt()
+ with scripted_in_process(_REASONING_STREAM) as handle, redis_store() as store:
+ first: Final = _replayed(_stream(handle.api_base(), text))
+ (first_id,) = first[0]
+ assert first[1:] == ("streamed response", (4,)), first
+ await_entry(store, str(first_id))
+ hit: Final = _replayed(_stream(handle.api_base(), text))
+ assert hit == first, hit
+ assert upstream.calls(handle.scenario_id) == 1
+
+
+@pytest.mark.usefixtures("sdk_cache")
+@pytest.mark.parametrize("variant", tuple(_UNVALIDATED_USAGE_DETAILS))
+def test_sync_stream_cache_hit_skips_cached_usage_details_that_fail_validation(variant: str) -> None:
+ upstream: Final = Upstream()
+ text: Final = prompt()
+ with scripted_in_process(_REASONING_STREAM) as handle, redis_store() as store:
+ first: Final = _replayed(_stream(handle.api_base(), text))
+ (first_id,) = first[0]
+ assert first[1:] == ("streamed response", (4,)), first
+ stale_id: Final = rewrite_usage_details(
+ store, await_entry(store, str(first_id)), _UNVALIDATED_USAGE_DETAILS[variant]
+ )
+ hit: Final = _replayed(_stream(handle.api_base(), text))
+ assert hit == (frozenset({stale_id}), "streamed response", (0,)), hit
+ assert upstream.calls(handle.scenario_id) == 1
diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py
index 2c3dae21058..e9d96f057ae 100644
--- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py
+++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py
@@ -11,6 +11,7 @@ content frame, restoring per-token cadence on cache hits.
import pytest
+import litellm
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_split_assembled_content_for_replay,
convert_to_streaming_response,
@@ -87,6 +88,71 @@ async def test_async_yields_multiple_content_chunks_with_lossless_join():
assert reassembled == text
+@pytest.mark.parametrize("content", ["", "Hello world! How are you?"], ids=["single_frame", "sliced"])
+@pytest.mark.asyncio
+async def test_async_and_sync_replays_carry_the_cached_provider_specific_fields(content):
+ provider_specific_fields = {"stop_reason": "refusal", "stop_details": {"category": "cyber"}}
+
+ def payload():
+ cached = _async_payload(content=content)
+ cached["choices"][0]["message"]["provider_specific_fields"] = provider_specific_fields
+ return cached
+
+ async_chunks = await _collect_async(payload())
+ sync_chunks = list(convert_to_streaming_response(payload()))
+
+ assert [c.choices[0].delta.provider_specific_fields for c in async_chunks] == [
+ c.choices[0].delta.provider_specific_fields for c in sync_chunks
+ ]
+ assert async_chunks[0].choices[0].delta.provider_specific_fields == provider_specific_fields
+
+
+@pytest.mark.parametrize("provider_fields", [{"thought_signatures": ["sig"]}, None], ids=["provider_fields", "text"])
+@pytest.mark.parametrize("replay", ["async", "sync"])
+@pytest.mark.asyncio
+async def test_a_replayed_stream_keeps_the_cached_completion_tokens_details(replay, provider_fields):
+ cached = _async_payload()
+ cached["choices"][0]["message"]["provider_specific_fields"] = provider_fields
+ cached["usage"]["completion_tokens_details"] = {"reasoning_tokens": 4, "text_tokens": 3}
+
+ chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached))
+ rebuilt = litellm.stream_chunk_builder(chunks=chunks)
+
+ assert rebuilt.choices[0].message.get("provider_specific_fields") == provider_fields
+ assert rebuilt.usage.completion_tokens_details.reasoning_tokens == 4
+ assert rebuilt.usage.completion_tokens_details.text_tokens == 3
+
+
+@pytest.mark.parametrize("cached_fields", [7, [1, 2], "", "x" * 5000], ids=["int", "list", "empty", "long"])
+@pytest.mark.asyncio
+async def test_the_async_replay_ignores_cached_provider_fields_that_are_not_an_object(cached_fields):
+ cached = _async_payload()
+ cached["choices"][0]["message"]["provider_specific_fields"] = cached_fields
+
+ chunks = await _collect_async(cached)
+
+ assert "".join(c.choices[0].delta.content or "" for c in chunks) == "Hello world! How are you?"
+ assert all(c.choices[0].delta.provider_specific_fields is None for c in chunks)
+
+
+@pytest.mark.parametrize(
+ "details",
+ [{"reasoning_tokens": "abc"}, {"reasoning_tokens": 1.5}, {"reasoning_tokens": {"a": 1}}],
+ ids=["text", "fraction", "object"],
+)
+@pytest.mark.parametrize("replay", ["async", "sync"])
+@pytest.mark.asyncio
+async def test_a_replayed_stream_skips_cached_completion_tokens_details_that_do_not_validate(replay, details):
+ cached = _async_payload()
+ cached["usage"]["completion_tokens_details"] = details
+
+ chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached))
+ rebuilt = litellm.stream_chunk_builder(chunks=chunks)
+
+ assert rebuilt.choices[0].message.content == "Hello world! How are you?"
+ assert rebuilt.usage.completion_tokens == 7
+
+
@pytest.mark.asyncio
async def test_async_finish_reason_only_on_last_chunk():
chunks = await _collect_async(_async_payload())
@@ -204,9 +270,7 @@ def test_sync_multi_chunk_and_lossless_join():
def test_sync_delta_and_choice_metadata_only_on_first_chunk():
- thinking_blocks = [
- {"type": "thinking", "thinking": "cached thinking", "signature": "sig"}
- ]
+ thinking_blocks = [{"type": "thinking", "thinking": "cached thinking", "signature": "sig"}]
payload = {
"id": "chatcmpl-test-sync",
"object": "chat.completion",
diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py
index 0172f0e9aeb..a604e61b9c2 100644
--- a/tests/unit/litellm_core_utils/test_streaming_handler.py
+++ b/tests/unit/litellm_core_utils/test_streaming_handler.py
@@ -732,6 +732,42 @@ def test_set_response_id_propagation_valid_to_invalid(
assert initialized_custom_stream_wrapper.response_id == "first-valid-id"
+@pytest.mark.parametrize(
+ "last_delta",
+ [
+ Delta(content="Hello"),
+ Delta(provider_specific_fields={"stop_reason": "refusal"}),
+ ],
+ ids=["content", "provider_specific_fields"],
+)
+def test_logged_response_keeps_received_model_id_when_the_finish_chunk_is_synthesized(
+ logging_obj: Logging, last_delta: Delta
+):
+ upstream_id: Final = "chatcmpl-upstream-id"
+ wrapper: Final = CustomStreamWrapper(
+ completion_stream=ModelResponseListIterator(
+ model_responses=[
+ ModelResponseStream(
+ id=upstream_id,
+ model="gpt-4o-mini",
+ choices=[StreamingChoices(index=0, delta=last_delta, finish_reason="content_filter")],
+ )
+ ]
+ ),
+ model="gpt-4o-mini",
+ logging_obj=logging_obj,
+ custom_llm_provider="openai",
+ )
+
+ emitted: Final = list(wrapper)
+ logged: Final = litellm.stream_chunk_builder(chunks=wrapper.chunks)
+
+ assert len(emitted) == 2
+ assert [chunk._hidden_params.get("received_model_id") for chunk in emitted] == [upstream_id, upstream_id]
+ assert logged is not None
+ assert logged._hidden_params.get("received_model_id") == upstream_id
+
+
@pytest.mark.asyncio
async def test_streaming_completion_start_time(logging_obj: Logging):
"""Test that the start time is set correctly"""
diff --git a/tests/unit/llms/anthropic/anthropic_docs_refusal.py b/tests/unit/llms/anthropic/anthropic_docs_refusal.py
new file mode 100644
index 00000000000..2bd75b4aab5
--- /dev/null
+++ b/tests/unit/llms/anthropic/anthropic_docs_refusal.py
@@ -0,0 +1,19 @@
+from typing import Final
+
+ANTHROPIC_DOCS_REFUSAL_URL: Final = "https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback#refusal-response"
+
+# Copied verbatim from ANTHROPIC_DOCS_REFUSAL_URL on 2026-10-12. Streams carry the same stop_details on message_delta.
+ANTHROPIC_DOCS_REFUSAL_MESSAGE: Final = {
+ "id": "msg_01XFUDYJgAACzvnptvVoYEL",
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-fable-5",
+ "content": [],
+ "stop_reason": "refusal",
+ "stop_details": {
+ "type": "refusal",
+ "category": "cyber",
+ "explanation": "This request was declined because it could enable cyber harm.",
+ },
+ "usage": {"input_tokens": 412, "output_tokens": 0},
+}
diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
index 7d1d2a3c434..a072010bfa9 100644
--- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
+++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
@@ -22,6 +22,7 @@ from litellm.types.llms.openai import (
)
from litellm.types.responses.main import OutputCodeInterpreterCall
from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream
+from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE
@pytest.mark.asyncio
@@ -3366,3 +3367,50 @@ def test_anthropic_streaming_invalid_key_maps_to_authentication_error() -> None:
assert exc_info.value.status_code == 401
assert exc_info.value.llm_provider == "anthropic"
assert route.calls.last.request.headers["x-api-key"] == "bad-anthropic-key"
+
+
+def _message_stream_events(stop_reason: str, stop_details: dict | None) -> tuple[dict, ...]:
+ return (
+ {
+ "type": "message_start",
+ "message": {
+ "id": ANTHROPIC_DOCS_REFUSAL_MESSAGE["id"],
+ "type": "message",
+ "role": "assistant",
+ "model": ANTHROPIC_DOCS_REFUSAL_MESSAGE["model"],
+ "content": [],
+ "usage": {"input_tokens": 412, "output_tokens": 0},
+ },
+ },
+ {
+ "type": "message_delta",
+ "delta": {"stop_reason": stop_reason, "stop_sequence": None, "stop_details": stop_details},
+ "usage": {"output_tokens": 0},
+ },
+ {"type": "message_stop"},
+ )
+
+
+def _assemble_stream(events: tuple[dict, ...]) -> litellm.ModelResponse:
+ iterator: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
+ chunks: Final = [iterator.chunk_parser(event) for event in events]
+ return litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}])
+
+
+def test_streamed_refusal_stop_details_reach_the_assembled_message():
+ stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"]
+ assembled: Final = _assemble_stream(_message_stream_events("refusal", stop_details))
+
+ assert assembled.choices[0].finish_reason == "content_filter"
+ provider_fields: Final = assembled.choices[0].message.provider_specific_fields
+ assert provider_fields["stop_reason"] == "refusal"
+ assert provider_fields["stop_details"] == stop_details
+
+
+def test_streamed_end_turn_carries_no_stop_details():
+ assembled: Final = _assemble_stream(_message_stream_events("end_turn", None))
+
+ assert assembled.choices[0].finish_reason == "stop"
+ provider_fields: Final = assembled.choices[0].message.provider_specific_fields or {}
+ assert "stop_details" not in provider_fields
+ assert "stop_reason" not in provider_fields
diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
index bbcc8629a10..5f4ecc4440b 100644
--- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
+++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
@@ -39,6 +39,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation im
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk
from litellm.types.utils import ChatCompletionToolCallChunk, ModelResponse, ServerToolUse, Usage
+from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE
def test_response_format_transformation_unit_test():
@@ -9082,3 +9083,93 @@ def test_completion_surfaces_response_citations_in_provider_specific_fields(resp
assert response.choices[0].message.provider_specific_fields["citations"] == [
[{**citation, "supported_text": "the grass is green"}]
]
+
+
+def _refusal_message_body(stop_reason: str, stop_details: dict | None) -> dict:
+ return {**ANTHROPIC_DOCS_REFUSAL_MESSAGE, "stop_reason": stop_reason, "stop_details": stop_details}
+
+
+def _transform_message_body(body: dict) -> litellm.ModelResponse:
+ return AnthropicConfig().transform_parsed_response(
+ completion_response=body,
+ raw_response=httpx.Response(status_code=200, headers={}),
+ model_response=litellm.ModelResponse(),
+ json_mode=False,
+ prefix_prompt=None,
+ )
+
+
+def test_refusal_stop_details_kept_in_provider_specific_fields():
+ stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"]
+ result: Final = _transform_message_body(ANTHROPIC_DOCS_REFUSAL_MESSAGE)
+
+ assert result.choices[0].finish_reason == "content_filter"
+ provider_fields: Final = result.choices[0].message.provider_specific_fields
+ assert provider_fields["stop_reason"] == "refusal"
+ assert provider_fields["stop_details"] == stop_details
+ assert result._hidden_params["provider_specific_fields"]["stop_details"] == stop_details
+
+
+def test_refusal_without_stop_details_keeps_the_native_stop_reason():
+ result: Final = _transform_message_body(_refusal_message_body("refusal", None))
+
+ provider_fields: Final = result.choices[0].message.provider_specific_fields
+ assert provider_fields["stop_reason"] == "refusal"
+ assert "stop_details" not in provider_fields
+
+
+def test_end_turn_response_carries_no_stop_details():
+ result: Final = _transform_message_body(_refusal_message_body("end_turn", None))
+
+ assert result.choices[0].finish_reason == "stop"
+ provider_fields: Final = result.choices[0].message.provider_specific_fields
+ assert "stop_details" not in provider_fields
+ assert "stop_reason" not in provider_fields
+
+
+def _json_mode_message_body(stop_reason: str, stop_details: dict | None) -> dict:
+ return {
+ **_refusal_message_body(stop_reason, stop_details),
+ "content": [
+ {"type": "tool_use", "id": "toolu_json", "name": RESPONSE_FORMAT_TOOL_NAME, "input": {"answer": "Hello"}}
+ ],
+ }
+
+
+def _transform_json_mode_body(body: dict) -> litellm.ModelResponse:
+ return AnthropicConfig().transform_parsed_response(
+ completion_response=body,
+ raw_response=httpx.Response(status_code=200, headers={}),
+ model_response=litellm.ModelResponse(),
+ json_mode=True,
+ prefix_prompt=None,
+ )
+
+
+def test_json_mode_refusal_keeps_stop_details_on_the_message():
+ stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"]
+ result: Final = _transform_json_mode_body(_json_mode_message_body("refusal", stop_details))
+
+ message: Final = result.choices[0].message
+ assert json.loads(message.content) == {"answer": "Hello"}
+ assert message.provider_specific_fields == {"stop_reason": "refusal", "stop_details": stop_details}
+
+
+def test_json_mode_tool_use_turn_adds_no_stop_fields_to_the_message():
+ result: Final = _transform_json_mode_body(_json_mode_message_body("tool_use", None))
+
+ message: Final = result.choices[0].message
+ assert json.loads(message.content) == {"answer": "Hello"}
+ assert not message.provider_specific_fields
+
+
+def test_non_dict_upstream_body_still_transforms_without_refusal_fields() -> None:
+ result = AnthropicConfig().transform_parsed_response(
+ completion_response=MagicMock(),
+ raw_response=httpx.Response(200, headers={}),
+ model_response=litellm.ModelResponse(),
+ json_mode=False,
+ prefix_prompt=None,
+ )
+ provider_specific_fields = result.choices[0].message.provider_specific_fields or {}
+ assert "stop_details" not in provider_specific_fields
diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
index 859b57c7cf3..8cb9ad9b73c 100644
--- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
+++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
@@ -4,7 +4,7 @@ import json
import uuid
from types import SimpleNamespace
from typing import Final
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
import httpx
import pytest
@@ -17,6 +17,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transfor
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
+from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE
ONE_PIXEL_PNG = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg=="
@@ -1689,3 +1690,23 @@ class TestBedrockAnthropicCombinedRegressions:
user_msg = result["messages"][0]
has_cache_control = any(isinstance(c, dict) and "cache_control" in c for c in user_msg["content"])
assert has_cache_control
+
+
+def test_refusal_stop_details_kept_through_bedrock_invoke_transform():
+ body: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE
+ result: Final = AmazonAnthropicClaudeConfig().transform_response(
+ model="anthropic.claude-sonnet-5-5-v1:0",
+ raw_response=httpx.Response(status_code=200, headers={}, json=body),
+ model_response=litellm.ModelResponse(),
+ logging_obj=MagicMock(),
+ request_data={},
+ messages=[{"role": "user", "content": "hi"}],
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+ assert result.choices[0].finish_reason == "content_filter"
+ provider_fields: Final = result.choices[0].message.provider_specific_fields
+ assert provider_fields["stop_reason"] == "refusal"
+ assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"]
diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py
index 56e76df25a4..ace24c8a3f4 100644
--- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py
+++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py
@@ -3,6 +3,7 @@ import copy
import json
from collections.abc import Mapping
from typing import Final
+from unittest.mock import MagicMock
import httpx
import pytest
@@ -19,6 +20,7 @@ from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation im
)
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.vertex_ai import VertexPartnerProvider
+from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE
@pytest.mark.parametrize(
@@ -1125,3 +1127,23 @@ def test_chat_flagged_model_replays_a_byte_identical_prefix_around_a_mid_convers
_assert_prefix_stable(requests)
assert [m["role"] for m in requests[1]["messages"]] == ["user", "assistant", "user", "system"]
assert [m["role"] for m in requests[2]["messages"]] == ["user", "assistant", "user", "system", "assistant", "user"]
+
+
+def test_refusal_stop_details_kept_through_vertex_anthropic_transform():
+ body: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE
+ result: Final = VertexAIAnthropicConfig().transform_response(
+ model="claude-sonnet-5-5",
+ raw_response=httpx.Response(status_code=200, headers={}, json=body),
+ model_response=litellm.ModelResponse(),
+ logging_obj=MagicMock(),
+ request_data={},
+ messages=[{"role": "user", "content": "hi"}],
+ optional_params={},
+ litellm_params={},
+ encoding=None,
+ )
+
+ assert result.choices[0].finish_reason == "content_filter"
+ provider_fields: Final = result.choices[0].message.provider_specific_fields
+ assert provider_fields["stop_reason"] == "refusal"
+ assert provider_fields["stop_details"] == ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"]
diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
index 973d6bc9f7b..88c5e2d2c1b 100644
--- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
+++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
@@ -16,6 +16,7 @@ from litellm.llms.anthropic.pass_through.stream_assembly import extract_sse_data
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
+from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE
async def _drain_tasks():
@@ -2951,3 +2952,58 @@ def test_get_user_from_metadata(end_user_id):
)
assert response == "test"
+
+
+def _anthropic_sse_stream(stop_reason: str, stop_details: dict | None) -> list[str]:
+ events: Final = (
+ (
+ "message_start",
+ {
+ "type": "message_start",
+ "message": {
+ "id": ANTHROPIC_DOCS_REFUSAL_MESSAGE["id"],
+ "type": "message",
+ "role": "assistant",
+ "model": ANTHROPIC_DOCS_REFUSAL_MESSAGE["model"],
+ "content": [],
+ "usage": {"input_tokens": 412, "output_tokens": 0},
+ },
+ },
+ ),
+ (
+ "message_delta",
+ {
+ "type": "message_delta",
+ "delta": {"stop_reason": stop_reason, "stop_sequence": None, "stop_details": stop_details},
+ "usage": {"output_tokens": 0},
+ },
+ ),
+ ("message_stop", {"type": "message_stop"}),
+ )
+ return [f"event: {name}\ndata: {json.dumps(data)}\n\n" for name, data in events]
+
+
+def test_streamed_messages_refusal_stop_details_on_the_logged_response(mock_logging_obj):
+ stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"]
+ logged: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response(
+ all_chunks=_anthropic_sse_stream("refusal", stop_details),
+ litellm_logging_obj=mock_logging_obj,
+ model="claude-sonnet-5-5",
+ )
+
+ assert logged.choices[0].finish_reason == "content_filter"
+ provider_fields: Final = logged.choices[0].message.provider_specific_fields
+ assert provider_fields["stop_reason"] == "refusal"
+ assert provider_fields["stop_details"] == stop_details
+
+
+def test_streamed_messages_end_turn_logs_no_stop_details(mock_logging_obj):
+ logged: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response(
+ all_chunks=_anthropic_sse_stream("end_turn", None),
+ litellm_logging_obj=mock_logging_obj,
+ model="claude-sonnet-5-5",
+ )
+
+ assert logged.choices[0].finish_reason == "stop"
+ provider_fields: Final = logged.choices[0].message.provider_specific_fields or {}
+ assert "stop_details" not in provider_fields
From b9849f35881782f137b750340963ceb1e64cb27d Mon Sep 17 00:00:00 2001
From: Kerry Lu
Date: Wed, 14 Oct 2026 14:43:01 -0700
Subject: [PATCH 38/47] fix(router): add fallback_on_unsupported_params so
param rejections can raise instead of falling back (#46350)
* fix(router): add fallback_on_unsupported_params so client-side param rejections can raise instead of falling back
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* style(router): wrap the fallback gate so ruff format passes
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(router): rename the flag to fallback_on_bad_request and cover provider 400s too
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(router): narrow the wire test responses so the type gate count does not grow
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* test(openai): swap the model map with setattr so the type gate count does not grow
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(router): import the exception classes directly so the type gate count does not grow
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(router): type the fallbacks list in fallback_event_handlers to pay for the new Router kwarg in the type gate
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(router): keep the LIT001 count flat on the typed fallbacks parameters
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* refactor(router): type the fallback stream wrapper weakref so the type gate stays level
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(router): make fallback_on_bad_request updatable at runtime so DB-stored router_settings apply it
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* chore(ui): regenerate schema.d.ts for UpdateRouterConfig.fallback_on_bad_request
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(router): rename the flag back to fallback_on_unsupported_params
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* revert(openai): drop the reasoning_effort=none client-side check, the Router flag covers the provider 400
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* ci: rerun docs checks now that the router_settings row is on litellm-docs main
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(router): only stop fallbacks on real 400s and cover the context window and content policy exclusions
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(router): move the no-fallback 400 check into is_bad_request_without_fallback and pin the 402 and Vertex 403 cases
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(router): keep falling back on router-raised 400s when fallback_on_unsupported_params is off
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
* fix(router): read llm_provider off the exception with getattr so the type gate stays flat
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
litellm/constants.py | 1 +
litellm/router.py | 18 +-
.../router_utils/fallback_event_handlers.py | 19 +-
litellm/types/router.py | 1 +
...router_unsupported_params_fallback_wire.py | 213 ++++++++++++++++++
.../test_fallback_event_handlers.py | 138 ++++++++++--
tests/unit/test_router_retry_policy_update.py | 15 ++
ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +
8 files changed, 380 insertions(+), 27 deletions(-)
create mode 100644 tests/integration/sdk/test_router_unsupported_params_fallback_wire.py
diff --git a/litellm/constants.py b/litellm/constants.py
index 4ac6adaba73..bbf6b8a4f04 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -60,6 +60,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(
"enable_tag_filtering",
"tag_routing_prefix",
"optional_pre_call_checks",
+ "fallback_on_unsupported_params",
}
)
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
diff --git a/litellm/router.py b/litellm/router.py
index 16257775d3f..b38b78b60d9 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -230,6 +230,7 @@ from litellm.router_utils.fallback_event_handlers import (
get_fallback_model_group_for_lookup_groups,
get_pre_routing_selection,
has_unattempted_fallback_target,
+ is_bad_request_without_fallback,
mid_stream_fallback_hop_kwargs,
mid_stream_fallback_snapshot_kwargs,
mid_stream_retry_kwargs,
@@ -1022,6 +1023,7 @@ class Router:
allowed_fails_policy: AllowedFailsPolicy | None = None, # set custom allowed fails policy
cooldown_time: float | None = None, # (seconds) time to cooldown a deployment after failure
disable_cooldowns: bool | None = None,
+ fallback_on_unsupported_params: bool = True,
routing_strategy: RoutingStrategyName = "simple-shuffle",
optional_pre_call_checks: OptionalPreCallChecks | None = None,
routing_strategy_args: dict = {}, # just for latency-based
@@ -1264,6 +1266,7 @@ class Router:
self.cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS
self.cooldown_cache = CooldownCache(cache=self.cache, default_cooldown_time=self.cooldown_time)
self.disable_cooldowns = disable_cooldowns
+ self.fallback_on_unsupported_params = fallback_on_unsupported_params
self.enable_health_check_routing = enable_health_check_routing
self.enable_weighted_failover = enable_weighted_failover
self.health_check_ignore_transient_errors = health_check_ignore_transient_errors
@@ -3270,7 +3273,7 @@ class Router:
wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks())
# weak, so the generator closing over it does not keep the wrapper out of
# refcount teardown and delay the `finally` that releases the deployment slot
- wrapper_ref: Final = weakref.ref(wrapped_response)
+ wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response)
return wrapped_response
@staticmethod
@@ -3875,7 +3878,7 @@ class Router:
wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks())
# weak, for the same reason as the async twin
- wrapper_ref: Final = weakref.ref(wrapped_response)
+ wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response)
return wrapped_response
async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
@@ -7728,7 +7731,15 @@ class Router:
hop_depth: Final = kwargs.get("fallback_depth")
nested_fallback_hop: Final = isinstance(hop_depth, int) and hop_depth > 0
- if disable_fallbacks is True or original_model_group is None or is_guardrail_intervention(e):
+ unsupported_params_without_fallback: Final = is_bad_request_without_fallback(
+ e, self.fallback_on_unsupported_params
+ )
+ if (
+ disable_fallbacks is True
+ or original_model_group is None
+ or is_guardrail_intervention(e)
+ or unsupported_params_without_fallback
+ ):
raise e
input_kwargs: Final = {
@@ -12802,6 +12813,7 @@ class Router:
"enable_weighted_failover",
"enable_tag_filtering",
"tag_routing_prefix",
+ "fallback_on_unsupported_params",
]
for var in vars_to_include:
diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py
index e1de57c91a0..4c393a7a93b 100644
--- a/litellm/router_utils/fallback_event_handlers.py
+++ b/litellm/router_utils/fallback_event_handlers.py
@@ -222,6 +222,16 @@ def has_unattempted_fallback_target(
return any((key := fallback_attempt_key(target)) is None or key not in attempted for target in fallback_model_group)
+def is_bad_request_without_fallback(exception: Exception, fallback_on_unsupported_params: bool) -> bool:
+ if fallback_on_unsupported_params or not isinstance(exception, litellm.BadRequestError):
+ return False
+ if isinstance(exception, litellm.UnsupportedParamsError):
+ return True
+ if isinstance(exception, (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError)):
+ return False
+ return exception.status_code == 400 and bool(getattr(exception, "llm_provider", ""))
+
+
def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool:
"""
Handles wildcard routing scenario
@@ -514,7 +524,7 @@ def fallback_lookup_groups(kwargs: Mapping[str, object], model_group: str | None
def _resolved_a_specific_chain(
- fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
+ fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape
result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract
) -> bool:
resolved, generic_idx = result
@@ -524,7 +534,7 @@ def _resolved_a_specific_chain(
def get_fallback_model_group_for_lookup_groups(
- fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract
+ fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape
lookup_groups: tuple[str, ...],
) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract
"""
@@ -538,7 +548,10 @@ def get_fallback_model_group_for_lookup_groups(
return next((result for result in results if result[0] is not None), (None, None))
-def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]:
+def get_fallback_model_group(
+ fallbacks: Sequence[Mapping[str, list[str]]], # mutable-ok: mirrors the fallbacks config shape
+ model_group: str,
+) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract
"""
Returns:
- fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"]
diff --git a/litellm/types/router.py b/litellm/types/router.py
index 273edf0f913..59b6798bfdc 100644
--- a/litellm/types/router.py
+++ b/litellm/types/router.py
@@ -188,6 +188,7 @@ class UpdateRouterConfig(LiteLLMBaseModel):
weights: RouterWeights | None = None
tag_routing_prefix: str | None = None
optional_pre_call_checks: OptionalPreCallChecks | None = None
+ fallback_on_unsupported_params: bool | None = None
model_config = ConfigDict(protected_namespaces=())
diff --git a/tests/integration/sdk/test_router_unsupported_params_fallback_wire.py b/tests/integration/sdk/test_router_unsupported_params_fallback_wire.py
new file mode 100644
index 00000000000..e7675fb64ee
--- /dev/null
+++ b/tests/integration/sdk/test_router_unsupported_params_fallback_wire.py
@@ -0,0 +1,213 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import Callable, Mapping
+from typing import Final
+
+import pytest
+from integration._support.wire import Reply, Request, Wire, wire_server
+
+import litellm
+from litellm import Router
+
+_MODEL: Final = "gpt-5-wire-unmapped"
+_API_KEY: Final = "synthetic-unsupported-params-key"
+_PROMPT: Final = "which deployment answers when the requested model rejects the request?"
+_CLIENTS: Final = ("sync", "async")
+_PRIMARY_TO_BACKUP: Final = [{"primary": ["backup"]}]
+_UNSUPPORTED_BY_PRIMARY: Final = {"reasoning_effort": "xhigh"}
+
+
+def _answers(text: str) -> Reply:
+ body: Final = {
+ "id": "chatcmpl-bad-request-wire",
+ "object": "chat.completion",
+ "created": 1700000000,
+ "model": _MODEL,
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+ }
+ return Reply(body=json.dumps(body).encode())
+
+
+_UPSTREAM_500: Final = Reply(status=500, body=json.dumps({"error": {"message": "overloaded"}}).encode())
+_UPSTREAM_400: Final = Reply(
+ status=400,
+ body=json.dumps({"error": {"message": "Unsupported value: temperature", "type": "invalid_request_error"}}).encode(),
+)
+_UPSTREAM_CONTEXT_WINDOW: Final = Reply(
+ status=400,
+ body=json.dumps(
+ {
+ "error": {
+ "message": "This model's maximum context length is 8192 tokens",
+ "type": "invalid_request_error",
+ "code": "context_length_exceeded",
+ }
+ }
+ ).encode(),
+)
+_UPSTREAM_CONTENT_POLICY: Final = Reply(
+ status=400,
+ body=json.dumps(
+ {
+ "error": {
+ "message": "Your request was rejected",
+ "type": "invalid_request_error",
+ "code": "content_policy_violation",
+ }
+ }
+ ).encode(),
+)
+_BACKUP_ANSWERS: Final = _answers("answered by the backup")
+
+
+def _peer(replies: Mapping[str, Reply]) -> Callable[[Request], Reply]:
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST", request.method
+ deployment, _, route = request.target.lstrip("/").partition("/")
+ assert route == "chat/completions", request.target
+ return replies[deployment]
+
+ return respond
+
+
+def _deployments_hit(wire: Wire) -> tuple[str, ...]:
+ return tuple(request.target.lstrip("/").partition("/")[0] for request in wire.drain())
+
+
+_ALL_DAY_FOR_TEAM_A: Final = (
+ {"start": "00:00", "end": "12:00", "timezone": "UTC", "team_ids": ["team-a"]},
+ {"start": "12:00", "end": "00:00", "timezone": "UTC", "team_ids": ["team-a"]},
+)
+
+
+def _router(wire: Wire, primary_model_info: Mapping[str, object] | None = None, **settings: object) -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": "primary",
+ "litellm_params": {"model": f"openai/{_MODEL}", "api_base": f"{wire.url}/primary", "api_key": _API_KEY},
+ "model_info": dict(primary_model_info or {}),
+ },
+ {
+ "model_name": "backup",
+ "litellm_params": {
+ "model": f"openai/{_MODEL}",
+ "api_base": f"{wire.url}/backup",
+ "api_key": _API_KEY,
+ "drop_params": True,
+ },
+ },
+ ],
+ num_retries=0,
+ disable_cooldowns=True,
+ **{"fallbacks": _PRIMARY_TO_BACKUP, **settings},
+ )
+
+
+_MESSAGES: Final = [{"role": "user", "content": _PROMPT}]
+
+
+def _content(response: litellm.ModelResponse) -> str:
+ choice: Final = response.choices[0]
+ assert isinstance(choice, litellm.Choices), choice
+ return str(choice.message.content)
+
+
+async def _acomplete(router: Router, **request: object) -> str:
+ response: Final = await router.acompletion(model="primary", messages=_MESSAGES, **request)
+ assert isinstance(response, litellm.ModelResponse), response
+ return _content(response)
+
+
+def _complete(client: str, router: Router, **request: object) -> str:
+ if client == "async":
+ return asyncio.run(_acomplete(router, **request))
+ response: Final = router.completion(model="primary", messages=_MESSAGES, **request)
+ assert isinstance(response, litellm.ModelResponse), response
+ return _content(response)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_client_side_param_rejection_reaches_the_caller_without_trying_fallbacks(client: str) -> None:
+ with wire_server(_peer({"backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(wire, fallback_on_unsupported_params=False)
+ with pytest.raises(litellm.UnsupportedParamsError, match="reasoning_effort"):
+ _complete(client, router, **_UNSUPPORTED_BY_PRIMARY)
+ assert _deployments_hit(wire) == ()
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_client_side_param_rejection_still_falls_back_by_default(client: str) -> None:
+ with wire_server(_peer({"backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(wire)
+ assert _complete(client, router, **_UNSUPPORTED_BY_PRIMARY) == "answered by the backup"
+ assert _deployments_hit(wire) == ("backup",)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_upstream_500_still_falls_back_when_bad_requests_raise(client: str) -> None:
+ with wire_server(_peer({"primary": _UPSTREAM_500, "backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(wire, fallback_on_unsupported_params=False)
+ assert _complete(client, router) == "answered by the backup"
+ assert _deployments_hit(wire) == ("primary", "backup")
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_router_raised_400_still_falls_back_when_bad_requests_raise(client: str) -> None:
+ with wire_server(_peer({"backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(
+ wire, primary_model_info={"access_windows": _ALL_DAY_FOR_TEAM_A}, fallback_on_unsupported_params=False
+ )
+ answer: Final = _complete(client, router, metadata={"user_api_key_team_id": "team-b"})
+ assert answer == "answered by the backup"
+ assert _deployments_hit(wire) == ("backup",)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_context_window_error_still_uses_its_own_fallbacks_when_bad_requests_raise(client: str) -> None:
+ with wire_server(_peer({"primary": _UPSTREAM_CONTEXT_WINDOW, "backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(
+ wire, fallback_on_unsupported_params=False, fallbacks=None, context_window_fallbacks=_PRIMARY_TO_BACKUP
+ )
+ assert _complete(client, router) == "answered by the backup"
+ assert _deployments_hit(wire) == ("primary", "backup")
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_content_policy_error_still_uses_its_own_fallbacks_when_bad_requests_raise(client: str) -> None:
+ with wire_server(_peer({"primary": _UPSTREAM_CONTENT_POLICY, "backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(
+ wire, fallback_on_unsupported_params=False, fallbacks=None, content_policy_fallbacks=_PRIMARY_TO_BACKUP
+ )
+ assert _complete(client, router) == "answered by the backup"
+ assert _deployments_hit(wire) == ("primary", "backup")
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_upstream_400_reaches_the_caller_without_trying_fallbacks(client: str) -> None:
+ with wire_server(_peer({"primary": _UPSTREAM_400, "backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(wire, fallback_on_unsupported_params=False)
+ with pytest.raises(litellm.BadRequestError, match="Unsupported value: temperature"):
+ _complete(client, router)
+ assert _deployments_hit(wire) == ("primary",)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_flag_applied_through_update_settings_stops_fallbacks_on_upstream_400(client: str) -> None:
+ with wire_server(_peer({"primary": _UPSTREAM_400, "backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(wire)
+ router.update_settings(fallback_on_unsupported_params=False)
+ with pytest.raises(litellm.BadRequestError, match="Unsupported value: temperature"):
+ _complete(client, router)
+ assert _deployments_hit(wire) == ("primary",)
+
+
+@pytest.mark.parametrize("client", _CLIENTS)
+def test_upstream_400_still_falls_back_by_default(client: str) -> None:
+ with wire_server(_peer({"primary": _UPSTREAM_400, "backup": _BACKUP_ANSWERS})) as wire:
+ router: Final = _router(wire)
+ assert _complete(client, router) == "answered by the backup"
+ assert _deployments_hit(wire) == ("primary", "backup")
diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py
index febb33aab7d..b353967df3a 100644
--- a/tests/unit/router_utils/test_fallback_event_handlers.py
+++ b/tests/unit/router_utils/test_fallback_event_handlers.py
@@ -29,6 +29,7 @@ from litellm.router_utils.fallback_event_handlers import (
fallback_attempt_key,
get_fallback_model_group,
get_pre_routing_selection,
+ is_bad_request_without_fallback,
log_failure_fallback_event,
log_success_fallback_event,
mid_stream_fallback_snapshot_kwargs,
@@ -1378,9 +1379,7 @@ class TestOrderedFallbackLookupGroups:
"smart-router",
"requested-model",
)
- assert fallback_lookup_groups({"metadata": {"model_group": []}}, "requested-model") == (
- "requested-model",
- )
+ assert fallback_lookup_groups({"metadata": {"model_group": []}}, "requested-model") == ("requested-model",)
def test_fallback_hop_resumes_the_original_groups_chain_last(self):
from litellm.router_utils.fallback_event_handlers import fallback_lookup_groups
@@ -1539,7 +1538,11 @@ def test_mid_stream_retry_kwargs_strips_what_the_retry_wrapper_pops_and_keeps_th
@pytest.mark.parametrize(
"kwargs,expected",
[
- pytest.param({"litellm_metadata": {"attempted_retries": 2}, "metadata": {"attempted_retries": 5}}, 2, id="litellm_metadata-wins"),
+ pytest.param(
+ {"litellm_metadata": {"attempted_retries": 2}, "metadata": {"attempted_retries": 5}},
+ 2,
+ id="litellm_metadata-wins",
+ ),
pytest.param({"metadata": {"attempted_retries": 1}}, 1, id="metadata-bucket"),
pytest.param({"litellm_metadata": {"attempted_retries": "2"}}, 0, id="string-is-not-a-count"),
pytest.param({"litellm_metadata": {"attempted_retries": -1}}, 0, id="negative-is-not-a-count"),
@@ -1566,8 +1569,12 @@ def test_record_retry_attempt_stamps_the_bucket_the_retry_wrapper_reads():
"kwargs,expected",
[
pytest.param({"litellm_metadata": {"attempted_retries": 1, "max_retries": 3}}, 3, id="committed-by-a-retry"),
- pytest.param({"litellm_metadata": {"attempted_retries": 0, "max_retries": 3}}, None, id="stamped-before-any-retry"),
- pytest.param({"litellm_metadata": {"attempted_retries": 1, "max_retries": "3"}}, None, id="string-is-not-a-budget"),
+ pytest.param(
+ {"litellm_metadata": {"attempted_retries": 0, "max_retries": 3}}, None, id="stamped-before-any-retry"
+ ),
+ pytest.param(
+ {"litellm_metadata": {"attempted_retries": 1, "max_retries": "3"}}, None, id="string-is-not-a-budget"
+ ),
pytest.param({"litellm_metadata": {"attempted_retries": 1}}, None, id="no-budget"),
pytest.param({}, None, id="no-bucket"),
],
@@ -1602,6 +1609,7 @@ def _vcr_outcome_gate(request, vcr):
yield
record_vcr_outcome(request, vcr)
+
@pytest.fixture(scope="function")
def setup_and_teardown():
"""
@@ -1622,6 +1630,7 @@ def setup_and_teardown():
loop.close()
asyncio.set_event_loop(None)
+
REFUSAL_RESPONSE: dict[str, Any] = {
"id": "msg_refusal",
"type": "message",
@@ -1647,9 +1656,11 @@ OK_RESPONSE: dict[str, Any] = {
"usage": {"input_tokens": 25, "output_tokens": 2},
}
+
def _sse(event: str, data: dict[str, Any]) -> bytes:
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
+
REFUSAL_STREAM_FRAMES: tuple[bytes, ...] = (
_sse("message_start", {"type": "message_start", "message": {**REFUSAL_RESPONSE, "stop_reason": None}}),
_sse(
@@ -1672,10 +1683,12 @@ OK_STREAM_FRAMES: tuple[bytes, ...] = (
_sse("message_stop", {"type": "message_stop"}),
)
+
def _split_frames_mid_data_line(frames: tuple[bytes, ...]) -> tuple[bytes, ...]:
"""Split each frame's data line in half, modeling a transport chunk boundary."""
return tuple(part for frame in frames for part in (frame[: len(frame) // 2], frame[len(frame) // 2 :]))
+
class _FrameStream(httpx.AsyncByteStream):
def __init__(self, frames: tuple[bytes, ...]) -> None:
self._frames = frames
@@ -1687,6 +1700,7 @@ class _FrameStream(httpx.AsyncByteStream):
async def aclose(self) -> None:
return None
+
class FakeAnthropicUpstream:
"""Intercepts the third-party transport (httpx.AsyncClient.send): refuses on fable
models, answers on others. The router deliberately does not forward caller-injected
@@ -1724,6 +1738,7 @@ class FakeAnthropicUpstream:
return patch("httpx.AsyncClient.send", new=_send)
+
FABLE_TIER = {
"model_name": "fable-tier",
"litellm_params": {"model": "anthropic/claude-fable-5", "api_key": "sk-test"},
@@ -1734,12 +1749,15 @@ OPUS_TARGET = {
"litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "sk-test"},
}
+
def _router(content_policy_fallbacks: list | None) -> Router:
return Router(model_list=[FABLE_TIER, OPUS_TARGET], content_policy_fallbacks=content_policy_fallbacks)
+
async def _collect(stream: AsyncIterator[bytes]) -> bytes:
return b"".join([chunk async for chunk in stream])
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response():
@@ -1756,6 +1774,7 @@ async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response
assert len(fake.calls) == 2
assert "claude-opus-5" in fake.calls[1]
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
@pytest.mark.parametrize(
@@ -1780,6 +1799,7 @@ async def test_non_streaming_refusal_passes_through_untouched(content_policy_fal
assert response.get("stop_details") == upstream_body.get("stop_details")
assert len(fake.calls) == 1
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_streaming_refusal_with_fallback_row_streams_fallback_frames():
@@ -1796,6 +1816,7 @@ async def test_streaming_refusal_with_fallback_row_streams_fallback_frames():
assert b"text_delta" in body
assert len(fake.calls) == 2
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_streaming_refusal_split_across_chunks_still_falls_back():
@@ -1812,6 +1833,7 @@ async def test_streaming_refusal_split_across_chunks_still_falls_back():
assert b"text_delta" in body
assert len(fake.calls) == 2
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_streaming_refusal_without_fallback_row_passes_frames_through():
@@ -1828,6 +1850,7 @@ async def test_streaming_refusal_without_fallback_row_passes_frames_through():
assert b"stop_details" in body
assert len(fake.calls) == 1
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata():
@@ -1861,6 +1884,7 @@ async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_i
assert b"text_delta" in body
assert len(fake.calls) == 2
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain():
@@ -1880,6 +1904,7 @@ async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain
assert b'"stop_reason": "refusal"' in body
assert len(fake.calls) == 1
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_tier_stamp_never_reaches_provider_bound_metadata():
@@ -1916,6 +1941,7 @@ async def test_tier_stamp_never_reaches_provider_bound_metadata():
for body in fake.bodies:
assert body.get("metadata") == {"user_id": "u1"}
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_record_pre_routing_selection_writes_only_the_internal_bucket():
"""The Anthropic request's own metadata field must never carry the tier stamp."""
@@ -1926,6 +1952,7 @@ def test_record_pre_routing_selection_writes_only_the_internal_bucket():
assert kwargs["litellm_metadata"] == {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-x"}
assert kwargs["metadata"] == {"user_id": "u1"}
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
@@ -1949,6 +1976,7 @@ async def test_generic_only_row_recovers_safeguard_refusal(stream):
assert len(fake.calls) == 2
assert "claude-opus-5" in fake.calls[1]
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_configured_content_policy_list_stays_authoritative_over_generic_rows():
@@ -1967,6 +1995,7 @@ async def test_configured_content_policy_list_stays_authoritative_over_generic_r
assert response["stop_reason"] == "refusal"
assert len(fake.calls) == 1
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_policy():
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"tier-group": ["opus-target"]}])
@@ -1976,6 +2005,7 @@ def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_po
assert router._refusal_fallback_available("router-group", {}) is False
assert router._refusal_fallback_available("router-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_chat_content_filter_gate_unchanged_by_generic_rows():
"""The generic-row arming is scoped to /v1/messages safeguard refusals; the chat surface's
@@ -1987,6 +2017,7 @@ def test_chat_content_filter_gate_unchanged_by_generic_rows():
assert router._should_raise_content_policy_error(model="fable-tier", response=response, kwargs={}) is False
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
@@ -2012,6 +2043,7 @@ async def test_disable_fallbacks_returns_the_refusal_instead_of_raising(stream):
assert body["stop_reason"] == "refusal"
assert len(fake.calls) == 1
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_disable_fallbacks_beats_a_content_policy_row_too():
@@ -2032,6 +2064,7 @@ async def test_disable_fallbacks_beats_a_content_policy_row_too():
assert response["stop_reason"] == "refusal"
assert len(fake.calls) == 1
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_refusal_gate_keys_on_pre_routing_tier_stamp():
router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}])
@@ -2059,6 +2092,7 @@ def test_refusal_gate_keys_on_pre_routing_tier_stamp():
is False
)
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_has_content_policy_fallback_default_fallbacks_arm():
router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}])
@@ -2066,6 +2100,7 @@ def test_has_content_policy_fallback_default_fallbacks_arm():
assert router._has_content_policy_fallback("any-group", {}) is True
assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested():
router = _router(content_policy_fallbacks=None)
@@ -2079,6 +2114,7 @@ def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested
) == ["backup-b"]
assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_refusal_gate_ignores_other_generic_call_types():
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
@@ -2269,9 +2305,7 @@ async def test_failed_fallbacks_raise_most_recent_exception(function_name):
model="gpt-3.5-turbo",
)
- request_kwargs: Dict[str, Any] = {
- "metadata": {"previous_models": ["gpt-3.5-turbo"]}
- }
+ request_kwargs: Dict[str, Any] = {"metadata": {"previous_models": ["gpt-3.5-turbo"]}}
if function_name == "_aembedding":
request_kwargs["input"] = "hello this is a test for run_async_fallback"
@@ -2323,27 +2357,19 @@ class CustomTestLogger(CustomLogger):
self.success_fallback_events = []
self.failure_fallback_events = []
- async def log_success_fallback_event(
- self, original_model_group, kwargs, original_exception
- ):
+ async def log_success_fallback_event(self, original_model_group, kwargs, original_exception):
print(
"in log_success_fallback_event for original_model_group: ",
original_model_group,
)
- self.success_fallback_events.append(
- (original_model_group, kwargs, original_exception)
- )
+ self.success_fallback_events.append((original_model_group, kwargs, original_exception))
- async def log_failure_fallback_event(
- self, original_model_group, kwargs, original_exception
- ):
+ async def log_failure_fallback_event(self, original_model_group, kwargs, original_exception):
print(
"in log_failure_fallback_event for original_model_group: ",
original_model_group,
)
- self.failure_fallback_events.append(
- (original_model_group, kwargs, original_exception)
- )
+ self.failure_fallback_events.append((original_model_group, kwargs, original_exception))
def create_test_router_2():
@@ -2458,3 +2484,73 @@ async def test_multiple_fallbacks(function_name, respx_mock: respx.MockRouter, m
print(result._hidden_params)
assert result._hidden_params["api_base"] == FAKE_OPENAI_API_BASE
+
+
+def _provider_error(status_code: int) -> httpx.Response:
+ return httpx.Response(status_code=status_code, request=httpx.Request(method="POST", url="https://provider.test"))
+
+
+@pytest.mark.parametrize(
+ ("exception", "stops_fallbacks"),
+ [
+ pytest.param(
+ litellm.BadRequestError(message="Unsupported value: temperature", model="m", llm_provider="openai"),
+ True,
+ id="provider 400",
+ ),
+ pytest.param(
+ litellm.UnsupportedParamsError(
+ message="reasoning_effort is not supported", model="m", llm_provider="openai"
+ ),
+ True,
+ id="client-side rejection",
+ ),
+ pytest.param(
+ litellm.UnsupportedParamsError(message="stop is not supported"),
+ True,
+ id="client-side rejection without a provider",
+ ),
+ pytest.param(
+ litellm.BadRequestError(
+ message="Deployment ptu is reserved for another team", model="ptu", llm_provider=""
+ ),
+ False,
+ id="router-raised 400 keeps falling back",
+ ),
+ pytest.param(
+ litellm.BadRequestError(
+ message="VertexAIException BadRequestError - 403",
+ model="m",
+ llm_provider="vertex_ai",
+ response=_provider_error(403),
+ ),
+ False,
+ id="vertex 403 wrapped as BadRequestError",
+ ),
+ pytest.param(
+ litellm.PaymentRequiredError(message="insufficient credits", model="m", llm_provider="openrouter"),
+ False,
+ id="402",
+ ),
+ pytest.param(
+ litellm.ContextWindowExceededError(message="maximum context length", model="m", llm_provider="openai"),
+ False,
+ id="context window keeps its own fallbacks",
+ ),
+ pytest.param(
+ litellm.ContentPolicyViolationError(message="content_policy_violation", model="m", llm_provider="openai"),
+ False,
+ id="content policy keeps its own fallbacks",
+ ),
+ pytest.param(
+ litellm.InternalServerError(message="overloaded", model="m", llm_provider="openai"),
+ False,
+ id="500 still falls back",
+ ),
+ ],
+)
+def test_is_bad_request_without_fallback_only_stops_on_a_real_400_when_the_flag_is_off(
+ exception: Exception, stops_fallbacks: bool
+) -> None:
+ assert is_bad_request_without_fallback(exception, fallback_on_unsupported_params=False) is stops_fallbacks
+ assert is_bad_request_without_fallback(exception, fallback_on_unsupported_params=True) is False
diff --git a/tests/unit/test_router_retry_policy_update.py b/tests/unit/test_router_retry_policy_update.py
index 48b5194a66e..918d215c89f 100644
--- a/tests/unit/test_router_retry_policy_update.py
+++ b/tests/unit/test_router_retry_policy_update.py
@@ -300,6 +300,21 @@ def test_update_settings_persists_retry_policy_dict():
assert router.retry_policy.TimeoutErrorRetries == 3
+def test_update_settings_fallback_on_unsupported_params_round_trips_through_get_settings():
+ router = _build_router()
+ assert router.get_settings()["fallback_on_unsupported_params"] is True
+
+ router.update_settings(fallback_on_unsupported_params=False)
+
+ assert router.fallback_on_unsupported_params is False
+ assert router.get_settings()["fallback_on_unsupported_params"] is False
+
+
+def test_update_router_config_keeps_fallback_on_unsupported_params_payload():
+ payload = UpdateRouterConfig(fallback_on_unsupported_params=False).model_dump(exclude_none=True)
+ assert payload["fallback_on_unsupported_params"] is False
+
+
def test_update_settings_accepts_retry_policy_object_unchanged():
"""A pre-built ``RetryPolicy`` instance must pass through verbatim so
callers that already constructed one (e.g. tests or programmatic
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 29d659ec1d9..1fd554e37a1 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -48111,6 +48111,8 @@ export interface components {
cooldown_time?: number | null;
/** Enable Tag Filtering */
enable_tag_filtering?: boolean | null;
+ /** Fallback On Unsupported Params */
+ fallback_on_unsupported_params?: boolean | null;
/** Fallbacks */
fallbacks?: {
[key: string]: unknown;
From 9f6379223c1d55b98f607b07683c2e731cda7587 Mon Sep 17 00:00:00 2001
From: Yassin Kortam
Date: Wed, 14 Oct 2026 14:50:55 -0700
Subject: [PATCH 39/47] feat(access-groups): per-server MCP tool permissions on
access groups (#46300)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
.../migration.sql | 4 +
.../litellm_proxy_extras/schema.prisma | 1 +
litellm/models/access_group.py | 2 +
.../mcp_server/auth/user_api_key_auth_mcp.py | 293 +++++++-
litellm/proxy/_lazy_openapi_snapshot.json | 51 ++
.../auth/agent_access_groups.py | 42 +-
litellm/proxy/auth/auth_checks.py | 141 ++--
.../access_group_endpoints.py | 19 +-
litellm/proxy/schema.prisma | 1 +
litellm/types/access_group.py | 4 +
schema.prisma | 1 +
.../auth/test_managed_agent_access.py | 74 +-
.../auth/test_user_api_key_auth_mcp.py | 693 ++++++++++++++++--
.../auth/test_agent_access_groups.py | 85 ++-
...st_auth_checks_object_access_and_lookup.py | 24 +
.../test_access_group_endpoints.py | 62 ++
.../AccessGroupsDetailsPage.test.tsx | 30 +-
.../_components/AccessGroupsDetailsPage.tsx | 32 +
.../AccessGroupsModal/AccessGroupBaseForm.tsx | 41 ++
.../AccessGroupEditModal.integration.test.tsx | 30 +
.../AccessGroupEditModal.tsx | 11 +-
.../AccessGroupCreateDialog.tsx | 6 +
.../access-group-create/mapper.test.ts | 18 +
.../_components/access-group-create/mapper.ts | 3 +
.../_components/access-group-create/schema.ts | 1 +
.../MCPToolPermissions.test.tsx | 44 ++
.../MCPToolPermissions.tsx | 21 +-
ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +
28 files changed, 1587 insertions(+), 159 deletions(-)
create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql
diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql
new file mode 100644
index 00000000000..35197a9e63c
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261013000000_access_group_mcp_tool_permissions/migration.sql
@@ -0,0 +1,4 @@
+SET LOCAL lock_timeout = '10s';
+
+-- AlterTable
+ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB;
diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
index ca58247a2c9..dcae56b5e71 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
@@ -1588,6 +1588,7 @@ model LiteLLM_AccessGroupTable {
// Resource memberships - explicit arrays per type
access_model_names String[] @default([])
access_mcp_server_ids String[] @default([])
+ mcp_tool_permissions Json?
access_agent_ids String[] @default([])
assigned_team_ids String[] @default([])
diff --git a/litellm/models/access_group.py b/litellm/models/access_group.py
index 513a60aa9ba..58b70c39661 100644
--- a/litellm/models/access_group.py
+++ b/litellm/models/access_group.py
@@ -5,6 +5,7 @@ Canonical definition for ``litellm_accessgrouptable``. Re-exported from
``litellm.proxy._types`` for backwards compatibility.
"""
+from collections.abc import Mapping, Sequence
from datetime import datetime
from litellm.types.llms.base import LiteLLMPydanticObjectBase
@@ -16,6 +17,7 @@ class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase):
description: str | None = None
access_model_names: list[str] = []
access_mcp_server_ids: list[str] = []
+ mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
access_agent_ids: list[str] = []
assigned_team_ids: list[str] = []
assigned_key_ids: list[str] = []
diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
index 18b7c279606..43537d48f5e 100644
--- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
@@ -1,3 +1,4 @@
+import functools
import re
import secrets
from collections.abc import Awaitable, Callable, Mapping, Sequence
@@ -42,6 +43,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
)
from litellm.proxy._types import (
UI_TEAM_ID,
+ LiteLLM_AccessGroupTable,
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTable,
ProxyException,
@@ -53,8 +55,10 @@ from litellm.proxy._types import (
user_api_key_has_admin_view,
)
from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
+ AgentAccessGroupCeiling,
CeilingResolver,
resolve_agent_access_group_ceiling,
+ resolve_managed_agent_ceilings,
)
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth
from litellm.proxy.agent_endpoints.auth.managed_authorization import managed_agent_policy
@@ -126,6 +130,22 @@ class UnloadableEntitlementError(Exception):
it places no ceiling; denying there would refuse MCP to every caller during a cold-cache fault."""
+@dataclass(frozen=True, slots=True)
+class KeyGroupTools:
+ """``owned``: tools the key's owned groups grant on a server, ``None`` when they place no bound.
+ ``owns_server``: whether an owned group grants the server at all, so an unbounded owned grant is told
+ apart from no owned grant. ``attached``: the bound every attached group, owned or not, places there,
+ applied only when the key holds no grant of its own on the server."""
+
+ owned: Sequence[str] | None
+ owns_server: bool
+ attached: Sequence[str] | None
+
+
+_NO_KEY_GROUP_TOOLS: Final = KeyGroupTools(owned=None, owns_server=False, attached=None)
+_UNREADABLE_KEY_GROUP_TOOLS: Final = KeyGroupTools(owned=(), owns_server=False, attached=())
+
+
def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: list[str] | None = None) -> list[str] | None:
"""Resolve the single MCP server name a cold-start passthrough bypass may
target. Delegates parsing to
@@ -251,6 +271,15 @@ def _agent_capped_servers(
)
+def _intersect_tool_ceiling(own: Sequence[str] | None, ceiling: Sequence[str] | None) -> Sequence[str] | None:
+ """The agent's own tool grants capped by its access-group ceiling, ``None`` only when neither bounds."""
+ if ceiling is None:
+ return own
+ if own is None:
+ return ceiling
+ return tuple(set(own) & set(ceiling))
+
+
def is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bool:
"""True when this auth is a keyless subject admitted by the gateway session / bridge user
path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``.
@@ -2363,16 +2392,177 @@ class MCPRequestHandler:
def _union_tool_grants(
direct: Sequence[str] | None,
via_toolsets: Sequence[str] | None,
+ via_access_groups: Sequence[str] | None = None,
) -> Sequence[str] | None:
- """Union of one level's direct tool grants and its toolset-granted tools on one server,
- ``None`` when neither source restricts (allow-all from this level). A direct grant
- containing ``MCP_ALL_TOOLS_WILDCARD`` makes the level unrestricted, so it returns
- ``None`` whatever the toolsets name."""
+ """Union of one level's direct tool grants, its toolset-granted tools and its unified
+ access-group tools on one server, ``None`` when no source restricts (allow-all from this
+ level). A direct grant containing ``MCP_ALL_TOOLS_WILDCARD`` makes the level unrestricted,
+ so it returns ``None`` whatever the toolsets or groups name."""
if direct is not None and MCP_ALL_TOOLS_WILDCARD in direct:
return None
- if direct is None and via_toolsets is None:
+ if direct is None and via_toolsets is None and via_access_groups is None:
return None
- return tuple({*(direct or ()), *(via_toolsets or ())})
+ return tuple({*(direct or ()), *(via_toolsets or ()), *(via_access_groups or ())})
+
+ @staticmethod
+ def access_group_tools_for_server(
+ groups: Sequence[LiteLLM_AccessGroupTable], server_id: str
+ ) -> Sequence[str] | None:
+ """Tools the unified access groups grant on ``server_id``: the union of each granting group's
+ ``mcp_tool_permissions`` entry for it. A granting group with no entry grants every tool, so
+ ``None`` wins. A group that does not grant the server contributes nothing, and an explicit
+ ``[]`` entry on a granting group grants no tools. ``None`` when no group grants the server."""
+ from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+ global_mcp_server_manager,
+ )
+
+ def granted_tools(group: LiteLLM_AccessGroupTable) -> Sequence[str] | None:
+ if not MCPRequestHandler._group_grants_server(group, server_id):
+ return None
+ permissions: Final = {
+ server: list(allowed) for server, allowed in (group.mcp_tool_permissions or {}).items()
+ }
+ tools: Final = global_mcp_server_manager.expand_tool_permissions(permissions)
+ return tools.get(server_id, (MCP_ALL_TOOLS_WILDCARD,))
+
+ grants: Final = tuple(granted_tools(group) for group in groups)
+ named: Final = tuple(grant for grant in grants if grant is not None)
+ if not named or any(MCP_ALL_TOOLS_WILDCARD in grant for grant in named):
+ return None
+ return tuple({tool for grant in named for tool in grant}) # comprehension-ok: flatten
+
+ @staticmethod
+ def _group_grants_server(group: LiteLLM_AccessGroupTable, server_id: str) -> bool:
+ from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+ global_mcp_server_manager,
+ )
+
+ return server_id in global_mcp_server_manager.expand_permission_list(list(group.access_mcp_server_ids))
+
+ @staticmethod
+ def _key_tool_scope(
+ direct: Sequence[str] | None,
+ via_toolsets: Sequence[str] | None,
+ groups: KeyGroupTools,
+ grants_server: bool,
+ ) -> Sequence[str] | None:
+ """The key's own tool scope on a server when it holds any grant there (a direct server grant with
+ or without a tool entry, a toolset or an owned group), so an unrestricted grant stays unrestricted.
+ Only a key with no grant of its own takes the bound its attached groups place."""
+ if grants_server or direct is not None or via_toolsets is not None or groups.owns_server:
+ return MCPRequestHandler._union_tool_grants(direct, via_toolsets, groups.owned)
+ return groups.attached
+
+ @staticmethod
+ def _group_authorizes_key(group: LiteLLM_AccessGroupTable, team_id: str | None, token: str | None) -> bool:
+ """Mirror of ``get_authorized_resources_from_key_access_groups``: a group a key merely names
+ binds only when the group lists the key's team or the key itself as an owner."""
+ return bool(team_id and team_id in group.assigned_team_ids) or bool(token and token in group.assigned_key_ids)
+
+ @staticmethod
+ def _loaded_group_tools_for_server(
+ groups: Sequence[LiteLLM_AccessGroupTable], requested: Sequence[str], server_id: str
+ ) -> Sequence[str] | None:
+ """Tool grants of the agent groups that loaded. A named agent group that did not load exists but
+ could not be read (a group that does not exist never reaches the ceiling), so the level grants no
+ tools on ``server_id`` rather than shedding a restriction while the fault lasts."""
+ if {group.access_group_id for group in groups} != set(requested):
+ return ()
+ return MCPRequestHandler.access_group_tools_for_server(groups, server_id)
+
+ @staticmethod
+ async def _key_direct_servers(
+ key_obj_perm: LiteLLM_ObjectPermissionTable | None, *, requires_fresh_policy: bool
+ ) -> AbstractSet[str]:
+ """Servers the key's object permission grants on its own: ``mcp_servers`` by id, name or alias and
+ legacy ``mcp_access_groups`` by the tag the server carries, exactly as server admission resolves
+ them. A tag never doubles as a server name, so a tag that collides with one grants nothing by itself."""
+ from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
+
+ if key_obj_perm is None:
+ return frozenset()
+ try:
+ tagged: Final = (
+ await MCPRequestHandler.get_mcp_servers_from_access_groups(
+ key_obj_perm.mcp_access_groups, requires_fresh_policy=requires_fresh_policy
+ )
+ if key_obj_perm.mcp_access_groups
+ else ()
+ )
+ except Exception as e: # noqa: BLE001 # a key naming tags we cannot resolve must deny, not widen
+ raise UnloadableEntitlementError(f"key access group tags unreadable: {e}") from e
+ return frozenset(global_mcp_server_manager.expand_permission_list(key_obj_perm.mcp_servers or [])) | frozenset(
+ tagged
+ )
+
+ @staticmethod
+ async def _key_access_group_tools_for_server(user_api_key_auth: UserAPIKeyAuth, server_id: str) -> KeyGroupTools:
+ """Tool bounds the key's attached groups place on ``server_id``. Attaching a group to a key admits
+ the group's servers (``_get_key_access_group_mcp_server_extras``), so an attached group the key
+ does not own still carries its own tool permissions onto the servers it admits, it just cannot
+ add tools to a grant the key already holds there. Only owned groups are grant sources. A named
+ group that does not exist grants and bounds nothing, one that exists but cannot be read denies."""
+ from litellm.proxy.auth.auth_checks import get_access_group_tables
+
+ if not user_api_key_auth.access_group_ids:
+ return _NO_KEY_GROUP_TOOLS
+ try:
+ groups: Final = await get_access_group_tables(
+ access_group_ids=user_api_key_auth.access_group_ids,
+ check_db_only=bool(user_api_key_auth.requires_fresh_policy),
+ unreadable="raise",
+ )
+ except Exception as e: # noqa: BLE001 # a key naming groups we cannot read must deny, not widen
+ if user_api_key_auth.requires_fresh_policy:
+ raise UnloadableEntitlementError(f"key access groups unreadable: {e}") from e
+ return _UNREADABLE_KEY_GROUP_TOOLS
+ owned: Final = tuple(
+ group
+ for group in groups
+ if MCPRequestHandler._group_authorizes_key(group, user_api_key_auth.team_id, user_api_key_auth.token)
+ )
+ return KeyGroupTools(
+ owned=MCPRequestHandler.access_group_tools_for_server(owned, server_id),
+ owns_server=any(MCPRequestHandler._group_grants_server(group, server_id) for group in owned),
+ attached=MCPRequestHandler.access_group_tools_for_server(groups, server_id),
+ )
+
+ @staticmethod
+ async def _team_access_group_tools_for_server(
+ user_api_key_auth: UserAPIKeyAuth, server_id: str
+ ) -> Sequence[str] | None:
+ from litellm.proxy.auth.auth_checks import get_access_group_tables, get_team_object
+ from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
+
+ if not user_api_key_auth.team_id or user_api_key_auth.team_id == UI_TEAM_ID or prisma_client is None:
+ return None
+ try:
+ team_obj: Final[LiteLLM_TeamTable | None] = await get_team_object(
+ team_id=user_api_key_auth.team_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ parent_otel_span=user_api_key_auth.parent_otel_span,
+ proxy_logging_obj=proxy_logging_obj,
+ check_db_only=bool(user_api_key_auth.requires_fresh_policy),
+ )
+ except Exception as e: # noqa: BLE001 # a team we cannot read must deny, not widen
+ raise UnloadableEntitlementError(f"team unreadable: {e}") from e
+ if not team_obj.access_group_ids:
+ return None
+ try:
+ groups: Final = await get_access_group_tables(
+ access_group_ids=team_obj.access_group_ids,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ check_db_only=bool(user_api_key_auth.requires_fresh_policy),
+ unreadable="raise",
+ )
+ except Exception as e: # noqa: BLE001 # team groups we cannot read grant no tools, they never widen
+ if user_api_key_auth.requires_fresh_policy:
+ raise UnloadableEntitlementError(f"team access groups unreadable: {e}") from e
+ return ()
+ return MCPRequestHandler.access_group_tools_for_server(groups, server_id)
@staticmethod
async def _key_object_permission_hydrated(
@@ -2493,7 +2683,17 @@ class MCPRequestHandler:
else None
)
- key_tools: Final = _as_list(MCPRequestHandler._union_tool_grants(key_direct_tools, key_toolset_tools))
+ key_group_tools: Final = await MCPRequestHandler._key_access_group_tools_for_server(
+ user_api_key_auth, server_id
+ )
+ key_direct_servers: Final = await MCPRequestHandler._key_direct_servers(
+ key_obj_perm, requires_fresh_policy=bool(user_api_key_auth.requires_fresh_policy)
+ )
+ key_tools: Final = _as_list(
+ MCPRequestHandler._key_tool_scope(
+ key_direct_tools, key_toolset_tools, key_group_tools, server_id in key_direct_servers
+ )
+ )
team_direct_tools: Final = (
global_mcp_server_manager.expand_tool_permissions(team_obj_perm.mcp_tool_permissions).get(server_id)
if team_obj_perm
@@ -2505,11 +2705,17 @@ class MCPRequestHandler:
team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(
team_obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy
)
- team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools)
+ team_group_tools: Final = await MCPRequestHandler._team_access_group_tools_for_server(
+ user_api_key_auth, server_id
+ )
+ team_tools: Final = MCPRequestHandler._union_tool_grants(
+ team_direct_tools, team_toolset_tools, team_group_tools
+ )
- # Apply same inheritance logic as get_allowed_mcp_servers
- if team_tools:
- if key_tools:
+ # Apply same inheritance logic as get_allowed_mcp_servers. ``is not None`` so an explicit
+ # empty allowlist stays a deny instead of reading as "no restriction".
+ if team_tools is not None:
+ if key_tools is not None:
# Both have restrictions → intersection
allowed_tools = list(set(team_tools) & set(key_tools))
else:
@@ -3525,6 +3731,9 @@ class MCPRequestHandler:
team_toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(
team_obj_perm, server_id, requires_fresh_policy=caller_auth.requires_fresh_policy
)
+ team_group_tools: Final = await MCPRequestHandler._team_access_group_tools_for_server(
+ caller_auth, server_id
+ )
except Exception as e: # noqa: BLE001 # an unresolved caller team must deny, not widen
verbose_logger.warning(
"MCP agent caller team tool ceiling unresolvable, denying tools on %r: %s", server_id, e
@@ -3535,7 +3744,9 @@ class MCPRequestHandler:
if team_obj_perm
else None
)
- team_tools: Final = MCPRequestHandler._union_tool_grants(team_direct_tools, team_toolset_tools)
+ team_tools: Final = MCPRequestHandler._union_tool_grants(
+ team_direct_tools, team_toolset_tools, team_group_tools
+ )
team_capped: Final = (
allowed_tools
if team_tools is None
@@ -3729,12 +3940,50 @@ class MCPRequestHandler:
global_mcp_server_manager,
)
+ ceilings: Final = await MCPRequestHandler._agent_access_group_ceilings(user_api_key_auth, resolve_ceiling)
+ if not ceilings:
+ return None
+ expanded: Final[tuple[frozenset[str], ...]] = tuple(
+ frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids)))
+ for ceiling in ceilings
+ )
+ return functools.reduce(lambda acc, servers: acc & servers, expanded[1:], expanded[0])
+
+ @staticmethod
+ async def _get_agent_access_group_tool_ceiling(
+ user_api_key_auth: UserAPIKeyAuth,
+ server_id: str,
+ resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
+ ) -> Sequence[str] | None:
+ """Ceiling the agent's attached unified access groups place on ``server_id`` tools, the tools axis
+ twin of ``_get_agent_access_group_server_ceiling``. None when the agent has no groups attached or
+ no attached group bounds tools on the server. The agent's own grants are intersected with it."""
+ ceilings: Final = await MCPRequestHandler._agent_access_group_ceilings(user_api_key_auth, resolve_ceiling)
+ if not ceilings:
+ return None
+ bounds: Final = tuple(
+ MCPRequestHandler._loaded_group_tools_for_server(ceiling.access_groups, ceiling.access_group_ids, server_id)
+ for ceiling in ceilings
+ )
+ return functools.reduce(_intersect_tool_ceiling, bounds[1:], bounds[0])
+
+ @staticmethod
+ async def _agent_access_group_ceilings(
+ user_api_key_auth: UserAPIKeyAuth,
+ resolve_ceiling: CeilingResolver,
+ ) -> tuple[AgentAccessGroupCeiling, ...]:
+ """A managed agent is bounded by the groups on its admitted policy, the same source server admission
+ reads, never by the registry row that policy may have superseded. Empty means no groups attached."""
if not user_api_key_auth.agent_id:
- return None
- ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id)
- if ceiling is None:
- return None
- return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids)))
+ return ()
+ managed: Final = managed_agent_policy(user_api_key_auth)
+ try:
+ if managed is not None:
+ return await resolve_managed_agent_ceilings(managed)
+ ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id)
+ except Exception as e: # noqa: BLE001 # agent groups we cannot read must deny, not widen
+ raise UnloadableEntitlementError(f"agent access groups unreadable: {e}") from e
+ return (ceiling,) if ceiling is not None else ()
@staticmethod
async def get_agent_tool_permissions_for_server(
@@ -3764,14 +4013,17 @@ class MCPRequestHandler:
if agent_object_permission is not None
else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
)
- if obj_perm is None:
- return None
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
+ group_tools: Final = await MCPRequestHandler._get_agent_access_group_tool_ceiling(
+ user_api_key_auth, server_id
+ )
+ if obj_perm is None:
+ return _as_list(group_tools)
direct_tools: Final = (
global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id)
if obj_perm.mcp_tool_permissions
@@ -3780,7 +4032,8 @@ class MCPRequestHandler:
toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(
obj_perm, server_id, requires_fresh_policy=user_api_key_auth.requires_fresh_policy
)
- agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
+ own_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
+ agent_tools: Final = _intersect_tool_ceiling(own_tools, group_tools)
return list(agent_tools) if agent_tools is not None else None
except Exception as e:
if managed_agent_policy(user_api_key_auth) is not None or isinstance(e, UnloadableEntitlementError):
diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json
index effb30fe8cb..225f289748e 100644
--- a/litellm/proxy/_lazy_openapi_snapshot.json
+++ b/litellm/proxy/_lazy_openapi_snapshot.json
@@ -686,6 +686,23 @@
}
],
"title": "Description"
+ },
+ "mcp_tool_permissions": {
+ "anyOf": [
+ {
+ "additionalProperties": {
+ "items": {
+ "type": "string"
+ },
+ "type": "array"
+ },
+ "type": "object"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "title": "Mcp Tool Permissions"
}
},
"required": [
@@ -867,6 +884,23 @@
],
"title": "Description"
},
+ "mcp_tool_permissions": {
+ "anyOf": [
+ {
+ "additionalProperties": {
+ "items": {
+ "type": "string"
+ },
+ "type": "array"
+ },
+ "type": "object"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "title": "Mcp Tool Permissions"
+ },
"updated_at": {
"format": "date-time",
"title": "Updated At",
@@ -995,6 +1029,23 @@
}
],
"title": "Description"
+ },
+ "mcp_tool_permissions": {
+ "anyOf": [
+ {
+ "additionalProperties": {
+ "items": {
+ "type": "string"
+ },
+ "type": "array"
+ },
+ "type": "object"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "title": "Mcp Tool Permissions"
}
},
"title": "AccessGroupUpdateRequest",
diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py
index db661fbea30..557d1c7bc0a 100644
--- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py
+++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py
@@ -25,11 +25,19 @@ class AgentAccessGroupCeiling:
models: frozenset[str]
mcp_server_ids: frozenset[str]
agent_ids: frozenset[str]
+ access_groups: tuple[LiteLLM_AccessGroupTable, ...] = ()
CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]]
+class _MissingGroup:
+ __slots__ = ()
+
+
+_MISSING: Final = _MissingGroup()
+
+
async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
@@ -38,7 +46,7 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
async def _load_access_group(access_group_id: str, *, check_db_only: bool = False) -> LoadedAccessGroup:
- from litellm.proxy.auth.auth_checks import get_access_object
+ from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
@@ -52,6 +60,8 @@ async def _load_access_group(access_group_id: str, *, check_db_only: bool = Fals
proxy_logging_obj=proxy_logging_obj,
check_db_only=check_db_only,
)
+ except AccessGroupNotFoundError:
+ raise
except HTTPException as e:
if check_db_only:
raise
@@ -68,23 +78,35 @@ async def resolve_agent_access_group_ceiling(
*,
check_db_only: bool = False,
) -> AgentAccessGroupCeiling | None:
- """``None`` when the agent has no access groups attached, so nothing is capped."""
- access_group_ids: Final = await load_access_group_ids(agent_id)
- if not access_group_ids:
+ """``None`` when the agent has no access groups attached, so nothing is capped. A named group that
+ provably does not exist is dropped from the ceiling, it grants and bounds nothing, while a group
+ that exists but could not be read stays named and loads as ``None``. The ceiling itself stays once
+ any group is named, so an agent whose every group is missing is still capped to nothing."""
+ named_ids: Final = await load_access_group_ids(agent_id)
+ if not named_ids:
return None
- loaded: Final = await asyncio.gather(
- *(
- _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id)
- for group_id in access_group_ids
- )
+ from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError
+
+ async def load(group_id: str) -> LoadedAccessGroup | _MissingGroup:
+ try:
+ return await (
+ _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id)
+ )
+ except AccessGroupNotFoundError:
+ return _MISSING
+
+ results: Final = await asyncio.gather(*(load(group_id) for group_id in named_ids))
+ access_group_ids: Final = tuple(
+ group_id for group_id, result in zip(named_ids, results, strict=True) if not isinstance(result, _MissingGroup)
)
- groups: Final = tuple(group for group in loaded if group is not None)
+ groups: Final = tuple(result for result in results if isinstance(result, LiteLLM_AccessGroupTable))
return AgentAccessGroupCeiling(
access_group_ids=access_group_ids,
models=frozenset(model for group in groups for model in group.access_model_names),
mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids),
agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids),
+ access_groups=groups,
)
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 5568ae2995c..1a12e1b21e0 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -3411,6 +3411,19 @@ class TeamNotFoundError(HTTPException):
super().__init__(status_code=404, detail=detail)
+class AccessGroupNotFoundError(HTTPException):
+ """The access group row is provably absent, as opposed to merely unreadable. ``get_access_object``
+ reports every failure as a 404, so a group that was never created and a database that would not
+ answer look the same to its callers. Policy code that must not treat a degraded read as a definitive
+ answer keys on this subclass. It stays a 404 with the same detail, so every other caller is unaffected."""
+
+ def __init__(self, access_group_id: str) -> None:
+ super().__init__(
+ status_code=404,
+ detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}."},
+ )
+
+
@log_db_metrics
async def _get_team_db_check(
team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None, *, use_writer: bool = False
@@ -3667,10 +3680,7 @@ async def get_access_object(
).find_unique(where={"access_group_id": access_group_id})
if response is None:
- raise HTTPException(
- status_code=404,
- detail={"error": f"Access group doesn't exist in db. Access group={access_group_id}."},
- )
+ raise AccessGroupNotFoundError(access_group_id=access_group_id)
_response: Final = LiteLLM_AccessGroupTable.model_validate(response.dict())
@@ -4532,6 +4542,56 @@ async def get_org_object_for_request(
return org
+async def get_access_group_tables(
+ access_group_ids: Sequence[str],
+ prisma_client: DatabaseClient | None = None,
+ user_api_key_cache: UserApiKeyCache | None = None,
+ proxy_logging_obj: ProxyLogging | None = None,
+ check_db_only: bool = False,
+ *,
+ unreadable: Literal["skip", "raise"] = "skip",
+) -> tuple[LiteLLM_AccessGroupTable, ...]:
+ """Load unified access groups by id (cache first, then DB). A group that provably does not exist is
+ skipped, it grants nothing. A group that exists but cannot be read is skipped too unless
+ ``check_db_only`` is set or ``unreadable`` is ``"raise"``, in which case the failure propagates so a
+ caller enforcing policy can deny instead of shedding a restriction it could not read."""
+ if not access_group_ids:
+ return ()
+
+ if prisma_client is None or user_api_key_cache is None:
+ from litellm.proxy.proxy_server import prisma_client as _prisma_client
+ from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj
+ from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache
+
+ prisma_client = prisma_client or _prisma_client # rebind-ok: lazy proxy_server fallback
+ user_api_key_cache = user_api_key_cache or _user_api_key_cache # rebind-ok: lazy proxy_server fallback
+ proxy_logging_obj = proxy_logging_obj or _proxy_logging_obj # rebind-ok: lazy proxy_server fallback
+
+ if user_api_key_cache is None:
+ return ()
+
+ async def load(ag_id: str) -> LiteLLM_AccessGroupTable | None:
+ try:
+ return await get_access_object(
+ access_group_id=ag_id,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ check_db_only=check_db_only,
+ )
+ except AccessGroupNotFoundError:
+ verbose_proxy_logger.debug("One of the requested access groups does not exist, it grants nothing")
+ return None
+ except Exception:
+ if check_db_only or unreadable == "raise":
+ raise
+ verbose_proxy_logger.debug("Could not fetch one of the requested access groups")
+ return None
+
+ loaded: Final = [await load(ag_id) for ag_id in access_group_ids]
+ return tuple(ag for ag in loaded if ag is not None)
+
+
async def _get_resources_from_access_groups(
access_group_ids: Sequence[str],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
@@ -4540,59 +4600,30 @@ async def _get_resources_from_access_groups(
proxy_logging_obj: ProxyLogging | None = None,
check_db_only: bool = False,
) -> list[str]:
- """
- Fetch access groups by their IDs (from cache or DB) and collect
- the specified resource field across all of them.
+ """Deduplicated ``resource_field`` values across the resolved access groups."""
+ groups: Final = await get_access_group_tables(
+ access_group_ids=access_group_ids,
+ prisma_client=prisma_client,
+ user_api_key_cache=user_api_key_cache,
+ proxy_logging_obj=proxy_logging_obj,
+ check_db_only=check_db_only,
+ )
+ return list(
+ {resource for ag in groups for resource in _access_group_field(ag, resource_field)} # comprehension-ok: flatten
+ )
- Args:
- access_group_ids: List of access group IDs to fetch
- resource_field: Which resource list to extract from each access group
- - "access_model_names": model names (for model access checks)
- - "access_mcp_server_ids": MCP server IDs (for MCP access checks)
- - "access_agent_ids": agent IDs (for agent access checks)
- prisma_client: Optional PrismaClient (lazy-imported from proxy_server if None)
- user_api_key_cache: Optional DualCache (lazy-imported from proxy_server if None)
- proxy_logging_obj: Optional ProxyLogging (lazy-imported from proxy_server if None)
- Returns:
- Deduplicated list of resource identifiers from all resolved access groups.
- """
- if not access_group_ids:
- return []
-
- # Lazy import to avoid circular imports
- if prisma_client is None or user_api_key_cache is None:
- from litellm.proxy.proxy_server import prisma_client as _prisma_client
- from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj
- from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache
-
- prisma_client = prisma_client or _prisma_client
- user_api_key_cache = user_api_key_cache or _user_api_key_cache
- proxy_logging_obj = proxy_logging_obj or _proxy_logging_obj
-
- if user_api_key_cache is None:
- return []
-
- resources: Final[list[str]] = []
- for ag_id in access_group_ids:
- try:
- ag = await get_access_object(
- access_group_id=ag_id,
- prisma_client=prisma_client,
- user_api_key_cache=user_api_key_cache,
- proxy_logging_obj=proxy_logging_obj,
- check_db_only=check_db_only,
- )
- resources.extend(getattr(ag, resource_field, []))
- except Exception:
- if check_db_only:
- raise
- verbose_proxy_logger.debug(
- "Could not fetch access group %s for resource field %s",
- ag_id,
- resource_field,
- )
- return list(set(resources))
+def _access_group_field(
+ group: LiteLLM_AccessGroupTable,
+ resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
+) -> list[str]:
+ match resource_field:
+ case "access_model_names":
+ return group.access_model_names
+ case "access_mcp_server_ids":
+ return group.access_mcp_server_ids
+ case "access_agent_ids":
+ return group.access_agent_ids
async def get_models_from_access_groups(
diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py
index 3f05c3a407a..5e0dbb3297f 100644
--- a/litellm/proxy/management_endpoints/access_group_endpoints.py
+++ b/litellm/proxy/management_endpoints/access_group_endpoints.py
@@ -2,7 +2,7 @@ import asyncio
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
-from typing import Final, Protocol
+from typing import TYPE_CHECKING, Final, Protocol
from fastapi import APIRouter, Depends, HTTPException, status
from typing_extensions import ReadOnly, TypedDict
@@ -45,6 +45,9 @@ from litellm.types.access_group import (
AccessGroupUpdateRequest,
)
+if TYPE_CHECKING:
+ from prisma import Json
+
router: Final = APIRouter(
tags=["access group management"],
)
@@ -57,6 +60,9 @@ class _AccessGroupRecord(Protocol):
@property
def access_mcp_server_ids(self) -> Sequence[str] | None: ...
+ @property
+ def mcp_tool_permissions(self) -> "Mapping[str, Sequence[str]] | Json | None": ...
+
@property
def access_agent_ids(self) -> Sequence[str] | None: ...
@@ -179,6 +185,12 @@ def _require_admin_view(user_api_key_dict: UserAPIKeyAuth) -> None:
)
+def _tool_permissions_json(value: Mapping[str, Sequence[str]] | None) -> "Json":
+ from prisma import Json # noqa: PLC0415 # local import: prisma may be ungenerated at module load in some tools
+
+ return Json(None) if value is None else Json({server: list(tools) for server, tools in value.items()})
+
+
@dataclass(frozen=True, slots=True)
class _ResourceNames:
mcp_servers: Mapping[str, str]
@@ -544,6 +556,7 @@ async def create_access_group(
"description": data.description,
"access_model_names": data.access_model_names or [],
"access_mcp_server_ids": data.access_mcp_server_ids or [],
+ "mcp_tool_permissions": _tool_permissions_json(data.mcp_tool_permissions),
"access_agent_ids": data.access_agent_ids or [],
"assigned_team_ids": data.assigned_team_ids or [],
"assigned_key_ids": data.assigned_key_ids or [],
@@ -648,7 +661,9 @@ async def update_access_group(
and value is None
):
value = []
- update_data[field] = value
+ update_data[field] = (
+ _tool_permissions_json(data.mcp_tool_permissions) if field == "mcp_tool_permissions" else value
+ )
# Initialize delta lists before the try block so they remain accessible
# for cache updates after the transaction, even if an error path is added later.
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index ca58247a2c9..dcae56b5e71 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -1588,6 +1588,7 @@ model LiteLLM_AccessGroupTable {
// Resource memberships - explicit arrays per type
access_model_names String[] @default([])
access_mcp_server_ids String[] @default([])
+ mcp_tool_permissions Json?
access_agent_ids String[] @default([])
assigned_team_ids String[] @default([])
diff --git a/litellm/types/access_group.py b/litellm/types/access_group.py
index b4bf345ecbe..d6d2a312f45 100644
--- a/litellm/types/access_group.py
+++ b/litellm/types/access_group.py
@@ -1,3 +1,4 @@
+from collections.abc import Mapping, Sequence
from datetime import datetime
from litellm.types.llms.base import LiteLLMBaseModel
@@ -8,6 +9,7 @@ class AccessGroupCreateRequest(LiteLLMBaseModel):
description: str | None = None
access_model_names: list[str] | None = None
access_mcp_server_ids: list[str] | None = None
+ mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
access_agent_ids: list[str] | None = None
assigned_team_ids: list[str] | None = None
assigned_key_ids: list[str] | None = None
@@ -18,6 +20,7 @@ class AccessGroupUpdateRequest(LiteLLMBaseModel):
description: str | None = None
access_model_names: list[str] | None = None
access_mcp_server_ids: list[str] | None = None
+ mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
access_agent_ids: list[str] | None = None
assigned_team_ids: list[str] | None = None
assigned_key_ids: list[str] | None = None
@@ -36,6 +39,7 @@ class AccessGroupResponse(LiteLLMBaseModel):
description: str | None = None
access_model_names: list[str]
access_mcp_server_ids: list[str]
+ mcp_tool_permissions: Mapping[str, Sequence[str]] | None = None
access_agent_ids: list[str]
assigned_team_ids: list[str]
assigned_key_ids: list[str]
diff --git a/schema.prisma b/schema.prisma
index ca58247a2c9..dcae56b5e71 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -1588,6 +1588,7 @@ model LiteLLM_AccessGroupTable {
// Resource memberships - explicit arrays per type
access_model_names String[] @default([])
access_mcp_server_ids String[] @default([])
+ mcp_tool_permissions Json?
access_agent_ids String[] @default([])
assigned_team_ids String[] @default([])
diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py
index 90403e5553f..1524aad5b66 100644
--- a/tests/unit/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py
+++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_managed_agent_access.py
@@ -7,7 +7,7 @@ from fastapi import HTTPException
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import mcp_server_manager
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
-from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_UserTable, UserAPIKeyAuth
+from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth
from litellm.proxy.auth import auth_checks
from litellm.types.agents import AgentResponse
from litellm.types.proxy.agent_identity import ManagedAgentContext
@@ -114,6 +114,41 @@ async def test_access_groups_cap_agent_servers_without_granting_new_ones(
assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == []
+@pytest.mark.asyncio
+async def test_managed_agent_tool_ceiling_reads_the_admitted_policy_not_the_registry(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from litellm.proxy._types import LiteLLM_AccessGroupTable
+ from litellm.proxy.common_utils import registry_read_through
+
+ groups: Final = {
+ "stale": LiteLLM_AccessGroupTable(
+ access_group_id="stale", access_group_name="Stale", access_mcp_server_ids=["slack"]
+ ),
+ "current": LiteLLM_AccessGroupTable(
+ access_group_id="current",
+ access_group_name="Current",
+ access_mcp_server_ids=["slack"],
+ mcp_tool_permissions={"slack": ["read"]},
+ ),
+ }
+
+ async def load_group(access_group_id: str, **_: object) -> LiteLLM_AccessGroupTable:
+ return groups[access_group_id]
+
+ monkeypatch.setattr(auth_checks, "get_access_object", load_group)
+ auth: Final = actor(("read", "write"))
+ assert auth.managed_agent_policy is not None
+ stale_registry_row: Final = auth.managed_agent_policy.model_copy(update={"access_group_ids": ["stale"]})
+ monkeypatch.setattr(
+ registry_read_through, "get_agent_with_read_through", AsyncMock(return_value=stale_registry_row)
+ )
+ auth.managed_agent_policy = auth.managed_agent_policy.model_copy(update={"access_group_ids": ["current"]})
+
+ assert tuple(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == ("slack",)
+ assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == ["read"]
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("change", ["tools", "servers", "disabled", "outage"])
async def test_delegated_mcp_revokes_warm_human_policy_before_tool_execution(
@@ -400,9 +435,9 @@ async def test_delegation_never_borrows_another_teams_server_or_tools(
team_id=name,
models=[],
members_with_roles=[{"user_id": "human", "role": "user"}],
- object_permission=permission if name == "selected" else LiteLLM_ObjectPermissionTable(
- object_permission_id="other-grant", mcp_servers=["slack", "linear"]
- ),
+ object_permission=permission
+ if name == "selected"
+ else LiteLLM_ObjectPermissionTable(object_permission_id="other-grant", mcp_servers=["slack", "linear"]),
)
for name in ("selected", "other")
}
@@ -439,7 +474,9 @@ async def test_managed_mcp_rejects_unavailable_authoritative_entitlements(
)
auth: Final = actor(None)
assert auth.managed_agent_policy is not None
- auth.managed_agent_policy = auth.managed_agent_policy.model_copy(update={"object_permission": permission.model_dump()})
+ auth.managed_agent_policy = auth.managed_agent_policy.model_copy(
+ update={"object_permission": permission.model_dump()}
+ )
auth.requires_fresh_policy = True
with pytest.raises(HTTPException) as failure:
await MCPRequestHandler.get_allowed_tools_for_server("slack", auth)
@@ -484,6 +521,7 @@ async def test_managed_agent_mcp_access_is_capped_at_the_invoking_callers_grants
"_apply_user_tool_ceiling",
AsyncMock(side_effect=lambda tools, _server_id, _auth: tools),
)
+ monkeypatch.setattr(auth_checks, "get_team_object", AsyncMock(return_value=LiteLLM_TeamTable(team_id="callers")))
auth: Final = actor(("read", "write"))
auth.agent_caller = AgentCaller(user_id="alice", team_id="callers")
@@ -496,25 +534,34 @@ async def test_managed_agent_mcp_access_is_capped_at_the_invoking_callers_grants
@pytest.mark.parametrize("fresh", [False, True])
@pytest.mark.parametrize("caller_kind", ["team", "user"])
async def test_caller_mcp_revocation_uses_fresh_policy(
- monkeypatch: pytest.MonkeyPatch, fresh: bool, caller_kind: str,
+ monkeypatch: pytest.MonkeyPatch,
+ fresh: bool,
+ caller_kind: str,
) -> None:
from litellm.proxy._types import LiteLLM_TeamTable
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key
from litellm.types.agents import AgentCaller
cached_permission: Final = LiteLLM_ObjectPermissionTable(
- object_permission_id="caller-permission", mcp_servers=["slack", "linear"],
+ object_permission_id="caller-permission",
+ mcp_servers=["slack", "linear"],
mcp_tool_permissions={"slack": ["read", "write"]},
)
current_permission: Final = LiteLLM_ObjectPermissionTable(
- object_permission_id="caller-permission", mcp_servers=["slack"],
+ object_permission_id="caller-permission",
+ mcp_servers=["slack"],
mcp_tool_permissions={"slack": ["read"]},
)
team: Final = LiteLLM_TeamTable(
- team_id="caller", object_permission_id="caller-permission", object_permission=current_permission,
+ team_id="caller",
+ object_permission_id="caller-permission",
+ object_permission=current_permission,
)
user: Final = LiteLLM_UserTable(
- user_id="caller", teams=[], object_permission_id="caller-permission", object_permission=current_permission,
+ user_id="caller",
+ teams=[],
+ object_permission_id="caller-permission",
+ object_permission=current_permission,
)
database: Final = MagicMock()
database.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team)
@@ -531,13 +578,16 @@ async def test_caller_mcp_revocation_uses_fresh_policy(
auth.agent_caller = AgentCaller(team_id="caller") if caller_kind == "team" else AgentCaller(user_id="caller")
assert set(await MCPRequestHandler.get_allowed_mcp_servers(auth)) == ({"slack"} if fresh else {"slack", "linear"})
- assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == (["read"] if fresh else ["read", "write"])
+ assert await MCPRequestHandler.get_allowed_tools_for_server("slack", auth) == (
+ ["read"] if fresh else ["read", "write"]
+ )
@pytest.mark.asyncio
@pytest.mark.parametrize("fresh", [False, True])
async def test_caller_team_outage_cannot_remove_authoritative_server_ceiling(
- monkeypatch: pytest.MonkeyPatch, fresh: bool,
+ monkeypatch: pytest.MonkeyPatch,
+ fresh: bool,
) -> None:
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.types.agents import AgentCaller
diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
index 078fff0266b..4e2597508fc 100644
--- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
@@ -1,6 +1,8 @@
import contextlib
import json
import os
+from collections.abc import Sequence
+from collections.abc import Set as AbstractSet
from datetime import datetime, timedelta, timezone
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@@ -10,6 +12,8 @@ from fastapi import HTTPException
from fastapi.testclient import TestClient
from starlette.datastructures import Headers
+from litellm.constants import MCP_ALL_TOOLS_WILDCARD
+from litellm.models.access_group import LiteLLM_AccessGroupTable
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
UnloadableEntitlementError,
@@ -18,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
)
from litellm.proxy._types import (
LiteLLM_ObjectPermissionTable,
+ LiteLLM_TeamTable,
ProxyException,
SpecialHeaders,
SpecialMCPServerNames,
@@ -26,17 +31,21 @@ from litellm.proxy._types import (
from litellm.types.agents import AgentCaller
-@pytest.mark.parametrize("headers,expected", [
- ({"Authorization": "Bearer caller.jwt.signature"}, "caller.jwt.signature"),
- ({"AUTHORIZATION": "bEaReR caller.jwt.signature"}, "caller.jwt.signature"),
- ({"authorization": "Bearer "}, ""),
- ({"authorization": "Basic credentials"}, None),
- ({"authorization-extra": "Bearer wrong.jwt.signature"}, None),
- ({"x-litellm-api-key": "sk-key"}, None),
- ({}, None),
-])
+@pytest.mark.parametrize(
+ "headers,expected",
+ [
+ ({"Authorization": "Bearer caller.jwt.signature"}, "caller.jwt.signature"),
+ ({"AUTHORIZATION": "bEaReR caller.jwt.signature"}, "caller.jwt.signature"),
+ ({"authorization": "Bearer "}, ""),
+ ({"authorization": "Basic credentials"}, None),
+ ({"authorization-extra": "Bearer wrong.jwt.signature"}, None),
+ ({"x-litellm-api-key": "sk-key"}, None),
+ ({}, None),
+ ],
+)
def test_incoming_guardrail_bearer_requires_exact_authorization_header(
- headers: dict[str, str], expected: str | None,
+ headers: dict[str, str],
+ expected: str | None,
) -> None:
token: Final = MCPRequestHandler.get_incoming_bearer_token(headers)
assert (token.get_secret_value() if token is not None else None) == expected
@@ -8372,7 +8381,9 @@ class TestUserSubjectTeamUnion:
assert servers == ["srv1"]
assert tools == ["add"]
assert resolve.await_args_list
- assert all(call.kwargs == {"toolset_ids": ["ts-1"], "requires_fresh_policy": True} for call in resolve.await_args_list)
+ assert all(
+ call.kwargs == {"toolset_ids": ["ts-1"], "requires_fresh_policy": True} for call in resolve.await_args_list
+ )
async def test_key_based_caller_uses_single_team_only(self):
"""A key-based caller (api_key set) with a team_id sees ONLY that team, even though the
@@ -9564,9 +9575,7 @@ class TestSessionBearerEgressScrub:
) -> None:
caller_key: Final = "sk-caller-admission-key-123"
upstream_token: Final = "Bearer real-upstream-token"
- oauth2_headers: Final = {
- "Authorization": caller_value if context == "oauth2_authorization" else upstream_token
- }
+ oauth2_headers: Final = {"Authorization": caller_value if context == "oauth2_authorization" else upstream_token}
raw_headers: Final = {
"x-litellm-api-key": f"Bearer {caller_key}",
"authorization": caller_value if context == "raw_authorization" else upstream_token,
@@ -9823,7 +9832,9 @@ class TestSessionBearerEgressScrub:
{},
)
- async def test_caller_admission_credential_falls_back_to_standard_headers_when_custom_header_is_absent(self) -> None:
+ async def test_caller_admission_credential_falls_back_to_standard_headers_when_custom_header_is_absent(
+ self,
+ ) -> None:
caller_value: Final = "Bearer sk-caller-admission-key-123"
headers: Final = Headers(
{
@@ -10243,7 +10254,11 @@ class TestGetUserObjectPermission:
prisma_client = MagicMock()
from litellm.proxy._types import LiteLLM_UserTable
- row = LiteLLM_UserTable(user_id="human", object_permission_id=user_row.object_permission_id) if user_row is not None else None
+ row = (
+ LiteLLM_UserTable(user_id="human", object_permission_id=user_row.object_permission_id)
+ if user_row is not None
+ else None
+ )
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=row)
return prisma_client
@@ -10259,7 +10274,10 @@ class TestGetUserObjectPermission:
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
- patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))),
+ patch(
+ "litellm.proxy.proxy_server.proxy_logging_obj",
+ MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())),
+ ),
patch(
"litellm.proxy.auth.auth_checks.get_object_permission",
new_callable=AsyncMock,
@@ -10286,7 +10304,10 @@ class TestGetUserObjectPermission:
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
- patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))),
+ patch(
+ "litellm.proxy.proxy_server.proxy_logging_obj",
+ MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())),
+ ),
patch("litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock) as mock_get_perm,
):
assert await MCPRequestHandler._get_user_object_permission(auth) is None
@@ -10308,7 +10329,10 @@ class TestGetUserObjectPermission:
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
- patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))),
+ patch(
+ "litellm.proxy.proxy_server.proxy_logging_obj",
+ MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())),
+ ),
):
assert await MCPRequestHandler._get_user_object_permission(auth) is None
@@ -10322,7 +10346,10 @@ class TestGetUserObjectPermission:
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
- patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))),
+ patch(
+ "litellm.proxy.proxy_server.proxy_logging_obj",
+ MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())),
+ ),
):
assert await MCPRequestHandler._get_user_object_permission(auth) is None
@@ -10339,7 +10366,10 @@ class TestGetUserObjectPermission:
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
- patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock()))),
+ patch(
+ "litellm.proxy.proxy_server.proxy_logging_obj",
+ MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())),
+ ),
patch(
"litellm.proxy.auth.auth_checks.get_object_permission",
new_callable=AsyncMock,
@@ -10759,7 +10789,9 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_UserTable
- current = LiteLLM_UserTable(user_id="catalog-user", organization_id="current-org", user_role="internal_user", teams=[])
+ current = LiteLLM_UserTable(
+ user_id="catalog-user", organization_id="current-org", user_role="internal_user", teams=[]
+ )
table = SimpleNamespace(find_unique=AsyncMock(return_value=current))
database = SimpleNamespace(writer_db=SimpleNamespace(litellm_usertable=table))
monkeypatch.setattr(proxy_server, "prisma_client", database)
@@ -10779,10 +10811,25 @@ async def test_catalog_refresh_reads_current_user_org_without_losing_resource_sc
@pytest.mark.parametrize("current_groups", [[], ["replacement-group"]])
async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session_scope(monkeypatch, current_groups):
permission = LiteLLM_ObjectPermissionTable(object_permission_id="current-policy", mcp_servers=["current-server"])
- current = UserAPIKeyAuth(object_permission=permission, object_permission_id="current-policy", team_id="new-team", org_id="new-org", project_id="new-project", user_id="new-owner", access_group_ids=current_groups)
+ current = UserAPIKeyAuth(
+ object_permission=permission,
+ object_permission_id="current-policy",
+ team_id="new-team",
+ org_id="new-org",
+ project_id="new-project",
+ user_id="new-owner",
+ access_group_ids=current_groups,
+ )
reload_key = AsyncMock(return_value=current)
monkeypatch.setattr(MCPRequestHandler, "_reload_admitted_key", reload_key)
- caller = UserAPIKeyAuth(api_key="owned-key-hash", team_id="old-team", org_id="old-org", project_id="old-project", user_id="old-owner", access_group_ids=["original-group"])
+ caller = UserAPIKeyAuth(
+ api_key="owned-key-hash",
+ team_id="old-team",
+ org_id="old-org",
+ project_id="old-project",
+ user_id="old-owner",
+ access_group_ids=["original-group"],
+ )
caller.via_virtual_key = True
caller.mcp_session_resource_server_id = "session-server"
caller.mcp_toolset_id = "session-toolset"
@@ -10790,7 +10837,12 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session
reload_key.assert_awaited_once_with("owned-key-hash", check_db_only=True)
assert refreshed.object_permission == permission
assert refreshed.object_permission_id == "current-policy"
- assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == ("new-team", "new-org", "new-project", "new-owner")
+ assert (refreshed.team_id, refreshed.org_id, refreshed.project_id, refreshed.user_id) == (
+ "new-team",
+ "new-org",
+ "new-project",
+ "new-owner",
+ )
assert refreshed.mcp_session_resource_server_id == "session-server"
assert refreshed.mcp_toolset_id == "session-toolset"
assert refreshed.via_virtual_key and refreshed.requires_fresh_policy
@@ -10820,33 +10872,40 @@ async def test_admission_request_body_serves_stashed_peek_callable():
)
assert await with_peek.body() == admission_body
- without_peek: Final = _admission_request(
- {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
- )
+ without_peek: Final = _admission_request({"type": "http", "method": "POST", "path": "/mcp", "headers": []})
assert await without_peek.body() == b"{}"
-@pytest.mark.parametrize("bearer,api_key,claims,explicit_key,expected", [
- (None, "sk-key", None, None, None),
- ("sk-key", "sk-key", None, None, None),
- ("sk-key", "sk-key", {"sub": "previous-jwt"}, "sk-key", None),
- ("caller.jwt.assertion", "sk-key", {"sub": "caller"}, None, "caller.jwt.assertion"),
- ("caller.jwt.assertion", None, {"sub": "caller"}, None, "caller.jwt.assertion"),
- ("upstream", "sk-key", None, "sk-key", "upstream"),
- ("llm_session_test", None, {"sub": "caller"}, None, None),
- ("llm_srefresh_test", None, {"sub": "caller"}, None, None),
- ("opaque-master", "litellm_proxy_master_key", None, None, None),
-])
+@pytest.mark.parametrize(
+ "bearer,api_key,claims,explicit_key,expected",
+ [
+ (None, "sk-key", None, None, None),
+ ("sk-key", "sk-key", None, None, None),
+ ("sk-key", "sk-key", {"sub": "previous-jwt"}, "sk-key", None),
+ ("caller.jwt.assertion", "sk-key", {"sub": "caller"}, None, "caller.jwt.assertion"),
+ ("caller.jwt.assertion", None, {"sub": "caller"}, None, "caller.jwt.assertion"),
+ ("upstream", "sk-key", None, "sk-key", "upstream"),
+ ("llm_session_test", None, {"sub": "caller"}, None, None),
+ ("llm_srefresh_test", None, {"sub": "caller"}, None, None),
+ ("opaque-master", "litellm_proxy_master_key", None, None, None),
+ ],
+)
def test_guardrail_bearer_preserves_identity_without_exposing_gateway_credentials(
- bearer: str | None, api_key: str | None, claims: dict[str, str] | None,
- explicit_key: str | None, expected: str | None,
+ bearer: str | None,
+ api_key: str | None,
+ claims: dict[str, str] | None,
+ explicit_key: str | None,
+ expected: str | None,
) -> None:
- headers: Final = Headers({
- **({"Authorization": f"Bearer {bearer}"} if bearer is not None else {}),
- **({"x-litellm-api-key": explicit_key} if explicit_key else {}),
- })
+ headers: Final = Headers(
+ {
+ **({"Authorization": f"Bearer {bearer}"} if bearer is not None else {}),
+ **({"x-litellm-api-key": explicit_key} if explicit_key else {}),
+ }
+ )
token: Final = MCPRequestHandler.get_guardrail_bearer_token(
- headers, UserAPIKeyAuth(api_key=api_key, jwt_claims=claims),
+ headers,
+ UserAPIKeyAuth(api_key=api_key, jwt_claims=claims),
)
assert (token.get_secret_value() if token is not None else None) == expected
@@ -10860,3 +10919,545 @@ def test_guardrail_bearer_respects_custom_admission_header(monkeypatch: pytest.M
UserAPIKeyAuth(api_key="litellm_proxy_master_key"),
)
assert token is None
+
+
+def _access_group(group_id, servers, tool_permissions=None, **owners):
+ from litellm.models.access_group import LiteLLM_AccessGroupTable
+
+ return LiteLLM_AccessGroupTable(
+ access_group_id=group_id,
+ access_group_name=group_id,
+ access_mcp_server_ids=list(servers),
+ mcp_tool_permissions=tool_permissions,
+ **owners,
+ )
+
+
+class TestAccessGroupToolsForServer:
+ """Tool grants of unified access groups on one server, before they meet the key/team ceilings."""
+
+ def test_no_group_grants_the_server_is_no_restriction(self):
+ groups = [_access_group("g", ["other"], {"other": ["x"], "srv": ["leak"]})]
+ assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") is None
+
+ def test_granting_group_without_entry_allows_every_tool(self):
+ groups = [_access_group("g", ["srv"])]
+ assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") is None
+
+ def test_granting_group_with_explicit_empty_list_grants_no_tools(self):
+ groups = [_access_group("g", ["srv"], {"srv": []})]
+ assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") == ()
+
+ def test_tools_union_across_granting_groups(self):
+ groups = [_access_group("a", ["srv"], {"srv": ["read"]}), _access_group("b", ["srv"], {"srv": ["write"]})]
+ assert sorted(MCPRequestHandler.access_group_tools_for_server(groups, "srv")) == ["read", "write"]
+
+ def test_unrestricted_granting_group_wins_the_union(self):
+ groups = [_access_group("a", ["srv"], {"srv": ["read"]}), _access_group("b", ["srv"])]
+ assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") is None
+
+ def test_non_granting_group_entry_does_not_widen(self):
+ groups = [_access_group("a", ["srv"], {"srv": []}), _access_group("b", ["other"], {"srv": ["write"]})]
+ assert MCPRequestHandler.access_group_tools_for_server(groups, "srv") == ()
+
+
+@pytest.mark.asyncio
+class TestAccessGroupToolPermissionsInheritance:
+ """Group tool allowlists reach the key, team and agent levels of get_allowed_tools_for_server."""
+
+ @staticmethod
+ def _team(access_group_ids):
+ from litellm.proxy._types import LiteLLM_TeamTable
+
+ return LiteLLM_TeamTable(team_id="team-1", access_group_ids=access_group_ids)
+
+ @staticmethod
+ def _groups_loader(groups):
+ by_id = {group.access_group_id: group for group in groups}
+
+ async def load(*, access_group_id, **_):
+ return by_id[access_group_id]
+
+ return AsyncMock(side_effect=load)
+
+ async def test_key_access_group_tool_allowlist_restricts_the_server(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"])
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]}, assigned_key_ids=[auth.token])]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ async def test_key_access_group_explicit_empty_list_denies_every_tool(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"])
+ groups = [_access_group("g", ["srv"], {"srv": []}, assigned_key_ids=[auth.token])]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_key_server_only_access_group_stays_unrestricted(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"])
+ groups = [_access_group("g", ["srv"], assigned_key_ids=[auth.token])]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None
+
+ async def test_key_object_permission_tools_intersect_with_group_tools(self):
+ auth = _key_auth_reaching("srv", tools=["read", "delete"], access_group_ids=["g"])
+ groups = [_access_group("g", ["srv"], {"srv": ["read", "write"]}, assigned_key_ids=[auth.token])]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)),
+ ):
+ allowed = await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)
+ assert sorted(allowed) == ["delete", "read", "write"]
+
+ async def test_team_access_group_tool_allowlist_caps_the_key(self):
+ auth = _key_auth_reaching("srv", team_id="team-1")
+ team = self._team(["g"])
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]})]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ async def test_team_access_group_explicit_empty_list_denies_even_with_key_tools(self):
+ auth = _key_auth_reaching("srv", tools=["read"], team_id="team-1")
+ team = self._team(["g"])
+ groups = [_access_group("g", ["srv"], {"srv": []})]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_team_and_key_group_tools_intersect(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", team_id="team-1", access_group_ids=["kg"])
+ team = self._team(["tg"])
+ groups = [
+ _access_group("kg", ["srv"], {"srv": ["read", "write"]}, assigned_team_ids=["team-1"]),
+ _access_group("tg", ["srv"], {"srv": ["write", "delete"]}),
+ ]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["write"]
+
+ async def test_agent_access_group_tool_allowlist_caps_an_agent_key(self):
+ auth = _key_auth_reaching("srv", agent_id="agent-1")
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]})]
+ agent = MagicMock(access_group_ids=["g"])
+ with (
+ _entitlement_fault_globals(_agent_prisma()),
+ patch(
+ "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
+ AsyncMock(return_value=agent),
+ ),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._groups_loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+
+class TestAccessGroupToolOwnershipAndFaults:
+ """Group tool grants bind only through groups that own the caller, fail closed when a named group
+ cannot be read, and cap rather than widen an agent's own grants. Every case runs the real
+ ``get_allowed_tools_for_server`` resolver with ``get_access_object`` as the only patched boundary."""
+
+ @staticmethod
+ def _loader(
+ groups: Sequence[LiteLLM_AccessGroupTable],
+ unreadable: AbstractSet[str] = frozenset(),
+ missing: AbstractSet[str] = frozenset(),
+ ) -> AsyncMock:
+ from fastapi import HTTPException
+
+ from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError
+
+ by_id = {group.access_group_id: group for group in groups}
+
+ async def load(*, access_group_id: str, **_: object) -> LiteLLM_AccessGroupTable:
+ if access_group_id in unreadable:
+ raise HTTPException(status_code=404, detail={"error": "Access group lookup failed. Error: db down"})
+ if access_group_id in missing:
+ raise AccessGroupNotFoundError(access_group_id=access_group_id)
+ return by_id[access_group_id]
+
+ return AsyncMock(side_effect=load)
+
+ async def test_unowned_key_group_contributes_no_tools(self):
+ auth = _key_auth_reaching("srv", tools=["read"], access_group_ids=["g"])
+ groups = [_access_group("g", ["srv"], {"srv": ["write"]})]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ async def test_key_group_owned_through_the_team_binds(self):
+ auth = _key_auth_reaching("srv", tools=["read"], team_id="team-1", access_group_ids=["g"])
+ groups = [_access_group("g", ["srv"], {"srv": ["write"]}, assigned_team_ids=["team-1"])]
+ team = LiteLLM_TeamTable(team_id="team-1")
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert sorted(await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)) == ["read", "write"]
+
+ async def test_unreadable_key_group_denies_the_server_tools(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"])
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"g"})),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_unreadable_key_group_denies_even_when_the_team_grants_the_server(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", team_id="team-1", access_group_ids=["g"])
+ team = LiteLLM_TeamTable(
+ team_id="team-1",
+ object_permission=LiteLLM_ObjectPermissionTable(
+ object_permission_id="op-team", mcp_servers=["srv"], mcp_tool_permissions={"srv": ["read", "write"]}
+ ),
+ )
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"g"})),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_unreadable_key_group_keeps_the_direct_allowlist(self):
+ auth = _key_auth_reaching("srv", tools=["read"], access_group_ids=["g"])
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"g"})),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ async def test_unreadable_team_group_denies_the_server_tools(self):
+ auth = _key_auth_reaching("srv", team_id="team-1")
+ team = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["tg"])
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"tg"})),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_nonexistent_key_group_leaves_the_team_grant_alone(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", team_id="team-1", access_group_ids=["typo-group"])
+ team = LiteLLM_TeamTable(
+ team_id="team-1",
+ object_permission=LiteLLM_ObjectPermissionTable(
+ object_permission_id="op-team", mcp_servers=["srv"], mcp_tool_permissions={"srv": ["read", "write"]}
+ ),
+ )
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], missing={"typo-group"})),
+ ):
+ assert sorted(await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)) == ["read", "write"]
+
+ @pytest.mark.parametrize("requires_fresh_policy", [False, True])
+ async def test_nonexistent_key_group_resolves_like_no_group_at_all(self, requires_fresh_policy):
+ without_group = _key_auth_reaching("srv", requires_fresh_policy=requires_fresh_policy)
+ with_typo = _key_auth_reaching(
+ "srv", requires_fresh_policy=requires_fresh_policy, access_group_ids=["typo-group"]
+ )
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], missing={"typo-group"})),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server(
+ "srv", with_typo
+ ) == await MCPRequestHandler.get_allowed_tools_for_server("srv", without_group)
+
+ async def test_nonexistent_key_group_does_not_lift_a_real_attached_bound(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["typo-group", "g"])
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]})]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups, missing={"typo-group"})),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_mcp_server_ids_from_access_groups", AsyncMock(return_value=["srv"])
+ ),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ async def test_nonexistent_team_group_resolves_like_a_team_without_groups(self):
+ auth = _key_auth_reaching("srv", team_id="team-1")
+ with_typo = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["typo-group"])
+ without_group = LiteLLM_TeamTable(team_id="team-1")
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], missing={"typo-group"})),
+ ):
+ with patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=with_typo)):
+ typo_tools = await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)
+ with patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=without_group)):
+ assert typo_tools == await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)
+
+ @pytest.mark.parametrize("group_tools, expected", [(["read"], ["read"]), ([], [])])
+ async def test_unowned_group_only_key_is_bounded_by_the_group_it_is_admitted_through(self, group_tools, expected):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g"])
+ groups = [_access_group("g", ["srv"], {"srv": group_tools})]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ patch(
+ "litellm.proxy.auth.auth_checks.get_mcp_server_ids_from_access_groups", AsyncMock(return_value=["srv"])
+ ),
+ ):
+ assert await MCPRequestHandler.get_allowed_mcp_servers(auth) == ["srv"]
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == expected
+
+ @pytest.mark.parametrize("key_owns_group", [True, False], ids=["owned-group", "attached-only-group"])
+ async def test_explicit_key_wildcard_is_not_narrowed_by_a_group_bound(self, key_owns_group):
+ auth = _key_auth_reaching("srv", tools=[MCP_ALL_TOOLS_WILDCARD], access_group_ids=["g"])
+ owners = {"assigned_key_ids": [auth.token]} if key_owns_group else {}
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]}, **owners)]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None
+
+ async def test_owned_group_without_tool_entry_keeps_the_key_unrestricted(self):
+ auth = UserAPIKeyAuth(api_key="sk-hash", user_id="u1", access_group_ids=["g", "h"])
+ groups = [
+ _access_group("g", ["srv"], assigned_key_ids=[auth.token]),
+ _access_group("h", ["srv"], {"srv": ["read"]}),
+ ]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None
+
+ async def test_direct_server_grant_without_tool_entry_is_not_narrowed_by_an_attached_group(self):
+ auth = _key_auth_reaching("srv", access_group_ids=["g"])
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]})]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None
+
+ async def test_legacy_tag_colliding_with_the_server_name_is_not_a_direct_grant(self):
+ auth = UserAPIKeyAuth(
+ api_key="sk-hash",
+ user_id="u1",
+ access_group_ids=["g"],
+ object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_access_groups=["srv"]),
+ )
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]})]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ async def test_legacy_tag_carried_by_the_server_keeps_the_key_unrestricted(self):
+ from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
+ from litellm.types.mcp import MCPTransport
+ from litellm.types.mcp_server.mcp_server_manager import MCPServer
+
+ auth = UserAPIKeyAuth(
+ api_key="sk-hash",
+ user_id="u1",
+ access_group_ids=["g"],
+ object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_access_groups=["tag"]),
+ )
+ srv = MCPServer(
+ server_id="srv",
+ name="srv",
+ url="https://srv.example.com",
+ transport=MCPTransport.http,
+ access_groups=["tag"],
+ )
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]})]
+ with (
+ _entitlement_fault_globals(),
+ patch.object(global_mcp_server_manager, "config_mcp_servers", {"srv": srv}),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) is None
+
+ async def test_duplicate_key_group_ids_do_not_deny(self):
+ auth = _key_auth_reaching("srv", tools=["read"], access_group_ids=["g", "g"])
+ groups = [_access_group("g", ["srv"], {"srv": ["write"]}, assigned_key_ids=[auth.token])]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert sorted(await MCPRequestHandler.get_allowed_tools_for_server("srv", auth)) == ["read", "write"]
+
+ async def test_duplicate_team_group_ids_do_not_deny(self):
+ auth = _key_auth_reaching("srv", team_id="team-1")
+ team = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["tg", "tg"])
+ groups = [_access_group("tg", ["srv"], {"srv": ["read"]})]
+ with (
+ _entitlement_fault_globals(),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ async def test_duplicate_agent_group_ids_do_not_deny(self):
+ auth = _key_auth_reaching("srv", agent_id="agent-1")
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]})]
+ agent = MagicMock(access_group_ids=["g", "g"])
+ with (
+ _entitlement_fault_globals(self._agent_prisma_with_own_tools(["read", "write"])),
+ patch(
+ "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
+ AsyncMock(return_value=agent),
+ ),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ @staticmethod
+ def _writer_down(*_: object, **__: object) -> None:
+ raise RuntimeError("writer down")
+
+ async def test_fresh_key_group_fault_denies_the_server_tools(self):
+ auth = _key_auth_reaching("srv", tools=["read"], access_group_ids=["g"]).model_copy(
+ update={"requires_fresh_policy": True}
+ )
+ assert auth.requires_fresh_policy is True
+ with (
+ _entitlement_fault_globals(),
+ patch.object(MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=None)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", AsyncMock(side_effect=self._writer_down)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_fresh_team_group_fault_denies_the_server_tools(self):
+ auth = _key_auth_reaching("srv", tools=["read"], team_id="team-1").model_copy(
+ update={"requires_fresh_policy": True}
+ )
+ assert auth.requires_fresh_policy is True
+ team = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["tg"])
+ with (
+ _entitlement_fault_globals(),
+ patch.object(MCPRequestHandler, "_get_user_object_permission", AsyncMock(return_value=None)),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", AsyncMock(side_effect=self._writer_down)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_team_lookup_fault_in_group_resolution_denies_the_server_tools(self):
+ auth = _key_auth_reaching("srv", tools=["read"], team_id="team-1")
+ with (
+ _entitlement_fault_globals(),
+ patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(side_effect=self._writer_down)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_agent_group_fault_denies_rather_than_shedding_the_agents_own_allowlist(self):
+ auth = _key_auth_reaching("srv", agent_id="agent-1")
+ agent = MagicMock(access_group_ids=["g"])
+ with (
+ _entitlement_fault_globals(self._agent_prisma_with_own_tools(["read"])),
+ patch(
+ "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
+ AsyncMock(return_value=agent),
+ ),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", AsyncMock(side_effect=self._writer_down)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_echoed_caller_team_group_tools_cap_the_agent_key(self):
+ from litellm.types.agents import AgentCaller
+
+ auth = _key_auth_reaching("srv", tools=["read", "write"], agent_id="agent-1")
+ auth.agent_caller = AgentCaller(team_id="team-1")
+ team = LiteLLM_TeamTable(team_id="team-1", access_group_ids=["tg"])
+ groups = [_access_group("tg", ["srv"], {"srv": ["read"]})]
+ with (
+ _entitlement_fault_globals(),
+ patch.object(MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None)),
+ patch("litellm.proxy.auth.auth_checks.get_team_object", AsyncMock(return_value=team)),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ capped = await MCPRequestHandler.apply_agent_caller_tool_ceiling(["read", "write"], "srv", auth)
+ assert sorted(capped) == ["read"]
+
+ @staticmethod
+ def _agent_prisma_with_own_tools(tools: Sequence[str]) -> MagicMock:
+ prisma_client = _agent_prisma(object_permission_id="op-agent")
+ prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
+ return_value=LiteLLM_ObjectPermissionTable(
+ object_permission_id="op-agent", mcp_servers=["srv"], mcp_tool_permissions={"srv": tools}
+ )
+ )
+ return prisma_client
+
+ async def test_agent_group_ceiling_caps_the_agents_own_tools(self):
+ auth = _key_auth_reaching("srv", agent_id="agent-1")
+ groups = [_access_group("g", ["srv"], {"srv": ["read"]})]
+ agent = MagicMock(access_group_ids=["g"])
+ with (
+ _entitlement_fault_globals(self._agent_prisma_with_own_tools(["read", "write"])),
+ patch(
+ "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
+ AsyncMock(return_value=agent),
+ ),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
+
+ async def test_agent_group_tools_outside_the_agents_own_grant_are_denied(self):
+ auth = _key_auth_reaching("srv", agent_id="agent-1")
+ groups = [_access_group("g", ["srv"], {"srv": ["write"]})]
+ agent = MagicMock(access_group_ids=["g"])
+ with (
+ _entitlement_fault_globals(self._agent_prisma_with_own_tools(["read"])),
+ patch(
+ "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
+ AsyncMock(return_value=agent),
+ ),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups)),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_unreadable_agent_group_denies_the_server_tools(self):
+ auth = _key_auth_reaching("srv", agent_id="agent-1")
+ agent = MagicMock(access_group_ids=["g"])
+ with (
+ _entitlement_fault_globals(_agent_prisma()),
+ patch(
+ "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
+ AsyncMock(return_value=agent),
+ ),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader([], unreadable={"g"})),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == []
+
+ async def test_nonexistent_agent_group_leaves_the_real_groups_tools_alone(self):
+ auth = _key_auth_reaching("srv", agent_id="agent-1")
+ agent = MagicMock(access_group_ids=["real", "typo-group"])
+ groups = [_access_group("real", ["srv"], {"srv": ["read"]})]
+ with (
+ _entitlement_fault_globals(_agent_prisma()),
+ patch(
+ "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through",
+ AsyncMock(return_value=agent),
+ ),
+ patch("litellm.proxy.auth.auth_checks.get_access_object", self._loader(groups, missing={"typo-group"})),
+ ):
+ assert await MCPRequestHandler.get_allowed_tools_for_server("srv", auth) == ["read"]
diff --git a/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py
index 08cceb0d967..cb9a63512d4 100644
--- a/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py
+++ b/tests/unit/proxy/agent_endpoints/auth/test_agent_access_groups.py
@@ -61,13 +61,11 @@ async def test_unknown_agent_has_no_ceiling():
@pytest.mark.asyncio
async def test_ceiling_is_the_union_of_every_attached_group():
- load_agent, load_group = _loaders(
- _agent(["g1", "g2"]),
- {
- "g1": _group("g1", models=("gpt-5",), mcp_servers=("mcp-a",), agents=("agent-b",)),
- "g2": _group("g2", models=("claude-sonnet",), mcp_servers=("mcp-b",), agents=("agent-c",)),
- },
- )
+ groups: Final = {
+ "g1": _group("g1", models=("gpt-5",), mcp_servers=("mcp-a",), agents=("agent-b",)),
+ "g2": _group("g2", models=("claude-sonnet",), mcp_servers=("mcp-b",), agents=("agent-c",)),
+ }
+ load_agent, load_group = _loaders(_agent(["g1", "g2"]), groups)
ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group)
@@ -76,12 +74,14 @@ async def test_ceiling_is_the_union_of_every_attached_group():
models=frozenset({"gpt-5", "claude-sonnet"}),
mcp_server_ids=frozenset({"mcp-a", "mcp-b"}),
agent_ids=frozenset({"agent-b", "agent-c"}),
+ access_groups=(groups["g1"], groups["g2"]),
)
@pytest.mark.asyncio
async def test_unloadable_group_contributes_nothing_but_the_ceiling_still_applies():
- load_agent, load_group = _loaders(_agent(["g1", "gone"]), {"g1": _group("g1", models=("gpt-5",))})
+ g1: Final = _group("g1", models=("gpt-5",))
+ load_agent, load_group = _loaders(_agent(["g1", "gone"]), {"g1": g1})
ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group)
@@ -90,6 +90,7 @@ async def test_unloadable_group_contributes_nothing_but_the_ceiling_still_applie
models=frozenset({"gpt-5"}),
mcp_server_ids=frozenset(),
agent_ids=frozenset(),
+ access_groups=(g1,),
)
@@ -109,7 +110,8 @@ async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted():
async def test_default_agent_loader_reads_the_attached_groups_from_the_registry():
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
- _, load_group = _loaders(None, {"g1": _group("g1", models=("gpt-5",))})
+ g1: Final = _group("g1", models=("gpt-5",))
+ _, load_group = _loaders(None, {"g1": g1})
global_agent_registry.register_agent(_agent(["g1"]))
try:
ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_access_group=load_group)
@@ -117,7 +119,11 @@ async def test_default_agent_loader_reads_the_attached_groups_from_the_registry(
global_agent_registry.deregister_agent("agent")
assert ceiling == AgentAccessGroupCeiling(
- access_group_ids=("g1",), models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), agent_ids=frozenset()
+ access_group_ids=("g1",),
+ models=frozenset({"gpt-5"}),
+ mcp_server_ids=frozenset(),
+ agent_ids=frozenset(),
+ access_groups=(g1,),
)
@@ -136,6 +142,61 @@ async def test_default_loader_treats_a_missing_group_as_unreadable(monkeypatch:
assert await _load_access_group("gone") is None
+@pytest.mark.asyncio
+async def test_default_loader_lets_a_provably_absent_group_surface(monkeypatch: pytest.MonkeyPatch):
+ from litellm.proxy import proxy_server
+ from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group
+ from litellm.proxy.auth import auth_checks
+ from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError
+
+ async def absent_group(*, access_group_id: str, **_: object) -> LiteLLM_AccessGroupTable:
+ raise AccessGroupNotFoundError(access_group_id=access_group_id)
+
+ monkeypatch.setattr(proxy_server, "prisma_client", object())
+ monkeypatch.setattr(auth_checks, "get_access_object", absent_group)
+
+ with pytest.raises(AccessGroupNotFoundError):
+ await _load_access_group("gone")
+
+
+@pytest.mark.asyncio
+async def test_ceiling_drops_a_group_that_does_not_exist_but_keeps_an_unreadable_one():
+ from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling
+ from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError
+
+ real = LiteLLM_AccessGroupTable(access_group_id="real", access_group_name="real", access_mcp_server_ids=["srv"])
+
+ async def ids(_agent_id: str) -> tuple[str, ...]:
+ return ("real", "typo-group", "unreadable")
+
+ async def load(group_id: str) -> LiteLLM_AccessGroupTable | None:
+ if group_id == "typo-group":
+ raise AccessGroupNotFoundError(access_group_id=group_id)
+ return real if group_id == "real" else None
+
+ ceiling = await resolve_agent_access_group_ceiling("agent-1", ids, load)
+ assert ceiling is not None
+ assert ceiling.access_group_ids == ("real", "unreadable")
+ assert ceiling.access_groups == (real,)
+ assert ceiling.mcp_server_ids == frozenset({"srv"})
+
+
+@pytest.mark.asyncio
+async def test_ceiling_stays_and_caps_to_nothing_when_every_named_group_does_not_exist():
+ from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling
+ from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError
+
+ async def ids(_agent_id: str) -> tuple[str, ...]:
+ return ("typo-group",)
+
+ async def load(group_id: str) -> LiteLLM_AccessGroupTable | None:
+ raise AccessGroupNotFoundError(access_group_id=group_id)
+
+ assert await resolve_agent_access_group_ceiling("agent-1", ids, load) == AgentAccessGroupCeiling(
+ access_group_ids=(), models=frozenset(), mcp_server_ids=frozenset(), agent_ids=frozenset(), access_groups=()
+ )
+
+
@pytest.mark.asyncio
async def test_default_loader_returns_nothing_without_a_db(monkeypatch: pytest.MonkeyPatch):
from litellm.proxy import proxy_server
@@ -159,7 +220,9 @@ async def test_authoritative_group_ceiling_propagates_policy_outages(
database: Final = MagicMock()
database.db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable"))
- database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable"))
+ database.writer_db.litellm_accessgrouptable.find_unique = AsyncMock(
+ side_effect=RuntimeError("database unavailable")
+ )
monkeypatch.setattr(proxy_server, "prisma_client", database)
monkeypatch.setattr(proxy_server, "user_api_key_cache", UserApiKeyCache())
if strict:
diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
index 25022621edc..1f8c85ce337 100644
--- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
+++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
@@ -8856,6 +8856,30 @@ async def test_get_access_object_db_fetch_returns_validated_access_group():
assert result.access_model_names == ["gpt-4"]
+@pytest.mark.asyncio
+async def test_get_access_object_distinguishes_absent_group_from_unreadable_row():
+ from fastapi import HTTPException
+
+ from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object
+
+ mock_cache = MagicMock()
+ mock_cache.async_get_cache = AsyncMock(return_value=None)
+
+ absent_client = MagicMock()
+ absent_client.db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=None)
+ with pytest.raises(AccessGroupNotFoundError) as absent_info:
+ await get_access_object("typo-group", absent_client, mock_cache)
+ assert absent_info.value.status_code == 404
+ assert "Access group doesn't exist in db" in str(absent_info.value.detail)
+
+ unreadable_client = MagicMock()
+ unreadable_client.db.litellm_accessgrouptable.find_unique = AsyncMock(side_effect=ConnectionError("db unreachable"))
+ with pytest.raises(HTTPException) as unreadable_info:
+ await get_access_object("real-group", unreadable_client, mock_cache)
+ assert unreadable_info.value.status_code == 404
+ assert not isinstance(unreadable_info.value, AccessGroupNotFoundError)
+
+
@pytest.mark.asyncio
async def test_get_team_object_by_alias_db_fetch_returns_cached_obj():
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
diff --git a/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py b/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py
index 13e57408bd0..5d7efe66449 100644
--- a/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py
+++ b/tests/unit/proxy/management_endpoints/test_access_group_endpoints.py
@@ -1602,3 +1602,65 @@ def test_update_access_group_response_carries_resolved_names(client_and_mocks):
assert body["access_mcp_servers"] == [{"id": "mcp-new", "name": "Linear"}]
assert body["access_agents"] == [{"id": "agent-a", "name": "support-bot"}]
assert body["access_mcp_server_ids"] == ["mcp-new"]
+
+
+def test_create_access_group_persists_mcp_tool_permissions_as_json(client_and_mocks):
+ """Tool allowlists are stored on the group as JSON, an explicit [] survives as a deny-all entry."""
+ client, _, mock_table, *_ = client_and_mocks
+ tool_permissions = {"mcp-1": ["read_file"], "mcp-2": []}
+
+ resp = client.post(
+ "/v1/access_group",
+ json={
+ "access_group_name": "tools",
+ "access_mcp_server_ids": ["mcp-1", "mcp-2"],
+ "mcp_tool_permissions": tool_permissions,
+ },
+ )
+
+ assert resp.status_code == 201, resp.text
+ written = mock_table.create.await_args.kwargs["data"]
+ assert written["mcp_tool_permissions"].data == tool_permissions
+
+
+def test_create_access_group_without_tool_permissions_writes_null(client_and_mocks):
+ client, _, mock_table, *_ = client_and_mocks
+
+ resp = client.post("/v1/access_group", json={"access_group_name": "servers-only", "access_mcp_server_ids": ["m"]})
+
+ assert resp.status_code == 201, resp.text
+ assert mock_table.create.await_args.kwargs["data"]["mcp_tool_permissions"].data is None
+
+
+@pytest.mark.parametrize("tool_permissions", [{"mcp-1": ["a", "b"]}, {"mcp-1": []}, None])
+def test_update_access_group_writes_mcp_tool_permissions(client_and_mocks, tool_permissions):
+ client, _, mock_table, *_ = client_and_mocks
+ mock_table.find_unique = AsyncMock(return_value=_make_access_group_record(access_group_id="ag-update"))
+
+ resp = client.put("/v1/access_group/ag-update", json={"mcp_tool_permissions": tool_permissions})
+
+ assert resp.status_code == 200, resp.text
+ written = mock_table.update.await_args.kwargs["data"]["mcp_tool_permissions"]
+ assert written.data == tool_permissions
+
+
+def test_update_access_group_leaves_tool_permissions_alone_when_not_sent(client_and_mocks):
+ client, _, mock_table, *_ = client_and_mocks
+ mock_table.find_unique = AsyncMock(return_value=_make_access_group_record(access_group_id="ag-update"))
+
+ resp = client.put("/v1/access_group/ag-update", json={"description": "x"})
+
+ assert resp.status_code == 200, resp.text
+ assert "mcp_tool_permissions" not in mock_table.update.await_args.kwargs["data"]
+
+
+def test_get_access_group_returns_mcp_tool_permissions(client_and_mocks):
+ client, _, mock_table, *_ = client_and_mocks
+ record = _make_access_group_record(access_group_id="ag-tools", access_mcp_server_ids=["mcp-1"])
+ record.dict = lambda: {**record.model_dump(), "mcp_tool_permissions": {"mcp-1": []}}
+ mock_table.find_unique = AsyncMock(return_value=record)
+
+ resp = client.get("/v1/access_group/ag-tools")
+
+ assert resp.status_code == 200, resp.text
+ assert resp.json()["mcp_tool_permissions"] == {"mcp-1": []}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.test.tsx
index aad63e979ac..91a58cef46e 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.test.tsx
@@ -1,6 +1,6 @@
import { useAccessGroupDetails } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails";
import { AccessGroupResponse } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups";
-import { screen } from "@testing-library/react";
+import { screen, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders } from "../../../../../tests/test-utils";
@@ -336,6 +336,34 @@ describe("AccessGroupDetail", () => {
expect(screen.getByText("No MCP servers assigned to this group")).toBeInTheDocument();
});
+
+ it("should list the allowed tools per server, with deny-all and unrestricted servers spelled out", async () => {
+ const user = userEvent.setup();
+ renderWith({
+ access_mcp_servers: [
+ { id: "mcp-1", name: "GitHub MCP" },
+ { id: "mcp-2", name: "Jira MCP" },
+ { id: "mcp-3", name: "Slack MCP" },
+ ],
+ mcp_tool_permissions: { "mcp-1": ["list_issues", "create_issue"], "mcp-2": [] },
+ });
+
+ await user.click(screen.getByRole("tab", { name: /MCP Servers/i }));
+
+ const list = screen.getByRole("region", { name: "MCP tool permissions" });
+ expect(within(list).getByText("list_issues, create_issue")).toBeInTheDocument();
+ expect(within(list).getByText("No tools")).toBeInTheDocument();
+ expect(within(list).getByText("All tools")).toBeInTheDocument();
+ });
+
+ it("should not render a tool permission list when the group has none", async () => {
+ const user = userEvent.setup();
+ renderWith({ mcp_tool_permissions: null });
+
+ await user.click(screen.getByRole("tab", { name: /MCP Servers/i }));
+
+ expect(screen.queryByRole("region", { name: "MCP tool permissions" })).not.toBeInTheDocument();
+ });
});
describe("Agents tab", () => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx
index 118a8655d5a..4f033a6ca02 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsDetailsPage.tsx
@@ -25,6 +25,37 @@ const MAX_PREVIEW = 5;
const shortId = (id: string) => (id.length > 20 ? `${id.slice(0, 10)}...${id.slice(-6)}` : id);
+const describeAllowedTools = (tools: readonly string[] | undefined): string => {
+ if (tools === undefined) {
+ return "All tools";
+ }
+ return tools.length === 0 ? "No tools" : tools.join(", ");
+};
+
+function MCPToolPermissionList({
+ servers,
+ toolPermissions,
+}: {
+ servers: readonly AccessGroupResource[];
+ toolPermissions: Readonly> | null;
+}) {
+ if (toolPermissions === null || servers.length === 0) {
+ return null;
+ }
+ return (
+
+ {servers.map(({ id, name }) => {
+ return (
+
+ {name ?? id}
+ {describeAllowedTools(toolPermissions[id])}
+
+ );
+ })}
+
+ );
+}
+
function ResourceList({ items, emptyMessage }: { items: readonly AccessGroupResource[]; emptyMessage: string }) {
if (items.length === 0) {
return {emptyMessage}
;
@@ -237,6 +268,7 @@ export function AccessGroupDetail({ accessGroupId, onBack }: AccessGroupDetailPr
+
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx
index 91e008de402..c66b7b1f1e3 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx
@@ -6,6 +6,8 @@ import { z } from "zod";
import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents";
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions";
import { ModelSelect } from "@/components/ModelSelect/ModelSelect";
import { FieldGroup } from "@/components/ui/field";
import { FormField } from "@/components/shared/form/FormField";
@@ -19,6 +21,7 @@ export const accessGroupFormSchema = z.object({
description: z.string(),
modelIds: z.array(z.string()),
mcpServerIds: z.array(z.string()),
+ mcpToolPermissions: z.record(z.string(), z.array(z.string())),
agentIds: z.array(z.string()),
});
@@ -34,6 +37,37 @@ interface AccessGroupBaseFormProps {
isNameDisabled?: boolean;
activeTab: string;
onTabChange: (tab: string) => void;
+ initialServerIds?: readonly string[];
+}
+
+export function AccessGroupMCPToolPermissions({
+ selectedServers,
+ toolPermissions,
+ onChange,
+ preserveUnrestrictedServers,
+}: {
+ selectedServers: string[];
+ toolPermissions: Record;
+ onChange: (next: Record) => void;
+ preserveUnrestrictedServers?: readonly string[];
+}) {
+ const { accessToken } = useAuthorized();
+
+ if (selectedServers.length === 0) {
+ return null;
+ }
+
+ return (
+
+
+
+ );
}
export function AccessGroupBaseForm({
@@ -41,6 +75,7 @@ export function AccessGroupBaseForm({
isNameDisabled = false,
activeTab,
onTabChange,
+ initialServerIds,
}: AccessGroupBaseFormProps) {
const { data: agentsData } = useAgents();
const { data: mcpServersData } = useMCPServers();
@@ -108,6 +143,12 @@ export function AccessGroupBaseForm({
/>
)}
+ form.setValue("mcpToolPermissions", next, { shouldDirty: true })}
+ preserveUnrestrictedServers={initialServerIds}
+ />
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx
index 7c4e218261f..ff21872c163 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx
@@ -3,6 +3,7 @@ import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"
import { fireEvent, renderWithProviders, screen, waitFor, within } from "../../../../../../tests/test-utils";
import { AccessGroupEditModal } from "./AccessGroupEditModal";
import { AccessGroupResponse } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups";
+import * as networking from "@/components/networking";
const mutate = vi.fn();
@@ -31,6 +32,8 @@ vi.mock("@/components/ModelSelect/ModelSelect", () => ({
),
}));
+vi.mock("@/components/networking");
+
vi.mock("@/lib/toast", () => ({
toast: { success: vi.fn(), fromError: vi.fn(), error: vi.fn() },
}));
@@ -44,6 +47,7 @@ const accessGroup: AccessGroupResponse = {
description: "Engineers",
access_model_names: ["gpt-4"],
access_mcp_server_ids: ["srv-1"],
+ mcp_tool_permissions: { "srv-1": ["read_file"] },
access_agent_ids: ["agent-1"],
assigned_team_ids: [],
assigned_key_ids: [],
@@ -67,6 +71,18 @@ const variables = () => mutate.mock.calls.at(-1)?.[0] as { accessGroupId: string
describe("AccessGroupEditModal submit payload", () => {
beforeEach(() => {
vi.clearAllMocks();
+ vi.mocked(networking.fetchMCPToolsets).mockResolvedValue([]);
+ vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]);
+ vi.mocked(networking.fetchMCPServers).mockResolvedValue([
+ { server_id: "srv-1", server_name: "Files", alias: "Files" },
+ ]);
+ vi.mocked(networking.listMCPTools).mockResolvedValue({
+ tools: [
+ { name: "read_file", description: "Read a file" },
+ { name: "write_file", description: "Write a file" },
+ ],
+ error: false,
+ });
});
it("sends exactly the antd payload for an untouched save", async () => {
@@ -169,6 +185,20 @@ describe("AccessGroupEditModal submit payload", () => {
expect(mutate).not.toHaveBeenCalled();
});
+ it("sends the edited tool checklist as mcp_tool_permissions, an emptied one as an explicit []", async () => {
+ const user = setup();
+ renderModal();
+ await screen.findByDisplayValue("Engineering");
+
+ await user.click(screen.getByRole("tab", { name: /MCP Servers/ }));
+ await user.click(await screen.findByRole("radio", { name: "Flat List" }));
+ await user.click(await screen.findByRole("checkbox", { name: "read_file" }));
+ await save(user);
+
+ await waitFor(() => expect(mutate).toHaveBeenCalled());
+ expect(variables().params.mcp_tool_permissions).toStrictEqual({ "srv-1": [] });
+ });
+
it("renders each selected MCP server as its own removable chip and drops one on remove", async () => {
const user = setup();
renderModal();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx
index 7370b3bc195..16bcf32d2ea 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.tsx
@@ -1,6 +1,7 @@
"use client";
import React, { useState } from "react";
+import { useFormState } from "react-hook-form";
import { toast } from "@/lib/toast";
import { useZodForm } from "@/lib/forms/useZodForm";
@@ -31,11 +32,13 @@ const toFormValues = (accessGroup: AccessGroupResponse): AccessGroupFormValues =
description: accessGroup.description ?? "",
modelIds: accessGroup.access_model_names ?? [],
mcpServerIds: accessGroup.access_mcp_server_ids ?? [],
+ mcpToolPermissions: accessGroup.mcp_tool_permissions ?? {},
agentIds: accessGroup.access_agent_ids ?? [],
});
function AccessGroupEditForm({ accessGroup, onCancel, onSuccess }: Omit) {
const form = useZodForm(accessGroupFormSchema, { defaultValues: toFormValues(accessGroup) });
+ const { dirtyFields } = useFormState({ control: form.control });
const editMutation = useEditAccessGroup();
const [activeTab, setActiveTab] = useState(GENERAL_TAB);
const [visitedTabs, setVisitedTabs] = useState>(new Set([GENERAL_TAB]));
@@ -52,6 +55,7 @@ function AccessGroupEditForm({ accessGroup, onCancel, onSuccess }: Omit event.preventDefault()}>
-
+