From fe8758f0556eb7efcffcc4b2c186b8ac5581ad4f Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 29 Sep 2026 04:54:24 +0000 Subject: [PATCH 1/6] refactor(stream)!: remove unused StreamingToolCallTracker 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 Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe --- aimux-stream/src/lib.rs | 5 - .../src/streaming_tool_call_tracker.rs | 429 ---------- .../tests/streaming_tool_call_tracker_test.rs | 787 ------------------ 3 files changed, 1221 deletions(-) delete mode 100644 aimux-stream/src/streaming_tool_call_tracker.rs delete mode 100644 aimux-stream/tests/streaming_tool_call_tracker_test.rs diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index e023a3204..65c368daf 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -6,12 +6,7 @@ pub mod lines; pub mod sse; -pub mod streaming_tool_call_tracker; // Re-export the most commonly used items. pub use lines::extract_lines; pub use sse::{SseError, SseEvent, SseStream}; -pub use streaming_tool_call_tracker::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, - ToolCallStreamPart, TrackerError, TypeValidation, -}; diff --git a/aimux-stream/src/streaming_tool_call_tracker.rs b/aimux-stream/src/streaming_tool_call_tracker.rs deleted file mode 100644 index be9ad8907..000000000 --- a/aimux-stream/src/streaming_tool_call_tracker.rs +++ /dev/null @@ -1,429 +0,0 @@ -//! Streaming tool call tracker. -//! -//! Rust translation of `@ai-sdk/provider-utils`'s `StreamingToolCallTracker` -//! (`packages/provider-utils/src/streaming-tool-call-tracker.ts`). -//! -//! Tracks streaming tool call state across multiple deltas from an -//! OpenAI-compatible chat completion stream: accumulates `arguments` string -//! fragments by `index`, emits `tool-input-start` / `tool-input-delta` / -//! `tool-input-end` / `tool-call` events, and finalizes any unfinished tool -//! calls on [`StreamingToolCallTracker::flush`]. -//! -//! Like the TS original, a tool call is *never* finalized before `flush` — a -//! parsable argument buffer can still be the prefix of a longer argument -//! string, so acting on it early would use truncated inputs (ai-sdk #13137). - -use serde_json::Value; -use thiserror::Error; - -/// Default upper bound accepted for a tool call `index`. -const DEFAULT_MAX_INDEX: usize = 1024; - -/// The `function` sub-object of a streaming tool call delta. -#[derive(Debug, Clone, Default)] -pub struct StreamingToolCallFunction { - pub name: Option, - pub arguments: Option, -} - -/// A streaming tool call delta — the `tool_calls[i]` entry of an OpenAI-style -/// streaming chunk. -/// -/// Use the builder methods ([`StreamingToolCallDelta::index`], etc.) to -/// construct one ergonomically. `arguments: null` (TS) maps to `None`; an -/// empty-string `arguments: ''` maps to `Some("")`. -#[derive(Debug, Clone, Default)] -pub struct StreamingToolCallDelta { - pub index: Option, - pub id: Option, - /// The `type` field. Named `r#type` because `type` is a reserved word. - pub r#type: Option, - pub function: Option, - /// Provider-specific payload carried alongside the standard fields. Used by - /// `extract_metadata` to pull out provider metadata (e.g. a Google thought - /// signature). Defaults to [`Value::Null`]. - pub extra: Value, -} - -impl StreamingToolCallDelta { - #[must_use] - pub fn new() -> Self { - Self::default() - } - - #[must_use] - pub fn index(mut self, index: usize) -> Self { - self.index = Some(index); - self - } - - #[must_use] - pub fn id(mut self, id: impl Into) -> Self { - self.id = Some(id.into()); - self - } - - /// Set the `type` field (named `tool_type` because `type` is reserved). - #[must_use] - pub fn tool_type(mut self, t: impl Into) -> Self { - self.r#type = Some(t.into()); - self - } - - #[must_use] - pub fn function_name(mut self, name: impl Into) -> Self { - self.function.get_or_insert_with(Default::default).name = Some(name.into()); - self - } - - /// Set the `function.arguments` fragment. Pass `""` for an explicit empty - /// fragment; omit the call entirely for `None` (TS `arguments: null`). - #[must_use] - pub fn arguments(mut self, args: impl Into) -> Self { - self.function.get_or_insert_with(Default::default).arguments = Some(args.into()); - self - } - - #[must_use] - pub fn extra(mut self, extra: Value) -> Self { - self.extra = extra; - self - } -} - -/// The stream parts emitted by [`StreamingToolCallTracker`]. -/// -/// Mirrors the subset of `LanguageModelV4StreamPart` the TS tracker enqueues. -/// Note that [`ToolCallStreamPart::ToolCall::input`] is the raw accumulated -/// argument *string* — the tracker does not parse it (matching the TS -/// behavior). -#[derive(Debug, Clone, PartialEq)] -pub enum ToolCallStreamPart { - /// Start of a tool call's input streaming. - ToolInputStart { id: String, tool_name: String }, - /// A delta of tool call input (a partial argument fragment). - ToolInputDelta { id: String, delta: String }, - /// End of a tool call's input streaming. - ToolInputEnd { id: String }, - /// A complete, finalized tool call. - ToolCall { - tool_call_id: String, - tool_name: String, - input: String, - provider_metadata: Option, - }, -} - -/// How to validate the `type` field on a new tool call delta. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum TypeValidation { - /// No validation (default). - #[default] - None, - /// Throw if `type` is present and not `"function"`. - IfPresent, - /// Throw if `type` is not exactly `"function"`. - Required, -} - -/// Errors raised while processing a tool call delta. -#[derive(Debug, Error, PartialEq)] -pub enum TrackerError { - #[error("Expected 'id' to be a string.")] - MissingId, - #[error("Expected 'function.name' to be a string.")] - MissingFunctionName, - #[error("Expected 'function' type.")] - InvalidType, - #[error("Tool call index out of range")] - IndexOutOfRange, -} - -struct TrackedToolCall { - id: String, - function_name: String, - arguments: String, - has_finished: bool, - metadata: Option, -} - -/// Extract provider metadata from a delta (the TS `extractMetadata` option). -type ExtractMetadataFn = Box Option>; -/// Build the `providerMetadata` for a finalized tool call (the TS -/// `buildToolCallProviderMetadata` option). -type BuildMetadataFn = Box) -> Option>; - -/// Tracks streaming tool call state across multiple deltas from an -/// OpenAI-compatible chat completion stream. -/// -/// Emitted [`ToolCallStreamPart`]s accumulate in an internal buffer; inspect -/// them with [`parts`](Self::parts) and reset between checks with -/// [`clear_parts`](Self::clear_parts) (the TS test's `parts.length = 0`). -/// -/// The type parameter `M` is the provider-metadata type. Use `()` (the -/// default) when no metadata handling is needed, or `serde_json::Value` (or -/// any `SharedV4ProviderMetadata`-like type) together with -/// [`with_extract_metadata`](Self::with_extract_metadata) / -/// [`with_build_provider_metadata`](Self::with_build_provider_metadata). -pub struct StreamingToolCallTracker { - tool_calls: Vec>>, - parts: Vec>, - // The TS `generateId` option: fallback for `toolCall.id ?? generateId()` - // when an incoming tool-call delta omits its id. - generate_id: Box String>, - type_validation: TypeValidation, - /// Upper bound accepted for a tool call `index`. Guards against a remote - /// index resizing `tool_calls` to a huge vector. Defaults to - /// [`DEFAULT_MAX_INDEX`] (1024). - max_index: usize, - extract_metadata: Option>, - build_provider_metadata: Option>, -} - -impl Default for StreamingToolCallTracker<()> { - fn default() -> Self { - Self::new() - } -} - -impl StreamingToolCallTracker { - /// Create a new tracker with no metadata handling and default settings. - #[must_use] - pub fn new() -> Self { - Self { - tool_calls: Vec::new(), - parts: Vec::new(), - generate_id: Box::new(|| "generated".to_string()), - type_validation: TypeValidation::None, - max_index: DEFAULT_MAX_INDEX, - extract_metadata: None, - build_provider_metadata: None, - } - } - - /// Set a custom id generator (the TS `generateId` option). - #[must_use] - pub fn with_generate_id String + 'static>(mut self, f: F) -> Self { - self.generate_id = Box::new(f); - self - } - - /// Set the `type` validation mode (the TS `typeValidation` option). - #[must_use] - pub fn with_type_validation(mut self, v: TypeValidation) -> Self { - self.type_validation = v; - self - } - - /// Set the maximum accepted tool call `index` (defaults to 1024). A delta - /// whose resolved `index` exceeds `max_index` returns - /// [`TrackerError::IndexOutOfRange`] instead of resizing `tool_calls` to - /// `index + 1` slots. - #[must_use] - pub fn with_max_index(mut self, max_index: usize) -> Self { - self.max_index = max_index; - self - } - - /// Set the metadata extractor (the TS `extractMetadata` option). Called - /// once when a new tool call is detected; the returned metadata is stored - /// on the tool call and passed to the builder at finalization. - #[must_use] - pub fn with_extract_metadata Option + 'static>( - mut self, - f: F, - ) -> Self { - self.extract_metadata = Some(Box::new(f)); - self - } - - /// Set the provider-metadata builder (the TS `buildToolCallProviderMetadata` - /// option). Receives the metadata previously extracted; if `None` is - /// returned, no `provider_metadata` is included in the `tool-call` event. - #[must_use] - pub fn with_build_provider_metadata) -> Option + 'static>( - mut self, - f: F, - ) -> Self { - self.build_provider_metadata = Some(Box::new(f)); - self - } - - /// Process a tool call delta from a streaming response chunk. Emits events - /// into the internal buffer. - /// - /// # Errors - /// - /// Returns `TrackerError::IndexOutOfRange` when the delta's tool index - /// exceeds the configured maximum. - pub fn process_delta(&mut self, delta: &StreamingToolCallDelta) -> Result<(), TrackerError> { - let index = delta.index.unwrap_or(self.tool_calls.len()); - // Guard against a remote `index` resizing `tool_calls` to a huge vector. - if index > self.max_index { - return Err(TrackerError::IndexOutOfRange); - } - let is_new = self - .tool_calls - .get(index) - .map(std::option::Option::is_none) - .unwrap_or(true); - if is_new { - self.process_new_tool_call(index, delta)?; - } else { - self.process_existing_tool_call(index, delta); - } - Ok(()) - } - - /// Finalize any unfinished tool calls. Should be called during the - /// stream's flush to ensure all tool calls are properly completed. Emits - /// `tool-input-end` and `tool-call` for each unfinished tool call. - pub fn flush(&mut self) { - for i in 0..self.tool_calls.len() { - let needs_finish = self - .tool_calls - .get(i) - .and_then(|o| o.as_ref()) - .is_some_and(|tc| !tc.has_finished); - if needs_finish { - self.finish_tool_call_at(i); - } - } - } - - /// The events emitted so far (accumulated across `process_delta`/`flush`). - #[must_use] - pub fn parts(&self) -> &[ToolCallStreamPart] { - &self.parts - } - - /// Clear the accumulated events (the TS test's `parts.length = 0`). - pub fn clear_parts(&mut self) { - self.parts.clear(); - } - - fn process_new_tool_call( - &mut self, - index: usize, - delta: &StreamingToolCallDelta, - ) -> Result<(), TrackerError> { - match self.type_validation { - TypeValidation::Required => { - if delta.r#type.as_deref() != Some("function") { - return Err(TrackerError::InvalidType); - } - } - TypeValidation::IfPresent => { - if delta.r#type.as_deref().is_some_and(|t| t != "function") { - return Err(TrackerError::InvalidType); - } - } - TypeValidation::None => {} - } - - let id = delta.id.clone().unwrap_or_else(|| (self.generate_id)()); - let function_name = delta - .function - .as_ref() - .and_then(|f| f.name.clone()) - .ok_or(TrackerError::MissingFunctionName)?; - - self.parts.push(ToolCallStreamPart::ToolInputStart { - id: id.clone(), - tool_name: function_name.clone(), - }); - - let metadata = self - .extract_metadata - .as_ref() - .and_then(|extract| extract(delta)); - - // TS: `toolCallDelta.function.arguments ?? ''`. - let arguments = delta - .function - .as_ref() - .and_then(|f| f.arguments.clone()) - .unwrap_or_default(); - - if index >= self.tool_calls.len() { - self.tool_calls.resize_with(index + 1, || None); - } - self.tool_calls[index] = Some(TrackedToolCall { - id: id.clone(), - function_name: function_name.clone(), - arguments: arguments.clone(), - has_finished: false, - metadata, - }); - - // Emit initial delta if arguments already present. - if !arguments.is_empty() { - self.parts.push(ToolCallStreamPart::ToolInputDelta { - id: id.clone(), - delta: arguments, - }); - } - - // Tool calls must not finalize before the stream ends (see #13137). - Ok(()) - } - - fn process_existing_tool_call(&mut self, index: usize, delta: &StreamingToolCallDelta) { - // TS: `toolCallDelta.function?.arguments != null`. - let new_args = match delta.function.as_ref().and_then(|f| f.arguments.as_ref()) { - Some(args) => args.clone(), - None => return, - }; - let id = { - let Some(Some(tool_call)) = self.tool_calls.get_mut(index) else { - return; - }; - if tool_call.has_finished { - return; - } - tool_call.arguments.push_str(&new_args); - tool_call.id.clone() - }; - self.parts.push(ToolCallStreamPart::ToolInputDelta { - id, - delta: new_args, - }); - } - - fn finish_tool_call_at(&mut self, index: usize) { - // Scope the mutable borrow of the tracked tool call, extracting the - // owned data we need to emit the final events. - let (id, function_name, arguments, metadata) = { - let Some(Some(tool_call)) = self.tool_calls.get_mut(index) else { - return; - }; - if tool_call.has_finished { - return; - } - ( - tool_call.id.clone(), - tool_call.function_name.clone(), - tool_call.arguments.clone(), - tool_call.metadata.take(), - ) - }; - if let Some(Some(tool_call)) = self.tool_calls.get_mut(index) { - tool_call.has_finished = true; - } - - self.parts - .push(ToolCallStreamPart::ToolInputEnd { id: id.clone() }); - - let provider_metadata = self - .build_provider_metadata - .as_ref() - .and_then(|build| build(metadata.as_ref())); - - self.parts.push(ToolCallStreamPart::ToolCall { - tool_call_id: id, - tool_name: function_name, - input: arguments, - provider_metadata, - }); - } -} diff --git a/aimux-stream/tests/streaming_tool_call_tracker_test.rs b/aimux-stream/tests/streaming_tool_call_tracker_test.rs deleted file mode 100644 index 25a9b7655..000000000 --- a/aimux-stream/tests/streaming_tool_call_tracker_test.rs +++ /dev/null @@ -1,787 +0,0 @@ -//! Rust translation of -//! `packages/provider-utils/src/streaming-tool-call-tracker.test.ts` (18 cases). -//! -//! The TS test collects emitted parts on a shared `controller.enqueue` sink and -//! resets it with `parts.length = 0` between assertions. The Rust equivalent -//! accumulates parts inside the tracker: inspect with `parts()` and reset with -//! `clear_parts()`. - -use aimux_stream::{ - StreamingToolCallDelta, StreamingToolCallTracker, ToolCallStreamPart, TrackerError, - TypeValidation, -}; -use serde_json::{Value, json}; - -/// Convenience builder alias. -fn d() -> StreamingToolCallDelta { - StreamingToolCallDelta::default() -} - -/// A tracker with no metadata handling (`M = ()`), matching the TS tests that -/// don't exercise the metadata hooks. Pinning `M = ()` here lets type -/// inference resolve `ToolCallStreamPart::ToolCall { provider_metadata: None }` -/// in the assertions. -fn new_tracker() -> StreamingToolCallTracker<()> { - StreamingToolCallTracker::<()>::new() -} - -fn tool_call_id(part: &ToolCallStreamPart) -> &str { - match part { - ToolCallStreamPart::ToolCall { tool_call_id, .. } => tool_call_id, - _ => panic!("expected tool-call part"), - } -} - -// == processDelta ========================================================== - -#[test] -fn single_tool_call_accumulated_across_multiple_deltas() { - // TS: "should handle a single tool call accumulated across multiple deltas" - let mut tracker = new_tracker(); - - // First delta: new tool call with id and name. - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("get_weather") - .arguments("{\"ci"), - ) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputStart { - id: "call_1".into(), - tool_name: "get_weather".into(), - }, - ToolCallStreamPart::ToolInputDelta { - id: "call_1".into(), - delta: "{\"ci".into(), - }, - ] - ); - - tracker.clear_parts(); - - // Second delta: more arguments. - tracker - .process_delta(&d().index(0).arguments("ty\": \"San")) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ToolCallStreamPart::ToolInputDelta { - id: "call_1".into(), - delta: "ty\": \"San".into(), - }] - ); - - tracker.clear_parts(); - - // Third delta: completes the JSON -- must not finalize before flush, since a - // parsable buffer can still be the prefix of longer arguments. - tracker - .process_delta(&d().index(0).arguments(" Francisco\"}")) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ToolCallStreamPart::ToolInputDelta { - id: "call_1".into(), - delta: " Francisco\"}".into(), - }] - ); - - tracker.clear_parts(); - - tracker.flush(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputEnd { - id: "call_1".into(), - }, - ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "get_weather".into(), - input: "{\"city\": \"San Francisco\"}".into(), - provider_metadata: None, - }, - ] - ); -} - -#[test] -fn full_tool_call_in_single_chunk() { - // TS: "should handle a full tool call in a single chunk" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("get_weather") - .arguments("{\"city\": \"London\"}"), - ) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputStart { - id: "call_1".into(), - tool_name: "get_weather".into(), - }, - ToolCallStreamPart::ToolInputDelta { - id: "call_1".into(), - delta: "{\"city\": \"London\"}".into(), - }, - ] - ); - - tracker.clear_parts(); - tracker.flush(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputEnd { - id: "call_1".into(), - }, - ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "get_weather".into(), - input: "{\"city\": \"London\"}".into(), - provider_metadata: None, - }, - ] - ); -} - -#[test] -fn not_finalize_when_argument_prefix_is_parsable_json() { - // TS: "should not finalize a tool call when its argument prefix is parsable JSON" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("search") - .arguments("{\"query\": \"test\"}"), - ) - .unwrap(); - - // The parsable prefix must not emit tool-input-end / tool-call. - let types: Vec<&str> = tracker - .parts() - .iter() - .map(|p| match p { - ToolCallStreamPart::ToolInputStart { .. } => "tool-input-start", - ToolCallStreamPart::ToolInputDelta { .. } => "tool-input-delta", - ToolCallStreamPart::ToolInputEnd { .. } => "tool-input-end", - ToolCallStreamPart::ToolCall { .. } => "tool-call", - }) - .collect(); - assert_eq!(types, vec!["tool-input-start", "tool-input-delta"]); - - tracker - .process_delta(&d().index(0).arguments(", \"limit\": 10}")) - .unwrap(); - - tracker.flush(); - - assert_eq!( - tracker.parts().last().unwrap(), - &ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "search".into(), - input: "{\"query\": \"test\"}, \"limit\": 10}".into(), - provider_metadata: None, - } - ); - - let tool_call_count = tracker - .parts() - .iter() - .filter(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .count(); - assert_eq!(tool_call_count, 1); -} - -#[test] -fn multiple_concurrent_tool_calls() { - // TS: "should handle multiple concurrent tool calls" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("get_weather") - .arguments(""), - ) - .unwrap(); - tracker - .process_delta( - &d().index(1) - .id("call_2") - .tool_type("function") - .function_name("get_time") - .arguments(""), - ) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputStart { - id: "call_1".into(), - tool_name: "get_weather".into(), - }, - ToolCallStreamPart::ToolInputStart { - id: "call_2".into(), - tool_name: "get_time".into(), - }, - ] - ); -} - -#[test] -fn skip_deltas_for_already_finished_tool_calls() { - // TS: "should skip deltas for already-finished tool calls" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{}"), - ) - .unwrap(); - - // Finalize via flush. - tracker.flush(); - tracker.clear_parts(); - - // Late delta for the same tool call. - tracker - .process_delta(&d().index(0).arguments("extra")) - .unwrap(); - - assert!(tracker.parts().is_empty()); -} - -#[test] -fn skip_delta_emission_when_arguments_null() { - // TS: "should skip delta emission when arguments are null". - // `arguments: null` maps to `None` (omitted on the builder); behaviorally - // identical to a `function` with no arguments field. - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments(""), - ) - .unwrap(); - - tracker.clear_parts(); - - // Delta with null arguments (no `.arguments(...)` call -> None). - tracker.process_delta(&d().index(0)).unwrap(); - - assert!(tracker.parts().is_empty()); -} - -#[test] -fn use_index_fallback_when_index_not_provided() { - // TS: "should use index fallback when index is not provided" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().id("call_1") - .tool_type("function") - .function_name("fn1") - .arguments("{}"), - ) - .unwrap(); - tracker - .process_delta( - &d().id("call_2") - .tool_type("function") - .function_name("fn2") - .arguments("{}"), - ) - .unwrap(); - - let starts: Vec<_> = tracker - .parts() - .iter() - .filter(|p| matches!(p, ToolCallStreamPart::ToolInputStart { .. })) - .cloned() - .collect(); - assert_eq!( - starts, - vec![ - ToolCallStreamPart::ToolInputStart { - id: "call_1".into(), - tool_name: "fn1".into(), - }, - ToolCallStreamPart::ToolInputStart { - id: "call_2".into(), - tool_name: "fn2".into(), - }, - ] - ); -} - -#[test] -fn generate_id_when_id_missing() { - // TS: id is missing → `toolCall.id ?? generateId()` generates a fallback id - // instead of throwing. The default generator returns "generated". - let mut tracker = new_tracker(); - tracker - .process_delta(&d().index(0).tool_type("function").function_name("fn")) - .unwrap(); - - assert_eq!( - tracker.parts(), - &[ToolCallStreamPart::ToolInputStart { - id: "generated".into(), - tool_name: "fn".into(), - }] - ); -} - -#[test] -fn throw_when_function_name_missing() { - // TS: "should throw when function.name is missing". - // `function: {}` (TS) maps to `function: None` here; both yield a missing - // name and the same error. - let mut tracker = new_tracker(); - let err = tracker - .process_delta(&d().index(0).id("call_1").tool_type("function")) - .unwrap_err(); - assert_eq!(err, TrackerError::MissingFunctionName); - assert_eq!(err.to_string(), "Expected 'function.name' to be a string."); -} - -// == typeValidation ======================================================== - -#[test] -fn no_validate_type_with_type_validation_none() { - // TS: "should not validate type with typeValidation: none" - let mut tracker = new_tracker().with_type_validation(TypeValidation::None); - // Should not throw even with a non-function type. - let result = tracker.process_delta( - &d().index(0) - .id("call_1") - .tool_type("custom") - .function_name("fn") - .arguments(""), - ); - assert!(result.is_ok()); -} - -#[test] -fn validate_type_when_present_with_type_validation_if_present() { - // TS: "should validate type when present with typeValidation: if-present" - let mut tracker = new_tracker().with_type_validation(TypeValidation::IfPresent); - - // Should throw for a non-function type. - let err = tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("custom") - .function_name("fn") - .arguments(""), - ) - .unwrap_err(); - assert_eq!(err, TrackerError::InvalidType); - assert_eq!(err.to_string(), "Expected 'function' type."); - - // Should not throw when type is null (absent). The first delta errored - // before creating any tool call, so index 0 is still new. - let result = - tracker.process_delta(&d().index(0).id("call_1").function_name("fn").arguments("")); - assert!(result.is_ok()); -} - -#[test] -fn require_function_type_with_type_validation_required() { - // TS: "should require function type with typeValidation: required" - let mut tracker = new_tracker().with_type_validation(TypeValidation::Required); - - // Should throw when type is null/undefined. - let err = tracker - .process_delta(&d().index(0).id("call_1").function_name("fn").arguments("")) - .unwrap_err(); - assert_eq!(err, TrackerError::InvalidType); - - // Should not throw for 'function' type. - let result = tracker.process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments(""), - ); - assert!(result.is_ok()); -} - -// == flush ================================================================= - -#[test] -fn finalize_unfinished_tool_calls_on_flush() { - // TS: "should finalize unfinished tool calls on flush" - let mut tracker = new_tracker(); - - // Start a tool call but don't complete it. - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{\"key\": \"val"), - ) - .unwrap(); - - tracker.clear_parts(); - - // Flush should finalize. - tracker.flush(); - - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputEnd { - id: "call_1".into(), - }, - ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{\"key\": \"val".into(), - provider_metadata: None, - }, - ] - ); -} - -#[test] -fn not_refinalize_already_finished_tool_calls() { - // TS: "should not re-finalize already finished tool calls" - let mut tracker = new_tracker(); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{}"), - ) - .unwrap(); - - // First flush finalizes the tool call. - tracker.flush(); - tracker.clear_parts(); - - tracker.flush(); - - // No events should be emitted since the tool call was already finished. - assert!(tracker.parts().is_empty()); -} - -// == metadata ============================================================= - -#[test] -fn extract_and_include_provider_metadata_in_tool_call_events() { - // TS: "should extract and include provider metadata in tool-call events" - let mut tracker = StreamingToolCallTracker::::new() - .with_extract_metadata(|delta| { - delta - .extra - .get("google")? - .get("thought_signature")? - .as_str() - .map(|s| json!({ "thoughtSignature": s })) - }) - .with_build_provider_metadata(|metadata| { - metadata - .as_ref() - .and_then(|m| m.get("thoughtSignature")) - .and_then(|s| s.as_str()) - .map(|s| json!({ "google": { "thoughtSignature": s } })) - }); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{}") - .extra(json!({ "google": { "thought_signature": "sig123" } })), - ) - .unwrap(); - - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!( - tool_call, - &ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{}".into(), - provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), - } - ); -} - -#[test] -fn include_provider_metadata_for_unfinished_tool_calls_finalized_in_flush() { - // TS: "should include provider metadata for unfinished tool calls finalized in flush" - let mut tracker = StreamingToolCallTracker::::new() - .with_extract_metadata(|_| Some(json!({ "custom": { "key": "value" } }))) - .with_build_provider_metadata(|metadata| { - metadata - .as_ref() - .map(|m| json!({ "provider": (*m).clone() })) - }); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{\"incomplete"), - ) - .unwrap(); - - tracker.clear_parts(); - - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!( - tool_call, - &ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{\"incomplete".into(), - provider_metadata: Some(json!({ "provider": { "custom": { "key": "value" } } })), - } - ); -} - -#[test] -fn not_include_provider_metadata_when_build_returns_none() { - // TS: "should not include providerMetadata when buildToolCallProviderMetadata returns undefined" - let mut tracker = StreamingToolCallTracker::::new() - .with_extract_metadata(|_| None) - .with_build_provider_metadata(|_| None); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{}"), - ) - .unwrap(); - - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!( - tool_call, - &ToolCallStreamPart::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{}".into(), - provider_metadata: None, - } - ); - // Mirrors the TS `not.toHaveProperty('providerMetadata')` assertion: the - // metadata is explicitly `None`, never `Some`. - assert!(matches!( - tool_call, - ToolCallStreamPart::ToolCall { - provider_metadata: None, - .. - } - )); -} - -// == generateId =========================================================== - -#[test] -fn use_custom_generate_id_for_tool_call_ids_when_id_missing_in_fallback() { - // TS: "should use custom generateId for tool call IDs when id is missing in fallback". - // The id is present (`call_1`), so the custom generator is NOT used; the - // original id is kept (mirrors the TS `toolCall.id ?? generateId()` path). - use std::sync::atomic::{AtomicUsize, Ordering}; - static CALLS: AtomicUsize = AtomicUsize::new(0); - - let mut tracker = StreamingToolCallTracker::<()>::new().with_generate_id(|| { - CALLS.fetch_add(1, Ordering::SeqCst); - "custom-id".to_string() - }); - - tracker - .process_delta( - &d().index(0) - .id("call_1") - .tool_type("function") - .function_name("fn") - .arguments("{\"key\": \"val"), - ) - .unwrap(); - - tracker.clear_parts(); - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!(tool_call_id(tool_call), "call_1"); - // The custom generator must not have been invoked. - assert_eq!(CALLS.load(Ordering::SeqCst), 0); -} - -#[test] -fn use_custom_generate_id_when_id_is_missing() { - // P2-05: when a tool-call delta omits its id, the configured `generateId` - // fallback is invoked (TS `toolCall.id ?? generateId()`). The generated id - // is used for the emitted ToolInputStart and the finalized ToolCall. - use std::sync::atomic::{AtomicUsize, Ordering}; - static CALLS: AtomicUsize = AtomicUsize::new(0); - - let mut tracker = StreamingToolCallTracker::<()>::new().with_generate_id(|| { - CALLS.fetch_add(1, Ordering::SeqCst); - "custom-id".to_string() - }); - - // No `.id(...)` — the generator must fill it in. - tracker - .process_delta( - &d().index(0) - .tool_type("function") - .function_name("fn") - .arguments("{\"key\": \"val"), - ) - .unwrap(); - - // The generator was invoked exactly once for the missing id. - assert_eq!(CALLS.load(Ordering::SeqCst), 1); - - // The emitted start part carries the generated id. - assert_eq!( - tracker.parts(), - &[ - ToolCallStreamPart::ToolInputStart { - id: "custom-id".into(), - tool_name: "fn".into(), - }, - ToolCallStreamPart::ToolInputDelta { - id: "custom-id".into(), - delta: "{\"key\": \"val".into(), - }, - ] - ); - - tracker.clear_parts(); - tracker.flush(); - - let tool_call = tracker - .parts() - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })) - .unwrap(); - assert_eq!(tool_call_id(tool_call), "custom-id"); - // No additional generator invocations during flush. - assert_eq!(CALLS.load(Ordering::SeqCst), 1); -} - -// == index bound (P1-10) ================================================== - -#[test] -fn reject_tool_call_index_above_max_index() { - // A remote `index` far above the cap must not resize `tool_calls` to a - // huge vector; it returns `IndexOutOfRange`. - let mut tracker = new_tracker().with_max_index(4); - - // index == max_index is accepted (boundary). - tracker - .process_delta( - &d().index(4) - .id("call_4") - .tool_type("function") - .function_name("fn") - .arguments(""), - ) - .unwrap(); - - // index == max_index + 1 is rejected. - let err = tracker - .process_delta( - &d().index(5) - .id("call_5") - .tool_type("function") - .function_name("fn") - .arguments(""), - ) - .unwrap_err(); - assert_eq!(err, TrackerError::IndexOutOfRange); - assert_eq!(err.to_string(), "Tool call index out of range"); -} - -#[test] -fn default_max_index_rejects_huge_index() { - // The default cap (1024) rejects an absurd index that would otherwise - // resize `tool_calls` to index+1 slots. - let mut tracker = new_tracker(); - let err = tracker - .process_delta( - &d().index(1_000_000) - .id("call_x") - .tool_type("function") - .function_name("fn") - .arguments(""), - ) - .unwrap_err(); - assert_eq!(err, TrackerError::IndexOutOfRange); -} From 6035201dbec549f5d77fa39390b070e3f1552f3f Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 29 Sep 2026 05:02:49 +0000 Subject: [PATCH 2/6] feat(stream)!: align StreamingToolCallTracker with the current AI SDK 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 Claude-Session: https://claude.ai/code/session_015EYDYDWYcPsDmjBjFeuDVe --- aimux-stream/src/lib.rs | 9 + .../src/streaming_tool_call_argument_state.rs | 116 ++ .../src/streaming_tool_call_tracker.rs | 638 ++++++++++ ...streaming_tool_call_argument_state_test.rs | 73 ++ .../tests/streaming_tool_call_tracker_test.rs | 1098 +++++++++++++++++ 5 files changed, 1934 insertions(+) create mode 100644 aimux-stream/src/streaming_tool_call_argument_state.rs create mode 100644 aimux-stream/src/streaming_tool_call_tracker.rs create mode 100644 aimux-stream/tests/streaming_tool_call_argument_state_test.rs create mode 100644 aimux-stream/tests/streaming_tool_call_tracker_test.rs diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index 65c368daf..d9078d8d8 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -6,7 +6,16 @@ pub mod lines; pub mod sse; +pub mod streaming_tool_call_argument_state; +pub mod streaming_tool_call_tracker; // Re-export the most commonly used items. pub use lines::extract_lines; pub use sse::{SseError, SseEvent, SseStream}; +pub use streaming_tool_call_argument_state::{ + StreamingToolCallArgumentState, starts_with_structured_value, +}; +pub use streaming_tool_call_tracker::{ + StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, + ToolCallStreamPart, TrackerError, TypeValidation, +}; diff --git a/aimux-stream/src/streaming_tool_call_argument_state.rs b/aimux-stream/src/streaming_tool_call_argument_state.rs new file mode 100644 index 000000000..db156e66d --- /dev/null +++ b/aimux-stream/src/streaming_tool_call_argument_state.rs @@ -0,0 +1,116 @@ +//! Incremental structure tracking for streamed tool-call arguments. +//! +//! Rust translation of `@ai-sdk/provider-utils`'s +//! `StreamingToolCallArgumentState` +//! (`packages/provider-utils/src/streaming-tool-call-argument-state.ts`). +//! +//! The state is intentionally *structural* rather than a JSON parse: a +//! currently parsable scalar can still be the prefix of a later value, but a +//! closed top-level `{…}` / `[…]` cannot be extended. + +#[derive(Debug, Clone, PartialEq, Eq)] +enum ArgumentStructure { + Undetermined, + Other, + Structured { + stack: Vec, + in_string: bool, + escaped: bool, + complete: bool, + }, +} + +/// Whether `value` (ignoring leading whitespace) starts with `{` or `[`. +#[must_use] +pub fn starts_with_structured_value(value: Option<&str>) -> bool { + value + .and_then(|v| v.trim_start().chars().next()) + .is_some_and(|c| c == '{' || c == '[') +} + +/// Incrementally tracks whether streamed tool-call arguments contain a +/// complete structured JSON value. +#[derive(Debug, Clone)] +pub struct StreamingToolCallArgumentState { + structure: ArgumentStructure, +} + +impl StreamingToolCallArgumentState { + /// Create a state seeded with `initial_value`. + #[must_use] + pub fn new(initial_value: &str) -> Self { + let mut state = Self { + structure: ArgumentStructure::Undetermined, + }; + state.append(initial_value); + state + } + + /// `true` once a top-level `{…}` / `[…]` value has been closed. + #[must_use] + pub fn has_complete_structured_value(&self) -> bool { + matches!( + self.structure, + ArgumentStructure::Structured { complete: true, .. } + ) + } + + /// Feed the next argument fragment. + pub fn append(&mut self, delta: &str) { + for character in delta.chars() { + match &mut self.structure { + ArgumentStructure::Undetermined => { + if character.is_whitespace() { + continue; + } + self.structure = if character == '{' || character == '[' { + ArgumentStructure::Structured { + stack: vec![character], + in_string: false, + escaped: false, + complete: false, + } + } else { + ArgumentStructure::Other + }; + } + ArgumentStructure::Other | ArgumentStructure::Structured { complete: true, .. } => { + } + ArgumentStructure::Structured { + stack, + in_string, + escaped, + complete, + } => { + if *in_string { + if *escaped { + *escaped = false; + } else if character == '\\' { + *escaped = true; + } else if character == '"' { + *in_string = false; + } + continue; + } + + match character { + '"' => *in_string = true, + '{' | '[' => stack.push(character), + '}' | ']' => { + let expected = if character == '}' { '{' } else { '[' }; + if stack.last() != Some(&expected) { + self.structure = ArgumentStructure::Other; + continue; + } + stack.pop(); + if stack.is_empty() { + *complete = true; + } + } + _ => {} + } + } + } + } + } +} diff --git a/aimux-stream/src/streaming_tool_call_tracker.rs b/aimux-stream/src/streaming_tool_call_tracker.rs new file mode 100644 index 000000000..9a8d7240e --- /dev/null +++ b/aimux-stream/src/streaming_tool_call_tracker.rs @@ -0,0 +1,638 @@ +//! Streaming tool call tracker. +//! +//! Rust translation of `@ai-sdk/provider-utils`'s `StreamingToolCallTracker` +//! (`packages/provider-utils/src/streaming-tool-call-tracker.ts`). +//! +//! Tracks streaming tool call state across the deltas of an OpenAI-compatible +//! chat completion stream: accumulates `arguments` fragments, emits +//! `tool-input-start` / `tool-input-delta` / `tool-input-end` / `tool-call` +//! parts, and finalizes unfinished calls on [`StreamingToolCallTracker::flush`]. +//! +//! Deltas are correlated to calls by wire `id`, `index` and function name +//! (see [`StreamingToolCallTracker`]'s resolution table), not by `index` +//! alone, so providers that reuse indices, repeat ids, drop ids on +//! continuations or send blank names are handled. +//! +//! Like the TS original, a call is *never* finalized before `flush`: a +//! parsable argument buffer can still be the prefix of a longer argument +//! string, so acting on it early would use truncated inputs (ai-sdk #13137). + +use std::collections::{HashMap, HashSet}; + +use serde_json::Value; +use thiserror::Error; + +use crate::streaming_tool_call_argument_state::{ + StreamingToolCallArgumentState, starts_with_structured_value, +}; + +/// Fallback id used when the id generator returns a blank string. +const FALLBACK_TOOL_CALL_ID: &str = "tool-call"; + +/// The `function` sub-object of a streaming tool call delta. +#[derive(Debug, Clone, Default)] +pub struct StreamingToolCallFunction { + pub name: Option, + pub arguments: Option, +} + +/// A streaming tool call delta — the `tool_calls[i]` entry of an OpenAI-style +/// streaming chunk. +/// +/// `arguments: null` (TS) maps to `None`; `arguments: ''` maps to `Some("")`. +#[derive(Debug, Clone, Default)] +pub struct StreamingToolCallDelta { + pub index: Option, + pub id: Option, + /// The `type` field. Named `r#type` because `type` is a reserved word. + pub r#type: Option, + pub function: Option, + /// Provider-specific payload carried alongside the standard fields, read + /// by the `extract_metadata` hook (e.g. a Google thought signature). + pub extra: Value, +} + +impl StreamingToolCallDelta { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + #[must_use] + pub fn index(mut self, index: usize) -> Self { + self.index = Some(index); + self + } + + #[must_use] + pub fn id(mut self, id: impl Into) -> Self { + self.id = Some(id.into()); + self + } + + /// Set the `type` field (named `tool_type` because `type` is reserved). + #[must_use] + pub fn tool_type(mut self, t: impl Into) -> Self { + self.r#type = Some(t.into()); + self + } + + #[must_use] + pub fn function_name(mut self, name: impl Into) -> Self { + self.function.get_or_insert_with(Default::default).name = Some(name.into()); + self + } + + /// Set the `function.arguments` fragment. Pass `""` for an explicit empty + /// fragment; omit the call for `None` (TS `arguments: null`). + #[must_use] + pub fn arguments(mut self, args: impl Into) -> Self { + self.function.get_or_insert_with(Default::default).arguments = Some(args.into()); + self + } + + #[must_use] + pub fn extra(mut self, extra: Value) -> Self { + self.extra = extra; + self + } +} + +/// The stream parts emitted by [`StreamingToolCallTracker`]. +/// +/// Mirrors the subset of `LanguageModelV4StreamPart` the TS tracker enqueues. +/// [`ToolCallStreamPart::ToolCall::input`] is the raw accumulated argument +/// *string*; the tracker does not parse it. +#[derive(Debug, Clone, PartialEq)] +pub enum ToolCallStreamPart { + /// Start of a tool call's input streaming. + ToolInputStart { id: String, tool_name: String }, + /// A partial argument fragment. + ToolInputDelta { id: String, delta: String }, + /// End of a tool call's input streaming. + ToolInputEnd { id: String }, + /// A complete, finalized tool call. + ToolCall { + tool_call_id: String, + tool_name: String, + input: String, + provider_metadata: Option, + }, +} + +/// How to validate the `type` field on a new tool call delta. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum TypeValidation { + /// No validation (default). + #[default] + None, + /// Error if `type` is present and not `"function"`. + IfPresent, + /// Error if `type` is not exactly `"function"`. + Required, +} + +/// Errors raised while processing a tool call delta. +#[derive(Debug, Error, PartialEq, Eq)] +pub enum TrackerError { + #[error("Expected 'function.name' to be a string.")] + MissingFunctionName, + #[error("Expected 'function' type.")] + InvalidType, + #[error("Failed to create a unique tool call ID.")] + IdExhausted, +} + +struct TrackedToolCall { + id: String, + index: Option, + sequence: usize, + function_name: String, + arguments: String, + argument_state: StreamingToolCallArgumentState, + has_finished: bool, + metadata: Option, +} + +enum ToolCallResolution { + Existing(usize), + New, + Ambiguous, +} + +type GenerateIdFn = Box String + Send + Sync>; +type ExtractMetadataFn = Box Option + Send + Sync>; +type BuildMetadataFn = Box) -> Option + Send + Sync>; + +/// Tracks streaming tool call state across multiple deltas. +/// +/// [`process_delta`](Self::process_delta) and [`flush`](Self::flush) return +/// the parts to forward downstream (the TS tracker enqueues them on a +/// controller instead). +/// +/// The type parameter `M` is the provider-metadata type; use `()` (the +/// default) when no metadata handling is needed. +/// +/// # Correlation +/// +/// | ID evidence | index/name evidence | start evidence | resolution | +/// | --- | --- | --- | --- | +/// | known | matching | any | matching call, new call, or ambiguity | +/// | known | conflicting | named | new call | +/// | unseen | matching | structured start | new call | +/// | unseen | matching | continuation | matching call or ambiguity | +/// | absent | matching | any | matching call, new call, or ambiguity | +/// | absent | absent | named | new call | +/// | absent | absent | unnamed | sole unfinished call, new call, or ambiguity | +/// +/// An *ambiguous* delta is dropped. +pub struct StreamingToolCallTracker { + tool_calls: Vec>, + tool_calls_by_id: HashMap>, + tool_calls_by_index: HashMap>, + used_tool_call_ids: HashSet, + next_generated_id_suffixes: HashMap, + generate_id: GenerateIdFn, + type_validation: TypeValidation, + extract_metadata: Option>, + build_provider_metadata: Option>, +} + +impl Default for StreamingToolCallTracker { + fn default() -> Self { + Self::new() + } +} + +impl StreamingToolCallTracker { + /// Create a tracker with no metadata handling and default settings. + /// + /// The default id generator returns a blank-free constant, so ids that + /// need generating become `tool-call`, `tool-call-1`, `tool-call-2`, … + #[must_use] + pub fn new() -> Self { + Self { + tool_calls: Vec::new(), + tool_calls_by_id: HashMap::new(), + tool_calls_by_index: HashMap::new(), + used_tool_call_ids: HashSet::new(), + next_generated_id_suffixes: HashMap::new(), + generate_id: Box::new(|| FALLBACK_TOOL_CALL_ID.to_string()), + type_validation: TypeValidation::None, + extract_metadata: None, + build_provider_metadata: None, + } + } + + /// Set the id generator (the TS `generateId` option). Blank or repeated + /// outputs are turned into usable unique ids. + #[must_use] + pub fn with_generate_id String + Send + Sync + 'static>(mut self, f: F) -> Self { + self.generate_id = Box::new(f); + self + } + + /// Set the `type` validation mode (the TS `typeValidation` option). + #[must_use] + pub fn with_type_validation(mut self, v: TypeValidation) -> Self { + self.type_validation = v; + self + } + + /// Set the metadata extractor (the TS `extractMetadata` option). Called + /// once when a new tool call is detected. + #[must_use] + pub fn with_extract_metadata< + F: Fn(&StreamingToolCallDelta) -> Option + Send + Sync + 'static, + >( + mut self, + f: F, + ) -> Self { + self.extract_metadata = Some(Box::new(f)); + self + } + + /// Set the provider-metadata builder (the TS + /// `buildToolCallProviderMetadata` option). If it returns `None`, the + /// `tool-call` part carries no `provider_metadata`. + #[must_use] + pub fn with_build_provider_metadata) -> Option + Send + Sync + 'static>( + mut self, + f: F, + ) -> Self { + self.build_provider_metadata = Some(Box::new(f)); + self + } + + /// Process a tool call delta from a streaming chunk and return the parts + /// it produces (possibly none). + /// + /// # Errors + /// + /// Returns [`TrackerError::InvalidType`] when `type` validation fails, + /// [`TrackerError::MissingFunctionName`] when a new call has no + /// (non-null) function name, and [`TrackerError::IdExhausted`] if no + /// unique id can be produced. + pub fn process_delta( + &mut self, + delta: &StreamingToolCallDelta, + ) -> Result>, TrackerError> { + let wire_name = delta.function.as_ref().and_then(|f| f.name.as_deref()); + let has_blank_name = wire_name.is_some_and(|n| n.trim().is_empty()); + let wire_id = non_blank(delta.id.as_deref()); + let name = non_blank(wire_name); + let index = delta.index; + let arguments = delta.function.as_ref().and_then(|f| f.arguments.as_deref()); + + let resolution = self.resolve_tool_call( + wire_id, + index, + name, + name.is_some() && starts_with_structured_value(arguments), + ); + + let mut parts = Vec::new(); + let call = match resolution { + ToolCallResolution::Ambiguous => return Ok(parts), + ToolCallResolution::New => { + // Blank names cannot start a usable call, but some providers + // repeat a blank name on continuations. Those were correlated + // above; only an unmatched blank-name delta is ignored. + if has_blank_name { + return Ok(parts); + } + self.process_new_tool_call(delta, wire_id, index, name, &mut parts)? + } + ToolCallResolution::Existing(call) => { + if let Some(wire_id) = wire_id { + self.associate_wire_id(call, wire_id); + } + self.process_existing_tool_call(call, arguments, &mut parts); + call + } + }; + + if let Some(index) = index { + self.tool_calls_by_index + .entry(index) + .or_default() + .insert(call); + } + Ok(parts) + } + + /// Finalize any unfinished tool calls and return the closing parts. Call + /// once when the stream ends. + pub fn flush(&mut self) -> Vec> { + // Index order is only reliable when every call has an index; for + // mixed streams keep insertion order. + let mut order: Vec = (0..self.tool_calls.len()).collect(); + if self.tool_calls.iter().all(|c| c.index.is_some()) { + order.sort_by_key(|&i| (self.tool_calls[i].index, self.tool_calls[i].sequence)); + } + + let mut parts = Vec::new(); + for call in order { + if !self.tool_calls[call].has_finished { + self.finish_tool_call(call, &mut parts); + } + } + parts + } + + fn resolve_tool_call( + &self, + wire_id: Option<&str>, + index: Option, + name: Option<&str>, + has_explicit_call_start: bool, + ) -> ToolCallResolution { + let indexed = index.and_then(|i| self.tool_calls_by_index.get(&i)); + let matching_indexed = self.filter_by_name(indexed, name); + + if let Some(wire_id) = wire_id { + if let Some(with_id) = self.tool_calls_by_id.get(wire_id) { + if index.is_some() { + let matching: Vec = matching_indexed + .iter() + .copied() + .filter(|c| with_id.contains(c)) + .collect(); + let resolved = self.resolve_matching(matching, has_explicit_call_start); + if !matches!(resolved, ToolCallResolution::New) { + return resolved; + } + + // A named delta with a distinct index starts a new call + // even when its wire id and name repeat: providers may + // reuse ids across parallel calls. + if name.is_some() { + return ToolCallResolution::New; + } + + // Conflicting labels on a continuation cannot be + // resolved safely. + if indexed.is_some() { + return ToolCallResolution::Ambiguous; + } + + return self.resolve_matching(sorted(with_id), false); + } + + if let Some(name) = name { + let matching: Vec = sorted(with_id) + .into_iter() + .filter(|&c| self.tool_calls[c].function_name == name) + .collect(); + return self.resolve_matching(matching, has_explicit_call_start); + } + + return self.resolve_matching(sorted(with_id), false); + } + + if !matching_indexed.is_empty() { + // A previously unseen id plus a named structured argument + // start is stronger evidence of a distinct call than a reused + // index/name; ids may still change on plain continuations. + return if has_explicit_call_start { + ToolCallResolution::New + } else { + self.resolve_matching(matching_indexed, false) + }; + } + + return ToolCallResolution::New; + } + + if indexed.is_some() { + // Repeated names are valid on continuations; a different name at + // the same index means a new call from a provider that reuses + // indices across parallel calls. + return self.resolve_matching(matching_indexed, has_explicit_call_start); + } + + if name.is_some() { + return ToolCallResolution::New; + } + + let unfinished: Vec = (0..self.tool_calls.len()) + .filter(|&c| !self.tool_calls[c].has_finished) + .collect(); + match unfinished.len() { + 0 => ToolCallResolution::New, + 1 => ToolCallResolution::Existing(unfinished[0]), + _ => ToolCallResolution::Ambiguous, + } + } + + fn filter_by_name(&self, calls: Option<&HashSet>, name: Option<&str>) -> Vec { + calls.map_or_else(Vec::new, |calls| { + sorted(calls) + .into_iter() + .filter(|&c| name.is_none_or(|n| self.tool_calls[c].function_name == n)) + .collect() + }) + } + + fn resolve_matching( + &self, + calls: Vec, + has_explicit_call_start: bool, + ) -> ToolCallResolution { + if calls.is_empty() { + return ToolCallResolution::New; + } + + if !has_explicit_call_start { + return if calls.len() == 1 { + ToolCallResolution::Existing(calls[0]) + } else { + ToolCallResolution::Ambiguous + }; + } + + // A repeated name can occur on continuations. A fresh structured + // argument prefix signals another call only once the matching call + // has completed its own structured payload. + let continuable: Vec = calls + .into_iter() + .filter(|&c| { + !self.tool_calls[c] + .argument_state + .has_complete_structured_value() + }) + .collect(); + match continuable.len() { + 0 => ToolCallResolution::New, + 1 => ToolCallResolution::Existing(continuable[0]), + _ => ToolCallResolution::Ambiguous, + } + } + + fn process_new_tool_call( + &mut self, + delta: &StreamingToolCallDelta, + wire_id: Option<&str>, + index: Option, + name: Option<&str>, + parts: &mut Vec>, + ) -> Result { + match self.type_validation { + TypeValidation::Required => { + if delta.r#type.as_deref() != Some("function") { + return Err(TrackerError::InvalidType); + } + } + TypeValidation::IfPresent => { + if delta.r#type.as_deref().is_some_and(|t| t != "function") { + return Err(TrackerError::InvalidType); + } + } + TypeValidation::None => {} + } + + let name = name.ok_or(TrackerError::MissingFunctionName)?; + let id = self.create_tool_call_id(wire_id)?; + + parts.push(ToolCallStreamPart::ToolInputStart { + id: id.clone(), + tool_name: name.to_string(), + }); + + let metadata = self + .extract_metadata + .as_ref() + .and_then(|extract| extract(delta)); + + let initial_arguments = delta + .function + .as_ref() + .and_then(|f| f.arguments.clone()) + .unwrap_or_default(); + + let call = self.tool_calls.len(); + self.tool_calls.push(TrackedToolCall { + id: id.clone(), + index, + sequence: call, + function_name: name.to_string(), + argument_state: StreamingToolCallArgumentState::new(&initial_arguments), + arguments: initial_arguments.clone(), + has_finished: false, + metadata, + }); + if let Some(wire_id) = wire_id { + self.associate_wire_id(call, wire_id); + } + + if !initial_arguments.is_empty() { + parts.push(ToolCallStreamPart::ToolInputDelta { + id, + delta: initial_arguments, + }); + } + + // Tool calls must not finalize before the stream ends (ai-sdk + // #13137); finalization happens in `flush`. + Ok(call) + } + + fn process_existing_tool_call( + &mut self, + call: usize, + arguments: Option<&str>, + parts: &mut Vec>, + ) { + let tool_call = &mut self.tool_calls[call]; + if tool_call.has_finished { + return; + } + if let Some(arguments) = arguments { + tool_call.argument_state.append(arguments); + tool_call.arguments.push_str(arguments); + parts.push(ToolCallStreamPart::ToolInputDelta { + id: tool_call.id.clone(), + delta: arguments.to_string(), + }); + } + } + + fn associate_wire_id(&mut self, call: usize, wire_id: &str) { + self.tool_calls_by_id + .entry(wire_id.to_string()) + .or_default() + .insert(call); + } + + fn create_tool_call_id(&mut self, wire_id: Option<&str>) -> Result { + if let Some(wire_id) = wire_id + && !self.used_tool_call_ids.contains(wire_id) + { + self.used_tool_call_ids.insert(wire_id.to_string()); + return Ok(wire_id.to_string()); + } + + let generated = non_blank(Some(&(self.generate_id)())) + .unwrap_or(FALLBACK_TOOL_CALL_ID) + .to_string(); + + if !self.used_tool_call_ids.contains(&generated) { + self.used_tool_call_ids.insert(generated.clone()); + return Ok(generated); + } + + // Resume after the last suffix checked for this generated value so + // deterministic generators stay bounded without rescanning occupied + // suffixes. + let initial_suffix = self + .next_generated_id_suffixes + .get(&generated) + .copied() + .unwrap_or(1); + let maximum_suffix = initial_suffix + self.used_tool_call_ids.len(); + for suffix in initial_suffix..=maximum_suffix { + let suffixed = format!("{generated}-{suffix}"); + if !self.used_tool_call_ids.contains(&suffixed) { + self.used_tool_call_ids.insert(suffixed.clone()); + self.next_generated_id_suffixes + .insert(generated, suffix + 1); + return Ok(suffixed); + } + } + + // Unreachable by the pigeonhole principle; guards the invariant + // instead of looping without bound. + Err(TrackerError::IdExhausted) + } + + fn finish_tool_call(&mut self, call: usize, parts: &mut Vec>) { + let tool_call = &mut self.tool_calls[call]; + tool_call.has_finished = true; + + parts.push(ToolCallStreamPart::ToolInputEnd { + id: tool_call.id.clone(), + }); + + let provider_metadata = self + .build_provider_metadata + .as_ref() + .and_then(|build| build(tool_call.metadata.as_ref())); + + parts.push(ToolCallStreamPart::ToolCall { + tool_call_id: tool_call.id.clone(), + tool_name: tool_call.function_name.clone(), + input: tool_call.arguments.clone(), + provider_metadata, + }); + } +} + +fn non_blank(value: Option<&str>) -> Option<&str> { + value.filter(|v| !v.trim().is_empty()) +} + +fn sorted(calls: &HashSet) -> Vec { + let mut calls: Vec = calls.iter().copied().collect(); + calls.sort_unstable(); + calls +} diff --git a/aimux-stream/tests/streaming_tool_call_argument_state_test.rs b/aimux-stream/tests/streaming_tool_call_argument_state_test.rs new file mode 100644 index 000000000..44cf754e2 --- /dev/null +++ b/aimux-stream/tests/streaming_tool_call_argument_state_test.rs @@ -0,0 +1,73 @@ +//! Port of `streaming-tool-call-argument-state.test.ts` from +//! `@ai-sdk/provider-utils`. + +use aimux_stream::{StreamingToolCallArgumentState, starts_with_structured_value}; + +#[test] +fn starts_with_structured_value_true() { + for value in ["{}", " {", "[]", "\n["] { + assert!(starts_with_structured_value(Some(value)), "{value:?}"); + } +} + +#[test] +fn starts_with_structured_value_false() { + assert!(!starts_with_structured_value(None)); + for value in ["", " ", "1", "\"value\""] { + assert!(!starts_with_structured_value(Some(value)), "{value:?}"); + } +} + +#[test] +fn tracks_a_structured_value_across_deltas() { + let mut state = StreamingToolCallArgumentState::new(" {\"value\":"); + assert!(!state.has_complete_structured_value()); + + state.append("1}"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn tracks_nested_objects_and_arrays() { + let state = StreamingToolCallArgumentState::new("[{\"value\":{\"items\":[1,2]}}]"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn ignores_structural_characters_inside_strings() { + let state = StreamingToolCallArgumentState::new("{\"value\":\"braces: } ] { [\"}"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn handles_escaped_quotes_across_deltas() { + let mut state = StreamingToolCallArgumentState::new("{\"value\":\"escaped quote: \\\""); + assert!(!state.has_complete_structured_value()); + + state.append(" still in string\"}"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn can_begin_after_an_empty_or_whitespace_only_delta() { + let mut state = StreamingToolCallArgumentState::new(" "); + + state.append("["); + assert!(!state.has_complete_structured_value()); + + state.append("]"); + assert!(state.has_complete_structured_value()); +} + +#[test] +fn does_not_treat_scalar_arguments_as_a_complete_structured_value() { + let state = StreamingToolCallArgumentState::new("12"); + assert!(!state.has_complete_structured_value()); +} + +#[test] +fn does_not_recover_mismatched_structures_as_complete() { + let mut state = StreamingToolCallArgumentState::new("{\"value\":]"); + state.append("}"); + assert!(!state.has_complete_structured_value()); +} diff --git a/aimux-stream/tests/streaming_tool_call_tracker_test.rs b/aimux-stream/tests/streaming_tool_call_tracker_test.rs new file mode 100644 index 000000000..7c55044c4 --- /dev/null +++ b/aimux-stream/tests/streaming_tool_call_tracker_test.rs @@ -0,0 +1,1098 @@ +//! Port of `streaming-tool-call-tracker.test.ts` from +//! `@ai-sdk/provider-utils`, case for case. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_stream::{ + StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, + ToolCallStreamPart, TrackerError, TypeValidation, +}; +use serde_json::{Value, json}; + +type Part = ToolCallStreamPart; + +/// Tracker plus the parts it has emitted so far (the TS `createCollector`). +struct Harness { + tracker: StreamingToolCallTracker, + parts: Vec>, +} + +impl Harness<()> { + fn new() -> Self { + Self::with(StreamingToolCallTracker::new()) + } + + fn with(tracker: StreamingToolCallTracker<()>) -> Self { + Self::with_meta(tracker) + } +} + +impl Harness { + fn with_meta(tracker: StreamingToolCallTracker) -> Self { + Self { + tracker, + parts: Vec::new(), + } + } + + fn delta(&mut self, delta: StreamingToolCallDelta) -> Result<(), TrackerError> { + let parts = self.tracker.process_delta(&delta)?; + self.parts.extend(parts); + Ok(()) + } + + fn flush(&mut self) { + let parts = self.tracker.flush(); + self.parts.extend(parts); + } + + fn clear(&mut self) { + self.parts.clear(); + } +} + +impl Harness { + /// `(id, name, input)` of every emitted `tool-call`, in order. + fn tool_calls(&self) -> Vec<(String, String, String)> { + self.parts + .iter() + .filter_map(|part| match part { + Part::ToolCall { + tool_call_id, + tool_name, + input, + .. + } => Some((tool_call_id.clone(), tool_name.clone(), input.clone())), + _ => None, + }) + .collect() + } +} + +/// A delta with an explicit `function` object (as every TS test passes one). +fn delta( + index: Option, + id: Option<&str>, + ty: Option<&str>, + name: Option<&str>, + arguments: Option<&str>, +) -> StreamingToolCallDelta { + StreamingToolCallDelta { + index, + id: id.map(str::to_string), + r#type: ty.map(str::to_string), + function: Some(StreamingToolCallFunction { + name: name.map(str::to_string), + arguments: arguments.map(str::to_string), + }), + extra: Value::Null, + } +} + +/// Full call start: index, id, `type: 'function'`, name and arguments. +fn start(index: usize, id: &str, name: &str, arguments: &str) -> StreamingToolCallDelta { + delta( + Some(index), + Some(id), + Some("function"), + Some(name), + Some(arguments), + ) +} + +/// Continuation carrying only arguments (plus an optional index). +fn cont(index: Option, arguments: &str) -> StreamingToolCallDelta { + delta(index, None, None, None, Some(arguments)) +} + +fn tc(id: &str, name: &str, input: &str) -> (String, String, String) { + (id.into(), name.into(), input.into()) +} + +fn input_start(id: &str, tool_name: &str) -> Part { + Part::ToolInputStart { + id: id.into(), + tool_name: tool_name.into(), + } +} + +fn input_delta(id: &str, delta: &str) -> Part { + Part::ToolInputDelta { + id: id.into(), + delta: delta.into(), + } +} + +fn input_end(id: &str) -> Part { + Part::ToolInputEnd { id: id.into() } +} + +fn tool_call(id: &str, name: &str, input: &str) -> Part { + Part::ToolCall { + tool_call_id: id.into(), + tool_name: name.into(), + input: input.into(), + provider_metadata: None, + } +} + +/// Deterministic id generator: `prefix-1`, `prefix-2`, … and a call counter. +fn counting_generator(ids: &'static [&'static str]) -> (impl Fn() -> String, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&calls); + ( + move || { + let n = counter.fetch_add(1, Ordering::SeqCst); + ids[n.min(ids.len() - 1)].to_string() + }, + calls, + ) +} + +mod process_delta { + use super::*; + + #[test] + fn single_tool_call_accumulated_across_multiple_deltas() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "get_weather", "{\"ci")).unwrap(); + assert_eq!( + h.parts, + vec![ + input_start("call_1", "get_weather"), + input_delta("call_1", "{\"ci"), + ] + ); + h.clear(); + + h.delta(cont(Some(0), "ty\": \"San")).unwrap(); + assert_eq!(h.parts, vec![input_delta("call_1", "ty\": \"San")]); + h.clear(); + + // Completing the JSON must not finalize before flush: a parsable + // buffer can still be the prefix of longer arguments. + h.delta(cont(Some(0), " Francisco\"}")).unwrap(); + assert_eq!(h.parts, vec![input_delta("call_1", " Francisco\"}")]); + h.clear(); + + h.flush(); + assert_eq!( + h.parts, + vec![ + input_end("call_1"), + tool_call("call_1", "get_weather", "{\"city\": \"San Francisco\"}"), + ] + ); + } + + #[test] + fn full_tool_call_in_a_single_chunk() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "get_weather", "{\"city\": \"London\"}")) + .unwrap(); + assert_eq!( + h.parts, + vec![ + input_start("call_1", "get_weather"), + input_delta("call_1", "{\"city\": \"London\"}"), + ] + ); + h.clear(); + + h.flush(); + assert_eq!( + h.parts, + vec![ + input_end("call_1"), + tool_call("call_1", "get_weather", "{\"city\": \"London\"}"), + ] + ); + } + + #[test] + fn does_not_finalize_when_argument_prefix_is_parsable_json() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "search", "{\"query\": \"test\"}")) + .unwrap(); + assert_eq!( + h.parts, + vec![ + input_start("call_1", "search"), + input_delta("call_1", "{\"query\": \"test\"}"), + ] + ); + + h.delta(cont(Some(0), ", \"limit\": 10}")).unwrap(); + h.flush(); + + assert_eq!( + h.parts.last(), + Some(&tool_call( + "call_1", + "search", + "{\"query\": \"test\"}, \"limit\": 10}" + )) + ); + assert_eq!(h.tool_calls().len(), 1); + } + + #[test] + fn multiple_concurrent_tool_calls() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "get_weather", "")).unwrap(); + h.delta(start(1, "call_2", "get_time", "")).unwrap(); + + assert_eq!( + h.parts, + vec![ + input_start("call_1", "get_weather"), + input_start("call_2", "get_time"), + ] + ); + } + + #[test] + fn non_zero_and_non_contiguous_indexes() { + let mut h = Harness::new(); + + h.delta(start(1, "call_1", "fn1", "{\"value\":1}")).unwrap(); + h.delta(start(3, "call_2", "fn2", "{\"value\":2}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "fn1", "{\"value\":1}"), + tc("call_2", "fn2", "{\"value\":2}"), + ] + ); + } + + #[test] + fn keeps_distinct_tool_calls_that_reuse_an_index() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{\"value\":1}")).unwrap(); + h.delta(start(0, "call_2", "fn", "{\"value\":2}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "fn", "{\"value\":1}"), + tc("call_2", "fn", "{\"value\":2}"), + ] + ); + } + + #[test] + fn continues_latest_call_when_index_is_omitted_after_starting_at() { + for index in [None, Some(7)] { + let mut h = Harness::new(); + + h.delta(delta( + index, + Some("call_1"), + Some("function"), + Some("fn"), + Some("{\"val"), + )) + .unwrap(); + h.delta(cont(None, "ue\":1}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "fn", "{\"value\":1}")], + "index {index:?}" + ); + } + } + + #[test] + fn uses_the_index_when_continuation_ids_are_empty() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{\"val")).unwrap(); + h.delta(delta( + Some(0), + Some(""), + Some("function"), + None, + Some("ue\":1}"), + )) + .unwrap(); + h.flush(); + + assert_eq!(h.tool_calls(), vec![tc("call_1", "fn", "{\"value\":1}")]); + } + + #[test] + fn skips_deltas_for_already_finished_tool_calls() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{}")).unwrap(); + h.flush(); + h.clear(); + + h.delta(cont(Some(0), "extra")).unwrap(); + assert_eq!(h.parts, vec![]); + } + + #[test] + fn skips_delta_emission_when_arguments_are_null() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "")).unwrap(); + h.clear(); + + h.delta(delta(Some(0), None, None, None, None)).unwrap(); + assert_eq!(h.parts, vec![]); + } + + #[test] + fn uses_index_fallback_when_index_is_not_provided() { + let mut h = Harness::new(); + + h.delta(delta( + None, + Some("call_1"), + Some("function"), + Some("fn1"), + Some("{}"), + )) + .unwrap(); + h.delta(delta( + None, + Some("call_2"), + Some("function"), + Some("fn2"), + Some("{}"), + )) + .unwrap(); + + let starts: Vec<&Part> = h + .parts + .iter() + .filter(|p| matches!(p, Part::ToolInputStart { .. })) + .collect(); + assert_eq!( + starts, + vec![&input_start("call_1", "fn1"), &input_start("call_2", "fn2")] + ); + } + + #[test] + fn generates_an_id_when_id_is_missing() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), + ); + + h.delta(delta( + Some(0), + None, + Some("function"), + Some("fn"), + Some("{}"), + )) + .unwrap(); + h.flush(); + + assert_eq!(h.tool_calls(), vec![tc("generated-id", "fn", "{}")]); + } + + #[test] + fn errors_when_function_name_is_missing() { + // TS covers `name: undefined` and `name: null`; both are `None`. + let mut h = Harness::new(); + + let result = h.delta(delta(Some(0), Some("call_1"), Some("function"), None, None)); + + assert_eq!(result, Err(TrackerError::MissingFunctionName)); + assert_eq!( + TrackerError::MissingFunctionName.to_string(), + "Expected 'function.name' to be a string." + ); + } + + #[test] + fn ignores_a_blank_function_name_without_preventing_prior_calls_from_finalizing() { + for name in ["", " "] { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "valid_tool", "{\"value\":1}")) + .unwrap(); + h.delta(start(1, "call_2", name, "{\"value\":2}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "valid_tool", "{\"value\":1}")], + "name {name:?}" + ); + } + } + + #[test] + fn retains_continuation_arguments_for_a_blank_name() { + // (blank name, continuation id, continuation index) + let cases = [ + ("", Some("call_1"), None), // blank name with a matching id + (" ", None, Some(0)), // whitespace name, matching index + ]; + for (name, id, index) in cases { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(delta(index, id, None, Some(name), Some("th\":\"a\"}"))) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")], + "name {name:?}" + ); + } + } + + #[test] + fn keeps_id_less_calls_distinct_when_an_index_is_reused_and_type_is_omitted() { + let (generate, _) = counting_generator(&["generated-1", "generated-2", "generated-3"]); + let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); + + h.delta(delta( + Some(0), + None, + None, + Some("read_file"), + Some("{\"path\":\"p0\"}"), + )) + .unwrap(); + h.delta(delta( + Some(0), + None, + None, + Some("write_file"), + Some("{\"path\":\"p1\"}"), + )) + .unwrap(); + h.delta(delta( + Some(0), + None, + None, + Some("read_file"), + Some("{\"path\":\"p2\"}"), + )) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("generated-1", "read_file", "{\"path\":\"p0\"}"), + tc("generated-2", "write_file", "{\"path\":\"p1\"}"), + tc("generated-3", "read_file", "{\"path\":\"p2\"}"), + ] + ); + } + + #[test] + fn keeps_complete_same_name_calls_distinct_with_reused_index() { + for id in [None, Some("dup")] { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), + ); + + for value in ["{\"value\":1}", "{\"value\":2}"] { + h.delta(delta( + Some(0), + id, + Some("function"), + Some("same_tool"), + Some(value), + )) + .unwrap(); + } + h.flush(); + + let calls = h.tool_calls(); + let inputs: Vec<&str> = calls.iter().map(|c| c.2.as_str()).collect(); + assert_eq!(inputs, vec!["{\"value\":1}", "{\"value\":2}"], "id {id:?}"); + let ids: std::collections::HashSet<&str> = calls.iter().map(|c| c.0.as_str()).collect(); + assert_eq!(ids.len(), 2, "id {id:?}"); + } + } + + #[test] + fn keeps_a_partial_same_name_call_distinct_with_reused_index() { + for id in [None, Some("dup")] { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), + ); + + for value in ["{\"value\":1}", "{\"value\":"] { + h.delta(delta( + Some(0), + id, + Some("function"), + Some("same_tool"), + Some(value), + )) + .unwrap(); + } + h.flush(); + + let calls = h.tool_calls(); + let inputs: Vec<&str> = calls.iter().map(|c| c.2.as_str()).collect(); + assert_eq!(inputs, vec!["{\"value\":1}", "{\"value\":"], "id {id:?}"); + let ids: std::collections::HashSet<&str> = calls.iter().map(|c| c.0.as_str()).collect(); + assert_eq!(ids.len(), 2, "id {id:?}"); + } + } + + #[test] + fn keeps_interleaved_same_name_calls_with_distinct_ids_and_reused_index_separate() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "same_tool", "{\"value\":")) + .unwrap(); + h.delta(start(0, "call_2", "same_tool", "{\"value\":2}")) + .unwrap(); + h.delta(delta(Some(0), Some("call_1"), None, None, Some("1}"))) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "same_tool", "{\"value\":1}"), + tc("call_2", "same_tool", "{\"value\":2}"), + ] + ); + } + + #[test] + fn ignores_an_index_only_continuation_after_the_index_is_reused() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "first", "{\"value\":1}")) + .unwrap(); + h.delta(start(0, "call_2", "second", "{\"value\":2}")) + .unwrap(); + h.delta(cont(Some(0), "{\"unattributed\":true}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "first", "{\"value\":1}"), + tc("call_2", "second", "{\"value\":2}"), + ] + ); + } + + #[test] + fn uses_the_index_when_continuation_ids_are_blank() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(delta(Some(0), Some(" "), None, None, Some("th\":\"a\"}"))) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")] + ); + } + + #[test] + fn generates_unique_ids_for_blank_and_repeated_ids() { + let (generate, _) = counting_generator(&["generated-1", "generated-2"]); + let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); + + h.delta(start(0, "", "read_file", "{}")).unwrap(); + h.delta(start(1, "dup", "read_file", "{}")).unwrap(); + h.delta(start(2, "dup", "write_file", "{}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("generated-1", "read_file", "{}"), + tc("dup", "read_file", "{}"), + tc("generated-2", "write_file", "{}"), + ] + ); + } + + #[test] + fn keeps_same_name_calls_with_repeated_ids_and_distinct_indices_separate() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), + ); + + h.delta(start(0, "dup", "same_tool", "{\"value\":0}")) + .unwrap(); + h.delta(start(1, "dup", "same_tool", "{\"value\":1}")) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("dup", "same_tool", "{\"value\":0}"), + tc("generated-id", "same_tool", "{\"value\":1}"), + ] + ); + } + + #[test] + fn preserves_nonblank_ids_and_function_names_exactly() { + let mut h = Harness::new(); + + h.delta(start(0, " spaced ", " same_tool ", "{\"value\":")) + .unwrap(); + h.delta(delta(Some(0), Some(" spaced "), None, None, Some("0}"))) + .unwrap(); + h.delta(start(1, "spaced", " same_tool ", "{\"value\":1}")) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc(" spaced ", " same_tool ", "{\"value\":0}"), + tc("spaced", " same_tool ", "{\"value\":1}"), + ] + ); + } + + #[test] + fn creates_bounded_unique_ids_when_generate_id_returns_duplicates() { + let (generate, calls) = counting_generator(&["generated-id"]); + let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); + + for (index, name) in [(0, "first"), (1, "second"), (2, "third")] { + h.delta(delta( + Some(index), + None, + Some("function"), + Some(name), + Some("{}"), + )) + .unwrap(); + } + h.flush(); + + assert_eq!(calls.load(Ordering::SeqCst), 3); + let ids: Vec = h.tool_calls().into_iter().map(|c| c.0).collect(); + assert_eq!( + ids, + vec!["generated-id", "generated-id-1", "generated-id-2"] + ); + } + + #[test] + fn creates_usable_ids_when_generate_id_returns_blank_values() { + let (generate, calls) = counting_generator(&[" "]); + let mut h = Harness::with(StreamingToolCallTracker::new().with_generate_id(generate)); + + for (index, name) in [(0, "first"), (1, "second")] { + h.delta(delta( + Some(index), + None, + Some("function"), + Some(name), + Some("{}"), + )) + .unwrap(); + } + h.flush(); + + assert_eq!(calls.load(Ordering::SeqCst), 2); + let ids: Vec = h.tool_calls().into_iter().map(|c| c.0).collect(); + assert_eq!(ids, vec!["tool-call", "tool-call-1"]); + } + + #[test] + fn ignores_unattributable_deltas_when_multiple_calls_are_active() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"path\":\"a\"}")) + .unwrap(); + h.delta(start(1, "call_2", "write_file", "{\"path\":\"b\"}")) + .unwrap(); + h.delta(cont(None, "{\"unattributed\":true}")).unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("call_1", "read_file", "{\"path\":\"a\"}"), + tc("call_2", "write_file", "{\"path\":\"b\"}"), + ] + ); + } + + #[test] + fn ignores_an_ambiguous_continuation_for_a_repeated_id() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "generated-id".to_string()), + ); + + h.delta(start(0, "dup", "read_file", "{\"path\":\"a\"}")) + .unwrap(); + h.delta(start(1, "dup", "write_file", "{\"path\":\"b\"}")) + .unwrap(); + h.delta(delta( + None, + Some("dup"), + None, + None, + Some("{\"unattributed\":true}"), + )) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![ + tc("dup", "read_file", "{\"path\":\"a\"}"), + tc("generated-id", "write_file", "{\"path\":\"b\"}"), + ] + ); + } + + #[test] + fn uses_a_matching_name_and_index_for_an_id_less_continuation() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_type_validation(TypeValidation::Required), + ); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(delta( + Some(0), + None, + None, + Some("read_file"), + Some("th\":\"a\"}"), + )) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")] + ); + } + + #[test] + fn uses_the_index_when_a_continuation_has_an_unexpected_id() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(delta( + Some(0), + Some("unexpected"), + None, + None, + Some("th\":\"a\"}"), + )) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")] + ); + } + + #[test] + fn continues_a_call_when_its_id_changes_but_index_and_name_match() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "read_file", "{\"pa")).unwrap(); + h.delta(start(0, "unexpected", "read_file", "th\":\"a\"}")) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "read_file", "{\"path\":\"a\"}")] + ); + } + + #[test] + fn continues_a_call_when_all_labels_repeat_after_a_parsable_argument_prefix() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "calculate", "1")).unwrap(); + h.delta(start(0, "call_1", "calculate", "2")).unwrap(); + h.flush(); + + assert_eq!(h.tool_calls(), vec![tc("call_1", "calculate", "12")]); + } + + #[test] + fn continues_a_structured_argument_when_repeated_labels_precede_a_nested_object() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "calculate", "{\"value\":")) + .unwrap(); + h.delta(start(0, "call_1", "calculate", "{\"nested\":true}}")) + .unwrap(); + h.flush(); + + assert_eq!( + h.tool_calls(), + vec![tc("call_1", "calculate", "{\"value\":{\"nested\":true}}")] + ); + } + + #[test] + fn emits_tool_calls_in_index_order() { + let mut h = Harness::new(); + + h.delta(start(1, "call_1", "second", "{}")).unwrap(); + h.delta(start(0, "call_0", "first", "{}")).unwrap(); + h.flush(); + + let names: Vec = h.tool_calls().into_iter().map(|c| c.1).collect(); + assert_eq!(names, vec!["first", "second"]); + } + + #[test] + fn preserves_insertion_order_when_calls_mix_present_and_omitted_indices() { + let mut h = Harness::new(); + + h.delta(delta( + None, + Some("call_without_index"), + Some("function"), + Some("first"), + Some("{}"), + )) + .unwrap(); + h.delta(start(0, "call_with_index", "second", "{}")) + .unwrap(); + h.flush(); + + let names: Vec = h.tool_calls().into_iter().map(|c| c.1).collect(); + assert_eq!(names, vec!["first", "second"]); + } + + #[test] + fn errors_when_function_name_is_missing_from_a_new_call() { + let mut h = Harness::new(); + + // `function: {}` — a function object with neither name nor arguments. + let result = h.delta(delta(Some(0), Some("call_1"), Some("function"), None, None)); + + assert_eq!(result, Err(TrackerError::MissingFunctionName)); + } +} + +mod type_validation { + use super::*; + + fn custom_type() -> StreamingToolCallDelta { + delta( + Some(0), + Some("call_1"), + Some("custom"), + Some("fn"), + Some(""), + ) + } + + fn no_type() -> StreamingToolCallDelta { + delta(Some(0), Some("call_1"), None, Some("fn"), Some("")) + } + + #[test] + fn does_not_validate_type_with_none() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_type_validation(TypeValidation::None), + ); + + assert_eq!(h.delta(custom_type()), Ok(())); + } + + #[test] + fn validates_type_when_present_with_if_present() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_type_validation(TypeValidation::IfPresent), + ); + + assert_eq!(h.delta(custom_type()), Err(TrackerError::InvalidType)); + assert_eq!( + TrackerError::InvalidType.to_string(), + "Expected 'function' type." + ); + + // A missing type is accepted. + assert_eq!(h.delta(no_type()), Ok(())); + } + + #[test] + fn requires_function_type_with_required() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_type_validation(TypeValidation::Required), + ); + + assert_eq!(h.delta(no_type()), Err(TrackerError::InvalidType)); + + assert_eq!(h.delta(start(0, "call_1", "fn", "")), Ok(())); + } +} + +mod flush { + use super::*; + + #[test] + fn finalizes_unfinished_tool_calls() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{\"key\": \"val")) + .unwrap(); + h.clear(); + + h.flush(); + + assert_eq!( + h.parts, + vec![ + input_end("call_1"), + tool_call("call_1", "fn", "{\"key\": \"val"), + ] + ); + } + + #[test] + fn does_not_re_finalize_already_finished_tool_calls() { + let mut h = Harness::new(); + + h.delta(start(0, "call_1", "fn", "{}")).unwrap(); + h.flush(); + h.clear(); + + h.flush(); + + assert_eq!(h.parts, vec![]); + } +} + +mod metadata { + use super::*; + + fn google_tracker() -> StreamingToolCallTracker { + StreamingToolCallTracker::new() + .with_extract_metadata(|delta| { + delta.extra["extra_content"]["google"]["thought_signature"] + .as_str() + .map(|sig| json!({ "thoughtSignature": sig })) + }) + .with_build_provider_metadata(|metadata| { + metadata + .and_then(|m| m.get("thoughtSignature")) + .map(|sig| json!({ "google": { "thoughtSignature": sig } })) + }) + } + + #[test] + fn extracts_and_includes_provider_metadata_in_tool_call_parts() { + let mut h = Harness::with_meta(google_tracker()); + + h.delta( + start(0, "call_1", "fn", "{}") + .extra(json!({ "extra_content": { "google": { "thought_signature": "sig123" } } })), + ) + .unwrap(); + h.flush(); + + let tool_call = h + .parts + .iter() + .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })); + assert_eq!( + tool_call, + Some(&ToolCallStreamPart::ToolCall { + tool_call_id: "call_1".into(), + tool_name: "fn".into(), + input: "{}".into(), + provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), + }) + ); + } + + #[test] + fn includes_provider_metadata_for_unfinished_tool_calls_finalized_in_flush() { + let mut h = Harness::with_meta( + StreamingToolCallTracker::new() + .with_extract_metadata(|_| Some(json!({ "custom": { "key": "value" } }))) + .with_build_provider_metadata(|metadata| { + metadata.map(|m| json!({ "provider": m })) + }), + ); + + h.delta(start(0, "call_1", "fn", "{\"incomplete")).unwrap(); + h.clear(); + + h.flush(); + + assert_eq!( + h.parts.last(), + Some(&ToolCallStreamPart::ToolCall { + tool_call_id: "call_1".into(), + tool_name: "fn".into(), + input: "{\"incomplete".into(), + provider_metadata: Some(json!({ "provider": { "custom": { "key": "value" } } })), + }) + ); + } + + #[test] + fn omits_provider_metadata_when_the_builder_returns_none() { + let mut h = Harness::with_meta( + StreamingToolCallTracker::<()>::new() + .with_extract_metadata(|_| None) + .with_build_provider_metadata(|_| None), + ); + + h.delta(start(0, "call_1", "fn", "{}")).unwrap(); + h.flush(); + + let tool_call = h + .parts + .iter() + .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })); + assert_eq!(tool_call, Some(&super::tool_call("call_1", "fn", "{}"))); + } +} + +mod generate_id { + use super::*; + + #[test] + fn keeps_the_wire_id_when_present() { + let mut h = Harness::with( + StreamingToolCallTracker::new().with_generate_id(|| "custom-id".to_string()), + ); + + h.delta(start(0, "call_1", "fn", "{\"key\": \"val")) + .unwrap(); + h.clear(); + + h.flush(); + + let ids: Vec = h.tool_calls().into_iter().map(|c| c.0).collect(); + assert_eq!(ids, vec!["call_1"]); + } +} From df4d09f0353b47ad56926f1c3a623e7a123e880e Mon Sep 17 00:00:00 2001 From: cunninghamcard-bit Date: Thu, 1 Oct 2026 16:09:31 +0000 Subject: [PATCH 3/6] feat(provider-utils)!: move the tool-call tracker next to its users and 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 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 (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 Claude-Session: https://claude.ai/code/session_01WGwanvo9WLU9HWD7sRR8TS --- CONTRIBUTING.md | 4 +- README.md | 8 +- aimux-provider-utils/Cargo.toml | 1 + aimux-provider-utils/src/lib.rs | 12 +- .../src/streaming_tool_call_argument_state.rs | 0 .../src/streaming_tool_call_tracker.rs | 119 +++++++++-------- ...streaming_tool_call_argument_state_test.rs | 2 +- .../tests/streaming_tool_call_tracker_test.rs | 124 +++++++++++++----- aimux-providers/src/openai/model.rs | 109 +++++---------- aimux-providers/src/openai/types.rs | 4 +- aimux-stream/src/lib.rs | 9 -- docs/PROJECT-OVERVIEW.md | 4 +- 12 files changed, 212 insertions(+), 184 deletions(-) rename {aimux-stream => aimux-provider-utils}/src/streaming_tool_call_argument_state.rs (100%) rename {aimux-stream => aimux-provider-utils}/src/streaming_tool_call_tracker.rs (87%) rename {aimux-stream => aimux-provider-utils}/tests/streaming_tool_call_argument_state_test.rs (95%) rename {aimux-stream => aimux-provider-utils}/tests/streaming_tool_call_tracker_test.rs (91%) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ab11942a7..5a38fd5a7 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -9,8 +9,8 @@ to set up a development environment, run the tests, and submit changes. aimux/ ├── aimux-core/ # Core abstractions: LanguageModel / Provider / Message / StreamPart ├── aimux-providers/ # provider implementations + cassettes (counts: docs/api/providers.md) -├── aimux-stream/ # SSE parsing, streamed tool-call tracking -├── aimux-provider-utils/ # HTTP utilities: retry, backoff, error parsing, API-key loading +├── aimux-stream/ # SSE parsing +├── aimux-provider-utils/ # HTTP utilities: retry, backoff, error parsing, API-key loading, streamed tool-call tracking ├── aimux-ffi/ # C ABI (opaque handle + JSON + push callback) for non-native bindings ├── bindings/ # Node, Python, Swift, Kotlin, Flutter, Go, C — share one Rust core ├── contract-tests/ # Shared JSON fixtures exercised across languages diff --git a/README.md b/README.md index 02ee0ed0d..32070b708 100644 --- a/README.md +++ b/README.md @@ -98,8 +98,8 @@ middleware, and telemetry per request). aimux/ ├── aimux-core # Core abstractions: LanguageModel / Provider / Message / StreamPart ├── aimux-providers # Provider implementations — registry-backed + typed (docs/api/providers.md) -├── aimux-stream # SSE parsing, streamed tool-call tracking -├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading +├── aimux-stream # SSE parsing +├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading, streamed tool-call tracking ├── aimux-ffi # C ABI (opaque handles + JSON results + owned aimux_error_t *) for non-native bindings └── tools/ # aimux-cli (cache probe) · aimux-replay · aimux-web (console) ``` @@ -122,8 +122,8 @@ cargo add aimux-core aimux-providers |-------|-------------|-----------| | `aimux-core` | Core abstractions: `LanguageModel` / `Provider` / `Message` / `StreamPart` | [crates.io](https://crates.io/crates/aimux-core) | | `aimux-providers` | Provider implementations — [registry-backed + typed](docs/api/providers.md) | [crates.io](https://crates.io/crates/aimux-providers) | -| `aimux-stream` | SSE parsing, streamed tool-call tracking | [crates.io](https://crates.io/crates/aimux-stream) | -| `aimux-provider-utils` | One-exchange HTTP helpers and typed response handlers | [crates.io](https://crates.io/crates/aimux-provider-utils) | +| `aimux-stream` | SSE parsing | [crates.io](https://crates.io/crates/aimux-stream) | +| `aimux-provider-utils` | One-exchange HTTP helpers, typed response handlers, streamed tool-call tracking | [crates.io](https://crates.io/crates/aimux-provider-utils) | | `aimux-ffi` | C ABI for non-native bindings | [crates.io](https://crates.io/crates/aimux-ffi) | **Node.js**: diff --git a/aimux-provider-utils/Cargo.toml b/aimux-provider-utils/Cargo.toml index 29a6049f2..bf427a093 100644 --- a/aimux-provider-utils/Cargo.toml +++ b/aimux-provider-utils/Cargo.toml @@ -16,6 +16,7 @@ reqwest = { workspace = true } http = "1" serde = { workspace = true } serde_json = { workspace = true } +thiserror = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index 679f24bef..fa33eb5df 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -3,7 +3,8 @@ //! Shared utilities for provider implementations. //! //! Provides one-exchange HTTP helpers, response handlers, API key loading, -//! header management, and URL utilities — the Rust equivalents of +//! header management, URL utilities and the streamed tool-call tracker for +//! the OpenAI chat-completions wire format — the Rust equivalents of //! `@ai-sdk/provider-utils`. Operation retry and timeout live in `aimux-core`. pub mod api_key; @@ -19,6 +20,8 @@ pub mod post_to_api; pub mod read_response_with_size_limit; pub mod response_handler; pub mod retry; +pub mod streaming_tool_call_argument_state; +pub mod streaming_tool_call_tracker; pub mod url; /// WebSocket client for realtime provider APIs (RFC-0028). Empty unless the /// `ws` feature is enabled. @@ -43,4 +46,11 @@ pub use response_handler::{ stream_error_api_call, }; pub use retry::RetryConfig; +pub use streaming_tool_call_argument_state::{ + StreamingToolCallArgumentState, starts_with_structured_value, +}; +pub use streaming_tool_call_tracker::{ + StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, TrackerError, + TypeValidation, +}; pub use url::{validate_base_url, without_trailing_slash, without_trailing_slash_opt}; diff --git a/aimux-stream/src/streaming_tool_call_argument_state.rs b/aimux-provider-utils/src/streaming_tool_call_argument_state.rs similarity index 100% rename from aimux-stream/src/streaming_tool_call_argument_state.rs rename to aimux-provider-utils/src/streaming_tool_call_argument_state.rs diff --git a/aimux-stream/src/streaming_tool_call_tracker.rs b/aimux-provider-utils/src/streaming_tool_call_tracker.rs similarity index 87% rename from aimux-stream/src/streaming_tool_call_tracker.rs rename to aimux-provider-utils/src/streaming_tool_call_tracker.rs index 9a8d7240e..ce9666dd8 100644 --- a/aimux-stream/src/streaming_tool_call_tracker.rs +++ b/aimux-provider-utils/src/streaming_tool_call_tracker.rs @@ -5,8 +5,9 @@ //! //! Tracks streaming tool call state across the deltas of an OpenAI-compatible //! chat completion stream: accumulates `arguments` fragments, emits -//! `tool-input-start` / `tool-input-delta` / `tool-input-end` / `tool-call` -//! parts, and finalizes unfinished calls on [`StreamingToolCallTracker::flush`]. +//! [`StreamPart::ToolInputStart`] / [`StreamPart::ToolInputDelta`] / +//! [`StreamPart::ToolInputEnd`] / [`StreamPart::ToolCall`] parts, and +//! finalizes unfinished calls on [`StreamingToolCallTracker::flush`]. //! //! Deltas are correlated to calls by wire `id`, `index` and function name //! (see [`StreamingToolCallTracker`]'s resolution table), not by `index` @@ -16,12 +17,22 @@ //! Like the TS original, a call is *never* finalized before `flush`: a //! parsable argument buffer can still be the prefix of a longer argument //! string, so acting on it early would use truncated inputs (ai-sdk #13137). +//! +//! This is a tool for the OpenAI chat-completions wire format only. Protocols +//! whose streams carry explicit tool-call boundaries (Anthropic content +//! blocks, Google complete `functionCall` parts, Bedrock content blocks, +//! Cohere `tool-call-*` events, the Responses API's output items) do not need +//! it. use std::collections::{HashMap, HashSet}; use serde_json::Value; use thiserror::Error; +use aimux_core::error::AiMuxError; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::ProviderMetadata; + use crate::streaming_tool_call_argument_state::{ StreamingToolCallArgumentState, starts_with_structured_value, }; @@ -98,28 +109,6 @@ impl StreamingToolCallDelta { } } -/// The stream parts emitted by [`StreamingToolCallTracker`]. -/// -/// Mirrors the subset of `LanguageModelV4StreamPart` the TS tracker enqueues. -/// [`ToolCallStreamPart::ToolCall::input`] is the raw accumulated argument -/// *string*; the tracker does not parse it. -#[derive(Debug, Clone, PartialEq)] -pub enum ToolCallStreamPart { - /// Start of a tool call's input streaming. - ToolInputStart { id: String, tool_name: String }, - /// A partial argument fragment. - ToolInputDelta { id: String, delta: String }, - /// End of a tool call's input streaming. - ToolInputEnd { id: String }, - /// A complete, finalized tool call. - ToolCall { - tool_call_id: String, - tool_name: String, - input: String, - provider_metadata: Option, - }, -} - /// How to validate the `type` field on a new tool call delta. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum TypeValidation { @@ -133,6 +122,9 @@ pub enum TypeValidation { } /// Errors raised while processing a tool call delta. +/// +/// The TS tracker throws `InvalidResponseDataError`; these convert into +/// [`AiMuxError::InvalidResponseData`]. #[derive(Debug, Error, PartialEq, Eq)] pub enum TrackerError { #[error("Expected 'function.name' to be a string.")] @@ -143,7 +135,13 @@ pub enum TrackerError { IdExhausted, } -struct TrackedToolCall { +impl From for AiMuxError { + fn from(error: TrackerError) -> Self { + AiMuxError::InvalidResponseData(error.to_string()) + } +} + +struct TrackedToolCall { id: String, index: Option, sequence: usize, @@ -151,7 +149,7 @@ struct TrackedToolCall { arguments: String, argument_state: StreamingToolCallArgumentState, has_finished: bool, - metadata: Option, + metadata: Option, } enum ToolCallResolution { @@ -161,17 +159,14 @@ enum ToolCallResolution { } type GenerateIdFn = Box String + Send + Sync>; -type ExtractMetadataFn = Box Option + Send + Sync>; -type BuildMetadataFn = Box) -> Option + Send + Sync>; +type ExtractMetadataFn = Box Option + Send + Sync>; +type BuildMetadataFn = Box) -> Option + Send + Sync>; /// Tracks streaming tool call state across multiple deltas. /// /// [`process_delta`](Self::process_delta) and [`flush`](Self::flush) return -/// the parts to forward downstream (the TS tracker enqueues them on a -/// controller instead). -/// -/// The type parameter `M` is the provider-metadata type; use `()` (the -/// default) when no metadata handling is needed. +/// the [`StreamPart`]s to forward downstream (the TS tracker enqueues them on +/// a controller instead). /// /// # Correlation /// @@ -186,25 +181,25 @@ type BuildMetadataFn = Box) -> Option + Send + Sync>; /// | absent | absent | unnamed | sole unfinished call, new call, or ambiguity | /// /// An *ambiguous* delta is dropped. -pub struct StreamingToolCallTracker { - tool_calls: Vec>, +pub struct StreamingToolCallTracker { + tool_calls: Vec, tool_calls_by_id: HashMap>, tool_calls_by_index: HashMap>, used_tool_call_ids: HashSet, next_generated_id_suffixes: HashMap, generate_id: GenerateIdFn, type_validation: TypeValidation, - extract_metadata: Option>, - build_provider_metadata: Option>, + extract_metadata: Option, + build_provider_metadata: Option, } -impl Default for StreamingToolCallTracker { +impl Default for StreamingToolCallTracker { fn default() -> Self { Self::new() } } -impl StreamingToolCallTracker { +impl StreamingToolCallTracker { /// Create a tracker with no metadata handling and default settings. /// /// The default id generator returns a blank-free constant, so ids that @@ -240,10 +235,11 @@ impl StreamingToolCallTracker { } /// Set the metadata extractor (the TS `extractMetadata` option). Called - /// once when a new tool call is detected. + /// once when a new tool call is detected; the value is kept with the call + /// and handed to the provider-metadata builder when the call finalizes. #[must_use] pub fn with_extract_metadata< - F: Fn(&StreamingToolCallDelta) -> Option + Send + Sync + 'static, + F: Fn(&StreamingToolCallDelta) -> Option + Send + Sync + 'static, >( mut self, f: F, @@ -254,9 +250,11 @@ impl StreamingToolCallTracker { /// Set the provider-metadata builder (the TS /// `buildToolCallProviderMetadata` option). If it returns `None`, the - /// `tool-call` part carries no `provider_metadata`. + /// `ToolCall` part carries no `provider_metadata`. #[must_use] - pub fn with_build_provider_metadata) -> Option + Send + Sync + 'static>( + pub fn with_build_provider_metadata< + F: Fn(Option<&Value>) -> Option + Send + Sync + 'static, + >( mut self, f: F, ) -> Self { @@ -276,7 +274,7 @@ impl StreamingToolCallTracker { pub fn process_delta( &mut self, delta: &StreamingToolCallDelta, - ) -> Result>, TrackerError> { + ) -> Result, TrackerError> { let wire_name = delta.function.as_ref().and_then(|f| f.name.as_deref()); let has_blank_name = wire_name.is_some_and(|n| n.trim().is_empty()); let wire_id = non_blank(delta.id.as_deref()); @@ -323,7 +321,7 @@ impl StreamingToolCallTracker { /// Finalize any unfinished tool calls and return the closing parts. Call /// once when the stream ends. - pub fn flush(&mut self) -> Vec> { + pub fn flush(&mut self) -> Vec { // Index order is only reliable when every call has an index; for // mixed streams keep insertion order. let mut order: Vec = (0..self.tool_calls.len()).collect(); @@ -475,7 +473,7 @@ impl StreamingToolCallTracker { wire_id: Option<&str>, index: Option, name: Option<&str>, - parts: &mut Vec>, + parts: &mut Vec, ) -> Result { match self.type_validation { TypeValidation::Required => { @@ -494,9 +492,13 @@ impl StreamingToolCallTracker { let name = name.ok_or(TrackerError::MissingFunctionName)?; let id = self.create_tool_call_id(wire_id)?; - parts.push(ToolCallStreamPart::ToolInputStart { + parts.push(StreamPart::ToolInputStart { id: id.clone(), tool_name: name.to_string(), + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, }); let metadata = self @@ -526,9 +528,10 @@ impl StreamingToolCallTracker { } if !initial_arguments.is_empty() { - parts.push(ToolCallStreamPart::ToolInputDelta { + parts.push(StreamPart::ToolInputDelta { id, delta: initial_arguments, + provider_metadata: None, }); } @@ -541,7 +544,7 @@ impl StreamingToolCallTracker { &mut self, call: usize, arguments: Option<&str>, - parts: &mut Vec>, + parts: &mut Vec, ) { let tool_call = &mut self.tool_calls[call]; if tool_call.has_finished { @@ -550,9 +553,10 @@ impl StreamingToolCallTracker { if let Some(arguments) = arguments { tool_call.argument_state.append(arguments); tool_call.arguments.push_str(arguments); - parts.push(ToolCallStreamPart::ToolInputDelta { + parts.push(StreamPart::ToolInputDelta { id: tool_call.id.clone(), delta: arguments.to_string(), + provider_metadata: None, }); } } @@ -605,12 +609,13 @@ impl StreamingToolCallTracker { Err(TrackerError::IdExhausted) } - fn finish_tool_call(&mut self, call: usize, parts: &mut Vec>) { + fn finish_tool_call(&mut self, call: usize, parts: &mut Vec) { let tool_call = &mut self.tool_calls[call]; tool_call.has_finished = true; - parts.push(ToolCallStreamPart::ToolInputEnd { + parts.push(StreamPart::ToolInputEnd { id: tool_call.id.clone(), + provider_metadata: None, }); let provider_metadata = self @@ -618,10 +623,16 @@ impl StreamingToolCallTracker { .as_ref() .and_then(|build| build(tool_call.metadata.as_ref())); - parts.push(ToolCallStreamPart::ToolCall { + parts.push(StreamPart::ToolCall { tool_call_id: tool_call.id.clone(), tool_name: tool_call.function_name.clone(), - input: tool_call.arguments.clone(), + // Raw argument text; Core parses it after `stream_text`. + input: Value::String(tool_call.arguments.clone()), + provider_executed: None, + dynamic: None, + thought_signature: None, + invalid: None, + error: None, provider_metadata, }); } diff --git a/aimux-stream/tests/streaming_tool_call_argument_state_test.rs b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs similarity index 95% rename from aimux-stream/tests/streaming_tool_call_argument_state_test.rs rename to aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs index 44cf754e2..fc05f4268 100644 --- a/aimux-stream/tests/streaming_tool_call_argument_state_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs @@ -1,7 +1,7 @@ //! Port of `streaming-tool-call-argument-state.test.ts` from //! `@ai-sdk/provider-utils`. -use aimux_stream::{StreamingToolCallArgumentState, starts_with_structured_value}; +use aimux_provider_utils::{StreamingToolCallArgumentState, starts_with_structured_value}; #[test] fn starts_with_structured_value_true() { diff --git a/aimux-stream/tests/streaming_tool_call_tracker_test.rs b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs similarity index 91% rename from aimux-stream/tests/streaming_tool_call_tracker_test.rs rename to aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs index 7c55044c4..238a0ecb9 100644 --- a/aimux-stream/tests/streaming_tool_call_tracker_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs @@ -1,35 +1,99 @@ //! Port of `streaming-tool-call-tracker.test.ts` from //! `@ai-sdk/provider-utils`, case for case. +//! +//! Tracker parts are projected to a local `Part` enum so the assertions stay +//! byte-for-byte comparable with the upstream expectations. use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use aimux_stream::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, - ToolCallStreamPart, TrackerError, TypeValidation, +use aimux_core::stream_part::StreamPart; +use aimux_provider_utils::{ + StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, TrackerError, + TypeValidation, }; use serde_json::{Value, json}; -type Part = ToolCallStreamPart; +/// The TS tracker's four events, projected out of [`StreamPart`] so the +/// assertions below read like the upstream suite. +#[derive(Debug, Clone, PartialEq)] +#[allow( + clippy::enum_variant_names, + reason = "named after the upstream stream parts" +)] +enum Part { + ToolInputStart { + id: String, + tool_name: String, + }, + ToolInputDelta { + id: String, + delta: String, + }, + ToolInputEnd { + id: String, + }, + ToolCall { + tool_call_id: String, + tool_name: String, + /// The raw argument text (`StreamPart::ToolCall.input` is a + /// `Value::String` from the tracker). + input: String, + provider_metadata: Option, + }, +} + +fn project(part: StreamPart) -> Part { + match part { + StreamPart::ToolInputStart { + id, + tool_name, + provider_executed: None, + dynamic: None, + title: None, + provider_metadata: None, + } => Part::ToolInputStart { id, tool_name }, + StreamPart::ToolInputDelta { + id, + delta, + provider_metadata: None, + } => Part::ToolInputDelta { id, delta }, + StreamPart::ToolInputEnd { + id, + provider_metadata: None, + } => Part::ToolInputEnd { id }, + StreamPart::ToolCall { + tool_call_id, + tool_name, + input: Value::String(input), + provider_executed: None, + dynamic: None, + thought_signature: None, + invalid: None, + error: None, + provider_metadata, + } => Part::ToolCall { + tool_call_id, + tool_name, + input, + provider_metadata, + }, + other => panic!("tracker emitted an unexpected part: {other:?}"), + } +} /// Tracker plus the parts it has emitted so far (the TS `createCollector`). -struct Harness { - tracker: StreamingToolCallTracker, - parts: Vec>, +struct Harness { + tracker: StreamingToolCallTracker, + parts: Vec, } -impl Harness<()> { +impl Harness { fn new() -> Self { Self::with(StreamingToolCallTracker::new()) } - fn with(tracker: StreamingToolCallTracker<()>) -> Self { - Self::with_meta(tracker) - } -} - -impl Harness { - fn with_meta(tracker: StreamingToolCallTracker) -> Self { + fn with(tracker: StreamingToolCallTracker) -> Self { Self { tracker, parts: Vec::new(), @@ -38,21 +102,19 @@ impl Harness { fn delta(&mut self, delta: StreamingToolCallDelta) -> Result<(), TrackerError> { let parts = self.tracker.process_delta(&delta)?; - self.parts.extend(parts); + self.parts.extend(parts.into_iter().map(project)); Ok(()) } fn flush(&mut self) { let parts = self.tracker.flush(); - self.parts.extend(parts); + self.parts.extend(parts.into_iter().map(project)); } fn clear(&mut self) { self.parts.clear(); } -} -impl Harness { /// `(id, name, input)` of every emitted `tool-call`, in order. fn tool_calls(&self) -> Vec<(String, String, String)> { self.parts @@ -992,7 +1054,7 @@ mod flush { mod metadata { use super::*; - fn google_tracker() -> StreamingToolCallTracker { + fn google_tracker() -> StreamingToolCallTracker { StreamingToolCallTracker::new() .with_extract_metadata(|delta| { delta.extra["extra_content"]["google"]["thought_signature"] @@ -1008,7 +1070,7 @@ mod metadata { #[test] fn extracts_and_includes_provider_metadata_in_tool_call_parts() { - let mut h = Harness::with_meta(google_tracker()); + let mut h = Harness::with(google_tracker()); h.delta( start(0, "call_1", "fn", "{}") @@ -1017,13 +1079,10 @@ mod metadata { .unwrap(); h.flush(); - let tool_call = h - .parts - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })); + let tool_call = h.parts.iter().find(|p| matches!(p, Part::ToolCall { .. })); assert_eq!( tool_call, - Some(&ToolCallStreamPart::ToolCall { + Some(&Part::ToolCall { tool_call_id: "call_1".into(), tool_name: "fn".into(), input: "{}".into(), @@ -1034,7 +1093,7 @@ mod metadata { #[test] fn includes_provider_metadata_for_unfinished_tool_calls_finalized_in_flush() { - let mut h = Harness::with_meta( + let mut h = Harness::with( StreamingToolCallTracker::new() .with_extract_metadata(|_| Some(json!({ "custom": { "key": "value" } }))) .with_build_provider_metadata(|metadata| { @@ -1049,7 +1108,7 @@ mod metadata { assert_eq!( h.parts.last(), - Some(&ToolCallStreamPart::ToolCall { + Some(&Part::ToolCall { tool_call_id: "call_1".into(), tool_name: "fn".into(), input: "{\"incomplete".into(), @@ -1060,8 +1119,8 @@ mod metadata { #[test] fn omits_provider_metadata_when_the_builder_returns_none() { - let mut h = Harness::with_meta( - StreamingToolCallTracker::<()>::new() + let mut h = Harness::with( + StreamingToolCallTracker::new() .with_extract_metadata(|_| None) .with_build_provider_metadata(|_| None), ); @@ -1069,10 +1128,7 @@ mod metadata { h.delta(start(0, "call_1", "fn", "{}")).unwrap(); h.flush(); - let tool_call = h - .parts - .iter() - .find(|p| matches!(p, ToolCallStreamPart::ToolCall { .. })); + let tool_call = h.parts.iter().find(|p| matches!(p, Part::ToolCall { .. })); assert_eq!(tool_call, Some(&super::tool_call("call_1", "fn", "{}"))); } } diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index 6445d52c2..107f984ee 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -20,7 +20,9 @@ use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::HttpRequest; +use aimux_provider_utils::{ + HttpRequest, StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, +}; use super::OpenAIConfig; use super::convert::{RequestBodyResult, build_request_body_with_warnings, parse_finish_reason}; @@ -176,15 +178,6 @@ fn convert_usage(usage: &UsageResponse, usage_raw: Option<&Value>) -> Usage { } } -// ── Tool-call accumulator (streaming) ──────────────────────────────────────── - -/// Accumulates a streamed tool call's id, name, and argument fragments. -struct ToolCallAccumulator { - id: String, - name: String, - arguments: String, -} - #[async_trait] impl LanguageModel for OpenAIModel { /// Provider identity for recording/routing. Uses `config.provider` (the @@ -491,9 +484,9 @@ pub async fn execute_stream( let mut response_metadata_emitted = false; let mut final_logprobs: Option = None; - // Tool-call accumulators keyed by OpenAI's `index` field. - let mut tool_calls: HashMap = HashMap::new(); - let mut tool_call_order: Vec = Vec::new(); + // Streamed tool calls, correlated by wire id, index and function name + // (the AI SDK's StreamingToolCallTracker) and finalized on flush. + let mut tool_calls = StreamingToolCallTracker::new(); // Process the first event (already peeked) then the rest. let mut event_iter = @@ -643,49 +636,30 @@ pub async fn execute_stream( reasoning_started = false; } for dtc in tool_call_deltas { - let idx = dtc.index; - let func = dtc.function.unwrap_or_default(); - - // New tool call: has id and/or name. - let is_new = !tool_calls.contains_key(&idx); - if is_new { - let id = dtc.id.unwrap_or_default(); - let name = func.name.unwrap_or_default(); - tool_calls.insert( - idx, - ToolCallAccumulator { - id: id.clone(), - name: name.clone(), - arguments: String::new(), - }, - ); - tool_call_order.push(idx); - yield Ok(StreamPart::ToolInputStart { - id, - tool_name: name, - provider_executed: None, - dynamic: None, - title: None, - provider_metadata: None, - }); - } - - // Argument delta. - // For new tool calls, skip the delta when - // arguments are empty (matches TS — the - // initial `""` is not emitted). For - // continuation chunks, always emit (even - // empty, matching TS). - if let Some(args) = func.arguments - && (!is_new || !args.is_empty()) - && let Some(acc) = tool_calls.get_mut(&idx) { - acc.arguments.push_str(&args); - yield Ok(StreamPart::ToolInputDelta { - id: acc.id.clone(), - delta: args, - provider_metadata: None, - }); + let delta = StreamingToolCallDelta { + index: dtc.index, + id: dtc.id, + r#type: None, + function: dtc.function.map(|f| StreamingToolCallFunction { + name: f.name, + arguments: f.arguments, + }), + extra: Value::Null, + }; + match tool_calls.process_delta(&delta) { + Ok(parts) => { + for part in parts { + yield Ok(part); } + } + // A malformed delta (new call without a + // function name) is invalid response + // data, as in the AI SDK; the stream ends. + Err(error) => { + yield Err(error.into()); + return; + } + } } } @@ -763,27 +737,10 @@ pub async fn execute_stream( }); } - // A parsable argument buffer can still be a prefix of a longer input. - // Match AI SDK's tracker by finalizing only when the stream flushes. - for &idx in &tool_call_order { - if let Some(acc) = tool_calls.get(&idx) { - yield Ok(StreamPart::ToolInputEnd { - id: acc.id.clone(), - provider_metadata: None, - }); - let input = Value::String(acc.arguments.clone()); - yield Ok(StreamPart::ToolCall { - tool_call_id: acc.id.clone(), - tool_name: acc.name.clone(), - input, - provider_executed: None, - dynamic: None, - thought_signature: None, - invalid: None, - error: None, - provider_metadata: None, - }); - } + // A parsable argument buffer can still be a prefix of a longer input: + // like the AI SDK's tracker, finalize only when the stream flushes. + for part in tool_calls.flush() { + yield Ok(part); } // Build provider metadata for the Finish part. diff --git a/aimux-providers/src/openai/types.rs b/aimux-providers/src/openai/types.rs index ea78af6ec..1a8c83400 100644 --- a/aimux-providers/src/openai/types.rs +++ b/aimux-providers/src/openai/types.rs @@ -156,8 +156,10 @@ pub struct Delta { #[derive(Debug, Deserialize)] pub struct DeltaToolCall { + /// Optional: some OpenAI-compatible providers omit it (the AI SDK schema + /// is `index: z.number().nullish()`); the tracker falls back to id/name. #[serde(default)] - pub index: usize, + pub index: Option, #[serde(default)] pub id: Option, #[serde(default)] diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index d9078d8d8..65c368daf 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -6,16 +6,7 @@ pub mod lines; pub mod sse; -pub mod streaming_tool_call_argument_state; -pub mod streaming_tool_call_tracker; // Re-export the most commonly used items. pub use lines::extract_lines; pub use sse::{SseError, SseEvent, SseStream}; -pub use streaming_tool_call_argument_state::{ - StreamingToolCallArgumentState, starts_with_structured_value, -}; -pub use streaming_tool_call_tracker::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, - ToolCallStreamPart, TrackerError, TypeValidation, -}; diff --git a/docs/PROJECT-OVERVIEW.md b/docs/PROJECT-OVERVIEW.md index a53a672eb..9a734971b 100644 --- a/docs/PROJECT-OVERVIEW.md +++ b/docs/PROJECT-OVERVIEW.md @@ -165,8 +165,8 @@ aimux/ │ ├── native protocols # standalone model + convert, handles provider-specific differences │ ├── OpenAI compatible # registry-backed: provider-registry.json + provider(name, ...) entry (RFC-0017 phase 4) │ └── modalities/search # voice / image / video / search implementations -├── aimux-stream # SSE parsing, streamed tool-call tracking -├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading +├── aimux-stream # SSE parsing +├── aimux-provider-utils # One-exchange HTTP helpers, response handlers, API-key loading, streamed tool-call tracking ├── aimux-ffi # C ABI (FFI infrastructure, shared by all bindings) └── bindings/ # 6 language bindings ├── node/ # napi-rs v3 + typed TS wrapper From 7bc200eae6cdb57f9af4ce58428b4a48e79ad0d3 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 04:49:27 +0000 Subject: [PATCH 4/6] test(provider-utils): note the None/undefined mapping in the argument-state port Carried over from the SSE test port commit so the two split branches leave the file identical. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WGwanvo9WLU9HWD7sRR8TS --- .../tests/streaming_tool_call_argument_state_test.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs index fc05f4268..ea6d87076 100644 --- a/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_argument_state_test.rs @@ -12,6 +12,7 @@ fn starts_with_structured_value_true() { #[test] fn starts_with_structured_value_false() { + // TS covers `undefined` and `null` separately; Rust has one `None`. assert!(!starts_with_structured_value(None)); for value in ["", " ", "1", "\"value\""] { assert!(!starts_with_structured_value(Some(value)), "{value:?}"); From 264ade36cb15d6801a443cac8cb084415604210c Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 04:49:27 +0000 Subject: [PATCH 5/6] docs(changelog): record the tracker move to aimux-provider-utils Describes only the final API: the tracker lives in aimux-provider-utils and emits StreamPart directly; the intermediate ToolCallStreamPart- returning API from earlier commits is not mentioned. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WGwanvo9WLU9HWD7sRR8TS --- CHANGELOG.md | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index fe45c4d4b..2fde8d081 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,39 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Removed `NdjsonStream` / `NdjsonError`: nothing in the workspace used them and the AI SDK has no counterpart. `tokio` is now a dev-dependency only and the unused direct `serde` dependency is dropped. +- `StreamingToolCallTracker`, `ToolCallStreamPart` and `TrackerError` are no + longer exported from `aimux-stream`; nothing in the workspace used them and + the AI SDK keeps the tracker in `@ai-sdk/provider-utils`, next to the + provider code that drives it. See `aimux-provider-utils` below. + +**Rust (aimux-provider-utils)** + +- Gains `StreamingToolCallTracker` (plus `StreamingToolCallDelta`, + `StreamingToolCallFunction`, `TypeValidation`, `TrackerError`, + `StreamingToolCallArgumentState` and `starts_with_structured_value`), + a port of the current `@ai-sdk/provider-utils` tracker. The previous + `aimux-stream` version was a port of an older, index-only tracker. 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. It emits `aimux_core::StreamPart` tool-input parts directly + (`ToolInputStart` / `ToolInputDelta` / `ToolInputEnd` / `ToolCall`); there + is no separate `ToolCallStreamPart` event type. `TrackerError` converts + into `AiMuxError::InvalidResponseData`. Metadata hooks are typed with + `serde_json::Value` / `ProviderMetadata`, and the builder closures must be + `Send + Sync`. + +**Rust (aimux-providers)** + +- The OpenAI chat-completions stream (`openai/model.rs`, which serves every + registry-backed provider) correlates `tool_calls` deltas with the tracker + instead of by `index` alone: deltas are matched by wire id, index and + function name; a continuation without an id follows its call; indices + reused across parallel calls stay distinct; ambiguous deltas are dropped; + a call whose delta carries no 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` (previously it started a + call with an empty name). `DeltaToolCall.index` is now `Option`. ## [0.5.0] - 2026-09-27 From b0fd923a3f312ee179c2e72a62e795ea44fb6cff Mon Sep 17 00:00:00 2001 From: chenhaonan Date: Fri, 9 Oct 2026 12:15:23 +0800 Subject: [PATCH 6/6] refactor(provider-utils)!: rewrite the tool-call tracker in idiomatic Rust Same behaviour as the upstream port (checked differentially against the previous implementation on ~1.1M random deltas), with Rust structure: - Correlate by linear scan over per-call wire ids and indices instead of id/index hash maps re-sorted on every delta; no per-delta allocations besides the emitted parts. - `finish(self)` consumes the tracker, so the never-true-before-flush `has_finished` state is gone and post-flush use cannot compile. - `StreamingToolCallDelta<'a>` borrows from the provider wire type and carries `provider_metadata`, replacing `StreamingToolCallFunction`, the untyped `extra` bag, the test-only builders and the two boxed metadata hooks. - `process` / `finish` append to a caller-owned buffer. - Drop the unreachable `TrackerError::IdExhausted`. - The argument state scans bytes and models completion as its own state. Co-authored-by: Amp Amp-Thread-ID: https://ampcode.com/threads/T-01a11e83-bf7b-75bb-900b-34851e6f12ee --- CHANGELOG.md | 30 +- aimux-provider-utils/src/lib.rs | 3 +- .../src/streaming_tool_call_argument_state.rs | 143 ++-- .../src/streaming_tool_call_tracker.rs | 747 ++++++------------ .../tests/streaming_tool_call_tracker_test.rs | 158 ++-- aimux-providers/src/openai/model.rs | 44 +- aimux-stream/src/lib.rs | 5 +- bindings/node/src/types/InputRecord.ts | 4 +- bindings/node/src/types/OutcomeRecord.ts | 6 +- 9 files changed, 414 insertions(+), 726 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fde8d081..1064907c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,19 +41,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 **Rust (aimux-provider-utils)** - Gains `StreamingToolCallTracker` (plus `StreamingToolCallDelta`, - `StreamingToolCallFunction`, `TypeValidation`, `TrackerError`, - `StreamingToolCallArgumentState` and `starts_with_structured_value`), - a port of the current `@ai-sdk/provider-utils` tracker. The previous - `aimux-stream` version was a port of an older, index-only tracker. 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. It emits `aimux_core::StreamPart` tool-input parts directly - (`ToolInputStart` / `ToolInputDelta` / `ToolInputEnd` / `ToolCall`); there - is no separate `ToolCallStreamPart` event type. `TrackerError` converts - into `AiMuxError::InvalidResponseData`. Metadata hooks are typed with - `serde_json::Value` / `ProviderMetadata`, and the builder closures must be - `Send + Sync`. + `TypeValidation`, `TrackerError`, `StreamingToolCallArgumentState` and + `starts_with_structured_value`), a port of the current + `@ai-sdk/provider-utils` tracker. The previous `aimux-stream` version was + a port of an older, index-only tracker. Deltas are correlated by wire + `id`, `index` and function name (ambiguous deltas are dropped), ids are + de-duplicated with `-N` suffixes, blank function names are ignored, and + `finish` orders calls by index only when every call has one. It emits + `aimux_core::StreamPart` tool-input parts directly (`ToolInputStart` / + `ToolInputDelta` / `ToolInputEnd` / `ToolCall`); there is no separate + `ToolCallStreamPart` event type. `StreamingToolCallDelta` borrows from the + provider's wire type and flattens upstream's `function` object; + `process` and `finish` append to a caller-owned `Vec`, and + `finish` consumes the tracker. The upstream `extractMetadata` / + `buildToolCallProviderMetadata` hooks are replaced by + `StreamingToolCallDelta::provider_metadata`, read from the delta that + starts a call. `TrackerError` converts into + `AiMuxError::InvalidResponseData`. **Rust (aimux-providers)** diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index fa33eb5df..c357f9dc2 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -50,7 +50,6 @@ pub use streaming_tool_call_argument_state::{ StreamingToolCallArgumentState, starts_with_structured_value, }; pub use streaming_tool_call_tracker::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, TrackerError, - TypeValidation, + StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, }; pub use url::{validate_base_url, without_trailing_slash, without_trailing_slash_opt}; diff --git a/aimux-provider-utils/src/streaming_tool_call_argument_state.rs b/aimux-provider-utils/src/streaming_tool_call_argument_state.rs index db156e66d..3de33b7f1 100644 --- a/aimux-provider-utils/src/streaming_tool_call_argument_state.rs +++ b/aimux-provider-utils/src/streaming_tool_call_argument_state.rs @@ -1,38 +1,37 @@ //! Incremental structure tracking for streamed tool-call arguments. //! -//! Rust translation of `@ai-sdk/provider-utils`'s -//! `StreamingToolCallArgumentState` -//! (`packages/provider-utils/src/streaming-tool-call-argument-state.ts`). -//! +//! Port of `@ai-sdk/provider-utils`'s `streaming-tool-call-argument-state.ts`. //! The state is intentionally *structural* rather than a JSON parse: a //! currently parsable scalar can still be the prefix of a later value, but a //! closed top-level `{…}` / `[…]` cannot be extended. -#[derive(Debug, Clone, PartialEq, Eq)] -enum ArgumentStructure { - Undetermined, - Other, - Structured { - stack: Vec, - in_string: bool, - escaped: bool, - complete: bool, - }, -} - /// Whether `value` (ignoring leading whitespace) starts with `{` or `[`. #[must_use] pub fn starts_with_structured_value(value: Option<&str>) -> bool { - value - .and_then(|v| v.trim_start().chars().next()) - .is_some_and(|c| c == '{' || c == '[') + value.is_some_and(|v| matches!(v.trim_start().as_bytes().first(), Some(b'{' | b'['))) } /// Incrementally tracks whether streamed tool-call arguments contain a /// complete structured JSON value. #[derive(Debug, Clone)] pub struct StreamingToolCallArgumentState { - structure: ArgumentStructure, + structure: Structure, +} + +#[derive(Debug, Clone)] +enum Structure { + /// Only whitespace so far. + Undetermined, + /// Not a structured value (a scalar, or a mismatched bracket). + Other, + /// Inside a top-level `{…}` / `[…]`; `stack` holds the open brackets. + Open { + stack: Vec, + in_string: bool, + escaped: bool, + }, + /// The top-level value has been closed. + Complete, } impl StreamingToolCallArgumentState { @@ -40,7 +39,7 @@ impl StreamingToolCallArgumentState { #[must_use] pub fn new(initial_value: &str) -> Self { let mut state = Self { - structure: ArgumentStructure::Undetermined, + structure: Structure::Undetermined, }; state.append(initial_value); state @@ -49,67 +48,67 @@ impl StreamingToolCallArgumentState { /// `true` once a top-level `{…}` / `[…]` value has been closed. #[must_use] pub fn has_complete_structured_value(&self) -> bool { - matches!( - self.structure, - ArgumentStructure::Structured { complete: true, .. } - ) + matches!(self.structure, Structure::Complete) } /// Feed the next argument fragment. pub fn append(&mut self, delta: &str) { - for character in delta.chars() { - match &mut self.structure { - ArgumentStructure::Undetermined => { - if character.is_whitespace() { - continue; - } - self.structure = if character == '{' || character == '[' { - ArgumentStructure::Structured { - stack: vec![character], - in_string: false, - escaped: false, - complete: false, - } - } else { - ArgumentStructure::Other + let mut rest = delta; + if matches!(self.structure, Structure::Undetermined) { + rest = rest.trim_start(); + match rest.as_bytes().first() { + None => return, + Some(b'{' | b'[') => { + self.structure = Structure::Open { + stack: Vec::new(), + in_string: false, + escaped: false, }; } - ArgumentStructure::Other | ArgumentStructure::Structured { complete: true, .. } => { + Some(_) => { + self.structure = Structure::Other; + return; } - ArgumentStructure::Structured { - stack, - in_string, - escaped, - complete, - } => { - if *in_string { - if *escaped { - *escaped = false; - } else if character == '\\' { - *escaped = true; - } else if character == '"' { - *in_string = false; - } - continue; - } + } + } + + let Structure::Open { + stack, + in_string, + escaped, + } = &mut self.structure + else { + return; + }; - match character { - '"' => *in_string = true, - '{' | '[' => stack.push(character), - '}' | ']' => { - let expected = if character == '}' { '{' } else { '[' }; - if stack.last() != Some(&expected) { - self.structure = ArgumentStructure::Other; - continue; - } - stack.pop(); - if stack.is_empty() { - *complete = true; - } - } - _ => {} + // Every structural character is ASCII and UTF-8 continuation bytes + // never collide with ASCII, so scanning bytes is exact. + for &byte in rest.as_bytes() { + if *in_string { + if *escaped { + *escaped = false; + } else if byte == b'\\' { + *escaped = true; + } else if byte == b'"' { + *in_string = false; + } + continue; + } + match byte { + b'"' => *in_string = true, + b'{' | b'[' => stack.push(byte), + b'}' | b']' => { + let opening = if byte == b'}' { b'{' } else { b'[' }; + if stack.pop() != Some(opening) { + self.structure = Structure::Other; + return; + } + if stack.is_empty() { + self.structure = Structure::Complete; + return; } } + _ => {} } } } diff --git a/aimux-provider-utils/src/streaming_tool_call_tracker.rs b/aimux-provider-utils/src/streaming_tool_call_tracker.rs index ce9666dd8..cb9fedd43 100644 --- a/aimux-provider-utils/src/streaming_tool_call_tracker.rs +++ b/aimux-provider-utils/src/streaming_tool_call_tracker.rs @@ -1,30 +1,31 @@ -//! Streaming tool call tracker. +//! Streaming tool call tracker for the OpenAI chat-completions wire format. //! -//! Rust translation of `@ai-sdk/provider-utils`'s `StreamingToolCallTracker` -//! (`packages/provider-utils/src/streaming-tool-call-tracker.ts`). +//! Port of `@ai-sdk/provider-utils`'s `streaming-tool-call-tracker.ts`. It +//! accumulates `tool_calls[i]` deltas, emits [`StreamPart::ToolInputStart`] / +//! [`StreamPart::ToolInputDelta`] as they arrive, and emits +//! [`StreamPart::ToolInputEnd`] / [`StreamPart::ToolCall`] for every call in +//! [`StreamingToolCallTracker::finish`]. A call is never finalized earlier: a +//! parsable argument buffer can still be the prefix of a longer argument +//! string (ai-sdk #13137). //! -//! Tracks streaming tool call state across the deltas of an OpenAI-compatible -//! chat completion stream: accumulates `arguments` fragments, emits -//! [`StreamPart::ToolInputStart`] / [`StreamPart::ToolInputDelta`] / -//! [`StreamPart::ToolInputEnd`] / [`StreamPart::ToolCall`] parts, and -//! finalizes unfinished calls on [`StreamingToolCallTracker::flush`]. +//! Deltas are correlated to calls by wire `id`, `index` and function name, +//! not by `index` alone, so providers that reuse indices, repeat ids, drop +//! ids on continuations or send blank names are handled: //! -//! Deltas are correlated to calls by wire `id`, `index` and function name -//! (see [`StreamingToolCallTracker`]'s resolution table), not by `index` -//! alone, so providers that reuse indices, repeat ids, drop ids on -//! continuations or send blank names are handled. +//! | ID evidence | index/name evidence | start evidence | resolution | +//! | --- | --- | --- | --- | +//! | known | matching | any | matching call, new call, or ambiguity | +//! | known | conflicting | named | new call | +//! | unseen | matching | structured start | new call | +//! | unseen | matching | continuation | matching call or ambiguity | +//! | absent | matching | any | matching call, new call, or ambiguity | +//! | absent | absent | named | new call | +//! | absent | absent | unnamed | sole call, new call, or ambiguity | //! -//! Like the TS original, a call is *never* finalized before `flush`: a -//! parsable argument buffer can still be the prefix of a longer argument -//! string, so acting on it early would use truncated inputs (ai-sdk #13137). +//! An ambiguous delta is dropped. //! -//! This is a tool for the OpenAI chat-completions wire format only. Protocols -//! whose streams carry explicit tool-call boundaries (Anthropic content -//! blocks, Google complete `functionCall` parts, Bedrock content blocks, -//! Cohere `tool-call-*` events, the Responses API's output items) do not need -//! it. - -use std::collections::{HashMap, HashSet}; +//! Protocols whose streams carry explicit tool-call boundaries (Anthropic, +//! Google, Bedrock, Cohere, the Responses API) do not need this. use serde_json::Value; use thiserror::Error; @@ -37,102 +38,48 @@ use crate::streaming_tool_call_argument_state::{ StreamingToolCallArgumentState, starts_with_structured_value, }; -/// Fallback id used when the id generator returns a blank string. +/// Id used when the id generator returns a blank string; also the output of +/// the default generator. const FALLBACK_TOOL_CALL_ID: &str = "tool-call"; -/// The `function` sub-object of a streaming tool call delta. -#[derive(Debug, Clone, Default)] -pub struct StreamingToolCallFunction { - pub name: Option, - pub arguments: Option, -} - -/// A streaming tool call delta — the `tool_calls[i]` entry of an OpenAI-style -/// streaming chunk. +/// One `tool_calls[i]` entry of a streaming chunk, borrowed from the +/// provider's own wire type. /// -/// `arguments: null` (TS) maps to `None`; `arguments: ''` maps to `Some("")`. +/// The upstream `function: { name, arguments }` object is flattened; a +/// missing `function` is the same as both fields being `None`. #[derive(Debug, Clone, Default)] -pub struct StreamingToolCallDelta { +pub struct StreamingToolCallDelta<'a> { pub index: Option, - pub id: Option, - /// The `type` field. Named `r#type` because `type` is a reserved word. - pub r#type: Option, - pub function: Option, - /// Provider-specific payload carried alongside the standard fields, read - /// by the `extract_metadata` hook (e.g. a Google thought signature). - pub extra: Value, -} - -impl StreamingToolCallDelta { - #[must_use] - pub fn new() -> Self { - Self::default() - } - - #[must_use] - pub fn index(mut self, index: usize) -> Self { - self.index = Some(index); - self - } - - #[must_use] - pub fn id(mut self, id: impl Into) -> Self { - self.id = Some(id.into()); - self - } - - /// Set the `type` field (named `tool_type` because `type` is reserved). - #[must_use] - pub fn tool_type(mut self, t: impl Into) -> Self { - self.r#type = Some(t.into()); - self - } - - #[must_use] - pub fn function_name(mut self, name: impl Into) -> Self { - self.function.get_or_insert_with(Default::default).name = Some(name.into()); - self - } - - /// Set the `function.arguments` fragment. Pass `""` for an explicit empty - /// fragment; omit the call for `None` (TS `arguments: null`). - #[must_use] - pub fn arguments(mut self, args: impl Into) -> Self { - self.function.get_or_insert_with(Default::default).arguments = Some(args.into()); - self - } - - #[must_use] - pub fn extra(mut self, extra: Value) -> Self { - self.extra = extra; - self - } + pub id: Option<&'a str>, + pub r#type: Option<&'a str>, + pub name: Option<&'a str>, + /// `None` is the wire `null` / absent; `Some("")` is an empty fragment. + pub arguments: Option<&'a str>, + /// Provider metadata for the `ToolCall` part (the upstream + /// `extractMetadata` hook). Read only from the delta that starts a call. + pub provider_metadata: Option, } -/// How to validate the `type` field on a new tool call delta. +/// How to validate the `type` field on the delta that starts a call. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum TypeValidation { - /// No validation (default). + /// No validation. #[default] None, - /// Error if `type` is present and not `"function"`. + /// Reject a `type` that is present and not `"function"`. IfPresent, - /// Error if `type` is not exactly `"function"`. + /// Reject a `type` that is not exactly `"function"`. Required, } -/// Errors raised while processing a tool call delta. -/// -/// The TS tracker throws `InvalidResponseDataError`; these convert into -/// [`AiMuxError::InvalidResponseData`]. +/// A delta the tracker cannot accept; the upstream +/// `InvalidResponseDataError`. #[derive(Debug, Error, PartialEq, Eq)] pub enum TrackerError { #[error("Expected 'function.name' to be a string.")] MissingFunctionName, #[error("Expected 'function' type.")] InvalidType, - #[error("Failed to create a unique tool call ID.")] - IdExhausted, } impl From for AiMuxError { @@ -141,56 +88,30 @@ impl From for AiMuxError { } } -struct TrackedToolCall { +struct ToolCall { id: String, - index: Option, - sequence: usize, - function_name: String, + name: String, arguments: String, argument_state: StreamingToolCallArgumentState, - has_finished: bool, - metadata: Option, + /// Index of the starting delta; orders [`StreamingToolCallTracker::finish`]. + index: Option, + /// Every wire id and index seen for this call. + wire_ids: Vec, + indices: Vec, + provider_metadata: Option, } -enum ToolCallResolution { +enum Resolution { Existing(usize), New, Ambiguous, } -type GenerateIdFn = Box String + Send + Sync>; -type ExtractMetadataFn = Box Option + Send + Sync>; -type BuildMetadataFn = Box) -> Option + Send + Sync>; - -/// Tracks streaming tool call state across multiple deltas. -/// -/// [`process_delta`](Self::process_delta) and [`flush`](Self::flush) return -/// the [`StreamPart`]s to forward downstream (the TS tracker enqueues them on -/// a controller instead). -/// -/// # Correlation -/// -/// | ID evidence | index/name evidence | start evidence | resolution | -/// | --- | --- | --- | --- | -/// | known | matching | any | matching call, new call, or ambiguity | -/// | known | conflicting | named | new call | -/// | unseen | matching | structured start | new call | -/// | unseen | matching | continuation | matching call or ambiguity | -/// | absent | matching | any | matching call, new call, or ambiguity | -/// | absent | absent | named | new call | -/// | absent | absent | unnamed | sole unfinished call, new call, or ambiguity | -/// -/// An *ambiguous* delta is dropped. +/// Tracks the tool calls of one streamed response. See the module docs. pub struct StreamingToolCallTracker { - tool_calls: Vec, - tool_calls_by_id: HashMap>, - tool_calls_by_index: HashMap>, - used_tool_call_ids: HashSet, - next_generated_id_suffixes: HashMap, - generate_id: GenerateIdFn, + calls: Vec, + generate_id: Box String + Send + Sync>, type_validation: TypeValidation, - extract_metadata: Option, - build_provider_metadata: Option, } impl Default for StreamingToolCallTracker { @@ -200,299 +121,226 @@ impl Default for StreamingToolCallTracker { } impl StreamingToolCallTracker { - /// Create a tracker with no metadata handling and default settings. - /// - /// The default id generator returns a blank-free constant, so ids that - /// need generating become `tool-call`, `tool-call-1`, `tool-call-2`, … + /// A tracker without `type` validation whose generated ids are + /// `tool-call`, `tool-call-1`, `tool-call-2`, … #[must_use] pub fn new() -> Self { Self { - tool_calls: Vec::new(), - tool_calls_by_id: HashMap::new(), - tool_calls_by_index: HashMap::new(), - used_tool_call_ids: HashSet::new(), - next_generated_id_suffixes: HashMap::new(), + calls: Vec::new(), generate_id: Box::new(|| FALLBACK_TOOL_CALL_ID.to_string()), type_validation: TypeValidation::None, - extract_metadata: None, - build_provider_metadata: None, } } - /// Set the id generator (the TS `generateId` option). Blank or repeated - /// outputs are turned into usable unique ids. + /// Set the generator for calls without a usable wire id. Blank or + /// repeated outputs are turned into unique ids. #[must_use] - pub fn with_generate_id String + Send + Sync + 'static>(mut self, f: F) -> Self { + pub fn with_generate_id(mut self, f: impl Fn() -> String + Send + Sync + 'static) -> Self { self.generate_id = Box::new(f); self } - /// Set the `type` validation mode (the TS `typeValidation` option). - #[must_use] - pub fn with_type_validation(mut self, v: TypeValidation) -> Self { - self.type_validation = v; - self - } - - /// Set the metadata extractor (the TS `extractMetadata` option). Called - /// once when a new tool call is detected; the value is kept with the call - /// and handed to the provider-metadata builder when the call finalizes. - #[must_use] - pub fn with_extract_metadata< - F: Fn(&StreamingToolCallDelta) -> Option + Send + Sync + 'static, - >( - mut self, - f: F, - ) -> Self { - self.extract_metadata = Some(Box::new(f)); - self - } - - /// Set the provider-metadata builder (the TS - /// `buildToolCallProviderMetadata` option). If it returns `None`, the - /// `ToolCall` part carries no `provider_metadata`. #[must_use] - pub fn with_build_provider_metadata< - F: Fn(Option<&Value>) -> Option + Send + Sync + 'static, - >( - mut self, - f: F, - ) -> Self { - self.build_provider_metadata = Some(Box::new(f)); + pub fn with_type_validation(mut self, validation: TypeValidation) -> Self { + self.type_validation = validation; self } - /// Process a tool call delta from a streaming chunk and return the parts - /// it produces (possibly none). + /// Process one delta, appending the parts it produces to `out`. /// /// # Errors /// - /// Returns [`TrackerError::InvalidType`] when `type` validation fails, - /// [`TrackerError::MissingFunctionName`] when a new call has no - /// (non-null) function name, and [`TrackerError::IdExhausted`] if no - /// unique id can be produced. - pub fn process_delta( + /// A delta that starts a call without a function name, or whose `type` + /// fails validation. + pub fn process( &mut self, - delta: &StreamingToolCallDelta, - ) -> Result, TrackerError> { - let wire_name = delta.function.as_ref().and_then(|f| f.name.as_deref()); - let has_blank_name = wire_name.is_some_and(|n| n.trim().is_empty()); - let wire_id = non_blank(delta.id.as_deref()); - let name = non_blank(wire_name); - let index = delta.index; - let arguments = delta.function.as_ref().and_then(|f| f.arguments.as_deref()); - - let resolution = self.resolve_tool_call( - wire_id, + delta: StreamingToolCallDelta<'_>, + out: &mut Vec, + ) -> Result<(), TrackerError> { + let StreamingToolCallDelta { index, - name, - name.is_some() && starts_with_structured_value(arguments), - ); - - let mut parts = Vec::new(); - let call = match resolution { - ToolCallResolution::Ambiguous => return Ok(parts), - ToolCallResolution::New => { - // Blank names cannot start a usable call, but some providers - // repeat a blank name on continuations. Those were correlated - // above; only an unmatched blank-name delta is ignored. - if has_blank_name { - return Ok(parts); - } - self.process_new_tool_call(delta, wire_id, index, name, &mut parts)? + id, + r#type, + name: wire_name, + arguments, + provider_metadata, + } = delta; + let wire_id = non_blank(id); + let name = non_blank(wire_name); + let explicit_start = name.is_some() && starts_with_structured_value(arguments); + + let call = match self.resolve(wire_id, index, name, explicit_start) { + Resolution::Ambiguous => return Ok(()), + // Blank names cannot start a usable call, but some providers + // repeat a blank name on continuations; those were correlated + // above, so only an unmatched blank-name delta is ignored. + Resolution::New if wire_name.is_some() && name.is_none() => return Ok(()), + Resolution::New => { + self.validate_type(r#type)?; + let name = name.ok_or(TrackerError::MissingFunctionName)?; + self.start( + wire_id, + index, + name, + arguments.unwrap_or(""), + provider_metadata, + out, + ) } - ToolCallResolution::Existing(call) => { - if let Some(wire_id) = wire_id { - self.associate_wire_id(call, wire_id); + Resolution::Existing(call) => { + let tool_call = &mut self.calls[call]; + if let Some(wire_id) = wire_id + && !tool_call.wire_ids.iter().any(|w| w == wire_id) + { + tool_call.wire_ids.push(wire_id.to_string()); + } + if let Some(arguments) = arguments { + tool_call.argument_state.append(arguments); + tool_call.arguments.push_str(arguments); + out.push(StreamPart::ToolInputDelta { + id: tool_call.id.clone(), + delta: arguments.to_string(), + provider_metadata: None, + }); } - self.process_existing_tool_call(call, arguments, &mut parts); call } }; - if let Some(index) = index { - self.tool_calls_by_index - .entry(index) - .or_default() - .insert(call); + let indices = &mut self.calls[call].indices; + if let Some(index) = index + && !indices.contains(&index) + { + indices.push(index); } - Ok(parts) + Ok(()) } - /// Finalize any unfinished tool calls and return the closing parts. Call - /// once when the stream ends. - pub fn flush(&mut self) -> Vec { - // Index order is only reliable when every call has an index; for - // mixed streams keep insertion order. - let mut order: Vec = (0..self.tool_calls.len()).collect(); - if self.tool_calls.iter().all(|c| c.index.is_some()) { - order.sort_by_key(|&i| (self.tool_calls[i].index, self.tool_calls[i].sequence)); + /// Finalize every call, appending `ToolInputEnd` and `ToolCall` parts to + /// `out`. Calls are ordered by index when every call has one, otherwise + /// by arrival. + pub fn finish(mut self, out: &mut Vec) { + if self.calls.iter().all(|c| c.index.is_some()) { + self.calls.sort_by_key(|c| c.index); // stable: ties keep arrival order } - - let mut parts = Vec::new(); - for call in order { - if !self.tool_calls[call].has_finished { - self.finish_tool_call(call, &mut parts); - } + for call in self.calls { + out.push(StreamPart::ToolInputEnd { + id: call.id.clone(), + provider_metadata: None, + }); + out.push(StreamPart::ToolCall { + tool_call_id: call.id, + tool_name: call.name, + // Raw argument text; Core parses it. + input: Value::String(call.arguments), + provider_executed: None, + dynamic: None, + thought_signature: None, + invalid: None, + error: None, + provider_metadata: call.provider_metadata, + }); } - parts } - fn resolve_tool_call( + /// Apply the correlation table in the module docs. + fn resolve( &self, wire_id: Option<&str>, index: Option, name: Option<&str>, - has_explicit_call_start: bool, - ) -> ToolCallResolution { - let indexed = index.and_then(|i| self.tool_calls_by_index.get(&i)); - let matching_indexed = self.filter_by_name(indexed, name); - - if let Some(wire_id) = wire_id { - if let Some(with_id) = self.tool_calls_by_id.get(wire_id) { - if index.is_some() { - let matching: Vec = matching_indexed - .iter() - .copied() - .filter(|c| with_id.contains(c)) - .collect(); - let resolved = self.resolve_matching(matching, has_explicit_call_start); - if !matches!(resolved, ToolCallResolution::New) { - return resolved; - } - - // A named delta with a distinct index starts a new call - // even when its wire id and name repeat: providers may - // reuse ids across parallel calls. - if name.is_some() { - return ToolCallResolution::New; - } - - // Conflicting labels on a continuation cannot be - // resolved safely. - if indexed.is_some() { - return ToolCallResolution::Ambiguous; - } - - return self.resolve_matching(sorted(with_id), false); - } - - if let Some(name) = name { - let matching: Vec = sorted(with_id) - .into_iter() - .filter(|&c| self.tool_calls[c].function_name == name) - .collect(); - return self.resolve_matching(matching, has_explicit_call_start); - } - - return self.resolve_matching(sorted(with_id), false); + explicit_start: bool, + ) -> Resolution { + let at_index = |c: &ToolCall| index.is_some_and(|i| c.indices.contains(&i)); + let named = |c: &ToolCall| name.is_none_or(|n| c.name == n); + let index_known = self.calls.iter().any(at_index); + + let Some(wire_id) = wire_id else { + if index_known { + // Repeated names are valid on continuations; a different name + // at the same index is a new call from a provider that reuses + // indices across parallel calls. + return self.matching(|c| at_index(c) && named(c), explicit_start); } - - if !matching_indexed.is_empty() { - // A previously unseen id plus a named structured argument - // start is stronger evidence of a distinct call than a reused - // index/name; ids may still change on plain continuations. - return if has_explicit_call_start { - ToolCallResolution::New - } else { - self.resolve_matching(matching_indexed, false) - }; + if name.is_some() { + return Resolution::New; } + return self.matching(|_| true, false); + }; - return ToolCallResolution::New; + let has_id = |c: &ToolCall| c.wire_ids.iter().any(|w| w == wire_id); + if !self.calls.iter().any(has_id) { + // A previously unseen id plus a named structured start is + // stronger evidence of a distinct call than a reused index/name; + // ids may still change on plain continuations. + if explicit_start { + return Resolution::New; + } + return self.matching(|c| at_index(c) && named(c), false); } - if indexed.is_some() { - // Repeated names are valid on continuations; a different name at - // the same index means a new call from a provider that reuses - // indices across parallel calls. - return self.resolve_matching(matching_indexed, has_explicit_call_start); + if index.is_none() { + return match name { + Some(name) => self.matching(|c| has_id(c) && c.name == name, explicit_start), + None => self.matching(has_id, false), + }; } + let resolution = self.matching(|c| has_id(c) && at_index(c) && named(c), explicit_start); + if !matches!(resolution, Resolution::New) { + return resolution; + } if name.is_some() { - return ToolCallResolution::New; + // A named delta at a distinct index starts a new call even when + // its id and name repeat: providers may reuse ids across + // parallel calls. + return Resolution::New; } - - let unfinished: Vec = (0..self.tool_calls.len()) - .filter(|&c| !self.tool_calls[c].has_finished) - .collect(); - match unfinished.len() { - 0 => ToolCallResolution::New, - 1 => ToolCallResolution::Existing(unfinished[0]), - _ => ToolCallResolution::Ambiguous, + if index_known { + // Conflicting labels on a continuation cannot be resolved safely. + return Resolution::Ambiguous; } + self.matching(has_id, false) } - fn filter_by_name(&self, calls: Option<&HashSet>, name: Option<&str>) -> Vec { - calls.map_or_else(Vec::new, |calls| { - sorted(calls) - .into_iter() - .filter(|&c| name.is_none_or(|n| self.tool_calls[c].function_name == n)) - .collect() - }) - } - - fn resolve_matching( - &self, - calls: Vec, - has_explicit_call_start: bool, - ) -> ToolCallResolution { - if calls.is_empty() { - return ToolCallResolution::New; - } - - if !has_explicit_call_start { - return if calls.len() == 1 { - ToolCallResolution::Existing(calls[0]) - } else { - ToolCallResolution::Ambiguous - }; + /// Resolve to the single call matching `predicate`. With an explicit + /// structured start, a call whose structured arguments are already + /// complete cannot continue, so it is not a candidate. + fn matching(&self, predicate: impl Fn(&ToolCall) -> bool, explicit_start: bool) -> Resolution { + let mut candidates = self.calls.iter().enumerate().filter(|(_, c)| { + predicate(c) && !(explicit_start && c.argument_state.has_complete_structured_value()) + }); + match (candidates.next(), candidates.next()) { + (None, _) => Resolution::New, + (Some((call, _)), None) => Resolution::Existing(call), + (Some(_), Some(_)) => Resolution::Ambiguous, } + } - // A repeated name can occur on continuations. A fresh structured - // argument prefix signals another call only once the matching call - // has completed its own structured payload. - let continuable: Vec = calls - .into_iter() - .filter(|&c| { - !self.tool_calls[c] - .argument_state - .has_complete_structured_value() - }) - .collect(); - match continuable.len() { - 0 => ToolCallResolution::New, - 1 => ToolCallResolution::Existing(continuable[0]), - _ => ToolCallResolution::Ambiguous, + fn validate_type(&self, r#type: Option<&str>) -> Result<(), TrackerError> { + let valid = match self.type_validation { + TypeValidation::None => true, + TypeValidation::IfPresent => r#type.is_none_or(|t| t == "function"), + TypeValidation::Required => r#type == Some("function"), + }; + if valid { + Ok(()) + } else { + Err(TrackerError::InvalidType) } } - fn process_new_tool_call( + fn start( &mut self, - delta: &StreamingToolCallDelta, wire_id: Option<&str>, index: Option, - name: Option<&str>, - parts: &mut Vec, - ) -> Result { - match self.type_validation { - TypeValidation::Required => { - if delta.r#type.as_deref() != Some("function") { - return Err(TrackerError::InvalidType); - } - } - TypeValidation::IfPresent => { - if delta.r#type.as_deref().is_some_and(|t| t != "function") { - return Err(TrackerError::InvalidType); - } - } - TypeValidation::None => {} - } - - let name = name.ok_or(TrackerError::MissingFunctionName)?; - let id = self.create_tool_call_id(wire_id)?; - - parts.push(StreamPart::ToolInputStart { + name: &str, + arguments: &str, + provider_metadata: Option, + out: &mut Vec, + ) -> usize { + let id = self.unique_id(wire_id); + out.push(StreamPart::ToolInputStart { id: id.clone(), tool_name: name.to_string(), provider_executed: None, @@ -500,150 +348,49 @@ impl StreamingToolCallTracker { title: None, provider_metadata: None, }); - - let metadata = self - .extract_metadata - .as_ref() - .and_then(|extract| extract(delta)); - - let initial_arguments = delta - .function - .as_ref() - .and_then(|f| f.arguments.clone()) - .unwrap_or_default(); - - let call = self.tool_calls.len(); - self.tool_calls.push(TrackedToolCall { - id: id.clone(), - index, - sequence: call, - function_name: name.to_string(), - argument_state: StreamingToolCallArgumentState::new(&initial_arguments), - arguments: initial_arguments.clone(), - has_finished: false, - metadata, - }); - if let Some(wire_id) = wire_id { - self.associate_wire_id(call, wire_id); - } - - if !initial_arguments.is_empty() { - parts.push(StreamPart::ToolInputDelta { - id, - delta: initial_arguments, - provider_metadata: None, - }); - } - - // Tool calls must not finalize before the stream ends (ai-sdk - // #13137); finalization happens in `flush`. - Ok(call) - } - - fn process_existing_tool_call( - &mut self, - call: usize, - arguments: Option<&str>, - parts: &mut Vec, - ) { - let tool_call = &mut self.tool_calls[call]; - if tool_call.has_finished { - return; - } - if let Some(arguments) = arguments { - tool_call.argument_state.append(arguments); - tool_call.arguments.push_str(arguments); - parts.push(StreamPart::ToolInputDelta { - id: tool_call.id.clone(), + if !arguments.is_empty() { + out.push(StreamPart::ToolInputDelta { + id: id.clone(), delta: arguments.to_string(), provider_metadata: None, }); } + self.calls.push(ToolCall { + id, + name: name.to_string(), + arguments: arguments.to_string(), + argument_state: StreamingToolCallArgumentState::new(arguments), + index, + wire_ids: wire_id.map(str::to_string).into_iter().collect(), + indices: Vec::new(), + provider_metadata, + }); + self.calls.len() - 1 } - fn associate_wire_id(&mut self, call: usize, wire_id: &str) { - self.tool_calls_by_id - .entry(wire_id.to_string()) - .or_default() - .insert(call); - } - - fn create_tool_call_id(&mut self, wire_id: Option<&str>) -> Result { + /// The wire id if unused, else the generated id, else the generated id + /// with the first free `-N` suffix. + fn unique_id(&self, wire_id: Option<&str>) -> String { + let used = |id: &str| self.calls.iter().any(|c| c.id == id); if let Some(wire_id) = wire_id - && !self.used_tool_call_ids.contains(wire_id) + && !used(wire_id) { - self.used_tool_call_ids.insert(wire_id.to_string()); - return Ok(wire_id.to_string()); - } - - let generated = non_blank(Some(&(self.generate_id)())) - .unwrap_or(FALLBACK_TOOL_CALL_ID) - .to_string(); - - if !self.used_tool_call_ids.contains(&generated) { - self.used_tool_call_ids.insert(generated.clone()); - return Ok(generated); + return wire_id.to_string(); } - - // Resume after the last suffix checked for this generated value so - // deterministic generators stay bounded without rescanning occupied - // suffixes. - let initial_suffix = self - .next_generated_id_suffixes - .get(&generated) - .copied() - .unwrap_or(1); - let maximum_suffix = initial_suffix + self.used_tool_call_ids.len(); - for suffix in initial_suffix..=maximum_suffix { - let suffixed = format!("{generated}-{suffix}"); - if !self.used_tool_call_ids.contains(&suffixed) { - self.used_tool_call_ids.insert(suffixed.clone()); - self.next_generated_id_suffixes - .insert(generated, suffix + 1); - return Ok(suffixed); - } + let generated = (self.generate_id)(); + let generated = non_blank(Some(&generated)).unwrap_or(FALLBACK_TOOL_CALL_ID); + if !used(generated) { + return generated.to_string(); } - - // Unreachable by the pigeonhole principle; guards the invariant - // instead of looping without bound. - Err(TrackerError::IdExhausted) - } - - fn finish_tool_call(&mut self, call: usize, parts: &mut Vec) { - let tool_call = &mut self.tool_calls[call]; - tool_call.has_finished = true; - - parts.push(StreamPart::ToolInputEnd { - id: tool_call.id.clone(), - provider_metadata: None, - }); - - let provider_metadata = self - .build_provider_metadata - .as_ref() - .and_then(|build| build(tool_call.metadata.as_ref())); - - parts.push(StreamPart::ToolCall { - tool_call_id: tool_call.id.clone(), - tool_name: tool_call.function_name.clone(), - // Raw argument text; Core parses it after `stream_text`. - input: Value::String(tool_call.arguments.clone()), - provider_executed: None, - dynamic: None, - thought_signature: None, - invalid: None, - error: None, - provider_metadata, - }); + // Ids are never released, so the first free suffix equals what the + // upstream resume-from-last-suffix search returns. + (1..) + .map(|suffix| format!("{generated}-{suffix}")) + .find(|id| !used(id)) + .expect("finitely many ids are in use") } } fn non_blank(value: Option<&str>) -> Option<&str> { value.filter(|v| !v.trim().is_empty()) } - -fn sorted(calls: &HashSet) -> Vec { - let mut calls: Vec = calls.iter().copied().collect(); - calls.sort_unstable(); - calls -} diff --git a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs index 238a0ecb9..82b6b2726 100644 --- a/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs +++ b/aimux-provider-utils/tests/streaming_tool_call_tracker_test.rs @@ -9,8 +9,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use aimux_core::stream_part::StreamPart; use aimux_provider_utils::{ - StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, TrackerError, - TypeValidation, + StreamingToolCallDelta, StreamingToolCallTracker, TrackerError, TypeValidation, }; use serde_json::{Value, json}; @@ -84,7 +83,8 @@ fn project(part: StreamPart) -> Part { /// Tracker plus the parts it has emitted so far (the TS `createCollector`). struct Harness { - tracker: StreamingToolCallTracker, + /// `None` once flushed: `finish` consumes the tracker. + tracker: Option, parts: Vec, } @@ -95,19 +95,25 @@ impl Harness { fn with(tracker: StreamingToolCallTracker) -> Self { Self { - tracker, + tracker: Some(tracker), parts: Vec::new(), } } - fn delta(&mut self, delta: StreamingToolCallDelta) -> Result<(), TrackerError> { - let parts = self.tracker.process_delta(&delta)?; + fn delta(&mut self, delta: StreamingToolCallDelta<'_>) -> Result<(), TrackerError> { + let mut parts = Vec::new(); + let tracker = self.tracker.as_mut().expect("delta after flush"); + let result = tracker.process(delta, &mut parts); self.parts.extend(parts.into_iter().map(project)); - Ok(()) + result } fn flush(&mut self) { - let parts = self.tracker.flush(); + let mut parts = Vec::new(); + self.tracker + .take() + .expect("flushed twice") + .finish(&mut parts); self.parts.extend(parts.into_iter().map(project)); } @@ -133,27 +139,30 @@ impl Harness { } /// A delta with an explicit `function` object (as every TS test passes one). -fn delta( +fn delta<'a>( index: Option, - id: Option<&str>, - ty: Option<&str>, - name: Option<&str>, - arguments: Option<&str>, -) -> StreamingToolCallDelta { + id: Option<&'a str>, + ty: Option<&'a str>, + name: Option<&'a str>, + arguments: Option<&'a str>, +) -> StreamingToolCallDelta<'a> { StreamingToolCallDelta { index, - id: id.map(str::to_string), - r#type: ty.map(str::to_string), - function: Some(StreamingToolCallFunction { - name: name.map(str::to_string), - arguments: arguments.map(str::to_string), - }), - extra: Value::Null, + id, + r#type: ty, + name, + arguments, + provider_metadata: None, } } /// Full call start: index, id, `type: 'function'`, name and arguments. -fn start(index: usize, id: &str, name: &str, arguments: &str) -> StreamingToolCallDelta { +fn start<'a>( + index: usize, + id: &'a str, + name: &'a str, + arguments: &'a str, +) -> StreamingToolCallDelta<'a> { delta( Some(index), Some(id), @@ -164,7 +173,7 @@ fn start(index: usize, id: &str, name: &str, arguments: &str) -> StreamingToolCa } /// Continuation carrying only arguments (plus an optional index). -fn cont(index: Option, arguments: &str) -> StreamingToolCallDelta { +fn cont(index: Option, arguments: &str) -> StreamingToolCallDelta<'_> { delta(index, None, None, None, Some(arguments)) } @@ -394,18 +403,6 @@ mod process_delta { assert_eq!(h.tool_calls(), vec![tc("call_1", "fn", "{\"value\":1}")]); } - #[test] - fn skips_deltas_for_already_finished_tool_calls() { - let mut h = Harness::new(); - - h.delta(start(0, "call_1", "fn", "{}")).unwrap(); - h.flush(); - h.clear(); - - h.delta(cont(Some(0), "extra")).unwrap(); - assert_eq!(h.parts, vec![]); - } - #[test] fn skips_delta_emission_when_arguments_are_null() { let mut h = Harness::new(); @@ -964,7 +961,7 @@ mod process_delta { mod type_validation { use super::*; - fn custom_type() -> StreamingToolCallDelta { + fn custom_type() -> StreamingToolCallDelta<'static> { delta( Some(0), Some("call_1"), @@ -974,7 +971,7 @@ mod type_validation { ) } - fn no_type() -> StreamingToolCallDelta { + fn no_type() -> StreamingToolCallDelta<'static> { delta(Some(0), Some("call_1"), None, Some("fn"), Some("")) } @@ -1036,100 +1033,49 @@ mod flush { ] ); } - - #[test] - fn does_not_re_finalize_already_finished_tool_calls() { - let mut h = Harness::new(); - - h.delta(start(0, "call_1", "fn", "{}")).unwrap(); - h.flush(); - h.clear(); - - h.flush(); - - assert_eq!(h.parts, vec![]); - } } mod metadata { use super::*; - fn google_tracker() -> StreamingToolCallTracker { - StreamingToolCallTracker::new() - .with_extract_metadata(|delta| { - delta.extra["extra_content"]["google"]["thought_signature"] - .as_str() - .map(|sig| json!({ "thoughtSignature": sig })) - }) - .with_build_provider_metadata(|metadata| { - metadata - .and_then(|m| m.get("thoughtSignature")) - .map(|sig| json!({ "google": { "thoughtSignature": sig } })) - }) - } + // The upstream `extractMetadata` / `buildToolCallProviderMetadata` hooks + // become `provider_metadata` on the delta that starts a call. #[test] - fn extracts_and_includes_provider_metadata_in_tool_call_parts() { - let mut h = Harness::with(google_tracker()); + fn includes_the_starting_delta_metadata_in_the_tool_call() { + let mut h = Harness::new(); - h.delta( - start(0, "call_1", "fn", "{}") - .extra(json!({ "extra_content": { "google": { "thought_signature": "sig123" } } })), - ) + h.delta(StreamingToolCallDelta { + provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), + ..start(0, "call_1", "fn", "{\"incomplete") + }) .unwrap(); h.flush(); - let tool_call = h.parts.iter().find(|p| matches!(p, Part::ToolCall { .. })); - assert_eq!( - tool_call, - Some(&Part::ToolCall { - tool_call_id: "call_1".into(), - tool_name: "fn".into(), - input: "{}".into(), - provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), - }) - ); - } - - #[test] - fn includes_provider_metadata_for_unfinished_tool_calls_finalized_in_flush() { - let mut h = Harness::with( - StreamingToolCallTracker::new() - .with_extract_metadata(|_| Some(json!({ "custom": { "key": "value" } }))) - .with_build_provider_metadata(|metadata| { - metadata.map(|m| json!({ "provider": m })) - }), - ); - - h.delta(start(0, "call_1", "fn", "{\"incomplete")).unwrap(); - h.clear(); - - h.flush(); - assert_eq!( h.parts.last(), Some(&Part::ToolCall { tool_call_id: "call_1".into(), tool_name: "fn".into(), input: "{\"incomplete".into(), - provider_metadata: Some(json!({ "provider": { "custom": { "key": "value" } } })), + provider_metadata: Some(json!({ "google": { "thoughtSignature": "sig123" } })), }) ); } #[test] - fn omits_provider_metadata_when_the_builder_returns_none() { - let mut h = Harness::with( - StreamingToolCallTracker::new() - .with_extract_metadata(|_| None) - .with_build_provider_metadata(|_| None), - ); + fn ignores_metadata_on_continuations() { + let mut h = Harness::new(); - h.delta(start(0, "call_1", "fn", "{}")).unwrap(); + h.delta(start(0, "call_1", "fn", "{")).unwrap(); + h.delta(StreamingToolCallDelta { + provider_metadata: Some(json!({ "custom": { "key": "value" } })), + ..cont(Some(0), "}") + }) + .unwrap(); h.flush(); - let tool_call = h.parts.iter().find(|p| matches!(p, Part::ToolCall { .. })); - assert_eq!(tool_call, Some(&super::tool_call("call_1", "fn", "{}"))); + assert_eq!(h.parts.last(), Some(&tool_call("call_1", "fn", "{}"))); } } diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index 107f984ee..b4267c866 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -20,9 +20,7 @@ use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::{ - HttpRequest, StreamingToolCallDelta, StreamingToolCallFunction, StreamingToolCallTracker, -}; +use aimux_provider_utils::{HttpRequest, StreamingToolCallDelta, StreamingToolCallTracker}; use super::OpenAIConfig; use super::convert::{RequestBodyResult, build_request_body_with_warnings, parse_finish_reason}; @@ -487,6 +485,7 @@ pub async fn execute_stream( // Streamed tool calls, correlated by wire id, index and function name // (the AI SDK's StreamingToolCallTracker) and finalized on flush. let mut tool_calls = StreamingToolCallTracker::new(); + let mut tool_parts = Vec::new(); // Process the first event (already peeked) then the rest. let mut event_iter = @@ -635,30 +634,24 @@ pub async fn execute_stream( }); reasoning_started = false; } - for dtc in tool_call_deltas { + for dtc in &tool_call_deltas { + let function = dtc.function.as_ref(); let delta = StreamingToolCallDelta { index: dtc.index, - id: dtc.id, - r#type: None, - function: dtc.function.map(|f| StreamingToolCallFunction { - name: f.name, - arguments: f.arguments, - }), - extra: Value::Null, + id: dtc.id.as_deref(), + name: function.and_then(|f| f.name.as_deref()), + arguments: function.and_then(|f| f.arguments.as_deref()), + ..Default::default() }; - match tool_calls.process_delta(&delta) { - Ok(parts) => { - for part in parts { - yield Ok(part); - } - } - // A malformed delta (new call without a - // function name) is invalid response - // data, as in the AI SDK; the stream ends. - Err(error) => { - yield Err(error.into()); - return; - } + // A malformed delta (new call without a + // function name) is invalid response data, + // as in the AI SDK; the stream ends. + if let Err(error) = tool_calls.process(delta, &mut tool_parts) { + yield Err(error.into()); + return; + } + for part in tool_parts.drain(..) { + yield Ok(part); } } } @@ -739,7 +732,8 @@ pub async fn execute_stream( // A parsable argument buffer can still be a prefix of a longer input: // like the AI SDK's tracker, finalize only when the stream flushes. - for part in tool_calls.flush() { + tool_calls.finish(&mut tool_parts); + for part in tool_parts.drain(..) { yield Ok(part); } diff --git a/aimux-stream/src/lib.rs b/aimux-stream/src/lib.rs index 65c368daf..ae249daf6 100644 --- a/aimux-stream/src/lib.rs +++ b/aimux-stream/src/lib.rs @@ -1,8 +1,7 @@ //! # aimux-stream //! -//! Low-level streaming primitives for SSE (Server-Sent Events) parsing and -//! streamed tool-call tracking, -//! used by provider implementations to decode model API response streams. +//! Low-level SSE (Server-Sent Events) parsing, used by provider +//! implementations to decode model API response streams. pub mod lines; pub mod sse; diff --git a/bindings/node/src/types/InputRecord.ts b/bindings/node/src/types/InputRecord.ts index c6e44df47..e187e6e3d 100644 --- a/bindings/node/src/types/InputRecord.ts +++ b/bindings/node/src/types/InputRecord.ts @@ -6,11 +6,11 @@ import type { JsonValue } from "./serde_json/JsonValue"; /** * ① 输入侧:完整调用参数,足以重建 generate_text 调用。 */ -export type InputRecord = { operation?: RecordingOperation, +export type InputRecord = { operation?: RecordingOperation, /** * 完整 prompt(消息数组,含 ContentPart::Image 等多模态)。 */ -prompt: Array, +prompt: Array, /** * 序列化的 CallOptions(abort_signal/call_id 已 serde skip); * headers/provider_options/body_overrides 已递归脱敏。 diff --git a/bindings/node/src/types/OutcomeRecord.ts b/bindings/node/src/types/OutcomeRecord.ts index 38777b88a..fb13110d6 100644 --- a/bindings/node/src/types/OutcomeRecord.ts +++ b/bindings/node/src/types/OutcomeRecord.ts @@ -5,16 +5,16 @@ import type { JsonValue } from "./serde_json/JsonValue"; /** * 最终结果摘要。 */ -export type OutcomeRecord = { +export type OutcomeRecord = { /** * Normalized decision result, independent of a provider's wire format. */ -decision_result?: JsonValue | null, status: OutcomeStatus, finish_reason: string | null, error: string | null, +decision_result?: JsonValue | null, status: OutcomeStatus, finish_reason: string | null, error: string | null, /** * Lossless structured domain error. In particular, `RetryError.errors` * keeps the complete attempt history rather than only its display text. */ -error_value: JsonValue | null, +error_value: JsonValue | null, /** * 序列化的 Usage。 */