fix(rate_limiter): attach v3 priority rate limit headers on /v1/messages (#37228)

* fix(rate_limiter): attach v3 priority rate limit headers on /v1/messages

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(playground): honor the Stream responses toggle for /v1/messages

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(rate_limiter): drop explanatory docstrings from v3 dict response tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/hooks/dynamic_rate_limiter_v3.py

@@ -32,6 +32,10 @@ from litellm.proxy.hooks.rate_limiter_utils import (
     resolve_llm_provider_for_rate_limit,
 )
 from litellm.proxy.utils import InternalUsageCache
+from litellm.router_utils.add_retry_fallback_headers import (
+    ensure_response_additional_headers,
+    response_has_hidden_params,
+)
 from litellm.types.router import ModelGroupInfo
 from litellm.types.utils import CallTypesLiteral
 
@@ -659,22 +663,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
                 data=data, user_api_key_dict=user_api_key_dict, response=response
             )
 
-            # Add additional priority-specific headers
-            if isinstance(response, ModelResponse):
+            if response_has_hidden_params(response):
                 priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)
-
-                # Get existing additional headers
-                additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {}
-
-                # Add priority information
+                additional_headers: Final = ensure_response_additional_headers(response)
                 additional_headers["x-litellm-priority"] = priority or "default"
                 additional_headers["x-litellm-rate-limiter-version"] = "v3"
 
-                # Update response
-                if not hasattr(response, "_hidden_params"):
-                    response._hidden_params = {}
-                response._hidden_params["additional_headers"] = additional_headers
-
             return response
 
         except Exception as e:

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -52,6 +52,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
     canonical_provider_batch_id,
 )
 from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
+from litellm.router_utils.add_retry_fallback_headers import (
+    ensure_response_additional_headers,
+    response_has_hidden_params,
+)
 from litellm.types.caching import RedisPipelineIncrementOperation
 from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
 from litellm.types.utils import (
@@ -4677,34 +4681,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         Post-call hook to update rate limit headers in the response.
         """
         try:
-            from pydantic import BaseModel
-
             stash: Final = get_request_stash()
             litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None
 
-            if litellm_proxy_rate_limit_response is not None:
-                # Update response headers
-                if hasattr(response, "_hidden_params"):
-                    _hidden_params = getattr(response, "_hidden_params")
-                else:
-                    _hidden_params = None
-
-                if _hidden_params is not None and (
-                    isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict)
-                ):
-                    if isinstance(_hidden_params, BaseModel):
-                        _hidden_params = _hidden_params.model_dump()
-
-                    _additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers(
-                        additional_headers=_hidden_params.get("additional_headers", {}) or {},
+            if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response):
+                additional_headers: Final = ensure_response_additional_headers(response)
+                additional_headers.update(
+                    self._merge_ratelimit_statuses_into_additional_headers(
+                        additional_headers={},
                         statuses=litellm_proxy_rate_limit_response["statuses"],
                     )
-
-                    setattr(
-                        response,
-                        "_hidden_params",
-                        {**_hidden_params, "additional_headers": _additional_headers},
-                    )
+                )
 
         except Exception as e:
             verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e)

litellm/router_utils/add_retry_fallback_headers.py

@@ -50,6 +50,12 @@ def prepare_response_for_header_attachment(response: object) -> object | None:
     return response
 
 
+def response_has_hidden_params(response: object) -> bool:
+    if isinstance(response, dict):
+        return "_hidden_params" in response
+    return hasattr(response, "_hidden_params")
+
+
 def ensure_response_additional_headers(response: object) -> dict[str, object]:
     hidden_params: Final = get_hidden_params_dict(response, create=isinstance(response, dict))
     _write_hidden_params(response, hidden_params)

tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py

@@ -1861,3 +1861,60 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch):
         )
     assert capacity_blocked.value.status_code == 429
     assert "Model capacity reached" in capacity_blocked.value.detail["error"]
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_hook_attaches_priority_headers_to_dict_response():
+    from litellm.proxy.hooks.parallel_request_limiter_v3 import (
+        RateLimitResponse,
+        RateLimitStatus,
+        get_or_create_request_stash,
+    )
+
+    handler = DynamicRateLimitHandler(internal_usage_cache=DualCache())
+    get_or_create_request_stash().rate_limit_response = RateLimitResponse(
+        overall_code="OK",
+        statuses=[
+            RateLimitStatus(
+                code="OK",
+                current_limit=75,
+                limit_remaining=74,
+                rate_limit_type="requests",
+                descriptor_key="priority_model",
+            )
+        ],
+    )
+    response = {
+        "id": "msg_123",
+        "type": "message",
+        "role": "assistant",
+        "content": [],
+        "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}},
+    }
+
+    await handler.async_post_call_success_hook(
+        data={"model": "anthropic-haiku"},
+        user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}),
+        response=response,
+    )
+
+    additional_headers = response["_hidden_params"]["additional_headers"]
+    assert additional_headers["x-litellm-attempted-retries"] == 0
+    assert additional_headers["x-ratelimit-priority_model-limit-requests"] == 75
+    assert additional_headers["x-ratelimit-priority_model-remaining-requests"] == 74
+    assert additional_headers["x-litellm-priority"] == "premium"
+    assert additional_headers["x-litellm-rate-limiter-version"] == "v3"
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_hook_leaves_raw_provider_dict_untouched():
+    handler = DynamicRateLimitHandler(internal_usage_cache=DualCache())
+    response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []}
+
+    await handler.async_post_call_success_hook(
+        data={"model": "anthropic-haiku"},
+        user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}),
+        response=response,
+    )
+
+    assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []}

tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py

@@ -6171,3 +6171,68 @@ async def test_success_hook_leaves_stash_untouched_for_non_batch_responses():
         data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5))
     )
     assert get_request_stash().batch_enqueued_reservation == reservation
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_hook_attaches_ratelimit_headers_to_dict_response():
+    from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus
+
+    handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
+    get_or_create_request_stash().rate_limit_response = RateLimitResponse(
+        overall_code="OK",
+        statuses=[
+            RateLimitStatus(
+                code="OK",
+                current_limit=100,
+                limit_remaining=99,
+                rate_limit_type="requests",
+                descriptor_key="model_saturation_check",
+            )
+        ],
+    )
+    response = {
+        "id": "msg_123",
+        "type": "message",
+        "role": "assistant",
+        "content": [],
+        "_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}},
+    }
+
+    await handler.async_post_call_success_hook(
+        data={"model": "anthropic-haiku"},
+        user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-dict-response")),
+        response=response,
+    )
+
+    additional_headers = response["_hidden_params"]["additional_headers"]
+    assert additional_headers["x-litellm-attempted-retries"] == 0
+    assert additional_headers["x-ratelimit-model_saturation_check-limit-requests"] == 100
+    assert additional_headers["x-ratelimit-model_saturation_check-remaining-requests"] == 99
+
+
+@pytest.mark.asyncio
+async def test_post_call_success_hook_leaves_raw_provider_dict_untouched():
+    from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus
+
+    handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
+    get_or_create_request_stash().rate_limit_response = RateLimitResponse(
+        overall_code="OK",
+        statuses=[
+            RateLimitStatus(
+                code="OK",
+                current_limit=100,
+                limit_remaining=99,
+                rate_limit_type="requests",
+                descriptor_key="model_saturation_check",
+            )
+        ],
+    )
+    response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []}
+
+    await handler.async_post_call_success_hook(
+        data={"model": "anthropic-haiku"},
+        user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-raw-dict")),
+        response=response,
+    )
+
+    assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []}

ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx

@@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import ChatUI from "./ChatUI";
 import * as fetchModelsModule from "@/components/llm_calls/fetch_models";
 import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion";
+import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages";
 
 vi.mock("@/components/llm_calls/fetch_models", () => ({
   fetchAvailableModels: vi.fn(),
@@ -14,6 +15,10 @@ vi.mock("@/components/llm_calls/chat_completion", () => ({
   makeOpenAIChatCompletionRequest: vi.fn().mockResolvedValue(undefined),
 }));
 
+vi.mock("../../llm_calls/anthropic_messages", () => ({
+  makeAnthropicMessagesRequest: vi.fn().mockResolvedValue(undefined),
+}));
+
 vi.mock("@/components/networking", () => ({
   tagListCall: vi.fn().mockResolvedValue({}),
   vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }),
@@ -32,6 +37,8 @@ beforeEach(() => {
 
 const CHAT_REQUEST_ARG_COUNT = 26;
 const STREAMING_ENABLED_ARG_INDEX = 25;
+const MESSAGES_REQUEST_ARG_COUNT = 19;
+const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18;
 
 async function openComboboxByPlaceholder(placeholder: string) {
   const user = userEvent.setup();
@@ -378,6 +385,52 @@ describe("ChatUI", () => {
     expect(requestArgs[STREAMING_ENABLED_ARG_INDEX]).toBe(false);
   });
 
+  it("should send the /v1/messages request non-streaming after Stream responses is unchecked", async () => {
+    const user = userEvent.setup();
+    render(
+      <ChatUI
+        accessToken="1234567890"
+        token="1234567890"
+        userRole="user"
+        userID="1234567890"
+        disabledPersonalKeyCreation={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Test Key")).toBeInTheDocument();
+    });
+
+    await selectComboboxOption("Select an endpoint", "/v1/messages");
+    await selectComboboxOption("Select a Model", "Model 1");
+
+    await user.click(await screen.findByTestId("model-settings-button"));
+
+    const streamingCheckbox = await screen.findByRole("checkbox", { name: /Stream responses/i });
+    expect(streamingCheckbox).toBeChecked();
+    await user.click(streamingCheckbox);
+
+    await waitFor(() => {
+      expect(screen.getByRole("checkbox", { name: /Stream responses/i })).not.toBeChecked();
+    });
+
+    const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)");
+    await act(async () => {
+      fireEvent.change(messageInput, { target: { value: "hello" } });
+    });
+    await act(async () => {
+      fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
+    });
+
+    await waitFor(() => {
+      expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1);
+    });
+
+    const requestArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0];
+    expect(requestArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT);
+    expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false);
+  });
+
   it("should force streaming in simplified mode even when the playground setting is off", async () => {
     sessionStorage.setItem("streamingEnabled", "false");
 

ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx

@@ -1025,6 +1025,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
             mcpServers,
             mcpServerToolRestrictions,
             mcpToolsets,
+            streamingEnabled,
           );
         } else if (endpointType === EndpointType.EMBEDDINGS) {
           await makeOpenAIEmbeddingsRequest(
@@ -1174,7 +1175,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
     return !model.mode || model.mode === "chat";
   };
 
-  const supportsStreamingToggle = endpointType === EndpointType.CHAT || endpointType === EndpointType.RESPONSES;
+  const supportsStreamingToggle =
+    endpointType === EndpointType.CHAT ||
+    endpointType === EndpointType.RESPONSES ||
+    endpointType === EndpointType.ANTHROPIC_MESSAGES;
   const modelsForEndpoint = useMemo(
     () => filterModelsForEndpoint(modelInfo, endpointType as EndpointType),
     [modelInfo, endpointType],

ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx

@@ -7,13 +7,27 @@ vi.mock("@/components/networking", () => ({
 }));
 
 const mockMessagesStream = vi.fn();
+const mockMessagesCreate = vi.fn();
 
 vi.mock("@anthropic-ai/sdk", () => ({
   default: vi.fn(function () {
-    return { messages: { stream: mockMessagesStream } };
+    return { messages: { stream: mockMessagesStream, create: mockMessagesCreate } };
   }),
 }));
 
+const NON_STREAMING_ARGS = [
+  undefined, // traceId
+  undefined, // vector_store_ids
+  undefined, // guardrails
+  undefined, // policies
+  undefined, // selectedMCPServers
+  undefined, // customBaseUrl
+  undefined, // mcpServers
+  undefined, // mcpServerToolRestrictions
+  undefined, // mcpToolsets
+  false, // streamingEnabled
+] as const;
+
 describe("anthropic_messages prompt cache usage", () => {
   const captureUsage = async (usage: Record<string, unknown>): Promise<TokenUsage> => {
     async function* mockStream() {
@@ -59,3 +73,53 @@ describe("anthropic_messages prompt cache usage", () => {
     expect(usageData.promptTokens).toBe(5000);
   });
 });
+
+describe("anthropic_messages non-streaming", () => {
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("sends stream:false through messages.create and renders the full reply at once", async () => {
+    mockMessagesCreate.mockResolvedValue({
+      content: [
+        { type: "thinking", thinking: "considering" },
+        { type: "text", text: "OK" },
+      ],
+      usage: { input_tokens: 12, output_tokens: 3, cache_read_input_tokens: 7 },
+    });
+    const updateTextUI = vi.fn();
+    const onReasoningContent = vi.fn();
+    const onUsageData = vi.fn();
+
+    await makeAnthropicMessagesRequest(
+      [{ role: "user", content: "Hello" }],
+      updateTextUI,
+      "claude-haiku-4-5",
+      "test-token",
+      undefined,
+      undefined,
+      onReasoningContent,
+      undefined,
+      onUsageData,
+      ...NON_STREAMING_ARGS,
+    );
+
+    expect(mockMessagesStream).not.toHaveBeenCalled();
+    expect(mockMessagesCreate).toHaveBeenCalledTimes(1);
+    expect(mockMessagesCreate.mock.calls[0][0]).toMatchObject({ model: "claude-haiku-4-5", stream: false });
+    expect(updateTextUI).toHaveBeenCalledWith("assistant", "OK", "claude-haiku-4-5");
+    expect(onReasoningContent).toHaveBeenCalledWith("considering");
+    const expectedUsage: TokenUsage = { completionTokens: 3, promptTokens: 12, totalTokens: 15, cacheReadTokens: 7 };
+    expect(onUsageData).toHaveBeenCalledWith(expectedUsage);
+  });
+
+  it("keeps streaming as the default when the flag is omitted", async () => {
+    async function* emptyStream() {}
+    mockMessagesStream.mockReturnValue(emptyStream());
+
+    await makeAnthropicMessagesRequest([{ role: "user", content: "Hello" }], vi.fn(), "claude-haiku-4-5", "test-token");
+
+    expect(mockMessagesCreate).not.toHaveBeenCalled();
+    expect(mockMessagesStream.mock.calls[0][0]).toMatchObject({ stream: true });
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx

@@ -7,6 +7,13 @@ import { getProxyBaseUrl } from "@/components/networking";
 import { toast } from "@/lib/toast";
 import { extractPromptCacheTokens } from "@/utils/promptCacheUsage";
 
+const toTokenUsage = (usage: Anthropic.Usage): TokenUsage => ({
+  completionTokens: usage.output_tokens,
+  promptTokens: usage.input_tokens,
+  totalTokens: usage.input_tokens + usage.output_tokens,
+  ...extractPromptCacheTokens(usage),
+});
+
 export async function makeAnthropicMessagesRequest(
   messages: MessageType[],
   updateTextUI: (role: string, delta: string, model?: string) => void,
@@ -26,6 +33,7 @@ export async function makeAnthropicMessagesRequest(
   mcpServers?: MCPServer[],
   mcpServerToolRestrictions?: Record<string, string[]>,
   mcpToolsets?: MCPToolset[],
+  streamingEnabled: boolean = true,
 ) {
   if (!accessToken) {
     throw new Error("Virtual Key is required");
@@ -58,7 +66,7 @@ export async function makeAnthropicMessagesRequest(
     const requestBody: any = {
       model: selectedModel,
       messages: messages.map((m) => ({ role: m.role, content: m.content })),
-      stream: true,
+      stream: streamingEnabled,
       max_tokens: 1024,
       // @ts-ignore - litellm specific parameter
       litellm_trace_id: traceId,
@@ -74,6 +82,20 @@ export async function makeAnthropicMessagesRequest(
     if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids;
     if (guardrails) requestBody.guardrails = guardrails;
     if (policies) requestBody.policies = policies;
+
+    if (!streamingEnabled) {
+      const message: Anthropic.Message = await client.messages.create({ ...requestBody, stream: false }, { signal });
+      for (const block of message.content) {
+        if (block.type === "text") {
+          updateTextUI("assistant", block.text, selectedModel);
+        } else if (block.type === "thinking" && onReasoningContent) {
+          onReasoningContent(block.thinking);
+        }
+      }
+      onUsageData?.(toTokenUsage(message.usage));
+      return;
+    }
+
     // Use the streaming helper method for cleaner async iteration
     // @ts-ignore - The SDK types might not include all litellm-specific parameters
     const stream = client.messages.stream(requestBody, { signal });
@@ -105,14 +127,7 @@ export async function makeAnthropicMessagesRequest(
 
       // Process usage data from message_delta events
       if (messageStreamEvent.type === "message_delta" && (messageStreamEvent as any).usage && onUsageData) {
-        const usage = (messageStreamEvent as any).usage;
-        const usageData: TokenUsage = {
-          completionTokens: usage.output_tokens,
-          promptTokens: usage.input_tokens,
-          totalTokens: usage.input_tokens + usage.output_tokens,
-          ...extractPromptCacheTokens(usage),
-        };
-        onUsageData(usageData);
+        onUsageData(toTokenUsage((messageStreamEvent as any).usage));
       }
     }
   } catch (error) {