From 30b6ee93d721d76752c297897183b1aee679db3c Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 9 Oct 2026 19:30:51 +0000 Subject: [PATCH] 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 | 2 +- pydantic_ai_slim/pydantic_ai/_agent_graph.py | 101 +++++++++-------- .../pydantic_ai/models/__init__.py | 5 +- .../pydantic_ai_harness/spend/README.md | 2 +- tests/harness/spend/test_spend.py | 104 +++++++++++++++++- tests/models/test_streamed_continuation.py | 20 ++-- tests/test_capability_hooks.py | 46 +++++++- 7 files changed, 217 insertions(+), 63 deletions(-) diff --git a/docs/harness/spend.md b/docs/harness/spend.md index 171150bf11..e38dfe5baf 100644 --- a/docs/harness/spend.md +++ b/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 diff --git a/pydantic_ai_slim/pydantic_ai/_agent_graph.py b/pydantic_ai_slim/pydantic_ai/_agent_graph.py index eb7614b6c0..39c1b2d47e 100644 --- a/pydantic_ai_slim/pydantic_ai/_agent_graph.py +++ b/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) diff --git a/pydantic_ai_slim/pydantic_ai/models/__init__.py b/pydantic_ai_slim/pydantic_ai/models/__init__.py index 26f0886fea..baa7c35591 100644 --- a/pydantic_ai_slim/pydantic_ai/models/__init__.py +++ b/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 diff --git a/src/pydantic_ai_harness/pydantic_ai_harness/spend/README.md b/src/pydantic_ai_harness/pydantic_ai_harness/spend/README.md index e030020683..78a71302eb 100644 --- a/src/pydantic_ai_harness/pydantic_ai_harness/spend/README.md +++ b/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 diff --git a/tests/harness/spend/test_spend.py b/tests/harness/spend/test_spend.py index 42b90c699f..f0f1ead661 100644 --- a/tests/harness/spend/test_spend.py +++ b/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.""" diff --git a/tests/models/test_streamed_continuation.py b/tests/models/test_streamed_continuation.py index d18bbb1cdd..6ada25fd93 100644 --- a/tests/models/test_streamed_continuation.py +++ b/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: diff --git a/tests/test_capability_hooks.py b/tests/test_capability_hooks.py index da60b343f5..516710611c 100644 --- a/tests/test_capability_hooks.py +++ b/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: