Record an interrupted stream's partial response in the billed-response ledger before wrap_model_request unwinds

So `SpendLimits` charges it, together with its `failed_attempts`, exactly once.

docs/harness/spend.md

@@ -82,7 +82,7 @@ SpendLimits(budgets=[Budget(window='month', scope=lambda ctx: ctx.deps.tenant_id
 
 Once a recorded response spends a window, `SpendLimits` refuses the next request.
 
-Not: that spend stays under the ceiling. The request that crosses the line completes, and concurrent runs can each pass the check before any of them records anything. A stream the caller abandons part-way never reaches the accounting hook, so its tokens are billed by the provider and invisible here. Treat this as a brake on a runaway loop, not as an accounting ledger; reconcile against the provider's own numbers if you need the second thing.
+Not: that spend stays under the ceiling. The request that crosses the line completes, and concurrent runs can each pass the check before any of them records anything. A stream the caller stops part-way is charged for the usage the provider had reported by then, which can fall short of what it bills, since many providers report output tokens only with a stream's last event. Treat this as a brake on a runaway loop, not as an accounting ledger; reconcile against the provider's own numbers if you need the second thing.
 
 ## Reading the numbers
 

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -252,12 +252,10 @@ def _apply_context_changes(changes: Sequence[tuple[ContextVar[Any], Any]]) -> No
         var.set(value)
 
 
-async def _resolve_interrupted_stream_state(
-    model: models.Model,
-    stream_error: BaseException,
-    partial: _messages.ModelResponse,
-) -> _messages.ModelResponseState:
-    """State to record for a streamed turn the consumer stopped, cancelling a leaked job when appropriate.
+def _interrupted_stream_response(
+    stream_error: BaseException, partial: _messages.ModelResponse
+) -> _messages.ModelResponse:
+    """The response to record for a streamed turn the consumer stopped.
 
     The composite treats every `aclose()` (which the handler teardown triggers) as a *detach*, so
     `partial.state` is `'suspended'` whenever the last segment is a still-pending job — regardless of
@@ -266,13 +264,12 @@ async def _resolve_interrupted_stream_state(
     - `GeneratorExit` is a walk-away detach (the consumer broke out of `run_stream`/`stream_text`). Mirror
       the non-streaming detach: keep `'suspended'` so the run is resumable, and leave the job alive.
     - any other exception is a genuine downstream failure. Mirror the non-streaming cancel-on-error policy:
-      force `'interrupted'` (non-resumable) and best-effort cancel the still-live job so it doesn't leak.
+      force `'interrupted'` (non-resumable); the graph best-effort cancels the still-live job so it doesn't leak.
     """
-    if isinstance(stream_error, GeneratorExit) and partial.state == 'suspended':
-        return 'suspended'
-    if partial.state == 'suspended':
-        await cancel_suspended_job(model, partial)
-    return 'interrupted'
+    detached = isinstance(stream_error, GeneratorExit) and partial.state == 'suspended'
+    response = replace(partial, state='suspended' if detached else 'interrupted')
+    fill_response_cost(response)
+    return response
 
 
 NEW_CONVERSATION: Literal['new'] = 'new'
@@ -1433,6 +1430,9 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         _handler_called = False
         _handler_usage_recorded = False
         _stream_cut_short = False
+        # Set when the consumer stops the stream by raising, before the graph task cancels the handler.
+        stream_error: BaseException | None = None
+        _interrupted_response: _messages.ModelResponse | None = None
         time_to_first_chunk: float | None = None
         accounted_responses: list[_messages.ModelResponse] = []
         before_model_request_context: list[tuple[ContextVar[Any], Any]] = []
@@ -1441,6 +1441,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
             req_ctx: ModelRequestContext,
         ) -> _messages.ModelResponse:
             nonlocal _handler_called, _handler_response, _handler_usage_recorded, _stream_cut_short, time_to_first_chunk
+            nonlocal _interrupted_response
             if _handler_called:
                 raise exceptions.UserError('`wrap_model_request` may call its handler only once')
             _handler_called = True
@@ -1464,25 +1465,34 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
             # into one continuous stream, so the whole chain is presented as a single
             # `AgentStream` and the model-request hooks wrap it once. The step is counted in
             # `ctx.state.usage.requests` when its response is committed, not here.
-            async with model_request_stream(req_ctx.model, request_context=req_ctx, run_context=run_context) as sr:
-                self._did_stream = True
-                agent_stream = self._build_agent_stream(ctx, sr, req_ctx.model_request_parameters)
-                agent_stream_holder.append(agent_stream)
-                stream_ready.set()
-                try:
-                    await stream_done.wait()
-                finally:
-                    if not stream_done.is_set():
-                        # `wrap_model_request` stopped the handler while the stream is still being consumed,
-                        # e.g. a parallel `InputGuardrail` blocking the prompt. Stop the consumer's pull before
-                        # the stream closes; it continues with the response `wrap_model_request` returns instead.
-                        _stream_cut_short = await agent_stream._abandon_model_stream(_wrap_response)  # pyright: ignore[reportPrivateUsage]
-                    # Report TTFT in a `finally` so it also lands when the consumer raises
-                    # mid-iteration and `_cancel_task(wrap_task)` injects CancelledError at
-                    # the `wait()` above, mirroring `InstrumentedModel.request_stream`. On
-                    # that cancelled path `finish` is never reached today (no metrics of any
-                    # kind are recorded), so this is symmetry rather than an observable fix.
-                    time_to_first_chunk = sr.time_to_first_chunk(request_start)
+            try:
+                async with model_request_stream(req_ctx.model, request_context=req_ctx, run_context=run_context) as sr:
+                    self._did_stream = True
+                    agent_stream = self._build_agent_stream(ctx, sr, req_ctx.model_request_parameters)
+                    agent_stream_holder.append(agent_stream)
+                    stream_ready.set()
+                    try:
+                        await stream_done.wait()
+                    finally:
+                        if not stream_done.is_set():
+                            # `wrap_model_request` stopped the handler while the stream is still being consumed,
+                            # e.g. a parallel `InputGuardrail` blocking the prompt. Stop the consumer's pull before
+                            # the stream closes; it continues with the response `wrap_model_request` returns instead.
+                            _stream_cut_short = await agent_stream._abandon_model_stream(_wrap_response)  # pyright: ignore[reportPrivateUsage]
+                        # Report TTFT in a `finally` so it also lands when the consumer raises
+                        # mid-iteration and `_cancel_task(wrap_task)` injects CancelledError at
+                        # the `wait()` above, mirroring `InstrumentedModel.request_stream`. On
+                        # that cancelled path `finish` is never reached today (no metrics of any
+                        # kind are recorded), so this is symmetry rather than an observable fix.
+                        time_to_first_chunk = sr.time_to_first_chunk(request_start)
+            except BaseException:
+                if stream_error is not None:
+                    # The consumer stopped the stream, so the graph task cancelled this one. Record what the
+                    # provider billed for so far before the `wrap_model_request` chain unwinds, so wrappers that
+                    # account for billed responses see it. The graph task commits this same response afterwards.
+                    _interrupted_response = _interrupted_stream_response(stream_error, agent_stream_holder[0].response)
+                    req_ctx._usage_response_ledger.responses.append(_interrupted_response)  # pyright: ignore[reportPrivateUsage]
+                raise
             # Streaming core errors surface in the consumer task, which cancels this wrap task;
             # `on_model_request_error` cannot recover an error after streaming has begun.
             response = sr.get()
@@ -1592,7 +1602,6 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
             return
 
         # Normal path: handler was called, stream is ready
-        stream_error: BaseException | None = None
         try:
             yield agent_stream_holder[0]
         except BaseException as exc:
@@ -1609,8 +1618,14 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
                     # We append directly rather than via `_append_response` to skip the usage-limit
                     # check; raising `UsageLimitExceeded` here would mask `stream_error`.
                     if agent_stream_holder:  # pragma: no branch
+                        partial = agent_stream_holder[0].response
                         await self._commit_interrupted_response(
-                            ctx, wrap_request_context.model, stream_error, agent_stream_holder[0].response
+                            ctx,
+                            wrap_request_context.model,
+                            partial,
+                            # The handler recorded the response it billed, unless `wrap_model_request`
+                            # had already stopped it before the consumer did.
+                            _interrupted_response or _interrupted_stream_response(stream_error, partial),
                         )
                 else:
                     try:
@@ -1647,18 +1662,18 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
     async def _commit_interrupted_response(
         ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[Any, Any]],
         model: models.Model,
-        stream_error: BaseException,
         partial: _messages.ModelResponse,
+        partial_response: _messages.ModelResponse,
     ) -> None:
-        """Record the response an interrupted stream produced so far, without checking usage limits."""
-        recorded_state = await _resolve_interrupted_stream_state(model, stream_error, partial)
-        partial_response = replace(
-            partial,
-            state=recorded_state,
-            run_id=ctx.state.run_id,
-            conversation_id=ctx.state.conversation_id,
-        )
-        fill_response_cost(partial_response)
+        """Commit the response an interrupted stream produced so far, without checking usage limits.
+
+        `partial_response` is `partial` as `_interrupted_stream_response` recorded it.
+        """
+        if partial.state == 'suspended' and partial_response.state == 'interrupted':
+            # A genuine failure rather than a walk-away: cancel the still-live job so it doesn't leak.
+            await cancel_suspended_job(model, partial)
+        partial_response.run_id = ctx.state.run_id
+        partial_response.conversation_id = ctx.state.conversation_id
         partial_response.workspace_ref = ctx.deps.workspace_ref
         _record_attempts_usage(ctx.state.usage, partial_response.failed_attempts)
         _usage_attribution.record_usage(ctx.state.usage, partial_response.usage)

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -407,8 +407,9 @@ class ModelRequestContext:
     def _usage_responses(self) -> tuple[ModelResponse, ...]:
         """The model responses whose usage was counted for this request.
 
-        Empty until the model responds. It includes responses a hook later replaced, and can hold more
-        than one response when a continued response was partly billed before an error hook recovered.
+        Empty until the model responds. It includes responses a hook later replaced and the partial response
+        of a stream the consumer stopped, and can hold more than one response when a continued response was
+        partly billed before an error hook recovered.
         Request contexts copied with `dataclasses.replace()` see the same responses.
 
         Private for now: read by the `Instrumentation` capability and by the Pydantic AI Harness's

src/pydantic_ai_harness/pydantic_ai_harness/spend/README.md

@@ -80,7 +80,7 @@ SpendLimits(budgets=[Budget(window='month', scope=lambda ctx: ctx.deps.tenant_id
 
 Once a recorded response spends a window, `SpendLimits` refuses the next request.
 
-Not: that spend stays under the ceiling. The request that crosses the line completes, and concurrent runs can each pass the check before any of them records anything. A stream the caller abandons part-way never reaches the accounting hook, so its tokens are billed by the provider and invisible here. Treat this as a brake on a runaway loop, not as an accounting ledger; reconcile against the provider's own numbers if you need the second thing.
+Not: that spend stays under the ceiling. The request that crosses the line completes, and concurrent runs can each pass the check before any of them records anything. A stream the caller stops part-way is charged for the usage the provider had reported by then, which can fall short of what it bills, since many providers report output tokens only with a stream's last event. Treat this as a brake on a runaway loop, not as an accounting ledger; reconcile against the provider's own numbers if you need the second thing.
 
 ## Reading the numbers
 

tests/harness/spend/test_spend.py

@@ -4,7 +4,8 @@ from __future__ import annotations
 
 import json
 import warnings
-from collections.abc import Awaitable, Callable, Mapping, Sequence
+from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
+from contextlib import asynccontextmanager
 from dataclasses import replace
 from datetime import UTC, datetime, timedelta
 from decimal import Decimal, InvalidOperation, localcontext
@@ -47,10 +48,12 @@ from pydantic_ai.messages import (
     ToolCallPart,
     ToolReturnPart,
 )
-from pydantic_ai.models import ModelRequestContext, ModelRequestParameters
+from pydantic_ai.models import ModelRequestContext, ModelRequestParameters, StreamedResponse
 from pydantic_ai.models.fallback import FallbackModel
 from pydantic_ai.models.function import AgentInfo, FunctionModel
 from pydantic_ai.models.test import TestModel
+from pydantic_ai.models.wrapper import WrapperModel
+from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import RunContext
 from pydantic_ai.usage import RequestUsage, RunUsage
 from pydantic_ai_harness import HarnessDeprecationWarning, spend
@@ -2649,6 +2652,103 @@ class TestContinuationAccrual:
         assert (await limits.status())[0].spent == Spent(usd=Decimal('5'), tokens=21, requests=1)
 
 
+class _AttemptsBeforeStream(WrapperModel):
+    """Streams the wrapped model's response as one a `FallbackModel` fell back to, after a billed rejected attempt."""
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncGenerator[StreamedResponse]:
+        async with super().request_stream(messages, model_settings, model_request_parameters, run_context) as stream:
+            stream.failed_attempts = [_attempt()]
+            yield stream
+
+
+async def _stop_streaming(agent: Agent[None, str], *, at: str) -> RunUsage:
+    """Stream the agent's first model response and raise once its text reaches `at`, returning the run's usage."""
+    async with agent.iter('go') as run:
+        with pytest.raises(RuntimeError, match='consumer stopped'):
+            async for node in run:  # pragma: no branch
+                if Agent.is_model_request_node(node):
+                    async with node.stream(run.ctx) as stream:
+                        async for text in stream.stream_text():  # pragma: no branch
+                            if at in text:
+                                raise RuntimeError('consumer stopped')
+    return run.usage
+
+
+class TestInterruptedStreams:
+    """A stream the consumer stops was billed for what the provider produced before it stopped."""
+
+    async def test_the_partial_response_is_charged_once(self):
+        store = _CountingStore()
+        model = ScriptedContinuationModel(
+            segments=[
+                StreamSegment(
+                    texts=['A', 'B'], state='complete', provider_response_id='a', input_tokens=7, output_tokens=3
+                )
+            ]
+        )
+        limits = SpendLimits[None](
+            budgets=[Budget(window='total')], store=store, price=lambda response: Decimal(response.usage.total_tokens)
+        )
+
+        usage = await _stop_streaming(Agent(model, deps_type=type(None), capabilities=[limits]), at='A')
+
+        assert usage == RunUsage(input_tokens=7, output_tokens=3, requests=1)
+        assert (await limits.status())[0].spent == Spent(usd=Decimal('10'), tokens=10, requests=1)
+        assert [[(entry.tokens, entry.requests) for entry in batch] for batch in store.batches] == [[(10, 1)]]
+
+    async def test_an_interrupted_chain_is_accrued_at_each_boundary(self):
+        """The segments a stopped chain completed are accrued at their boundaries, adding up to what core counted."""
+        store = _CountingStore()
+        model = ScriptedContinuationModel(
+            segments=[
+                StreamSegment(
+                    texts=['A'], state='suspended', provider_response_id='a', input_tokens=7, output_tokens=3
+                ),
+                StreamSegment(
+                    texts=['B', 'C'], state='complete', provider_response_id='b', input_tokens=8, output_tokens=3
+                ),
+            ]
+        )
+        limits = SpendLimits[None](
+            budgets=[Budget(window='total')], store=store, price=lambda response: Decimal(response.usage.total_tokens)
+        )
+
+        usage = await _stop_streaming(Agent(model, deps_type=type(None), capabilities=[limits]), at='B')
+
+        assert usage.total_tokens == 21
+        assert (await limits.status())[0].spent == Spent(usd=Decimal('21'), tokens=21, requests=1)
+        assert [[(entry.tokens, entry.requests) for entry in batch] for batch in store.batches] == [
+            [(10, 1)],
+            [(11, 0)],
+        ]
+
+    async def test_the_rejected_attempts_before_it_are_charged_too(self):
+        model = _AttemptsBeforeStream(
+            ScriptedContinuationModel(
+                segments=[
+                    StreamSegment(
+                        texts=['A', 'B'], state='complete', provider_response_id='a', input_tokens=7, output_tokens=3
+                    )
+                ]
+            )
+        )
+        limits = SpendLimits[None](
+            budgets=[Budget(window='total')], price=lambda response: Decimal(response.usage.total_tokens)
+        )
+
+        usage = await _stop_streaming(Agent(model, deps_type=type(None), capabilities=[limits]), at='A')
+
+        assert usage.total_tokens == 120
+        assert (await limits.status())[0].spent == Spent(usd=Decimal('120'), tokens=120, requests=2)
+
+
 class TestDeprecatedStore:
     """A store written against the released `SpendStore` keeps working, and says what it costs."""
 

tests/models/test_streamed_continuation.py

@@ -25,7 +25,7 @@ from typing import Any, Literal
 import pytest
 
 from pydantic_ai import Agent, capture_run_messages
-from pydantic_ai._agent_graph import _resolve_interrupted_stream_state  # pyright: ignore[reportPrivateUsage]
+from pydantic_ai._agent_graph import _interrupted_stream_response  # pyright: ignore[reportPrivateUsage]
 from pydantic_ai._run_context import RunContext
 from pydantic_ai.capabilities import AbstractCapability, Hooks
 from pydantic_ai.exceptions import SkipModelRequest, UnexpectedModelBehavior, UsageLimitExceeded, UserError
@@ -1333,7 +1333,7 @@ async def test_iter_node_stream_early_break_records_suspended() -> None:
 
     Unlike the `run_stream` early break (which records via the composite's own detach path), exiting a
     `node.stream(...)` context sends `GeneratorExit` into the graph's model-request stream, where
-    `_resolve_interrupted_stream_state` records the trailing response as resumable `'suspended'` and leaves
+    `_interrupted_stream_response` records the trailing response as resumable `'suspended'` and leaves
     the server-side job alive — the graph-level detach branch.
     """
     model = _ScriptedModel(
@@ -1364,23 +1364,19 @@ async def test_iter_node_stream_early_break_records_suspended() -> None:
     assert messages[-1].provider_response_id == 'job1'
 
 
-async def test_resolve_interrupted_stream_state_detach_vs_error() -> None:
-    """`_resolve_interrupted_stream_state`: a `GeneratorExit` walk-away with a suspended partial is a
-    resumable detach (`'suspended'`, job left alive); any other error is a genuine failure (`'interrupted'`,
-    job best-effort cancelled).
+def test_interrupted_stream_response_detach_vs_error() -> None:
+    """`_interrupted_stream_response`: a `GeneratorExit` walk-away with a suspended partial is a
+    resumable detach (`'suspended'`); any other error is a genuine failure (`'interrupted'`), whose
+    still-live job the graph then best-effort cancels.
 
     Unit-tested directly on the graph helper: whether the `GeneratorExit`-vs-error branch fires depends on
     stream-teardown ordering (which handle is closed first) that the `run_stream`/`iter` early-break tests
     above don't deterministically drive, so the branch itself is pinned here.
     """
-    model = _ScriptedModel(segments=[])
     suspended = _suspended(texts=['partial '], provider_response_id='job1', input_tokens=1, output_tokens=1)
 
-    assert await _resolve_interrupted_stream_state(model, GeneratorExit(), suspended) == 'suspended'
-    assert model.cancelled == []  # detach leaves the job alive
-
-    assert await _resolve_interrupted_stream_state(model, RuntimeError('boom'), suspended) == 'interrupted'
-    assert model.cancelled == [suspended]  # genuine error cancels the still-live job
+    assert _interrupted_stream_response(GeneratorExit(), suspended).state == 'suspended'
+    assert _interrupted_stream_response(RuntimeError('boom'), suspended).state == 'interrupted'
 
 
 async def test_cancel_through_wrapper_model_delegates() -> None:

tests/test_capability_hooks.py

@@ -21,7 +21,7 @@ import pytest
 from pydantic import BaseModel, ValidationError
 
 from pydantic_ai._run_context import RunContext
-from pydantic_ai.agent import Agent
+from pydantic_ai.agent import Agent, capture_run_messages
 from pydantic_ai.capabilities import (
     CombinedCapability,
     ToolSearch,
@@ -86,7 +86,7 @@ from .capability_models import (
     tool_calling_stream_function,
 )
 from .conftest import IsDatetime, IsStr
-from .continuation_utils import ScriptedContinuationModel, scripted_response
+from .continuation_utils import ScriptedContinuationModel, StreamSegment, scripted_response
 
 _SEARCH_TOOLS_NAME = ToolSearch.function_tool_name
 
@@ -554,6 +554,48 @@ class TestModelRequestHooks:
         assert len(observed_responses) == 1
         assert observed_responses[0] is result.all_messages()[-1]
 
+    async def test_usage_ledger_records_a_stream_the_consumer_stopped(self):
+        """The partial response is recorded before `wrap_model_request` unwinds, and is the one committed to history."""
+        observed_responses: list[ModelResponse] = []
+
+        class ObserveUsage(AbstractCapability[Any]):
+            async def wrap_model_request(
+                self,
+                ctx: RunContext[Any],
+                *,
+                request_context: ModelRequestContext,
+                handler: Any,
+            ) -> ModelResponse:
+                try:
+                    return await handler(request_context)
+                finally:
+                    observed_responses.extend(request_context._usage_responses)  # pyright: ignore[reportPrivateUsage]
+
+        model = ScriptedContinuationModel(
+            segments=[
+                StreamSegment(
+                    texts=['A', 'B'], state='complete', provider_response_id='a', input_tokens=7, output_tokens=3
+                )
+            ]
+        )
+        agent = Agent(model, capabilities=[ObserveUsage()])
+
+        with capture_run_messages() as messages:
+            async with agent.iter('hello') as run:
+                with pytest.raises(RuntimeError, match='stop'):
+                    async for node in run:  # pragma: no branch
+                        if Agent.is_model_request_node(node):
+                            async with node.stream(run.ctx) as stream:
+                                async for _ in stream:  # pragma: no branch
+                                    raise RuntimeError('stop')
+
+        partial = messages[-1]
+        assert isinstance(partial, ModelResponse)
+        assert partial.state == 'interrupted'
+        assert len(observed_responses) == 1
+        assert observed_responses[0] is partial
+        assert run.usage == RunUsage(input_tokens=7, output_tokens=3, requests=1)
+
     async def test_wrapper_recovery_keeps_continuation_usage(self):
         class RecoverAfterRejection(AbstractCapability[Any]):
             async def wrap_model_request(self, ctx: RunContext[Any], *, request_context: Any, handler: Any) -> Any: