diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 6c8680d0..0bf1fd58 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -127,6 +127,9 @@ impl Agent { if let Some(ref tx) = deps.sse_tx { scheduler.set_sse_sender(tx.clone()); } + if let Some(ref interceptor) = deps.http_interceptor { + scheduler.set_http_interceptor(Arc::clone(interceptor)); + } let scheduler = Arc::new(scheduler); Self { diff --git a/src/agent/routine.rs b/src/agent/routine.rs index 7fa56d7d..4cf691be 100644 --- a/src/agent/routine.rs +++ b/src/agent/routine.rs @@ -175,6 +175,11 @@ pub enum RoutineAction { /// Max reasoning iterations (default: 10). #[serde(default = "default_max_iterations")] max_iterations: u32, + /// Tool names pre-authorized for `Always`-approval tools (e.g. destructive + /// shell commands, cross-channel messaging). `UnlessAutoApproved` tools are + /// automatically permitted in routine jobs without listing them here. + #[serde(default)] + tool_permissions: Vec, }, } @@ -186,6 +191,19 @@ fn default_max_iterations() -> u32 { 10 } +/// Parse a `tool_permissions` JSON array into a `Vec`. +pub fn parse_tool_permissions(value: &serde_json::Value) -> Vec { + value + .get("tool_permissions") + .and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default() +} + impl RoutineAction { /// The string tag stored in the DB action_type column. pub fn type_tag(&self) -> &'static str { @@ -248,10 +266,12 @@ impl RoutineAction { .and_then(|v| v.as_u64()) .unwrap_or(default_max_iterations() as u64) as u32; + let tool_permissions = parse_tool_permissions(&config); Ok(RoutineAction::FullJob { title, description, max_iterations, + tool_permissions, }) } other => Err(RoutineError::UnknownActionType { @@ -276,10 +296,12 @@ impl RoutineAction { title, description, max_iterations, + tool_permissions, } => serde_json::json!({ "title": title, "description": description, "max_iterations": max_iterations, + "tool_permissions": tool_permissions, }), } } @@ -450,12 +472,13 @@ mod tests { title: "Deploy review".to_string(), description: "Review and deploy pending changes".to_string(), max_iterations: 5, + tool_permissions: vec!["shell".to_string()], }; let json = action.to_config_json(); let parsed = RoutineAction::from_db("full_job", json).expect("parse full_job"); assert!( - matches!(parsed, RoutineAction::FullJob { title, max_iterations, .. } - if title == "Deploy review" && max_iterations == 5) + matches!(parsed, RoutineAction::FullJob { title, max_iterations, tool_permissions, .. } + if title == "Deploy review" && max_iterations == 5 && tool_permissions == vec!["shell".to_string()]) ); } diff --git a/src/agent/routine_engine.rs b/src/agent/routine_engine.rs index 75970c4f..bc5508d5 100644 --- a/src/agent/routine_engine.rs +++ b/src/agent/routine_engine.rs @@ -28,6 +28,7 @@ use crate::config::RoutineConfig; use crate::db::Database; use crate::error::RoutineError; use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider}; +use crate::tools::ApprovalContext; use crate::workspace::Workspace; /// The routine execution engine. @@ -180,6 +181,9 @@ impl RoutineEngine { } /// Fire a routine manually (from tool call or CLI). + /// + /// Bypasses cooldown checks (those only apply to cron/event triggers). + /// Still enforces enabled check and concurrent run limit. pub async fn fire_manual(&self, routine_id: Uuid) -> Result { let routine = self .store @@ -327,7 +331,19 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) title, description, max_iterations, - } => execute_full_job(&ctx, &routine, &run, title, description, *max_iterations).await, + tool_permissions, + } => { + execute_full_job( + &ctx, + &routine, + &run, + title, + description, + *max_iterations, + tool_permissions, + ) + .await + } }; // Decrement running count @@ -418,6 +434,7 @@ async fn execute_full_job( title: &str, description: &str, max_iterations: u32, + tool_permissions: &[String], ) -> Result<(RunStatus, Option, Option), RoutineError> { let scheduler = ctx .scheduler @@ -426,10 +443,31 @@ async fn execute_full_job( reason: "scheduler not available".to_string(), })?; + // Set the message tool's default channel/target from the routine's notify config + // so the LLM can send results without triggering cross-channel approval. + // TODO: This mutates shared global state and can race with concurrent jobs. + // Move notify config into JobContext metadata and apply per-job instead. + if let Some(channel) = &routine.notify.channel { + scheduler + .tools() + .set_message_tool_context(Some(channel.clone()), Some(routine.notify.user.clone())) + .await; + } + let metadata = serde_json::json!({ "max_iterations": max_iterations }); + // Build approval context: UnlessAutoApproved tools are auto-approved for routines; + // Always tools require explicit listing in tool_permissions. + let approval_context = ApprovalContext::autonomous_with_tools(tool_permissions.iter().cloned()); + let job_id = scheduler - .dispatch_job(&routine.user_id, title, description, Some(metadata)) + .dispatch_job_with_context( + &routine.user_id, + title, + description, + Some(metadata), + approval_context, + ) .await .map_err(|e| RoutineError::JobDispatchFailed { reason: format!("failed to dispatch job: {e}"), diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 17ffc644..99386d7f 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -18,7 +18,7 @@ use crate::error::{Error, JobError}; use crate::hooks::HookRegistry; use crate::llm::LlmProvider; use crate::safety::SafetyLayer; -use crate::tools::ToolRegistry; +use crate::tools::{ApprovalContext, ToolRegistry}; /// Message to send to a worker. #[derive(Debug)] @@ -56,6 +56,8 @@ pub struct Scheduler { hooks: Arc, /// SSE broadcast sender for live job event streaming. sse_tx: Option>, + /// HTTP interceptor for trace recording/replay (propagated to workers). + http_interceptor: Option>, /// Running jobs (main LLM-driven jobs). jobs: Arc>>, /// Running sub-tasks (tool executions, background tasks). @@ -82,6 +84,7 @@ impl Scheduler { store, hooks, sse_tx: None, + http_interceptor: None, jobs: Arc::new(RwLock::new(HashMap::new())), subtasks: Arc::new(RwLock::new(HashMap::new())), } @@ -92,6 +95,14 @@ impl Scheduler { self.sse_tx = Some(tx); } + /// Set the HTTP interceptor for trace recording/replay. + pub fn set_http_interceptor( + &mut self, + interceptor: Arc, + ) { + self.http_interceptor = Some(interceptor); + } + /// Create, persist, and schedule a job in one shot. /// /// This is the preferred entry point for dispatching new jobs. It: @@ -108,6 +119,41 @@ impl Scheduler { title: &str, description: &str, metadata: Option, + ) -> Result { + self.dispatch_job_inner(user_id, title, description, metadata, None) + .await + } + + /// Dispatch a job with an explicit approval context for autonomous execution. + /// + /// Same as `dispatch_job`, but the worker will use the given `ApprovalContext` + /// to determine which tools are pre-approved (instead of blocking all non-`Never` tools). + pub async fn dispatch_job_with_context( + &self, + user_id: &str, + title: &str, + description: &str, + metadata: Option, + approval_context: ApprovalContext, + ) -> Result { + self.dispatch_job_inner( + user_id, + title, + description, + metadata, + Some(approval_context), + ) + .await + } + + /// Shared implementation for `dispatch_job` and `dispatch_job_with_context`. + async fn dispatch_job_inner( + &self, + user_id: &str, + title: &str, + description: &str, + metadata: Option, + approval_context: Option, ) -> Result { let job_id = self .context_manager @@ -132,12 +178,21 @@ impl Scheduler { })?; } - self.schedule(job_id).await?; + self.schedule_with_context(job_id, approval_context).await?; Ok(job_id) } /// Schedule a job for execution. pub async fn schedule(&self, job_id: Uuid) -> Result<(), JobError> { + self.schedule_with_context(job_id, None).await + } + + /// Schedule a job with an optional approval context. + async fn schedule_with_context( + &self, + job_id: Uuid, + approval_context: Option, + ) -> Result<(), JobError> { // Hold write lock for the entire check-insert sequence to prevent // TOCTOU races where two concurrent calls both pass the checks. { @@ -181,6 +236,8 @@ impl Scheduler { timeout: self.config.job_timeout, use_planning: self.config.use_planning, sse_tx: self.sse_tx.clone(), + approval_context, + http_interceptor: self.http_interceptor.clone(), }; let worker = Worker::new(job_id, deps); @@ -257,11 +314,14 @@ impl Scheduler { let context_manager = self.context_manager.clone(); let safety = self.safety.clone(); + // TODO: propagate parent job's ApprovalContext here when subtasks + // are used in autonomous/routine paths (currently only used in tests). tokio::spawn(async move { let result = Self::execute_tool_task( tools, context_manager, safety, + None, tool_parent_id, &tool_name, params, @@ -390,6 +450,7 @@ impl Scheduler { tools: Arc, context_manager: Arc, safety: Arc, + approval_context: Option, job_id: Uuid, tool_name: &str, params: serde_json::Value, @@ -413,7 +474,10 @@ impl Scheduler { .into()); } - if tool.requires_approval(¶ms).is_required() { + let requirement = tool.requires_approval(¶ms); + let blocked = + ApprovalContext::is_blocked_or_default(&approval_context, tool_name, requirement); + if blocked { return Err(crate::error::ToolError::AuthRequired { name: tool_name.to_string(), } @@ -617,6 +681,11 @@ impl Scheduler { #[cfg(test)] mod tests { + use super::*; + use crate::config::SafetyConfig; + use crate::safety::SafetyLayer; + use crate::tools::{ApprovalRequirement, Tool, ToolError, ToolOutput}; + #[test] fn test_scheduler_creation() { // Would need to mock dependencies for proper testing @@ -627,4 +696,202 @@ mod tests { // This test would need mock dependencies. // For now just verify the empty case doesn't panic. } + + /// A tool that returns `UnlessAutoApproved`. + struct SoftApprovalTool; + + #[async_trait::async_trait] + impl Tool for SoftApprovalTool { + fn name(&self) -> &str { + "soft_gate" + } + fn description(&self) -> &str { + "needs soft approval" + } + 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::text( + "soft_ok", + std::time::Instant::now().elapsed(), + )) + } + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { + ApprovalRequirement::UnlessAutoApproved + } + fn requires_sanitization(&self) -> bool { + false + } + } + + /// A tool that returns `Always`. + struct HardApprovalTool; + + #[async_trait::async_trait] + impl Tool for HardApprovalTool { + fn name(&self) -> &str { + "hard_gate" + } + fn description(&self) -> &str { + "needs hard approval" + } + 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::text( + "hard_ok", + std::time::Instant::now().elapsed(), + )) + } + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { + ApprovalRequirement::Always + } + fn requires_sanitization(&self) -> bool { + false + } + } + + async fn setup_tools_and_job() -> ( + Arc, + Arc, + Arc, + Uuid, + ) { + let registry = ToolRegistry::new(); + registry.register(Arc::new(SoftApprovalTool)).await; + registry.register(Arc::new(HardApprovalTool)).await; + + let cm = Arc::new(ContextManager::new(5)); + let job_id = cm.create_job("test", "approval test").await.unwrap(); + cm.update_context(job_id, |ctx| ctx.transition_to(JobState::InProgress, None)) + .await + .unwrap() + .unwrap(); + + let safety = Arc::new(SafetyLayer::new(&SafetyConfig { + max_output_length: 100_000, + injection_check_enabled: false, + })); + + (Arc::new(registry), cm, safety, job_id) + } + + #[tokio::test] + async fn test_execute_tool_task_blocks_without_context() { + let (tools, cm, safety, job_id) = setup_tools_and_job().await; + + // Without approval context, UnlessAutoApproved is blocked + let result = Scheduler::execute_tool_task( + tools.clone(), + cm.clone(), + safety.clone(), + None, + job_id, + "soft_gate", + serde_json::json!({}), + ) + .await; + assert!( + result.is_err(), + "soft_gate should be blocked without context" + ); + + // Always is also blocked + let result = Scheduler::execute_tool_task( + tools, + cm, + safety, + None, + job_id, + "hard_gate", + serde_json::json!({}), + ) + .await; + assert!( + result.is_err(), + "hard_gate should be blocked without context" + ); + } + + #[tokio::test] + async fn test_execute_tool_task_autonomous_unblocks_soft() { + let (tools, cm, safety, job_id) = setup_tools_and_job().await; + + // Autonomous context auto-approves UnlessAutoApproved + let result = Scheduler::execute_tool_task( + tools.clone(), + cm.clone(), + safety.clone(), + Some(ApprovalContext::autonomous()), + job_id, + "soft_gate", + serde_json::json!({}), + ) + .await; + assert!( + result.is_ok(), + "soft_gate should pass with autonomous context" + ); + + // But still blocks Always + let result = Scheduler::execute_tool_task( + tools, + cm, + safety, + Some(ApprovalContext::autonomous()), + job_id, + "hard_gate", + serde_json::json!({}), + ) + .await; + assert!( + result.is_err(), + "hard_gate should still be blocked without explicit permission" + ); + } + + #[tokio::test] + async fn test_execute_tool_task_autonomous_with_permissions() { + let (tools, cm, safety, job_id) = setup_tools_and_job().await; + + // Autonomous context with explicit permission for hard_gate + let ctx = ApprovalContext::autonomous_with_tools(["hard_gate".to_string()]); + + let result = Scheduler::execute_tool_task( + tools.clone(), + cm.clone(), + safety.clone(), + Some(ctx.clone()), + job_id, + "soft_gate", + serde_json::json!({}), + ) + .await; + assert!(result.is_ok(), "soft_gate should pass"); + + let result = Scheduler::execute_tool_task( + tools, + cm, + safety, + Some(ctx), + job_id, + "hard_gate", + serde_json::json!({}), + ) + .await; + assert!( + result.is_ok(), + "hard_gate should pass with explicit permission" + ); + } } diff --git a/src/agent/worker.rs b/src/agent/worker.rs index f5aa32b3..e3fa11e7 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -19,7 +19,7 @@ use crate::llm::{ }; use crate::safety::SafetyLayer; use crate::tools::rate_limiter::RateLimitResult; -use crate::tools::{ToolRegistry, redact_params}; +use crate::tools::{ApprovalContext, ToolRegistry, redact_params}; /// Shared dependencies for worker execution. /// @@ -37,6 +37,12 @@ pub struct WorkerDeps { pub use_planning: bool, /// SSE broadcast sender for live job event streaming to the web gateway. pub sse_tx: Option>, + /// Approval context for tool execution. When `None`, all non-`Never` tools are + /// blocked (legacy behavior). When `Some`, the context determines which tools + /// are pre-approved for autonomous execution. + pub approval_context: Option, + /// HTTP interceptor for trace recording/replay (propagated to JobContext). + pub http_interceptor: Option>, } /// Worker that executes a single job. @@ -246,6 +252,9 @@ Report when the job is complete or if you encounter issues you cannot resolve."# // Already in a terminal state (e.g. execution_loop // called mark_completed itself). } + Ok(JobState::Completed) => { + // execution_loop already called mark_completed. + } Ok(JobState::Stuck) => { // execution_loop marked this as stuck (e.g. "plan // completed but work remains"); leave for self-repair. @@ -353,11 +362,13 @@ Report when the job is complete or if you encounter issues you cannot resolve."# if let Some(ref plan) = plan { self.execute_plan(rx, reasoning, reason_ctx, plan).await?; - // If the plan marked the job terminal, we're done. Only fall - // through to the direct selection loop if the plan was - // interrupted or explicitly left the job in-progress. + // If the plan marked the job completed, terminal, or stuck, we're + // done. Only fall through to the direct selection loop if the + // plan was interrupted or explicitly left the job in-progress. if let Ok(ctx) = self.context_manager().get_context(self.job_id).await - && (ctx.state.is_terminal() || ctx.state == JobState::Stuck) + && (ctx.state.is_terminal() + || ctx.state == JobState::Stuck + || ctx.state == JobState::Completed) { return Ok(()); } @@ -671,8 +682,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."# name: tool_name.to_string(), })?; - // Tools requiring approval are blocked in autonomous jobs - if tool.requires_approval(params).is_required() { + // Check approval: use context-aware check if available, else block all non-Never tools + let requirement = tool.requires_approval(params); + let blocked = + ApprovalContext::is_blocked_or_default(&deps.approval_context, tool_name, requirement); + if blocked { return Err(crate::error::ToolError::AuthRequired { name: tool_name.to_string(), } @@ -680,7 +694,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } // Fetch job context early so we have the real user_id for hooks and rate limiting - let job_ctx = deps.context_manager.get_context(job_id).await?; + let mut job_ctx = deps.context_manager.get_context(job_id).await?; + // Propagate http_interceptor for trace recording/replay + if job_ctx.http_interceptor.is_none() { + job_ctx.http_interceptor = deps.http_interceptor.clone(); + } // Check per-tool rate limit before running hooks or executing (cheaper check first) if let Some(config) = tool.rate_limit_config() @@ -1298,6 +1316,8 @@ mod tests { timeout: Duration::from_secs(30), use_planning: false, sse_tx: None, + approval_context: None, + http_interceptor: None, }; Worker::new(job_id, deps) @@ -1496,4 +1516,199 @@ mod tests { "Missing tool should produce an error, not a panic" ); } + + /// Verify that calling mark_completed on an already-Completed job returns + /// an error (Completed → Completed is an invalid state transition). + #[tokio::test] + async fn test_mark_completed_twice_returns_error() { + let worker = make_worker(vec![]).await; + + // Transition to InProgress first (required by state machine) + worker + .context_manager() + .update_context(worker.job_id, |ctx| { + ctx.transition_to(JobState::InProgress, None) + }) + .await + .unwrap() + .unwrap(); + + // First mark_completed should succeed + worker.mark_completed().await.unwrap(); + + // Verify state is Completed + let ctx = worker + .context_manager() + .get_context(worker.job_id) + .await + .unwrap(); + assert_eq!(ctx.state, JobState::Completed); + + // Second mark_completed should fail (Completed → Completed is invalid) + let result = worker.mark_completed().await; + assert!( + result.is_err(), + "Completed → Completed transition should be rejected by state machine" + ); + } + + /// Build a Worker with the given approval context. + async fn make_worker_with_approval( + tools: Vec>, + approval_context: Option, + ) -> Worker { + let registry = ToolRegistry::new(); + for t in tools { + registry.register(t).await; + } + + let cm = Arc::new(crate::context::ContextManager::new(5)); + let job_id = cm.create_job("test", "test job").await.unwrap(); + + let deps = WorkerDeps { + context_manager: cm, + llm: Arc::new(StubLlm), + safety: Arc::new(SafetyLayer::new(&SafetyConfig { + max_output_length: 100_000, + injection_check_enabled: false, + })), + tools: Arc::new(registry), + store: None, + hooks: Arc::new(crate::hooks::HookRegistry::new()), + timeout: Duration::from_secs(30), + use_planning: false, + sse_tx: None, + approval_context, + http_interceptor: None, + }; + + Worker::new(job_id, deps) + } + + /// A tool that requires approval (UnlessAutoApproved). + struct ApprovalTool; + + #[async_trait::async_trait] + impl Tool for ApprovalTool { + fn name(&self) -> &str { + "needs_approval" + } + fn description(&self) -> &str { + "Tool requiring approval" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) + } + async fn execute( + &self, + _params: serde_json::Value, + _ctx: &crate::context::JobContext, + ) -> Result { + Ok(ToolOutput::text( + "approved", + std::time::Instant::now().elapsed(), + )) + } + fn requires_approval( + &self, + _params: &serde_json::Value, + ) -> crate::tools::ApprovalRequirement { + crate::tools::ApprovalRequirement::UnlessAutoApproved + } + fn requires_sanitization(&self) -> bool { + false + } + } + + /// A tool that always requires approval. + struct AlwaysApprovalTool; + + #[async_trait::async_trait] + impl Tool for AlwaysApprovalTool { + fn name(&self) -> &str { + "always_approval" + } + fn description(&self) -> &str { + "Tool always requiring approval" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {}}) + } + async fn execute( + &self, + _params: serde_json::Value, + _ctx: &crate::context::JobContext, + ) -> Result { + Ok(ToolOutput::text( + "always", + std::time::Instant::now().elapsed(), + )) + } + fn requires_approval( + &self, + _params: &serde_json::Value, + ) -> crate::tools::ApprovalRequirement { + crate::tools::ApprovalRequirement::Always + } + fn requires_sanitization(&self) -> bool { + false + } + } + + #[tokio::test] + async fn test_approval_context_unblocks_unless_auto_approved() { + // Without approval context, UnlessAutoApproved is blocked + let worker_blocked = make_worker_with_approval(vec![Arc::new(ApprovalTool)], None).await; + let result = worker_blocked + .execute_tool("needs_approval", &serde_json::json!({})) + .await; + assert!( + result.is_err(), + "Should be blocked without approval context" + ); + + // With autonomous approval context, UnlessAutoApproved is allowed + let worker_allowed = make_worker_with_approval( + vec![Arc::new(ApprovalTool)], + Some(crate::tools::ApprovalContext::autonomous()), + ) + .await; + let result = worker_allowed + .execute_tool("needs_approval", &serde_json::json!({})) + .await; + assert!(result.is_ok(), "Should be allowed with autonomous context"); + } + + #[tokio::test] + async fn test_approval_context_blocks_always_unless_permitted() { + // Autonomous context without tool_permissions blocks Always tools + let worker_blocked = make_worker_with_approval( + vec![Arc::new(AlwaysApprovalTool)], + Some(crate::tools::ApprovalContext::autonomous()), + ) + .await; + let result = worker_blocked + .execute_tool("always_approval", &serde_json::json!({})) + .await; + assert!( + result.is_err(), + "Always tool should be blocked without permission" + ); + + // Autonomous context with tool_permissions allows Always tools + let worker_allowed = make_worker_with_approval( + vec![Arc::new(AlwaysApprovalTool)], + Some(crate::tools::ApprovalContext::autonomous_with_tools([ + "always_approval".to_string(), + ])), + ) + .await; + let result = worker_allowed + .execute_tool("always_approval", &serde_json::json!({})) + .await; + assert!( + result.is_ok(), + "Always tool should be allowed with permission" + ); + } } diff --git a/src/tools/builtin/message.rs b/src/tools/builtin/message.rs index 532b41e4..9e37da6c 100644 --- a/src/tools/builtin/message.rs +++ b/src/tools/builtin/message.rs @@ -207,26 +207,10 @@ impl Tool for MessageTool { } } - fn requires_approval(&self, params: &serde_json::Value) -> ApprovalRequirement { - // Require approval when sending to a different channel than the default - // (cross-channel messages are more sensitive) - let param_channel = params.get("channel").and_then(|v| v.as_str()); - if let Some(channel) = param_channel { - // Check if it differs from the default channel - let default_channel = self - .default_channel - .read() - .unwrap_or_else(|e| e.into_inner()); - if let Some(default) = default_channel.as_ref() - && channel != default - { - return ApprovalRequirement::Always; - } - // No default set - require approval for explicit channel selection - return ApprovalRequirement::Always; - } - // No channel specified in params - uses default, less risky - ApprovalRequirement::UnlessAutoApproved + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { + // Message tool only delivers to channels the user has configured + // (TUI, Telegram, Slack, web gateway, etc.) via ChannelManager::broadcast. + ApprovalRequirement::Never } fn rate_limit_config(&self) -> Option { @@ -533,53 +517,17 @@ mod tests { ); } - // ── Multi-thread runtime safety tests ───────────────────────────── - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn requires_approval_no_channel_multi_thread() { + #[test] + fn requires_approval_always_never() { + // Message tool only sends to user-owned channels, so never needs approval. let tool = MessageTool::new(Arc::new(ChannelManager::new())); - // No channel set, no channel param - should not panic in multi-thread runtime - let result = tool.requires_approval(&serde_json::json!({"content": "hello"})); - assert_eq!(result, ApprovalRequirement::UnlessAutoApproved); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn requires_approval_with_context_multi_thread() { - let tool = MessageTool::new(Arc::new(ChannelManager::new())); - tool.set_context(Some("signal".to_string()), Some("+1234567890".to_string())) - .await; - - // No channel param - uses default, less risky - let result = tool.requires_approval(&serde_json::json!({"content": "hello"})); - assert_eq!(result, ApprovalRequirement::UnlessAutoApproved); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn requires_approval_cross_channel_multi_thread() { - let tool = MessageTool::new(Arc::new(ChannelManager::new())); - tool.set_context(Some("signal".to_string()), Some("+1234567890".to_string())) - .await; - - // Different channel than default requires approval - let result = tool.requires_approval(&serde_json::json!({ - "content": "hello", - "channel": "telegram" - })); - assert_eq!(result, ApprovalRequirement::Always); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn requires_approval_same_channel_explicit_multi_thread() { - let tool = MessageTool::new(Arc::new(ChannelManager::new())); - tool.set_context(Some("signal".to_string()), Some("+1234567890".to_string())) - .await; - - // Explicit channel that matches default still returns Always - // (existing behavior: any explicit channel param triggers Always) - let result = tool.requires_approval(&serde_json::json!({ - "content": "hello", - "channel": "signal" - })); - assert_eq!(result, ApprovalRequirement::Always); + assert_eq!( + tool.requires_approval(&serde_json::json!({"content": "hello"})), + ApprovalRequirement::Never, + ); + assert_eq!( + tool.requires_approval(&serde_json::json!({"content": "hi", "channel": "telegram"})), + ApprovalRequirement::Never, + ); } } diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index 4931e5b8..23f170f9 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -32,7 +32,8 @@ pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTo pub use message::MessageTool; pub use restart::RestartTool; pub use routine::{ - RoutineCreateTool, RoutineDeleteTool, RoutineHistoryTool, RoutineListTool, RoutineUpdateTool, + RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool, RoutineListTool, + RoutineUpdateTool, }; pub use secrets_tools::{SecretDeleteTool, SecretListTool}; pub use shell::ShellTool; diff --git a/src/tools/builtin/routine.rs b/src/tools/builtin/routine.rs index 6a0abce9..59a57e0c 100644 --- a/src/tools/builtin/routine.rs +++ b/src/tools/builtin/routine.rs @@ -1,10 +1,11 @@ //! LLM-facing tools for managing routines. //! -//! Five tools let the agent manage routines conversationally: +//! Six tools let the agent manage routines conversationally: //! - `routine_create` - Create a new routine //! - `routine_list` - List all routines with status //! - `routine_update` - Modify or toggle a routine //! - `routine_delete` - Remove a routine +//! - `routine_fire` - Manually trigger a routine //! - `routine_history` - View past runs use std::sync::Arc; @@ -20,7 +21,7 @@ use crate::agent::routine::{ use crate::agent::routine_engine::RoutineEngine; use crate::context::JobContext; use crate::db::Database; -use crate::tools::tool::{Tool, ToolError, ToolOutput, require_str}; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput, require_str}; // ==================== routine_create ==================== @@ -93,6 +94,19 @@ impl Tool for RoutineCreateTool { "cooldown_secs": { "type": "integer", "description": "Minimum seconds between fires (default: 300)" + }, + "tool_permissions": { + "type": "array", + "items": { "type": "string" }, + "description": "Tool names pre-authorized for Always-approval tools in full_job mode (e.g. ['shell']). UnlessAutoApproved tools are automatically permitted in routines." + }, + "notify_channel": { + "type": "string", + "description": "Channel to send results to (e.g. 'telegram', 'slack', 'tui'). Sets the default channel for message tool calls in routine jobs." + }, + "notify_user": { + "type": "string", + "description": "User/target to notify (e.g. username, chat ID). Defaults to 'default'." } }, "required": ["name", "trigger_type", "prompt"] @@ -192,11 +206,15 @@ impl Tool for RoutineCreateTool { context_paths, max_tokens: 4096, }, - "full_job" => RoutineAction::FullJob { - title: name.to_string(), - description: prompt.to_string(), - max_iterations: 10, - }, + "full_job" => { + let tool_permissions = crate::agent::routine::parse_tool_permissions(¶ms); + RoutineAction::FullJob { + title: name.to_string(), + description: prompt.to_string(), + max_iterations: 10, + tool_permissions, + } + } other => { return Err(ToolError::InvalidParameters(format!( "unknown action_type: {other}" @@ -229,7 +247,18 @@ impl Tool for RoutineCreateTool { max_concurrent: 1, dedup_window: None, }, - notify: NotifyConfig::default(), + notify: NotifyConfig { + channel: params + .get("notify_channel") + .and_then(|v| v.as_str()) + .map(String::from), + user: params + .get("notify_user") + .and_then(|v| v.as_str()) + .unwrap_or("default") + .to_string(), + ..NotifyConfig::default() + }, last_run_at: None, next_fire_at: next_fire, run_count: 0, @@ -533,6 +562,82 @@ impl Tool for RoutineDeleteTool { } } +// ==================== routine_fire ==================== + +pub struct RoutineFireTool { + store: Arc, + engine: Arc, +} + +impl RoutineFireTool { + pub fn new(store: Arc, engine: Arc) -> Self { + Self { store, engine } + } +} + +#[async_trait] +impl Tool for RoutineFireTool { + fn name(&self) -> &str { + "routine_fire" + } + + fn description(&self) -> &str { + "Manually trigger a routine to run immediately, bypassing schedule, trigger type, and cooldown." + } + + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { + // Firing a routine can dispatch a full_job with pre-authorized Always-gated tools, + // so this is a meaningful escalation that warrants auto-approval gating. + ApprovalRequirement::UnlessAutoApproved + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Name of the routine to fire" + } + }, + "required": ["name"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let name = require_str(¶ms, "name")?; + + let routine = self + .store + .get_routine_by_name(&ctx.user_id, name) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("DB error: {e}")))? + .ok_or_else(|| ToolError::ExecutionFailed(format!("routine '{}' not found", name)))?; + + let run_id = self.engine.fire_manual(routine.id).await.map_err(|e| { + ToolError::ExecutionFailed(format!("failed to fire routine '{}': {e}", name)) + })?; + + let result = serde_json::json!({ + "name": name, + "run_id": run_id.to_string(), + "status": "fired", + }); + + Ok(ToolOutput::success(result, start.elapsed())) + } + + fn requires_sanitization(&self) -> bool { + false + } +} + // ==================== routine_history ==================== pub struct RoutineHistoryTool { diff --git a/src/tools/mod.rs b/src/tools/mod.rs index cd225bd1..d379d474 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -25,6 +25,6 @@ pub use builder::{ pub use rate_limiter::RateLimiter; pub use registry::ToolRegistry; pub use tool::{ - ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, ToolRateLimitConfig, - redact_params, validate_tool_schema, + ApprovalContext, ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, + ToolRateLimitConfig, redact_params, validate_tool_schema, }; diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 5809305e..4f98a30b 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -63,6 +63,7 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[ "routine_list", "routine_update", "routine_delete", + "routine_fire", "routine_history", "skill_list", "skill_search", @@ -423,8 +424,8 @@ impl ToolRegistry { engine: Arc, ) { use crate::tools::builtin::{ - RoutineCreateTool, RoutineDeleteTool, RoutineHistoryTool, RoutineListTool, - RoutineUpdateTool, + RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool, + RoutineListTool, RoutineUpdateTool, }; self.register_sync(Arc::new(RoutineCreateTool::new( Arc::clone(&store), @@ -439,8 +440,12 @@ impl ToolRegistry { Arc::clone(&store), Arc::clone(&engine), ))); + self.register_sync(Arc::new(RoutineFireTool::new( + Arc::clone(&store), + Arc::clone(&engine), + ))); self.register_sync(Arc::new(RoutineHistoryTool::new(store))); - tracing::info!("Registered 5 routine management tools"); + tracing::info!("Registered 6 routine management tools"); } /// Register message tool for sending messages to channels. diff --git a/src/tools/schema_validator.rs b/src/tools/schema_validator.rs index f4aa0968..8da0b613 100644 --- a/src/tools/schema_validator.rs +++ b/src/tools/schema_validator.rs @@ -582,7 +582,14 @@ mod tests { "enum": ["lightweight", "full_job"], "description": "Execution mode" }, - "cooldown_secs": { "type": "integer", "description": "Min seconds between fires" } + "cooldown_secs": { "type": "integer", "description": "Min seconds between fires" }, + "tool_permissions": { + "type": "array", + "items": { "type": "string" }, + "description": "Pre-authorized tools for full_job mode" + }, + "notify_channel": { "type": "string", "description": "Channel for message tool" }, + "notify_user": { "type": "string", "description": "User/target to notify" } }, "required": ["name", "trigger_type", "prompt"] }), @@ -619,6 +626,16 @@ mod tests { "required": ["name"] }), ), + ( + "routine_fire", + serde_json::json!({ + "type": "object", + "properties": { + "name": { "type": "string", "description": "Routine name" } + }, + "required": ["name"] + }), + ), ( "routine_history", serde_json::json!({ diff --git a/src/tools/tool.rs b/src/tools/tool.rs index 78728cf3..2e1b5183 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -28,6 +28,62 @@ impl ApprovalRequirement { } } +/// Approval context for autonomous tool execution (routines, background jobs). +/// +/// Interactive sessions don't use this type — they rely on session-level +/// auto-approve lists managed by the UI. This enum models only the autonomous +/// case where no interactive user is present. +#[derive(Debug, Clone)] +pub enum ApprovalContext { + /// Autonomous job with no interactive user. `UnlessAutoApproved` tools are + /// pre-approved. `Always` tools are blocked unless listed in `allowed_tools`. + Autonomous { + /// Tool names that are pre-authorized even for `Always` approval. + allowed_tools: std::collections::HashSet, + }, +} + +impl ApprovalContext { + /// Create an autonomous context with no extra tool permissions. + pub fn autonomous() -> Self { + Self::Autonomous { + allowed_tools: std::collections::HashSet::new(), + } + } + + /// Create an autonomous context with specific tools pre-authorized. + pub fn autonomous_with_tools(tools: impl IntoIterator) -> Self { + Self::Autonomous { + allowed_tools: tools.into_iter().collect(), + } + } + + /// Check whether a tool invocation is blocked in this context. + pub fn is_blocked(&self, tool_name: &str, requirement: ApprovalRequirement) -> bool { + match self { + Self::Autonomous { allowed_tools } => match requirement { + ApprovalRequirement::Never => false, + ApprovalRequirement::UnlessAutoApproved => false, + ApprovalRequirement::Always => !allowed_tools.contains(tool_name), + }, + } + } + + /// Check whether a tool is blocked given an optional context. + /// + /// When `None`, falls back to legacy behavior: all non-`Never` tools are blocked. + pub fn is_blocked_or_default( + context: &Option, + tool_name: &str, + requirement: ApprovalRequirement, + ) -> bool { + match context { + Some(ctx) => ctx.is_blocked(tool_name, requirement), + None => requirement.is_required(), + } + } +} + /// Per-tool rate limit configuration for built-in tool invocations. /// /// Controls how many times a tool can be invoked per user, per time window. @@ -733,4 +789,69 @@ mod tests { assert!(errors[0].contains("headers.items")); assert!(errors[0].contains("\"missing_field\"")); } + + #[test] + fn test_approval_context_autonomous_allows_unless_auto_approved() { + let ctx = ApprovalContext::autonomous(); + assert!(!ctx.is_blocked("shell", ApprovalRequirement::Never)); + assert!(!ctx.is_blocked("shell", ApprovalRequirement::UnlessAutoApproved)); + assert!(ctx.is_blocked("shell", ApprovalRequirement::Always)); + } + + #[test] + fn test_approval_context_autonomous_with_tools_allows_always() { + let ctx = + ApprovalContext::autonomous_with_tools(["shell".to_string(), "message".to_string()]); + assert!(!ctx.is_blocked("shell", ApprovalRequirement::Always)); + assert!(!ctx.is_blocked("message", ApprovalRequirement::Always)); + assert!(ctx.is_blocked("http", ApprovalRequirement::Always)); + } + + #[test] + fn test_approval_context_never_is_not_blocked() { + let ctx = ApprovalContext::autonomous(); + assert!(!ctx.is_blocked("any_tool", ApprovalRequirement::Never)); + } + + #[test] + fn test_is_blocked_or_default_with_none_uses_legacy() { + // None context: all non-Never tools are blocked + assert!(!ApprovalContext::is_blocked_or_default( + &None, + "any", + ApprovalRequirement::Never + )); + assert!(ApprovalContext::is_blocked_or_default( + &None, + "any", + ApprovalRequirement::UnlessAutoApproved + )); + assert!(ApprovalContext::is_blocked_or_default( + &None, + "any", + ApprovalRequirement::Always + )); + } + + #[test] + fn test_is_blocked_or_default_with_some_delegates() { + let ctx = Some(ApprovalContext::autonomous_with_tools( + ["shell".to_string()], + )); + assert!(!ApprovalContext::is_blocked_or_default( + &ctx, + "shell", + ApprovalRequirement::Always + )); + assert!(ApprovalContext::is_blocked_or_default( + &ctx, + "other", + ApprovalRequirement::Always + )); + assert!(!ApprovalContext::is_blocked_or_default( + &ctx, + "any", + ApprovalRequirement::UnlessAutoApproved + )); + } } diff --git a/tests/e2e_advanced_traces.rs b/tests/e2e_advanced_traces.rs index cd9d0326..92dd81f4 100644 --- a/tests/e2e_advanced_traces.rs +++ b/tests/e2e_advanced_traces.rs @@ -252,7 +252,123 @@ mod advanced { } // ----------------------------------------------------------------------- - // 6. Prompt injection resilience + // 6. Routine news digest (end-to-end: create, fire, verify message) + // + // Exercises the full routine execution stack: + // routine_create → routine_fire → RoutineEngine::fire_manual → + // Scheduler::dispatch_job_with_context → Worker (autonomous) → + // http + memory_write + message (broadcast to test channel) + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn routine_news_digest() { + use ironclaw::llm::recording::{HttpExchange, HttpExchangeRequest, HttpExchangeResponse}; + + let trace = LlmTrace::from_file(format!("{FIXTURES}/routine_news_digest.json")).unwrap(); + + // Mock HTTP response for the news API call made by the routine worker. + let http_exchanges = vec![HttpExchange { + request: HttpExchangeRequest { + method: "GET".to_string(), + url: "https://news-api.example.com/v1/tech/headlines".to_string(), + headers: Vec::new(), + body: None, + }, + response: HttpExchangeResponse { + status: 200, + headers: vec![( + "content-type".to_string(), + "application/json".to_string(), + )], + body: serde_json::json!({ + "headlines": [ + {"title": "Rust 2026 Edition", "summary": "async closures, generator syntax"}, + {"title": "WASM Component Model 1.0", "summary": "cross-language interop"}, + {"title": "NEAR AI Agent Framework", "summary": "on-chain identity"} + ] + }) + .to_string(), + }, + }]; + + let rig = TestRigBuilder::new() + .with_trace(trace.clone()) + .with_routines() + .with_http_exchanges(http_exchanges) + .build() + .await; + + // Turn 1: Create the routine (manual trigger, full_job, message+http pre-authorized). + rig.send_message( + "Set up a morning tech news routine with manual trigger \ + and full_job mode. Pre-authorize the message and http tools.", + ) + .await; + let r1 = rig.wait_for_responses(1, TIMEOUT).await; + assert!(!r1.is_empty(), "Turn 1: no response"); + let t1 = r1[0].content.to_lowercase(); + assert!( + t1.contains("routine") || t1.contains("created"), + "Turn 1: expected routine/created, got: {t1}" + ); + + // Turn 2: Fire the routine. This dispatches a full_job through the scheduler. + // The routine worker runs autonomously and consumes TraceLlm steps for + // http, memory_write, and message tool calls. The http tool uses the + // ReplayingHttpInterceptor to return the mock news API response. + rig.send_message("Fire it now.").await; + + // Wait for: + // - response 2: main conversation reply ("fired the routine") + // - response 3: message tool broadcast from routine worker ("Tech News Digest: ...") + // The routine worker runs asynchronously, so we wait for 3 total responses. + let responses = rig.wait_for_responses(3, Duration::from_secs(15)).await; + + // Find the main conversation reply (from turn 2) by content, since + // the routine worker runs asynchronously and may interleave messages. + let fire_reply = responses.iter().find(|r| { + let c = r.content.to_lowercase(); + c.contains("fired") || c.contains("running") + }); + assert!( + fire_reply.is_some(), + "Turn 2: expected fired/running, got: {:?}", + responses.iter().map(|r| &r.content).collect::>() + ); + + // The routine worker runs autonomously: http → memory_write → message. + // The message tool broadcasts to the test channel, proving the full + // chain executed successfully (including ApprovalContext allowing the + // http and message tools in autonomous mode). + let message_broadcast = responses.iter().find(|r| { + r.content.contains("Tech News Digest") + || r.content.contains("Rust 2026") + || r.content.contains("WASM Component Model") + }); + assert!( + message_broadcast.is_some(), + "Routine worker should have broadcast a message. Got: {:?}", + responses.iter().map(|r| &r.content).collect::>() + ); + + // Verify main conversation tools were called. + let started = rig.tool_calls_started(); + for tool in &["routine_create", "routine_fire"] { + assert!( + started.iter().any(|s| s == *tool), + "{tool} not called: {started:?}" + ); + } + + // Main conversation tools should have succeeded. + let completed = rig.tool_calls_completed(); + crate::support::assertions::assert_all_tools_succeeded(&completed); + + rig.shutdown(); + } + + // ----------------------------------------------------------------------- + // 7. Prompt injection resilience // ----------------------------------------------------------------------- #[tokio::test] diff --git a/tests/fixtures/llm_traces/advanced/routine_news_digest.json b/tests/fixtures/llm_traces/advanced/routine_news_digest.json new file mode 100644 index 00000000..4c98b49f --- /dev/null +++ b/tests/fixtures/llm_traces/advanced/routine_news_digest.json @@ -0,0 +1,140 @@ +{ + "model_name": "advanced-routine-news-digest", + "expects": { + "tools_used": ["routine_create", "routine_fire", "http", "memory_write", "message"], + "all_tools_succeeded": true, + "min_responses": 2 + }, + "turns": [ + { + "user_input": "Set up a morning tech news routine with manual trigger and full_job mode. Pre-authorize the message and http tools.", + "steps": [ + { + "request_hint": { "last_user_message_contains": "routine" }, + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_routine_1", + "name": "routine_create", + "arguments": { + "name": "morning-tech-news", + "description": "Fetch tech news via HTTP, write digest to memory, send summary", + "trigger_type": "manual", + "prompt": "Fetch the latest tech news from the API, write a digest to workspace memory, then send a summary message to the user.", + "action_type": "full_job", + "tool_permissions": ["message", "http"], + "cooldown_secs": 60, + "notify_channel": "test", + "notify_user": "default" + } + } + ], + "input_tokens": 120, + "output_tokens": 60 + } + }, + { + "response": { + "type": "text", + "content": "Created the **morning-tech-news** routine with manual trigger and full_job mode. The `message` and `http` tools are pre-authorized.", + "input_tokens": 200, + "output_tokens": 50 + } + } + ] + }, + { + "user_input": "Fire it now.", + "steps": [ + { + "request_hint": { "last_user_message_contains": "Fire" }, + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_fire_1", + "name": "routine_fire", + "arguments": { + "name": "morning-tech-news" + } + } + ], + "input_tokens": 250, + "output_tokens": 30 + } + }, + { + "response": { + "type": "text", + "content": "Fired the **morning-tech-news** routine. The job is running now.", + "input_tokens": 300, + "output_tokens": 40 + } + }, + { + "_comment": "Steps below are consumed by the routine worker (spawned async by routine_fire). The worker hits the same TraceLlm sequentially.", + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_rw_http", + "name": "http", + "arguments": { + "method": "GET", + "url": "https://news-api.example.com/v1/tech/headlines" + } + } + ], + "input_tokens": 100, + "output_tokens": 20 + } + }, + { + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_rw_mw", + "name": "memory_write", + "arguments": { + "content": "# Tech News Digest - 2026-03-05\n\n1. **Rust 2026 Edition** - async closures, generator syntax\n2. **WASM Component Model 1.0** - cross-language interop\n3. **NEAR AI Agent Framework** - on-chain identity", + "target": "routines/morning-tech-news/digest-2026-03-05.md", + "append": false + } + } + ], + "input_tokens": 150, + "output_tokens": 50 + } + }, + { + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_rw_msg", + "name": "message", + "arguments": { + "content": "Tech News Digest:\n- Rust 2026 Edition released\n- WASM Component Model 1.0 finalized\n- NEAR AI Agent Framework launched", + "channel": "test", + "target": "default" + } + } + ], + "input_tokens": 200, + "output_tokens": 30 + } + }, + { + "response": { + "type": "text", + "content": "Done. Digest written and summary sent.", + "input_tokens": 250, + "output_tokens": 20 + } + } + ] + } + ] +} diff --git a/tests/support/test_rig.rs b/tests/support/test_rig.rs index 4b1939f9..430d9182 100644 --- a/tests/support/test_rig.rs +++ b/tests/support/test_rig.rs @@ -27,6 +27,8 @@ use crate::support::metrics::{ToolInvocation, TraceMetrics}; use crate::support::test_channel::TestChannel; use crate::support::trace_llm::{LlmTrace, TraceLlm}; +use ironclaw::llm::recording::{HttpExchange, ReplayingHttpInterceptor}; + // --------------------------------------------------------------------------- // TestChannelHandle -- wraps Arc as Box // --------------------------------------------------------------------------- @@ -362,6 +364,8 @@ pub struct TestRigBuilder { llm: Option>, max_tool_iterations: usize, injection_check: bool, + enable_routines: bool, + http_exchanges: Vec, extra_tools: Vec>, } @@ -373,6 +377,8 @@ impl TestRigBuilder { llm: None, max_tool_iterations: 10, injection_check: false, + enable_routines: false, + http_exchanges: Vec::new(), extra_tools: Vec::new(), } } @@ -411,6 +417,23 @@ impl TestRigBuilder { self } + /// Enable the routines system so the scheduler is wired with a `RoutineEngine`, + /// allowing routine jobs to actually execute. Routine tools are always registered + /// but require the engine to dispatch jobs. + pub fn with_routines(mut self) -> Self { + self.enable_routines = true; + self + } + + /// Add pre-recorded HTTP exchanges for the `ReplayingHttpInterceptor`. + /// + /// When set, all `http` tool calls will return these responses in order + /// instead of making real network requests. + pub fn with_http_exchanges(mut self, exchanges: Vec) -> Self { + self.http_exchanges = exchanges; + self + } + /// Build the test rig, creating a real agent and spawning it in the background. /// /// Uses `AppBuilder::build_all()` to get the same component set as the real @@ -422,6 +445,17 @@ impl TestRigBuilder { use ironclaw::channels::ChannelManager; use ironclaw::db::libsql::LibSqlBackend; + // Destructure self up front to avoid partial-move issues. + let TestRigBuilder { + trace, + llm, + max_tool_iterations, + injection_check, + enable_routines, + http_exchanges: explicit_http_exchanges, + extra_tools, + } = self; + // 1. Create temp dir + libSQL database + run migrations. let temp_dir = tempfile::tempdir().expect("failed to create temp dir"); let db_path = temp_dir.path().join("test_rig.db"); @@ -440,24 +474,23 @@ impl TestRigBuilder { let _ = std::fs::create_dir_all(&skills_dir); let _ = std::fs::create_dir_all(&installed_skills_dir); let mut config = Config::for_testing(db_path, skills_dir, installed_skills_dir); - config.agent.max_tool_iterations = self.max_tool_iterations; - config.safety.injection_check_enabled = self.injection_check; + config.agent.max_tool_iterations = max_tool_iterations; + config.safety.injection_check_enabled = injection_check; // 3. Create SessionManager + LogBroadcaster. let session = Arc::new(SessionManager::new(SessionConfig::default())); let log_broadcaster = Arc::new(LogBroadcaster::new()); // 4. Create TraceLlm + InstrumentedLlm, extract HTTP exchanges for replay. - let http_exchanges = self - .trace + let trace_http_exchanges = trace .as_ref() .map(|t| t.http_exchanges.clone()) .unwrap_or_default(); let mut trace_llm_ref: Option> = None; - let base_llm: Arc = if let Some(llm) = self.llm { + let base_llm: Arc = if let Some(llm) = llm { llm - } else if let Some(trace) = self.trace { + } else if let Some(trace) = trace { let tlm = Arc::new(TraceLlm::from_trace(trace)); trace_llm_ref = Some(Arc::clone(&tlm)); tlm @@ -536,7 +569,7 @@ impl TestRigBuilder { } // Register any extra test-specific tools. - for tool in self.extra_tools { + for tool in extra_tools { components.tools.register(tool).await; } } @@ -560,12 +593,19 @@ impl TestRigBuilder { hooks: components.hooks, cost_guard: components.cost_guard, sse_tx: None, - http_interceptor: if http_exchanges.is_empty() { - None - } else { - Some(Arc::new( - ironclaw::llm::recording::ReplayingHttpInterceptor::new(http_exchanges), - )) + http_interceptor: { + // Prefer explicit exchanges from with_http_exchanges(), fall back to trace. + let exchanges = if explicit_http_exchanges.is_empty() { + trace_http_exchanges + } else { + explicit_http_exchanges + }; + if exchanges.is_empty() { + None + } else { + Some(Arc::new(ReplayingHttpInterceptor::new(exchanges)) + as Arc) + } }, }; @@ -576,14 +616,30 @@ impl TestRigBuilder { channel_manager.add(Box::new(handle)).await; let channels = Arc::new(channel_manager); + // 7b. Register message tool so routines can send messages to channels. + deps.tools + .register_message_tools(Arc::clone(&channels)) + .await; + // 8. Create Agent. + let routine_config = if enable_routines { + Some(ironclaw::config::RoutineConfig { + enabled: true, + cron_check_interval_secs: 60, + max_concurrent_routines: 3, + default_cooldown_secs: 300, + max_lightweight_tokens: 4096, + }) + } else { + None + }; let agent = Agent::new( components.config.agent.clone(), deps, channels, None, // heartbeat_config None, // hygiene_config - None, // routine_config + routine_config, None, // context_manager None, // session_manager ); @@ -604,7 +660,7 @@ impl TestRigBuilder { channel: test_channel, instrumented_llm: instrumented, start_time: Instant::now(), - max_tool_iterations: self.max_tool_iterations, + max_tool_iterations, agent_handle: Some(agent_handle), db: db_ref, workspace: workspace_ref,