feat: parallel tool execution for read-only tools
When the model returns multiple tool_use blocks in a single response,
read-only tools (read_file, glob_search, grep_search, WebFetch,
WebSearch, LSP, Git*) now execute concurrently via std:🧵:scope
instead of sequentially.
Architecture:
- Add ToolCall/ToolResult types to ToolExecutor trait
- Add execute_batch() with default sequential fallback
- CliToolExecutor overrides execute_batch to classify tools as
parallel-safe or sequential, then runs parallel-safe tools
concurrently via thread scopes
- run_turn refactored into 3 phases:
1. Pre-hooks + permission checks (sequential, hooks may mutate state)
2. Tool execution (batch, parallel for read-only tools)
3. Post-hooks + session updates (sequential, preserves ordering)
Safety guarantees:
- Pre/post hooks always run sequentially
- Permission checks complete before any tool executes
- Tool results pushed to session in original model order
- Fallback to sequential for single-tool batches
This commit is contained in:
parent
5bb811aa5e
commit
0ff3882ca3
|
|
@ -61,6 +61,40 @@ pub trait ApiClient {
|
||||||
/// Trait implemented by tool dispatchers that execute model-requested tools.
|
/// Trait implemented by tool dispatchers that execute model-requested tools.
|
||||||
pub trait ToolExecutor {
|
pub trait ToolExecutor {
|
||||||
fn execute(&mut self, tool_name: &str, input: &str) -> Result<String, ToolError>;
|
fn execute(&mut self, tool_name: &str, input: &str) -> Result<String, ToolError>;
|
||||||
|
|
||||||
|
/// Execute a batch of tool calls, potentially in parallel.
|
||||||
|
/// Returns results in the same order as the input calls.
|
||||||
|
/// The default implementation executes sequentially via `execute`.
|
||||||
|
/// Override this to provide parallel execution for read-only tools.
|
||||||
|
fn execute_batch(&mut self, calls: Vec<ToolCall>) -> Vec<ToolResult> {
|
||||||
|
calls
|
||||||
|
.into_iter()
|
||||||
|
.map(|call| {
|
||||||
|
let result = self.execute(&call.tool_name, &call.input);
|
||||||
|
ToolResult {
|
||||||
|
tool_use_id: call.tool_use_id,
|
||||||
|
tool_name: call.tool_name,
|
||||||
|
result,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A single tool call to execute.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct ToolCall {
|
||||||
|
pub tool_use_id: String,
|
||||||
|
pub tool_name: String,
|
||||||
|
pub input: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The result of executing a tool call.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub struct ToolResult {
|
||||||
|
pub tool_use_id: String,
|
||||||
|
pub tool_name: String,
|
||||||
|
pub result: Result<String, ToolError>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Error returned when a tool invocation fails locally.
|
/// Error returned when a tool invocation fails locally.
|
||||||
|
|
@ -415,11 +449,24 @@ where
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
for (tool_use_id, tool_name, input) in pending_tool_uses {
|
// Phase 1: Pre-hooks and permission checks (sequential).
|
||||||
let pre_hook_result = self.run_pre_tool_use_hook(&tool_name, &input);
|
// Hooks may mutate state and must run in order.
|
||||||
|
struct PendingTool {
|
||||||
|
tool_use_id: String,
|
||||||
|
tool_name: String,
|
||||||
|
effective_input: String,
|
||||||
|
pre_hook_messages: Vec<String>,
|
||||||
|
allowed: bool,
|
||||||
|
deny_reason: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut pending = Vec::with_capacity(pending_tool_uses.len());
|
||||||
|
for (tool_use_id, tool_name, input) in &pending_tool_uses {
|
||||||
|
let pre_hook_result = self.run_pre_tool_use_hook(tool_name, input);
|
||||||
let effective_input = pre_hook_result
|
let effective_input = pre_hook_result
|
||||||
.updated_input()
|
.updated_input()
|
||||||
.map_or_else(|| input.clone(), ToOwned::to_owned);
|
.map_or_else(|| input.clone(), ToOwned::to_owned);
|
||||||
|
let pre_hook_messages = pre_hook_result.messages().to_vec();
|
||||||
let permission_context = PermissionContext::new(
|
let permission_context = PermissionContext::new(
|
||||||
pre_hook_result.permission_override(),
|
pre_hook_result.permission_override(),
|
||||||
pre_hook_result.permission_reason().map(ToOwned::to_owned),
|
pre_hook_result.permission_reason().map(ToOwned::to_owned),
|
||||||
|
|
@ -448,66 +495,112 @@ where
|
||||||
}
|
}
|
||||||
} else if let Some(prompt) = prompter.as_mut() {
|
} else if let Some(prompt) = prompter.as_mut() {
|
||||||
self.permission_policy.authorize_with_context(
|
self.permission_policy.authorize_with_context(
|
||||||
&tool_name,
|
tool_name,
|
||||||
&effective_input,
|
&effective_input,
|
||||||
&permission_context,
|
&permission_context,
|
||||||
Some(*prompt),
|
Some(*prompt),
|
||||||
)
|
)
|
||||||
} else {
|
} else {
|
||||||
self.permission_policy.authorize_with_context(
|
self.permission_policy.authorize_with_context(
|
||||||
&tool_name,
|
tool_name,
|
||||||
&effective_input,
|
&effective_input,
|
||||||
&permission_context,
|
&permission_context,
|
||||||
None,
|
None,
|
||||||
)
|
)
|
||||||
};
|
};
|
||||||
|
|
||||||
let result_message = match permission_outcome {
|
match permission_outcome {
|
||||||
PermissionOutcome::Allow => {
|
PermissionOutcome::Allow => {
|
||||||
self.record_tool_started(iterations, &tool_name);
|
pending.push(PendingTool {
|
||||||
let (mut output, mut is_error) =
|
tool_use_id: tool_use_id.clone(),
|
||||||
match self.tool_executor.execute(&tool_name, &effective_input) {
|
tool_name: tool_name.clone(),
|
||||||
Ok(output) => (output, false),
|
effective_input,
|
||||||
Err(error) => (error.to_string(), true),
|
pre_hook_messages,
|
||||||
};
|
allowed: true,
|
||||||
output = merge_hook_feedback(pre_hook_result.messages(), output, false);
|
deny_reason: None,
|
||||||
|
});
|
||||||
let post_hook_result = if is_error {
|
|
||||||
self.run_post_tool_use_failure_hook(
|
|
||||||
&tool_name,
|
|
||||||
&effective_input,
|
|
||||||
&output,
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
self.run_post_tool_use_hook(
|
|
||||||
&tool_name,
|
|
||||||
&effective_input,
|
|
||||||
&output,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
};
|
|
||||||
if post_hook_result.is_denied()
|
|
||||||
|| post_hook_result.is_failed()
|
|
||||||
|| post_hook_result.is_cancelled()
|
|
||||||
{
|
|
||||||
is_error = true;
|
|
||||||
}
|
|
||||||
output = merge_hook_feedback(
|
|
||||||
post_hook_result.messages(),
|
|
||||||
output,
|
|
||||||
post_hook_result.is_denied()
|
|
||||||
|| post_hook_result.is_failed()
|
|
||||||
|| post_hook_result.is_cancelled(),
|
|
||||||
);
|
|
||||||
|
|
||||||
ConversationMessage::tool_result(tool_use_id, tool_name, output, is_error)
|
|
||||||
}
|
}
|
||||||
PermissionOutcome::Deny { reason } => ConversationMessage::tool_result(
|
PermissionOutcome::Deny { reason } => {
|
||||||
tool_use_id,
|
pending.push(PendingTool {
|
||||||
tool_name,
|
tool_use_id: tool_use_id.clone(),
|
||||||
merge_hook_feedback(pre_hook_result.messages(), reason, true),
|
tool_name: tool_name.clone(),
|
||||||
|
effective_input: String::new(),
|
||||||
|
pre_hook_messages,
|
||||||
|
allowed: false,
|
||||||
|
deny_reason: Some(reason),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Phase 2: Execute allowed tools (batch, may run in parallel).
|
||||||
|
let allowed_calls: Vec<ToolCall> = pending
|
||||||
|
.iter()
|
||||||
|
.filter(|p| p.allowed)
|
||||||
|
.map(|p| {
|
||||||
|
self.record_tool_started(iterations, &p.tool_name);
|
||||||
|
ToolCall {
|
||||||
|
tool_use_id: p.tool_use_id.clone(),
|
||||||
|
tool_name: p.tool_name.clone(),
|
||||||
|
input: p.effective_input.clone(),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let batch_results = self.tool_executor.execute_batch(allowed_calls);
|
||||||
|
let mut batch_index = 0;
|
||||||
|
|
||||||
|
// Phase 3: Post-hooks and session updates (sequential, original order).
|
||||||
|
for p in &pending {
|
||||||
|
let result_message = if p.allowed {
|
||||||
|
let batch_result = &batch_results[batch_index];
|
||||||
|
batch_index += 1;
|
||||||
|
let (mut output, mut is_error) = match &batch_result.result {
|
||||||
|
Ok(output) => (output.clone(), false),
|
||||||
|
Err(error) => (error.to_string(), true),
|
||||||
|
};
|
||||||
|
output = merge_hook_feedback(&p.pre_hook_messages, output, false);
|
||||||
|
|
||||||
|
let post_hook_result = if is_error {
|
||||||
|
self.run_post_tool_use_failure_hook(
|
||||||
|
&p.tool_name,
|
||||||
|
&p.effective_input,
|
||||||
|
&output,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
self.run_post_tool_use_hook(
|
||||||
|
&p.tool_name,
|
||||||
|
&p.effective_input,
|
||||||
|
&output,
|
||||||
|
false,
|
||||||
|
)
|
||||||
|
};
|
||||||
|
if post_hook_result.is_denied()
|
||||||
|
|| post_hook_result.is_failed()
|
||||||
|
|| post_hook_result.is_cancelled()
|
||||||
|
{
|
||||||
|
is_error = true;
|
||||||
|
}
|
||||||
|
output = merge_hook_feedback(
|
||||||
|
post_hook_result.messages(),
|
||||||
|
output,
|
||||||
|
post_hook_result.is_denied()
|
||||||
|
|| post_hook_result.is_failed()
|
||||||
|
|| post_hook_result.is_cancelled(),
|
||||||
|
);
|
||||||
|
|
||||||
|
ConversationMessage::tool_result(
|
||||||
|
p.tool_use_id.clone(),
|
||||||
|
p.tool_name.clone(),
|
||||||
|
output,
|
||||||
|
is_error,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
ConversationMessage::tool_result(
|
||||||
|
p.tool_use_id.clone(),
|
||||||
|
p.tool_name.clone(),
|
||||||
|
merge_hook_feedback(&p.pre_hook_messages, p.deny_reason.clone().unwrap_or_default(), true),
|
||||||
true,
|
true,
|
||||||
),
|
)
|
||||||
};
|
};
|
||||||
self.session
|
self.session
|
||||||
.push_message(result_message.clone())
|
.push_message(result_message.clone())
|
||||||
|
|
|
||||||
|
|
@ -123,8 +123,8 @@ pub use lsp_discovery::{
|
||||||
};
|
};
|
||||||
pub use conversation::{
|
pub use conversation::{
|
||||||
auto_compaction_threshold_from_env, ApiClient, ApiRequest, AssistantEvent, AutoCompactionEvent,
|
auto_compaction_threshold_from_env, ApiClient, ApiRequest, AssistantEvent, AutoCompactionEvent,
|
||||||
ConversationRuntime, PromptCacheEvent, RuntimeError, StaticToolExecutor, ToolError,
|
ConversationRuntime, PromptCacheEvent, RuntimeError, StaticToolExecutor, ToolCall, ToolError,
|
||||||
ToolExecutor, TurnSummary,
|
ToolExecutor, ToolResult, TurnSummary,
|
||||||
};
|
};
|
||||||
pub use file_ops::{
|
pub use file_ops::{
|
||||||
edit_file, edit_file_in_workspace, glob_search, glob_search_in_workspace, grep_search,
|
edit_file, edit_file_in_workspace, glob_search, glob_search_in_workspace, grep_search,
|
||||||
|
|
|
||||||
|
|
@ -14218,6 +14218,145 @@ impl ToolExecutor for CliToolExecutor {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn execute_batch(&mut self, calls: Vec<runtime::ToolCall>) -> Vec<runtime::ToolResult> {
|
||||||
|
if calls.len() <= 1 {
|
||||||
|
return calls
|
||||||
|
.into_iter()
|
||||||
|
.map(|call| {
|
||||||
|
let result = self.execute(&call.tool_name, &call.input);
|
||||||
|
runtime::ToolResult {
|
||||||
|
tool_use_id: call.tool_use_id,
|
||||||
|
tool_name: call.tool_name,
|
||||||
|
result,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Tools that are safe to run in parallel because they only read
|
||||||
|
/// state and dispatch through the stateless tool registry.
|
||||||
|
const PARALLEL_SAFE_TOOLS: &[&str] = &[
|
||||||
|
"read_file",
|
||||||
|
"glob_search",
|
||||||
|
"grep_search",
|
||||||
|
"WebFetch",
|
||||||
|
"WebSearch",
|
||||||
|
"ToolSearch",
|
||||||
|
"Skill",
|
||||||
|
"LSP",
|
||||||
|
"GitStatus",
|
||||||
|
"GitDiff",
|
||||||
|
"GitLog",
|
||||||
|
"GitShow",
|
||||||
|
"GitBlame",
|
||||||
|
];
|
||||||
|
|
||||||
|
let emit_output = self.emit_output;
|
||||||
|
let mut results: Vec<Option<runtime::ToolResult>> = vec![None; calls.len()];
|
||||||
|
let mut parallel_calls: Vec<(usize, String, String, String)> = Vec::new();
|
||||||
|
let mut sequential_indices: Vec<usize> = Vec::new();
|
||||||
|
|
||||||
|
// Classify calls as parallel-safe or sequential
|
||||||
|
for (i, call) in calls.iter().enumerate() {
|
||||||
|
if self
|
||||||
|
.allowed_tools
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|allowed| !allowed.contains(&call.tool_name))
|
||||||
|
{
|
||||||
|
results[i] = Some(runtime::ToolResult {
|
||||||
|
tool_use_id: call.tool_use_id.clone(),
|
||||||
|
tool_name: call.tool_name.clone(),
|
||||||
|
result: Err(ToolError::new(format!(
|
||||||
|
"tool `{}` is not enabled by the current --allowedTools setting",
|
||||||
|
call.tool_name
|
||||||
|
))),
|
||||||
|
});
|
||||||
|
} else if PARALLEL_SAFE_TOOLS.contains(&call.tool_name.as_str())
|
||||||
|
&& !self.tool_registry.has_runtime_tool(&call.tool_name)
|
||||||
|
{
|
||||||
|
parallel_calls.push((
|
||||||
|
i,
|
||||||
|
call.tool_use_id.clone(),
|
||||||
|
call.tool_name.clone(),
|
||||||
|
call.input.clone(),
|
||||||
|
));
|
||||||
|
} else {
|
||||||
|
sequential_indices.push(i);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute parallel-safe tools concurrently
|
||||||
|
if !parallel_calls.is_empty() {
|
||||||
|
let registry = self.tool_registry.clone();
|
||||||
|
let parallel_results: Vec<(usize, String, String, Result<String, ToolError>)> =
|
||||||
|
std::thread::scope(|s| {
|
||||||
|
let mut handles = Vec::new();
|
||||||
|
for (idx, tool_use_id, tool_name, input) in ¶llel_calls {
|
||||||
|
let registry = ®istry;
|
||||||
|
let tool_use_id = tool_use_id.clone();
|
||||||
|
let tool_name = tool_name.clone();
|
||||||
|
let input = input.clone();
|
||||||
|
let idx = *idx;
|
||||||
|
handles.push(s.spawn(move || {
|
||||||
|
let value = serde_json::from_str(&input)
|
||||||
|
.map_err(|error| ToolError::new(format!("invalid tool input JSON: {error}")));
|
||||||
|
let result = match value {
|
||||||
|
Ok(v) => registry
|
||||||
|
.execute(&tool_name, &v)
|
||||||
|
.map_err(ToolError::new),
|
||||||
|
Err(e) => Err(e),
|
||||||
|
};
|
||||||
|
(idx, tool_use_id, tool_name, result)
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
handles
|
||||||
|
.into_iter()
|
||||||
|
.map(|h| h.join().unwrap_or_else(|_| {
|
||||||
|
(
|
||||||
|
0,
|
||||||
|
String::new(),
|
||||||
|
String::new(),
|
||||||
|
Err(ToolError::new("parallel thread panicked")),
|
||||||
|
)
|
||||||
|
}))
|
||||||
|
.collect()
|
||||||
|
});
|
||||||
|
|
||||||
|
for (idx, tool_use_id, tool_name, result) in parallel_results {
|
||||||
|
if emit_output {
|
||||||
|
let output_str = match &result {
|
||||||
|
Ok(o) => o.clone(),
|
||||||
|
Err(e) => e.to_string(),
|
||||||
|
};
|
||||||
|
let is_error = result.is_err();
|
||||||
|
let markdown = format_tool_result(&tool_name, &output_str, is_error);
|
||||||
|
self.renderer
|
||||||
|
.stream_markdown(&markdown, &mut io::stdout())
|
||||||
|
.map_err(|error| ToolError::new(error.to_string()))
|
||||||
|
.ok();
|
||||||
|
}
|
||||||
|
results[idx] = Some(runtime::ToolResult {
|
||||||
|
tool_use_id,
|
||||||
|
tool_name,
|
||||||
|
result,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute sequential tools one at a time
|
||||||
|
for idx in sequential_indices {
|
||||||
|
let call = &calls[idx];
|
||||||
|
let result = self.execute(&call.tool_name, &call.input);
|
||||||
|
results[idx] = Some(runtime::ToolResult {
|
||||||
|
tool_use_id: call.tool_use_id.clone(),
|
||||||
|
tool_name: call.tool_name.clone(),
|
||||||
|
result,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
results.into_iter().map(|r| r.unwrap()).collect()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn permission_policy(
|
fn permission_policy(
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue