From 2a86810d7367b53ccd8c74b77452eb9a1c370291 Mon Sep 17 00:00:00 2001 From: xuezhixin Date: Wed, 5 Aug 2026 11:44:59 +0800 Subject: [PATCH 1/3] feat(tools): harden read-only bash classifier - fd redirects: treat `>&N`/`>&-` as read-only fd duplication, but `>&word` (non-digit, e.g. `echo x >& /tmp/f`) as a write redirect. Same distinction applied in split_compound_segments and has_background_operator so `2>&1` is not split or misread as a background job. - globs: allow unquoted globs; writes (`rm *.txt`, `cp *.x /d`) are still caught by the blocked-command list and `echo *.x > f` by the write-redirect check. - find -exec: recursively classify the embedded command so `find -exec grep {} \;` passes and `find -exec rm {} \;` is blocked; `-delete` stays destructive. Handles `\;` and `+` terminators. Closes #428 --- crates/aish-tools/src/bash/read_only.rs | 126 +++++++++++++++++++++--- 1 file changed, 110 insertions(+), 16 deletions(-) diff --git a/crates/aish-tools/src/bash/read_only.rs b/crates/aish-tools/src/bash/read_only.rs index 92555c60..6c45b015 100644 --- a/crates/aish-tools/src/bash/read_only.rs +++ b/crates/aish-tools/src/bash/read_only.rs @@ -111,6 +111,12 @@ fn split_compound_segments(command: &str) -> Vec { if !in_single_quote && !in_double_quote { if ch == '&' { + // `>&N` / `<&N` duplicate a file descriptor; the `&` is part of + // the redirect token, not a command separator. + if current.ends_with('>') || current.ends_with('<') { + current.push(ch); + continue; + } if chars.peek() == Some(&'&') { chars.next(); segments.push(current.clone()); @@ -153,27 +159,39 @@ fn has_background_operator(command: &str) -> bool { let mut in_single_quote = false; let mut in_double_quote = false; let mut chars = command.chars().peekable(); + let mut prev: Option = None; while let Some(ch) = chars.next() { if ch == '\'' && !in_double_quote { in_single_quote = !in_single_quote; + prev = Some(ch); continue; } if ch == '"' && !in_single_quote { in_double_quote = !in_double_quote; + prev = Some(ch); continue; } if ch == '\\' && !in_single_quote { chars.next(); + prev = None; continue; } if !in_single_quote && !in_double_quote && ch == '&' { + // `>&N` / `<&N` duplicate a file descriptor (e.g. `2>&1`, `<&2`); + // the `&` is part of the redirect, not a background operator. + if prev == Some('>') || prev == Some('<') { + prev = Some(ch); + continue; + } if chars.peek() == Some(&'&') { chars.next(); + prev = Some('&'); continue; } return true; } + prev = Some(ch); } false @@ -188,10 +206,18 @@ fn non_readonly_segment_reason(segment: &str) -> Option { }) { return Some("command or process substitution".into()); } - if scan_outside_quotes(segment, |window| window.starts_with('*')) { - return Some("unquoted glob".into()); - } if scan_outside_quotes(segment, |window| { + // `>&N` (e.g. `2>&1`) duplicates a file descriptor and `>&-` closes + // it; both are read-only. But `>&word` where `word` is NOT a digit + // or `-` redirects stdout to a file (e.g. `echo x >& /tmp/f` writes + // /tmp/f in bash), which IS a write — let it fall through to the `>` + // check below instead of blanket-allowing every `>&`. + if window.starts_with(">&") { + let after = window[2..].chars().next(); + if matches!(after, Some(c) if c.is_ascii_digit()) || after == Some('-') { + return false; + } + } if let Some(rest) = window.strip_prefix(">>") { return !rest.trim_start().starts_with("/dev/null"); } @@ -242,15 +268,10 @@ fn non_readonly_segment_reason(segment: &str) -> Option { } } - if base == "find" - && tokens.iter().any(|token| { - matches!( - token.to_ascii_lowercase().as_str(), - "-delete" | "-exec" | "-execdir" | "-ok" | "-okdir" - ) - }) - { - return Some("find mutating or command execution".into()); + if base == "find" { + if let Some(reason) = find_action_not_readonly(&tokens) { + return Some(reason); + } } if base == "curl" && curl_writes_or_mutates(segment) { @@ -265,6 +286,44 @@ fn non_readonly_segment_reason(segment: &str) -> Option { None } +/// Judge `find` action flags. `-delete` is inherently destructive; `-exec`, +/// `-execdir`, `-ok` and `-okdir` run an embedded command whose read-only +/// status is determined by recursively classifying it. This lets +/// `find ... -exec grep/wc/head {} \;` through while still blocking +/// `find ... -exec rm {} \;`. +fn find_action_not_readonly(tokens: &[String]) -> Option { + const EXEC_FLAGS: [&str; 4] = ["-exec", "-execdir", "-ok", "-okdir"]; + let mut index = 0; + while index < tokens.len() { + let flag = tokens[index].to_ascii_lowercase(); + if flag == "-delete" { + return Some("find -delete".into()); + } + if EXEC_FLAGS.contains(&flag.as_str()) { + // Collect the embedded command up to the terminator `\;` or `+`. + let mut end = index + 1; + while end < tokens.len() { + let term = normalize_shell_word(&tokens[end]); + if term == ";" || term == "+" { + break; + } + end += 1; + } + let embedded = tokens[index + 1..end].join(" "); + if embedded.trim().is_empty() { + return Some("find -exec without command".into()); + } + if !matches!(classify(&embedded), ReadOnlyVerdict::ReadOnly) { + return Some("find -exec non-read-only command".into()); + } + index = end + 1; + } else { + index += 1; + } + } + None +} + fn curl_writes_or_mutates(segment: &str) -> bool { let compact = segment .chars() @@ -698,8 +757,13 @@ mod tests { } #[test] - fn blocks_unquoted_glob() { - assert_not_read_only("ls *.txt"); + fn allows_unquoted_glob() { + // Glob expansion does not mutate the filesystem; write commands + // (`rm *.txt`, `cp *.x /d`) are still caught by the BLOCKED list and + // redirects (`echo *.x > f`) by the write-redirect check. + assert_read_only("ls *.txt"); + assert_read_only("grep -r foo *.md"); + assert_read_only("cat *.log"); } #[test] @@ -727,11 +791,41 @@ mod tests { fn blocks_trailing_background_operator() { assert_not_read_only("sleep 9999 &"); } - #[test] - fn blocks_find_exec_and_ok() { + fn fd_redirects_are_read_only() { + // `2>&1`, `>&2`, `<&2` duplicate file descriptors; they are read-only + // and must not be misclassified as background jobs or write redirects. + assert_read_only("ls 2>&1"); + assert_read_only("cat foo 2>&1"); + assert_read_only("grep bar file 2>&1 | head"); + assert_read_only("echo hi >&2"); + assert_read_only("echo hi >&-"); + assert_read_only("cmd <&2"); + // Real background operators and write redirects are still caught. + assert_not_read_only("sleep 10 &"); + assert_not_read_only("cmd > file &"); + assert_not_read_only("echo hi > file"); + // `>&word` where `word` is NOT a digit redirects stdout to a file + // (verified in bash: `echo x >& /tmp/f` writes the file). These must + // NOT slip through the `>&N` allowance. + assert_not_read_only("echo x >& /tmp/f"); + assert_not_read_only("echo x >&/tmp/f"); + assert_not_read_only("echo x >& f"); + } + + #[test] + fn blocks_find_exec_mutating_and_allows_read_only() { + // -exec/-ok/-execdir running a mutating command -> blocked assert_not_read_only("find . -exec rm {} \\;"); assert_not_read_only("find . -ok rm {} \\;"); + assert_not_read_only("find . -execdir mv {} /tmp \\;"); + // -exec/-ok running a read-only command -> allowed + assert_read_only("find . -exec grep foo {} \\;"); + assert_read_only("find . -exec wc -l {} \\;"); + assert_read_only("find . -exec head -n 1 {} \\;"); + assert_read_only("find . -ok grep foo {} \\;"); + // -delete stays destructive + assert_not_read_only("find . -name '*.tmp' -delete"); } #[test] From 4ca691cc2aba9fac860cf1f628adb98c91b8e2e5 Mon Sep 17 00:00:00 2001 From: xuezhixin Date: Wed, 5 Aug 2026 11:45:20 +0800 Subject: [PATCH 2/3] feat(llm): execute Agent sub-agent calls in parallel When a tool-call batch is entirely `Agent` calls and len > 1, dispatch them concurrently via futures::join_all; all other batches stay sequential to preserve stop-on-first-short-circuit semantics. Each sub-agent is an isolated SubSession (own history, tools, cancellation, counters) so there is no shared mutable parent state under the concurrent execute_tool(&self) calls. - session.rs: extract process_tool_call_result (shared by both loops); add run_tool_calls with the parallel/sequential branch. - agent_tool.rs: wrap each sub-agent run in catch_unwind so a panicking sub-session degrades to ToolResult::error instead of unwinding through join_all and aborting its siblings. - app.rs: replace the AtomicBool sub-agent flag with an AtomicU32 active count (saturating decrement; ignore sub-agent OpEnd) so the spinner survives a fast-finishing sibling. - prompt.rs: steer the model to fan out multiple independent Agent calls in one response. - ApprovalMemory migrated to parking_lot::Mutex (no poisoning) so an unwound sub-agent cannot poison a lock shared via the parent and cascade a panic. Tests: join_all+zip preserves tool_call id order; concurrency probe proves real overlap (peak >= 2). Closes #427 --- Cargo.lock | 1 + crates/aish-llm/Cargo.toml | 1 + crates/aish-llm/src/session.rs | 521 ++++++++++++------ crates/aish-shell/src/app.rs | 35 +- .../aish-tools/src/agent_tool/agent_tool.rs | 53 +- crates/aish-tools/src/agent_tool/prompt.rs | 11 +- 6 files changed, 437 insertions(+), 185 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 43d3a82b..da82b9a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -128,6 +128,7 @@ dependencies = [ "chrono", "futures", "langfuse-ergonomic", + "parking_lot", "rand 0.9.4", "reqwest", "serde", diff --git a/crates/aish-llm/Cargo.toml b/crates/aish-llm/Cargo.toml index 763b9dc7..7fe51d99 100644 --- a/crates/aish-llm/Cargo.toml +++ b/crates/aish-llm/Cargo.toml @@ -16,6 +16,7 @@ tokio.workspace = true tracing.workspace = true uuid.workspace = true futures.workspace = true +parking_lot.workspace = true bytes.workspace = true sha2.workspace = true base64.workspace = true diff --git a/crates/aish-llm/src/session.rs b/crates/aish-llm/src/session.rs index 938e4c4d..579493c6 100644 --- a/crates/aish-llm/src/session.rs +++ b/crates/aish-llm/src/session.rs @@ -58,7 +58,7 @@ pub struct LlmSession { /// Session-scoped approval memory. When set, commands the user approved /// with "remember" skip confirmation (and the sandbox preflight) for the /// rest of the session on the same host. - approval_memory: Option>>, + approval_memory: Option>>, /// Callback invoked when the tool-call iteration limit is reached. /// Receives the current iteration count and returns true to reset and continue. iteration_limit_callback: Option bool + Send + Sync>>, @@ -239,7 +239,7 @@ impl LlmSession { /// Install session-scoped approval memory. When present, commands approved /// with "remember" skip confirmation (and sandbox preflight) for the rest /// of the session. - pub fn set_approval_memory(&mut self, memory: Arc>) { + pub fn set_approval_memory(&mut self, memory: Arc>) { self.approval_memory = Some(memory); } @@ -844,89 +844,17 @@ impl LlmSession { messages.push(chat_msg); } - // Execute each tool call and append results - for tc in &tool_calls { - let result = self.execute_tool(tc).await; - let short_circuit = is_short_circuit_result(&result); - let output = result.output.clone(); - - // Log tool call span to Langfuse - if let (Some(ref langfuse), Some(ref tid)) = (&self.langfuse, &trace_id) { - langfuse - .span_tool_call(tid, &tc.name, &tc.arguments, &output, 0) - .await; - } - if short_circuit { - self.emit_event(LlmEvent { - event_type: LlmEventType::GenerationEnd, - data: serde_json::json!({}), - timestamp: now_timestamp(), - metadata: None, - }); - self.emit_event(LlmEvent { - event_type: LlmEventType::OpEnd, - data: serde_json::json!({"reason": "short_circuit"}), - timestamp: now_timestamp(), - metadata: None, - }); - // Sub-agent cancel is shown by the shell as `shell.interrupted` - // (same as Ctrl+C); do not return the tool string as AI body. - let suppress_body = result.meta.as_ref().is_some_and(|meta| { - matches!( - meta.get("reason").and_then(|v| v.as_str()), - Some("sub_agent_cancelled") | Some("user_cancelled") - ) - }); - let text = if self.security_notice_callback.is_some() || suppress_body { - String::new() - } else { - output - }; - let new_messages = messages[initial_len..].to_vec(); - return Ok(crate::types::ProcessResult { text, new_messages }); - } - - // Track consecutive failures for early termination. - // short_circuit results (security blocked) are excluded. - if result.ok { - consecutive_failures = 0; - } else { - consecutive_failures += 1; - } - if consecutive_failures >= MAX_CONSECUTIVE_FAILURES { - tracing::warn!( - consecutive_failures, - "Too many consecutive tool failures, stopping loop" - ); - self.emit_event(LlmEvent { - event_type: LlmEventType::Error, - data: serde_json::json!({ - "error": format!( - "Stopped: {} consecutive tool failures", - consecutive_failures - ), - "consecutive_failures": consecutive_failures, - }), - timestamp: now_timestamp(), - metadata: None, - }); - self.emit_event(LlmEvent { - event_type: LlmEventType::OpEnd, - data: serde_json::json!({"reason": "consecutive_failures"}), - timestamp: now_timestamp(), - metadata: None, - }); - messages.push(ChatMessage::tool_result(&tc.id, output)); - let text = format!( - "Stopped after {} consecutive tool execution failures. \ - Please check your connection and retry.", - consecutive_failures - ); - let new_messages = messages[initial_len..].to_vec(); - return Ok(crate::types::ProcessResult { text, new_messages }); - } - - messages.push(ChatMessage::tool_result(&tc.id, output)); + if let Some(pr) = self + .run_tool_calls( + &tool_calls, + &mut messages, + &mut consecutive_failures, + &trace_id, + initial_len, + ) + .await + { + return Ok(pr); } // Trim old tool-call rounds to prevent unbounded growth @@ -1242,89 +1170,17 @@ impl LlmSession { } messages.push(assistant_msg); - // Execute tools - for tc in &tool_calls { - let result = self.execute_tool(tc).await; - let short_circuit = is_short_circuit_result(&result); - let output = result.output.clone(); - - // Log tool call span to Langfuse - if let (Some(ref langfuse), Some(ref tid)) = (&self.langfuse, &trace_id) { - langfuse - .span_tool_call(tid, &tc.name, &tc.arguments, &output, 0) - .await; - } - if short_circuit { - self.emit_event(LlmEvent { - event_type: LlmEventType::GenerationEnd, - data: serde_json::json!({}), - timestamp: now_timestamp(), - metadata: None, - }); - self.emit_event(LlmEvent { - event_type: LlmEventType::OpEnd, - data: serde_json::json!({"reason": "short_circuit"}), - timestamp: now_timestamp(), - metadata: None, - }); - // Sub-agent cancel is shown by the shell as `shell.interrupted` - // (same as Ctrl+C); do not return the tool string as AI body. - let suppress_body = result.meta.as_ref().is_some_and(|meta| { - matches!( - meta.get("reason").and_then(|v| v.as_str()), - Some("sub_agent_cancelled") | Some("user_cancelled") - ) - }); - let text = if self.security_notice_callback.is_some() || suppress_body { - String::new() - } else { - output - }; - let new_messages = messages[initial_len..].to_vec(); - return Ok(crate::types::ProcessResult { text, new_messages }); - } - - // Track consecutive failures for early termination. - // short_circuit results (security blocked) are excluded. - if result.ok { - consecutive_failures = 0; - } else { - consecutive_failures += 1; - } - if consecutive_failures >= MAX_CONSECUTIVE_FAILURES { - tracing::warn!( - consecutive_failures, - "Too many consecutive tool failures, stopping loop" - ); - self.emit_event(LlmEvent { - event_type: LlmEventType::Error, - data: serde_json::json!({ - "error": format!( - "Stopped: {} consecutive tool failures", - consecutive_failures - ), - "consecutive_failures": consecutive_failures, - }), - timestamp: now_timestamp(), - metadata: None, - }); - self.emit_event(LlmEvent { - event_type: LlmEventType::OpEnd, - data: serde_json::json!({"reason": "consecutive_failures"}), - timestamp: now_timestamp(), - metadata: None, - }); - messages.push(ChatMessage::tool_result(&tc.id, output)); - let text = format!( - "Stopped after {} consecutive tool execution failures. \ - Please check your connection and retry.", - consecutive_failures - ); - let new_messages = messages[initial_len..].to_vec(); - return Ok(crate::types::ProcessResult { text, new_messages }); - } - - messages.push(ChatMessage::tool_result(&tc.id, output)); + if let Some(pr) = self + .run_tool_calls( + &tool_calls, + &mut messages, + &mut consecutive_failures, + &trace_id, + initial_len, + ) + .await + { + return Ok(pr); } // Smart-trim old tool outputs to prevent unbounded growth @@ -1523,6 +1379,164 @@ impl LlmSession { ToolResult::error(format!("Unknown tool: {}", tool_call.name)) } } + /// Post-process one tool result inside the tool-calling loop: log the + /// Langfuse span, honor short-circuit (security block / cancel), track + /// consecutive failures, and append the `tool_result` message. + /// + /// Returns `Some(ProcessResult)` when the caller must return immediately + /// (short-circuit or failure threshold reached); `None` to keep looping. + async fn process_tool_call_result( + &self, + tc: &ToolCall, + result: ToolResult, + messages: &mut Vec, + consecutive_failures: &mut usize, + trace_id: &Option, + initial_len: usize, + ) -> Option { + let short_circuit = is_short_circuit_result(&result); + let output = result.output.clone(); + + // Log tool call span to Langfuse + if let (Some(ref langfuse), Some(ref tid)) = (&self.langfuse, trace_id) { + langfuse + .span_tool_call(tid, &tc.name, &tc.arguments, &output, 0) + .await; + } + if short_circuit { + self.emit_event(LlmEvent { + event_type: LlmEventType::GenerationEnd, + data: serde_json::json!({}), + timestamp: now_timestamp(), + metadata: None, + }); + self.emit_event(LlmEvent { + event_type: LlmEventType::OpEnd, + data: serde_json::json!({"reason": "short_circuit"}), + timestamp: now_timestamp(), + metadata: None, + }); + // Sub-agent cancel is shown by the shell as `shell.interrupted` + // (same as Ctrl+C); do not return the tool string as AI body. + let suppress_body = result.meta.as_ref().is_some_and(|meta| { + matches!( + meta.get("reason").and_then(|v| v.as_str()), + Some("sub_agent_cancelled") | Some("user_cancelled") + ) + }); + let text = if self.security_notice_callback.is_some() || suppress_body { + String::new() + } else { + output + }; + let new_messages = messages[initial_len..].to_vec(); + return Some(crate::types::ProcessResult { text, new_messages }); + } + + // Track consecutive failures for early termination. + // short_circuit results (security blocked) are excluded. + if result.ok { + *consecutive_failures = 0; + } else { + *consecutive_failures += 1; + } + if *consecutive_failures >= MAX_CONSECUTIVE_FAILURES { + tracing::warn!( + consecutive_failures = *consecutive_failures, + "Too many consecutive tool failures, stopping loop" + ); + self.emit_event(LlmEvent { + event_type: LlmEventType::Error, + data: serde_json::json!({ + "error": format!( + "Stopped: {} consecutive tool failures", + *consecutive_failures + ), + "consecutive_failures": *consecutive_failures, + }), + timestamp: now_timestamp(), + metadata: None, + }); + self.emit_event(LlmEvent { + event_type: LlmEventType::OpEnd, + data: serde_json::json!({"reason": "consecutive_failures"}), + timestamp: now_timestamp(), + metadata: None, + }); + messages.push(ChatMessage::tool_result(&tc.id, output)); + let text = format!( + "Stopped after {} consecutive tool execution failures. \ + Please check your connection and retry.", + *consecutive_failures + ); + let new_messages = messages[initial_len..].to_vec(); + return Some(crate::types::ProcessResult { text, new_messages }); + } + + messages.push(ChatMessage::tool_result(&tc.id, output)); + None + } + + /// Execute a batch of tool calls and post-process every result. + /// + /// When every call targets the `Agent` sub-agent tool and there is more + /// than one, the calls run **concurrently** — each `Agent` spawns an + /// isolated `SubSession` (independent tool registry, cancellation, and + /// event handling), so there are no shared mutable side effects between + /// them. All other batches execute **sequentially**, preserving the + /// original "stop on first short-circuit" semantics. + async fn run_tool_calls( + &self, + tool_calls: &[ToolCall], + messages: &mut Vec, + consecutive_failures: &mut usize, + trace_id: &Option, + initial_len: usize, + ) -> Option { + let parallel = tool_calls.len() > 1 && tool_calls.iter().all(|tc| tc.name == "Agent"); + + if parallel { + tracing::info!( + count = tool_calls.len(), + "executing sub-agent tool calls in parallel" + ); + let results: Vec = + futures::future::join_all(tool_calls.iter().map(|tc| self.execute_tool(tc))).await; + for (tc, result) in tool_calls.iter().zip(results) { + if let Some(pr) = self + .process_tool_call_result( + tc, + result, + messages, + consecutive_failures, + trace_id, + initial_len, + ) + .await + { + return Some(pr); + } + } + } else { + for tc in tool_calls { + let result = self.execute_tool(tc).await; + if let Some(pr) = self + .process_tool_call_result( + tc, + result, + messages, + consecutive_failures, + trace_id, + initial_len, + ) + .await + { + return Some(pr); + } + } + } + None + } /// Create an isolated subsession that shares the LLM client credentials /// and confirmation callback but has independent event handling, cancellation, @@ -1585,7 +1599,7 @@ impl LlmSession { if memory_command.is_some_and(|command| { self.approval_memory .as_ref() - .is_some_and(|memory| memory.lock().unwrap().is_allowed(command)) + .is_some_and(|memory| memory.lock().is_allowed(command)) }) { self.emit_audit(AuditEvent::security_decision( chrono::Utc::now(), @@ -1634,7 +1648,7 @@ impl LlmSession { if matches!(choice, ApprovalChoice::RememberSession) { if let Some(memory) = &self.approval_memory { if let Some(command) = memory_command { - memory.lock().unwrap().remember(command); + memory.lock().remember(command); } } } @@ -3244,7 +3258,7 @@ mod tests { session.set_confirmation_callback(std::sync::Arc::new( |_ctx: &PreflightSecurityContext| ApprovalChoice::RememberSession, )); - session.set_approval_memory(std::sync::Arc::new(std::sync::Mutex::new( + session.set_approval_memory(std::sync::Arc::new(parking_lot::Mutex::new( ApprovalMemory::new(), ))); @@ -3277,7 +3291,7 @@ mod tests { session.set_confirmation_callback(std::sync::Arc::new( |_ctx: &PreflightSecurityContext| ApprovalChoice::ReplyToAi, )); - session.set_approval_memory(std::sync::Arc::new(std::sync::Mutex::new( + session.set_approval_memory(std::sync::Arc::new(parking_lot::Mutex::new( ApprovalMemory::new(), ))); @@ -3292,4 +3306,167 @@ mod tests { // Preflight still ran once (memory miss → confirm → reply). assert_eq!(counter.load(std::sync::atomic::Ordering::SeqCst), 1); } + /// Tool that echoes a `tag` from its args so parallel calls can be told + /// apart and matched back to their originating call id. + struct ArgsEchoTool; + + impl Tool for ArgsEchoTool { + fn name(&self) -> &str { + "Agent" + } + fn description(&self) -> &str { + "echo args tag" + } + fn parameters(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) + } + fn execute(&self, args: serde_json::Value) -> crate::types::ToolResult { + let tag = args.get("tag").and_then(|v| v.as_str()).unwrap_or("?"); + crate::types::ToolResult::success(format!("echo:{tag}")) + } + } + + #[tokio::test] + async fn parallel_agent_calls_all_execute_and_keep_id_order() { + use crate::agents::{mock_text_response, mock_tool_call_response}; + + let mut session = LlmSession::new("http://localhost", "key", "model", None, None); + session.set_context_budget_policy(ContextBudgetPolicy { + enabled: false, + ..Default::default() + }); + // Turn 1: two independent Agent calls in one batch; turn 2: plain text. + session.set_test_chat_responses(vec![ + Ok(mock_tool_call_response(&[ + ("c1", "Agent", r#"{"tag":"alpha"}"#), + ("c2", "Agent", r#"{"tag":"beta"}"#), + ])), + Ok(mock_text_response("all done")), + ]); + session.register_tool(Box::new(ArgsEchoTool)); + + let result = session + .process_input( + &ChatMessage::user("spawn two agents"), + &[], + Some("system"), + false, + ) + .await + .expect("process_input should succeed"); + + // Both sub-agent calls executed, and each tool_result carries the tag + // from its own args — proving results were matched back to the correct + // call id rather than collapsed or swapped. + let tool_outputs: Vec = result + .new_messages + .iter() + .filter(|m| m.role == "tool") + .map(|m| { + m.content + .as_ref() + .and_then(|c| c.to_text()) + .unwrap_or_default() + .to_string() + }) + .collect(); + assert_eq!( + tool_outputs.len(), + 2, + "both agent calls must produce a tool result" + ); + assert_eq!( + tool_outputs, + vec!["echo:alpha".to_string(), "echo:beta".to_string()], + "join_all + zip must preserve tool_call order: {tool_outputs:?}" + ); + } + use std::future::Future; + use std::pin::Pin; + /// Tool that tracks live concurrency: increments a counter on entry, + /// sleeps long enough to force overlap, decrements on exit, and records + /// the peak number of simultaneously-active calls. This makes parallel + /// vs sequential execution directly observable (peak >= 2 means the + /// calls actually overlapped; peak == 1 means they ran one after another). + struct ConcurrencyProbe { + inflight: std::sync::Arc, + peak: std::sync::Arc, + } + + impl Tool for ConcurrencyProbe { + fn name(&self) -> &str { + "Agent" + } + fn description(&self) -> &str { + "concurrency probe" + } + fn parameters(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) + } + fn execute(&self, _args: serde_json::Value) -> crate::types::ToolResult { + crate::types::ToolResult::success("probe") + } + fn execute_async<'a>( + &'a self, + _args: serde_json::Value, + ) -> Pin + Send + 'a>> { + let inflight = self.inflight.clone(); + let peak = self.peak.clone(); + Box::pin(async move { + use std::sync::atomic::Ordering; + let cur = inflight.fetch_add(1, Ordering::SeqCst) + 1; + let mut known = peak.load(Ordering::SeqCst); + while cur > known { + match peak.compare_exchange(known, cur, Ordering::SeqCst, Ordering::SeqCst) { + Ok(_) => break, + Err(actual) => known = actual, + } + } + tokio::time::sleep(std::time::Duration::from_millis(60)).await; + inflight.fetch_sub(1, Ordering::SeqCst); + crate::types::ToolResult::success("probe") + }) + } + } + + #[tokio::test] + async fn parallel_agent_batch_runs_concurrently() { + use crate::agents::{mock_text_response, mock_tool_call_response}; + use std::sync::atomic::Ordering; + + let inflight = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); + let peak = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)); + + let mut session = LlmSession::new("http://localhost", "key", "model", None, None); + session.set_context_budget_policy(ContextBudgetPolicy { + enabled: false, + ..Default::default() + }); + session.set_test_chat_responses(vec![ + Ok(mock_tool_call_response(&[ + ("c1", "Agent", "{}"), + ("c2", "Agent", "{}"), + ("c3", "Agent", "{}"), + ])), + Ok(mock_text_response("done")), + ]); + session.register_tool(Box::new(ConcurrencyProbe { + inflight: inflight.clone(), + peak: peak.clone(), + })); + + session + .process_input(&ChatMessage::user("x"), &[], Some("sys"), false) + .await + .expect("process_input should succeed"); + + // join_all runs the probes concurrently, so at least two must overlap. + // If `run_tool_calls` regresses to sequential `for` (or `parallel` is + // forced false), peak stays 1 and this assertion fails. + let observed = peak.load(Ordering::SeqCst); + assert!( + observed >= 2, + "expected concurrent sub-agent execution (peak >= 2), got peak = {observed}" + ); + } } diff --git a/crates/aish-shell/src/app.rs b/crates/aish-shell/src/app.rs index d0120c38..cb9131fe 100644 --- a/crates/aish-shell/src/app.rs +++ b/crates/aish-shell/src/app.rs @@ -1,5 +1,5 @@ use std::io::{self, Read, Write}; -use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU32, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Instant; @@ -674,7 +674,7 @@ pub struct AishShell { inline_ai: Option>, /// Session-scoped approval memory shared with the LLM session. Kept on the /// shell so slash commands (e.g. `/forget-approvals`) can reset it. - approval_memory: Arc>, + approval_memory: Arc>, } impl AishShell { @@ -1258,8 +1258,8 @@ impl AishShell { let compaction_active_ref = compaction_active.clone(); let compaction_notice_shown = Arc::new(AtomicBool::new(false)); let compaction_notice_shown_ref = compaction_notice_shown.clone(); - let sub_agent_ui_active = Arc::new(AtomicBool::new(false)); - let sub_agent_ui_active_ref = sub_agent_ui_active.clone(); + let sub_agent_active_count = Arc::new(AtomicU32::new(0)); + let sub_agent_active_count_ref = sub_agent_active_count.clone(); let sub_agent_animation = Arc::new(SubAgentThinkingAnimation::new()); let sub_agent_animation_ref = sub_agent_animation.clone(); @@ -1358,9 +1358,16 @@ impl AishShell { compaction_notice_shown_ref.store(false, Ordering::SeqCst); animation_ref.start(&t("shell.status.thinking")); } + LlmEventType::OpEnd if crate::llm_event_ui::sub_agent_llm_event(&event) => { + // A sub-agent OpEnd marks one sub-session finishing, not the + // parent turn end. Ignore it so a fast-finishing sibling does + // not reset the active count (or stop the spinner) while other + // sub-agents are still running. The count is decremented via + // ToolExecutionEnd's saturating subtraction instead. + } LlmEventType::OpEnd => { // Operation ends — stop animation and show timing - sub_agent_ui_active_ref.store(false, Ordering::SeqCst); + sub_agent_active_count_ref.store(0, Ordering::SeqCst); sub_agent_animation_ref.stop(); animation_ref.stop(); let ttft = *ttft_value_ref.lock().unwrap(); @@ -1417,7 +1424,7 @@ impl AishShell { reasoning_buf_ref.lock().unwrap().clear(); reasoning_frame_ref.store(0, Ordering::SeqCst); renderer_ref.lock().unwrap().reset(); - if !sub_agent_ui_active_ref.load(Ordering::SeqCst) { + if sub_agent_active_count_ref.load(Ordering::SeqCst) == 0 { animation_ref.start(&t("shell.status.thinking")); } } @@ -1573,7 +1580,7 @@ impl AishShell { } LlmEventType::ToolExecutionStart => { if crate::llm_event_ui::is_parent_agent_spawn_tool_event(&event) { - sub_agent_ui_active_ref.store(true, Ordering::SeqCst); + sub_agent_active_count_ref.fetch_add(1, Ordering::SeqCst); animation_ref.stop(); sub_agent_animation_ref.stop(); clear_reasoning(); @@ -1669,7 +1676,13 @@ impl AishShell { } LlmEventType::ToolExecutionEnd => { if crate::llm_event_ui::is_parent_agent_spawn_tool_event(&event) { - sub_agent_ui_active_ref.store(false, Ordering::SeqCst); + // Decrement; saturating so a stray End (no matching Start, + // e.g. after an OpEnd reset) cannot underflow the count. + sub_agent_active_count_ref + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |v| { + Some(v.saturating_sub(1)) + }) + .ok(); } // Stop progress spinner started at ToolExecutionStart. animation_ref.stop(); @@ -1883,8 +1896,8 @@ impl AishShell { // Session-scoped approval memory: when the user approves a command with // "remember", equivalent commands (same host + normalized text) skip // confirmation — and the sandbox preflight — for the rest of the session. - let approval_memory: Arc> = - Arc::new(Mutex::new(aish_llm::ApprovalMemory::new())); + let approval_memory: Arc> = + Arc::new(parking_lot::Mutex::new(aish_llm::ApprovalMemory::new())); llm_session.set_approval_memory(approval_memory.clone()); if let Some(ref audit) = audit_store { @@ -3423,7 +3436,7 @@ impl AishShell { /// so previously "remembered" commands prompt for confirmation again. fn handle_forget_approvals(&mut self) { let count = { - let mut memory = self.approval_memory.lock().unwrap(); + let mut memory = self.approval_memory.lock(); let n = memory.len(); memory.clear(); n diff --git a/crates/aish-tools/src/agent_tool/agent_tool.rs b/crates/aish-tools/src/agent_tool/agent_tool.rs index afdc3d43..b19730b1 100644 --- a/crates/aish-tools/src/agent_tool/agent_tool.rs +++ b/crates/aish-tools/src/agent_tool/agent_tool.rs @@ -125,7 +125,15 @@ impl Tool for AgentTool { args: serde_json::Value, session: &'a LlmSession, ) -> Pin + Send + 'a>> { - Box::pin(async move { + use futures::FutureExt; + // Wrap the sub-agent run in `catch_unwind` so a panicking sub-session + // degrades to a `ToolResult::error` instead of unwinding through the + // caller. This matches the framework's `Tool::execute_async` contract + // (types.rs) and is essential under parallel `join_all` execution, + // where an uncontained panic would abort every concurrent sibling and + // the whole turn (e.g. a poisoned `Arc` shared via the parent + // session cascades a `.lock().unwrap()` panic to all siblings). + let inner = std::panic::AssertUnwindSafe(async move { let (_description, prompt, subagent_type) = match Self::validate_args(&args) { Ok(v) => v, Err(err) => return err, @@ -162,6 +170,21 @@ impl Tool for AgentTool { } Self::spawn_result_to_tool_result(result) + }); + Box::pin(async move { + match inner.catch_unwind().await { + Ok(result) => result, + Err(payload) => { + let message = if let Some(s) = payload.downcast_ref::<&str>() { + s.to_string() + } else if let Some(s) = payload.downcast_ref::() { + s.clone() + } else { + "Agent sub-agent execution panicked".to_string() + }; + ToolResult::error(format!("Error: {}", message)) + } + } }) } } @@ -214,4 +237,32 @@ mod tests { assert!(!names.contains(&"command-diagnose")); assert!(!names.contains(&"diagnose")); } + #[test] + fn description_encourages_parallel_independent_agents() { + let tool = AgentTool::new(); + // The tool prompt must steer the model toward emitting multiple Agent + // calls in one response — that is what unlocks concurrent execution. + let desc = tool.description(); + assert!(desc.contains("Parallelize"), "missing parallel guidance"); + assert!( + desc.contains("this single response"), + "must tell the model to emit calls in one response" + ); + assert!( + desc.contains("concurrently"), + "must state the calls run concurrently" + ); + assert!( + desc.contains("I/O-bound"), + "must explain why read-only tasks parallelize well" + ); + assert!( + desc.contains("scan"), + "parallel guidance must cover file-scan scenarios" + ); + assert!( + desc.contains("are healthy"), + "parallel guidance must cover multi-service diagnosis scenarios" + ); + } } diff --git a/crates/aish-tools/src/agent_tool/prompt.rs b/crates/aish-tools/src/agent_tool/prompt.rs index 0db84c8d..e95e85cf 100644 --- a/crates/aish-tools/src/agent_tool/prompt.rs +++ b/crates/aish-tools/src/agent_tool/prompt.rs @@ -39,7 +39,16 @@ pub const USAGE_SECTION: &str = "\ thorough), and whether the sub-agent must stay read-only. Default to quick or medium unless the \ user asked for exhaustive coverage; do not expand scope to \"everywhere\" on your own. - When the sub-agent finishes, only its final conclusion is returned here; summarize for the user if needed. -- Launch multiple agents in one turn when their tasks are independent. +- **Parallelize independent work — decide for the user.** Users rarely say \"parallel\"; when \ +their request naturally covers several independent areas, YOU decompose it and emit **multiple \ +`Agent` calls in this single response** so they run concurrently. Read-only investigation, \ +environment checks, system diagnosis, and file scans are I/O-bound and independent — parallel \ +sub-agents finish in a fraction of the time, so default to fanning out. Pick the right type per \ +area: \"explain the frontend, backend, and database\" → 3× explore; \"check the auth, payment, \ +and logging services are healthy\" → 3× troubleshoot; \"scan src/, tests/, and configs/ for TODO \ +markers\" → 3× explore; \"audit env vars, disk usage, and open ports\" → troubleshoot or explore. \ +Do not wait for one to return before issuing the next, and do not make the user ask for parallelism. \ +Keep dependent tasks (a later one needs an earlier result) sequential. - If you delegate research to a sub-agent, do not duplicate the same searches in this session."; pub fn parameters(subagent_types: &[String]) -> serde_json::Value { From 23d6357dbc3427bbd630a16e789d0658629e16ed Mon Sep 17 00:00:00 2001 From: xuezhixin Date: Wed, 5 Aug 2026 13:30:38 +0800 Subject: [PATCH 3/3] fix(llm): pair every tool_call with a tool_result on short-circuit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A short-circuit / failure-threshold return from the tool loop left the remaining assistant tool_calls without a matching tool_result. The shell persists new_messages verbatim, so the next turn would send the provider an assistant message with dangling tool_call_ids (rejected by OpenAI / Anthropic, which require a tool_result for every tool_call). - process_tool_call_result: append the short-circuiting call's own tool_result before returning (previously skipped). - run_tool_calls: on early return, backfill tool_results for every remaining call — the already-computed result in the parallel branch, a synthetic "[skipped ...]" in the sequential branch (those calls never executed). - Add regression test: a parallel Agent batch whose first call short-circuits must still produce a tool_result for every tool_call_id. Also drop two wording-coupled assertions in the Agent prompt test (scan / are healthy) that broke on example rewording without any behavior change; the remaining four still cover the parallelism contract. Addresses CodeRabbit review on #429. --- crates/aish-llm/src/session.rs | 139 +++++++++++++++++- .../aish-tools/src/agent_tool/agent_tool.rs | 8 - 2 files changed, 131 insertions(+), 16 deletions(-) diff --git a/crates/aish-llm/src/session.rs b/crates/aish-llm/src/session.rs index 579493c6..98dcab96 100644 --- a/crates/aish-llm/src/session.rs +++ b/crates/aish-llm/src/session.rs @@ -1416,6 +1416,11 @@ impl LlmSession { timestamp: now_timestamp(), metadata: None, }); + // Record this call's tool_result before returning so the assistant + // message never carries a dangling tool_call_id — providers reject + // an assistant tool_call with no following tool_result. Remaining + // calls in the batch are backfilled by `run_tool_calls`. + messages.push(ChatMessage::tool_result(&tc.id, output.clone())); // Sub-agent cancel is shown by the shell as `shell.interrupted` // (same as Ctrl+C); do not return the tool string as AI body. let suppress_body = result.meta.as_ref().is_some_and(|meta| { @@ -1502,11 +1507,11 @@ impl LlmSession { ); let results: Vec = futures::future::join_all(tool_calls.iter().map(|tc| self.execute_tool(tc))).await; - for (tc, result) in tool_calls.iter().zip(results) { + for idx in 0..tool_calls.len() { if let Some(pr) = self .process_tool_call_result( - tc, - result, + &tool_calls[idx], + results[idx].clone(), messages, consecutive_failures, trace_id, @@ -1514,15 +1519,32 @@ impl LlmSession { ) .await { - return Some(pr); + // A short-circuit / failure threshold stops the loop + // here, but the assistant message already lists every + // tool_call. Append the remaining (already-computed) + // results so no tool_call_id is left without a + // tool_result — providers reject an assistant tool_call + // that has no matching tool_result. + for (rest_tc, rest_result) in + tool_calls[idx + 1..].iter().zip(&results[idx + 1..]) + { + messages.push(ChatMessage::tool_result( + &rest_tc.id, + rest_result.output.clone(), + )); + } + return Some(crate::types::ProcessResult { + text: pr.text, + new_messages: messages[initial_len..].to_vec(), + }); } } } else { - for tc in tool_calls { - let result = self.execute_tool(tc).await; + for idx in 0..tool_calls.len() { + let result = self.execute_tool(&tool_calls[idx]).await; if let Some(pr) = self .process_tool_call_result( - tc, + &tool_calls[idx], result, messages, consecutive_failures, @@ -1531,7 +1553,19 @@ impl LlmSession { ) .await { - return Some(pr); + // Backfill the remaining (never-executed) tool calls with + // a synthetic tool_result so the assistant message carries + // no dangling tool_call_id (provider API invariant). + for rest_tc in &tool_calls[idx + 1..] { + messages.push(ChatMessage::tool_result( + &rest_tc.id, + "[skipped: tool execution stopped after a blocking result]", + )); + } + return Some(crate::types::ProcessResult { + text: pr.text, + new_messages: messages[initial_len..].to_vec(), + }); } } } @@ -3469,4 +3503,93 @@ mod tests { "expected concurrent sub-agent execution (peak >= 2), got peak = {observed}" ); } + /// Agent tool that short-circuits (`sub_agent_cancelled`) when its args + /// carry `"cancel": true`, else succeeds. Exercises the parallel + /// short-circuit path and the tool_result backfill. + struct ShortCircuitAgent; + + impl Tool for ShortCircuitAgent { + fn name(&self) -> &str { + "Agent" + } + fn description(&self) -> &str { + "short-circuit probe" + } + fn parameters(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) + } + fn execute(&self, args: serde_json::Value) -> crate::types::ToolResult { + if args + .get("cancel") + .and_then(|v| v.as_bool()) + .unwrap_or(false) + { + crate::types::ToolResult { + ok: false, + output: "cancelled".into(), + meta: Some(serde_json::json!({ + "dispatch_status": "short_circuit", + "reason": "sub_agent_cancelled", + })), + } + } else { + crate::types::ToolResult::success("ok") + } + } + } + + #[tokio::test] + async fn parallel_short_circuit_pairs_every_tool_call() { + use crate::agents::mock_tool_call_response; + use std::collections::HashSet; + + let mut session = LlmSession::new("http://localhost", "key", "model", None, None); + session.set_context_budget_policy(ContextBudgetPolicy { + enabled: false, + ..Default::default() + }); + // Three Agent calls in one batch; the first short-circuits. join_all + // still runs all three, so the loop hits the short-circuit at idx 0. + // Without the backfill, c2/c3 would be dangling tool_call_ids in the + // persisted assistant message (providers reject that next turn). + session.set_test_chat_responses(vec![Ok(mock_tool_call_response(&[ + ("c1", "Agent", r#"{"cancel":true}"#), + ("c2", "Agent", "{}"), + ("c3", "Agent", "{}"), + ]))]); + session.register_tool(Box::new(ShortCircuitAgent)); + + let result = session + .process_input(&ChatMessage::user("x"), &[], Some("sys"), false) + .await + .expect("process_input should succeed"); + + // Every assistant tool_call_id must have a matching tool_result. + let mut expected: HashSet = HashSet::new(); + for m in &result.new_messages { + if m.role == "assistant" { + if let Some(calls) = &m.tool_calls { + for c in calls { + expected.insert(c.id.clone()); + } + } + } + } + let mut answered: HashSet = HashSet::new(); + for m in &result.new_messages { + if m.role == "tool" { + if let Some(id) = &m.tool_call_id { + answered.insert(id.clone()); + } + } + } + assert!( + !expected.is_empty(), + "test setup error: no assistant tool_calls found" + ); + assert_eq!( + expected, answered, + "every tool_call_id must have a tool_result" + ); + } } diff --git a/crates/aish-tools/src/agent_tool/agent_tool.rs b/crates/aish-tools/src/agent_tool/agent_tool.rs index b19730b1..74bb5eda 100644 --- a/crates/aish-tools/src/agent_tool/agent_tool.rs +++ b/crates/aish-tools/src/agent_tool/agent_tool.rs @@ -256,13 +256,5 @@ mod tests { desc.contains("I/O-bound"), "must explain why read-only tasks parallelize well" ); - assert!( - desc.contains("scan"), - "parallel guidance must cover file-scan scenarios" - ); - assert!( - desc.contains("are healthy"), - "parallel guidance must cover multi-service diagnosis scenarios" - ); } }