Skip to content

fix(providers): retry streaming requests dropped before first body bytes #10779

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
vincenzopalazzo wants to merge 6 commits into from
Closed

fix(providers): retry streaming requests dropped before first body bytes #10779

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
146a780
fix(providers): retry streaming requests dropped before first body bytes
vincenzopalazzo Jul 29, 2026
0b78310
fix(providers): pre-read first stream chunk in remaining providers
vincenzopalazzo Aug 10, 2026
2361850
fix: use rmcp MetaObject in tool response test
vincenzopalazzo Aug 10, 2026
def8410
fix(providers): pre-read first chunk on Anthropic stream routes
vincenzopalazzo Aug 10, 2026
a08ce85
fix(providers): pre-read first chunk on custom SSE providers
vincenzopalazzo Aug 10, 2026
3a225da
fix(providers): pre-read first chunk on Ollama and Vertex AI
vincenzopalazzo Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion crates/goose-provider-types/src/conversation/message.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
1 change: 1 addition & 0 deletions crates/goose-providers/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
13 changes: 7 additions & 6 deletions crates/goose-providers/src/anthropic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;
Expand Down Expand Up @@ -182,23 +181,25 @@ 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)?
.streaming(true)
.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);
Expand Down
93 changes: 54 additions & 39 deletions crates/goose-providers/src/databricks.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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::<Value>(&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)
Expand All @@ -701,14 +680,50 @@ 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::<Value>(&error_text).ok();
return Err(map_http_error_to_provider_error(
status,
json_payload,
&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::<Value>(&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| {
Expand All @@ -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)
}
}

Expand Down
75 changes: 43 additions & 32 deletions crates/goose-providers/src/databricks_v2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -25,7 +24,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::{
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};
use crate::retry::ProviderRetry;
use crate::retry::{
Expand Down Expand Up @@ -201,23 +203,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?;
Comment thread
vincenzopalazzo marked this conversation as resolved.
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(
Expand All @@ -240,23 +245,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(
Expand All @@ -277,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);
Expand Down
16 changes: 11 additions & 5 deletions crates/goose-providers/src/google.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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};
Expand Down Expand Up @@ -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);
Expand Down
Loading