feat: enhance responsiveness by implementing abort checks in streaming API calls
This commit is contained in:
@@ -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);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -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}");
|
||||
}
|
||||
};
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user