Skip to content

Commit ff76250

Browse files
committed
fix(ccr): make continuation-usage publication additive for batch paths
set_pending_ccr_continuation_usage now appends to a list instead of replacing, so multiple handle_response calls in a batch (e.g. BatchResultProcessor.process_results) accumulate rather than overwrite. consume sums the list and returns a single 4-tuple, preserving the existing consumer interface. Fixes the CHANGES_REQUESTED from @JerrettDavis: with two CCR-bearing batch items, item 2 was overwriting item 1's continuation overhead, so cost_tracker recorded only the last item's dropped rounds. Added test_batch_two_ccr_items_both_reach_cost_view: two set_pending calls + one emit, asserting both items' overhead reaches cost_tracker.record_tokens (sum), while record_request stays on the base outcome (cost-only fold, unchanged). 14/14 pass.
1 parent b4f479b commit ff76250

2 files changed

Lines changed: 67 additions & 10 deletions

File tree

headroom/proxy/outcome.py

Lines changed: 22 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -65,24 +65,36 @@ def consume_pending_proactive_retrieval() -> tuple[int, int] | None:
6565
# The payload is cache-split (uncached_input, cache_read, cache_write, output) so
6666
# cost_with_headroom prices continuation cache at the right rate rather than
6767
# folding cached tokens into uncached input at full list price.
68-
_pending_ccr_continuation_usage: ContextVar[tuple[int, int, int, int] | None] = ContextVar(
69-
"headroom_ccr_continuation_pending", default=None
70-
)
68+
_pending_ccr_continuation_usage: ContextVar[
69+
list[tuple[int, int, int, int]] | None
70+
] = ContextVar("headroom_ccr_continuation_pending", default=None)
7171

7272

7373
def set_pending_ccr_continuation_usage(
7474
value: tuple[int, int, int, int] | None,
7575
) -> None:
76-
"""Bind the dropped continuation rounds' (uncached_in, cache_read, cache_write, output) to this request."""
77-
_pending_ccr_continuation_usage.set(value)
76+
"""Append the dropped continuation rounds' (uncached_in, cache_read, cache_write, output) to this request.
77+
78+
Additive: each call appends to a list so batch processing (multiple
79+
handle_response calls before a single emit) accumulates all items.
80+
consume_pending_ccr_continuation_usage sums the list.
81+
"""
82+
if value is None:
83+
return
84+
items = _pending_ccr_continuation_usage.get()
85+
if items is None:
86+
items = []
87+
_pending_ccr_continuation_usage.set(items)
88+
items.append(value)
7889

7990

8091
def consume_pending_ccr_continuation_usage() -> tuple[int, int, int, int] | None:
81-
"""Read and clear the pending continuation usage for this request."""
82-
value = _pending_ccr_continuation_usage.get()
83-
if value is not None:
84-
_pending_ccr_continuation_usage.set(None)
85-
return value
92+
"""Read, sum, and clear all pending continuation usage for this request."""
93+
items = _pending_ccr_continuation_usage.get()
94+
_pending_ccr_continuation_usage.set(None)
95+
if not items:
96+
return None
97+
return tuple(sum(group) for group in zip(*items))
8698

8799

88100
def clear_pending_outcome_side_channels() -> None:

tests/test_ccr_continuation_cost_fold.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -487,3 +487,48 @@ async def main():
487487
assert metrics_rec["uncached_input_tokens"] == U2["input_tokens"]
488488
assert metrics_rec["cache_read_tokens"] == U2["cache_read_input_tokens"]
489489
assert metrics_rec["output_tokens"] == U2["output_tokens"]
490+
491+
492+
# --- batch path: multiple CCR items accumulate, not overwrite ------------
493+
494+
495+
def test_batch_two_ccr_items_both_reach_cost_view():
496+
"""Two CCR-bearing batch items: both items' continuation overhead
497+
reaches cost_tracker.record_tokens, not just the last.
498+
499+
Simulates BatchResultProcessor: multiple handle_response calls
500+
(each set_pending), then one emit_request_outcome (one consume).
501+
"""
502+
handler = _Handler()
503+
handler.cost_tracker = _CostTracker()
504+
base_out, base_uin, base_cr, base_cw = 800, 110_000, 30_000, 3_000
505+
item1 = (205_000, 80_000, 11_000, 1_100)
506+
item2 = (100_000, 50_000, 5_000, 500)
507+
508+
outcome = RequestOutcome(
509+
request_id="req-batch", provider="anthropic", model="m", status_code=200,
510+
original_tokens=base_uin, optimized_tokens=base_uin, output_tokens=base_out,
511+
tokens_saved=0, attempted_input_tokens=base_uin,
512+
uncached_input_tokens=base_uin, cache_read_tokens=base_cr, cache_write_tokens=base_cw,
513+
)
514+
515+
async def main():
516+
consume_pending_ccr_continuation_usage()
517+
set_pending_ccr_continuation_usage(item1)
518+
set_pending_ccr_continuation_usage(item2)
519+
await emit_request_outcome(handler, outcome)
520+
return consume_pending_ccr_continuation_usage()
521+
522+
leftover = asyncio.run(main())
523+
524+
assert handler.cost_tracker.recorded, "cost_tracker.record_tokens was not called"
525+
ct = handler.cost_tracker.recorded[0]
526+
# Both items accumulated into cost_tracker, not just the last
527+
assert ct["uncached_tokens"] == base_uin + item1[0] + item2[0]
528+
assert ct["cache_read_tokens"] == base_cr + item1[1] + item2[1]
529+
assert ct["cache_write_tokens"] == base_cw + item1[2] + item2[2]
530+
assert ct["output_tokens"] == base_out + item1[3] + item2[3]
531+
# record_request stays on base outcome (cost-only fold)
532+
assert handler.metrics.requested[0]["uncached_input_tokens"] == base_uin
533+
# Consume cleared after emit
534+
assert leftover is None

0 commit comments

Comments
 (0)