diff --git a/CHANGELOG.md b/CHANGELOG.md index 621117f82..ea313a12b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Breaking +- Provider results carry request and response information as the AI SDK does + (`aimux-core`). `GenerateResult` has `request: Option` and + `response: Option`; `ResponseInfo` gains `id` and holds the + timestamp, model id, headers and body. `StreamResult` has `request` and + `response: Option`. The flat `request_body` and + `response_headers` fields are removed and `response` is optional; the call + layer fills a missing id, timestamp and model id. The user-facing results and + their JSON are unchanged. + - The prompt a provider receives is modelled by role (`aimux-core`), as the AI SDK's `LanguageModelV4Message`. `LanguageModelPrompt` is `Vec`: `System { content: String }`, @@ -102,8 +111,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `InvalidPromptError` and `InvalidResponseDataError`. - Text results carry the AI SDK's core `ContentPart` (`aimux-core`), eleven - kinds, instead of the provider's `GenerateContent`: `GenerateTextResult` - and `StreamTextResultAggregated` `content` is `Vec` + kinds, instead of the provider's `GenerateContent`: `StepResult.content` + (and so `GenerateTextResult.content`) is `Vec` (`TextContent` is removed; Node exports `ContentPart`, Python keeps the `TextContent` model name). New kinds: `ToolError(TypedToolError)` and `ToolApprovalResponse(ToolApprovalResponseOutput)`. A provider tool result @@ -174,11 +183,41 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `CallOptions.call_id` and `CallOptions::for_step` are removed (`aimux-core`); the V4 call options have no call id. The recording - correlation id is `recording_context.call_id`, and `HttpRequest::new` takes - it from there. Provider code that built `HttpRequest { call_id: + correlation id is the recording context's `call_id`, and `HttpRequest::new` + takes it from there. Provider code that built `HttpRequest { call_id: options.call_id.clone(), .. }` uses `HttpRequest::new(url, headers, options)`. The JSON is unchanged (both were internal and never serialized). +- `CallOptions.max_retries`, `CallOptions.timeout` and + `CallOptions.session_id` are removed (`aimux-core`), as in the V4 call + options. They are settings of the `generate_text` / `stream_text` + operation: `GenerateTextOptions` keeps all three, and the JSON bindings + are unchanged. The operation runs its model calls inside + `operation::OperationSettings::scope`. The router and MoA take the child + retry count from `OperationSettings::current()`, and `TraceLayer` takes + the session id from it. Code that calls `do_generate` / `do_stream` on a + router, MoA or trace layer directly wraps the call in + `OperationSettings { .. }.scope(..)`. Recordings keep the three keys in + `input.options`, and replay restores them. + +- `generate_text` accepts a provider tool result without its call only for a + provider tool that sets `supportsDeferredResults: true` (new optional + `ProviderTool.supports_deferred_results`), as in the AI SDK + (`convert-language-model-content.ts:96-103`); before, any provider tool of + that name was enough. `stream_text` passes such a result on without its + input, as `streamText` does (`stream-language-model-call.ts:760-798`). The + flag is not sent to the provider. + +- `CallOptions.recording_context` and `DecisionCallOptions.recording_context` + are removed (`aimux-core`): the V4 call options carry no recording context. + `generate_text`, `stream_text` and `decide` run the model call in + `recording::RecordingContext::scope`, and the stream in + `RecordingContext::scope_stream`; the router and MoA scope each child step + in a child context. `HttpRequest::new` and the trace layer read + `RecordingContext::current()`. Code that set the field on call options + wraps the call in `RecordingContext::scope(Some(context), ..)`. The JSON is + unchanged (the field was never serialized). + - `ProviderTool.id` is a `types::NamespacedKind`, the V4 `${string}.${string}` identifier: an id without a `.` (for example `"web_search"` instead of `"openai.web_search"`) no longer deserializes, @@ -196,6 +235,104 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 [docs/api/providers.md](docs/api/providers.md); Rust also gains `provider_names()`. +- `RetryError`'s text names the last attempt's message as the AI SDK's + `getErrorMessage` does: `Failed after 3 attempts. Last error: Too Many + Requests`, not `... Last error: API call error: HTTP 429: Too Many + Requests`. The HTTP status stays a field of the attempt error. + +- Video poll settings belong to the operation, as the AI SDK's + `generateVideo` `poll` option (`aimux-core`, `aimux-providers`). + `VideoModel::poll_config()` and `VideoPollConfig` are removed; unset + `VideoCallOptions::poll` fields use the AI SDK defaults (5000 ms interval, + 600000 ms timeout). `RunwaymlConfig` loses `poll_interval`, `timeout`, + `with_poll_interval` and `with_timeout` (RunwayML used 2 s / 300 s); pass + `poll` per call instead. Each status check, its retries included, now races + the time left in the poll budget and a hanging check fails with `Timeout` + "Video generation timed out after {timeoutMs}ms." as upstream. + +- Call-level usage and finish reasons follow the AI SDK (`aimux-core`). The + results of `generate_text` / `stream_text` / `generate_object` and the + `TextStreamPart` finish parts report usage as `LanguageModelUsage` + (`input_tokens`, `input_token_details { no_cache_tokens, + cache_read_tokens, cache_write_tokens }`, `output_tokens`, + `output_token_details { text_tokens, reasoning_tokens }`, `total_tokens`, + `raw`) and the finish reason as the unified string (`"stop"`, + `"tool-calls"`, ...) with `raw_finish_reason` beside it. Usage is summed + over steps with `add_language_model_usage`. Provider results + (`GenerateResult`, `StreamPart`) keep `Usage` and `FinishReason`. + Recordings store the call-level usage for both the generate and the + stream path. + +- `GenerateTextResult` is built from steps, as the AI SDK's + (`aimux-core`). It holds `steps: Vec`, the summed usage and the + parsed output; everything else is an accessor: `final_step()`, `content()`, + `text()`, `files()`, `sources()`, `tool_calls()`, `static_tool_calls()`, + `dynamic_tool_calls()`, `tool_results()` (and the static/dynamic + variants), `finish_reason()`, `raw_finish_reason()`, `usage()`, + `warnings()`, `response_messages()`, `output()`, and the deprecated + `reasoning()`, `reasoning_text()`, `total_usage()`, `request()`, + `response()`, `provider_metadata()` that the AI SDK also deprecates. + `StepResult` carries the step's call id, model, content, finish reason, + usage, performance, warnings, request, response (with its messages) and + provider metadata. The `raw` provider result is removed, and + `reasoning_text` is `None` when there is no reasoning. The JSON form is a + snapshot: it writes the derived values next to `steps`, `total_usage` and + `output`, and only those three are read back. + +- Structured output is the `output` option, as the AI SDK's `Output` + (`aimux-core`). `GenerateTextOptions::response_format` is replaced by + `output: Option` (`Text`, `Object`, `Array`, `Choice`, `Json`); the + parsed and schema-checked value is `GenerateTextResult::output()`. + `generate_object` is deprecated, as upstream, and is a wrapper over + `generate_text` with `Output::Object`; its `GenerateObjectResult` has + `object`, `reasoning`, `finish_reason`, `usage`, `warnings`, `request`, + `response` (`GenerateObjectResponse`: id, timestamp, model id, headers, + body) and `provider_metadata`, and no `raw`. It no longer repairs + malformed JSON before parsing. + +- `TextStreamPart` is the AI SDK's call-level stream part, its own enum and + no longer `StreamPart` (`aimux-core`). The stream opens with + `Start` and `StartStep { request, warnings }` and closes with + `FinishStep { response, usage, performance, finish_reason, + raw_finish_reason, provider_metadata }` and + `Finish { finish_reason, raw_finish_reason, total_usage }`; the provider's + `StreamStart` and `ResponseMetadata` parts are not forwarded. `TextDelta` + and `ReasoningDelta` carry `text` (was `delta`). Aborting the call ends the + stream with `Abort { reason }`. `StreamTextResult::consume()` returns a + `GenerateTextResult`; `StreamTextResultAggregated` is removed. + +- `stream_text` drops an empty `TextDelta` that carries no provider + metadata, as the AI SDK does (`aimux-core`); providers emit one, e.g. the + opening `"content": ""` chunk of OpenAI chat. + +- `TextStreamPart` carries the AI SDK's tool parts (`aimux-core`, + stream-language-model-call.ts:760-792). `ToolResult` is a + `TypedToolResult` (`tool_call_id`, `tool_name`, `input?` from the call it + answers, `output`, `provider_executed`, `provider_metadata?`, `dynamic?`, + `preliminary?`) instead of the provider's `ToolResult`; a provider result + marked `is_error` arrives as `ToolError` (`TypedToolError`, with + `error`), and an invalid client tool call is followed by a `ToolError` + holding its error message and `dynamic: true` + (stream-language-model-call.ts:709-725). `ToolOutputDenied` and + `ToolApprovalResponse` are added to complete the AI SDK's set; aimux's + stream does not produce them yet. Unlike upstream, the provider's `preliminary` flag is kept, so + a preliminary result stays out of the step content, as on the generated + path. A consumed stream's step content records the tool results, tool + errors and approval responses as they arrive (stream-text.ts:1529-1538, + :1648-1650), so an invalid call's tool error is inline after its call + there, where `generate_text` puts it after the content, as upstream does. + +- `generate_text` fails with `Tool call {id} not found.` when the model + returns a provider tool result whose call is not in the response and no + provider tool of that name was passed (`aimux-core`, + convert-language-model-content.ts:87-104). A result for a passed provider + tool is kept as a deferred result. The stream passes such results through, + without `input`, as upstream does. + +- The response messages write a tool approval request's + `input_schema_input` only when it differs from the call's input under the + AI SDK's `isDeepEqualData` (`aimux-core`, to-response-messages.ts:133-146). + **Rust (aimux-stream)** - `SseStream` is now a thin adapter over the `sse-stream` crate (WHATWG @@ -269,6 +406,28 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `InvalidResponseData` (previously it started a call with an empty name). `DeltaToolCall.index` is now `Option`. +**All bindings** + +- Typed provider tools in Java, Kotlin, Swift, Python and Flutter carry + optional `supportsDeferredResults`, omitted from JSON when absent + (provider-utils@5.0.55 src/types/tool.ts:316-326). + +- The typed results and stream parts follow the core's AI SDK shapes + (Go, Java, Kotlin, Swift, Flutter, Node, Python). `GenerateTextResult` + gains `steps`, `tool_results` and `output` and loses `raw`; its `content` + holds the final step's parsed parts; `usage` / `total_usage` are + `LanguageModelUsage` and `finish_reason` the unified string. + `GenerateObjectResult` loses `raw` and `raw_finish_reason`. + `consume_stream_text` returns a `GenerateTextResult` + (`StreamTextResultAggregated` is removed). The options' `response_format` + is replaced by `output`. Stream parts: `TextDelta` / `ReasoningDelta` carry + `text`; `StreamStart` and `ResponseMetadata` give way to `Start`, + `StartStep`, `FinishStep`, `Finish { finish_reason, raw_finish_reason, + total_usage }` and `Abort`; `ToolResult` carries `output` and the call's + `input` (was `result`, `is_error`), and `ToolError`, `ToolOutputDenied` and + `ToolApprovalResponse` are decoded. Tool-call + repair reads the top-level `tool_calls` only. + ### Fixed - Gated aimux-core's ts-rs exports behind `ts-export`, so plain `cargo test` @@ -300,6 +459,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 fixes for released versions ship from a maintenance branch. `ROADMAP.md` defers the ops protocol schema and L0 passthrough and no longer fixes C1 at nine exports. +- `docs/aisdk-architecture-alignment.md` and RFC-0031 + (`docs/ai-sdk-request-pipeline.md`), against ai 7.0.127: object generation + is the deprecated wrapper over `output` (no separate object spans), async + jobs follow `generateVideo`'s start/status retries and poll budget, the + result convenience fields are marked deprecated as upstream marks them, + the response-message rules gain the empty-reasoning, invalid-input, + `inputSchemaInput`, tool-message order and tool-error encoding rows, and + the remaining differences are listed: stream retries, the first-chunk + timeout's start, `additionalRetryableError` and per-batch idempotency + keys. ## [0.5.0] - 2026-09-27 diff --git a/README.md b/README.md index 542a7ebd7..ba5627084 100644 --- a/README.md +++ b/README.md @@ -184,8 +184,8 @@ let result = stream_text( let mut stream = result.stream; while let Some(part) = stream.next().await { match part? { - StreamPart::TextDelta { delta, .. } => print!("{}", delta), - StreamPart::Finish { .. } => println!("\n[done]"), + TextStreamPart::TextDelta { text, .. } => print!("{}", text), + TextStreamPart::Finish { .. } => println!("\n[done]"), _ => {} } } diff --git a/aimux-core/src/composite.rs b/aimux-core/src/composite.rs index 363936992..a11e5712f 100644 --- a/aimux-core/src/composite.rs +++ b/aimux-core/src/composite.rs @@ -8,7 +8,8 @@ //! //! This module holds the pieces both share: //! - [`ChildModel`] — the child-handle shape (`Arc`). -//! - `step_options` — the call options of one child step. +//! - `step_context` — the recording context of one child step. +//! - `child_max_retries` — the operation's retry count for a child call. //! - [`add_usage`] — per-field `Usage` accumulation. //! - [`extract_text`] — pull text out of a `GenerateContent` list. //! - [`build_aggregator_prompt`] — assemble the MoA aggregator prompt. @@ -21,7 +22,8 @@ use std::sync::Arc; use crate::language_model::LanguageModel; use crate::language_model_message::{LanguageModelMessage, LanguageModelPrompt}; -use crate::options::CallOptions; +use crate::operation::OperationSettings; +use crate::recording::RecordingContext; use crate::result::GenerateContent; use crate::types::{InputTokenUsage, OutputTokenUsage, Usage}; @@ -30,18 +32,20 @@ use crate::types::{InputTokenUsage, OutputTokenUsage, Usage}; /// dropped (shared by registry + composite). pub type ChildModel = Arc; -/// Clone `options` for one composite child step (Router child, MoA -/// reference or aggregator): the recording context, when present, is -/// replaced by a child context labeled `step`, so the child's exchanges and -/// retry attempts are recorded as their own step instead of blending into the -/// parent operation's. -pub(crate) fn step_options(options: &CallOptions, step: impl Into) -> CallOptions { - let mut child = options.clone(); - child.recording_context = options - .recording_context - .as_ref() - .map(|context| context.child(step.into())); - child +/// The recording context of one composite child step (Router child, MoA +/// reference or aggregator): a child of the current context labeled `step`, +/// so the child's exchanges and retry attempts are recorded as their own step +/// instead of blending into the parent operation's. The child call runs in +/// [`RecordingContext::scope`] of it. `None` when nothing is recorded. +pub(crate) fn step_context(step: impl Into) -> Option { + RecordingContext::current().map(|context| context.child(step)) +} + +/// The retry count of a child call: the `max_retries` of the operation the +/// composite runs in ([`OperationSettings::current`]). `None` outside an +/// operation, or when the operation sets none, leaves the child's default. +pub(crate) fn child_max_retries() -> Option { + OperationSettings::current().and_then(|settings| settings.max_retries) } /// Add two `Usage` values field-by-field. `Usage::raw` (provider-opaque) is diff --git a/aimux-core/src/decision_model.rs b/aimux-core/src/decision_model.rs index 51c940c88..36a867904 100644 --- a/aimux-core/src/decision_model.rs +++ b/aimux-core/src/decision_model.rs @@ -153,9 +153,6 @@ impl DecisionQuestion { #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[cfg_attr(feature = "ts-export", ts(export))] pub struct DecisionCallOptions { - #[serde(skip)] - #[ts(skip)] - pub recording_context: Option, pub state: DecisionState, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub images: Vec, @@ -180,7 +177,6 @@ impl DecisionCallOptions { state, images: Vec::new(), questions, - recording_context: None, abort_signal: None, max_retries: None, timeout: None, @@ -534,7 +530,7 @@ pub trait DecisionModel: Send + Sync { /// retry exhaustion, timeout or caller abort errors. pub async fn decide( model: &dyn DecisionModel, - mut options: DecisionCallOptions, + options: DecisionCallOptions, ) -> Result { options.validate()?; let capabilities = model.capabilities(); @@ -614,7 +610,6 @@ pub async fn decide( let context = crate::recording::recorder().map(|recorder| { let ctx = crate::recording::RecordingContext::new(crate::recording::new_call_id(), recorder); - options.recording_context = Some(ctx.clone()); ctx.recorder.record_decision_input( &ctx.call_id, &options, @@ -626,37 +621,40 @@ pub async fn decide( ctx }); let result = timeout::run( - retries.retry(|| async { - if let Some(ctx) = &options.recording_context { - let _ = ctx.start_attempt(); - } - let result = model.do_decide(&options).await?; - if result.rounding != capabilities.rounding { - return Err(AiMuxError::InvalidResponseData( - "decision rounding does not match provider capabilities".into(), - )); - } - result.validate(&options)?; - if capabilities.returns_distributions - && result.answers.values().any(|answer| { - matches!( - answer, - DecisionAnswer::Choice { - probabilities: None, - .. - } | DecisionAnswer::Score { - probabilities: None, - .. - } - ) - }) - { - return Err(AiMuxError::InvalidResponseData( - "decision provider omitted promised distributions".into(), - )); - } - Ok(result) - }), + crate::recording::RecordingContext::scope( + context.clone(), + retries.retry(|| async { + if let Some(ctx) = &context { + let _ = ctx.start_attempt(); + } + let result = model.do_decide(&options).await?; + if result.rounding != capabilities.rounding { + return Err(AiMuxError::InvalidResponseData( + "decision rounding does not match provider capabilities".into(), + )); + } + result.validate(&options)?; + if capabilities.returns_distributions + && result.answers.values().any(|answer| { + matches!( + answer, + DecisionAnswer::Choice { + probabilities: None, + .. + } | DecisionAnswer::Score { + probabilities: None, + .. + } + ) + }) + { + return Err(AiMuxError::InvalidResponseData( + "decision provider omitted promised distributions".into(), + )); + } + Ok(result) + }), + ), abort_signal.as_ref(), timeout, ) diff --git a/aimux-core/src/embedding_model.rs b/aimux-core/src/embedding_model.rs index a00620efa..63471e1bb 100644 --- a/aimux-core/src/embedding_model.rs +++ b/aimux-core/src/embedding_model.rs @@ -35,7 +35,10 @@ pub struct EmbeddingCallOptions { #[ts(skip)] pub abort_signal: Option, - /// Per-call retry override. `None` uses the model default. + /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables + /// retrying. `None` uses 2, the AI SDK default, unless the model's + /// `retry_config()` sets another count (model-level retry settings are + /// removed in the target design, docs/aisdk-architecture-alignment.md D14). pub max_retries: Option, /// Per-call operation timeout. diff --git a/aimux-core/src/error.rs b/aimux-core/src/error.rs index 75ee1a866..25458f9d8 100644 --- a/aimux-core/src/error.rs +++ b/aimux-core/src/error.rs @@ -124,12 +124,31 @@ impl RetryError { } } +/// The text the AI SDK's `getErrorMessage` gives for an attempt error +/// (`@ai-sdk/provider` `errors/get-error-message.ts`), which the retry +/// wrapper puts after "Last error:" +/// (`@ai-sdk/provider-utils` `retry-with-exponential-backoff.ts:97`). +/// +/// For an API failure that is the provider's message alone: the AI SDK's +/// `APICallError.message` never carries the HTTP status, so neither does this +/// text. (JavaScript's `Error.toString()` also prefixes the class name, such as +/// `AI_APICallError: `; Aimux errors have no JavaScript class name, so only +/// the message part is reproduced.) Every other variant is its own message. +fn retry_attempt_message(error: &AiMuxError) -> String { + match error { + AiMuxError::ApiCall(api) if !api.message.is_empty() => api.message.clone(), + AiMuxError::ApiCall(api) => api.to_string(), + other => other.to_string(), + } +} + impl std::fmt::Display for RetryError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let attempts = self.errors.len(); let Some(last_error) = self.errors.last() else { return f.write_str("Retry failed without recorded errors"); }; + let last_error = retry_attempt_message(last_error); match self.reason { RetryErrorReason::MaxRetriesExceeded => { write!( @@ -175,6 +194,9 @@ pub enum AiMuxError { #[error("JSON parse error: {0}")] JsonParse(String), + #[error("{0}")] + NoOutputGenerated(String), + #[error("invalid response data: {0}")] InvalidResponseData(String), diff --git a/aimux-core/src/generate.rs b/aimux-core/src/generate.rs index 2906b783b..d85a6e344 100644 --- a/aimux-core/src/generate.rs +++ b/aimux-core/src/generate.rs @@ -4,14 +4,18 @@ //! 1. Convert user prompt (`ModelPrompt`) to `LanguageModelPrompt` (provider-facing). //! 2. Build `CallOptions` from user-facing options. //! 3. Call `LanguageModel::do_generate` / `do_stream`. -//! 4. Extract / wrap the result for the user. +//! 4. Record the model call as a [`StepResult`] and aggregate the steps into +//! a [`GenerateTextResult`], as the AI SDK's `generateText` and +//! `streamText` do (ai@7.0.127 `generate-text/generate-text.ts`, +//! `generate-text/stream-text.ts`). Aimux runs one step per call. use std::collections::HashMap; use std::pin::Pin; +use std::sync::Mutex; use std::time::Instant; use futures::{Stream, StreamExt}; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; use serde_json::Value; use tracing::Instrument; use ts_rs::TS; @@ -20,18 +24,52 @@ use crate::error::AiMuxError; use crate::language_model::LanguageModel; use crate::language_model_message::convert_to_language_model_prompt; use crate::message::{ModelMessage, ModelPrompt}; -use crate::options::{CallOptions, ResponseFormat, ToolChoice, prepare_tool_choice}; +use crate::operation::OperationSettings; +use crate::options::{CallOptions, ToolChoice, prepare_tool_choice}; +use crate::output::Output; use crate::parse_tool_call::{ToolCallRepair, parse_tool_call}; use crate::result::{ - GenerateContent, GenerateResult, GeneratedFile, ReasoningFileOutput, ReasoningPart, Source, - StreamResult, StreamTextResultAggregated, ToolApprovalRequestOutput, + ContentPart, GenerateContent, GenerateResult, GeneratedFile, ReasoningFileOutput, + ReasoningPart, Source, StreamResult, ToolApprovalRequestOutput, TypedToolError, + TypedToolResult, +}; +use crate::shared::{RequestInfo, SharedHeaders, SharedProviderOptions}; +use crate::step_result::{ + StepModel, StepResponse, StepResponseMetadata, StepResult, StepResultPerformance, +}; +use crate::stream_part::{StreamPart, StreamPartPayload, TextStreamPart}; +use crate::tool::{Tool, ToolCall, ToolResult}; +use crate::types::{ + FinishReasonUnified, LanguageModelUsage, ProviderMetadata, ReasoningEffort, Warning, + add_language_model_usage, as_language_model_usage, }; -use crate::shared::SharedProviderOptions; -use crate::stream_part::{StreamPart, TextStreamPart}; -use crate::tool::Tool; -use crate::types::{FinishReason, ProviderMetadata, ReasoningEffort, Usage, Warning}; use crate::{AbortSignal, retry, timeout}; +fn response_metadata_defaults( + mut metadata: crate::types::ResponseMetadata, + model_id: &str, +) -> crate::types::ResponseMetadata { + metadata.id.get_or_insert_with(|| { + use rand::Rng; + rand::thread_rng() + .sample_iter(rand::distributions::Alphanumeric) + .take(16) + .map(char::from) + .collect() + }); + metadata + .timestamp + .get_or_insert_with(crate::util::rfc3339_now); + metadata.model_id.get_or_insert_with(|| model_id.to_owned()); + metadata +} + +/// Milliseconds between two instants, the unit of the AI SDK's step timing +/// (`util/now.ts` reads `performance.now()`). +fn ms_between(start: Instant, end: Instant) -> f64 { + end.saturating_duration_since(start).as_secs_f64() * 1000.0 +} + /// Matches the AI SDK's `isOutputChunk`: only chunks containing model output /// start or reset the first/chunk output timers. fn is_output_chunk(part: &StreamPart) -> bool { @@ -72,7 +110,12 @@ pub struct GenerateTextOptions { pub top_k: Option, pub presence_penalty: Option, pub frequency_penalty: Option, - pub response_format: Option, + /// The structured output to generate, the AI SDK's `output` + /// (`generate-text/output.ts`). It sets the model call's + /// `response_format` and parses the final text into + /// [`GenerateTextResult::output`]. `None` sends no response format and + /// takes the text as the output. + pub output: Option, #[ts(type = "number | null")] pub seed: Option, pub tools: Option>, @@ -85,7 +128,10 @@ pub struct GenerateTextOptions { pub instructions: Option, /// Per-call request body overrides (deep-merged). See RFC-0017. pub body_overrides: Option, - /// Per-call retry count override. `None` = provider default, `Some(0)` = disable. + /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables + /// retrying. `None` uses 2, the AI SDK default, unless the model's + /// `retry_config()` sets another count (model-level retry settings are + /// removed in the target design, docs/aisdk-architecture-alignment.md D14). pub max_retries: Option, /// Per-call timeout configuration (total / first-chunk / chunk idle). pub timeout: Option, @@ -112,7 +158,19 @@ pub struct GenerateTextOptions { } impl GenerateTextOptions { + /// The operation settings of these options: the retry count, the timeout + /// and the session id, which the model call options do not carry. + #[must_use] + pub fn operation_settings(&self) -> OperationSettings { + OperationSettings { + max_retries: self.max_retries, + timeout: self.timeout, + session_id: self.session_id.clone(), + } + } + /// Build the provider-facing `CallOptions` from user-facing options. + /// The operation settings ([`Self::operation_settings`]) are left out. #[must_use] pub fn into_call_options( self, @@ -131,7 +189,9 @@ impl GenerateTextOptions { top_k: self.top_k, presence_penalty: self.presence_penalty, frequency_penalty: self.frequency_penalty, - response_format: self.response_format, + // `responseFormat: await output?.responseFormat` + // (generate-text.ts:1048, stream-language-model-call.ts:361). + response_format: self.output.as_ref().map(Output::response_format), seed: self.seed, tools: self.tools, tool_choice, @@ -139,11 +199,7 @@ impl GenerateTextOptions { provider_options: self.provider_options, reasoning: self.reasoning, body_overrides: self.body_overrides, - max_retries: self.max_retries, - timeout: self.timeout, - session_id: self.session_id, abort_signal: self.abort_signal, - recording_context: None, include_raw_chunks: self.include_raw_chunks, } } @@ -153,88 +209,435 @@ impl GenerateTextOptions { // User-facing result types // ───────────────────────────────────────────────────────────────────────────── -/// Result of `generate_text` (user-facing). -#[derive(Debug, Serialize, Deserialize, TS)] -#[cfg_attr(feature = "ts-export", ts(export))] +/// Result of `generate_text` and of a consumed `stream_text`: the AI SDK's +/// `GenerateTextResult` (ai@7.0.127 `generate-text/generate-text-result.ts`). +/// +/// It stores the steps, the usage summed over them and the parsed output, as +/// `DefaultGenerateTextResult` does (`generate-text.ts:1691-1808`); every +/// other value is derived from the steps by an accessor. The serialized form +/// is a snapshot that writes the derived values next to the stored ones, so +/// a binding reads `text`, `tool_calls`, `usage` and the rest directly. Only +/// `steps`, `total_usage` and `output` are read back. +#[derive(Debug, Clone, Deserialize)] +#[serde(try_from = "GenerateTextResultFacts")] pub struct GenerateTextResult { - /// Ordered generated content with parsed tool calls (the core - /// [`ContentPart`](crate::result::ContentPart)). - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub content: Vec, - /// The generated text (concatenated from all text content parts). - pub text: String, - /// Tool calls requested by the model. - pub tool_calls: Vec, - /// Why generation stopped. - pub finish_reason: FinishReason, - /// Token usage. - pub usage: Usage, - /// Warnings from the provider. - pub warnings: Vec, - /// Raw provider result (for advanced use). - pub raw: GenerateResult, - // ── M7: top-level aggregation (extracted from `raw.content`) ── - /// Reasoning / thinking segments from the model. - #[serde(default)] - pub reasoning: Vec, - /// Concatenated reasoning text (convenience for `reasoning.iter().map(text).join("")`). - #[serde(default)] - pub reasoning_text: String, - /// Sources / citations (search-preview models). - #[serde(default)] - pub sources: Vec, - /// Files generated by the model. - #[serde(default)] - pub files: Vec, - /// Assistant messages ready to append to the prompt for the next turn - /// (solves the multi-turn "manually build assistant message" footgun). - #[serde(default)] - pub response_messages: Vec, - /// The raw provider-specific finish reason string (e.g. "stop", - /// "end_turn", "safety"). Useful when `finish_reason.unified` is `Other`. - #[serde(default)] - pub raw_finish_reason: Option, - /// Provider-specific metadata (e.g. Anthropic cache info). Mirrored from - /// `raw.provider_metadata` for top-level convenience. - #[serde(default)] - pub provider_metadata: Option, - /// Response metadata (id, timestamp, model_id). Mirrored from - /// `raw.response` for top-level convenience. - #[serde(default)] - pub response: crate::types::ResponseMetadata, - /// Total token usage across all steps. In single-step mode (aimux's - /// default), `total_usage` equals `usage`. Provided for AI SDK parity. - #[serde(default)] - pub total_usage: Usage, + steps: Vec, + total_usage: LanguageModelUsage, + output: OutputState, +} + +/// The output of a result. +#[derive(Debug, Clone)] +pub(crate) enum OutputState { + /// Parsed by the output specification. + Parsed(Value), + /// Not parsed: `generate_text` parses only a final step that finished + /// with `stop`, or with text and not with tool calls + /// (`generate-text.ts:1599-1612`). + NotParsed, + /// Parsing a streamed final step failed. Only reading the output fails, + /// as only `StreamTextResult.output` rejects upstream + /// (`stream-text.ts:3419-3439`). + Failed(AiMuxError), } -/// Result of `generate_object` (user-facing, M12). The parsed JSON object plus -/// convenience fields from the underlying `generate_text` call. -#[derive(Debug, Serialize, Deserialize, TS)] +/// The text output, used when no output is configured (`output ?? text()`). +static TEXT_OUTPUT: Output = Output::Text {}; + +impl GenerateTextResult { + /// Aggregate the steps of a call with the usage summed over them + /// (`generate-text.ts:1526-1544`). + pub(crate) fn from_steps(steps: Vec, output: OutputState) -> Self { + let total_usage = steps + .iter() + .fold(LanguageModelUsage::default(), |total, step| { + add_language_model_usage(&total, &step.usage) + }); + Self::with_total_usage(steps, total_usage, output) + } + + fn with_total_usage( + steps: Vec, + total_usage: LanguageModelUsage, + output: OutputState, + ) -> Self { + assert!(!steps.is_empty(), "a text result has at least one step"); + Self { + steps, + total_usage, + output, + } + } + + /// Every step of the call, in order. Aimux runs one step per call. + #[must_use] + pub fn steps(&self) -> &[StepResult] { + &self.steps + } + + /// The last step. + #[must_use] + pub fn final_step(&self) -> &StepResult { + self.steps + .last() + .expect("a text result has at least one step") + } + + /// The content of every step. + #[must_use] + pub fn content(&self) -> Vec<&ContentPart> { + self.steps.iter().flat_map(|step| &step.content).collect() + } + + /// The text of the last step. + #[must_use] + pub fn text(&self) -> String { + self.final_step().text() + } + + /// The reasoning of the last step in output form. + #[deprecated(note = "use `final_step().reasoning()`, as the AI SDK deprecates \ + `GenerateTextResult.reasoning` for `finalStep.reasoning`")] + #[must_use] + pub fn reasoning(&self) -> Vec { + self.final_step().reasoning_outputs() + } + + /// The reasoning text of the last step, `None` when it is empty. + #[deprecated( + note = "use `final_step().reasoning_text()`, as the AI SDK deprecates \ + `GenerateTextResult.reasoningText` for `finalStep.reasoningText`" + )] + #[must_use] + pub fn reasoning_text(&self) -> Option { + self.final_step().reasoning_text() + } + + /// The files of every step. + #[must_use] + pub fn files(&self) -> Vec<&GeneratedFile> { + self.steps.iter().flat_map(StepResult::files).collect() + } + + /// The sources of every step. + #[must_use] + pub fn sources(&self) -> Vec<&Source> { + self.steps.iter().flat_map(StepResult::sources).collect() + } + + /// The tool calls of every step. + #[must_use] + pub fn tool_calls(&self) -> Vec<&ToolCall> { + self.steps.iter().flat_map(StepResult::tool_calls).collect() + } + + /// The tool calls of tools known in advance, from every step. + #[must_use] + pub fn static_tool_calls(&self) -> Vec<&ToolCall> { + self.steps + .iter() + .flat_map(StepResult::static_tool_calls) + .collect() + } + + /// The tool calls of dynamic tools, from every step. + #[must_use] + pub fn dynamic_tool_calls(&self) -> Vec<&ToolCall> { + self.steps + .iter() + .flat_map(StepResult::dynamic_tool_calls) + .collect() + } + + /// The tool results of every step. + #[must_use] + pub fn tool_results(&self) -> Vec<&TypedToolResult> { + self.steps + .iter() + .flat_map(StepResult::tool_results) + .collect() + } + + /// The tool results of tools known in advance, from every step. + #[must_use] + pub fn static_tool_results(&self) -> Vec<&TypedToolResult> { + self.steps + .iter() + .flat_map(StepResult::static_tool_results) + .collect() + } + + /// The tool results of dynamic tools, from every step. + #[must_use] + pub fn dynamic_tool_results(&self) -> Vec<&TypedToolResult> { + self.steps + .iter() + .flat_map(StepResult::dynamic_tool_results) + .collect() + } + + /// Why the last step finished. + #[must_use] + pub fn finish_reason(&self) -> FinishReasonUnified { + self.final_step().finish_reason + } + + /// The provider's own finish reason of the last step. + #[must_use] + pub fn raw_finish_reason(&self) -> Option<&str> { + self.final_step().raw_finish_reason.as_deref() + } + + /// The usage summed over every step. + #[must_use] + pub fn usage(&self) -> &LanguageModelUsage { + &self.total_usage + } + + /// The usage summed over every step. + #[deprecated(note = "use `usage()`, as the AI SDK deprecates \ + `GenerateTextResult.totalUsage` for `usage`")] + #[must_use] + pub fn total_usage(&self) -> &LanguageModelUsage { + &self.total_usage + } + + /// The warnings of every step. + #[must_use] + pub fn warnings(&self) -> Vec<&Warning> { + self.steps.iter().flat_map(|step| &step.warnings).collect() + } + + /// The request of the last step. + #[deprecated(note = "use `final_step().request`, as the AI SDK deprecates \ + `GenerateTextResult.request` for `finalStep.request`")] + #[must_use] + pub fn request(&self) -> &RequestInfo { + &self.final_step().request + } + + /// The response of the last step. + #[deprecated(note = "use `final_step().response`, as the AI SDK deprecates \ + `GenerateTextResult.response` for `finalStep.response`")] + #[must_use] + pub fn response(&self) -> &StepResponse { + &self.final_step().response + } + + /// The response messages of every step, ready to append to the prompt of + /// the next call. + #[must_use] + pub fn response_messages(&self) -> Vec<&ModelMessage> { + self.steps + .iter() + .flat_map(|step| &step.response.messages) + .collect() + } + + /// The provider metadata of the last step. + #[deprecated( + note = "use `final_step().provider_metadata`, as the AI SDK deprecates \ + `GenerateTextResult.providerMetadata` for `finalStep.providerMetadata`" + )] + #[must_use] + pub fn provider_metadata(&self) -> Option<&ProviderMetadata> { + self.final_step().provider_metadata.as_ref() + } + + /// The output parsed by [`GenerateTextOptions::output`] (the text when + /// no output is configured). + /// + /// # Errors + /// + /// `NoOutputGenerated` ("No output generated.", the AI SDK's + /// `NoOutputGeneratedError`) when the output was not parsed; for a + /// streamed result, the parse error when parsing failed. + pub fn output(&self) -> Result<&Value, AiMuxError> { + match &self.output { + OutputState::Parsed(value) => Ok(value), + OutputState::NotParsed => { + Err(AiMuxError::NoOutputGenerated("No output generated.".into())) + } + OutputState::Failed(error) => Err(error.clone()), + } + } + + pub(crate) fn steps_mut(&mut self) -> &mut [StepResult] { + &mut self.steps + } + + fn snapshot(&self) -> GenerateTextResultSnapshot<'_> { + let final_step = self.final_step(); + GenerateTextResultSnapshot { + content: self.content(), + text: final_step.text(), + reasoning: final_step.reasoning_outputs(), + reasoning_text: final_step.reasoning_text(), + files: self.files(), + sources: self.sources(), + tool_calls: self.tool_calls(), + tool_results: self.tool_results(), + finish_reason: final_step.finish_reason, + raw_finish_reason: final_step.raw_finish_reason.as_deref(), + usage: &self.total_usage, + total_usage: &self.total_usage, + warnings: self.warnings(), + request: &final_step.request, + response: &final_step.response, + response_messages: self.response_messages(), + provider_metadata: final_step.provider_metadata.as_ref(), + steps: &self.steps, + output: match &self.output { + OutputState::Parsed(value) => Some(value), + OutputState::NotParsed | OutputState::Failed(_) => None, + }, + } + } +} + +/// The serialized [`GenerateTextResult`]: the stored values and the values +/// the AI SDK's result derives from them (`generate-text.ts:1714-1808`). +/// `finalStep` and the static/dynamic subsets are left out: they are the +/// last entry of `steps` and a filter on `dynamic`. +#[derive(Serialize, TS)] +#[ts(rename = "GenerateTextResult")] +#[cfg_attr(feature = "ts-export", ts(export))] +struct GenerateTextResultSnapshot<'a> { + content: Vec<&'a ContentPart>, + text: String, + /// Deprecated upstream: the last step's reasoning. + reasoning: Vec, + /// Deprecated upstream: the last step's reasoning text. + reasoning_text: Option, + files: Vec<&'a GeneratedFile>, + sources: Vec<&'a Source>, + tool_calls: Vec<&'a ToolCall>, + tool_results: Vec<&'a TypedToolResult>, + finish_reason: FinishReasonUnified, + raw_finish_reason: Option<&'a str>, + usage: &'a LanguageModelUsage, + /// Deprecated upstream: equal to `usage`. + total_usage: &'a LanguageModelUsage, + warnings: Vec<&'a Warning>, + /// Deprecated upstream: the last step's request. + request: &'a RequestInfo, + /// Deprecated upstream: the last step's response. + response: &'a StepResponse, + response_messages: Vec<&'a ModelMessage>, + /// Deprecated upstream: the last step's provider metadata. + provider_metadata: Option<&'a ProviderMetadata>, + steps: &'a [StepResult], + /// Absent when the output was not parsed. + #[serde(skip_serializing_if = "Option::is_none")] + #[ts(optional)] + output: Option<&'a Value>, +} + +impl Serialize for GenerateTextResult { + fn serialize(&self, serializer: S) -> Result { + self.snapshot().serialize(serializer) + } +} + +/// The values a serialized [`GenerateTextResult`] is read back from. +#[derive(Deserialize)] +struct GenerateTextResultFacts { + steps: Vec, + total_usage: LanguageModelUsage, + #[serde(default, deserialize_with = "present_value")] + output: Option, +} + +/// A present `output` is `Some`, JSON `null` included. +fn present_value<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + Value::deserialize(deserializer).map(Some) +} + +impl TryFrom for GenerateTextResult { + type Error = String; + + fn try_from(facts: GenerateTextResultFacts) -> Result { + if facts.steps.is_empty() { + return Err("a text result has at least one step".into()); + } + Ok(Self::with_total_usage( + facts.steps, + facts.total_usage, + facts + .output + .map_or(OutputState::NotParsed, OutputState::Parsed), + )) + } +} + +/// Parse the output of a generated final step: only a step that finished with +/// `stop`, or with text and not with tool calls, is parsed, and a parse +/// failure fails the call (`generate-text.ts:1599-1612`). +fn generate_output(step: &StepResult, output: Option<&Output>) -> Result { + let text = step.text(); + if step.finish_reason == FinishReasonUnified::Stop + || (step.finish_reason != FinishReasonUnified::ToolCalls && !text.is_empty()) + { + let value = output + .unwrap_or(&TEXT_OUTPUT) + .parse_complete_output(&text)?; + Ok(OutputState::Parsed(value)) + } else { + Ok(OutputState::NotParsed) + } +} + +/// Parse the output of a streamed final step: always parsed, and a failure +/// only fails reading the output (`stream-text.ts:3419-3435`). +fn stream_output(step: &StepResult, output: Option<&Output>) -> OutputState { + match output + .unwrap_or(&TEXT_OUTPUT) + .parse_complete_output(&step.text()) + { + Ok(value) => OutputState::Parsed(value), + Err(error) => OutputState::Failed(error), + } +} + +/// Result of the deprecated [`generate_object`]: the AI SDK's +/// `GenerateObjectResult` (`generate-object/generate-object-result.ts`). +#[derive(Debug, Clone, Serialize, Deserialize, TS)] #[cfg_attr(feature = "ts-export", ts(export))] pub struct GenerateObjectResult { - /// The parsed JSON object returned by the model. + /// The generated object, as parsed by the output. pub object: Value, - /// Why generation stopped. - pub finish_reason: FinishReason, - /// Raw provider-specific finish reason string. + /// The reasoning text, `None` when there is none. + #[serde(default)] + pub reasoning: Option, + /// Why the model call finished. + pub finish_reason: FinishReasonUnified, + /// Token usage of the model call. + pub usage: LanguageModelUsage, + /// Warnings of the model call. #[serde(default)] - pub raw_finish_reason: Option, - /// Token usage. - pub usage: Usage, - /// Warnings from the provider. pub warnings: Vec, - /// Concatenated reasoning text (if the model produced reasoning/thinking). + /// The request of the model call. #[serde(default)] - pub reasoning: Option, - /// Provider-specific metadata (e.g. Anthropic cache info). + pub request: RequestInfo, + /// The response of the model call, without messages. + pub response: GenerateObjectResponse, + /// Provider metadata of the model call. #[serde(default)] pub provider_metadata: Option, - /// Response metadata (id, timestamp, model_id). +} + +/// `GenerateObjectResult.response`: the AI SDK's +/// `Omit`. +#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)] +#[cfg_attr(feature = "ts-export", ts(export))] +pub struct GenerateObjectResponse { + pub id: String, + pub timestamp: String, + pub model_id: String, + #[serde(default)] + pub headers: Option, #[serde(default)] - pub response: crate::types::ResponseMetadata, - /// The full `generate_text` result (for advanced use). - pub raw: GenerateTextResult, + pub body: Option, } /// Result of `stream_text` (user-facing). @@ -245,6 +648,9 @@ pub struct StreamTextResult { pub request_body: Option, /// Response headers. pub response_headers: Option>, + call_id: String, + model: StepModel, + output: Option, } impl std::fmt::Debug for StreamTextResult { @@ -253,243 +659,196 @@ impl std::fmt::Debug for StreamTextResult { .field("stream", &"") .field("request_body", &self.request_body) .field("response_headers", &self.response_headers) + .field("call_id", &self.call_id) + .field("model", &self.model) + .field("output", &self.output) .finish() } } impl StreamTextResult { - /// Consume the stream and collect all text deltas into a single `String`. + /// Consume the stream and return the text of the final step. /// /// # Errors /// - /// Returns the stream's error part or transport error, propagated from the - /// underlying model stream. + /// The errors of [`consume`](Self::consume). pub async fn text(self) -> Result { - let mut result = String::new(); - let mut saw_output = false; - let mut saw_finish = false; - let mut provider_error: Option = None; - let mut stream = self.stream; - while let Some(part) = stream.next().await { - let part = match part { - Ok(part) => part, - // A recoverable frame error is data on this stream (core keeps - // the stream alive past it); fold it into the same first-error - // slot as a provider error event and keep consuming so the - // trailing Finish still lands (usage, recording). - Err(error) if error.is_recoverable_stream_error() => { - if provider_error.is_none() { - provider_error = Some(error); - } - continue; - } - Err(error) => return Err(error), - }; - if is_output_chunk(&part) { - saw_output = true; - } - match part { - StreamPart::TextDelta { delta, .. } => result.push_str(&delta), - StreamPart::Finish { .. } => { - saw_finish = true; - break; - } - // Keep consuming after a provider error event: a trailing - // Finish still carries usage for the recording layer. The - // error is returned once the stream reaches its end. - StreamPart::Error { error } if provider_error.is_none() => { - provider_error = Some(error); - } - _ => {} - } - } - if let Some(error) = provider_error { - return Err(error); - } - // AI SDK semantics: an incomplete stream with no output at all is an - // error; an incomplete stream with partial output retains the partial - // result. - if !saw_finish && !saw_output { - return Err(AiMuxError::InvalidResponseData( - "No output generated. The model stream ended without a finish chunk.".into(), - )); - } - Ok(result) + self.consume().await.map(|result| result.text()) } - /// Consume the full stream and return an aggregated result (M11). + /// Consume the stream and return the result, as the AI SDK's + /// `StreamTextResult` promises resolve once the stream has finished. /// - /// Unlike [`text`](Self::text) (which returns only the concatenated text), - /// this collects reasoning, tool calls, sources, files, usage, finish - /// reason, and builds `response_messages` for the next turn. + /// Each step is recorded from the parts between its `StartStep` and + /// `FinishStep` (the event processor of `stream-text.ts:1480-1700`); the + /// total usage comes from `Finish`. /// /// # Errors /// - /// Returns the stream's error part or transport error, propagated from the - /// underlying model stream. - pub async fn consume(self) -> Result { - let mut text = String::new(); - let mut tool_calls: Vec = Vec::new(); - let mut sources: Vec = Vec::new(); - let mut files: Vec = Vec::new(); - let mut warnings: Vec = Vec::new(); - let mut finish_reason = FinishReason { - unified: crate::types::FinishReasonUnified::Stop, - raw: None, - }; - let mut raw_finish_reason: Option = None; - let mut usage = Usage::default(); - let mut finish_provider_metadata: Option = None; - let mut response: Option = None; + /// The first error part or recoverable stream error, once the stream has + /// ended; a transport error, timeout or abort error at once. A stream that + /// recorded no step fails with the abort error when it was aborted and + /// with `NoOutputGenerated` ("No output generated. Check the stream for + /// errors.") otherwise (`stream-text.ts:1724-1736`). + pub async fn consume(self) -> Result { + let Self { + stream, + call_id, + model, + output, + .. + } = self; + let mut steps: Vec = Vec::new(); let mut rm = crate::response_messages::ResponseMessageBuilder::new(); - - let mut saw_output = false; - let mut saw_finish = false; - let mut provider_error: Option = None; - let mut stream = self.stream; + let mut request = RequestInfo::default(); + let mut warnings: Vec = Vec::new(); + let mut total_usage: Option = None; + let mut aborted: Option = None; + let mut first_error: Option = None; + let mut stream = stream; while let Some(part) = stream.next().await { let part = match part { Ok(part) => part, // A recoverable frame error is data on this stream (core keeps // the stream alive past it); fold it into the same first-error - // slot as a provider error event and keep consuming so the - // trailing Finish still lands (usage, recording). + // slot as an error part and keep consuming so the trailing + // Finish still lands (usage, recording). Err(error) if error.is_recoverable_stream_error() => { - if provider_error.is_none() { - provider_error = Some(error); - } + first_error.get_or_insert(error); continue; } Err(error) => return Err(error), }; - if is_output_chunk(&part) { - saw_output = true; - } match part { - StreamPart::TextStart { + TextStreamPart::StartStep { + request: step_request, + warnings: step_warnings, + } => { + rm = crate::response_messages::ResponseMessageBuilder::new(); + request = step_request; + warnings = step_warnings; + } + TextStreamPart::TextStart { id, provider_metadata, } => rm.text_start(id, provider_metadata), - StreamPart::TextDelta { + TextStreamPart::TextDelta { id, - delta, + text, provider_metadata, - } => { - text.push_str(&delta); - rm.text_delta(id, &delta, provider_metadata); - } - StreamPart::TextEnd { + } => rm.text_delta(id, &text, provider_metadata), + TextStreamPart::TextEnd { id, provider_metadata, } => rm.text_end(id, provider_metadata), - StreamPart::ReasoningStart { + TextStreamPart::ReasoningStart { id, provider_metadata, } => rm.reasoning_start(id, provider_metadata), - StreamPart::ReasoningDelta { + TextStreamPart::ReasoningDelta { id, - delta, + text, provider_metadata, - } => rm.reasoning_delta(id, &delta, provider_metadata), - StreamPart::ReasoningEnd { + } => rm.reasoning_delta(id, &text, provider_metadata), + TextStreamPart::ReasoningEnd { id, provider_metadata, } => rm.reasoning_end(id, provider_metadata), - StreamPart::ToolCall(call) => { - rm.tool_call(&call); - tool_calls.push(call); - } - StreamPart::ToolResult(result) => rm.tool_result(result), - StreamPart::Source(source) => { - rm.source(source.clone()); - sources.push(source); - } - StreamPart::File(file) => { - rm.file(&file, false); - files.push(file); - } - StreamPart::ReasoningFile(output) => rm.file(&output.file, true), - StreamPart::Custom { + TextStreamPart::ToolCall(call) => rm.tool_call(&call), + // Tool results (not preliminary ones), tool errors and + // approval responses enter the step content as they arrive + // (stream-text.ts:1529-1538, :1648-1650). + TextStreamPart::ToolResult(result) => rm.typed_tool_result(result), + TextStreamPart::ToolError(error) => rm.tool_error(error), + TextStreamPart::ToolApprovalResponse(response) => rm.approval_response(response), + TextStreamPart::Source(source) => rm.source(source), + TextStreamPart::File(file) => rm.file(&file, false), + TextStreamPart::ReasoningFile(output) => rm.file(&output.file, true), + TextStreamPart::Custom { kind, provider_metadata, } => rm.custom(kind, provider_metadata), - StreamPart::ToolApprovalRequest(approval) => rm.approval(&approval), - StreamPart::StreamStart { warnings: w, .. } => { - warnings = w; + TextStreamPart::ToolApprovalRequest(approval) => rm.approval(&approval), + TextStreamPart::FinishStep { + response, + usage, + performance, + finish_reason, + raw_finish_reason, + provider_metadata, + } => { + let assembled = std::mem::take(&mut rm).finish(); + steps.push(StepResult { + call_id: call_id.clone(), + step_number: u32::try_from(steps.len()).unwrap_or(u32::MAX), + model: model.clone(), + content: assembled.content, + finish_reason, + raw_finish_reason, + usage, + performance, + warnings: std::mem::take(&mut warnings), + request: std::mem::take(&mut request), + response: StepResponse { + id: response.id, + timestamp: response.timestamp, + model_id: response.model_id, + headers: response.headers, + body: None, + messages: assembled.messages, + }, + provider_metadata, + }); } - StreamPart::Finish { - finish_reason: fr, - usage: u, - provider_metadata: pm, + TextStreamPart::Finish { + total_usage: usage, .. } => { - raw_finish_reason = fr.raw.clone(); - finish_reason = fr; - usage = u.clone(); - finish_provider_metadata = pm; - saw_finish = true; + total_usage = Some(usage); + break; + } + // stream-text.ts:1725-1757: an abort after a recorded step + // resolves the result (null total usage without `finish`); + // only a stream that recorded no step fails with the abort + // (:1729-1740). stream_text itself yields FinishStep and + // Finish back to back, so this order only comes from a + // caller-supplied stream. + TextStreamPart::Abort { reason } => { + aborted = Some(AiMuxError::Aborted( + reason.unwrap_or_else(|| "request aborted".into()), + )); break; } - // Keep consuming after a provider error event: a trailing - // Finish still carries usage for the recording layer. The - // error is returned once the stream reaches its end. - StreamPart::Error { error } if provider_error.is_none() => { - provider_error = Some(error); + // Keep consuming after an error part: a trailing Finish still + // carries usage for the recording layer. The error is + // returned once the stream reaches its end. + TextStreamPart::Error { error } => { + first_error.get_or_insert(error); } - StreamPart::ResponseMetadata(meta) => response = Some(meta), - _ => {} + TextStreamPart::ToolInputStart { .. } + | TextStreamPart::ToolInputDelta { .. } + | TextStreamPart::ToolInputEnd { .. } + | TextStreamPart::ToolOutputDenied(_) + | TextStreamPart::Start {} + | TextStreamPart::Raw { .. } => {} } } - if let Some(error) = provider_error { + if let Some(error) = first_error { return Err(error); } - if !saw_finish { - // AI SDK semantics: an incomplete stream with no output at all is - // an error; an incomplete stream with partial output retains the - // partial result but must not claim a normal stop. - if !saw_output { - return Err(AiMuxError::InvalidResponseData( - "No output generated. The model stream ended without a finish chunk.".into(), - )); - } - finish_reason = FinishReason { - unified: crate::types::FinishReasonUnified::Other, - raw: None, - }; + if steps.is_empty() { + return Err(aborted.unwrap_or_else(|| { + AiMuxError::NoOutputGenerated( + "No output generated. Check the stream for errors.".into(), + ) + })); } - - let crate::response_messages::ResponseMessages { - messages: response_messages, - content, - reasoning, - } = rm.finish(); - - let reasoning_text = reasoning - .iter() - .filter_map(|r| match r { - ReasoningPart::Text(r) => Some(r.text.as_str()), - ReasoningPart::File(_) => None, - }) - .collect::>() - .join(""); - - Ok(StreamTextResultAggregated { - content, - text, - reasoning, - reasoning_text, - tool_calls, - sources, - files, - finish_reason, - raw_finish_reason, - total_usage: usage.clone(), - usage, - warnings, - provider_metadata: finish_provider_metadata, - response, - response_messages, - }) + let output = stream_output(steps.last().expect("checked above"), output.as_ref()); + Ok(GenerateTextResult::with_total_usage( + steps, + total_usage.unwrap_or_default(), + output, + )) } } @@ -519,7 +878,7 @@ fn provider_id(model: &dyn LanguageModel) -> Option<&str> { /// GenerateTextOptions::default(), /// ).await?; /// -/// println!("{}", result.text); +/// println!("{}", result.text()); /// # Ok(()) /// # } /// ``` @@ -527,7 +886,8 @@ fn provider_id(model: &dyn LanguageModel) -> Option<&str> { /// # Errors /// /// Returns the error produced by the model's `do_generate` call (provider API, -/// network, or abort). +/// network, or abort), and the output's parse error when the final step's +/// text does not parse. pub async fn generate_text( model: &dyn LanguageModel, prompt: impl Into, @@ -546,18 +906,39 @@ pub(crate) async fn generate_text_from_language_model_prompt( lm_prompt: crate::language_model_message::LanguageModelPrompt, messages: Vec, options: GenerateTextOptions, +) -> Result { + let settings = options.operation_settings(); + settings + .clone() + .scope(generate_in_operation( + model, lm_prompt, messages, options, settings, + )) + .await +} + +// The operation body runs inside its settings' scope, where the model, the +// composite children and the recorder read them. +async fn generate_in_operation( + model: &dyn LanguageModel, + lm_prompt: crate::language_model_message::LanguageModelPrompt, + messages: Vec, + options: GenerateTextOptions, + settings: OperationSettings, ) -> Result { let repair_tool_call = options.repair_tool_call.clone(); let tools = options.tools.clone(); let operation_instructions = options.instructions.clone(); + let output = options.output.clone(); + if let Some(output) = &output { + output.validate()?; + } // 2. Build CallOptions. - let mut call_options = options.into_call_options(lm_prompt); - let operation_timeout = - timeout::OperationTimeout::new(call_options.timeout.unwrap_or_default())?; + let call_options = options.into_call_options(lm_prompt); + let operation_timeout = timeout::OperationTimeout::new(settings.timeout.unwrap_or_default())?; let abort_signal = call_options.abort_signal.clone(); let retries = retry::prepare_retries( - call_options.max_retries, + settings.max_retries, model.retry_config(), abort_signal.clone(), ); @@ -568,7 +949,6 @@ pub(crate) async fn generate_text_from_language_model_prompt( let context = crate::recording::recorder().map(|recorder| { let ctx = crate::recording::RecordingContext::new(crate::recording::new_call_id(), recorder); - call_options.recording_context = Some(ctx.clone()); ctx.recorder.record_input( &ctx.call_id, &call_options, @@ -593,10 +973,7 @@ pub(crate) async fn generate_text_from_language_model_prompt( // and store lookups are cheap (once-locked env check + two RwLock // reads) when no store / inferer is registered. if let (Some((session_id, source)), Some(store)) = ( - crate::session::resolve_session_id( - call_options.session_id.as_deref(), - &call_options.prompt, - ), + crate::session::resolve_session_id(settings.session_id.as_deref(), &call_options.prompt), crate::session::session_store(), ) { let call = store.append(&session_id, call_id.as_deref(), source); @@ -609,6 +986,8 @@ pub(crate) async fn generate_text_from_language_model_prompt( // 3. Call provider, inside the "generate" span (RFC-0014 §4.1 span tree). // The http layer's `http_request` span nests under this one. + // `stepStartTimestampMs` (generate-text.ts:1023) is taken before the + // retry loop, so the response time includes retries (:1068). let started = Instant::now(); let span = tracing::info_span!( target: "aimux_core::generate", @@ -617,12 +996,15 @@ pub(crate) async fn generate_text_from_language_model_prompt( model = %model.model_id(), modality = "text", ); - let operation = retries.retry(|| { - if let Some(context) = &context { - let _ = context.start_attempt(); - } - do_generate_with_logging(model, &call_options, span.clone(), started) - }); + let operation = crate::recording::RecordingContext::scope( + context.clone(), + retries.retry(|| { + if let Some(context) = &context { + let _ = context.start_attempt(); + } + do_generate_with_logging(model, &call_options, span.clone(), started) + }), + ); let result = match timeout::run(operation, abort_signal.as_ref(), operation_timeout).await { Ok(r) => r, Err(e) => { @@ -632,51 +1014,152 @@ pub(crate) async fn generate_text_from_language_model_prompt( return Err(e); } }; - // 4. Extract text, tool calls, reasoning, sources, files from content. - let mut text = String::new(); - let mut tool_calls = Vec::new(); - let mut sources = Vec::new(); - let mut files = Vec::new(); - // Build the assistant response message content parts in parallel. - let mut rm = crate::response_messages::ResponseMessageBuilder::new(); - for content in &result.content { - if let GenerateContent::ToolCall(call) = content { - tool_calls.push( - parse_tool_call( - call.clone(), - tools.as_deref(), - repair_tool_call.as_ref(), - &messages, - operation_instructions.as_deref(), - ) - .await, - ); + let response_time_ms = ms_between(started, Instant::now()); + + // 4. Parse tool calls, build the step and parse the output. Parsing, + // repair and file resolution are part of the user operation: the + // outcome is recorded once this phase has completed within the same + // deadline/cancellation scope as the provider call. + let step = async { + let mut tool_calls = Vec::new(); + for content in &result.content { + if let GenerateContent::ToolCall(call) = content { + tool_calls.push( + parse_tool_call( + call.clone(), + tools.as_deref(), + repair_tool_call.as_ref(), + &messages, + operation_instructions.as_deref(), + ) + .await, + ); + } } - } - if let Some(error) = tool_choice_violation( - call_options.tool_choice.as_ref(), - tool_calls.iter().map(|call| call.tool_name.as_str()), - result.finish_reason.unified, - model, - || result.content.clone(), - ) { - if let (Some(rec), Some(call_id)) = (&recorder, &call_id) { - rec.record_outcome( - call_id, - &crate::recording::OutcomeRecord::from_error(&error), - ); + if let Some(error) = tool_choice_violation( + call_options.tool_choice.as_ref(), + tool_calls.iter().map(|call| call.tool_name.as_str()), + result.finish_reason.unified, + model, + || result.content.clone(), + ) { + return Err(error); } - return Err(error); + // The step time ends after tool handling and before content + // conversion (generate-text.ts:1382, :1437-1446). + let step_time_ms = ms_between(started, Instant::now()); + let content = convert_generated_content( + &result, + &tool_calls, + tools.as_deref(), + abort_signal.as_ref(), + operation_timeout, + ) + .await?; + Ok::<_, AiMuxError>(build_generated_step( + model, + call_id.clone(), + result, + content, + response_time_ms, + step_time_ms, + )) + } + .await; + let outcome = step.and_then(|step| { + let output = generate_output(&step, output.as_ref())?; + Ok(GenerateTextResult::from_steps(vec![step], output)) + }); + if let (Some(rec), Some(call_id)) = (&recorder, &call_id) { + rec.record_outcome( + call_id, + &match &outcome { + Ok(result) => crate::recording::OutcomeRecord::success( + result.finish_reason(), + &result.final_step().usage, + ), + Err(error) => crate::recording::OutcomeRecord::from_error(error), + }, + ); } + outcome +} + +/// The tool error that follows a client tool call that could not be parsed +/// (`stream-language-model-call.ts:709-725`): the call's input, its error +/// message, and `dynamic`. +pub(crate) fn invalid_tool_call_error(call: &ToolCall) -> Option { + (call.invalid == Some(true) && call.provider_executed != Some(true)).then(|| TypedToolError { + tool_call_id: call.tool_call_id.clone(), + tool_name: call.tool_name.clone(), + input: Some(call.input.clone()), + error: Value::String( + call.error + .as_ref() + .map_or_else(|| "unknown error".to_owned(), ToString::to_string), + ), + provider_executed: None, + provider_metadata: None, + dynamic: Some(true), + }) +} + +/// A provider's tool result as a stream part +/// (`stream-language-model-call.ts:760-792`): a tool error when the provider +/// marks it as one, otherwise a tool result, with the input of the call it +/// answers (`None` when the call is not in this stream). +/// +/// Upstream drops the provider's `preliminary` flag here; aimux keeps it so +/// that, as on the generated path, a preliminary result does not enter the +/// step content. +fn provider_tool_result_part(result: ToolResult, input: Option) -> TextStreamPart { + if result.is_error == Some(true) { + TextStreamPart::ToolError(TypedToolError { + tool_call_id: result.tool_call_id, + tool_name: result.tool_name, + input, + error: result.result, + provider_executed: Some(true), + provider_metadata: result.provider_metadata, + dynamic: result.dynamic, + }) + } else { + TextStreamPart::ToolResult(TypedToolResult { + tool_call_id: result.tool_call_id, + tool_name: result.tool_name, + input, + output: result.result, + provider_executed: Some(true), + provider_metadata: result.provider_metadata, + dynamic: result.dynamic, + preliminary: result.preliminary, + }) + } +} + +/// The content of a generated model call with parsed tool calls and +/// resolved files, and its response messages (the AI SDK's +/// `convertLanguageModelContent` and `toResponseMessages`, +/// `generate-text.ts:1437-1455`). +/// +/// A provider tool result whose call is not in this response is a deferred +/// result. It is accepted only for a provider tool of that name that sets +/// `supportsDeferredResults`, and otherwise fails with `Tool call {id} not +/// found.` (`convert-language-model-content.ts:87-104`). +async fn convert_generated_content( + result: &GenerateResult, + tool_calls: &[ToolCall], + tools: Option<&[Tool]>, + abort_signal: Option<&AbortSignal>, + operation_timeout: timeout::OperationTimeout, +) -> Result { + let mut rm = crate::response_messages::ResponseMessageBuilder::new(); for content in &result.content { match content { GenerateContent::Text { - text: t, + text, provider_metadata, - } => { - text.push_str(t); - rm.text(t, provider_metadata.as_ref()); - } + } => rm.text(text, provider_metadata.as_ref()), GenerateContent::ToolCall(call) => { let parsed = tool_calls .iter() @@ -685,24 +1168,20 @@ pub(crate) async fn generate_text_from_language_model_prompt( rm.tool_call(parsed); } GenerateContent::Reasoning(reasoning) => rm.reasoning(reasoning), - GenerateContent::Source(source) => { - rm.source(source.clone()); - sources.push(source.clone()); - } + GenerateContent::Source(source) => rm.source(source.clone()), GenerateContent::File(file) => { let file = timeout::run( - resolve_file(file.clone(), abort_signal.as_ref()), - abort_signal.as_ref(), + resolve_file(file.clone(), abort_signal), + abort_signal, operation_timeout, ) .await?; rm.file(&file, false); - files.push(file); } GenerateContent::ReasoningFile(file) => { let file = timeout::run( - resolve_file(file.clone(), abort_signal.as_ref()), - abort_signal.as_ref(), + resolve_file(file.clone(), abort_signal), + abort_signal, operation_timeout, ) .await?; @@ -726,109 +1205,120 @@ pub(crate) async fn generate_text_from_language_model_prompt( reason: None, is_automatic: None, signature: None, + input_schema_input: None, }); } - GenerateContent::ToolResult(result) => rm.tool_result(result.clone()), + GenerateContent::ToolResult(result) => { + let has_call = tool_calls + .iter() + .any(|call| call.tool_call_id == result.tool_call_id); + let deferred_allowed = tools.unwrap_or_default().iter().any(|tool| { + matches!(tool, Tool::Provider(tool) + if tool.name == result.tool_name + && tool.supports_deferred_results == Some(true)) + }); + if !has_call && !deferred_allowed { + return Err(AiMuxError::Other(format!( + "Tool call {} not found.", + result.tool_call_id + ))); + } + rm.tool_result(result.clone()); + } } } // A tool call that could not be parsed gets a tool error, after the // model's content (`generate-text.ts:1303-1323`, // `convert-language-model-content.ts:208-214`). - for call in &tool_calls { - if call.invalid == Some(true) - && call.dynamic == Some(true) - && call.provider_executed != Some(true) - { - rm.tool_error(crate::result::TypedToolError { - tool_call_id: call.tool_call_id.clone(), - tool_name: call.tool_name.clone(), - input: Some(call.input.clone()), - error: Value::String( - call.error - .as_ref() - .map_or_else(|| "unknown error".to_owned(), ToString::to_string), - ), - provider_executed: None, - provider_metadata: None, - dynamic: Some(true), - }); + for call in tool_calls.iter().filter(|call| call.dynamic == Some(true)) { + if let Some(error) = invalid_tool_call_error(call) { + rm.tool_error(error); } } + Ok(rm.finish()) +} - let crate::response_messages::ResponseMessages { - messages: response_messages, - content, - reasoning, - } = rm.finish(); - let reasoning_text = reasoning - .iter() - .filter_map(|r| match r { - ReasoningPart::Text(r) => Some(r.text.as_str()), - ReasoningPart::File(_) => None, - }) - .collect::>() - .join(""); - - // Parsing and optional repair are part of the user operation. Record a - // successful outcome only after that phase has completed within the same - // deadline/cancellation scope as the provider call. - if let (Some(rec), Some(call_id)) = (&recorder, &call_id) { - rec.record_outcome( - call_id, - &crate::recording::OutcomeRecord::from_generate_result(&result), - ); - } - - // Extract fields before moving `result` into `raw`. - let raw_finish_reason = result.finish_reason.raw.clone(); - let provider_metadata = result.provider_metadata.clone(); - let response = result.response.clone().map(Into::into).unwrap_or_default(); - let usage = result.usage.clone(); - - Ok(GenerateTextResult { - content, - text, - tool_calls, - finish_reason: result.finish_reason.clone(), - usage: usage.clone(), - warnings: result.warnings.clone(), - raw: result, - reasoning, - reasoning_text, - sources, - files, - response_messages, - raw_finish_reason, +/// The step of a generated model call (`generate-text.ts:1382-1401`, +/// `:1457-1490`). A generated call has no first-output timing. +fn build_generated_step( + model: &dyn LanguageModel, + call_id: Option, + result: GenerateResult, + content: crate::response_messages::ResponseMessages, + response_time_ms: f64, + step_time_ms: f64, +) -> StepResult { + let GenerateResult { + finish_reason, + usage, + warnings, provider_metadata, + request, response, - total_usage: usage, - }) + .. + } = result; + let response = response.unwrap_or_default(); + let metadata = response_metadata_defaults( + crate::types::ResponseMetadata { + id: response.id, + timestamp: response.timestamp, + model_id: response.model_id, + }, + model.model_id(), + ); + let usage = as_language_model_usage(&usage); + let performance = StepResultPerformance::new(&usage, response_time_ms, step_time_ms, None, &[]); + StepResult { + call_id: call_id.unwrap_or_else(crate::recording::new_call_id), + step_number: 0, + model: StepModel { + provider: model.provider().to_owned(), + model_id: model.model_id().to_owned(), + }, + content: content.content, + finish_reason: finish_reason.unified, + raw_finish_reason: finish_reason.raw, + usage, + performance, + warnings, + request: request.unwrap_or_default(), + response: StepResponse { + id: metadata.id.unwrap_or_default(), + timestamp: metadata.timestamp.unwrap_or_default(), + model_id: metadata.model_id.unwrap_or_default(), + headers: response.headers, + body: response.body, + messages: content.messages, + }, + provider_metadata, + } } -/// Generate a structured JSON object (M12). Uses `generate_text` with -/// `response_format: Json`, then parses the model output into a JSON value. +/// Generate a structured JSON value: a deprecated wrapper over +/// [`generate_text`] with [`GenerateTextOptions::output`], as the AI SDK +/// deprecates `generateObject` for `generateText({ output })` +/// (`generate-object/generate-object.ts:120`). /// -/// The schema is passed via `options.response_format = ResponseFormat::Json { -/// schema: Some(schema), .. }`. JSON repair (`fix_json`) is applied before -/// parsing to handle truncated or slightly malformed model output. No schema -/// validation is performed (weak validation — relies on the model following -/// the schema; use retries for robustness). +/// `options.output` describes the value; without one the JSON output +/// (`Output::Json`, any JSON value) applies. The object is the result's +/// output, so the model's text must parse and match the schema, with the +/// output's errors. /// /// # Example /// ``` +/// # #![allow(deprecated)] /// # use aimux_core::prelude::*; -/// # use aimux_core::options::ResponseFormat; /// # async fn example(model: &dyn aimux_core::LanguageModel) -> Result<(), aimux_core::error::AiMuxError> { /// let result = generate_object( /// model, /// "Extract: John is 25 years old.", /// GenerateTextOptions { -/// response_format: Some(ResponseFormat::Json { -/// schema: Some(serde_json::json!({ +/// output: Some(Output::Object { +/// schema: serde_json::json!({ /// "type": "object", /// "properties": { "name": { "type": "string" }, "age": { "type": "number" } }, /// "required": ["name", "age"] -/// })), +/// }), /// name: None, /// description: None, /// }), @@ -842,46 +1332,40 @@ pub(crate) async fn generate_text_from_language_model_prompt( /// /// # Errors /// -/// Propagates `generate_text` failures; additionally returns `JsonParse` when -/// the (repaired) model output is not valid JSON. +/// Propagates `generate_text` failures, including the output's parse +/// errors; `NoOutputGenerated` when the model call finished without a +/// parsed output (for example with tool calls). +#[deprecated( + note = "use `generate_text` with `GenerateTextOptions::output`, as the AI SDK \ + deprecates `generateObject` for `generateText({ output })`" +)] pub async fn generate_object( model: &dyn LanguageModel, prompt: impl Into, - options: GenerateTextOptions, + mut options: GenerateTextOptions, ) -> Result { - let text_result = generate_text(model, prompt, options).await?; - let repaired = crate::json_repair::fix_json(&text_result.text); - let object: Value = serde_json::from_str(&repaired).map_err(|e| { - AiMuxError::JsonParse(format!( - "generateObject: model output is not valid JSON after repair: {e}" - )) - })?; - let raw_finish_reason = text_result.raw_finish_reason.clone(); - let finish_reason = text_result.finish_reason.clone(); - let usage = text_result.usage.clone(); - let warnings = text_result.warnings.clone(); - let reasoning = if text_result.reasoning_text.is_empty() { - None - } else { - Some(text_result.reasoning_text.clone()) - }; - let provider_metadata = text_result.raw.provider_metadata.clone(); - let response = text_result - .raw - .response - .clone() - .map(Into::into) - .unwrap_or_default(); + options.output.get_or_insert(Output::Json { + name: None, + description: None, + }); + let result = generate_text(model, prompt, options).await?; + let object = result.output()?.clone(); + let step = result.final_step(); Ok(GenerateObjectResult { object, - finish_reason, - raw_finish_reason, - usage, - warnings, - reasoning, - provider_metadata, - response, - raw: text_result, + reasoning: step.reasoning_text(), + finish_reason: step.finish_reason, + usage: step.usage.clone(), + warnings: step.warnings.clone(), + request: step.request.clone(), + response: GenerateObjectResponse { + id: step.response.id.clone(), + timestamp: step.response.timestamp.clone(), + model_id: step.response.model_id.clone(), + headers: step.response.headers.clone(), + body: step.response.body.clone(), + }, + provider_metadata: step.provider_metadata.clone(), }) } @@ -902,36 +1386,59 @@ pub async fn generate_object( /// /// let mut stream = result.stream; /// while let Some(part) = stream.next().await { -/// if let StreamPart::TextDelta { delta, .. } = part? { -/// print!("{}", delta); +/// if let TextStreamPart::TextDelta { text, .. } = part? { +/// print!("{}", text); /// } /// } /// # Ok(()) /// # } /// ``` /// +/// The stream is framed as the AI SDK's `fullStream` +/// (`stream-text.ts:1855`, `:2830-2840`, `:3015-3135`): `Start`, then +/// `StartStep` on the step's first chunk, the content parts, `FinishStep` and +/// `Finish`. Aborting ends it with an `Abort` part (`:1860-1882`). +/// /// # Errors /// /// Returns setup failures from the model's `do_stream` call. Provider error -/// events surface as `Ok(StreamPart::Error { .. })`; transport failures, -/// caller cancellation, and Core timeouts surface as `Err(AiMuxError)` stream -/// items. +/// events surface as `Ok(TextStreamPart::Error { .. })`; transport failures +/// and Core timeouts surface as `Err(AiMuxError)` stream items. pub async fn stream_text( model: &dyn LanguageModel, prompt: impl Into, options: GenerateTextOptions, +) -> Result { + let settings = options.operation_settings(); + settings + .clone() + .scope(stream_in_operation(model, prompt.into(), options, settings)) + .await +} + +// The stream setup runs inside the operation settings' scope, like +// `generate_in_operation`. The pump task that drives the returned stream +// calls no model and does not read them. +async fn stream_in_operation( + model: &dyn LanguageModel, + prompt: ModelPrompt, + options: GenerateTextOptions, + settings: OperationSettings, ) -> Result { // 1. Convert user prompt to provider-facing prompt. let repair_tool_call = options.repair_tool_call.clone(); let tools = options.tools.clone(); let operation_instructions = options.instructions.clone(); - let (messages, instructions) = - standardize_prompt(prompt.into(), operation_instructions.as_deref())?; + let output = options.output.clone(); + if let Some(output) = &output { + output.validate()?; + } + let (messages, instructions) = standardize_prompt(prompt, operation_instructions.as_deref())?; let lm_prompt = convert_to_language_model_prompt(&messages, instructions, provider_id(model))?; // 2. Build CallOptions. - let mut call_options = options.into_call_options(lm_prompt); - let stream_timeout = call_options.timeout.unwrap_or_default(); + let call_options = options.into_call_options(lm_prompt); + let stream_timeout = settings.timeout.unwrap_or_default(); let operation_timeout = timeout::OperationTimeout::new(stream_timeout)?; // Armed at operation start: providers await their first SSE event inside // do_stream, so the first-chunk budget must already be counting there — @@ -943,7 +1450,7 @@ pub async fn stream_text( .transpose()?; let abort_signal = call_options.abort_signal.clone(); let retries = retry::prepare_retries( - call_options.max_retries, + settings.max_retries, model.retry_config(), abort_signal.clone(), ); @@ -954,7 +1461,6 @@ pub async fn stream_text( let context = crate::recording::recorder().map(|recorder| { let ctx = crate::recording::RecordingContext::new(crate::recording::new_call_id(), recorder); - call_options.recording_context = Some(ctx.clone()); ctx.recorder.record_input( &ctx.call_id, &call_options, @@ -971,6 +1477,13 @@ pub async fn stream_text( }); let call_id = context.as_ref().map(|c| c.call_id.clone()); let recorder = context.as_ref().map(|c| c.recorder.clone()); + let step_call_id = call_id + .clone() + .unwrap_or_else(crate::recording::new_call_id); + let step_model = StepModel { + provider: model.provider().to_owned(), + model_id: model.model_id().to_owned(), + }; // 2b. Session grouping (RFC-0024): explicit session_id first, opt-in // prompt-prefix inference as fallback. Recorded before the call so @@ -979,10 +1492,7 @@ pub async fn stream_text( // and store lookups are cheap (once-locked env check + two RwLock // reads) when no store / inferer is registered. if let (Some((session_id, source)), Some(store)) = ( - crate::session::resolve_session_id( - call_options.session_id.as_deref(), - &call_options.prompt, - ), + crate::session::resolve_session_id(settings.session_id.as_deref(), &call_options.prompt), crate::session::session_store(), ) { let call = store.append(&session_id, call_id.as_deref(), source); @@ -994,6 +1504,9 @@ pub async fn stream_text( } // 3. Call provider, inside the "generate" span (RFC-0014 §4.1 span tree). + // The step starts before the retry loop (`stepStartTimestampMs`, + // stream-text.ts:2399); each attempt restarts the model call's clock + // (`callStartTimestampMs`, stream-language-model-call.ts:348). let span = tracing::info_span!( target: "aimux_core::generate", "generate", @@ -1001,15 +1514,23 @@ pub async fn stream_text( model = %model.model_id(), modality = "text", ); - let operation = retries.retry(|| { - if let Some(context) = &context { - let _ = context.start_attempt(); - } - // Call-level retry boundary: once do_stream returns a stream, retry - // no longer recreates the provider request; the returned stream is - // consumed by the stream-step timeout logic below. - model.do_stream(&call_options).instrument(span.clone()) - }); + let step_started = Instant::now(); + let attempt_started = Mutex::new(step_started); + let operation = crate::recording::RecordingContext::scope( + context.clone(), + retries.retry(|| { + if let Some(context) = &context { + let _ = context.start_attempt(); + } + *attempt_started + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now(); + // Call-level retry boundary: once do_stream returns a stream, retry + // no longer recreates the provider request; the returned stream is + // consumed by the stream-step timeout logic below. + model.do_stream(&call_options).instrument(span.clone()) + }), + ); let setup_deadline = timeout::min_deadline(operation_timeout.deadline(), first_chunk_deadline); let result: StreamResult = match timeout::run_until(operation, abort_signal.as_ref(), setup_deadline).await { @@ -1021,6 +1542,9 @@ pub async fn stream_text( return Err(e); } }; + let call_started = *attempt_started + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); // 解构避免部分 move(result 各字段去向不同)。 let StreamResult { stream, @@ -1040,105 +1564,124 @@ pub async fn stream_text( // for the consumer's latency. The channel is unbounded so the pump never // blocks on the consumer (which would put that latency back into the // timers); the AI SDK's `streamText` consumes eagerly for the same - // reason, and buffering is bounded by the response itself. + // reason, and buffering is bounded by the response itself. Each item is + // stamped with its arrival time, which the step's performance uses. // // Spawned even with no abort signal or deadlines armed: the pump is also // the stream's terminal fuse — `Finish` and non-recoverable errors must // end the stream under every configuration, or a provider that keeps the // connection open after `Finish` hangs a consumer reading to end-of-stream. - let file_abort_signal = abort_signal.clone(); - let stream: Pin> + Send>> = { - let mut stream = stream; - let chunk_ms = stream_timeout.chunk_ms; - let pump_abort = abort_signal.clone(); - let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); - let pump = tokio::spawn(async move { - let mut chunk_deadline = first_chunk_deadline; + // The pump runs in its own task, which does not inherit the recording + // scope; the stream carries it into every poll. + let mut provider_stream = + crate::recording::RecordingContext::scope_stream(context.clone(), stream); + let chunk_ms = stream_timeout.chunk_ms; + let pump_abort = abort_signal.clone(); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let pump = tokio::spawn(async move { + let mut chunk_deadline = first_chunk_deadline; - loop { - let next = tokio::select! { - biased; - // Dropping the provider stream cancels the in-flight HTTP - // exchange; the consumer side reports the abort error. - () = timeout::wait_for_abort(pump_abort.as_ref()) => break, - () = timeout::wait_for_deadline(operation_deadline) => { - let _ = tx.send(Err(operation_deadline - .expect("operation deadline future only resolves for a deadline") - .error())); - break; - } - () = timeout::wait_for_deadline(chunk_deadline) => { - let _ = tx.send(Err(chunk_deadline - .expect("chunk deadline future only resolves for a deadline") - .error())); + loop { + let next = tokio::select! { + biased; + // Dropping the provider stream cancels the in-flight HTTP + // exchange; the consumer side reports the abort. + () = timeout::wait_for_abort(pump_abort.as_ref()) => break, + () = timeout::wait_for_deadline(operation_deadline) => { + let _ = tx.send((Instant::now(), Err(operation_deadline + .expect("operation deadline future only resolves for a deadline") + .error()))); + break; + } + () = timeout::wait_for_deadline(chunk_deadline) => { + let _ = tx.send((Instant::now(), Err(chunk_deadline + .expect("chunk deadline future only resolves for a deadline") + .error()))); + break; + } + item = provider_stream.next() => item, + }; + + let Some(item) = next else { break }; + let arrived = Instant::now(); + let terminal = is_terminal_item(&item); + if let Ok(part) = &item + && is_output_chunk(part) + { + chunk_deadline = match chunk_ms + .map(|duration_ms| timeout::TimeoutDeadline::from_now("Chunk", duration_ms)) + .transpose() + { + Ok(deadline) => deadline, + Err(error) => { + let _ = tx.send((arrived, Err(error))); break; } - item = stream.next() => item, }; + } - let Some(item) = next else { break }; - let terminal = is_terminal_item(&item); - if let Ok(part) = &item - && is_output_chunk(part) - { - chunk_deadline = match chunk_ms - .map(|duration_ms| timeout::TimeoutDeadline::from_now("Chunk", duration_ms)) - .transpose() - { - Ok(deadline) => deadline, - Err(error) => { - let _ = tx.send(Err(error)); - break; - } - }; - } - - // A closed receiver means the consumer dropped the stream. - if tx.send(item).is_err() || terminal { - break; - } + // A closed receiver means the consumer dropped the stream. + if tx.send((arrived, item)).is_err() || terminal { + break; } - }); - let pump = AbortOnDrop(pump); - Box::pin(async_stream::stream! { + } + }); + let pump = AbortOnDrop(pump); + + let file_abort_signal = abort_signal.clone(); + let tool_choice = call_options.tool_choice.clone(); + let step_provider = model.provider().to_owned(); + let step_model_id = model.model_id().to_owned(); + let response_metadata = + response_metadata_defaults(crate::types::ResponseMetadata::default(), model.model_id()); + let stream_response_headers = response_headers.clone(); + let step_request = RequestInfo { + body: request_body.clone(), + }; + let stream: Pin> + Send>> = Box::pin( + async_stream::stream! { let _pump = pump; + let mut step = StreamStep::new(call_started, response_metadata); + let mut tool_calls: HashMap = HashMap::new(); + let mut content = StreamedContent::default(); + yield Ok(TextStreamPart::Start {}); loop { let next = tokio::select! { biased; () = timeout::wait_for_abort(abort_signal.as_ref()) => { - yield Err(AiMuxError::from_abort_signal( + let error = AiMuxError::from_abort_signal( abort_signal.as_ref().expect("abort future only resolves for a signal"), - )); - break; + ); + yield Ok(TextStreamPart::Abort { reason: Some(error.to_string()) }); + return; } item = rx.recv() => item, }; - let Some(item) = next else { break }; - // Mirror the pump's fuse here so the stream ends right after - // the terminal item, ahead of a later abort or the channel - // close racing it. - let terminal = is_terminal_item(&item); - yield item; - if terminal { - break; - } - } - }) - }; - let mut stream = stream; - let tool_choice = call_options.tool_choice.clone(); - let provider = model.provider().to_owned(); - let model_id = model.model_id().to_owned(); - let stream: Pin> + Send>> = Box::pin( - async_stream::stream! { - let mut tool_calls = HashMap::new(); - let mut content = StreamedContent::default(); - while let Some(item) = stream.next().await { - if let Ok(part) = &item { - content.observe(part); + let Some((arrived, item)) = next else { break }; + let part = match item { + Ok(part) => part, + Err(error) => { + let terminal = !error.is_recoverable_stream_error(); + yield Err(error); + if terminal { return; } + continue; + } + }; + // `model-call-start` only carries the warnings; every other + // chunk first opens the step (stream-text.ts:2822-2840). + if !matches!(part, StreamPart::StreamStart { .. }) { + if !step.started { + step.started = true; + yield Ok(TextStreamPart::StartStep { + request: step_request.clone(), + warnings: step.warnings.clone(), + }); + } + step.record_chunk(&part, arrived); + content.observe(&part); } - match item { - Ok(StreamPart::ToolCall(raw)) => { + match part.split_payload() { + StreamPartPayload::ToolCall(raw) => { let parsed = parse_tool_call( raw, tools.as_deref(), @@ -1147,69 +1690,101 @@ pub async fn stream_text( operation_instructions.as_deref(), ).await; tool_calls.insert(parsed.tool_call_id.clone(), parsed.clone()); - yield Ok(StreamPart::ToolCall(parsed)); - } - Ok(StreamPart::File(file)) => { - match timeout::run(resolve_file(file, file_abort_signal.as_ref()), file_abort_signal.as_ref(), operation_timeout).await { - Ok(file) => yield Ok(StreamPart::File(file)), - Err(error) => yield Ok(StreamPart::Error { error }), - } - } - Ok(StreamPart::ReasoningFile(file)) => { - match timeout::run(resolve_file(file, file_abort_signal.as_ref()), file_abort_signal.as_ref(), operation_timeout).await { - Ok(file) => yield Ok(StreamPart::ReasoningFile(ReasoningFileOutput { - provider_metadata: file.provider_metadata.clone(), file, - })), - Err(error) => yield Ok(StreamPart::Error { error }), + let tool_error = invalid_tool_call_error(&parsed); + yield Ok(TextStreamPart::ToolCall(parsed)); + if let Some(error) = tool_error { + yield Ok(TextStreamPart::ToolError(error)); } } - Ok(StreamPart::ToolApprovalRequest(approval)) => { + StreamPartPayload::ToolApprovalRequest(approval) => { if let Some(call) = tool_calls.get(&approval.tool_call_id) { - yield Ok(StreamPart::ToolApprovalRequest(ToolApprovalRequestOutput { + yield Ok(TextStreamPart::ToolApprovalRequest(ToolApprovalRequestOutput { approval_id: approval.approval_id, tool_call: call.clone(), - reason: None, is_automatic: None, signature: None, + reason: None, is_automatic: None, signature: None, input_schema_input: None, })); } else { - yield Ok(StreamPart::Error { error: AiMuxError::ToolCallNotFoundForApproval { + yield Ok(TextStreamPart::Error { error: AiMuxError::ToolCallNotFoundForApproval { tool_call_id: approval.tool_call_id, approval_id: approval.approval_id, }}); } } - // `stream-language-model-call.ts:644-687`: a response that - // breaks an enforced tool choice ends with an error and - // an `error` finish reason (the raw reason is kept). - Ok(StreamPart::Finish { finish_reason, usage, provider_metadata }) => { - let violation = tool_choice_violation_error( - tool_choice.as_ref(), - tool_calls.values().map(|call: &crate::tool::ToolCall| call.tool_name.as_str()), - finish_reason.unified, - &provider, - &model_id, - || std::mem::take(&mut content.content), - ); - match violation { - None => yield Ok(StreamPart::Finish { finish_reason, usage, provider_metadata }), - Some(error) => { - yield Ok(StreamPart::Error { error }); - yield Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: crate::types::FinishReasonUnified::Error, - raw: finish_reason.raw, - }, - usage, - provider_metadata, - }); - } + StreamPartPayload::ReasoningFile(file) => { + match timeout::run(resolve_file(file, file_abort_signal.as_ref()), file_abort_signal.as_ref(), operation_timeout).await { + Ok(file) => yield Ok(TextStreamPart::ReasoningFile(ReasoningFileOutput { + provider_metadata: file.provider_metadata.clone(), file, + })), + Err(error) => yield Ok(TextStreamPart::Error { error }), } } - Ok(part) => yield Ok(part.map_payloads(|_| unreachable!("matched above"), |_| unreachable!("matched above")) - .map_reasoning_file(|file| ReasoningFileOutput { provider_metadata: file.provider_metadata.clone(), file })), - Err(error) => yield Err(error), + StreamPartPayload::Other(part) => match part { + StreamPart::StreamStart { warnings } => step.warnings = warnings, + StreamPart::ResponseMetadata(metadata) => step.response_metadata(metadata), + // `stream-language-model-call.ts:644-687`: a response that + // breaks an enforced tool choice finishes the model call + // with an `error` finish reason (the raw reason is kept), + // followed by an error part. + StreamPart::Finish { mut finish_reason, usage, provider_metadata } => { + let violation = tool_choice_violation_error( + tool_choice.as_ref(), + tool_calls.values().map(|call| call.tool_name.as_str()), + finish_reason.unified, + &step_provider, + &step_model_id, + || std::mem::take(&mut content.content), + ); + if violation.is_some() { + finish_reason.unified = FinishReasonUnified::Error; + } + step.finish(finish_reason, &usage, provider_metadata, arrived); + if let Some(error) = violation { + yield Ok(TextStreamPart::Error { error }); + } + } + StreamPart::Error { error } => { + step.terminal = true; + step.finish_reason = FinishReasonUnified::Error; + yield Ok(TextStreamPart::Error { error }); + } + StreamPart::ToolResult(result) => { + let input = tool_calls.get(&result.tool_call_id).map(|call| call.input.clone()); + yield Ok(provider_tool_result_part(result, input)); + } + StreamPart::File(file) => { + match timeout::run(resolve_file(file, file_abort_signal.as_ref()), file_abort_signal.as_ref(), operation_timeout).await { + Ok(file) => yield Ok(TextStreamPart::File(file)), + Err(error) => yield Ok(TextStreamPart::Error { error }), + } + } + // An empty text delta without provider metadata is not + // forwarded (stream-text.ts:2881-2891); providers emit + // one, e.g. OpenAI chat's opening `"content": ""`. + StreamPart::TextDelta { delta, provider_metadata: None, .. } if delta.is_empty() => {} + part => { + if let Some(part) = part + .widen::() + .into_text_stream_part() + { + yield Ok(part); + } + } + }, } } + // An incomplete model stream without output is an error, not an + // empty step; one with partial output keeps it + // (stream-text.ts:3015-3030). + if !step.terminal && !step.saw_output { + yield Ok(TextStreamPart::Error { error: AiMuxError::NoOutputGenerated( + "No output generated. The model stream ended without a finish chunk.".into(), + ) }); + return; + } + let (finish_step, finish) = step.into_parts(stream_response_headers, step_started); + yield Ok(finish_step); + yield Ok(finish); }, ); - // 录制开启时才包装(终结时写 outcome + 传输封闭);关闭时零成本透传。 + let stream = crate::recording::RecordingOutcomeStream::new( stream, recorder.clone(), @@ -1222,9 +1797,137 @@ pub async fn stream_text( stream: Box::pin(stream), request_body, response_headers, + call_id: step_call_id, + model: step_model, + output, }) } +/// The state of the streamed step: what `stream-text.ts:2736-2780` keeps +/// per step, plus the output timing of `stream-language-model-call.ts:441-460`. +struct StreamStep { + started: bool, + warnings: Vec, + /// A finish or error chunk arrived. + terminal: bool, + saw_output: bool, + finish_reason: FinishReasonUnified, + raw_finish_reason: Option, + usage: LanguageModelUsage, + provider_metadata: Option, + response: crate::types::ResponseMetadata, + call_started: Instant, + first_output_ms: Option, + previous_output: Option, + chunk_gaps_ms: Vec, + /// The model call's performance, known once its finish chunk arrived. + model_call_performance: Option, +} + +impl StreamStep { + fn new(call_started: Instant, response: crate::types::ResponseMetadata) -> Self { + Self { + started: false, + warnings: Vec::new(), + terminal: false, + saw_output: false, + finish_reason: FinishReasonUnified::Other, + raw_finish_reason: None, + usage: LanguageModelUsage::default(), + provider_metadata: None, + response, + call_started, + first_output_ms: None, + previous_output: None, + chunk_gaps_ms: Vec::new(), + model_call_performance: None, + } + } + + /// Output timing by arrival time (`stream-language-model-call.ts:450-460`). + fn record_chunk(&mut self, part: &StreamPart, arrived: Instant) { + if !is_output_chunk(part) { + return; + } + self.saw_output = true; + if self.first_output_ms.is_none() { + self.first_output_ms = Some(ms_between(self.call_started, arrived)); + } else if let Some(previous) = self.previous_output { + self.chunk_gaps_ms.push(ms_between(previous, arrived)); + } + self.previous_output = Some(arrived); + } + + /// `model-call-response-metadata` (stream-text.ts:2962-2968). + fn response_metadata(&mut self, metadata: crate::types::ResponseMetadata) { + if metadata.id.is_some() { + self.response.id = metadata.id; + } + if metadata.timestamp.is_some() { + self.response.timestamp = metadata.timestamp; + } + if metadata.model_id.is_some() { + self.response.model_id = metadata.model_id; + } + } + + /// The finish chunk: usage and the model call's performance + /// (`stream-language-model-call.ts:593-624`, stream-text.ts:2971-2983). + fn finish( + &mut self, + finish_reason: crate::types::FinishReason, + usage: &crate::types::Usage, + provider_metadata: Option, + arrived: Instant, + ) { + self.terminal = true; + self.finish_reason = finish_reason.unified; + self.raw_finish_reason = finish_reason.raw; + self.usage = as_language_model_usage(usage); + self.provider_metadata = provider_metadata; + self.model_call_performance = Some(StepResultPerformance::new( + &self.usage, + ms_between(self.call_started, arrived), + 0.0, + self.first_output_ms, + &self.chunk_gaps_ms, + )); + } + + /// The `finish-step` and `finish` parts (stream-text.ts:3032-3048, + /// :3128-3134). Without a finish chunk the model call's performance + /// stays zero (`createModelCallPerformance`, :2755-2767). + fn into_parts( + self, + headers: Option, + step_started: Instant, + ) -> (TextStreamPart, TextStreamPart) { + let mut performance = self.model_call_performance.unwrap_or_default(); + performance.step_time_ms = ms_between(step_started, Instant::now()); + let total_usage = add_language_model_usage(&LanguageModelUsage::default(), &self.usage); + ( + TextStreamPart::FinishStep { + response: StepResponseMetadata { + id: self.response.id.unwrap_or_default(), + timestamp: self.response.timestamp.unwrap_or_default(), + model_id: self.response.model_id.unwrap_or_default(), + headers, + }, + usage: self.usage, + performance, + finish_reason: self.finish_reason, + raw_finish_reason: self.raw_finish_reason.clone(), + provider_metadata: self.provider_metadata, + }, + TextStreamPart::Finish { + finish_reason: self.finish_reason, + raw_finish_reason: self.raw_finish_reason, + total_usage, + }, + ) + } +} + /// The AI SDK's tool choice check after a model call (`generate-text.ts:1155-1176`, /// `stream-language-model-call.ts:644-662`): a `required` tool choice needs /// some tool call, a `tool` choice a call to that tool. Every parsed call @@ -1394,6 +2097,11 @@ impl StreamedContent { } } +#[cfg(test)] +mod pipeline_tests; +#[cfg(test)] +mod tests; + /// A malformed individual frame does not prove that the transport is dead: /// SSE framing lets a later event still be valid. `Finish` and transport/Core /// errors are terminal so the stream is not polled past a dead source. @@ -1485,17 +2193,15 @@ pub(crate) fn split_prompt( // ───────────────────────────────────────────────────────────────────────────── use crate::openai_output::{ - ChatCompletion, ChatCompletionFunction, ChatCompletionStream, ChatCompletionToolCall, - OpenAiStreamOptions, parsed_tool_call_arguments, to_chat_completion, to_chat_completion_stream, + ChatCompletion, ChatCompletionStream, OpenAiStreamOptions, to_chat_completion_stream, }; /// Generate text and return the result as an OpenAI Chat Completion. /// /// This is the OpenAI-compatible equivalent of [`generate_text`]. Internally -/// it calls `generate_text`, then converts the [`GenerateResult`] into a -/// [`ChatCompletion`] via [`to_chat_completion`]. -/// Tool calls reflect Core parsing and any configured repair; provider -/// response metadata, usage, and non-tool content are preserved from `raw`. +/// it calls `generate_text`, then converts the result with +/// [`generate_text_result_to_chat_completion`]. Tool calls reflect Core +/// parsing and any configured repair. /// /// Works with **any** provider (OpenAI, Anthropic, Google, …) — the output is /// always standard OpenAI Chat Completions JSON. @@ -1535,45 +2241,21 @@ pub async fn generate_text_as_openai( /// Convert a [`GenerateTextResult`] into an OpenAI [`ChatCompletion`]. /// -/// The pure half of [`generate_text_as_openai`]: tool calls come from -/// `result.tool_calls` (so they reflect Core parsing and any repair already -/// applied to the result), everything else from `result.raw` via -/// [`to_chat_completion`]. Hosts that repair a result themselves (RFC-0035) -/// convert it with this afterwards. +/// The pure half of [`generate_text_as_openai`]: the message is built from +/// the content of the result's steps, so tool calls reflect Core parsing and +/// any repair already applied to the result; the finish reason, id, model and +/// logprobs come from the last step and the usage is the result's total. +/// Hosts that repair a result themselves (RFC-0035) convert it with this +/// afterwards. /// -/// `model_id` is the fallback for `ChatCompletion.model` when the provider -/// response carries no model id. +/// `model_id` is the fallback for `ChatCompletion.model` when the response +/// carries no model id. #[must_use] pub fn generate_text_result_to_chat_completion( result: &GenerateTextResult, model_id: &str, ) -> ChatCompletion { - let mut completion = to_chat_completion(&result.raw, model_id); - if let Some(choice) = completion.choices.first_mut() { - choice.message.tool_calls = if result.tool_calls.is_empty() { - None - } else { - Some( - result - .tool_calls - .iter() - .map(|tool_call| ChatCompletionToolCall { - id: tool_call.tool_call_id.clone(), - tool_type: "function".to_string(), - function: ChatCompletionFunction { - name: tool_call.tool_name.clone(), - arguments: parsed_tool_call_arguments( - &tool_call.input, - tool_call.invalid, - tool_call.error.as_ref(), - ), - }, - }) - .collect(), - ) - }; - } - completion + crate::openai_output::text_result_to_chat_completion(result, model_id) } /// Stream text and return the result as a stream of OpenAI Chat Completion chunks. diff --git a/aimux-core/src/generate/pipeline_tests.rs b/aimux-core/src/generate/pipeline_tests.rs new file mode 100644 index 000000000..7d265d87d --- /dev/null +++ b/aimux-core/src/generate/pipeline_tests.rs @@ -0,0 +1,468 @@ +//! RFC-0031 §13 acceptance for the streaming call pipeline: the chunk timer +//! (items 16-19c), producer-side timing, the stream's terminal fuse, the +//! first-event peek and its retry, and no timer state after a drop +//! (21a-21c). Ported from the master branch's +//! `tests/stream_timeout_matrix_test.rs` and `generate.rs` stream tests, +//! adjusted to the AI SDK's full stream (`Start`, `StartStep`, ...). +//! +//! The chunk timer follows `streamText` (ai@7.0.127 `stream-text.ts`): it is +//! armed and reset only by output chunks (`isOutputChunk`, :182-225, :2842- +//! 2857) and reports `Chunk timeout of {ms}ms exceeded` (:2257-2266). The +//! first-chunk timer is armed at operation start rather than when the +//! response stream begins (upstream :2228-2243), because aimux providers +//! peek the first event inside `do_stream` (RFC-0031 §8.3). + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use async_trait::async_trait; +use futures::StreamExt; + +use super::*; +use crate::error::ApiCallError; +use crate::options::TimeoutConfiguration; +use crate::types::{FinishReason, Usage}; + +#[derive(Clone, Copy)] +enum Behavior { + /// One text delta, then the stream ends. + Normal, + /// A text delta at t=0, another part at t=600ms (a delta when `true`, + /// a non-output `TextStart` otherwise), then silence. + Probe(bool), + /// A finish part, then silence. + FinishThenPending, + /// A provider error part and a finish part, then silence. + ProviderErrorThenFinish, + /// A malformed frame, then a text delta. + ParseErrorThenDelta, + /// A retryable transport error, then silence. + TransportErrorThenPending, + /// `do_stream` never returns (a 200-then-silence server whose first + /// event the provider awaits inside `do_stream`). + NeverReturns, + /// The first `do_stream` fails with a retryable 429. + RetryableFirstError, + /// `do_stream` fails with a non-retryable 400. + NonRetryableFirstError, +} + +struct PipelineModel { + behavior: Behavior, + calls: AtomicUsize, +} + +impl PipelineModel { + fn new(behavior: Behavior) -> Self { + Self { + behavior, + calls: AtomicUsize::new(0), + } + } +} + +fn delta(text: &str) -> StreamPart { + StreamPart::TextDelta { + id: "1".into(), + delta: text.into(), + provider_metadata: None, + } +} + +fn finish() -> StreamPart { + StreamPart::Finish { + finish_reason: FinishReason { + unified: FinishReasonUnified::Error, + raw: None, + }, + usage: Usage::default(), + provider_metadata: None, + } +} + +fn api_error(status: u16, retryable: bool) -> AiMuxError { + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + is_retryable: retryable, + ..ApiCallError::new( + "failed", + "https://example.test/stream", + serde_json::json!({}), + ) + })) +} + +fn stream_result( + stream: impl Stream> + Send + 'static, +) -> StreamResult { + StreamResult { + stream: Box::pin(stream), + request: None, + response: None, + } +} + +#[async_trait] +impl LanguageModel for PipelineModel { + fn provider(&self) -> &str { + "mock.pipeline" + } + + fn model_id(&self) -> &str { + "pipeline-1" + } + + async fn do_generate(&self, _options: &CallOptions) -> Result { + unimplemented!("streaming-only mock") + } + + async fn do_stream(&self, _options: &CallOptions) -> Result { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst); + match self.behavior { + Behavior::Normal => Ok(stream_result(futures::stream::iter([Ok(delta("hello"))]))), + Behavior::Probe(second_is_output) => Ok(stream_result(async_stream::stream! { + yield Ok(delta("first")); + tokio::time::sleep(Duration::from_millis(600)).await; + if second_is_output { + yield Ok(delta("second")); + } else { + yield Ok(StreamPart::TextStart { id: "1".into(), provider_metadata: None }); + } + futures::future::pending::<()>().await; + })), + Behavior::FinishThenPending => Ok(stream_result( + futures::stream::iter([Ok(delta("done")), Ok(finish())]) + .chain(futures::stream::pending()), + )), + Behavior::ProviderErrorThenFinish => Ok(stream_result( + futures::stream::iter([ + Ok(StreamPart::Error { + error: AiMuxError::Other("provider error".into()), + }), + Ok(finish()), + ]) + .chain(futures::stream::pending()), + )), + Behavior::ParseErrorThenDelta => Ok(stream_result(futures::stream::iter([ + Err(AiMuxError::JsonParse("malformed SSE data".into())), + Ok(delta("after-error")), + ]))), + Behavior::TransportErrorThenPending => Ok(stream_result( + futures::stream::iter([Err(api_error(502, true))]) + .chain(futures::stream::pending()), + )), + Behavior::NeverReturns => std::future::pending().await, + Behavior::RetryableFirstError if attempt == 0 => Err(api_error(429, true)), + Behavior::RetryableFirstError => { + Ok(stream_result(futures::stream::iter([Ok(delta("hello"))]))) + } + Behavior::NonRetryableFirstError => Err(api_error(400, false)), + } + } +} + +fn options(timeout: TimeoutConfiguration) -> GenerateTextOptions { + GenerateTextOptions { + max_retries: Some(0), + timeout: Some(timeout), + ..GenerateTextOptions::default() + } +} + +fn chunk_options() -> GenerateTextOptions { + options(TimeoutConfiguration { + chunk_ms: Some(1_000), + ..TimeoutConfiguration::default() + }) +} + +/// The text of each `TextDelta` among `parts`. +fn texts(parts: &[Result]) -> Vec { + parts + .iter() + .filter_map(|part| match part { + Ok(TextStreamPart::TextDelta { text, .. }) => Some(text.clone()), + _ => None, + }) + .collect() +} + +/// Reads the next `TextDelta`'s text, skipping framing parts. +async fn next_text( + stream: &mut (impl Stream> + Unpin), +) -> String { + loop { + match stream.next().await.expect("stream must not end") { + Ok(TextStreamPart::TextDelta { text, .. }) => return text, + Ok(_) => {} + Err(error) => panic!("unexpected error: {error}"), + } + } +} + +/// Drives the probe stream to the chunk timeout and returns when it fired. +async fn time_to_chunk_timeout(second_part_is_output: bool) -> Duration { + let model = PipelineModel::new(Behavior::Probe(second_part_is_output)); + let mut result = stream_text(&model, "hello", chunk_options()).await.unwrap(); + let start = tokio::time::Instant::now(); + loop { + match result.stream.next().await.expect("stream must not end") { + Ok(_) => {} + Err(AiMuxError::Timeout(message)) => { + assert_eq!(message, "Chunk timeout of 1000ms exceeded"); + return start.elapsed(); + } + Err(other) => panic!("unexpected error: {other}"), + } + } +} + +/// §13-17, stream-text.ts:182-225: a `text-start` is not an output chunk, so +/// the timeout fires 1000ms after the only output at t=0, not at t=1600. +#[tokio::test(start_paused = true)] +async fn non_output_part_does_not_reset_the_chunk_timer() { + assert_eq!( + time_to_chunk_timeout(false).await, + Duration::from_millis(1_000) + ); +} + +/// §13-17, stream-text.ts:2842-2857: the second delta at t=600 re-arms the +/// deadline to t=1600. +#[tokio::test(start_paused = true)] +async fn output_resets_the_chunk_timer() { + assert_eq!( + time_to_chunk_timeout(true).await, + Duration::from_millis(1_600) + ); +} + +/// §13-18: the timers run on the producer side. The second delta arrives at +/// t=600, inside the budget; a consumer that polls again only at t=2000 +/// still receives it, and the next poll sees the timeout that fired at +/// t=1600. +#[tokio::test(start_paused = true)] +async fn consumer_delay_does_not_count_against_the_chunk_timer() { + let model = PipelineModel::new(Behavior::Probe(true)); + let mut result = stream_text(&model, "hello", chunk_options()).await.unwrap(); + assert_eq!(next_text(&mut result.stream).await, "first"); + + tokio::time::sleep(Duration::from_millis(2_000)).await; + + assert_eq!(next_text(&mut result.stream).await, "second"); + match result.stream.next().await.expect("timeout") { + Err(AiMuxError::Timeout(message)) => { + assert_eq!(message, "Chunk timeout of 1000ms exceeded"); + } + other => panic!("expected the chunk timeout, got {other:?}"), + } + assert!(result.stream.next().await.is_none()); +} + +/// §13-21b: the deadlines live in the pump task, which is aborted when the +/// returned stream is dropped; no timer is left to fire later and touch the +/// caller's signal. +#[tokio::test(start_paused = true)] +async fn a_dropped_stream_leaves_no_armed_timer() { + let caller = AbortSignal::new(); + let model = PipelineModel::new(Behavior::Probe(false)); + let options = GenerateTextOptions { + abort_signal: Some(caller.clone()), + ..options(TimeoutConfiguration { + total_ms: Some(5_000), + chunk_ms: Some(1_000), + ..TimeoutConfiguration::default() + }) + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + assert_eq!(next_text(&mut result.stream).await, "first"); + drop(result); + + tokio::time::sleep(Duration::from_secs(60)).await; + assert!(!caller.is_aborted()); +} + +/// §8.3: the event the provider peeked inside `do_stream` is not lost. +#[tokio::test] +async fn stream_setup_preserves_the_peeked_first_event() { + let model = PipelineModel::new(Behavior::Normal); + let result = stream_text(&model, "hello", GenerateTextOptions::default()) + .await + .unwrap(); + let parts: Vec<_> = result.stream.collect().await; + assert_eq!(texts(&parts), ["hello"]); +} + +/// §13-18: the first-chunk budget measures the arrival of the first output, +/// not when the consumer first polls. +#[tokio::test] +async fn a_late_consumer_still_receives_a_first_chunk_delivered_on_time() { + let model = PipelineModel::new(Behavior::Normal); + let result = stream_text( + &model, + "hello", + options(TimeoutConfiguration { + first_chunk_ms: Some(1), + ..TimeoutConfiguration::default() + }), + ) + .await + .unwrap(); + + tokio::time::sleep(Duration::from_millis(20)).await; + + let parts: Vec<_> = result.stream.collect().await; + assert!(parts.iter().all(Result::is_ok), "{parts:?}"); + assert_eq!(texts(&parts), ["hello"]); +} + +/// The model's `finish` ends the stream even when the provider keeps the +/// connection open: no total timeout fires afterwards. +#[tokio::test(start_paused = true)] +async fn the_finish_part_ends_the_stream() { + let model = PipelineModel::new(Behavior::FinishThenPending); + let result = stream_text( + &model, + "hello", + options(TimeoutConfiguration { + total_ms: Some(50), + ..TimeoutConfiguration::default() + }), + ) + .await + .unwrap(); + let parts: Vec<_> = result.stream.collect().await; + let kinds: Vec = parts + .iter() + .map(|part| part.as_ref().map(super::tests::kind).unwrap_or_default()) + .collect(); + assert_eq!( + kinds, + ["Start", "StartStep", "TextDelta", "FinishStep", "Finish"] + ); +} + +/// A provider's error part does not end the stream before its `finish`. +#[tokio::test] +async fn a_provider_error_part_does_not_hide_the_finish() { + let model = PipelineModel::new(Behavior::ProviderErrorThenFinish); + let result = stream_text(&model, "hello", GenerateTextOptions::default()) + .await + .unwrap(); + let parts: Vec<_> = result.stream.collect().await; + let kinds: Vec = parts + .iter() + .map(|part| part.as_ref().map(super::tests::kind).unwrap_or_default()) + .collect(); + assert_eq!( + kinds, + ["Start", "StartStep", "Error", "FinishStep", "Finish"] + ); + let Some(Ok(TextStreamPart::Finish { finish_reason, .. })) = parts.last() else { + unreachable!() + }; + assert_eq!(*finish_reason, FinishReasonUnified::Error); +} + +/// A malformed frame is reported and the stream goes on, with a total +/// timeout or an abort signal armed alike. +#[tokio::test] +async fn the_stream_continues_after_a_malformed_frame() { + for options in [ + options(TimeoutConfiguration { + total_ms: Some(1_000), + ..TimeoutConfiguration::default() + }), + GenerateTextOptions { + abort_signal: Some(AbortSignal::new()), + ..GenerateTextOptions::default() + }, + ] { + let model = PipelineModel::new(Behavior::ParseErrorThenDelta); + let result = stream_text(&model, "hello", options).await.unwrap(); + let parts: Vec<_> = result.stream.collect().await; + let error = parts + .iter() + .position(|part| matches!(part, Err(AiMuxError::JsonParse(_)))); + let text = parts.iter().position( + |part| matches!(part, Ok(TextStreamPart::TextDelta { text, .. }) if text == "after-error"), + ); + assert!( + error.is_some() && text.is_some() && error < text, + "{parts:?}" + ); + } +} + +/// A transport error ends the stream instead of waiting on the dead source. +#[tokio::test] +async fn a_transport_error_ends_the_stream() { + let model = PipelineModel::new(Behavior::TransportErrorThenPending); + let mut result = stream_text(&model, "hello", GenerateTextOptions::default()) + .await + .unwrap(); + loop { + match result.stream.next().await.expect("the error part") { + Ok(_) => {} + Err(error) => { + assert!(matches!(error, AiMuxError::ApiCall(_)), "{error:?}"); + break; + } + } + } + assert!(result.stream.next().await.is_none()); +} + +/// §8.3: with `first_chunk_ms` alone, a provider awaiting its first event +/// inside `do_stream` is bounded. +#[tokio::test] +async fn the_first_chunk_timeout_bounds_stream_setup() { + let model = PipelineModel::new(Behavior::NeverReturns); + let Err(error) = stream_text( + &model, + "hello", + options(TimeoutConfiguration { + first_chunk_ms: Some(5), + ..TimeoutConfiguration::default() + }), + ) + .await + else { + panic!("expected a timeout"); + }; + assert!( + matches!(error, AiMuxError::Timeout(ref message) if message == "First chunk timeout of 5ms exceeded"), + "got {error:?}" + ); +} + +/// A retryable error from `do_stream` (the provider's first event, e.g. a +/// 429) is retried before the stream is returned (stream-text.ts:2409-2420 +/// wraps the model call in `prepareRetries`' retry). +#[tokio::test(start_paused = true)] +async fn a_retryable_first_stream_error_is_retried() { + let model = PipelineModel::new(Behavior::RetryableFirstError); + let options = GenerateTextOptions { + max_retries: Some(2), + ..GenerateTextOptions::default() + }; + let result = stream_text(&model, "hello", options).await.unwrap(); + assert_eq!(model.calls.load(Ordering::SeqCst), 2); + let parts: Vec<_> = result.stream.collect().await; + assert_eq!(texts(&parts), ["hello"]); +} + +/// A non-retryable error from `do_stream` is returned once, unchanged. +#[tokio::test] +async fn a_non_retryable_first_stream_error_is_returned_unchanged() { + let model = PipelineModel::new(Behavior::NonRetryableFirstError); + let options = GenerateTextOptions { + max_retries: Some(2), + ..GenerateTextOptions::default() + }; + let Err(error) = stream_text(&model, "hello", options).await else { + panic!("expected the error"); + }; + assert!(matches!(error, AiMuxError::ApiCall(ref detail) if detail.status_code == Some(400))); + assert_eq!(model.calls.load(Ordering::SeqCst), 1); +} diff --git a/aimux-core/src/generate/tests.rs b/aimux-core/src/generate/tests.rs new file mode 100644 index 000000000..104c53879 --- /dev/null +++ b/aimux-core/src/generate/tests.rs @@ -0,0 +1,1434 @@ +//! Call-layer behaviour of `generate_text` / `stream_text`, checked against +//! the AI SDK sources (ai@7.0.127) cited on each test. + +use std::sync::Mutex; + +use async_trait::async_trait; +use futures::StreamExt; +use serde_json::json; + +use super::*; +use crate::result::GenerateResult; +use crate::shared::{ResponseInfo, StreamResponseInfo}; +use crate::tool::RawToolCall; +use crate::types::{FinishReason, InputTokenUsage, OutputTokenUsage, Usage}; + +/// A model that replays scripted results. +pub(super) struct ScriptedModel { + pub generate: Mutex>, + pub streams: Mutex>>>, + /// Keep each scripted stream open after its last part. + pub hang_after_script: bool, +} + +impl ScriptedModel { + pub fn generating(result: GenerateResult) -> Self { + Self { + generate: Mutex::new(Some(result)), + streams: Mutex::new(Vec::new()), + hang_after_script: false, + } + } + + pub fn streaming(parts: Vec>) -> Self { + Self { + generate: Mutex::new(None), + streams: Mutex::new(vec![parts]), + hang_after_script: false, + } + } +} + +#[async_trait] +impl LanguageModel for ScriptedModel { + fn provider(&self) -> &str { + "mock.chat" + } + + fn model_id(&self) -> &str { + "mock-model" + } + + async fn do_generate(&self, _options: &CallOptions) -> Result { + Ok(self + .generate + .lock() + .unwrap() + .clone() + .expect("a scripted generate result")) + } + + async fn do_stream(&self, _options: &CallOptions) -> Result { + let parts = self.streams.lock().unwrap().remove(0); + let stream = futures::stream::iter(parts); + let stream: Pin> + Send>> = + if self.hang_after_script { + Box::pin(stream.chain(futures::stream::pending())) + } else { + Box::pin(stream) + }; + Ok(StreamResult { + stream, + request: Some(RequestInfo { + body: Some(json!({ "model": "mock-model" })), + }), + response: Some(StreamResponseInfo { + headers: Some(HashMap::from([( + "x-request-id".to_string(), + "req-1".to_string(), + )])), + }), + }) + } +} + +pub(super) fn usage(input: u32, cache_read: u32, output: u32, reasoning: u32) -> Usage { + Usage { + input_tokens: InputTokenUsage { + total: Some(input), + no_cache: Some(input - cache_read), + cache_read: Some(cache_read), + cache_write: None, + }, + output_tokens: OutputTokenUsage { + total: Some(output), + text: Some(output - reasoning), + reasoning: Some(reasoning), + }, + raw: Some(serde_json::Map::from_iter([( + "prompt_tokens".to_string(), + json!(input), + )])), + } +} + +fn finish(unified: FinishReasonUnified, raw: &str) -> FinishReason { + FinishReason { + unified, + raw: Some(raw.to_string()), + } +} + +pub(super) fn text_stream( + texts: &[&str], + reason: FinishReasonUnified, +) -> Vec> { + let mut parts = vec![ + Ok(StreamPart::StreamStart { + warnings: vec![Warning::Other { + message: "careful".into(), + }], + }), + Ok(StreamPart::ResponseMetadata( + crate::types::ResponseMetadata { + id: Some("resp-1".into()), + timestamp: Some("2026-01-01T00:00:00.000Z".into()), + model_id: Some("mock-model-2026".into()), + }, + )), + Ok(StreamPart::TextStart { + id: "t1".into(), + provider_metadata: None, + }), + ]; + for text in texts { + parts.push(Ok(StreamPart::TextDelta { + id: "t1".into(), + delta: (*text).into(), + provider_metadata: None, + })); + } + parts.push(Ok(StreamPart::TextEnd { + id: "t1".into(), + provider_metadata: None, + })); + parts.push(Ok(StreamPart::Finish { + finish_reason: finish(reason, "end_turn"), + usage: usage(10, 4, 6, 2), + provider_metadata: None, + })); + parts +} + +pub(super) async fn collect(result: StreamTextResult) -> Vec { + result + .stream + .map(|part| part.expect("no transport error")) + .collect() + .await +} + +pub(super) fn kind(part: &TextStreamPart) -> String { + let value = serde_json::to_value(part).unwrap(); + value + .as_object() + .and_then(|object| object.keys().next().cloned()) + .unwrap_or_else(|| value.as_str().unwrap_or_default().to_string()) +} + +fn generated(content: Vec, reason: FinishReasonUnified) -> GenerateResult { + GenerateResult { + content, + finish_reason: finish(reason, "stop"), + usage: usage(10, 4, 6, 2), + warnings: vec![Warning::Other { + message: "careful".into(), + }], + provider_metadata: Some(indexmap::IndexMap::from([( + "mock".to_string(), + serde_json::Map::from_iter([("k".to_string(), json!(1))]), + )])), + request: Some(RequestInfo { + body: Some(json!({ "model": "mock-model" })), + }), + response: Some(ResponseInfo { + id: Some("resp-1".into()), + timestamp: None, + model_id: None, + headers: None, + body: Some(json!({ "id": "resp-1" })), + }), + } +} + +fn text_content(text: &str) -> GenerateContent { + GenerateContent::Text { + text: text.into(), + provider_metadata: None, + } +} + +/// stream-text-result.ts:581-607 and stream-text.ts:1855, :2830-2840, +/// :3032-3048, :3128-3134: `start`, `start-step` on the step's first chunk +/// (the provider's `stream-start` only carries warnings), the content, +/// `finish-step` with the step's response, usage, performance and finish +/// reason, then `finish` with the total usage. Text deltas carry `text` +/// (stream-text-result.ts:436-441). +#[tokio::test] +async fn the_stream_is_framed_as_the_ai_sdk_full_stream() { + let model = ScriptedModel::streaming(text_stream(&["Hel", "lo"], FinishReasonUnified::Stop)); + let parts = collect( + stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(), + ) + .await; + let kinds: Vec = parts.iter().map(kind).collect(); + assert_eq!( + kinds, + [ + "Start", + "StartStep", + "TextStart", + "TextDelta", + "TextDelta", + "TextEnd", + "FinishStep", + "Finish" + ] + ); + + let TextStreamPart::StartStep { request, warnings } = &parts[1] else { + unreachable!() + }; + assert_eq!(request.body, Some(json!({ "model": "mock-model" }))); + assert_eq!( + serde_json::to_value(warnings).unwrap(), + json!([{ "Other": { "message": "careful" } }]) + ); + + assert_eq!( + serde_json::to_value(&parts[3]).unwrap(), + json!({ "TextDelta": { "id": "t1", "text": "Hel" } }) + ); + + let TextStreamPart::FinishStep { + response, + usage, + performance, + finish_reason, + raw_finish_reason, + provider_metadata, + } = &parts[6] + else { + unreachable!() + }; + assert_eq!(response.id, "resp-1"); + assert_eq!(response.model_id, "mock-model-2026"); + assert_eq!(response.timestamp, "2026-01-01T00:00:00.000Z"); + assert_eq!( + response.headers.as_ref().unwrap()["x-request-id"], + "req-1".to_string() + ); + assert_eq!(*finish_reason, FinishReasonUnified::Stop); + assert_eq!(raw_finish_reason.as_deref(), Some("end_turn")); + assert!(provider_metadata.is_none()); + // usage.ts:81-102: the step usage keeps the provider's raw usage. + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.input_token_details.cache_read_tokens, Some(4)); + assert_eq!(usage.input_token_details.no_cache_tokens, Some(6)); + assert_eq!(usage.output_token_details.reasoning_tokens, Some(2)); + assert_eq!(usage.total_tokens, Some(16)); + assert!(usage.raw.is_some()); + // stream-language-model-call.ts:597-624: a finished streamed call has + // the streaming rates and the first-output timing. + assert!(performance.time_to_first_output_ms.is_some()); + assert!(performance.output_tokens_per_second.is_some()); + assert!(performance.input_tokens_per_second.is_some()); + assert_eq!( + performance + .time_between_output_chunks_ms + .map(|stats| stats.min >= 0.0), + Some(true), + "two output chunks give one gap" + ); + assert!(performance.step_time_ms >= performance.response_time_ms); + + let TextStreamPart::Finish { + finish_reason, + raw_finish_reason, + total_usage, + } = &parts[7] + else { + unreachable!() + }; + assert_eq!(*finish_reason, FinishReasonUnified::Stop); + assert_eq!(raw_finish_reason.as_deref(), Some("end_turn")); + // usage.ts:122-155: the sum carries no raw usage. + assert_eq!(total_usage.total_tokens, Some(16)); + assert!(total_usage.raw.is_none()); + assert_eq!( + serde_json::to_value(&parts[7]).unwrap()["Finish"]["finish_reason"], + json!("stop") + ); +} + +/// stream-text.ts:1480-1700: the consumed stream records one step from the +/// parts between `start-step` and `finish-step`; text, usage and the response +/// messages derive from it as in `generateText`. +#[tokio::test] +async fn a_consumed_stream_aggregates_its_step() { + let model = ScriptedModel::streaming(text_stream(&["Hel", "lo"], FinishReasonUnified::Stop)); + let result = stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap() + .consume() + .await + .unwrap(); + assert_eq!(result.steps().len(), 1); + let step = result.final_step(); + assert_eq!(step.step_number, 0); + assert!(!step.call_id.is_empty()); + assert_eq!(step.model.provider, "mock.chat"); + assert_eq!(step.model.model_id, "mock-model"); + assert_eq!(step.text(), "Hello"); + assert_eq!(step.response.id, "resp-1"); + assert_eq!(step.response.body, None); + assert_eq!(step.request.body, Some(json!({ "model": "mock-model" }))); + assert_eq!(step.warnings.len(), 1); + assert_eq!(result.text(), "Hello"); + assert_eq!(result.finish_reason(), FinishReasonUnified::Stop); + assert_eq!(result.raw_finish_reason(), Some("end_turn")); + assert_eq!(result.usage().total_tokens, Some(16)); + assert_eq!(result.output().unwrap(), &json!("Hello")); + assert_eq!(result.response_messages().len(), 1); +} + +/// generate-text.ts:1382-1401 and :1691-1808: a generated step has the +/// response-time rates but no streaming timing; the result's `usage` is the +/// total over the steps (no raw usage), the step's usage keeps it; the +/// response falls back to a generated id, the current time and the model's +/// id (:1058-1064). +#[tokio::test] +async fn generate_text_derives_its_result_from_the_step() { + let call = RawToolCall { + tool_call_id: "call-1".into(), + tool_name: "weather".into(), + input: r#"{"city":"Paris"}"#.into(), + provider_executed: None, + dynamic: None, + provider_metadata: None, + }; + let model = ScriptedModel::generating(generated( + vec![ + text_content("Checking "), + GenerateContent::ToolCall(call), + text_content("the weather."), + ], + FinishReasonUnified::ToolCalls, + )); + let options = GenerateTextOptions { + tools: Some(vec![Tool::Function(crate::tool::FunctionTool::new( + "weather", + json!({ + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"], + }), + ))]), + ..GenerateTextOptions::default() + }; + let result = generate_text(&model, "hi", options).await.unwrap(); + let step = result.final_step(); + assert_eq!(step.response.id, "resp-1"); + assert_eq!(step.response.model_id, "mock-model"); + assert!(!step.response.timestamp.is_empty()); + assert_eq!(step.response.body, Some(json!({ "id": "resp-1" }))); + assert_eq!(step.performance.time_to_first_output_ms, None); + assert_eq!(step.performance.output_tokens_per_second, None); + assert_eq!(step.performance.time_between_output_chunks_ms, None); + assert!(step.usage.raw.is_some()); + + assert_eq!(result.text(), "Checking the weather."); + assert_eq!(result.content().len(), 3); + assert_eq!(result.tool_calls().len(), 1); + assert_eq!(result.tool_calls()[0].input, json!({ "city": "Paris" })); + assert_eq!(result.static_tool_calls().len(), 1); + assert!(result.dynamic_tool_calls().is_empty()); + assert_eq!(result.usage().input_tokens, Some(10)); + assert!(result.usage().raw.is_none()); + assert_eq!(result.warnings().len(), 1); + assert_eq!(result.response_messages().len(), 1); + #[allow(deprecated)] + { + assert_eq!(result.total_usage(), result.usage()); + assert_eq!(result.provider_metadata().unwrap()["mock"]["k"], json!(1)); + assert_eq!(result.response().id, "resp-1"); + assert_eq!(result.reasoning_text(), None); + } + // generate-text.ts:1599-1612: a tool-calls step is not parsed, and + // reading the output is a NoOutputGeneratedError. + assert!(matches!( + result.output(), + Err(AiMuxError::NoOutputGenerated(message)) if message == "No output generated." + )); +} + +/// The serialized result writes the derived values next to the stored ones +/// and reads back only `steps`, `total_usage` and `output`. The finish reason +/// is a string and the usage is the flat `LanguageModelUsage` +/// (usage.ts:10-68). +#[tokio::test] +async fn the_result_snapshot_round_trips() { + let model = ScriptedModel::generating(generated( + vec![text_content("Hello")], + FinishReasonUnified::Stop, + )); + let result = generate_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(); + let wire = serde_json::to_value(&result).unwrap(); + assert_eq!(wire["text"], json!("Hello")); + assert_eq!(wire["finish_reason"], json!("stop")); + assert_eq!(wire["raw_finish_reason"], json!("stop")); + assert_eq!(wire["usage"]["input_tokens"], json!(10)); + assert_eq!( + wire["usage"]["input_token_details"], + json!({ "no_cache_tokens": 6, "cache_read_tokens": 4, "cache_write_tokens": null }) + ); + assert_eq!(wire["usage"]["total_tokens"], json!(16)); + assert_eq!(wire["output"], json!("Hello")); + assert_eq!(wire["steps"].as_array().unwrap().len(), 1); + assert!(wire.get("raw").is_none()); + + let back: GenerateTextResult = serde_json::from_value(wire.clone()).unwrap(); + assert_eq!(serde_json::to_value(&back).unwrap(), wire); + + let mut no_steps = wire; + no_steps["steps"] = json!([]); + assert!(serde_json::from_value::(no_steps).is_err()); +} + +/// output.ts and generate-text.ts:1599-1612: the output's response format +/// goes to the model; a stop step is parsed and a parse failure fails +/// `generateText`. +#[tokio::test] +async fn generate_text_parses_the_output() { + let object = Output::Object { + schema: json!({ + "type": "object", + "properties": { "name": { "type": "string" } }, + "required": ["name"], + }), + name: None, + description: None, + }; + let options = || GenerateTextOptions { + output: Some(object.clone()), + ..GenerateTextOptions::default() + }; + let model = ScriptedModel::generating(generated( + vec![text_content(r#"{"name":"Ada"}"#)], + FinishReasonUnified::Stop, + )); + let result = generate_text(&model, "hi", options()).await.unwrap(); + assert_eq!(result.output().unwrap(), &json!({ "name": "Ada" })); + + let model = ScriptedModel::generating(generated( + vec![text_content("not json")], + FinishReasonUnified::Stop, + )); + let error = generate_text(&model, "hi", options()).await.unwrap_err(); + assert!( + matches!(&error, AiMuxError::JsonParse(message) + if message.starts_with("No object generated: could not parse the response.")), + "{error:?}" + ); + + // A length step with text is parsed too (`finishReason !== 'tool-calls' + // && text.length > 0`). + let model = ScriptedModel::generating(generated( + vec![text_content(r#"{"name":"Bo"}"#)], + FinishReasonUnified::Length, + )); + let result = generate_text(&model, "hi", options()).await.unwrap(); + assert_eq!(result.output().unwrap(), &json!({ "name": "Bo" })); +} + +/// stream-text.ts:3419-3439: a streamed result always parses its final text, +/// and a failure only fails reading `output`. +#[tokio::test] +async fn a_stream_output_failure_only_fails_the_output() { + let model = ScriptedModel::streaming(text_stream(&["not json"], FinishReasonUnified::Stop)); + let result = stream_text( + &model, + "hi", + GenerateTextOptions { + output: Some(Output::Json { + name: None, + description: None, + }), + ..GenerateTextOptions::default() + }, + ) + .await + .unwrap() + .consume() + .await + .unwrap(); + assert_eq!(result.text(), "not json"); + assert!(matches!(result.output(), Err(AiMuxError::JsonParse(_)))); + assert!( + serde_json::to_value(&result) + .unwrap() + .get("output") + .is_none() + ); +} + +/// stream-text.ts:1860-1882: an aborted stream ends with an `abort` part +/// carrying the abort reason's message, and the result fails with the abort +/// when no step finished (:1724-1736). +#[tokio::test] +async fn an_aborted_stream_ends_with_an_abort_part() { + let mut model = ScriptedModel::streaming(vec![ + Ok(StreamPart::StreamStart { + warnings: Vec::new(), + }), + Ok(StreamPart::TextDelta { + id: "t1".into(), + delta: "partial".into(), + provider_metadata: None, + }), + ]); + model.hang_after_script = true; + let signal = AbortSignal::new(); + let result = stream_text( + &model, + "hi", + GenerateTextOptions { + abort_signal: Some(signal.clone()), + ..GenerateTextOptions::default() + }, + ) + .await + .unwrap(); + let mut stream = result.stream; + let mut kinds = Vec::new(); + while let Some(part) = stream.next().await { + let part = part.unwrap(); + kinds.push(kind(&part)); + if matches!(part, TextStreamPart::TextDelta { .. }) { + signal.abort(); + } + if let TextStreamPart::Abort { reason } = &part { + assert_eq!(reason.as_deref(), Some("request aborted")); + } + } + assert_eq!(kinds, ["Start", "StartStep", "TextDelta", "Abort"]); + + let mut model = ScriptedModel::streaming(vec![Ok(StreamPart::StreamStart { + warnings: Vec::new(), + })]); + model.hang_after_script = true; + let signal = AbortSignal::new(); + let result = stream_text( + &model, + "hi", + GenerateTextOptions { + abort_signal: Some(signal.clone()), + ..GenerateTextOptions::default() + }, + ) + .await + .unwrap(); + signal.abort(); + assert!(matches!( + result.consume().await, + Err(AiMuxError::Aborted(message)) if message == "request aborted" + )); +} + +/// stream-text.ts:3015-3030: a model stream that ends without a finish chunk +/// and without output ends with a NoOutputGeneratedError part and no +/// `finish-step`; with partial output the step finishes with `other`. +#[tokio::test] +async fn an_incomplete_stream_without_output_is_an_error() { + let model = ScriptedModel::streaming(vec![Ok(StreamPart::StreamStart { + warnings: Vec::new(), + })]); + let parts = collect( + stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(), + ) + .await; + let kinds: Vec = parts.iter().map(kind).collect(); + assert_eq!(kinds, ["Start", "Error"]); + assert!(matches!( + &parts[1], + TextStreamPart::Error { error: AiMuxError::NoOutputGenerated(message) } + if message == "No output generated. The model stream ended without a finish chunk." + )); + + let model = ScriptedModel::streaming(vec![Ok(StreamPart::TextDelta { + id: "t1".into(), + delta: "partial".into(), + provider_metadata: None, + })]); + let parts = collect( + stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(), + ) + .await; + let kinds: Vec = parts.iter().map(kind).collect(); + assert_eq!( + kinds, + ["Start", "StartStep", "TextDelta", "FinishStep", "Finish"] + ); + let TextStreamPart::FinishStep { + finish_reason, + usage, + performance, + .. + } = &parts[3] + else { + unreachable!() + }; + assert_eq!(*finish_reason, FinishReasonUnified::Other); + assert_eq!(usage, &LanguageModelUsage::default()); + // stream-text.ts:2755-2767: without a finish chunk the model call's + // performance stays zero. + assert_eq!(performance.response_time_ms, 0.0); + assert_eq!(performance.time_to_first_output_ms, None); +} + +/// generate-object.ts:120 deprecates `generateObject` for +/// `generateText({ output })`; the wrapper returns the AI SDK's +/// `GenerateObjectResult` fields (generate-object-result.ts:17-55). +#[tokio::test] +#[allow(deprecated)] +async fn generate_object_wraps_the_output() { + let model = ScriptedModel::generating(generated( + vec![text_content(r#"{"name":"Ada"}"#)], + FinishReasonUnified::Stop, + )); + let result = generate_object(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(); + assert_eq!(result.object, json!({ "name": "Ada" })); + assert_eq!(result.finish_reason, FinishReasonUnified::Stop); + assert_eq!(result.usage.input_tokens, Some(10)); + assert_eq!(result.response.id, "resp-1"); + assert_eq!(result.response.body, Some(json!({ "id": "resp-1" }))); + assert_eq!(result.reasoning, None); + let wire = serde_json::to_value(&result).unwrap(); + assert!(wire.get("raw").is_none()); + assert!(wire.get("raw_finish_reason").is_none()); + + let model = ScriptedModel::generating(generated( + vec![text_content("[1,2")], + FinishReasonUnified::Stop, + )); + assert!(matches!( + generate_object(&model, "hi", GenerateTextOptions::default()).await, + Err(AiMuxError::JsonParse(_)) + )); +} + +/// stream-text.ts:2881-2891: a `text-delta` is enqueued only when its text is +/// non-empty or it carries provider metadata. +#[tokio::test] +async fn an_empty_text_delta_is_forwarded_only_with_provider_metadata() { + let metadata = indexmap::IndexMap::from([( + "mock".to_string(), + serde_json::Map::from_iter([("itemId".to_string(), json!("msg_1"))]), + )]); + let mut parts = text_stream(&["", "hi"], FinishReasonUnified::Stop); + parts.insert( + 4, + Ok(StreamPart::TextDelta { + id: "t1".into(), + delta: String::new(), + provider_metadata: Some(metadata.clone()), + }), + ); + let model = ScriptedModel::streaming(parts); + let result = stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(); + let deltas: Vec<(String, bool)> = collect(result) + .await + .into_iter() + .filter_map(|part| match part { + TextStreamPart::TextDelta { + text, + provider_metadata, + .. + } => Some((text, provider_metadata.is_some())), + _ => None, + }) + .collect(); + assert_eq!( + deltas, + vec![(String::new(), true), ("hi".to_string(), false)] + ); +} + +fn provider_result(id: &str, name: &str, result: Value, is_error: Option) -> ToolResult { + ToolResult { + tool_call_id: id.into(), + tool_name: name.into(), + result, + is_error, + preliminary: None, + dynamic: Some(true), + provider_metadata: None, + } +} + +/// convert-language-model-content.ts:87-104: a provider tool result whose +/// call is not in the response fails with `Tool call {id} not found.` unless +/// it belongs to a provider tool of that name with `supportsDeferredResults`. +#[tokio::test] +async fn generate_text_rejects_a_tool_result_without_its_call() { + let deferred = GenerateContent::ToolResult(provider_result( + "srv-0", + "code_execution", + json!({ "stdout": "2" }), + None, + )); + let run = |tools: Option>| { + let model = ScriptedModel::generating(generated( + vec![text_content("hi"), deferred.clone()], + FinishReasonUnified::Stop, + )); + async move { + let options = GenerateTextOptions { + tools, + ..GenerateTextOptions::default() + }; + generate_text(&model, "hi", options).await + } + }; + let code_execution = |supports_deferred_results| { + Tool::Provider(crate::tool::ProviderTool { + id: crate::types::NamespacedKind::new("anthropic.code_execution_20250825").unwrap(), + name: "code_execution".into(), + args: serde_json::Map::new(), + supports_deferred_results, + }) + }; + + for tools in [None, Some(vec![code_execution(None)])] { + let error = run(tools).await.unwrap_err(); + assert_eq!(error.to_string(), "Tool call srv-0 not found."); + } + let result = run(Some(vec![code_execution(Some(true))])).await.unwrap(); + assert!(result.final_step().content.iter().any(|part| matches!( + part, + ContentPart::ToolResult(result) if result.tool_call_id == "srv-0" + ))); +} + +fn provider_call(id: &str, name: &str, input: &str) -> StreamPart { + StreamPart::ToolCall(RawToolCall { + tool_call_id: id.into(), + tool_name: name.into(), + input: input.into(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: None, + }) +} + +/// Provider tool calls, their results (one an error) and a deferred result, +/// before the step's `finish`. +fn tool_stream() -> Vec> { + let mut parts = text_stream(&["hi"], FinishReasonUnified::ToolCalls); + let finish = parts.pop().unwrap(); + parts.extend([ + Ok(provider_call("srv-1", "web_search", r#"{"q":"rust"}"#)), + Ok(StreamPart::ToolResult(provider_result( + "srv-1", + "web_search", + json!([{ "url": "https://example.com" }]), + None, + ))), + Ok(provider_call( + "srv-2", + "code_execution", + r#"{"code":"1/0"}"#, + )), + Ok(StreamPart::ToolResult(provider_result( + "srv-2", + "code_execution", + json!({ "error_code": "unavailable" }), + Some(true), + ))), + Ok(StreamPart::ToolResult(provider_result( + "srv-0", + "code_execution", + json!({ "stdout": "2" }), + None, + ))), + ]); + parts.push(finish); + parts +} + +/// stream-language-model-call.ts:760-792: a provider `tool-result` becomes a +/// `tool-result` (or, with `isError`, a `tool-error`) part with +/// `providerExecuted: true`, the chunk's `dynamic` and the input of its call +/// (`undefined` for a deferred result, whose call is not in the stream). +#[tokio::test] +async fn stream_tool_results_and_errors_follow_the_ai_sdk() { + let model = ScriptedModel::streaming(tool_stream()); + let parts = collect( + stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(), + ) + .await; + let kinds: Vec = parts.iter().map(kind).collect(); + assert_eq!( + kinds, + [ + "Start", + "StartStep", + "TextStart", + "TextDelta", + "TextEnd", + "ToolCall", + "ToolResult", + "ToolCall", + "ToolError", + "ToolResult", + "FinishStep", + "Finish" + ] + ); + assert_eq!( + serde_json::to_value(&parts[6]).unwrap(), + json!({ "ToolResult": { + "tool_call_id": "srv-1", + "tool_name": "web_search", + "input": { "q": "rust" }, + "output": [{ "url": "https://example.com" }], + "provider_executed": true, + "dynamic": true, + } }) + ); + assert_eq!( + serde_json::to_value(&parts[8]).unwrap(), + json!({ "ToolError": { + "tool_call_id": "srv-2", + "tool_name": "code_execution", + "input": { "code": "1/0" }, + "error": { "error_code": "unavailable" }, + "provider_executed": true, + "dynamic": true, + } }) + ); + assert_eq!( + serde_json::to_value(&parts[9]).unwrap(), + json!({ "ToolResult": { + "tool_call_id": "srv-0", + "tool_name": "code_execution", + "output": { "stdout": "2" }, + "provider_executed": true, + "dynamic": true, + } }) + ); +} + +/// The step content of a consumed stream keeps the provider's results and +/// error results (stream-text.ts:1529-1538 records `tool-error` parts, +/// :1648-1650 non-preliminary `tool-result` parts). +#[tokio::test] +async fn a_consumed_stream_keeps_provider_tool_results_and_errors() { + let model = ScriptedModel::streaming(tool_stream()); + let result = stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap() + .consume() + .await + .unwrap(); + let results: Vec<(&str, &str)> = result + .final_step() + .content + .iter() + .filter_map(|part| match part { + ContentPart::ToolResult(result) => Some((result.tool_call_id.as_str(), "result")), + ContentPart::ToolError(error) => Some((error.tool_call_id.as_str(), "error")), + _ => None, + }) + .collect(); + assert_eq!( + results, + [("srv-1", "result"), ("srv-2", "error"), ("srv-0", "result")] + ); +} + +/// stream-language-model-call.ts:709-725: a client tool call that could not +/// be parsed is followed by a `tool-error` with the call's input, +/// `getErrorMessage(error)` and `dynamic: true`; an invalid provider-executed +/// call is not. +#[tokio::test] +async fn an_invalid_tool_call_is_followed_by_a_tool_error() { + let mut parts = text_stream(&["hi"], FinishReasonUnified::ToolCalls); + let finish = parts.pop().unwrap(); + parts.extend([ + Ok(StreamPart::ToolCall(RawToolCall { + tool_call_id: "call-1".into(), + tool_name: "weather".into(), + input: r#"{"city":"Paris"}"#.into(), + provider_executed: None, + dynamic: None, + provider_metadata: None, + })), + Ok(StreamPart::ToolCall(RawToolCall { + tool_call_id: "srv-1".into(), + tool_name: "web_search".into(), + input: "{".into(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: None, + })), + ]); + parts.push(finish); + let model = ScriptedModel::streaming(parts); + let parts = collect( + stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(), + ) + .await; + let kinds: Vec = parts.iter().map(kind).collect(); + assert_eq!( + kinds, + [ + "Start", + "StartStep", + "TextStart", + "TextDelta", + "TextEnd", + "ToolCall", + "ToolError", + "ToolCall", + "FinishStep", + "Finish" + ] + ); + let TextStreamPart::ToolCall(invalid) = &parts[5] else { + unreachable!() + }; + assert_eq!(invalid.invalid, Some(true)); + let TextStreamPart::ToolError(error) = &parts[6] else { + unreachable!() + }; + assert_eq!(error.tool_call_id, "call-1"); + assert_eq!(error.tool_name, "weather"); + assert_eq!(error.input, Some(json!({ "city": "Paris" }))); + assert_eq!( + error.error, + json!(invalid.error.as_ref().unwrap().to_string()) + ); + assert_eq!(error.provider_executed, None); + assert_eq!(error.dynamic, Some(true)); + let TextStreamPart::ToolCall(provider_call) = &parts[7] else { + unreachable!() + }; + assert_eq!(provider_call.invalid, Some(true)); +} + +/// stream-text.ts:1529-1538 records a `tool-error` part in the step content +/// where it arrives, so a consumed stream keeps an invalid call's tool error +/// right after the call (the generated path appends it after the content, +/// generate-text.ts:1303-1323). The response messages replay it in a tool +/// message as error text (to-response-messages.ts:158-225). +#[tokio::test] +async fn a_consumed_stream_keeps_an_invalid_calls_tool_error_in_place() { + let mut parts = text_stream(&["hi"], FinishReasonUnified::ToolCalls); + let finish = parts.pop().unwrap(); + // The call goes before the text, so "in place" differs from "appended". + parts.insert( + 1, + Ok(StreamPart::ToolCall(RawToolCall { + tool_call_id: "call-1".into(), + tool_name: "weather".into(), + input: r#"{"city":"Paris"}"#.into(), + provider_executed: None, + dynamic: None, + provider_metadata: None, + })), + ); + parts.push(finish); + let model = ScriptedModel::streaming(parts); + let result = stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap() + .consume() + .await + .unwrap(); + let step = result.final_step(); + let kinds: Vec<&str> = step + .content + .iter() + .map(|part| match part { + ContentPart::ToolCall(_) => "ToolCall", + ContentPart::ToolError(_) => "ToolError", + ContentPart::Text { .. } => "Text", + other => panic!("unexpected content {other:?}"), + }) + .collect(); + assert_eq!(kinds, ["ToolCall", "ToolError", "Text"]); + let ContentPart::ToolError(error) = &step.content[1] else { + unreachable!() + }; + assert_eq!(error.tool_call_id, "call-1"); + assert_eq!(error.dynamic, Some(true)); + let [ + ModelMessage::Assistant { .. }, + ModelMessage::Tool { content, .. }, + ] = &step.response.messages[..] + else { + panic!("{:#?}", step.response.messages) + }; + let [crate::content::ToolContentPart::ToolResult(result)] = &content[..] else { + panic!("{content:#?}") + }; + assert_eq!(result.tool_call_id, "call-1"); + assert!(matches!( + result.output, + crate::content::ToolResultOutput::ErrorText { .. } + )); +} + +/// The `type` of each AI SDK `TextStreamPart`, in the union's order +/// (stream-text-result.ts:581-607). +const UPSTREAM_STREAM_PART_TYPES: [&str; 26] = [ + "text-start", + "text-end", + "text-delta", + "reasoning-start", + "reasoning-end", + "reasoning-delta", + "custom", + "tool-input-start", + "tool-input-end", + "tool-input-delta", + "source", + "file", + "reasoning-file", + "tool-call", + "tool-result", + "tool-error", + "tool-output-denied", + "tool-approval-request", + "tool-approval-response", + "start-step", + "finish-step", + "start", + "finish", + "abort", + "error", + "raw", +]; + +/// `TextStream` → `text-stream`. +fn kebab(name: &str) -> String { + let mut out = String::new(); + for (index, c) in name.chars().enumerate() { + if c.is_ascii_uppercase() { + if index > 0 { + out.push('-'); + } + out.push(c.to_ascii_lowercase()); + } else { + out.push(c); + } + } + out +} + +/// NP2-5: `TextStreamPart` has exactly the AI SDK's 26 part types, in the +/// union's order (stream-text-result.ts:581-607). The variant tags are read +/// from the exported TypeScript type. +#[test] +fn the_stream_part_set_is_the_ai_sdks() { + let ts = ::inline(&ts_rs::Config::default()); + let tags: Vec = ts + .match_indices("{ \"") + .filter_map(|(index, _)| { + let rest = &ts[index + 3..]; + let tag = &rest[..rest.find('"')?]; + tag.starts_with(|c: char| c.is_ascii_uppercase()) + .then(|| kebab(tag)) + }) + .collect(); + assert_eq!(tags, UPSTREAM_STREAM_PART_TYPES); +} + +/// NP2-5 stream state: text and reasoning segments keep the position of +/// their start part while deltas update them by id, interleaved or not +/// (stream-text.ts:1540-1600); the response message replays them in that +/// order. +#[tokio::test] +async fn interleaved_segments_keep_their_start_positions() { + let mut parts = text_stream(&[], FinishReasonUnified::Stop); + let finish = parts.pop().unwrap(); + let _text_end = parts.pop().unwrap(); + let text = |id: &str, delta: &str| { + Ok(StreamPart::TextDelta { + id: id.into(), + delta: delta.into(), + provider_metadata: None, + }) + }; + parts.extend([ + Ok(StreamPart::TextStart { + id: "t2".into(), + provider_metadata: None, + }), + text("t2", "B"), + text("t1", "A"), + Ok(StreamPart::ReasoningStart { + id: "r1".into(), + provider_metadata: None, + }), + Ok(StreamPart::ReasoningDelta { + id: "r1".into(), + delta: "think".into(), + provider_metadata: None, + }), + text("t1", "a"), + Ok(StreamPart::TextEnd { + id: "t2".into(), + provider_metadata: None, + }), + Ok(StreamPart::TextEnd { + id: "t1".into(), + provider_metadata: None, + }), + Ok(StreamPart::ReasoningEnd { + id: "r1".into(), + provider_metadata: None, + }), + finish, + ]); + let model = ScriptedModel::streaming(parts); + let result = stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap() + .consume() + .await + .unwrap(); + let content: Vec<(&str, &str)> = result + .final_step() + .content + .iter() + .map(|part| match part { + ContentPart::Text { text, .. } => ("text", text.as_str()), + ContentPart::Reasoning(reasoning) => ("reasoning", reasoning.text.as_str()), + other => panic!("unexpected content {other:?}"), + }) + .collect(); + assert_eq!( + content, + [("text", "Aa"), ("text", "B"), ("reasoning", "think")] + ); + let messages = serde_json::to_value(result.response_messages()).unwrap(); + let replayed: Vec = messages[0]["content"] + .as_array() + .unwrap() + .iter() + .map(|part| { + part.as_object().unwrap()["text"] + .as_str() + .unwrap() + .to_owned() + }) + .collect(); + assert_eq!(replayed, ["Aa", "B", "think"]); +} + +/// NP2-5 stream state: a preliminary tool result reaches the stream but not +/// the step content (stream-text.ts:1648-1650; alignment doc rule 6). +#[tokio::test] +async fn a_preliminary_tool_result_stays_out_of_the_step_content() { + let mut parts = text_stream(&["hi"], FinishReasonUnified::Stop); + let finish = parts.pop().unwrap(); + let mut preliminary = provider_result("srv-1", "image_generation", json!("partial"), None); + preliminary.preliminary = Some(true); + parts.extend([ + Ok(provider_call("srv-1", "image_generation", "{}")), + Ok(StreamPart::ToolResult(preliminary)), + Ok(StreamPart::ToolResult(provider_result( + "srv-1", + "image_generation", + json!("final"), + None, + ))), + finish, + ]); + let model = ScriptedModel::streaming(parts.clone()); + let streamed = collect( + stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(), + ) + .await; + let outputs: Vec<(Value, Option)> = streamed + .iter() + .filter_map(|part| match part { + TextStreamPart::ToolResult(result) => Some((result.output.clone(), result.preliminary)), + _ => None, + }) + .collect(); + assert_eq!( + outputs, + [(json!("partial"), Some(true)), (json!("final"), None)] + ); + + let model = ScriptedModel::streaming(parts); + let result = stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap() + .consume() + .await + .unwrap(); + let kept: Vec<&Value> = result + .final_step() + .content + .iter() + .filter_map(|part| match part { + ContentPart::ToolResult(result) => Some(&result.output), + _ => None, + }) + .collect(); + assert_eq!(kept, [&json!("final")]); +} + +/// NP2-5 message replay: the same model output gives the same response +/// messages whether it was generated or streamed (generate-text.ts:1437-1455 +/// and stream-text.ts build both with `toResponseMessages`). +#[tokio::test] +async fn generated_and_streamed_output_replay_the_same_messages() { + let generated_model = ScriptedModel::generating(generated( + vec![ + GenerateContent::Reasoning(crate::result::ReasoningOutput { + text: "think".into(), + provider_metadata: None, + }), + text_content("hi"), + GenerateContent::ToolCall(RawToolCall { + tool_call_id: "srv-1".into(), + tool_name: "web_search".into(), + input: r#"{"q":"rust"}"#.into(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: None, + }), + GenerateContent::ToolResult(provider_result( + "srv-1", + "web_search", + json!([{ "url": "https://example.com" }]), + None, + )), + ], + FinishReasonUnified::Stop, + )); + let generated = generate_text(&generated_model, "hi", GenerateTextOptions::default()) + .await + .unwrap(); + + let mut parts = text_stream(&["hi"], FinishReasonUnified::Stop); + let finish = parts.pop().unwrap(); + parts.splice( + 2..2, + [ + Ok(StreamPart::ReasoningStart { + id: "r1".into(), + provider_metadata: None, + }), + Ok(StreamPart::ReasoningDelta { + id: "r1".into(), + delta: "think".into(), + provider_metadata: None, + }), + Ok(StreamPart::ReasoningEnd { + id: "r1".into(), + provider_metadata: None, + }), + ], + ); + parts.extend([ + Ok(provider_call("srv-1", "web_search", r#"{"q":"rust"}"#)), + Ok(StreamPart::ToolResult(provider_result( + "srv-1", + "web_search", + json!([{ "url": "https://example.com" }]), + None, + ))), + finish, + ]); + let streamed_model = ScriptedModel::streaming(parts); + let streamed = stream_text(&streamed_model, "hi", GenerateTextOptions::default()) + .await + .unwrap() + .consume() + .await + .unwrap(); + + assert_eq!( + serde_json::to_value(streamed.response_messages()).unwrap(), + serde_json::to_value(generated.response_messages()).unwrap() + ); +} + +/// stream-text.ts:1725-1757: an abort after a step finished keeps the step; +/// the result resolves with the recorded step and, without a `finish` part, +/// the null total usage. Only a stream that recorded no step fails with the +/// abort (:1729-1740; `an_aborted_stream_ends_with_an_abort_part`). +#[tokio::test] +async fn an_abort_after_a_finished_step_keeps_the_step() { + let model = ScriptedModel::streaming(text_stream(&["Hel", "lo"], FinishReasonUnified::Stop)); + let mut result = stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(); + let mut parts: Vec = + std::mem::replace(&mut result.stream, Box::pin(futures::stream::empty())) + .map(|part| part.expect("no transport error")) + .collect() + .await; + // stream_text itself never lets an abort in here: Finish directly + // follows FinishStep. + assert!(matches!(parts.pop(), Some(TextStreamPart::Finish { .. }))); + assert!(matches!( + parts.last(), + Some(TextStreamPart::FinishStep { .. }) + )); + parts.push(TextStreamPart::Abort { + reason: Some("request aborted".into()), + }); + result.stream = Box::pin(futures::stream::iter(parts.into_iter().map(Ok))); + + let result = result.consume().await.expect("the finished step is kept"); + assert_eq!(result.steps().len(), 1); + assert_eq!(result.text(), "Hello"); + assert_eq!(result.usage(), &LanguageModelUsage::default()); +} + +/// stream-language-model-call.ts:432, :705-707, :764: both calls with a +/// repeated id are forwarded, and a provider result for that id carries the +/// input of the last one. +#[tokio::test] +async fn a_repeated_tool_call_id_resolves_to_the_last_call_in_the_stream() { + let mut parts = text_stream(&["hi"], FinishReasonUnified::ToolCalls); + let finish = parts.pop().unwrap(); + parts.extend([ + Ok(provider_call("srv-1", "web_search", r#"{"q":"first"}"#)), + Ok(provider_call("srv-1", "web_search", r#"{"q":"second"}"#)), + Ok(StreamPart::ToolResult(provider_result( + "srv-1", + "web_search", + json!(["hit"]), + None, + ))), + ]); + parts.push(finish); + let model = ScriptedModel::streaming(parts); + let parts = collect( + stream_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(), + ) + .await; + + let inputs: Vec<&Value> = parts + .iter() + .filter_map(|part| match part { + TextStreamPart::ToolCall(call) => Some(&call.input), + _ => None, + }) + .collect(); + assert_eq!( + inputs, + [&json!({ "q": "first" }), &json!({ "q": "second" })] + ); + + let result = parts + .iter() + .find(|part| matches!(part, TextStreamPart::ToolResult(_))) + .expect("the provider result"); + assert_eq!( + serde_json::to_value(result).unwrap()["ToolResult"]["input"], + json!({ "q": "second" }) + ); +} + +/// convert-language-model-content.ts:74-84: the generated content maps each +/// `tool-call` part to the first parsed call with its id, so both parts of a +/// repeated id show the first call. +#[tokio::test] +async fn a_repeated_tool_call_id_maps_to_the_first_call_when_generated() { + let call = |input: &str| { + GenerateContent::ToolCall(RawToolCall { + tool_call_id: "srv-1".into(), + tool_name: "web_search".into(), + input: input.into(), + provider_executed: Some(true), + dynamic: Some(true), + provider_metadata: None, + }) + }; + let model = ScriptedModel::generating(generated( + vec![call(r#"{"q":"first"}"#), call(r#"{"q":"second"}"#)], + FinishReasonUnified::ToolCalls, + )); + let result = generate_text(&model, "hi", GenerateTextOptions::default()) + .await + .unwrap(); + let inputs: Vec<&Value> = result + .final_step() + .content + .iter() + .filter_map(|part| match part { + ContentPart::ToolCall(call) => Some(&call.input), + _ => None, + }) + .collect(); + assert_eq!(inputs, [&json!({ "q": "first" }), &json!({ "q": "first" })]); +} diff --git a/aimux-core/src/image_model.rs b/aimux-core/src/image_model.rs index 13cd7628d..e41ab3fe8 100644 --- a/aimux-core/src/image_model.rs +++ b/aimux-core/src/image_model.rs @@ -112,7 +112,10 @@ pub struct ImageCallOptions { #[ts(skip)] pub abort_signal: Option, - /// Per-call retry override. `None` uses the model default. + /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` disables + /// retrying. `None` uses 2, the AI SDK default, unless the model's + /// `retry_config()` sets another count (model-level retry settings are + /// removed in the target design, docs/aisdk-architecture-alignment.md D14). pub max_retries: Option, /// Per-call operation timeout. diff --git a/aimux-core/src/lib.rs b/aimux-core/src/lib.rs index b5ef7ca3f..60dc4d30d 100644 --- a/aimux-core/src/lib.rs +++ b/aimux-core/src/lib.rs @@ -39,7 +39,9 @@ pub mod moa; pub mod model_catalogue; pub mod model_id; pub mod openai_output; +pub mod operation; pub mod options; +pub mod output; pub mod parse_tool_call; pub mod provider; pub mod recording; @@ -53,6 +55,7 @@ pub mod search_model; pub mod session; pub mod shared; pub mod speech_model; +pub mod step_result; pub mod stream_part; mod timeout; pub mod tool; @@ -81,9 +84,11 @@ pub mod prelude { }; pub use crate::error::{AiMuxError, ApiCallError}; pub use crate::files_model::{Files, UploadFileCallOptions, UploadFileResult}; + #[allow(deprecated)] + pub use crate::generate::generate_object; pub use crate::generate::{ - GenerateObjectResult, GenerateTextOptions, GenerateTextResult, StreamTextResult, - generate_object, generate_text, generate_text_as_openai, + GenerateObjectResponse, GenerateObjectResult, GenerateTextOptions, GenerateTextResult, + StreamTextResult, generate_text, generate_text_as_openai, generate_text_result_to_chat_completion, stream_text, stream_text_as_openai, }; pub use crate::image_model::{ImageCallOptions, ImageModel, ImageResult, generate_image}; @@ -105,6 +110,7 @@ pub mod prelude { pub use crate::options::{ CallOptions, LanguageModelV4ToolChoice, ResponseFormat, ToolChoice, prepare_tool_choice, }; + pub use crate::output::Output; pub use crate::parse_tool_call::{ ToolCallRepair, ToolCallRepairContext, ToolCallRepairReply, apply_tool_call_repair, apply_tool_call_repair_to_result, tool_call_repair_context, tool_call_repair_inputs, @@ -114,9 +120,9 @@ pub mod prelude { RerankingCallOptions, RerankingModel, RerankingResult, rerank, }; pub use crate::result::{ - ContentPart, GenerateResponse, GenerateResult, GeneratedFile, RawToolApprovalRequest, - ReasoningOutput, ReasoningPart, Source, StreamResponse, StreamResult, - ToolApprovalRequestOutput, ToolApprovalResponseOutput, TypedToolError, TypedToolResult, + ContentPart, GenerateResult, GeneratedFile, RawToolApprovalRequest, ReasoningOutput, + ReasoningPart, Source, StreamResult, ToolApprovalRequestOutput, ToolApprovalResponseOutput, + ToolOutputDenied, TypedToolError, TypedToolResult, }; pub use crate::router::{ FallbackPolicy, Router, RouterConfig, RouterModel, RuleRouter, WeightedRouter, @@ -134,6 +140,10 @@ pub mod prelude { provider_namespace, }; pub use crate::speech_model::{SpeechCallOptions, SpeechModel, SpeechResult, generate_speech}; + pub use crate::step_result::{ + OutputChunkTimingStats, StepModel, StepResponse, StepResponseMetadata, StepResult, + StepResultPerformance, + }; pub use crate::stream_part::{StreamPart, TextStreamPart}; pub use crate::tool::{ FunctionTool, FunctionToolInputExample, ProviderTool, RawToolCall, Tool, ToolCall, @@ -145,12 +155,12 @@ pub mod prelude { TranscriptionStreamResult, stream_transcribe, transcribe, }; pub use crate::types::{ - FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage, ReasoningEffort, - Usage, Warning, + FinishReason, FinishReasonUnified, InputTokenDetails, InputTokenUsage, LanguageModelUsage, + OutputTokenDetails, OutputTokenUsage, ReasoningEffort, Usage, Warning, }; pub use crate::video_model::{ - VideoCallOptions, VideoModel, VideoOperationStart, VideoOperationStatus, VideoPollConfig, - VideoPollOptions, VideoResult, generate_video, + VideoCallOptions, VideoModel, VideoOperationStart, VideoOperationStatus, VideoPollOptions, + VideoResult, generate_video, }; } diff --git a/aimux-core/src/moa.rs b/aimux-core/src/moa.rs index b298808b8..fc1c8929c 100644 --- a/aimux-core/src/moa.rs +++ b/aimux-core/src/moa.rs @@ -14,11 +14,12 @@ use futures::{StreamExt, future::join_all}; use serde::Deserialize; use crate::composite::{ - ChildModel, add_usage, build_aggregator_prompt, extract_text, step_options, + ChildModel, add_usage, build_aggregator_prompt, child_max_retries, extract_text, step_context, }; use crate::error::AiMuxError; use crate::language_model::LanguageModel; use crate::options::CallOptions; +use crate::recording::RecordingContext; use crate::result::{GenerateResult, StreamResult}; use crate::retry; use crate::stream_part::StreamPart; @@ -110,6 +111,15 @@ impl MoaModel { o } + /// The recording context of the aggregator step. + fn aggregator_context(&self) -> Option { + step_context(format!( + "moa.aggregator:{}/{}", + self.aggregator.provider(), + self.aggregator.model_id() + )) + } + /// Fan out references in parallel (non-streaming), accumulate usage, and /// collect `(model_id, text)` for the successful ones. Failures are handled /// per `fail_mode`. Returns the reference texts, accumulated usage, and any @@ -124,24 +134,24 @@ impl MoaModel { } let ref_opts = self.reference_options(options); let results = join_all(self.references.iter().enumerate().map(|(i, m)| { - let child_opts = step_options( - &ref_opts, - format!("moa.ref[{i}]:{}/{}", m.provider(), m.model_id()), + let context = step_context(format!("moa.ref[{i}]:{}/{}", m.provider(), m.model_id())); + let ref_opts = &ref_opts; + let retries = retry::prepare_retries( + child_max_retries(), + m.retry_config(), + ref_opts.abort_signal.clone(), ); async move { - let retries = retry::prepare_retries( - child_opts.max_retries, - m.retry_config(), - child_opts.abort_signal.clone(), - ); - retries - .retry(|| { - if let Some(context) = &child_opts.recording_context { + RecordingContext::scope( + context.clone(), + retries.retry(|| { + if let Some(context) = &context { let _ = context.start_attempt(); } - m.do_generate(&child_opts) - }) - .await + m.do_generate(ref_opts) + }), + ) + .await } })) .await; @@ -211,26 +221,21 @@ impl LanguageModel for MoaModel { // 3. Run the aggregator. let retries = retry::prepare_retries( - agg_opts.max_retries, + child_max_retries(), self.aggregator.retry_config(), agg_opts.abort_signal.clone(), ); - let agg_opts = step_options( - &agg_opts, - format!( - "moa.aggregator:{}/{}", - self.aggregator.provider(), - self.aggregator.model_id() - ), - ); - let mut agg = retries - .retry(|| { - if let Some(context) = &agg_opts.recording_context { + let context = self.aggregator_context(); + let mut agg = RecordingContext::scope( + context.clone(), + retries.retry(|| { + if let Some(context) = &context { let _ = context.start_attempt(); } self.aggregator.do_generate(&agg_opts) - }) - .await?; + }), + ) + .await?; // 4. Fold reference usage + drop warnings into the aggregator result. agg.usage = add_usage(agg.usage, &ref_usage); @@ -256,27 +261,22 @@ impl LanguageModel for MoaModel { // usage onto its Finish. We swallow the aggregator's StreamStart // (we've already emitted ours). let retries = retry::prepare_retries( - agg_opts.max_retries, + child_max_retries(), self.aggregator.retry_config(), agg_opts.abort_signal.clone(), ); - let agg_opts = step_options( - &agg_opts, - format!( - "moa.aggregator:{}/{}", - self.aggregator.provider(), - self.aggregator.model_id() - ), - ); - let agg = retries - .retry(|| { - if let Some(context) = &agg_opts.recording_context { + let context = self.aggregator_context(); + let agg = RecordingContext::scope( + context.clone(), + retries.retry(|| { + if let Some(context) = &context { let _ = context.start_attempt(); } self.aggregator.do_stream(&agg_opts) - }) - .await?; - let mut agg_stream = agg.stream; + }), + ) + .await?; + let mut agg_stream = RecordingContext::scope_stream(context, agg.stream); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: drop_warnings }); @@ -305,7 +305,7 @@ impl LanguageModel for MoaModel { } }; - // RFC-0022 §3.4: return None for request_body/response_headers. The + // RFC-0022 §3.4: return None for request/response. The // aggregator's request body is a synthesized prompt (references // spliced in), not the user's original — exposing it would mislead // cache probing (RFC-0015 fingerprinting). @@ -323,7 +323,8 @@ mod tests { use crate::language_model_message::{LanguageModelMessage, TextPart, UserPart}; use crate::options::LanguageModelV4ToolChoice; use crate::result::GenerateContent; - use crate::result::{GenerateResponse, Source}; + use crate::result::Source; + use crate::shared::ResponseInfo; use crate::stream_part::StreamPart; use crate::tool::RawToolCall; use crate::types::{FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage}; @@ -371,7 +372,7 @@ mod tests { usage: self.usage.clone(), warnings: vec![], provider_metadata: None, - response: Some(GenerateResponse { + response: Some(ResponseInfo { model_id: Some(self.name.into()), ..Default::default() }), @@ -585,7 +586,7 @@ mod tests { .await .unwrap(); - assert_eq!(result.text, "final"); + assert_eq!(result.text(), "final"); assert_eq!(ref_a_calls.load(Ordering::SeqCst), 2); assert_eq!(ref_b_calls.load(Ordering::SeqCst), 1); assert_eq!(aggregator_calls.load(Ordering::SeqCst), 2); diff --git a/aimux-core/src/openai_output.rs b/aimux-core/src/openai_output.rs index 7835a5a67..4b632f339 100644 --- a/aimux-core/src/openai_output.rs +++ b/aimux-core/src/openai_output.rs @@ -1,6 +1,6 @@ //! OpenAI Chat Completions output format. //! -//! Converts aimux's internal types ([`GenerateResult`] / [`StreamPart`]) into +//! Converts aimux's internal types ([`GenerateResult`] / [`TextStreamPart`]) into //! standard OpenAI Chat Completions structures ([`ChatCompletion`] / //! [`ChatCompletionChunk`]). This lets any provider (OpenAI, Anthropic, Google, //! …) be consumed via the OpenAI wire format. @@ -9,7 +9,7 @@ //! //! ```text //! provider.do_generate() → GenerateResult → to_chat_completion() → ChatCompletion -//! provider.do_stream() → Stream → to_chat_completion_stream() → Stream +//! stream_text() → Stream → to_chat_completion_stream() → Stream //! ``` //! //! The conversion is a post-processing step — it does not modify the @@ -34,12 +34,16 @@ use serde_json::Value; use ts_rs::TS; use crate::error::AiMuxError; +use crate::generate::GenerateTextResult; use crate::parse_tool_call::original_tool_call_text; -use crate::result::{GenerateContent, GenerateResult, GeneratedFile, ReasoningOutput, Source}; +use crate::result::{ + ContentPart, GenerateContent, GenerateResult, GeneratedFile, ReasoningOutput, Source, + TypedToolError, TypedToolResult, +}; use crate::shared::GeneratedFileData; -use crate::stream_part::{StreamPart, TextStreamPart}; -use crate::tool::{ToolCall, ToolResult}; -use crate::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; +use crate::stream_part::TextStreamPart; +use crate::tool::ToolCall; +use crate::types::{FinishReasonUnified, LanguageModelUsage, as_language_model_usage}; // ───────────────────────────────────────────────────────────────────────────── // Non-streaming response types @@ -227,38 +231,166 @@ impl Default for OpenAiStreamOptions { /// precedence. #[must_use] pub fn to_chat_completion(result: &GenerateResult, model: &str) -> ChatCompletion { + let content = result.content.iter().map(|item| match item { + GenerateContent::Text { text, .. } => CompletionContent::Text(text), + GenerateContent::Reasoning(ReasoningOutput { text, .. }) => { + CompletionContent::Reasoning(text) + } + // The provider's raw argument text passes through verbatim; OpenAI's + // wire format requires a JSON object even when the model emitted no + // arguments at all. + GenerateContent::ToolCall(call) => CompletionContent::ToolCall { + id: &call.tool_call_id, + name: &call.tool_name, + arguments: if call.input.trim().is_empty() { + "{}".to_string() + } else { + call.input.clone() + }, + }, + GenerateContent::Source(source) => CompletionContent::Source(source), + GenerateContent::File(file) => CompletionContent::File(file), + GenerateContent::ToolResult(result) => CompletionContent::ToolResult { + tool_name: &result.tool_name, + output: &result.result, + is_error: result.is_error == Some(true), + }, + GenerateContent::Custom { .. } + | GenerateContent::ReasoningFile(_) + | GenerateContent::ToolApprovalRequest(_) => CompletionContent::None, + }); + let response = result.response.as_ref(); + build_chat_completion( + content, + finish_reason_to_openai( + result.finish_reason.unified, + result.finish_reason.raw.as_deref(), + ), + result.provider_metadata.as_ref(), + response.and_then(|response| response.id.clone()), + response + .and_then(|response| response.model_id.clone()) + .unwrap_or_else(|| model.to_string()), + usage_to_openai(&as_language_model_usage(&result.usage)), + ) +} + +/// Convert a [`GenerateTextResult`] into an OpenAI [`ChatCompletion`]: the +/// content of every step, the last step's finish reason, id, model and +/// logprobs, and the usage summed over the steps. Tool calls carry the +/// Core-parsed (and possibly repaired) arguments. +#[must_use] +pub(crate) fn text_result_to_chat_completion( + result: &GenerateTextResult, + model: &str, +) -> ChatCompletion { + let content = result.content().into_iter().map(|item| match item { + ContentPart::Text { text, .. } => CompletionContent::Text(text), + ContentPart::Reasoning(ReasoningOutput { text, .. }) => CompletionContent::Reasoning(text), + ContentPart::ToolCall(call) => CompletionContent::ToolCall { + id: &call.tool_call_id, + name: &call.tool_name, + arguments: parsed_tool_call_arguments(&call.input, call.invalid, call.error.as_ref()), + }, + ContentPart::Source(source) => CompletionContent::Source(source), + ContentPart::File { file, .. } => CompletionContent::File(file), + ContentPart::ToolResult(result) => CompletionContent::ToolResult { + tool_name: &result.tool_name, + output: &result.output, + is_error: false, + }, + // A provider's tool error renders as the generated path renders it; + // the error of a call that could not be parsed has no Chat + // Completions form. + ContentPart::ToolError(error) if error.provider_executed == Some(true) => { + CompletionContent::ToolResult { + tool_name: &error.tool_name, + output: &error.error, + is_error: true, + } + } + ContentPart::Custom { .. } + | ContentPart::ReasoningFile(_) + | ContentPart::ToolError(_) + | ContentPart::ToolApprovalRequest(_) + | ContentPart::ToolApprovalResponse(_) => CompletionContent::None, + }); + let final_step = result.final_step(); + let model = if final_step.response.model_id.is_empty() { + model.to_string() + } else { + final_step.response.model_id.clone() + }; + build_chat_completion( + content, + finish_reason_to_openai( + final_step.finish_reason, + final_step.raw_finish_reason.as_deref(), + ), + final_step.provider_metadata.as_ref(), + Some(final_step.response.id.clone()).filter(|id| !id.is_empty()), + model, + usage_to_openai(result.usage()), + ) +} + +/// A content item as the non-streaming converter maps it. +enum CompletionContent<'a> { + Text(&'a str), + Reasoning(&'a str), + ToolCall { + id: &'a str, + name: &'a str, + arguments: String, + }, + Source(&'a Source), + File(&'a GeneratedFile), + /// A provider-executed tool result or error. + ToolResult { + tool_name: &'a str, + output: &'a serde_json::Value, + is_error: bool, + }, + /// No Chat Completions equivalent. + None, +} + +fn build_chat_completion<'a>( + content: impl Iterator>, + finish_reason: Option, + provider_metadata: Option<&crate::types::ProviderMetadata>, + id: Option, + model: String, + usage: ChatCompletionUsage, +) -> ChatCompletion { let mut content_text = String::new(); let mut reasoning_text = String::new(); let mut tool_calls = Vec::new(); let mut annotations = Vec::new(); - for item in &result.content { + for item in content { match item { - GenerateContent::Text { text, .. } => { + CompletionContent::Text(text) => { content_text.push_str(text); } - GenerateContent::Reasoning(ReasoningOutput { text, .. }) => { + CompletionContent::Reasoning(text) => { reasoning_text.push_str(text); } - GenerateContent::ToolCall(call) => { - // The provider's raw argument text passes through verbatim; - // OpenAI's wire format requires a JSON object even when the - // model emitted no arguments at all. - let arguments = if call.input.trim().is_empty() { - "{}".to_string() - } else { - call.input.clone() - }; + CompletionContent::ToolCall { + id, + name, + arguments, + } => { tool_calls.push(ChatCompletionToolCall { - id: call.tool_call_id.clone(), + id: id.to_string(), tool_type: "function".to_string(), function: ChatCompletionFunction { - name: call.tool_name.clone(), + name: name.to_string(), arguments, }, }); } - GenerateContent::Source(Source::Url { url, title, .. }) => { + CompletionContent::Source(Source::Url { url, title, .. }) => { // Map to OpenAI url_citation annotation. let mut ann = serde_json::json!({ "type": "url_citation", @@ -276,7 +408,7 @@ pub fn to_chat_completion(result: &GenerateResult, model: &str) -> ChatCompletio } content_text.push_str(url); } - GenerateContent::File(GeneratedFile { + CompletionContent::File(GeneratedFile { data, media_type, .. }) => { if !content_text.is_empty() { @@ -284,27 +416,23 @@ pub fn to_chat_completion(result: &GenerateResult, model: &str) -> ChatCompletio } content_text.push_str(&file_data_to_text(data, media_type)); } - GenerateContent::Custom { .. } - | GenerateContent::Source(Source::Document { .. }) - | GenerateContent::ReasoningFile(_) - | GenerateContent::ToolApprovalRequest(_) => {} - GenerateContent::ToolResult(ToolResult { + CompletionContent::Source(Source::Document { .. }) | CompletionContent::None => {} + CompletionContent::ToolResult { tool_name, - result, + output, is_error, - .. - }) => { + } => { // Degraded mapping: provider-executed tool result as text. if !content_text.is_empty() { content_text.push('\n'); } - let prefix = if is_error.unwrap_or(false) { + let prefix = if is_error { format!("[tool error: {tool_name}] ") } else { format!("[tool result: {tool_name}] ") }; content_text.push_str(&prefix); - content_text.push_str(&result.to_string()); + content_text.push_str(&output.to_string()); } } } @@ -336,34 +464,16 @@ pub fn to_chat_completion(result: &GenerateResult, model: &str) -> ChatCompletio }, }; - // finish_reason - let finish_reason = finish_reason_to_openai(&result.finish_reason); - // logprobs from provider_metadata - let logprobs = result - .provider_metadata - .as_ref() + let logprobs = provider_metadata .and_then(|pm| pm.get("openai")) .and_then(|o| o.get("logprobs")) .cloned(); - // id / model / created - let id = result - .response - .as_ref() - .and_then(|response| response.id.clone()) - .unwrap_or_else(|| format!("chatcmpl-{}", random_id())); - let model = result - .response - .as_ref() - .and_then(|response| response.model_id.clone()) - .unwrap_or_else(|| model.to_string()); - let created = now_unix(); - ChatCompletion { - id, + id: id.unwrap_or_else(|| format!("chatcmpl-{}", random_id())), object: "chat.completion".to_string(), - created, + created: now_unix(), model, choices: vec![ChatCompletionChoice { index: 0, @@ -371,13 +481,13 @@ pub fn to_chat_completion(result: &GenerateResult, model: &str) -> ChatCompletio finish_reason, logprobs, }], - usage: usage_to_openai(&result.usage), + usage, system_fingerprint: None, } } // ───────────────────────────────────────────────────────────────────────────── -// Streaming conversion: Stream → Stream +// Streaming conversion: Stream → Stream // ───────────────────────────────────────────────────────────────────────────── /// A streaming OpenAI Chat Completions result. @@ -394,7 +504,7 @@ impl std::fmt::Debug for ChatCompletionStream { } } -/// Convert a `StreamPart` stream into a `ChatCompletionChunk` stream. +/// Convert a `TextStreamPart` stream into a `ChatCompletionChunk` stream. /// /// Uses a stateful converter that maintains tool-call indices and accumulates /// the final usage / finish_reason. The output follows OpenAI SSE conventions: @@ -445,7 +555,7 @@ pub fn to_chat_completion_stream( } // After Finish, stop (but we already emitted the final chunk). - if matches!(part, StreamPart::Finish { .. }) { + if matches!(part, TextStreamPart::Finish { .. }) { return; } } @@ -473,8 +583,8 @@ struct StreamState { next_tool_index: u32, /// Whether each tool_call_id has had its opening chunk emitted. tool_call_opened: std::collections::HashSet, - final_usage: Option, - final_finish_reason: Option, + final_usage: Option, + final_finish_reason: Option<(FinishReasonUnified, Option)>, finish_emitted: bool, } @@ -536,7 +646,7 @@ impl StreamState { Some(chunk) } - /// Process a single StreamPart, returning zero or more output chunks. + /// Process a single TextStreamPart, returning zero or more output chunks. fn process_part( &mut self, part: &TextStreamPart, @@ -546,27 +656,26 @@ impl StreamState { let mut chunks = Vec::new(); match part { - StreamPart::StreamStart { .. } => { + TextStreamPart::Start {} => { if let Some(c) = self.ensure_started() { chunks.push(c); } } - StreamPart::ResponseMetadata(ResponseMetadata { id, model_id, .. }) => { - if let Some(id) = id { - self.id = id.clone(); - } - if let Some(m) = model_id { - self.model = m.clone(); + TextStreamPart::FinishStep { response, .. } => { + // Response metadata arrives at step end; emitted chunks keep their identity. + if !self.started { + self.id.clone_from(&response.id); + self.model.clone_from(&response.model_id); } } - StreamPart::TextStart { .. } => { + TextStreamPart::TextStart { .. } => { if let Some(c) = self.ensure_started() { chunks.push(c); } } - StreamPart::TextDelta { delta, .. } => { + TextStreamPart::TextDelta { text: delta, .. } => { if let Some(c) = self.ensure_started() { chunks.push(c); } @@ -582,10 +691,10 @@ impl StreamState { }]; chunks.push(chunk); } - StreamPart::TextEnd { .. } => {} + TextStreamPart::TextEnd { .. } => {} - StreamPart::ReasoningStart { .. } => {} - StreamPart::ReasoningDelta { delta, .. } => { + TextStreamPart::ReasoningStart { .. } => {} + TextStreamPart::ReasoningDelta { text: delta, .. } => { if include_reasoning { if let Some(c) = self.ensure_started() { chunks.push(c); @@ -603,9 +712,9 @@ impl StreamState { chunks.push(chunk); } } - StreamPart::ReasoningEnd { .. } => {} + TextStreamPart::ReasoningEnd { .. } => {} - StreamPart::ToolInputStart { id, tool_name, .. } => { + TextStreamPart::ToolInputStart { id, tool_name, .. } => { if let Some(c) = self.ensure_started() { chunks.push(c); } @@ -649,7 +758,7 @@ impl StreamState { }]; chunks.push(chunk); } - StreamPart::ToolInputDelta { id, delta, .. } => { + TextStreamPart::ToolInputDelta { id, delta, .. } => { // Ensure started (shouldn't happen without Start, but be safe). if let Some(c) = self.ensure_started() { chunks.push(c); @@ -697,9 +806,9 @@ impl StreamState { }]; chunks.push(chunk); } - StreamPart::ToolInputEnd { .. } => {} + TextStreamPart::ToolInputEnd { .. } => {} - StreamPart::ToolCall(ToolCall { + TextStreamPart::ToolCall(ToolCall { tool_call_id, tool_name, input, @@ -788,14 +897,39 @@ impl StreamState { } } - StreamPart::ToolResult(ToolResult { - tool_name, result, .. + TextStreamPart::ToolResult(TypedToolResult { + tool_name, output, .. }) => { // Degraded: provider-executed tool result as content. if let Some(c) = self.ensure_started() { chunks.push(c); } - let text = format!("[tool result: {tool_name}] {result}"); + let text = format!("[tool result: {tool_name}] {output}"); + let mut chunk = self.base_chunk(); + chunk.choices = vec![ChatCompletionChunkChoice { + index: 0, + delta: ChatCompletionDelta { + content: Some(text), + ..Default::default() + }, + finish_reason: None, + logprobs: None, + }]; + chunks.push(chunk); + } + + TextStreamPart::ToolError(TypedToolError { + tool_name, + error, + provider_executed: Some(true), + .. + }) => { + // Degraded: provider-executed tool error as content, as the + // generated path renders it. + if let Some(c) = self.ensure_started() { + chunks.push(c); + } + let text = format!("[tool error: {tool_name}] {error}"); let mut chunk = self.base_chunk(); chunk.choices = vec![ChatCompletionChunkChoice { index: 0, @@ -809,7 +943,7 @@ impl StreamState { chunks.push(chunk); } - StreamPart::File(GeneratedFile { + TextStreamPart::File(GeneratedFile { data, media_type, .. }) => { if let Some(c) = self.ensure_started() { @@ -829,7 +963,7 @@ impl StreamState { chunks.push(chunk); } - StreamPart::Source(Source::Url { url, .. }) => { + TextStreamPart::Source(Source::Url { url, .. }) => { if let Some(c) = self.ensure_started() { chunks.push(c); } @@ -849,30 +983,35 @@ impl StreamState { } } - StreamPart::Finish { + TextStreamPart::Finish { finish_reason, - usage, - .. + raw_finish_reason, + total_usage, } => { - self.final_usage = Some(usage.clone()); - self.final_finish_reason = Some(finish_reason.clone()); + self.final_usage = Some(total_usage.clone()); + self.final_finish_reason = Some((*finish_reason, raw_finish_reason.clone())); // Emit the final chunk here. if let Some(c) = self.final_chunk_impl(include_usage) { chunks.push(c); } } - StreamPart::Error { error } => { + TextStreamPart::Error { error } => { if let Some(chunk) = self.error_chunk(error) { chunks.push(chunk); } } - StreamPart::Raw { .. } - | StreamPart::Source(Source::Document { .. }) - | StreamPart::Custom { .. } - | StreamPart::ReasoningFile(_) - | StreamPart::ToolApprovalRequest(_) => { /* no OpenAI chat equivalent */ } + TextStreamPart::Raw { .. } + | TextStreamPart::StartStep { .. } + | TextStreamPart::Abort { .. } + | TextStreamPart::Source(Source::Document { .. }) + | TextStreamPart::Custom { .. } + | TextStreamPart::ReasoningFile(_) + | TextStreamPart::ToolError(_) + | TextStreamPart::ToolOutputDenied(_) + | TextStreamPart::ToolApprovalRequest(_) + | TextStreamPart::ToolApprovalResponse(_) => { /* no OpenAI chat equivalent */ } } chunks @@ -917,7 +1056,7 @@ impl StreamState { let finish_reason_str = self .final_finish_reason .as_ref() - .map(finish_reason_to_openai) + .map(|(unified, raw)| finish_reason_to_openai(*unified, raw.as_deref())) .unwrap_or(Some("stop".to_string())); let usage = if include_usage { @@ -956,29 +1095,29 @@ pub const DONE_FRAME: &str = "data: [DONE]\n\n"; // Internal helpers // ───────────────────────────────────────────────────────────────────────────── -/// Map an aimux [`FinishReason`] to an OpenAI `finish_reason` string. -fn finish_reason_to_openai(reason: &FinishReason) -> Option { - match reason.unified { +/// Map a finish reason to an OpenAI `finish_reason` string. +fn finish_reason_to_openai(unified: FinishReasonUnified, raw: Option<&str>) -> Option { + match unified { FinishReasonUnified::Stop => Some("stop".to_string()), FinishReasonUnified::Length => Some("length".to_string()), FinishReasonUnified::ContentFilter => Some("content_filter".to_string()), FinishReasonUnified::ToolCalls => Some("tool_calls".to_string()), FinishReasonUnified::Error => Some("stop".to_string()), - FinishReasonUnified::Other => reason.raw.clone().or(Some("stop".to_string())), + FinishReasonUnified::Other => Some(raw.unwrap_or("stop").to_string()), } } -/// Convert aimux [`Usage`] to OpenAI [`ChatCompletionUsage`]. +/// Convert a [`LanguageModelUsage`] to OpenAI [`ChatCompletionUsage`]. /// /// This is the inverse of `convert_usage` in /// `aimux-providers/src/openai/model.rs`. -fn usage_to_openai(usage: &Usage) -> ChatCompletionUsage { - let prompt_tokens = usage.input_tokens.total.unwrap_or(0); - let completion_tokens = usage.output_tokens.total.unwrap_or(0); +fn usage_to_openai(usage: &LanguageModelUsage) -> ChatCompletionUsage { + let prompt_tokens = usage.input_tokens.unwrap_or(0); + let completion_tokens = usage.output_tokens.unwrap_or(0); - let cached_tokens = usage.input_tokens.cache_read.unwrap_or(0); - let cache_write = usage.input_tokens.cache_write; - let reasoning_tokens = usage.output_tokens.reasoning; + let cached_tokens = usage.input_token_details.cache_read_tokens.unwrap_or(0); + let cache_write = usage.input_token_details.cache_write_tokens; + let reasoning_tokens = usage.output_token_details.reasoning_tokens; ChatCompletionUsage { prompt_tokens, @@ -1019,7 +1158,7 @@ fn now_unix() -> u64 { // Rejected repairs retain the original call. Recover its text from the // original error, never from the replacement's validation failure. -/// Render a Core-parsed `StreamPart::ToolCall`'s arguments as OpenAI-compatible +/// Render a Core-parsed `TextStreamPart::ToolCall`'s arguments as OpenAI-compatible /// wire text: the provider's raw argument text verbatim for an invalid call, /// compact JSON of the parsed value otherwise. /// @@ -1065,7 +1204,8 @@ fn random_id() -> String { mod tests { use super::*; use crate::result::GenerateResult; - use crate::result::{GenerateResponse, ReasoningOutput}; + use crate::result::ReasoningOutput; + use crate::shared::ResponseInfo; use crate::tool::RawToolCall; use crate::types::{ FinishReason, FinishReasonUnified, InputTokenUsage, OutputTokenUsage, Usage, @@ -1083,7 +1223,7 @@ mod tests { usage: Usage::default(), warnings: vec![], provider_metadata: None, - response: Some(GenerateResponse { + response: Some(ResponseInfo { id: Some("chatcmpl-test123".to_string()), timestamp: None, model_id: Some("gpt-4o".to_string()), @@ -1261,7 +1401,7 @@ mod tests { assert_eq!(parsed.invalid, Some(true), "{raw_input}"); let parts: Vec> = - vec![Ok(StreamPart::ToolCall(ToolCall { + vec![Ok(TextStreamPart::ToolCall(ToolCall { tool_call_id: parsed.tool_call_id, tool_name: parsed.tool_name, input: parsed.input, @@ -1386,8 +1526,10 @@ mod tests { (FinishReasonUnified::Error, "stop"), ]; for (unified, expected) in cases { - let fr = FinishReason { unified, raw: None }; - assert_eq!(finish_reason_to_openai(&fr).as_deref(), Some(expected)); + assert_eq!( + finish_reason_to_openai(unified, None).as_deref(), + Some(expected) + ); } } @@ -1405,32 +1547,29 @@ mod tests { #[tokio::test] async fn test_stream_text_only() { let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextStart { + Ok(TextStreamPart::Start {}), + Ok(TextStreamPart::TextStart { id: "0".to_string(), provider_metadata: None, }), - Ok(StreamPart::TextDelta { + Ok(TextStreamPart::TextDelta { id: "0".to_string(), - delta: "Hello ".to_string(), + text: "Hello ".to_string(), provider_metadata: None, }), - Ok(StreamPart::TextDelta { + Ok(TextStreamPart::TextDelta { id: "0".to_string(), - delta: "world".to_string(), + text: "world".to_string(), provider_metadata: None, }), - Ok(StreamPart::TextEnd { + Ok(TextStreamPart::TextEnd { id: "0".to_string(), provider_metadata: None, }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, + Ok(TextStreamPart::Finish { + finish_reason: FinishReasonUnified::Stop, + raw_finish_reason: None, + total_usage: LanguageModelUsage::default(), }), ]; @@ -1462,8 +1601,8 @@ mod tests { #[tokio::test] async fn test_stream_tool_calls() { let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::ToolInputStart { + Ok(TextStreamPart::Start {}), + Ok(TextStreamPart::ToolInputStart { id: "call_1".to_string(), tool_name: "get_weather".to_string(), provider_executed: None, @@ -1471,21 +1610,21 @@ mod tests { title: None, provider_metadata: None, }), - Ok(StreamPart::ToolInputDelta { + Ok(TextStreamPart::ToolInputDelta { id: "call_1".to_string(), delta: "{\"city\":".to_string(), provider_metadata: None, }), - Ok(StreamPart::ToolInputDelta { + Ok(TextStreamPart::ToolInputDelta { id: "call_1".to_string(), delta: "\"Tokyo\"}".to_string(), provider_metadata: None, }), - Ok(StreamPart::ToolInputEnd { + Ok(TextStreamPart::ToolInputEnd { id: "call_1".to_string(), provider_metadata: None, }), - Ok(StreamPart::ToolCall(ToolCall { + Ok(TextStreamPart::ToolCall(ToolCall { tool_call_id: "call_1".to_string(), tool_name: "get_weather".to_string(), input: json!({ "city": "Tokyo" }), @@ -1495,13 +1634,10 @@ mod tests { error: None, provider_metadata: None, })), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::ToolCalls, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, + Ok(TextStreamPart::Finish { + finish_reason: FinishReasonUnified::ToolCalls, + raw_finish_reason: None, + total_usage: LanguageModelUsage::default(), }), ]; @@ -1563,7 +1699,7 @@ mod tests { #[tokio::test] async fn test_stream_complete_call_backfills_arguments_absent_from_deltas() { let parts: Vec> = vec![ - Ok(StreamPart::ToolInputStart { + Ok(TextStreamPart::ToolInputStart { id: "call_1".to_string(), tool_name: "fetch".to_string(), provider_executed: Some(true), @@ -1571,11 +1707,11 @@ mod tests { title: None, provider_metadata: None, }), - Ok(StreamPart::ToolInputEnd { + Ok(TextStreamPart::ToolInputEnd { id: "call_1".to_string(), provider_metadata: None, }), - Ok(StreamPart::ToolCall(ToolCall { + Ok(TextStreamPart::ToolCall(ToolCall { tool_call_id: "call_1".to_string(), tool_name: "fetch".to_string(), input: json!({ "url": "https://example.com" }), @@ -1585,13 +1721,10 @@ mod tests { error: None, provider_metadata: None, })), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, + Ok(TextStreamPart::Finish { + finish_reason: FinishReasonUnified::Stop, + raw_finish_reason: None, + total_usage: LanguageModelUsage::default(), }), ]; @@ -1623,8 +1756,8 @@ mod tests { #[tokio::test] async fn test_stream_multiple_tool_calls_index() { let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::ToolInputStart { + Ok(TextStreamPart::Start {}), + Ok(TextStreamPart::ToolInputStart { id: "call_a".to_string(), tool_name: "tool_a".to_string(), provider_executed: None, @@ -1632,7 +1765,7 @@ mod tests { title: None, provider_metadata: None, }), - Ok(StreamPart::ToolInputStart { + Ok(TextStreamPart::ToolInputStart { id: "call_b".to_string(), tool_name: "tool_b".to_string(), provider_executed: None, @@ -1640,13 +1773,10 @@ mod tests { title: None, provider_metadata: None, }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::ToolCalls, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, + Ok(TextStreamPart::Finish { + finish_reason: FinishReasonUnified::ToolCalls, + raw_finish_reason: None, + total_usage: LanguageModelUsage::default(), }), ]; @@ -1680,32 +1810,29 @@ mod tests { #[tokio::test] async fn test_stream_reasoning() { let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::ReasoningStart { + Ok(TextStreamPart::Start {}), + Ok(TextStreamPart::ReasoningStart { id: "reasoning-0".to_string(), provider_metadata: None, }), - Ok(StreamPart::ReasoningDelta { + Ok(TextStreamPart::ReasoningDelta { id: "reasoning-0".to_string(), - delta: "Hmm...".to_string(), + text: "Hmm...".to_string(), provider_metadata: None, }), - Ok(StreamPart::ReasoningEnd { + Ok(TextStreamPart::ReasoningEnd { id: "reasoning-0".to_string(), provider_metadata: None, }), - Ok(StreamPart::TextDelta { + Ok(TextStreamPart::TextDelta { id: "0".to_string(), - delta: "Answer".to_string(), + text: "Answer".to_string(), provider_metadata: None, }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage::default(), - provider_metadata: None, + Ok(TextStreamPart::Finish { + finish_reason: FinishReasonUnified::Stop, + raw_finish_reason: None, + total_usage: LanguageModelUsage::default(), }), ]; @@ -1759,19 +1886,16 @@ mod tests { }; let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { + Ok(TextStreamPart::Start {}), + Ok(TextStreamPart::TextDelta { id: "0".to_string(), - delta: "Hi".to_string(), + text: "Hi".to_string(), provider_metadata: None, }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage, - provider_metadata: None, + Ok(TextStreamPart::Finish { + finish_reason: FinishReasonUnified::Stop, + raw_finish_reason: None, + total_usage: as_language_model_usage(&usage), }), ]; @@ -1795,21 +1919,18 @@ mod tests { #[tokio::test] async fn test_stream_no_usage_when_disabled() { let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: None, - }, - usage: Usage { + Ok(TextStreamPart::Start {}), + Ok(TextStreamPart::Finish { + finish_reason: FinishReasonUnified::Stop, + raw_finish_reason: None, + total_usage: as_language_model_usage(&Usage { input_tokens: InputTokenUsage { total: Some(42), ..Default::default() }, output_tokens: OutputTokenUsage::default(), raw: None, - }, - provider_metadata: None, + }), }), ]; @@ -1831,13 +1952,13 @@ mod tests { #[tokio::test] async fn test_stream_error() { let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { + Ok(TextStreamPart::Start {}), + Ok(TextStreamPart::TextDelta { id: "0".to_string(), - delta: "partial".to_string(), + text: "partial".to_string(), provider_metadata: None, }), - Ok(StreamPart::Error { + Ok(TextStreamPart::Error { error: AiMuxError::ApiCall(Box::new(crate::error::ApiCallError::new( "something went wrong", "https://example.test/v1", @@ -1880,20 +2001,17 @@ mod tests { #[tokio::test] async fn recoverable_frame_error_does_not_truncate_chat_completion_stream() { let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), + Ok(TextStreamPart::Start {}), Err(AiMuxError::JsonParse("malformed SSE data".into())), - Ok(StreamPart::TextDelta { + Ok(TextStreamPart::TextDelta { id: "0".to_string(), - delta: "after-error".to_string(), + text: "after-error".to_string(), provider_metadata: None, }), - Ok(StreamPart::Finish { - finish_reason: FinishReason { - unified: FinishReasonUnified::Stop, - raw: Some("stop".into()), - }, - usage: Usage::default(), - provider_metadata: None, + Ok(TextStreamPart::Finish { + finish_reason: FinishReasonUnified::Stop, + raw_finish_reason: Some("stop".into()), + total_usage: LanguageModelUsage::default(), }), ]; @@ -1926,10 +2044,10 @@ mod tests { async fn test_stream_completes_without_finish() { // Stream ends without Finish — should still emit a final chunk. let parts: Vec> = vec![ - Ok(StreamPart::StreamStart { warnings: vec![] }), - Ok(StreamPart::TextDelta { + Ok(TextStreamPart::Start {}), + Ok(TextStreamPart::TextDelta { id: "0".to_string(), - delta: "Hello".to_string(), + text: "Hello".to_string(), provider_metadata: None, }), ]; diff --git a/aimux-core/src/operation.rs b/aimux-core/src/operation.rs new file mode 100644 index 000000000..956b308f8 --- /dev/null +++ b/aimux-core/src/operation.rs @@ -0,0 +1,103 @@ +//! The settings of the user operation that a model call runs in. +//! +//! The AI SDK keeps `maxRetries` and `timeout` on `generateText` / +//! `streamText` (`RequestOptions`, ai@7.0.127 `prompt/request-options.ts:104-127`) +//! and reports them on the operation's start event +//! (`generate-text/generate-text-events.ts:54-62`). The model call options +//! carry neither (`LanguageModelV4CallOptions`, +//! @ai-sdk/provider@4.0.21 `language-model/v4/language-model-v4-call-options.ts:8-138`). +//! Aimux's session id (RFC-0024) is an operation setting too. +//! +//! `generate_text` and `stream_text` run their model calls inside +//! [`OperationSettings::scope`]. Code below the model interface that needs a +//! setting reads [`OperationSettings::current`]: the router and MoA take the +//! child retry count from it, the trace layer the session id, and the recorder +//! writes the settings into the input record (docs/aisdk-architecture-alignment.md +//! §3.5, §4.5). + +use std::future::Future; + +use serde::{Deserialize, Serialize}; + +use crate::options::TimeoutConfiguration; + +tokio::task_local! { + static CURRENT: OperationSettings; +} + +/// Operation settings that are not part of the model call. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct OperationSettings { + /// Maximum number of retries, the AI SDK's `maxRetries`; `Some(0)` + /// disables retrying. `None` uses the default. + #[serde(default)] + pub max_retries: Option, + /// Timeout configuration (total / step / first-chunk / chunk idle). + #[serde(default)] + pub timeout: Option, + /// Explicit session id (RFC-0024), for grouping consecutive calls. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session_id: Option, +} + +impl OperationSettings { + /// The settings of the operation the current task runs in, or `None` + /// outside an operation (a direct `do_generate` / `do_stream` call). + #[must_use] + pub fn current() -> Option { + CURRENT.try_with(Clone::clone).ok() + } + + /// Run `future` with these settings as the current operation settings. + /// + /// `generate_text` and `stream_text` call it themselves. A caller that + /// drives a model directly uses it to give composite children a retry + /// count or the trace layer a session id. A task spawned inside the + /// future does not inherit the settings. + pub async fn scope(self, future: F) -> F::Output { + CURRENT.scope(self, future).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn the_settings_are_current_only_inside_their_scope() { + assert_eq!(OperationSettings::current(), None); + let settings = OperationSettings { + max_retries: Some(0), + timeout: None, + session_id: Some("s-1".into()), + }; + let inside = settings + .clone() + .scope(async { OperationSettings::current() }) + .await; + assert_eq!(inside, Some(settings)); + assert_eq!(OperationSettings::current(), None); + } + + #[test] + fn the_settings_serialize_as_the_call_options_fields_did() { + // The recorded input keeps the keys it had when the three settings + // were `CallOptions` fields: `max_retries` and `timeout` always, + // `session_id` only when set. + let json = serde_json::to_value(OperationSettings::default()).unwrap(); + assert_eq!( + json, + serde_json::json!({ "max_retries": null, "timeout": null }) + ); + let back: OperationSettings = serde_json::from_value(serde_json::json!({ + "max_retries": 1, + "timeout": { "total_ms": 5, "step_ms": null, "first_chunk_ms": null, "chunk_ms": null }, + "session_id": "s", + "temperature": 0.5 + })) + .unwrap(); + assert_eq!(back.max_retries, Some(1)); + assert_eq!(back.timeout.unwrap().total_ms, Some(5)); + assert_eq!(back.session_id.as_deref(), Some("s")); + } +} diff --git a/aimux-core/src/options.rs b/aimux-core/src/options.rs index 295f38843..3553c7858 100644 --- a/aimux-core/src/options.rs +++ b/aimux-core/src/options.rs @@ -73,6 +73,10 @@ pub struct TimeoutConfiguration { /// This is the **provider-facing** options struct. Users interact with /// `GenerateTextOptions` (user-facing) which is converted to `CallOptions` /// by the `generate_text` / `stream_text` functions. +/// +/// The retry count, the timeout and the session id are settings of the +/// operation (`OperationSettings`), not call options, as in the AI SDK +/// (`language-model-v4-call-options.ts:8-138` has none of them). #[derive(Debug, Clone, Default, Serialize, Deserialize, TS)] #[cfg_attr(feature = "ts-export", ts(export))] pub struct CallOptions { @@ -129,23 +133,6 @@ pub struct CallOptions { /// `null` values delete the corresponding key. See RFC-0017. pub body_overrides: Option, - /// Per-call retry count override. `None` uses the provider's configured - /// Core operation retry. `Some(0)` disables retries. - pub max_retries: Option, - - /// Per-call timeout configuration (total / first-chunk / chunk idle). - /// `None` = no timeouts (provider defaults still apply at the HTTP layer). - pub timeout: Option, - - /// Session identifier, for grouping consecutive calls into a session - /// (observability, see RFC-0024). Explicit values take precedence; when - /// `None` and the optional session inferer is enabled, one may be inferred. - /// Orthogonal to RFC-0019 session-affinity headers: this field is for - /// local grouping, while session headers in `headers` are for upstream - /// routing — both may share an id value but travel different paths. - #[serde(skip_serializing_if = "Option::is_none")] - pub session_id: Option, - /// Abort signal for cancelling the call. /// /// Runtime handle — never crosses the JSON boundary (bindings that only @@ -154,12 +141,6 @@ pub struct CallOptions { #[ts(skip)] pub abort_signal: Option, - /// RFC-0023 recording context (R7 快照绑定;internal;never serialized)。 - /// 层 A 入口构造一次,provider 复制到 `HttpRequest`。None = 不录制。 - #[serde(skip)] - #[ts(skip)] - pub recording_context: Option, - /// Emit raw provider stream chunks as `StreamPart::Raw` (debugging aid). /// When `Some(true)`, streaming providers yield one `Raw` part per JSON /// SSE event, carrying the **parsed JSON payload** of the event, emitted @@ -198,11 +179,7 @@ impl CallOptions { provider_options: None, reasoning: None, body_overrides: None, - max_retries: None, - timeout: None, - session_id: None, abort_signal: None, - recording_context: None, include_raw_chunks: None, } } diff --git a/aimux-core/src/output.rs b/aimux-core/src/output.rs new file mode 100644 index 000000000..5ee691af7 --- /dev/null +++ b/aimux-core/src/output.rs @@ -0,0 +1,447 @@ +//! Structured output of a text generation: the AI SDK's `Output` +//! specifications (ai@7.0.127 `generate-text/output.ts`). +//! +//! An [`Output`] decides the `response_format` sent to the model and how the +//! final step's text is parsed into the result's `output`. `generate_text` and +//! `stream_text` take it as [`crate::generate::GenerateTextOptions::output`]; +//! without one the text output applies and `output` is the text. +//! +//! Parse failures follow the AI SDK's `NoObjectGeneratedError` messages +//! ("No object generated: could not parse the response." and "No object +//! generated: response did not match schema."), carried by the Aimux variants +//! for the same two cases: [`AiMuxError::JsonParse`] for text that is not +//! JSON and [`AiMuxError::InvalidResponseData`] for JSON of the wrong shape. + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; +use ts_rs::TS; + +use crate::error::AiMuxError; +use crate::options::ResponseFormat; + +/// What the final text of a generation is parsed into. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)] +#[cfg_attr(feature = "ts-export", ts(export))] +pub enum Output { + /// The text itself (`Output.text()`). + Text {}, + /// A JSON object matching `schema` (`Output.object()`). + Object { + schema: Value, + #[serde(default, skip_serializing_if = "Option::is_none")] + name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + description: Option, + }, + /// An array of elements matching `element` (`Output.array()`). The model + /// is asked for `{"elements": [...]}`. + Array { + element: Value, + #[serde(default, skip_serializing_if = "Option::is_none")] + min_items: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + max_items: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + description: Option, + }, + /// One of `options` (`Output.choice()`). The model is asked for + /// `{"result": "