mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
272d31797e
* chore: remove dead code (LlmEvaluator, chunk_by_paragraphs, bundled channel installer, Reasoning::safety) Delete unused code flagged in #648: - evaluation/success.rs: delete LlmEvaluator struct/impl, remove #[allow(dead_code)] from RuleBasedEvaluator methods - workspace/chunker.rs: delete chunk_by_paragraphs() and its tests (zero production callers) - extensions/manager.rs: delete install_bundled_channel_from_artifacts() (hot-activation never shipped) - llm/reasoning.rs: remove unused safety field from Reasoning struct; cascade removal through ContextCompactor, HeartbeatRunner, LlmSoftwareBuilder, and all callers Closes #648 [skip-regression-check] Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: move RuleBasedEvaluator into test module to fix dead_code warning RuleBasedEvaluator has no production callers -- it was only used in tests of itself. Moving it into #[cfg(test)] eliminates the clippy dead_code error that broke CI. [skip-regression-check] Co-Authored-By: Claude Opus 4.6 <[email protected]> --------- Co-authored-by: Claude Opus 4.6 <[email protected]>
1766 lines
66 KiB
Rust
1766 lines
66 KiB
Rust
//! Per-job worker execution.
|
|
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use tokio::sync::mpsc;
|
|
use tokio::task::JoinSet;
|
|
use uuid::Uuid;
|
|
|
|
use crate::agent::scheduler::WorkerMessage;
|
|
use crate::agent::task::TaskOutput;
|
|
use crate::channels::web::types::SseEvent;
|
|
use crate::context::{ContextManager, JobState};
|
|
use crate::db::Database;
|
|
use crate::error::Error;
|
|
use crate::hooks::HookRegistry;
|
|
use crate::llm::{
|
|
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolCall,
|
|
ToolSelection,
|
|
};
|
|
use crate::safety::SafetyLayer;
|
|
use crate::tools::rate_limiter::RateLimitResult;
|
|
use crate::tools::{ApprovalContext, ToolRegistry, redact_params};
|
|
|
|
/// Shared dependencies for worker execution.
|
|
///
|
|
/// This bundles the dependencies that are shared across all workers,
|
|
/// reducing the number of arguments to `Worker::new`.
|
|
#[derive(Clone)]
|
|
pub struct WorkerDeps {
|
|
pub context_manager: Arc<ContextManager>,
|
|
pub llm: Arc<dyn LlmProvider>,
|
|
pub safety: Arc<SafetyLayer>,
|
|
pub tools: Arc<ToolRegistry>,
|
|
pub store: Option<Arc<dyn Database>>,
|
|
pub hooks: Arc<HookRegistry>,
|
|
pub timeout: Duration,
|
|
pub use_planning: bool,
|
|
/// SSE broadcast sender for live job event streaming to the web gateway.
|
|
pub sse_tx: Option<tokio::sync::broadcast::Sender<SseEvent>>,
|
|
/// Approval context for tool execution. When `None`, all non-`Never` tools are
|
|
/// blocked (legacy behavior). When `Some`, the context determines which tools
|
|
/// are pre-approved for autonomous execution.
|
|
pub approval_context: Option<ApprovalContext>,
|
|
/// HTTP interceptor for trace recording/replay (propagated to JobContext).
|
|
pub http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
|
|
}
|
|
|
|
/// Worker that executes a single job.
|
|
pub struct Worker {
|
|
job_id: Uuid,
|
|
deps: WorkerDeps,
|
|
}
|
|
|
|
/// Result of a tool execution with metadata for context building.
|
|
struct ToolExecResult {
|
|
result: Result<String, Error>,
|
|
}
|
|
|
|
impl Worker {
|
|
/// Create a new worker for a specific job.
|
|
pub fn new(job_id: Uuid, deps: WorkerDeps) -> Self {
|
|
Self { job_id, deps }
|
|
}
|
|
|
|
// Convenience accessors to avoid deps.field everywhere
|
|
fn context_manager(&self) -> &Arc<ContextManager> {
|
|
&self.deps.context_manager
|
|
}
|
|
|
|
fn llm(&self) -> &Arc<dyn LlmProvider> {
|
|
&self.deps.llm
|
|
}
|
|
|
|
fn safety(&self) -> &Arc<SafetyLayer> {
|
|
&self.deps.safety
|
|
}
|
|
|
|
fn tools(&self) -> &Arc<ToolRegistry> {
|
|
&self.deps.tools
|
|
}
|
|
|
|
fn store(&self) -> Option<&Arc<dyn Database>> {
|
|
self.deps.store.as_ref()
|
|
}
|
|
|
|
fn timeout(&self) -> Duration {
|
|
self.deps.timeout
|
|
}
|
|
|
|
fn use_planning(&self) -> bool {
|
|
self.deps.use_planning
|
|
}
|
|
|
|
/// Fire-and-forget persistence of job status.
|
|
fn persist_status(&self, status: JobState, reason: Option<String>) {
|
|
if let Some(store) = self.store() {
|
|
let store = store.clone();
|
|
let job_id = self.job_id;
|
|
tokio::spawn(async move {
|
|
if let Err(e) = store
|
|
.update_job_status(job_id, status, reason.as_deref())
|
|
.await
|
|
{
|
|
tracing::warn!("Failed to persist status for job {}: {}", job_id, e);
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
/// Fire-and-forget persistence of a job event and SSE broadcast.
|
|
fn log_event(&self, event_type: &str, data: serde_json::Value) {
|
|
let job_id = self.job_id;
|
|
|
|
// Persist to DB
|
|
if let Some(store) = self.store() {
|
|
let store = store.clone();
|
|
let et = event_type.to_string();
|
|
let d = data.clone();
|
|
tokio::spawn(async move {
|
|
if let Err(e) = store.save_job_event(job_id, &et, &d).await {
|
|
tracing::warn!("Failed to persist event for job {}: {}", job_id, e);
|
|
}
|
|
});
|
|
}
|
|
|
|
// Broadcast SSE for live web UI updates
|
|
if let Some(ref tx) = self.deps.sse_tx {
|
|
let job_id_str = job_id.to_string();
|
|
let event = match event_type {
|
|
"message" => Some(SseEvent::JobMessage {
|
|
job_id: job_id_str,
|
|
role: data
|
|
.get("role")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("assistant")
|
|
.to_string(),
|
|
content: data
|
|
.get("content")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
}),
|
|
"tool_use" => Some(SseEvent::JobToolUse {
|
|
job_id: job_id_str,
|
|
tool_name: data
|
|
.get("tool_name")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("unknown")
|
|
.to_string(),
|
|
input: data
|
|
.get("input")
|
|
.cloned()
|
|
.unwrap_or(serde_json::Value::Null),
|
|
}),
|
|
"tool_result" => Some(SseEvent::JobToolResult {
|
|
job_id: job_id_str,
|
|
tool_name: data
|
|
.get("tool_name")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("unknown")
|
|
.to_string(),
|
|
output: data
|
|
.get("output")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
}),
|
|
"status" => Some(SseEvent::JobStatus {
|
|
job_id: job_id_str,
|
|
message: data
|
|
.get("message")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("")
|
|
.to_string(),
|
|
}),
|
|
"result" => Some(SseEvent::JobResult {
|
|
job_id: job_id_str,
|
|
status: data
|
|
.get("status")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("completed")
|
|
.to_string(),
|
|
session_id: data
|
|
.get("session_id")
|
|
.and_then(|v| v.as_str())
|
|
.map(|s| s.to_string()),
|
|
}),
|
|
_ => None,
|
|
};
|
|
if let Some(event) = event {
|
|
let _ = tx.send(event);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Run the worker until the job is complete or stopped.
|
|
pub async fn run(self, mut rx: mpsc::Receiver<WorkerMessage>) -> Result<(), Error> {
|
|
tracing::info!("Worker starting for job {}", self.job_id);
|
|
|
|
// Wait for start signal
|
|
match rx.recv().await {
|
|
Some(WorkerMessage::Start) => {}
|
|
Some(WorkerMessage::Stop) | None => {
|
|
tracing::debug!("Worker for job {} stopped before starting", self.job_id);
|
|
return Ok(());
|
|
}
|
|
Some(WorkerMessage::Ping) | Some(WorkerMessage::UserMessage(_)) => {}
|
|
}
|
|
|
|
// Get job context
|
|
let job_ctx = self.context_manager().get_context(self.job_id).await?;
|
|
|
|
// Create reasoning engine
|
|
let reasoning = Reasoning::new(self.llm().clone());
|
|
|
|
// Build initial reasoning context (tool definitions refreshed each iteration in execution_loop)
|
|
let mut reason_ctx = ReasoningContext::new().with_job(&job_ctx.description);
|
|
|
|
// Add system message
|
|
reason_ctx.messages.push(ChatMessage::system(format!(
|
|
r#"You are an autonomous agent working on a job.
|
|
|
|
Job: {}
|
|
Description: {}
|
|
|
|
You have access to tools to complete this job. Plan your approach and execute tools as needed.
|
|
You may request multiple tools at once if they can be executed in parallel.
|
|
Report when the job is complete or if you encounter issues you cannot resolve."#,
|
|
job_ctx.title, job_ctx.description
|
|
)));
|
|
|
|
// Main execution loop with timeout
|
|
let result = tokio::time::timeout(self.timeout(), async {
|
|
self.execution_loop(&mut rx, &reasoning, &mut reason_ctx)
|
|
.await
|
|
})
|
|
.await;
|
|
|
|
match result {
|
|
Ok(Ok(())) => {
|
|
tracing::info!("Worker for job {} completed successfully", self.job_id);
|
|
// Only mark completed if still in an active, non-stuck state.
|
|
// The execution_loop may have already called mark_completed or
|
|
// mark_stuck (e.g. "plan completed but work remains").
|
|
let current_state = self
|
|
.context_manager()
|
|
.get_context(self.job_id)
|
|
.await
|
|
.map(|ctx| ctx.state);
|
|
match current_state {
|
|
Ok(state) if state.is_terminal() => {
|
|
// Already in a terminal state (e.g. execution_loop
|
|
// called mark_completed itself).
|
|
}
|
|
Ok(JobState::Completed) => {
|
|
// execution_loop already called mark_completed.
|
|
}
|
|
Ok(JobState::Stuck) => {
|
|
// execution_loop marked this as stuck (e.g. "plan
|
|
// completed but work remains"); leave for self-repair.
|
|
tracing::info!(
|
|
"Job {} returned Ok but is Stuck — leaving for self-repair",
|
|
self.job_id
|
|
);
|
|
}
|
|
Ok(_) => {
|
|
self.mark_completed().await?;
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
job_id = %self.job_id,
|
|
"Failed to get job context, cannot mark as completed: {}", e
|
|
);
|
|
}
|
|
}
|
|
}
|
|
Ok(Err(e)) => {
|
|
tracing::error!("Worker for job {} failed: {}", self.job_id, e);
|
|
self.mark_failed(&e.to_string()).await?;
|
|
}
|
|
Err(_) => {
|
|
tracing::warn!("Worker for job {} timed out", self.job_id);
|
|
self.mark_stuck("Execution timeout").await?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn execution_loop(
|
|
&self,
|
|
rx: &mut mpsc::Receiver<WorkerMessage>,
|
|
reasoning: &Reasoning,
|
|
reason_ctx: &mut ReasoningContext,
|
|
) -> Result<(), Error> {
|
|
const MAX_WORKER_ITERATIONS: usize = 500;
|
|
let max_iterations = self
|
|
.context_manager()
|
|
.get_context(self.job_id)
|
|
.await
|
|
.ok()
|
|
.and_then(|ctx| ctx.metadata.get("max_iterations").and_then(|v| v.as_u64()))
|
|
.unwrap_or(50) as usize;
|
|
let max_iterations = max_iterations.min(MAX_WORKER_ITERATIONS);
|
|
let mut iteration = 0;
|
|
const MAX_CONSECUTIVE_RATE_LIMITS: usize = 10;
|
|
let mut consecutive_rate_limits = 0usize;
|
|
const MAX_TOOL_INTENT_NUDGES: u32 = 2;
|
|
let mut consecutive_tool_intent_nudges: u32 = 0;
|
|
|
|
// Initial tool definitions for planning (will be refreshed in loop)
|
|
reason_ctx.available_tools = self.tools().tool_definitions().await;
|
|
|
|
// Generate plan if planning is enabled
|
|
let plan = if self.use_planning() {
|
|
match reasoning.plan(reason_ctx).await {
|
|
Ok(p) => {
|
|
tracing::info!(
|
|
"Created plan for job {}: {} actions, {:.0}% confidence",
|
|
self.job_id,
|
|
p.actions.len(),
|
|
p.confidence * 100.0
|
|
);
|
|
|
|
// Add plan to context as assistant message
|
|
reason_ctx.messages.push(ChatMessage::assistant(format!(
|
|
"I've created a plan to accomplish this goal: {}\n\nSteps:\n{}",
|
|
p.goal,
|
|
p.actions
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, a)| format!("{}. {} - {}", i + 1, a.tool_name, a.reasoning))
|
|
.collect::<Vec<_>>()
|
|
.join("\n")
|
|
)));
|
|
|
|
self.log_event("message", serde_json::json!({
|
|
"role": "assistant",
|
|
"content": format!("Plan: {}\n\n{}", p.goal,
|
|
p.actions.iter().enumerate()
|
|
.map(|(i, a)| format!("{}. {} - {}", i + 1, a.tool_name, a.reasoning))
|
|
.collect::<Vec<_>>().join("\n"))
|
|
}));
|
|
|
|
Some(p)
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
"Planning failed for job {}, falling back to direct selection: {}",
|
|
self.job_id,
|
|
e
|
|
);
|
|
None
|
|
}
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// If we have a plan, execute it. Two exit paths:
|
|
// 1. Plan ran to completion → job is Completed or needs continuation
|
|
// (check state and only fall through if not terminal)
|
|
// 2. Plan was interrupted by UserMessage → fall through to direct loop
|
|
if let Some(ref plan) = plan {
|
|
self.execute_plan(rx, reasoning, reason_ctx, plan).await?;
|
|
|
|
// If the plan marked the job completed, terminal, or stuck, we're
|
|
// done. Only fall through to the direct selection loop if the
|
|
// plan was interrupted or explicitly left the job in-progress.
|
|
if let Ok(ctx) = self.context_manager().get_context(self.job_id).await
|
|
&& (ctx.state.is_terminal()
|
|
|| ctx.state == JobState::Stuck
|
|
|| ctx.state == JobState::Completed)
|
|
{
|
|
return Ok(());
|
|
}
|
|
}
|
|
|
|
// Direct tool selection loop (also used as fallback after plan interruption)
|
|
loop {
|
|
// Check for stop signal and injected user messages
|
|
while let Ok(msg) = rx.try_recv() {
|
|
match msg {
|
|
WorkerMessage::Stop => {
|
|
tracing::debug!("Worker for job {} received stop signal", self.job_id);
|
|
return Ok(());
|
|
}
|
|
WorkerMessage::Ping => {
|
|
tracing::trace!("Worker for job {} received ping", self.job_id);
|
|
}
|
|
WorkerMessage::Start => {}
|
|
WorkerMessage::UserMessage(content) => {
|
|
tracing::info!(
|
|
job_id = %self.job_id,
|
|
"Worker received follow-up user message"
|
|
);
|
|
reason_ctx.messages.push(ChatMessage::user(&content));
|
|
self.log_event(
|
|
"message",
|
|
serde_json::json!({
|
|
"role": "user",
|
|
"content": content,
|
|
}),
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check for cancellation
|
|
if let Ok(ctx) = self.context_manager().get_context(self.job_id).await
|
|
&& ctx.state == JobState::Cancelled
|
|
{
|
|
tracing::info!("Worker for job {} detected cancellation", self.job_id);
|
|
return Ok(());
|
|
}
|
|
|
|
iteration += 1;
|
|
if iteration > max_iterations {
|
|
self.mark_stuck("Maximum iterations exceeded").await?;
|
|
return Ok(());
|
|
}
|
|
|
|
// Refresh tool definitions so newly built tools become visible
|
|
reason_ctx.available_tools = self.tools().tool_definitions().await;
|
|
|
|
// Select next tool(s) to use, with rate-limit retry.
|
|
let selections = match reasoning.select_tools(reason_ctx).await {
|
|
Ok(s) => s,
|
|
Err(crate::error::LlmError::RateLimited { retry_after, .. }) => {
|
|
consecutive_rate_limits += 1;
|
|
let wait = retry_after.unwrap_or(Duration::from_secs(5));
|
|
tracing::warn!(
|
|
job_id = %self.job_id,
|
|
wait_secs = wait.as_secs(),
|
|
attempt = consecutive_rate_limits,
|
|
"LLM rate limited during tool selection, backing off"
|
|
);
|
|
if consecutive_rate_limits >= MAX_CONSECUTIVE_RATE_LIMITS {
|
|
self.mark_stuck("Persistent rate limiting").await?;
|
|
return Ok(());
|
|
}
|
|
self.log_event(
|
|
"status",
|
|
serde_json::json!({
|
|
"message": format!("Rate limited, retrying in {}s ({}/{})...",
|
|
wait.as_secs(), consecutive_rate_limits, MAX_CONSECUTIVE_RATE_LIMITS),
|
|
}),
|
|
);
|
|
tokio::time::sleep(wait).await;
|
|
continue;
|
|
}
|
|
Err(e) => return Err(e.into()),
|
|
};
|
|
|
|
if selections.is_empty() {
|
|
// No tools from select_tools, ask LLM directly (may still return tool calls)
|
|
let respond_output = match reasoning.respond_with_tools(reason_ctx).await {
|
|
Ok(o) => o,
|
|
Err(crate::error::LlmError::RateLimited { retry_after, .. }) => {
|
|
consecutive_rate_limits += 1;
|
|
let wait = retry_after.unwrap_or(Duration::from_secs(5));
|
|
tracing::warn!(
|
|
job_id = %self.job_id,
|
|
wait_secs = wait.as_secs(),
|
|
attempt = consecutive_rate_limits,
|
|
"LLM rate limited during respond_with_tools, backing off"
|
|
);
|
|
if consecutive_rate_limits >= MAX_CONSECUTIVE_RATE_LIMITS {
|
|
self.mark_stuck("Persistent rate limiting").await?;
|
|
return Ok(());
|
|
}
|
|
self.log_event(
|
|
"status",
|
|
serde_json::json!({
|
|
"message": format!("Rate limited, retrying in {}s ({}/{})...",
|
|
wait.as_secs(), consecutive_rate_limits, MAX_CONSECUTIVE_RATE_LIMITS),
|
|
}),
|
|
);
|
|
tokio::time::sleep(wait).await;
|
|
continue;
|
|
}
|
|
Err(e) => return Err(e.into()),
|
|
};
|
|
|
|
match respond_output.result {
|
|
RespondResult::Text(response) => {
|
|
// Check for explicit completion phrases. Use word-boundary
|
|
// aware checks to avoid false positives like "incomplete",
|
|
// "not done", or "unfinished". Only the LLM's own response
|
|
// (not tool output) can trigger this.
|
|
if crate::util::llm_signals_completion(&response) {
|
|
self.mark_completed().await?;
|
|
return Ok(());
|
|
}
|
|
|
|
// Add assistant response to context
|
|
reason_ctx.messages.push(ChatMessage::assistant(&response));
|
|
|
|
self.log_event(
|
|
"message",
|
|
serde_json::json!({
|
|
"role": "assistant",
|
|
"content": response,
|
|
}),
|
|
);
|
|
|
|
// Nudge the LLM if it expressed tool intent without calling tools
|
|
let signals_intent = !reason_ctx.available_tools.is_empty()
|
|
&& crate::llm::llm_signals_tool_intent(&response);
|
|
if signals_intent && consecutive_tool_intent_nudges < MAX_TOOL_INTENT_NUDGES
|
|
{
|
|
consecutive_tool_intent_nudges += 1;
|
|
tracing::info!(
|
|
job_id = %self.job_id,
|
|
"LLM expressed tool intent without calling a tool, nudging"
|
|
);
|
|
reason_ctx
|
|
.messages
|
|
.push(ChatMessage::user(crate::llm::TOOL_INTENT_NUDGE));
|
|
} else if !signals_intent {
|
|
consecutive_tool_intent_nudges = 0;
|
|
if iteration > 3 && iteration % 5 == 0 {
|
|
// Generic fallback nudge
|
|
reason_ctx.messages.push(ChatMessage::user(
|
|
"Are you stuck? Do you need help completing this job?",
|
|
));
|
|
}
|
|
}
|
|
}
|
|
RespondResult::ToolCalls {
|
|
tool_calls,
|
|
content,
|
|
} => {
|
|
consecutive_tool_intent_nudges = 0;
|
|
// Model returned tool calls - execute them
|
|
tracing::debug!(
|
|
"Job {} respond_with_tools returned {} tool calls",
|
|
self.job_id,
|
|
tool_calls.len()
|
|
);
|
|
|
|
if let Some(ref text) = content {
|
|
self.log_event(
|
|
"message",
|
|
serde_json::json!({
|
|
"role": "assistant",
|
|
"content": text,
|
|
}),
|
|
);
|
|
}
|
|
|
|
// Add assistant message with tool_calls (OpenAI protocol)
|
|
reason_ctx
|
|
.messages
|
|
.push(ChatMessage::assistant_with_tool_calls(
|
|
content,
|
|
tool_calls.clone(),
|
|
));
|
|
|
|
// Convert ToolCalls to ToolSelections and execute in parallel
|
|
let selections: Vec<ToolSelection> = tool_calls
|
|
.iter()
|
|
.map(|tc| ToolSelection {
|
|
tool_name: tc.name.clone(),
|
|
parameters: tc.arguments.clone(),
|
|
reasoning: String::new(),
|
|
alternatives: vec![],
|
|
tool_call_id: tc.id.clone(),
|
|
})
|
|
.collect();
|
|
|
|
let results = self.execute_tools_parallel(&selections).await;
|
|
for (selection, result) in selections.iter().zip(results) {
|
|
self.process_tool_result(reason_ctx, selection, result.result)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
consecutive_tool_intent_nudges = 0;
|
|
|
|
// Record the assistant tool_calls message so that tool_result
|
|
// messages have a matching parent (prevents orphaned rewrites).
|
|
let tool_calls: Vec<ToolCall> = selections
|
|
.iter()
|
|
.map(|s| ToolCall {
|
|
id: s.tool_call_id.clone(),
|
|
name: s.tool_name.clone(),
|
|
arguments: s.parameters.clone(),
|
|
})
|
|
.collect();
|
|
reason_ctx
|
|
.messages
|
|
.push(ChatMessage::assistant_with_tool_calls(None, tool_calls));
|
|
|
|
if selections.len() == 1 {
|
|
// Single tool: execute directly
|
|
let selection = &selections[0];
|
|
tracing::debug!(
|
|
"Job {} selecting tool: {} - {}",
|
|
self.job_id,
|
|
selection.tool_name,
|
|
selection.reasoning
|
|
);
|
|
|
|
let result = self
|
|
.execute_tool(&selection.tool_name, &selection.parameters)
|
|
.await;
|
|
|
|
self.process_tool_result(reason_ctx, selection, result)
|
|
.await?;
|
|
} else {
|
|
// Multiple tools: execute in parallel
|
|
tracing::debug!(
|
|
"Job {} executing {} tools in parallel",
|
|
self.job_id,
|
|
selections.len()
|
|
);
|
|
|
|
let results = self.execute_tools_parallel(&selections).await;
|
|
|
|
// Process all results
|
|
for (selection, result) in selections.iter().zip(results) {
|
|
self.process_tool_result(reason_ctx, selection, result.result)
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Reset rate-limit counter after a successful iteration (all LLM
|
|
// calls succeeded). Placed here so alternating success/fail between
|
|
// select_tools and respond_with_tools cannot bypass the cap.
|
|
consecutive_rate_limits = 0;
|
|
|
|
// Small delay between iterations
|
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
}
|
|
}
|
|
|
|
/// Execute multiple tools in parallel using a JoinSet.
|
|
///
|
|
/// Each task is tagged with its original index so results are returned
|
|
/// in the same order as `selections`, regardless of completion order.
|
|
async fn execute_tools_parallel(&self, selections: &[ToolSelection]) -> Vec<ToolExecResult> {
|
|
let count = selections.len();
|
|
|
|
// Short-circuit for single tool: execute directly without JoinSet overhead
|
|
if count <= 1 {
|
|
let mut results = Vec::with_capacity(count);
|
|
for selection in selections {
|
|
let result = Self::execute_tool_inner(
|
|
&self.deps,
|
|
self.job_id,
|
|
&selection.tool_name,
|
|
&selection.parameters,
|
|
)
|
|
.await;
|
|
results.push(ToolExecResult { result });
|
|
}
|
|
return results;
|
|
}
|
|
|
|
let mut join_set = JoinSet::new();
|
|
|
|
for (idx, selection) in selections.iter().enumerate() {
|
|
let deps = self.deps.clone();
|
|
let job_id = self.job_id;
|
|
let tool_name = selection.tool_name.clone();
|
|
let params = selection.parameters.clone();
|
|
join_set.spawn(async move {
|
|
let result = Self::execute_tool_inner(&deps, job_id, &tool_name, ¶ms).await;
|
|
(idx, ToolExecResult { result })
|
|
});
|
|
}
|
|
|
|
// Collect and reorder by original index
|
|
let mut results: Vec<Option<ToolExecResult>> = (0..count).map(|_| None).collect();
|
|
while let Some(join_result) = join_set.join_next().await {
|
|
match join_result {
|
|
Ok((idx, exec_result)) => results[idx] = Some(exec_result),
|
|
Err(e) => {
|
|
if e.is_panic() {
|
|
tracing::error!("Tool execution task panicked: {}", e);
|
|
} else {
|
|
tracing::error!("Tool execution task cancelled: {}", e);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fill any panicked slots with error results
|
|
results
|
|
.into_iter()
|
|
.enumerate()
|
|
.map(|(i, opt)| {
|
|
opt.unwrap_or_else(|| ToolExecResult {
|
|
result: Err(crate::error::ToolError::ExecutionFailed {
|
|
name: selections[i].tool_name.clone(),
|
|
reason: "Task failed during execution".to_string(),
|
|
}
|
|
.into()),
|
|
})
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Inner tool execution logic that can be called from both single and parallel paths.
|
|
async fn execute_tool_inner(
|
|
deps: &WorkerDeps,
|
|
job_id: Uuid,
|
|
tool_name: &str,
|
|
params: &serde_json::Value,
|
|
) -> Result<String, Error> {
|
|
let tool =
|
|
deps.tools
|
|
.get(tool_name)
|
|
.await
|
|
.ok_or_else(|| crate::error::ToolError::NotFound {
|
|
name: tool_name.to_string(),
|
|
})?;
|
|
|
|
// Check approval: use context-aware check if available, else block all non-Never tools
|
|
let requirement = tool.requires_approval(params);
|
|
let blocked =
|
|
ApprovalContext::is_blocked_or_default(&deps.approval_context, tool_name, requirement);
|
|
if blocked {
|
|
return Err(crate::error::ToolError::AuthRequired {
|
|
name: tool_name.to_string(),
|
|
}
|
|
.into());
|
|
}
|
|
|
|
// Fetch job context early so we have the real user_id for hooks and rate limiting
|
|
let mut job_ctx = deps.context_manager.get_context(job_id).await?;
|
|
// Propagate http_interceptor for trace recording/replay
|
|
if job_ctx.http_interceptor.is_none() {
|
|
job_ctx.http_interceptor = deps.http_interceptor.clone();
|
|
}
|
|
|
|
// Check per-tool rate limit before running hooks or executing (cheaper check first)
|
|
if let Some(config) = tool.rate_limit_config()
|
|
&& let RateLimitResult::Limited { retry_after, .. } = deps
|
|
.tools
|
|
.rate_limiter()
|
|
.check_and_record(&job_ctx.user_id, tool_name, &config)
|
|
.await
|
|
{
|
|
return Err(crate::error::ToolError::RateLimited {
|
|
name: tool_name.to_string(),
|
|
retry_after: Some(retry_after),
|
|
}
|
|
.into());
|
|
}
|
|
|
|
// Run BeforeToolCall hook
|
|
let params = {
|
|
use crate::hooks::{HookError, HookEvent, HookOutcome};
|
|
let hook_params = redact_params(params, tool.sensitive_params());
|
|
let event = HookEvent::ToolCall {
|
|
tool_name: tool_name.to_string(),
|
|
parameters: hook_params,
|
|
user_id: job_ctx.user_id.clone(),
|
|
context: format!("job:{}", job_id),
|
|
};
|
|
match deps.hooks.run(&event).await {
|
|
Err(HookError::Rejected { reason }) => {
|
|
return Err(crate::error::ToolError::ExecutionFailed {
|
|
name: tool_name.to_string(),
|
|
reason: format!("Blocked by hook: {}", reason),
|
|
}
|
|
.into());
|
|
}
|
|
Err(err) => {
|
|
return Err(crate::error::ToolError::ExecutionFailed {
|
|
name: tool_name.to_string(),
|
|
reason: format!("Blocked by hook failure mode: {}", err),
|
|
}
|
|
.into());
|
|
}
|
|
Ok(HookOutcome::Continue {
|
|
modified: Some(new_params),
|
|
}) => serde_json::from_str(&new_params).unwrap_or_else(|e| {
|
|
tracing::warn!(
|
|
tool = %tool_name,
|
|
"Hook returned non-JSON modification for ToolCall, ignoring: {}",
|
|
e
|
|
);
|
|
params.clone()
|
|
}),
|
|
_ => params.clone(),
|
|
}
|
|
};
|
|
if job_ctx.state == JobState::Cancelled {
|
|
return Err(crate::error::ToolError::ExecutionFailed {
|
|
name: tool_name.to_string(),
|
|
reason: "Job is cancelled".to_string(),
|
|
}
|
|
.into());
|
|
}
|
|
|
|
// Validate tool parameters
|
|
let validation = deps.safety.validator().validate_tool_params(¶ms);
|
|
if !validation.is_valid {
|
|
let details = validation
|
|
.errors
|
|
.iter()
|
|
.map(|e| format!("{}: {}", e.field, e.message))
|
|
.collect::<Vec<_>>()
|
|
.join("; ");
|
|
return Err(crate::error::ToolError::InvalidParameters {
|
|
name: tool_name.to_string(),
|
|
reason: format!("Invalid tool parameters: {}", details),
|
|
}
|
|
.into());
|
|
}
|
|
|
|
// Redact sensitive parameter values (e.g. secret_save's "value") before
|
|
// they touch any observability or audit path.
|
|
let safe_params = redact_params(¶ms, tool.sensitive_params());
|
|
tracing::debug!(
|
|
tool = %tool_name,
|
|
params = %safe_params,
|
|
job = %job_id,
|
|
"Tool call started"
|
|
);
|
|
|
|
// Execute with per-tool timeout and timing
|
|
let tool_timeout = tool.execution_timeout();
|
|
let start = std::time::Instant::now();
|
|
let result = tokio::time::timeout(tool_timeout, async {
|
|
tool.execute(params.clone(), &job_ctx).await
|
|
})
|
|
.await;
|
|
let elapsed = start.elapsed();
|
|
|
|
match &result {
|
|
Ok(Ok(output)) => {
|
|
let result_str = serde_json::to_string(&output.result)
|
|
.unwrap_or_else(|_| "<serialize error>".to_string());
|
|
tracing::debug!(
|
|
tool = %tool_name,
|
|
elapsed_ms = elapsed.as_millis() as u64,
|
|
result = %result_str,
|
|
"Tool call succeeded"
|
|
);
|
|
}
|
|
Ok(Err(e)) => {
|
|
tracing::debug!(
|
|
tool = %tool_name,
|
|
elapsed_ms = elapsed.as_millis() as u64,
|
|
error = %e,
|
|
"Tool call failed"
|
|
);
|
|
}
|
|
Err(_) => {
|
|
tracing::debug!(
|
|
tool = %tool_name,
|
|
elapsed_ms = elapsed.as_millis() as u64,
|
|
timeout_secs = tool_timeout.as_secs(),
|
|
"Tool call timed out"
|
|
);
|
|
}
|
|
}
|
|
|
|
// Record action in memory and get the ActionRecord for persistence
|
|
let action = match &result {
|
|
Ok(Ok(output)) => {
|
|
let output_str = serde_json::to_string_pretty(&output.result)
|
|
.ok()
|
|
.map(|s| deps.safety.sanitize_tool_output(tool_name, &s).content);
|
|
match deps
|
|
.context_manager
|
|
.update_memory(job_id, |mem| {
|
|
let rec = mem.create_action(tool_name, safe_params.clone()).succeed(
|
|
output_str.clone(),
|
|
output.result.clone(),
|
|
elapsed,
|
|
);
|
|
mem.record_action(rec.clone());
|
|
rec
|
|
})
|
|
.await
|
|
{
|
|
Ok(rec) => Some(rec),
|
|
Err(e) => {
|
|
tracing::warn!(job_id = %job_id, tool = tool_name, "Failed to record action in memory: {e}");
|
|
None
|
|
}
|
|
}
|
|
}
|
|
Ok(Err(e)) => {
|
|
match deps
|
|
.context_manager
|
|
.update_memory(job_id, |mem| {
|
|
let rec = mem
|
|
.create_action(tool_name, safe_params.clone())
|
|
.fail(e.to_string(), elapsed);
|
|
mem.record_action(rec.clone());
|
|
rec
|
|
})
|
|
.await
|
|
{
|
|
Ok(rec) => Some(rec),
|
|
Err(e) => {
|
|
tracing::warn!(job_id = %job_id, tool = tool_name, "Failed to record action in memory: {e}");
|
|
None
|
|
}
|
|
}
|
|
}
|
|
Err(_) => {
|
|
match deps
|
|
.context_manager
|
|
.update_memory(job_id, |mem| {
|
|
let rec = mem
|
|
.create_action(tool_name, safe_params.clone())
|
|
.fail("Execution timeout", elapsed);
|
|
mem.record_action(rec.clone());
|
|
rec
|
|
})
|
|
.await
|
|
{
|
|
Ok(rec) => Some(rec),
|
|
Err(e) => {
|
|
tracing::warn!(job_id = %job_id, tool = tool_name, "Failed to record action in memory: {e}");
|
|
None
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
// Persist action to database (fire-and-forget)
|
|
if let (Some(action), Some(store)) = (action, deps.store.clone()) {
|
|
tokio::spawn(async move {
|
|
if let Err(e) = store.save_action(job_id, &action).await {
|
|
tracing::warn!("Failed to persist action for job {}: {}", job_id, e);
|
|
}
|
|
});
|
|
}
|
|
|
|
// Handle the result
|
|
let output = result
|
|
.map_err(|_| crate::error::ToolError::Timeout {
|
|
name: tool_name.to_string(),
|
|
timeout: tool_timeout,
|
|
})?
|
|
.map_err(|e| crate::error::ToolError::ExecutionFailed {
|
|
name: tool_name.to_string(),
|
|
reason: e.to_string(),
|
|
})?;
|
|
|
|
// Return result as string
|
|
serde_json::to_string_pretty(&output.result).map_err(|e| {
|
|
crate::error::ToolError::ExecutionFailed {
|
|
name: tool_name.to_string(),
|
|
reason: format!("Failed to serialize result: {}", e),
|
|
}
|
|
.into()
|
|
})
|
|
}
|
|
|
|
/// Process a tool execution result and add it to the reasoning context.
|
|
async fn process_tool_result(
|
|
&self,
|
|
reason_ctx: &mut ReasoningContext,
|
|
selection: &ToolSelection,
|
|
result: Result<String, Error>,
|
|
) -> Result<bool, Error> {
|
|
self.log_event(
|
|
"tool_use",
|
|
serde_json::json!({
|
|
"tool_name": selection.tool_name,
|
|
"input": crate::agent::agent_loop::truncate_for_preview(
|
|
&selection.parameters.to_string(), 500),
|
|
}),
|
|
);
|
|
|
|
match result {
|
|
Ok(output) => {
|
|
// Sanitize output
|
|
let sanitized = self
|
|
.safety()
|
|
.sanitize_tool_output(&selection.tool_name, &output);
|
|
|
|
// Add to context
|
|
let wrapped = self.safety().wrap_for_llm(
|
|
&selection.tool_name,
|
|
&sanitized.content,
|
|
sanitized.was_modified,
|
|
);
|
|
|
|
reason_ctx.messages.push(ChatMessage::tool_result(
|
|
&selection.tool_call_id,
|
|
&selection.tool_name,
|
|
wrapped,
|
|
));
|
|
|
|
self.log_event("tool_result", serde_json::json!({
|
|
"tool_name": selection.tool_name,
|
|
"success": true,
|
|
"output": crate::agent::agent_loop::truncate_for_preview(&sanitized.content, 500),
|
|
}));
|
|
|
|
// Tool output never drives job completion. A malicious tool could
|
|
// emit "TASK_COMPLETE" to force premature completion. Only the LLM's
|
|
// own structured response (in execution_loop) can mark a job done.
|
|
Ok(false)
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
"Tool {} failed for job {}: {}",
|
|
selection.tool_name,
|
|
self.job_id,
|
|
e
|
|
);
|
|
|
|
// Record failure for self-repair tracking
|
|
if let Some(store) = self.store() {
|
|
let store = store.clone();
|
|
let tool_name = selection.tool_name.clone();
|
|
let error_msg = e.to_string();
|
|
tokio::spawn(async move {
|
|
if let Err(db_err) = store.record_tool_failure(&tool_name, &error_msg).await
|
|
{
|
|
tracing::warn!("Failed to record tool failure: {}", db_err);
|
|
}
|
|
});
|
|
}
|
|
|
|
self.log_event(
|
|
"tool_result",
|
|
serde_json::json!({
|
|
"tool_name": selection.tool_name,
|
|
"success": false,
|
|
"output": format!("Error: {}", e),
|
|
}),
|
|
);
|
|
|
|
reason_ctx.messages.push(ChatMessage::tool_result(
|
|
&selection.tool_call_id,
|
|
&selection.tool_name,
|
|
format!("Error: {}", e),
|
|
));
|
|
|
|
Ok(false)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Execute a pre-generated plan.
|
|
async fn execute_plan(
|
|
&self,
|
|
rx: &mut mpsc::Receiver<WorkerMessage>,
|
|
reasoning: &Reasoning,
|
|
reason_ctx: &mut ReasoningContext,
|
|
plan: &ActionPlan,
|
|
) -> Result<(), Error> {
|
|
for (i, action) in plan.actions.iter().enumerate() {
|
|
// Check for stop signal and injected user messages
|
|
while let Ok(msg) = rx.try_recv() {
|
|
match msg {
|
|
WorkerMessage::Stop => {
|
|
tracing::debug!(
|
|
"Worker for job {} received stop signal during plan execution",
|
|
self.job_id
|
|
);
|
|
return Ok(());
|
|
}
|
|
WorkerMessage::Ping => {
|
|
tracing::trace!("Worker for job {} received ping", self.job_id);
|
|
}
|
|
WorkerMessage::Start => {}
|
|
WorkerMessage::UserMessage(content) => {
|
|
tracing::info!(
|
|
job_id = %self.job_id,
|
|
"User message received during plan execution, abandoning plan"
|
|
);
|
|
reason_ctx.messages.push(ChatMessage::user(&content));
|
|
self.log_event(
|
|
"message",
|
|
serde_json::json!({
|
|
"role": "user",
|
|
"content": content,
|
|
}),
|
|
);
|
|
self.log_event(
|
|
"status",
|
|
serde_json::json!({
|
|
"message": "Plan interrupted by user message, re-evaluating...",
|
|
}),
|
|
);
|
|
// Return Ok to break out of plan; caller falls through to
|
|
// the direct selection loop for LLM re-evaluation.
|
|
return Ok(());
|
|
}
|
|
}
|
|
}
|
|
|
|
tracing::debug!(
|
|
"Job {} executing planned action {}/{}: {} - {}",
|
|
self.job_id,
|
|
i + 1,
|
|
plan.actions.len(),
|
|
action.tool_name,
|
|
action.reasoning
|
|
);
|
|
|
|
// Create a synthetic ToolSelection for process_tool_result.
|
|
// Plan actions don't originate from an LLM tool_call response so
|
|
// there is no real tool_call_id; generate a unique one.
|
|
let selection = ToolSelection {
|
|
tool_name: action.tool_name.clone(),
|
|
parameters: action.parameters.clone(),
|
|
reasoning: action.reasoning.clone(),
|
|
alternatives: vec![],
|
|
tool_call_id: format!("plan_{}_{}", self.job_id, i),
|
|
};
|
|
|
|
// Record the assistant tool_calls message so that the tool_result
|
|
// has a matching parent (prevents orphaned rewrites).
|
|
reason_ctx
|
|
.messages
|
|
.push(ChatMessage::assistant_with_tool_calls(
|
|
None,
|
|
vec![ToolCall {
|
|
id: selection.tool_call_id.clone(),
|
|
name: selection.tool_name.clone(),
|
|
arguments: selection.parameters.clone(),
|
|
}],
|
|
));
|
|
|
|
// Execute the planned tool
|
|
let result = self
|
|
.execute_tool(&action.tool_name, &action.parameters)
|
|
.await;
|
|
|
|
// Process the result
|
|
let completed = self
|
|
.process_tool_result(reason_ctx, &selection, result)
|
|
.await?;
|
|
|
|
if completed {
|
|
return Ok(());
|
|
}
|
|
|
|
// Small delay between actions
|
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
}
|
|
|
|
// Plan completed, check with LLM if job is done
|
|
reason_ctx.messages.push(ChatMessage::user(
|
|
"All planned actions have been executed. Is the job complete? If not, what else needs to be done?",
|
|
));
|
|
|
|
let response = reasoning.respond(reason_ctx).await?;
|
|
reason_ctx.messages.push(ChatMessage::assistant(&response));
|
|
|
|
if crate::util::llm_signals_completion(&response) {
|
|
self.mark_completed().await?;
|
|
} else {
|
|
// Job not complete — return Ok without marking terminal so the
|
|
// caller falls through to the direct selection loop for continuation.
|
|
tracing::info!(
|
|
"Job {} plan completed but work remains, falling back to direct selection",
|
|
self.job_id
|
|
);
|
|
self.log_event(
|
|
"status",
|
|
serde_json::json!({
|
|
"message": "Plan completed but job needs more work, continuing...",
|
|
}),
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn execute_tool(
|
|
&self,
|
|
tool_name: &str,
|
|
params: &serde_json::Value,
|
|
) -> Result<String, Error> {
|
|
Self::execute_tool_inner(&self.deps, self.job_id, tool_name, params).await
|
|
}
|
|
|
|
async fn mark_completed(&self) -> Result<(), Error> {
|
|
self.context_manager()
|
|
.update_context(self.job_id, |ctx| {
|
|
ctx.transition_to(
|
|
JobState::Completed,
|
|
Some("Job completed successfully".to_string()),
|
|
)
|
|
})
|
|
.await?
|
|
.map_err(|s| crate::error::JobError::ContextError {
|
|
id: self.job_id,
|
|
reason: s,
|
|
})?;
|
|
|
|
self.log_event(
|
|
"result",
|
|
serde_json::json!({
|
|
"status": "completed",
|
|
"success": true,
|
|
"message": "Job completed successfully",
|
|
}),
|
|
);
|
|
self.persist_status(
|
|
JobState::Completed,
|
|
Some("Job completed successfully".to_string()),
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
async fn mark_failed(&self, reason: &str) -> Result<(), Error> {
|
|
self.context_manager()
|
|
.update_context(self.job_id, |ctx| {
|
|
ctx.transition_to(JobState::Failed, Some(reason.to_string()))
|
|
})
|
|
.await?
|
|
.map_err(|s| crate::error::JobError::ContextError {
|
|
id: self.job_id,
|
|
reason: s,
|
|
})?;
|
|
|
|
self.log_event(
|
|
"result",
|
|
serde_json::json!({
|
|
"status": "failed",
|
|
"success": false,
|
|
"message": format!("Execution failed: {}", reason),
|
|
}),
|
|
);
|
|
self.persist_status(JobState::Failed, Some(reason.to_string()));
|
|
Ok(())
|
|
}
|
|
|
|
async fn mark_stuck(&self, reason: &str) -> Result<(), Error> {
|
|
self.context_manager()
|
|
.update_context(self.job_id, |ctx| ctx.mark_stuck(reason))
|
|
.await?
|
|
.map_err(|s| crate::error::JobError::ContextError {
|
|
id: self.job_id,
|
|
reason: s,
|
|
})?;
|
|
|
|
self.log_event(
|
|
"result",
|
|
serde_json::json!({
|
|
"status": "stuck",
|
|
"success": false,
|
|
"message": format!("Job stuck: {}", reason),
|
|
}),
|
|
);
|
|
self.persist_status(JobState::Stuck, Some(reason.to_string()));
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Convert a TaskOutput to a string result for tool execution.
|
|
impl From<TaskOutput> for Result<String, Error> {
|
|
fn from(output: TaskOutput) -> Self {
|
|
serde_json::to_string_pretty(&output.result).map_err(|e| {
|
|
crate::error::ToolError::ExecutionFailed {
|
|
name: "task".to_string(),
|
|
reason: format!("Failed to serialize result: {}", e),
|
|
}
|
|
.into()
|
|
})
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use crate::llm::ToolSelection;
|
|
use crate::util::llm_signals_completion;
|
|
|
|
use super::*;
|
|
use crate::config::SafetyConfig;
|
|
use crate::context::JobContext;
|
|
use crate::llm::{
|
|
CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest,
|
|
ToolCompletionResponse,
|
|
};
|
|
use crate::safety::SafetyLayer;
|
|
use crate::tools::{Tool, ToolError, ToolOutput};
|
|
|
|
/// A test tool that sleeps for a configurable duration before returning.
|
|
struct SlowTool {
|
|
tool_name: String,
|
|
delay: Duration,
|
|
}
|
|
|
|
#[async_trait::async_trait]
|
|
impl Tool for SlowTool {
|
|
fn name(&self) -> &str {
|
|
&self.tool_name
|
|
}
|
|
fn description(&self) -> &str {
|
|
"Test tool with configurable delay"
|
|
}
|
|
fn parameters_schema(&self) -> serde_json::Value {
|
|
serde_json::json!({"type": "object", "properties": {}})
|
|
}
|
|
async fn execute(
|
|
&self,
|
|
_params: serde_json::Value,
|
|
_ctx: &JobContext,
|
|
) -> Result<ToolOutput, ToolError> {
|
|
let start = std::time::Instant::now();
|
|
tokio::time::sleep(self.delay).await;
|
|
Ok(ToolOutput::text(
|
|
format!("done_{}", self.tool_name),
|
|
start.elapsed(),
|
|
))
|
|
}
|
|
fn requires_sanitization(&self) -> bool {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Stub LLM provider (never called in these tests).
|
|
struct StubLlm;
|
|
|
|
#[async_trait::async_trait]
|
|
impl LlmProvider for StubLlm {
|
|
fn model_name(&self) -> &str {
|
|
"stub"
|
|
}
|
|
fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) {
|
|
(rust_decimal::Decimal::ZERO, rust_decimal::Decimal::ZERO)
|
|
}
|
|
async fn complete(
|
|
&self,
|
|
_req: CompletionRequest,
|
|
) -> Result<CompletionResponse, crate::error::LlmError> {
|
|
unimplemented!("stub")
|
|
}
|
|
async fn complete_with_tools(
|
|
&self,
|
|
_req: ToolCompletionRequest,
|
|
) -> Result<ToolCompletionResponse, crate::error::LlmError> {
|
|
unimplemented!("stub")
|
|
}
|
|
}
|
|
|
|
/// Build a Worker wired to a ToolRegistry containing the given tools.
|
|
async fn make_worker(tools: Vec<Arc<dyn Tool>>) -> Worker {
|
|
let registry = ToolRegistry::new();
|
|
for t in tools {
|
|
registry.register(t).await;
|
|
}
|
|
|
|
let cm = Arc::new(crate::context::ContextManager::new(5));
|
|
let job_id = cm.create_job("test", "test job").await.unwrap();
|
|
|
|
let deps = WorkerDeps {
|
|
context_manager: cm,
|
|
llm: Arc::new(StubLlm),
|
|
safety: Arc::new(SafetyLayer::new(&SafetyConfig {
|
|
max_output_length: 100_000,
|
|
injection_check_enabled: false,
|
|
})),
|
|
tools: Arc::new(registry),
|
|
store: None,
|
|
hooks: Arc::new(crate::hooks::HookRegistry::new()),
|
|
timeout: Duration::from_secs(30),
|
|
use_planning: false,
|
|
sse_tx: None,
|
|
approval_context: None,
|
|
http_interceptor: None,
|
|
};
|
|
|
|
Worker::new(job_id, deps)
|
|
}
|
|
|
|
#[test]
|
|
fn test_tool_selection_preserves_call_id() {
|
|
let selection = ToolSelection {
|
|
tool_name: "memory_search".to_string(),
|
|
parameters: serde_json::json!({"query": "test"}),
|
|
reasoning: "Need to search memory".to_string(),
|
|
alternatives: vec![],
|
|
tool_call_id: "call_abc123".to_string(),
|
|
};
|
|
|
|
assert_eq!(selection.tool_call_id, "call_abc123");
|
|
assert_ne!(
|
|
selection.tool_call_id, "tool_call_id",
|
|
"tool_call_id must not be the hardcoded placeholder string"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_completion_positive_signals() {
|
|
assert!(llm_signals_completion("The job is complete."));
|
|
assert!(llm_signals_completion(
|
|
"I have completed the task successfully."
|
|
));
|
|
assert!(llm_signals_completion("The task is done."));
|
|
assert!(llm_signals_completion("The task is finished."));
|
|
assert!(llm_signals_completion(
|
|
"All steps are complete and verified."
|
|
));
|
|
assert!(llm_signals_completion(
|
|
"I've done all the work. The work is done."
|
|
));
|
|
assert!(llm_signals_completion(
|
|
"Successfully completed the migration."
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn test_completion_negative_signals_block_false_positives() {
|
|
// These contain completion keywords but also negation, should NOT trigger.
|
|
assert!(!llm_signals_completion("The task is not complete yet."));
|
|
assert!(!llm_signals_completion("This is not done."));
|
|
assert!(!llm_signals_completion("The work is incomplete."));
|
|
assert!(!llm_signals_completion(
|
|
"The migration is not yet finished."
|
|
));
|
|
assert!(!llm_signals_completion("The job isn't done yet."));
|
|
assert!(!llm_signals_completion("This remains unfinished."));
|
|
}
|
|
|
|
#[test]
|
|
fn test_completion_does_not_match_bare_substrings() {
|
|
// Bare words embedded in other text should NOT trigger completion.
|
|
assert!(!llm_signals_completion(
|
|
"I need to complete more work first."
|
|
));
|
|
assert!(!llm_signals_completion(
|
|
"Let me finish the remaining steps."
|
|
));
|
|
assert!(!llm_signals_completion(
|
|
"I'm done analyzing, now let me fix it."
|
|
));
|
|
assert!(!llm_signals_completion(
|
|
"I completed step 1 but step 2 remains."
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn test_completion_tool_output_injection() {
|
|
// A malicious tool output echoed by the LLM should not trigger
|
|
// completion unless it forms a genuine completion phrase.
|
|
assert!(!llm_signals_completion("TASK_COMPLETE"));
|
|
assert!(!llm_signals_completion("JOB_DONE"));
|
|
assert!(!llm_signals_completion(
|
|
"The tool returned: TASK_COMPLETE signal"
|
|
));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_parallel_speedup() {
|
|
// 3 tools each sleeping 200ms should finish in roughly 200ms (parallel),
|
|
// not ~600ms (sequential).
|
|
let tools: Vec<Arc<dyn Tool>> = (0..3)
|
|
.map(|i| {
|
|
Arc::new(SlowTool {
|
|
tool_name: format!("slow_{}", i),
|
|
delay: Duration::from_millis(200),
|
|
}) as Arc<dyn Tool>
|
|
})
|
|
.collect();
|
|
|
|
let worker = make_worker(tools).await;
|
|
|
|
let selections: Vec<ToolSelection> = (0..3)
|
|
.map(|i| ToolSelection {
|
|
tool_name: format!("slow_{}", i),
|
|
parameters: serde_json::json!({}),
|
|
reasoning: String::new(),
|
|
alternatives: vec![],
|
|
tool_call_id: format!("call_{}", i),
|
|
})
|
|
.collect();
|
|
|
|
let start = std::time::Instant::now();
|
|
let results = worker.execute_tools_parallel(&selections).await;
|
|
let elapsed = start.elapsed();
|
|
|
|
assert_eq!(results.len(), 3);
|
|
for r in &results {
|
|
assert!(r.result.is_ok(), "Tool should succeed");
|
|
}
|
|
// Parallel should complete well under the sequential 600ms threshold.
|
|
// Use a generous bound (800ms) to avoid flaky failures on slow CI runners,
|
|
// while still proving parallelism (sequential would be >= 600ms on any machine).
|
|
assert!(
|
|
elapsed < Duration::from_millis(800),
|
|
"Parallel execution took {:?}, expected < 800ms (sequential would be ~600ms)",
|
|
elapsed
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_result_ordering_preserved() {
|
|
// Tools with different delays finish in different order.
|
|
// Results must be returned in the original request order.
|
|
let tools: Vec<Arc<dyn Tool>> = vec![
|
|
Arc::new(SlowTool {
|
|
tool_name: "tool_a".into(),
|
|
delay: Duration::from_millis(300),
|
|
}),
|
|
Arc::new(SlowTool {
|
|
tool_name: "tool_b".into(),
|
|
delay: Duration::from_millis(100),
|
|
}),
|
|
Arc::new(SlowTool {
|
|
tool_name: "tool_c".into(),
|
|
delay: Duration::from_millis(200),
|
|
}),
|
|
];
|
|
|
|
let worker = make_worker(tools).await;
|
|
|
|
let selections = vec![
|
|
ToolSelection {
|
|
tool_name: "tool_a".into(),
|
|
parameters: serde_json::json!({}),
|
|
reasoning: String::new(),
|
|
alternatives: vec![],
|
|
tool_call_id: "call_a".into(),
|
|
},
|
|
ToolSelection {
|
|
tool_name: "tool_b".into(),
|
|
parameters: serde_json::json!({}),
|
|
reasoning: String::new(),
|
|
alternatives: vec![],
|
|
tool_call_id: "call_b".into(),
|
|
},
|
|
ToolSelection {
|
|
tool_name: "tool_c".into(),
|
|
parameters: serde_json::json!({}),
|
|
reasoning: String::new(),
|
|
alternatives: vec![],
|
|
tool_call_id: "call_c".into(),
|
|
},
|
|
];
|
|
|
|
let results = worker.execute_tools_parallel(&selections).await;
|
|
|
|
// Results must be in same order as selections, not completion order.
|
|
assert!(results[0].result.as_ref().unwrap().contains("done_tool_a"));
|
|
assert!(results[1].result.as_ref().unwrap().contains("done_tool_b"));
|
|
assert!(results[2].result.as_ref().unwrap().contains("done_tool_c"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_missing_tool_produces_error_not_panic() {
|
|
// If a tool doesn't exist, the result slot should contain an error.
|
|
let worker = make_worker(vec![]).await;
|
|
|
|
let selections = vec![ToolSelection {
|
|
tool_name: "nonexistent_tool".into(),
|
|
parameters: serde_json::json!({}),
|
|
reasoning: String::new(),
|
|
alternatives: vec![],
|
|
tool_call_id: "call_x".into(),
|
|
}];
|
|
|
|
let results = worker.execute_tools_parallel(&selections).await;
|
|
assert_eq!(results.len(), 1);
|
|
assert!(
|
|
results[0].result.is_err(),
|
|
"Missing tool should produce an error, not a panic"
|
|
);
|
|
}
|
|
|
|
/// Verify that calling mark_completed on an already-Completed job returns
|
|
/// an error (Completed → Completed is an invalid state transition).
|
|
#[tokio::test]
|
|
async fn test_mark_completed_twice_returns_error() {
|
|
let worker = make_worker(vec![]).await;
|
|
|
|
// Transition to InProgress first (required by state machine)
|
|
worker
|
|
.context_manager()
|
|
.update_context(worker.job_id, |ctx| {
|
|
ctx.transition_to(JobState::InProgress, None)
|
|
})
|
|
.await
|
|
.unwrap()
|
|
.unwrap();
|
|
|
|
// First mark_completed should succeed
|
|
worker.mark_completed().await.unwrap();
|
|
|
|
// Verify state is Completed
|
|
let ctx = worker
|
|
.context_manager()
|
|
.get_context(worker.job_id)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(ctx.state, JobState::Completed);
|
|
|
|
// Second mark_completed should fail (Completed → Completed is invalid)
|
|
let result = worker.mark_completed().await;
|
|
assert!(
|
|
result.is_err(),
|
|
"Completed → Completed transition should be rejected by state machine"
|
|
);
|
|
}
|
|
|
|
/// Build a Worker with the given approval context.
|
|
async fn make_worker_with_approval(
|
|
tools: Vec<Arc<dyn Tool>>,
|
|
approval_context: Option<crate::tools::ApprovalContext>,
|
|
) -> Worker {
|
|
let registry = ToolRegistry::new();
|
|
for t in tools {
|
|
registry.register(t).await;
|
|
}
|
|
|
|
let cm = Arc::new(crate::context::ContextManager::new(5));
|
|
let job_id = cm.create_job("test", "test job").await.unwrap();
|
|
|
|
let deps = WorkerDeps {
|
|
context_manager: cm,
|
|
llm: Arc::new(StubLlm),
|
|
safety: Arc::new(SafetyLayer::new(&SafetyConfig {
|
|
max_output_length: 100_000,
|
|
injection_check_enabled: false,
|
|
})),
|
|
tools: Arc::new(registry),
|
|
store: None,
|
|
hooks: Arc::new(crate::hooks::HookRegistry::new()),
|
|
timeout: Duration::from_secs(30),
|
|
use_planning: false,
|
|
sse_tx: None,
|
|
approval_context,
|
|
http_interceptor: None,
|
|
};
|
|
|
|
Worker::new(job_id, deps)
|
|
}
|
|
|
|
/// A tool that requires approval (UnlessAutoApproved).
|
|
struct ApprovalTool;
|
|
|
|
#[async_trait::async_trait]
|
|
impl Tool for ApprovalTool {
|
|
fn name(&self) -> &str {
|
|
"needs_approval"
|
|
}
|
|
fn description(&self) -> &str {
|
|
"Tool requiring approval"
|
|
}
|
|
fn parameters_schema(&self) -> serde_json::Value {
|
|
serde_json::json!({"type": "object", "properties": {}})
|
|
}
|
|
async fn execute(
|
|
&self,
|
|
_params: serde_json::Value,
|
|
_ctx: &crate::context::JobContext,
|
|
) -> Result<ToolOutput, crate::tools::ToolError> {
|
|
Ok(ToolOutput::text(
|
|
"approved",
|
|
std::time::Instant::now().elapsed(),
|
|
))
|
|
}
|
|
fn requires_approval(
|
|
&self,
|
|
_params: &serde_json::Value,
|
|
) -> crate::tools::ApprovalRequirement {
|
|
crate::tools::ApprovalRequirement::UnlessAutoApproved
|
|
}
|
|
fn requires_sanitization(&self) -> bool {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// A tool that always requires approval.
|
|
struct AlwaysApprovalTool;
|
|
|
|
#[async_trait::async_trait]
|
|
impl Tool for AlwaysApprovalTool {
|
|
fn name(&self) -> &str {
|
|
"always_approval"
|
|
}
|
|
fn description(&self) -> &str {
|
|
"Tool always requiring approval"
|
|
}
|
|
fn parameters_schema(&self) -> serde_json::Value {
|
|
serde_json::json!({"type": "object", "properties": {}})
|
|
}
|
|
async fn execute(
|
|
&self,
|
|
_params: serde_json::Value,
|
|
_ctx: &crate::context::JobContext,
|
|
) -> Result<ToolOutput, crate::tools::ToolError> {
|
|
Ok(ToolOutput::text(
|
|
"always",
|
|
std::time::Instant::now().elapsed(),
|
|
))
|
|
}
|
|
fn requires_approval(
|
|
&self,
|
|
_params: &serde_json::Value,
|
|
) -> crate::tools::ApprovalRequirement {
|
|
crate::tools::ApprovalRequirement::Always
|
|
}
|
|
fn requires_sanitization(&self) -> bool {
|
|
false
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_approval_context_unblocks_unless_auto_approved() {
|
|
// Without approval context, UnlessAutoApproved is blocked
|
|
let worker_blocked = make_worker_with_approval(vec![Arc::new(ApprovalTool)], None).await;
|
|
let result = worker_blocked
|
|
.execute_tool("needs_approval", &serde_json::json!({}))
|
|
.await;
|
|
assert!(
|
|
result.is_err(),
|
|
"Should be blocked without approval context"
|
|
);
|
|
|
|
// With autonomous approval context, UnlessAutoApproved is allowed
|
|
let worker_allowed = make_worker_with_approval(
|
|
vec![Arc::new(ApprovalTool)],
|
|
Some(crate::tools::ApprovalContext::autonomous()),
|
|
)
|
|
.await;
|
|
let result = worker_allowed
|
|
.execute_tool("needs_approval", &serde_json::json!({}))
|
|
.await;
|
|
assert!(result.is_ok(), "Should be allowed with autonomous context");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_approval_context_blocks_always_unless_permitted() {
|
|
// Autonomous context without tool_permissions blocks Always tools
|
|
let worker_blocked = make_worker_with_approval(
|
|
vec![Arc::new(AlwaysApprovalTool)],
|
|
Some(crate::tools::ApprovalContext::autonomous()),
|
|
)
|
|
.await;
|
|
let result = worker_blocked
|
|
.execute_tool("always_approval", &serde_json::json!({}))
|
|
.await;
|
|
assert!(
|
|
result.is_err(),
|
|
"Always tool should be blocked without permission"
|
|
);
|
|
|
|
// Autonomous context with tool_permissions allows Always tools
|
|
let worker_allowed = make_worker_with_approval(
|
|
vec![Arc::new(AlwaysApprovalTool)],
|
|
Some(crate::tools::ApprovalContext::autonomous_with_tools([
|
|
"always_approval".to_string(),
|
|
])),
|
|
)
|
|
.await;
|
|
let result = worker_allowed
|
|
.execute_tool("always_approval", &serde_json::json!({}))
|
|
.await;
|
|
assert!(
|
|
result.is_ok(),
|
|
"Always tool should be allowed with permission"
|
|
);
|
|
}
|
|
}
|