[Bugfix][Responses API] Complete cache-write token accounting (#60973)

Signed-off-by: Kay Yan <kay.yan@daocloud.io>
Co-authored-by: Codex <noreply@openai.com>

tests/entrypoints/openai/responses/test_parsable_context_unit.py

@@ -169,6 +169,7 @@ def _make_request_output(
     text: str = "Hello, world!",
     token_ids: Sequence[int] = (1, 2, 3),
     finish_reason: str = "stop",
+    num_cache_creation_tokens: int | None = None,
 ) -> RequestOutput:
     return RequestOutput(
         request_id="test",
@@ -186,6 +187,7 @@ def _make_request_output(
             )
         ],
         finished=True,
+        num_cache_creation_tokens=num_cache_creation_tokens,
     )
 
 
@@ -425,6 +427,14 @@ def test_reasoning_tokens_counted_per_round():
     assert ctx.num_reasoning_tokens == 3
 
 
+def test_cache_creation_tokens_counted_per_round():
+    ctx = _make_context(None)
+
+    for count, expected in [(None, 0), (0, 0), (3, 3), (4, 7)]:
+        ctx.append_output(_make_request_output(num_cache_creation_tokens=count))
+        assert ctx.num_cache_creation_tokens == expected
+
+
 def test_reasoning_tokens_zero_without_parser():
     ctx = _make_context(None)
 

tests/entrypoints/unit_tests/test_context.py

@@ -20,6 +20,7 @@ def create_mock_request_output(
     output_token_ids=None,
     num_cached_tokens=0,
     finished=True,
+    num_cache_creation_tokens=None,
 ):
     """Helper function to create a mock RequestOutput object for testing."""
     outputs = []
@@ -44,6 +45,7 @@ def create_mock_request_output(
         outputs=outputs,
         finished=finished,
         num_cached_tokens=num_cached_tokens,
+        num_cache_creation_tokens=num_cache_creation_tokens,
     )
 
 
@@ -117,7 +119,8 @@ def make_harmony_context(
     return context, fake_parser
 
 
-def test_single_turn_token_counting():
+@pytest.mark.parametrize("num_cache_creation_tokens", [None, 0, 3])
+def test_single_turn_token_counting(num_cache_creation_tokens):
     """Test token counting behavior for a single turn."""
     # Create a context
     context, _ = make_harmony_context()
@@ -127,6 +130,7 @@ def test_single_turn_token_counting():
         prompt_token_ids=[1, 2, 3, 4, 5],  # 5 prompt tokens
         output_token_ids=[6, 7, 8],  # 3 output tokens
         num_cached_tokens=2,  # 2 cached tokens
+        num_cache_creation_tokens=num_cache_creation_tokens,
     )
 
     # Append the output to the context
@@ -136,6 +140,7 @@ def test_single_turn_token_counting():
     assert context.num_prompt_tokens == 5
     assert context.num_output_tokens == 3
     assert context.num_cached_tokens == 2
+    assert context.num_cache_creation_tokens == (num_cache_creation_tokens or 0)
     assert context.num_tool_output_tokens == 0  # No tool tokens in first turn
 
     # Verify internal state tracking
@@ -398,6 +403,7 @@ async def test_streaming_multi_turn_token_counting():
     num_prompt_tokens = [3, 8, 13]
     num_output_tokens = [3, 3, 2]
     num_cached_tokens = [0, 3, 8]
+    num_cache_creation_tokens = [3, 4, 2]
 
     # Simulate three turns of conversation:
     # Turn 1: stream tokens one by one, then finish the message
@@ -411,6 +417,7 @@ async def test_streaming_multi_turn_token_counting():
             prompt_token_ids=[1, 2, 3],  # 3 prompt tokens
             output_token_ids=[101],  # Single token
             num_cached_tokens=num_cached_tokens[0],
+            num_cache_creation_tokens=num_cache_creation_tokens[0],
             finished=False,  # Not end of message yet
         )
     )
@@ -419,6 +426,7 @@ async def test_streaming_multi_turn_token_counting():
     context.append_output(
         create_mock_request_output(
             output_token_ids=[102],
+            num_cache_creation_tokens=num_cache_creation_tokens[0],
             finished=False,
         )
     )
@@ -427,6 +435,7 @@ async def test_streaming_multi_turn_token_counting():
     context.append_output(
         create_mock_request_output(
             output_token_ids=[103],
+            num_cache_creation_tokens=num_cache_creation_tokens[0],
             finished=True,  # End of message
         )
     )
@@ -435,6 +444,7 @@ async def test_streaming_multi_turn_token_counting():
     assert context.num_prompt_tokens == 3  # Initial prompt tokens
     assert context.num_output_tokens == 3  # Three output tokens
     assert context.num_cached_tokens == 0
+    assert context.num_cache_creation_tokens == 3
     assert context.num_tool_output_tokens == 0  # No tool output in first turn
     assert context.first_tok_of_message is True  # Ready for next message
 
@@ -454,6 +464,7 @@ async def test_streaming_multi_turn_token_counting():
             ],  # 8 tokens (includes previous)
             output_token_ids=[201],
             num_cached_tokens=num_cached_tokens[1],  # Some tokens cached
+            num_cache_creation_tokens=num_cache_creation_tokens[1],
             finished=False,
         )
     )
@@ -463,6 +474,7 @@ async def test_streaming_multi_turn_token_counting():
     context.append_output(
         create_mock_request_output(
             output_token_ids=[202],
+            num_cache_creation_tokens=num_cache_creation_tokens[1],
             finished=False,
         )
     )
@@ -471,6 +483,7 @@ async def test_streaming_multi_turn_token_counting():
     context.append_output(
         create_mock_request_output(
             output_token_ids=[203],
+            num_cache_creation_tokens=num_cache_creation_tokens[1],
             finished=True,  # End of reasoning message
         )
     )
@@ -480,6 +493,7 @@ async def test_streaming_multi_turn_token_counting():
     assert context.num_output_tokens == 3 + 3  # First turn + second turn
     assert context.num_reasoning_tokens == 3  # All tokens in analysis channel
     assert context.num_cached_tokens == 3  # Cached tokens from second turn
+    assert context.num_cache_creation_tokens == 3 + 4
 
     # Formula: this turn prompt tokens - last turn prompt - last turn output
     expected_tool_tokens = 8 - 3 - 3  # = 2
@@ -505,6 +519,7 @@ async def test_streaming_multi_turn_token_counting():
             ],  # 13 tokens
             output_token_ids=[301],
             num_cached_tokens=num_cached_tokens[2],  # More cached tokens
+            num_cache_creation_tokens=num_cache_creation_tokens[2],
             finished=False,
         )
     )
@@ -512,6 +527,7 @@ async def test_streaming_multi_turn_token_counting():
     context.append_output(
         create_mock_request_output(
             output_token_ids=[302],
+            num_cache_creation_tokens=num_cache_creation_tokens[2],
             finished=True,
         )
     )
@@ -523,6 +539,7 @@ async def test_streaming_multi_turn_token_counting():
     assert context.num_cached_tokens == sum(
         num_cached_tokens
     )  # Accumulated cached tokens
+    assert context.num_cache_creation_tokens == sum(num_cache_creation_tokens)
 
     # Additional tool tokens from third turn
     # Formula: this turn prompt - last turn prompt - last turn output

vllm/entrypoints/openai/responses/context.py

@@ -315,6 +315,7 @@ class ParsableContext(ConversationContext):
         self.num_prompt_tokens = 0
         self.num_output_tokens = 0
         self.num_cached_tokens = 0
+        self.num_cache_creation_tokens = 0
         self.num_reasoning_tokens = 0
         # not implemented yet for ParsableContext
         self.all_turn_metrics: list[TurnMetrics] = []
@@ -347,6 +348,7 @@ class ParsableContext(ConversationContext):
     def append_output(self, output: RequestOutput) -> None:
         self.num_prompt_tokens = len(output.prompt_token_ids or [])
         self.num_cached_tokens = output.num_cached_tokens or 0
+        self.num_cache_creation_tokens += output.num_cache_creation_tokens or 0
         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
@@ -651,6 +653,7 @@ class HarmonyContext(ConversationContext):
         self.num_prompt_tokens = 0
         self.num_output_tokens = 0
         self.num_cached_tokens = 0
+        self.num_cache_creation_tokens = 0
         self.num_reasoning_tokens = 0
         self.num_tool_output_tokens = 0
 
@@ -773,6 +776,8 @@ class HarmonyContext(ConversationContext):
             self.num_cached_tokens += num_cached_token
             self.current_turn_metrics.cached_input_tokens = num_cached_token
 
+        self.num_cache_creation_tokens += output.num_cache_creation_tokens or 0
+
     def _update_decode_token_usage(self, output: RequestOutput) -> int:
         """Update token usage statistics for the decode phase of generation.
 

vllm/entrypoints/openai/responses/serving.py

@@ -881,7 +881,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),
+                cache_write_tokens=context.num_cache_creation_tokens,
                 cached_tokens=num_cached_tokens,
                 input_tokens_per_turn=[
                     turn.input_tokens for turn in context.all_turn_metrics