Lesson 6: Operating Streams

Capturing usage when numbers arrive last

Token counts, and therefore cost, usually arrive in the final chunk of a stream. This creates three practical problems.

You cannot know the cost until the stream ends, so a pre-flight budget check must work from estimates. Counting input tokens before sending is reliable; output is a reservation you make and reconcile afterwards.

A cancelled or failed stream may never deliver usage. The tokens were still generated and you are still billed for them, so recording nothing understates your spend, and the gap grows with every disconnect.

Estimate when the real number is unavailable.

@dataclass
class StreamUsage:
    input_tokens: int
    output_tokens: int
    is_estimated: bool
    completed: bool


async def stream_with_usage(
    events: AsyncIterator[StreamEvent], *, input_tokens: int
) -> tuple[AsyncIterator[StreamEvent], Callable[[], StreamUsage]]:
    """Wrap a stream so usage is recorded even when it ends early."""
    state = {"fragments": 0, "usage": None, "completed": False}

    async def wrapped() -> AsyncIterator[StreamEvent]:
        try:
            async for event in events:
                if event.kind == "content":
                    state["fragments"] += 1
                elif event.kind == "usage":
                    state["usage"] = event.usage
                elif event.kind == "finish":
                    state["completed"] = True
                yield event
        finally:
            record_usage(current_usage())      # runs on cancellation too

    def current_usage() -> StreamUsage:
        reported = state["usage"]
        if reported is not None:
            return StreamUsage(
                reported.input_tokens, reported.output_tokens,
                is_estimated=False, completed=state["completed"],
            )
        return StreamUsage(
            input_tokens,
            estimate_output_tokens(state["fragments"]),
            is_estimated=True,
            completed=state["completed"],
        )

    return wrapped(), current_usage

The finally block is what makes this work under cancellation, which is exactly the disconnect case where usage would otherwise be lost.

Record is_estimated rather than hiding it. When reconciling against the provider's invoice, knowing which figures are measured and which are inferred is the difference between finding a discrepancy and being confused by one.

Tracing a stream as one logical span

A normal request is one event with a start, an end, and a duration. A stream is a long-lived operation with internal structure, and logging it as one line at the end loses everything interesting.

Log it as one span with several recorded points.

async def traced_stream(
    events: AsyncIterator[StreamEvent], *, request_id: str, model: str
) -> AsyncIterator[StreamEvent]:
    """Record one span covering the whole stream, with phase timings."""
    started = time.monotonic()
    first_token_at: float | None = None
    fragments = 0
    outcome = "unknown"

    try:
        async for event in events:
            if event.kind == "content":
                if first_token_at is None:
                    first_token_at = time.monotonic()
                fragments += 1
            elif event.kind == "finish":
                outcome = event.finish_reason or "stop"
            yield event
        if outcome == "unknown":
            outcome = "ended_without_finish"
    except asyncio.CancelledError:
        outcome = "client_disconnected"
        raise
    except ProviderError:
        outcome = "provider_error"
        raise
    finally:
        logger.info(
            "stream_completed",
            extra={
                "request_id": request_id,
                "model": model,
                "outcome": outcome,
                "ttft_ms": _ms(started, first_token_at),
                "total_ms": _ms(started, time.monotonic()),
                "fragments": fragments,
            },
        )

Four fields do most of the work. ttft_ms is the perceived speed metric from Lesson 3. total_ms is the throughput metric. outcome distinguishes a clean finish from a disconnect, a provider error, and a stream that simply stopped, and the proportions between those four tell you whether the system is healthy. fragments gives a rough size independent of token accounting.

Structured logging with an extra dictionary rather than a formatted message means these are queryable fields. The full treatment of structured logging, correlation identifiers, and tracing across async boundaries belongs to the observability module, and this is the shape it will build on.

[IMAGE PROMPT M8-7
Purpose: Show a stream as one logical span with internal timing points and possible outcomes, rather than a single end-of-request log line.
Visual type: Annotated span timeline with branching outcomes.
Prompt: A clean educational diagram showing one long horizontal bar labelled "one stream span" running left to right. Marked points along the bar from left: a start marker labelled "request sent", a marker roughly one tenth along labelled "first content chunk", a series of small tick marks across the middle labelled "content fragments", and an end marker. Beneath the bar, two brackets: a short one from start to the first content marker labelled "ttft_ms", and a full-width one labelled "total_ms". At the right end, four short branch arrows fan out to labelled outcome boxes reading "stop: finished normally", "length: hit token limit", "client_disconnected", and "provider_error". Beneath the outcome boxes, a note reads "usage arrives in the final chunk, and may never arrive". Above the bar, a small annotation reads "logged once, with fields, not one line per chunk".
Required elements: A single span bar with the stated internal markers, two labelled timing brackets, four branching outcome boxes, the usage note, the logging annotation.
Style: Clean educational illustration, professional, uncluttered, high contrast, flat vector.
Layout: Horizontal span across the upper half, brackets beneath it, outcomes fanning out at the right end.
Text labels: "one stream span", "request sent", "first content chunk", "content fragments", "ttft_ms", "total_ms", "stop: finished normally", "length: hit token limit", "client_disconnected", "provider_error", "usage arrives in the final chunk, and may never arrive", "logged once, with fields, not one line per chunk".
Aspect ratio: 16:9
Accessibility: Label every marker and outcome in text, and distinguish the two brackets by span length and label rather than colour.
Avoid: Vendor observability tool imagery, screenshots, decorative icons, tiny text, watermarks.
Alt text: Timeline of a single stream span showing request start, first content chunk, content fragments, brackets for time to first token and total duration, and four possible outcomes including client disconnect and provider error.
END IMAGE PROMPT]

Reconnection, resumption, and buffering

Reconnection. SSE clients reconnect automatically after a dropped connection, and browsers send a Last-Event-ID header carrying the id of the last event they received, if your events had ids.

Whether you can usefully resume is a harder question. Regenerating from scratch is wasteful and produces different text, since generation is not deterministic. Genuine resumption requires storing the completed portion server-side, keyed by a request identifier, so a reconnecting client can be sent what it missed and then continue.

That is real machinery, and it is worth building only when reconnection is common, such as on mobile networks. For most services the honest answer is to surface the failure and offer a retry rather than pretending the interruption did not happen.

Buffering for smooth rendering. Provider chunks are uneven, arriving in bursts and pauses. Rendering each fragment the instant it arrives produces visible stutter.

A small buffer that releases text at a steadier rate reads better, and the cost is a small addition to time-to-first-token. Keep the delay in the tens of milliseconds, because a buffer large enough to be noticeable has given back the benefit you streamed for.

This is a client concern rather than a server one. Your API should forward fragments as they arrive and let the client decide how to render them, since the client knows about frame rates and the server does not.

Concept check. Your streaming endpoint works in local development. In production, users report that responses appear all at once after fifteen seconds. Nothing in your code changed. What are the two most likely causes?

Answer

An intermediary buffering the response. A reverse proxy or CDN between your service and the user is collecting the whole response before forwarding it, which is invisible locally because nothing sits in the path. For nginx, the X-Accel-Buffering: no header addresses it, and a CDN may need the route excluded from buffering or caching entirely.

Something in your own code materialising the stream before yielding. A list comprehension over the upstream events, a join into a single string, or validating the complete response before forwarding all consume the entire stream first. This behaves identically to buffering and is worth checking before blaming infrastructure.

Both produce exactly the symptom described, and the way to tell them apart is to hit the service directly, bypassing the proxy. If it streams correctly there, the problem is in the path rather than the code.