refactor(streaming)!: replace stream_with_chunking with single-task stream()

[ajbozarth]

Pull Request

Issue

Fixes #1440

Description

Replaces the two-task stream_with_chunking() + StreamChunkingResult with a single-task stream()/Streamer primitive consumed by a plain async for, and factors the inline chunk-boundary bookkeeping into a stateful Chunker. Builds on the POC in #1409.

stream() returns a Streamer driven by one async generator on the caller's task — there is no background orchestration task. It is consumed with async for, ideally inside async with, which provides a guaranteed cleanup contract (aclose()/__aexit__) that cancels an abandoned generation on an early break, an exception, or external cancellation (e.g. an asyncio.wait_for timeout) — the terminal events fire and the thunk is finalized on every path. Typed StreamEvents are emitted through the streaming_event plugin hook rather than a result.events() iterator, and acomplete() is removed.

Key changes:

• stream() / Streamer — single async generator on the caller's task; terminal state (failed_early, failure_reason, streaming_failures, full_text, final_validations, mot) lives on the thin Streamer handle.
• Chunker — new stateful driver that wraps a stateless ChunkingStrategy and holds only the pending fragment between deltas; delta-invariant (any slicing yields the same chunks as one split() over the whole text).
• Chunking strategy renames — SentenceChunker/WordChunker/ParagraphChunker → SentenceChunking/WordChunking/ParagraphChunking, freeing the Chunker name for the driver. String aliases (chunking="sentence", etc.) are unchanged.
• ModelOutputThunk iterator API — __aiter__/__anext__ wrap astream() in the async-iterator protocol with a single-consumer guard, plus aclose()/async with cleanup. astream() itself is unchanged.
• Telemetry — with the whole run on one task, the cross-task span-reattachment machinery and the STREAMING_ORCHESTRATION_START/_END hooks are removed; the span is renamed stream_with_chunking → stream; CompletedEvent moves to the streaming_event hook.
• Docs and examples — tutorials, how-to, tracing, and the streaming examples updated to the new API; added a multi-stream events example (multi_stream_events.py) and a raw MOT async-iterator example (async-iterator.py); added a migration guide at docs/dev/migrate-streaming-v0.8.md.

Breaking change, no deprecation shim: stream_with_chunking(), StreamChunkingResult, and the ...Chunker strategy names are removed. Migration guide: docs/dev/migrate-streaming-v0.8.md.

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Rewrote test/stdlib/test_streaming.py for the new API and added the Chunker delta-invariance suite to test/stdlib/test_chunking.py. Retargeted the telemetry span/metrics tests and the hook-call-site tests to the new topology (stream span roots the chat span; events via the streaming_event hook). Added the mocked integration twin for the streaming span-topology test so it is covered in the fast tier, not only behind the slow e2e path. Verified ruff, ruff format, and mypy clean; ran the tutorial example code against Ollama to confirm the documented sample output shapes.

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

• Component
• Requirement
• Sampling Strategy
• Tool

[AngeloDanducci]

Generally LGTM, a few review comments. I'll leave the larger design discussion to Jake since I think he's been more in the loop on the streaming refactor.

[AngeloDanducci]

IIUC there is a subtle change here - moving the completed event from streaming_end to the CompletedEvent changes full_text_length in cases of exceptions.

Previously it was payload.full_text_length = len(accumulated) which is now len(ev.full_text).

Not a bug per se but it looks like finish_streaming_span still uses payload.full_text_length = len(accumulated) so we may end up with two different lengths for what I think should be the same metric.

[planetf1]

Adding to this: it's actually two different attribute names, not one name with two values. tracing_plugins.py sets mellea.streaming.full_text_length (lines 374 and 401), while finish_streaming_span in tracing.py:841 sets mellea.full_text_length — no streaming. segment. So a dashboard or query keyed on one name won't pick up data recorded under the other at all, on top of the value mismatch you already flagged.

[ajbozarth]

actually full_text_length= len(accumulated) was the bug here, I fixed it in c6a130a so the value match now

I also updated the attr name to mellea.streaming.full_text_length along with all the attrs on that span

[AngeloDanducci]

Do we need this anymore after the refactor? I don't see it being used anywhere.

[ajbozarth]

nice catch, removed in c6a130a

[planetf1]

Nice simplification overall — dropping the cross-task span-reattachment machinery in particular removes some genuinely hairy code. Four things below are worth fixing before merge; none call the core design into question, they're all narrow and fixable.

[planetf1]

self._finalized flips to True here, before the try/finally that awaits self._mot.aclose() and then fires CompletedEvent/STREAMING_END. If a CancelledError lands during either of those two awaits inside the finally block, STREAMING_END never finishes firing — and since the guard is already True, nothing (not even aclose() calling _finalize() again) will retry it. That leaks the stream span (never popped from the in-flight registry, OTel token never detached) and the metrics plugins that key off STREAMING_END never run. I reproduced this with a slow streaming_event subscriber plus an external cancellation — STREAMING_END fired 0 times. The old code had an explicit shield against exactly this.

asyncio.shield() around the whole cleanup body fixes it — I tested this pattern in isolation and confirmed the shielded work keeps running to completion even though the caller sees CancelledError:

async def _finalize(self, *, success=False, error=None, full_text_length=0) -> None:
    if self._finalized:
        return
    self._finalized = True
    await asyncio.shield(self._do_finalize(success, error, full_text_length))

async def _do_finalize(self, success, error, full_text_length) -> None:
    try:
        await self._mot.aclose()
    finally:
        await _emit_event(...)
        if has_plugins(HookType.STREAMING_END):
            ...

[ajbozarth]

I had Claude draft a explanation for this after we put a simpler fix in c6a130a:

Good catch on the teardown window — fixed, though not with shield(). shield() runs its body in a new task, but the stream span's OTel context token is task-affine (finish_streaming_span → _detach_token must detach on the task that attached, per the start_application_span docstring). Shielding would detach on the wrong task and corrupt the context on the normal path — trading a rare leak for a guaranteed bug, and re-introducing the cross-task boundary this refactor removed.

Instead I nested a try/finally so STREAMING_END fires even if the CompletedEvent emission is cancelled or raises, all on the caller's task — the token detaches correctly and the CancelledError still propagates, so an asyncio.wait_for timeout still surfaces as TimeoutError. The single-cancel path (the one the PR advertises) was already clean; this closes the double-cancel-mid-teardown window you reproduced with the slow subscriber.

On the raising-subscriber variant: the STREAMING_END subscribers are ours and mode-isolated (StreamMetricsPlugin is FIRE_AND_FORGET, so cpex never propagates its errors; the span-closer is our own code), and a third-party streaming_event hook that raises is the subscriber's own bug — I'd rather it surface than be silently swallowed. So no blanket except BaseException like the old orchestrator had.

[planetf1]

Thanks — the nested finally now reaches STREAMING_END if cancellation interrupts _mot.aclose() or CompletedEvent emission. However, _finalized is set before teardown awaits begin. If cancellation arrives while invoke_hook(STREAMING_END, ...) is in flight, terminal dispatch is interrupted; later aclose() or _finalize() calls return without retrying it.

Reproduced against _finalize() with a blocking terminal dispatcher. Existing cancellation tests cover only a single cancellation and do not exercise cancellation during terminal-hook delivery.

Suggest adding a regression test and ensuring STREAMING_END completes reliably under cancellation during teardown.

[planetf1]

Two more from a second pass, plus one I've folded into the existing thread on tracing_plugins.py:395 as a reply rather than a new comment since it's building on the same spot.

[jakelorocco]

Absent my other comments about the multi-chat streaming, I think it's unclear what the chunk actually contains here. Could you please add either an in docstring example or say that it returns string segments of the model output thunk?

[ajbozarth]

I added a clarifying sentence in c6a130a

[jakelorocco]

I did some tests with this multi stream events example, and I'm worried that it shows our current approach is actually difficult to utilize.

A few thoughts:

1. It is quite difficult to correlate a stream, the events, and the place those events need to go. For instance, I could get the stream id from the streamer and give it to some tui consumer. But I would have to have the plugin route to some queue / map that the tui consumer can grab events with that stream id from.
2. What is the preferred pattern if a user is creating and constantly streaming new mots? I feel like that pattern would either require multiple plugins that basically do the same thing (which might be a performance issue); and/or a complicated pattern to route the events to the correct consumer.
3. Is it possible to just have the streamer return events as well as mot chunks (or a different function of the streamer that can return the events) allowing the user to choose what to consume?

Additionally, I wonder if we could just register handlers directly to a streamer instead of having to go through the plugin system. Or maybe there's a way to utilize the plugin system to do something similar? After trying to utilize the plugin system to do some things beyond the example below, I'm just worried that it's quite convoluted to handle events from streamers.

[ajbozarth]

I actually went back and forth on how this should work a few items before settling on this implementation. I am open to further design discussions around it at or after Monday scrum.

If I'm understanding you correctly though this isn't an issue with the current implementation as it stands but the design we chose to implement. If that's the case we could choose to deal with this in a follow up PR before the next release. If you disagree and want to block merge of this PR on that design discussion we can do that too.

Even if we redesign, these hooks would need to stick around for the telemetry at least.

[jakelorocco]

Are we ensuring anywhere that the final mot == what was actually received in chunks from the model even if the chunker edits things during streaming?

I also think that if split drops whitespace, isn't it possible for .find to fail here, messing up _pending?

[ajbozarth]

I dug into this and part of it was already addressed in a fix suggested by @planetf1 above, I'll include Claudes explanation:

Good instinct — the mutation-via-find() case is now a hard error: Chunker.feed() raises ValueError if split() returns a chunk that isn't a verbatim substring of the buffered text, so a normalizing/rewriting strategy fails loudly instead of silently re-emitting. On the mot question: the final mot can't diverge from what the model produced, because mot.value is accumulated from the raw deltas directly and the chunker never writes back to it — chunking only shapes the consumer-facing async for stream. The only place chunk text feeds back into state is full_text on early exit (accumulated[:emitted_end], located by find()), and that's covered by the same substring guard.

[jakelorocco]

Should this get similar handling as mots to prevent multi-consumers?

[ajbozarth]

I actually asked this myself during my own self review, I'll let Claude explain:

Streamer doesn't need the MOT-style guard, because the failure mode it protects against can't occur here. MOT's __aiter__ returns self and re-arms, so a second async for would re-drive and split the stream — hence its explicit guard. Streamer's __aiter__ returns the same async generator object (self._gen), which is single-consumer by construction: once the first async for exhausts it, a second one just gets StopAsyncIteration immediately and yields nothing — no split, no re-run, full_text/mot intact (verified). So a second iteration is inert rather than dangerous. It's silent rather than loud (you get [], not an error); happy to add an explicit raise if you'd prefer the louder contract, but there's no correctness risk either way.

[ajbozarth]

@AngeloDanducci @planetf1 @jakelorocco I believe I had addressed all your review (except the open design point @jakelorocco raise) and left response on every comment. If you could re-review and mark them as resolved.

[jakelorocco]

lgtm; a few minor issues and discussed verbally about next steps

[jakelorocco]

Should we also enforce that the chunker always advances? or if not, that it doesn't produce empty chunks? This doesn't seem to impact our chunking strategies but maybe we should include this check to ensure third party chunkers don't fail in these modes?

[jakelorocco]

I think this marker should be elsewhere / deleted, not on the class?

[jakelorocco]

Can you please add detail here that completed_normally doesn't prevent the final validation from raising an exception, etc...?