fix(openrouter): preserve token counts from streaming usage chunks (#21011)

* docs: add reference to example_openai_endpoint repo for self-hosting fake OpenAI proxy (#21006)

- Updated benchmarks.md with a section on setting up fake OpenAI endpoints
- Updated load_test.md to mention the self-hosted option
- Updated load_test_advanced.md with a tip box about the example repo

Reference: https://github.com/BerriAI/example_openai_endpoint

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

* MCP fixes

* fix(oldteams.tsx): show policies when creating

* fix(proxy/_types.py): ensure mcp rest endpoints can be called by virtual key

ensures UI works with virtual key testing mcp endpoints

* refactor: migrate get object permissions table logic to happen in user api key auth - allows functions to trust user api key object they receive has what they need

* fix(rest_endpoints.py): filter for allowed tools based on what key has access to

* fix(mcp_server_manager.py): ensure only allowed MCP's are returned to the user, via rest endpoints

* Guardrails - add toxic/abusive content filter guardrails

* fix(streaming): preserve usage data from post-finish_reason chunks in OpenAI-compatible streaming

Fixes #16112

OpenRouter and other OpenAI-compatible providers send a usage chunk after
the finish_reason='stop' chunk when stream_options.include_usage is True.
The OpenAIChatCompletionStreamingHandler.chunk_parser() was not passing
the usage field to ModelResponseStream, causing real token counts from the
provider to be lost and falling back to inaccurate estimates.

* fix: resolve merge conflict in test file

- Fix typo in test method name (extra space)
- Move test_prompt_cache_key_in_optional_params to its own class

---------

Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>

litellm/llms/openai/chat/gpt_transformation.py

@@ -772,12 +772,15 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig):
 class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator):
     def chunk_parser(self, chunk: dict) -> ModelResponseStream:
         try:
-            return ModelResponseStream(
-                id=chunk["id"],
-                object="chat.completion.chunk",
-                created=chunk.get("created"),
-                model=chunk.get("model"),
-                choices=chunk.get("choices", []),
-            )
+            kwargs = {
+                "id": chunk["id"],
+                "object": "chat.completion.chunk",
+                "created": chunk.get("created"),
+                "model": chunk.get("model"),
+                "choices": chunk.get("choices", []),
+            }
+            if "usage" in chunk and chunk["usage"] is not None:
+                kwargs["usage"] = chunk["usage"]
+            return ModelResponseStream(**kwargs)
         except Exception as e:
             raise e

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -71,9 +71,7 @@ try:
     from mcp.shared.tool_name_validation import (
         validate_tool_name,  # pyright: ignore[reportAssignmentType]
     )
-    from mcp.shared.tool_name_validation import (
-        SEP_986_URL,
-    )
+    from mcp.shared.tool_name_validation import SEP_986_URL
 except ImportError:
     from pydantic import BaseModel
 

tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py

@@ -9,7 +9,10 @@ import pytest
 
 sys.path.insert(0, os.path.abspath("../../../../.."))
 
-from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
+from litellm.llms.openai.chat.gpt_transformation import (
+    OpenAIGPTConfig,
+    OpenAIChatCompletionStreamingHandler,
+)
 
 
 class TestOpenAIGPTConfig:
@@ -136,6 +139,75 @@ class TestGetOptionalParamsIntegration:
         assert regular_params.get("user") == "my-end-user"
         assert responses_params.get("user") == "my-end-user"
 
+
+class TestOpenAIChatCompletionStreamingHandler:
+    """Tests for OpenAIChatCompletionStreamingHandler.chunk_parser()"""
+
+    def test_chunk_parser_preserves_usage(self):
+        """
+        Test that chunk_parser preserves the usage field from streaming chunks.
+
+        """
+        handler = OpenAIChatCompletionStreamingHandler(
+            streaming_response=None, sync_stream=True
+        )
+
+        usage_chunk = {
+            "id": "gen-123",
+            "created": 1234567890,
+            "model": "openai/gpt-4o-mini",
+            "object": "chat.completion.chunk",
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"role": "assistant", "content": ""},
+                    "finish_reason": None,
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 13797,
+                "completion_tokens": 350,
+                "total_tokens": 14147,
+            },
+        }
+
+        result = handler.chunk_parser(usage_chunk)
+
+        assert result.usage is not None
+        assert result.usage.prompt_tokens == 13797
+        assert result.usage.completion_tokens == 350
+        assert result.usage.total_tokens == 14147
+
+    def test_chunk_parser_without_usage(self):
+        """Test that chunk_parser works normally for chunks without usage."""
+        handler = OpenAIChatCompletionStreamingHandler(
+            streaming_response=None, sync_stream=True
+        )
+
+        chunk = {
+            "id": "gen-123",
+            "created": 1234567890,
+            "model": "openai/gpt-4o-mini",
+            "object": "chat.completion.chunk",
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"role": "assistant", "content": "Hello"},
+                    "finish_reason": None,
+                }
+            ],
+        }
+
+        result = handler.chunk_parser(chunk)
+
+        assert result.id == "gen-123"
+        assert result.choices[0].delta.content == "Hello"
+        assert not hasattr(result, "usage") or result.usage is None
+
+
+class TestPromptCacheKeyIntegration:
+    """Tests for prompt_cache_key support"""
+
     def test_prompt_cache_key_in_optional_params(self):
         """Test that 'prompt_cache_key' flows through get_optional_params for OpenAI models."""
         from litellm.utils import get_optional_params