diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 323de68d6116..1b562b2ddba0 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -2463,8 +2463,21 @@ impl Agent { // reasoning without hiding final-only non-streaming thoughts. let mut surfaced_thinking_in_turn = false; - while let Some(next) = stream.next().await { - if is_token_cancelled(&cancel_token) || exit_chat { + loop { + let next = if let Some(cancel_token) = &cancel_token { + tokio::select! { + biased; + _ = cancel_token.cancelled() => break, + next = stream.next() => next, + } + } else { + stream.next().await + }; + let Some(next) = next else { + break; + }; + + if exit_chat { break; } diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 52670dd2eb06..06eb12a93af2 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -361,6 +361,7 @@ pub(crate) async fn stream_response_from_provider( // Clone owned data to move into the async stream let system_prompt = system_prompt.to_owned(); + let session_id = session_id.to_owned(); let tools = tools.to_owned(); let toolshim_tools = toolshim_tools.to_owned(); let provider = provider.clone(); @@ -372,7 +373,7 @@ pub(crate) async fn stream_response_from_provider( let request_started = std::time::Instant::now(); debug!("WAITING_LLM_STREAM_START"); let stream_result = crate::session_context::with_session_id( - Some(session_id.to_string()), + Some(session_id.clone()), provider.stream( &model_config, system_prompt.as_str(), @@ -397,6 +398,66 @@ pub(crate) async fn stream_response_from_provider( }; Ok(Box::pin(try_stream! { + if !provider.manages_own_context() { + let retry_config = provider.retry_config().transient_only(); + let mut attempts = 0; + + loop { + match stream.next().await { + None => break, + Some(Ok(item)) => { + stream = Box::pin( + futures::stream::once(std::future::ready(Ok(item))).chain(stream), + ); + break; + } + Some(Err(error)) + if goose_providers::retry::should_retry(&error, &retry_config) + && attempts < retry_config.max_retries => + { + attempts += 1; + let delay = match &error { + ProviderError::RateLimitExceeded { + retry_delay: Some(provider_delay), + .. + } => *provider_delay, + _ => retry_config.delay_for_attempt(attempts), + }; + warn!( + "Provider stream failed before its first item, retrying ({}/{}): {:?}", + attempts, retry_config.max_retries, error + ); + + let skip_backoff = std::env::var("GOOSE_PROVIDER_SKIP_BACKOFF") + .unwrap_or_default() + .parse::() + .unwrap_or(false); + if !skip_backoff { + tokio::time::sleep(delay).await; + } + + stream = match crate::session_context::with_session_id( + Some(session_id.clone()), + provider.stream( + &model_config, + system_prompt.as_str(), + messages_for_provider.messages(), + &tools, + ), + ) + .await + { + Ok(stream) => stream, + Err(error) => { + Err(enhance_model_error(error, &provider, config.toolshim).await)? + } + }; + } + Some(Err(error)) => Err(error)?, + } + } + } + if config.toolshim { // Toolshim mode: accumulate the full response before processing // so that tool-use markers spanning multiple chunks are detected @@ -1789,4 +1850,219 @@ mod tests { ); assert!(stats.elapsed_ms.expect("elapsed_ms must be filled") >= 100); } + + type TestStreamItem = Result<(Option, Option), ProviderError>; + + struct SequencedProvider { + responses: Mutex, ProviderError>>>, + calls: Arc, + manages_context: bool, + } + + impl SequencedProvider { + fn new(responses: Vec, ProviderError>>) -> Self { + Self { + responses: Mutex::new(responses.into()), + calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + manages_context: false, + } + } + + fn managing_context(mut self) -> Self { + self.manages_context = true; + self + } + } + + #[async_trait] + impl Provider for SequencedProvider { + fn get_name(&self) -> &str { + "sequenced" + } + + fn retry_config(&self) -> goose_providers::retry::RetryConfig { + goose_providers::retry::RetryConfig::new(2, 0, 1.0, 0) + } + + fn manages_own_context(&self) -> bool { + self.manages_context + } + + async fn stream( + &self, + _model_config: &ModelConfig, + _system: &str, + _messages: &[Message], + _tools: &[Tool], + ) -> Result { + self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let response = self.responses.lock().unwrap().pop_front().unwrap(); + response.map(|items| Box::pin(futures::stream::iter(items)) as MessageStream) + } + } + + fn successful_item() -> TestStreamItem { + Ok((Some(Message::assistant().with_text("ok")), None)) + } + + fn transient_error() -> ProviderError { + ProviderError::NetworkError("stream closed".into()) + } + + async fn stream_for_test(provider: Arc) -> MessageStream { + stream_response_from_provider( + provider, + ModelConfig::new("test-model"), + "session", + "system", + &[Message::user().with_text("hi")], + &[], + &[], + ) + .await + .unwrap() + } + + #[tokio::test] + async fn first_item_transient_error_retries() { + let provider = Arc::new(SequencedProvider::new(vec![ + Ok(vec![Err(transient_error())]), + Ok(vec![successful_item()]), + ])); + let calls = provider.calls.clone(); + let mut stream = stream_for_test(provider).await; + + assert_eq!( + stream + .next() + .await + .unwrap() + .unwrap() + .0 + .unwrap() + .as_concat_text(), + "ok" + ); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn first_item_non_transient_error_does_not_retry() { + let provider = Arc::new(SequencedProvider::new(vec![Ok(vec![Err( + ProviderError::ContextLengthExceeded("too long".into()), + )])])); + let calls = provider.calls.clone(); + let mut stream = stream_for_test(provider).await; + + assert!(matches!( + stream.next().await.unwrap(), + Err(ProviderError::ContextLengthExceeded(_)) + )); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn first_item_retry_exhaustion_returns_last_error() { + let provider = Arc::new(SequencedProvider::new(vec![ + Ok(vec![Err(transient_error())]), + Ok(vec![Err(transient_error())]), + Ok(vec![Err(transient_error())]), + ])); + let calls = provider.calls.clone(); + let mut stream = stream_for_test(provider).await; + + assert!(matches!( + stream.next().await.unwrap(), + Err(ProviderError::NetworkError(_)) + )); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn provider_managing_context_does_not_retry() { + let provider = Arc::new( + SequencedProvider::new(vec![Ok(vec![Err(transient_error())])]).managing_context(), + ); + let calls = provider.calls.clone(); + let mut stream = stream_for_test(provider).await; + + assert!(matches!( + stream.next().await.unwrap(), + Err(ProviderError::NetworkError(_)) + )); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn replacement_stream_creation_error_does_not_retry() { + let provider = Arc::new(SequencedProvider::new(vec![ + Ok(vec![Err(transient_error())]), + Err(ProviderError::ServerError("unavailable".into())), + ])); + let calls = provider.calls.clone(); + let mut stream = stream_for_test(provider).await; + + assert!(matches!( + stream.next().await.unwrap(), + Err(ProviderError::ServerError(_)) + )); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn error_after_first_item_does_not_retry() { + let provider = Arc::new(SequencedProvider::new(vec![Ok(vec![ + successful_item(), + Err(transient_error()), + ])])); + let calls = provider.calls.clone(); + let mut stream = stream_for_test(provider).await; + + assert!(stream.next().await.unwrap().is_ok()); + assert!(matches!( + stream.next().await.unwrap(), + Err(ProviderError::NetworkError(_)) + )); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn empty_stream_is_not_retried() { + let provider = Arc::new(SequencedProvider::new(vec![Ok(vec![])])); + let calls = provider.calls.clone(); + let mut stream = stream_for_test(provider).await; + + assert!(stream.next().await.is_none()); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + struct PendingProvider; + + #[async_trait] + impl Provider for PendingProvider { + fn get_name(&self) -> &str { + "pending" + } + + async fn stream( + &self, + _model_config: &ModelConfig, + _system: &str, + _messages: &[Message], + _tools: &[Tool], + ) -> Result { + Ok(Box::pin(futures::stream::pending())) + } + } + + #[tokio::test] + async fn pending_first_item_does_not_block_stream_creation() { + let result = tokio::time::timeout( + Duration::from_secs(1), + stream_for_test(Arc::new(PendingProvider)), + ) + .await; + + assert!(result.is_ok()); + } }