Skip to content

refactor: replace stream_with_chunking with single-task stream() #1440

Closed

Description

@ajbozarth

Replace the two-task stream_with_chunking() + StreamChunkingResult with a single-task stream()/Streamer primitive consumed by a plain async for, and factor the inline chunk-boundary bookkeeping into a Chunker object. Chunking stays shared across validators in this issue; per-validator chunking and chunk-state ownership are a follow-up (#1441).

Builds on the POC in PR #1409.

Scope

Already prototyped (POC in PR #1409)

  • stream() / Streamer — single generator on the caller's task, no background orchestration task. Durable state is the thin Streamer handle (failed_early, failure_reason, streaming_failures, full_text, final_validations, terminal mot).
  • ModelOutputThunk.__aiter__ / __anext__ — wraps astream() in the async-iterator protocol so plain and validated streaming share one loop.
  • Event emission via STREAMING_EVENT hook — full Phase 1 vocabulary (QuickCheckEvent, ChunkEvent, StreamingDoneEvent, FullValidationEvent, ErrorEvent, CompletedEvent) fired uniformly on the single task.

New in this issue

  • Chunker — see design below. Factors the chunk-boundary state currently inline in the orchestrator into a reusable object.
  • Rename the Phase 1 chunking strategies — SentenceChunker → SentenceChunking, WordChunker → WordChunking, ParagraphChunker → ParagraphChunking (the ChunkingStrategy ABC keeps its name). This frees Chunker for the new stateful object above. Rationale: feat(stdlib): add ChunkingStrategy ABC and built-in chunkers #899 specified these as ...Chunking; they shipped as ...Chunker in feat(stdlib): add ChunkingStrategy ABC and built-in chunkers #923 with no recorded rationale (silent implementation choice). ...Chunker (an agent noun) reads as "the thing that chunks," which fits the new stateful driver better than a stateless strategy — so the strategies are restored to their ...Chunking spec and Chunker names the driver. Note: these names are exported (__all__) and used across docs/examples/streaming/, tests, and docs — this is a breaking rename of public names, with all call sites updated (consistent with this issue's hard break, below).
  • __aiter__ single-consumer guard — __aiter__/__anext__ make MOT look like a normal async iterator, but a second consumer splits the same stream rather than starting an independent one. Add a real guard: mark the thunk consumed on first iteration, raise if a second consumer starts. (Full parallel-consumer support — buffering, per-consumer cursors, late-joiner semantics — is out of scope unless a concrete need appears; the guard is enough for now.)
  • Events example — an example demonstrating an events plugin that consumes STREAMING_EVENT and surfaces streaming events for multiple parallel streams. The plugin may live in the example or in src depending on implementation. This is the primary user-facing deliverable for event consumption (hooks are the interface; the example is how users learn to consume them).
  • Hard break, no deprecation shim — stream() replaces stream_with_chunking() outright. A shim would keep the old orchestrator alive alongside the new one, reintroducing the complexity this issue removes; the API is new enough to break directly.

Constraints

This issue's chunking scope

Exactly one Chunker on the Streamer, fed by the single chunking= param; all validators validate the same chunks (as today). Per-validator chunking — each Requirement carrying its own strategy, and the resulting chunk-state ownership question — is #1441. The Chunker carries its own small state, so that follow-up relocates/multiplies it without an internal reshape.

Chunker class

Factors the chunk-boundary bookkeeping that currently lives inline in _orchestrate_streaming / _drive (accumulated scanning, prev_chunk_count, withheld-fragment handling) into a dedicated object. It is the stateful counterpart to a ChunkingStrategy: the strategy is stateless "how to split," the Chunker holds "how far through this stream we are."

Shape

  • __init__(strategy: ChunkingStrategy) — holds the (stateless) strategy.
  • feed(delta: str) -> list[str] — takes one stream delta, returns any new complete chunks, holding the trailing fragment between calls.
  • A flush/finalize method for the trailing fragment at stream end.

State division

State Owner Notes
accumulated (full raw text) Streamer one raw stream even when N validators later; needed for full_text regardless
pending fragment (text since last boundary) Chunker small (≤ one chunk); all the Chunker needs
full_text (validated-emitted output) Streamer validation concern, not chunking

The Chunker holds only the pending fragment, not the full accumulated text. The Streamer keeps accumulated for its own full_text needs and feeds deltas to the Chunker.

Implementation approach — reuse split(), don't reimplement boundaries

Chunker.feed() should call the existing stateless ChunkingStrategy.split() on the pending fragment, not reimplement boundary detection:

feed(delta):
    self._pending += delta
    chunks = self._strategy.split(self._pending)
    self._pending = <tail after last returned chunk>
    return chunks

split() already withholds the trailing fragment correctly, so feeding it the pending fragment reuses the exact boundary logic that already works. The only new logic is computing the tail carried forward.

Behavioral-change risk: chunking on the delta, not full accumulated

Today split() is called on the full accumulated every delta (O(n) re-scan). Moving to incremental delta input is a real behavioral change, not a pure refactor. Risk is confined to boundaries that span the delta seam:

  • Sentence — boundary is punctuation + whitespace; the . and the following space can arrive in different deltas.
  • Word — trailing partial word split across deltas.
  • Paragraph — \n{2,} is greedy; the newlines can split across deltas, and a prefix (\n\n) must not be emitted as a complete boundary before a possible third \n arrives.

Reusing split() on the accumulated pending fragment (rather than delta-native detection) sidesteps most of this, because split()'s withhold-the-tail contract already handles incomplete trailing boundaries.

Required test: delta-invariance — for any way a text is sliced into deltas, feeding the slices to Chunker yields the same chunks as one split(full_text). Fixtures: 1 delta, per-character deltas, random splits, plus the three seam cases above. The test/core/test_astream_mock.py mock gives deterministic delta control.

Activity

  1. self-assigned this
    on Jul 24, 2026
  2. added
    enhancementNew feature or request
    area/stdlibCore abstractions: Context, MOT, SamplingStrategy, formatters, serialization
    area/streamingStreaming chunks, events, per-chunk validation
    on Jul 24, 2026
  3. ajbozarth commented on Jul 31, 2026

    @ajbozarth
    MemberAuthor

    stream() lazy start — behavior change, needs sign-off

    Wiring the #1440 refactor in surfaced a trace-nesting bug: the generation span becomes the parent of the streaming span, so the parent closes before its child and the OTel context unwinds out of order — a broken trace. The clean fix requires changing when a stream starts, which is outside the original design scope, so flagging before proceeding.

    Today: await stream(...) immediately starts generation, fires generation_pre_call, and raises any setup error — all before the first async for.

    Proposed: stream() returns a ready handle but starts nothing. The first async for starts generation, fires the hook, and surfaces setup errors.

    Today Proposed
    Generation starts at await stream(...) at first async for
    generation_pre_call fires at await stream(...) at first async for
    Setup/backend errors raise from await stream(...) from first async for
    A stream created but never iterated still generates (incurs cost) never runs

    Why it's arguably better: "iterating starts the work" matches the single-task design; an unconsumed stream costs nothing; and it removes the lifecycle bug rather than masking it.

    Trade-off: callers can no longer start a stream and do other work while the model warms up — generation now waits for the loop. For the usual s = await stream(...); async for ... in s this is negligible.

    Ask: OK to make stream() lazy? Clean fix for the trace bug, better contract, but a behavior change beyond #1440's original scope.

  4. ajbozarth commented on Jul 31, 2026

    @ajbozarth
    MemberAuthor

    @jakelorocco @planetf1 are you ok with the above change? Given it was not part of my design pitch during the call I wanted to confirm before making the change. I'll also raise this during Monday scrum for a decision.

  5. jakelorocco commented on Aug 3, 2026

    @jakelorocco
    Contributor

    I do like the current behavior of immediately starting generation while allowing other work to be done elsewhere (which is how regular generation works now).

    I remember we had talked about reorganizing the calls / spans so that the streaming span could totally encompass the generation span. Is that not option?

  6. ajbozarth commented on Aug 3, 2026

    @ajbozarth
    MemberAuthor

    @jakelorocco The other option would be to mirror what I did with out current code, which is to move the start hook out of the iterator into the start of the function that creates the iterator. The issue with that is it would require a user to always async for ... in s otherwise the already open Otel span will never close, also if they decided to do anything between the s = await stream(...) and async for ... in s the span would already be open. IMHO this behavior though similar to how it is currently is not great.

  7. planetf1 commented on Aug 4, 2026

    @planetf1
    Contributor

    Thinking about it more — I don't think it has to be either/or. What both suggestions were missing was a close that actually happens, whether or not iteration finishes. If someone just breaks out of the loop early, a plain async generator won't reliably clean itself up - the leak concern. But if the stream is something you explicitly open and close (async with) rather than just loop over, we get both: it stays responsive, and the span closes properly regardless of how the loop ends. That also makes it the outer span wrapping generation - which Jake mentioned.

    Claude says (implementation notes, checked against current code):

    • stream() needs to become sync, with generation moved into __aenter__ — so streaming starts unambiguously at the point you enter the async with, not buried inside the generator. async with await stream(...) would otherwise reproduce the same ordering bug.
    • This restores behaviour main already has (STREAMING_START fires before generation starts) — feat: validated streaming refactor POC + Phase 2 epic breakdown #1409's POC regressed it by moving the start inside the generator.
    • Drop "guaranteed" — closes on every normal exit path (finish, break, error, cancellation), but wrap cleanup in the same except BaseException the current code uses, since cleanup awaits are themselves cancellable.
    • Reordering fixes the mis-parenting on 3.11 too — span creation doesn't take an explicit parent, so it always follows whatever's ambient at creation time regardless of Python version. What's 3.12+-only is the fuller nesting (generation properly nested under streaming); the mis-parenting fix itself applies everywhere.

    Follow-up issue needed to fix other incomplete spans — e.g. STREAMING_END today only fires from acomplete(), not from the orchestrator's finally in streaming.py, so an abandoned stream already leaks that span on main, independent of this refactor.

  8. added 2 commits that reference this issue on Aug 11, 2026
    24cc933
    1b7942c
  9. added a commit that references this issue on Aug 18, 2026
    058e3dd
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area/stdlibCore abstractions: Context, MOT, SamplingStrategy, formatters, serializationarea/streamingStreaming chunks, events, per-chunk validationenhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions