From 7ec07f2c792e7f677757f897c80a3da8424fa8fd Mon Sep 17 00:00:00 2001 From: fre$h Date: Thu, 26 Mar 2026 14:43:59 -0400 Subject: [PATCH] fix: prevent Ollama provider from hanging on tool-calling requests (#7723) Signed-off-by: fre Signed-off-by: fresh3nough Signed-off-by: Douwe Osinga Co-authored-by: Douwe Osinga --- Cargo.lock | 23 ++-- crates/goose/Cargo.toml | 2 + crates/goose/src/providers/ollama.rs | 197 +++++++++++++++++++++++++-- 3 files changed, 204 insertions(+), 18 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1fc6b5ad83..8014cbae07 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4342,6 +4342,7 @@ dependencies = [ "base64 0.22.1", "blake3", "byteorder", + "bytes", "candle-core", "candle-nn", "candle-transformers", @@ -4358,6 +4359,7 @@ dependencies = [ "fs2", "futures", "goose-test-support", + "http 1.4.0", "ignore", "include_dir", "indexmap 2.13.0", @@ -5787,9 +5789,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libredox" -version = "0.1.14" +version = "0.1.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1744e39d1d6a9948f4f388969627434e31128196de472883b39f148769bfe30a" +checksum = "7ddbf48fd451246b1f8c2610bd3b4ac0cc6e149d89832867093ab69a17194f08" dependencies = [ "bitflags 2.11.0", "libc", @@ -6344,9 +6346,9 @@ dependencies = [ [[package]] name = "num-conv" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf97ec579c3c42f953ef76dbf8d55ac91fb219dde70e49aa4a6b7d74e9919050" +checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967" [[package]] name = "num-derive" @@ -6960,9 +6962,9 @@ dependencies = [ [[package]] name = "pctx_executor" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52c6f6e9ae409b20157d7f60623e32de6627b7a80a159e3ee5c1472c4f36ba64" +checksum = "459b68c42775a94e0e160537e644f9a8d4ab7041fa724976b08382eb386836e1" dependencies = [ "deno_core", "deno_resolver", @@ -6984,13 +6986,14 @@ dependencies = [ [[package]] name = "pctx_registry" -version = "0.1.0" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a49c948ffc8c07357e76b2e0008503c9fafaf49d91fd7e00e672bbd7aabd157" +checksum = "2f5a53bf73ba98a5352e788391f47c33af096c696c64e130110b85cc51336e20" dependencies = [ "deno_error", "pctx_config", "rmcp", + "serde", "serde_json", "thiserror 2.0.18", "tokio", @@ -11297,9 +11300,9 @@ checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" [[package]] name = "unicode-segmentation" -version = "1.12.0" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" +checksum = "da36089a805484bcccfffe0739803392c8298778a2d2f09febf76fac5ad9025b" [[package]] name = "unicode-width" diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index c1194c178a..5df8168fd8 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -215,6 +215,8 @@ env-lock = { workspace = true } rmcp = { workspace = true, features = ["transport-streamable-http-server"] } opentelemetry_sdk = { workspace = true, features = ["testing"] } goose-test-support = { path = "../goose-test-support" } +bytes.workspace = true +http.workspace = true [[example]] name = "agent" diff --git a/crates/goose/src/providers/ollama.rs b/crates/goose/src/providers/ollama.rs index 84bc18676f..5678a1da9a 100644 --- a/crates/goose/src/providers/ollama.rs +++ b/crates/goose/src/providers/ollama.rs @@ -68,14 +68,28 @@ fn resolve_ollama_num_ctx(model_config: &ModelConfig) -> Option { } fn apply_ollama_options(payload: &mut Value, model_config: &ModelConfig) { - let Some(limit) = resolve_ollama_num_ctx(model_config) else { - return; - }; - if let Some(obj) = payload.as_object_mut() { - let options = obj.entry("options").or_insert_with(|| json!({})); - if let Some(options_obj) = options.as_object_mut() { - options_obj.insert("num_ctx".to_string(), json!(limit)); + // Ollama does not support stream_options; remove it to prevent hangs. + obj.remove("stream_options"); + + // Convert max_completion_tokens / max_tokens to Ollama's options.num_predict. + // Reasoning models emit max_completion_tokens; non-reasoning models emit max_tokens. + let max_tokens = obj + .remove("max_completion_tokens") + .or_else(|| obj.remove("max_tokens")); + if let Some(max_tokens) = max_tokens { + let options = obj.entry("options").or_insert_with(|| json!({})); + if let Some(options_obj) = options.as_object_mut() { + options_obj.entry("num_predict").or_insert(max_tokens); + } + } + + // Apply num_ctx from context limit settings. + if let Some(limit) = resolve_ollama_num_ctx(model_config) { + let options = obj.entry("options").or_insert_with(|| json!({})); + if let Some(options_obj) = options.as_object_mut() { + options_obj.insert("num_ctx".to_string(), json!(limit)); + } } } } @@ -300,9 +314,49 @@ impl Provider for OllamaProvider { } } +/// Per-chunk timeout for Ollama streaming responses. +/// If no new raw SSE data arrives within this duration, the connection is considered dead. +const OLLAMA_CHUNK_TIMEOUT_SECS: u64 = 30; + +/// Wraps a line stream with a per-item timeout at the raw SSE level. +/// This detects dead connections without false-positive stalls during long +/// tool-call generations where response_to_streaming_message_ollama buffers. +fn with_line_timeout( + stream: impl futures::Stream> + Unpin + Send + 'static, + timeout_secs: u64, +) -> std::pin::Pin> + Send>> { + let timeout = Duration::from_secs(timeout_secs); + Box::pin(try_stream! { + let mut stream = stream; + + // Allow time-to-first-token to be governed by the request timeout. + // Only enforce per-chunk timeout after first SSE line arrives. + match stream.next().await { + Some(first_item) => yield first_item?, + None => return, + } + loop { + match tokio::time::timeout(timeout, stream.next()).await { + Ok(Some(item)) => yield item?, + Ok(None) => break, + Err(_) => { + Err::<(), anyhow::Error>(anyhow::anyhow!( + "Ollama stream stalled: no data received for {}s. \ + This may indicate the model is overwhelmed by the request payload. \ + Try a smaller model or reduce the number of tools.", + timeout_secs + ))?; + } + } + } + }) +} + /// Ollama-specific streaming handler with XML tool call fallback. /// Uses the Ollama format module which buffers text when XML tool calls are detected, /// preventing duplicate content from being emitted to the UI. +/// Timeout is applied at the raw SSE line level via with_line_timeout so that +/// buffering inside response_to_streaming_message_ollama does not cause false stalls. fn stream_ollama(response: Response, mut log: RequestLog) -> Result { let stream = response.bytes_stream().map_err(std::io::Error::other); @@ -311,8 +365,10 @@ fn stream_ollama(response: Response, mut log: RequestLog) -> Result)]); + let model_config = ModelConfig::new("llama3.1") + .unwrap() + .with_max_tokens(Some(4096)); + let messages = vec![crate::conversation::message::Message::user().with_text("hi")]; + + let mut payload = create_request( + &model_config, + "You are a helpful assistant.", + &messages, + &[], + &ImageFormat::OpenAi, + true, + ) + .unwrap(); + + apply_ollama_options(&mut payload, &model_config); + + assert!( + payload.get("stream_options").is_none(), + "stream_options should be removed for Ollama" + ); + assert!( + payload.get("max_tokens").is_none(), + "max_tokens should be removed for Ollama" + ); + assert!( + payload.get("max_completion_tokens").is_none(), + "max_completion_tokens should be removed for Ollama" + ); + assert_eq!( + payload["options"]["num_predict"], 4096, + "max_tokens should be moved to options.num_predict" + ); + assert_eq!(payload["stream"], true, "stream field should be preserved"); + } + + #[tokio::test] + async fn test_stream_ollama_timeout_on_stall() { + use std::convert::Infallible; + + let (tx, rx) = tokio::sync::mpsc::channel::>(1); + tx.send(Ok(bytes::Bytes::from( + "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"index\":0}],\ + \"model\":\"test\",\"object\":\"chat.completion.chunk\",\"created\":0}\n", + ))) + .await + .unwrap(); + let stream = tokio_stream::wrappers::ReceiverStream::new(rx); + let body = reqwest::Body::wrap_stream(stream); + let response = http::Response::builder().status(200).body(body).unwrap(); + let response: reqwest::Response = response.into(); + + let log = RequestLog::start( + &ModelConfig::new("test").unwrap(), + &json!({"model": "test"}), + ) + .unwrap(); + + let mut msg_stream = stream_ollama(response, log).unwrap(); + + let result = + tokio::time::timeout(Duration::from_secs(OLLAMA_CHUNK_TIMEOUT_SECS + 5), async { + let mut last_err = None; + while let Some(item) = msg_stream.next().await { + if let Err(e) = item { + last_err = Some(e); + break; + } + } + last_err + }) + .await; + + match result { + Ok(Some(err)) => { + let err_msg = err.to_string(); + assert!( + err_msg.contains("stream stalled"), + "Expected stall timeout error, got: {}", + err_msg + ); + } + Ok(None) => panic!("Expected timeout error but stream completed normally"), + Err(_) => panic!("Outer timeout elapsed -- per-chunk timeout did not fire"), + } + + drop(tx); + } + #[test] fn test_ollama_retry_config_is_transient_only() { let config = RetryConfig::new(