From 96792d09138deecd5bc9eb5a746e040ea6f49beb Mon Sep 17 00:00:00 2001 From: Arsh Goyal Date: Mon, 26 Oct 2026 23:33:37 +0530 Subject: [PATCH 01/10] test(streaming): pin that a stream stuck on one chunk closes its provider stream and logs one failure (#38753) * fix(streaming): stop re-raising and re-logging a failed stream on every next() A stream that raised keeps its provider iterator alive, so a consumer that iterates past the error re-detected, re-raised and re-logged the same failure on every call, flooding failure callbacks for one request. Latch the failure and re-raise it without logging again. * fix(streaming): drop the retained traceback before re-raising a failed stream * test(streaming): count failures through a real logging callback and require a clean stop on later reads * test(streaming): require the stuck stream to close early and release its socket The provider stream now records a GeneratorExit and caps its chunks, so draining it or detecting the loop late fails. Two wire cases check a sync and an async stuck stream release their TCP session. * test(streaming): avoid failure recorder name collision --------- Co-authored-by: Devin AI Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: moe-berri --- ...est_sync_stream_connection_release_wire.py | 33 +++++ .../test_streaming_handler.py | 118 ++++++++++++++++++ 2 files changed, 151 insertions(+) diff --git a/tests/integration/sdk/test_sync_stream_connection_release_wire.py b/tests/integration/sdk/test_sync_stream_connection_release_wire.py index a8b49d285c3..21774df66d0 100644 --- a/tests/integration/sdk/test_sync_stream_connection_release_wire.py +++ b/tests/integration/sdk/test_sync_stream_connection_release_wire.py @@ -212,3 +212,36 @@ def test_sync_stream_that_raises_mid_stream_releases_the_connection() -> None: assert _open_sessions_settling_to(wire, at_most=0) == 0, ( "a stream that raised mid-way left its TCP session open" ) + + +_STUCK: Final = _Shape( + "hosted_vllm", + "hosted_vllm/stuck-model", + _OPENAI_OPENING, + (_openai_delta("the model is stuck on this", None),) * (10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT), +) + + +def _still_generating_peer(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=(_STUCK.opening, *_STUCK.closing), pause_between_chunks=0.01) + + +def test_sync_stream_stuck_repeating_one_chunk_releases_the_connection() -> None: + with wire_server(_still_generating_peer, keep_alive=True) as wire: + stuck: Final = litellm.completion(**_request_kwargs(_STUCK, wire)) + with pytest.raises(Exception, match="repeating the same chunk"): + for _ in stuck: + pass + assert _open_sessions_settling_to(wire, at_most=0) == 0, "a stream stuck on one chunk left its TCP session open" + + +@pytest.mark.asyncio +async def test_async_stream_stuck_repeating_one_chunk_releases_the_connection() -> None: + with wire_server(_still_generating_peer, keep_alive=True) as wire: + stuck: Final = await litellm.acompletion(**_request_kwargs(_STUCK, wire)) + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in stuck: + pass + assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=0) == 0, ( + "a stream stuck on one chunk 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 c13cea51183..cc1fa1b8803 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,6 +1,7 @@ import asyncio import datetime import json +import queue import sys import threading import time @@ -2819,6 +2820,123 @@ def test_raise_on_model_repetition_tolerates_empty_choices( wrapper.raise_on_model_repetition() +_STUCK_CHUNK: Final = "the model is stuck on this" +_STUCK_PROVIDER_CHUNK_CAP: Final = 10 * litellm.REPEATED_STREAMING_CHUNK_LIMIT + + +class _StuckProvider: + def __init__(self) -> None: + self.closed_early = threading.Event() + self.chunks_sent = 0 + + def _next_chunk(self) -> ModelResponseStream: + self.chunks_sent += 1 + return _make_chunk(_STUCK_CHUNK) + + def sync_stream(self) -> Iterator[ModelResponseStream]: + try: + yield from (self._next_chunk() for _ in range(_STUCK_PROVIDER_CHUNK_CAP)) + except GeneratorExit: + self.closed_early.set() + raise + + async def async_stream(self) -> AsyncIterator[ModelResponseStream]: + try: + for _ in range(_STUCK_PROVIDER_CHUNK_CAP): + yield self._next_chunk() + except GeneratorExit: + self.closed_early.set() + raise + + +class _StuckStreamFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: Final = queue.SimpleQueue[str]() + + def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"])) + + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.failures.put(str(cast(Mapping[str, object], kwargs)["exception"])) + + +def _stuck_stream_wrapper( + provider_stream: Iterator[ModelResponseStream] | AsyncIterator[ModelResponseStream], + call_type: Literal["completion", "acompletion"], + recorder: _StuckStreamFailureRecorder, +) -> CustomStreamWrapper: + logging_obj: Final = Logging( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type=call_type, + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id=f"stuck-stream-{call_type}", + function_id=f"stuck-stream-{call_type}", + dynamic_failure_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="test-model", + optional_params={}, + litellm_params=get_litellm_params(custom_llm_provider="openai", acompletion=call_type == "acompletion"), + ) + return CustomStreamWrapper( + completion_stream=provider_stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + +@pytest.mark.asyncio +async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): + """https://github.com/BerriAI/litellm/issues/19428 and https://github.com/BerriAI/litellm/issues/13786""" + provider: Final = _StuckProvider() + recorder: Final = _StuckStreamFailureRecorder() + wrapper: Final = _stuck_stream_wrapper(provider.async_stream(), "acompletion", recorder) + + with pytest.raises(Exception, match="repeating the same chunk"): + async for _ in wrapper: + pass + assert provider.closed_early.is_set() + assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1 + chunks_sent_at_failure: Final = provider.chunks_sent + + with pytest.raises(StopAsyncIteration): + await wrapper.__anext__() + await wrapper.wait_for_failure_logging(timeout=5) + + assert provider.chunks_sent == chunks_sent_at_failure + assert "repeating the same chunk" in recorder.failures.get_nowait() + assert recorder.failures.empty() + + +def test_sync_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure(): + """Sync counterpart of the async test above""" + provider: Final = _StuckProvider() + recorder: Final = _StuckStreamFailureRecorder() + wrapper: Final = _stuck_stream_wrapper(provider.sync_stream(), "completion", recorder) + + with pytest.raises(Exception, match="repeating the same chunk"): + for _ in wrapper: + pass + assert provider.closed_early.is_set() + assert provider.chunks_sent <= litellm.REPEATED_STREAMING_CHUNK_LIMIT + 1 + chunks_sent_at_failure: Final = provider.chunks_sent + assert "repeating the same chunk" in recorder.failures.get(timeout=5) + + with pytest.raises(StopIteration): + next(wrapper) + + assert provider.chunks_sent == chunks_sent_at_failure + with pytest.raises(queue.Empty): + recorder.failures.get(timeout=0.5) + + def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): """ Test that provider-reported usage from a post-finish_reason chunk From be3a7be1db6c4ee38d0d61957a2ae5179446c874 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 14:12:11 -0400 Subject: [PATCH 02/10] fix(ui): restore clipboard copying in key and log details (#48425) --- ui/litellm-dashboard/eslint-suppressions.json | 5 - .../AuditLogDrawer/AuditLogDrawer.test.tsx | 23 +- .../audit/AuditLogDrawer/AuditLogDrawer.tsx | 19 +- .../logs/detail/DrawerHeader.test.tsx | 25 ++ .../components/logs/detail/DrawerHeader.tsx | 6 +- .../components/logs/detail/InputCard.test.tsx | 42 +-- .../src/components/logs/detail/InputCard.tsx | 5 +- .../LogDetailContent.integration.test.tsx | 41 +++ .../logs/detail/LogDetailContent.tsx | 6 +- .../logs/detail/LogDetailsDrawer.test.tsx | 22 ++ .../logs/detail/LogDetailsDrawer.tsx | 6 +- .../logs/detail/OutputCard.test.tsx | 35 +-- .../src/components/logs/detail/OutputCard.tsx | 5 +- .../logs/detail/RealtimePrettyView.test.tsx | 22 +- .../logs/detail/RealtimePrettyView.tsx | 3 +- .../src/components/shared/CopyButton.test.tsx | 55 +--- .../src/components/shared/CopyButton.tsx | 12 +- .../src/utils/dataUtils.test.ts | 277 +++++++----------- ui/litellm-dashboard/src/utils/dataUtils.ts | 23 +- 19 files changed, 287 insertions(+), 345 deletions(-) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 38d4c435e9a..32777b43504 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2175,11 +2175,6 @@ "count": 1 } }, - "src/utils/dataUtils.test.ts": { - "max-nested-callbacks": { - "count": 1 - } - }, "tests/CreateKeyPage.expiredToken.test.tsx": { "@typescript-eslint/no-require-imports": { "count": 3 diff --git a/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.test.tsx b/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.test.tsx index da158fb8450..8ac99760c01 100644 --- a/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.test.tsx @@ -1,10 +1,16 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { describe, it, expect, vi, beforeEach } from "vitest"; import userEvent from "@testing-library/user-event"; -import { render, screen, waitFor, within } from "@testing-library/react"; +import { render, screen, within } from "@testing-library/react"; import moment from "moment"; import { AuditLogDrawer } from "./AuditLogDrawer"; import { AuditLogEntry } from "../AuditLogsTableColumns"; +vi.mock("@/utils/dataUtils", async (importOriginal) => ({ + ...(await importOriginal()), + copyToClipboard: vi.fn(), +})); + vi.mock("../../../common_components/DefaultProxyAdminTag", () => ({ default: ({ userId }: { userId: string }) => {userId}, })); @@ -120,16 +126,17 @@ describe("AuditLogDrawer", () => { expect(screen.queryByText(/"spend"/)).not.toBeInTheDocument(); }); - it("should copy the json of a block to the clipboard", async () => { + it.each([true, false])("confirms copied JSON only when copying succeeds (%s)", async (success) => { const user = userEvent.setup(); - const writeText = vi.fn().mockResolvedValue(undefined); - Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true }); - Object.defineProperty(window, "isSecureContext", { value: true, configurable: true }); - + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockResolvedValue(success); render(); - await user.click(within(blockNamed("Before")).getByTitle("Copy JSON")); - await waitFor(() => expect(writeText).toHaveBeenCalledWith(JSON.stringify({ max_budget: 10 }, null, 2))); + const button = within(blockNamed("Before")).getByRole("button", { name: "Copy JSON" }); + await user.click(button); + + expect(copyToClipboard).toHaveBeenCalledWith(JSON.stringify({ max_budget: 10 }, null, 2)); + expect(button.querySelector(".lucide-check") !== null).toBe(success); }); describe("width", () => { diff --git a/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx b/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx index fae3c35b80c..eaddf1e507f 100644 --- a/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx +++ b/ui/litellm-dashboard/src/components/logs/audit/AuditLogDrawer/AuditLogDrawer.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { Check, Copy } from "lucide-react"; import { useState, useCallback } from "react"; import moment from "moment"; @@ -27,25 +28,9 @@ function CopyableJsonBlock({ label, value }: { label: string; value: Record { - try { - const text = JSON.stringify(value, null, 2); - if (navigator.clipboard && window.isSecureContext) { - await navigator.clipboard.writeText(text); - } else { - const el = document.createElement("textarea"); - el.value = text; - el.style.position = "fixed"; - el.style.opacity = "0"; - document.body.appendChild(el); - el.focus(); - el.select(); - document.execCommand("copy"); - document.body.removeChild(el); - } + if (await copyToClipboard(JSON.stringify(value, null, 2))) { setCopied(true); setTimeout(() => setCopied(false), 2000); - } catch (e) { - console.error("Copy failed:", e); } }, [value]); diff --git a/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx index 564862acccd..b24667c6a2e 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.test.tsx @@ -6,6 +6,12 @@ import { describe, expect, it, vi } from "vitest"; import { render } from "../../../../tests/test-utils"; import type { LogEntry } from "../types"; import { DrawerHeader } from "./DrawerHeader"; +import { copyToClipboard } from "@/utils/dataUtils"; + +vi.mock("@/utils/dataUtils", async (importOriginal) => ({ + ...(await importOriginal()), + copyToClipboard: vi.fn(), +})); const logEntry = (overrides: Partial): LogEntry => ({ @@ -118,3 +124,22 @@ describe("DrawerHeader back link", () => { expect(screen.queryByRole("button", { name: /^Back to/ })).not.toBeInTheDocument(); }); }); + +describe("DrawerHeader copying", () => { + it.each([true, false])("confirms both ids only when copying succeeds (%s)", async (success) => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockReset().mockResolvedValue(success); + renderHeader(logEntry({ request_id: "request-copy", litellm_call_id: "call-copy" }), false); + + const requestButton = screen.getByRole("button", { name: "Copy Request ID" }); + const callButton = screen.getByRole("button", { name: "Copy x-litellm-call-id" }); + await user.click(requestButton); + await user.click(callButton); + + expect(copyToClipboard).toHaveBeenNthCalledWith(1, "request-copy"); + expect(copyToClipboard).toHaveBeenNthCalledWith(2, "call-copy"); + expect(requestButton).toHaveAccessibleName(success ? "Copied!" : "Copy Request ID"); + expect(callButton).toHaveAccessibleName(success ? "Copied!" : "Copy x-litellm-call-id"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx b/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx index 8ec615d444d..312a0797f1d 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/DrawerHeader.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { type ReactNode, useState } from "react"; import { ArrowLeft, Check, ChevronDown, ChevronUp, Copy, X } from "lucide-react"; import moment from "moment"; @@ -183,12 +184,9 @@ function CopyableId({ const [copied, setCopied] = useState(false); const handleCopy = async () => { - try { - await navigator.clipboard.writeText(value); + if (await copyToClipboard(value)) { setCopied(true); setTimeout(() => setCopied(false), 1200); - } catch { - /* clipboard unavailable in non-secure contexts */ } }; diff --git a/ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx index dfd9dd54dd7..87105d3f48f 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/InputCard.test.tsx @@ -1,12 +1,15 @@ import React from "react"; -import { render, screen, waitFor, fireEvent, act } from "@testing-library/react"; +import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { InputCard } from "./InputCard"; import { ParsedMessage } from "./prettyMessagesTypes"; +import { copyToClipboard } from "@/utils/dataUtils"; +import { toast } from "@/lib/toast"; + +vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() })); describe("InputCard", () => { - const mockWriteText = vi.fn().mockResolvedValue(undefined); const mockMessages: ParsedMessage[] = [ { role: "user", @@ -16,13 +19,7 @@ describe("InputCard", () => { beforeEach(() => { vi.clearAllMocks(); - Object.defineProperty(navigator, "clipboard", { - value: { - writeText: mockWriteText, - }, - writable: true, - configurable: true, - }); + vi.mocked(copyToClipboard).mockResolvedValue(true); }); it("should render the InputCard component", () => { @@ -91,27 +88,16 @@ describe("InputCard", () => { expect(screen.getByText(/Cost: \$0\.001500/)).toBeInTheDocument(); }); - it("should copy last message content when copy button is clicked", async () => { - const messages: ParsedMessage[] = [ - { - role: "user", - content: "Copy this text", - }, - ]; + it.each([true, false])("delegates copying and feedback to the helper (%s)", async (success) => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockResolvedValue(success); + render(); - render(); + await user.click(screen.getByRole("button", { name: /copy/i })); - const copyButton = screen.getByRole("button", { name: /copy/i }); - - expect(copyButton).toBeInTheDocument(); - - await act(async () => { - fireEvent.click(copyButton); - }); - - await waitFor(() => { - expect(mockWriteText).toHaveBeenCalledWith("Copy this text"); - }); + expect(copyToClipboard).toHaveBeenCalledWith("Copy this text", "Input copied"); + expect(toast.success).not.toHaveBeenCalled(); }); it("should toggle collapse state when header is clicked", async () => { diff --git a/ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx b/ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx index d367c69f263..0b0bde85b91 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/InputCard.tsx @@ -4,7 +4,7 @@ */ import { useState } from "react"; -import { toast } from "@/lib/toast"; +import { copyToClipboard } from "@/utils/dataUtils"; import { ParsedMessage } from "./prettyMessagesTypes"; import { SectionHeader } from "./SectionHeader"; import { CollapsibleMessage } from "./CollapsibleMessage"; @@ -32,8 +32,7 @@ export function InputCard({ messages, promptTokens, inputCost }: InputCardProps) const handleCopy = () => { const content = lastMessage?.content || ""; - navigator.clipboard.writeText(content); - toast.success("Input copied"); + void copyToClipboard(content, "Input copied"); }; return ( diff --git a/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.integration.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.integration.test.tsx index 7f39a0d8fe5..21ecf19235e 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.integration.test.tsx @@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; import { GuardrailJumpLink, LogDetailContent } from "./LogDetailContent"; import type { LogEntry } from "../types"; +import { toast } from "@/lib/toast"; vi.mock("./sections/GuardrailViewer/GuardrailViewer", () => ({ default: ({ data }: { data: unknown }) =>
{JSON.stringify(data)}
, @@ -295,6 +296,46 @@ describe("LogDetailContent", () => { expect(screen.queryByRole("tab", { name: "Request" })).not.toBeInTheDocument(); }); + it.each([ + ["Request", true], + ["Response", true], + ["Metadata", true], + ["Request", false], + ] as const)("copies %s through the fallback and reports its outcome (%s)", async (section, success) => { + const user = userEvent.setup(); + const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard"); + const execCommandDescriptor = Object.getOwnPropertyDescriptor(document, "execCommand"); + const writeText = vi.fn<(value: string) => boolean>().mockReturnValue(success); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + Object.defineProperty(document, "execCommand", { + configurable: true, + value: (command: string) => + command === "copy" && writeText((document.activeElement as HTMLTextAreaElement).value), + }); + vi.mocked(toast.success).mockClear(); + const log = createLogEntry(); + const expected = { Request: log.messages, Response: log.response, Metadata: log.metadata }[section]; + + try { + render(); + if (section !== "Metadata") { + await user.click(screen.getByRole("tab", { name: "JSON", exact: true })); + await user.click(screen.getByRole("tab", { name: section, exact: true })); + } + const label = section === "Metadata" ? "Copy Metadata" : "Copy JSON"; + const button = screen.getByRole("button", { name: label, exact: true }); + await user.click(button); + + expect(writeText).toHaveBeenCalledWith(JSON.stringify(expected, null, 2)); + expect(button).toHaveAccessibleName(success ? "Copied!" : label); + expect(toast.success).toHaveBeenCalledTimes(success ? 1 : 0); + } finally { + if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor); + if (execCommandDescriptor) Object.defineProperty(document, "execCommand", execCommandDescriptor); + else Reflect.deleteProperty(document, "execCommand"); + } + }); + it.each(["object", "serialized", "messages only", "null captures"])( "preserves request inspection and copying for classifier logs with %s data", async (shape) => { diff --git a/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx index f6c63582938..f740b6c586a 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/LogDetailContent.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { useState } from "react"; import { Check, ChevronDown, ChevronRight, CircleAlert, Copy, Info } from "lucide-react"; import moment from "moment"; @@ -304,12 +305,9 @@ function CopyButton({ const [copied, setCopied] = useState(false); const handleCopy = async () => { - try { - await navigator.clipboard.writeText(getText()); + if (await copyToClipboard(getText())) { setCopied(true); setTimeout(() => setCopied(false), 1200); - } catch { - /* clipboard unavailable in non-secure contexts */ } }; diff --git a/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx index 1321a4a152a..36ab70ecf2a 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.test.tsx @@ -1,3 +1,5 @@ +import userEvent from "@testing-library/user-event"; +import { copyToClipboard } from "@/utils/dataUtils"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; @@ -6,6 +8,11 @@ import { sessionSpendLogsCall } from "../../networking"; import { LogEntry } from "../types"; import { AutoRouterModelGroupsProvider } from "@/components/shared/table_cells"; +vi.mock("@/utils/dataUtils", async (importOriginal) => ({ + ...(await importOriginal()), + copyToClipboard: vi.fn(), +})); + vi.mock("../../networking", () => ({ sessionSpendLogsCall: vi.fn(), })); @@ -201,3 +208,18 @@ describe("LogDetailsDrawer session sidebar auto-router icon", () => { expect(directRow.querySelector(".lucide-waypoints")).toBeNull(); }); }); + +describe("LogDetailsDrawer copying", () => { + it.each([true, false])("confirms a session id only when copying succeeds (%s)", async (success) => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockReset().mockResolvedValue(success); + renderSessionDrawer(); + + const button = await screen.findByRole("button", { name: "Copy trace id" }); + await user.click(button); + + expect(copyToClipboard).toHaveBeenCalledWith("session-1"); + expect(button.querySelector(".lucide-check") !== null).toBe(success); + }); +}); diff --git a/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx index 48b89ee8ac1..23ad629d572 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/LogDetailsDrawer.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; import { useEffect, useMemo, useState } from "react"; import { Bot, Check, Copy, Sparkles, Wrench } from "lucide-react"; import { Sheet, SheetContent, SheetTitle } from "@/components/ui/sheet"; @@ -289,12 +290,9 @@ export function LogDetailsDrawer({ const handleCopyLeftPanelId = async () => { if (!leftPanelId) return; - try { - await navigator.clipboard.writeText(leftPanelId); + if (await copyToClipboard(leftPanelId)) { setCopiedLeftPanelId(true); setTimeout(() => setCopiedLeftPanelId(false), 1200); - } catch { - /* clipboard unavailable in non-secure contexts */ } }; diff --git a/ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx index 3dda7c5fa65..7b4c1a987a1 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/OutputCard.test.tsx @@ -1,12 +1,15 @@ import React from "react"; -import { render, screen, waitFor, fireEvent, act } from "@testing-library/react"; +import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { OutputCard } from "./OutputCard"; import { ParsedMessage } from "./prettyMessagesTypes"; +import { copyToClipboard } from "@/utils/dataUtils"; +import { toast } from "@/lib/toast"; + +vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() })); describe("OutputCard", () => { - const mockWriteText = vi.fn().mockResolvedValue(undefined); const mockMessage: ParsedMessage = { role: "assistant", content: "This is a test response", @@ -14,13 +17,7 @@ describe("OutputCard", () => { beforeEach(() => { vi.clearAllMocks(); - Object.defineProperty(navigator, "clipboard", { - value: { - writeText: mockWriteText, - }, - writable: true, - configurable: true, - }); + vi.mocked(copyToClipboard).mockResolvedValue(true); }); it("should render the OutputCard component", () => { @@ -53,20 +50,16 @@ describe("OutputCard", () => { expect(screen.getByText(/Cost: \$0\.002500/)).toBeInTheDocument(); }); - it("should copy message content when copy button is clicked", async () => { + it.each([true, false])("delegates copying and feedback to the helper (%s)", async (success) => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + vi.mocked(copyToClipboard).mockResolvedValue(success); render(); - const copyButton = screen.getByRole("button", { name: /copy/i }); + await user.click(screen.getByRole("button", { name: /copy/i })); - expect(copyButton).toBeInTheDocument(); - - await act(async () => { - fireEvent.click(copyButton); - }); - - await waitFor(() => { - expect(mockWriteText).toHaveBeenCalledWith("This is a test response"); - }); + expect(copyToClipboard).toHaveBeenCalledWith("This is a test response", "Output copied"); + expect(toast.success).not.toHaveBeenCalled(); }); it("should not copy when message is null and copy button is clicked", async () => { @@ -80,7 +73,7 @@ describe("OutputCard", () => { await user.click(copyButton); await waitFor(() => { - expect(mockWriteText).not.toHaveBeenCalled(); + expect(copyToClipboard).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx b/ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx index 7d0ec62ffe2..7daa8b02546 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/OutputCard.tsx @@ -4,7 +4,7 @@ */ import { useState } from "react"; -import { toast } from "@/lib/toast"; +import { copyToClipboard } from "@/utils/dataUtils"; import { COLOR_BORDER } from "./constants"; import { ParsedMessage } from "./prettyMessagesTypes"; import { SectionHeader } from "./SectionHeader"; @@ -22,8 +22,7 @@ export function OutputCard({ message, completionTokens, outputCost }: OutputCard const handleCopy = () => { if (!message) return; - navigator.clipboard.writeText(message.content || ""); - toast.success("Output copied"); + void copyToClipboard(message.content || "", "Output copied"); }; return ( diff --git a/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx b/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx index 201c9e74ce8..d5280c35983 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.test.tsx @@ -1,9 +1,12 @@ import React from "react"; import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import { copyToClipboard } from "@/utils/dataUtils"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { RealtimePrettyView, isRealtimeResponse } from "./RealtimePrettyView"; +vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn().mockResolvedValue(true) })); + const sampleRealtimeResponse = { usage: { total_tokens: 587, @@ -142,15 +145,20 @@ describe("isRealtimeResponse", () => { }); describe("RealtimePrettyView", () => { - const mockWriteText = vi.fn().mockResolvedValue(undefined); - beforeEach(() => { vi.clearAllMocks(); - Object.defineProperty(navigator, "clipboard", { - value: { writeText: mockWriteText }, - writable: true, - configurable: true, - }); + }); + + it("copies all transcript turns when the Clipboard API is unavailable", async () => { + const user = userEvent.setup(); + Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); + render(); + + await user.click(screen.getByRole("button", { name: /copy/i })); + + expect(copyToClipboard).toHaveBeenCalledWith( + "assistant: Hello! How's your day going?\nassistant: I'm here to help with information and general questions.", + ); }); it("should render the component successfully", () => { diff --git a/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx b/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx index 82f0d3e7fb2..04f0a178462 100644 --- a/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx +++ b/ui/litellm-dashboard/src/components/logs/detail/RealtimePrettyView.tsx @@ -1,3 +1,4 @@ +import { copyToClipboard } from "@/utils/dataUtils"; /** * RealtimePrettyView - Structured pretty view for OpenAI Realtime API logs * Displays session config, conversation turns, and token usage @@ -276,7 +277,7 @@ function ConversationCard({ (r.output || []).flatMap((o) => (o.content || []).map((c) => `${o.role}: ${c.transcript || c.text || ""}`)), ) .join("\n"); - navigator.clipboard.writeText(transcripts); + void copyToClipboard(transcripts); }; return ( diff --git a/ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx b/ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx index c2c2ee199b4..311b536d89c 100644 --- a/ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/CopyButton.test.tsx @@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; import type { ComponentProps } from "react"; import { copyToClipboard } from "@/utils/dataUtils"; -import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import { renderWithProviders, screen } from "../../../tests/test-utils"; import CopyButton from "./CopyButton"; vi.mock("@/utils/dataUtils", () => ({ copyToClipboard: vi.fn() })); @@ -13,13 +13,9 @@ vi.mock("lucide-react", async (importOriginal) => ({ Check: (props: ComponentProps<"svg">) => , })); -const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard"); - afterEach(() => { vi.useRealTimers(); vi.resetAllMocks(); - if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor); - else Reflect.deleteProperty(navigator, "clipboard"); }); describe("CopyButton", () => { @@ -28,55 +24,37 @@ describe("CopyButton", () => { expect(screen.queryByRole("button")).not.toBeInTheDocument(); }); - it("shows the confirmation checkmark only after a successful write", async () => { + it("shows the icon confirmation only after copying succeeds", async () => { const user = userEvent.setup(); - const writeText = vi.fn().mockResolvedValue(undefined); - Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true }); + const pendingCopy = Promise.withResolvers(); + vi.mocked(copyToClipboard).mockReturnValue(pendingCopy.promise); + renderWithProviders(); - renderWithProviders(); + await user.click(screen.getByRole("button", { name: "Copy value" })); - const button = screen.getByRole("button", { name: "Copy value" }); + expect(copyToClipboard).toHaveBeenCalledWith("secret-value", "Value copied"); expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument(); + expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument(); - await user.click(button); - - expect(writeText).toHaveBeenCalledWith("secret-value"); - expect(await screen.findByRole("img", { name: "Copied icon" })).toBeInTheDocument(); + await act(async () => pendingCopy.resolve(true)); + expect(screen.getByRole("img", { name: "Copied icon" })).toBeInTheDocument(); }); - it("does not show the checkmark when the clipboard write is rejected", async () => { + it("keeps the icon unchanged when copying fails", async () => { const user = userEvent.setup(); - const writeText = vi.fn().mockRejectedValue(new Error("permission denied")); - Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true }); - + vi.mocked(copyToClipboard).mockResolvedValue(false); renderWithProviders(); - const button = screen.getByRole("button", { name: "Copy value" }); - await user.click(button); + await user.click(screen.getByRole("button", { name: "Copy value" })); - await waitFor(() => expect(writeText).toHaveBeenCalledWith("secret-value")); + expect(copyToClipboard).toHaveBeenCalledWith("secret-value", "Copied"); expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument(); expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument(); }); - it("does not show the checkmark when the clipboard API is unavailable", async () => { - const user = userEvent.setup(); - Object.defineProperty(navigator, "clipboard", { value: undefined, configurable: true }); - - renderWithProviders(); - - const button = screen.getByRole("button", { name: "Copy value" }); - await user.click(button); - - expect(screen.queryByRole("img", { name: "Copied icon" })).not.toBeInTheDocument(); - expect(screen.getByRole("img", { name: "Copy icon" })).toBeInTheDocument(); - expect(copyToClipboard).not.toHaveBeenCalled(); - }); - it("keeps the default icon styling, custom classes, title and 1200ms confirmation", async () => { vi.useFakeTimers(); - const writeText = vi.fn().mockResolvedValue(undefined); - Object.defineProperty(navigator, "clipboard", { value: { writeText }, configurable: true }); + vi.mocked(copyToClipboard).mockResolvedValue(true); renderWithProviders( , ); @@ -87,8 +65,7 @@ describe("CopyButton", () => { expect(screen.getByRole("img", { name: "Copy icon" })).toHaveClass("size-3"); await act(async () => fireEvent.click(button)); - expect(writeText).toHaveBeenCalledWith("request-id"); - expect(copyToClipboard).not.toHaveBeenCalled(); + expect(copyToClipboard).toHaveBeenCalledWith("request-id", "Copied"); act(() => vi.advanceTimersByTime(1199)); expect(screen.getByRole("img", { name: "Copied icon" })).toBeInTheDocument(); act(() => vi.advanceTimersByTime(1)); diff --git a/ui/litellm-dashboard/src/components/shared/CopyButton.tsx b/ui/litellm-dashboard/src/components/shared/CopyButton.tsx index 57ae86a26d0..225154e4bcf 100644 --- a/ui/litellm-dashboard/src/components/shared/CopyButton.tsx +++ b/ui/litellm-dashboard/src/components/shared/CopyButton.tsx @@ -54,17 +54,7 @@ const CopyButton: React.FC = ({ if (value == null || (!value && !isAction)) return null; const handleCopy = async () => { - if (isAction) { - setCopied(await copyToClipboard(value, copiedLabel)); - return; - } - if (!navigator.clipboard) return; - try { - await navigator.clipboard.writeText(value); - setCopied(true); - } catch { - setCopied(false); - } + setCopied(await copyToClipboard(value, copiedLabel)); }; const button = ( diff --git a/ui/litellm-dashboard/src/utils/dataUtils.test.ts b/ui/litellm-dashboard/src/utils/dataUtils.test.ts index af249ca601a..23b425b47e0 100644 --- a/ui/litellm-dashboard/src/utils/dataUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/dataUtils.test.ts @@ -1,5 +1,6 @@ // @vitest-environment jsdom +import "@testing-library/jest-dom/vitest"; import { describe, it, expect, beforeEach, vi, afterEach } from "vitest"; import { copyToClipboard, @@ -16,8 +17,6 @@ const mockToast = vi.mocked(toast); describe("dataUtils", () => { beforeEach(() => { vi.clearAllMocks(); - // Reset document.execCommand mock - delete (document as any).execCommand; }); afterEach(() => { @@ -134,194 +133,114 @@ describe("dataUtils", () => { }); describe("copyToClipboard", () => { - describe("when Clipboard API is available (HTTPS scenario)", () => { - beforeEach(() => { - // Mock modern Clipboard API - Object.assign(navigator, { - clipboard: { - writeText: vi.fn().mockResolvedValue(undefined), - }, - }); - }); - - it("should use navigator.clipboard.writeText when available", async () => { - const result = await copyToClipboard("test text"); - - expect(navigator.clipboard.writeText).toHaveBeenCalledWith("test text"); - expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard"); - expect(result).toBe(true); - }); - - it("should use custom message when provided", async () => { - await copyToClipboard("test text", "Custom message"); - - expect(mockToast.success).toHaveBeenCalledWith("Custom message"); - }); - - it("should return false for null/undefined text", async () => { - expect(await copyToClipboard(null)).toBe(false); - expect(await copyToClipboard(undefined)).toBe(false); - expect(navigator.clipboard.writeText).not.toHaveBeenCalled(); - }); - - it("should fall back to execCommand when clipboard API fails", async () => { - // Make clipboard API fail - navigator.clipboard.writeText = vi.fn().mockRejectedValue(new Error("Permission denied")); - - // Mock successful execCommand - document.execCommand = vi.fn().mockReturnValue(true); - - // Mock DOM methods - const mockTextArea = { - value: "", - style: {} as Record, - setAttribute: vi.fn(), - focus: vi.fn(), - select: vi.fn(), - }; - document.createElement = vi.fn().mockReturnValue(mockTextArea); - document.body.appendChild = vi.fn(); - document.body.removeChild = vi.fn(); - - const result = await copyToClipboard("test text"); - - expect(navigator.clipboard.writeText).toHaveBeenCalledWith("test text"); - expect(document.execCommand).toHaveBeenCalledWith("copy"); - expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard"); - expect(result).toBe(true); + const clipboardDescriptor = Object.getOwnPropertyDescriptor(navigator, "clipboard"); + const execCommandDescriptor = Object.getOwnPropertyDescriptor(document, "execCommand"); + const writeText = vi.fn<(text: string) => Promise>(); + const execCommand = vi.fn<(command: string) => boolean>(); + const selectedText = vi.fn(() => { + const element = document.activeElement; + return element instanceof HTMLTextAreaElement + ? element.value.slice(element.selectionStart, element.selectionEnd) + : undefined; + }); + const trigger = document.createElement("button"); + const text = "Request α\nsecond line"; + + beforeEach(() => { + writeText.mockResolvedValue(undefined); + execCommand.mockImplementation(() => { + selectedText(); + return true; }); + Object.defineProperty(navigator, "clipboard", { configurable: true, value: { writeText } }); + Object.defineProperty(document, "execCommand", { configurable: true, value: execCommand }); + vi.spyOn(console, "error").mockImplementation(() => {}); + document.body.appendChild(trigger); + trigger.focus(); }); - describe("when Clipboard API is not available (HTTP scenario)", () => { - beforeEach(() => { - // Mock HTTP scenario - no clipboard API - Object.assign(navigator, { - clipboard: undefined, - }); - - // Mock successful execCommand - document.execCommand = vi.fn().mockReturnValue(true); - - // Mock DOM methods - const mockTextArea = { - value: "", - style: {} as Record, - setAttribute: vi.fn(), - focus: vi.fn(), - select: vi.fn(), - }; - document.createElement = vi.fn().mockReturnValue(mockTextArea); - document.body.appendChild = vi.fn(); - document.body.removeChild = vi.fn(); - }); - - it("should fall back to execCommand when clipboard API is not available", async () => { - const result = await copyToClipboard("test text"); - - expect(document.createElement).toHaveBeenCalledWith("textarea"); - expect(document.execCommand).toHaveBeenCalledWith("copy"); - expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard"); - expect(result).toBe(true); - }); - - it("should set textarea properties correctly", async () => { - const mockTextArea = { - value: "", - style: {} as Record, - setAttribute: vi.fn(), - focus: vi.fn(), - select: vi.fn(), - }; - document.createElement = vi.fn().mockReturnValue(mockTextArea); - - await copyToClipboard("test text"); - - expect(mockTextArea.value).toBe("test text"); - /* eslint-disable jest-dom/prefer-to-have-style -- the subject is a plain mock object, not a DOM node, so toHaveStyle cannot read it */ - expect(mockTextArea.style.position).toBe("fixed"); - expect(mockTextArea.style.left).toBe("-999999px"); - expect(mockTextArea.style.top).toBe("-999999px"); - /* eslint-enable jest-dom/prefer-to-have-style */ - expect(mockTextArea.setAttribute).toHaveBeenCalledWith("readonly", ""); - expect(mockTextArea.focus).toHaveBeenCalled(); - expect(mockTextArea.select).toHaveBeenCalled(); - }); - - it("should handle execCommand failure", async () => { - document.execCommand = vi.fn().mockReturnValue(false); - - const result = await copyToClipboard("test text"); - - expect(document.execCommand).toHaveBeenCalledWith("copy"); - expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard"); - expect(result).toBe(false); - }); - - it("should handle DOM manipulation errors", async () => { - document.createElement = vi.fn().mockImplementation(() => { - throw new Error("DOM error"); - }); - - const result = await copyToClipboard("test text"); - - expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard"); - expect(result).toBe(false); - }); - - it("should clean up textarea element after successful copy", async () => { - const mockTextArea = { - value: "", - style: {} as Record, - setAttribute: vi.fn(), - focus: vi.fn(), - select: vi.fn(), - }; - document.createElement = vi.fn().mockReturnValue(mockTextArea); - - await copyToClipboard("test text"); - - expect(document.body.appendChild).toHaveBeenCalledWith(mockTextArea); - expect(document.body.removeChild).toHaveBeenCalledWith(mockTextArea); - }); + afterEach(() => { + document.body.replaceChildren(); + if (clipboardDescriptor) Object.defineProperty(navigator, "clipboard", clipboardDescriptor); + else Reflect.deleteProperty(navigator, "clipboard"); + if (execCommandDescriptor) Object.defineProperty(document, "execCommand", execCommandDescriptor); + else Reflect.deleteProperty(document, "execCommand"); }); - describe("edge cases", () => { - beforeEach(() => { - // Mock scenario where navigator exists but clipboard is null - Object.assign(navigator, { - clipboard: null, - }); + it("confirms a modern clipboard write only after it succeeds", async () => { + const pendingWrite = Promise.withResolvers(); + writeText.mockReturnValue(pendingWrite.promise); + const result = copyToClipboard(text, "Request copied"); - document.execCommand = vi.fn().mockReturnValue(true); - const mockTextArea = { - value: "", - style: {} as Record, - setAttribute: vi.fn(), - focus: vi.fn(), - select: vi.fn(), - }; - document.createElement = vi.fn().mockReturnValue(mockTextArea); - document.body.appendChild = vi.fn(); - document.body.removeChild = vi.fn(); + expect(writeText).toHaveBeenCalledWith(text); + expect(mockToast.success).not.toHaveBeenCalled(); + pendingWrite.resolve(); + + expect(await result).toBe(true); + expect(mockToast.success).toHaveBeenCalledWith("Request copied"); + expect(execCommand).not.toHaveBeenCalled(); + expect(trigger).toHaveFocus(); + }); + + it.each([null, undefined, ""])("does not copy absent content %s", async (value) => { + expect(await copyToClipboard(value)).toBe(false); + expect(writeText).not.toHaveBeenCalled(); + expect(execCommand).not.toHaveBeenCalled(); + expect(mockToast.success).not.toHaveBeenCalled(); + }); + + it.each([undefined, null, {}])("copies the full selection when the Clipboard API is %s", async (clipboard) => { + Object.defineProperty(navigator, "clipboard", { configurable: true, value: clipboard }); + + expect(await copyToClipboard(text)).toBe(true); + + expect(execCommand).toHaveBeenCalledWith("copy"); + expect(selectedText).toHaveReturnedWith(text); + expect(mockToast.success).toHaveBeenCalledWith("Copied to clipboard"); + expect(document.querySelector("textarea")).not.toBeInTheDocument(); + expect(trigger).toHaveFocus(); + }); + + it("falls back after a rejected modern write", async () => { + writeText.mockRejectedValue(new Error("Permission denied")); + + expect(await copyToClipboard(text, "Request copied")).toBe(true); + + expect(writeText).toHaveBeenCalledWith(text); + expect(execCommand).toHaveBeenCalledWith("copy"); + expect(selectedText).toHaveReturnedWith(text); + expect(mockToast.success).toHaveBeenCalledExactlyOnceWith("Request copied"); + expect(document.querySelector("textarea")).not.toBeInTheDocument(); + expect(trigger).toHaveFocus(); + }); + + it.each(["false", "throw"])("cleans up and reports failure when the fallback returns %s", async (outcome) => { + Object.defineProperty(navigator, "clipboard", { configurable: true, value: undefined }); + execCommand.mockImplementation(() => { + selectedText(); + if (outcome === "throw") throw new Error("Copy denied"); + return false; }); - it("should handle navigator.clipboard being null", async () => { - const result = await copyToClipboard("test text"); + expect(await copyToClipboard(text)).toBe(false); - expect(document.execCommand).toHaveBeenCalledWith("copy"); - expect(result).toBe(true); + expect(selectedText).toHaveReturnedWith(text); + expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard"); + expect(mockToast.success).not.toHaveBeenCalled(); + expect(document.querySelector("textarea")).not.toBeInTheDocument(); + expect(trigger).toHaveFocus(); + }); + + it("reports failure when the fallback cannot create its selection", async () => { + Object.defineProperty(navigator, "clipboard", { configurable: true, value: undefined }); + vi.spyOn(document, "createElement").mockImplementationOnce(() => { + throw new Error("DOM error"); }); - it("should handle navigator.clipboard.writeText being undefined", async () => { - Object.assign(navigator, { - clipboard: {}, // clipboard exists but writeText doesn't - }); - - const result = await copyToClipboard("test text"); - - expect(document.execCommand).toHaveBeenCalledWith("copy"); - expect(result).toBe(true); - }); + expect(await copyToClipboard(text)).toBe(false); + expect(mockToast.fromError).toHaveBeenCalledWith("Failed to copy to clipboard"); + expect(mockToast.success).not.toHaveBeenCalled(); + expect(trigger).toHaveFocus(); }); }); }); diff --git a/ui/litellm-dashboard/src/utils/dataUtils.ts b/ui/litellm-dashboard/src/utils/dataUtils.ts index c813a69b999..a05b8d6b0c2 100644 --- a/ui/litellm-dashboard/src/utils/dataUtils.ts +++ b/ui/litellm-dashboard/src/utils/dataUtils.ts @@ -103,19 +103,20 @@ const fallbackCopyToClipboard = (text: string, messageText: string): boolean => textArea.setAttribute("readonly", ""); document.body.appendChild(textArea); - textArea.focus(); - textArea.select(); + try { + textArea.focus(); + textArea.select(); - const successful = document.execCommand("copy"); - document.body.removeChild(textArea); - if (previouslyFocused instanceof HTMLElement) previouslyFocused.focus(); - - if (successful) { - toast.success(messageText); - return true; - } else { - throw new Error("execCommand failed"); + if (document.execCommand("copy")) { + toast.success(messageText); + return true; + } + } finally { + textArea.remove(); + if (previouslyFocused instanceof HTMLElement) previouslyFocused.focus(); } + toast.fromError("Failed to copy to clipboard"); + return false; } catch (err) { toast.fromError("Failed to copy to clipboard"); console.error("Failed to copy: ", err); From 4c674e8e3f0758531308b68f4d42399811cf3b8c Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 14:18:23 -0400 Subject: [PATCH 03/10] fix(mcp): report oversized SSE catalog errors (#48424) --- .../mcp_server/faults/list_outcomes.py | 46 +++++++++++-- .../mcp_server/rest_endpoints.py | 2 +- .../test_mcp_client.py | 56 ++++++++++++++++ .../mcp_server/faults/test_list_outcomes.py | 65 ++++++++++++++++++- 4 files changed, 163 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py index a3e33929f6d..e148e28e826 100644 --- a/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py +++ b/litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py @@ -3,8 +3,8 @@ The aggregate listing deliberately keeps serving the healthy subset when one server fails, but a failed server must contribute a classified outcome instead of silently shrinking the list: an empty contribution with no signal makes a broken upstream indistinguishable from a healthy server with no -tools. Outcomes carry only machine fields (category and status code) so nothing from an upstream -body crosses the trust boundary; classification is total, so any exception out of a server fetch +tools. Outcomes carry machine fields and controlled diagnostics so nothing from an upstream body +crosses the trust boundary; classification is total, so any exception out of a server fetch becomes an outcome, never a second failure. """ @@ -21,8 +21,10 @@ from urllib.request import parse_http_list, parse_keqv_list import anyio import httpx import httpx2 +from mcp import MCPError +from mcp.types import CONNECTION_CLOSED from mcp.types import Tool as MCPTool -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, PositiveInt, TypeAdapter from typing_extensions import assert_never from litellm._logging import verbose_logger @@ -65,6 +67,16 @@ class ServerListFault(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) tag: ListFaultCategory status_code: int | None = None + sse_event_limit_bytes: PositiveInt | None = None + + @property + def public_message(self) -> str | None: + if self.sse_event_limit_bytes is None: + return None + return ( + f"Upstream MCP SSE event exceeded the {self.sse_event_limit_bytes} byte per-event limit. " + "Paginate tools/list responses so each event stays below the limit." + ) ServerOutcome: TypeAlias = ServerListOk | ServerListFault @@ -155,6 +167,18 @@ def raise_classified_list_failure( raise MCPServerListError(classify_list_exception(exc), server_name) from exc +def _sse_event_limit(exc: BaseException) -> int | None: + """The SDK replaces decoder exceptions with CONNECTION_CLOSED and a fixed diagnostic.""" + if not isinstance(exc, (httpx2.SSEError, MCPError)): + return None + if isinstance(exc, MCPError) and exc.error.code != CONNECTION_CLOSED: + return None + prefix: Final = "SSE stream failed: " if isinstance(exc, MCPError) else "" + message: Final = exc.error.message if isinstance(exc, MCPError) else str(exc) + matched: Final = re.fullmatch(prefix + r"Server-sent event exceeded the ([1-9][0-9]{0,18}) byte limit\.", message) + return int(matched[1]) if matched is not None else None + + def classify_list_exception(exc: BaseException) -> ServerListFault: """Classify a per-server listing failure into exactly one outcome. Total: an exception this function cannot recognize is the gateway's own fault (``internal``), never a re-raise.""" @@ -179,6 +203,11 @@ def classify_list_exception(exc: BaseException) -> ServerListFault: response: Final = _find_upstream_response(exc) if response is not None: return ServerListFault(tag="upstream_error", status_code=response.status_code) + sse_event_limit: Final = next( + (limit for error in iter_exception_tree(exc) if (limit := _sse_event_limit(error)) is not None), None + ) + if sse_event_limit is not None: + return ServerListFault(tag="upstream_error", sse_event_limit_bytes=sse_event_limit) if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)): return ServerListFault(tag="timeout") if isinstance(exc, (httpx.TransportError, httpx2.TransportError)): @@ -295,7 +324,7 @@ def listing_failed_outright(outcomes: Mapping[str, ServerOutcome]) -> bool: def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: """The client-visible form of one outcome, for the tools/list result ``_meta`` and the REST - response: category plus status code only, never upstream prose or URLs.""" + response: category, status code and controlled diagnostics, never upstream prose or URLs.""" match outcome.tag: case "ok": return {"status": "ok", "tool_count": outcome.tool_count} @@ -303,6 +332,15 @@ def outcome_wire_value(outcome: ServerOutcome) -> dict[str, object]: return { "status": outcome.tag, **({"http_status": outcome.status_code} if outcome.status_code is not None else {}), + **( + { + "reason": "sse_event_too_large", + "limit_bytes": outcome.sse_event_limit_bytes, + "message": outcome.public_message, + } + if outcome.sse_event_limit_bytes is not None + else {} + ), } case _: assert_never(outcome.tag) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index ee063b81ae4..62387b7cc28 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -957,7 +957,7 @@ if MCP_AVAILABLE: status_code=list_fault_http_status(fault), detail={ "error": fault.tag, - "message": f"Failed to list tools from server {get_server_prefix(server)}", + "message": fault.public_message or f"Failed to list tools from server {get_server_prefix(server)}", }, ) from e except Exception as e: diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py index b4485bcc5f1..60600afab84 100644 --- a/tests/unit/experimental_mcp_client/test_mcp_client.py +++ b/tests/unit/experimental_mcp_client/test_mcp_client.py @@ -1921,6 +1921,62 @@ async def test_list_tools_classifies_upstream_status_for_post_failures(failure_m assert rate_limit is None +@pytest.mark.asyncio +@pytest.mark.parametrize("paginated", (False, True)) +async def test_sse_catalog_limit_reports_the_failure_and_allows_smaller_pages(paginated: bool) -> None: + from mcp.types import Tool + + limit: Final = 1024 + tools: Final = tuple( + Tool(name=f"tool_{index}", description="private-upstream-secret " * 25, input_schema={"type": "object"}) + for index in range(2) + ) + catalog: Final = tuple(tool.model_dump(by_alias=True, exclude_none=True) for tool in tools) + assert len(json.dumps({"tools": catalog}).encode()) > limit + + def respond(request: httpx2.Request) -> httpx2.Response: + if request.method != "POST": + return httpx2.Response(200) + payload: Final = _JSONRPC_MESSAGE_ADAPTER.validate_json(request.content) + assert isinstance(payload, (JSONRPCRequest, JSONRPCNotification)) + if isinstance(payload, JSONRPCNotification): + return httpx2.Response(202) + if payload.method == "initialize": + return httpx2.Response(200, json={ + "jsonrpc": "2.0", "id": payload.id, "result": { + "protocolVersion": LATEST_HANDSHAKE_VERSION, + "capabilities": {"tools": {}}, "serverInfo": {"name": "catalog", "version": "1"}, + }, + }) + assert payload.method == "tools/list" + cursor: Final = (payload.params or {}).get("cursor") + page: Final = catalog[1:] if cursor == "next-page" else catalog[:1] if paginated else catalog + result: Final = { + "tools": page, + **({"nextCursor": "next-page"} if paginated and cursor is None else {}), + } + event: Final = f"data: {json.dumps({'jsonrpc': '2.0', 'id': payload.id, 'result': result})}\n\n" + assert (len(event.encode()) < limit) == paginated + return httpx2.Response(200, headers={"content-type": "text/event-stream"}, content=event) + + def transport_context(url: str, http_client: httpx2.AsyncClient) -> _TransportContext: + return streamable_http_client(url, http_client=http_client, max_sse_event_size=limit) + + client: Final = _MockTransportClient( + respond, server_url="https://upstream.example/mcp", transport_context=transport_context, + ) + if paginated: + assert await client.list_tools(raise_on_error=True) == list(tools) + return + with pytest.raises(MCPError) as caught: + await client.list_tools(raise_on_error=True) + fault: Final = classify_list_exception(caught.value) + assert fault == ServerListFault(tag="upstream_error", sse_event_limit_bytes=limit) + assert list_fault_http_status(fault) == 502 + assert fault.public_message is not None + assert "private-upstream-secret" not in fault.public_message + + @pytest.mark.asyncio async def test_http_status_check_allows_auth_refresh_before_rejecting() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.client_credentials import ( diff --git a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index f6902c778e5..fdae575276d 100644 --- a/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -9,9 +9,10 @@ if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 from exceptiongroup import BaseExceptionGroup import httpx +import httpx2 import pytest from mcp import MCPError -from mcp.types import ErrorData +from mcp.types import CONNECTION_CLOSED, ErrorData from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, @@ -102,6 +103,68 @@ def test_unknown_exception_is_internal(): assert classify_list_exception(ValueError("who knows")).tag == "internal" +@pytest.mark.parametrize("sdk_wrapped", (False, True)) +@pytest.mark.parametrize("nested", (False, True)) +def test_sse_event_limit_is_a_safe_upstream_fault(sdk_wrapped: bool, nested: bool) -> None: + limit: Final = 1024 + decoder_message: Final = f"Server-sent event exceeded the {limit} byte limit." + failure: Final = ( + MCPError(code=CONNECTION_CLOSED, message=f"SSE stream failed: {decoder_message}") + if sdk_wrapped else httpx2.SSEError(decoder_message) + ) + wrapped: Final = BaseExceptionGroup("private-upstream-secret", [failure]) if nested else failure + fault: Final = classify_list_exception(wrapped) + message: Final = ( + f"Upstream MCP SSE event exceeded the {limit} byte per-event limit. " + "Paginate tools/list responses so each event stays below the limit." + ) + assert fault == ServerListFault(tag="upstream_error", sse_event_limit_bytes=limit) + assert fault.public_message == message + assert list_fault_http_status(fault) == 502 + assert outcome_wire_value(fault) == { + "status": "upstream_error", "reason": "sse_event_too_large", "limit_bytes": limit, "message": message, + } + + +@pytest.mark.parametrize( + "failure", + ( + MCPError(code=CONNECTION_CLOSED, message="Connection closed"), + MCPError(code=-32603, message="SSE stream failed: Server-sent event exceeded the 1024 byte limit."), + MCPError(code=CONNECTION_CLOSED, message="SSE stream failed: Server-sent event exceeded the 0 byte limit."), + MCPError(code=CONNECTION_CLOSED, message="SSE stream failed: Server-sent event exceeded the -1 byte limit."), + MCPError(code=CONNECTION_CLOSED, message=( + "SSE stream failed: Server-sent event exceeded the 10000000000000000000 byte limit." + )), + MCPError(code=CONNECTION_CLOSED, message=( + "SSE stream failed: Server-sent event exceeded the 1024 byte limit. private-secret" + )), + httpx2.SSEError("private-secret Server-sent event exceeded the 1024 byte limit."), + httpx2.SSEError("Expected response with content type 'text/event-stream', got 'private-secret'."), + RuntimeError("SSE stream failed: Server-sent event exceeded the 1024 byte limit."), + ), +) +def test_unrelated_errors_do_not_claim_an_sse_event_limit(failure: Exception) -> None: + fault: Final = classify_list_exception(failure) + assert fault.sse_event_limit_bytes is None + assert fault.public_message is None + assert "private-secret" not in str(outcome_wire_value(fault)) + + +@pytest.mark.parametrize("status_code", (401, 403)) +def test_upstream_auth_takes_precedence_over_sse_event_limit(status_code: int) -> None: + request: Final = httpx2.Request("POST", "https://upstream.example/mcp") + response: Final = httpx2.Response(status_code, request=request) + failure: Final = BaseExceptionGroup("upstream errors", [ + httpx2.SSEError("Server-sent event exceeded the 1024 byte limit."), + httpx2.HTTPStatusError("private-secret", request=request, response=response), + ]) + fault: Final = classify_list_exception(failure) + assert fault.tag == ("auth_required" if status_code == 401 else "forbidden") + assert list_fault_http_status(fault) == status_code + assert fault.sse_event_limit_bytes is None + + def test_wire_value_carries_no_prose(): fault = ServerListFault(tag="upstream_error", status_code=500) assert outcome_wire_value(fault) == {"status": "upstream_error", "http_status": 500} From 1c022d1d06dba1cb42f9190c7b2641a1402bd2d9 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Mon, 26 Oct 2026 11:31:26 -0700 Subject: [PATCH 04/10] chore(cost-map): add openai gpt-6.1-terra from the pricing page (#48429) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 102 ++++++++++++++++++ model_prices_and_context_window.json | 102 ++++++++++++++++++ 2 files changed, 204 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9a7d88b8eda..0a85d8a00cd 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -78652,5 +78652,107 @@ "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", "supports_web_search": true + }, + "gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_12hr": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-07, + "cache_creation_input_token_cost_priority": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "cache_read_input_token_cost_above_272k_tokens_batches": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-07, + "cache_read_input_token_cost_batches": 2.5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, + "cache_read_input_token_cost_priority": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "input_cost_per_token_above_272k_tokens_batches": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_priority": 4e-06, + "input_cost_per_token_batches": 5e-07, + "input_cost_per_token_flex": 5e-07, + "input_cost_per_token_priority": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "output_cost_per_token_above_272k_tokens_batches": 3.75e-06, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.5e-05, + "output_cost_per_token_batches": 2.5e-06, + "output_cost_per_token_flex": 2.5e-06, + "output_cost_per_token_priority": 1e-05, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6.1-terra": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9a7d88b8eda..0a85d8a00cd 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -78652,5 +78652,107 @@ "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", "supports_web_search": true + }, + "gpt-6.1-terra": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_12hr": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_flex": 1.25e-06, + "cache_creation_input_token_cost_above_12hr_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_creation_input_token_cost_flex": 6.25e-07, + "cache_creation_input_token_cost_priority": 2.5e-06, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_above_272k_tokens": 1e-07, + "cache_read_input_token_cost_above_272k_tokens_batches": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 5e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-07, + "cache_read_input_token_cost_batches": 2.5e-08, + "cache_read_input_token_cost_flex": 2.5e-08, + "cache_read_input_token_cost_priority": 1e-07, + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 2e-06, + "input_cost_per_token_above_272k_tokens_batches": 1e-06, + "input_cost_per_token_above_272k_tokens_flex": 1e-06, + "input_cost_per_token_above_272k_tokens_priority": 4e-06, + "input_cost_per_token_batches": 5e-07, + "input_cost_per_token_flex": 5e-07, + "input_cost_per_token_priority": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_above_272k_tokens": 7.5e-06, + "output_cost_per_token_above_272k_tokens_batches": 3.75e-06, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.5e-05, + "output_cost_per_token_batches": 2.5e-06, + "output_cost_per_token_flex": 2.5e-06, + "output_cost_per_token_priority": 1e-05, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6.1-terra": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } From 38f318e73659c5ae73c7ce26f387e1933ba210e9 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 26 Oct 2026 11:44:30 -0700 Subject: [PATCH 05/10] fix(otel): decide per-team Arize sampling from the trace id so every pod agrees (#47468) * fix(otel): decide per-team Arize sampling from the trace id so every pod agrees * test(otel): cover per-team sampling across pods, workers, endpoints and restarts * fix(otel): remember only a failed tree's sampling verdict so a later failed call answers to the error rate A successful tree's verdict follows from the trace id alone, so remembering it only made a later failed call of the same trace follow the success rate on the pod that served both, while other pods judged it at the error rate * test(otel): flush every worker before the restart cell kills it and cover a failed call after a dropped one * fix(otel): decide team sampling per request, not per trace Calls sharing a caller's trace id are separate requests, keyed by their root span. Each is decided by the trace hash at the rate for its own outcome, and only its own late spans follow its verdict, so no pod applies an earlier call's verdict to a later call of the trace * fix(otel): forget a span's request root when it ends, not when newer spans start The root of each open span was kept in a first-in-first-out table of 16384 entries, so a span two levels below the root that stayed open while the worker started more spans than that lost its entry and ended under its parent's id. Its request then never closed and was held from every sampled destination until shutdown. Each open span now keeps its entry until it ends, and a span starting under a parent that has already ended reads the root from the parent span itself * fix(otel): decide a request as its spans end once it falls out of the open span table A request whose span was pushed out of the table of open spans, because the worker held more than 16384 open at once, ended that span under another key and was held from every sampled destination until shutdown. Pushing a span out now drops its request's open count too, so the request is decided as each of its spans ends, the way one pushed out of the open request table already is * test(otel): annotate the team sampling test locals as Final * test(otel): build the caller's remote parent in one helper instead of rebinding it per trace * fix(otel): keep each open span's request until it ends, bounded by the open requests The span table had its own bound of 16384, so a worker holding more spans open than that pushed out the oldest, and a failure in such a span no longer reached its request. Each request now holds the ids of its open spans, and a span's entry in the lookup goes when the span ends or when its request is pushed out of the open request table, the same bound main already applies. A failure anywhere in a tracked request reaches its verdict however many spans are open --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../integrations/otel/model/destination.py | 2 +- .../integrations/otel/plumbing/providers.py | 160 ++-- .../test_otel_team_sampling_across_pods.py | 723 ++++++++++++++++++ .../otel/test_otel_v2_destinations.py | 403 ++++++++-- 4 files changed, 1180 insertions(+), 108 deletions(-) create mode 100644 tests/integration/observability/test_otel_team_sampling_across_pods.py diff --git a/litellm/integrations/otel/model/destination.py b/litellm/integrations/otel/model/destination.py index 4226b4361dd..e760ba3600b 100644 --- a/litellm/integrations/otel/model/destination.py +++ b/litellm/integrations/otel/model/destination.py @@ -42,7 +42,7 @@ class OtelDestination(LiteLLMBaseModel): success_sampling_rate: float | None = Field( default=None, description=( - "Share of the request trees forwarded, 0.0..1.0, drawn once per request when its root span ends; " + "Share of the request trees forwarded, 0.0..1.0, picked by a hash of the trace id; " "``None`` forwards every one." ), ) diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 898eab5715b..fc1e9b7099d 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -1,13 +1,14 @@ """Provider / exporter factory + the Baggage span processor.""" import queue -import random import threading import time +import weakref from collections import OrderedDict from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from hashlib import sha256 from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from opentelemetry import _logs, baggage, metrics, trace from opentelemetry._logs import Logger, LoggerProvider, NoOpLoggerProvider @@ -486,15 +487,15 @@ def is_llm_call_span(span: ReadableSpan) -> bool: return GenAI.OPERATION_NAME in attributes and MCP.METHOD_NAME not in attributes -#: Request trees held back for a sampled destination until the last span of the trace +#: Request trees held back for a sampled destination until the last span of the request #: still open ends, the verdicts kept for the spans that start after that (a post-call -#: database write), and the spans in flight per trace. All bounded, so a span that never -#: ends, or a flood of requests, cannot hold spans for ever: the oldest tree is decided -#: on what it has. +#: database write), the spans in flight per request, and the request each open span +#: belongs to. All bounded, so a span that never ends, or a flood of requests, cannot +#: hold spans for ever: the oldest tree is decided on what it has. _MAX_PENDING_TREES: Final = 1024 _MAX_PENDING_SPANS_PER_TREE: Final = 512 _MAX_REMEMBERED_VERDICTS: Final = 4096 -_MAX_OPEN_TRACES: Final = 16384 +_MAX_OPEN_REQUESTS: Final = 16384 class _PendingTree: @@ -512,25 +513,35 @@ def _is_sampled(destination: "OtelDestination") -> bool: return destination.success_sampling_rate is not None or destination.error_sampling_rate is not None -#: A verdict is one request tree's draw for one destination at its rates: the exporter +#: A request is its trace id and the id of its root, the span of it whose parent is +#: remote or missing. Calls sharing a caller's trace are separate requests, each decided +#: on its own outcome. +_RequestKey: TypeAlias = tuple[int, int] + +#: A verdict is one request's keep or drop for one destination at its rates: the exporter #: ``cache_key`` leaves the rates out, so two destinations sharing an exporter with -#: different rates would otherwise share one draw. -_VerdictKey = tuple[int, object, float | None, float | None] +#: different rates would otherwise share one verdict. +_VerdictKey: TypeAlias = tuple[int, int, object, float | None, float | None] -def _verdict_key(trace_id: int, destination: "OtelDestination") -> _VerdictKey: - return (trace_id, destination.cache_key(), destination.success_sampling_rate, destination.error_sampling_rate) +def _verdict_key(request: _RequestKey, destination: "OtelDestination") -> _VerdictKey: + return (*request, destination.cache_key(), destination.success_sampling_rate, destination.error_sampling_rate) -def _keeps_tree(destination: "OtelDestination", failed: bool, draw: Callable[[], float]) -> bool: - """Whether one draw against the destination's rate keeps a request tree, read the - way the legacy Arize callback reads ``arize_success_sampling_rate`` / - ``arize_error_sampling_rate``: a tree with a failed span answers to the error - rate, an unset rate keeps everything and ``0.0`` keeps nothing.""" +def _keeps_tree(destination: "OtelDestination", failed: bool, trace_id: int) -> bool: + """Whether the destination's rate keeps a request tree of ``trace_id``: a tree with + a failed span answers to ``arize_error_sampling_rate``, the rest to + ``arize_success_sampling_rate``, an unset rate keeps everything and ``0.0`` keeps + nothing. + + The first eight bytes of a SHA-256 of the trace id are read against the rate, so + every process deciding a request of the trace with the same outcome reaches the same + verdict.""" rate: Final = destination.error_sampling_rate if failed else destination.success_sampling_rate if rate is None: return True - return rate > 0.0 and draw() <= rate + digest: Final = sha256(trace_id.to_bytes(16, "big")).digest() + return int.from_bytes(digest[:8], "big") < round(rate * 2**64) def _in_scope(span: ReadableSpan, scope: "OtelSpanScope") -> bool: @@ -667,10 +678,8 @@ class TenantFanOutSpanProcessor(SpanProcessor): default_capture_by_backend: Mapping[str, str] = MappingProxyType({}), pending_drains: int = _MAX_PENDING_DRAINS, drain_pool: _DrainPool | None = None, - sampling_draw: Callable[[], float] | None = None, ) -> None: self._operator_sinks: Final = operator_sinks - self._draw: Final = sampling_draw if sampling_draw is not None else random.random self._excluded_db_systems: Final = excluded_db_systems self._default_capture: Final = default_capture self._default_capture_by_backend: Final = default_capture_by_backend @@ -682,24 +691,32 @@ class TenantFanOutSpanProcessor(SpanProcessor): self._processors: OrderedDict[object, SpanProcessor] = OrderedDict() # mutable-ok: bounded LRU self._retired: OrderedDict[int, SpanProcessor] = OrderedDict() # mutable-ok: drains as exports finish self._exporting: dict[int, int] = {} # mutable-ok: per-processor in-flight export count - self._pending: OrderedDict[int, _PendingTree] = OrderedDict() # mutable-ok: bounded, by trace id - self._open: OrderedDict[int, int] = OrderedDict() # mutable-ok: bounded, spans in flight by trace id + self._pending: OrderedDict[_RequestKey, _PendingTree] = OrderedDict() # mutable-ok: bounded, by request + self._open: OrderedDict[_RequestKey, set[int]] = OrderedDict() # mutable-ok: bounded, open span ids by request + self._roots: dict[int, _RequestKey] = {} # mutable-ok: holds only the spans ``_open`` holds, request by span id + self._span_roots: weakref.WeakKeyDictionary[SDKSpan, int] = weakref.WeakKeyDictionary() self._verdicts: OrderedDict[_VerdictKey, bool] = OrderedDict() # mutable-ok: bounded LRU self._drain: Final = drain_pool if drain_pool is not None else _DrainPool(capacity=pending_drains) def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None: context: Final = span.context - trace_id: Final = context.trace_id if context is not None else 0 with self._lock: - self._open[trace_id] = self._open.get(trace_id, 0) + 1 - while len(self._open) > _MAX_OPEN_TRACES: - self._open.popitem(last=False) + request: Final = self._request_of(span, trace.get_current_span(parent_context)) + if context is not None: + self._span_roots[span] = request[1] + self._roots[context.span_id] = request + self._open.setdefault(request, set()).add(context.span_id) + while len(self._open) > _MAX_OPEN_REQUESTS: + for forgotten in self._open.popitem(last=False)[1]: + self._roots.pop(forgotten, None) def on_end(self, span: ReadableSpan) -> None: suppressed: Final = suppressed_backends() attributes: Final = span.attributes or _NO_ATTRIBUTES - context: Final = span.context - trace_id: Final = context.trace_id if context is not None else 0 + with self._lock: + request: Final = self._request_of(span) + if span.context is not None: + self._roots.pop(span.context.span_id, None) failed: Final = span.status.status_code is StatusCode.ERROR for destination in request_destinations(): if ( @@ -716,11 +733,32 @@ class TenantFanOutSpanProcessor(SpanProcessor): if not _is_sampled(destination): self._forward(span, destination) continue - for held_span, held_destination in self._route(trace_id, span, destination, failed): + for held_span, held_destination in self._route(request, span, destination, failed): self._forward(held_span, held_destination) - for held_span, held_destination in self._settle(trace_id, failed): + for held_span, held_destination in self._settle(request, span, failed): self._forward(held_span, held_destination) + def _request_of(self, span: ReadableSpan, parent_span: Span | None = None) -> _RequestKey: + """The request ``span`` belongs to, read under ``_lock``: the request recorded while + it is open, else the span itself when its parent is remote or missing, else its + parent's request, recorded while the parent is open or on ``parent_span`` for a + child starting under a parent that has ended, else the parent itself.""" + context: Final = span.context + if context is None: + return (0, 0) + recorded: Final = self._roots.get(context.span_id) + if recorded is not None: + return recorded + parent: Final = span.parent + if parent is None or parent.is_remote: + return (context.trace_id, context.span_id) + parent_request: Final = self._roots.get(parent.span_id) + if parent_request is not None: + return parent_request + if isinstance(parent_span, SDKSpan): + return (context.trace_id, self._span_roots.get(parent_span, parent.span_id)) + return (context.trace_id, parent.span_id) + def _default_capture_for(self, destination: "OtelDestination") -> str: backend: Final = destination.callback_name return ( @@ -743,14 +781,14 @@ class TenantFanOutSpanProcessor(SpanProcessor): self._release(processor) def _route( - self, trace_id: int, span: ReadableSpan, destination: "OtelDestination", failed: bool + self, request: _RequestKey, span: ReadableSpan, destination: "OtelDestination", failed: bool ) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": """The spans to forward now for ``span`` on a sampled ``destination``: the span itself when its tree's verdict is already known and keeps it, nothing while the - tree is held for the rest of its trace, or what the bounds made the fan-out decide + tree is held for the rest of its request, or what the bounds made the fan-out decide early (a tree that outgrew its cap is decided on what it has, and so is the oldest tree once too many are waiting, so a span that never ends holds nothing back for - ever). A failed span marks its tree before any such early draw. Once shutdown + ever). A failed span marks its tree before any such early decision. Once shutdown has started flushing, nothing is held: a span ending then is decided at once, since no later flush would reach it. @@ -758,58 +796,62 @@ class TenantFanOutSpanProcessor(SpanProcessor): another thread between the two would leave this span in a fresh tree. """ with self._lock: - verdict: Final = self._verdicts.get(_verdict_key(trace_id, destination)) + verdict: Final = self._verdicts.get(_verdict_key(request, destination)) if verdict is not None: return ((span, destination),) if verdict else () - tree: Final = self._pending.setdefault(trace_id, _PendingTree()) + tree: Final = self._pending.setdefault(request, _PendingTree()) tree.held.append((span, destination)) if failed: tree.failed = True if self._draining or len(tree.held) >= _MAX_PENDING_SPANS_PER_TREE: - return self._decide_locked(trace_id) + return self._decide_locked(request) if len(self._pending) > _MAX_PENDING_TREES: oldest: Final = next(iter(self._pending)) return self._decide_locked(oldest) return () - def _decide(self, trace_id: int) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + def _decide(self, request: _RequestKey) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": with self._lock: - return self._decide_locked(trace_id) + return self._decide_locked(request) - def _settle(self, trace_id: int, failed: bool) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": - """What the trace's held tree exports once the span that just ended was the last - of its trace still open. The model-call span ends in the post-call callback, - after the server span, so a tree is decided when its trace goes quiet rather than + def _settle( + self, request: _RequestKey, span: ReadableSpan, failed: bool + ) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + """What the request's held tree exports once the span that just ended was the last + of its request still open. The model-call span ends in the post-call callback, + after the server span, so a tree is decided when its request goes quiet rather than when its root ends, and a model call that fails late still answers to the error - rate. A trace whose starts this fan-out never saw (attached mid-flight, or forgotten - by the bound) is decided as each of its spans ends. + rate. A request whose starts this fan-out never saw (attached mid-flight, or + forgotten by the bound) is decided as each of its spans ends. """ with self._lock: - tree: Final = self._pending.get(trace_id) + tree: Final = self._pending.get(request) if tree is not None and failed: tree.failed = True - still_open: Final = self._open.pop(trace_id, 1) - 1 - if still_open > 0: - self._open[trace_id] = still_open + still_open: Final = self._open.pop(request, set()) + if span.context is not None: + still_open.discard(span.context.span_id) + if still_open: + self._open[request] = still_open return () - return self._decide_locked(trace_id) + return self._decide_locked(request) - def _decide_locked(self, trace_id: int) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": - """Draw once per destination for the held tree; the spans the draws keep, and a - verdict remembered for the spans of this trace that start later.""" - tree: Final = self._pending.pop(trace_id, None) + def _decide_locked(self, request: _RequestKey) -> "tuple[tuple[ReadableSpan, OtelDestination], ...]": + """Decide once per destination for the held tree; the spans the verdicts keep, and + a verdict remembered for the spans of this request that start later.""" + tree: Final = self._pending.pop(request, None) if tree is None: return () - verdicts: dict[_VerdictKey, bool] = {} # mutable-ok: one draw per destination of this tree + verdicts: dict[_VerdictKey, bool] = {} # mutable-ok: one verdict per destination of this tree for _, destination in tree.held: - key = _verdict_key(trace_id, destination) + key = _verdict_key(request, destination) if key not in verdicts: - verdicts[key] = _keeps_tree(destination, tree.failed, self._draw) + verdicts[key] = _keeps_tree(destination, tree.failed, request[0]) self._verdicts[key] = verdicts[key] while len(self._verdicts) > _MAX_REMEMBERED_VERDICTS: self._verdicts.popitem(last=False) return tuple( - (span, destination) for span, destination in tree.held if verdicts[_verdict_key(trace_id, destination)] + (span, destination) for span, destination in tree.held if verdicts[_verdict_key(request, destination)] ) def _operator_already_writes( @@ -850,8 +892,8 @@ class TenantFanOutSpanProcessor(SpanProcessor): with self._lock: self._draining = True undecided: Final = tuple(self._pending) - for trace_id in undecided: - for held_span, held_destination in self._decide(trace_id): + for request in undecided: + for held_span, held_destination in self._decide(request): self._forward(held_span, held_destination) deadline: Final = time.monotonic() + self._drain_seconds with self._lock: diff --git a/tests/integration/observability/test_otel_team_sampling_across_pods.py b/tests/integration/observability/test_otel_team_sampling_across_pods.py new file mode 100644 index 00000000000..dc7b295febd --- /dev/null +++ b/tests/integration/observability/test_otel_team_sampling_across_pods.py @@ -0,0 +1,723 @@ +"""A team's Arize sampling rate decides once per trace, so every pod and worker serving a call of it agrees. + +Two proxies share one database, as two pods behind a load balancer do: the first runs one worker, the second +two. Calls of one trace carry the caller's ``traceparent`` and land on different pods and workers; the team's +Arize space must receive every call of a trace or none of them. +""" + +import asyncio +import json +import re +import uuid +from collections.abc import Callable, Generator, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass, replace +from itertools import chain, product +from pathlib import Path +from typing import Final, Literal, TypeAlias +from urllib.parse import urlparse + +import httpx +import pytest +import yaml +from integration._support.client import ( + Scenario, + eventually, + gateway_from_environment, + list_value, + object_value, + until_asserted, +) +from integration._support.database import scratch_database +from integration._support.otlp_sink import Span, SpanSinks, configure_sink, recorded_requests, recorded_spans +from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +MODEL: Final = "team-sampling" +SERVER: Final = 2 +CALL_ID: Final = "litellm.call_id" +OPERATION: Final = "gen_ai.operation.name" +SPACE: Final = "space-sampled" +CREDENTIALS: Final = {"arize_space_key": SPACE, "arize_api_key": "arize-sampled"} +FAILING: Final = re.compile(rb"sampling-fail") +JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +TRACES: Final = 16 +Endpoint: TypeAlias = Literal["chat", "responses", "messages"] +Verdicts: TypeAlias = Mapping[str, tuple[bool, ...]] + + +def _chat_reply(stream: bool) -> Reply: + usage: Final[JsonValue] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9} + if not stream: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-sampling", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": usage, + } + ).encode() + ) + chunk: Final = {"id": "chatcmpl-sampling", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"} + deltas: Final[tuple[JsonValue, ...]] = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {**chunk, "choices": [], "usage": usage}, + ) + return Reply( + content_type="text/event-stream", + chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"), + ) + + +def _responses_reply(stream: bool) -> Reply: + response: Final[Mapping[str, JsonValue]] = { + "id": "resp_sampling", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": "msg_sampling", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9}, + } + if not stream: + return Reply(body=json.dumps(response).encode()) + events: Final[tuple[Mapping[str, JsonValue], ...]] = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**response, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_sampling", + "output_index": 0, + "content_index": 0, + "delta": "ok", + }, + {"type": "response.completed", "sequence_number": 2, "response": dict(response)}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events), + ) + + +def _upstream(request: Request) -> Reply: + if FAILING.search(request.body) is not None: + return Reply(status=500, body=b'{"error":{"message":"scripted upstream failure","type":"server_error"}}') + stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True + if request.target.endswith("/responses"): + return _responses_reply(stream) + return _chat_reply(stream) + + +@dataclass(frozen=True, slots=True) +class Sent: + trace_id: str + parent_span_id: str + call_id: str + + +@dataclass(frozen=True, slots=True) +class Call: + path: str + body: Mapping[str, JsonValue] + headers: Mapping[str, str] + sent: Sent + expected: int + + +def _call( + key: str, trace_id: str | None, endpoint: Endpoint = "chat", stream: bool = False, content: str = "hello" +) -> Call: + parent: Final = uuid.uuid4().hex[:16] + call_id: Final = str(uuid.uuid4()) + messages: Final[JsonValue] = [{"role": "user", "content": content}] + path, body = { + "chat": ("/v1/chat/completions", {"model": MODEL, "messages": messages, "stream": stream}), + "responses": ("/v1/responses", {"model": MODEL, "input": content, "stream": stream}), + "messages": ("/v1/messages", {"model": MODEL, "max_tokens": 16, "messages": messages, "stream": stream}), + }[endpoint] + headers: Final = { + "Authorization": f"Bearer {key}", + "x-litellm-call-id": call_id, + "connection": "close", + **({} if trace_id is None else {"traceparent": f"00-{trace_id}-{parent}-01"}), + } + expected: Final = 500 if FAILING.search(content.encode()) else 200 + return Call(path, body, headers, Sent(trace_id or "", parent if trace_id else "", call_id), expected) + + +@dataclass(frozen=True, slots=True) +class Pods: + """``first`` serves with one worker, ``second`` with two.""" + + first: OwnedProxy + second: OwnedProxy + scenario: Scenario + sinks: SpanSinks + + def team_key(self, variables: Mapping[str, str]) -> str: + team: Final = self.scenario.team() + attached: Final = self.first.gateway.request( + "POST", f"/team/{team}/callback", {"callback_name": "arize", "callback_vars": dict(variables)} + ) + assert attached.status_code == 200, attached.text + return self.scenario.key(team_id=team) + + def send(self, pod: OwnedProxy, call: Call) -> Sent: + with pod.gateway.client.stream("POST", call.path, json=call.body, headers=call.headers) as response: + text: Final = response.read().decode() + assert response.status_code == call.expected, text + return call.sent + + async def send_async(self, client: httpx.AsyncClient, call: Call) -> Sent: + async with client.stream("POST", call.path, json=call.body, headers=call.headers) as response: + text: Final = (await response.aread()).decode() + assert response.status_code == call.expected, text + return call.sent + + def on_each_worker(self, pod: OwnedProxy, call: Callable[[], Call]) -> tuple[Sent, ...]: + """One call served by each worker of ``pod``, the others suspended while it runs.""" + workers: Final = sorted(pod.worker_pids()) + assert len(workers) == pod.gateway.workers, workers + sent: Final[list[Sent]] = [] # mutable-ok: one call per worker, in worker order + for serving in workers: + with _suspended_all_but(pod, workers, serving): + sent.append(self.send(pod, call())) + return tuple(sent) + + def fence(self, key: str) -> tuple[Sent, ...]: + """Calls of an unsampled team on the sampled team's exporter, one per worker of both pods. + + Each worker sends spans to that exporter in the order they were handed to it, so all of these + arriving means every call any worker decided before them has been sent too. + """ + return (self.send(self.first, _call(key, uuid.uuid4().hex)),) + self.on_each_worker( + self.second, lambda: _call(key, uuid.uuid4().hex) + ) + + +@contextmanager +def _suspended_all_but(pod: OwnedProxy, workers: Sequence[int], serving: int) -> Generator[None]: + if len(workers) == 1: + yield + return + others: Final = [pid for pid in workers if pid != serving] + with pod.suspended(others[0]): + yield + + +def _space_id(request: Mapping[str, JsonValue]) -> JsonValue: + return {name.lower(): value for name, value in object_value(request.get("headers") or {}).items()}.get( + "arize-space-id" + ) + + +def _space_span_ids(url: str) -> frozenset[str]: + """Span ids the Arize sink received in export requests sent with the team's space id.""" + ours: Final = ( + list_value(request.get("span_ids") or []) for request in recorded_requests(url) if _space_id(request) == SPACE + ) + return frozenset(span_id for span_id in chain.from_iterable(ours) if isinstance(span_id, str)) + + +def _team_spans(pods: Pods, since: int) -> tuple[Span, ...]: + owned: Final = _space_span_ids(pods.sinks.arize) + return tuple(span for span in recorded_spans(pods.sinks.arize, since)[1] if span["span_id"] in owned) + + +def _delivered(spans: tuple[Span, ...], sent: Sequence[Sent]) -> frozenset[str]: + """Call ids of ``sent`` whose server span is among ``spans``.""" + by_parent: Final = {(call.trace_id, call.parent_span_id): call.call_id for call in sent} + return frozenset( + by_parent[(span["trace_id"], span["parent_span_id"])] + for span in spans + if span["kind"] == SERVER and (span["trace_id"], span["parent_span_id"]) in by_parent + ) + + +def _await_logged(pods: Pods, since: int, sent: Sequence[Sent]) -> None: + """Wait until every call of ``sent`` has its model span in the operator's unsampled copy.""" + expected: Final = {call.call_id for call in sent} + eventually( + lambda: ( + { + str(span["attributes"][CALL_ID]) + for span in recorded_spans(pods.sinks.operator, since)[1] + if OPERATION in span["attributes"] and CALL_ID in span["attributes"] + } + & expected + ), + lambda logged: logged == expected, + seconds=60, + ) + + +@dataclass(frozen=True, slots=True) +class Cursors: + operator: int + arize: int + + +def _cursors(pods: Pods) -> Cursors: + return Cursors(recorded_spans(pods.sinks.operator)[0], recorded_spans(pods.sinks.arize)[0]) + + +def _await_flushed(pods: Pods, cursors: Cursors, fence_key: str) -> None: + """Wait until every worker of both pods has sent the team's exporter everything it decided so far.""" + fences: Final = pods.fence(fence_key) + + def arrived() -> None: + delivered: Final = _delivered(_team_spans(pods, cursors.arize), fences) + assert {fence.call_id for fence in fences} <= delivered, "the unsampled team's calls have not arrived yet" + + until_asserted(arrived, seconds=30) + + +def _settled( + pods: Pods, + cursors: Cursors, + by_trace: Mapping[str, Sequence[Sent]], + fence_key: str, + holds: Callable[[tuple[bool, ...]], bool] = lambda kept: len(set(kept)) == 1, +) -> Verdicts: + """Per trace, which of its calls reached the team's space, once every worker has sent what it decided + and ``holds`` is true of every trace.""" + calls: Final = tuple(chain.from_iterable(by_trace.values())) + _await_logged(pods, cursors.operator, calls) + _await_flushed(pods, cursors, fence_key) + + def check() -> Verdicts: + delivered: Final = _delivered(_team_spans(pods, cursors.arize), calls) + verdicts: Final = { + trace_id: tuple(call.call_id in delivered for call in sent) for trace_id, sent in by_trace.items() + } + broken: Final = {trace_id: kept for trace_id, kept in verdicts.items() if not holds(kept)} + assert broken == {}, f"{len(broken)} of {len(verdicts)} traces were split (per call, kept): {broken}" + return verdicts + + return until_asserted(check, seconds=30) + + +def _trace_ids(cell: str, count: int = TRACES) -> tuple[str, ...]: + """Fixed RFC 4122 shaped trace ids, like the uuid4 hex a caller mints for ``traceparent``.""" + return tuple(uuid.uuid5(uuid.NAMESPACE_URL, f"team-sampling/{cell}/{index}").hex for index in range(count)) + + +def _kept(verdicts: Verdicts) -> int: + return sum(1 for kept in verdicts.values() if all(kept)) + + +@pytest.mark.parametrize( + ("variables", "content"), + [({"arize_success_sampling_rate": "0.5"}, "hello"), ({"arize_error_sampling_rate": "0.5"}, "sampling-fail")], + ids=["success-rate", "error-rate"], +) +def test_each_trace_reaches_the_team_space_whole_or_not_at_all( + pods: Pods, variables: Mapping[str, str], content: str +) -> None: + """The ticket's case: three calls per trace load balanced over the pods.""" + sampled: Final = pods.team_key({**CREDENTIALS, **variables}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple( + pods.send(pod, _call(sampled, trace_id, content=content)) for pod in (pods.first, pods.second, pods.first) + ) + for trace_id in _trace_ids("-".join(sorted(variables))) + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES, f"a rate of 0.5 kept {_kept(verdicts)} of {TRACES} traces" + + +@pytest.mark.parametrize("client", ["sync", "async"]) +def test_every_endpoint_follows_the_verdict_of_its_trace(pods: Pods, client: str) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + plan: Final[tuple[tuple[OwnedProxy, Endpoint, bool], ...]] = ( + (pods.first, "chat", False), + (pods.second, "chat", True), + (pods.first, "responses", False), + (pods.second, "messages", True), + ) + + async def send_trace(trace_id: str) -> tuple[Sent, ...]: + clients: Final = { + id(pod): httpx.AsyncClient(base_url=str(pod.gateway.client.base_url), timeout=15, trust_env=False) + for pod in (pods.first, pods.second) + } + try: + return tuple( + await asyncio.gather( + *( + pods.send_async(clients[id(pod)], _call(sampled, trace_id, endpoint, stream)) + for pod, endpoint, stream in plan + ) + ) + ) + finally: + for opened in clients.values(): + await opened.aclose() + + by_trace: Final = { + trace_id: ( + tuple(pods.send(pod, _call(sampled, trace_id, endpoint, stream)) for pod, endpoint, stream in plan) + if client == "sync" + else asyncio.run(send_trace(trace_id)) + ) + for trace_id in _trace_ids(f"endpoints-{client}") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_the_workers_of_one_pod_agree(pods: Pods) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: pods.on_each_worker(pods.second, lambda trace_id=trace_id: _call(sampled, trace_id)) + for trace_id in _trace_ids("workers") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_replaced_worker_reaches_the_verdict_its_predecessor_reached(pods: Pods) -> None: + """A worker that restarts forgets the verdicts it remembered; the calls it serves next for the same + traces still answer the way the first ones did.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + fence_key: Final = pods.team_key(CREDENTIALS) + traces: Final = _trace_ids("replaced") + before: Final = { + trace_id: pods.on_each_worker(pods.second, lambda t=trace_id: _call(sampled, t)) for trace_id in traces + } + _await_logged(pods, cursors.operator, tuple(chain.from_iterable(before.values()))) + _await_flushed(pods, cursors, fence_key) + for victim in sorted(pods.second.worker_pids()): + pods.second.kill_worker(victim) + by_trace: Final = { + trace_id: before[trace_id] + pods.on_each_worker(pods.second, lambda t=trace_id: _call(sampled, t)) + for trace_id in traces + } + + verdicts: Final = _settled(pods, cursors, by_trace, fence_key) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_failed_call_is_kept_wherever_a_successful_call_of_its_trace_is(pods: Pods) -> None: + """One call of the trace fails on one pod and answers to the error rate, another succeeds on the other + pod and answers to the success rate: with the error rate above the success rate, no trace keeps its + successful call and drops its failed one.""" + sampled: Final = pods.team_key( + {**CREDENTIALS, "arize_success_sampling_rate": "0.3", "arize_error_sampling_rate": "0.7"} + ) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: ( + pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")), + pods.send(pods.second, _call(sampled, trace_id)), + ) + for trace_id in _trace_ids("mixed", 32) + } + + verdicts: Final = _settled( + pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept[0] or not kept[1] + ) + + assert any(kept == (True, False) for kept in verdicts.values()), "no trace sat between the two rates" + assert any(kept == (True, True) for kept in verdicts.values()), "no trace sat under both rates" + + +def test_a_failed_call_answers_to_the_error_rate_on_a_pod_that_dropped_its_trace(pods: Pods) -> None: + """A pod that already dropped a successful call of the trace still keeps a later failed call of it at an + unset error rate, as the pod that only served the failed call does.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.3"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("failed-after-dropped") + succeeded: Final = {trace_id: pods.send(pods.first, _call(sampled, trace_id)) for trace_id in traces} + _await_logged(pods, cursors.operator, tuple(succeeded.values())) + by_trace: Final = { + trace_id: ( + succeeded[trace_id], + pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")), + pods.send(pods.second, _call(sampled, trace_id, content="sampling-fail")), + ) + for trace_id in traces + } + + verdicts: Final = _settled( + pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept[1] and kept[2] + ) + + assert any(not kept[0] for kept in verdicts.values()), "no successful call was dropped before its trace failed" + + +def test_a_successful_call_answers_to_the_success_rate_on_a_pod_that_kept_its_failed_call(pods: Pods) -> None: + """A pod that already kept a failed call of the trace drops a later successful call of it at a zero + success rate, as the pod that only served the successful call does.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.0"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("succeeded-after-kept") + failed: Final = { + trace_id: pods.send(pods.first, _call(sampled, trace_id, content="sampling-fail")) for trace_id in traces + } + _await_logged(pods, cursors.operator, tuple(failed.values())) + by_trace: Final = { + trace_id: ( + failed[trace_id], + pods.send(pods.first, _call(sampled, trace_id)), + pods.send(pods.second, _call(sampled, trace_id)), + ) + for trace_id in traces + } + + _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept == (True, False, False)) + + +def test_concurrent_calls_of_one_trace_on_one_worker_each_answer_to_their_own_outcome(pods: Pods) -> None: + """A failed and a successful call of the trace in flight at once on the single worker of the first pod + are decided apart: the failed one is kept at an unset error rate, the successful one dropped at a zero + success rate.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.0"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("concurrent-mixed") + plan: Final = tuple(product(traces, ("sampling-fail", "hello"))) + + def send(job: tuple[str, str]) -> Sent: + return pods.send(pods.first, _call(sampled, job[0], content=job[1])) + + with ThreadPoolExecutor(max_workers=8) as pool: + sent: Final = tuple(pool.map(send, plan)) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS), holds=lambda kept: kept == (True, False)) + + +@pytest.mark.parametrize( + ("variables", "expected"), + [({"arize_success_sampling_rate": "1.0"}, TRACES), ({"arize_success_sampling_rate": "0.0"}, 0), ({}, TRACES)], + ids=["one", "zero", "unset"], +) +def test_a_rate_at_either_end_keeps_every_trace_or_none( + pods: Pods, variables: Mapping[str, str], expected: int +) -> None: + sampled: Final = pods.team_key({**CREDENTIALS, **variables}) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple(pods.send(pod, _call(sampled, trace_id)) for pod in (pods.first, pods.second)) + for trace_id in _trace_ids(f"edge-{'-'.join(f'{k}={v}' for k, v in sorted(variables.items()))}") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert _kept(verdicts) == expected + + +def test_a_request_without_a_traceparent_is_sampled_whole_at_the_rate(pods: Pods) -> None: + """With no caller trace the proxy starts one per request: each request's server and model spans travel + together and the rate still keeps some requests and drops others.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + sent: Final = tuple( + pods.send(pod, _call(sampled, None)) for _, pod in product(range(12), (pods.first, pods.second)) + ) + _await_logged(pods, cursors.operator, sent) + trace_of: Final = { + str(span["attributes"][CALL_ID]): span["trace_id"] + for span in recorded_spans(pods.sinks.operator, cursors.operator)[1] + if CALL_ID in span["attributes"] + } + fences: Final = pods.fence(pods.team_key(CREDENTIALS)) + + def check() -> Mapping[str, tuple[bool, bool]]: + spans: Final = _team_spans(pods, cursors.arize) + assert {fence.call_id for fence in fences} <= _delivered(spans, fences), ( + "the unsampled team's calls have not arrived yet" + ) + kinds: Final = { + call.call_id: ( + any(span["kind"] == SERVER for span in spans if span["trace_id"] == trace_of[call.call_id]), + any(OPERATION in span["attributes"] for span in spans if span["trace_id"] == trace_of[call.call_id]), + ) + for call in sent + } + partial: Final = {call_id: kept for call_id, kept in kinds.items() if kept[0] != kept[1]} + assert partial == {}, f"requests reached the space without their whole tree (server, model): {partial}" + return kinds + + kept: Final = sum(1 for server, _ in until_asserted(check, seconds=30).values() if server) + assert 0 < kept < len(sent), f"a rate of 0.5 kept {kept} of {len(sent)} requests" + + +def test_a_key_level_destination_follows_the_verdict_of_its_trace(pods: Pods) -> None: + team: Final = pods.scenario.team() + sampled: Final = pods.scenario.key( + team_id=team, + metadata={ + "logging": [ + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {**CREDENTIALS, "arize_success_sampling_rate": "0.5"}, + } + ] + }, + ) + cursors: Final = _cursors(pods) + by_trace: Final = { + trace_id: tuple(pods.send(pod, _call(sampled, trace_id)) for pod in (pods.first, pods.second, pods.first)) + for trace_id in _trace_ids("key-level") + } + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < TRACES + + +def test_a_concurrent_burst_keeps_every_trace_whole(pods: Pods) -> None: + """Four calls per trace, all in flight at once across both pods.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("burst", 24) + plan: Final = tuple(product(traces, (pods.first, pods.second, pods.first, pods.second))) + + def send(job: tuple[str, OwnedProxy]) -> Sent: + return pods.send(job[1], _call(sampled, job[0])) + + with ThreadPoolExecutor(max_workers=16) as pool: + sent: Final = list(pool.map(send, plan)) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < len(traces) + + +def test_a_stalled_team_space_during_a_burst_keeps_serving_and_each_trace_whole(pods: Pods) -> None: + """The team's Arize endpoint stops answering while a burst of calls runs over both pods: every call is + still served, and once the endpoint answers again each trace arrives whole or not at all.""" + sampled: Final = pods.team_key({**CREDENTIALS, "arize_success_sampling_rate": "0.5"}) + cursors: Final = _cursors(pods) + traces: Final = _trace_ids("stalled") + plan: Final = tuple(product(traces, (pods.first, pods.second))) + + def send(job: tuple[str, OwnedProxy]) -> Sent: + return pods.send(job[1], _call(sampled, job[0])) + + configure_sink(pods.sinks.arize, paused=True) + try: + with ThreadPoolExecutor(max_workers=8) as pool: + sent: Final = tuple(pool.map(send, plan)) + finally: + configure_sink(pods.sinks.arize, paused=False) + by_trace: Final = {trace_id: tuple(call for call in sent if call.trace_id == trace_id) for trace_id in traces} + + verdicts: Final = _settled(pods, cursors, by_trace, pods.team_key(CREDENTIALS)) + + assert 0 < _kept(verdicts) < len(traces) + + +def _write_config(directory: Path, sinks: SpanSinks, provider: Wire) -> Path: + shared: Final = object_value( + JSON.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + ) + config: Final = { + **shared, + "model_list": [ + *list_value(shared.get("model_list") or []), + { + "model_name": MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": provider.url + "/v1", + }, + }, + ], + "litellm_settings": { + **object_value(shared.get("litellm_settings") or {}), + "callbacks": ["langfuse_otel"], + "otel_tenant_destination_mode": "additive", + "provider_url_destination_allowed_hosts": [urlparse(sinks.operator).netloc, urlparse(sinks.arize).netloc], + }, + "router_settings": {**object_value(shared.get("router_settings") or {}), "num_retries": 0}, + "general_settings": {**object_value(shared.get("general_settings") or {}), "user_api_key_cache_ttl": 2}, + } + path: Final = directory / "team-sampling.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@contextmanager +def _two_pods(provider: Wire, sinks: SpanSinks, directory: Path) -> Generator[Pods]: + environment: Final = { + "LITELLM_OTEL_V2": "1", + "OTEL_BSP_SCHEDULE_DELAY": "300", + "MAPPER_NAMES": "genai,openinference", + "LANGFUSE_HOST": sinks.operator, + "LANGFUSE_PUBLIC_KEY": "pk-lf-operator", + "LANGFUSE_SECRET_KEY": "sk-lf-operator", + "ARIZE_SPACE_KEY": "space-operator", + "ARIZE_API_KEY": "arize-operator", + "ARIZE_HTTP_ENDPOINT": sinks.arize + "/v1/traces", + } + removed: Final = ("LITELLM_OTEL_TENANT_DESTINATION_MODE", "OTEL_TRACES_SAMPLER", "OTEL_TRACES_SAMPLER_ARG") + config: Final = _write_config(directory, sinks, provider) + for pod in ("first", "second"): + (directory / pod).mkdir() + with gateway_from_environment() as shared, scratch_database("team_sampling") as database_url: + overrides: Final = {**environment, "DATABASE_URL": database_url} + with ( + owned_proxy_process( + shared, directory / "first", overrides, config=config, remove_environment=removed + ) as first, + owned_proxy_process( + shared, + directory / "second", + overrides, + config=config, + remove_environment=removed, + workers=2, + healthcheck=WorkerHealthcheck.LENIENT, + ) as second, + replace(first.gateway, database_url=database_url).scenario() as scenario, + ): + eventually(second.worker_pids, lambda pids: len(pids) == 2, seconds=120) + yield Pods(first, second, scenario, sinks) + + +@pytest.fixture(scope="module") +def provider() -> Iterator[Wire]: + with wire_server(_upstream) as wire: + yield wire + + +@pytest.fixture(scope="module") +def pods(provider: Wire, audit_sinks: SpanSinks, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Pods]: + with _two_pods(provider, audit_sinks, tmp_path_factory.mktemp("team-sampling")) as started: + yield started diff --git a/tests/unit/integrations/otel/test_otel_v2_destinations.py b/tests/unit/integrations/otel/test_otel_v2_destinations.py index c7a057cca5c..782442aad22 100644 --- a/tests/unit/integrations/otel/test_otel_v2_destinations.py +++ b/tests/unit/integrations/otel/test_otel_v2_destinations.py @@ -4,6 +4,7 @@ import asyncio import contextvars import threading import time +import uuid from base64 import b64encode from collections.abc import Mapping from dataclasses import replace @@ -14,10 +15,12 @@ from typing import Final import pytest from opentelemetry import trace as trace_api +from opentelemetry.context import Context from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.id_generator import RandomIdGenerator from opentelemetry.trace import Status, StatusCode import litellm @@ -1159,12 +1162,36 @@ class TestFanOut: ARIZE_TEAM_PARAMS = {"arize_space_id": "space-team", "arize_api_key": "key-team"} +def caller_context(trace_id: int) -> Context: + return trace_api.set_span_in_context( + trace_api.NonRecordingSpan( + trace_api.SpanContext(trace_id=trace_id, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01)) + ) + ) + + def arize_destination(**rates: str) -> OtelDestination: destination = destination_for("arize", {**ARIZE_TEAM_PARAMS, **rates}) assert destination is not None, "the fixture must resolve for the test to mean anything" return destination +#: Trace ids whose SHA-256 reads 0.40 and 0.80 of the way through its first eight bytes, +#: so a rate of 0.5 keeps the first and drops the second. +KEPT_AT_HALF: Final = 4 +DROPPED_AT_HALF: Final = 5 + + +class _FixedTraceId(RandomIdGenerator): + """Every trace the provider starts gets ``trace_id``.""" + + def __init__(self, trace_id: int) -> None: + self._trace_id = trace_id + + def generate_trace_id(self) -> int: + return self._trace_id + + class TestTenantSampling: """A team's ``arize_success_sampling_rate`` and ``arize_error_sampling_rate`` decide how much of its traffic the fan-out delivers to its space, the way they already decide what @@ -1173,14 +1200,17 @@ class TestTenantSampling: NAMES = {"POST /v1/chat/completions", "auth /v1/chat/completions", "chat gpt-4"} @staticmethod - def _fan_out(dest_exporter, draw=None, global_exporter=None): - provider = TracerProvider() + def _fan_out( + dest_exporter: InMemorySpanExporter, + trace_id: int | None = None, + global_exporter: InMemorySpanExporter | None = None, + ) -> TracerProvider: + provider: Final = TracerProvider() if trace_id is None else TracerProvider(id_generator=_FixedTraceId(trace_id)) if global_exporter is not None: provider.add_span_processor(_OverriddenBackendFilter(SimpleSpanProcessor(global_exporter), "arize")) - kwargs = {"processor_factory": lambda _d: SimpleSpanProcessor(dest_exporter)} - if draw is not None: - kwargs["sampling_draw"] = draw - provider.add_span_processor(TenantFanOutSpanProcessor(**kwargs)) + provider.add_span_processor( + TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) + ) return provider @staticmethod @@ -1224,7 +1254,7 @@ class TestTenantSampling: def test_an_unset_rate_exports_everything(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.99) + provider: Final = self._fan_out(dest_exporter, trace_id=DROPPED_AT_HALF) self._run(provider, (arize_destination(),)) @@ -1300,10 +1330,10 @@ class TestTenantSampling: def test_a_span_that_ends_after_the_root_follows_the_requests_verdict(self): """A post-call database write that starts after the whole request tree has ended - goes where the rest of the tree went, without a second draw.""" + goes where the rest of the tree went, without a second verdict.""" kept_exporter, dropped_exporter = InMemorySpanExporter(), InMemorySpanExporter() - kept = self._fan_out(kept_exporter, draw=lambda: 0.1) - dropped = self._fan_out(dropped_exporter, draw=lambda: 0.9) + kept: Final = self._fan_out(kept_exporter, trace_id=KEPT_AT_HALF) + dropped: Final = self._fan_out(dropped_exporter, trace_id=DROPPED_AT_HALF) destinations = (arize_destination(arize_success_sampling_rate="0.5"),) def late_tree(provider): @@ -1327,14 +1357,263 @@ class TestTenantSampling: } assert dropped_exporter.get_finished_spans() == () + def test_a_span_that_ends_after_a_failed_root_follows_the_failed_requests_verdict(self): + """A post-call database write after a failed request goes where that request went, + though on its own it would answer to the success rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + root_context: Final = trace_api.set_span_in_context(root) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=root_context): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "postgres INSERT LiteLLM_SpendLogs", + } + + def test_a_failed_call_after_a_dropped_call_of_its_trace_still_answers_to_the_error_rate(self): + """A process that dropped a successful call of the trace keeps a later failed call of + it at an unset error rate, as a process that only served the failed call does.""" + served_both, served_failed = InMemorySpanExporter(), InMemorySpanExporter() + destinations: Final = (arize_destination(arize_success_sampling_rate="0.5"),) + both: Final = self._fan_out(served_both, trace_id=DROPPED_AT_HALF) + + self._run(both, destinations) + self._run(both, destinations, failed=True) + self._run(self._fan_out(served_failed, trace_id=DROPPED_AT_HALF), destinations, failed=True) + + assert [sorted(s.name for s in exporter.get_finished_spans()) for exporter in (served_both, served_failed)] == [ + sorted(self.NAMES) + ] * 2 + + def test_a_successful_call_after_a_kept_failed_call_of_its_trace_answers_to_the_success_rate(self): + """A process that kept a failed call of the trace drops a later successful call of it + at a zero success rate, as a process that only served the successful call does.""" + served_both, served_success = InMemorySpanExporter(), InMemorySpanExporter() + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + both: Final = self._fan_out(served_both, trace_id=KEPT_AT_HALF) + + self._run(both, destinations, failed=True) + self._run(both, destinations) + self._run(self._fan_out(served_success, trace_id=KEPT_AT_HALF), destinations) + + assert len(served_both.get_finished_spans()) == len(self.NAMES), "only the failed call" + assert all(s.name in self.NAMES for s in served_both.get_finished_spans()) + assert served_success.get_finished_spans() == () + + def test_a_late_child_that_fails_after_its_request_was_dropped_is_dropped_with_it(self): + """A post-call write that fails after its successful request was dropped goes where + the request went, rather than reaching the space on its own at the error rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + root_context: Final = trace_api.set_span_in_context(root) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=root_context) as write: + write.set_status(Status(StatusCode.ERROR, "connection reset")) + + in_fresh_context(run) + + assert dest_exporter.get_finished_spans() == () + + def test_a_failure_deep_in_a_request_marks_the_whole_request(self): + """A database span that fails under the auth span puts the whole request, root + included, on the error rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("auth /v1/chat/completions"): + with tracer.start_as_current_span("postgres SELECT LiteLLM_VerificationToken") as lookup: + lookup.set_status(Status(StatusCode.ERROR, "connection reset")) + with tracer.start_as_current_span("chat gpt-4"): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + *self.NAMES, + "postgres SELECT LiteLLM_VerificationToken", + } + + def test_a_span_two_levels_down_that_outlives_many_other_spans_still_closes_its_request(self): + """A classifier call under the route span that stays open while the worker serves + more requests than it tracks at once still counts against its own request, so the + request is decided once it ends, on its own outcome.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def other_request(): + set_request_destinations(()) + with get_tracer(provider, "litellm").start_as_current_span("POST /v1/chat/completions"): + pass + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + with tracer.start_as_current_span("route gpt-4"): + with tracer.start_as_current_span("chat classifier"): + for _ in range(otel_providers._MAX_OPEN_REQUESTS + 1): + contextvars.Context().run(other_request) + with tracer.start_as_current_span("chat gpt-4"): + pass + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "chat classifier", + "chat gpt-4", + } + + def test_a_failure_two_levels_down_reaches_its_request_however_many_spans_are_open(self): + """A classifier call that fails while another request holds more spans open than + the worker tracks requests marks its own request, which is then kept at the error + rate as a whole.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + others: Final = contextvars.Context() + + def open_other_request(): + set_request_destinations(()) + tracer: Final = get_tracer(provider, "litellm") + root: Final = tracer.start_span("POST /v1/chat/completions") + parent: Final = trace_api.set_span_in_context(root) + return ( + root, + *(tracer.start_span("chat flood", context=parent) for _ in range(otel_providers._MAX_OPEN_REQUESTS)), + ) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("route gpt-4"): + with tracer.start_as_current_span("chat classifier") as classifier: + held_open: Final = others.run(open_other_request) + classifier.set_status(Status(StatusCode.ERROR, "classifier timed out")) + return held_open + + held_open: Final = in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "chat classifier", + } + for span in held_open: + span.end() + + def test_the_fan_out_forgets_a_span_once_it_ends_and_a_request_once_it_is_pushed_out(self): + """The request lookup holds only spans still open, and only those of the requests + the fan-out still counts, so neither ended spans nor a flood of open requests grow + it past the open request bound.""" + fan_out: Final = TenantFanOutSpanProcessor( + processor_factory=lambda _d: SimpleSpanProcessor(InMemorySpanExporter()) + ) + provider: Final = TracerProvider() + provider.add_span_processor(fan_out) + + def run(): + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions"): + with tracer.start_as_current_span("chat gpt-4"): + pass + assert (fan_out._roots, dict(fan_out._open)) == ({}, {}) + root: Final = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) + tracer.start_span("route gpt-4", context=trace_api.set_span_in_context(root)) + return tuple( + tracer.start_span(f"POST {index}", context=trace_api.Context()) + for index in range(otel_providers._MAX_OPEN_REQUESTS) + ) + + held_open: Final = in_fresh_context(run) + + assert len(fan_out._roots) == otel_providers._MAX_OPEN_REQUESTS + for span in held_open: + span.end() + assert (fan_out._roots, dict(fan_out._open)) == ({}, {}) + + def test_a_span_starting_under_an_ended_span_below_the_root_follows_its_request(self): + """A post-call write started under the route span after its failed request ended + is kept with the request at the error rate, rather than decided on its own at the + success rate.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + with tracer.start_as_current_span("POST /v1/chat/completions") as root: + with tracer.start_as_current_span("route gpt-4") as route: + route_context: Final = trace_api.set_span_in_context(route) + root.set_status(Status(StatusCode.ERROR, "upstream 500")) + with tracer.start_as_current_span("postgres INSERT LiteLLM_SpendLogs", context=route_context): + pass + + in_fresh_context(run) + + assert {s.name for s in dest_exporter.get_finished_spans()} == { + "POST /v1/chat/completions", + "route gpt-4", + "postgres INSERT LiteLLM_SpendLogs", + } + + def test_concurrent_calls_of_one_trace_each_answer_to_their_own_outcome(self): + """Two calls of the caller's trace in flight at once on one process, one failing, + are decided as two processes serving one each would decide them.""" + dest_exporter: Final = InMemorySpanExporter() + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.0", arize_error_sampling_rate="1.0"),) + remote: Final = trace_api.SpanContext( + trace_id=KEPT_AT_HALF, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01) + ) + + def run(): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + caller: Final = trace_api.set_span_in_context(trace_api.NonRecordingSpan(remote)) + with tracer.start_as_current_span("POST failed", context=caller) as failed: + with tracer.start_as_current_span("POST succeeded", context=caller): + pass + failed.set_status(Status(StatusCode.ERROR, "upstream 500")) + + in_fresh_context(run) + + assert [s.name for s in dest_exporter.get_finished_spans()] == ["POST failed"] + def test_a_late_span_whose_verdict_was_forgotten_is_decided_on_its_own_not_stranded(self): """Once enough other requests have been decided to evict a request's verdict, a span of it that starts late (a post-call database write) opens and closes its own count, - so it is decided on a draw of its own as soon as it ends instead of waiting in a tree - nothing would ever close.""" + so it is decided on its own as soon as it ends instead of waiting in a tree nothing + would ever close.""" dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.1) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1353,18 +1632,17 @@ class TestTenantSampling: assert names[0] == "POST /v1/chat/completions" assert names[-1] == "postgres INSERT LiteLLM_SpendLogs", "exported as it ended, not held for a root" - def test_a_callers_traceparent_neither_steers_the_draw_nor_hides_the_root(self): - """The draw is random, not read from the trace id a caller may have chosen, and the - server span under a remote parent still closes the tree.""" + def test_a_callers_traceparent_decides_by_its_trace_id_and_still_closes_the_tree(self): + """A request continuing a caller's trace is decided on that trace id, and the server + span under a remote parent still closes the tree.""" dest_exporter = InMemorySpanExporter() - draws = iter([0.9, 0.1]) - provider = self._fan_out(dest_exporter, draw=lambda: next(draws)) + provider: Final = self._fan_out(dest_exporter) destinations = (arize_destination(arize_success_sampling_rate="0.5"),) def run(): set_request_destinations(destinations) tracer = get_tracer(provider, "litellm") - for trace_id in (1, 2): + for trace_id in (DROPPED_AT_HALF, KEPT_AT_HALF): remote = trace_api.SpanContext( trace_id=trace_id, span_id=1, is_remote=True, trace_flags=trace_api.TraceFlags(0x01) ) @@ -1377,12 +1655,12 @@ class TestTenantSampling: kept = dest_exporter.get_finished_spans() assert [s.name for s in kept] == ["chat gpt-4", "POST /v1/chat/completions"] - assert all(s.context.trace_id == 2 for s in kept), "the second request, whose draw of 0.1 passed" + assert all(s.context.trace_id == KEPT_AT_HALF for s in kept) def test_a_root_that_never_ends_holds_the_tree_only_up_to_the_bound(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1397,8 +1675,8 @@ class TestTenantSampling: def test_a_flood_of_waiting_trees_decides_the_oldest_on_what_it_has(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) @@ -1413,21 +1691,21 @@ class TestTenantSampling: assert [s.name for s in dest_exporter.get_finished_spans()] == ["chat 0"] - def test_a_flood_of_open_traces_forgets_the_oldest_and_decides_it_as_its_spans_end(self): - """The fan-out counts open spans for a bounded number of traces, a bound far above - what one process keeps in flight, since every trace of the provider is counted, not - only the sampled ones. A trace pushed out of that count is decided on what it has as + def test_a_flood_of_open_requests_forgets_the_oldest_and_decides_it_as_its_spans_end(self): + """The fan-out counts open spans for a bounded number of requests, a bound far above + what one process keeps in flight, since every request of the provider is counted, not + only the sampled ones. A request pushed out of that count is decided on what it has as its next span ends, root still open, rather than held until a bound gets to it.""" dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) - destinations = (arize_destination(arize_success_sampling_rate="0.5"),) + provider: Final = self._fan_out(dest_exporter) + destinations: Final = (arize_destination(arize_success_sampling_rate="1.0"),) def run(): set_request_destinations(destinations) tracer = get_tracer(provider, "litellm") root = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) tracer.start_span("auth /v1/chat/completions", context=trace_api.set_span_in_context(root)).end() - for index in range(otel_providers._MAX_OPEN_TRACES): + for index in range(otel_providers._MAX_OPEN_REQUESTS): tracer.start_span(f"POST {index}", context=trace_api.Context()) assert dest_exporter.get_finished_spans() == (), "still held before the count is forgotten" tracer.start_span("chat gpt-4", context=trace_api.set_span_in_context(root)).end() @@ -1443,14 +1721,12 @@ class TestTenantSampling: def test_shutdown_decides_what_is_still_held(self): dest_exporter = InMemorySpanExporter() - fan_out = TenantFanOutSpanProcessor( - processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter), sampling_draw=lambda: 0.0 - ) + fan_out: Final = TenantFanOutSpanProcessor(processor_factory=lambda _d: SimpleSpanProcessor(dest_exporter)) provider = TracerProvider() provider.add_span_processor(fan_out) def run(): - set_request_destinations((arize_destination(arize_success_sampling_rate="0.5"),)) + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) tracer = get_tracer(provider, "litellm") with tracer.start_as_current_span("POST /v1/chat/completions"): with tracer.start_as_current_span("chat gpt-4"): @@ -1476,14 +1752,12 @@ class TestTenantSampling: straggler.end() super().on_end(span) - fan_out = TenantFanOutSpanProcessor( - processor_factory=lambda _d: _EndsAStragglerOnExport(dest_exporter), sampling_draw=lambda: 0.0 - ) + fan_out: Final = TenantFanOutSpanProcessor(processor_factory=lambda _d: _EndsAStragglerOnExport(dest_exporter)) provider = TracerProvider() provider.add_span_processor(fan_out) def run(): - set_request_destinations((arize_destination(arize_success_sampling_rate="0.5"),)) + set_request_destinations((arize_destination(arize_success_sampling_rate="1.0"),)) tracer = get_tracer(provider, "litellm") root = tracer.start_span("POST /v1/chat/completions", context=trace_api.Context()) tracer.start_span("chat gpt-4", context=trace_api.set_span_in_context(root)).end() @@ -1495,20 +1769,20 @@ class TestTenantSampling: assert {s.name for s in dest_exporter.get_finished_spans()} == {"chat gpt-4", "embed ada"} - def test_a_zero_rate_drops_even_a_draw_of_zero(self): + def test_a_zero_rate_drops_a_trace_a_half_rate_keeps(self): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.0) + provider: Final = self._fan_out(dest_exporter, trace_id=KEPT_AT_HALF) self._run(provider, (arize_destination(arize_success_sampling_rate="0.0"),)) assert dest_exporter.get_finished_spans() == () - def test_the_draw_is_compared_to_the_rate_the_way_the_legacy_callback_compares_it(self): + def test_the_trace_id_is_read_against_the_rate(self): dropped_exporter, kept_exporter = InMemorySpanExporter(), InMemorySpanExporter() - destinations = (arize_destination(arize_success_sampling_rate="0.2"),) + destinations: Final = (arize_destination(arize_success_sampling_rate="0.5"),) - self._run(self._fan_out(dropped_exporter, draw=lambda: 0.3), destinations) - self._run(self._fan_out(kept_exporter, draw=lambda: 0.2), destinations) + self._run(self._fan_out(dropped_exporter, trace_id=DROPPED_AT_HALF), destinations) + self._run(self._fan_out(kept_exporter, trace_id=KEPT_AT_HALF), destinations) assert dropped_exporter.get_finished_spans() == () assert {s.name for s in kept_exporter.get_finished_spans()} == self.NAMES @@ -1527,6 +1801,39 @@ class TestTenantSampling: assert all(names == self.NAMES for names in by_trace.values()) assert 0 < len(by_trace) < 64 + @pytest.mark.parametrize( + ("rates", "failed"), + [({"arize_success_sampling_rate": "0.5"}, False), ({"arize_error_sampling_rate": "0.5"}, True)], + ids=["success-rate", "error-rate"], + ) + def test_two_processes_serving_one_trace_reach_the_same_verdict(self, rates: Mapping[str, str], failed: bool): + """Two pods behind a load balancer each serve one call of a trace the caller + continued with its ``traceparent``: the team's space gets both calls or neither, + and a rate of 0.5 over UUID-shaped trace ids keeps some traces and drops others.""" + exporters: Final = (InMemorySpanExporter(), InMemorySpanExporter()) + pods: Final = tuple(self._fan_out(exporter) for exporter in exporters) + destinations: Final = (arize_destination(**rates),) + trace_ids: Final = tuple(uuid.uuid5(uuid.NAMESPACE_URL, f"pod/{index}").int for index in range(32)) + + def run(provider: TracerProvider): + set_request_destinations(destinations) + tracer: Final = get_tracer(provider, "litellm") + for trace_id in trace_ids: + with tracer.start_as_current_span("POST /v1/chat/completions", context=caller_context(trace_id)): + with tracer.start_as_current_span("chat gpt-4") as span: + if failed: + span.set_status(Status(StatusCode.ERROR, "upstream 500")) + + for provider in pods: + in_fresh_context(run, provider) + + first, second = ( + {s.context.trace_id for s in exporter.get_finished_spans() if s.context is not None} + for exporter in exporters + ) + assert first == second + assert 0 < len(first) < len(trace_ids) + def test_in_additive_mode_the_operators_own_copy_is_not_sampled(self, monkeypatch): """The rates are the team's setting for the team's space; the operator's backbone keeps every span.""" @@ -1547,7 +1854,7 @@ class TestTenantSampling: assert {s.name for s in dest_exporter.get_finished_spans()} == self.NAMES - def test_two_views_of_one_exporter_each_draw_at_their_own_rate(self): + def test_two_views_of_one_exporter_each_answer_to_their_own_rate(self): """Two destinations to the same space with different rates share an exporter, not a verdict: the 1.0 view gets the tree once and the 0.0 view adds nothing.""" dest_exporter = InMemorySpanExporter() @@ -1582,7 +1889,7 @@ class TestTenantSampling: @pytest.mark.parametrize("bad", ["abc", "-0.1", "1.5", "nan"]) def test_an_unusable_rate_exports_rather_than_dropping(self, bad): dest_exporter = InMemorySpanExporter() - provider = self._fan_out(dest_exporter, draw=lambda: 0.99) + provider: Final = self._fan_out(dest_exporter, trace_id=DROPPED_AT_HALF) self._run(provider, (arize_destination(arize_success_sampling_rate=bad),)) From 1ecc782986162c8359d38eb6b22a213fd7c8a0e0 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 14:49:50 -0400 Subject: [PATCH 06/10] fix(vertex_ai): reuse rerank credential cache (#48430) --- litellm/utils.py | 7 +- .../test_vertex_ai_rerank_transformation.py | 86 +++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index e2820f9c682..4315ae426cf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9571,7 +9571,7 @@ class ProviderConfigManager: return get_nvidia_nim_rerank_config(model) elif litellm.LlmProviders.VERTEX_AI == provider: - return litellm.VertexAIRerankConfig() + return ProviderConfigManager._get_vertex_ai_rerank_config_cached() elif litellm.LlmProviders.FIREWORKS_AI == provider: return litellm.FireworksAIRerankConfig() elif litellm.LlmProviders.VOYAGE == provider: @@ -9590,6 +9590,11 @@ class ProviderConfigManager: return get_dashscope_family_rerank_config(provider.value) return litellm.CohereRerankConfig() + @staticmethod + @lru_cache(maxsize=1) + def _get_vertex_ai_rerank_config_cached() -> BaseRerankConfig: + return litellm.VertexAIRerankConfig() + @staticmethod def get_provider_decisions_config(model: str, provider: LlmProviders) -> BaseDecisionsConfig | None: if provider == LlmProviders.PERPLEXITY: diff --git a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index 630b2e1eb34..8abd012f9d8 100644 --- a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -5,13 +5,19 @@ Based on the test patterns from other rerank providers and the current Vertex AI import json import os +from datetime import datetime, timedelta +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +import responses +import respx +import litellm from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig from litellm.types.rerank import RerankResponse +from litellm.utils import ProviderConfigManager class TestVertexAIRerankTransform: @@ -646,3 +652,83 @@ class TestVertexAIRerankTransform: second_call = mock_ensure_access_token.call_args_list[1] assert first_call.kwargs["credentials"] == "path/to/credentials.json" assert second_call.kwargs["credentials"] == "path/to/credentials.json" + + +@pytest.fixture +def clear_vertex_rerank_cache(request: pytest.FixtureRequest) -> None: + ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear() + request.addfinalizer( + ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("clear_vertex_rerank_cache") +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_rerank_reuses_scoped_credentials_and_refreshes_expired_tokens( + is_async: bool, +) -> None: + now: Final = datetime(2025, 1, 1) + credentials: Final = tuple( + json.dumps( + { + "type": "authorized_user", + "client_id": identity, + "client_secret": "synthetic-client-secret", + "refresh_token": f"synthetic-refresh-{identity}", + } + ) + for identity in ("first-principal", "second-principal") + ) + with ( + patch("litellm.disable_aiohttp_transport", True), + patch("google.auth._helpers.utcnow", return_value=now) as clock, + responses.RequestsMock(assert_all_requests_are_fired=False) as token_http, + respx.mock as rank_http, + ): + for token in ("initial", "other-credential", "other-project", "renewed"): + token_http.post( + "https://oauth2.googleapis.com/token", + json={"access_token": token, "expires_in": 3600, "token_type": "Bearer"}, + ) + rank_http.post( + url__regex=r"https://discoveryengine\.googleapis\.com/v1/projects/[^/]+/" + r"locations/global/rankingConfigs/default_ranking_config:rank" + ).respond(200, json={"records": [{"id": "0", "score": 0.9}]}) + + async def rank( + credential: str, project: str, expected_token: str, refresh_count: int + ) -> None: + params: Final = { + "model": "vertex_ai/semantic-ranker-default@latest", + "query": f"Query for {expected_token}", + "documents": ["A relevant document"], + "vertex_credentials": credential, + "vertex_project": project, + "caching": False, + "num_retries": 0, + } + result: Final = ( + await litellm.arerank(**params) + if is_async + else litellm.rerank(**params) + ) + assert isinstance(result, RerankResponse) + assert result.results[0]["index"] == 0 + assert len(token_http.calls) == refresh_count + request: Final = rank_http.calls.last.request + assert request.headers["Authorization"] == f"Bearer {expected_token}" + assert request.headers["X-Goog-User-Project"] == project + assert str(request.url) == ( + f"https://discoveryengine.googleapis.com/v1/projects/{project}/" + "locations/global/rankingConfigs/default_ranking_config:rank" + ) + assert json.loads(request.content)["query"] == params["query"] + + await rank(credentials[0], "first-project", "initial", 1) + await rank(credentials[0], "first-project", "initial", 1) + await rank(credentials[1], "first-project", "other-credential", 2) + await rank(credentials[0], "second-project", "other-project", 3) + await rank(credentials[0], "first-project", "initial", 3) + clock.return_value = now + timedelta(hours=2) + await rank(credentials[0], "first-project", "renewed", 4) From 85871f791c12ca381b53dad8cbb7c8a47076ef4c Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 26 Oct 2026 14:50:31 -0400 Subject: [PATCH 07/10] fix(ui): render string inputs in general settings (#48431) --- .../general_settings.integration.test.tsx | 36 +++++++++++++++++++ .../_components/general_settings.tsx | 9 +++++ 2 files changed, 45 insertions(+) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx index 35434830d98..d7dfb8151fc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.integration.test.tsx @@ -80,6 +80,42 @@ describe("GeneralSettings General tab", () => { vi.mocked(deleteConfigFieldSetting).mockClear(); }); + it.each([ + ["maximum_spend_logs_retention_period", null, "7d"], + ["maximum_health_check_retention_period", "30d", "14d"], + ["maximum_daily_tag_spend_retention_period", undefined, "90d"], + ["maximum_spend_logs_cleanup_run_budget", "5m", "10m"], + ["maximum_spend_logs_cleanup_batch_timeout", "", "45s"], + ["mcp_client_id_header", "x-client-id", "X-MCP-Client.ID"], + ])("displays, updates and resets the String setting %s", async (fieldName, initialValue, updatedValue) => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + ...SETTINGS_FIXTURE, + { + field_name: fieldName, + field_type: "String", + field_value: initialValue, + field_description: "string setting", + stored_in_db: initialValue != null, + field_default_value: null, + }, + ]); + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("tab", { name: "General" })); + const input = await screen.findByRole("textbox", { name: fieldName }); + expect(input).toHaveValue(initialValue ?? ""); + fireEvent.change(input, { target: { value: updatedValue } }); + const row = screen.getByRole("row", { name: new RegExp(fieldName) }); + await user.click(within(row).getByRole("button", { name: "Update" })); + expect(vi.mocked(updateConfigFieldSetting).mock.calls).toEqual([["token", fieldName, updatedValue]]); + expect(within(row).getByText("In DB")).toBeInTheDocument(); + const resetIcon = row.querySelectorAll("td")[3].querySelector("svg"); + expect(resetIcon).not.toBeNull(); + await user.click(resetIcon as unknown as Element); + expect(vi.mocked(deleteConfigFieldSetting).mock.calls).toEqual([["token", fieldName]]); + expect(input).toHaveValue(""); + }); + it("updates max_ui_session_budget with its own value, not the value at its filtered index", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx index 5d28adb7ef3..67712d78a13 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/router-settings/_components/general_settings.tsx @@ -68,6 +68,15 @@ const SettingValueEditor: React.FC<{ /> ); } + if (setting.field_type === "String") { + return ( + onChange(setting.field_name, event.target.value)} + /> + ); + } if (setting.field_type === "Integer") { return ( Date: Mon, 26 Oct 2026 12:07:58 -0700 Subject: [PATCH 08/10] fix(guardrails): exclude litellm_metadata from pass-through guardrail scan text (#48276) * fix(guardrails): exclude litellm_metadata from pass-through guardrail scan text The pass-through guardrail translation json-dumps the request body as the scan text when no request_fields are configured. The unified guardrail hook injects litellm_metadata into that body first, so the dump carried the caller's key metadata. Exclude it alongside metadata and litellm_logging_obj; vendors still receive it through request_data. * test(guardrails): type the pass-through scan-text tests per AGENTS.md --- .../guardrail_translation/handler.py | 2 +- .../test_passthrough_guardrail_scan_text.py | 153 ++++++++++++++++++ .../guardrail_translation/test_handler.py | 16 ++ .../test_unified_guardrail.py | 33 ++++ 4 files changed, 203 insertions(+), 1 deletion(-) create mode 100644 tests/integration/observability/test_passthrough_guardrail_scan_text.py diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 3ce8d7ddad5..dea8bcb73a4 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -90,7 +90,7 @@ class PassThroughEndpointHandler(BaseTranslation): payload_to_check: Final = { k: v for k, v in without_server_streaming_classification(data).items() - if not k.startswith("_") and k not in ("metadata", "litellm_logging_obj") + if not k.startswith("_") and k not in ("metadata", "litellm_metadata", "litellm_logging_obj") } verbose_proxy_logger.debug("PassThroughEndpointHandler: Using full payload for guardrail") return safe_dumps(payload_to_check) diff --git a/tests/integration/observability/test_passthrough_guardrail_scan_text.py b/tests/integration/observability/test_passthrough_guardrail_scan_text.py new file mode 100644 index 00000000000..7437f696d4c --- /dev/null +++ b/tests/integration/observability/test_passthrough_guardrail_scan_text.py @@ -0,0 +1,153 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.anthropic_sse import message_json +from integration._support.client import Gateway +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +GENERIC: Final = "passthrough-scan-generic" +AIRS: Final = "passthrough-scan-airs" +OPENAI_ROUTE: Final = "/openai/chat/completions" +ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" +ROUTE_BODIES: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + OPENAI_ROUTE: {"model": "gpt-5.6", "messages": [{"role": "user", "content": "Say hi in one word"}]}, + ANTHROPIC_ROUTE: { + "model": "claude-sonnet-5-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "Say hi in one word"}], + }, + } +) +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) + + +def _generic_verdict(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _airs_verdict(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "allow", "category": "benign", "scan_id": "s", "report_id": "r"}).encode()) + + +def _openai_reply(_request: Request) -> Reply: + return chat_reply("chatcmpl-passthrough-scan", "gpt-5.6", "Hi", stream=False) + + +def _anthropic_reply(_request: Request) -> Reply: + return Reply(body=message_json("msg_passthrough_scan", "claude-sonnet-5-5", "Hi")) + + +def _parsed_texts(texts: JsonValue, vendor_body: dict[str, JsonValue]) -> list[JsonValue]: + assert isinstance(texts, list), vendor_body + return [json.loads(text) if isinstance(text, str) else text for text in texts] + + +def _scanned_payloads(guardrail: str, vendor_body: dict[str, JsonValue]) -> list[JsonValue]: + if guardrail == GENERIC: + return _parsed_texts(vendor_body["texts"], vendor_body) + contents: Final = vendor_body["contents"] + assert isinstance(contents, list), vendor_body + return _parsed_texts([content.get("prompt") for content in contents if isinstance(content, dict)], vendor_body) + + +@dataclass(frozen=True, slots=True) +class PassThroughRig: + owned: OwnedProxy + vendors: Mapping[str, Wire] + upstreams: Mapping[str, Wire] + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[PassThroughRig]: + directory: Final = tmp_path_factory.mktemp("passthrough-guardrail-scan") + with ExitStack() as stack: + vendors: Final = MappingProxyType( + { + GENERIC: stack.enter_context(wire_server(_generic_verdict)), + AIRS: stack.enter_context(wire_server(_airs_verdict)), + } + ) + upstreams: Final = MappingProxyType( + { + OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), + ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)), + } + ) + config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + path: Final = directory / "passthrough-guardrail-scan.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "guardrails": [ + { + "guardrail_name": GENERIC, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": vendors[GENERIC].url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": AIRS, + "litellm_params": { + "guardrail": "panw_prisma_airs", + "mode": "pre_call", + "default_on": False, + "api_base": vendors[AIRS].url, + "api_key": "synthetic-airs-key", + "profile_name": "integration", + }, + }, + ], + } + ) + ) + environment: Final = { + "OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url, + "OPENAI_API_KEY": "synthetic-openai-key", + "ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + } + owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) + yield PassThroughRig(owned, vendors, upstreams) + + +@pytest.mark.parametrize("guardrail", [GENERIC, AIRS]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE]) +def test_pass_through_guardrail_scans_the_client_body_without_the_callers_key_metadata( + rig: PassThroughRig, route: str, guardrail: str +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + team: Final = scenario.team() + user: Final = scenario.user( + user_email=f"passthrough-scan-{uuid.uuid4().hex}@example.com", user_role="internal_user", teams=[team] + ) + key: Final = scenario.key(user_id=user, team_id=team) + body: Final = {**ROUTE_BODIES[route], "guardrails": [guardrail]} + + response: Final = gateway.request("POST", route, body, key=key) + + assert response.status_code == 200, response.text + scanned: Final = [ + _scanned_payloads(guardrail, _JSON.validate_json(request.body)) + for request in rig.vendors[guardrail].drain() + ] + assert scanned == [[body]], scanned + assert len(rig.upstreams[route].drain()) == 1 diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 0dba0d7b90f..7cbabb95931 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -19,6 +19,22 @@ def test_full_payload_guardrail_text_excludes_the_server_streaming_marker(): assert json.loads(text) == body, text +def test_full_payload_guardrail_text_excludes_the_injected_litellm_metadata(): + body: Final = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "guardrails": ["g"]} + litellm_metadata: Final = { + "user_api_key_token": "hashed-key", + "user_api_key_user_email": "someone@example.com", + "user_api_key_user_role": "proxy_admin", + } + + text: Final = PassThroughEndpointHandler()._extract_text_for_guardrail( + {**body, "litellm_metadata": litellm_metadata}, + None, + ) + + assert json.loads(text) == body, text + + @pytest.mark.parametrize( "marker", [ diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 3efbe6475b2..a132ce8eada 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3,6 +3,7 @@ import asyncio import contextlib import io +import json import logging from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator from types import SimpleNamespace @@ -35,6 +36,7 @@ from litellm.llms.openai.chat.guardrail_translation.handler import ( from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) +from litellm.llms.pass_through.guardrail_translation.handler import PassThroughEndpointHandler from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) @@ -3053,3 +3055,34 @@ async def test_pre_call_hook_skips_attachment_scan_when_guardrail_has_none(): ) assert len(guardrail.apply_calls) == 1 + + +@pytest.mark.parametrize("hook", ["async_pre_call_hook", "async_moderation_hook"]) +@pytest.mark.asyncio +async def test_pass_through_guardrail_text_is_the_client_body_without_the_injected_key_metadata( + monkeypatch: pytest.MonkeyPatch, hook: str +) -> None: + _patch_translation_mappings(monkeypatch, {CallTypes.pass_through: PassThroughEndpointHandler}) + guardrail: Final = RecordingGuardrail() + body: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Say hi in one word"}], + "guardrails": ["recording-guardrail"], + } + data: Final = {**body, "guardrail_to_apply": guardrail} + user_api_key_dict: Final = UserAPIKeyAuth( + token="hashed-key", user_email="someone@example.com", team_id="team-1", user_role="proxy_admin" + ) + + if hook == "async_pre_call_hook": + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=CallTypes.pass_through.value + ) + else: + await UnifiedLLMGuardrails().async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=CallTypes.pass_through.value + ) + + texts: Final = [call["inputs"]["texts"] for call in guardrail.apply_calls] + assert [[json.loads(text) for text in call_texts] for call_texts in texts] == [[body]], texts + assert data["litellm_metadata"]["user_api_key_user_email"] == "someone@example.com", data["litellm_metadata"] From dab39d7402cb1e772045c5a702abe74f69f66970 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:09:14 -0700 Subject: [PATCH 09/10] fix(azure): set azure_ai/Microsoft-Decision-1 max_input_tokens to 32768 (#48434) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 3 ++- model_prices_and_context_window.json | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0a85d8a00cd..af2b4655d6f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -76439,9 +76439,10 @@ "azure_ai/Microsoft-Decision-1": { "input_cost_per_token": 4.2e-08, "litellm_provider": "azure_ai", + "max_input_tokens": 32768, "mode": "evaluation", "output_cost_per_token": 0.0, - "source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'" + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" }, "gpt-rosalind-discovery": { "cache_read_input_token_cost": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0a85d8a00cd..af2b4655d6f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -76439,9 +76439,10 @@ "azure_ai/Microsoft-Decision-1": { "input_cost_per_token": 4.2e-08, "litellm_provider": "azure_ai", + "max_input_tokens": 32768, "mode": "evaluation", "output_cost_per_token": 0.0, - "source": "https://prices.azure.com/api/retail/prices?$filter=productName eq 'Microsoft Decision Models'" + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" }, "gpt-rosalind-discovery": { "cache_read_input_token_cost": 5e-07, From c69e82f538690c40f16e7aa6519c7e05579a55ab Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:18:19 -0700 Subject: [PATCH 10/10] chore(cost-map): add openai batch and ultrafast 12-hour cache write rates from the pricing page (#48433) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++++++++ model_prices_and_context_window.json | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index af2b4655d6f..6ec1bd29beb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -36840,10 +36840,14 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_12hr": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, "cache_creation_input_token_cost_above_12hr_priority": 5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -36931,8 +36935,10 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -37014,8 +37020,10 @@ "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_12hr": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, @@ -37186,8 +37194,10 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_12hr": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, @@ -37268,8 +37278,10 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -37349,8 +37361,10 @@ "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_12hr": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_priority": 1e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, @@ -75046,10 +75060,14 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index af2b4655d6f..6ec1bd29beb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -36840,10 +36840,14 @@ "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_12hr": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 0.0001, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 0.0003, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-05, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-05, "cache_creation_input_token_cost_above_12hr_priority": 5e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 0.00015, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, @@ -36931,8 +36935,10 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -37014,8 +37020,10 @@ "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_12hr": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_above_12hr_batches": 1.25e-07, "cache_creation_input_token_cost_above_12hr_flex": 1.25e-07, "cache_creation_input_token_cost_above_12hr_priority": 5e-07, "cache_creation_input_token_cost_decisions": 1.25e-07, @@ -37186,8 +37194,10 @@ "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_12hr": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 1e-05, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 4e-05, + "cache_creation_input_token_cost_above_12hr_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_priority": 2e-05, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, @@ -37268,8 +37278,10 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -37349,8 +37361,10 @@ "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_12hr": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-06, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-07, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-06, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-07, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-07, "cache_creation_input_token_cost_above_12hr_priority": 1e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, @@ -75046,10 +75060,14 @@ "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_batches": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_flex": 5e-06, "cache_creation_input_token_cost_above_12hr_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens_ultrafast": 6e-05, + "cache_creation_input_token_cost_above_12hr_batches": 2.5e-06, "cache_creation_input_token_cost_above_12hr_flex": 2.5e-06, "cache_creation_input_token_cost_above_12hr_priority": 1e-05, + "cache_creation_input_token_cost_above_12hr_ultrafast": 3e-05, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_creation_input_token_cost_above_272k_tokens_batches": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06,