Skip to content

feat(stream)!: align SseStream and StreamingToolCallTracker with the AI SDK - #201

Closed
cunninghamcard-bit wants to merge 9 commits into
arcships:masterfrom
cunninghamcard-bit:stream/align-ai-sdk
Closed

feat(stream)!: align SseStream and StreamingToolCallTracker with the AI SDK#201
cunninghamcard-bit wants to merge 9 commits into
arcships:masterfrom
cunninghamcard-bit:stream/align-ai-sdk

Conversation

@cunninghamcard-bit

@cunninghamcard-bit cunninghamcard-bit commented Sep 29, 2026

Copy link
Copy Markdown
Contributor

Summary

Aligns aimux-stream with the AI SDK's current implementations. The old code was a port of an older, index-only tracker and a hand-rolled SSE parser; nothing in the workspace used the tracker or NdjsonStream, and the only consumer of SseStream is aimux-provider-utils/src/response_handler.rs.

  • StreamingToolCallTracker is rewritten to match @ai-sdk/provider-utils (5.0.51): deltas are correlated by wire id, index and function name, ambiguous deltas are dropped, ids are de-duplicated with bounded suffixes, blank function names are ignored, and flush orders calls by index only when every call has one. Adds StreamingToolCallArgumentState / starts_with_structured_value.
  • SseStream is rewritten as a line-based parser matching eventsource-parser (3.1.1), the parser behind parseJsonEventStream: \n / \r / \r\n in any mix, BOM stripping, field-parsing rules, id with U+0000 ignored, retry digits only. Exceeding max_event_size now yields SseError::FrameTooLarge and ends the stream instead of skipping the frame.
  • NdjsonStream / NdjsonError are removed (unused, no AI SDK counterpart). tokio becomes a dev-dependency and the unused direct serde dependency is dropped.

Full API changes are in the [Unreleased] section of CHANGELOG.md.

Type of change

  • Bug fix (non-breaking)
  • New feature (non-breaking)
  • Breaking change
  • New provider
  • New / updated binding
  • Docs / RFC

Checklist

  • cargo fmt --all -- --check passes
  • cargo clippy --workspace --all-targets -- -D warnings passes
  • cargo test --workspace passes (171 test binaries, 3856 tests, 0 failed)
  • Added or updated cassette tests / contract fixtures for behavior changes — no cassette changes; the tests are ports of the upstream suites (see below)
  • Updated relevant docs (CHANGELOG.md, README.md, CONTRIBUTING.md, docs/PROJECT-OVERVIEW.md)
  • Commit messages follow Conventional Commits

Notes for reviewers

Tests are ported, not written. streaming_tool_call_tracker_test.rs and the argument-state test are case-for-case ports of the upstream tests at the pinned tag (@ai-sdk/provider-utils@5.0.51). sse_test.rs is a port of eventsource-parser v3.1.1's parse.test.ts and stream.test.ts (MIT), driven by the same fixtures; test/multibyte.ts is checked in as JSON. Test names follow the upstream titles. What could not be ported (onComment counts, reset(), onError payloads) and the adaptations (retry is reported on the dispatched event; maxBufferSize maps to with_max_event_size) are listed in the header of sse_test.rs. Four aimux-only tests remain: strict UTF-8, byte-level chunk splitting, transport errors, and the default limit. sha2 (already used by aimux-providers) is a dev-dependency for the upstream SHA-256 check.

Parity was checked by running upstream, not by reading it. The same delta sequences were fed to the upstream JS tracker and the Rust tracker (identical output for the cases tried), and a few SSE inputs to eventsource-parser. Some behaviours that look odd are upstream's own, for example an id-less continuation with a reused index and name is appended to the existing call. They are kept on purpose.

Known, intentional differences from upstream

  • Invalid UTF-8 is a strict SseError::Utf8 that discards the rest of that event; upstream decodes lossily.
  • The default max_event_size is 1 MiB; upstream's default is unbounded.
  • Rust trim() and JS trim() differ on U+FEFF and U+0085; not reachable with real wire data, left as is.
  • The default tool-call id generator is the constant tool-call (as before this change), not a random id.

Consumer impact: response_handler.rs matches on SseError::Stream and turns every other error into JsonParse. A FrameTooLarge now ends the stream after that one error, where it used to skip the frame and continue.

The first 5 commits are the rewrite; the last two are a clippy fix (collapsible_match fails the CI -D warnings gate on stable) and the test port.

🤖 Generated with Claude Code

claude and others added 9 commits September 29, 2026 04:54
@claude
refactor(stream)!: remove unused StreamingToolCallTracker
7b0fbef 
No crate in the workspace referenced the tracker; providers accumulate
streamed tool-call deltas themselves. Drops the module, its 787-line test
file and the re-exported companion types. Breaking for aimux-stream's
public Rust API; noted in CHANGELOG under Unreleased.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe
@claude
feat(stream)!: align StreamingToolCallTracker with the current AI SDK
d72cfc3 
Re-adds the tracker as a port of the current
@ai-sdk/provider-utils streaming-tool-call-tracker.ts instead of the
older index-only version removed in the previous commit.

- Correlate deltas by wire id, index and function name; drop ambiguous ones
- Unique tool-call ids with bounded suffixes; blank generator output falls
  back to "tool-call"
- Ignore unmatched blank function names; retain blank-name continuations
- Flush in index order only when every call has an index
- New StreamingToolCallArgumentState / starts_with_structured_value port
- process_delta and flush return the emitted parts instead of buffering

Tests are ported case for case from the upstream suites (45 tracker cases,
9 argument-state cases). No caller in the workspace uses the tracker yet.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe
@claude
feat(stream)!: align SseStream with eventsource-parser
b17aa5d 
The AI SDK's parseJsonEventStream pipes through eventsource-parser. Replace
the blank-line frame splitter with its line-based state machine:

- \n, \r and \r\n line endings in any mix; a trailing \r waits for a
  possible \n in the next chunk
- strip a UTF-8 BOM at the start of the stream
- a field line without ':' has an empty value; empty event -> None;
  id containing U+0000 ignored; retry accepted only when all ASCII digits
- exceeding the buffer limit (event data + partial line) is fatal and ends
  the stream, as maxBufferSize does upstream

Kept from aimux: strict UTF-8 decoding after reassembly (an invalid line
now drops only the event being built) and a default 1 MiB bound.

extract_lines already matched upstream; the [DONE]-skipping JSON layer in
aimux-provider-utils already matches parseJsonEventStream.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe
@claude
refactor(stream)!: remove unused NdjsonStream
4e06914 
Nothing in the workspace used NdjsonStream / NdjsonError and the AI SDK's
provider-utils has no counterpart. Also move tokio to dev-dependencies (only
the tests use it), drop the unused direct serde dependency, and update the
crate description in README / PROJECT-OVERVIEW.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe
@claude
fix(stream): discard the rest of an SSE event after a UTF-8 error
eef5195 
A line that failed UTF-8 decoding reset the event, but the lines after it
(up to the blank line) were then dispatched as a separate fragment event,
producing a second downstream error. Mark the event poisoned until the next
blank line so one bad event yields exactly one SseError::Utf8. Also update
the aimux-stream description in CONTRIBUTING.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe
chenhaonan @claude
fix(stream): collapse the SSE retry arm into a match guard
0644161 
clippy::collapsible_match rejects the nested `if` under the CI
`-D warnings` gate (stable 1.98); behaviour is unchanged.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
chenhaonan @claude
test(stream): replace hand-written SSE tests with eventsource-parser'…
d47746f 
…s suite

sse_test.rs was a behaviour list written from scratch, and its header
claimed the AI SDK has no SSE parser tests. The parser behind
parseJsonEventStream, eventsource-parser, ships its own suite. Port
test/parse.test.ts and test/stream.test.ts (v3.1.1, MIT) case for case,
driven by the same fixtures; test/multibyte.ts is checked in as JSON
and read with include_str!.

Cases with no Rust counterpart (onComment counts, reset(), onError
payloads) are listed in the file header, as are the adaptations
(retry is reported on the event, maxBufferSize maps to
with_max_event_size). Four aimux-only tests remain: strict UTF-8,
byte-level chunk splitting, transport errors, and the default limit.

The four unit tests inside sse.rs are dropped; the port covers them.
sha2 (already used by aimux-providers) is a dev-dependency for the
upstream SHA-256 check on the 4.8M-character event.

The tracker and argument-state tests were already case-for-case ports
of the pinned upstream tests; only a comment is added.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
@cunninghamcard-bit @claude
feat(stream)!: decode SSE with the sse-stream crate
20b2b50 
Replace aimux-stream's hand-written SSE parser with a thin adapter over
sse-stream 0.3 (WHATWG event-stream parsing), the way the AI SDK's
parseJsonEventStream wraps eventsource-parser. The adapter keeps the AI SDK
dispatch rules: only blocks with a data line are dispatched, an empty
`event:` is None, and `retry` rides on its event.

- No event size limit any more, as upstream: `with_max_event_size` and
  `SseError::FrameTooLarge` are removed. Real recordings of OpenAI Responses
  and Gemini image generation carry 2-3 MB events; the 1 MiB cap dropped
  them, and after the previous commit ended the stream on the first one.
  Both recordings now decode completely (15/15 and 5/5 events).
- `SseError::Stream` carries the transport error as its source instead of a
  String; `Utf8` holds a `str::Utf8Error`; `Decode` covers the rest. Decoder
  errors end the stream after being reported.
- The ported eventsource-parser suite runs unchanged against the adapter
  (29 cases); the 6 `maxBufferSize` cases go with the limit.
- provider-utils: the `SseError::Stream` arm formats the source error.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WGwanvo9WLU9HWD7sRR8TS
@cunninghamcard-bit @claude
feat(provider-utils)!: move the tool-call tracker next to its users a…
ff2517c 
…nd wire the OpenAI stream to it

The AI SDK keeps StreamingToolCallTracker in @ai-sdk/provider-utils, where
every provider that speaks the OpenAI chat-completions wire format imports
it. aimux had the tracker in aimux-stream (the eventsource-parser role),
with its own ToolCallStreamPart event type and no callers, while
openai/model.rs kept an index-only accumulator. This commit puts the
pieces where the AI SDK has them.

aimux-stream
- StreamingToolCallTracker, StreamingToolCallArgumentState,
  starts_with_structured_value and ToolCallStreamPart move out; the crate
  is SSE decoding only. serde_json becomes a dev-dependency (test
  fixture); the `tool-calls` keyword is dropped.

aimux-provider-utils
- streaming_tool_call_tracker.rs / streaming_tool_call_argument_state.rs
  land here, logic unchanged. The tracker emits aimux_core::StreamPart
  ToolInputStart / ToolInputDelta / ToolInputEnd / ToolCall directly, so
  there is no second event type (RFC-0036: ToolCallStreamPart must not be a
  second public protocol). The generic metadata parameter becomes
  serde_json::Value / ProviderMetadata; TrackerError converts into
  AiMuxError::InvalidResponseData. The ported upstream tests (45 + 9)
  move with it, projecting StreamPart back to the upstream event shape.

aimux-providers
- openai/model.rs replaces its HashMap<usize, ToolCallAccumulator> with the
  tracker: tool_calls deltas are correlated by wire id, index and function
  name instead of index alone; id-less continuations follow their call;
  indices reused across parallel calls stay distinct; ambiguous deltas are
  dropped; a call without a wire id gets a generated `tool-call` /
  `tool-call-N` id instead of an empty string; a new call without a
  function name ends the stream with InvalidResponseData (AI SDK
  behaviour) instead of starting a call with an empty name.
- DeltaToolCall.index is Option<usize> (the AI SDK schema is
  `index: z.number().nullish()`).

Every registry-backed provider goes through openai/model.rs, so this is
the one plug point until S4-2 folds the mistral and xai chat copies into
the same implementation.

Verification: cargo test -p aimux-stream -p aimux-provider-utils; cargo test
-p aimux-providers (129 suites, 3041 passed, 0 failed, including the
cassette replays); cargo clippy -p aimux-stream -p aimux-provider-utils
-p aimux-providers --all-targets -- -D warnings; cargo fmt --all -- --check.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WGwanvo9WLU9HWD7sRR8TS
@cunninghamcard-bit cunninghamcard-bit mentioned this pull request Oct 2, 2026
Draft
7 of 12 tasks
@cunninghamcard-bit

cunninghamcard-bit commented Oct 2, 2026

Copy link
Copy Markdown
Contributor Author

Superseded by #204, rebased onto #203 for stacking

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@cunninghamcard-bit @claude