From 5a6ccc589281b4302791cc301e258adcbcc5fead Mon Sep 17 00:00:00 2001 From: yzong-rh Date: Wed, 16 Sep 2026 17:52:34 -0400 Subject: [PATCH] [Bugfix] Add Responses cache_write_tokens for API compat and CC parity (#57222) Signed-off-by: Yifan Zong Co-authored-by: Flora Feng <4florafeng@gmail.com> --- .../entrypoints/openai/responses/test_serving_responses.py | 7 +++++++ tests/entrypoints/unit_tests/test_context.py | 5 +++++ vllm/entrypoints/openai/responses/context.py | 3 +++ vllm/entrypoints/openai/responses/protocol.py | 1 + vllm/entrypoints/openai/responses/serving.py | 1 + 5 files changed, 17 insertions(+) diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index a445774fcc..0f361c1c6a 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -18,6 +18,9 @@ from openai.types.responses import ( ResponseTextConfig, ResponseTextDeltaEvent, ) +from openai.types.responses import ( + ResponseUsage as OpenAIResponseUsage, +) from openai.types.responses.response_format_text_json_schema_config import ( ResponseFormatTextJSONSchemaConfig, ) @@ -975,6 +978,7 @@ async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch): outputs=[completion], finished=True, num_cached_tokens=0, + num_cache_creation_tokens=3, ) context.append_output(req_output) @@ -995,6 +999,9 @@ async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch): ) assert response.usage.output_tokens_details.reasoning_tokens == 1 + assert response.usage.input_tokens_details.cache_write_tokens == 3 + + OpenAIResponseUsage.model_validate(response.usage.model_dump(mode="json")) class TestExtractAllowedToolsFromMcpRequests: diff --git a/tests/entrypoints/unit_tests/test_context.py b/tests/entrypoints/unit_tests/test_context.py index 8c3d1467ed..a5b1afe8c5 100644 --- a/tests/entrypoints/unit_tests/test_context.py +++ b/tests/entrypoints/unit_tests/test_context.py @@ -682,6 +682,7 @@ def create_simple_context_output( prompt="Test prompt", prompt_token_ids=None, num_cached_tokens=0, + num_cache_creation_tokens=None, logprobs=None, finished=True, ): @@ -707,6 +708,7 @@ def create_simple_context_output( ], finished=finished, num_cached_tokens=num_cached_tokens, + num_cache_creation_tokens=num_cache_creation_tokens, ) @@ -827,6 +829,7 @@ def test_simple_context_token_counting(): token_ids=[10, 11], prompt_token_ids=[1, 2, 3, 4, 5], num_cached_tokens=2, + num_cache_creation_tokens=3, ) ) context.append_output( @@ -835,12 +838,14 @@ def test_simple_context_token_counting(): token_ids=[12], prompt_token_ids=[1, 2, 3, 4, 5], num_cached_tokens=2, + num_cache_creation_tokens=3, ) ) assert context.num_prompt_tokens == 5 assert context.num_output_tokens == 3 # 2 + 1 assert context.num_cached_tokens == 2 + assert context.num_cache_creation_tokens == 3 def test_simple_context_final_output(): diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 077f78c02b..3414bd6d82 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -189,6 +189,7 @@ class SimpleContext(ConversationContext): self.num_prompt_tokens = 0 self.num_output_tokens = 0 self.num_cached_tokens = 0 + self.num_cache_creation_tokens = 0 # todo num_reasoning_tokens is not implemented yet. self.num_reasoning_tokens = 0 # not implemented yet for SimpleContext @@ -204,6 +205,8 @@ class SimpleContext(ConversationContext): raise ValueError("SimpleContext only supports RequestOutput.") self.num_prompt_tokens = len(output.prompt_token_ids or []) self.num_cached_tokens = output.num_cached_tokens or 0 + if output.num_cache_creation_tokens is not None: + self.num_cache_creation_tokens = output.num_cache_creation_tokens self.num_output_tokens += len(output.outputs[0].token_ids or []) if output.kv_transfer_params is not None: self.kv_transfer_params = output.kv_transfer_params diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 201e46ed04..ec28d2823b 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -80,6 +80,7 @@ _INT64_MAX = 2**63 - 1 class InputTokensDetails(OpenAIBaseModel): cached_tokens: int + cache_write_tokens: int input_tokens_per_turn: list[int] = Field(default_factory=list) cached_tokens_per_turn: list[int] = Field(default_factory=list) diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 2c61205e22..bce54338c4 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -837,6 +837,7 @@ class OpenAIServingResponses(GenerateBaseServing): output_tokens=num_generated_tokens, total_tokens=num_prompt_tokens + num_generated_tokens, input_tokens_details=InputTokensDetails( + cache_write_tokens=getattr(context, "num_cache_creation_tokens", 0), cached_tokens=num_cached_tokens, input_tokens_per_turn=[ turn.input_tokens for turn in context.all_turn_metrics