fix(middleware): preserve serving-hop usage iterations

.changelog/serving-hop-usage-iterations.md

@@ -0,0 +1 @@
+fix(middleware): preserve serving-hop usage iterations during streaming refusal fallback

src/anthropic/lib/middleware/_fallbacks.py

@@ -45,6 +45,7 @@
 log: logging.Logger = logging.getLogger("anthropic.lib.middleware")
 
 _MESSAGES_PATH = "/v1/messages"
+SAMPLING_ITERATION_TYPES = ("message", "fallback_message")
 
 CAPACITY_ERROR_STATUSES = (429, 529)
 
@@ -954,19 +955,23 @@ def feed(self, sse: ServerSentEvent) -> list[bytes]:
                     stop_details = _as_dict(delta.get("stop_details"))
                     if stop_details is not None:
                         stop_details.setdefault("recommended_model", None)
-                # Terminal hop. Replace iterations, don't append: this hop's own
-                # message_delta self-reports its iterations without a `model` (a
-                # fresh non-fallback request doesn't know it served a chain).
-                # Server-side `fallbacks` relabels the whole chain instead —
-                # refused hops as `message`, the serving hop as
-                # `fallback_message` — so the recorded chain replaces the
-                # self-report, with this hop's own entry relabeled as the
-                # `fallback_message` completer.
+                # Earlier tool-loop and compaction entries must survive the chain splice.
                 usage = _as_dict(event.get("usage")) or {}
-                usage["iterations"] = [
-                    *splice.iterations,
-                    _serving_iteration_entry(usage, splice.model),
-                ]
+                reported = usage.get("iterations")
+                entries: list[dict[str, Any]] = []
+                if isinstance(reported, list):
+                    entries = cast("list[dict[str, Any]]", reported)
+                for entry in entries:
+                    entry["model"] = entry.get("model") or splice.model
+                for entry in reversed(entries):
+                    if entry.get("type") in SAMPLING_ITERATION_TYPES:
+                        entry["type"] = "fallback_message"
+                        break
+                else:
+                    # Without a sampling record, use the hop total alone: it already
+                    # includes any reported compaction or advisor tokens.
+                    entries = [_to_iteration_usage("fallback_message", splice.model, usage)]
+                usage["iterations"] = [*splice.iterations, *entries]
                 event["usage"] = usage
                 _forward_input_transformations(event, self._suppressed_start_message())
                 return [*frames, _emit("message_delta", event)]
@@ -1251,17 +1256,6 @@ def _declined_iteration_entries(refusal: _Refusal, model_label: str) -> list[dic
     return [_to_iteration_usage("message", model_label, refusal.usage)]
 
 
-def _serving_iteration_entry(delta_usage: dict[str, Any], model: str) -> dict[str, Any]:
-    """The serving hop's `fallback_message` completer entry: its own
-    self-reported iteration relabeled, or one built from its delta usage."""
-    reported = delta_usage.get("iterations")
-    if isinstance(reported, list) and reported:
-        last = _as_dict(cast("List[Any]", reported)[-1])
-        if last is not None:
-            return {**last, "type": "fallback_message", "model": last.get("model") or model}
-    return _to_iteration_usage("fallback_message", model, delta_usage)
-
-
 class _BlockTracker:
     """Block bookkeeping for one stream of the splice: accumulates each content
     block from its deltas (for the continuation prefill), shifts wire indices

tests/lib/streaming/test_beta_messages.py

@@ -5,7 +5,7 @@
 import base64
 import pathlib
 import threading
-from typing import Any, Set, Dict, List, Tuple, TypeVar, cast
+from typing import Any, Set, Dict, List, Tuple, TypeVar, cast, get_args
 
 import httpx2
 import pytest
@@ -14,6 +14,7 @@
 from anthropic import Anthropic, AsyncAnthropic
 from anthropic._utils import assert_overloads_in_sync, assert_signatures_in_sync
 from anthropic._compat import PYDANTIC_V1, get_model_fields
+from anthropic.lib.middleware import BetaFallbackState, BetaRefusalFallbackMiddleware, _fallbacks
 from anthropic.types.beta.beta_message import BetaMessage
 from anthropic.lib.streaming._beta_types import (
     BetaInputJsonEvent,
@@ -23,10 +24,21 @@
 from anthropic.types.beta.beta_tool_param import BetaToolParam
 from anthropic.resources.messages.messages import DEPRECATED_MODELS
 from anthropic.lib.streaming._beta_messages import TRACKS_TOOL_INPUT, BetaMessageStream, BetaAsyncMessageStream
+from anthropic.types.beta.beta_iterations_usage import BetaIterationsUsageItem
 from anthropic.types.beta.beta_message_delta_usage import BetaMessageDeltaUsage
 from anthropic.types.beta.beta_raw_message_delta_event import Delta as BetaRawMessageDelta, BetaRawMessageDeltaEvent
 
 from .helpers import get_response, to_async_iter
+from ..test_refusal_fallback_streaming import (
+    PARAMS,
+    STREAM_A,
+    STREAM_B,
+    FALLBACKS,
+    FALLBACK_MODEL,
+    sse_response,
+    make_sync_client,
+    make_async_client,
+)
 
 base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
 api_key = "my-anthropic-api-key"
@@ -1096,3 +1108,156 @@ def test_tracks_tool_input_type_alias_is_up_to_date() -> None:
             f"ContentBlock type {block_type.__name__} has an input property, "
             f"but is not included in TRACKS_TOOL_INPUT. You probably need to update the TRACKS_TOOL_INPUT type alias."
         )
+
+
+def iteration(kind: str, input_tokens: int, output_tokens: int, **extra: Any) -> dict[str, Any]:
+    return {
+        "type": kind,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cache_read_input_tokens": 0,
+        "cache_creation_input_tokens": 0,
+        **extra,
+    }
+
+
+FIRST = iteration(
+    "message",
+    210_001,
+    8,
+    model=FALLBACK_MODEL,
+    cache_read_input_tokens=17,
+    cache_creation_input_tokens=23,
+    cache_creation={"ephemeral_5m_input_tokens": 11, "ephemeral_1h_input_tokens": 12},
+    future_detail={"source": "synthetic"},
+)
+FINAL = iteration("message", 31, 547)
+COMPACTION = iteration("compaction", 205_000, 64)
+ADVISOR = iteration("advisor_message", 50, 12, model="claude-sonnet-4-6")
+COMPLETER = {**FINAL, "type": "fallback_message", "model": FALLBACK_MODEL}
+AGGREGATE = {**COMPLETER, "cache_creation": None}
+LABELED_FINAL = {**FINAL, "model": FALLBACK_MODEL}
+LABELED_COMPACTION = {**COMPACTION, "model": FALLBACK_MODEL}
+
+USAGE_CASES = [
+    pytest.param(None, [AGGREGATE], id="absent-breakdown"),
+    pytest.param([], [AGGREGATE], id="empty-breakdown"),
+    pytest.param([FINAL], [COMPLETER], id="single-sampling"),
+    pytest.param([COMPLETER], [COMPLETER], id="reported-fallback-sampling"),
+    # The delta's aggregate already includes these partial records (5 input / 7 output).
+    pytest.param([iteration("compaction", 5, 7)], [AGGREGATE], id="compaction-only-breakdown"),
+    pytest.param([iteration("advisor_message", 5, 7)], [AGGREGATE], id="advisor-only-breakdown"),
+    pytest.param(
+        [iteration("compaction", 5, 7), iteration("advisor_message", 3, 2)],
+        [AGGREGATE],
+        id="non-sampling-breakdown",
+    ),
+    pytest.param([FIRST, FINAL], [FIRST, COMPLETER], id="server-tool-loop"),
+    pytest.param(
+        [{key: value for key, value in FIRST.items() if key != "model"}, FINAL],
+        [FIRST, COMPLETER],
+        id="unlabeled-first-sampling",
+    ),
+    pytest.param([FIRST, FINAL, FINAL], [FIRST, LABELED_FINAL, COMPLETER], id="multiple-sampling"),
+    pytest.param(
+        [FIRST, COMPACTION, ADVISOR, FINAL],
+        [FIRST, LABELED_COMPACTION, ADVISOR, COMPLETER],
+        id="mixed-breakdown",
+    ),
+    pytest.param([FIRST, FINAL, COMPACTION], [FIRST, COMPLETER, LABELED_COMPACTION], id="trailing-compaction"),
+    pytest.param(
+        [{**ADVISOR, "model": None}, FINAL],
+        [{**ADVISOR, "model": FALLBACK_MODEL}, COMPLETER],
+        id="unlabeled-advisor",
+    ),
+    pytest.param(
+        [{**FINAL, "model": "claude-opus-4-8-canonical"}],
+        [{**COMPLETER, "model": "claude-opus-4-8-canonical"}],
+        id="reported-model",
+    ),
+]
+
+
+def serving_with_iterations(iterations: list[dict[str, Any]] | None) -> str:
+    """Keep the synthetic serving wire, supplying an independently specified usage breakdown."""
+    lines: list[str] = []
+    for line in STREAM_B.splitlines(keepends=True):
+        if line.startswith("data: "):
+            event = json.loads(line[6:])
+            if event["type"] == "message_delta":
+                if iterations is None:
+                    event["usage"].pop("iterations", None)
+                else:
+                    event["usage"]["iterations"] = iterations
+                line = "data: " + json.dumps(event) + "\n"
+        lines.append(line)
+    return "".join(lines)
+
+
+@pytest.mark.respx(base_url=base_url)
+@pytest.mark.parametrize("iterations, expected", USAGE_CASES)
+def test_serving_usage_retains_each_iteration(
+    respx_mock: MockRouter, iterations: list[dict[str, Any]] | None, expected: list[dict[str, Any]]
+) -> None:
+    respx_mock.post("/v1/messages").mock(
+        side_effect=[sse_response(STREAM_A), sse_response(serving_with_iterations(iterations))]
+    )
+    with make_sync_client(middleware=[BetaRefusalFallbackMiddleware(FALLBACKS)]) as client, BetaFallbackState():
+        with client.beta.messages.stream(**PARAMS) as stream:
+            events = list(stream)
+            message = stream.get_final_message()
+
+    assert len(respx_mock.calls) == 2
+    request = cast(httpx2.Request, respx_mock.calls[1][0])
+    assert request.method == "POST"
+    assert request.url.path == "/v1/messages"
+    assert "fallback-credit-2026-07-01" in request.headers["anthropic-beta"]
+    assert json.loads(request.content)["model"] == FALLBACK_MODEL
+    delta = next(event for event in events if event.type == "message_delta")
+    assert delta.usage.iterations is not None
+    assert [entry.to_dict() for entry in delta.usage.iterations[1:]] == expected
+    assert message.usage.iterations is not None
+    assert [entry.to_dict() for entry in message.usage.iterations[1:]] == expected
+    assert len([event for event in events if event.type == "message_start"]) == 1
+    assert len([event for event in events if event.type == "message_stop"]) == 1
+    assert message.stop_reason == "end_turn"
+    assert message.usage.output_tokens == 547
+
+
+@pytest.mark.respx(base_url=base_url)
+@pytest.mark.parametrize("iterations, expected", USAGE_CASES)
+async def test_async_serving_usage_retains_each_iteration(
+    respx_mock: MockRouter, iterations: list[dict[str, Any]] | None, expected: list[dict[str, Any]]
+) -> None:
+    respx_mock.post("/v1/messages").mock(
+        side_effect=[sse_response(STREAM_A), sse_response(serving_with_iterations(iterations))]
+    )
+    with BetaFallbackState():
+        async with make_async_client(middleware=[BetaRefusalFallbackMiddleware(FALLBACKS)]) as client:
+            async with client.beta.messages.stream(**PARAMS) as stream:
+                events = [event async for event in stream]
+                message = await stream.get_final_message()
+
+    assert len(respx_mock.calls) == 2
+    request = cast(httpx2.Request, respx_mock.calls[1][0])
+    assert request.method == "POST"
+    assert request.url.path == "/v1/messages"
+    assert "fallback-credit-2026-07-01" in request.headers["anthropic-beta"]
+    assert json.loads(request.content)["model"] == FALLBACK_MODEL
+    delta = next(event for event in events if event.type == "message_delta")
+    assert delta.usage.iterations is not None
+    assert [entry.to_dict() for entry in delta.usage.iterations[1:]] == expected
+    assert message.usage.iterations is not None
+    assert [entry.to_dict() for entry in message.usage.iterations[1:]] == expected
+    assert len([event for event in events if event.type == "message_start"]) == 1
+    assert len([event for event in events if event.type == "message_stop"]) == 1
+    assert message.stop_reason == "end_turn"
+    assert message.usage.output_tokens == 547
+
+
+def test_serving_iteration_types_track_generated_union() -> None:
+    variants = get_args(get_args(BetaIterationsUsageItem)[0])
+    generated_types = {get_args(get_model_fields(variant)["type"].annotation)[0] for variant in variants}
+    assert generated_types == {*_fallbacks.SAMPLING_ITERATION_TYPES, "compaction", "advisor_message"}, (
+        "Update _HopReader when the generated usage iteration variants change."
+    )