From 7721a181aa24db7bb80656b5abd28ea76ad36ab7 Mon Sep 17 00:00:00 2001 From: Jake Barnby Date: Mon, 5 Oct 2026 17:59:09 +1300 Subject: [PATCH 1/3] fix(comfy): time out training at the deadline when ComfyUI is slow wait_prompt cut each history poll's request timeout to whatever was left of the training deadline. When ComfyUI was too slow to answer the last poll in time, reqwest's own timeout fired first and the job failed with "error sending request for url (.../history/)" instead of "training timed out". Under load the in-flight poll is the usual place the deadline lands, which is how a_graph_that_never_finishes_times_out failed. Polls now keep their own five-second request timeout and the whole poll runs under timeout_at(deadline), so whichever of the deadline or a real request failure comes first is reported as what it is. cancel_and_wait bounds its polls the same way. The timeout tests now mount a ComfyUI that marks the prompt terminal once it is cancelled, so they assert the cancellation was confirmed and no longer spend thirty seconds polling a graph that never ends. Co-Authored-By: Claude Opus 5.5 --- runner/zone_comfy/src/train.rs | 152 +++++++++++++++++++++++---------- 1 file changed, 109 insertions(+), 43 deletions(-) diff --git a/runner/zone_comfy/src/train.rs b/runner/zone_comfy/src/train.rs index 748c0968..83719499 100644 --- a/runner/zone_comfy/src/train.rs +++ b/runner/zone_comfy/src/train.rs @@ -20,6 +20,7 @@ const MIN_WEIGHT_BYTES: usize = 10_000; const MANIFEST_VERSION: u32 = 1; const CANCEL_TIMEOUT: Duration = Duration::from_secs(30); const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); +const TIMED_OUT: &str = "training timed out"; #[derive(Debug, Deserialize)] pub struct TrainConfig { @@ -897,25 +898,17 @@ async fn wait_prompt( let deadline = tokio::time::Instant::now() + timeout; loop { if tokio::time::Instant::now() >= deadline { - let cleanup = cancel_and_wait(client, config, prompt, false).await; - return Err(WaitFailure { - error: TrainError::Failed("training timed out".into()), - cleanup, - }); + return Err(abandon(client, config, prompt, timed_out()).await); } - let history = match history(client, config, prompt, deadline).await { - Ok(history) => history, - Err(error) => { - let cleanup = cancel_and_wait(client, config, prompt, false).await; - return Err(WaitFailure { error, cleanup }); - } + let history = match tokio::time::timeout_at(deadline, history(client, config, prompt)).await + { + Ok(Ok(history)) => history, + Ok(Err(error)) => return Err(abandon(client, config, prompt, error).await), + Err(_) => return Err(abandon(client, config, prompt, timed_out()).await), }; let entry = match exact_history(&history, prompt) { Ok(entry) => entry, - Err(error) => { - let cleanup = cancel_and_wait(client, config, prompt, false).await; - return Err(WaitFailure { error, cleanup }); - } + Err(error) => return Err(abandon(client, config, prompt, error).await), }; if let Some(entry) = entry { match train_prompt_complete(entry) { @@ -964,20 +957,30 @@ async fn prompt_progress( parse_progress(&response.bytes().await.ok()?) } +async fn abandon( + client: &reqwest::Client, + config: &Config, + prompt: Uuid, + error: TrainError, +) -> WaitFailure { + let cleanup = cancel_and_wait(client, config, prompt, false).await; + WaitFailure { error, cleanup } +} + +fn timed_out() -> TrainError { + TrainError::Failed(TIMED_OUT.into()) +} + async fn history( client: &reqwest::Client, config: &Config, prompt: Uuid, - deadline: tokio::time::Instant, ) -> Result { - let timeout = deadline - .saturating_duration_since(tokio::time::Instant::now()) - .min(REQUEST_TIMEOUT); authorize( config, client .get(format!("{}/history/{prompt}", config.base_url)) - .timeout(timeout), + .timeout(REQUEST_TIMEOUT), ) .send() .await @@ -1069,7 +1072,8 @@ async fn cancel_and_wait( if tokio::time::Instant::now() >= deadline { return false; } - if let Ok(history) = history(client, config, prompt, deadline).await + if let Ok(Ok(history)) = + tokio::time::timeout_at(deadline, history(client, config, prompt)).await && let Ok(Some(entry)) = exact_history(&history, prompt) && train_prompt_terminal(entry) { @@ -1284,6 +1288,8 @@ fn authorize(config: &Config, request: reqwest::RequestBuilder) -> reqwest::Requ mod tests { use super::*; use crate::recipe::{Recipe, RecipeCatalog}; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use wiremock::matchers::{method, path, query_param}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -1635,26 +1641,65 @@ mod tests { ); } + #[derive(Clone)] + struct RunningUntilCancelled { + prompt: Uuid, + delay: Duration, + cancelled: Arc, + } + + impl RunningUntilCancelled { + async fn mount(server: &MockServer, prompt: Uuid, delay: Duration) { + let running = Self { + prompt, + delay, + cancelled: Default::default(), + }; + Mock::given(method("GET")) + .and(path(format!("/history/{prompt}"))) + .respond_with(running.clone()) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path(format!("/api/jobs/{prompt}/cancel"))) + .respond_with(Cancel(running.cancelled)) + .mount(server) + .await; + } + } + + impl wiremock::Respond for RunningUntilCancelled { + fn respond(&self, _request: &wiremock::Request) -> ResponseTemplate { + if self.cancelled.load(Ordering::SeqCst) { + return ResponseTemplate::new(200).set_body_json(json!({ + self.prompt.to_string(): {"status": {"status_str": "error"}} + })); + } + ResponseTemplate::new(200) + .set_body_json(json!({ + self.prompt.to_string(): {"status": {"status_str": "running"}} + })) + .set_delay(self.delay) + } + } + + struct Cancel(Arc); + + impl wiremock::Respond for Cancel { + fn respond(&self, _request: &wiremock::Request) -> ResponseTemplate { + self.0.store(true, Ordering::SeqCst); + ResponseTemplate::new(200).set_body_json(json!({"cancelled": true})) + } + } + #[tokio::test] async fn a_graph_that_never_finishes_times_out() { let server = MockServer::start().await; let prompt = Uuid::new_v4(); - uploads(&server).await; - queues(&server, json!({"prompt_id": prompt, "number": 1})).await; - Mock::given(method("GET")) - .and(path(format!("/history/{prompt}"))) - .respond_with( - ResponseTemplate::new(200).set_body_json( - json!({prompt.to_string(): {"status": {"status_str": "running"}}}), - ), - ) - .mount(&server) - .await; + RunningUntilCancelled::mount(&server, prompt, Duration::ZERO).await; - let work = dataset(); - let client = reqwest::Client::new(); let failure = wait_prompt( - &client, + &reqwest::Client::new(), &config(&server), prompt, Duration::from_millis(120), @@ -1664,11 +1709,35 @@ mod tests { .await .unwrap_err(); assert!( - matches!(&failure.error, TrainError::Failed(message) if message.contains("timed out")), + matches!(&failure.error, TrainError::Failed(message) if message == TIMED_OUT), "{}", failure.error ); - drop(work); + assert!(failure.cleanup, "the timed-out graph was cancelled"); + } + + #[tokio::test] + async fn a_comfyui_too_slow_to_answer_still_times_out_at_the_deadline() { + let server = MockServer::start().await; + let prompt = Uuid::new_v4(); + RunningUntilCancelled::mount(&server, prompt, Duration::from_secs(30)).await; + + let failure = wait_prompt( + &reqwest::Client::new(), + &config(&server), + prompt, + Duration::from_millis(200), + None, + None, + ) + .await + .unwrap_err(); + assert!( + matches!(&failure.error, TrainError::Failed(message) if message == TIMED_OUT), + "{}", + failure.error + ); + assert!(failure.cleanup, "the timed-out graph was cancelled"); } #[tokio::test] @@ -2182,10 +2251,7 @@ mod tests { let artifact = "zone-lora-progress"; Mock::given(method("GET")) .and(path(format!("/history/{prompt}"))) - .respond_with(RunningThenDone( - std::sync::atomic::AtomicUsize::new(0), - prompt, - )) + .respond_with(RunningThenDone(AtomicUsize::new(0), prompt)) .mount(&server) .await; Mock::given(method("GET")) @@ -2213,11 +2279,11 @@ mod tests { assert_eq!(update, TrainProgress::new(12, 400)); } - struct RunningThenDone(std::sync::atomic::AtomicUsize, Uuid); + struct RunningThenDone(AtomicUsize, Uuid); impl wiremock::Respond for RunningThenDone { fn respond(&self, _request: &wiremock::Request) -> ResponseTemplate { - let seen = self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let seen = self.0.fetch_add(1, Ordering::SeqCst); let status = if seen == 0 { json!({"status_str": "running"}) } else { From df3416b809d22e61af04562a75de956f7ab259d2 Mon Sep 17 00:00:00 2001 From: Jake Barnby Date: Mon, 5 Oct 2026 18:04:19 +1300 Subject: [PATCH 2/3] test(chat): renew the catalog test's lease like a generation does catalog_pages_all_references_without_exposing_result_bodies_or_other_chats acquired a 30-second lease and never renewed it, then made several hundred store round trips paging the catalog seven characters at a time. It takes about a second on its own, but with every test binary sharing one Postgres under load it took 39 seconds and the final delete_message failed with LeaseLost. Production never holds a lease without keep_alive (session.rs pairs every acquire with it), so the test now does the same and also asserts the renewal never lost the lease. What it asserts about paging, redaction, scoping and invalidation is unchanged. Co-Authored-By: Claude Opus 5.5 --- runner/zone_server/tests/chat_context_store_tests.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/runner/zone_server/tests/chat_context_store_tests.rs b/runner/zone_server/tests/chat_context_store_tests.rs index 1fb41d5c..8a58e5e1 100644 --- a/runner/zone_server/tests/chat_context_store_tests.rs +++ b/runner/zone_server/tests/chat_context_store_tests.rs @@ -780,6 +780,7 @@ async fn preview_does_not_initialize_mcp_in_real_application_state() { async fn catalog_pages_all_references_without_exposing_result_bodies_or_other_chats() { let (pool, store, chat, workspace) = fixture().await; let lease = store.acquire(Uuid::new_v4(), LIFETIME).await.unwrap(); + let mut guard = store.keep_alive(lease.clone(), LIFETIME).unwrap(); let (turn, user) = begin(&store, &lease).await; for index in 0..40 { let call = format!("call-{index}"); @@ -887,6 +888,11 @@ async fn catalog_pages_all_references_without_exposing_result_bodies_or_other_ch .is_err(), "Deletion invalidates the snapshot rather than shifting pages" ); + assert!( + !guard.is_lost(), + "renewal kept the lease through the whole catalog walk" + ); + guard.stop().await; store.release(&lease).await.unwrap(); sqlx::query("DELETE FROM chats WHERE id=$1") .bind(chat) From 5935179f0c910b5eed9b38be91df8b17b1317499 Mon Sep 17 00:00:00 2001 From: Jake Barnby Date: Mon, 5 Oct 2026 15:53:35 +1300 Subject: [PATCH 3/3] refactor(chat): box the route a chat turn carries Clippy 1.98 flags Routing as a large enum: the Chat variant's Route makes every Routing at least 856 bytes, 216 more than the next largest variant, and clippy -D warnings fails zone_server on main. Boxing the route keeps the enum near the media variants' size; it is built once per turn, so the allocation costs nothing that matters. Co-Authored-By: Claude Opus 5.5 --- runner/zone_server/src/ws/chat.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/runner/zone_server/src/ws/chat.rs b/runner/zone_server/src/ws/chat.rs index 70cc1d5d..df188670 100644 --- a/runner/zone_server/src/ws/chat.rs +++ b/runner/zone_server/src/ws/chat.rs @@ -388,7 +388,7 @@ enum Routing { Video(crate::config::ComfyUiConfig), Audio(crate::config::ComfyUiConfig), Upscale(crate::config::ComfyUiConfig), - Chat(chats::ChatRow, Route), + Chat(chats::ChatRow, Box), } /// What a provider reported when a turn failed, remedied for a coding agent's @@ -2495,7 +2495,7 @@ async fn handle_send_message( return Ok(()); } _ = session.guard.lost() => { return Err(OWNERSHIP_LOST.into()); } - result = prepare_chat(state, stream, chat_id, workspace_id, user_id, content, metadata.as_ref(), chat, route, web_search_requested) => result?, + result = prepare_chat(state, stream, chat_id, workspace_id, user_id, content, metadata.as_ref(), chat, *route, web_search_requested) => result?, }; handle_chat_generation(state, stream, chat_id, workspace_id, user_id, preparation, &mut request, &mut session, &mut jobs).await } @@ -2666,7 +2666,9 @@ async fn prepare_message( } crate::services::image_intent::GenerationIntent::Audio => Routing::Audio(image_config), crate::services::image_intent::GenerationIntent::Upscale => Routing::Upscale(image_config), - crate::services::image_intent::GenerationIntent::Chat => Routing::Chat(chat, route), + crate::services::image_intent::GenerationIntent::Chat => { + Routing::Chat(chat, Box::new(route)) + } }) } @@ -4569,7 +4571,7 @@ mod tests { QUESTION, None, chat, - route, + *route, false, ) .await