feat(routines): approval context for autonomous job execution (#577)

* feat(routines): add approval context for autonomous job execution

Routines and background jobs were unable to use any tools that required
approval (file ops, shell, message, http), making them effectively
useless. This adds an ApprovalContext system that lets autonomous jobs
pre-authorize tools at dispatch time.

- Add ApprovalContext enum with Autonomous variant that auto-approves
  UnlessAutoApproved tools and optionally pre-authorizes Always tools
- Add tool_permissions field to RoutineAction::FullJob for pre-authorizing
  Always-gated tools (e.g. destructive shell, cross-channel messaging)
- Add Scheduler::dispatch_job_with_context() to thread approval context
  through to workers
- Set message tool default channel/target from routine NotifyConfig
  so routines can send results without cross-channel approval
- Fix Completed→Completed state transition error in worker (plan marks
  job completed, then direct loop or outer run() tries again)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* test(routines): add E2E trace for routine news digest workflow

Add a 3-turn trace fixture and test that exercises:
- Turn 1: routine_create with full_job mode and tool_permissions
- Turn 2: Simulated digest workflow with echo + memory_write
- Turn 3: Verification via memory_search

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(test): wire RoutineEngine into test rig for routine_create E2E

- Add `with_routines()` to TestRigBuilder that passes a RoutineConfig
  to Agent::new, enabling routine tool registration during agent startup
- Add Turn 2 (routine_list) to the trace to verify routine persistence
  in the database after routine_create
- Fix formatting issues flagged by CI (cargo fmt)

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* refactor(scheduler): deduplicate dispatch_job and dispatch_job_with_context

Extract shared logic into private `dispatch_job_inner` to prevent
divergence when dispatch behavior changes in the future.

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat(routines): add routine_fire tool and real E2E routine execution test

- Add `routine_fire` tool that calls `RoutineEngine::fire_manual` to
  trigger a routine on demand. Registered alongside the other 5 routine
  tools (now 6 total).

- Rewrite the routine_news_digest E2E trace to exercise the full
  execution stack end-to-end:
  1. routine_create (manual trigger, full_job, tool_permissions: [message])
  2. routine_fire → RoutineEngine → Scheduler::dispatch_job_with_context
     → autonomous Worker consuming TraceLlm steps
  3. Worker calls echo → memory_write → message (broadcast to test channel)
  4. Test verifies the message broadcast arrived, proving ApprovalContext
     correctly allowed the Always-approval message tool

- Register message tools in TestRig so routines can send messages to
  the test channel via channel_manager.broadcast().

[skip-regression-check]

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* feat(routines): wire HttpInterceptor through scheduler for routine worker http calls

Propagate http_interceptor from AgentDeps → Scheduler → WorkerDeps → JobContext
so that routine workers (and any scheduler-dispatched workers) can use the
ReplayingHttpInterceptor for mock HTTP responses during tests.

Changes:
- Add http_interceptor field to Scheduler and WorkerDeps
- Set job_ctx.http_interceptor in Worker before tool execution
- Add with_http_exchanges() builder method to TestRigBuilder
- Replace echo tool with http tool in routine_news_digest trace
- Test now exercises real http tool with mock response → memory_write → message

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address review comments from Copilot on PR #577

- Extract `ApprovalContext::is_blocked_or_default()` helper to deduplicate
  approval check logic in worker.rs and scheduler.rs
- Extract `parse_tool_permissions()` helper to deduplicate JSON array
  parsing in routine.rs and builtin/routine.rs
- Fix test name: `test_mark_completed_twice_does_not_error` →
  `test_mark_completed_twice_returns_error` (matches actual behavior)
- Fix ApprovalContext doc comment to clarify it only models autonomous mode
- Fix flaky index-based assertion in routine_news_digest test — now uses
  content-based search instead of fixed position
- Fix stale comment: echo → http in routine test header
- Add TODO for subtask approval context propagation (latent, not in
  active code paths)
- Add TODO for global message tool context race in routine_engine

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* style: fix formatting in is_blocked_or_default test

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* refactor(test_rig): destructure self in build() to avoid partial-move fragility

Destructure TestRigBuilder at the top of build() instead of accessing
self.* fields after moving self.http_exchanges. While the prior code
compiled (remaining fields are Copy), it was fragile and would break
if any non-Copy field were added.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* docs: clarify that routine_fire bypasses cooldown

Manual fires are explicitly user-initiated and intentionally bypass
cooldown checks (which only apply to automated cron/event triggers).
Updated tool description and fire_manual docstring to make this clear.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix(routines): fix message tool approval in routine context

Two fixes for message tool failures in autonomous routine jobs:

1. MessageTool::requires_approval() now returns UnlessAutoApproved when
   the explicit channel param matches the default channel (was Always,
   causing "requires authentication" errors for routine workers).

2. routine_create tool now accepts notify_channel and notify_user params,
   wired into NotifyConfig. Without these, routines had channel: None,
   so set_message_tool_context was never called, causing "No channel
   specified" errors.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* refactor(message): remove approval requirement from message tool

The message tool only sends to user-owned channels via
ChannelManager::broadcast (TUI, Telegram, Slack, web gateway, etc.).
It cannot reach arbitrary external services, so approval adds friction
with no security benefit. This also eliminates the routine context
errors entirely since approval is never checked.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address review comments — routine_fire approval + test rename

- routine_fire now returns UnlessAutoApproved since firing a routine
  can dispatch a full_job with pre-authorized Always-gated tools
- Rename test_approval_context_never_always_passes to
  test_approval_context_never_is_not_blocked for clarity

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address review nits — update stale docs and comments

- Remove 'message' from tool_permissions example (no longer Always)
- Reword message tool approval comment for accuracy
- Clarify with_routines() docstring re: tool registration vs engine wiring

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-03-07 05:21:58 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 5c2ba44f12
commit ae89a52ac2
15 changed files with 1168 additions and 113 deletions
+3
View File
@@ -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 {
+25 -2
View File
@@ -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<String>,
},
}
@@ -186,6 +191,19 @@ fn default_max_iterations() -> u32 {
10
}
/// Parse a `tool_permissions` JSON array into a `Vec<String>`.
pub fn parse_tool_permissions(value: &serde_json::Value) -> Vec<String> {
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()])
);
}
+40 -2
View File
@@ -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<Uuid, RoutineError> {
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<String>, Option<i32>), 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}"),
+270 -3
View File
@@ -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<HookRegistry>,
/// SSE broadcast sender for live job event streaming.
sse_tx: Option<tokio::sync::broadcast::Sender<SseEvent>>,
/// HTTP interceptor for trace recording/replay (propagated to workers).
http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
/// Running jobs (main LLM-driven jobs).
jobs: Arc<RwLock<HashMap<Uuid, ScheduledJob>>>,
/// 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<dyn crate::llm::recording::HttpInterceptor>,
) {
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<serde_json::Value>,
) -> Result<Uuid, JobError> {
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<serde_json::Value>,
approval_context: ApprovalContext,
) -> Result<Uuid, JobError> {
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<serde_json::Value>,
approval_context: Option<ApprovalContext>,
) -> Result<Uuid, JobError> {
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<ApprovalContext>,
) -> 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<ToolRegistry>,
context_manager: Arc<ContextManager>,
safety: Arc<SafetyLayer>,
approval_context: Option<ApprovalContext>,
job_id: Uuid,
tool_name: &str,
params: serde_json::Value,
@@ -413,7 +474,10 @@ impl Scheduler {
.into());
}
if tool.requires_approval(&params).is_required() {
let requirement = tool.requires_approval(&params);
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<ToolOutput, ToolError> {
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<ToolOutput, ToolError> {
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<ToolRegistry>,
Arc<ContextManager>,
Arc<SafetyLayer>,
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"
);
}
}
+223 -8
View File
@@ -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<tokio::sync::broadcast::Sender<SseEvent>>,
/// Approval context for tool execution. When `None`, all non-`Never` tools are
/// blocked (legacy behavior). When `Some`, the context determines which tools
/// are pre-approved for autonomous execution.
pub approval_context: Option<ApprovalContext>,
/// HTTP interceptor for trace recording/replay (propagated to JobContext).
pub http_interceptor: Option<Arc<dyn crate::llm::recording::HttpInterceptor>>,
}
/// Worker that executes a single job.
@@ -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<Arc<dyn Tool>>,
approval_context: Option<crate::tools::ApprovalContext>,
) -> Worker {
let registry = ToolRegistry::new();
for t in tools {
registry.register(t).await;
}
let cm = Arc::new(crate::context::ContextManager::new(5));
let job_id = cm.create_job("test", "test job").await.unwrap();
let deps = WorkerDeps {
context_manager: cm,
llm: Arc::new(StubLlm),
safety: Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: false,
})),
tools: Arc::new(registry),
store: None,
hooks: Arc::new(crate::hooks::HookRegistry::new()),
timeout: Duration::from_secs(30),
use_planning: false,
sse_tx: None,
approval_context,
http_interceptor: None,
};
Worker::new(job_id, deps)
}
/// A tool that requires approval (UnlessAutoApproved).
struct ApprovalTool;
#[async_trait::async_trait]
impl Tool for ApprovalTool {
fn name(&self) -> &str {
"needs_approval"
}
fn description(&self) -> &str {
"Tool requiring approval"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({"type": "object", "properties": {}})
}
async fn execute(
&self,
_params: serde_json::Value,
_ctx: &crate::context::JobContext,
) -> Result<ToolOutput, crate::tools::ToolError> {
Ok(ToolOutput::text(
"approved",
std::time::Instant::now().elapsed(),
))
}
fn requires_approval(
&self,
_params: &serde_json::Value,
) -> crate::tools::ApprovalRequirement {
crate::tools::ApprovalRequirement::UnlessAutoApproved
}
fn requires_sanitization(&self) -> bool {
false
}
}
/// A tool that always requires approval.
struct AlwaysApprovalTool;
#[async_trait::async_trait]
impl Tool for AlwaysApprovalTool {
fn name(&self) -> &str {
"always_approval"
}
fn description(&self) -> &str {
"Tool always requiring approval"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({"type": "object", "properties": {}})
}
async fn execute(
&self,
_params: serde_json::Value,
_ctx: &crate::context::JobContext,
) -> Result<ToolOutput, crate::tools::ToolError> {
Ok(ToolOutput::text(
"always",
std::time::Instant::now().elapsed(),
))
}
fn requires_approval(
&self,
_params: &serde_json::Value,
) -> crate::tools::ApprovalRequirement {
crate::tools::ApprovalRequirement::Always
}
fn requires_sanitization(&self) -> bool {
false
}
}
#[tokio::test]
async fn test_approval_context_unblocks_unless_auto_approved() {
// Without approval context, UnlessAutoApproved is blocked
let worker_blocked = make_worker_with_approval(vec![Arc::new(ApprovalTool)], None).await;
let result = worker_blocked
.execute_tool("needs_approval", &serde_json::json!({}))
.await;
assert!(
result.is_err(),
"Should be blocked without approval context"
);
// With autonomous approval context, UnlessAutoApproved is allowed
let worker_allowed = make_worker_with_approval(
vec![Arc::new(ApprovalTool)],
Some(crate::tools::ApprovalContext::autonomous()),
)
.await;
let result = worker_allowed
.execute_tool("needs_approval", &serde_json::json!({}))
.await;
assert!(result.is_ok(), "Should be allowed with autonomous context");
}
#[tokio::test]
async fn test_approval_context_blocks_always_unless_permitted() {
// Autonomous context without tool_permissions blocks Always tools
let worker_blocked = make_worker_with_approval(
vec![Arc::new(AlwaysApprovalTool)],
Some(crate::tools::ApprovalContext::autonomous()),
)
.await;
let result = worker_blocked
.execute_tool("always_approval", &serde_json::json!({}))
.await;
assert!(
result.is_err(),
"Always tool should be blocked without permission"
);
// Autonomous context with tool_permissions allows Always tools
let worker_allowed = make_worker_with_approval(
vec![Arc::new(AlwaysApprovalTool)],
Some(crate::tools::ApprovalContext::autonomous_with_tools([
"always_approval".to_string(),
])),
)
.await;
let result = worker_allowed
.execute_tool("always_approval", &serde_json::json!({}))
.await;
assert!(
result.is_ok(),
"Always tool should be allowed with permission"
);
}
}
+15 -67
View File
@@ -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<ToolRateLimitConfig> {
@@ -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,
);
}
}
+2 -1
View File
@@ -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;
+113 -8
View File
@@ -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(&params);
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<dyn Database>,
engine: Arc<RoutineEngine>,
}
impl RoutineFireTool {
pub fn new(store: Arc<dyn Database>, engine: Arc<RoutineEngine>) -> 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<ToolOutput, ToolError> {
let start = std::time::Instant::now();
let name = require_str(&params, "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 {
+2 -2
View File
@@ -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,
};
+8 -3
View File
@@ -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<crate::agent::routine_engine::RoutineEngine>,
) {
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.
+18 -1
View File
@@ -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!({
+121
View File
@@ -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<String>,
},
}
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<Item = String>) -> 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<Self>,
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
));
}
}
+117 -1
View File
@@ -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::<Vec<_>>()
);
// 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::<Vec<_>>()
);
// 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]
@@ -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
}
}
]
}
]
}
+71 -15
View File
@@ -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<TestChannel> as Box<dyn Channel>
// ---------------------------------------------------------------------------
@@ -362,6 +364,8 @@ pub struct TestRigBuilder {
llm: Option<Arc<dyn LlmProvider>>,
max_tool_iterations: usize,
injection_check: bool,
enable_routines: bool,
http_exchanges: Vec<HttpExchange>,
extra_tools: Vec<Arc<dyn Tool>>,
}
@@ -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<HttpExchange>) -> 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<Arc<TraceLlm>> = None;
let base_llm: Arc<dyn LlmProvider> = if let Some(llm) = self.llm {
let base_llm: Arc<dyn LlmProvider> = 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<dyn ironclaw::llm::recording::HttpInterceptor>)
}
},
};
@@ -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,