diff --git a/crates/agentic-server-core/src/executor/inference.rs b/crates/agentic-server-core/src/executor/inference.rs index 2281c97d..1fba4ee9 100644 --- a/crates/agentic-server-core/src/executor/inference.rs +++ b/crates/agentic-server-core/src/executor/inference.rs @@ -64,11 +64,10 @@ fn drain_complete_utf8_lines(buffer: &mut Vec) -> ExecutorResult Ok(lines) } -async fn response_text_limited(resp: reqwest::Response) -> ExecutorResult { +async fn response_text_limited(resp: reqwest::Response, chunk_timeout: Duration) -> ExecutorResult { let mut stream = resp.bytes_stream(); let mut body = Vec::new(); - while let Some(chunk) = stream.next().await { - let chunk = chunk.map_err(ExecutorError::NetworkError)?; + while let Some(chunk) = next_chunk(&mut stream, chunk_timeout).await? { if chunk.len() > MAX_EXECUTOR_RESPONSE_BYTES.saturating_sub(body.len()) { return Err(ExecutorError::StreamError(format!( "upstream response exceeded {MAX_EXECUTOR_RESPONSE_BYTES} bytes" @@ -85,13 +84,15 @@ async fn response_text_limited(resp: reqwest::Response) -> ExecutorResult, forwarded_headers: Option<&reqwest::header::HeaderMap>, + chunk_timeout: Duration, ) -> ExecutorResult { let mut headers = forwarded_headers.cloned().unwrap_or_default(); headers @@ -120,7 +121,7 @@ pub(super) async fn send_request( let headers = processed_response_headers(resp.headers()); // Log and discard any error reading the error body — the status code // is the primary signal; an empty body is acceptable here. - let body = response_text_limited(resp) + let body = response_text_limited(resp, chunk_timeout) .await .inspect_err(|error| tracing::debug!(%error, "failed to read bounded error response body")) .unwrap_or_default(); @@ -143,9 +144,9 @@ pub(super) async fn fetch_response_json( client: &reqwest::Client, auth: Option<&str>, ) -> ExecutorResult { - let resp = send_request(client, url, upstream_json, auth, None).await?; + let resp = send_request(client, url, upstream_json, auth, None, Duration::ZERO).await?; // Preserve the reqwest::Error as the typed source (NetworkError). - response_text_limited(resp).await + response_text_limited(resp, Duration::ZERO).await } /// Makes a non-streaming HTTP POST with caller-supplied upstream headers. @@ -155,9 +156,9 @@ pub(super) async fn fetch_response_json_with_headers( client: &reqwest::Client, headers: &reqwest::header::HeaderMap, ) -> ExecutorResult<(String, http::HeaderMap)> { - let resp = send_request(client, url, upstream_json, None, Some(headers)).await?; + let resp = send_request(client, url, upstream_json, None, Some(headers), Duration::ZERO).await?; let response_headers = processed_response_headers(resp.headers()); - let body = response_text_limited(resp).await?; + let body = response_text_limited(resp, Duration::ZERO).await?; Ok((body, response_headers)) } @@ -178,7 +179,7 @@ pub fn call_inference( chunk_timeout: Duration, ) -> impl Stream> + Send + 'static { stream! { - let resp = match send_request(&client, &url, upstream_json, auth.as_deref(), None).await { + let resp = match send_request(&client, &url, upstream_json, auth.as_deref(), None, chunk_timeout).await { Ok(r) => r, Err(e) => { yield Err(e); return; } }; @@ -436,9 +437,16 @@ mod tests { #[tokio::test] async fn non_success_response_discards_a_cumulative_oversized_body() { let (url, server) = oversized_body_server(StatusCode::BAD_GATEWAY).await; - let error = send_request(&reqwest::Client::new(), &url, "{}".to_owned(), None, None) - .await - .expect_err("non-success response must fail"); + let error = send_request( + &reqwest::Client::new(), + &url, + "{}".to_owned(), + None, + None, + Duration::ZERO, + ) + .await + .expect_err("non-success response must fail"); let ExecutorError::LLMRequest { status, body, .. } = error else { panic!("expected upstream request error"); diff --git a/crates/agentic-server-core/src/executor/messages_stream.rs b/crates/agentic-server-core/src/executor/messages_stream.rs index efbb2ea4..a0e6e657 100644 --- a/crates/agentic-server-core/src/executor/messages_stream.rs +++ b/crates/agentic-server-core/src/executor/messages_stream.rs @@ -61,6 +61,7 @@ pub async fn run_messages_stream( first_body, None, Some(upstream.headers()), + exec_ctx.streaming_timeout, ) .await?; let response_headers = processed_response_headers(first_response.headers()); @@ -83,6 +84,7 @@ pub async fn run_messages_stream( body, None, Some(upstream.headers()), + exec_ctx.streaming_timeout, ) .await { diff --git a/crates/agentic-server-core/src/executor/request.rs b/crates/agentic-server-core/src/executor/request.rs index b29ae12b..5ffa9ed7 100644 --- a/crates/agentic-server-core/src/executor/request.rs +++ b/crates/agentic-server-core/src/executor/request.rs @@ -66,7 +66,8 @@ pub struct ExecutionContext { pub messages_gateway_tools: GatewayToolMap, /// Base URL for the LLM backend, e.g. `"http://localhost:8000"`. pub llm_base_url: String, - /// Maximum wait time for the next SSE chunk. `Duration::ZERO` disables the timeout. + /// Maximum wait for the next body chunk in a streaming request, including HTTP error bodies. + /// Applies after response headers arrive; `Duration::ZERO` disables the timeout. /// Sourced from the `STREAMING_CHUNK_TIMEOUT_S` environment variable, defaulting to /// [`DEFAULT_STREAMING_TIMEOUT`] when unset or unparseable. pub streaming_timeout: Duration, diff --git a/crates/agentic-server-core/tests/upstream_error_timeout_test.rs b/crates/agentic-server-core/tests/upstream_error_timeout_test.rs new file mode 100644 index 00000000..10887a8a --- /dev/null +++ b/crates/agentic-server-core/tests/upstream_error_timeout_test.rs @@ -0,0 +1,202 @@ +//! Streaming error bodies share the configured idle timeout and retain HTTP errors. + +use std::convert::Infallible; +use std::sync::Arc; +use std::time::Duration; + +use agentic_core::executor::ExecutorError; +use agentic_core::executor::inference::call_inference; +use axum::Router; +use axum::body::Body; +use axum::response::Response; +use axum::routing::post; +use bytes::Bytes; +use futures::{Stream, StreamExt}; +use http::StatusCode; +use tokio::net::TcpListener; + +struct Server(tokio::task::JoinHandle<()>); + +impl Drop for Server { + fn drop(&mut self) { + self.0.abort(); + } +} + +async fn error_server(status: StatusCode, body: F) -> (String, Server) +where + F: Fn() -> S + Clone + Send + Sync + 'static, + S: Stream> + Send + 'static, +{ + let app = Router::new().route( + "/v1/responses", + post(move || { + let stream = body(); + async move { + Response::builder() + .status(status) + .header("content-type", "text/plain; charset=utf-8") + .header("retry-after", "7") + .header("x-request-id", "upstream-error") + .header("connection", "x-private-hop") + .header("x-private-hop", "discard") + .body(Body::from_stream(stream)) + .unwrap() + } + }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/v1/responses", listener.local_addr().unwrap()); + let server = Server(tokio::spawn(async move { axum::serve(listener, app).await.unwrap() })); + (url, server) +} + +async fn read_error(url: String, timeout: Duration) -> Vec> { + call_inference("{}".to_owned(), url, Arc::new(reqwest::Client::new()), None, timeout) + .collect() + .await +} + +fn assert_error(mut result: Vec>, expected_status: StatusCode, expected_body: &str) { + assert_eq!(result.len(), 1, "a failed request emits exactly one error"); + let ExecutorError::LLMRequest { status, body, headers } = result.pop().unwrap().unwrap_err() else { + panic!("expected the original HTTP error"); + }; + assert_eq!(status, expected_status); + assert_eq!(body, expected_body); + assert_eq!(headers["retry-after"], "7"); + assert_eq!(headers["x-request-id"], "upstream-error"); + assert_eq!(headers["content-type"], "text/plain; charset=utf-8"); + assert!(!headers.contains_key("x-private-hop")); + assert!(!headers.contains_key("connection")); +} + +async fn stalled_error(partial: bool, status: StatusCode) { + let (url, _server) = error_server(status, move || { + async_stream::stream! { + if partial { yield Ok(Bytes::from_static(b"partial error")); } + std::future::pending::<()>().await; + } + }) + .await; + let result = tokio::time::timeout(Duration::from_secs(2), read_error(url, Duration::from_millis(50))) + .await + .expect("the configured idle timeout must bound an upstream error body"); + assert_error(result, status, ""); +} + +#[tokio::test] +async fn streaming_error_timeout_bounds_an_empty_429_body() { + stalled_error(false, StatusCode::TOO_MANY_REQUESTS).await; +} + +#[tokio::test] +async fn streaming_error_timeout_bounds_a_partial_503_body() { + stalled_error(true, StatusCode::SERVICE_UNAVAILABLE).await; +} + +#[tokio::test] +async fn streaming_error_timeout_resets_for_each_chunk() { + let (url, _server) = error_server(StatusCode::TOO_MANY_REQUESTS, || { + async_stream::stream! { + for chunk in ["rate ", "limited ", "雪"] { + tokio::time::sleep(Duration::from_millis(80)).await; + yield Ok(Bytes::from(chunk)); + } + } + }) + .await; + let result = tokio::time::timeout(Duration::from_secs(3), read_error(url, Duration::from_millis(200))) + .await + .unwrap(); + assert_error(result, StatusCode::TOO_MANY_REQUESTS, "rate limited 雪"); +} + +#[tokio::test] +async fn streaming_error_timeout_zero_still_allows_delayed_bodies() { + let (url, _server) = error_server(StatusCode::BAD_GATEWAY, || { + async_stream::stream! { + tokio::time::sleep(Duration::from_millis(120)).await; + yield Ok(Bytes::from_static(b"upstream unavailable")); + } + }) + .await; + let result = tokio::time::timeout(Duration::from_secs(2), read_error(url, Duration::ZERO)) + .await + .unwrap(); + assert_error(result, StatusCode::BAD_GATEWAY, "upstream unavailable"); +} + +#[tokio::test] +async fn streaming_error_timeout_retains_empty_completed_bodies() { + let (url, _server) = error_server(StatusCode::BAD_GATEWAY, futures::stream::empty).await; + let result = read_error(url, Duration::from_millis(50)).await; + assert_error(result, StatusCode::BAD_GATEWAY, ""); +} + +#[tokio::test] +async fn streaming_error_timeout_preserves_byte_limit_and_utf8_policy() { + for (bytes, expected) in [ + (vec![b'x'; 1024 * 1024], "x".repeat(1024 * 1024)), + (vec![b'x'; 1024 * 1024 + 1], String::new()), + (vec![0xff], String::new()), + (b"{malformed json".to_vec(), "{malformed json".to_owned()), + ] { + let (url, _server) = error_server(StatusCode::BAD_GATEWAY, move || { + futures::stream::iter( + bytes + .chunks(4096) + .map(|chunk| Ok(Bytes::copy_from_slice(chunk))) + .collect::>(), + ) + }) + .await; + let result = tokio::time::timeout(Duration::from_secs(3), read_error(url, Duration::from_millis(200))) + .await + .unwrap(); + assert_error(result, StatusCode::BAD_GATEWAY, &expected); + } +} + +#[tokio::test] +async fn streaming_error_timeout_cancellation_releases_the_upstream_body() { + use std::sync::atomic::{AtomicBool, Ordering}; + struct Dropped(Arc); + impl Drop for Dropped { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + let dropped = Arc::new(AtomicBool::new(false)); + let route_dropped = Arc::clone(&dropped); + let started = Arc::new(tokio::sync::Notify::new()); + let route_started = Arc::clone(&started); + let (url, _server) = error_server(StatusCode::TOO_MANY_REQUESTS, move || { + let guard = Dropped(Arc::clone(&route_dropped)); + let started = Arc::clone(&route_started); + async_stream::stream! { + let _guard = guard; + started.notify_one(); + yield Ok(Bytes::from_static(b"partial")); + std::future::pending::<()>().await; + } + }) + .await; + let mut read = Box::pin(read_error(url, Duration::ZERO)); + tokio::time::timeout(Duration::from_secs(2), async { + tokio::select! { + () = started.notified() => {}, + _ = &mut read => panic!("the stalled request must remain pending"), + } + }) + .await + .expect("upstream body starts before cancellation"); + drop(read); + tokio::time::timeout(Duration::from_secs(2), async { + while !dropped.load(Ordering::SeqCst) { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("cancelling the read must release the upstream body without waiting for a timer"); +} diff --git a/crates/agentic-server/tests/streaming_error_timeout_test.rs b/crates/agentic-server/tests/streaming_error_timeout_test.rs new file mode 100644 index 00000000..d511da98 --- /dev/null +++ b/crates/agentic-server/tests/streaming_error_timeout_test.rs @@ -0,0 +1,364 @@ +//! Upstream error-body timeouts through public transports and durable recovery. +#[allow(dead_code)] +mod common; + +use std::convert::Infallible; +use std::fmt::Write as _; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use agentic_core::executor::ExecutionContext; +use axum::body::{Body, Bytes}; +use axum::response::{IntoResponse, Response}; +use axum::routing::post; +use axum::{Json, Router}; +use futures::{SinkExt, StreamExt}; +use http::StatusCode; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tokio::sync::Mutex; +use tokio_tungstenite::{connect_async, tungstenite::Message}; + +struct Server(tokio::task::JoinHandle<()>); +impl Drop for Server { + fn drop(&mut self) { + self.0.abort(); + } +} + +struct BodyGuard(Arc); +impl Drop for BodyGuard { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } +} + +async fn wait_for_body_drop(dropped: &AtomicUsize) { + tokio::time::timeout(Duration::from_secs(2), async { + while dropped.load(Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("the upstream body must be released"); +} + +async fn spawn(app: Router) -> (String, Server) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + ( + url, + Server(tokio::spawn(async move { axum::serve(listener, app).await.unwrap() })), + ) +} + +fn stalled_response(dropped: Arc) -> Response { + let guard = BodyGuard(dropped); + let body = futures::stream::once(std::future::ready(Ok::<_, Infallible>(Bytes::from_static( + b"partial diagnostic", + )))) + .chain(futures::stream::once(async move { + let _guard = guard; + std::future::pending::>().await + })); + Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .header("content-type", "text/plain") + .header("retry-after", "7") + .header("x-request-id", "req_stalled") + .body(Body::from_stream(body)) + .unwrap() +} + +fn successful_response() -> Value { + json!({"id":"resp_upstream", "object":"response", "model":"test-model", "status":"completed", + "output":[{"id":"msg_success", "type":"message", "role":"assistant", "status":"completed", + "content":[{"type":"output_text", "text":"recovered 雪"}]}]}) +} + +fn messages_request() -> Value { + json!({"model":"test-model", "max_tokens":64, "stream":true, + "messages":[{"role":"user", "content":"search"}], + "tools":[{"name":"web_search", "input_schema":{"type":"object"}}]}) +} + +#[tokio::test] +async fn messages_streaming_error_timeout_preserves_initial_http_error() { + let dropped = Arc::new(AtomicUsize::new(0)); + let route_dropped = Arc::clone(&dropped); + let requests = Arc::new(AtomicUsize::new(0)); + let route_requests = Arc::clone(&requests); + let (upstream_url, _upstream) = spawn(Router::new().route( + "/v1/messages", + post(move || { + route_requests.fetch_add(1, Ordering::SeqCst); + let dropped = Arc::clone(&route_dropped); + async move { stalled_response(dropped) } + }), + )) + .await; + let mut state = common::test_state(&common::test_config(&upstream_url)); + Arc::make_mut(&mut state.exec_ctx).streaming_timeout = Duration::from_millis(50); + let (gateway_url, gateway) = common::spawn_gateway(state).await; + let _gateway = Server(gateway); + let response = reqwest::Client::builder() + .timeout(Duration::from_secs(2)) + .build() + .unwrap() + .post(format!("{gateway_url}/v1/messages")) + .json(&messages_request()) + .send() + .await + .expect("a stalled upstream error must not hang the Messages handler"); + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(response.headers()["retry-after"], "7"); + assert_eq!(response.headers()["x-request-id"], "req_stalled"); + assert_eq!(response.headers()["content-type"], "text/plain"); + assert!(response.bytes().await.unwrap().is_empty()); + wait_for_body_drop(&dropped).await; + assert_eq!(requests.load(Ordering::SeqCst), 1); +} + +#[derive(Clone, Copy)] +enum Transport { + Http, + WebSocket, +} + +#[allow( + clippy::too_many_lines, + reason = "keeps failure, retry, restart, and actual continuation in one fixture" +)] +async fn response_recovery(transport: Transport) { + let requests = Arc::new(Mutex::new(Vec::::new())); + let route_requests = Arc::clone(&requests); + let dropped = Arc::new(AtomicUsize::new(0)); + let route_dropped = Arc::clone(&dropped); + let (upstream_url, _upstream) = spawn(Router::new().route( + "/v1/responses", + post(move |Json(request): Json| { + let requests = Arc::clone(&route_requests); + let dropped = Arc::clone(&route_dropped); + async move { + let first = { + let mut requests = requests.lock().await; + requests.push(request); + requests.len() == 1 + }; + if first { + stalled_response(dropped) + } else { + Json(successful_response()).into_response() + } + } + }), + )) + .await; + let directory = tempfile::tempdir().unwrap(); + let mut config = common::test_config(&upstream_url); + config.db_url = Some(format!("sqlite://{}", directory.path().join("history.db").display())); + let mut context = ExecutionContext::from_config(&config).await.unwrap(); + context.streaming_timeout = Duration::from_millis(50); + let context = Arc::new(context); + let mut state = common::test_state(&config); + state.exec_ctx = Arc::clone(&context); + let (gateway_url, gateway) = common::spawn_gateway(state).await; + let mut gateway = Server(gateway); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(2)) + .build() + .unwrap(); + let request = json!({"model":"test-model", "input":"failed input", "store":true, "stream":true}); + let error = match transport { + Transport::Http => { + let response = client + .post(format!("{gateway_url}/v1/responses")) + .json(&request) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response + .text() + .await + .expect("a stalled error must terminate the Responses stream"); + assert!(body.ends_with("data: [DONE]\n\n")); + let events: Vec = body + .lines() + .filter_map(|line| line.strip_prefix("data: ")) + .filter(|data| *data != "[DONE]") + .map(|data| serde_json::from_str(data).unwrap()) + .collect(); + assert_eq!(events.len(), 1); + events[0].clone() + } + Transport::WebSocket => { + let (mut socket, _) = connect_async(format!("{}/v1/responses", gateway_url.replace("http:", "ws:"))) + .await + .unwrap(); + let mut request = request; + request["type"] = json!("response.create"); + request["stream_id"] = json!("stalled-request"); + socket.send(Message::Text(request.to_string().into())).await.unwrap(); + let event = tokio::time::timeout(Duration::from_secs(2), socket.next()) + .await + .expect("a stalled error must terminate the WebSocket lane") + .unwrap() + .unwrap(); + let Message::Text(text) = event else { + panic!("expected an error event"); + }; + let error: Value = serde_json::from_str(&text).unwrap(); + assert_eq!(error["stream_id"], "stalled-request"); + socket.close(None).await.unwrap(); + error + } + }; + assert_eq!(error["type"], "error"); + assert_eq!(error["status"], 429); + wait_for_body_drop(&dropped).await; + assert_eq!(requests.lock().await.len(), 1, "no inference retry after an HTTP error"); + let pool = context.storage_pool().unwrap(); + for table in ["responses", "items"] { + let count: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {table}")) + .fetch_one(pool) + .await + .unwrap(); + assert_eq!(count, 0, "failed requests cannot persist {table}"); + } + let success: Value = client + .post(format!("{gateway_url}/v1/responses")) + .json(&json!({"model":"test-model", "input":"successful input", "store":true})) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(success["output"][0]["content"][0]["text"], "recovered 雪"); + gateway.0.abort(); + let _ = (&mut gateway.0).await; + context.storage_pool().unwrap().close().await; + drop(gateway); + drop(context); + let context = Arc::new(ExecutionContext::from_config(&config).await.unwrap()); + let mut state = common::test_state(&config); + state.exec_ctx = Arc::clone(&context); + let (gateway_url, gateway) = common::spawn_gateway(state).await; + let _gateway = Server(gateway); + let response: Value = client + .post(format!("{gateway_url}/v1/responses")) + .json(&json!({"model":"test-model", "input":"continue", "previous_response_id":success["id"], "store":true})) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(response["status"], "completed"); + let requests = requests.lock().await; + assert_eq!(requests.len(), 3); + let downstream = requests[2]["input"].to_string(); + assert!(downstream.contains("successful input")); + assert!(downstream.contains("recovered 雪")); + assert!(downstream.contains("continue")); + assert!(!downstream.contains("failed input")); + assert!(!downstream.contains("partial diagnostic")); + context.storage_pool().unwrap().close().await; +} + +#[tokio::test] +async fn responses_http_streaming_error_timeout_allows_durable_recovery() { + response_recovery(Transport::Http).await; +} + +#[tokio::test] +async fn responses_websocket_streaming_error_timeout_allows_durable_recovery() { + response_recovery(Transport::WebSocket).await; +} + +#[tokio::test] +async fn messages_streaming_error_timeout_bounds_a_later_tool_round() { + let requests = Arc::new(Mutex::new(Vec::::new())); + let route_requests = Arc::clone(&requests); + let dropped = Arc::new(AtomicUsize::new(0)); + let route_dropped = Arc::clone(&dropped); + let searches = Arc::new(AtomicUsize::new(0)); + let route_searches = Arc::clone(&searches); + let events = [ + json!({"type":"message_start", "message":{"id":"msg_tool", "type":"message", "role":"assistant", "content":[], + "model":"test-model", "stop_reason":null, "usage":{"input_tokens":2,"output_tokens":0}}}), + json!({"type":"content_block_start", "index":0, "content_block":{"type":"tool_use", "id":"call_search", "name":"web_search", "input":{}}}), + json!({"type":"content_block_delta", "index":0, "delta":{"type":"input_json_delta", "partial_json":"{\"query\":\"weather\"}"}}), + json!({"type":"content_block_stop", "index":0}), + json!({"type":"message_delta", "delta":{"stop_reason":"tool_use", "stop_sequence":null}, "usage":{"output_tokens":3}}), + json!({"type":"message_stop"}), + ]; + let mut body = String::new(); + for event in events { + write!(body, "event: {}\ndata: {event}\n\n", event["type"].as_str().unwrap()).unwrap(); + } + let app = Router::new().route("/v1/messages", post(move |Json(request): Json| { + let requests = Arc::clone(&route_requests); + let dropped = Arc::clone(&route_dropped); + let body = body.clone(); + async move { + let first = { let mut requests = requests.lock().await; requests.push(request); requests.len() == 1 }; + if first { Response::builder().header("content-type", "text/event-stream").body(Body::from(body)).unwrap() } + else { stalled_response(dropped) } + } + })).route("/v1/search", axum::routing::get(move || { + route_searches.fetch_add(1, Ordering::SeqCst); + async {Json(json!({"results":{"web":[{"url":"https://example.com/weather", "title":"weather", "description":"sunny proof"}]}}))} + })); + let (upstream_url, _upstream) = spawn(app).await; + let directory = tempfile::tempdir().unwrap(); + let mut config = common::test_config(&upstream_url); + config.db_url = Some(format!("sqlite://{}", directory.path().join("history.db").display())); + config.tools.web_search.api_key = Some("local-search-key".to_owned()); + config.tools.web_search.base_url = Some(upstream_url); + let mut context = ExecutionContext::from_config(&config).await.unwrap(); + context.streaming_timeout = Duration::from_millis(50); + let context = Arc::new(context); + let mut state = common::test_state(&config); + state.exec_ctx = Arc::clone(&context); + let (gateway_url, gateway) = common::spawn_gateway(state).await; + let _gateway = Server(gateway); + let response = reqwest::Client::builder() + .timeout(Duration::from_secs(3)) + .build() + .unwrap() + .post(format!("{gateway_url}/v1/messages")) + .json(&messages_request()) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response + .text() + .await + .expect("a later stalled upstream error must terminate the Messages stream"); + let events: Vec = body + .lines() + .filter_map(|line| line.strip_prefix("data: ")) + .map(|data| serde_json::from_str(data).unwrap()) + .collect(); + assert_eq!(events.first().unwrap()["type"], "message_start"); + assert_eq!(events.last().unwrap()["type"], "error"); + assert_eq!(events.iter().filter(|event| event["type"] == "error").count(), 1); + assert!(!events.iter().any(|event| event["type"] == "message_stop")); + wait_for_body_drop(&dropped).await; + assert_eq!(searches.load(Ordering::SeqCst), 1); + let requests = requests.lock().await; + assert_eq!(requests.len(), 2); + let continuation = requests[1]["messages"].to_string(); + assert!(continuation.contains("call_search")); + assert!(continuation.contains("sunny proof")); + assert!(continuation.contains("tool_result")); + context.storage_pool().unwrap().close().await; +}