diff --git a/src/agent/routine.rs b/src/agent/routine.rs index 72226502..2dee6333 100644 --- a/src/agent/routine.rs +++ b/src/agent/routine.rs @@ -207,7 +207,7 @@ impl Trigger { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub enum RoutineAction { - /// Single LLM call, no tools. Cheap and fast. + /// Single LLM call (optionally with tools). Cheap and fast. Lightweight { /// The prompt sent to the LLM. prompt: String, @@ -217,6 +217,14 @@ pub enum RoutineAction { /// Max output tokens (default: 4096). #[serde(default = "default_max_tokens")] max_tokens: u32, + /// Enable tool access (default: false for backward compatibility). + /// When true, the LLM can call tools during execution. + /// Tools requiring approval are automatically filtered out. + #[serde(default)] + use_tools: bool, + /// Max tool call rounds (default: 3). Only used when use_tools is true. + #[serde(default = "default_max_tool_rounds")] + max_tool_rounds: u32, }, /// Full multi-turn worker job with tool access. FullJob { @@ -243,6 +251,19 @@ fn default_max_iterations() -> u32 { 10 } +fn default_max_tool_rounds() -> u32 { + 3 +} + +/// Hard upper bound for max_tool_rounds to prevent runaway loops and cost explosion. +pub(crate) const MAX_TOOL_ROUNDS_LIMIT: u32 = 20; + +/// Clamp max_tool_rounds to [1, MAX_TOOL_ROUNDS_LIMIT]. +/// Accepts u64 to avoid truncation before clamping. +fn clamp_max_tool_rounds(value: u64) -> u32 { + value.clamp(1, MAX_TOOL_ROUNDS_LIMIT as u64) as u32 +} + /// Parse a `tool_permissions` JSON array into a `Vec`. pub fn parse_tool_permissions(value: &serde_json::Value) -> Vec { value @@ -290,10 +311,22 @@ impl RoutineAction { .get("max_tokens") .and_then(|v| v.as_u64()) .unwrap_or(default_max_tokens() as u64) as u32; + let use_tools = config + .get("use_tools") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let max_tool_rounds = clamp_max_tool_rounds( + config + .get("max_tool_rounds") + .and_then(|v| v.as_u64()) + .unwrap_or(default_max_tool_rounds() as u64), + ); Ok(RoutineAction::Lightweight { prompt, context_paths, max_tokens, + use_tools, + max_tool_rounds, }) } "full_job" => { @@ -339,10 +372,14 @@ impl RoutineAction { prompt, context_paths, max_tokens, + use_tools, + max_tool_rounds, } => serde_json::json!({ "prompt": prompt, "context_paths": context_paths, "max_tokens": max_tokens, + "use_tools": use_tools, + "max_tool_rounds": max_tool_rounds, }), RoutineAction::FullJob { title, @@ -504,7 +541,8 @@ pub fn next_cron_fire( #[cfg(test)] mod tests { use crate::agent::routine::{ - RoutineAction, RoutineGuardrails, RunStatus, Trigger, content_hash, next_cron_fire, + MAX_TOOL_ROUNDS_LIMIT, RoutineAction, RoutineGuardrails, RunStatus, Trigger, content_hash, + next_cron_fire, }; #[test] @@ -554,11 +592,13 @@ mod tests { prompt: "Check PRs".to_string(), context_paths: vec!["context/priorities.md".to_string()], max_tokens: 2048, + use_tools: false, + max_tool_rounds: 3, }; let json = action.to_config_json(); let parsed = RoutineAction::from_db("lightweight", json).expect("parse lightweight"); assert!( - matches!(parsed, RoutineAction::Lightweight { prompt, context_paths, max_tokens } + matches!(parsed, RoutineAction::Lightweight { prompt, context_paths, max_tokens, .. } if prompt == "Check PRs" && context_paths.len() == 1 && max_tokens == 2048) ); } @@ -695,4 +735,77 @@ mod tests { ); assert_eq!(Trigger::Manual.type_tag(), "manual"); } + + #[test] + fn test_action_lightweight_backward_compat_no_use_tools() { + // Simulate old DB record without use_tools field + let json = serde_json::json!({ + "prompt": "old routine", + "context_paths": [], + "max_tokens": 4096 + }); + let parsed = RoutineAction::from_db("lightweight", json).expect("parse lightweight"); + assert!( + matches!(parsed, RoutineAction::Lightweight { use_tools, max_tool_rounds, .. } + if !use_tools && max_tool_rounds == 3), + "missing use_tools should default to false, max_tool_rounds to 3" + ); + } + + #[test] + fn test_max_tool_rounds_clamped_to_upper_bound() { + let json = serde_json::json!({ + "prompt": "test", + "use_tools": true, + "max_tool_rounds": 9999 + }); + let parsed = RoutineAction::from_db("lightweight", json).expect("parse"); + match parsed { + RoutineAction::Lightweight { + max_tool_rounds, .. + } => { + assert_eq!( + max_tool_rounds, MAX_TOOL_ROUNDS_LIMIT, + "should clamp to MAX_TOOL_ROUNDS_LIMIT" + ); + } + _ => panic!("expected Lightweight"), + } + } + + #[test] + fn test_max_tool_rounds_clamped_to_lower_bound() { + let json = serde_json::json!({ + "prompt": "test", + "use_tools": true, + "max_tool_rounds": 0 + }); + let parsed = RoutineAction::from_db("lightweight", json).expect("parse"); + match parsed { + RoutineAction::Lightweight { + max_tool_rounds, .. + } => { + assert_eq!(max_tool_rounds, 1, "should clamp 0 to 1"); + } + _ => panic!("expected Lightweight"), + } + } + + #[test] + fn test_max_tool_rounds_normal_value_passes_through() { + let json = serde_json::json!({ + "prompt": "test", + "use_tools": true, + "max_tool_rounds": 10 + }); + let parsed = RoutineAction::from_db("lightweight", json).expect("parse"); + match parsed { + RoutineAction::Lightweight { + max_tool_rounds, .. + } => { + assert_eq!(max_tool_rounds, 10, "normal value should pass through"); + } + _ => panic!("expected Lightweight"), + } + } } diff --git a/src/agent/routine_engine.rs b/src/agent/routine_engine.rs index b10021ef..a973437a 100644 --- a/src/agent/routine_engine.rs +++ b/src/agent/routine_engine.rs @@ -459,7 +459,20 @@ async fn execute_routine(ctx: EngineContext, routine: Routine, run: RoutineRun) prompt, context_paths, max_tokens, - } => execute_lightweight(&ctx, &routine, prompt, context_paths, *max_tokens).await, + use_tools, + max_tool_rounds, + } => { + execute_lightweight( + &ctx, + &routine, + prompt, + context_paths, + *max_tokens, + *use_tools, + *max_tool_rounds, + ) + .await + } RoutineAction::FullJob { title, description, @@ -670,6 +683,8 @@ async fn execute_lightweight( prompt: &str, context_paths: &[String], max_tokens: u32, + use_tools: bool, + max_tool_rounds: u32, ) -> Result<(RunStatus, Option, Option), RoutineError> { // Load context from workspace let mut context_parts = Vec::new(); @@ -732,14 +747,15 @@ async fn execute_lightweight( Err(_) => max_tokens, }; - // If tools are enabled, use the tool execution loop; otherwise, single LLM call - if ctx.config.lightweight_tools_enabled { + // If tools are enabled (both globally and per-routine), use the tool execution loop + if use_tools && ctx.config.lightweight_tools_enabled { execute_lightweight_with_tools( ctx, routine, &system_prompt, &full_prompt, effective_max_tokens, + max_tool_rounds, ) .await } else { @@ -783,24 +799,12 @@ async fn execute_lightweight_no_tools( reason: e.to_string(), })?; - let content = response.content.trim(); - let tokens_used = Some((response.input_tokens + response.output_tokens) as i32); - - // Empty content guard - if content.is_empty() { - return if response.finish_reason == FinishReason::Length { - Err(RoutineError::TruncatedResponse) - } else { - Err(RoutineError::EmptyResponse) - }; - } - - // Check for the "nothing to do" sentinel - if content == "ROUTINE_OK" || content.contains("ROUTINE_OK") { - return Ok((RunStatus::Ok, None, tokens_used)); - } - - Ok((RunStatus::Attention, Some(content.to_string()), tokens_used)) + handle_text_response( + &response.content, + response.finish_reason, + response.input_tokens, + response.output_tokens, + ) } /// Handle a text-only LLM response in lightweight routine execution. @@ -850,6 +854,7 @@ async fn execute_lightweight_with_tools( system_prompt: &str, full_prompt: &str, effective_max_tokens: u32, + max_tool_rounds: u32, ) -> Result<(RunStatus, Option, Option), RoutineError> { let mut messages = if system_prompt.is_empty() { vec![ChatMessage::user(full_prompt)] @@ -860,7 +865,9 @@ async fn execute_lightweight_with_tools( ] }; - let max_iterations = ctx.config.lightweight_max_iterations.min(5); + let max_iterations = max_tool_rounds + .min(ctx.config.lightweight_max_iterations) + .min(5); let mut iteration = 0; let mut total_input_tokens = 0; let mut total_output_tokens = 0; @@ -906,7 +913,10 @@ async fn execute_lightweight_with_tools( ); } else { // Tool-enabled iteration - let tool_defs = ctx.tools.tool_definitions().await; + let tool_defs = ctx + .tools + .tool_definitions_excluding(ROUTINE_TOOL_DENYLIST) + .await; let request = ToolCompletionRequest::new(messages.clone(), tool_defs) .with_max_tokens(effective_max_tokens) @@ -972,12 +982,33 @@ async fn execute_lightweight_with_tools( } } +/// Tools that must never be callable from lightweight routines. +/// +/// These tools pose autonomy-escalation risks: a routine could self-replicate, +/// modify its own triggers/prompts, delete other routines, or restart the agent. +const ROUTINE_TOOL_DENYLIST: &[&str] = &[ + "routine_create", + "routine_update", + "routine_delete", + "routine_fire", + "restart", +]; + /// Execute a single tool for a lightweight routine. async fn execute_routine_tool( ctx: &EngineContext, job_ctx: &JobContext, tc: &ToolCall, ) -> Result> { + // Block tools that pose autonomy-escalation risks + if ROUTINE_TOOL_DENYLIST.contains(&tc.name.as_str()) { + return Err(format!( + "Tool '{}' is not available in lightweight routines", + tc.name + ) + .into()); + } + // Check if tool exists let tool = ctx .tools @@ -1283,6 +1314,36 @@ mod tests { } } + #[test] + fn test_routine_tool_denylist_blocks_self_management_tools() { + let denylisted = vec![ + "routine_create", + "routine_update", + "routine_delete", + "routine_fire", + "restart", + ]; + for tool in &denylisted { + assert!( + super::ROUTINE_TOOL_DENYLIST.contains(tool), + "Tool '{}' should be in ROUTINE_TOOL_DENYLIST", + tool + ); + } + } + + #[test] + fn test_routine_tool_denylist_allows_safe_tools() { + let allowed = vec!["echo", "time", "json", "http", "memory_search", "shell"]; + for tool in &allowed { + assert!( + !super::ROUTINE_TOOL_DENYLIST.contains(tool), + "Tool '{}' should NOT be in ROUTINE_TOOL_DENYLIST", + tool + ); + } + } + #[test] fn test_empty_response_handling() { // Simulate the empty content guard logic diff --git a/src/testing/mod.rs b/src/testing/mod.rs index d2078a80..33702e67 100644 --- a/src/testing/mod.rs +++ b/src/testing/mod.rs @@ -1067,6 +1067,8 @@ mod tests { prompt: "Check status".to_string(), context_paths: vec![], max_tokens: 500, + use_tools: false, + max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(60), @@ -1198,6 +1200,8 @@ mod tests { prompt: "test".to_string(), context_paths: vec![], max_tokens: 100, + use_tools: false, + max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: std::time::Duration::from_secs(0), diff --git a/src/tools/builtin/routine.rs b/src/tools/builtin/routine.rs index 573c3c60..43e2add7 100644 --- a/src/tools/builtin/routine.rs +++ b/src/tools/builtin/routine.rs @@ -104,6 +104,14 @@ impl Tool for RoutineCreateTool { "enum": ["lightweight", "full_job"], "description": "Execution mode: 'lightweight' (single LLM call, default) or 'full_job' (multi-turn with tools)" }, + "use_tools": { + "type": "boolean", + "description": "Enable tool access in lightweight mode (default: false). Only safe tools (no approval required) are available. Ignored for full_job mode." + }, + "max_tool_rounds": { + "type": "integer", + "description": "Max tool call rounds in lightweight mode (default: 3). Only used when use_tools is true." + }, "cooldown_secs": { "type": "integer", "description": "Minimum seconds between fires (default: 300)" @@ -262,11 +270,24 @@ impl Tool for RoutineCreateTool { }) .unwrap_or_default(); + let use_tools = params + .get("use_tools") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let max_tool_rounds = params + .get("max_tool_rounds") + .and_then(|v| v.as_u64()) + .map(|v| v.clamp(1, crate::agent::routine::MAX_TOOL_ROUNDS_LIMIT as u64) as u32) + .unwrap_or(3); + let action = match action_type { "lightweight" => RoutineAction::Lightweight { prompt: prompt.to_string(), context_paths, max_tokens: 4096, + use_tools, + max_tool_rounds, }, "full_job" => { let tool_permissions = crate::agent::routine::parse_tool_permissions(¶ms); diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 7054eea3..bee8cf27 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -23,7 +23,7 @@ use crate::tools::builtin::{ ToolUpgradeTool, WriteFileTool, }; use crate::tools::rate_limiter::RateLimiter; -use crate::tools::tool::{Tool, ToolDomain}; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain}; use crate::tools::wasm::{ Capabilities, OAuthRefreshConfig, ResourceLimits, SharedCredentialRegistry, WasmError, WasmStorageError, WasmToolRuntime, WasmToolStore, WasmToolWrapper, @@ -278,6 +278,38 @@ impl ToolRegistry { .collect() } + /// Get tool definitions excluding specific tools by name. + /// + /// Used by lightweight routines to filter out denylisted and approval-gated tools + /// so the LLM only sees tools it is actually allowed to call. + pub async fn tool_definitions_excluding(&self, deny: &[&str]) -> Vec { + let empty_params = serde_json::Value::Object(serde_json::Map::new()); + let mut defs: Vec = self + .tools + .read() + .await + .values() + .filter(|tool| { + // Exclude denylisted tools + if deny.contains(&tool.name()) { + return false; + } + // Exclude tools that require approval + matches!( + tool.requires_approval(&empty_params), + ApprovalRequirement::Never + ) + }) + .map(|tool| ToolDefinition { + name: tool.name().to_string(), + description: tool.description().to_string(), + parameters: tool.parameters_schema(), + }) + .collect(); + defs.sort_unstable_by(|a, b| a.name.cmp(&b.name)); + defs + } + /// Register development tools for building software. /// /// These tools provide shell access, file operations, and code editing diff --git a/tests/e2e_routine_heartbeat.rs b/tests/e2e_routine_heartbeat.rs index f245e656..f5a28c25 100644 --- a/tests/e2e_routine_heartbeat.rs +++ b/tests/e2e_routine_heartbeat.rs @@ -61,6 +61,8 @@ mod tests { prompt: prompt.to_string(), context_paths: vec![], max_tokens: 1000, + use_tools: false, + max_tool_rounds: 3, }, guardrails: RoutineGuardrails { cooldown: Duration::from_secs(0),