Skip to content

Recipe: Stream Cost & Time Budgets

Goal

Enforce cost-based and time-based limits during streaming to control spend in budget-sensitive applications.

Cost budget

Use the iteration_complete event to track accumulated cost and break when a threshold is exceeded. Wrap the stream in contextlib.aclosing so that the break closes the generator, which cancels the request behind it (see Stopping a stream early):

import asyncio
import contextlib

from doc_processing_toolkit import Agent

class GroupedStreamPrinter:
    def __init__(self):
        self.current_type = None
        self.parts = []

    def push(self, event_type, content):
        if content is None:
            return

        if self.current_type is None:
            self.current_type = event_type
        elif event_type != self.current_type:
            self.flush()
            self.current_type = event_type

        self.parts.append(content)

    def flush(self):
        if not self.parts or self.current_type is None:
            self.current_type = None
            self.parts = []
            return

        print(f"\n[{self.current_type}]\n{''.join(self.parts)}\n---", flush=True)
        self.current_type = None
        self.parts = []

async def budget_limited_stream(agent: Agent, message: str, max_cost_usd: float = 0.10):
    cost_so_far = 0.0
    printer = GroupedStreamPrinter()

    async with contextlib.aclosing(agent.stream(message)) as stream:
        async for event in stream:
            if event.event_type in {"text_delta", "thinking_delta"}:
                printer.push(event.event_type, event.content)
            elif event.event_type == "tool_call_result":
                # A tool that spent money (speech, images, a sub-agent) reports it here.
                cost_so_far += event.metadata.get("tool_reported_cost", 0)
            elif event.event_type == "iteration_complete":
                printer.flush()
                cost_so_far += event.metadata.get("cost", 0)
                if cost_so_far > max_cost_usd:
                    print(f"\nBudget limit reached: ${cost_so_far:.4f}")
                    break
            elif event.event_type == "error":
                printer.flush()
                print(f"\nError: {event.content}")
                break

    printer.flush()

When does cost update?

Cost information is available in the iteration_complete event, which fires after each iteration of the tool-calling loop. For single-turn calls this fires once at the end. Multi-turn tool-calling agents emit it after every iteration. Each event covers every model request of the tool loop since the previous one, so the running sum is the run's request cost up to that event. Before 0.5.1 the request that asked for a round's tools was in no event, and this sum under-counted every tool run. The complete event's result carries the exact total, including a request after the last event and a run that emitted none (a forced tool_choice exit, for example). Its cost is the model's token cost only: a tool that spends money of its own reports it on its tool_call_result event as tool_reported_cost (see Tool-reported spend), so a budget that reads only iteration_complete misses it.

Budget checks are post-iteration by design

Breaking on iteration_complete inside aclosing stops the next iteration from starting: the run is cancelled at that point, between tool cycles, so no model/tool exchange is cut off midway. A bare break without aclosing may not close the generator, and then the run keeps going and keeps billing.

A failed or aborted turn still spent money

The error event's metadata carries the partial input_tokens, output_tokens, cost and iterations recovered from the failure — including a hook abort, which reports aborted: True and its reason. Add event.metadata["cost"] to your running total there rather than dropping the turn, or a budget tracker under-counts every run that ends badly. A terminal complete event always follows the error event, so a finally-style teardown can stay bound to complete. See Streaming.

Time budget

Check elapsed wall-clock time on each event:

async def time_limited_stream(agent: Agent, message: str, max_seconds: float = 60.0):
    loop = asyncio.get_running_loop()
    start = loop.time()

    async with contextlib.aclosing(agent.stream(message)) as stream:
        async for event in stream:
            elapsed = loop.time() - start
            if elapsed > max_seconds:
                print(f"\nTime limit reached: {elapsed:.1f}s")
                break
            if event.event_type == "text_delta":
                print(event.content, end="", flush=True)

Combined budget

Check both cost and time together for defense-in-depth:

async def guarded_stream(
    agent: Agent,
    message: str,
    max_cost: float = 0.10,
    max_seconds: float = 120.0,
):
    loop = asyncio.get_running_loop()
    start = loop.time()
    cost_so_far = 0.0
    printer = GroupedStreamPrinter()

    async with contextlib.aclosing(agent.stream(message)) as stream:
        async for event in stream:
            # Check time budget
            if loop.time() - start > max_seconds:
                printer.flush()
                print("\nTime budget exceeded")
                break

            # Track cost budget
            if event.event_type in {"text_delta", "thinking_delta"}:
                printer.push(event.event_type, event.content)
            elif event.event_type == "tool_call_result":
                # A tool that spent money (speech, images, a sub-agent) reports it here.
                cost_so_far += event.metadata.get("tool_reported_cost", 0)
            elif event.event_type == "iteration_complete":
                printer.flush()
                cost_so_far += event.metadata.get("cost", 0)
                if cost_so_far > max_cost:
                    print(f"\nCost budget exceeded: ${cost_so_far:.4f}")
                    break

    printer.flush()

Provider billing on cancellation

Closing the stream cancels the request and waits at most 5 s for it to stop. What the run had already spent is not lost: agent.last_abandoned_stream holds it as an AgentResult dict with "abandoned": True, and an AgentSession banks it into its totals. Its cost is the run's spend so far, not an increment, so read it in place of your running total rather than adding it. On a run with local tools it covers the completed iterations only, and the request in flight at the close is not included. See What a closed stream reports.

When a stream is cancelled by closing it:

Provider Billing behavior
Fireworks Stops generation; avoids billing for ungenerated tokens
OpenAI Billed for tokens generated before cancellation
Anthropic Billed for tokens generated before cancellation
Google Billed for tokens generated before cancellation

The key savings come from tokens that would have been generated after cancellation.

What to validate

  • Cost reported in iteration_complete matches expected magnitude
  • Time budget triggers within a reasonable margin of max_seconds
  • Breaking from the stream does not raise exceptions
  • After a budget break, agent.last_abandoned_stream carries "abandoned": True and the spend so far
  • result.cost in the complete event (if reached) matches tracked cost