From 0d6f558b2bd0282a7a1695f7680ab1d1c6142579 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Mon, 13 Jul 2026 09:44:28 +0700 Subject: [PATCH] feat: enhance responsiveness by implementing abort checks in streaming API calls --- src/app/runtime/actions/mod.rs | 64 +++++++++++++++++++--------------- src/app/subagent/engine.rs | 45 +++++++++++++++++++----- src/app/workflow/engine.rs | 2 +- 3 files changed, 73 insertions(+), 38 deletions(-) diff --git a/src/app/runtime/actions/mod.rs b/src/app/runtime/actions/mod.rs index bce9b5f..7d2376e 100644 --- a/src/app/runtime/actions/mod.rs +++ b/src/app/runtime/actions/mod.rs @@ -1166,7 +1166,11 @@ fn run_agent_turn( let token_estimate = total_chars / 4; let max_wire_tokens = tc.context_window; - let wire_msgs = if crate::app::runtime::shortsend::should_shape(token_estimate, max_wire_tokens, prev_shaped) { + // Skip message compaction if abort was requested — the non-streaming + // LLM call for summarization would block without checking abort_flag. + let wire_msgs = if !tc.abort_flag.load(std::sync::atomic::Ordering::SeqCst) + && crate::app::runtime::shortsend::should_shape(token_estimate, max_wire_tokens, prev_shaped) + { prev_shaped = true; let compacted = crate::app::runtime::shortsend::shape_messages(&msgs, token_estimate, max_wire_tokens, false, Some(&tc.client)); @@ -1240,42 +1244,44 @@ fn run_agent_turn( let (response, final_usage) = match result { Ok((msg, u)) => (msg, u.or(usage)), Err(e) => { + // If abort was requested, return immediately. if tc.abort_flag.load(std::sync::atomic::Ordering::SeqCst) || e.to_string().contains("aborted") { if let Ok(mut q) = events_q.lock() { q.push_back(TurnEvent::Error("Generation aborted by user".to_string())); } return Ok(()); } - match tc.client.chat_with_tools_non_streaming(&wire_msgs, Some(tc.tdefs.clone())) { - Ok((msg, usage_fb)) => (msg, usage_fb), - Err(api_err) => { - let todo_path = tc.ctx.session_dir.join("todo.md"); - let mut has_unfinished = false; - if let Ok(todo_text) = std::fs::read_to_string(&todo_path) { - if todo_text.lines().any(|l| l.trim_start().starts_with("- [ ]")) { - has_unfinished = true; - } - } - if has_unfinished { - todo_retry_count += 1; - if todo_retry_count > MAX_TODO_RETRIES { - anyhow::bail!( - "exhausted {MAX_TODO_RETRIES} todo-retries — giving up on unfinished tasks. \ - Edit todo.md manually or ask me to focus on specific items.", - ); - } - if let Ok(mut q) = events_q.lock() { - q.push_back(TurnEvent::SystemNote { - kind: "task_retry".to_string(), - message: format!("Network/API error: {api_err}. Auto-retrying to finish tasks... (retry {todo_retry_count}/{MAX_TODO_RETRIES})"), - }); - } - std::thread::sleep(std::time::Duration::from_secs(5)); - continue; - } - return Err(api_err); + // Streaming-only: no non-streaming fallback. + // Non-streaming blocks up to 1 minute without checking + // abort_flag, making cancellation unresponsive. + // If the API supports streaming (which it must), this + // path handles transient errors via the retry loop below. + let api_err = e; + let todo_path = tc.ctx.session_dir.join("todo.md"); + let mut has_unfinished = false; + if let Ok(todo_text) = std::fs::read_to_string(&todo_path) { + if todo_text.lines().any(|l| l.trim_start().starts_with("- [ ]")) { + has_unfinished = true; } } + if has_unfinished { + todo_retry_count += 1; + if todo_retry_count > MAX_TODO_RETRIES { + anyhow::bail!( + "exhausted {MAX_TODO_RETRIES} todo-retries — giving up on unfinished tasks. \ + Edit todo.md manually or ask me to focus on specific items.", + ); + } + if let Ok(mut q) = events_q.lock() { + q.push_back(TurnEvent::SystemNote { + kind: "task_retry".to_string(), + message: format!("Network/API error: {api_err}. Auto-retrying to finish tasks... (retry {todo_retry_count}/{MAX_TODO_RETRIES})"), + }); + } + std::thread::sleep(std::time::Duration::from_secs(5)); + continue; + } + return Err(api_err); } }; diff --git a/src/app/subagent/engine.rs b/src/app/subagent/engine.rs index d6b37a3..d91a6bb 100644 --- a/src/app/subagent/engine.rs +++ b/src/app/subagent/engine.rs @@ -278,10 +278,10 @@ fn generate_workspace_tree(roots: &[std::path::PathBuf]) -> String { /// /// Flow: inject system prompt (with workspace tree if available) → for each /// step: resolve provider config, build an LLM client, call -/// `chat_with_tools_non_streaming`, process tool calls (gated against both -/// the allowlist and Harness-style content safety checks) or collect text -/// output → send `SubagentEvent`s on `tx` → break on first text-only -/// (non-empty) response. +/// `chat_with_tools_streaming` (with abort check per SSE event), process +/// tool calls (gated against both the allowlist and Harness-style content +/// safety checks) or collect text output → send `SubagentEvent`s on `tx` → +/// break on first text-only (non-empty) response. /// /// Why: runs synchronously on a dedicated thread so the main async event /// loop is not blocked. Tool gating prevents restricted, risky, or @@ -333,15 +333,44 @@ pub fn run_subagent(ctx: &SubagentContext, tx: &mpsc::Sender) -> anyhow::bail!("subagent aborted by parent at step {step}"); } - // Use the structured tool-calling API so the LLM can request tools with - // proper arguments, exactly like the main agent does. - let (response, _usage) = match client.chat_with_tools_non_streaming(&messages, tdefs_opt.clone()) { + // Use streaming API so the abort flag is checked per SSE event, + // making the subagent responsive to cancellation even during an + // LLM call (non-streaming would block for 10-30s unchecked). + let stream_result = client.chat_with_tools_streaming( + &messages, + tdefs_opt.clone(), + Some(0.7), + Some(4096), + |_event| -> bool { + // Check abort on every SSE event for responsive cancellation. + if ctx.abort_flag.as_ref().is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) { + return false; // signals provider to abort + } + // We don't stream tokens to the UI for subagents — just + // need the assembled message at the end. + true + }, + ); + + let (response, _usage) = match stream_result { Ok(result) => result, Err(e) => { + let is_abort = ctx.abort_flag.as_ref().is_some_and(|f| f.load(std::sync::atomic::Ordering::SeqCst)) + || e.to_string().contains("aborted"); let _ = tx.blocking_send(SubagentEvent::StepFailed { step, - error: e.to_string(), + error: if is_abort { + "subagent aborted by user".to_string() + } else { + e.to_string() + }, }); + if is_abort { + anyhow::bail!("subagent aborted by parent at step {step}"); + } + // No non-streaming fallback — API must support streaming. + // Non-streaming calls block for up to 1 min without checking + // abort_flag, making cancellation unresponsive. anyhow::bail!("subagent call failed at step {step}: {e}"); } }; diff --git a/src/app/workflow/engine.rs b/src/app/workflow/engine.rs index 9357164..fcf2e59 100644 --- a/src/app/workflow/engine.rs +++ b/src/app/workflow/engine.rs @@ -237,7 +237,7 @@ fn spawn_single_agent( let _ = done_tx.send(run_subagent(&bg_ctx, &bg_tx)); }); - let poll_interval = Duration::from_millis(500); + let poll_interval = Duration::from_millis(200); let result = if let Some(timeout) = timeout_ms { let deadline = Duration::from_millis(timeout); let mut elapsed = Duration::ZERO;