diff --git a/CHANGELOG.md b/CHANGELOG.md index 44aa653..62ad1c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Added +- `DeltaPart::Thinking { text }` variant + `on_thinking_delta` observer event + (`ThinkingDeltaContext`): reasoning-model tokens (Claude extended-thinking, + DeepSeek-R1, OpenAI o-series, Gemini 2.5+) are now routed as their own + stream kind instead of being dropped or misrouted into text. Stream-only — + reasoning is not accumulated into the `Message` and does not reach + `ResponseContext.text`; consume it via `on_thinking_delta` or the raw + `IndexedDelta(Thinking)` stream event. Anthropic parses `thinking_delta` + (and emits an empty Thinking delta for `redacted_thinking`); OpenAI parses + `reasoning_content` (aliased to `reasoning`); Gemini parses per-part + `thought: true` flags. An empty `delta` signals redacted reasoning (render + a placeholder). For Gemini, `GeminiClientBuilder::include_thoughts(true)` + opts into `generationConfig.thinkingConfig.includeThoughts` on the request + side — opt-in because the Gemini API rejects `thinkingConfig` with + `400 INVALID_ARGUMENT` on non-reasoning models; defaults to `false`. - `DisplayHint` advisory rendering hint on `ToolOutput` (`with_hint()` builder), threaded through `ToolDispatchResult` and `ToolPostContext` so presentation layers (TUI, headless console) can render a tool result by the tool's own @@ -35,8 +49,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. `RequestOptions` (carrying `tool_constraint`) applied to every provider call. Default is `RequestOptions::default()` (no constraint), reproducing prior behavior. -- `StreamHandler::with_request_options(opts)` builder: the same option applied - to the handler's stream-open call. +- `StreamHandler::passthrough()` constructor + `passthrough_default()` — a + no-resilience handler (no retries, no timeouts, no fallback) used as the + engine default when no handler is configured. +- `HandlerEvent` enum (`stream::handler`): events yielded by the new + stream-based `StreamHandler::stream_turn`. Variants: `Stream(StreamEvent)` + for raw provider events, `AttemptReset` on retry, `Fallback { message, + stop_reason }` on non-streaming fallback. +- `StreamRequest` struct (`api`): bundles `(messages, system, tools)` into a + single parameter for `ApiClient` methods. Replaces the positional + `(Vec, Option, Option>)` parameter lists on + `stream_messages`, `create_message`, and their `_with_options` variants. + Builders: `new`, `system`, `system_opt`, `tools`, `tools_opt`. +- `GeminiClientBuilder::include_thoughts(bool)` builder: opt into Gemini's + `thinkingConfig.includeThoughts` for reasoning-capable models (2.5 Pro/Flash, + Gemini 3). Defaults to `false` — the Gemini API rejects `thinkingConfig` + with `400 INVALID_ARGUMENT` on non-reasoning models, so the caller must opt + in once they know their model supports thinking. - `ContextContributor` trait and `ContributorContext<'a>` (`engine::contributor` module): a write-side hook at the turn boundary. Implementors return an optional `Message` that the loop appends to the conversation before the next @@ -53,10 +82,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. - `LoopError::RateLimitEscalation { attempts, retry_after }` variant, recoverable, raised when the stream handler exhausts rate-limit retries on a model and escalates to the circuit breaker. -- `StreamHandlerError::RateLimitEscalation { attempts, retry_after, prior }` +- `StreamHandlerError::RateLimitEscalation { attempts, retry_after }` variant. After `RateLimitConfig::fallback_after_retries` rate-limit retries, - `stream_turn` returns this instead of looping indefinitely or falling back to - the same model's non-streaming endpoint. + `stream_turn` yields this as a stream error instead of looping + indefinitely or falling back to the same model's non-streaming endpoint. - Rate-limit backoff sleeps are now clamped to the turn's `total_stream_timeout` so a large `Retry-After` cannot overrun the turn budget. - The engine routes `RateLimitEscalation` to @@ -135,6 +164,30 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. ### Changed +- **Breaking:** `stream::DeltaPart` is now `#[non_exhaustive]`. Every + `match` on `DeltaPart` in downstream code must include a `_ =>` wildcard arm + — handling `Thinking` explicitly is optional but does not replace the + wildcard requirement. Same-crate matches are unaffected. Future variant + additions (e.g. `Image`, `Audio`) will arrive non-breaking. +- **Breaking:** `ApiClient::stream_messages`, `stream_messages_with_options`, + `create_message`, and `create_message_with_options` now take a single + `StreamRequest` parameter instead of positional `(messages, system, tools)`. + Every `impl ApiClient` must update its signatures. +- **Breaking:** `StreamHandler::stream_turn` now returns + `impl Stream>` instead of + `Future>`. Callers must drive the stream and + accumulate events themselves. The engine's `stream_turn` does this internally + and fires observer callbacks (`on_text_delta`, `on_thinking_delta`, + `text_streamer`) per event — configuring a `StreamHandler` for resilience no + longer drops real-time observability. +- **Breaking:** `StreamCapable::stream_handler` now returns `&StreamHandler` + (not `Option<&StreamHandler>`). When no handler is configured, returns a + shared `StreamHandler::passthrough_default()` (no-resilience default). +- **Breaking:** `StreamHandlerError::RateLimitEscalation` lost its `prior` + field. The variant is now `{ attempts, retry_after }`. +- **Breaking:** `StreamHandler::stream_turn` now takes `options: + RequestOptions` as an explicit parameter. `StreamHandler::with_request_options` + is removed — use `BareLoop::set_request_options` instead. - **Breaking:** `ToolOutput`, `ToolDispatchResult`, and `ToolPostContext` are now `#[non_exhaustive]`, matching `DisplayHint`. Downstream code that constructs these via struct literal must switch to the named constructors @@ -194,6 +247,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. honored at the next turn boundary; now it aborts the remaining calls in the batch. +### Removed + +- `StreamTurnResult` (the handler no longer accumulates; the engine assembles + the result from the event stream). +- `StreamHandler::with_request_options` builder (options now flow via + `stream_turn`'s parameter; configure via `BareLoop::set_request_options`). +- `StreamHandlerError::RateLimitEscalation.prior: StreamOutcome` field (never + read by any consumer). + ## [0.1.0] - 2025-07-01 Initial crates.io release. diff --git a/Cargo.toml b/Cargo.toml index 04f8eb5..e55c922 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,7 +28,7 @@ uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" reqwest = { version = "0.13", default-features = false, features = ["json", "stream", "rustls"], optional = true } -async-stream = { version = "0.3", optional = true } +async-stream = "0.3" httpdate = { version = "1", optional = true } jsonschema = { version = "0.30", optional = true } @@ -44,7 +44,7 @@ tool_health = [] tool_shield = ["tool_health"] # Providers -providers = ["dep:reqwest", "dep:async-stream", "dep:httpdate"] +providers = ["dep:reqwest", "dep:httpdate"] openai = ["providers"] anthropic = ["providers"] ollama = ["providers", "openai"] diff --git a/src/api.rs b/src/api.rs index a2a7865..34dd566 100644 --- a/src/api.rs +++ b/src/api.rs @@ -15,6 +15,83 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +/// A request to an LLM provider — the conversation, system prompt, and tools. +/// +/// Bundles the three fields that every [`ApiClient`] method takes together, +/// so method signatures stay short (`&self, request, ...`) instead of +/// repeating `messages, system, tools` at every call site. +/// +/// Construct via `StreamRequest::new(messages)` and chain `.system(...)` +/// / `.tools(...)` as needed, or build with struct-literal syntax. +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::api::StreamRequest; +/// +/// let req = StreamRequest::new(vec![Message::user("hi")]) +/// .system("be brief") +/// .tools(vec![search_tool]); +/// let stream = client.stream_messages(req); +/// ``` +#[derive(Debug, Clone)] +pub struct StreamRequest { + /// The conversation history. Owned because the request body must be built + /// from owned data — callers clone the full history each turn + /// (O(n) in messages). + pub messages: Vec, + /// An optional system prompt to prepend. + pub system: Option, + /// Optional tool definitions the model may invoke. + pub tools: Option>, +} + +impl StreamRequest { + /// Create a request with messages and no system prompt or tools. + #[must_use] + pub fn new(messages: Vec) -> Self { + Self { + messages, + system: None, + tools: None, + } + } + + /// Set the system prompt. + #[must_use] + pub fn system(mut self, system: impl Into) -> Self { + self.system = Some(system.into()); + self + } + + /// Set the system prompt from an existing [`Option`]. + /// + /// Convenience for call sites that already hold the prompt as an `Option` + /// (e.g. forwarded from another field). `None` leaves it unset. + #[must_use] + pub fn system_opt(mut self, system: Option) -> Self { + self.system = system; + self + } + + /// Set the tool definitions. + #[must_use] + pub fn tools(mut self, tools: Vec) -> Self { + self.tools = Some(tools); + self + } + + /// Set the tool definitions from an existing [`Option>`]. + /// + /// Convenience for call sites that already hold tools as an `Option`. + /// `None` leaves them unset. + #[must_use] + pub fn tools_opt(mut self, tools: Option>) -> Self { + self.tools = tools; + self + } +} + /// Interface for API clients that communicate with LLM providers. /// /// Defines the contract for both streaming and non-streaming @@ -54,9 +131,7 @@ use std::sync::Arc; /// /// fn stream_messages( /// &self, -/// messages: Vec, -/// system: Option, -/// tools: Option>, +/// request: StreamRequest, /// ) -> Pin> + Send + 'static>> { /// // Clone data from &self, then build and return a stream /// let model = self.model.clone(); @@ -67,9 +142,7 @@ use std::sync::Arc; /// /// fn create_message( /// &self, -/// messages: Vec, -/// system: Option, -/// tools: Option>, +/// request: StreamRequest, /// ) -> Pin> + Send + '_>> { /// // Non-streaming fallback /// todo!() @@ -118,9 +191,9 @@ pub trait ApiClient: Send + Sync { /// Stream messages from the LLM provider. /// - /// Sends the conversation history (`messages`), an optional - /// `system` prompt, and optional [`ToolSchema`] definitions to the - /// LLM, returning a `'static` [`Stream`] of [`StreamEvent`]s. + /// Sends the [`StreamRequest`] (conversation history, optional system + /// prompt, optional tool definitions) to the LLM, returning a `'static` + /// [`Stream`] of [`StreamEvent`]s. /// /// Called by the agent's turn-processing loop for every LLM /// interaction. The stream produces events such as @@ -136,55 +209,32 @@ pub trait ApiClient: Send + Sync { /// required data from `&self` before constructing the stream. No /// references to `&self` may be captured. /// - /// # Parameters - /// - /// - `messages` — The conversation history as a [`Vec`]. - /// Takes ownership because the request body must be built from owned - /// data for the returned `+ '_` future; callers (e.g. - /// [`BareLoop`](crate::engine::BareLoop)) clone the full history each - /// turn — O(n) in the number of messages. - /// - `system` — An optional system prompt to prepend. - /// - `tools` — Optional tool definitions the model may invoke. - /// /// # Returns /// /// A pinned, boxed stream of [`Result`]. fn stream_messages( &self, - messages: Vec, - system: Option, - tools: Option>, + request: StreamRequest, ) -> Pin> + Send + 'static>>; /// Non-streaming message request (fallback). /// - /// Sends the same parameters as [`stream_messages`](ApiClient::stream_messages) - /// but returns a single [`serde_json::Value`] instead of a stream. Useful - /// for simple one-shot queries where streaming overhead isn't needed, - /// or as a fallback when the provider does not support streaming. + /// Sends the same [`StreamRequest`] as + /// [`stream_messages`](ApiClient::stream_messages) but returns a single + /// [`serde_json::Value`] instead of a stream. Useful for simple one-shot + /// queries where streaming overhead isn't needed, or as a fallback when + /// the provider does not support streaming. /// /// Called by utility code that needs a complete response in one shot, /// such as token estimation probes or health checks. /// - /// # Parameters - /// - /// - `messages` — The conversation history as a [`Vec`]. - /// Takes ownership because the request body must be built from owned - /// data for the returned `+ '_` future; callers (e.g. - /// [`BareLoop`](crate::engine::BareLoop)) clone the full history each - /// turn — O(n) in the number of messages. - /// - `system` — An optional system prompt to prepend. - /// - `tools` — Optional tool definitions the model may invoke. - /// /// # Returns /// /// A pinned, boxed future resolving to the raw JSON response value /// from the provider, or an [`ApiError`] if the request fails. fn create_message( &self, - messages: Vec, - system: Option, - tools: Option>, + request: StreamRequest, ) -> Pin> + Send + '_>>; /// Streaming variant that honors [`RequestOptions`](crate::structured::RequestOptions). @@ -198,13 +248,11 @@ pub trait ApiClient: Send + Sync { /// this to inject the schema. fn stream_messages_with_options( &self, - messages: Vec, - system: Option, - tools: Option>, + request: StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { if options.response_format.is_none() { - return self.stream_messages(messages, system, tools); + return self.stream_messages(request); } Box::pin(futures::stream::once(async { Err(ApiError::config( @@ -227,13 +275,11 @@ pub trait ApiClient: Send + Sync { /// on the extracted payload. fn create_message_with_options( &self, - messages: Vec, - system: Option, - tools: Option>, + request: StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { if options.response_format.is_none() { - return self.create_message(messages, system, tools); + return self.create_message(request); } Box::pin(async { Err(ApiError::config( @@ -317,9 +363,7 @@ mod tests { fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { // Return a simple stream with one event let events: Vec> = vec![ @@ -347,9 +391,7 @@ mod tests { fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { @@ -369,7 +411,11 @@ mod tests { #[tokio::test] async fn test_mock_client_stream() { let client = MockClient::new("test-model"); - let stream = client.stream_messages(vec![Message::user("Hi")], None, None); + let stream = client.stream_messages(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }); let events: Vec<_> = stream.collect().await; assert_eq!(events.len(), 4); @@ -395,13 +441,63 @@ mod tests { async fn test_mock_client_create_message() { let client = MockClient::new("test-model"); let result = client - .create_message(vec![Message::user("Hi")], None, None) + .create_message(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }) .await; assert!(result.is_ok()); let json = result.unwrap(); assert!(json.get("content").is_some()); } + #[test] + fn stream_request_new_defaults() { + let req = StreamRequest::new(vec![Message::user("hi")]); + assert_eq!(req.messages.len(), 1); + assert!(req.system.is_none()); + assert!(req.tools.is_none()); + } + + #[test] + fn stream_request_system_builder() { + let req = StreamRequest::new(vec![]).system("be brief"); + assert_eq!(req.system.as_deref(), Some("be brief")); + } + + #[test] + fn stream_request_system_opt_builder() { + let req = StreamRequest::new(vec![]).system_opt(Some("x".into())); + assert_eq!(req.system.as_deref(), Some("x")); + let req = StreamRequest::new(vec![]).system_opt(None); + assert!(req.system.is_none()); + } + + #[test] + fn stream_request_tools_builder() { + let tools = vec![crate::tool::ToolSchema { + tool: "search".into(), + description: "Search".into(), + input_schema: serde_json::json!({"type": "object"}), + }]; + let req = StreamRequest::new(vec![]).tools(tools); + assert_eq!(req.tools.as_ref().unwrap().len(), 1); + } + + #[test] + fn stream_request_tools_opt_builder() { + let tools = vec![crate::tool::ToolSchema { + tool: "search".into(), + description: "Search".into(), + input_schema: serde_json::json!({"type": "object"}), + }]; + let req = StreamRequest::new(vec![]).tools_opt(Some(tools)); + assert_eq!(req.tools.as_ref().unwrap().len(), 1); + let req = StreamRequest::new(vec![]).tools_opt(None); + assert!(req.tools.is_none()); + } + #[test] fn test_boxed_client() { let client: BoxedApiClient = Box::new(MockClient::new("boxed")); diff --git a/src/capabilities.rs b/src/capabilities.rs index 8e04221..6ba4565 100644 --- a/src/capabilities.rs +++ b/src/capabilities.rs @@ -171,8 +171,14 @@ pub trait Compactable { /// that need to control streaming behaviour (timeouts, retries, fallback /// to non-streaming mode). pub trait StreamCapable { - /// Returns the stream handler, if resilient streaming is configured. - fn stream_handler(&self) -> Option<&StreamHandler>; + /// Returns the stream handler. + /// + /// Always returns a handler — when no resilient handler is configured, + /// returns a shared reference to [`StreamHandler::passthrough_default`] + /// (a no-resilience handler that yields the raw provider stream with no + /// retries, timeouts, or fallback). The engine's `stream_turn` always + /// routes through a handler; this never returns `None`. + fn stream_handler(&self) -> &StreamHandler; } /// Capability to run bidirectional hooks that can block actions. diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 330b283..6261178 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -931,46 +931,90 @@ impl BareLoop { async fn do_stream(&mut self) -> Result<(Message, Option, StreamStopReason), LoopError> { match self.stream_turn().await { Ok((msg, usage, stop)) => { - self.managers.fallback.record_model_success(); - let (in_tok, out_tok) = Self::usage_tokens(usage.as_ref()); - self.managers.observers().on_stream_success(&StreamContext { - turn: self.budget.total_turns, - model: self.client.model(), - input_tokens: in_tok, - output_tokens: out_tok, - }); + self.record_stream_success(usage.as_ref()); Ok((msg, usage, stop)) } - Err(e) => { - let tripped = if matches!(e, LoopError::RateLimitEscalation { .. }) { - self.managers.fallback.record_model_failure() - } else { - self.managers.fallback.record_api_failure() - }; - if tripped { - let from = self.client.model(); - if let Some(to) = self.managers.fallback.fallback_model() { - tracing::warn!(from = %from, to = %to, "fallback manager tripped"); - self.managers - .observers() - .on_fallback(&FallbackContext { from, to }); - } - } + Err(e) => Err(self.record_stream_failure(e)), + } + } + /// Record a successful stream completion. + /// + /// Fires [`on_stream_success`](crate::observer::LoopObserver::on_stream_success) + /// with this turn's token counts and tells the + /// [`FallbackManager`](crate::fallback::FallbackManager) the current model + /// is healthy (so a transient failure earlier in the session doesn't keep + /// the circuit breaker tripped forever). + /// + /// Called from [`do_stream`](Self::do_stream) on the `Ok` branch only; + /// has no return value because the caller already holds the successful + /// `(Message, Option, StreamStopReason)` and just needs the + /// side-effects. + fn record_stream_success(&mut self, usage: Option<&Usage>) { + self.managers.fallback.record_model_success(); + let (in_tok, out_tok) = Self::usage_tokens(usage); + + // TODO: fire_stream_success + self.managers.observers().on_stream_success(&StreamContext { + turn: self.budget.total_turns, + model: self.client.model(), + input_tokens: in_tok, + output_tokens: out_tok, + }); + } + + /// Record a stream failure and return the error to propagate. + /// + /// Distinguishes [`LoopError::RateLimitEscalation`] (which trips the + /// model circuit breaker via + /// [`record_model_failure`](crate::fallback::FallbackManager::record_model_failure)) + /// from other stream errors (which only count as a generic API failure + /// via + /// [`record_api_failure`](crate::fallback::FallbackManager::record_api_failure)). + /// When the breaker trips and a fallback model is configured, fires + /// [`on_fallback`](crate::observer::LoopObserver::on_fallback). + /// + /// Then fires + /// [`on_stream_failure`](crate::observer::LoopObserver::on_stream_failure) + /// (regardless of breaker outcome), sets the terminal + /// [`LoopState::Failed`], and returns the original error so the caller + /// can propagate it from [`do_stream`](Self::do_stream). + /// + /// `LoopError::Cancelled` is **not** routed through here — cancellation + /// is a clean termination that sets [`LoopState::Cancelled`] without + /// tripping the breaker or firing `on_stream_failure`. The caller + /// ([`run_turn_body`](Self::run_turn_body)) handles cancellation before + /// reaching [`do_stream`](Self::do_stream). + fn record_stream_failure(&mut self, e: LoopError) -> LoopError { + let tripped = if matches!(e, LoopError::RateLimitEscalation { .. }) { + self.managers.fallback.record_model_failure() + } else { + self.managers.fallback.record_api_failure() + }; + + if tripped { + let from = self.client.model(); + if let Some(to) = self.managers.fallback.fallback_model() { + tracing::warn!(from = %from, to = %to, "fallback manager tripped"); self.managers .observers() - .on_stream_failure(&StreamFailureContext { - turn: self.budget.total_turns, - model: self.client.model(), - error: e.clone(), - }); - - self.state = LoopState::Failed { - error: e.to_string(), - }; - Err(e) + .on_fallback(&FallbackContext { from, to }); } } + + // TODO: fire_stream_failure + self.managers + .observers() + .on_stream_failure(&StreamFailureContext { + turn: self.budget.total_turns, + model: self.client.model(), + error: e.clone(), + }); + + self.state = LoopState::Failed { + error: e.to_string(), + }; + e } /// Fold this turn's token counts into the running session totals on @@ -1146,6 +1190,7 @@ impl BareLoop { turn_out: u64, turn_start: Instant, ) -> TurnResult { + // TODO: fire_complete_turn self.finish_turn(turn_in, turn_out, turn_start.elapsed()); self.state = LoopState::Completed { summary: text.clone(), @@ -1299,6 +1344,7 @@ impl BareLoop { self.budget.total_turns = self.budget.total_turns.saturating_add(1); if tool_calls.is_empty() { + // TODO: complete turn return Ok(self.complete_session(text, turn_in, turn_out, turn_start)); } @@ -1392,6 +1438,7 @@ impl crate::engine::loop_core::Loop for BareLoop { biased; () = cancel.notified() => { self.state = LoopState::Cancelled; + // TODO: fire_turn_cancelled self.managers.observers().on_turn_end(&TurnEndContext { turn: current_turn, success: false, @@ -1658,9 +1705,7 @@ mod tests { fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { let mut guard = crate::error::recover_guard(self.responses.lock()); @@ -1677,9 +1722,7 @@ mod tests { fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(json!({"content": []})) }) } @@ -1969,7 +2012,7 @@ mod tests { let result = agent.run("Hi").await; assert!(result.is_err()); match result.unwrap_err() { - LoopError::Api(msg) => assert!(msg.contains("No more mock responses")), + LoopError::Api(msg) => assert!(msg.contains("No more mock responses"), "got: {msg}"), other => panic!("Expected Api error, got: {other}"), } } @@ -2193,6 +2236,36 @@ mod tests { ); } + #[tokio::test] + async fn test_text_streamer_fires_when_stream_handler_configured() { + // Regression: when a StreamHandler is attached, the engine must still + // fire text_streamer / on_text_delta for each streamed text delta. The + // handler path used to bypass observers entirely. + let client = MockClient::new("test-model"); + client.add_text_response("via handler"); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + agent.set_stream_handler(StreamHandler::new()); + + let received = Arc::new(Mutex::new(Vec::new())); + let buf = Arc::clone(&received); + agent.set_text_streamer(Arc::new(move |delta: &str| { + crate::error::recover_guard(buf.lock()).push(delta.to_string()); + })); + + let result = agent.run("Hi").await.unwrap(); + assert!(result.success); + + let received = crate::error::recover_guard(received.lock()); + assert!( + !received.is_empty(), + "streamer should fire even with a StreamHandler configured" + ); + assert!( + received.join("").contains("via handler"), + "got: {received:?}", + ); + } + #[tokio::test] async fn test_text_streamer_none_works() { let client = MockClient::new("test-model"); @@ -3117,9 +3190,7 @@ mod tests { fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { let rx = crate::error::recover_guard(self.rx.lock()) @@ -3130,9 +3201,7 @@ mod tests { fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Err(ApiError::api("not implemented")) }) } @@ -3336,9 +3405,7 @@ mod tests { fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn futures::stream::Stream> @@ -3351,9 +3418,7 @@ mod tests { fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn std::future::Future> @@ -3819,9 +3884,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { Box::pin(futures::stream::once(async { @@ -3833,9 +3896,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(json!({})) }) } @@ -4029,11 +4090,10 @@ mod tests { fn stream_messages( &self, - messages: Vec, - _system: Option, - _tools: Option>, + request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { + let messages = request.messages; crate::error::recover_guard(self.seen.lock()).push(messages); let mut guard = crate::error::recover_guard(self.responses.lock()); if let Some(events) = guard.pop_front() { @@ -4048,12 +4108,11 @@ mod tests { fn stream_messages_with_options( &self, - messages: Vec, - _system: Option, - _tools: Option>, + request: crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { + let messages = request.messages; crate::error::recover_guard(self.seen.lock()).push(messages); crate::error::recover_guard(self.seen_options.lock()).push(options); let mut guard = crate::error::recover_guard(self.responses.lock()); @@ -4069,9 +4128,7 @@ mod tests { fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(json!({"content": []})) }) } @@ -4450,4 +4507,158 @@ mod tests { "GoalReminder (cadence 1) should have injected the goal text as a System message" ); } + + #[tokio::test] + async fn test_on_thinking_delta_fires_per_thinking_delta() { + struct ThinkingRecorder { + deltas: Arc>>, + } + impl crate::observer::LoopObserver for ThinkingRecorder { + fn name(&self) -> &'static str { + "thinking-recorder" + } + fn on_thinking_delta(&self, ctx: &crate::observer::ThinkingDeltaContext) { + crate::error::recover_guard(self.deltas.lock()).push((ctx.turn, ctx.delta.clone())); + } + } + + let client = MockClient::new("test-model"); + let events = vec![ + StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: "msg-1".into(), + role: "assistant".into(), + model: "test-model".into(), + }, + }), + StreamEvent::PartStart(PartStart { + index: 1, + part: None, + }), + StreamEvent::IndexedDelta(IndexedDelta { + index: 1, + delta: DeltaPart::Thinking { + text: "First reasoning".into(), + }, + }), + StreamEvent::IndexedDelta(IndexedDelta { + index: 1, + delta: DeltaPart::Thinking { + text: " chunk".into(), + }, + }), + StreamEvent::PartStop, + StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::text("ignored")), + }), + StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::Text { + text: "final answer".into(), + }, + }), + StreamEvent::PartStop, + StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some("end_turn".into()), + }, + usage: None, + }), + StreamEvent::MessageStop, + ]; + client.add_events(events); + + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + let captured = Arc::new(Mutex::new(Vec::new())); + let recorder = Arc::new(ThinkingRecorder { + deltas: Arc::clone(&captured), + }); + agent.register_observer(recorder as Arc); + + let result = agent.run("Hi").await.unwrap(); + assert!(result.success); + + let captured = crate::error::recover_guard(captured.lock()); + assert_eq!( + captured.len(), + 2, + "one on_thinking_delta per Thinking delta" + ); + let joined: String = captured.iter().map(|(_, d)| d.as_str()).collect(); + assert_eq!(joined, "First reasoning chunk"); + assert_eq!(captured[0].0, 0, "turn number matches budget.total_turns"); + } + + #[tokio::test] + async fn test_on_thinking_delta_independent_of_text_delta() { + struct MixedRecorder { + text_calls: Arc>, + thinking_calls: Arc>, + } + impl crate::observer::LoopObserver for MixedRecorder { + fn name(&self) -> &'static str { + "mixed-recorder" + } + fn on_text_delta(&self, _ctx: &crate::observer::TextDeltaContext) { + *crate::error::recover_guard(self.text_calls.lock()) += 1; + } + fn on_thinking_delta(&self, _ctx: &crate::observer::ThinkingDeltaContext) { + *crate::error::recover_guard(self.thinking_calls.lock()) += 1; + } + } + + let client = MockClient::new("test-model"); + let events = vec![ + StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: "msg-1".into(), + role: "assistant".into(), + model: "test-model".into(), + }, + }), + StreamEvent::IndexedDelta(IndexedDelta { + index: 1, + delta: DeltaPart::Thinking { + text: "reasoning".into(), + }, + }), + StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::Text { + text: "answer".into(), + }, + }), + StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some("end_turn".into()), + }, + usage: None, + }), + StreamEvent::MessageStop, + ]; + client.add_events(events); + + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + let text_calls = Arc::new(Mutex::new(0usize)); + let thinking_calls = Arc::new(Mutex::new(0usize)); + let recorder = Arc::new(MixedRecorder { + text_calls: Arc::clone(&text_calls), + thinking_calls: Arc::clone(&thinking_calls), + }); + agent.register_observer(recorder as Arc); + + agent.run("Hi").await.unwrap(); + + assert_eq!( + *crate::error::recover_guard(text_calls.lock()), + 1, + "text callback fires once (for the Text delta)" + ); + assert_eq!( + *crate::error::recover_guard(thinking_calls.lock()), + 1, + "thinking callback fires once (for the Thinking delta)" + ); + } } diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index d887fc0..3bb9a31 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -1139,9 +1139,7 @@ mod tests { } fn stream_messages( &self, - _history: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream> @@ -1153,9 +1151,7 @@ mod tests { } fn create_message( &self, - _history: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Err(ApiError::http("not implemented")) }) diff --git a/src/engine/bare/stream.rs b/src/engine/bare/stream.rs index 0285252..0c3a497 100644 --- a/src/engine/bare/stream.rs +++ b/src/engine/bare/stream.rs @@ -1,30 +1,34 @@ //! Streaming — send the conversation to the LLM API and accumulate the response. //! -//! When a [`StreamHandler`](crate::stream::handler::StreamHandler) is configured, -//! delegates to it for resilient streaming (retry, timeout, fallback). Otherwise, -//! uses basic inline logic. +//! Streaming always routes through a [`StreamHandler`](crate::stream::handler::StreamHandler). +//! When no handler is configured, the engine uses +//! [`StreamHandler::passthrough_default`](crate::stream::handler::StreamHandler::passthrough_default) +//! — a no-resilience handler that yields the raw provider stream with no retries, +//! timeouts, or fallback (equivalent to the pre-redesign inline path). Configuring +//! a handler via [`set_stream_handler()`](BareLoop::set_stream_handler) opts into +//! retry, timeout, fallback, and rate-limit handling. use super::{ ApiClient, BareLoop, LoopError, Message, StreamAccumulator, StreamEvent, StreamStopReason, Usage, }; use crate::capabilities::StreamCapable; -use crate::observer::TextDeltaContext; -use crate::stream::handler::{StreamHandler, StreamHandlerError}; +use crate::observer::{TextDeltaContext, ThinkingDeltaContext}; +use crate::stream::handler::{HandlerEvent, StreamHandlerError}; use futures::StreamExt; impl BareLoop { /// Stream one turn from the API, accumulating the response. /// - /// When a [`StreamHandler`] is configured (via - /// [`set_stream_handler()`](BareLoop::set_stream_handler)), delegates - /// to the handler for resilient streaming with retry, timeout, and - /// fallback capabilities. Otherwise, uses the basic inline logic - /// with no retries. + /// Always routes through a [`StreamHandler`] — when none is configured, + /// [`passthrough_default`](crate::stream::handler::StreamHandler::passthrough_default) + /// is used (no retries, no timeouts, no fallback). Configure a handler via + /// [`set_stream_handler()`](BareLoop::set_stream_handler) to opt into + /// resilient streaming. /// /// Sends the current conversation history to the LLM API via - /// [`ApiClient::stream_messages`] and uses a [`StreamAccumulator`] - /// to collect the events into a single [`Message`]. + /// [`ApiClient::stream_messages_with_options`] and uses a + /// [`StreamAccumulator`] to collect the events into a single [`Message`]. /// /// Also captures the stop reason (e.g. `end_turn`, `tool_call`) and /// token [`Usage`] from the stream's final `MessageDelta` event. @@ -40,116 +44,98 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`LoopError::Api`] if any stream event is an error. When a - /// [`StreamHandler`](crate::stream::handler::StreamHandler) is configured, - /// may also return [`LoopError::Cancelled`] if the handler's cancel-aware - /// `select!` fires mid-stream. The inline path does not check cancellation - /// itself — that is handled by the `select!` in `process_turn`, which drops - /// this future if cancelled. + /// Returns [`LoopError::Api`] if any stream event is an error. May also + /// return [`LoopError::Cancelled`] if the handler's cancel-aware `select!` + /// fires mid-stream. pub(super) async fn stream_turn( &self, ) -> Result<(Message, Option, StreamStopReason), LoopError> { - // Delegate to StreamHandler if configured. - if let Some(handler) = self.managers.stream_handler() { - return self.stream_turn_via_handler(handler).await; - } - - // Inline streaming (no handler). - let system = self.config.system_prompt.clone(); - let tool_schemas = self.build_tool_schemas(); - // Clone the conversation history for the API request. The `ApiClient` - // trait requires `'static` streams (it takes ownership of the - // messages), so a clone is unavoidable here. The in-memory clone is - // O(n) in the number of messages but is typically dwarfed by the - // cost of serialising the messages into an HTTP request body. For - // very long sessions (>200 turns with large tool outputs), consider - // enabling auto-compaction to bound the history size. - let mut stream = self.client.stream_messages_with_options( - self.conversation.clone(), - system, - tool_schemas, + let handler = self.managers.stream_handler(); + let mut stream = handler.stream_turn( + &*self.client, + crate::api::StreamRequest::new(self.conversation.clone()) + .system_opt(self.config.system_prompt.clone()) + .tools_opt(self.build_tool_schemas()), self.request_options.clone(), + &self.cancelled, ); + let mut accumulator = StreamAccumulator::new(); let mut stop_reason = StreamStopReason::EndTurn; - loop { - let event_result = stream.next().await; - match event_result { - Some(Ok(event)) => { - // Match the text delta once, then notify both the raw - // text_streamer callback (if set) and every registered - // observer. The match is hoisted out of the - // `if let Some(streamer)` guard so observers fire even - // when no streamer is set. - if let StreamEvent::IndexedDelta(indexed_delta) = &event - && let crate::stream::DeltaPart::Text { text } = &indexed_delta.delta - { - if let Some(ref streamer) = self.text_streamer { - streamer(text); - } - self.managers.observers().on_text_delta(&TextDeltaContext { - turn: self.budget.total_turns, - delta: text.clone(), - }); - } - if let StreamEvent::MessageDelta(delta) = &event - && let Some(ref reason_str) = delta.delta.stop_reason - { - stop_reason = - StreamStopReason::from_api_str(reason_str).unwrap_or(stop_reason); - } - accumulator - .process(&event) - .map_err(|e| LoopError::Api(format!("stream accumulation error: {e}")))?; + while let Some(result) = stream.next().await { + match result.map_err(Self::map_handler_error)? { + HandlerEvent::Stream(ev) => { + self.dispatch_stream_event(&ev, &mut accumulator, &mut stop_reason)?; + } + HandlerEvent::AttemptReset => { + accumulator = StreamAccumulator::new(); + stop_reason = StreamStopReason::EndTurn; } - Some(Err(api_error)) => { - return Err(LoopError::Api(api_error.to_string())); + HandlerEvent::Fallback { + message, + stop_reason: fallback_stop_reason, + } => { + return Ok((message, None, fallback_stop_reason)); } - None => break, } } let usage = accumulator.usage().copied(); - let message = accumulator.build(); - - Ok((message, usage, stop_reason)) + Ok((accumulator.build(), usage, stop_reason)) } - /// Stream one turn via the [`StreamHandler`]. + /// Dispatch one stream event: fire observer callbacks + /// (`text_streamer` + `on_text_delta` for [`DeltaPart::Text`], + /// `on_thinking_delta` for [`DeltaPart::Thinking`]), extract the stop + /// reason from [`MessageDelta`] events, then fold the event into the + /// accumulator. /// - /// Delegates streaming to the handler, which manages retries, - /// timeouts, and fallback to non-streaming. Maps the handler's - /// result/error types back to the `(Message, Option, - /// StreamStopReason)` tuple expected by the run loop. + /// Shared by all `HandlerEvent::Stream` events regardless of whether the + /// source is the passthrough handler or a configured resilient handler. /// /// # Errors /// - /// Maps [`StreamHandlerError`] variants to the appropriate - /// [`LoopError`] variants: - /// - [`Cancelled`](StreamHandlerError::Cancelled) → [`LoopError::Cancelled`] - /// - [`InitFailed`](StreamHandlerError::InitFailed) → [`LoopError::Api`] - /// - [`StreamFailed`](StreamHandlerError::StreamFailed) → [`LoopError::Api`] - /// - [`FallbackFailed`](StreamHandlerError::FallbackFailed) → [`LoopError::Api`] - /// - [`RateLimitEscalation`](StreamHandlerError::RateLimitEscalation) → - /// [`LoopError::RateLimitEscalation`] - async fn stream_turn_via_handler( + /// Returns [`LoopError::Api`] if the event cannot be accumulated (e.g. + /// malformed tool-call JSON in a `PartStop` boundary). + /// + /// [`DeltaPart::Text`]: crate::stream::DeltaPart::Text + /// [`DeltaPart::Thinking`]: crate::stream::DeltaPart::Thinking + fn dispatch_stream_event( &self, - handler: &StreamHandler, - ) -> Result<(Message, Option, StreamStopReason), LoopError> { - let system = self.config.system_prompt.clone(); - let tool_schemas = self.build_tool_schemas(); - let result = handler - .stream_turn( - &*self.client, - self.conversation.clone(), - system, - tool_schemas, - &self.cancelled, - ) - .await - .map_err(Self::map_handler_error)?; - Ok((result.message, result.usage, result.stop_reason)) + event: &StreamEvent, + accumulator: &mut StreamAccumulator, + stop_reason: &mut StreamStopReason, + ) -> Result<(), LoopError> { + if let StreamEvent::IndexedDelta(d) = event + && let crate::stream::DeltaPart::Text { text } = &d.delta + { + if let Some(streamer) = &self.text_streamer { + streamer(text); + } + self.managers.observers().on_text_delta(&TextDeltaContext { + turn: self.budget.total_turns, + delta: text.clone(), + }); + } + if let StreamEvent::IndexedDelta(d) = event + && let crate::stream::DeltaPart::Thinking { text } = &d.delta + { + self.managers + .observers() + .on_thinking_delta(&ThinkingDeltaContext { + turn: self.budget.total_turns, + delta: text.clone(), + }); + } + if let StreamEvent::MessageDelta(d) = event + && let Some(reason_str) = &d.delta.stop_reason + { + *stop_reason = StreamStopReason::from_api_str(reason_str).unwrap_or(*stop_reason); + } + accumulator + .process(event) + .map_err(|e| LoopError::Api(format!("stream accumulation error: {e}"))) } /// Map a [`StreamHandlerError`] to an [`LoopError`]. @@ -176,7 +162,6 @@ impl BareLoop { StreamHandlerError::RateLimitEscalation { attempts, retry_after, - prior: _, } => LoopError::RateLimitEscalation { attempts, retry_after, @@ -189,7 +174,6 @@ impl BareLoop { mod tests { use super::*; use crate::api::error::ApiError; - use crate::stream::handler::{DetectedRateLimit, RateLimitKind, StreamOutcome}; // Minimal ApiClient so the `BareLoop` associated fn is callable. struct StubClient; @@ -199,9 +183,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -209,9 +191,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + '_>, > { @@ -221,20 +201,10 @@ mod tests { #[test] fn map_handler_error_escalation() { - let prior = StreamOutcome::RateLimited { - detail: DetectedRateLimit { - kind: RateLimitKind::RateLimited, - retry_after: Some(std::time::Duration::from_secs(12)), - message: "slow down".to_string(), - }, - has_partial_data: false, - events_processed: 0, - }; let mapped = BareLoop::::map_handler_error(StreamHandlerError::RateLimitEscalation { attempts: 3, retry_after: Some(std::time::Duration::from_secs(12)), - prior, }); match mapped { LoopError::RateLimitEscalation { diff --git a/src/observer.rs b/src/observer.rs index 5f36584..e95a1cc 100644 --- a/src/observer.rs +++ b/src/observer.rs @@ -13,6 +13,7 @@ //! - [`StreamContext`] / [`StreamFailureContext`] — stream success/failure //! - [`ResponseContext`] — model response text and usage //! - [`TextDeltaContext`] — incremental text chunk while streaming +//! - [`ThinkingDeltaContext`] — incremental reasoning chunk while streaming //! - [`ToolCallReceivedContext`] — tool call accumulated, before dispatch //! - [`ToolPreContext`] / [`ToolPostContext`] — tool dispatch lifecycle //! - [`CompactedContext`] — context window compaction @@ -43,8 +44,8 @@ pub mod context; pub use context::{ CompactedContext, ConvergenceDetectedContext, FallbackContext, LoopDetectedContext, ModelSwitchedContext, ResponseContext, SessionEndContext, SessionStartContext, StreamContext, - StreamFailureContext, TextDeltaContext, ToolCallReceivedContext, ToolPostContext, - ToolPreContext, TurnEndContext, TurnStartContext, + StreamFailureContext, TextDeltaContext, ThinkingDeltaContext, ToolCallReceivedContext, + ToolPostContext, ToolPreContext, TurnEndContext, TurnStartContext, }; // ================================================== // LoopObserver Trait @@ -125,6 +126,17 @@ pub trait LoopObserver: Send + Sync { /// must not parse, render, or perform I/O synchronously, since it runs on /// the stream-ingestion task. /// + /// # Retry caveat (handler path) + /// + /// When a [`StreamHandler`](crate::stream::handler::StreamHandler) is + /// configured, this callback fires for every event of every attempt — + /// including partial events from a failed attempt that got cut off + /// mid-stream. An observer that concatenates `delta` across calls will, + /// under the handler path only, see duplicated or truncated-then-restarted + /// fragments after a retry. Consumers that need only committed output must + /// buffer until [`on_response`](Self::on_response) (or + /// [`on_turn_end`](Self::on_turn_end)). + /// /// # Examples /// /// ```no_run @@ -144,6 +156,20 @@ pub trait LoopObserver: Send + Sync { /// ``` fn on_text_delta(&self, _ctx: &TextDeltaContext) {} + /// Called for each incremental reasoning ("thinking") chunk. + /// + /// Fired per `DeltaPart::Thinking` during streaming, symmetric to + /// [`on_text_delta`](Self::on_text_delta). Reasoning is distinct from + /// visible assistant text; do not concatenate it with `on_text_delta` / + /// [`on_response`](Self::on_response) output. Redacted reasoning arrives + /// as an empty `delta` (render a placeholder, not the empty string). + /// + /// Inherits the same retry caveat as [`on_text_delta`](Self::on_text_delta): + /// under a configured [`StreamHandler`](crate::stream::handler::StreamHandler), + /// partial events from a failed attempt fire here too. Buffer until + /// [`on_turn_end`](Self::on_turn_end) if you need only committed reasoning. + fn on_thinking_delta(&self, _ctx: &ThinkingDeltaContext) {} + /// Called when the engine has accumulated a tool call and is about to dispatch it. /// /// Fires once per call, after the streaming response is accumulated and @@ -348,6 +374,17 @@ impl ObserverHost { } } + /// Dispatch [`LoopObserver::on_thinking_delta`] to all observers. + /// + /// Iterates registered observers in registration order. Called once per + /// streamed reasoning delta, so each observer's `on_thinking_delta` must + /// be cheap. + pub fn on_thinking_delta(&self, ctx: &ThinkingDeltaContext) { + for obs in &self.observers { + obs.on_thinking_delta(ctx); + } + } + /// Dispatch [`LoopObserver::on_tool_pre`] to all observers. /// /// Iterates registered observers in registration order. @@ -708,4 +745,60 @@ mod tests { input: serde_json::Value::Null, }); } + + #[test] + fn host_dispatches_on_thinking_delta_to_all_observers() { + struct ThinkingRecorder { + deltas: std::sync::Mutex>, + } + impl LoopObserver for ThinkingRecorder { + fn name(&self) -> &'static str { + "thinking-recorder" + } + fn on_thinking_delta(&self, ctx: &ThinkingDeltaContext) { + crate::error::recover_guard(self.deltas.lock()).push(ctx.delta.clone()); + } + } + + let obs1 = Arc::new(ThinkingRecorder { + deltas: std::sync::Mutex::new(Vec::new()), + }); + let obs2 = Arc::new(ThinkingRecorder { + deltas: std::sync::Mutex::new(Vec::new()), + }); + let mut host = ObserverHost::new(); + host.register(Arc::clone(&obs1) as Arc); + host.register(Arc::clone(&obs2) as Arc); + + host.on_thinking_delta(&ThinkingDeltaContext { + turn: 2, + delta: "reasoning".into(), + }); + + assert_eq!( + crate::error::recover_guard(obs1.deltas.lock()).clone(), + vec!["reasoning".to_string()] + ); + assert_eq!( + crate::error::recover_guard(obs2.deltas.lock()).clone(), + vec!["reasoning".to_string()] + ); + } + + #[test] + fn on_thinking_delta_default_is_noop() { + struct NoopObserver; + impl LoopObserver for NoopObserver { + fn name(&self) -> &'static str { + "noop" + } + } + let obs = NoopObserver; + // A bare observer that doesn't override on_thinking_delta must compile + // and not panic when the method is called. + obs.on_thinking_delta(&ThinkingDeltaContext { + turn: 0, + delta: String::new(), + }); + } } diff --git a/src/observer/context.rs b/src/observer/context.rs index 798ed91..560e941 100644 --- a/src/observer/context.rs +++ b/src/observer/context.rs @@ -223,6 +223,44 @@ pub struct TextDeltaContext { pub delta: String, } +/// Context for [`LoopObserver::on_thinking_delta`](crate::observer::LoopObserver::on_thinking_delta). +/// +/// Carries one incremental reasoning ("thinking") chunk. Parallel in shape to +/// [`TextDeltaContext`]: same fields, same lifetime semantics (concatenate in +/// arrival order per turn to reconstruct the full reasoning trace). Reasoning +/// is distinct from the assistant's visible text and is never included in +/// [`ResponseContext`]. +/// +/// # Redacted reasoning +/// +/// An empty `delta` signals redacted reasoning (e.g. Anthropic +/// `redacted_thinking`) — the provider withheld the content. Consumers should +/// render a placeholder ("reasoning redacted"), not the empty string. +/// +/// # Example +/// +/// ``` +/// use loopctl::observer::ThinkingDeltaContext; +/// +/// let ctx = ThinkingDeltaContext { turn: 0, delta: "considering options…".to_string() }; +/// assert_eq!(ctx.turn, 0); +/// assert_eq!(ctx.delta, "considering options…"); +/// ``` +#[derive(Debug, Clone)] +pub struct ThinkingDeltaContext { + /// Turn number (0-indexed), matching `on_turn_start` / `on_turn_end`. + /// + /// Same value [`TextDeltaContext::turn`] carries for the same turn, so an + /// observer can interleave thinking and text deltas correctly. + pub turn: usize, + + /// The incremental reasoning chunk. Concatenate in arrival order per turn. + /// + /// Empty string for redacted/encrypted reasoning (e.g. Anthropic + /// `redacted_thinking`) — render a placeholder, not the empty string. + pub delta: String, +} + /// Context for [`LoopObserver::on_tool_call_received`](crate::observer::LoopObserver::on_tool_call_received). /// /// Fired once per tool call after the streaming response has been accumulated diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index 62042a9..10ecd81 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -225,39 +225,28 @@ impl ApiClient for AnthropicClient { fn stream_messages( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { - self.stream_messages_with_options( - messages, - system, - tools, - crate::structured::RequestOptions::default(), - ) + self.stream_messages_with_options(request, crate::structured::RequestOptions::default()) } fn create_message( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { - self.create_message_with_options( - messages, - system, - tools, - crate::structured::RequestOptions::default(), - ) + self.create_message_with_options(request, crate::structured::RequestOptions::default()) } fn stream_messages_with_options( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = build_request_body( @@ -296,11 +285,14 @@ impl ApiClient for AnthropicClient { fn create_message_with_options( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = build_request_body( @@ -897,6 +889,7 @@ impl SseReader { /// - Emitting [`PartStop`] when parts finish. /// - Emitting the final [`MessageDelta`] with stop reason and usage. #[derive(Default)] +#[allow(clippy::struct_excessive_bools)] struct StreamEmitter { /// Whether [`MessageStart`] has been emitted for the current stream. /// @@ -936,6 +929,23 @@ struct StreamEmitter { /// [`DeltaPart::InputJson`]: crate::stream::DeltaPart::InputJson current_tool_index: Option, + /// Whether a thinking/reasoning content block is currently open. + /// + /// Anthropic signals the start of a thinking block with + /// `content_block_start` (`type: "thinking"` or `"redacted_thinking"`) + /// and its end with `content_block_stop`. This flag tracks the open state + /// so the matching `content_block_stop` emits exactly one + /// [`StreamEvent::PartStop`]. + thinking_part_open: bool, + + /// Index of the thinking block. + /// + /// Most recently opened thinking bolock by + /// `content_block_start`, used to route subsequent `thinking_delta` + /// fragments. Mirrors [`current_tool_index`](Self::current_tool_index) + /// for the reasoning lane. + thinking_index: Option, + /// Whether the terminal stop signal has been processed. /// /// Set by either `message_delta` (which carries the stop reason and @@ -1064,6 +1074,20 @@ impl StreamEmitter { part: Some(MessagePart::text("")), })); } + Some("thinking" | "redacted_thinking") => { + self.thinking_part_open = true; + self.thinking_index = Some(index); + self.push(StreamEvent::PartStart(PartStart { index, part: None })); + + if matches!(block_type, Some("redacted_thinking")) { + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index, + delta: DeltaPart::Thinking { + text: String::new(), + }, + })); + } + } _ => {} } } @@ -1113,6 +1137,21 @@ impl StreamEmitter { })); } } + Some("thinking_delta") => { + let text = v + .pointer("/delta/thinking") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + if !text.is_empty() { + // Use the index from the corresponding content_block_start. + let thinking_index = self.thinking_index.unwrap_or(TEXT_PART_INDEX); + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: thinking_index, + delta: DeltaPart::Thinking { text }, + })); + } + } _ => {} } } @@ -1130,6 +1169,10 @@ impl StreamEmitter { if self.text_part_open { self.text_part_open = false; self.push(StreamEvent::PartStop); + } else if self.thinking_part_open { + self.thinking_part_open = false; + self.thinking_index = None; + self.push(StreamEvent::PartStop); } else if self.tool_parts_open > 0 { self.tool_parts_open = self.tool_parts_open.saturating_sub(1); self.current_tool_index = None; @@ -1191,12 +1234,12 @@ impl StreamEmitter { /// Handle a `message_stop` event. /// /// Marks the stream finished, closes any content blocks still marked - /// open (one [`StreamEvent::PartStop`] for an open text block, then - /// one per open tool block, with both counters reset), and emits the - /// terminal [`StreamEvent::MessageStop`] that consumers rely on to - /// know the stream is complete. The closing `PartStop`s are pushed - /// first so the event order matches the documented protocol - /// (`PartStop* → MessageStop`). + /// open (one [`StreamEvent::PartStop`] for an open thinking block, one + /// for an open text block, then one per open tool block, with all + /// counters reset), and emits the terminal [`StreamEvent::MessageStop`] + /// that consumers rely on to know the stream is complete. The closing + /// `PartStop`s are pushed first so the event order matches the + /// documented protocol (`PartStop* → MessageStop`). /// /// Setting `finished` here is what suppresses the synthetic /// [`StreamEvent::MessageStop`] in [`finish`](Self::finish), so a @@ -1204,12 +1247,17 @@ impl StreamEmitter { /// terminal event appended after the SSE stream ends. fn on_message_stop(&mut self) { self.finished = true; + if self.thinking_part_open { + self.push(StreamEvent::PartStop); + } if self.text_part_open { self.push(StreamEvent::PartStop); } for _ in 0..self.tool_parts_open { self.push(StreamEvent::PartStop); } + self.thinking_part_open = false; + self.thinking_index = None; self.tool_parts_open = 0; self.text_part_open = false; self.push(StreamEvent::MessageStop); @@ -2269,4 +2317,163 @@ mod tests { assert_eq!(messages[0]["content"], "first"); assert_eq!(messages[2]["content"], "third"); } + + #[test] + fn emitter_thinking_delta_emits_thinking_variant() { + let mut em = StreamEmitter::default(); + + // Start a thinking block at index 0. + em.on_block_start(Some(serde_json::json!({ + "index": 0, + "content_block": {"type": "thinking"} + }))); + em.drain(); + + // Thinking delta. + em.on_block_delta(Some(serde_json::json!({ + "delta": {"type": "thinking_delta", "thinking": "reasoning here"} + }))); + let events = em.drain(); + + assert_eq!(events.len(), 1); + match &events[0] { + StreamEvent::IndexedDelta(d) => match &d.delta { + DeltaPart::Thinking { text } => assert_eq!(text, "reasoning here"), + other => panic!("expected Thinking, got {other:?}"), + }, + other => panic!("expected IndexedDelta, got {other:?}"), + } + } + + #[test] + fn emitter_signature_delta_is_ignored() { + let mut em = StreamEmitter::default(); + + // Start a thinking block + emit a thinking delta. + em.on_block_start(Some(serde_json::json!({ + "index": 0, + "content_block": {"type": "thinking"} + }))); + em.on_block_delta(Some(serde_json::json!({ + "delta": {"type": "thinking_delta", "thinking": "visible reasoning"} + }))); + em.drain(); + + // Now send a signature_delta — must NOT emit any additional event. + em.on_block_delta(Some(serde_json::json!({ + "delta": {"type": "signature_delta", "signature": "opaque_base64_blob"} + }))); + let events = em.drain(); + + assert!( + events.is_empty(), + "signature_delta must not emit any events: got {events:?}" + ); + } + + #[test] + fn emitter_redacted_thinking_emits_empty_delta() { + let mut em = StreamEmitter::default(); + + // Start a redacted_thinking block — should emit PartStart + one + // empty Thinking delta (the placeholder convention). + em.on_block_start(Some(serde_json::json!({ + "index": 0, + "content_block": {"type": "redacted_thinking"} + }))); + let events = em.drain(); + + // PartStart + one empty Thinking delta. + assert_eq!(events.len(), 2, "expected PartStart + empty Thinking delta"); + assert!(matches!(events[0], StreamEvent::PartStart(_))); + match &events[1] { + StreamEvent::IndexedDelta(d) => match &d.delta { + DeltaPart::Thinking { text } => { + assert!(text.is_empty(), "redacted thinking → empty text"); + } + other => panic!("expected Thinking, got {other:?}"), + }, + other => panic!("expected IndexedDelta, got {other:?}"), + } + assert!(em.thinking_part_open, "thinking_part_open set"); + } + + #[test] + fn emitter_thinking_block_stop_closes_part() { + let mut em = StreamEmitter::default(); + + em.on_block_start(Some(serde_json::json!({ + "index": 0, + "content_block": {"type": "thinking"} + }))); + em.on_block_delta(Some(serde_json::json!({ + "delta": {"type": "thinking_delta", "thinking": "reasoning"} + }))); + em.drain(); + + // Block stop must close the thinking part: emit one PartStop and + // reset the tracking fields. + em.on_block_stop(None); + let events = em.drain(); + + assert_eq!(events.len(), 1, "exactly one PartStop"); + assert!(matches!(events[0], StreamEvent::PartStop)); + assert!(!em.thinking_part_open, "thinking_part_open reset"); + assert!(em.thinking_index.is_none(), "thinking_index cleared"); + } + + #[test] + fn emitter_message_stop_closes_open_thinking_block() { + // If a thinking block is still open when `message_stop` arrives + // (e.g. a stream that ended without a final `content_block_stop`), + // `on_message_stop` must defensively emit a PartStop for it before + // the terminal MessageStop. Without it the open block leaks and the + // downstream accumulator never finalizes the part boundary. + let mut em = StreamEmitter::default(); + + em.on_block_start(Some(serde_json::json!({ + "index": 0, + "content_block": {"type": "thinking"} + }))); + em.on_block_delta(Some(serde_json::json!({ + "delta": {"type": "thinking_delta", "thinking": "reasoning"} + }))); + em.drain(); + assert!( + em.thinking_part_open, + "test precondition: thinking lane open" + ); + + em.on_message_stop(); + let events = em.drain(); + + // Expected: one PartStop (thinking) then one MessageStop. + assert!( + events.iter().any(|e| matches!(e, StreamEvent::PartStop)), + "message_stop must emit a PartStop for the open thinking block" + ); + assert!( + events.iter().any(|e| matches!(e, StreamEvent::MessageStop)), + "message_stop must emit the terminal MessageStop" + ); + assert!( + !em.thinking_part_open, + "thinking_part_open must be reset by message_stop" + ); + assert!( + em.thinking_index.is_none(), + "thinking_index must be cleared by message_stop" + ); + + // Ordering: PartStop before MessageStop. + let stop_idx = events + .iter() + .position(|e| matches!(e, StreamEvent::PartStop)) + .expect("PartStop present"); + let msg_stop_idx = events + .iter() + .position(|e| matches!(e, StreamEvent::MessageStop)) + .expect("MessageStop present"); + assert!(stop_idx < msg_stop_idx, "PartStop must precede MessageStop"); + } } diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 3e1dea6..27513f1 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -46,6 +46,7 @@ const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta const DEFAULT_MODEL: &str = "gemini-2.0-flash"; const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; +const THINKING_PART_INDEX: usize = 1; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_mins(2); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb @@ -87,6 +88,17 @@ pub struct GeminiClient { /// [`FallbackManager`](crate::fallback::FallbackManager) trips to a /// fallback model. model: std::sync::Mutex, + + /// Whether to request thought summaries from reasoning-capable models. + /// + /// When `true`, every request body gets + /// `generationConfig.thinkingConfig.includeThoughts = true`. The Gemini + /// API rejects `thinkingConfig` with `400 INVALID_ARGUMENT` on models + /// that don't support thinking, so this is opt-in — the caller must know + /// their model is reasoning-capable (e.g. Gemini 2.5 Pro/Flash, Gemini 3) + /// before enabling it. Defaults to `false`. Set via + /// [`GeminiClientBuilder::include_thoughts`]. + include_thoughts: bool, } impl GeminiClient { @@ -223,16 +235,20 @@ impl ApiClient for GeminiClient { fn stream_messages( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let body = build_request_body( &messages, system.as_deref(), tools.as_deref(), None, &ToolConstraint::None, + self.include_thoughts, ); let url = self.stream_url(); let http = self.http.clone(); @@ -258,16 +274,20 @@ impl ApiClient for GeminiClient { fn create_message( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let body = build_request_body( &messages, system.as_deref(), tools.as_deref(), None, &ToolConstraint::None, + self.include_thoughts, ); let url = self.generate_url(); @@ -290,11 +310,14 @@ impl ApiClient for GeminiClient { fn stream_messages_with_options( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let rf = options.response_format.as_ref(); let body = build_request_body( &messages, @@ -302,6 +325,7 @@ impl ApiClient for GeminiClient { tools.as_deref(), rf, &options.tool_constraint, + self.include_thoughts, ); let url = self.stream_url(); let http = self.http.clone(); @@ -327,11 +351,14 @@ impl ApiClient for GeminiClient { fn create_message_with_options( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let response_format = options.response_format.as_ref(); let body = build_request_body( &messages, @@ -339,6 +366,7 @@ impl ApiClient for GeminiClient { tools.as_deref(), response_format, &options.tool_constraint, + self.include_thoughts, ); let url = self.generate_url(); Box::pin(async move { @@ -405,6 +433,15 @@ pub struct GeminiClientBuilder { /// Separate from the total timeout so a slow-connecting server can be /// detected faster. Defaults to 10 seconds. connect_timeout: Duration, + + /// Whether to opt into thought summaries from reasoning-capable models. + /// + /// When `true`, every request gets + /// `generationConfig.thinkingConfig.includeThoughts = true`. The Gemini + /// API rejects `thinkingConfig` on non-reasoning models with + /// `400 INVALID_ARGUMENT`, so this is `false` by default and the caller + /// must opt in once they know their model supports thinking. + include_thoughts: bool, } impl Default for GeminiClientBuilder { @@ -415,6 +452,7 @@ impl Default for GeminiClientBuilder { model: DEFAULT_MODEL.into(), timeout: DEFAULT_REQUEST_TIMEOUT, connect_timeout: DEFAULT_CONNECT_TIMEOUT, + include_thoughts: false, } } } @@ -473,6 +511,24 @@ impl GeminiClientBuilder { self } + /// Opt into thought summaries from reasoning-capable models. + /// + /// When enabled, every request body gets + /// `generationConfig.thinkingConfig.includeThoughts = true`, which asks + /// Gemini to surface its reasoning alongside the visible answer. Only + /// reasoning-capable models (e.g. Gemini 2.5 Pro/Flash, Gemini 3) honor + /// this — non-reasoning models reject `thinkingConfig` with + /// `400 INVALID_ARGUMENT`, so this defaults to `false`. + /// + /// The response-side parser routes `thought: true` parts to + /// [`DeltaPart::Thinking`] regardless of this flag — it's purely the + /// request-side opt-in. + #[must_use] + pub fn include_thoughts(mut self, enabled: bool) -> Self { + self.include_thoughts = enabled; + self + } + /// Build the client. /// /// # Errors @@ -493,6 +549,7 @@ impl GeminiClientBuilder { api_key, base_url: self.base_url, model: std::sync::Mutex::new(self.model), + include_thoughts: self.include_thoughts, }) } } @@ -506,10 +563,14 @@ impl GeminiClientBuilder { /// Unlike OpenAI/Anthropic, Gemini puts the model in the URL, not the /// request body. Each [`Message`] is serialized via [`convert_message`]. /// +/// `generationConfig` is injected only when it has something to carry: +/// `thinkingConfig.includeThoughts = true` when `include_thoughts` is set, +/// and/or `responseMimeType` + `responseJsonSchema` when `response_format` +/// is set. When neither applies, `generationConfig` is omitted entirely. +/// /// Tool-call constraint: -/// - When `response_format` is set, injects `generationConfig` for the -/// structured-output path and suppresses `tools`; `tool_constraint` is -/// ignored in that case. +/// - When `response_format` is set, suppresses `tools`; `tool_constraint` +/// is ignored in that case. /// - Otherwise, when `tool_constraint` is `Strict`, each /// `functionDeclaration`'s `parameters` is tightened /// (`additionalProperties: false`, full `required`) via [`convert_tools`]. @@ -521,6 +582,7 @@ fn build_request_body( tools: Option<&[ToolSchema]>, response_format: Option<&crate::structured::ResponseFormat>, tool_constraint: &ToolConstraint, + include_thoughts: bool, ) -> Value { let (non_system, effective_system) = super::fold_system_messages(messages, system); let contents: Vec = non_system.iter().map(|m| convert_message(m)).collect(); @@ -544,15 +606,20 @@ fn build_request_body( ); } - if let Some(rf) = response_format { - obj.insert( - "generationConfig".into(), - serde_json::json!({ - "responseMimeType": "application/json", - "responseJsonSchema": rf.schema, - }), + let mut generation_config = serde_json::Map::new(); + if include_thoughts { + generation_config.insert( + "thinkingConfig".into(), + serde_json::json!({ "includeThoughts": true }), ); } + if let Some(rf) = response_format { + generation_config.insert("responseMimeType".into(), "application/json".into()); + generation_config.insert("responseJsonSchema".into(), rf.schema.clone()); + } + if !generation_config.is_empty() { + obj.insert("generationConfig".into(), Value::Object(generation_config)); + } } body @@ -733,8 +800,11 @@ impl SseReader { /// Gemini's SSE format is simpler than Anthropic's: each `data:` line /// is a complete JSON object with `candidates[0].content.parts[]` for /// text/function-call data and `candidates[0].finishReason` for the -/// stop reason. +/// stop reason. Reasoning models (Gemini 2.5+) interleave thought parts +/// with regular parts in the same `parts[]` array, flagged by a +/// `thought: true` boolean on each thought part. #[derive(Default)] +#[allow(clippy::struct_excessive_bools)] struct StreamEmitter { /// Whether [`StreamEvent::MessageStart`] has been emitted for the /// current stream. @@ -745,6 +815,22 @@ struct StreamEmitter { /// chunks as content only. started: bool, + /// Whether the text content part is currently open. + /// + /// The emitter opens a text part with [`StreamEvent::PartStart`] on the + /// first non-empty text fragment and tracks the open state so + /// [`extract_finish_reason`](Self::extract_finish_reason) emits exactly + /// one [`StreamEvent::PartStop`] to close it. + text_part_open: bool, + + /// Whether the reasoning (thinking) content part is currently open. + /// + /// Reasoning models flag thought parts with `thought: true`. The emitter + /// opens a thinking part on the first non-empty thought fragment and + /// closes it in [`extract_finish_reason`](Self::extract_finish_reason), + /// symmetric to the text lane. + thinking_part_open: bool, + /// Whether the terminal stop signal has been processed. /// /// Set by [`finish`](Self::finish) when it appends the synthetic @@ -786,86 +872,182 @@ impl StreamEmitter { })); } - self.extract_text(json); + self.extract_parts(json); self.extract_function_call(json); self.extract_finish_reason(json); } - /// Extract a text delta from the chunk and emit an `IndexedDelta` event. + /// Extract text and thought deltas from the chunk, routing by the + /// `thought: true` flag on each part. /// - /// Reads `candidates[0].content.parts[0].text`. If the text is non-empty, - /// pushes a [`DeltaPart::Text`](crate::stream::DeltaPart::Text) event at - /// index 0. Does nothing if the path is absent or the text is empty. - fn extract_text(&mut self, json: &Value) { - if let Some(text) = json - .pointer("/candidates/0/content/parts/0/text") - .and_then(Value::as_str) - && !text.is_empty() - { - self.push(StreamEvent::IndexedDelta(IndexedDelta { - index: TEXT_PART_INDEX, - delta: DeltaPart::Text { - text: text.to_string(), - }, - })); + /// Gemini reasoning models (2.5+) interleave thought parts with regular + /// text parts in `candidates[0].content.parts[]`. Each part carries a + /// `thought` boolean: `true` for reasoning (routed to + /// [`DeltaPart::Thinking`]), absent or `false` for visible text (routed + /// to [`DeltaPart::Text`]). Opening a lane after the other has been + /// streaming emits a [`PartStop`](StreamEvent::PartStop) for the previous + /// lane first — the downstream accumulator keys on a single active index, + /// so failing to close would let deltas for one lane clobber the other's + /// buffered state. Both lanes are closed by + /// [`extract_finish_reason`](Self::extract_finish_reason). + /// + /// The `functionCall` part is skipped here — it's handled by + /// [`extract_function_call`](Self::extract_function_call). + fn extract_parts(&mut self, json: &Value) { + let Some(parts) = json + .pointer("/candidates/0/content/parts") + .and_then(Value::as_array) + else { + return; + }; + + for part in parts { + // Function-call parts are handled by extract_function_call. + if part.get("functionCall").is_some() { + continue; + } + + let Some(text) = part.get("text").and_then(Value::as_str) else { + continue; + }; + + if text.is_empty() { + continue; + } + + let is_thought = part + .get("thought") + .and_then(Value::as_bool) + .unwrap_or(false); + + if is_thought { + if self.text_part_open { + self.text_part_open = false; + self.push(StreamEvent::PartStop); + } + if !self.thinking_part_open { + self.thinking_part_open = true; + self.push(StreamEvent::PartStart(PartStart { + index: THINKING_PART_INDEX, + part: None, + })); + } + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: THINKING_PART_INDEX, + delta: DeltaPart::Thinking { + text: text.to_string(), + }, + })); + } else { + if self.thinking_part_open { + self.thinking_part_open = false; + self.push(StreamEvent::PartStop); + } + if !self.text_part_open { + self.text_part_open = true; + self.push(StreamEvent::PartStart(PartStart { + index: TEXT_PART_INDEX, + part: Some(MessagePart::text("")), + })); + } + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: TEXT_PART_INDEX, + delta: DeltaPart::Text { + text: text.to_string(), + }, + })); + } } } /// Extract a function (tool) call from the chunk and emit the /// corresponding part-start and input-json events. /// - /// Reads `candidates[0].content.parts[0].functionCall` for the tool - /// `name` and `args`. Emits a [`PartStart`](StreamEvent::PartStart) + /// Scans `candidates[0].content.parts` for the first part containing a + /// `functionCall` field. Emits a [`PartStart`](StreamEvent::PartStart) /// with a [`ToolCall`](crate::message::MessagePart::ToolCall) part, /// followed by an [`InputJson`](crate::stream::DeltaPart::InputJson) /// delta carrying the serialized arguments. Does nothing if no function /// call is present in the chunk. fn extract_function_call(&mut self, json: &Value) { - if let Some(func_call) = json.pointer("/candidates/0/content/parts/0/functionCall") { - let name = func_call - .pointer("/name") - .and_then(Value::as_str) - .unwrap_or("") - .to_string(); - let args = func_call.pointer("/args").cloned().unwrap_or(Value::Null); - let args_str = serde_json::to_string(&args).unwrap_or_default(); - - self.push(StreamEvent::PartStart(PartStart { - index: TEXT_PART_INDEX, - part: Some(MessagePart::ToolCall { - id: String::new(), - name, - input: args, - }), - })); - self.push(StreamEvent::IndexedDelta(IndexedDelta { - index: TEXT_PART_INDEX, - delta: DeltaPart::InputJson { - partial_json: args_str, - }, - })); + let Some(parts) = json + .pointer("/candidates/0/content/parts") + .and_then(Value::as_array) + else { + return; + }; + + let Some(func_call) = parts.iter().find_map(|p| p.get("functionCall")) else { + return; + }; + let name = func_call + .pointer("/name") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + let args = func_call.pointer("/args").cloned().unwrap_or(Value::Null); + let args_str = serde_json::to_string(&args).unwrap_or_default(); + + if self.thinking_part_open { + self.thinking_part_open = false; + self.push(StreamEvent::PartStop); + } + + if self.text_part_open { + self.text_part_open = false; + self.push(StreamEvent::PartStop); } + + self.push(StreamEvent::PartStart(PartStart { + index: TEXT_PART_INDEX, + part: Some(MessagePart::ToolCall { + id: String::new(), + name, + input: args, + }), + })); + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: TEXT_PART_INDEX, + delta: DeltaPart::InputJson { + partial_json: args_str, + }, + })); } /// Extract the finish reason from the chunk and emit stop events. /// /// Reads `candidates[0].finishReason` (e.g. `"STOP"`, `"MAX_TOKENS"`). - /// When present, emits [`PartStop`](StreamEvent::PartStop) followed by a + /// When present, closes any open thinking part and text part with + /// [`PartStop`](StreamEvent::PartStop), then emits a /// [`MessageDelta`](StreamEvent::MessageDelta) carrying the mapped - /// [`StreamStopReason`]. Does nothing if the chunk has no finish reason. + /// [`StreamStopReason`]. No-op on a second `finishReason` chunk (the + /// `finished` guard defends against proxies/gateways that re-emit). fn extract_finish_reason(&mut self, json: &Value) { + if self.finished { + return; + } let Some(reason) = json .pointer("/candidates/0/finishReason") .and_then(Value::as_str) else { return; }; + self.finished = true; let stop = match reason { "MAX_TOKENS" => StreamStopReason::MaxTokens, _ => StreamStopReason::EndTurn, }; - self.push(StreamEvent::PartStop); + if self.thinking_part_open { + self.thinking_part_open = false; + self.push(StreamEvent::PartStop); + } + + if self.text_part_open { + self.text_part_open = false; + self.push(StreamEvent::PartStop); + } + self.push(StreamEvent::MessageDelta(MessageDelta { delta: MessageDeltaPayload { stop_reason: Some(stop.to_api_str().into()), @@ -922,7 +1104,7 @@ mod tests { #[test] fn request_body_user_text() { let msgs = vec![Message::user("hello")]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); let contents = body["contents"].as_array().unwrap(); assert_eq!(contents.len(), 1); @@ -934,7 +1116,14 @@ mod tests { #[test] fn request_body_includes_system_instruction() { let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, Some("be brief"), None, None, &ToolConstraint::None); + let body = build_request_body( + &msgs, + Some("be brief"), + None, + None, + &ToolConstraint::None, + false, + ); let sys = &body["systemInstruction"]; assert!(sys.is_object()); @@ -944,7 +1133,7 @@ mod tests { #[test] fn request_body_no_system_instruction_when_none() { let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); assert!(body.get("systemInstruction").is_none()); } @@ -954,14 +1143,14 @@ mod tests { Role::Assistant, vec![MessagePart::text("hello")], )]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); assert_eq!(body["contents"][0]["role"], "model"); } #[test] fn request_body_user_role() { let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); assert_eq!(body["contents"][0]["role"], "user"); } @@ -975,7 +1164,7 @@ mod tests { input: serde_json::json!({"msg": "hi"}), }], )]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); let parts = body["contents"][0]["parts"].as_array().unwrap(); assert_eq!(parts[0]["functionCall"]["name"], "echo"); @@ -992,7 +1181,7 @@ mod tests { is_error: None, }], )]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); let parts = body["contents"][0]["parts"].as_array().unwrap(); assert_eq!(parts[0]["functionResponse"]["name"], "call_1"); @@ -1010,7 +1199,14 @@ mod tests { description: "Search the web".into(), input_schema: serde_json::json!({"type": "object"}), }]; - let body = build_request_body(&msgs, None, Some(&tools), None, &ToolConstraint::None); + let body = build_request_body( + &msgs, + None, + Some(&tools), + None, + &ToolConstraint::None, + false, + ); let tools_arr = body["tools"].as_array().unwrap(); assert_eq!(tools_arr.len(), 1); @@ -1023,7 +1219,7 @@ mod tests { #[test] fn request_body_no_tools_when_none() { let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); assert!(body.get("tools").is_none()); } @@ -1034,7 +1230,7 @@ mod tests { Message::new(Role::Assistant, vec![MessagePart::text("hi")]), Message::user("bye"), ]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); let contents = body["contents"].as_array().unwrap(); assert_eq!(contents.len(), 3); @@ -1224,10 +1420,289 @@ mod tests { ); } + #[test] + fn emitter_thought_part_routes_to_thinking_variant() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{ + "text": "reasoning here", + "thought": true + }] + } + }] + })); + let events = em.drain(); + let delta = events + .iter() + .find_map(|e| match e { + StreamEvent::IndexedDelta(d) => Some(d.clone()), + _ => None, + }) + .expect("expected at least one IndexedDelta"); + assert!( + matches!(delta.delta, DeltaPart::Thinking { .. }), + "thought:true part must route to Thinking, got {:?}", + delta.delta + ); + if let DeltaPart::Thinking { text } = delta.delta { + assert_eq!(text, "reasoning here"); + } + } + + #[test] + fn emitter_thought_part_emits_part_start_at_thinking_index() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{ + "text": "hmm", + "thought": true + }] + } + }] + })); + let events = em.drain(); + let start = events + .iter() + .find_map(|e| match e { + StreamEvent::PartStart(p) => Some(p.clone()), + _ => None, + }) + .expect("expected a PartStart for the thinking lane"); + assert_eq!( + start.index, THINKING_PART_INDEX, + "thought part must open at the thinking index" + ); + } + + #[test] + fn emitter_thought_part_does_not_open_text_lane() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{ + "text": "private reasoning", + "thought": true + }] + } + }] + })); + let events = em.drain(); + // No DeltaPart::Text event must appear for a thought:true part. + assert!( + !events.iter().any(|e| matches!( + e, + StreamEvent::IndexedDelta(IndexedDelta { + delta: DeltaPart::Text { .. }, + .. + }) + )), + "thought:true part must not produce a Text delta" + ); + } + + #[test] + fn emitter_text_part_does_not_open_thinking_lane() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{ + "text": "visible answer" + }] + } + }] + })); + let events = em.drain(); + assert!( + !events.iter().any(|e| matches!( + e, + StreamEvent::IndexedDelta(IndexedDelta { + delta: DeltaPart::Thinking { .. }, + .. + }) + )), + "thought field absent must not produce a Thinking delta" + ); + assert!( + events.iter().any(|e| matches!( + e, + StreamEvent::IndexedDelta(IndexedDelta { + delta: DeltaPart::Text { .. }, + .. + }) + )), + "plain text part must produce a Text delta" + ); + } + + #[test] + fn emitter_thought_and_text_parts_interleave() { + // Gemini interleaves thought and text parts in the same parts[] array. + // When the lane changes the emitter closes the previous lane with a + // PartStop before opening the next, so the downstream accumulator + // (single active index) never sees one lane's deltas clobber the + // other's buffered state. + let mut em = StreamEmitter::default(); + em.started = true; + + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [ + {"text": "step 1", "thought": true}, + {"text": "answer", "thought": false} + ] + } + }] + })); + let events = em.drain(); + let deltas: Vec<&IndexedDelta> = events + .iter() + .filter_map(|e| match e { + StreamEvent::IndexedDelta(d) => Some(d), + _ => None, + }) + .collect(); + assert_eq!(deltas.len(), 2); + assert!( + matches!(deltas[0].delta, DeltaPart::Thinking { ref text } if text == "step 1"), + "first delta must be the thinking fragment, got {:?}", + deltas[0].delta + ); + assert_eq!(deltas[0].index, THINKING_PART_INDEX); + assert!( + matches!(deltas[1].delta, DeltaPart::Text { ref text } if text == "answer"), + "second delta must be the text fragment, got {:?}", + deltas[1].delta + ); + assert_eq!(deltas[1].index, TEXT_PART_INDEX); + // The lane switch between the two deltas must emit exactly one + // PartStop for the thinking lane before the text PartStart. + assert_eq!( + events + .iter() + .filter(|e| matches!(e, StreamEvent::PartStop)) + .count(), + 1, + "lane switch must close the thinking lane" + ); + } + + #[test] + fn emitter_text_to_thought_lane_switch_emits_part_stop() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": {"parts": [{"text": "visible", "thought": false}]} + }] + })); + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": {"parts": [{"text": "thinking…", "thought": true}]} + }] + })); + let events = em.drain(); + assert_eq!( + events + .iter() + .filter(|e| matches!(e, StreamEvent::PartStop)) + .count(), + 1, + "lane switch must close the text lane" + ); + } + + #[test] + fn emitter_finish_closes_thinking_lane() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{"text": "thinking…", "thought": true}] + } + }] + })); + em.drain(); + em.process_chunk(&serde_json::json!({ + "candidates": [{"finishReason": "STOP"}] + })); + let events = em.drain(); + // The thinking lane was open; finish must close it with a PartStop. + assert!( + events.iter().any(|e| matches!(e, StreamEvent::PartStop)), + "finish must emit PartStop for the open thinking lane" + ); + } + + #[test] + fn request_body_includes_thinking_config_when_opted_in() { + // Opt-in: includeThoughts injected only when the caller asked for it. + let msgs = vec![Message::user("hi")]; + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, true); + + assert_eq!( + body["generationConfig"]["thinkingConfig"]["includeThoughts"], true, + "includeThoughts must be injected when include_thoughts=true" + ); + } + + #[test] + fn request_body_omits_thinking_config_by_default() { + // Default: includeThoughts NOT injected (non-reasoning models reject + // it with 400 INVALID_ARGUMENT). generationConfig should be absent + // entirely when there's nothing else to put in it either. + let msgs = vec![Message::user("hi")]; + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); + + assert!( + body.get("generationConfig").is_none(), + "generationConfig must be omitted when include_thoughts=false and no response_format" + ); + } + + #[test] + fn request_body_thinking_config_composes_with_response_format() { + // Both thinkingConfig and responseJsonSchema land under generationConfig. + let msgs = vec![Message::user("hi")]; + let rf = crate::structured::ResponseFormat::new( + "result", + serde_json::json!({"type": "object", "properties": {"x": {"type": "string"}}}), + ); + let body = build_request_body(&msgs, None, None, Some(&rf), &ToolConstraint::None, true); + + assert_eq!( + body["generationConfig"]["thinkingConfig"]["includeThoughts"], + true + ); + assert_eq!( + body["generationConfig"]["responseMimeType"], + "application/json" + ); + assert_eq!( + body["generationConfig"]["responseJsonSchema"], + serde_json::json!({"type": "object", "properties": {"x": {"type": "string"}}}) + ); + } + #[test] fn emitter_finish_reason_end_turn() { let mut em = StreamEmitter::default(); em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"text": "hi"}]}}] + })); + em.drain(); em.process_chunk(&serde_json::json!({ "candidates": [{"finishReason": "STOP"}] })); @@ -1247,6 +1722,10 @@ mod tests { fn emitter_finish_reason_max_tokens() { let mut em = StreamEmitter::default(); em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"text": "hi"}]}}] + })); + em.drain(); em.process_chunk(&serde_json::json!({ "candidates": [{"finishReason": "MAX_TOKENS"}] })); @@ -1261,6 +1740,162 @@ mod tests { } } + #[test] + fn emitter_duplicate_finish_reason_is_noop() { + // A second finishReason chunk (e.g. from a proxy/gateway that re-emits) + // must not produce duplicate PartStop / MessageDelta events. + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"text": "hi"}]}}] + })); + em.drain(); + em.process_chunk(&serde_json::json!({ + "candidates": [{"finishReason": "STOP"}] + })); + let first = em.drain(); + let first_deltas = first + .iter() + .filter(|e| matches!(e, StreamEvent::MessageDelta(_))) + .count(); + let first_stops = first + .iter() + .filter(|e| matches!(e, StreamEvent::PartStop)) + .count(); + + em.process_chunk(&serde_json::json!({ + "candidates": [{"finishReason": "STOP"}] + })); + let second = em.drain(); + + assert_eq!(first_deltas, 1, "first finishReason emits one MessageDelta"); + assert_eq!(first_stops, 1, "first finishReason emits one PartStop"); + assert!( + second.is_empty(), + "second finishReason must produce no events, got {second:?}" + ); + } + + #[test] + fn emitter_function_call_after_text_closes_text_lane() { + // When a functionCall arrives after the text lane has been streaming, + // the emitter must close the text lane with a PartStop before opening + // the tool part. Both reuse TEXT_PART_INDEX; without the close the + // accumulator would clobber the text buffer with tool state. + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": {"parts": [{"text": "calling tool", "thought": false}]} + }] + })); + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{"functionCall": {"name": "search", "args": {"q": "rust"}}}] + } + }] + })); + let events = em.drain(); + + // Exactly one PartStop for the text lane (between the text delta and + // the tool PartStart). + let stops = events + .iter() + .filter(|e| matches!(e, StreamEvent::PartStop)) + .count(); + assert_eq!(stops, 1, "text lane must be closed before tool PartStart"); + + // Ordering: TextDelta → PartStop → tool PartStart. + let text_delta_idx = events + .iter() + .position(|e| { + matches!( + e, + StreamEvent::IndexedDelta(IndexedDelta { + delta: DeltaPart::Text { .. }, + .. + }) + ) + }) + .expect("Text delta present"); + let stop_idx = events + .iter() + .position(|e| matches!(e, StreamEvent::PartStop)) + .expect("PartStop present"); + let tool_start_idx = events + .iter() + .position(|e| { + matches!( + e, + StreamEvent::PartStart(PartStart { + part: Some(MessagePart::ToolCall { .. }), + .. + }) + ) + }) + .expect("Tool PartStart present"); + assert!(text_delta_idx < stop_idx, "text delta before PartStop"); + assert!(stop_idx < tool_start_idx, "PartStop before tool PartStart"); + } + + #[test] + fn emitter_function_call_at_nonzero_part_index() { + // A functionCall may appear at parts[1] (or later) when Gemini + // interleaves a thought part with a tool call in the same chunk. + // The emitter must scan all parts, not just parts[0]. + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [ + {"text": "reasoning about the call", "thought": true}, + {"functionCall": {"name": "search", "args": {"q": "rust"}}} + ] + } + }] + })); + let events = em.drain(); + assert!( + events.iter().any(|e| matches!( + e, + StreamEvent::PartStart(PartStart { + part: Some(MessagePart::ToolCall { name, .. }), + .. + }) if name == "search" + )), + "functionCall at parts[1] must be found and emitted" + ); + } + + #[test] + fn emitter_no_function_call_does_nothing() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [ + {"text": "just text"}, + {"text": "more text"} + ] + } + }] + })); + let events = em.drain(); + assert!( + !events.iter().any(|e| matches!( + e, + StreamEvent::PartStart(PartStart { + part: Some(MessagePart::ToolCall { .. }), + .. + }) + )), + "no functionCall in any part must not emit a tool PartStart" + ); + } + #[test] fn emitter_finish_emits_message_stop_if_needed() { let mut em = StreamEmitter::default(); @@ -1359,9 +1994,6 @@ mod tests { #[test] fn builder_timeouts_applied_on_build() { - // Verify the build succeeds — reqwest validates the configuration - // internally. If timeout/connect_timeout were somehow invalid, - // .build() would return an error. let client = GeminiClient::builder() .api_key("test-key") .timeout(Duration::from_mins(3)) @@ -1370,6 +2002,28 @@ mod tests { assert!(client.is_ok(), "build should succeed with valid timeouts"); } + #[test] + fn builder_include_thoughts_defaults_false() { + let client = GeminiClient::builder().api_key("test-key").build().unwrap(); + assert!( + !client.include_thoughts, + "include_thoughts must default to false (non-reasoning models reject thinkingConfig)" + ); + } + + #[test] + fn builder_include_thoughts_true_propagates_to_client() { + let client = GeminiClient::builder() + .api_key("test-key") + .include_thoughts(true) + .build() + .unwrap(); + assert!( + client.include_thoughts, + "include_thoughts(true) must propagate to the built client" + ); + } + #[tokio::test] async fn sse_reader_take_line_splits_on_newline() { let mut reader = SseReader { @@ -1400,8 +2054,6 @@ mod tests { #[tokio::test] async fn sse_reader_next_data_malformed_returns_none() { - // Malformed JSON data should be logged (H4) and the reader - // continues looking for the next valid data line. let chunk = "data: not valid json\n\ndata: {\"ok\":true}\n\n"; let stream = futures::stream::iter(vec![Ok::(chunk.to_string())]); let mut reader = SseReader { @@ -1444,7 +2096,7 @@ mod tests { "result", serde_json::json!({"type": "object", "properties": {"x": {"type": "string"}}}), ); - let body = build_request_body(&msgs, None, None, Some(&rf), &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, Some(&rf), &ToolConstraint::None, false); assert_eq!( body["generationConfig"]["responseMimeType"], @@ -1457,12 +2109,15 @@ mod tests { } #[test] - fn request_body_response_format_absent_when_none() { + fn request_body_generation_config_omitted_when_nothing_applies() { + // Without response_format AND include_thoughts=false, generationConfig + // has nothing to carry, so it must be absent entirely. let msgs = vec![Message::user("hi")]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); + assert!( body.get("generationConfig").is_none(), - "generationConfig should be absent when no response_format" + "generationConfig must be omitted when there's nothing to put in it" ); } @@ -1482,6 +2137,7 @@ mod tests { Some(&[caller_tool]), Some(&rf), &ToolConstraint::None, + false, ); assert!( @@ -1533,7 +2189,14 @@ mod tests { "properties": {"msg": {"type": "string"}} }), }]; - let body = build_request_body(&msgs, None, Some(&tools), None, &ToolConstraint::Strict); + let body = build_request_body( + &msgs, + None, + Some(&tools), + None, + &ToolConstraint::Strict, + false, + ); let decls = body["tools"][0]["functionDeclarations"].as_array().unwrap(); assert_eq!(decls.len(), 1); @@ -1553,7 +2216,14 @@ mod tests { description: "Echo".into(), input_schema: serde_json::json!({"type": "object"}), }]; - let body = build_request_body(&msgs, None, Some(&tools), None, &ToolConstraint::None); + let body = build_request_body( + &msgs, + None, + Some(&tools), + None, + &ToolConstraint::None, + false, + ); let decls = body["tools"][0]["functionDeclarations"].as_array().unwrap(); assert_eq!( @@ -1581,6 +2251,7 @@ mod tests { Some(&[caller_tool]), Some(&rf), &ToolConstraint::Strict, + false, ); assert!( @@ -1592,15 +2263,20 @@ mod tests { #[test] fn gemini_strict_does_not_emit_tool_config() { - // Strict constrains shape, not selection — toolConfig must not - // appear under Strict. let msgs = vec![Message::user("hi")]; let tools = vec![ToolSchema { tool: "echo".into(), description: "Echo".into(), input_schema: serde_json::json!({"type": "object"}), }]; - let body = build_request_body(&msgs, None, Some(&tools), None, &ToolConstraint::Strict); + let body = build_request_body( + &msgs, + None, + Some(&tools), + None, + &ToolConstraint::Strict, + false, + ); assert!( body.get("toolConfig").is_none(), @@ -1621,7 +2297,7 @@ mod tests { system_role_msg("stay on task"), Message::assistant("working"), ]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); let contents = body["contents"].as_array().expect("contents is an array"); assert_eq!(contents.len(), 2, "system-role message is filtered out"); @@ -1640,10 +2316,15 @@ mod tests { #[test] fn request_body_system_role_merges_with_caller_system() { - // Both caller `system` and an inline Role::System message present → - // systemInstruction carries both (caller prompt first). let msgs = vec![Message::user("hi"), system_role_msg("reminder")]; - let body = build_request_body(&msgs, Some("be brief"), None, None, &ToolConstraint::None); + let body = build_request_body( + &msgs, + Some("be brief"), + None, + None, + &ToolConstraint::None, + false, + ); let sys_text = body["systemInstruction"]["parts"][0]["text"] .as_str() @@ -1664,14 +2345,13 @@ mod tests { #[test] fn request_body_system_role_preserves_message_order() { - // Folding must not reorder the remaining (non-system) messages. let msgs = vec![ Message::user("first"), system_role_msg("mid reminder"), Message::assistant("second"), Message::user("third"), ]; - let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None); + let body = build_request_body(&msgs, None, None, None, &ToolConstraint::None, false); let contents = body["contents"].as_array().expect("contents is an array"); let roles: Vec<&str> = contents .iter() diff --git a/src/provider/grammar.rs b/src/provider/grammar.rs index af50e91..f1e97e7 100644 --- a/src/provider/grammar.rs +++ b/src/provider/grammar.rs @@ -178,9 +178,8 @@ mod tests { let mut total = 0usize; for prompt in corpus { let stream = client.stream_messages_with_options( - vec![crate::message::Message::user(prompt)], - None, - Some(schemas.clone()), + crate::api::StreamRequest::new(vec![crate::message::Message::user(prompt)]) + .tools_opt(Some(schemas.clone())), opts.clone(), ); let events: Vec> = diff --git a/src/provider/openai.rs b/src/provider/openai.rs index c6d9e85..e0e8bef 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -49,6 +49,7 @@ const DEFAULT_MODEL: &str = "gpt-4o"; const SSE_DONE: &str = "[DONE]"; const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; +const THINKING_PART_INDEX: usize = 1; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_mins(2); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb @@ -222,10 +223,13 @@ impl ApiClient for OpenAiClient { fn stream_messages( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let model = crate::error::recover_guard(self.model.lock()).clone(); let body = RequestBody::build( &model, @@ -262,10 +266,13 @@ impl ApiClient for OpenAiClient { fn create_message( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let model = crate::error::recover_guard(self.model.lock()).clone(); let body = RequestBody::build( &model, @@ -298,11 +305,14 @@ impl ApiClient for OpenAiClient { fn stream_messages_with_options( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = RequestBody::build( @@ -340,11 +350,14 @@ impl ApiClient for OpenAiClient { fn create_message_with_options( &self, - messages: Vec, - system: Option, - tools: Option>, + request: crate::api::StreamRequest, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { + let crate::api::StreamRequest { + messages, + system, + tools, + } = request; let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = RequestBody::build( @@ -934,6 +947,8 @@ struct OpenAiChoice { #[derive(Deserialize)] struct OpenAiDelta { content: Option, + #[serde(alias = "reasoning")] + reasoning_content: Option, tool_calls: Option>, } @@ -966,6 +981,7 @@ struct OpenAiToolCallFunction { /// Splitting this out from the `try_stream!` macro body makes the /// translation logic testable without a live network connection. #[derive(Default)] +#[allow(clippy::struct_excessive_bools)] struct StreamEmitter { /// Whether [`StreamEvent::MessageStart`] has been emitted for the /// current stream. @@ -984,6 +1000,14 @@ struct StreamEmitter { /// emits exactly one [`StreamEvent::PartStop`] to close it. text_part_open: bool, + /// Whether the reasoning (thinking) content part is currently open. + /// + /// Reasoning models stream reasoning via `delta.reasoning_content` (or its + /// alias `delta.reasoning`). The emitter opens a thinking part on the + /// first non-empty fragment and closes it in `process_finish`, symmetric + /// to the text lane. + thinking_part_open: bool, + /// Number of tool-call parts currently open. /// /// Each `delta.tool_calls` entry with a `function` field opens a new @@ -1044,16 +1068,24 @@ impl StreamEmitter { } } - /// Translate a single delta object into text and/or tool-call events. + /// Translate a single delta object into text and/or reasoning events. /// /// If the delta carries non-empty `content`, emits a `PartStart` (on the /// first text delta) followed by `IndexedDelta(Text)` events. If it - /// carries `tool_calls`, delegates each to + /// carries non-empty `reasoning_content`, the same shape is emitted for + /// the reasoning lane. Switching lanes emits a `PartStop` for the previous + /// lane first — the downstream accumulator keys on a single active index, + /// so leaving both open would let one lane's deltas clobber the other's + /// buffered state. If it carries `tool_calls`, delegates each to /// [`process_tool_call`](Self::process_tool_call). fn process_delta(&mut self, delta: &OpenAiDelta) { if let Some(text) = &delta.content && !text.is_empty() { + if self.thinking_part_open { + self.thinking_part_open = false; + self.push(StreamEvent::PartStop); + } if !self.text_part_open { self.text_part_open = true; self.push(StreamEvent::PartStart(PartStart { @@ -1067,6 +1099,28 @@ impl StreamEmitter { })); } + if let Some(reasoning) = &delta.reasoning_content + && !reasoning.is_empty() + { + if self.text_part_open { + self.text_part_open = false; + self.push(StreamEvent::PartStop); + } + if !self.thinking_part_open { + self.thinking_part_open = true; + self.push(StreamEvent::PartStart(PartStart { + index: THINKING_PART_INDEX, + part: None, + })); + } + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: THINKING_PART_INDEX, + delta: DeltaPart::Thinking { + text: reasoning.clone(), + }, + })); + } + if let Some(tool_calls) = &delta.tool_calls { for tc in tool_calls { self.process_tool_call(tc); @@ -1133,6 +1187,11 @@ impl StreamEmitter { self.push(StreamEvent::PartStop); } + // Close any open reasoning (thinking) part. + if self.thinking_part_open { + self.push(StreamEvent::PartStop); + } + // Close each open tool-call part. for _ in 0..self.open_tool_count { self.push(StreamEvent::PartStop); @@ -2100,4 +2159,190 @@ mod tests { assert_eq!(value["role"], "system"); assert_eq!(value["content"], "stay on task"); } + + #[test] + fn openai_delta_reasoning_content_emits_thinking() { + let mut em = StreamEmitter::default(); + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"o1","choices":[{"delta":{"reasoning_content":"thinking…"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + let events = em.drain(); + + // Expect at least one IndexedDelta carrying Thinking. + let thinking = events.iter().find_map(|e| match e { + StreamEvent::IndexedDelta(d) => match &d.delta { + DeltaPart::Thinking { text } => Some(text.clone()), + _ => None, + }, + _ => None, + }); + assert_eq!(thinking.as_deref(), Some("thinking…")); + assert!( + events + .iter() + .any(|e| matches!(e, StreamEvent::PartStart(_))), + "a PartStart should fire for the reasoning lane" + ); + } + + #[test] + fn openai_delta_reasoning_alias_works() { + let mut em = StreamEmitter::default(); + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"o1","choices":[{"delta":{"reasoning":"via alias"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + let events = em.drain(); + + let thinking = events.iter().find_map(|e| match e { + StreamEvent::IndexedDelta(d) => match &d.delta { + DeltaPart::Thinking { text } => Some(text.clone()), + _ => None, + }, + _ => None, + }); + assert_eq!( + thinking.as_deref(), + Some("via alias"), + "#[serde(alias = \"reasoning\")] must accept the field" + ); + } + + #[test] + fn openai_reasoning_does_not_open_text_part() { + let mut em = StreamEmitter::default(); + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"o1","choices":[{"delta":{"reasoning_content":"reasoning only"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + let events = em.drain(); + + // No Text delta should be emitted from a reasoning-only chunk. + let has_text = events.iter().any(|e| { + matches!( + e, + StreamEvent::IndexedDelta(d) if matches!(d.delta, DeltaPart::Text { .. }) + ) + }); + assert!(!has_text, "reasoning-only chunk must not emit Text deltas"); + assert!(!em.text_part_open, "reasoning must not set text_part_open"); + } + + #[test] + fn openai_reasoning_and_text_interleave() { + let mut em = StreamEmitter::default(); + + // Chunk 1: reasoning. + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"o1","choices":[{"delta":{"reasoning_content":"think"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + + // Chunk 2: text. + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"o1","choices":[{"delta":{"content":"answer"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + let events = em.drain(); + + // Both lanes should have fired — one Thinking, one Text. + let has_thinking = events.iter().any(|e| { + matches!( + e, + StreamEvent::IndexedDelta(d) if matches!(d.delta, DeltaPart::Thinking { .. }) + ) + }); + let has_text = events.iter().any(|e| { + matches!( + e, + StreamEvent::IndexedDelta(d) if matches!(d.delta, DeltaPart::Text { .. }) + ) + }); + assert!(has_thinking, "a Thinking delta should have fired"); + assert!(has_text, "a Text delta should have fired"); + + // The lane switch from reasoning to text must emit exactly one + // PartStop for the reasoning lane before the text PartStart. Without + // it the downstream accumulator (which keys on a single active index) + // would let text deltas clobber the reasoning lane's state. + let lane_switch_stops = events + .iter() + .filter(|e| matches!(e, StreamEvent::PartStop)) + .count(); + assert_eq!( + lane_switch_stops, 1, + "lane switch must close the reasoning lane with one PartStop" + ); + } + + #[test] + fn openai_combined_content_and_reasoning_in_one_delta() { + // A single delta object can carry both content and reasoning_content. + // process_delta handles content first (closing thinking if open), then + // reasoning (closing text if open). Lock in the expected sequence: + // PartStart(text) → TextDelta → PartStop → PartStart(thinking) → + // ThinkingDelta. + let mut em = StreamEmitter::default(); + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"o1","choices":[{"delta":{"content":"answer","reasoning_content":"why"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + let events = em.drain(); + + let has_text = events.iter().any(|e| { + matches!( + e, + StreamEvent::IndexedDelta(d) if matches!(d.delta, DeltaPart::Text { ref text } if text == "answer") + ) + }); + let has_thinking = events.iter().any(|e| { + matches!( + e, + StreamEvent::IndexedDelta(d) if matches!(d.delta, DeltaPart::Thinking { ref text } if text == "why") + ) + }); + assert!(has_text, "text delta must fire"); + assert!(has_thinking, "thinking delta must fire"); + + // The reasoning lane opens after the text lane, so the text lane must + // be closed with exactly one PartStop between them. + let stops = events + .iter() + .filter(|e| matches!(e, StreamEvent::PartStop)) + .count(); + assert_eq!(stops, 1, "exactly one PartStop for the text lane"); + + // Ordering: TextDelta → PartStop → ThinkingDelta. + let text_idx = events + .iter() + .position(|e| { + matches!( + e, + StreamEvent::IndexedDelta(d) if matches!(d.delta, DeltaPart::Text { .. }) + ) + }) + .expect("text delta present"); + let stop_idx = events + .iter() + .position(|e| matches!(e, StreamEvent::PartStop)) + .expect("PartStop present"); + let thinking_idx = events + .iter() + .rposition(|e| { + matches!( + e, + StreamEvent::IndexedDelta(d) if matches!(d.delta, DeltaPart::Thinking { .. }) + ) + }) + .expect("thinking delta present"); + assert!(text_idx < stop_idx, "text delta before PartStop"); + assert!(stop_idx < thinking_idx, "PartStop before thinking delta"); + } } diff --git a/src/reflection/llm.rs b/src/reflection/llm.rs index 0217559..ffc88f2 100644 --- a/src/reflection/llm.rs +++ b/src/reflection/llm.rs @@ -293,9 +293,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream> @@ -307,21 +305,22 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(serde_json::json!({})) }) } fn create_message_with_options( &self, - messages: Vec, - system: Option, - _tools: Option>, + request: crate::api::StreamRequest, _options: RequestOptions, ) -> Pin> + Send + '_>> { + let crate::api::StreamRequest { + messages, + system, + tools: _, + } = request; let user = messages.first().and_then(|m| { m.parts.first().and_then(|p| match p { MessagePart::Text { text } => Some(text.clone()), @@ -345,9 +344,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream> @@ -359,18 +356,14 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { Box::pin(async { Ok(serde_json::json!({})) }) } fn create_message_with_options( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, _options: RequestOptions, ) -> Pin> + Send + '_>> { @@ -387,9 +380,7 @@ mod tests { } fn stream_messages( &self, - m: Vec, - s: Option, - t: Option>, + request: crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream> @@ -397,22 +388,18 @@ mod tests { + 'static, >, > { - self.0.stream_messages(m, s, t) + self.0.stream_messages(request) } fn create_message( &self, - m: Vec, - s: Option, - t: Option>, + request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { - self.0.create_message(m, s, t) + self.0.create_message(request) } fn create_message_with_options( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, _options: RequestOptions, ) -> Pin> + Send + '_>> { diff --git a/src/runtime.rs b/src/runtime.rs index c9dad6d..9ff1cbd 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -692,8 +692,10 @@ impl crate::capabilities::Compactable for LoopRuntime { } impl crate::capabilities::StreamCapable for LoopRuntime { - fn stream_handler(&self) -> Option<&StreamHandler> { - self.stream_handler.as_ref() + fn stream_handler(&self) -> &StreamHandler { + self.stream_handler + .as_ref() + .unwrap_or(StreamHandler::passthrough_default()) } } @@ -814,9 +816,13 @@ mod tests { } #[test] - fn test_stream_handler_defaults_to_none() { + fn test_stream_handler_defaults_to_passthrough() { let runtime = LoopRuntime::new(); - assert!(runtime.stream_handler().is_none()); + let handler = runtime.stream_handler(); + assert_eq!( + handler.timeout_config().total_stream_timeout, + std::time::Duration::MAX + ); } #[test] @@ -846,14 +852,27 @@ mod tests { } #[test] - fn test_set_stream_handler_returns_some() { + fn test_set_stream_handler_overrides_passthrough() { use crate::stream::handler::StreamHandler; let handler = StreamHandler::new(); let mut runtime = LoopRuntime::new(); - assert!(runtime.stream_handler().is_none()); + assert_eq!( + runtime + .stream_handler() + .timeout_config() + .total_stream_timeout, + std::time::Duration::MAX + ); runtime.set_stream_handler(handler); - assert!(runtime.stream_handler().is_some()); + // After set_stream_handler, the production defaults apply (15 min total). + assert_eq!( + runtime + .stream_handler() + .timeout_config() + .total_stream_timeout, + std::time::Duration::from_mins(15) + ); } #[test] diff --git a/src/stream.rs b/src/stream.rs index 51dfdaf..3818ec5 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -402,11 +402,18 @@ pub struct IndexedDelta { /// the API sends. Consumers should append the payload to the /// appropriate in-progress part. /// +/// `#[non_exhaustive]` so the framework can add content kinds (e.g. +/// `Image`/`Audio`) in a later minor release without breaking downstream +/// `match`es. Downstream code that matches on `DeltaPart` MUST include a +/// wildcard arm. +/// /// # Variants /// /// - [`Text`](Self::Text) — Append to the text buffer for text parts. /// - [`ToolCall`](Self::ToolCall) — Append to the JSON buffer for tool-call parts. /// - [`InputJson`](Self::InputJson) — Append to the JSON buffer (raw string form). +/// - [`Thinking`](Self::Thinking) — Append to the reasoning buffer for separate +/// display; not part of the assistant's visible text. /// /// # Example /// @@ -424,10 +431,15 @@ pub struct IndexedDelta { /// json_buf.push_str(s); /// } /// } +/// DeltaPart::Thinking { .. } => { +/// // Reasoning is delivered via on_thinking_delta; not accumulated here. +/// } +/// _ => {} /// } /// ``` #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type")] +#[non_exhaustive] pub enum DeltaPart { /// Text delta — append to the existing text content. /// @@ -480,6 +492,29 @@ pub enum DeltaPart { /// [`PartStop`](StreamEvent::PartStop) is received. partial_json: String, }, + + /// Thinking/reasoning delta — append to a reasoning buffer for separate + /// display. NOT part of the assistant's visible text. + /// + /// Emitted by reasoning models (Claude extended-thinking, DeepSeek-R1, + /// OpenAI o-series). Stream-only: the [`StreamAccumulator`] does NOT carry + /// reasoning into the built [`Message`]; consume + /// it via + /// [`on_thinking_delta`](crate::observer::LoopObserver::on_thinking_delta). + /// An empty `text` signals redacted reasoning (e.g. Anthropic + /// `redacted_thinking`); render a placeholder rather than the empty string. + /// + /// Serialized as `"type":"thinking_delta"` with a `"text"` field. + #[serde(rename = "thinking_delta")] + Thinking { + /// The reasoning text fragment to append. + /// + /// Concatenate in arrival order per turn to reconstruct the full + /// reasoning trace. Empty string when the reasoning is redacted + /// (the provider withheld the content); consumers should render a + /// placeholder, not the empty string. + text: String, + }, } // ================================================== @@ -1029,6 +1064,10 @@ impl StreamAccumulator { self.current_tool_input.push_str(s); } } + // Reasoning is not part of the assistant Message. Drop it + // here; observers already received it via on_thinking_delta + // upstream (engine/bare/stream.rs). + DeltaPart::Thinking { .. } => {} } Ok(()) } @@ -1535,7 +1574,6 @@ mod tests { assert_eq!(acc.usage().unwrap().input_tokens, 100); assert_eq!(acc.usage().unwrap().output_tokens, 50); - // Second delta overwrites usage. acc.process(&StreamEvent::MessageDelta(MessageDelta { delta: MessageDeltaPayload { stop_reason: Some("max_tokens".into()), @@ -1546,4 +1584,67 @@ mod tests { assert_eq!(acc.usage().unwrap().input_tokens, 200); assert_eq!(acc.usage().unwrap().output_tokens, 75); } + + #[test] + fn accumulator_drops_thinking_not_into_text() { + let mut acc = StreamAccumulator::new(); + acc.process(&StreamEvent::PartStart(PartStart { + index: 1, + part: None, + })) + .unwrap(); + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 1, + delta: DeltaPart::Thinking { + text: "reasoning here".into(), + }, + })) + .unwrap(); + acc.process(&StreamEvent::PartStop).unwrap(); + let msg = acc.build(); + let text = msg.parts.iter().find_map(|p| match p { + MessagePart::Text { text } => Some(text.as_str()), + _ => None, + }); + assert!( + !text.unwrap_or("").contains("reasoning here"), + "reasoning must not leak into the message text: {text:?}" + ); + } + + #[test] + fn deltapart_thinking_serde_roundtrip() { + let delta = DeltaPart::Thinking { text: "hmm".into() }; + let json = serde_json::to_string(&delta).unwrap(); + assert_eq!(json, r#"{"type":"thinking_delta","text":"hmm"}"#); + let parsed: DeltaPart = serde_json::from_str(&json).unwrap(); + match &parsed { + DeltaPart::Thinking { text } => assert_eq!(text, "hmm"), + other => panic!("expected Thinking, got {other:?}"), + } + } + + #[test] + fn deltapart_thinking_empty_text_roundtrip() { + let delta = DeltaPart::Thinking { + text: String::new(), + }; + let json = serde_json::to_string(&delta).unwrap(); + let parsed: DeltaPart = serde_json::from_str(&json).unwrap(); + match &parsed { + DeltaPart::Thinking { text } => assert_eq!(text, ""), + other => panic!("expected Thinking with empty text, got {other:?}"), + } + } + + #[test] + fn deltapart_thinking_match_compiles() { + let delta = DeltaPart::Thinking { text: "x".into() }; + let result = match &delta { + DeltaPart::Text { text } => format!("text:{text}"), + DeltaPart::Thinking { text } => format!("thinking:{text}"), + DeltaPart::ToolCall { .. } | DeltaPart::InputJson { .. } => "other".into(), + }; + assert_eq!(result, "thinking:x"); + } } diff --git a/src/stream/handler.rs b/src/stream/handler.rs index f810a31..0544823 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -38,10 +38,11 @@ use crate::api::ApiClient; use crate::api::error::{ApiError, http_status_is_overload}; use crate::cancel::CancelSignal; use crate::message::Message; -use crate::stream::{StreamAccumulator, StreamEvent, StreamStopReason, Usage}; -use crate::tool::ToolSchema; +use crate::stream::{StreamAccumulator, StreamEvent, StreamStopReason}; use futures::StreamExt; +use futures::stream::Stream; use std::fmt; +use std::pin::Pin; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -742,7 +743,7 @@ enum EventPoll { /// Read-only diagnostic context used to build timeout and error outcomes. /// -/// Snapshotted once per loop iteration in [`StreamHandler::process_events`] and +/// Snapshotted once per loop iteration in [`StreamHandler::stream_turn`] and /// handed to [`StreamHandler::next_event`], which needs progress/elapsed data to /// populate [`StreamOutcome`] fields when it short-circuits. struct EventDiagnostics { @@ -759,7 +760,7 @@ struct EventDiagnostics { /// /// Read via [`Instant::elapsed`] when building /// [`StreamOutcome::TotalTimeout`]'s `duration` field. Captured once at - /// the top of [`process_events`](StreamHandler::process_events) rather + /// the top of [`stream_turn`](Self::stream_turn) rather /// than per event so the reported duration is the full stream lifetime, /// not the time since the most recent event. stream_start: Instant, @@ -767,7 +768,7 @@ struct EventDiagnostics { /// Whether partial content has been accumulated. /// /// Recomputed each loop iteration (the whole [`EventDiagnostics`] is - /// rebuilt per iteration in [`process_events`](StreamHandler::process_events)) + /// rebuilt per iteration in [`stream_turn`](Self::stream_turn)) /// from the accumulator's current part count: `true` once at least one /// usable event has been received. Flows into the `has_partial_data` flag /// on [`StreamOutcome::TotalTimeout`], [`StreamOutcome::EventTimeout`], @@ -781,7 +782,7 @@ impl EventDiagnostics { /// /// Snapshots the current diagnostic state — partial-data flag, events /// processed so far, and elapsed time since [`stream_start`](Self::stream_start) - /// — into a `TotalTimeout` outcome. Used by [`process_events`](StreamHandler::process_events) + /// — into a `TotalTimeout` outcome. Used by [`stream_turn`](Self::stream_turn) /// when the turn's `total_stream_timeout` deadline fires (both at the /// top-of-loop check and inside the per-event `select!`). fn total_timeout(&self) -> StreamOutcome { @@ -796,7 +797,7 @@ impl EventDiagnostics { /// /// Carries the partial-data flag plus the caller-supplied /// `consecutive_timeouts` count (this method does not track the counter - /// itself — `process_events` owns it and passes the current value in). + /// itself — `stream_turn` owns it and passes the current value in). /// Used once the per-event timeout crosses /// [`max_consecutive_timeouts`](StreamTimeoutConfig::max_consecutive_timeouts). fn event_timeout(&self, consecutive_timeouts: u32) -> StreamOutcome { @@ -813,7 +814,7 @@ impl EventDiagnostics { /// parsed `Retry-After` and current progress; otherwise wraps it as a /// generic [`StreamOutcome::InitFailed`] with `attempts: 1` (this is the /// per-event error path, not the init-retry path, so the attempt counter - /// isn't meaningful here). Used by `process_events` when the stream yields + /// isn't meaningful here). Used by `stream_turn` when the stream yields /// an `Err` event. fn api_error_outcome(&self, error: &crate::api::error::ApiError) -> StreamHandlerError { if let Some(detail) = DetectedRateLimit::detect(error) { @@ -1059,8 +1060,6 @@ pub enum StreamHandlerError { attempts: u32, /// Last server-advised `Retry-After` hint, after clamping. `None` when no header sent. retry_after: Option, - /// The rate-limit outcome that triggered escalation. - prior: StreamOutcome, }, } @@ -1082,7 +1081,6 @@ impl fmt::Display for StreamHandlerError { Self::RateLimitEscalation { attempts, retry_after, - prior: _, } => write!( f, "rate-limit escalation after {attempts} retries (retry-after {retry_after:?})" @@ -1193,8 +1191,8 @@ pub struct StreamHandler { /// /// Drives the initial-event / per-event / total-stream deadlines, the /// consecutive-timeout escalation threshold, and the - /// non-streaming-fallback toggle. Read on every turn in `stream_turn` and - /// on each event poll. + /// non-streaming-fallback toggle. Read on every turn in + /// [`stream_turn`](Self::stream_turn) and on each event poll. timeout_config: StreamTimeoutConfig, /// Retry configuration for stream-initialization failures. @@ -1220,16 +1218,6 @@ pub struct StreamHandler { /// than risking a 429. `None` (the default) means reactive-only handling: /// no pre-throttling, server 429s still handled via `rate_limit_config`. rate_limiter: Option>, - - /// Per-turn request options applied to every stream-open call. - /// - /// Carries [`tool_constraint`](crate::structured::ToolConstraint) for - /// constrained tool-call decoding. Default is - /// [`RequestOptions::default`](crate::structured::RequestOptions::default) - /// (no constraint), reproducing the prior unconstrained behavior. - /// Passed to [`ApiClient::stream_messages_with_options`] at the stream-open - /// call site. Set via [`with_request_options`](StreamHandler::with_request_options). - request_options: crate::structured::RequestOptions, } impl fmt::Debug for StreamHandler { @@ -1239,7 +1227,6 @@ impl fmt::Debug for StreamHandler { .field("retry_config", &self.retry_config) .field("rate_limit_config", &self.rate_limit_config) .field("rate_limiter", &self.rate_limiter) - .field("request_options", &self.request_options) .finish() } } @@ -1251,6 +1238,74 @@ impl Default for StreamHandler { } impl StreamHandler { + /// Create a no-resilience handler — yields the raw provider stream with + /// no retries, no timeouts, no fallback, no rate-limit handling. + /// + /// Used as the default when no `StreamHandler` is configured: the engine + /// always routes streaming through a handler, and `passthrough()` is what + /// you get when you haven't opted into resilience. Behaves like the + /// pre-redesign inline streaming path — one attempt, no retry, surface + /// the underlying [`ApiClient`] errors directly. + /// + /// Also useful as an explicit baseline for callers who want to start + /// fresh and reconfigure only the fields they care about: + /// + /// ```rust,no_run + /// # use loopctl::stream::handler::{StreamHandler, StreamTimeoutConfig}; + /// let handler = StreamHandler::passthrough() + /// .with_config( + /// StreamTimeoutConfig { + /// total_stream_timeout: std::time::Duration::from_secs(60), + /// ..Default::default() + /// }, + /// Default::default(), + /// ); + /// ``` + /// + /// Returned by [`passthrough_default`](Self::passthrough_default) as a + /// shared static, so [`StreamCapable::stream_handler`](crate::capabilities::StreamCapable::stream_handler) + /// can return `&Self` (never `Option<&Self>`). + #[must_use] + pub fn passthrough() -> Self { + // Duration::MAX overflows Instant::now() + it; stream_turn maps that + // overflow to None (no deadline), so this never fires spuriously. + const NEVER_TIME_OUT: Duration = Duration::MAX; + Self { + timeout_config: StreamTimeoutConfig { + initial_event_timeout: NEVER_TIME_OUT, + per_event_timeout: NEVER_TIME_OUT, + total_stream_timeout: NEVER_TIME_OUT, + // validate() rejects 0; value is irrelevant since timeouts never fire. + max_consecutive_timeouts: 1, + progress_interval: NEVER_TIME_OUT, + fallback_to_non_streaming: false, + }, + retry_config: StreamRetryConfig { + max_retries: 0, + ..Default::default() + }, + rate_limit_config: RateLimitConfig { + max_retries: 0, + fallback_after_retries: 0, + ..Default::default() + }, + rate_limiter: None, + } + } + + /// Return a shared reference to the default passthrough handler. + /// + /// Constructed lazily on first access (via `std::sync::OnceLock`) and + /// shared by all callers that don't configure their own [`StreamHandler`]. + /// Used by + /// [`StreamCapable::stream_handler`](crate::capabilities::StreamCapable::stream_handler) + /// to always return `&Self` regardless of configuration. + #[must_use] + pub fn passthrough_default() -> &'static Self { + static PASSTHROUGH: std::sync::OnceLock = std::sync::OnceLock::new(); + PASSTHROUGH.get_or_init(Self::passthrough) + } + /// Create a handler with default configuration. /// /// The defaults are suitable for production LLM API usage: @@ -1273,7 +1328,6 @@ impl StreamHandler { retry_config: StreamRetryConfig::default(), rate_limit_config: RateLimitConfig::default(), rate_limiter: None, - request_options: crate::structured::RequestOptions::default(), } } @@ -1303,20 +1357,6 @@ impl StreamHandler { self } - /// Set the per-turn [`RequestOptions`](crate::structured::RequestOptions) - /// applied to every stream-open call. Consuming builder. - /// - /// Default is [`RequestOptions::default`](crate::structured::RequestOptions::default) - /// (no constraint), which reproduces prior behavior. Set - /// [`tool_constraint`](crate::structured::ToolConstraint) to - /// [`Strict`](crate::structured::ToolConstraint::Strict) for constrained - /// tool-call decoding. - #[must_use] - pub fn with_request_options(mut self, options: crate::structured::RequestOptions) -> Self { - self.request_options = options; - self - } - /// Returns a reference to the timeout configuration. /// /// Read-only access to the [`StreamTimeoutConfig`] stored on the handler. @@ -1368,8 +1408,8 @@ impl StreamHandler { /// Attach a per-provider rate limiter (proactive throttle). /// - /// When set, every `stream_turn` attempt waits for a token before opening - /// the stream, spacing requests to the limiter's `requests_per_minute` + /// When set, every `stream_turn` attempt waits for a token before + /// opening the stream, spacing requests to the limiter's `requests_per_minute` /// ceiling. `None` (the default) disables proactive throttling — the /// handler then relies purely on the reactive 429 handling. /// @@ -1396,148 +1436,173 @@ impl StreamHandler { // Runtime methods // ================================================== - /// Stream one complete turn with retry, timeout, and fallback. - /// - /// Primary entry point for resilient streaming. It - /// orchestrates the full lifecycle: - /// - /// 1. Opens a stream via [`ApiClient::stream_messages_with_options`] - /// (gated by the proactive [`RateLimiter`](crate::stream::rate_limit::RateLimiter), - /// if attached). - /// 2. Processes events with per-event and total timeouts. - /// 3. On transient transport errors, retries with exponential backoff; - /// on 429/503/529 responses, backs off under the rate-limit budget - /// and escalates to the model circuit breaker once that budget is - /// exhausted. - /// 4. If all retries fail and - /// [`fallback_to_non_streaming`](StreamTimeoutConfig::fallback_to_non_streaming) - /// is enabled, falls back to [`ApiClient::create_message`]. + /// Drive one turn as a stream of [`HandlerEvent`]s. /// - /// The returned [`StreamTurnResult`] carries the accumulated message, - /// token usage, stop reason, and timing regardless of which path - /// produced the result (streaming or fallback). + /// The engine consumes this stream directly: real stream events flow to + /// observers and the engine's accumulator (identical to the inline + /// streaming path), [`HandlerEvent::AttemptReset`] signals a retry so the + /// engine can discard partial state, and [`HandlerEvent::Fallback`] + /// delivers the non-streaming fallback message when retries exhaust. /// - /// # Parameters - /// - /// - `client` — The LLM API client to stream from. - /// - `conversation` — The conversation history to send. - /// - `system` — An optional system prompt. - /// - `tool_schemas` — Optional tool definitions. - /// - `cancel` — Shared cancellation signal. + /// The handler keeps its retry/rate-limit/timeout/fallback machinery — + /// this method is the retry loop reshaped as a stream. Events from failed + /// attempts are yielded as `HandlerEvent::Stream` to the engine before the + /// attempt fails; the engine then receives `AttemptReset` and discards + /// them. Consumers that want only committed output must reset their state + /// on `AttemptReset`. /// /// # Errors /// - /// Returns [`StreamHandlerError`] on unrecoverable failure: - /// - [`InitFailed`](StreamHandlerError::InitFailed) — stream could not - /// be opened after all retries. - /// - [`StreamFailed`](StreamHandlerError::StreamFailed) — stream - /// failed mid-event with a timeout or API error. - /// - [`FallbackFailed`](StreamHandlerError::FallbackFailed) — both - /// streaming and non-streaming fallback failed. - /// - [`RateLimitEscalation`](StreamHandlerError::RateLimitEscalation) — - /// rate-limit retries on this model were exhausted; escalate to a - /// fallback model. - /// - [`Cancelled`](StreamHandlerError::Cancelled) — cancellation - /// signal fired. - pub async fn stream_turn( - &self, - client: &C, - conversation: Vec, - system: Option, - tool_schemas: Option>, - cancel: &Arc, - ) -> Result { - let total_deadline = Some( - Instant::now() - .checked_add(self.timeout_config.total_stream_timeout) - .unwrap_or(Instant::now()), - ); + /// Item errors carry the same [`StreamHandlerError`] variants the + /// pre-redesign `stream_turn` returned: cancellation, timeout, transport + /// retry exhaustion, rate-limit escalation, and non-streaming fallback + /// failure. + /// A successful turn ends with `None` from the stream (or after + /// [`HandlerEvent::Fallback`]). + pub fn stream_turn<'a, C: ApiClient>( + &'a self, + client: &'a C, + request: crate::api::StreamRequest, + options: crate::structured::RequestOptions, + cancel: &'a Arc, + ) -> Pin> + Send + 'a>> { + let crate::api::StreamRequest { + messages: conversation, + system, + tools: tool_schemas, + } = request; + // `checked_add` returns `None` when the timeout (e.g. `Duration::MAX` + // in the passthrough handler) would overflow `Instant`. Treat that as + // "no total deadline" (`None`) rather than wrapping to `Instant::now()` + // (which would immediately fire a spurious TotalTimeout). + let total_deadline = Instant::now().checked_add(self.timeout_config.total_stream_timeout); let max_attempts = self.retry_config.max_retries.saturating_add(1); - let mut rate_limit_retries: u32 = 0; - let mut transport_attempts: u32 = 0; - loop { - if let Some(deadline) = total_deadline - && Instant::now() >= deadline - { - return Err(StreamHandlerError::InitFailed( - StreamOutcome::TotalTimeout { - has_partial_data: false, - events_processed: 0, - duration: self.timeout_config.total_stream_timeout, - }, - )); - } - match self - .try_stream( - client, - conversation.clone(), - system.clone(), - tool_schemas.clone(), - cancel, - total_deadline, - ) - .await - { - Ok(result) => return Ok(result), - Err(StreamHandlerError::Cancelled) => { - return Err(StreamHandlerError::Cancelled); + + Box::pin(async_stream::try_stream! { + let mut rate_limit_retries: u32 = 0; + let mut transport_attempts: u32 = 0; + let mut attempt: u32 = 0; + // Shadow accumulator tracks partial-data presence for diagnostics. + let mut shadow = StreamAccumulator::new(); + + 'outer: loop { + if let Some(deadline) = total_deadline + && Instant::now() >= deadline + { + Err(StreamHandlerError::InitFailed( + StreamOutcome::TotalTimeout { + has_partial_data: !shadow.peek_parts().is_empty(), + events_processed: 0, + duration: self.timeout_config.total_stream_timeout, + }, + ))?; } - Err(e) => { - let last_stream_outcome = carried_outcome(&e); - - // Rate-limit retries draw on their own budget - // (RateLimitConfig), independent of the transport retry - // budget below — a rate-limit storm must not exhaust - // transport retries (nor vice versa). - if let Some(StreamOutcome::RateLimited { detail, .. }) = &last_stream_outcome { - match self.rate_limit_retry(detail, &mut rate_limit_retries, total_deadline) - { - RateLimitRetry::Escalate { - attempts, - retry_after, - } => { - return Err(StreamHandlerError::RateLimitEscalation { - attempts, - retry_after, - prior: last_stream_outcome.clone().unwrap_or( - StreamOutcome::RateLimited { - detail: detail.clone(), - has_partial_data: false, - events_processed: 0, - }, - ), - }); - } - RateLimitRetry::HardStop => return Err(e), - RateLimitRetry::Retry(delay) => { - sleep_cancellable(delay, cancel).await?; - continue; + + attempt = attempt.saturating_add(1); + if attempt > 1 { + shadow = StreamAccumulator::new(); + yield HandlerEvent::AttemptReset; + } + + self.gate_on_rate_limit(client, cancel, total_deadline).await?; + let request = crate::api::StreamRequest { + messages: conversation.clone(), + system: system.clone(), + tools: tool_schemas.clone(), + }; + let mut stream = + client.stream_messages_with_options(request, options.clone()); + + let mut consecutive_timeouts: usize = 0; + let mut events_processed: u64 = 0; + let stream_start = Instant::now(); + + loop { + let diagnostics = EventDiagnostics { + events_processed, + stream_start, + has_partial_data: !shadow.peek_parts().is_empty(), + }; + match self.next_event( + &mut stream, + cancel, + &mut consecutive_timeouts, + total_deadline, + &diagnostics, + ).await { + Ok(Some(event)) => { + events_processed = events_processed.saturating_add(1); + consecutive_timeouts = 0; + if let Err(e) = shadow.process(&event) { + Err(StreamHandlerError::StreamFailed( + StreamOutcome::InitFailed { + attempts: 1, + last_error: e.to_string(), + }, + ))?; } + yield HandlerEvent::Stream(event); } - } + Ok(None) => return, + Err(err) => { + let last_stream_outcome = carried_outcome(&err); + + if let Some(StreamOutcome::RateLimited { detail, .. }) = &last_stream_outcome { + let rate_limit_decision = self.rate_limit_retry( + detail, + &mut rate_limit_retries, + total_deadline, + ); + match rate_limit_decision { + RateLimitRetry::Escalate { attempts, retry_after } => { + Err(StreamHandlerError::RateLimitEscalation { + attempts, + retry_after, + })?; + return; + } + RateLimitRetry::HardStop => { + Err(err)?; + return; + } + RateLimitRetry::Retry(delay) => { + sleep_cancellable(delay, cancel).await?; + continue 'outer; + } + } + } - // Non-rate-limit errors consume the transport retry budget. - if transport_attempts >= max_attempts.saturating_sub(1) { - if self.timeout_config.fallback_to_non_streaming { - return self - .fallback_non_streaming( - client, - conversation, - system, - tool_schemas, - cancel, - last_stream_outcome, - ) - .await; + if transport_attempts >= max_attempts.saturating_sub(1) { + if self.timeout_config.fallback_to_non_streaming { + let request = crate::api::StreamRequest { + messages: conversation.clone(), + system: system.clone(), + tools: tool_schemas.clone(), + }; + let (message, fallback_stop_reason) = self.fallback_non_streaming( + client, + request, + cancel, + last_stream_outcome.clone(), + ).await?; + yield HandlerEvent::Fallback { + message, + stop_reason: fallback_stop_reason, + }; + return; + } + Err(err)?; + return; + } + let delay = self.retry_config.base_delay(transport_attempts); + transport_attempts = transport_attempts.saturating_add(1); + sleep_cancellable(delay, cancel).await?; + continue 'outer; } - return Err(e); } - let delay = self.retry_config.base_delay(transport_attempts); - transport_attempts = transport_attempts.saturating_add(1); - sleep_cancellable(delay, cancel).await?; } } - } + }) } /// Decide how to handle a rate-limit failure on the current model. @@ -1578,32 +1643,9 @@ impl StreamHandler { /// [`ApiClient::stream_messages_with_options`] carrying the handler's /// [`RequestOptions`](crate::structured::RequestOptions), and processes /// all events with timeout and cancellation support. Called inside the - /// retry loop in [`stream_turn`](StreamHandler::stream_turn), so a retried - /// 429 re-gates on the rate limiter. + /// retry loop in [`stream_turn`](Self::stream_turn), so a + /// retried 429 re-gates on the rate limiter. /// - /// # Errors - /// - /// Returns [`StreamHandlerError`] if the stream fails or times out. - async fn try_stream( - &self, - client: &C, - conversation: Vec, - system: Option, - tool_schemas: Option>, - cancel: &Arc, - total_deadline: Option, - ) -> Result { - self.gate_on_rate_limit(client, cancel, total_deadline) - .await?; - let stream = client.stream_messages_with_options( - conversation, - system, - tool_schemas, - self.request_options.clone(), - ); - self.process_events(stream, cancel, total_deadline).await - } - /// Wait for a rate-limit token before opening the stream. /// /// When a [`RateLimiter`](crate::stream::rate_limit::RateLimiter) is @@ -1614,11 +1656,12 @@ impl StreamHandler { /// the turn's remaining `total_deadline` so the gate cannot overrun the /// turn budget; if the deadline has already elapsed the gate proceeds /// rather than sleeping (the downstream per-event/total-timeout checks in - /// [`process_events`](Self::process_events) report the expiry). A no-op + /// [`stream_turn`](Self::stream_turn) report the expiry). A no-op /// when no limiter is attached. /// - /// Fires per attempt (called from [`try_stream`](Self::try_stream), which - /// runs inside the retry loop), so a retried 429 re-respects the budget. + /// Fires per attempt (called from + /// [`stream_turn`](Self::stream_turn), which runs the retry + /// loop), so a retried 429 re-respects the budget. /// Cancel-safe: a turn stuck waiting for tokens is still user-cancellable. /// /// # Errors @@ -1665,68 +1708,6 @@ impl StreamHandler { } } - /// Process events from an open stream with per-event and total timeouts. - /// - /// Reads events from the stream, accumulating them into a - /// [`Message`]. Applies per-event timeouts and an overall deadline. - /// Checks cancellation between events. - /// - /// # Errors - /// - /// Returns [`StreamHandlerError::StreamFailed`] on timeout or API error, - /// or [`StreamHandlerError::Cancelled`] if the cancel signal fires. - async fn process_events( - &self, - mut stream: S, - cancel: &Arc, - total_deadline: Option, - ) -> Result - where - S: futures::Stream> - + Unpin, - { - let mut accumulator = StreamAccumulator::new(); - let mut stop_reason = StreamStopReason::EndTurn; - let mut consecutive_timeouts: usize = 0; - let mut events_processed: u64 = 0; - let stream_start = Instant::now(); - - loop { - let diagnostics = EventDiagnostics { - events_processed, - stream_start, - has_partial_data: !accumulator.peek_parts().is_empty(), - }; - let Some(event) = self - .next_event( - &mut stream, - cancel, - &mut consecutive_timeouts, - total_deadline, - &diagnostics, - ) - .await? - else { - break; - }; - events_processed = events_processed.saturating_add(1); - consecutive_timeouts = 0; - Self::apply_event(&event, &mut accumulator, &mut stop_reason)?; - } - - let usage = accumulator.usage().copied(); - let message = accumulator.build(); - let elapsed = stream_start.elapsed(); - - Ok(StreamTurnResult { - message, - usage, - stop_reason, - from_fallback: false, - elapsed, - }) - } - /// Wait for the next stream event, enforcing the total deadline and /// per-event timeout. /// @@ -1766,7 +1747,12 @@ impl StreamHandler { let event_result = tokio::select! { event = stream.next() => EventPoll::Next(event), () = cancel.notified() => return Err(StreamHandlerError::Cancelled), - () = tokio::time::sleep_until(event_deadline.into()) => EventPoll::TimedOut, + () = async { + match event_deadline { + Some(d) => tokio::time::sleep_until(d.into()).await, + None => std::future::pending::<()>().await, + } + } => EventPoll::TimedOut, () = Self::total_deadline_future(total_deadline) => { return Err(StreamHandlerError::StreamFailed(diagnostics.total_timeout())); } @@ -1790,23 +1776,22 @@ impl StreamHandler { } } - /// The deadline for the next stream event. + /// The deadline for the next stream event, or `None` if per-event + /// timeouts are disabled (the timeout is `Duration::MAX`, which + /// overflows `Instant::now() + it`). /// /// Uses [`initial_event_timeout`](StreamTimeoutConfig::initial_event_timeout) /// before any event has arrived (the model may need time to begin /// generating), then switches to - /// [`per_event_timeout`](StreamTimeoutConfig::per_event_timeout) once events - /// are flowing. Falls back to "now" if the addition overflows. - fn event_deadline(&self, events_processed: u64) -> Instant { + /// [`per_event_timeout`](StreamTimeoutConfig::per_event_timeout) once + /// events are flowing. + fn event_deadline(&self, events_processed: u64) -> Option { let base_timeout = if events_processed == 0 { self.timeout_config.initial_event_timeout } else { self.timeout_config.per_event_timeout }; - - Instant::now() - .checked_add(base_timeout) - .unwrap_or(Instant::now()) + Instant::now().checked_add(base_timeout) } /// Whether the total-stream deadline has already passed. @@ -1844,33 +1829,6 @@ impl StreamHandler { } } - /// Fold one event into the accumulator, tracking stop reason. - /// - /// # Errors - /// - /// Returns [`StreamHandlerError::StreamFailed`] if the event cannot be - /// accumulated. - fn apply_event( - event: &StreamEvent, - accumulator: &mut StreamAccumulator, - stop_reason: &mut StreamStopReason, - ) -> Result<(), StreamHandlerError> { - if let StreamEvent::MessageDelta(delta) = event - && let Some(ref reason_str) = delta.delta.stop_reason - { - *stop_reason = StreamStopReason::from_api_str(reason_str).unwrap_or(*stop_reason); - } - if let Err(e) = accumulator.process(event) { - return Err(StreamHandlerError::StreamFailed( - StreamOutcome::InitFailed { - attempts: 1, - last_error: e.to_string(), - }, - )); - } - Ok(()) - } - /// Fall back to non-streaming message creation. /// /// Called when streaming fails (timeout, retries exhausted) and @@ -1885,20 +1843,16 @@ impl StreamHandler { async fn fallback_non_streaming( &self, client: &C, - conversation: Vec, - system: Option, - tool_schemas: Option>, + request: crate::api::StreamRequest, cancel: &Arc, stream_outcome: Option, - ) -> Result { + ) -> Result<(Message, StreamStopReason), StreamHandlerError> { if cancel.is_cancelled() { return Err(StreamHandlerError::Cancelled); } - let start = Instant::now(); - let result = tokio::select! { - res = client.create_message(conversation, system, tool_schemas) => res, + res = client.create_message(request) => res, () = cancel.notified() => { return Err(StreamHandlerError::Cancelled); } @@ -1916,20 +1870,12 @@ impl StreamHandler { .find_map(|p| p.get("text").and_then(|t| t.as_str()).map(String::from)) }) .unwrap_or_default(); - let stop_reason = value .get("stop_reason") .and_then(|r| r.as_str()) .and_then(StreamStopReason::from_api_str) .unwrap_or(StreamStopReason::EndTurn); - - Ok(StreamTurnResult { - message: Message::assistant(&text), - usage: None, - stop_reason, - from_fallback: true, - elapsed: start.elapsed(), - }) + Ok((Message::assistant(&text), stop_reason)) } Err(e) => Err(StreamHandlerError::FallbackFailed { stream_outcome: stream_outcome.unwrap_or(StreamOutcome::InitFailed { @@ -1942,72 +1888,139 @@ impl StreamHandler { } } -/// Result of a successful streaming turn via [`StreamHandler::stream_turn`]. +/// Event yielded by [`StreamHandler::stream_turn`]. /// -/// Carries the accumulated message, token usage, and stop reason, -/// alongside metadata about how the turn was completed (normal streaming -/// vs. non-streaming fallback). +/// The engine drives the handler's turn as a stream of these events: real +/// stream events flow through to observers and the engine's accumulator; +/// retry boundaries and the non-streaming fallback are surfaced as +/// first-class signals so the engine can react (reset state, swap in the +/// fallback message). /// -/// # Example -/// -/// ```rust,ignore -/// let result = handler.stream_turn(client, messages, system, tools, cancel).await?; -/// if result.from_fallback { -/// eprintln!("Warning: fell back to non-streaming"); -/// } -/// println!("Response: {:?}", result.message); -/// ``` +/// See [`StreamHandler::stream_turn`] for the contract. #[derive(Debug, Clone)] #[non_exhaustive] -pub struct StreamTurnResult { - /// The fully accumulated assistant message. - /// - /// On the streaming path, built by the accumulator from every received - /// event — so it carries the full part structure (text blocks, tool - /// calls, images). On the non-streaming fallback path, built from just - /// the first text part of the JSON response, so non-text parts are lost - /// when [`from_fallback`](Self::from_fallback) is `true`. - pub message: Message, - - /// Token counts for this turn, if reported by the provider. - /// - /// Populated on the streaming path from the provider's final usage event. - /// `None` on the non-streaming fallback path (the JSON extraction does - /// not parse usage), and `None` on either path when the provider did not - /// report usage. - pub usage: Option, - - /// Why the model stopped generating. - /// - /// [`EndTurn`](StreamStopReason::EndTurn), - /// [`ToolCall`](StreamStopReason::ToolCall), etc. — parsed from the - /// provider's stop signal. On the non-streaming fallback path, defaults to - /// [`EndTurn`](StreamStopReason::EndTurn) when the response carries no - /// parseable stop reason. - pub stop_reason: StreamStopReason, - - /// Whether the result came from a non-streaming fallback. - /// - /// `false` on the normal streaming path; `true` when streaming exhausted - /// its retries and [`fallback_to_non_streaming`](StreamTimeoutConfig::fallback_to_non_streaming) - /// routed the turn through [`ApiClient::create_message`]. Callers can use - /// this to downgrade trust in the result (the fallback path loses non-text - /// parts and usage) or surface a warning to the user. - pub from_fallback: bool, - - /// Wall-clock time spent on this turn. - /// - /// Measured from the turn's stream-open (or fallback dispatch) to the - /// completed result, inclusive of any retries and backoff sleeps along - /// the way — so it reflects the *true* latency the turn incurred, not - /// just the final successful attempt's duration. - pub elapsed: Duration, +pub enum HandlerEvent { + /// A raw event from the provider's stream (text delta, thinking delta, + /// tool call, etc.). + /// + /// The engine forwards these to observers (`on_text_delta`, + /// `on_thinking_delta`, `text_streamer`) exactly like the inline path, + /// then feeds them to its own [`StreamAccumulator`]. + Stream(StreamEvent), + + /// The handler is starting a new attempt after a retry decision. + /// + /// Fired before the first `Stream` event of attempts 2, 3, … (never on + /// the first attempt). The engine must reset any per-attempt state — + /// including its [`StreamAccumulator`] — so events from the failed + /// attempt are discarded rather than concatenated with the retry's + /// events. Observers that concatenate deltas per turn must reset too. + AttemptReset, + + /// Streaming retries are exhausted and the non-streaming fallback + /// succeeded. + /// + /// Carries the final message and stop reason extracted from the + /// non-streaming + /// [`create_message`](crate::api::ApiClient::create_message) JSON + /// response. The engine should stop accumulating and use these directly — + /// the streaming accumulator's partial state from failed attempts is + /// irrelevant on this path. + /// + /// Always the last event before the stream ends (when the fallback path + /// is taken). + Fallback { + /// The fallback assistant message (text extracted from the JSON + /// response's first text part). + message: Message, + /// Stop reason parsed from the JSON response’s `stop_reason` field, + /// defaulting to [`EndTurn`](StreamStopReason::EndTurn) when absent + /// or unrecognized. + stop_reason: StreamStopReason, + }, } #[cfg(test)] mod tests { use super::*; + /// Test-only result shape matching the old `StreamTurnResult`, used by + /// [`StreamHandler::drive_turn`] to keep existing tests' assertions working. + #[derive(Debug)] + #[allow(dead_code)] + struct DriveResult { + message: Message, + usage: Option, + stop_reason: StreamStopReason, + from_fallback: bool, + } + + impl StreamHandler { + /// Drive `stream_turn` to completion and return the assembled + /// `(Message, Option, StreamStopReason, from_fallback)` tuple — + /// the same shape `stream_turn` used to return directly. Mirrors how the + /// engine consumes the stream: events accumulate, `AttemptReset` + /// discards partial state, `Fallback` short-circuits with the fallback + /// message. + async fn drive_turn( + &self, + client: &C, + request: crate::api::StreamRequest, + cancel: &Arc, + ) -> Result { + let mut stream = self.stream_turn( + client, + request, + crate::structured::RequestOptions::default(), + cancel, + ); + let mut accumulator = StreamAccumulator::new(); + let mut stop_reason = StreamStopReason::EndTurn; + let mut from_fallback = false; + while let Some(item) = stream.next().await { + match item? { + HandlerEvent::Stream(ev) => { + if let StreamEvent::MessageDelta(delta) = &ev + && let Some(ref reason_str) = delta.delta.stop_reason + { + stop_reason = + StreamStopReason::from_api_str(reason_str).unwrap_or(stop_reason); + } + accumulator.process(&ev).map_err(|e| { + StreamHandlerError::StreamFailed(StreamOutcome::InitFailed { + attempts: 1, + last_error: e.to_string(), + }) + })?; + } + HandlerEvent::AttemptReset => { + accumulator = StreamAccumulator::new(); + stop_reason = StreamStopReason::EndTurn; + } + HandlerEvent::Fallback { + message, + stop_reason: fallback_stop_reason, + } => { + from_fallback = true; + return Ok(DriveResult { + message, + usage: None, + stop_reason: fallback_stop_reason, + from_fallback, + }); + } + } + } + let usage = accumulator.usage().copied(); + Ok(DriveResult { + message: accumulator.build(), + usage, + stop_reason, + from_fallback, + }) + } + } + #[test] fn timeout_config_default_values() { let config = StreamTimeoutConfig::default(); @@ -2019,6 +2032,28 @@ mod tests { assert!(config.fallback_to_non_streaming); } + #[test] + fn passthrough_sets_no_resilience_config() { + let h = StreamHandler::passthrough(); + assert_eq!(h.timeout_config().initial_event_timeout, Duration::MAX); + assert_eq!(h.timeout_config().per_event_timeout, Duration::MAX); + assert_eq!(h.timeout_config().total_stream_timeout, Duration::MAX); + assert!(!h.timeout_config().fallback_to_non_streaming); + assert_eq!(h.retry_config().max_retries, 0); + assert_eq!(h.rate_limit_config().max_retries, 0); + assert_eq!(h.rate_limit_config().fallback_after_retries, 0); + } + + #[test] + fn passthrough_default_returns_shared_static() { + let a = StreamHandler::passthrough_default(); + let b = StreamHandler::passthrough_default(); + assert!( + std::ptr::eq(a, b), + "passthrough_default must return the same static" + ); + } + #[test] fn timeout_config_custom_values() { let config = StreamTimeoutConfig { @@ -2166,19 +2201,9 @@ mod tests { #[test] fn error_rate_limit_escalation_display() { - let prior = StreamOutcome::RateLimited { - detail: DetectedRateLimit { - kind: RateLimitKind::RateLimited, - retry_after: Some(Duration::from_secs(5)), - message: "slow down".to_string(), - }, - has_partial_data: false, - events_processed: 0, - }; let err = StreamHandlerError::RateLimitEscalation { attempts: 4, retry_after: Some(Duration::from_secs(5)), - prior, }; let s = err.to_string(); assert!(s.contains("rate-limit escalation"), "got: {s}"); @@ -2395,37 +2420,10 @@ mod tests { assert!(config.validate().is_ok()); } - #[test] - fn stream_turn_result_fields() { - let result = StreamTurnResult { - message: Message::assistant("hello"), - usage: Some(Usage::new(10, 5)), - stop_reason: StreamStopReason::EndTurn, - from_fallback: false, - elapsed: Duration::from_millis(100), - }; - assert!(!result.from_fallback); - assert_eq!(result.stop_reason, StreamStopReason::EndTurn); - assert!(result.usage.is_some()); - } - - #[test] - fn stream_turn_result_fallback_flag() { - let result = StreamTurnResult { - message: Message::assistant("fallback text"), - usage: None, - stop_reason: StreamStopReason::EndTurn, - from_fallback: true, - elapsed: Duration::from_millis(200), - }; - assert!(result.from_fallback); - assert!(result.usage.is_none()); - } - use crate::api::error::ApiError; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, - PartStart, StreamEvent, + PartStart, StreamEvent, Usage, }; fn happy_stream_events() -> Vec> { @@ -2458,142 +2456,6 @@ mod tests { ] } - fn event_stream( - events: Vec>, - ) -> std::pin::Pin< - Box> + Send + 'static>, - > { - Box::pin(futures::stream::iter(events)) - } - - #[tokio::test] - async fn process_events_happy_path() { - let handler = StreamHandler::new(); - let cancel = Arc::new(CancelSignal::new()); - let stream = event_stream(happy_stream_events()); - - let result = handler - .process_events(stream, &cancel, None) - .await - .expect("should succeed"); - - assert!(!result.from_fallback); - assert_eq!(result.stop_reason, StreamStopReason::EndTurn); - assert!(result.elapsed > Duration::ZERO); - } - - #[tokio::test] - async fn process_events_api_error_mid_stream() { - let handler = StreamHandler::new(); - let cancel = Arc::new(CancelSignal::new()); - let events = vec![ - Ok(StreamEvent::MessageStart(MessageStart { - message: MessageMetadata { - id: "msg_test".to_string(), - role: "assistant".to_string(), - model: "test-model".to_string(), - }, - })), - Err(ApiError::api("connection lost")), - ]; - let stream = event_stream(events); - - let err = handler - .process_events(stream, &cancel, None) - .await - .expect_err("should fail on API error"); - - match err { - StreamHandlerError::StreamFailed(outcome) => { - let s = outcome.to_string(); - assert!(s.contains("connection lost"), "unexpected: {s}"); - } - other => panic!("expected StreamFailed, got: {other}"), - } - } - - #[tokio::test] - async fn process_events_total_timeout() { - // Use a very short total timeout to trigger it immediately. - let handler = StreamHandler::new().with_config( - StreamTimeoutConfig { - total_stream_timeout: Duration::from_millis(1), - per_event_timeout: Duration::from_mins(5), - initial_event_timeout: Duration::from_mins(2), - max_consecutive_timeouts: 10, - progress_interval: Duration::from_secs(30), - fallback_to_non_streaming: false, - }, - StreamRetryConfig::default(), - ); - let cancel = Arc::new(CancelSignal::new()); - - // Set a total deadline in the past so it triggers immediately. - let deadline = Some( - Instant::now() - .checked_sub(Duration::from_secs(1)) - .unwrap_or(Instant::now()), - ); - - // Use a stream that never produces events (pending forever). - let pending_stream: std::pin::Pin< - Box> + Send + 'static>, - > = Box::pin(futures::stream::pending()); - - let err = handler - .process_events(pending_stream, &cancel, deadline) - .await - .expect_err("should fail on total timeout"); - - match err { - StreamHandlerError::StreamFailed(StreamOutcome::TotalTimeout { .. }) => {} - other => panic!("expected StreamFailed(TotalTimeout), got: {other}"), - } - } - - #[tokio::test] - async fn process_events_cancelled() { - let handler = StreamHandler::new().with_config( - StreamTimeoutConfig { - per_event_timeout: Duration::from_mins(5), - ..Default::default() - }, - StreamRetryConfig::default(), - ); - let cancel = Arc::new(CancelSignal::new()); - cancel.cancel(); - - let pending_stream: std::pin::Pin< - Box> + Send + 'static>, - > = Box::pin(futures::stream::pending()); - - let err = handler - .process_events(pending_stream, &cancel, None) - .await - .expect_err("should fail on cancellation"); - - assert!( - matches!(err, StreamHandlerError::Cancelled), - "expected Cancelled, got: {err}" - ); - } - - #[tokio::test] - async fn process_events_empty_stream() { - // A stream that ends immediately (None) should produce an - // empty-but-successful result. - let handler = StreamHandler::new(); - let cancel = Arc::new(CancelSignal::new()); - let stream = event_stream(vec![]); - - let result = handler - .process_events(stream, &cancel, None) - .await - .expect("empty stream should succeed"); - - assert!(!result.from_fallback); - } - struct HandlerMock { create_error: Option, create_response: Option, @@ -2628,9 +2490,7 @@ mod tests { fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -2640,9 +2500,7 @@ mod tests { fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + '_>, > { @@ -2670,12 +2528,10 @@ mod tests { let client = HandlerMock::new().with_text_response("fallback works"); let cancel = Arc::new(CancelSignal::new()); - let result = handler + let (message, stop_reason) = handler .fallback_non_streaming( &client, - vec![], - None, - None, + crate::api::StreamRequest::new(vec![]), &cancel, Some(StreamOutcome::InitFailed { last_error: "stream failed".to_string(), @@ -2685,8 +2541,19 @@ mod tests { .await .expect("fallback should succeed"); - assert!(result.from_fallback); - assert_eq!(result.stop_reason, StreamStopReason::EndTurn); + // The fallback returns a Message built from the first text part of the + // non-streaming JSON response, plus the stop_reason from the JSON. + let text: String = message + .parts + .iter() + .filter_map(|p| match p { + crate::stream::MessagePart::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(); + assert!(text.contains("fallback works"), "got: {text:?}"); + // HandlerMock::with_text_response sets stop_reason: "end_turn". + assert_eq!(stop_reason, StreamStopReason::EndTurn); } #[tokio::test] @@ -2703,7 +2570,12 @@ mod tests { cancel.cancel(); let err = handler - .fallback_non_streaming(&client, vec![], None, None, &cancel, None) + .fallback_non_streaming( + &client, + crate::api::StreamRequest::new(vec![]), + &cancel, + None, + ) .await .expect_err("should fail on cancellation"); @@ -2728,9 +2600,7 @@ mod tests { let err = handler .fallback_non_streaming( &client, - vec![], - None, - None, + crate::api::StreamRequest::new(vec![]), &cancel, Some(StreamOutcome::InitFailed { last_error: "stream timeout".to_string(), @@ -2759,6 +2629,133 @@ mod tests { } } + #[tokio::test] + async fn stream_turn_yields_handler_event_stream_per_event() { + // Each raw stream event must arrive as a HandlerEvent::Stream, in + // arrival order, when driving the handler as a stream. + let handler = StreamHandler::new(); + let client = HandlerMock::new().with_text_response("hello"); + let cancel = Arc::new(CancelSignal::new()); + + let mut stream = handler.stream_turn( + &client, + crate::api::StreamRequest::new(vec![]), + crate::structured::RequestOptions::default(), + &cancel, + ); + let mut saw_stream_events = 0; + let mut saw_attempt_reset = false; + let mut saw_fallback = false; + while let Some(item) = stream.next().await { + match item.expect("stream item ok") { + HandlerEvent::Stream(_) => saw_stream_events += 1, + HandlerEvent::AttemptReset => saw_attempt_reset = true, + HandlerEvent::Fallback { .. } => saw_fallback = true, + } + } + assert!(saw_stream_events > 0, "should yield Stream events"); + assert!(!saw_attempt_reset, "happy path must not emit AttemptReset"); + assert!(!saw_fallback, "happy path must not emit Fallback"); + } + + /// Mock that fails its first streaming attempt with a transport error, + /// then succeeds on the second. Used by the AttemptReset test to verify + /// the handler emits `AttemptReset` before the retried attempt's events. + struct RetryingMock { + attempts: Arc, + } + + impl ApiClient for RetryingMock { + fn model(&self) -> String { + "retry-test".to_string() + } + fn base_url(&self) -> String { + "retry-test".to_string() + } + fn set_model(&self, _: &str) -> bool { + false + } + fn stream_messages( + &self, + request: crate::api::StreamRequest, + ) -> Pin> + Send + 'static>> { + self.stream_messages_with_options(request, crate::structured::RequestOptions::default()) + } + fn create_message( + &self, + request: crate::api::StreamRequest, + ) -> Pin< + Box> + Send + '_>, + > { + self.create_message_with_options(request, crate::structured::RequestOptions::default()) + } + fn stream_messages_with_options( + &self, + _request: crate::api::StreamRequest, + _options: crate::structured::RequestOptions, + ) -> Pin> + Send + 'static>> { + use std::sync::atomic::Ordering; + let n = self.attempts.fetch_add(1, Ordering::SeqCst); + if n == 0 { + // First attempt: one event then transport error. + Box::pin(futures::stream::iter(vec![ + Ok(StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: String::new(), + role: "assistant".into(), + model: String::new(), + }, + })), + Err(ApiError::api("transient")), + ])) + } else { + // Second attempt: clean happy-path events. + Box::pin(futures::stream::iter(happy_stream_events())) + } + } + fn create_message_with_options( + &self, + _request: crate::api::StreamRequest, + _options: crate::structured::RequestOptions, + ) -> Pin< + Box> + Send + '_>, + > { + Box::pin(async { Ok(serde_json::Value::Null) }) + } + fn extract_structured(&self, _: &serde_json::Value) -> serde_json::Value { + serde_json::Value::Null + } + } + + #[tokio::test] + async fn stream_turn_yields_attempt_reset_on_retry() { + // A retried transport failure must emit AttemptReset before the + // retried attempt's events. Engine uses this to discard partial state. + use std::sync::atomic::AtomicUsize; + + let attempts = Arc::new(AtomicUsize::new(0)); + let handler = StreamHandler::new(); + let client = RetryingMock { attempts }; + let cancel = Arc::new(CancelSignal::new()); + + let mut stream = handler.stream_turn( + &client, + crate::api::StreamRequest::new(vec![]), + crate::structured::RequestOptions::default(), + &cancel, + ); + let mut saw_attempt_reset = false; + while let Some(item) = stream.next().await { + if let HandlerEvent::AttemptReset = item.expect("stream item ok") { + saw_attempt_reset = true; + } + } + assert!( + saw_attempt_reset, + "second attempt must be preceded by AttemptReset" + ); + } + #[tokio::test] async fn stream_turn_happy_path() { let handler = StreamHandler::new(); @@ -2766,7 +2763,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let result = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("stream_turn should succeed"); @@ -2782,7 +2779,7 @@ mod tests { cancel.cancel(); let err = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("should fail on cancellation"); @@ -2811,9 +2808,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -2823,9 +2818,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -2852,7 +2845,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("should fail when streaming errors and fallback is disabled"); @@ -2866,6 +2859,126 @@ mod tests { } } + /// Mock that always fails streaming but succeeds on non-streaming + /// `create_message`. Used by the fallback regression test to verify the + /// handler yields `HandlerEvent::Fallback` with the message and stop_reason + /// extracted from the JSON response. + struct StreamingFailingFallbackMock; + impl ApiClient for StreamingFailingFallbackMock { + fn model(&self) -> String { + "fallback-test".to_string() + } + fn base_url(&self) -> String { + "fallback-test".to_string() + } + fn set_model(&self, _: &str) -> bool { + false + } + fn stream_messages( + &self, + _request: crate::api::StreamRequest, + ) -> Pin> + Send + 'static>> { + // Always fail streaming — forces the handler to retry, then + // fall back to create_message. + Box::pin(futures::stream::once(async { + Err(ApiError::api("stream down")) + })) + } + fn create_message( + &self, + _request: crate::api::StreamRequest, + ) -> Pin< + Box> + Send + '_>, + > { + Box::pin(async { + Ok(serde_json::json!({ + "content": [{"type": "text", "text": "fallback answer"}], + "stop_reason": "max_tokens", + })) + }) + } + fn stream_messages_with_options( + &self, + _request: crate::api::StreamRequest, + _options: crate::structured::RequestOptions, + ) -> Pin> + Send + 'static>> { + Box::pin(futures::stream::once(async { + Err(ApiError::api("stream down")) + })) + } + fn create_message_with_options( + &self, + _request: crate::api::StreamRequest, + _options: crate::structured::RequestOptions, + ) -> Pin< + Box> + Send + '_>, + > { + Box::pin(async { + Ok(serde_json::json!({ + "content": [{"type": "text", "text": "fallback answer"}], + "stop_reason": "max_tokens", + })) + }) + } + fn extract_structured(&self, _: &serde_json::Value) -> serde_json::Value { + serde_json::Value::Null + } + } + + #[tokio::test] + async fn drive_turn_returns_fallback_message_and_stop_reason() { + // When streaming fails after retries and fallback is enabled, the + // engine should see the fallback message and the stop_reason from the + // non-streaming JSON response (not the streaming accumulator's stale + // values). Regression test for an earlier bug where Fallback dropped + // stop_reason. + let handler = StreamHandler::new().with_config( + StreamTimeoutConfig { + fallback_to_non_streaming: true, + ..Default::default() + }, + StreamRetryConfig { + max_retries: 0, + ..Default::default() + }, + ); + let client = StreamingFailingFallbackMock; + let cancel = Arc::new(CancelSignal::new()); + + let result = handler + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) + .await + .expect("fallback should succeed"); + + assert!( + result.from_fallback, + "result should be marked from_fallback" + ); + // Stop reason comes from the JSON's "max_tokens", not the streaming + // default "end_turn". + assert_eq!( + result.stop_reason, + StreamStopReason::MaxTokens, + "fallback stop_reason must come from the JSON response" + ); + let text: String = result + .message + .parts + .iter() + .filter_map(|p| match p { + crate::stream::MessagePart::Text { text } => Some(text.clone()), + _ => None, + }) + .collect(); + assert!( + text.contains("fallback answer"), + "fallback message text, got {text:?}" + ); + // Usage is None on the fallback path (non-streaming JSON doesn't + // reliably carry token counts). + assert!(result.usage.is_none()); + } + #[test] fn rate_limit_config_default_values() { let cfg = RateLimitConfig::default(); @@ -3085,69 +3198,6 @@ mod tests { assert_eq!(count, 3); } - #[test] - fn apply_event_extracts_stop_reason_and_accumulates() { - let mut accumulator = StreamAccumulator::new(); - let mut stop_reason = StreamStopReason::EndTurn; - - StreamHandler::apply_event( - &StreamEvent::MessageStart(MessageStart { - message: MessageMetadata { - id: "msg".to_string(), - role: "assistant".to_string(), - model: "m".to_string(), - }, - }), - &mut accumulator, - &mut stop_reason, - ) - .expect("MessageStart should accumulate"); - StreamHandler::apply_event( - &StreamEvent::PartStart(PartStart { - index: 0, - part: Some(crate::stream::MessagePart::text("")), - }), - &mut accumulator, - &mut stop_reason, - ) - .expect("PartStart should accumulate"); - - let result = StreamHandler::apply_event( - &StreamEvent::MessageDelta(MessageDelta { - delta: MessageDeltaPayload { - stop_reason: Some("end_turn".to_string()), - }, - usage: None, - }), - &mut accumulator, - &mut stop_reason, - ); - assert!(result.is_ok()); - assert_eq!(stop_reason, StreamStopReason::EndTurn); - - let message = accumulator.build(); - assert_eq!(message.role, crate::message::Role::Assistant); - } - - #[test] - fn apply_event_unknown_stop_reason_keeps_prior() { - let mut accumulator = StreamAccumulator::new(); - let mut stop_reason = StreamStopReason::ToolCall; - - StreamHandler::apply_event( - &StreamEvent::MessageDelta(MessageDelta { - delta: MessageDeltaPayload { - stop_reason: Some("not_a_real_reason".to_string()), - }, - usage: None, - }), - &mut accumulator, - &mut stop_reason, - ) - .expect("apply_event should not error on an out-of-context delta"); - assert_eq!(stop_reason, StreamStopReason::ToolCall); - } - #[test] fn detected_rate_limit_from_structured_variant() { let err = ApiError::RateLimit { @@ -3209,117 +3259,6 @@ mod tests { } } - #[tokio::test] - async fn process_events_rate_limit_mid_stream() { - let handler = StreamHandler::new(); - let cancel = Arc::new(CancelSignal::new()); - let events = vec![ - Ok(StreamEvent::MessageStart(MessageStart { - message: MessageMetadata { - id: "msg_test".to_string(), - role: "assistant".to_string(), - model: "test-model".to_string(), - }, - })), - Err(ApiError::RateLimit { - retry_after: Some(Duration::from_secs(4)), - message: "slow down".into(), - }), - ]; - let stream = event_stream(events); - - let err = handler - .process_events(stream, &cancel, None) - .await - .expect_err("rate-limit error should fail the stream"); - - match err { - StreamHandlerError::StreamFailed(StreamOutcome::RateLimited { - detail, - has_partial_data, - .. - }) => { - assert_eq!(detail.kind, RateLimitKind::RateLimited); - assert_eq!(detail.retry_after, Some(Duration::from_secs(4))); - assert!(!has_partial_data, "MessageStart carries no parts"); - } - other => panic!("expected RateLimited outcome, got: {other:?}"), - } - } - - #[tokio::test] - async fn process_events_rate_limit_with_partial_data() { - let handler = StreamHandler::new(); - let cancel = Arc::new(CancelSignal::new()); - let events = vec![ - Ok(StreamEvent::MessageStart(MessageStart { - message: MessageMetadata { - id: "msg_test".to_string(), - role: "assistant".to_string(), - model: "test-model".to_string(), - }, - })), - Ok(StreamEvent::PartStart(PartStart { - index: 0, - part: Some(crate::stream::MessagePart::text("partial")), - })), - Ok(StreamEvent::IndexedDelta(IndexedDelta { - index: 0, - delta: DeltaPart::Text { - text: " response".into(), - }, - })), - Ok(StreamEvent::PartStop), - Err(ApiError::http_with_status(503, "overloaded")), - ]; - let stream = event_stream(events); - - let err = handler - .process_events(stream, &cancel, None) - .await - .expect_err("503 should fail the stream"); - - match err { - StreamHandlerError::StreamFailed(StreamOutcome::RateLimited { - detail, - has_partial_data, - events_processed, - }) => { - assert_eq!(detail.kind, RateLimitKind::Overloaded); - assert!(has_partial_data, "PartStart+delta should count as partial"); - assert!(events_processed >= 2, "at least 2 events processed"); - } - other => panic!("expected RateLimited outcome, got: {other:?}"), - } - } - - #[tokio::test] - async fn process_events_non_rate_error_still_init_failed() { - let handler = StreamHandler::new(); - let cancel = Arc::new(CancelSignal::new()); - let events: Vec> = vec![ - Ok(StreamEvent::MessageStart(MessageStart { - message: MessageMetadata { - id: "msg_test".to_string(), - role: "assistant".to_string(), - model: "test-model".to_string(), - }, - })), - Err(ApiError::api("connection reset")), - ]; - let stream = event_stream(events); - - let err = handler - .process_events(stream, &cancel, None) - .await - .expect_err("should fail"); - - match err { - StreamHandlerError::StreamFailed(StreamOutcome::InitFailed { .. }) => {} - other => panic!("expected InitFailed for non-rate error, got: {other:?}"), - } - } - #[test] fn stream_outcome_rate_limited_display() { let outcome = StreamOutcome::RateLimited { @@ -3370,9 +3309,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3380,9 +3317,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + '_>, > { @@ -3511,7 +3446,7 @@ mod tests { let start = Instant::now(); for _ in 0..3 { handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("turn should succeed"); } @@ -3536,7 +3471,7 @@ mod tests { // First turn consumes the only token. let cancel = Arc::new(CancelSignal::new()); handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("first turn should succeed"); @@ -3545,7 +3480,7 @@ mod tests { cancel2.cancel(); let start = Instant::now(); let err = handler - .stream_turn(&client, vec![], None, None, &cancel2) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel2) .await .expect_err("should be cancelled, not hang for 60s"); let elapsed = start.elapsed(); @@ -3572,14 +3507,14 @@ mod tests { // First turn consumes the only token. handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("first turn should succeed"); // Second turn: bucket empty, wait capped at 50ms, then proceeds. let start = Instant::now(); let result = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await; let elapsed = start.elapsed(); @@ -3606,9 +3541,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3626,9 +3559,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3662,7 +3593,7 @@ mod tests { let start = Instant::now(); let result = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("second attempt should succeed"); let elapsed = start.elapsed(); @@ -3685,9 +3616,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3700,9 +3629,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3724,25 +3651,17 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("should escalate, not succeed"); match err { StreamHandlerError::RateLimitEscalation { attempts, retry_after, - prior, } => { // fallback_after_retries == 2, so escalation fires on the 3rd hit. assert_eq!(attempts, 3); assert_eq!(retry_after, Some(Duration::from_millis(1))); - match prior { - StreamOutcome::RateLimited { detail, .. } => { - assert_eq!(detail.kind, RateLimitKind::RateLimited); - assert_eq!(detail.retry_after, Some(Duration::from_millis(1))); - } - other => panic!("prior should be RateLimited, got {other:?}"), - } } other => panic!("expected RateLimitEscalation, got {other:?}"), } @@ -3767,9 +3686,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3786,9 +3703,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3825,7 +3740,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("should escalate after the rate-limit budget, not fall through"); match err { @@ -3849,9 +3764,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3864,9 +3777,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3889,7 +3800,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("hard-stop should fail the turn"); match err { @@ -3917,9 +3828,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -3937,9 +3846,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -3963,7 +3870,7 @@ mod tests { }; let cancel = Arc::new(CancelSignal::new()); handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("first call should succeed after one rate-limit retry"); @@ -3973,7 +3880,7 @@ mod tests { attempts: AtomicUsize::new(0), }; handler - .stream_turn(&client2, vec![], None, None, &cancel) + .drive_turn(&client2, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect("second call should not see leaked rate-limit state"); } @@ -3989,9 +3896,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -4001,9 +3906,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -4030,7 +3933,7 @@ mod tests { let cancel = Arc::new(CancelSignal::new()); let err = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("transport errors should fail"); assert!( @@ -4050,9 +3953,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -4065,9 +3966,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -4105,7 +4004,7 @@ mod tests { let start = Instant::now(); let err = handler - .stream_turn(&client, vec![], None, None, &cancel) + .drive_turn(&client, crate::api::StreamRequest::new(vec![]), &cancel) .await .expect_err("tight total timeout should fail the turn"); let elapsed = start.elapsed(); @@ -4131,9 +4030,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box> + Send + 'static>, > { @@ -4143,9 +4040,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> std::pin::Pin< Box< dyn std::future::Future> @@ -4174,7 +4069,11 @@ mod tests { let start = Instant::now(); let err = handler - .stream_turn(&AlwaysFailingMock, vec![], None, None, &cancel) + .drive_turn( + &AlwaysFailingMock, + crate::api::StreamRequest::new(vec![]), + &cancel, + ) .await .expect_err("should return Cancelled, not hang for 60s"); let elapsed = start.elapsed(); diff --git a/src/structured.rs b/src/structured.rs index 9da6263..790e5f0 100644 --- a/src/structured.rs +++ b/src/structured.rs @@ -584,8 +584,13 @@ pub async fn request_structured, ) -> Result { let opts = RequestOptions::new().response_format(ResponseFormat::from_type::()); + let request = crate::api::StreamRequest { + messages, + system, + tools: None, + }; let raw = client - .create_message_with_options(messages, system, None, opts) + .create_message_with_options(request, opts) .await .map_err(StructuredError::Api)?; let value = client.extract_structured(&raw); @@ -729,9 +734,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream< @@ -744,9 +747,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn Future> @@ -762,9 +763,8 @@ mod tests { async fn default_client_rejects_response_format() { let client = PlainMockClient; let opts = RequestOptions::new().response_format(ResponseFormat::from_type::()); - let result = client - .create_message_with_options(vec![], None, None, opts) - .await; + let request = crate::api::StreamRequest::new(vec![]); + let result = client.create_message_with_options(request, opts).await; assert!( result.is_err(), "client without structured-output support should reject response_format" @@ -780,9 +780,8 @@ mod tests { async fn default_client_delegates_empty_options() { let client = PlainMockClient; let opts = RequestOptions::new(); - let result = client - .create_message_with_options(vec![], None, None, opts) - .await; + let request = crate::api::StreamRequest::new(vec![]); + let result = client.create_message_with_options(request, opts).await; assert!(result.is_ok(), "empty options should delegate normally"); } @@ -793,9 +792,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream< @@ -808,9 +805,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn Future> @@ -822,9 +817,7 @@ mod tests { } fn create_message_with_options( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, _options: RequestOptions, ) -> Pin< Box< @@ -858,9 +851,7 @@ mod tests { } fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn futures::Stream< @@ -873,9 +864,7 @@ mod tests { } fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin< Box< dyn Future> @@ -887,9 +876,7 @@ mod tests { } fn create_message_with_options( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, _options: RequestOptions, ) -> Pin< Box< diff --git a/src/testing.rs b/src/testing.rs index 245d1bf..851e172 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -635,20 +635,18 @@ impl ApiClient for MockApiClient { /// /// ```rust /// # tokio::runtime::Runtime::new().unwrap().block_on(async { - /// use loopctl::api::ApiClient; + /// use loopctl::api::{ApiClient, StreamRequest}; /// use loopctl::testing::MockApiClient; /// /// let client = MockApiClient::new("test-model").with_text_response("Hi!"); - /// let stream = client.stream_messages(vec![], None, None); + /// let stream = client.stream_messages(StreamRequest::new(vec![])); /// let events: Vec<_> = futures::StreamExt::collect(stream).await; /// assert!(events.len() >= 4); /// # }); /// ``` fn stream_messages( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + 'static>> { if let Some(ref err) = self.error { let err = err.clone(); @@ -715,26 +713,23 @@ impl ApiClient for MockApiClient { /// future resolves to an [`ApiError`] instead, bypassing the /// response queue entirely. /// - /// The `_messages`, `_system`, and `_tools` parameters are accepted - /// for trait compatibility but ignored. + /// The `_request` parameter is accepted for trait compatibility but ignored. /// /// # Example /// /// ```rust /// # tokio::runtime::Runtime::new().unwrap().block_on(async { - /// use loopctl::api::ApiClient; + /// use loopctl::api::{ApiClient, StreamRequest}; /// use loopctl::testing::MockApiClient; /// /// let client = MockApiClient::new("test-model").with_text_response("Hi!"); - /// let result = client.create_message(vec![], None, None).await; + /// let result = client.create_message(StreamRequest::new(vec![])).await; /// assert_eq!(result.unwrap()["content"][0]["text"], "Hi!"); /// # }); /// ``` fn create_message( &self, - _messages: Vec, - _system: Option, - _tools: Option>, + _request: crate::api::StreamRequest, ) -> Pin> + Send + '_>> { if let Some(ref err) = self.error { let err = err.clone(); @@ -1451,7 +1446,11 @@ mod tests { #[tokio::test] async fn test_mock_client_default_response() { let client = MockApiClient::new("test-model"); - let stream = client.stream_messages(vec![Message::user("Hi")], None, None); + let stream = client.stream_messages(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }); let events: Vec<_> = stream.collect().await; assert!(events.len() >= 4); assert!(events[0].is_ok()); @@ -1460,7 +1459,11 @@ mod tests { #[tokio::test] async fn test_mock_client_custom_text() { let client = MockApiClient::new("test-model").with_text_response("Custom response"); - let stream = client.stream_messages(vec![Message::user("Hi")], None, None); + let stream = client.stream_messages(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }); let events: Vec<_> = stream.collect().await; let has_text = events.iter().any(|e| { if let Ok(StreamEvent::IndexedDelta(delta)) = e @@ -1480,7 +1483,11 @@ mod tests { "echo", json!({"message": "hi"}), ); - let stream = client.stream_messages(vec![Message::user("Hi")], None, None); + let stream = client.stream_messages(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }); let events: Vec<_> = stream.collect().await; let has_tool_use = events.iter().any(|e| { @@ -1506,7 +1513,11 @@ mod tests { #[tokio::test] async fn test_mock_client_error() { let client = MockApiClient::new("test-model").with_error("API error"); - let stream = client.stream_messages(vec![Message::user("Hi")], None, None); + let stream = client.stream_messages(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }); let events: Vec<_> = stream.collect().await; assert_eq!(events.len(), 1); assert!(events[0].is_err()); @@ -1527,7 +1538,11 @@ mod tests { }, ]); - let stream1 = client.stream_messages(vec![Message::user("Hi")], None, None); + let stream1 = client.stream_messages(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }); let events1: Vec<_> = stream1.collect().await; let has_first = events1.iter().any(|e| { if let Ok(StreamEvent::IndexedDelta(delta)) = e @@ -1539,7 +1554,11 @@ mod tests { }); assert!(has_first); - let stream2 = client.stream_messages(vec![Message::user("Hi")], None, None); + let stream2 = client.stream_messages(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }); let events2: Vec<_> = stream2.collect().await; let has_second = events2.iter().any(|e| { if let Ok(StreamEvent::IndexedDelta(delta)) = e @@ -1556,7 +1575,11 @@ mod tests { async fn test_mock_client_create_message() { let client = MockApiClient::new("test-model").with_text_response("Hello!"); let result = client - .create_message(vec![Message::user("Hi")], None, None) + .create_message(crate::api::StreamRequest { + messages: vec![Message::user("Hi")], + system: None, + tools: None, + }) .await; assert!(result.is_ok()); let json = result.unwrap();