fix(ccr): scope the continuation-usage fold to the cost view

The fold added the dropped CCR continuation rounds onto the RequestOutcome, but
three consumers of the outcome funnel are not cost surfaces:

- metrics.record_request forwards cache_read_tokens to SavingsTracker, which
  credits it as cache *savings* (cache_savings_usd -> SavingsSnapshot
  .total_savings and `headroom doctor`), so retrieval overhead RAISED the
  reported savings: 80k continuation cache reads on a 1k-read turn moved
  cache_savings_usd from $0.0027 to $0.2187 in a synthetic check.
- the output-shaper counterfactual reads outcome.output_tokens as this turn's
  completion length; folded continuation output biased the per-arm estimate
  (the recorder saw 1180 tokens for an 80-token turn).
- the RequestLog row and the PERF line describe the single round the client
  received.

Move the fold to the cost_tracker.record_tokens call, which is where the billed
input breakdown behind cost_with_headroom_usd lands and where _costs is appended
for budget enforcement. Cost accounting keeps the original fix's intent; the
savings, log and PERF surfaces stay on the final round.

Also:
- clear both pending side channels in the caller's context after the shielded
  emit. asyncio.shield runs the funnel in a child task with a copied context, so
  its consume_* cleared the copy, and a second outcome emitted from the same
  context re-booked the same continuation usage and proactive drawback.
- keep the original_tokens floor on the Responses-path usage recompute. Falling
  to 0 priced that round's input at nothing, because the folded continuation
  tokens keep record_tokens' "no breakdown -> use tokens_sent" fallback from
  firing.
- cover the Responses recompute end-to-end (initial round 50/10, final 60/5) and
  pin each non-cost surface against re-folding.

Author: nangsontay

headroom/proxy/cost.py

@@ -671,15 +671,18 @@ def _entry_number(entry: Any, attr: str) -> int | float:
     summary["compression"]["net_tokens_saved"] = net_tokens_saved
 
     # Attribution view only. The dropped CCR continuation rounds' usage is now
-    # folded into the billed totals at the outcome funnel
-    # (set_pending_ccr_continuation_usage -> uncached_input_tokens /
-    # cache_read_tokens / cache_write_tokens, which cost_with_headroom prices, and
-    # output_tokens for the budget/metrics total). So the dropped rounds ARE
-    # inside cost_with_headroom (their input side). This block surfaces them
-    # again purely as a breakdown; do NOT add ccr_overhead into a second cost
+    # folded into the cost view at the outcome funnel
+    # (set_pending_ccr_continuation_usage -> the cost_tracker.record_tokens
+    # cache_read / cache_write / uncached args that cost_with_headroom prices,
+    # plus output_tokens for the budget). So the dropped rounds ARE inside
+    # cost_with_headroom (their input side). This block surfaces them again
+    # purely as a token breakdown; do NOT add ccr_overhead into a second cost
     # total, that double-counts. (The earlier "already inside cost_with_headroom"
     # claim was inaccurate: handle_response returns only the final round, so
     # without the funnel fold the dropped rounds were never costed at all.)
+    # Note the funnel deliberately leaves metrics / SavingsTracker / RequestLog
+    # on the final round's numbers — continuation cache reads must not be
+    # credited as cache savings.
     summary["cost"]["retrieval_cost_usd"] = round(
         estimate_cost_usd(primary_model, tokens_retrieved + overhead_tokens), 6
     )

headroom/proxy/handlers/openai.py

@@ -5255,12 +5255,22 @@ async def api_call_fn(
                             # outcome carries the initial round's usage (captured
                             # above, before CCR): it under-counts the final round
                             # and, combined with the continuation-cost fold,
-                            # double-books the initial round. Absent fields
-                            # default to 0 so a partial final payload cannot
-                            # leave the stale initial usage in place.
+                            # double-books the initial round. A partial final
+                            # payload must not leave the stale initial usage in
+                            # place either, so absent fields are re-derived, not
+                            # kept: output/cached fall to 0, and the input side
+                            # keeps the same ``original_tokens`` floor the
+                            # pre-CCR capture above used — dropping to 0 there
+                            # would price this round's input at nothing, since
+                            # the folded continuation tokens keep
+                            # ``record_tokens``' "no breakdown -> use
+                            # tokens_sent" fallback from firing.
                             _final_usage = resp_json.get("usage") or {}
                             if isinstance(_final_usage, dict):
-                                total_input_tokens = _usage_int(_final_usage.get("input_tokens"))
+                                total_input_tokens = _usage_int(
+                                    _final_usage.get("input_tokens"),
+                                    original_tokens,
+                                )
                                 output_tokens = _usage_int(_final_usage.get("output_tokens"))
                                 _final_details = _final_usage.get("input_tokens_details") or {}
                                 if isinstance(_final_details, dict):

headroom/proxy/outcome.py

@@ -85,6 +85,21 @@ def consume_pending_ccr_continuation_usage() -> tuple[int, int, int, int] | None
     return value
 
 
+def clear_pending_outcome_side_channels() -> None:
+    """Drop both pending side-channel values in the CALLER's context.
+
+    ``emit_request_outcome`` consumes them itself, but production reaches it
+    through ``asyncio.shield`` (``HeadroomProxy._record_request_outcome``), which
+    runs the funnel in a child task holding a *copy* of the caller's context —
+    so the funnel's ``set(None)`` never reaches the caller. Any call site that
+    shields must clear here afterwards; otherwise a second outcome emitted from
+    the same context (a batch item, a re-driven turn) books the same continuation
+    usage and proactive drawback a second time.
+    """
+    _pending_proactive_retrieval.set(None)
+    _pending_ccr_continuation_usage.set(None)
+
+
 @dataclass(frozen=True)
 class RequestOutcome:
     """Immutable, value-equal snapshot of a completed request.
@@ -391,8 +406,6 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
     and is awaitable-compatible. We could lift this to a typing.Protocol
     if/when another contract surface emerges, but YAGNI.
     """
-    import dataclasses
-
     from headroom.copilot_auth import consume_request_routed_to_copilot
     from headroom.proxy.cost import _summarize_transforms
     from headroom.proxy.models import RequestLog
@@ -417,7 +430,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
     _pending_proactive = consume_pending_proactive_retrieval()
 
     # Consume (read + clear) the dropped CCR continuation rounds' usage for this
-    # request. Folded into the billed totals only in the success section below —
+    # request. Folded into the cost view only in the success section below —
     # a >=500 short-circuit must never book continuation tokens that, by the
     # account-after-success rule in handle_response, were never billed.
     _pending_continuation = consume_pending_ccr_continuation_usage()
@@ -440,24 +453,28 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
         if _proactive_rec is not None:
             _proactive_rec(*_pending_proactive)
 
-    # Fold the dropped continuation rounds' real billed usage into the cost-
-    # relevant token totals. handle_response returns only the final round, so
-    # without this the intermediate rounds' usage never reaches cost_with_headroom
-    # (it sat in the /stats ccr_overhead side-channel but was deliberately
-    # excluded from the total). Each bucket lands in the field cost_with_headroom
-    # prices: uncached input at list, cache_read at the discounted rate,
-    # cache_write at the premium rate, and output into the budget/metrics total
-    # (cost_with_headroom itself is input-only). Client-facing response.usage is
-    # unchanged; only the internal billed totals move.
-    if _pending_continuation is not None:
-        _ct_uin, _ct_cr, _ct_cw, _ct_out = _pending_continuation
-        outcome = dataclasses.replace(
-            outcome,
-            output_tokens=outcome.output_tokens + _ct_out,
-            uncached_input_tokens=outcome.uncached_input_tokens + _ct_uin,
-            cache_read_tokens=outcome.cache_read_tokens + _ct_cr,
-            cache_write_tokens=outcome.cache_write_tokens + _ct_cw,
-        )
+    # The dropped continuation rounds were really billed, so they belong in the
+    # cost view — and ONLY there, which is why they are added at the
+    # ``cost_tracker.record_tokens`` call below instead of onto ``outcome``:
+    #
+    # * ``metrics.record_request`` forwards ``cache_read_tokens`` to
+    #   SavingsTracker, which credits it as cache *savings*
+    #   (``_estimate_cache_savings_usd`` -> ``lifetime.cache_savings_usd`` ->
+    #   ``SavingsSnapshot.total_savings``). Folding there would let CCR retrieval
+    #   overhead RAISE the reported savings.
+    # * the output-shaper estimator below reads ``outcome.output_tokens`` as the
+    #   observed completion length of *this* turn; continuation output would bias
+    #   the per-arm counterfactual.
+    # * the RequestLog row and the PERF line describe the one round the client
+    #   received, so they stay on the final round's numbers.
+    #
+    # ``cost_tracker.record_tokens`` is where the billed input breakdown lands
+    # (``_api_{uncached,cache_read,cache_write}_by_model`` -> the per-category
+    # pricing behind ``cost_with_headroom_usd``) and where ``_costs`` is appended
+    # for budget enforcement, so the fold reaches both of the surfaces that are
+    # meant to reflect real spend. ``/stats`` ``cost.ccr_overhead`` keeps the
+    # token-level attribution view.
+    _ct_uin, _ct_cr, _ct_cw, _ct_out = _pending_continuation or (0, 0, 0, 0)
 
     # Output-shaping savings ledger (counterfactual estimator). The shaper
     # tags each request's (arm, stratum) onto ``transforms_applied``; feed the
@@ -513,19 +530,23 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
         tool_search_saved=tool_search_saved,
     )
 
-    # 2. Cost tracker (optional).
+    # 2. Cost tracker (optional). The dropped CCR continuation rounds are folded
+    #    in here — see the comment above the ``_ct_*`` unpack. The 5m/1h split is
+    #    left alone: the provider does not report a TTL for a continuation round,
+    #    so only the cache-write total moves and the TTL attribution view stays
+    #    honest about what it actually observed.
     cost_tracker = getattr(handler, "cost_tracker", None)
     if cost_tracker is not None:
         cost_tracker.record_tokens(
             outcome.model,
             outcome.tokens_saved,
             outcome.optimized_tokens,
-            cache_read_tokens=outcome.cache_read_tokens,
-            cache_write_tokens=outcome.cache_write_tokens,
+            cache_read_tokens=outcome.cache_read_tokens + _ct_cr,
+            cache_write_tokens=outcome.cache_write_tokens + _ct_cw,
             cache_write_5m_tokens=outcome.cache_write_5m_tokens,
             cache_write_1h_tokens=outcome.cache_write_1h_tokens,
-            uncached_tokens=outcome.uncached_input_tokens,
-            output_tokens=outcome.output_tokens,
+            uncached_tokens=outcome.uncached_input_tokens + _ct_uin,
+            output_tokens=outcome.output_tokens + _ct_out,
         )
 
     # 3. Per-request log (optional). The ``client`` outcome field is

headroom/proxy/server.py

@@ -1940,7 +1940,10 @@ async def _record_request_outcome(self, outcome: RequestOutcome) -> None:
         See ``docs/superpowers/specs/P0-proxy-pipeline-audit.md`` for the
         divergence catalog this funnel collapses.
         """
-        from headroom.proxy.outcome import emit_request_outcome
+        from headroom.proxy.outcome import (
+            clear_pending_outcome_side_channels,
+            emit_request_outcome,
+        )
 
         # Shielded because four call sites are `finally:` blocks inside streaming
         # async generators (streaming.py:1611, :1859, :2069, openai.py:8614). A
@@ -1953,7 +1956,14 @@ async def _record_request_outcome(self, outcome: RequestOutcome) -> None:
         # The shield does not swallow the cancellation — the await below still
         # raises CancelledError, so generator teardown propagates exactly as
         # before. It only keeps the bookkeeping from being torn in half.
-        await asyncio.shield(emit_request_outcome(self, outcome))
+        try:
+            await asyncio.shield(emit_request_outcome(self, outcome))
+        finally:
+            # The shielded funnel runs in a child task, which holds a COPY of this
+            # context: its consume_* cleared the copy, not ours. Clear here as well,
+            # or a second outcome emitted from this same context would re-book the
+            # same CCR continuation usage / proactive drawback.
+            clear_pending_outcome_side_channels()
 
     async def _next_request_id(self) -> str:
         """Generate unique request ID."""