diff --git a/.claude/commands/add-sse-event.md b/.claude/commands/add-sse-event.md index 7215a48e..23f47a08 100644 --- a/.claude/commands/add-sse-event.md +++ b/.claude/commands/add-sse-event.md @@ -64,7 +64,7 @@ If the event needs custom UI (cards, badges, etc.), add styles. Follow the exist Identify where in the backend this event should be triggered. Common locations: - `src/agent/agent_loop.rs` - During message processing or tool execution -- `src/agent/worker.rs` - During job execution +- `src/worker/job.rs` - During job execution - `src/agent/heartbeat.rs` - During periodic execution Use the existing pattern: diff --git a/CLAUDE.md b/CLAUDE.md index 1b454e21..f7c0b403 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -99,7 +99,8 @@ src/ │ └── job_manager.rs # Container lifecycle (create, stop, cleanup) │ ├── worker/ # Runs inside Docker containers -│ ├── runtime.rs # Worker execution loop (tool calls, LLM) +│ ├── container.rs # Container worker runtime (ContainerDelegate + shared agentic loop) +│ ├── job.rs # Background job worker (JobDelegate + shared agentic loop) │ ├── claude_bridge.rs # Claude Code bridge (spawns claude CLI) │ └── proxy_llm.rs # LlmProvider that proxies through orchestrator │ diff --git a/COVERAGE_PLAN.md b/COVERAGE_PLAN.md index c9d7d73b..af5f872c 100644 --- a/COVERAGE_PLAN.md +++ b/COVERAGE_PLAN.md @@ -63,12 +63,12 @@ These files account for the vast majority of the coverage gap: | `src/main.rs` | 740 | 522 | 29.4% | 485 | | `src/channels/web/handlers/jobs.rs` | 513 | 456 | 11.1% | 430 | | `src/tools/builder/core.rs` | 524 | 456 | 13.0% | 429 | -| `src/agent/worker.rs` | 1,078 | 467 | 56.7% | 413 | +| `src/worker/job.rs` | 1,078 | 467 | 56.7% | 413 | | `src/channels/web/handlers/chat.rs` | 564 | 417 | 26.1% | 388 | | `src/tools/wasm/wrapper.rs` | 1,005 | 436 | 56.6% | 385 | | `src/channels/signal.rs` | 1,814 | 472 | 74.0% | 381 | | `src/tools/mcp/auth.rs` | 472 | 378 | 19.9% | 354 | -| `src/worker/runtime.rs` | 350 | 330 | 5.7% | 312 | +| `src/worker/container.rs` | 350 | 330 | 5.7% | 312 | | `src/tools/builtin/job.rs` | 1,014 | 359 | 64.6% | 308 | | `src/cli/mcp.rs` | 322 | 319 | 0.9% | 302 | | `src/cli/oauth_defaults.rs` | 730 | 335 | 54.1% | 298 | @@ -346,7 +346,7 @@ Test slash commands through the agent loop. ### Trace: Worker Multi-Turn Execution -**Covers:** `agent/worker.rs` (+413 lines), `agent/agent_loop.rs` (+207 lines) +**Covers:** `worker/job.rs` (+413 lines), `agent/agent_loop.rs` (+207 lines) Test multi-turn tool calling, error recovery, and completion flows. @@ -769,7 +769,7 @@ HTTP proxy for container network access. - `test_proxy_connect_tunnel` -- HTTPS CONNECT method handling - `test_proxy_logging` -- request/response logging -### `src/worker/runtime.rs` -- 5.7% -> 95% (+312 lines) +### `src/worker/container.rs` -- 5.7% -> 95% (+312 lines) Worker execution loop (runs inside containers). diff --git a/FEATURE_PARITY.md b/FEATURE_PARITY.md index 61e32b0c..634131fc 100644 --- a/FEATURE_PARITY.md +++ b/FEATURE_PARITY.md @@ -46,7 +46,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O | Bonjour/mDNS discovery | ✅ | ❌ | | | Tailscale integration | ✅ | ❌ | | | Health check endpoints | ✅ | ✅ | /api/health + /api/gateway/status + /healthz + /readyz, with channel-backed readiness probes | -| `doctor` diagnostics | ✅ | ❌ | | +| `doctor` diagnostics | ✅ | 🚧 | 16 checks: settings, LLM, DB, embeddings, routines, gateway, MCP, skills, secrets, service, Docker daemon, tunnel binaries | | Agent event broadcast | ✅ | 🚧 | SSE broadcast manager exists (SseManager) but tool/job-state events not fully wired | | Channel health monitor | ✅ | ❌ | Auto-restart with configurable interval | | Presence system | ✅ | ❌ | Beacons on connect, system presence for agents | @@ -175,7 +175,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O | `message send` | ✅ | ❌ | P2 | Send to channels | | `browser` | ✅ | ❌ | P3 | Browser automation | | `sandbox` | ✅ | ✅ | - | WASM sandbox | -| `doctor` | ✅ | ❌ | P2 | Diagnostics | +| `doctor` | ✅ | 🚧 | P2 | 16 subsystem checks | | `logs` | ✅ | ❌ | P3 | Query logs | | `update` | ✅ | ❌ | P3 | Self-update | | `completion` | ✅ | ✅ | - | Shell completion | diff --git a/src/agent/CLAUDE.md b/src/agent/CLAUDE.md index 386492d0..e55c9591 100644 --- a/src/agent/CLAUDE.md +++ b/src/agent/CLAUDE.md @@ -14,7 +14,8 @@ Core agent logic. This is the most complex subsystem — read this before workin | `session_manager.rs` | Lifecycle: create/lookup sessions, map external thread IDs to internal UUIDs, prune stale sessions, manage undo managers. | | `router.rs` | Routes explicit `/commands` to `MessageIntent`. Natural language bypasses the router entirely. | | `scheduler.rs` | Parallel job scheduling. Maintains `jobs` map (full LLM-driven) and `subtasks` map (tool-exec/background). | -| `worker.rs` | Per-job execution for background scheduler jobs: calls LLM, runs tools, handles the reasoning loop. Distinct from `dispatcher.rs`. | +| *(moved to `src/worker/job.rs`)* | Per-job execution now lives in `src/worker/job.rs` as `JobDelegate`, using the shared `run_agentic_loop()` engine. | +| `agentic_loop.rs` | Shared agentic loop engine: `run_agentic_loop()`, `LoopDelegate` trait, `LoopOutcome`, `LoopSignal`, `TextAction`. All three execution paths (chat, job, container) delegate to this. | | `compaction.rs` | Context window management: summarize old turns, write to workspace daily log, trim context. Three strategies. | | `context_monitor.rs` | Detects memory pressure. Suggests `CompactionStrategy` based on usage level. | | `self_repair.rs` | Detects stuck jobs and broken tools, attempts recovery. | @@ -49,26 +50,28 @@ Session (per user) ## Agentic Loop (dispatcher.rs) -The `dispatcher.rs` module handles **direct conversational turns** (user messages processed inline by the main agent). Background scheduler jobs use `worker.rs` instead — these are two separate execution paths. +All three execution paths (chat, job, container) now use the shared `run_agentic_loop()` engine in `agentic_loop.rs`, each providing their own `LoopDelegate` implementation: + +- **`ChatDelegate`** (`dispatcher.rs`) — conversational turns, tool approval, skill context injection +- **`JobDelegate`** (`src/worker/job.rs`) — background scheduler jobs, planning support, completion detection +- **`ContainerDelegate`** (`src/worker/container.rs`) — Docker container worker, sequential tool exec, HTTP event streaming ``` -run_agentic_loop() [dispatcher.rs — conversational turns] - 1. Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.) - 2. Detect group chat from metadata; exclude MEMORY.md if group chat - 3. Select active skills (keyword/pattern scoring against message content) - 4. Build skill context block (injected before user message) - 5. LLM call → text response OR tool calls - 6. If tool calls: - a. Check tool approval (session auto-approvals, pending approval queue) - b. Execute tools (parallel via JoinSet) - c. Sanitize results through SafetyLayer - d. Feed results back → goto 5 - 7. Return AgenticLoopResult::Response or NeedApproval +run_agentic_loop(delegate, reasoning, reason_ctx, config) + 1. Check signals (stop/cancel) via delegate.check_signals() + 2. Pre-LLM hook via delegate.before_llm_call() + 3. LLM call via delegate.call_llm() + 4. If text response → delegate.handle_text_response() → Continue or Return + 5. If tool calls → delegate.execute_tool_calls() → Continue or Return + 6. Post-iteration hook via delegate.after_iteration() + 7. Repeat until LoopOutcome returned or max_iterations reached ``` -**Tool approval:** Tools flagged `requires_approval` pause the loop and return `NeedApproval`. The web gateway stores the `PendingApproval` in session state and sends an `approval_needed` SSE event. The user's approval/deny resumes the loop. +**Tool approval:** Tools flagged `requires_approval` pause the loop — `ChatDelegate` returns `LoopOutcome::NeedApproval(pending)`. The web gateway stores the `PendingApproval` in session state and sends an `approval_needed` SSE event. The user's approval/deny resumes the loop. -**worker.rs vs dispatcher.rs:** `dispatcher.rs` runs the agentic loop for user-initiated conversational turns (holds session lock, tracks turns). `worker.rs` is spawned by the `Scheduler` for background jobs created via `CreateJob` / `/job` — it runs independently of the session and has its own LLM reasoning loop with planning support (`use_planning` flag). +**Shared tool execution:** `tools/execute.rs` provides `execute_tool_with_safety()` (validate → timeout → execute → serialize) and `process_tool_result()` (sanitize → wrap → ChatMessage), used by all three delegates. + +**ChatDelegate vs JobDelegate:** `ChatDelegate` runs for user-initiated conversational turns (holds session lock, tracks turns). `JobDelegate` is spawned by the `Scheduler` for background jobs created via `CreateJob` / `/job` — it runs independently of the session and has planning support (`use_planning` flag). ## Command Routing (router.rs) diff --git a/src/agent/agentic_loop.rs b/src/agent/agentic_loop.rs new file mode 100644 index 00000000..0e5bef9d --- /dev/null +++ b/src/agent/agentic_loop.rs @@ -0,0 +1,587 @@ +//! Unified agentic loop engine. +//! +//! Provides a single implementation of the core LLM call → tool execution → +//! result processing → context update → repeat cycle. Three consumers +//! (chat dispatcher, job worker, container runtime) customize behavior +//! via the `LoopDelegate` trait. + +use async_trait::async_trait; + +use crate::agent::session::PendingApproval; +use crate::error::Error; +use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult}; + +/// Signal from the delegate indicating how the loop should proceed. +pub enum LoopSignal { + /// Continue normally. + Continue, + /// Stop the loop gracefully. + Stop, + /// Inject a user message into context and continue. + InjectMessage(String), +} + +/// Outcome of a text response from the LLM. +pub enum TextAction { + /// Return this as the final loop result. + Return(LoopOutcome), + /// Continue the loop (text was handled but loop should proceed). + Continue, +} + +/// Final outcome of the agentic loop. +pub enum LoopOutcome { + /// Completed with a text response. + Response(String), + /// Loop was stopped by a signal. + Stopped, + /// Max iterations exceeded. + MaxIterations, + /// A tool requires user approval before continuing (chat delegate only). + NeedApproval(Box), +} + +/// Configuration for the agentic loop. +pub struct AgenticLoopConfig { + pub max_iterations: usize, + pub enable_tool_intent_nudge: bool, + pub max_tool_intent_nudges: u32, +} + +impl Default for AgenticLoopConfig { + fn default() -> Self { + Self { + max_iterations: 50, + enable_tool_intent_nudge: true, + max_tool_intent_nudges: 2, + } + } +} + +/// Strategy trait — each consumer implements this to customize I/O and lifecycle. +/// +/// The shared loop calls these methods at well-defined points. Consumers +/// implement only the behavior that differs between chat, job, and container +/// contexts. The loop itself handles the common logic: tool intent nudge, +/// iteration counting, tool definition refresh, and the respond → execute → process cycle. +/// +/// # `Send + Sync` requirement +/// +/// This trait requires `Send + Sync` because the loop accepts `&dyn LoopDelegate`. +/// Delegates using borrowed references (e.g. `ChatDelegate<'a>`) must ensure all +/// borrowed fields are `Send + Sync`. This is a load-bearing constraint: if a +/// delegate needs to be spawned into a detached task, it must use `Arc`-based +/// ownership instead of borrows (as `JobDelegate` and `ContainerDelegate` do). +#[async_trait] +pub trait LoopDelegate: Send + Sync { + /// Called at the start of each iteration. Check for external signals + /// (cancellation, user messages, stop requests). + async fn check_signals(&self) -> LoopSignal; + + /// Called before the LLM call. Allows the delegate to refresh tool + /// definitions, enforce cost guards, or inject messages. + /// Return `Some(outcome)` to break the loop early. + async fn before_llm_call( + &self, + reason_ctx: &mut ReasoningContext, + iteration: usize, + ) -> Option; + + /// Call the LLM and return the result. Delegates own the LLM call + /// to handle consumer-specific concerns (rate limiting, auto-compaction, + /// cost tracking, force_text mode). + async fn call_llm( + &self, + reasoning: &Reasoning, + reason_ctx: &mut ReasoningContext, + iteration: usize, + ) -> Result; + + /// Handle a text-only response from the LLM. + /// Return `TextAction::Return` to exit the loop, `TextAction::Continue` to proceed. + async fn handle_text_response( + &self, + text: &str, + reason_ctx: &mut ReasoningContext, + ) -> TextAction; + + /// Execute tool calls and add results to context. + /// Return `Some(outcome)` to break the loop (e.g. approval needed). + async fn execute_tool_calls( + &self, + tool_calls: Vec, + content: Option, + reason_ctx: &mut ReasoningContext, + ) -> Result, Error>; + + /// Called when the LLM expresses tool intent without actually calling a tool. + /// Delegates can use this to emit events or log the nudge for observability. + async fn on_tool_intent_nudge(&self, _text: &str, _reason_ctx: &mut ReasoningContext) {} + + /// Called after each successful iteration (no error, no early return). + async fn after_iteration(&self, _iteration: usize) {} +} + +/// Run the unified agentic loop. +/// +/// This is the single implementation used by all three consumers (chat, job, container). +/// The `delegate` provides consumer-specific behavior via the `LoopDelegate` trait. +pub async fn run_agentic_loop( + delegate: &dyn LoopDelegate, + reasoning: &Reasoning, + reason_ctx: &mut ReasoningContext, + config: &AgenticLoopConfig, +) -> Result { + let mut consecutive_tool_intent_nudges: u32 = 0; + + for iteration in 1..=config.max_iterations { + // Check for external signals (stop, cancellation, user messages) + match delegate.check_signals().await { + LoopSignal::Continue => {} + LoopSignal::Stop => return Ok(LoopOutcome::Stopped), + LoopSignal::InjectMessage(msg) => { + reason_ctx.messages.push(ChatMessage::user(&msg)); + } + } + + // Pre-LLM call hook (cost guard, tool refresh, iteration limit nudge) + if let Some(outcome) = delegate.before_llm_call(reason_ctx, iteration).await { + return Ok(outcome); + } + + // Call LLM + let output = delegate.call_llm(reasoning, reason_ctx, iteration).await?; + + match output.result { + RespondResult::Text(text) => { + // Tool intent nudge: if the LLM says "let me search..." without + // actually calling a tool, inject a nudge message. + if config.enable_tool_intent_nudge + && !reason_ctx.available_tools.is_empty() + && !reason_ctx.force_text + && consecutive_tool_intent_nudges < config.max_tool_intent_nudges + && crate::llm::llm_signals_tool_intent(&text) + { + consecutive_tool_intent_nudges += 1; + tracing::info!( + iteration, + "LLM expressed tool intent without calling a tool, nudging" + ); + delegate.on_tool_intent_nudge(&text, reason_ctx).await; + reason_ctx.messages.push(ChatMessage::assistant(&text)); + reason_ctx + .messages + .push(ChatMessage::user(crate::llm::TOOL_INTENT_NUDGE)); + delegate.after_iteration(iteration).await; + continue; + } + + // Reset nudge counter since we got a non-intent text response + if !crate::llm::llm_signals_tool_intent(&text) { + consecutive_tool_intent_nudges = 0; + } + + match delegate.handle_text_response(&text, reason_ctx).await { + TextAction::Return(outcome) => return Ok(outcome), + TextAction::Continue => {} + } + } + RespondResult::ToolCalls { + tool_calls, + content, + } => { + consecutive_tool_intent_nudges = 0; + + if let Some(outcome) = delegate + .execute_tool_calls(tool_calls, content, reason_ctx) + .await? + { + return Ok(outcome); + } + } + } + + delegate.after_iteration(iteration).await; + } + + Ok(LoopOutcome::MaxIterations) +} + +/// Truncate a string for log/status previews. +/// +/// `max` is a byte budget. The result is truncated at the last valid char +/// boundary at or before `max` bytes, so it is always valid UTF-8. +pub fn truncate_for_preview(s: &str, max: usize) -> String { + if s.len() <= max { + s.to_string() + } else { + let end = crate::util::floor_char_boundary(s, max); + format!("{}...", &s[..end]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::llm::{RespondOutput, TokenUsage, ToolCall}; + use crate::testing::StubLlm; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::sync::Mutex; + + fn stub_reasoning() -> Reasoning { + Reasoning::new(Arc::new(StubLlm::default())) + } + + fn zero_usage() -> TokenUsage { + TokenUsage { + input_tokens: 0, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + } + } + + fn text_output(text: &str) -> RespondOutput { + RespondOutput { + result: RespondResult::Text(text.to_string()), + usage: zero_usage(), + } + } + + fn tool_calls_output(calls: Vec) -> RespondOutput { + RespondOutput { + result: RespondResult::ToolCalls { + tool_calls: calls, + content: None, + }, + usage: zero_usage(), + } + } + + /// Configurable mock delegate for testing run_agentic_loop. + struct MockDelegate { + signal: Mutex, + llm_responses: Mutex>, + tool_exec_count: AtomicUsize, + tool_exec_outcome: Mutex>, + iterations_seen: Mutex>, + early_exit: Mutex>, + nudge_count: AtomicUsize, + } + + impl MockDelegate { + fn new(responses: Vec) -> Self { + Self { + signal: Mutex::new(LoopSignal::Continue), + llm_responses: Mutex::new(responses), + tool_exec_count: AtomicUsize::new(0), + tool_exec_outcome: Mutex::new(None), + iterations_seen: Mutex::new(Vec::new()), + early_exit: Mutex::new(None), + nudge_count: AtomicUsize::new(0), + } + } + + fn with_signal(mut self, signal: LoopSignal) -> Self { + self.signal = Mutex::new(signal); + self + } + + fn with_early_exit(mut self, iteration: usize, outcome: LoopOutcome) -> Self { + self.early_exit = Mutex::new(Some((iteration, outcome))); + self + } + } + + #[async_trait] + impl LoopDelegate for MockDelegate { + async fn check_signals(&self) -> LoopSignal { + let mut sig = self.signal.lock().await; + std::mem::replace(&mut *sig, LoopSignal::Continue) + } + + async fn before_llm_call( + &self, + _reason_ctx: &mut ReasoningContext, + iteration: usize, + ) -> Option { + let mut guard = self.early_exit.lock().await; + let should_take = guard + .as_ref() + .is_some_and(|(target, _)| *target == iteration); + if should_take { + guard.take().map(|(_, o)| o) + } else { + None + } + } + + async fn call_llm( + &self, + _reasoning: &Reasoning, + _reason_ctx: &mut ReasoningContext, + _iteration: usize, + ) -> Result { + let mut responses = self.llm_responses.lock().await; + if responses.is_empty() { + panic!("MockDelegate: no more LLM responses queued"); + } + Ok(responses.remove(0)) + } + + async fn handle_text_response( + &self, + text: &str, + _reason_ctx: &mut ReasoningContext, + ) -> TextAction { + TextAction::Return(LoopOutcome::Response(text.to_string())) + } + + async fn execute_tool_calls( + &self, + _tool_calls: Vec, + _content: Option, + reason_ctx: &mut ReasoningContext, + ) -> Result, crate::error::Error> { + self.tool_exec_count.fetch_add(1, Ordering::SeqCst); + reason_ctx + .messages + .push(ChatMessage::user("tool result stub")); + let outcome = self.tool_exec_outcome.lock().await.take(); + Ok(outcome) + } + + async fn on_tool_intent_nudge(&self, _text: &str, _reason_ctx: &mut ReasoningContext) { + self.nudge_count.fetch_add(1, Ordering::SeqCst); + } + + async fn after_iteration(&self, iteration: usize) { + self.iterations_seen.lock().await.push(iteration); + } + } + + // --- Tests --- + + #[tokio::test] + async fn test_text_response_returns_immediately() { + let delegate = MockDelegate::new(vec![text_output("Hello, world!")]); + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + let config = AgenticLoopConfig::default(); + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + match outcome { + LoopOutcome::Response(text) => assert_eq!(text, "Hello, world!"), + _ => panic!("Expected LoopOutcome::Response"), + } + // after_iteration is NOT called when handle_text_response returns Return + // (the loop exits before reaching after_iteration). + assert!(delegate.iterations_seen.lock().await.is_empty()); + } + + #[tokio::test] + async fn test_tool_call_then_text_response() { + let tool_call = ToolCall { + id: "call_1".to_string(), + name: "echo".to_string(), + arguments: serde_json::json!({}), + }; + let delegate = MockDelegate::new(vec![ + tool_calls_output(vec![tool_call]), + text_output("Done!"), + ]); + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + let config = AgenticLoopConfig::default(); + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + match outcome { + LoopOutcome::Response(text) => assert_eq!(text, "Done!"), + _ => panic!("Expected LoopOutcome::Response"), + } + assert_eq!(delegate.tool_exec_count.load(Ordering::SeqCst), 1); + // after_iteration called for iteration 1 (tool call), but not 2 + // (text response exits before after_iteration). + assert_eq!(*delegate.iterations_seen.lock().await, vec![1]); + } + + #[tokio::test] + async fn test_stop_signal_exits_immediately() { + let delegate = + MockDelegate::new(vec![text_output("unreachable")]).with_signal(LoopSignal::Stop); + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + let config = AgenticLoopConfig::default(); + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + assert!(matches!(outcome, LoopOutcome::Stopped)); + assert!(delegate.iterations_seen.lock().await.is_empty()); + } + + #[tokio::test] + async fn test_inject_message_adds_user_message() { + let delegate = MockDelegate::new(vec![text_output("Got it")]) + .with_signal(LoopSignal::InjectMessage("injected prompt".to_string())); + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + let config = AgenticLoopConfig::default(); + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + assert!(matches!(outcome, LoopOutcome::Response(_))); + assert!( + ctx.messages + .iter() + .any(|m| m.role == crate::llm::Role::User && m.content.contains("injected prompt")), + "Injected message should appear in context" + ); + } + + #[tokio::test] + async fn test_max_iterations_reached() { + struct ContinueDelegate; + + #[async_trait] + impl LoopDelegate for ContinueDelegate { + async fn check_signals(&self) -> LoopSignal { + LoopSignal::Continue + } + async fn before_llm_call( + &self, + _: &mut ReasoningContext, + _: usize, + ) -> Option { + None + } + async fn call_llm( + &self, + _: &Reasoning, + _: &mut ReasoningContext, + _: usize, + ) -> Result { + Ok(text_output("still working")) + } + async fn handle_text_response( + &self, + _: &str, + ctx: &mut ReasoningContext, + ) -> TextAction { + ctx.messages.push(ChatMessage::assistant("still working")); + TextAction::Continue + } + async fn execute_tool_calls( + &self, + _: Vec, + _: Option, + _: &mut ReasoningContext, + ) -> Result, crate::error::Error> { + Ok(None) + } + } + + let delegate = ContinueDelegate; + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + let config = AgenticLoopConfig { + max_iterations: 3, + ..Default::default() + }; + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + assert!(matches!(outcome, LoopOutcome::MaxIterations)); + let assistant_count = ctx + .messages + .iter() + .filter(|m| m.role == crate::llm::Role::Assistant) + .count(); + assert_eq!(assistant_count, 3); + } + + #[tokio::test] + async fn test_tool_intent_nudge_fires_and_caps() { + let delegate = MockDelegate::new(vec![ + text_output("Let me search for that file"), + text_output("Let me search for that file"), + text_output("Let me search for that file"), + ]); + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + ctx.available_tools.push(crate::llm::ToolDefinition { + name: "search".to_string(), + description: "Search files".to_string(), + parameters: serde_json::json!({"type": "object"}), + }); + let config = AgenticLoopConfig { + max_iterations: 10, + enable_tool_intent_nudge: true, + max_tool_intent_nudges: 2, + }; + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + assert!(matches!(outcome, LoopOutcome::Response(_))); + assert_eq!(delegate.nudge_count.load(Ordering::SeqCst), 2); + let nudge_messages = ctx + .messages + .iter() + .filter(|m| { + m.role == crate::llm::Role::User + && m.content.contains("you did not include any tool calls") + }) + .count(); + assert_eq!( + nudge_messages, 2, + "Should have exactly 2 nudge messages in context" + ); + } + + #[tokio::test] + async fn test_before_llm_call_early_exit() { + let delegate = MockDelegate::new(vec![text_output("unreachable")]) + .with_early_exit(1, LoopOutcome::Stopped); + let reasoning = stub_reasoning(); + let mut ctx = ReasoningContext::new(); + let config = AgenticLoopConfig::default(); + + let outcome = run_agentic_loop(&delegate, &reasoning, &mut ctx, &config) + .await + .unwrap(); + + assert!(matches!(outcome, LoopOutcome::Stopped)); + assert!(delegate.iterations_seen.lock().await.is_empty()); + } + + #[test] + fn test_truncate_short_string_unchanged() { + assert_eq!(truncate_for_preview("hello", 10), "hello"); + } + + #[test] + fn test_truncate_long_string_adds_ellipsis() { + let result = truncate_for_preview("hello world", 5); + assert_eq!(result, "hello..."); + } + + #[test] + fn test_truncate_multibyte_safe() { + let result = truncate_for_preview("café", 4); + assert_eq!(result, "caf..."); + } +} diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index b1678d89..18121086 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -10,12 +10,16 @@ use tokio::task::JoinSet; use uuid::Uuid; use crate::agent::Agent; -use crate::agent::context_monitor::{ContextBreakdown, estimate_text_tokens}; use crate::agent::session::{PendingApproval, Session, ThreadState}; use crate::channels::{IncomingMessage, StatusUpdate}; use crate::context::JobContext; use crate::error::Error; -use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult}; +use async_trait::async_trait; + +use crate::agent::agentic_loop::{ + AgenticLoopConfig, LoopDelegate, LoopOutcome, LoopSignal, TextAction, +}; +use crate::llm::{ChatMessage, Reasoning, ReasoningContext}; use crate::tools::redact_params; /// Result of the agentic loop execution. @@ -134,9 +138,6 @@ impl Agent { reasoning = reasoning.with_skill_context(ctx); } - // Build context with messages that we'll mutate during the loop - let mut context_messages = initial_messages; - // Create a JobContext for tool execution (chat doesn't have a real job) let mut job_ctx = JobContext::with_user(&message.user_id, "chat", "Interactive chat session"); @@ -155,721 +156,62 @@ impl Agent { let cached_prompt_no_tools = reasoning.build_system_prompt_with_tools(&[]); let max_tool_iterations = self.config.max_tool_iterations; - // Force a text-only response on the last iteration to guarantee termination - // instead of hard-erroring. The penultimate iteration also gets a nudge - // message so the LLM knows it should wrap up. let force_text_at = max_tool_iterations; let nudge_at = max_tool_iterations.saturating_sub(1); - let mut iteration = 0; - const MAX_TOOL_INTENT_NUDGES: u32 = 2; - let mut consecutive_tool_intent_nudges: u32 = 0; - loop { - iteration += 1; - // Hard ceiling one past the forced-text iteration (should never be reached - // since force_text_at guarantees a text response, but kept as a safety net). - if iteration > max_tool_iterations + 1 { - return Err(crate::error::LlmError::InvalidResponse { - provider: "agent".to_string(), - reason: format!("Exceeded maximum tool iterations ({max_tool_iterations})"), - } - .into()); + + let delegate = ChatDelegate { + agent: self, + session: session.clone(), + thread_id, + message, + job_ctx, + active_skills, + cached_prompt, + cached_prompt_no_tools, + nudge_at, + force_text_at, + user_tz, + }; + + let mut reason_ctx = ReasoningContext::new() + .with_messages(initial_messages) + .with_tools(initial_tool_defs) + .with_system_prompt(delegate.cached_prompt.clone()) + .with_metadata({ + let mut m = std::collections::HashMap::new(); + m.insert("thread_id".to_string(), thread_id.to_string()); + m + }); + + let loop_config = AgenticLoopConfig { + // Hard ceiling: one past force_text_at (safety net). + max_iterations: max_tool_iterations + 1, + enable_tool_intent_nudge: true, + max_tool_intent_nudges: 2, + }; + + let outcome = crate::agent::agentic_loop::run_agentic_loop( + &delegate, + &reasoning, + &mut reason_ctx, + &loop_config, + ) + .await?; + + match outcome { + LoopOutcome::Response(text) => Ok(AgenticLoopResult::Response(text)), + LoopOutcome::Stopped => Err(crate::error::JobError::ContextError { + id: thread_id, + reason: "Interrupted".to_string(), } - - // Check if interrupted - { - let sess = session.lock().await; - if let Some(thread) = sess.threads.get(&thread_id) - && thread.state == ThreadState::Interrupted - { - return Err(crate::error::JobError::ContextError { - id: thread_id, - reason: "Interrupted".to_string(), - } - .into()); - } + .into()), + LoopOutcome::MaxIterations => Err(crate::error::LlmError::InvalidResponse { + provider: "agent".to_string(), + reason: format!("Exceeded maximum tool iterations ({max_tool_iterations})"), } - - // Enforce cost guardrails before the LLM call - if let Err(limit) = self.cost_guard().check_allowed().await { - return Err(crate::error::LlmError::InvalidResponse { - provider: "agent".to_string(), - reason: limit.to_string(), - } - .into()); - } - - // Inject a nudge message when approaching the iteration limit so the - // LLM is aware it should produce a final answer on the next turn. - if iteration == nudge_at { - context_messages.push(ChatMessage::system( - "You are approaching the tool call limit. \ - Provide your best final answer on the next response \ - using the information you have gathered so far. \ - Do not call any more tools.", - )); - } - - let force_text = iteration >= force_text_at; - - // Refresh tool definitions each iteration so newly built tools become visible - let tool_defs = self.tools().tool_definitions().await; - - // Apply trust-based tool attenuation if skills are active. - let tool_defs = if !active_skills.is_empty() { - let result = crate::skills::attenuate_tools(&tool_defs, &active_skills); - tracing::info!( - min_trust = %result.min_trust, - tools_available = result.tools.len(), - tools_removed = result.removed_tools.len(), - removed = ?result.removed_tools, - explanation = %result.explanation, - "Tool attenuation applied" - ); - result.tools - } else { - tool_defs - }; - - // Call LLM with current context; force_text drops tools to guarantee a - // text response on the final iteration. The pre-built system prompt - // avoids rebuilding the same ~1,500-token string each iteration. - let mut context = ReasoningContext::new() - .with_messages(context_messages.clone()) - .with_tools(tool_defs) - .with_system_prompt(if force_text { - cached_prompt_no_tools.clone() - } else { - cached_prompt.clone() - }) - .with_metadata({ - let mut m = std::collections::HashMap::new(); - m.insert("thread_id".to_string(), thread_id.to_string()); - m - }); - context.force_text = force_text; - - if force_text { - tracing::info!( - iteration, - "Forcing text-only response (iteration limit reached)" - ); - } - - // Pre-prompt context diagnostics: log token breakdown before LLM call - { - let breakdown = ContextBreakdown::analyze(&context_messages); - let system_prompt_tokens = - estimate_text_tokens(context.system_prompt.as_deref().unwrap_or("")); - let total_tokens = breakdown.total_tokens + system_prompt_tokens; - tracing::debug!( - iteration, - messages = breakdown.message_count, - total_tokens, - system_prompt_tokens, - system_msg_tokens = breakdown.system_tokens, - user_tokens = breakdown.user_tokens, - assistant_tokens = breakdown.assistant_tokens, - tool_tokens = breakdown.tool_tokens, - tools_available = context.available_tools.len(), - force_text, - "Pre-prompt context diagnostics" - ); - } - - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Thinking("Calling LLM...".into()), - &message.metadata, - ) - .await; - - let output = match reasoning.respond_with_tools(&context).await { - Ok(output) => output, - Err(crate::error::LlmError::ContextLengthExceeded { used, limit }) => { - tracing::warn!( - used, - limit, - iteration, - "Context length exceeded, compacting messages and retrying" - ); - - // Compact: keep system messages + last user message + current turn - context_messages = compact_messages_for_retry(&context_messages); - - // Rebuild context with compacted messages, reusing cached prompt - let mut retry_context = ReasoningContext::new() - .with_messages(context_messages.clone()) - .with_tools(if force_text { - Vec::new() - } else { - context.available_tools.clone() - }) - .with_metadata(context.metadata.clone()); - retry_context.force_text = force_text; - retry_context.system_prompt = context.system_prompt.clone(); - - reasoning - .respond_with_tools(&retry_context) - .await - .map_err(|retry_err| { - tracing::error!( - original_used = used, - original_limit = limit, - retry_error = %retry_err, - "Retry after auto-compaction also failed" - ); - // Propagate the actual retry error so callers see the real failure - crate::error::Error::from(retry_err) - })? - } - Err(e) => return Err(e.into()), - }; - - // Record cost and track token usage - let model_name = self.llm().active_model_name(); - let read_discount = self.llm().cache_read_discount(); - let write_multiplier = self.llm().cache_write_multiplier(); - let call_cost = self - .cost_guard() - .record_llm_call( - &model_name, - output.usage.input_tokens, - output.usage.output_tokens, - output.usage.cache_read_input_tokens, - output.usage.cache_creation_input_tokens, - read_discount, - write_multiplier, - Some(self.llm().cost_per_token()), - ) - .await; - tracing::debug!( - "LLM call used {} input + {} output tokens (${:.6})", - output.usage.input_tokens, - output.usage.output_tokens, - call_cost, - ); - - match output.result { - RespondResult::Text(text) => { - // Nudge the LLM if it expressed tool intent without calling tools. - // This is common with non-Anthropic models (e.g. GLM-5 via NEAR AI) - // that output "Let me search…" but don't issue tool_calls. - if !force_text - && !context.available_tools.is_empty() - && consecutive_tool_intent_nudges < MAX_TOOL_INTENT_NUDGES - && crate::llm::llm_signals_tool_intent(&text) - { - consecutive_tool_intent_nudges += 1; - tracing::info!( - iteration, - "LLM expressed tool intent without calling a tool, nudging" - ); - context_messages.push(ChatMessage::assistant(&text)); - context_messages.push(ChatMessage::user(crate::llm::TOOL_INTENT_NUDGE)); - continue; - } - - // Strip internal "[Called tool ...]" text that can leak when - // provider flattening (e.g. NEAR AI) converts tool_calls to - // plain text and the LLM echoes it back. - let sanitized = strip_internal_tool_call_text(&text); - return Ok(AgenticLoopResult::Response(sanitized)); - } - RespondResult::ToolCalls { - tool_calls, - content, - } => { - consecutive_tool_intent_nudges = 0; - // Add the assistant message with tool_calls to context. - // OpenAI protocol requires this before tool-result messages. - context_messages.push(ChatMessage::assistant_with_tool_calls( - content, - tool_calls.clone(), - )); - - // Execute tools and add results to context - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::Thinking(format!( - "Executing {} tool(s)...", - tool_calls.len() - )), - &message.metadata, - ) - .await; - - // Record tool calls in the thread with sensitive params redacted. - // Look up each tool's sensitive_params before acquiring the session lock. - { - let mut redacted_args: Vec = - Vec::with_capacity(tool_calls.len()); - for tc in &tool_calls { - let safe = if let Some(tool) = self.tools().get(&tc.name).await { - redact_params(&tc.arguments, tool.sensitive_params()) - } else { - tc.arguments.clone() - }; - redacted_args.push(safe); - } - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) - && let Some(turn) = thread.last_turn_mut() - { - for (tc, safe_args) in tool_calls.iter().zip(redacted_args) { - turn.record_tool_call(&tc.name, safe_args); - } - } - } - - // === Phase 1: Preflight (sequential) === - // Walk tool_calls checking approval and hooks. Classify - // each tool as Rejected (by hook) or Runnable. Stop at the - // first tool that needs approval. - // - // Outcomes are indexed by original tool_calls position so - // Phase 3 can emit results in the correct order. - enum PreflightOutcome { - /// Hook rejected/blocked this tool; contains the error message. - Rejected(String), - /// Tool passed preflight and will be executed. - Runnable, - } - let mut preflight: Vec<(crate::llm::ToolCall, PreflightOutcome)> = Vec::new(); - let mut runnable: Vec<(usize, crate::llm::ToolCall)> = Vec::new(); - let mut approval_needed: Option<( - usize, - crate::llm::ToolCall, - Arc, - )> = None; - - for (idx, original_tc) in tool_calls.iter().enumerate() { - let mut tc = original_tc.clone(); - - // Fetch the tool upfront so we can redact sensitive params - // before they touch hooks or approval display. - let tool_opt = self.tools().get(&tc.name).await; - let sensitive = tool_opt - .as_ref() - .map(|t| t.sensitive_params()) - .unwrap_or(&[]); - - // Hook: BeforeToolCall (runs before approval so hooks can - // modify parameters — approval is checked on final params). - // Hooks receive redacted params so sensitive values are not - // exposed to hook handlers or their logs. - let hook_params = redact_params(&tc.arguments, sensitive); - let event = crate::hooks::HookEvent::ToolCall { - tool_name: tc.name.clone(), - parameters: hook_params, - user_id: message.user_id.clone(), - context: "chat".to_string(), - }; - match self.hooks().run(&event).await { - Err(crate::hooks::HookError::Rejected { reason }) => { - preflight.push(( - tc, - PreflightOutcome::Rejected(format!( - "Tool call rejected by hook: {}", - reason - )), - )); - continue; // skip to next tool (not infinite: using for loop) - } - Err(err) => { - preflight.push(( - tc, - PreflightOutcome::Rejected(format!( - "Tool call blocked by hook policy: {}", - err - )), - )); - continue; - } - Ok(crate::hooks::HookOutcome::Continue { - modified: Some(new_params), - }) => match serde_json::from_str::(&new_params) { - Ok(mut parsed) => { - // Restore original sensitive param values so a hook - // cannot overwrite them (they were sent as [REDACTED]). - if let Some(obj) = parsed.as_object_mut() { - for key in sensitive { - if let Some(orig_val) = original_tc.arguments.get(*key) - { - obj.insert((*key).to_string(), orig_val.clone()); - } - } - } - tc.arguments = parsed; - } - Err(e) => { - tracing::warn!( - tool = %tc.name, - "Hook returned non-JSON modification for ToolCall, ignoring: {}", - e - ); - } - }, - _ => {} - } - - // Check if tool requires approval on the final (post-hook) - // parameters. Skipped when auto_approve_tools is set. - if !self.config.auto_approve_tools - && let Some(tool) = tool_opt - { - use crate::tools::ApprovalRequirement; - let needs_approval = match tool.requires_approval(&tc.arguments) { - ApprovalRequirement::Never => false, - ApprovalRequirement::UnlessAutoApproved => { - let sess = session.lock().await; - !sess.is_tool_auto_approved(&tc.name) - } - ApprovalRequirement::Always => true, - }; - - if needs_approval { - approval_needed = Some((idx, tc, tool)); - break; // remaining tools are deferred - } - } - - let preflight_idx = preflight.len(); - preflight.push((tc.clone(), PreflightOutcome::Runnable)); - runnable.push((preflight_idx, tc)); - } - - // === Phase 2: Parallel execution === - // Execute runnable tools and slot results back by preflight - // index so Phase 3 can iterate in original order. - let mut exec_results: Vec>> = - (0..preflight.len()).map(|_| None).collect(); - - if runnable.len() <= 1 { - // Single tool (or none): execute inline - for (pf_idx, tc) in &runnable { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolStarted { - name: tc.name.clone(), - }, - &message.metadata, - ) - .await; - - let result = self - .execute_chat_tool(&tc.name, &tc.arguments, &job_ctx) - .await; - - let disp_tool = self.tools().get(&tc.name).await; - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::tool_completed( - tc.name.clone(), - &result, - &tc.arguments, - disp_tool.as_deref(), - ), - &message.metadata, - ) - .await; - - exec_results[*pf_idx] = Some(result); - } - } else { - // Multiple tools: execute in parallel via JoinSet - let mut join_set = JoinSet::new(); - - for (pf_idx, tc) in &runnable { - let pf_idx = *pf_idx; - 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 par_tool = tools.get(&tc.name).await; - let _ = channels - .send_status( - &channel, - StatusUpdate::tool_completed( - tc.name.clone(), - &result, - &tc.arguments, - par_tool.as_deref(), - ), - &metadata, - ) - .await; - - (pf_idx, result) - }); - } - - while let Some(join_result) = join_set.join_next().await { - match join_result { - Ok((pf_idx, result)) => { - exec_results[pf_idx] = Some(result); - } - Err(e) => { - if e.is_panic() { - tracing::error!("Chat tool execution task panicked: {}", e); - } else { - tracing::error!( - "Chat tool execution task cancelled: {}", - e - ); - } - } - } - } - - // Fill panicked slots with error results - for (runnable_idx, (pf_idx, tc)) in runnable.iter().enumerate() { - if exec_results[*pf_idx].is_none() { - tracing::error!( - tool = %tc.name, - runnable_idx, - "Filling failed task slot with error" - ); - exec_results[*pf_idx] = - Some(Err(crate::error::ToolError::ExecutionFailed { - name: tc.name.clone(), - reason: "Task failed during execution".to_string(), - } - .into())); - } - } - } - - // === Phase 3: Post-flight (sequential, in original order) === - // Process all results — both hook rejections and execution - // results — in the original tool_calls order. Auth intercept - // is deferred until after every result is recorded. - let mut deferred_auth: Option = None; - - for (pf_idx, (tc, outcome)) in preflight.into_iter().enumerate() { - match outcome { - PreflightOutcome::Rejected(error_msg) => { - // Record hook rejection in thread - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) - && let Some(turn) = thread.last_turn_mut() - { - turn.record_tool_error(error_msg.clone()); - } - } - context_messages - .push(ChatMessage::tool_result(&tc.id, &tc.name, error_msg)); - } - PreflightOutcome::Runnable => { - // Retrieve the execution result for this slot - let tool_result = - exec_results[pf_idx].take().unwrap_or_else(|| { - Err(crate::error::ToolError::ExecutionFailed { - name: tc.name.clone(), - reason: "No result available".to_string(), - } - .into()) - }); - - // Detect image generation sentinel in tool output - // (only from image tools — avoids parsing all tool outputs) - let is_image_sentinel = if let Ok(ref output) = tool_result - && matches!(tc.name.as_str(), "image_generate" | "image_edit") - { - if let Ok(sentinel) = - serde_json::from_str::(output) - && sentinel.get("type").and_then(|v| v.as_str()) - == Some("image_generated") - { - let data_url = sentinel - .get("data") - .and_then(|v| v.as_str()) - .unwrap_or_default() - .to_string(); - let path = sentinel - .get("path") - .and_then(|v| v.as_str()) - .map(String::from); - // Skip broadcasting if data_url is empty to avoid - // sending a broken ImageGenerated SSE event. - if data_url.is_empty() { - tracing::warn!( - "Image generation sentinel has empty data URL, skipping broadcast" - ); - } else { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ImageGenerated { data_url, path }, - &message.metadata, - ) - .await; - } - true - } else { - false - } - } else { - false - }; - - // Send ToolResult preview (skip for image sentinels to avoid - // broadcasting multi-MB base64 data as a preview) - if !is_image_sentinel - && let Ok(ref output) = tool_result - && !output.is_empty() - { - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ToolResult { - name: tc.name.clone(), - preview: output.clone(), - }, - &message.metadata, - ) - .await; - } - - // Check for auth awaiting — defer the return - // until all results are recorded. - if deferred_auth.is_none() - && let Some((ext_name, instructions)) = - check_auth_required(&tc.name, &tool_result) - { - let auth_data = parse_auth_result(&tool_result); - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) { - thread.enter_auth_mode(ext_name.clone()); - } - } - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::AuthRequired { - extension_name: ext_name, - instructions: Some(instructions.clone()), - auth_url: auth_data.auth_url, - setup_url: auth_data.setup_url, - }, - &message.metadata, - ) - .await; - deferred_auth = Some(instructions); - } - - // Stash full output so subsequent tools can reference it - if let Ok(ref output) = tool_result { - job_ctx - .tool_output_stash - .write() - .await - .insert(tc.id.clone(), output.clone()); - } - - // Sanitize and add tool result to context - let is_tool_error = tool_result.is_err(); - let result_content = match tool_result { - Ok(output) => { - let sanitized = - self.safety().sanitize_tool_output(&tc.name, &output); - self.safety().wrap_for_llm( - &tc.name, - &sanitized.content, - sanitized.was_modified, - ) - } - Err(e) => format!("Tool '{}' failed: {}", tc.name, e), - }; - - // Record sanitized result in thread so messages() - // and persist_tool_calls() use cleaned content. - { - let mut sess = session.lock().await; - if let Some(thread) = sess.threads.get_mut(&thread_id) - && let Some(turn) = thread.last_turn_mut() - { - if is_tool_error { - turn.record_tool_error(result_content.clone()); - } else { - turn.record_tool_result(serde_json::json!( - result_content - )); - } - } - } - - context_messages.push(ChatMessage::tool_result( - &tc.id, - &tc.name, - result_content, - )); - } - } - } - - // Return auth response after all results are recorded - if let Some(instructions) = deferred_auth { - return Ok(AgenticLoopResult::Response(instructions)); - } - - // Handle approval if a tool needed it - if let Some((approval_idx, tc, tool)) = approval_needed { - // Show redacted params in the approval UI — the user already knows - // the sensitive value (they provided it); showing it again is - // unnecessary and creates a leakage path through channel logs. - let display_params = redact_params(&tc.arguments, tool.sensitive_params()); - let pending = PendingApproval { - request_id: Uuid::new_v4(), - tool_name: tc.name.clone(), - parameters: tc.arguments.clone(), - display_parameters: display_params, - description: tool.description().to_string(), - tool_call_id: tc.id.clone(), - context_messages: context_messages.clone(), - deferred_tool_calls: tool_calls[approval_idx + 1..].to_vec(), - user_timezone: Some(user_tz.name().to_string()), - }; - - return Ok(AgenticLoopResult::NeedApproval { pending }); - } - } + .into()), + LoopOutcome::NeedApproval(pending) => { + Ok(AgenticLoopResult::NeedApproval { pending: *pending }) } } } @@ -885,11 +227,660 @@ impl Agent { } } +/// Delegate for the chat (dispatcher) context. +/// +/// Implements `LoopDelegate` to customize the shared agentic loop for +/// interactive chat sessions with the full 3-phase tool execution +/// (preflight → parallel exec → post-flight), approval flow, hooks, +/// auth intercept, and cost tracking. +struct ChatDelegate<'a> { + agent: &'a Agent, + session: Arc>, + thread_id: Uuid, + message: &'a IncomingMessage, + job_ctx: JobContext, + active_skills: Vec, + cached_prompt: String, + cached_prompt_no_tools: String, + nudge_at: usize, + force_text_at: usize, + user_tz: chrono_tz::Tz, +} + +#[async_trait] +impl<'a> LoopDelegate for ChatDelegate<'a> { + async fn check_signals(&self) -> LoopSignal { + let sess = self.session.lock().await; + if let Some(thread) = sess.threads.get(&self.thread_id) + && thread.state == ThreadState::Interrupted + { + return LoopSignal::Stop; + } + LoopSignal::Continue + } + + async fn before_llm_call( + &self, + reason_ctx: &mut ReasoningContext, + iteration: usize, + ) -> Option { + // Inject a nudge message when approaching the iteration limit so the + // LLM is aware it should produce a final answer on the next turn. + if iteration == self.nudge_at { + reason_ctx.messages.push(ChatMessage::system( + "You are approaching the tool call limit. \ + Provide your best final answer on the next response \ + using the information you have gathered so far. \ + Do not call any more tools.", + )); + } + + let force_text = iteration >= self.force_text_at; + + // Refresh tool definitions each iteration so newly built tools become visible + let tool_defs = self.agent.tools().tool_definitions().await; + + // Apply trust-based tool attenuation if skills are active. + let tool_defs = if !self.active_skills.is_empty() { + let result = crate::skills::attenuate_tools(&tool_defs, &self.active_skills); + tracing::info!( + min_trust = %result.min_trust, + tools_available = result.tools.len(), + tools_removed = result.removed_tools.len(), + removed = ?result.removed_tools, + explanation = %result.explanation, + "Tool attenuation applied" + ); + result.tools + } else { + tool_defs + }; + + // Update context for this iteration + reason_ctx.available_tools = tool_defs; + reason_ctx.system_prompt = Some(if force_text { + self.cached_prompt_no_tools.clone() + } else { + self.cached_prompt.clone() + }); + reason_ctx.force_text = force_text; + + if force_text { + tracing::info!( + iteration, + "Forcing text-only response (iteration limit reached)" + ); + } + + let _ = self + .agent + .channels + .send_status( + &self.message.channel, + StatusUpdate::Thinking("Calling LLM...".into()), + &self.message.metadata, + ) + .await; + + None + } + + async fn call_llm( + &self, + reasoning: &Reasoning, + reason_ctx: &mut ReasoningContext, + iteration: usize, + ) -> Result { + // Enforce cost guardrails before the LLM call + if let Err(limit) = self.agent.cost_guard().check_allowed().await { + return Err(crate::error::LlmError::InvalidResponse { + provider: "agent".to_string(), + reason: limit.to_string(), + } + .into()); + } + + let output = match reasoning.respond_with_tools(reason_ctx).await { + Ok(output) => output, + Err(crate::error::LlmError::ContextLengthExceeded { used, limit }) => { + tracing::warn!( + used, + limit, + iteration, + "Context length exceeded, compacting messages and retrying" + ); + + // Compact messages in place and retry + reason_ctx.messages = compact_messages_for_retry(&reason_ctx.messages); + + // When force_text, clear tools to further reduce token count + if reason_ctx.force_text { + reason_ctx.available_tools.clear(); + } + + reasoning + .respond_with_tools(reason_ctx) + .await + .map_err(|retry_err| { + tracing::error!( + original_used = used, + original_limit = limit, + retry_error = %retry_err, + "Retry after auto-compaction also failed" + ); + crate::error::Error::from(retry_err) + })? + } + Err(e) => return Err(e.into()), + }; + + // Record cost and track token usage + let model_name = self.agent.llm().active_model_name(); + let read_discount = self.agent.llm().cache_read_discount(); + let write_multiplier = self.agent.llm().cache_write_multiplier(); + let call_cost = self + .agent + .cost_guard() + .record_llm_call( + &model_name, + output.usage.input_tokens, + output.usage.output_tokens, + output.usage.cache_read_input_tokens, + output.usage.cache_creation_input_tokens, + read_discount, + write_multiplier, + Some(self.agent.llm().cost_per_token()), + ) + .await; + tracing::debug!( + "LLM call used {} input + {} output tokens (${:.6})", + output.usage.input_tokens, + output.usage.output_tokens, + call_cost, + ); + + Ok(output) + } + + async fn handle_text_response( + &self, + text: &str, + _reason_ctx: &mut ReasoningContext, + ) -> TextAction { + // Strip internal "[Called tool ...]" text that can leak when + // provider flattening (e.g. NEAR AI) converts tool_calls to + // plain text and the LLM echoes it back. + let sanitized = strip_internal_tool_call_text(text); + TextAction::Return(LoopOutcome::Response(sanitized)) + } + + async fn execute_tool_calls( + &self, + tool_calls: Vec, + content: Option, + reason_ctx: &mut ReasoningContext, + ) -> Result, Error> { + // Add the assistant message with tool_calls to context. + // OpenAI protocol requires this before tool-result messages. + reason_ctx + .messages + .push(ChatMessage::assistant_with_tool_calls( + content, + tool_calls.clone(), + )); + + // Execute tools and add results to context + let _ = self + .agent + .channels + .send_status( + &self.message.channel, + StatusUpdate::Thinking(format!("Executing {} tool(s)...", tool_calls.len())), + &self.message.metadata, + ) + .await; + + // Record tool calls in the thread with sensitive params redacted. + { + let mut redacted_args: Vec = Vec::with_capacity(tool_calls.len()); + for tc in &tool_calls { + let safe = if let Some(tool) = self.agent.tools().get(&tc.name).await { + redact_params(&tc.arguments, tool.sensitive_params()) + } else { + tc.arguments.clone() + }; + redacted_args.push(safe); + } + let mut sess = self.session.lock().await; + if let Some(thread) = sess.threads.get_mut(&self.thread_id) + && let Some(turn) = thread.last_turn_mut() + { + for (tc, safe_args) in tool_calls.iter().zip(redacted_args) { + turn.record_tool_call(&tc.name, safe_args); + } + } + } + + // === Phase 1: Preflight (sequential) === + // Walk tool_calls checking approval and hooks. Classify + // each tool as Rejected (by hook) or Runnable. Stop at the + // first tool that needs approval. + enum PreflightOutcome { + Rejected(String), + Runnable, + } + let mut preflight: Vec<(crate::llm::ToolCall, PreflightOutcome)> = Vec::new(); + let mut runnable: Vec<(usize, crate::llm::ToolCall)> = Vec::new(); + let mut approval_needed: Option<( + usize, + crate::llm::ToolCall, + Arc, + )> = None; + + for (idx, original_tc) in tool_calls.iter().enumerate() { + let mut tc = original_tc.clone(); + + let tool_opt = self.agent.tools().get(&tc.name).await; + let sensitive = tool_opt + .as_ref() + .map(|t| t.sensitive_params()) + .unwrap_or(&[]); + + // Hook: BeforeToolCall + let hook_params = redact_params(&tc.arguments, sensitive); + let event = crate::hooks::HookEvent::ToolCall { + tool_name: tc.name.clone(), + parameters: hook_params, + user_id: self.message.user_id.clone(), + context: "chat".to_string(), + }; + match self.agent.hooks().run(&event).await { + Err(crate::hooks::HookError::Rejected { reason }) => { + preflight.push(( + tc, + PreflightOutcome::Rejected(format!( + "Tool call rejected by hook: {}", + reason + )), + )); + continue; + } + Err(err) => { + preflight.push(( + tc, + PreflightOutcome::Rejected(format!( + "Tool call blocked by hook policy: {}", + err + )), + )); + continue; + } + Ok(crate::hooks::HookOutcome::Continue { + modified: Some(new_params), + }) => match serde_json::from_str::(&new_params) { + Ok(mut parsed) => { + if let Some(obj) = parsed.as_object_mut() { + for key in sensitive { + if let Some(orig_val) = original_tc.arguments.get(*key) { + obj.insert((*key).to_string(), orig_val.clone()); + } + } + } + tc.arguments = parsed; + } + Err(e) => { + tracing::warn!( + tool = %tc.name, + "Hook returned non-JSON modification for ToolCall, ignoring: {}", + e + ); + } + }, + _ => {} + } + + // Check if tool requires approval + if !self.agent.config.auto_approve_tools + && let Some(tool) = tool_opt + { + use crate::tools::ApprovalRequirement; + let needs_approval = match tool.requires_approval(&tc.arguments) { + ApprovalRequirement::Never => false, + ApprovalRequirement::UnlessAutoApproved => { + let sess = self.session.lock().await; + !sess.is_tool_auto_approved(&tc.name) + } + ApprovalRequirement::Always => true, + }; + + if needs_approval { + approval_needed = Some((idx, tc, tool)); + break; + } + } + + let preflight_idx = preflight.len(); + preflight.push((tc.clone(), PreflightOutcome::Runnable)); + runnable.push((preflight_idx, tc)); + } + + // === Phase 2: Parallel execution === + let mut exec_results: Vec>> = + (0..preflight.len()).map(|_| None).collect(); + + if runnable.len() <= 1 { + for (pf_idx, tc) in &runnable { + let _ = self + .agent + .channels + .send_status( + &self.message.channel, + StatusUpdate::ToolStarted { + name: tc.name.clone(), + }, + &self.message.metadata, + ) + .await; + + let result = self + .agent + .execute_chat_tool(&tc.name, &tc.arguments, &self.job_ctx) + .await; + + let disp_tool = self.agent.tools().get(&tc.name).await; + let _ = self + .agent + .channels + .send_status( + &self.message.channel, + StatusUpdate::tool_completed( + tc.name.clone(), + &result, + &tc.arguments, + disp_tool.as_deref(), + ), + &self.message.metadata, + ) + .await; + + exec_results[*pf_idx] = Some(result); + } + } else { + let mut join_set = JoinSet::new(); + + for (pf_idx, tc) in &runnable { + let pf_idx = *pf_idx; + let tools = self.agent.tools().clone(); + let safety = self.agent.safety().clone(); + let channels = self.agent.channels.clone(); + let job_ctx = self.job_ctx.clone(); + let tc = tc.clone(); + let channel = self.message.channel.clone(); + let metadata = self.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 par_tool = tools.get(&tc.name).await; + let _ = channels + .send_status( + &channel, + StatusUpdate::tool_completed( + tc.name.clone(), + &result, + &tc.arguments, + par_tool.as_deref(), + ), + &metadata, + ) + .await; + + (pf_idx, result) + }); + } + + while let Some(join_result) = join_set.join_next().await { + match join_result { + Ok((pf_idx, result)) => { + exec_results[pf_idx] = Some(result); + } + Err(e) => { + if e.is_panic() { + tracing::error!("Chat tool execution task panicked: {}", e); + } else { + tracing::error!("Chat tool execution task cancelled: {}", e); + } + } + } + } + + // Fill panicked slots with error results + for (pf_idx, tc) in runnable.iter() { + if exec_results[*pf_idx].is_none() { + tracing::error!( + tool = %tc.name, + "Filling failed task slot with error" + ); + exec_results[*pf_idx] = Some(Err(crate::error::ToolError::ExecutionFailed { + name: tc.name.clone(), + reason: "Task failed during execution".to_string(), + } + .into())); + } + } + } + + // === Phase 3: Post-flight (sequential, in original order) === + let mut deferred_auth: Option = None; + + for (pf_idx, (tc, outcome)) in preflight.into_iter().enumerate() { + match outcome { + PreflightOutcome::Rejected(error_msg) => { + { + let mut sess = self.session.lock().await; + if let Some(thread) = sess.threads.get_mut(&self.thread_id) + && let Some(turn) = thread.last_turn_mut() + { + turn.record_tool_error(error_msg.clone()); + } + } + reason_ctx + .messages + .push(ChatMessage::tool_result(&tc.id, &tc.name, error_msg)); + } + PreflightOutcome::Runnable => { + let tool_result = exec_results[pf_idx].take().unwrap_or_else(|| { + Err(crate::error::ToolError::ExecutionFailed { + name: tc.name.clone(), + reason: "No result available".to_string(), + } + .into()) + }); + + // Detect image generation sentinel + let is_image_sentinel = if let Ok(ref output) = tool_result + && matches!(tc.name.as_str(), "image_generate" | "image_edit") + { + if let Ok(sentinel) = serde_json::from_str::(output) + && sentinel.get("type").and_then(|v| v.as_str()) + == Some("image_generated") + { + let data_url = sentinel + .get("data") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + let path = sentinel + .get("path") + .and_then(|v| v.as_str()) + .map(String::from); + if data_url.is_empty() { + tracing::warn!( + "Image generation sentinel has empty data URL, skipping broadcast" + ); + } else { + let _ = self + .agent + .channels + .send_status( + &self.message.channel, + StatusUpdate::ImageGenerated { data_url, path }, + &self.message.metadata, + ) + .await; + } + true + } else { + false + } + } else { + false + }; + + // Send ToolResult preview + if !is_image_sentinel + && let Ok(ref output) = tool_result + && !output.is_empty() + { + let _ = self + .agent + .channels + .send_status( + &self.message.channel, + StatusUpdate::ToolResult { + name: tc.name.clone(), + preview: output.clone(), + }, + &self.message.metadata, + ) + .await; + } + + // Check for auth awaiting + if deferred_auth.is_none() + && let Some((ext_name, instructions)) = + check_auth_required(&tc.name, &tool_result) + { + let auth_data = parse_auth_result(&tool_result); + { + let mut sess = self.session.lock().await; + if let Some(thread) = sess.threads.get_mut(&self.thread_id) { + thread.enter_auth_mode(ext_name.clone()); + } + } + let _ = self + .agent + .channels + .send_status( + &self.message.channel, + StatusUpdate::AuthRequired { + extension_name: ext_name, + instructions: Some(instructions.clone()), + auth_url: auth_data.auth_url, + setup_url: auth_data.setup_url, + }, + &self.message.metadata, + ) + .await; + deferred_auth = Some(instructions); + } + + // Stash full output so subsequent tools can reference it + if let Ok(ref output) = tool_result { + self.job_ctx + .tool_output_stash + .write() + .await + .insert(tc.id.clone(), output.clone()); + } + + // Sanitize and add tool result to context + let is_tool_error = tool_result.is_err(); + let result_content = match tool_result { + Ok(output) => { + let sanitized = + self.agent.safety().sanitize_tool_output(&tc.name, &output); + self.agent.safety().wrap_for_llm( + &tc.name, + &sanitized.content, + sanitized.was_modified, + ) + } + Err(e) => format!("Tool '{}' failed: {}", tc.name, e), + }; + + // Record sanitized result in thread + { + let mut sess = self.session.lock().await; + if let Some(thread) = sess.threads.get_mut(&self.thread_id) + && let Some(turn) = thread.last_turn_mut() + { + if is_tool_error { + turn.record_tool_error(result_content.clone()); + } else { + turn.record_tool_result(serde_json::json!(result_content)); + } + } + } + + reason_ctx.messages.push(ChatMessage::tool_result( + &tc.id, + &tc.name, + result_content, + )); + } + } + } + + // Return auth response after all results are recorded + if let Some(instructions) = deferred_auth { + return Ok(Some(LoopOutcome::Response(instructions))); + } + + // Handle approval if a tool needed it + if let Some((approval_idx, tc, tool)) = approval_needed { + let display_params = redact_params(&tc.arguments, tool.sensitive_params()); + let pending = PendingApproval { + request_id: Uuid::new_v4(), + tool_name: tc.name.clone(), + parameters: tc.arguments.clone(), + display_parameters: display_params, + description: tool.description().to_string(), + tool_call_id: tc.id.clone(), + context_messages: reason_ctx.messages.clone(), + deferred_tool_calls: tool_calls[approval_idx + 1..].to_vec(), + user_timezone: Some(self.user_tz.name().to_string()), + }; + + return Ok(Some(LoopOutcome::NeedApproval(Box::new(pending)))); + } + + Ok(None) + } +} + /// Execute a chat tool without requiring `&Agent`. /// /// This standalone function enables parallel invocation from spawned JoinSet -/// tasks, which cannot borrow `&self`. It replicates the logic from -/// `Agent::execute_chat_tool`. +/// tasks, which cannot borrow `&self`. Delegates to the shared +/// `execute_tool_with_safety` pipeline. pub(super) async fn execute_chat_tool_standalone( tools: &crate::tools::ToolRegistry, safety: &crate::safety::SafetyLayer, @@ -897,91 +888,7 @@ pub(super) async fn execute_chat_tool_standalone( params: &serde_json::Value, job_ctx: &crate::context::JobContext, ) -> Result { - let tool = tools - .get(tool_name) - .await - .ok_or_else(|| crate::error::ToolError::NotFound { - name: tool_name.to_string(), - })?; - - // Validate tool parameters - let validation = safety.validator().validate_tool_params(params); - if !validation.is_valid { - let details = validation - .errors - .iter() - .map(|e| format!("{}: {}", e.field, e.message)) - .collect::>() - .join("; "); - return Err(crate::error::ToolError::InvalidParameters { - name: tool_name.to_string(), - reason: format!("Invalid tool parameters: {}", details), - } - .into()); - } - - let safe_params = redact_params(params, tool.sensitive_params()); - tracing::debug!( - tool = %tool_name, - params = %safe_params, - "Tool call started" - ); - - // Execute with per-tool timeout - let timeout = tool.execution_timeout(); - let start = std::time::Instant::now(); - let result = tokio::time::timeout(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(|_| "".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 = timeout.as_secs(), - "Tool call timed out" - ); - } - } - - let result = result - .map_err(|_| crate::error::ToolError::Timeout { - name: tool_name.to_string(), - timeout, - })? - .map_err(|e| crate::error::ToolError::ExecutionFailed { - name: tool_name.to_string(), - reason: e.to_string(), - })?; - - serde_json::to_string_pretty(&result.result).map_err(|e| { - crate::error::ToolError::ExecutionFailed { - name: tool_name.to_string(), - reason: format!("Failed to serialize result: {}", e), - } - .into() - }) + crate::tools::execute::execute_tool_with_safety(tools, safety, tool_name, params, job_ctx).await } /// Parsed auth result fields for emitting StatusUpdate::AuthRequired. diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 895a551a..de2434be 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -11,6 +11,7 @@ //! - Context compaction for long conversations mod agent_loop; +pub mod agentic_loop; mod attachments; mod commands; pub mod compaction; @@ -22,7 +23,7 @@ pub mod job_monitor; mod router; pub mod routine; pub mod routine_engine; -mod scheduler; +pub(crate) mod scheduler; mod self_repair; pub mod session; mod session_manager; @@ -30,8 +31,8 @@ pub mod submission; pub mod task; mod thread_ops; pub mod undo; -pub mod worker; +pub use crate::worker::{Worker, WorkerDeps}; pub(crate) use agent_loop::truncate_for_preview; pub use agent_loop::{Agent, AgentDeps}; pub use compaction::{CompactionResult, ContextCompactor}; @@ -47,4 +48,3 @@ pub use session_manager::SessionManager; pub use submission::{Submission, SubmissionParser, SubmissionResult}; pub use task::{Task, TaskContext, TaskHandler, TaskOutput}; pub use undo::{Checkpoint, UndoManager}; -pub use worker::{Worker, WorkerDeps}; diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 85f3f6eb..971842e5 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -9,7 +9,6 @@ use tokio::task::JoinHandle; use uuid::Uuid; use crate::agent::task::{Task, TaskContext, TaskOutput}; -use crate::agent::worker::{Worker, WorkerDeps}; use crate::channels::web::types::SseEvent; use crate::config::AgentConfig; use crate::context::{ContextManager, JobContext, JobState}; @@ -19,6 +18,7 @@ use crate::hooks::HookRegistry; use crate::llm::LlmProvider; use crate::safety::SafetyLayer; use crate::tools::{ApprovalContext, ToolRegistry}; +use crate::worker::job::{Worker, WorkerDeps}; /// Message to send to a worker. #[derive(Debug)] @@ -462,6 +462,9 @@ impl Scheduler { } /// Execute a single tool as a subtask. + /// + /// Performs scheduler-specific checks (approval, cancellation) then + /// delegates to the shared `execute_tool_with_safety` pipeline. async fn execute_tool_task( tools: Arc, context_manager: Arc, @@ -473,7 +476,7 @@ impl Scheduler { ) -> Result { let start = std::time::Instant::now(); - // Get the tool + // Get the tool for approval check let tool = tools.get(tool_name).await.ok_or_else(|| { Error::Tool(crate::error::ToolError::NotFound { name: tool_name.to_string(), @@ -490,6 +493,7 @@ impl Scheduler { .into()); } + // Scheduler-specific approval check let requirement = tool.requires_approval(¶ms); let blocked = ApprovalContext::is_blocked_or_default(&approval_context, tool_name, requirement); @@ -500,41 +504,23 @@ impl Scheduler { .into()); } - // Validate tool parameters - let validation = safety.validator().validate_tool_params(¶ms); - if !validation.is_valid { - let details = validation - .errors - .iter() - .map(|e| format!("{}: {}", e.field, e.message)) - .collect::>() - .join("; "); - return Err(crate::error::ToolError::InvalidParameters { + // Delegate to shared tool execution pipeline + let output_str = crate::tools::execute::execute_tool_with_safety( + &tools, &safety, tool_name, ¶ms, &job_ctx, + ) + .await?; + + // Parse back to Value for TaskOutput; this should be infallible given + // `execute_tool_with_safety` uses `serde_json::to_string_pretty`, but if it + // ever fails we surface a clear error instead of silently changing types. + let result_value: serde_json::Value = serde_json::from_str(&output_str).map_err(|e| { + Error::Tool(crate::error::ToolError::ExecutionFailed { name: tool_name.to_string(), - reason: format!("Invalid tool parameters: {}", details), - } - .into()); - } + reason: format!("Failed to parse tool output as JSON: {}", e), + }) + })?; - // Execute with per-tool timeout - let tool_timeout = tool.execution_timeout(); - let result = - tokio::time::timeout(tool_timeout, async { tool.execute(params, &job_ctx).await }) - .await - .map_err(|_| { - Error::Tool(crate::error::ToolError::Timeout { - name: tool_name.to_string(), - timeout: tool_timeout, - }) - })? - .map_err(|e| { - Error::Tool(crate::error::ToolError::ExecutionFailed { - name: tool_name.to_string(), - reason: e.to_string(), - }) - })?; - - Ok(TaskOutput::new(result.result, start.elapsed())) + Ok(TaskOutput::new(result_value, start.elapsed())) } /// Stop a running job. diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index c987b826..e7f526e3 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -852,19 +852,12 @@ impl Agent { // Sanitize tool result, then record the cleaned version in the // thread. Must happen before auth intercept check which may return early. let is_tool_error = tool_result.is_err(); - let result_content = match &tool_result { - Ok(output) => { - let sanitized = self - .safety() - .sanitize_tool_output(&pending.tool_name, output); - self.safety().wrap_for_llm( - &pending.tool_name, - &sanitized.content, - sanitized.was_modified, - ) - } - Err(e) => format!("Error: {}", e), - }; + let (result_content, _) = crate::tools::execute::process_tool_result( + self.safety(), + &pending.tool_name, + &pending.tool_call_id, + &tool_result, + ); // Record sanitized result in thread { @@ -1104,17 +1097,12 @@ impl Agent { // Sanitize first, then record the cleaned version in thread. // Must happen before auth detection which may set deferred_auth. let is_deferred_error = deferred_result.is_err(); - let deferred_content = match &deferred_result { - Ok(output) => { - let sanitized = self.safety().sanitize_tool_output(&tc.name, output); - self.safety().wrap_for_llm( - &tc.name, - &sanitized.content, - sanitized.was_modified, - ) - } - Err(e) => format!("Error: {}", e), - }; + let (deferred_content, _) = crate::tools::execute::process_tool_result( + self.safety(), + &tc.name, + &tc.id, + &deferred_result, + ); // Record sanitized result in thread { diff --git a/src/cli/doctor.rs b/src/cli/doctor.rs index c46f4863..aa47b6bf 100644 --- a/src/cli/doctor.rs +++ b/src/cli/doctor.rs @@ -7,6 +7,7 @@ use std::path::PathBuf; use crate::bootstrap::ironclaw_base_dir; +use crate::settings::Settings; /// Run all diagnostic checks and print results. pub async fn run_doctor_command() -> anyhow::Result<()> { @@ -15,14 +16,35 @@ pub async fn run_doctor_command() -> anyhow::Result<()> { let mut passed = 0u32; let mut failed = 0u32; + let mut skipped = 0u32; - // ── Configuration checks ────────────────────────────────── + // Load settings once for checks that need them. + let settings = Settings::load(); + + // ── Settings & core config ───────────────────────────────── + + check( + "Settings file", + check_settings_file(), + &mut passed, + &mut failed, + &mut skipped, + ); check( "NEAR AI session", check_nearai_session().await, &mut passed, &mut failed, + &mut skipped, + ); + + check( + "LLM configuration", + check_llm_config(&settings), + &mut passed, + &mut failed, + &mut skipped, ); check( @@ -30,6 +52,7 @@ pub async fn run_doctor_command() -> anyhow::Result<()> { check_database().await, &mut passed, &mut failed, + &mut skipped, ); check( @@ -37,15 +60,75 @@ pub async fn run_doctor_command() -> anyhow::Result<()> { check_workspace_dir(), &mut passed, &mut failed, + &mut skipped, + ); + + // ── Subsystem configuration checks ───────────────────────── + + check( + "Embeddings", + check_embeddings(&settings), + &mut passed, + &mut failed, + &mut skipped, + ); + + check( + "Routines config", + check_routines_config(), + &mut passed, + &mut failed, + &mut skipped, + ); + + check( + "Gateway config", + check_gateway_config(&settings), + &mut passed, + &mut failed, + &mut skipped, + ); + + check( + "MCP servers", + check_mcp_config().await, + &mut passed, + &mut failed, + &mut skipped, + ); + + check( + "Skills", + check_skills().await, + &mut passed, + &mut failed, + &mut skipped, + ); + + check( + "Secrets", + check_secrets(&settings), + &mut passed, + &mut failed, + &mut skipped, + ); + + check( + "Service", + check_service_installed(), + &mut passed, + &mut failed, + &mut skipped, ); // ── External binary checks ──────────────────────────────── check( - "Docker", - check_binary("docker", &["--version"]), + "Docker daemon", + check_docker_daemon().await, &mut passed, &mut failed, + &mut skipped, ); check( @@ -53,6 +136,7 @@ pub async fn run_doctor_command() -> anyhow::Result<()> { check_binary("cloudflared", &["--version"]), &mut passed, &mut failed, + &mut skipped, ); check( @@ -60,6 +144,7 @@ pub async fn run_doctor_command() -> anyhow::Result<()> { check_binary("ngrok", &["version"]), &mut passed, &mut failed, + &mut skipped, ); check( @@ -67,12 +152,13 @@ pub async fn run_doctor_command() -> anyhow::Result<()> { check_binary("tailscale", &["version"]), &mut passed, &mut failed, + &mut skipped, ); // ── Summary ─────────────────────────────────────────────── println!(); - println!(" {passed} passed, {failed} failed"); + println!(" {passed} passed, {failed} failed, {skipped} skipped"); if failed > 0 { println!("\n Some checks failed. This is normal if you don't use those features."); @@ -83,7 +169,7 @@ pub async fn run_doctor_command() -> anyhow::Result<()> { // ── Individual checks ─────────────────────────────────────── -fn check(name: &str, result: CheckResult, passed: &mut u32, failed: &mut u32) { +fn check(name: &str, result: CheckResult, passed: &mut u32, failed: &mut u32, skipped: &mut u32) { match result { CheckResult::Pass(detail) => { *passed += 1; @@ -94,6 +180,7 @@ fn check(name: &str, result: CheckResult, passed: &mut u32, failed: &mut u32) { println!(" [FAIL] {name}: {detail}"); } CheckResult::Skip(reason) => { + *skipped += 1; println!(" [skip] {name}: {reason}"); } } @@ -105,6 +192,29 @@ enum CheckResult { Skip(String), } +// ── Settings file ─────────────────────────────────────────── + +fn check_settings_file() -> CheckResult { + let path = Settings::default_path(); + if !path.exists() { + return CheckResult::Pass("no settings file (defaults will be used)".into()); + } + + match std::fs::read_to_string(&path) { + Ok(data) => match serde_json::from_str::(&data) { + Ok(_) => CheckResult::Pass(format!("valid ({})", path.display())), + Err(e) => CheckResult::Fail(format!( + "settings.json is malformed: {}. Fix or delete {}", + e, + path.display() + )), + }, + Err(e) => CheckResult::Fail(format!("cannot read {}: {}", path.display(), e)), + } +} + +// ── NEAR AI session ───────────────────────────────────────── + async fn check_nearai_session() -> CheckResult { // Check if session file exists let session_path = crate::config::llm::default_session_path(); @@ -129,6 +239,27 @@ async fn check_nearai_session() -> CheckResult { } } +// ── LLM configuration ────────────────────────────────────── + +fn check_llm_config(settings: &Settings) -> CheckResult { + match crate::llm::LlmConfig::resolve(settings) { + Ok(config) => { + // Show the model for the active backend, not always nearai.model. + let model = if let Some(ref bedrock) = config.bedrock { + &bedrock.model + } else if let Some(ref provider) = config.provider { + &provider.model + } else { + &config.nearai.model + }; + CheckResult::Pass(format!("backend={}, model={}", config.backend, model)) + } + Err(e) => CheckResult::Fail(format!("LLM config error: {e}")), + } +} + +// ── Database ──────────────────────────────────────────────── + async fn check_database() -> CheckResult { let backend = std::env::var("DATABASE_BACKEND") .ok() @@ -192,6 +323,8 @@ async fn try_pg_connect() -> Result<(), String> { Err("postgres feature not compiled in".into()) } +// ── Workspace directory ───────────────────────────────────── + fn check_workspace_dir() -> CheckResult { let dir = ironclaw_base_dir(); @@ -206,6 +339,222 @@ fn check_workspace_dir() -> CheckResult { } } +// ── Embeddings ────────────────────────────────────────────── + +fn check_embeddings(settings: &Settings) -> CheckResult { + match crate::config::EmbeddingsConfig::resolve(settings) { + Ok(config) => { + if !config.enabled { + return CheckResult::Skip("disabled (set EMBEDDING_ENABLED=true)".into()); + } + let has_creds = match config.provider.as_str() { + "openai" => config.openai_api_key().is_some(), + "nearai" => { + // NearAiEmbeddings uses SessionManager::get_token() which + // only returns session tokens, NOT NEARAI_API_KEY + // (src/workspace/embeddings.rs:309, src/llm/session.rs:132). + let session_path = crate::config::llm::default_session_path(); + session_path.exists() + && std::fs::read_to_string(&session_path) + .map(|s| !s.trim().is_empty()) + .unwrap_or(false) + } + "ollama" => true, // local, no creds needed + _ => config.openai_api_key().is_some(), + }; + if has_creds { + CheckResult::Pass(format!( + "provider={}, model={}", + config.provider, config.model + )) + } else { + let hint = match config.provider.as_str() { + "nearai" => "run `ironclaw onboard` to create a session", + _ => "set OPENAI_API_KEY", + }; + CheckResult::Fail(format!( + "provider={} but credentials missing ({})", + config.provider, hint + )) + } + } + Err(e) => CheckResult::Fail(format!("config error: {e}")), + } +} + +// ── Routines config ───────────────────────────────────────── + +fn check_routines_config() -> CheckResult { + match crate::config::RoutineConfig::resolve() { + Ok(config) => { + if config.enabled { + CheckResult::Pass(format!( + "enabled (interval={}s, max_concurrent={})", + config.cron_check_interval_secs, config.max_concurrent_routines + )) + } else { + CheckResult::Skip("disabled".into()) + } + } + Err(e) => CheckResult::Fail(format!("config error: {e}")), + } +} + +// ── Gateway config ────────────────────────────────────────── + +fn check_gateway_config(settings: &Settings) -> CheckResult { + // Use the same resolve() path as runtime so invalid env values + // (e.g. GATEWAY_PORT=abc) are caught here too. + match crate::config::ChannelsConfig::resolve(settings) { + Ok(channels) => match channels.gateway { + Some(gw) => { + if gw.auth_token.is_some() { + CheckResult::Pass(format!( + "enabled at {}:{} (auth token set)", + gw.host, gw.port + )) + } else { + CheckResult::Pass(format!( + "enabled at {}:{} (no auth token — random token will be generated)", + gw.host, gw.port + )) + } + } + None => CheckResult::Skip("disabled (GATEWAY_ENABLED=false)".into()), + }, + Err(e) => CheckResult::Fail(format!("config error: {e}")), + } +} + +// ── MCP servers ───────────────────────────────────────────── + +async fn check_mcp_config() -> CheckResult { + match crate::tools::mcp::config::load_mcp_servers().await { + Ok(file) => { + let servers: Vec<_> = file.enabled_servers().collect(); + if servers.is_empty() { + return CheckResult::Skip("no MCP servers configured".into()); + } + + let mut invalid = Vec::new(); + for server in &servers { + if let Err(e) = server.validate() { + invalid.push(format!("{}: {}", server.name, e)); + } + } + + if invalid.is_empty() { + CheckResult::Pass(format!("{} server(s) configured, all valid", servers.len())) + } else { + CheckResult::Fail(format!( + "{} server(s), {} invalid: {}", + servers.len(), + invalid.len(), + invalid.join("; ") + )) + } + } + Err(e) => { + // Distinguish no config from corrupted config + let msg = e.to_string(); + if msg.contains("not found") || msg.contains("No such file") { + CheckResult::Skip("no MCP config file".into()) + } else { + CheckResult::Fail(format!("config error: {e}")) + } + } + } +} + +// ── Skills ────────────────────────────────────────────────── + +async fn check_skills() -> CheckResult { + let user_dir = ironclaw_base_dir().join("skills"); + let installed_dir = ironclaw_base_dir().join("installed_skills"); + + let mut registry = crate::skills::SkillRegistry::new(user_dir.clone()); + registry = registry.with_installed_dir(installed_dir); + + // discover_all() returns loaded skill names (not warnings). + let _loaded_names = registry.discover_all().await; + + let count = registry.count(); + if count == 0 { + return CheckResult::Skip("no skills discovered".into()); + } + + CheckResult::Pass(format!("{count} skill(s) loaded")) +} + +// ── Secrets ───────────────────────────────────────────────── + +fn check_secrets(settings: &Settings) -> CheckResult { + match settings.secrets_master_key_source { + crate::settings::KeySource::Keychain => { + CheckResult::Pass("master key source: OS keychain".into()) + } + crate::settings::KeySource::Env => { + if std::env::var("SECRETS_MASTER_KEY").is_ok() { + CheckResult::Pass("master key source: env var (set)".into()) + } else { + CheckResult::Fail( + "master key source: env var but SECRETS_MASTER_KEY not set".into(), + ) + } + } + crate::settings::KeySource::None => { + CheckResult::Skip("secrets not configured (run `ironclaw onboard`)".into()) + } + } +} + +// ── Service ───────────────────────────────────────────────── + +fn check_service_installed() -> CheckResult { + if cfg!(target_os = "macos") { + let plist = + dirs::home_dir().map(|h| h.join("Library/LaunchAgents/com.ironclaw.daemon.plist")); + match plist { + Some(path) if path.exists() => { + CheckResult::Pass(format!("launchd plist installed ({})", path.display())) + } + Some(_) => CheckResult::Skip("not installed (run `ironclaw service install`)".into()), + None => CheckResult::Skip("cannot determine home directory".into()), + } + } else if cfg!(target_os = "linux") { + let unit = dirs::home_dir().map(|h| h.join(".config/systemd/user/ironclaw.service")); + match unit { + Some(path) if path.exists() => { + CheckResult::Pass(format!("systemd unit installed ({})", path.display())) + } + Some(_) => CheckResult::Skip("not installed (run `ironclaw service install`)".into()), + None => CheckResult::Skip("cannot determine home directory".into()), + } + } else { + CheckResult::Skip("service management not supported on this platform".into()) + } +} + +// ── Docker daemon ─────────────────────────────────────────── + +async fn check_docker_daemon() -> CheckResult { + let detection = crate::sandbox::check_docker().await; + match detection.status { + crate::sandbox::DockerStatus::Available => CheckResult::Pass("running".into()), + crate::sandbox::DockerStatus::NotInstalled => CheckResult::Skip(format!( + "not installed. {}", + detection.platform.install_hint() + )), + crate::sandbox::DockerStatus::NotRunning => CheckResult::Fail(format!( + "installed but not running. {}", + detection.platform.start_hint() + )), + crate::sandbox::DockerStatus::Disabled => CheckResult::Skip("sandbox disabled".into()), + } +} + +// ── External binary ───────────────────────────────────────── + fn check_binary(name: &str, args: &[&str]) -> CheckResult { match std::process::Command::new(name) .args(args) @@ -273,6 +622,193 @@ mod tests { } } + #[test] + fn check_settings_file_handles_missing() { + // Settings::default_path() might or might not exist, but must not panic + let result = check_settings_file(); + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[test] + fn check_llm_config_does_not_panic() { + let settings = Settings::default(); + let result = check_llm_config(&settings); + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[test] + fn check_routines_config_does_not_panic() { + let result = check_routines_config(); + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[test] + fn check_gateway_config_does_not_panic() { + let settings = Settings::default(); + let result = check_gateway_config(&settings); + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[test] + fn check_embeddings_does_not_panic() { + let settings = Settings::default(); + let result = check_embeddings(&settings); + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[test] + fn check_secrets_none_returns_skip() { + let settings = Settings::default(); + match check_secrets(&settings) { + CheckResult::Skip(msg) => { + assert!( + msg.contains("not configured"), + "expected 'not configured' in skip message, got: {msg}" + ); + } + other => panic!( + "expected Skip for default settings, got: {}", + format_result(&other) + ), + } + } + + #[test] + fn check_service_installed_does_not_panic() { + let result = check_service_installed(); + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[tokio::test] + async fn check_docker_daemon_does_not_panic() { + let result = check_docker_daemon().await; + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[tokio::test] + async fn check_mcp_config_does_not_panic() { + let result = check_mcp_config().await; + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[tokio::test] + async fn check_skills_does_not_panic() { + let result = check_skills().await; + match result { + CheckResult::Pass(_) | CheckResult::Fail(_) | CheckResult::Skip(_) => {} + } + } + + #[test] + fn check_llm_config_shows_nearai_model_for_nearai_backend() { + let _guard = crate::config::helpers::ENV_MUTEX.lock().expect("env mutex"); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::remove_var("LLM_BACKEND"); + } + let settings = Settings::default(); + match check_llm_config(&settings) { + CheckResult::Pass(msg) => { + assert!( + msg.contains("backend=nearai"), + "expected nearai backend, got: {msg}" + ); + // Must NOT show a bedrock or registry model when backend is nearai + assert!( + !msg.contains("anthropic.claude"), + "should not show bedrock model for nearai backend: {msg}" + ); + } + other => panic!( + "expected Pass for default LLM config, got: {}", + format_result(&other) + ), + } + } + + #[test] + fn check_embeddings_disabled_by_default_returns_skip() { + let _guard = crate::config::helpers::ENV_MUTEX.lock().expect("env mutex"); + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("EMBEDDING_ENABLED"); + } + let settings = Settings::default(); + match check_embeddings(&settings) { + CheckResult::Skip(msg) => { + assert!( + msg.contains("disabled"), + "expected 'disabled' in skip message, got: {msg}" + ); + } + other => panic!( + "expected Skip for disabled embeddings, got: {}", + format_result(&other) + ), + } + } + + #[test] + fn check_routines_enabled_by_default() { + let _guard = crate::config::helpers::ENV_MUTEX.lock().expect("env mutex"); + // SAFETY: Under ENV_MUTEX. + unsafe { + std::env::remove_var("ROUTINES_ENABLED"); + } + match check_routines_config() { + CheckResult::Pass(msg) => { + assert!( + msg.contains("enabled"), + "routines should be enabled by default, got: {msg}" + ); + } + other => panic!( + "expected Pass for default routines, got: {}", + format_result(&other) + ), + } + } + + #[test] + fn check_secrets_env_without_var_returns_fail() { + let settings = Settings { + secrets_master_key_source: crate::settings::KeySource::Env, + ..Default::default() + }; + match check_secrets(&settings) { + CheckResult::Fail(msg) => { + assert!( + msg.contains("SECRETS_MASTER_KEY not set"), + "expected mention of missing env var, got: {msg}" + ); + } + CheckResult::Pass(_) => { + // If SECRETS_MASTER_KEY happens to be set in the environment, + // Pass is correct — don't fail the test. + } + other => panic!( + "expected Fail or Pass for env key source, got: {}", + format_result(&other) + ), + } + } + fn format_result(r: &CheckResult) -> String { match r { CheckResult::Pass(s) => format!("Pass({s})"), diff --git a/src/tools/execute.rs b/src/tools/execute.rs new file mode 100644 index 00000000..7c82d7ff --- /dev/null +++ b/src/tools/execute.rs @@ -0,0 +1,391 @@ +//! Shared tool execution pipeline. +//! +//! Provides a single implementation of the validate → timeout → execute → serialize +//! pipeline used by all agentic loop consumers (chat, job, container) and the +//! scheduler's subtask execution. + +use crate::context::JobContext; +use crate::error::Error; +use crate::llm::ChatMessage; +use crate::safety::SafetyLayer; +use crate::tools::{ToolRegistry, redact_params}; + +/// Execute a tool with safety checks: lookup → validate → timeout → execute → serialize. +/// +/// This is the single canonical implementation of tool execution. All consumers +/// (chat dispatcher, job worker, container runtime, scheduler subtasks) use this +/// function instead of maintaining their own copies. +pub async fn execute_tool_with_safety( + tools: &ToolRegistry, + safety: &SafetyLayer, + tool_name: &str, + params: &serde_json::Value, + job_ctx: &JobContext, +) -> Result { + let tool = tools + .get(tool_name) + .await + .ok_or_else(|| crate::error::ToolError::NotFound { + name: tool_name.to_string(), + })?; + + // Validate tool parameters + let validation = safety.validator().validate_tool_params(params); + if !validation.is_valid { + let details = validation + .errors + .iter() + .map(|e| format!("{}: {}", e.field, e.message)) + .collect::>() + .join("; "); + return Err(crate::error::ToolError::InvalidParameters { + name: tool_name.to_string(), + reason: format!("Invalid tool parameters: {}", details), + } + .into()); + } + + let safe_params = redact_params(params, tool.sensitive_params()); + tracing::debug!( + tool = %tool_name, + params = %safe_params, + "Tool call started" + ); + + // Execute with per-tool timeout + let timeout = tool.execution_timeout(); + let start = std::time::Instant::now(); + let result = tokio::time::timeout(timeout, async { + tool.execute(params.clone(), job_ctx).await + }) + .await; + let elapsed = start.elapsed(); + + match &result { + Ok(Ok(output)) => { + let result_size = serde_json::to_string(&output.result) + .map(|s| s.len()) + .unwrap_or(0); + tracing::debug!( + tool = %tool_name, + elapsed_ms = elapsed.as_millis() as u64, + result_size_bytes = result_size, + "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 = timeout.as_secs(), + "Tool call timed out" + ); + } + } + + let result = result + .map_err(|_| crate::error::ToolError::Timeout { + name: tool_name.to_string(), + timeout, + })? + .map_err(|e| crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: e.to_string(), + })?; + + serde_json::to_string_pretty(&result.result).map_err(|e| { + crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: format!("Failed to serialize result: {}", e), + } + .into() + }) +} + +/// Process a tool result into a `ChatMessage::tool_result` with safety sanitization. +/// +/// On success: sanitize → wrap → ChatMessage::tool_result. +/// On error: format error → ChatMessage::tool_result. +/// +/// Returns the content string and the ChatMessage. +pub fn process_tool_result( + safety: &SafetyLayer, + tool_name: &str, + tool_call_id: &str, + result: &Result, +) -> (String, ChatMessage) { + let content = match result { + Ok(output) => { + let sanitized = safety.sanitize_tool_output(tool_name, output); + safety.wrap_for_llm(tool_name, &sanitized.content, sanitized.was_modified) + } + Err(e) => format!("Error: {}", e), + }; + let message = ChatMessage::tool_result(tool_call_id, tool_name, content.clone()); + (content, message) +} + +/// Execute a tool with safety checks, returning a string error (for container runtime). +/// +/// This is a thin wrapper around `execute_tool_with_safety` that converts +/// `Error` to `String` for the container runtime's simpler error model. +pub async fn execute_tool_simple( + tools: &ToolRegistry, + safety: &SafetyLayer, + tool_name: &str, + params: &serde_json::Value, + job_ctx: &JobContext, +) -> Result { + execute_tool_with_safety(tools, safety, tool_name, params, job_ctx) + .await + .map_err(|e| e.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tools::tool::{Tool, ToolError, ToolOutput}; + use std::sync::Arc; + use std::time::Duration; + + struct EchoTool; + + #[async_trait::async_trait] + impl Tool for EchoTool { + fn name(&self) -> &str { + "echo" + } + fn description(&self) -> &str { + "Echoes input" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) + } + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + Ok(ToolOutput::success(params, Duration::default())) + } + fn requires_sanitization(&self) -> bool { + false + } + } + + struct FailTool; + + #[async_trait::async_trait] + impl Tool for FailTool { + fn name(&self) -> &str { + "fail_tool" + } + fn description(&self) -> &str { + "Always fails" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) + } + async fn execute( + &self, + _: serde_json::Value, + _: &JobContext, + ) -> Result { + Err(ToolError::ExecutionFailed( + "intentional failure".to_string(), + )) + } + fn requires_sanitization(&self) -> bool { + false + } + } + + struct SlowTool; + + #[async_trait::async_trait] + impl Tool for SlowTool { + fn name(&self) -> &str { + "slow_tool" + } + fn description(&self) -> &str { + "Sleeps forever" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) + } + async fn execute( + &self, + _: serde_json::Value, + _: &JobContext, + ) -> Result { + tokio::time::sleep(Duration::from_secs(60)).await; + unreachable!() + } + fn execution_timeout(&self) -> Duration { + Duration::from_millis(50) + } + fn requires_sanitization(&self) -> bool { + false + } + } + + fn test_safety() -> SafetyLayer { + SafetyLayer::new(&crate::config::SafetyConfig { + max_output_length: 100_000, + injection_check_enabled: false, + }) + } + + fn test_job_ctx() -> JobContext { + JobContext::default() + } + + async fn registry_with(tools: Vec>) -> ToolRegistry { + let registry = ToolRegistry::new(); + for tool in tools { + registry.register(tool).await; + } + registry + } + + #[tokio::test] + async fn test_execute_success() { + let registry = registry_with(vec![Arc::new(EchoTool)]).await; + let safety = test_safety(); + let params = serde_json::json!({"message": "hello"}); + + let result = + execute_tool_with_safety(®istry, &safety, "echo", ¶ms, &test_job_ctx()).await; + + assert!(result.is_ok(), "Echo tool should succeed"); + let output = result.unwrap(); + assert!( + output.contains("hello"), + "Output should contain the echoed input" + ); + } + + #[tokio::test] + async fn test_execute_missing_tool() { + let registry = registry_with(vec![]).await; + let safety = test_safety(); + + let result = execute_tool_with_safety( + ®istry, + &safety, + "nonexistent", + &serde_json::json!({}), + &test_job_ctx(), + ) + .await; + + assert!(result.is_err(), "Missing tool should return error"); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("nonexistent") || err.contains("not found"), + "Error should mention the tool: {}", + err + ); + } + + #[tokio::test] + async fn test_execute_tool_failure() { + let registry = registry_with(vec![Arc::new(FailTool)]).await; + let safety = test_safety(); + + let result = execute_tool_with_safety( + ®istry, + &safety, + "fail_tool", + &serde_json::json!({}), + &test_job_ctx(), + ) + .await; + + assert!(result.is_err(), "FailTool should return error"); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("intentional failure"), + "Error should contain the failure reason: {}", + err + ); + } + + #[tokio::test] + async fn test_execute_tool_timeout() { + let registry = registry_with(vec![Arc::new(SlowTool)]).await; + let safety = test_safety(); + + let start = std::time::Instant::now(); + let result = execute_tool_with_safety( + ®istry, + &safety, + "slow_tool", + &serde_json::json!({}), + &test_job_ctx(), + ) + .await; + let elapsed = start.elapsed(); + + assert!(result.is_err(), "SlowTool should timeout"); + let err = result.unwrap_err().to_string(); + assert!( + err.to_lowercase().contains("timeout") || err.to_lowercase().contains("timed out"), + "Error should mention timeout: {}", + err + ); + assert!( + elapsed < Duration::from_secs(1), + "Should timeout quickly, not wait 60s" + ); + } + + #[test] + fn test_process_tool_result_success() { + let safety = test_safety(); + let result: Result = Ok("tool output data".to_string()); + + let (content, message) = process_tool_result(&safety, "echo", "call_1", &result); + + assert!( + content.contains("tool_output"), + "Content should be XML-wrapped: {}", + content + ); + assert!( + content.contains("tool output data"), + "Content should contain the output: {}", + content + ); + assert_eq!(message.role, crate::llm::Role::Tool); + assert_eq!(message.name.as_deref(), Some("echo")); + } + + #[test] + fn test_process_tool_result_error() { + let safety = test_safety(); + let result: Result = Err("something went wrong".to_string()); + + let (content, message) = process_tool_result(&safety, "echo", "call_1", &result); + + assert!( + content.contains("Error:"), + "Error content should start with 'Error:': {}", + content + ); + assert!( + content.contains("something went wrong"), + "Error content should contain the message: {}", + content + ); + assert_eq!(message.role, crate::llm::Role::Tool); + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index d379d474..833d278b 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -9,6 +9,7 @@ pub mod builder; pub mod builtin; +pub mod execute; pub mod mcp; pub mod rate_limiter; pub mod schema_validator; diff --git a/src/util.rs b/src/util.rs index 0ac7b69d..866f623c 100644 --- a/src/util.rs +++ b/src/util.rs @@ -24,7 +24,7 @@ pub fn floor_char_boundary(s: &str, pos: usize) -> usize { pub fn llm_signals_completion(response: &str) -> bool { let lower = response.to_lowercase(); - // Superset of phrases from agent/worker.rs and worker/runtime.rs. + // Superset of phrases from worker/job.rs and worker/container.rs. let positive_phrases = [ "job is complete", "job is done", diff --git a/src/worker/container.rs b/src/worker/container.rs new file mode 100644 index 00000000..0b7f41d0 --- /dev/null +++ b/src/worker/container.rs @@ -0,0 +1,539 @@ +//! Worker runtime: the main execution loop inside a container. +//! +//! Reuses the existing `Reasoning` and `SafetyLayer` infrastructure but +//! connects to the orchestrator for LLM calls instead of calling APIs directly. +//! Streams real-time events (message, tool_use, tool_result, result) through +//! the orchestrator's job event pipeline for UI visibility. +//! +//! Uses the shared `AgenticLoop` engine via `ContainerDelegate`. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::agent::agentic_loop::{ + AgenticLoopConfig, LoopDelegate, LoopOutcome, LoopSignal, TextAction, truncate_for_preview, +}; +use crate::config::SafetyConfig; +use crate::context::JobContext; +use crate::error::WorkerError; +use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext}; +use crate::safety::SafetyLayer; +use crate::tools::ToolRegistry; +use crate::tools::execute::{execute_tool_simple, process_tool_result}; +use crate::worker::api::{CompletionReport, JobEventPayload, StatusUpdate, WorkerHttpClient}; +use crate::worker::proxy_llm::ProxyLlmProvider; + +/// Configuration for the worker runtime. +pub struct WorkerConfig { + pub job_id: Uuid, + pub orchestrator_url: String, + pub max_iterations: u32, + pub timeout: Duration, +} + +impl Default for WorkerConfig { + fn default() -> Self { + Self { + job_id: Uuid::nil(), + orchestrator_url: String::new(), + max_iterations: 50, + timeout: Duration::from_secs(600), + } + } +} + +/// The worker runtime runs inside a Docker container. +/// +/// It connects to the orchestrator over HTTP, fetches its job description, +/// then runs a tool execution loop until the job is complete. Events are +/// streamed to the orchestrator so the UI can show real-time progress. +pub struct WorkerRuntime { + config: WorkerConfig, + client: Arc, + llm: Arc, + safety: Arc, + tools: Arc, + /// Credentials fetched from the orchestrator, injected into child processes + /// via `Command::envs()` rather than mutating the global process environment. + /// + /// Wrapped in `Arc` to avoid deep-cloning the map on every tool invocation. + extra_env: Arc>, +} + +impl WorkerRuntime { + /// Create a new worker runtime. + /// + /// Reads `IRONCLAW_WORKER_TOKEN` from the environment for auth. + pub fn new(config: WorkerConfig) -> Result { + let client = Arc::new(WorkerHttpClient::from_env( + config.orchestrator_url.clone(), + config.job_id, + )?); + + let llm: Arc = Arc::new(ProxyLlmProvider::new( + Arc::clone(&client), + "proxied".to_string(), + )); + + let safety = Arc::new(SafetyLayer::new(&SafetyConfig { + max_output_length: 100_000, + injection_check_enabled: true, + })); + + let tools = Arc::new(ToolRegistry::new()); + // Register only container-safe tools + tools.register_container_tools(); + + Ok(Self { + config, + client, + llm, + safety, + tools, + extra_env: Arc::new(HashMap::new()), + }) + } + + /// Run the worker until the job is complete or an error occurs. + pub async fn run(mut self) -> Result<(), WorkerError> { + tracing::info!("Worker starting for job {}", self.config.job_id); + + // Fetch job description from orchestrator + let job = self.client.get_job().await?; + + tracing::info!( + "Received job: {} - {}", + job.title, + truncate_for_preview(&job.description, 100) + ); + + // Fetch credentials and store them for injection into child processes + // via Command::envs() (avoids unsafe std::env::set_var in multi-threaded runtime). + let credentials = self.client.fetch_credentials().await?; + { + let mut env_map = HashMap::new(); + for cred in &credentials { + env_map.insert(cred.env_var.clone(), cred.value.clone()); + } + self.extra_env = Arc::new(env_map); + } + if !credentials.is_empty() { + tracing::info!( + "Fetched {} credential(s) for child process injection", + credentials.len() + ); + } + + // Report that we're starting + self.client + .report_status(&StatusUpdate { + state: "in_progress".to_string(), + message: Some("Worker started, beginning execution".to_string()), + iteration: 0, + }) + .await?; + + // Create reasoning engine + let reasoning = Reasoning::new(self.llm.clone()); + + // Build initial context + let mut reason_ctx = ReasoningContext::new().with_job(&job.description); + + reason_ctx.messages.push(ChatMessage::system(format!( + r#"You are an autonomous agent running inside a Docker container. + +Job: {} +Description: {} + +You have tools for shell commands, file operations, and code editing. +Work independently to complete this job. Report when done."#, + job.title, job.description + ))); + + // Load tool definitions + reason_ctx.available_tools = self.tools.tool_definitions().await; + + // Shared iteration tracker — read after the loop to report accurate counts. + let iteration_tracker = Arc::new(Mutex::new(0u32)); + + // Run with timeout using the shared agentic loop + let result = tokio::time::timeout(self.config.timeout, async { + let delegate = ContainerDelegate { + client: self.client.clone(), + safety: self.safety.clone(), + tools: self.tools.clone(), + extra_env: self.extra_env.clone(), + last_output: Mutex::new(String::new()), + iteration_tracker: iteration_tracker.clone(), + }; + + let config = AgenticLoopConfig { + max_iterations: self.config.max_iterations as usize, + enable_tool_intent_nudge: true, + max_tool_intent_nudges: 2, + }; + + crate::agent::agentic_loop::run_agentic_loop( + &delegate, + &reasoning, + &mut reason_ctx, + &config, + ) + .await + }) + .await; + + let iterations = *iteration_tracker.lock().await; + + match result { + Ok(Ok(LoopOutcome::Response(output))) => { + tracing::info!("Worker completed job {} successfully", self.config.job_id); + self.post_event( + "result", + serde_json::json!({ + "success": true, + "message": truncate_for_preview(&output, 2000), + }), + ) + .await; + self.client + .report_complete(&CompletionReport { + success: true, + message: Some(output), + iterations, + }) + .await?; + } + Ok(Ok(LoopOutcome::MaxIterations)) => { + let msg = format!("max iterations ({}) exceeded", self.config.max_iterations); + tracing::warn!("Worker failed for job {}: {}", self.config.job_id, msg); + self.post_event( + "result", + serde_json::json!({ + "success": false, + "message": format!("Execution failed: {}", msg), + }), + ) + .await; + self.client + .report_complete(&CompletionReport { + success: false, + message: Some(format!("Execution failed: {}", msg)), + iterations, + }) + .await?; + } + Ok(Ok(LoopOutcome::Stopped | LoopOutcome::NeedApproval(_))) => { + tracing::info!("Worker for job {} stopped", self.config.job_id); + self.client + .report_complete(&CompletionReport { + success: false, + message: Some("Execution stopped".to_string()), + iterations, + }) + .await?; + } + Ok(Err(e)) => { + tracing::error!("Worker failed for job {}: {}", self.config.job_id, e); + self.post_event( + "result", + serde_json::json!({ + "success": false, + "message": format!("Execution failed: {}", e), + }), + ) + .await; + self.client + .report_complete(&CompletionReport { + success: false, + message: Some(format!("Execution failed: {}", e)), + iterations, + }) + .await?; + } + Err(_) => { + tracing::warn!("Worker timed out for job {}", self.config.job_id); + self.post_event( + "result", + serde_json::json!({ + "success": false, + "message": "Execution timed out", + }), + ) + .await; + self.client + .report_complete(&CompletionReport { + success: false, + message: Some("Execution timed out".to_string()), + iterations, + }) + .await?; + } + } + + Ok(()) + } + + /// Post a job event to the orchestrator (fire-and-forget). + async fn post_event(&self, event_type: &str, data: serde_json::Value) { + self.client + .post_event(&JobEventPayload { + event_type: event_type.to_string(), + data, + }) + .await; + } +} + +/// Container delegate: implements `LoopDelegate` for the Docker container context. +/// +/// Tools execute sequentially. Events are posted to the orchestrator via HTTP. +/// Completion is detected via `llm_signals_completion()`. +struct ContainerDelegate { + client: Arc, + safety: Arc, + tools: Arc, + extra_env: Arc>, + /// Tracks the last successful tool output for the final response. + last_output: Mutex, + /// Tracks the current iteration — shared with the outer `run` method so + /// `CompletionReport` can include accurate iteration counts. + iteration_tracker: Arc>, +} + +impl ContainerDelegate { + async fn post_event(&self, event_type: &str, data: serde_json::Value) { + self.client + .post_event(&JobEventPayload { + event_type: event_type.to_string(), + data, + }) + .await; + } + + /// Poll the orchestrator for a follow-up prompt. If one is available, + /// inject it as a user message into the reasoning context. + async fn poll_and_inject_prompt(&self, reason_ctx: &mut ReasoningContext) { + match self.client.poll_prompt().await { + Ok(Some(prompt)) => { + tracing::info!( + "Received follow-up prompt: {}", + truncate_for_preview(&prompt.content, 100) + ); + self.post_event( + "message", + serde_json::json!({ + "role": "user", + "content": truncate_for_preview(&prompt.content, 2000), + }), + ) + .await; + reason_ctx.messages.push(ChatMessage::user(&prompt.content)); + } + Ok(None) => {} + Err(e) => { + tracing::debug!("Failed to poll for prompt: {}", e); + } + } + } +} + +#[async_trait] +impl LoopDelegate for ContainerDelegate { + async fn check_signals(&self) -> LoopSignal { + // Container runtime has no stop signals — the orchestrator manages lifecycle. + LoopSignal::Continue + } + + async fn before_llm_call( + &self, + reason_ctx: &mut ReasoningContext, + iteration: usize, + ) -> Option { + let iteration = iteration as u32; + *self.iteration_tracker.lock().await = iteration; + + // Report progress every 5 iterations + if iteration % 5 == 1 { + let _ = self + .client + .report_status(&StatusUpdate { + state: "in_progress".to_string(), + message: Some(format!("Iteration {}", iteration)), + iteration, + }) + .await; + } + + // Poll for follow-up prompts from the user + self.poll_and_inject_prompt(reason_ctx).await; + + // Refresh tools (in case WASM tools were built) + reason_ctx.available_tools = self.tools.tool_definitions().await; + + None + } + + async fn call_llm( + &self, + reasoning: &Reasoning, + reason_ctx: &mut ReasoningContext, + _iteration: usize, + ) -> Result { + // Container uses respond_with_tools (which may return either text or tool calls) + reasoning + .respond_with_tools(reason_ctx) + .await + .map_err(Into::into) + } + + async fn handle_text_response( + &self, + text: &str, + reason_ctx: &mut ReasoningContext, + ) -> TextAction { + self.post_event( + "message", + serde_json::json!({ + "role": "assistant", + "content": truncate_for_preview(text, 2000), + }), + ) + .await; + + // Check for completion + if crate::util::llm_signals_completion(text) { + let last = self.last_output.lock().await; + let output = if last.is_empty() { + text.to_string() + } else { + last.clone() + }; + return TextAction::Return(LoopOutcome::Response(output)); + } + + reason_ctx.messages.push(ChatMessage::assistant(text)); + TextAction::Continue + } + + async fn execute_tool_calls( + &self, + tool_calls: Vec, + content: Option, + reason_ctx: &mut ReasoningContext, + ) -> Result, crate::error::Error> { + if let Some(ref text) = content { + self.post_event( + "message", + serde_json::json!({ + "role": "assistant", + "content": truncate_for_preview(text, 2000), + }), + ) + .await; + } + + // Add assistant message with tool_calls (OpenAI protocol) + reason_ctx + .messages + .push(ChatMessage::assistant_with_tool_calls( + content, + tool_calls.clone(), + )); + + // Execute tools sequentially (container context — no parallel execution) + for tc in tool_calls { + self.post_event( + "tool_use", + serde_json::json!({ + "tool_name": tc.name, + "input": truncate_for_preview(&tc.arguments.to_string(), 500), + }), + ) + .await; + + let job_ctx = JobContext { + extra_env: self.extra_env.clone(), + ..Default::default() + }; + + let result = + execute_tool_simple(&self.tools, &self.safety, &tc.name, &tc.arguments, &job_ctx) + .await; + + self.post_event( + "tool_result", + serde_json::json!({ + "tool_name": tc.name, + "output": match &result { + Ok(output) => truncate_for_preview(output, 2000), + Err(e) => format!("Error: {}", truncate_for_preview(e, 500)), + }, + "success": result.is_ok(), + }), + ) + .await; + + if let Ok(ref output) = result { + *self.last_output.lock().await = output.clone(); + } + + // Use shared result processing + let (_, message) = process_tool_result(&self.safety, &tc.name, &tc.id, &result); + reason_ctx.messages.push(message); + } + + Ok(None) + } + + async fn on_tool_intent_nudge(&self, text: &str, _reason_ctx: &mut ReasoningContext) { + self.post_event( + "message", + serde_json::json!({ + "role": "assistant", + "content": truncate_for_preview(text, 2000), + "nudge": true, + }), + ) + .await; + } + + async fn after_iteration(&self, _iteration: usize) { + // Brief pause between iterations + tokio::time::sleep(Duration::from_millis(100)).await; + } +} + +#[cfg(test)] +mod tests { + use crate::agent::agentic_loop::truncate_for_preview; + + #[test] + fn test_truncate_within_limit() { + assert_eq!(truncate_for_preview("hello", 10), "hello"); + } + + #[test] + fn test_truncate_at_limit() { + assert_eq!(truncate_for_preview("hello", 5), "hello"); + } + + #[test] + fn test_truncate_beyond_limit() { + let result = truncate_for_preview("hello world", 5); + assert_eq!(result, "hello..."); + } + + #[test] + fn test_truncate_multibyte_safe() { + // "é" is 2 bytes in UTF-8; slicing at byte 1 would panic without safety + let result = truncate_for_preview("é is fancy", 1); + // Should truncate to 0 chars (can't fit "é" in 1 byte) + assert_eq!(result, "..."); + } +} diff --git a/src/agent/worker.rs b/src/worker/job.rs similarity index 72% rename from src/agent/worker.rs rename to src/worker/job.rs index 5f6901d7..fd7fcd12 100644 --- a/src/agent/worker.rs +++ b/src/worker/job.rs @@ -1,12 +1,21 @@ -//! Per-job worker execution. +//! Job worker execution via the shared `AgenticLoop`. +//! +//! Replaces `src/agent/worker.rs` with a `JobDelegate` that implements +//! `LoopDelegate`. The `Worker` struct and `WorkerDeps` remain as the +//! public API consumed by `scheduler.rs`. use std::sync::Arc; use std::time::Duration; +use async_trait::async_trait; use tokio::sync::mpsc; use tokio::task::JoinSet; use uuid::Uuid; +use crate::agent::agentic_loop::{ + AgenticLoopConfig, LoopDelegate, LoopOutcome, LoopSignal, TextAction, run_agentic_loop, + truncate_for_preview, +}; use crate::agent::scheduler::WorkerMessage; use crate::agent::task::TaskOutput; use crate::channels::web::types::SseEvent; @@ -19,6 +28,7 @@ use crate::llm::{ ToolSelection, }; use crate::safety::SafetyLayer; +use crate::tools::execute::process_tool_result; use crate::tools::rate_limiter::RateLimitResult; use crate::tools::{ApprovalContext, ToolRegistry, redact_params}; @@ -72,6 +82,7 @@ impl Worker { &self.deps.llm } + #[allow(dead_code)] fn safety(&self) -> &Arc { &self.deps.safety } @@ -242,24 +253,15 @@ Report when the job is complete or if you encounter issues you cannot resolve."# 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(state) if state.is_terminal() => {} + Ok(JobState::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 @@ -304,11 +306,6 @@ Report when the job is complete or if you encounter issues you cannot resolve."# .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; @@ -359,16 +356,10 @@ Report when the job is complete or if you encounter issues you cannot resolve."# 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 we have a plan, execute it. 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 @@ -378,282 +369,36 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } } - // 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, - }), - ); - } - } - } + // Build the delegate and run the shared agentic loop + let delegate = JobDelegate { + worker: self, + rx: tokio::sync::Mutex::new(rx), + consecutive_rate_limits: std::sync::atomic::AtomicUsize::new(0), + }; - // 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(()); - } + let config = AgenticLoopConfig { + max_iterations, + enable_tool_intent_nudge: true, + max_tool_intent_nudges: 2, + }; - iteration += 1; - if iteration > max_iterations { + let outcome = run_agentic_loop(&delegate, reasoning, reason_ctx, &config).await?; + + match outcome { + LoopOutcome::Response(_) => { + // Completion was already handled in handle_text_response via mark_completed + } + LoopOutcome::MaxIterations => { self.mark_failed("Maximum iterations exceeded: job hit the iteration cap") .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_failed("Persistent rate limiting: exceeded retry limit") - .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_failed("Persistent rate limiting: exceeded retry limit") - .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()), - }; - - // Track token usage from LLM call against the job budget. - // NOTE: select_tools() also makes LLM calls but doesn't expose - // TokenUsage; only respond_with_tools() usage is tracked here. - let total_tokens = respond_output.usage.total() as u64; - if total_tokens > 0 - && let Err(msg) = self - .context_manager() - .update_context(self.job_id, |ctx| ctx.add_tokens(total_tokens)) - .await? - { - self.mark_failed(&msg).await?; - return Ok(()); - } - - 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 = 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 = 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?; - } - } + LoopOutcome::Stopped => { + // Stop signal handled — nothing more to do } - - // 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; + LoopOutcome::NeedApproval(_) => {} } + + Ok(()) } /// Execute multiple tools in parallel using a JoinSet. @@ -833,8 +578,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# .into()); } - // Redact sensitive parameter values (e.g. secret_save's "value") before - // they touch any observability or audit path. + // Redact sensitive parameter values before they touch any observability or audit path. let safe_params = redact_params(¶ms, tool.sensitive_params()); tracing::debug!( tool = %tool_name, @@ -854,12 +598,13 @@ Report when the job is complete or if you encounter issues you cannot resolve."# match &result { Ok(Ok(output)) => { - let result_str = serde_json::to_string(&output.result) - .unwrap_or_else(|_| "".to_string()); + let result_size = serde_json::to_string(&output.result) + .map(|s| s.len()) + .unwrap_or(0); tracing::debug!( tool = %tool_name, elapsed_ms = elapsed.as_millis() as u64, - result = %result_str, + result_size_bytes = result_size, "Tool call succeeded" ); } @@ -978,51 +723,47 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } /// Process a tool execution result and add it to the reasoning context. - async fn process_tool_result( + async fn process_tool_result_job( &self, reason_ctx: &mut ReasoningContext, selection: &ToolSelection, result: Result, - ) -> Result { + ) -> Result<(), Error> { self.log_event( "tool_use", serde_json::json!({ "tool_name": selection.tool_name, - "input": crate::agent::agent_loop::truncate_for_preview( + "input": truncate_for_preview( &selection.parameters.to_string(), 500), }), ); - match result { - Ok(output) => { - // Sanitize output + // Use shared result processing for sanitize → wrap → ChatMessage. + // The wrapped content (XML tags) goes into reason_ctx for the LLM. + // The raw sanitized content goes into events/SSE for human-readable UI. + let (_wrapped, message) = process_tool_result( + &self.deps.safety, + &selection.tool_name, + &selection.tool_call_id, + &result, + ); + reason_ctx.messages.push(message); + + match &result { + Ok(raw_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, + .deps + .safety + .sanitize_tool_output(&selection.tool_name, raw_output); + self.log_event( + "tool_result", + serde_json::json!({ + "tool_name": selection.tool_name, + "success": true, + "output": truncate_for_preview(&sanitized.content, 500), + }), ); - - 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) + Ok(()) } Err(e) => { tracing::warn!( @@ -1050,17 +791,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."# serde_json::json!({ "tool_name": selection.tool_name, "success": false, - "output": format!("Error: {}", e), + "output": truncate_for_preview(&format!("Error: {}", e), 500), }), ); - reason_ctx.messages.push(ChatMessage::tool_result( - &selection.tool_call_id, - &selection.tool_name, - format!("Error: {}", e), - )); - - Ok(false) + Ok(()) } } } @@ -1107,8 +842,6 @@ Report when the job is complete or if you encounter issues you cannot resolve."# "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(()); } } @@ -1123,9 +856,6 @@ Report when the job is complete or if you encounter issues you cannot resolve."# 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(), @@ -1134,8 +864,6 @@ Report when the job is complete or if you encounter issues you cannot resolve."# 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( @@ -1147,21 +875,13 @@ Report when the job is complete or if you encounter issues you cannot resolve."# }], )); - // 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) + self.process_tool_result_job(reason_ctx, &selection, result) .await?; - if completed { - return Ok(()); - } - - // Small delay between actions tokio::time::sleep(Duration::from_millis(100)).await; } @@ -1176,8 +896,6 @@ Report when the job is complete or if you encounter issues you cannot resolve."# 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 @@ -1275,6 +993,343 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } } +/// Job delegate: implements `LoopDelegate` for the background job context. +/// +/// Handles: signal channel (stop/ping/user messages), cancellation checks, +/// rate-limit retry, parallel tool execution, DB persistence, SSE broadcasting. +struct JobDelegate<'a> { + worker: &'a Worker, + rx: tokio::sync::Mutex<&'a mut mpsc::Receiver>, + /// Tracks consecutive rate-limit errors to fail fast instead of burning iterations. + consecutive_rate_limits: std::sync::atomic::AtomicUsize, +} + +impl<'a> JobDelegate<'a> { + const MAX_CONSECUTIVE_RATE_LIMITS: usize = 10; + + /// Handle a rate-limit error: back off, increment counter, and fail fast + /// if the provider remains rate-limited for too many consecutive attempts. + async fn handle_rate_limit( + &self, + retry_after: Option, + context: &str, + ) -> Result { + use std::sync::atomic::Ordering::Relaxed; + + let count = self.consecutive_rate_limits.fetch_add(1, Relaxed) + 1; + let wait = retry_after.unwrap_or(Duration::from_secs(5)); + tracing::warn!( + job_id = %self.worker.job_id, + wait_secs = wait.as_secs(), + attempt = count, + "LLM rate limited during {}, backing off", + context, + ); + + if count >= Self::MAX_CONSECUTIVE_RATE_LIMITS { + self.worker + .mark_failed("Persistent rate limiting: exceeded retry limit") + .await?; + return Err(crate::error::LlmError::RateLimited { + provider: "rate-limit-exhausted".to_string(), + retry_after: None, + } + .into()); + } + + self.worker.log_event( + "status", + serde_json::json!({ + "message": format!( + "Rate limited, retrying in {}s... ({}/{})", + wait.as_secs(), count, Self::MAX_CONSECUTIVE_RATE_LIMITS + ), + }), + ); + tokio::time::sleep(wait).await; + + Ok(crate::llm::RespondOutput { + result: RespondResult::Text(String::new()), + usage: crate::llm::TokenUsage::default(), + }) + } +} + +#[async_trait] +impl<'a> LoopDelegate for JobDelegate<'a> { + async fn check_signals(&self) -> LoopSignal { + // Drain the entire message channel, prioritizing Stop over user messages. + // Scope the lock so it's dropped before any .await below. + let mut stop_requested = false; + let mut first_user_message: Option = None; + { + let mut rx = self.rx.lock().await; + while let Ok(msg) = rx.try_recv() { + match msg { + WorkerMessage::Stop => { + tracing::debug!( + "Worker for job {} received stop signal", + self.worker.job_id + ); + stop_requested = true; + } + WorkerMessage::Ping => { + tracing::trace!("Worker for job {} received ping", self.worker.job_id); + } + WorkerMessage::Start => {} + WorkerMessage::UserMessage(content) => { + tracing::info!( + job_id = %self.worker.job_id, + "Worker received follow-up user message" + ); + self.worker.log_event( + "message", + serde_json::json!({ + "role": "user", + "content": content, + }), + ); + // Keep only the first user message; subsequent ones will be + // picked up on the next iteration's drain. + if first_user_message.is_none() { + first_user_message = Some(content); + } + } + } + } + } // MutexGuard dropped here, before the cancellation .await + + // Stop takes priority over user messages + if stop_requested { + return LoopSignal::Stop; + } + + if let Some(content) = first_user_message { + return LoopSignal::InjectMessage(content); + } + + // Check for terminal or non-progressing state. The loop should stop when the + // job has been cancelled, failed, stuck, or already completed — not just the + // three states that `is_terminal()` covers (Accepted/Failed/Cancelled). + if let Ok(ctx) = self + .worker + .context_manager() + .get_context(self.worker.job_id) + .await + && matches!( + ctx.state, + JobState::Cancelled + | JobState::Failed + | JobState::Stuck + | JobState::Completed + | JobState::Submitted + | JobState::Accepted + ) + { + tracing::info!( + "Worker for job {} detected terminal state {:?}", + self.worker.job_id, + ctx.state, + ); + return LoopSignal::Stop; + } + + LoopSignal::Continue + } + + async fn before_llm_call( + &self, + reason_ctx: &mut ReasoningContext, + _iteration: usize, + ) -> Option { + // Refresh tool definitions so newly built tools become visible + reason_ctx.available_tools = self.worker.tools().tool_definitions().await; + None + } + + async fn call_llm( + &self, + reasoning: &Reasoning, + reason_ctx: &mut ReasoningContext, + _iteration: usize, + ) -> Result { + // Try select_tools first, fall back to respond_with_tools + match reasoning.select_tools(reason_ctx).await { + Ok(s) if !s.is_empty() => { + // Reset counter after a successful LLM call + self.consecutive_rate_limits + .store(0, std::sync::atomic::Ordering::Relaxed); + let tool_calls: Vec = selections_to_tool_calls(&s); + return Ok(crate::llm::RespondOutput { + result: RespondResult::ToolCalls { + tool_calls, + content: None, + }, + usage: crate::llm::TokenUsage::default(), + }); + } + Ok(_) => {} // empty selections, fall through + Err(crate::error::LlmError::RateLimited { retry_after, .. }) => { + return self.handle_rate_limit(retry_after, "tool selection").await; + } + Err(e) => return Err(e.into()), + }; + + // Fall back to respond_with_tools + match reasoning.respond_with_tools(reason_ctx).await { + Ok(output) => { + // Reset counter after a successful LLM call + self.consecutive_rate_limits + .store(0, std::sync::atomic::Ordering::Relaxed); + + // Track token usage against the job budget. + // NOTE: select_tools() also makes LLM calls but doesn't expose + // TokenUsage; only respond_with_tools() usage is tracked here. + let total_tokens = output.usage.total() as u64; + if total_tokens > 0 + && let Err(msg) = self + .worker + .context_manager() + .update_context(self.worker.job_id, |ctx| ctx.add_tokens(total_tokens)) + .await? + { + self.worker.mark_failed(&msg).await?; + } + + Ok(output) + } + Err(crate::error::LlmError::RateLimited { retry_after, .. }) => { + self.handle_rate_limit(retry_after, "respond_with_tools") + .await + } + Err(e) => Err(e.into()), + } + } + + async fn handle_text_response( + &self, + text: &str, + reason_ctx: &mut ReasoningContext, + ) -> TextAction { + // Empty text from rate-limit backoff retry — skip processing and let the + // loop proceed to the next iteration which will re-call the LLM. + if text.is_empty() { + return TextAction::Continue; + } + + // Check for explicit completion + if crate::util::llm_signals_completion(text) { + if let Err(e) = self.worker.mark_completed().await { + tracing::warn!( + "Failed to mark job {} as completed: {}", + self.worker.job_id, + e + ); + } + return TextAction::Return(LoopOutcome::Response(text.to_string())); + } + + // Add assistant response to context + reason_ctx.messages.push(ChatMessage::assistant(text)); + + self.worker.log_event( + "message", + serde_json::json!({ + "role": "assistant", + "content": text, + }), + ); + + TextAction::Continue + } + + async fn execute_tool_calls( + &self, + tool_calls: Vec, + content: Option, + reason_ctx: &mut ReasoningContext, + ) -> Result, crate::error::Error> { + if let Some(ref text) = content { + self.worker.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 to ToolSelections + let selections: Vec = 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(); + + // Execute tools (parallel for multiple, direct for single) + if selections.len() == 1 { + let selection = &selections[0]; + let result = self + .worker + .execute_tool(&selection.tool_name, &selection.parameters) + .await; + self.worker + .process_tool_result_job(reason_ctx, selection, result) + .await?; + } else { + let results = self.worker.execute_tools_parallel(&selections).await; + for (selection, result) in selections.iter().zip(results) { + self.worker + .process_tool_result_job(reason_ctx, selection, result.result) + .await?; + } + } + + Ok(None) + } + + async fn on_tool_intent_nudge(&self, text: &str, _reason_ctx: &mut ReasoningContext) { + self.worker.log_event( + "message", + serde_json::json!({ + "role": "assistant", + "content": truncate_for_preview(text, 2000), + "nudge": true, + }), + ); + } + + async fn after_iteration(&self, _iteration: usize) { + // Small delay between iterations + tokio::time::sleep(Duration::from_millis(100)).await; + } +} + +/// Convert `ToolSelection`s to `ToolCall`s. +fn selections_to_tool_calls(selections: &[ToolSelection]) -> Vec { + selections + .iter() + .map(|s| ToolCall { + id: s.tool_call_id.clone(), + name: s.tool_name.clone(), + arguments: s.parameters.clone(), + }) + .collect() +} + /// Convert a TaskOutput to a string result for tool execution. impl From for Result { fn from(output: TaskOutput) -> Self { @@ -1291,7 +1346,6 @@ impl From for Result { #[cfg(test)] mod tests { use crate::llm::ToolSelection; - use crate::util::llm_signals_completion; use super::*; use crate::config::SafetyConfig; @@ -1301,7 +1355,7 @@ mod tests { ToolCompletionResponse, }; use crate::safety::SafetyLayer; - use crate::tools::{Tool, ToolError, ToolOutput}; + use crate::tools::{Tool, ToolError as ToolExecError, ToolOutput}; /// A test tool that sleeps for a configurable duration before returning. struct SlowTool { @@ -1324,7 +1378,7 @@ mod tests { &self, _params: serde_json::Value, _ctx: &JobContext, - ) -> Result { + ) -> Result { let start = std::time::Instant::now(); tokio::time::sleep(self.delay).await; Ok(ToolOutput::text( @@ -1409,70 +1463,11 @@ mod tests { ); } - #[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" - )); - } + // Completion detection tests live in src/util.rs (the canonical location). + // See: test_completion_signals, test_completion_negative, etc. #[tokio::test] async fn test_parallel_speedup() { - // 3 tools each sleeping 200ms should finish in roughly 200ms (parallel), - // not ~600ms (sequential). let tools: Vec> = (0..3) .map(|i| { Arc::new(SlowTool { @@ -1502,9 +1497,6 @@ mod tests { 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)", @@ -1514,8 +1506,6 @@ mod tests { #[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> = vec![ Arc::new(SlowTool { tool_name: "tool_a".into(), @@ -1559,7 +1549,6 @@ mod tests { 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")); @@ -1567,7 +1556,6 @@ mod tests { #[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 { @@ -1586,13 +1574,10 @@ mod tests { ); } - /// 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| { @@ -1602,10 +1587,8 @@ mod tests { .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) @@ -1613,7 +1596,6 @@ mod tests { .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(), @@ -1726,7 +1708,6 @@ mod tests { #[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!({})) @@ -1736,7 +1717,6 @@ mod tests { "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()), @@ -1750,7 +1730,6 @@ mod tests { #[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()), @@ -1764,7 +1743,6 @@ mod tests { "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([ diff --git a/src/worker/mod.rs b/src/worker/mod.rs index dce75b3d..c6028b96 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -26,13 +26,15 @@ pub mod api; pub mod claude_bridge; +pub mod container; +pub mod job; pub mod proxy_llm; -pub mod runtime; pub use api::WorkerHttpClient; pub use claude_bridge::ClaudeBridgeRuntime; +pub use container::WorkerRuntime; +pub use job::{Worker, WorkerDeps}; pub use proxy_llm::ProxyLlmProvider; -pub use runtime::WorkerRuntime; /// Run the Worker subcommand (inside Docker containers). pub async fn run_worker( @@ -46,7 +48,7 @@ pub async fn run_worker( orchestrator_url ); - let config = runtime::WorkerConfig { + let config = container::WorkerConfig { job_id, orchestrator_url: orchestrator_url.to_string(), max_iterations, diff --git a/src/worker/runtime.rs b/src/worker/runtime.rs deleted file mode 100644 index 677a4cf8..00000000 --- a/src/worker/runtime.rs +++ /dev/null @@ -1,570 +0,0 @@ -//! Worker runtime: the main execution loop inside a container. -//! -//! Reuses the existing `Reasoning` and `SafetyLayer` infrastructure but -//! connects to the orchestrator for LLM calls instead of calling APIs directly. -//! Streams real-time events (message, tool_use, tool_result, result) through -//! the orchestrator's job event pipeline for UI visibility. - -use std::collections::HashMap; -use std::sync::Arc; -use std::time::Duration; - -use uuid::Uuid; - -use crate::config::SafetyConfig; -use crate::context::JobContext; -use crate::error::WorkerError; -use crate::llm::{ - ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection, -}; -use crate::safety::SafetyLayer; -use crate::tools::ToolRegistry; -use crate::worker::api::{CompletionReport, JobEventPayload, StatusUpdate, WorkerHttpClient}; -use crate::worker::proxy_llm::ProxyLlmProvider; - -/// Configuration for the worker runtime. -pub struct WorkerConfig { - pub job_id: Uuid, - pub orchestrator_url: String, - pub max_iterations: u32, - pub timeout: Duration, -} - -impl Default for WorkerConfig { - fn default() -> Self { - Self { - job_id: Uuid::nil(), - orchestrator_url: String::new(), - max_iterations: 50, - timeout: Duration::from_secs(600), - } - } -} - -/// The worker runtime runs inside a Docker container. -/// -/// It connects to the orchestrator over HTTP, fetches its job description, -/// then runs a tool execution loop until the job is complete. Events are -/// streamed to the orchestrator so the UI can show real-time progress. -pub struct WorkerRuntime { - config: WorkerConfig, - client: Arc, - llm: Arc, - safety: Arc, - tools: Arc, - /// Credentials fetched from the orchestrator, injected into child processes - /// via `Command::envs()` rather than mutating the global process environment. - /// - /// Wrapped in `Arc` to avoid deep-cloning the map on every tool invocation. - extra_env: Arc>, -} - -impl WorkerRuntime { - /// Create a new worker runtime. - /// - /// Reads `IRONCLAW_WORKER_TOKEN` from the environment for auth. - pub fn new(config: WorkerConfig) -> Result { - let client = Arc::new(WorkerHttpClient::from_env( - config.orchestrator_url.clone(), - config.job_id, - )?); - - let llm: Arc = Arc::new(ProxyLlmProvider::new( - Arc::clone(&client), - "proxied".to_string(), - )); - - let safety = Arc::new(SafetyLayer::new(&SafetyConfig { - max_output_length: 100_000, - injection_check_enabled: true, - })); - - let tools = Arc::new(ToolRegistry::new()); - // Register only container-safe tools - tools.register_container_tools(); - - Ok(Self { - config, - client, - llm, - safety, - tools, - extra_env: Arc::new(HashMap::new()), - }) - } - - /// Run the worker until the job is complete or an error occurs. - pub async fn run(mut self) -> Result<(), WorkerError> { - tracing::info!("Worker starting for job {}", self.config.job_id); - - // Fetch job description from orchestrator - let job = self.client.get_job().await?; - - tracing::info!( - "Received job: {} - {}", - job.title, - truncate(&job.description, 100) - ); - - // Fetch credentials and store them for injection into child processes - // via Command::envs() (avoids unsafe std::env::set_var in multi-threaded runtime). - let credentials = self.client.fetch_credentials().await?; - { - let mut env_map = HashMap::new(); - for cred in &credentials { - env_map.insert(cred.env_var.clone(), cred.value.clone()); - } - self.extra_env = Arc::new(env_map); - } - if !credentials.is_empty() { - tracing::info!( - "Fetched {} credential(s) for child process injection", - credentials.len() - ); - } - - // Report that we're starting - self.client - .report_status(&StatusUpdate { - state: "in_progress".to_string(), - message: Some("Worker started, beginning execution".to_string()), - iteration: 0, - }) - .await?; - - // Create reasoning engine - let reasoning = Reasoning::new(self.llm.clone()) - .with_model_name(self.llm.active_model_name()); - - // Build initial context - let mut reason_ctx = ReasoningContext::new().with_job(&job.description); - - reason_ctx.messages.push(ChatMessage::system(format!( - r#"You are an autonomous agent running inside a Docker container. - -Job: {} -Description: {} - -You have tools for shell commands, file operations, and code editing. -Work independently to complete this job. Report when done."#, - job.title, job.description - ))); - - // Run with timeout - let result = tokio::time::timeout(self.config.timeout, async { - self.execution_loop(&reasoning, &mut reason_ctx).await - }) - .await; - - match result { - Ok(Ok(output)) => { - tracing::info!("Worker completed job {} successfully", self.config.job_id); - self.post_event( - "result", - serde_json::json!({ - "success": true, - "message": truncate(&output, 2000), - }), - ) - .await; - self.client - .report_complete(&CompletionReport { - success: true, - message: Some(output), - iterations: 0, - }) - .await?; - } - Ok(Err(e)) => { - tracing::error!("Worker failed for job {}: {}", self.config.job_id, e); - self.post_event( - "result", - serde_json::json!({ - "success": false, - "message": format!("Execution failed: {}", e), - }), - ) - .await; - self.client - .report_complete(&CompletionReport { - success: false, - message: Some(format!("Execution failed: {}", e)), - iterations: 0, - }) - .await?; - } - Err(_) => { - tracing::warn!("Worker timed out for job {}", self.config.job_id); - self.post_event( - "result", - serde_json::json!({ - "success": false, - "message": "Execution timed out", - }), - ) - .await; - self.client - .report_complete(&CompletionReport { - success: false, - message: Some("Execution timed out".to_string()), - iterations: 0, - }) - .await?; - } - } - - Ok(()) - } - - async fn execution_loop( - &self, - reasoning: &Reasoning, - reason_ctx: &mut ReasoningContext, - ) -> Result { - let max_iterations = self.config.max_iterations; - let mut last_output = String::new(); - const MAX_TOOL_INTENT_NUDGES: u32 = 2; - let mut consecutive_tool_intent_nudges: u32 = 0; - - // Load tool definitions - reason_ctx.available_tools = self.tools.tool_definitions().await; - - for iteration in 1..=max_iterations { - // Report progress - if iteration % 5 == 1 { - let _ = self - .client - .report_status(&StatusUpdate { - state: "in_progress".to_string(), - message: Some(format!("Iteration {}", iteration)), - iteration, - }) - .await; - } - - // Poll for follow-up prompts from the user - self.poll_and_inject_prompt(reason_ctx).await; - - // Refresh tools (in case WASM tools were built) - reason_ctx.available_tools = self.tools.tool_definitions().await; - - // Ask the LLM what to do next - let selections = reasoning.select_tools(reason_ctx).await.map_err(|e| { - WorkerError::ExecutionFailed { - reason: format!("tool selection failed: {}", e), - } - })?; - - if selections.is_empty() { - // No tools selected, try direct response - let respond_result = - reasoning - .respond_with_tools(reason_ctx) - .await - .map_err(|e| WorkerError::ExecutionFailed { - reason: format!("respond_with_tools failed: {}", e), - })?; - - match respond_result.result { - RespondResult::Text(response) => { - self.post_event( - "message", - serde_json::json!({ - "role": "assistant", - "content": truncate(&response, 2000), - }), - ) - .await; - - if crate::util::llm_signals_completion(&response) { - if last_output.is_empty() { - last_output = response.clone(); - } - return Ok(last_output); - } - reason_ctx.messages.push(ChatMessage::assistant(&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!( - "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; - } - } - RespondResult::ToolCalls { - tool_calls, - content, - } => { - consecutive_tool_intent_nudges = 0; - if let Some(ref text) = content { - self.post_event( - "message", - serde_json::json!({ - "role": "assistant", - "content": truncate(text, 2000), - }), - ) - .await; - } - - // Add assistant message with tool_calls (OpenAI protocol) - reason_ctx - .messages - .push(ChatMessage::assistant_with_tool_calls( - content, - tool_calls.clone(), - )); - - for tc in tool_calls { - self.post_event( - "tool_use", - serde_json::json!({ - "tool_name": tc.name, - "input": truncate(&tc.arguments.to_string(), 500), - }), - ) - .await; - - let result = self.execute_tool(&tc.name, &tc.arguments).await; - - self.post_event( - "tool_result", - serde_json::json!({ - "tool_name": tc.name, - "output": match &result { - Ok(output) => truncate(output, 2000), - Err(e) => format!("Error: {}", truncate(e, 500)), - }, - "success": result.is_ok(), - }), - ) - .await; - - if let Ok(ref output) = result { - last_output = output.clone(); - } - let selection = ToolSelection { - tool_name: tc.name.clone(), - parameters: tc.arguments.clone(), - reasoning: String::new(), - alternatives: vec![], - tool_call_id: tc.id.clone(), - }; - self.process_result(reason_ctx, &selection, result); - } - } - } - } else { - consecutive_tool_intent_nudges = 0; - // Execute selected tools - for selection in &selections { - self.post_event( - "tool_use", - serde_json::json!({ - "tool_name": selection.tool_name, - "input": truncate(&selection.parameters.to_string(), 500), - }), - ) - .await; - - let result = self - .execute_tool(&selection.tool_name, &selection.parameters) - .await; - - self.post_event( - "tool_result", - serde_json::json!({ - "tool_name": selection.tool_name, - "output": match &result { - Ok(output) => truncate(output, 2000), - Err(e) => format!("Error: {}", truncate(e, 500)), - }, - "success": result.is_ok(), - }), - ) - .await; - - if let Ok(ref output) = result { - last_output = output.clone(); - } - - let completed = self.process_result(reason_ctx, selection, result); - if completed { - return Ok(last_output); - } - } - } - - // Brief pause between iterations - tokio::time::sleep(Duration::from_millis(100)).await; - } - - Err(WorkerError::ExecutionFailed { - reason: format!("max iterations ({}) exceeded", max_iterations), - }) - } - - async fn execute_tool( - &self, - tool_name: &str, - params: &serde_json::Value, - ) -> Result { - let tool = match self.tools.get(tool_name).await { - Some(t) => t, - None => return Err(format!("tool '{}' not found", tool_name)), - }; - - let ctx = JobContext { - extra_env: self.extra_env.clone(), - ..Default::default() - }; - - // Validate params - let validation = self.safety.validator().validate_tool_params(params); - if !validation.is_valid { - let details = validation - .errors - .iter() - .map(|e| format!("{}: {}", e.field, e.message)) - .collect::>() - .join("; "); - return Err(format!("invalid parameters: {}", details)); - } - - // Execute with per-tool timeout - let tool_timeout = tool.execution_timeout(); - let result = tokio::time::timeout(tool_timeout, tool.execute(params.clone(), &ctx)).await; - - match result { - Ok(Ok(output)) => serde_json::to_string_pretty(&output.result) - .map_err(|e| format!("serialization error: {}", e)), - Ok(Err(e)) => Err(e.to_string()), - Err(_) => Err("tool execution timed out".to_string()), - } - } - - /// Process a tool result into the reasoning context. Returns true if the job is complete. - fn process_result( - &self, - reason_ctx: &mut ReasoningContext, - selection: &ToolSelection, - result: Result, - ) -> bool { - match result { - Ok(output) => { - let sanitized = self - .safety - .sanitize_tool_output(&selection.tool_name, &output); - 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, - )); - - // Tool output should never signal job completion. Only the LLM's - // natural language response should decide when a job is done. A - // tool could return text containing "TASK_COMPLETE" in its output - // (e.g. from file contents) and trigger a false positive. - false - } - Err(e) => { - tracing::warn!("Tool {} failed: {}", selection.tool_name, e); - reason_ctx.messages.push(ChatMessage::tool_result( - &selection.tool_call_id, - &selection.tool_name, - format!("Error: {}", e), - )); - false - } - } - } - - /// Post a job event to the orchestrator (fire-and-forget). - async fn post_event(&self, event_type: &str, data: serde_json::Value) { - self.client - .post_event(&JobEventPayload { - event_type: event_type.to_string(), - data, - }) - .await; - } - - /// Poll the orchestrator for a follow-up prompt. If one is available, - /// inject it as a user message into the reasoning context. - async fn poll_and_inject_prompt(&self, reason_ctx: &mut ReasoningContext) { - match self.client.poll_prompt().await { - Ok(Some(prompt)) => { - tracing::info!( - "Received follow-up prompt: {}", - truncate(&prompt.content, 100) - ); - self.post_event( - "message", - serde_json::json!({ - "role": "user", - "content": truncate(&prompt.content, 2000), - }), - ) - .await; - reason_ctx.messages.push(ChatMessage::user(&prompt.content)); - } - Ok(None) => {} - Err(e) => { - tracing::debug!("Failed to poll for prompt: {}", e); - } - } - } -} - -fn truncate(s: &str, max: usize) -> String { - if s.len() <= max { - s.to_string() - } else { - let end = crate::util::floor_char_boundary(s, max); - format!("{}...", &s[..end]) - } -} - -#[cfg(test)] -mod tests { - use crate::worker::runtime::truncate; - - #[test] - fn test_truncate_within_limit() { - assert_eq!(truncate("hello", 10), "hello"); - } - - #[test] - fn test_truncate_at_limit() { - assert_eq!(truncate("hello", 5), "hello"); - } - - #[test] - fn test_truncate_beyond_limit() { - let result = truncate("hello world", 5); - assert_eq!(result, "hello..."); - } - - #[test] - fn test_truncate_multibyte_safe() { - // "é" is 2 bytes in UTF-8; slicing at byte 1 would panic without safety - let result = truncate("é is fancy", 1); - // Should truncate to 0 chars (can't fit "é" in 1 byte) - assert_eq!(result, "..."); - } -}