mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 15:40:18 +00:00
* fix: parallelize tool call execution via JoinSet (#219) When the LLM returns multiple tool_calls in a single response, they were executed sequentially. This change makes both the worker and dispatcher paths concurrent using tokio::task::JoinSet, so N independent tool calls complete in ~max(latency) instead of sum(latency). Worker path: migrate execute_tools_parallel from join_all to JoinSet and route the respond_with_tools branch through the same parallel path. Dispatcher path: restructure the while-idx loop into three phases — preflight (sequential approval/hook checks), parallel execution via JoinSet, and sequential post-flight processing (session recording, auth detection, sanitization). Also fixes a pre-existing infinite loop bug where hook rejection used `continue` inside a `while idx` loop, skipping `idx += 1` and retrying the same rejected tool forever. Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address PR review — ordered results, deferred auth, dedup standalone fn - Fix auth early return skipping unrecorded tool results: defer auth response until after all results in the batch are recorded in session history and context_messages (both dispatcher and thread_ops paths) - Fix tool results appearing out of order: collect Phase 1 hook rejections indexed by original position, merge with Phase 2 execution results, and emit all in Phase 3 in original tool_calls order - Deduplicate execute_chat_tool: Agent method now delegates to the standalone function instead of duplicating 90 lines of logic - Fix benchmark compilation: add missing session_manager arg to Agent::new Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: rustfmt alignment for CI compatibility Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address second round of PR review comments - Distinguish JoinError panic vs cancellation in log messages and error reasons across all 3 files (dispatcher, thread_ops, worker) - Simplify deferred_auth from Option<(String, String)> to Option<String> since only the instructions string is used - Add single-tool short-circuit in worker execute_tools_parallel to avoid JoinSet overhead for the common single-tool case Co-Authored-By: Claude Opus 4.6 <[email protected]> --------- Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
9906190de7
commit
bfe393eb38
+196
-69
@@ -6,11 +6,14 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::task::JoinSet;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::Agent;
|
||||
use crate::agent::compaction::ContextCompactor;
|
||||
use crate::agent::dispatcher::{AgenticLoopResult, check_auth_required, parse_auth_result};
|
||||
use crate::agent::dispatcher::{
|
||||
AgenticLoopResult, check_auth_required, execute_chat_tool_standalone, parse_auth_result,
|
||||
};
|
||||
use crate::agent::session::{PendingApproval, Session, ThreadState};
|
||||
use crate::agent::submission::SubmissionResult;
|
||||
use crate::channels::{IncomingMessage, StatusUpdate};
|
||||
@@ -785,9 +788,17 @@ impl Agent {
|
||||
.await;
|
||||
}
|
||||
|
||||
let mut deferred_queue = std::collections::VecDeque::from(deferred_tool_calls);
|
||||
while let Some(tc) = deferred_queue.pop_front() {
|
||||
// Re-check approval for each deferred tool call
|
||||
// === Phase 1: Preflight (sequential) ===
|
||||
// Walk deferred tools checking approval. Collect runnable
|
||||
// tools; stop at the first that needs approval.
|
||||
let mut runnable: Vec<crate::llm::ToolCall> = Vec::new();
|
||||
let mut approval_needed: Option<(
|
||||
usize,
|
||||
crate::llm::ToolCall,
|
||||
Arc<dyn crate::tools::Tool>,
|
||||
)> = None;
|
||||
|
||||
for (idx, tc) in deferred_tool_calls.iter().enumerate() {
|
||||
if let Some(tool) = self.tools().get(&tc.name).await
|
||||
&& tool.requires_approval()
|
||||
{
|
||||
@@ -801,73 +812,142 @@ impl Agent {
|
||||
};
|
||||
|
||||
if !is_auto_approved {
|
||||
let new_pending = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
tool_name: tc.name.clone(),
|
||||
parameters: tc.arguments.clone(),
|
||||
description: tool.description().to_string(),
|
||||
tool_call_id: tc.id.clone(),
|
||||
context_messages: context_messages.clone(),
|
||||
deferred_tool_calls: deferred_queue.iter().cloned().collect(),
|
||||
};
|
||||
|
||||
let request_id = new_pending.request_id;
|
||||
let tool_name = new_pending.tool_name.clone();
|
||||
let description = new_pending.description.clone();
|
||||
let parameters = new_pending.parameters.clone();
|
||||
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.await_approval(new_pending);
|
||||
}
|
||||
}
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Awaiting approval".into()),
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
return Ok(SubmissionResult::NeedApproval {
|
||||
request_id,
|
||||
tool_name,
|
||||
description,
|
||||
parameters,
|
||||
});
|
||||
approval_needed = Some((idx, tc.clone(), tool));
|
||||
break; // remaining tools stay deferred
|
||||
}
|
||||
}
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolStarted {
|
||||
name: tc.name.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
runnable.push(tc.clone());
|
||||
}
|
||||
|
||||
let deferred_result = self
|
||||
.execute_chat_tool(&tc.name, &tc.arguments, &job_ctx)
|
||||
.await;
|
||||
// === Phase 2: Parallel execution ===
|
||||
let exec_results: Vec<(crate::llm::ToolCall, Result<String, Error>)> = if runnable.len()
|
||||
<= 1
|
||||
{
|
||||
// Single tool (or none): execute inline
|
||||
let mut results = Vec::new();
|
||||
for tc in &runnable {
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolStarted {
|
||||
name: tc.name.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolCompleted {
|
||||
name: tc.name.clone(),
|
||||
success: deferred_result.is_ok(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
let result = self
|
||||
.execute_chat_tool(&tc.name, &tc.arguments, &job_ctx)
|
||||
.await;
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolCompleted {
|
||||
name: tc.name.clone(),
|
||||
success: result.is_ok(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
results.push((tc.clone(), result));
|
||||
}
|
||||
results
|
||||
} else {
|
||||
// Multiple tools: execute in parallel via JoinSet
|
||||
let mut join_set = JoinSet::new();
|
||||
let runnable_count = runnable.len();
|
||||
|
||||
for (spawn_idx, tc) in runnable.iter().enumerate() {
|
||||
let tools = self.tools().clone();
|
||||
let safety = self.safety().clone();
|
||||
let channels = self.channels.clone();
|
||||
let job_ctx = job_ctx.clone();
|
||||
let tc = tc.clone();
|
||||
let channel = message.channel.clone();
|
||||
let metadata = message.metadata.clone();
|
||||
|
||||
join_set.spawn(async move {
|
||||
let _ = channels
|
||||
.send_status(
|
||||
&channel,
|
||||
StatusUpdate::ToolStarted {
|
||||
name: tc.name.clone(),
|
||||
},
|
||||
&metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
let result = execute_chat_tool_standalone(
|
||||
&tools,
|
||||
&safety,
|
||||
&tc.name,
|
||||
&tc.arguments,
|
||||
&job_ctx,
|
||||
)
|
||||
.await;
|
||||
|
||||
let _ = channels
|
||||
.send_status(
|
||||
&channel,
|
||||
StatusUpdate::ToolCompleted {
|
||||
name: tc.name.clone(),
|
||||
success: result.is_ok(),
|
||||
},
|
||||
&metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
(spawn_idx, tc, result)
|
||||
});
|
||||
}
|
||||
|
||||
// Collect and reorder by original index
|
||||
let mut ordered: Vec<Option<(crate::llm::ToolCall, Result<String, Error>)>> =
|
||||
(0..runnable_count).map(|_| None).collect();
|
||||
while let Some(join_result) = join_set.join_next().await {
|
||||
match join_result {
|
||||
Ok((idx, tc, result)) => {
|
||||
ordered[idx] = Some((tc, result));
|
||||
}
|
||||
Err(e) => {
|
||||
if e.is_panic() {
|
||||
tracing::error!("Deferred tool execution task panicked: {}", e);
|
||||
} else {
|
||||
tracing::error!("Deferred tool execution task cancelled: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fill panicked slots with error results
|
||||
ordered
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, opt)| {
|
||||
opt.unwrap_or_else(|| {
|
||||
let tc = runnable[i].clone();
|
||||
let err: Error = crate::error::ToolError::ExecutionFailed {
|
||||
name: tc.name.clone(),
|
||||
reason: "Task failed during execution".to_string(),
|
||||
}
|
||||
.into();
|
||||
(tc, Err(err))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
|
||||
// === Phase 3: Post-flight (sequential, in original order) ===
|
||||
// Process all results before any conditional return so every
|
||||
// tool result is recorded in the session audit trail.
|
||||
let mut deferred_auth: Option<String> = None;
|
||||
|
||||
for (tc, deferred_result) in exec_results {
|
||||
if let Ok(ref output) = deferred_result
|
||||
&& !output.is_empty()
|
||||
{
|
||||
@@ -897,9 +977,10 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
// Auth detection for deferred tools
|
||||
if let Some((ext_name, instructions)) =
|
||||
check_auth_required(&tc.name, &deferred_result)
|
||||
// Auth detection — defer return until all results are recorded
|
||||
if deferred_auth.is_none()
|
||||
&& let Some((ext_name, instructions)) =
|
||||
check_auth_required(&tc.name, &deferred_result)
|
||||
{
|
||||
self.handle_auth_intercept(
|
||||
&session,
|
||||
@@ -910,7 +991,7 @@ impl Agent {
|
||||
instructions.clone(),
|
||||
)
|
||||
.await;
|
||||
return Ok(SubmissionResult::response(instructions));
|
||||
deferred_auth = Some(instructions);
|
||||
}
|
||||
|
||||
let deferred_content = match deferred_result {
|
||||
@@ -928,6 +1009,52 @@ impl Agent {
|
||||
context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content));
|
||||
}
|
||||
|
||||
// Return auth response after all results are recorded
|
||||
if let Some(instructions) = deferred_auth {
|
||||
return Ok(SubmissionResult::response(instructions));
|
||||
}
|
||||
|
||||
// Handle approval if a tool needed it
|
||||
if let Some((approval_idx, tc, tool)) = approval_needed {
|
||||
let new_pending = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
tool_name: tc.name.clone(),
|
||||
parameters: tc.arguments.clone(),
|
||||
description: tool.description().to_string(),
|
||||
tool_call_id: tc.id.clone(),
|
||||
context_messages: context_messages.clone(),
|
||||
deferred_tool_calls: deferred_tool_calls[approval_idx + 1..].to_vec(),
|
||||
};
|
||||
|
||||
let request_id = new_pending.request_id;
|
||||
let tool_name = new_pending.tool_name.clone();
|
||||
let description = new_pending.description.clone();
|
||||
let parameters = new_pending.parameters.clone();
|
||||
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.await_approval(new_pending);
|
||||
}
|
||||
}
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Awaiting approval".into()),
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
return Ok(SubmissionResult::NeedApproval {
|
||||
request_id,
|
||||
tool_name,
|
||||
description,
|
||||
parameters,
|
||||
});
|
||||
}
|
||||
|
||||
// Continue the agentic loop (a tool was already executed this turn)
|
||||
let result = self
|
||||
.run_agentic_loop(message, session.clone(), thread_id, context_messages, true)
|
||||
|
||||
Reference in New Issue
Block a user