feat: enhance responsiveness by implementing abort checks in streaming API calls

This commit is contained in:
asepharyana
2026-07-13 14:39:39 +07:00
parent 8388a83af0
commit 0d6f558b2b
3 changed files with 73 additions and 38 deletions
+35 -29
View File
@@ -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);
}
};
+37 -8
View File
@@ -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<SubagentEvent>) ->
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}");
}
};
+1 -1
View File
@@ -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;