From 146a780ceddd16273673d1bb056d26b6be9d5b14 Mon Sep 17 00:00:00 2001 From: Vincenzo Palazzo Date: Wed, 29 Jul 2026 10:06:24 +0200 Subject: [PATCH 1/6] fix(providers): retry streaming requests dropped before first body bytes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The with_retry scope in the OpenAI and OpenAI-compatible providers ended the moment response headers arrived, so an SSE connection cut before delivering any body (reqwest's "error decoding response body") surfaced as a mid-stream item error where retry is structurally unreachable, aborting the whole turn. Await the first body chunk inside the with_retry closure and prepend it to the reconstructed body stream, so pre-first-token failures — the common case with long-lived SSE connections — are transparently retried with the existing backoff. should_retry(NetworkError) is no longer dead code for this failure mode. Adds a regression test with a raw TCP server that drops the first connection after 200 headers and serves a full SSE response on retry. --- Cargo.lock | 1 + crates/goose-providers/Cargo.toml | 1 + crates/goose-providers/src/openai.rs | 208 ++++++++++++++++-- .../goose-providers/src/openai_compatible.rs | 73 +++++- 4 files changed, 262 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ba9bc3b1fd4f..10596e28c468 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5225,6 +5225,7 @@ dependencies = [ "anyhow", "async-stream", "async-trait", + "bytes", "chrono", "env-lock", "futures", diff --git a/crates/goose-providers/Cargo.toml b/crates/goose-providers/Cargo.toml index ab231235b947..7eadc676ff3d 100644 --- a/crates/goose-providers/Cargo.toml +++ b/crates/goose-providers/Cargo.toml @@ -54,6 +54,7 @@ pkcs1 = { version = "0.7.5", default-features = false, features = ["pkcs8", "std pkcs8 = { version = "0.10", default-features = false, features = ["alloc", "std"], optional = true } sec1 = { version = "0.7", default-features = false, features = ["der", "pkcs8", "std"], optional = true } include_dir = { workspace = true } +bytes = { workspace = true } [dev-dependencies] test-case = { workspace = true } diff --git a/crates/goose-providers/src/openai.rs b/crates/goose-providers/src/openai.rs index b256ca989da8..776e7ea0d609 100644 --- a/crates/goose-providers/src/openai.rs +++ b/crates/goose-providers/src/openai.rs @@ -17,7 +17,8 @@ use crate::formats::openai_responses::{ }; use crate::images::ImageFormat; use crate::openai_compatible::{ - handle_response_openai_compat, handle_status, stream_openai_compat, stream_responses_compat, + first_body_chunk, handle_response_openai_compat, handle_status, + stream_openai_compat_with_prefix, stream_responses_compat_with_prefix, }; use crate::request_log::{start_log, LoggerHandleExt}; use crate::thinking::ThinkingEffort; @@ -301,9 +302,11 @@ impl OpenAiProvider { payload: serde_json::Value, ) -> Result { let mut log = start_log(model_config, &payload)?; - let response = self + // For streaming, pre-read the first body chunk inside the retry scope so a + // 200 OK whose connection dies before any body bytes still retries. + let (response, first_chunk) = self .with_retry(|| async { - handle_status( + let mut response = handle_status( self.api_client .request(&Self::map_base_path( &self.base_path, @@ -315,14 +318,20 @@ impl OpenAiProvider { .response_post(&payload) .await?, ) - .await + .await?; + let first_chunk = if self.supports_streaming { + Some(first_body_chunk(&mut response).await?) + } else { + None + }; + Ok((response, first_chunk)) }) .await .inspect_err(|e| { let _ = log.error(e); })?; if self.supports_streaming { - stream_responses_compat(response, log) + stream_responses_compat_with_prefix(response, first_chunk, log) } else { let json: serde_json::Value = response.json().await.map_err(|e| { ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e)) @@ -792,16 +801,25 @@ impl Provider for OpenAiProvider { let payload = self.sanitize_request_for_compat(payload, model_config); let mut log = start_log(model_config, &payload)?; - let response = self + // For streaming, pre-read the first body chunk inside the retry scope so a + // 200 OK whose connection dies before any body bytes still retries. + let (response, first_chunk) = self .with_retry(|| async { - let resp = self - .api_client - .request(&self.base_path) - .model_headers(model_config)? - .streaming(self.supports_streaming) - .response_post(&payload) - .await?; - handle_status(resp).await + let mut response = handle_status( + self.api_client + .request(&self.base_path) + .model_headers(model_config)? + .streaming(self.supports_streaming) + .response_post(&payload) + .await?, + ) + .await?; + let first_chunk = if self.supports_streaming { + Some(first_body_chunk(&mut response).await?) + } else { + None + }; + Ok((response, first_chunk)) }) .await .inspect_err(|e| { @@ -809,7 +827,7 @@ impl Provider for OpenAiProvider { })?; if self.supports_streaming { - stream_openai_compat(response, log) + stream_openai_compat_with_prefix(response, first_chunk, log) } else { let json: serde_json::Value = response.json().await.map_err(|e| { ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e)) @@ -1702,4 +1720,164 @@ mod tests { assert_eq!(payload["stream_options"], json!({"include_usage": true})); assert_eq!(payload["messages"].as_array().unwrap().len(), 1); } + + #[tokio::test] + async fn stream_retries_when_connection_drops_before_first_body_bytes() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + use tokio_stream::StreamExt; + + async fn read_request(socket: &mut TcpStream) { + let mut buf = Vec::new(); + let mut chunk = [0u8; 4096]; + loop { + let n = socket.read(&mut chunk).await.unwrap(); + if n == 0 { + break; + } + buf.extend_from_slice(&chunk[..n]); + if buf.windows(4).any(|w| w == b"\r\n\r\n") { + break; + } + } + } + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let server = tokio::spawn(async move { + // First attempt: 200 headers promise a 1024-byte body, then the + // connection drops without delivering any of it — consuming the + // body fails with reqwest's "error decoding response body". + let (mut socket, _) = listener.accept().await.unwrap(); + read_request(&mut socket).await; + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: 1024\r\n\r\n", + ) + .await + .unwrap(); + socket.shutdown().await.unwrap(); + + // Retried attempt: a complete SSE response. + let (mut socket, _) = listener.accept().await.unwrap(); + read_request(&mut socket).await; + let body = concat!( + "data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1,", + "\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,", + "\"delta\":{\"role\":\"assistant\",\"content\":\"hello\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n", + ); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{}", + body.len(), + body + ); + socket.write_all(response.as_bytes()).await.unwrap(); + socket.shutdown().await.unwrap(); + }); + + let provider = OpenAiProvider { + api_client: ApiClient::new_with_tls(format!("http://{addr}"), AuthMethod::NoAuth, None) + .unwrap(), + ..make_provider("openai") + }; + + let model_config = ModelConfig::new("gpt-4o"); + let messages = vec![Message::user().with_text("hi")]; + let mut stream = provider + .stream(&model_config, "system", &messages, &[]) + .await + .expect("dropped pre-body connection should be transparently retried"); + + let mut text = String::new(); + while let Some(item) = stream.next().await { + let (message, _usage) = item.expect("stream item should be ok"); + if let Some(message) = message { + text.push_str(&message.as_concat_text()); + } + } + server.await.unwrap(); + assert_eq!(text, "hello"); + } + + #[tokio::test] + async fn stream_retries_when_connection_closes_cleanly_before_first_body_bytes() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + use tokio_stream::StreamExt; + + async fn read_request(socket: &mut TcpStream) { + let mut buf = Vec::new(); + let mut chunk = [0u8; 4096]; + loop { + let n = socket.read(&mut chunk).await.unwrap(); + if n == 0 { + break; + } + buf.extend_from_slice(&chunk[..n]); + if buf.windows(4).any(|w| w == b"\r\n\r\n") { + break; + } + } + } + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let server = tokio::spawn(async move { + // First attempt: connection-close framing with zero body bytes — + // a clean EOF, so chunk() returns Ok(None) rather than an error. + // An empty body is not a valid streaming response and must be + // retried like any other pre-first-token failure. + let (mut socket, _) = listener.accept().await.unwrap(); + read_request(&mut socket).await; + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n", + ) + .await + .unwrap(); + socket.shutdown().await.unwrap(); + + let (mut socket, _) = listener.accept().await.unwrap(); + read_request(&mut socket).await; + let body = concat!( + "data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":1,", + "\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,", + "\"delta\":{\"role\":\"assistant\",\"content\":\"hello\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: [DONE]\n\n", + ); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\n\r\n{}", + body.len(), + body + ); + socket.write_all(response.as_bytes()).await.unwrap(); + socket.shutdown().await.unwrap(); + }); + + let provider = OpenAiProvider { + api_client: ApiClient::new_with_tls(format!("http://{addr}"), AuthMethod::NoAuth, None) + .unwrap(), + ..make_provider("openai") + }; + + let model_config = ModelConfig::new("gpt-4o"); + let messages = vec![Message::user().with_text("hi")]; + let mut stream = provider + .stream(&model_config, "system", &messages, &[]) + .await + .expect("clean pre-body EOF should be transparently retried"); + + let mut text = String::new(); + while let Some(item) = stream.next().await { + let (message, _usage) = item.expect("stream item should be ok"); + if let Some(message) = message { + text.push_str(&message.as_concat_text()); + } + } + server.await.unwrap(); + assert_eq!(text, "hello"); + } } diff --git a/crates/goose-providers/src/openai_compatible.rs b/crates/goose-providers/src/openai_compatible.rs index f40d756da26c..ebb965860ce2 100644 --- a/crates/goose-providers/src/openai_compatible.rs +++ b/crates/goose-providers/src/openai_compatible.rs @@ -2,11 +2,14 @@ use crate::conversation::token_usage::{CostSource, ProviderUsage}; use crate::images::ImageFormat; use anyhow::Error; use async_stream::try_stream; +use bytes::Bytes; +use futures::Stream; use futures::TryStreamExt; use reqwest::Response; #[cfg(test)] use reqwest::StatusCode; use serde_json::Value; +use std::pin::Pin; use tokio::pin; use tokio_stream::StreamExt; use tokio_util::codec::{FramedRead, LinesCodec}; @@ -106,9 +109,11 @@ impl OpenAiCompatibleProvider { ) -> Result { let mut log = start_log(model_config, &payload)?; let path = format!("{}chat/completions", self.completions_prefix); - let response = self + // For streaming, pre-read the first body chunk inside the retry scope so a + // 200 OK whose connection dies before any body bytes still retries. + let (response, first_chunk) = self .with_retry(|| async { - handle_status( + let mut response = handle_status( self.api_client .request(&path) .model_headers(model_config)? @@ -116,14 +121,20 @@ impl OpenAiCompatibleProvider { .response_post(&payload) .await?, ) - .await + .await?; + let first_chunk = if self.supports_streaming { + Some(first_body_chunk(&mut response).await?) + } else { + None + }; + Ok((response, first_chunk)) }) .await .inspect_err(|e| { let _ = log.error(e); })?; if self.supports_streaming { - stream_openai_compat(response, log) + stream_openai_compat_with_prefix(response, first_chunk, log) } else { let json = response.json().await.map_err(|e| { ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e)) @@ -232,11 +243,53 @@ pub use super::http_status::{ // Legacy alias kept for callers that haven't migrated their import path yet. pub use super::http_status::handle_response as handle_response_openai_compat; +/// Reads the first chunk of a streaming response body. +/// +/// A `200 OK` whose connection dies before delivering any body bytes only +/// fails here, while the caller is still inside its `with_retry` scope — +/// afterwards the failure would surface as a mid-stream item error where +/// retry is unreachable. A clean EOF is also an error: an empty body is +/// never a valid answer to a streaming request, and `Ok(None)` here would +/// otherwise complete the retry scope with a stream that yields nothing. +pub async fn first_body_chunk(response: &mut Response) -> Result { + response + .chunk() + .await + .map_err(ProviderError::stream_decode_error)? + .ok_or_else(|| { + ProviderError::stream_decode_error("stream ended before the first body chunk") + }) +} + +fn body_stream_with_prefix( + response: Response, + first_chunk: Option, +) -> Pin> + Send>> { + Box::pin(async_stream::stream! { + if let Some(chunk) = first_chunk { + yield Ok(chunk); + } + let body = response.bytes_stream().map_err(std::io::Error::other); + tokio::pin!(body); + while let Some(item) = body.next().await { + yield item; + } + }) +} + pub fn stream_openai_compat( response: Response, + log: Option>, +) -> Result { + stream_openai_compat_with_prefix(response, None, log) +} + +pub fn stream_openai_compat_with_prefix( + response: Response, + first_chunk: Option, mut log: Option>, ) -> Result { - let stream = response.bytes_stream().map_err(std::io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); @@ -258,9 +311,17 @@ pub fn stream_openai_compat( pub fn stream_responses_compat( response: Response, + log: Option>, +) -> Result { + stream_responses_compat_with_prefix(response, None, log) +} + +pub fn stream_responses_compat_with_prefix( + response: Response, + first_chunk: Option, mut log: Option>, ) -> Result { - let stream = response.bytes_stream().map_err(std::io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); From 0b78310eac26642f91e492598437c30ca445355f Mon Sep 17 00:00:00 2001 From: Vincenzo Palazzo Date: Mon, 10 Aug 2026 20:23:27 +0200 Subject: [PATCH 2/6] fix(providers): pre-read first stream chunk in remaining providers Apply the same first-body-chunk-inside-retry pattern to Databricks, OpenRouter, GitHub Copilot, NanoGPT, Tetrate, and Bedrock mantle so a 200 OK that dies before any body bytes is still retried. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/goose-providers/src/databricks.rs | 93 ++++++++++++--------- crates/goose-providers/src/databricks_v2.rs | 51 ++++++----- crates/goose/Cargo.toml | 1 + crates/goose/src/providers/bedrock.rs | 14 ++-- crates/goose/src/providers/githubcopilot.rs | 33 +++++--- crates/goose/src/providers/nanogpt.rs | 25 +++--- crates/goose/src/providers/openrouter.rs | 52 ++++++------ crates/goose/src/providers/tetrate.rs | 35 ++++---- 8 files changed, 174 insertions(+), 130 deletions(-) diff --git a/crates/goose-providers/src/databricks.rs b/crates/goose-providers/src/databricks.rs index eb8a4380aa56..75b4e1e2386c 100644 --- a/crates/goose-providers/src/databricks.rs +++ b/crates/goose-providers/src/databricks.rs @@ -24,8 +24,8 @@ pub use crate::formats::databricks::DATABRICKS_PROVIDER_NAME; use crate::formats::openai_responses::create_responses_request; use crate::model::ModelConfig; use crate::openai_compatible::{ - handle_status, map_http_error_to_provider_error, sanitize_url, stream_openai_compat, - stream_responses_compat, + first_body_chunk, handle_status, map_http_error_to_provider_error, sanitize_url, + stream_openai_compat_with_prefix, stream_responses_compat_with_prefix, }; use crate::request_log::{start_log, LoggerHandleExt}; use crate::retry::ProviderRetry; @@ -594,24 +594,27 @@ impl Provider for DatabricksProvider { let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { let payload_clone = payload.clone(); - let resp = self - .api_client - .request(&path) - .model_headers(model_config)? - .streaming(true) - .response_post(&payload_clone) - .await?; - handle_status(resp).await + let mut response = handle_status( + self.api_client + .request(&path) + .model_headers(model_config)? + .streaming(true) + .response_post(&payload_clone) + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - stream_responses_compat(response, log) + stream_responses_compat_with_prefix(response, first_chunk, log) } else { let format_model_config; let request_model_config = if Self::is_claude_model(effective_model_name) @@ -663,31 +666,7 @@ impl Provider for DatabricksProvider { let mut log = start_log(model_config, &payload)?; let response = self .with_retry(|| async { - let resp = self - .api_client - .request(&path) - .model_headers(model_config)? - .streaming(true) - .response_post(&payload) - .await?; - if !resp.status().is_success() { - let status = resp.status(); - let url = sanitize_url(resp.url().as_str()); - let error_text = crate::http_status::read_error_body(resp) - .await - .unwrap_or_default(); - - let json_payload = serde_json::from_str::(&error_text).ok(); - return Err(map_http_error_to_provider_error(status, json_payload, &url)); - } - Ok(resp) - }) - .await; - - let response = match response { - Err(e) if e.to_string().contains("stream_options") => { - payload.as_object_mut().unwrap().remove("stream_options"); - self.with_retry(|| async { + let mut response = { let resp = self .api_client .request(&path) @@ -701,6 +680,7 @@ impl Provider for DatabricksProvider { let error_text = crate::http_status::read_error_body(resp) .await .unwrap_or_default(); + let json_payload = serde_json::from_str::(&error_text).ok(); return Err(map_http_error_to_provider_error( status, @@ -708,7 +688,42 @@ impl Provider for DatabricksProvider { &url, )); } - Ok(resp) + resp + }; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) + }) + .await; + + let (response, first_chunk) = match response { + Err(e) if e.to_string().contains("stream_options") => { + payload.as_object_mut().unwrap().remove("stream_options"); + self.with_retry(|| async { + let mut response = { + let resp = self + .api_client + .request(&path) + .model_headers(model_config)? + .streaming(true) + .response_post(&payload) + .await?; + if !resp.status().is_success() { + let status = resp.status(); + let url = sanitize_url(resp.url().as_str()); + let error_text = crate::http_status::read_error_body(resp) + .await + .unwrap_or_default(); + let json_payload = serde_json::from_str::(&error_text).ok(); + return Err(map_http_error_to_provider_error( + status, + json_payload, + &url, + )); + } + resp + }; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { @@ -722,7 +737,7 @@ impl Provider for DatabricksProvider { Ok(resp) => resp, }; - stream_openai_compat(response, log) + stream_openai_compat_with_prefix(response, first_chunk, log) } } diff --git a/crates/goose-providers/src/databricks_v2.rs b/crates/goose-providers/src/databricks_v2.rs index 0557676527a8..fb5b20c887cd 100644 --- a/crates/goose-providers/src/databricks_v2.rs +++ b/crates/goose-providers/src/databricks_v2.rs @@ -25,7 +25,10 @@ use crate::errors::ProviderError; use crate::formats::anthropic; use crate::formats::openai_responses; use crate::model::ModelConfig; -use crate::openai_compatible::{handle_status, stream_openai_compat, stream_responses_compat}; +use crate::openai_compatible::{ + first_body_chunk, handle_status, stream_openai_compat_with_prefix, + stream_responses_compat_with_prefix, +}; use crate::request_log::{start_log, LoggerHandleExt}; use crate::retry::ProviderRetry; use crate::retry::{ @@ -201,23 +204,26 @@ impl DatabricksV2Provider { payload["stream"] = Value::Bool(true); let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { - let resp = self - .api_client - .request("ai-gateway/openai/v1/responses") - .model_headers(model_config)? - .streaming(true) - .response_post(&payload) - .await?; - handle_status(resp).await + let mut response = handle_status( + self.api_client + .request("ai-gateway/openai/v1/responses") + .model_headers(model_config)? + .streaming(true) + .response_post(&payload) + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - stream_responses_compat(response, log) + stream_responses_compat_with_prefix(response, first_chunk, log) } async fn stream_mlflow_chat_completions( @@ -240,23 +246,26 @@ impl DatabricksV2Provider { } let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { - let resp = self - .api_client - .request("ai-gateway/mlflow/v1/chat/completions") - .model_headers(model_config)? - .streaming(true) - .response_post(&payload) - .await?; - handle_status(resp).await + let mut response = handle_status( + self.api_client + .request("ai-gateway/mlflow/v1/chat/completions") + .model_headers(model_config)? + .streaming(true) + .response_post(&payload) + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - stream_openai_compat(response, log) + stream_openai_compat_with_prefix(response, first_chunk, log) } async fn stream_anthropic_messages( diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index cef1da9d872c..2bdae46c9e65 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -98,6 +98,7 @@ thiserror = { workspace = true } futures = { workspace = true } dirs = { workspace = true } reqwest = { workspace = true, features = ["json", "cookies", "gzip", "brotli", "deflate", "zstd", "charset", "http2", "stream", "blocking", "multipart", "system-proxy"] } +bytes = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "macros", "sync", "time", "fs", "process", "net", "signal", "io-std", "io-util"] } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/goose/src/providers/bedrock.rs b/crates/goose/src/providers/bedrock.rs index efd42d25b485..dea8e2c6e9dc 100644 --- a/crates/goose/src/providers/bedrock.rs +++ b/crates/goose/src/providers/bedrock.rs @@ -4,7 +4,9 @@ use super::base::{ ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata, DEFAULT_CONNECT_TIMEOUT_SECS, DEFAULT_PROVIDER_TIMEOUT_SECS, }; -use super::openai_compatible::{handle_status, stream_responses_compat}; +use super::openai_compatible::{ + first_body_chunk, handle_status, stream_responses_compat_with_prefix, +}; use super::retry::{ProviderRetry, RetryConfig}; use crate::conversation::message::Message; use crate::session_context::SESSION_ID_HEADER; @@ -244,7 +246,7 @@ impl BedrockProvider { &self, session_id: Option<&str>, payload: &Value, - ) -> Result { + ) -> Result<(reqwest::Response, Option), ProviderError> { let region = self.region.as_deref().ok_or_else(|| { ProviderError::Authentication( "AWS region is required for Bedrock mantle endpoint".to_string(), @@ -284,7 +286,9 @@ impl BedrockProvider { ) .await?; - handle_status(response).await + let mut response = handle_status(response).await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) } /// Build the request inputs shared by [`Self::converse`] and @@ -808,14 +812,14 @@ impl Provider for BedrockProvider { payload["stream"] = Value::Bool(true); let mut log = start_log(model_config, &payload).map_err(anyhow::Error::from)?; - let response = self + let (response, first_chunk) = self .with_retry(|| self.post_mantle_streaming(session_id_opt, &payload)) .await .inspect_err(|e| { let _ = log.error(e); })?; - return stream_responses_compat(response, log); + return stream_responses_compat_with_prefix(response, first_chunk, log); } let model_name = model_config.model_name.clone(); diff --git a/crates/goose/src/providers/githubcopilot.rs b/crates/goose/src/providers/githubcopilot.rs index 81bb6ad858e2..a17b901ba065 100644 --- a/crates/goose/src/providers/githubcopilot.rs +++ b/crates/goose/src/providers/githubcopilot.rs @@ -2,7 +2,8 @@ use crate::config::paths::Paths; use crate::providers::api_client::{ApiClient, AuthMethod}; use crate::providers::oauth_device_flow::{run_device_flow, DeviceFlowConfig, RequestEncoding}; use crate::providers::openai_compatible::{ - handle_status, stream_openai_compat, stream_responses_compat, + first_body_chunk, handle_status, stream_openai_compat_with_prefix, + stream_responses_compat_with_prefix, }; use crate::providers::private_file::write_private_file; use anyhow::{anyhow, Result}; @@ -423,11 +424,11 @@ impl GithubCopilotProvider { let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { let mut payload_clone = payload.clone(); - let resp = self - .post( + let mut response = handle_status( + self.post( model_config, "responses", is_user_initiated, @@ -435,15 +436,18 @@ impl GithubCopilotProvider { has_images, true, ) - .await?; - handle_status(resp).await + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - stream_responses_compat(response, log) + stream_responses_compat_with_prefix(response, first_chunk, log) } #[allow(clippy::too_many_arguments)] @@ -471,11 +475,11 @@ impl GithubCopilotProvider { )?; let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { let mut payload_clone = payload.clone(); - let resp = self - .post( + let mut response = handle_status( + self.post( model_config, "chat/completions", is_user_initiated, @@ -483,15 +487,18 @@ impl GithubCopilotProvider { has_images, true, ) - .await?; - handle_status(resp).await + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - stream_openai_compat(response, log) + stream_openai_compat_with_prefix(response, first_chunk, log) } else { let payload = create_request( model_config, diff --git a/crates/goose/src/providers/nanogpt.rs b/crates/goose/src/providers/nanogpt.rs index 2ced5d7b2893..0e2207a80470 100644 --- a/crates/goose/src/providers/nanogpt.rs +++ b/crates/goose/src/providers/nanogpt.rs @@ -1,6 +1,6 @@ use super::api_client::{ApiClient, AuthMethod}; use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata}; -use super::openai_compatible::{handle_status, stream_openai_compat}; +use super::openai_compatible::{first_body_chunk, handle_status, stream_openai_compat_with_prefix}; use super::retry::ProviderRetry; use crate::conversation::message::Message; use anyhow::Result; @@ -192,23 +192,26 @@ impl Provider for NanoGptProvider { let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { - let resp = self - .api_client - .request("chat/completions") - .model_headers(model_config)? - .streaming(true) - .response_post(&payload) - .await?; - handle_status(resp).await + let mut response = handle_status( + self.api_client + .request("chat/completions") + .model_headers(model_config)? + .streaming(true) + .response_post(&payload) + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - stream_openai_compat(response, log) + stream_openai_compat_with_prefix(response, first_chunk, log) } } diff --git a/crates/goose/src/providers/openrouter.rs b/crates/goose/src/providers/openrouter.rs index 01e1c9bc96a3..4c20a87fc1f5 100644 --- a/crates/goose/src/providers/openrouter.rs +++ b/crates/goose/src/providers/openrouter.rs @@ -7,7 +7,7 @@ use std::collections::HashMap; use super::api_client::{ApiClient, AuthMethod}; use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata}; -use super::openai_compatible::{handle_status, stream_openai_compat}; +use super::openai_compatible::{first_body_chunk, handle_status, stream_openai_compat_with_prefix}; use super::retry::ProviderRetry; use crate::conversation::message::Message; use crate::providers::formats::openrouter as openrouter_format; @@ -79,16 +79,19 @@ impl OpenRouterProvider { &self, model_config: &ModelConfig, payload: &Value, - ) -> Result { + ) -> Result<(reqwest::Response, Option), ProviderError> { self.with_retry(|| async { - let resp = self - .api_client - .request("api/v1/chat/completions") - .model_headers(model_config)? - .streaming(true) - .response_post(payload) - .await?; - handle_status(resp).await + let mut response = handle_status( + self.api_client + .request("api/v1/chat/completions") + .model_headers(model_config)? + .streaming(true) + .response_post(payload) + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await } @@ -296,22 +299,23 @@ impl Provider for OpenRouterProvider { let mut log = start_log(model_config, &payload)?; - let response = match self.post_chat_completions(model_config, &payload).await { - // Mandatory-reasoning endpoints reject the disable request, so - // downgrade to the lowest effort they all accept and retry once. - Err(error) if sent_reasoning_disable && is_mandatory_reasoning_error(&error) => { - let _ = log.error(&error); - payload["reasoning"] = json!({ "effort": "low" }); - log = start_log(model_config, &payload)?; - self.post_chat_completions(model_config, &payload).await + let (response, first_chunk) = + match self.post_chat_completions(model_config, &payload).await { + // Mandatory-reasoning endpoints reject the disable request, so + // downgrade to the lowest effort they all accept and retry once. + Err(error) if sent_reasoning_disable && is_mandatory_reasoning_error(&error) => { + let _ = log.error(&error); + payload["reasoning"] = json!({ "effort": "low" }); + log = start_log(model_config, &payload)?; + self.post_chat_completions(model_config, &payload).await + } + result => result, } - result => result, - } - .inspect_err(|e| { - let _ = log.error(e); - })?; + .inspect_err(|e| { + let _ = log.error(e); + })?; - stream_openai_compat(response, log) + stream_openai_compat_with_prefix(response, first_chunk, log) } } diff --git a/crates/goose/src/providers/tetrate.rs b/crates/goose/src/providers/tetrate.rs index 0524a10e3822..6c8e01618098 100644 --- a/crates/goose/src/providers/tetrate.rs +++ b/crates/goose/src/providers/tetrate.rs @@ -1,8 +1,8 @@ use super::api_client::{ApiClient, AuthMethod}; use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata}; use super::openai_compatible::{ - handle_response_openai_compat, handle_status, map_http_error_to_provider_error, - stream_openai_compat, + first_body_chunk, handle_response_openai_compat, handle_status, + map_http_error_to_provider_error, stream_openai_compat_with_prefix, }; use super::retry::ProviderRetry; use crate::config::signup_tetrate::TETRATE_DEFAULT_MODEL; @@ -148,20 +148,20 @@ impl Provider for TetrateProvider { let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { - let resp = self - .api_client - .request("v1/chat/completions") - .model_headers(model_config)? - .streaming(true) - .response_post(&payload) - .await?; - let resp = handle_status(resp) - .await - .map_err(Self::enrich_credits_error)?; + let mut response = handle_status( + self.api_client + .request("v1/chat/completions") + .model_headers(model_config)? + .streaming(true) + .response_post(&payload) + .await?, + ) + .await + .map_err(Self::enrich_credits_error)?; - let is_json = resp + let is_json = response .headers() .get(reqwest::header::CONTENT_TYPE) .and_then(|v| v.to_str().ok()) @@ -169,7 +169,7 @@ impl Provider for TetrateProvider { .is_some_and(|v| v.contains("json")); if is_json { - let body = goose_providers::http_status::read_error_body(resp) + let body = goose_providers::http_status::read_error_body(response) .await .unwrap_or_default(); if let Ok(payload) = serde_json::from_str::(&body) { @@ -186,14 +186,15 @@ impl Provider for TetrateProvider { ))); } - Ok(resp) + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - stream_openai_compat(response, log) + stream_openai_compat_with_prefix(response, first_chunk, log) } /// Fetch supported models from Tetrate Agent Router Service API From 2361850b0261b2f9ad7ccf8a883221a1e02c434c Mon Sep 17 00:00:00 2001 From: Vincenzo Palazzo Date: Mon, 10 Aug 2026 20:36:24 +0200 Subject: [PATCH 3/6] fix: use rmcp MetaObject in tool response test rmcp 3 renamed Meta to MetaObject; update the sanitization test so the crate builds under --all-targets. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/goose-provider-types/src/conversation/message.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/goose-provider-types/src/conversation/message.rs b/crates/goose-provider-types/src/conversation/message.rs index 51493b1ecd63..13523f021821 100644 --- a/crates/goose-provider-types/src/conversation/message.rs +++ b/crates/goose-provider-types/src/conversation/message.rs @@ -1580,7 +1580,7 @@ mod tests { ); let mut result = CallToolResult::success(vec![text, image]); result.structured_content = Some(serde_json::json!({"safe": "世界"})); - result.meta = Some(rmcp::model::Meta(object!({"source": "test"}))); + result.meta = Some(rmcp::model::MetaObject(object!({"source": "test"}))); let expected = result.clone(); let content = MessageContentBlock::tool_response("tool-1", Ok(result)); From def841041fe9154e6cb723512c953632b2e8d342 Mon Sep 17 00:00:00 2001 From: Vincenzo Palazzo Date: Mon, 10 Aug 2026 20:47:10 +0200 Subject: [PATCH 4/6] fix(providers): pre-read first chunk on Anthropic stream routes Apply first_body_chunk inside with_retry for Databricks v2 Anthropic messages and the Anthropic provider stream path, using the shared body_stream_with_prefix helper. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/goose-providers/src/anthropic.rs | 13 +++++----- crates/goose-providers/src/databricks_v2.rs | 26 ++++++++++--------- .../goose-providers/src/openai_compatible.rs | 2 +- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/crates/goose-providers/src/anthropic.rs b/crates/goose-providers/src/anthropic.rs index 1c82eb606596..35f872c0f298 100644 --- a/crates/goose-providers/src/anthropic.rs +++ b/crates/goose-providers/src/anthropic.rs @@ -9,7 +9,6 @@ use async_trait::async_trait; use futures::TryStreamExt; use reqwest::StatusCode; use serde_json::Value; -use std::io; use tokio::pin; use tokio_util::io::StreamReader; @@ -19,7 +18,7 @@ use super::formats::anthropic::{ create_request_for_model, response_to_streaming_message, AnthropicFormatOptions, ANTHROPIC_PROVIDER_NAME, }; -use super::openai_compatible::handle_status; +use super::openai_compatible::{body_stream_with_prefix, first_body_chunk, handle_status}; use super::retry::ProviderRetry; use crate::conversation::message::Message; use crate::model::ModelConfig; @@ -182,9 +181,9 @@ impl AnthropicProvider { )?; payload["stream"] = Value::Bool(true); let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { - handle_status( + let mut response = handle_status( self.api_client .request("v1/messages") .model_headers(model_config)? @@ -192,13 +191,15 @@ impl AnthropicProvider { .response_post(&payload) .await?, ) - .await + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - let stream = response.bytes_stream().map_err(io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let reader = StreamReader::new(stream); let framed = tokio_util::codec::FramedRead::new(reader, tokio_util::codec::LinesCodec::new()).map_err(anyhow::Error::from); diff --git a/crates/goose-providers/src/databricks_v2.rs b/crates/goose-providers/src/databricks_v2.rs index fb5b20c887cd..28e06c9c9a40 100644 --- a/crates/goose-providers/src/databricks_v2.rs +++ b/crates/goose-providers/src/databricks_v2.rs @@ -7,7 +7,6 @@ use async_trait::async_trait; use futures::TryStreamExt; use serde::Serialize; use serde_json::Value; -use std::io; use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio::pin; @@ -26,7 +25,7 @@ use crate::formats::anthropic; use crate::formats::openai_responses; use crate::model::ModelConfig; use crate::openai_compatible::{ - first_body_chunk, handle_status, stream_openai_compat_with_prefix, + body_stream_with_prefix, first_body_chunk, handle_status, stream_openai_compat_with_prefix, stream_responses_compat_with_prefix, }; use crate::request_log::{start_log, LoggerHandleExt}; @@ -286,23 +285,26 @@ impl DatabricksV2Provider { payload["stream"] = Value::Bool(true); let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { - let resp = self - .api_client - .request("ai-gateway/anthropic/v1/messages") - .model_headers(model_config)? - .streaming(true) - .response_post(&payload) - .await?; - handle_status(resp).await + let mut response = handle_status( + self.api_client + .request("ai-gateway/anthropic/v1/messages") + .model_headers(model_config)? + .streaming(true) + .response_post(&payload) + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - let stream = response.bytes_stream().map_err(io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); diff --git a/crates/goose-providers/src/openai_compatible.rs b/crates/goose-providers/src/openai_compatible.rs index ebb965860ce2..76ed521f0ac1 100644 --- a/crates/goose-providers/src/openai_compatible.rs +++ b/crates/goose-providers/src/openai_compatible.rs @@ -261,7 +261,7 @@ pub async fn first_body_chunk(response: &mut Response) -> Result, ) -> Pin> + Send>> { From a08ce855dd3917b8353f7818e5654c57566a1fe6 Mon Sep 17 00:00:00 2001 From: Vincenzo Palazzo Date: Mon, 10 Aug 2026 20:59:27 +0200 Subject: [PATCH 5/6] fix(providers): pre-read first chunk on custom SSE providers Apply first_body_chunk inside with_retry for Google, ChatGPT Codex, Gemini OAuth, and Kimi Code stream paths using body_stream_with_prefix. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/goose-providers/src/google.rs | 16 +++++++++++----- crates/goose/src/providers/chatgpt_codex.rs | 12 ++++++++---- crates/goose/src/providers/gemini_oauth.rs | 12 ++++++++---- crates/goose/src/providers/kimicode.rs | 12 ++++++------ 4 files changed, 33 insertions(+), 19 deletions(-) diff --git a/crates/goose-providers/src/google.rs b/crates/goose-providers/src/google.rs index bc9e7eab7e1e..24fa4ac92808 100644 --- a/crates/goose-providers/src/google.rs +++ b/crates/goose-providers/src/google.rs @@ -2,7 +2,10 @@ use crate::api_client::{ApiClient, AuthMethod}; use crate::base::MessageStream; use crate::conversation::message::Message; use crate::errors::ProviderError; -use crate::openai_compatible::{handle_status, map_http_error_to_provider_error, sanitize_url}; +use crate::openai_compatible::{ + body_stream_with_prefix, first_body_chunk, handle_status, map_http_error_to_provider_error, + sanitize_url, +}; use crate::retry::ProviderRetry; use crate::base::{ConfigKey, Provider, ProviderMetadata}; @@ -15,7 +18,6 @@ use async_trait::async_trait; use futures::TryStreamExt; use rmcp::model::Tool; use serde_json::Value; -use std::io; use tokio::pin; use tokio_stream::StreamExt; use tokio_util::codec::{FramedRead, LinesCodec}; @@ -192,14 +194,18 @@ impl Provider for GoogleProvider { )?; let mut log = start_log(model_config, &payload)?; - let response = self - .with_retry(|| async { self.post_stream(model_config, &payload).await }) + let (response, first_chunk) = self + .with_retry(|| async { + let mut response = self.post_stream(model_config, &payload).await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) + }) .await .inspect_err(|e| { let _ = log.error(e); })?; - let stream = response.bytes_stream().map_err(io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); diff --git a/crates/goose/src/providers/chatgpt_codex.rs b/crates/goose/src/providers/chatgpt_codex.rs index 3622898cd2f0..b0d0ce3d6f6f 100644 --- a/crates/goose/src/providers/chatgpt_codex.rs +++ b/crates/goose/src/providers/chatgpt_codex.rs @@ -5,7 +5,9 @@ use crate::providers::base::{ ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata, DEFAULT_CONNECT_TIMEOUT_SECS, DEFAULT_PROVIDER_TIMEOUT_SECS, }; -use crate::providers::openai_compatible::handle_status; +use crate::providers::openai_compatible::{ + body_stream_with_prefix, first_body_chunk, handle_status, +}; use crate::providers::private_file::write_private_file; use crate::providers::retry::ProviderRetry; use anyhow::{anyhow, Result}; @@ -1001,14 +1003,16 @@ impl Provider for ChatGptCodexProvider { .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; payload["stream"] = serde_json::Value::Bool(true); - let response = self + let (response, first_chunk) = self .with_retry(|| async { let payload_clone = payload.clone(); - self.post_streaming(&payload_clone).await + let mut response = self.post_streaming(&payload_clone).await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await?; - let stream = response.bytes_stream().map_err(io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); diff --git a/crates/goose/src/providers/gemini_oauth.rs b/crates/goose/src/providers/gemini_oauth.rs index 39bdb7f437bb..9dfcd0a32511 100644 --- a/crates/goose/src/providers/gemini_oauth.rs +++ b/crates/goose/src/providers/gemini_oauth.rs @@ -10,6 +10,7 @@ use crate::providers::google::GOOGLE_DOC_URL; use crate::providers::private_file::write_private_file; use goose_providers::errors::ProviderError; use goose_providers::model::ModelConfig; +use goose_providers::openai_compatible::{body_stream_with_prefix, first_body_chunk}; use goose_providers::request_log::{start_log, LoggerHandleExt}; const GEMINI_OAUTH_DEFAULT_MODEL: &str = "gemini-3-flash-preview"; @@ -27,7 +28,6 @@ use rmcp::model::Tool; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use sha2::Digest; -use std::io; use std::net::SocketAddr; use std::path::PathBuf; use std::sync::{Arc, LazyLock}; @@ -992,14 +992,18 @@ impl Provider for GeminiOAuthProvider { let payload = create_request(model_config, system, messages, tools)?; let mut log = start_log(model_config, &payload)?; - let response = self - .with_retry(|| async { self.post_stream(&model_config.model_name, &payload).await }) + let (response, first_chunk) = self + .with_retry(|| async { + let mut response = self.post_stream(&model_config.model_name, &payload).await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) + }) .await .inspect_err(|e| { let _ = log.error(e); })?; - let stream = response.bytes_stream().map_err(io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); diff --git a/crates/goose/src/providers/kimicode.rs b/crates/goose/src/providers/kimicode.rs index b36a8f4b33ce..d0c3c2750673 100644 --- a/crates/goose/src/providers/kimicode.rs +++ b/crates/goose/src/providers/kimicode.rs @@ -9,7 +9,6 @@ use reqwest::header::{HeaderMap, HeaderValue}; use reqwest::Client; use serde::{Deserialize, Serialize}; use serde_json::Value; -use std::io; use std::time::Duration as StdDuration; use tokio::pin; use tokio_util::io::StreamReader; @@ -25,7 +24,7 @@ use super::oauth_device_flow::{ refresh_device_flow_token, run_device_flow, DeviceFlowConfig, DeviceFlowTokenRefreshError, DeviceFlowTokens, RequestEncoding, }; -use super::openai_compatible::handle_status; +use super::openai_compatible::{body_stream_with_prefix, first_body_chunk, handle_status}; use super::retry::ProviderRetry; use crate::conversation::message::Message; use futures::future::BoxFuture; @@ -436,17 +435,18 @@ impl Provider for KimiCodeProvider { let mut log = start_log(model_config, &payload) .map_err(|e| ProviderError::RequestFailed(e.to_string()))?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { - let resp = self.post(&payload).await?; - handle_status(resp).await + let mut response = handle_status(self.post(&payload).await?).await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - let stream = response.bytes_stream().map_err(io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); From 3a225da465829b25be4dfa3dbfc28150ff0251f6 Mon Sep 17 00:00:00 2001 From: Vincenzo Palazzo Date: Mon, 10 Aug 2026 21:16:02 +0200 Subject: [PATCH 6/6] fix(providers): pre-read first chunk on Ollama and Vertex AI Apply first_body_chunk inside the retry path for Ollama streaming and GCP Vertex AI send_request_with_retry, replaying via body_stream_with_prefix. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/goose-providers/src/ollama.rs | 28 ++++++++++--------- crates/goose/src/providers/gcpvertexai.rs | 33 +++++++++++++++++------ 2 files changed, 41 insertions(+), 20 deletions(-) diff --git a/crates/goose-providers/src/ollama.rs b/crates/goose-providers/src/ollama.rs index c1de5a336370..fa72c8285c30 100644 --- a/crates/goose-providers/src/ollama.rs +++ b/crates/goose-providers/src/ollama.rs @@ -1,6 +1,6 @@ use super::api_client::ApiClient; use super::base::{ConfigKey, MessageStream, Provider, ProviderMetadata}; -use super::openai_compatible::handle_status; +use super::openai_compatible::{body_stream_with_prefix, first_body_chunk, handle_status}; use super::retry::{ProviderRetry, RetryConfig}; use crate::api_client::{AuthMethod, TlsConfig}; use crate::base::ProviderDescriptor; @@ -422,22 +422,25 @@ impl Provider for OllamaProvider { apply_ollama_options(&mut payload, &self.options, model_config); let mut log = start_log(model_config, &payload)?; - let response = self + let (response, first_chunk) = self .with_retry(|| async { - let resp = self - .api_client - .request("v1/chat/completions") - .model_headers(model_config)? - .streaming(true) - .response_post(&payload) - .await?; - handle_status(resp).await + let mut response = handle_status( + self.api_client + .request("v1/chat/completions") + .model_headers(model_config)? + .streaming(true) + .response_post(&payload) + .await?, + ) + .await?; + let first_chunk = first_body_chunk(&mut response).await?; + Ok((response, Some(first_chunk))) }) .await .inspect_err(|e| { let _ = log.error(e); })?; - stream_ollama(response, self.options.chunk_timeout_secs, log) + stream_ollama(response, first_chunk, self.options.chunk_timeout_secs, log) } async fn fetch_supported_models(&self) -> Result, ProviderError> { @@ -513,10 +516,11 @@ fn with_line_timeout( /// buffering inside response_to_streaming_message_ollama does not cause false stalls. fn stream_ollama( response: Response, + first_chunk: Option, chunk_timeout: u64, mut log: Option>, ) -> Result { - let stream = response.bytes_stream().map_err(std::io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); diff --git a/crates/goose/src/providers/gcpvertexai.rs b/crates/goose/src/providers/gcpvertexai.rs index 07f1aa218396..f66475280cbd 100644 --- a/crates/goose/src/providers/gcpvertexai.rs +++ b/crates/goose/src/providers/gcpvertexai.rs @@ -1,4 +1,3 @@ -use std::io; use std::time::Duration; use anyhow::Result; @@ -27,7 +26,9 @@ use crate::providers::formats::gcpvertexai::{ DEFAULT_MODEL, KNOWN_MODELS, }; use crate::providers::gcpauth::GcpAuth; -use crate::providers::openai_compatible::{map_http_error_to_provider_error, sanitize_url}; +use crate::providers::openai_compatible::{ + body_stream_with_prefix, first_body_chunk, map_http_error_to_provider_error, sanitize_url, +}; use crate::providers::retry::RetryConfig; use goose_providers::errors::ProviderError; use goose_providers::http_status::read_error_body; @@ -284,7 +285,7 @@ impl GcpVertexAIProvider { model: &ModelConfig, url: Url, payload: &Value, - ) -> Result { + ) -> Result<(reqwest::Response, Option), ProviderError> { let mut rate_limit_attempts = 0; let mut overloaded_attempts = 0; let mut last_error = None; @@ -376,7 +377,23 @@ impl GcpVertexAIProvider { }); sleep(self.retry_config.delay_for_attempt(overloaded_attempts)).await; } else if status == StatusCode::OK { - return Ok(response); + // Pre-read the first body chunk inside the retry loop so a 200 + // whose connection dies before any body bytes still retries. + let mut response = response; + match first_body_chunk(&mut response).await { + Ok(first_chunk) => return Ok((response, Some(first_chunk))), + Err(e) => { + rate_limit_attempts += 1; + if rate_limit_attempts > max_retries { + return Err(e); + } + tracing::warn!( + "pre-body stream failure (attempt {rate_limit_attempts}/{max_retries}): {e}" + ); + last_error = Some(e); + sleep(self.retry_config.delay_for_attempt(rate_limit_attempts)).await; + } + } } else if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { if !retried_auth { retried_auth = true; @@ -407,7 +424,7 @@ impl GcpVertexAIProvider { payload: &Value, context: &RequestContext, location: &str, - ) -> Result { + ) -> Result<(reqwest::Response, Option), ProviderError> { let url = self .build_request_url(model, context.provider(), location, true) .map_err(|e| ProviderError::RequestFailed(e.to_string()))?; @@ -420,7 +437,7 @@ impl GcpVertexAIProvider { model: &ModelConfig, payload: &Value, context: &RequestContext, - ) -> Result { + ) -> Result<(reqwest::Response, Option), ProviderError> { let result = self .post_stream_with_location(model, payload, context, &self.location) .await; @@ -628,14 +645,14 @@ impl Provider for GcpVertexAIProvider { let mut log = start_log(model_config, &request)?; - let response = self + let (response, first_chunk) = self .post_stream(model_config, &request, &context) .await .inspect_err(|e| { let _ = log.error(e); })?; - let stream = response.bytes_stream().map_err(io::Error::other); + let stream = body_stream_with_prefix(response, first_chunk); let context_clone = context.clone(); Ok(Box::pin(try_stream! {