From 75d0d2bd7ad091fcd81e2ee7cc173e9918d12a33 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Wed, 11 Feb 2026 23:57:28 -0300 Subject: [PATCH] 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 * 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 Co-authored-by: Cursor Agent --- .../llms/openai/chat/gpt_transformation.py | 17 +++-- .../mcp_server/mcp_server_manager.py | 4 +- .../chat/test_openai_gpt_transformation.py | 74 ++++++++++++++++++- 3 files changed, 84 insertions(+), 11 deletions(-) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 1636890707..5b9840d95b 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/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 diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 62a06468e9..e003841ac0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/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 diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index c0695bf358..e6ab199168 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/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