From 5fcb4d3c034c8694fb32ba1dc86feec397b18429 Mon Sep 17 00:00:00 2001 From: "ilblackdragon@gmail.com" Date: Tue, 10 Mar 2026 02:01:58 -0700 Subject: [PATCH] feat: add tool execution idempotency cache When the LLM re-requests the same idempotent tool with identical args (common during self-repair recovery, stuck job retries, or chat-mode retry loops), the cache returns the previous result instantly without re-executing. Design improvements over the closed PR #204: - Single global LRU cache (lru crate) instead of per-job HashMap with O(N) manual eviction - Global cap of 2000 entries prevents memory leaks from chat-path ephemeral job IDs that never get invalidated - TTL expiry (30min) checked on read; LRU eviction handles the rest - Job invalidation still runs on worker completion for prompt cleanup Idempotent tools: echo, time, json, memory_read/search/tree, read_file, list_dir, list_jobs, job_status. Closes #204 Co-Authored-By: Claude Opus 4.6 --- src/agent/agent_loop.rs | 4 + src/agent/dispatcher.rs | 67 +++++++-- src/agent/scheduler.rs | 7 + src/agent/thread_ops.rs | 2 + src/agent/worker.rs | 57 +++++++- src/main.rs | 3 + src/testing.rs | 3 + src/tools/builtin/echo.rs | 4 + src/tools/builtin/file.rs | 8 ++ src/tools/builtin/job.rs | 8 ++ src/tools/builtin/json.rs | 4 + src/tools/builtin/memory.rs | 12 ++ src/tools/builtin/time.rs | 4 + src/tools/idempotency.rs | 270 ++++++++++++++++++++++++++++++++++++ src/tools/mod.rs | 1 + src/tools/tool.rs | 11 ++ tests/support/test_rig.rs | 3 + 17 files changed, 453 insertions(+), 15 deletions(-) create mode 100644 src/tools/idempotency.rs diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 15853f14..7f90c109 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -29,6 +29,7 @@ use crate::llm::LlmProvider; use crate::safety::SafetyLayer; use crate::skills::SkillRegistry; use crate::tools::ToolRegistry; +use crate::tools::idempotency::ToolIdempotencyCache; use crate::workspace::Workspace; /// Collapse a tool output string into a single-line preview for display. @@ -81,6 +82,8 @@ pub struct AgentDeps { pub transcription: Option>, /// Document text extraction middleware for PDF, DOCX, PPTX, etc. pub document_extraction: Option>, + /// Idempotency cache for tool executions. + pub idempotency_cache: ToolIdempotencyCache, } /// The main agent that coordinates all components. @@ -130,6 +133,7 @@ impl Agent { deps.tools.clone(), deps.store.clone(), deps.hooks.clone(), + deps.idempotency_cache.clone(), ); if let Some(ref tx) = deps.sse_tx { scheduler.set_sse_sender(tx.clone()); diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 99feed9d..f981b25b 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -15,6 +15,7 @@ use crate::channels::{IncomingMessage, StatusUpdate}; use crate::context::JobContext; use crate::error::Error; use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult}; +use crate::tools::idempotency::ToolIdempotencyCache; use crate::tools::redact_params; /// Result of the agentic loop execution. @@ -569,6 +570,7 @@ impl Agent { let pf_idx = *pf_idx; let tools = self.tools().clone(); let safety = self.safety().clone(); + let idempotency_cache = self.deps.idempotency_cache.clone(); let channels = self.channels.clone(); let job_ctx = job_ctx.clone(); let tc = tc.clone(); @@ -589,6 +591,7 @@ impl Agent { let result = execute_chat_tool_standalone( &tools, &safety, + &idempotency_cache, &tc.name, &tc.arguments, &job_ctx, @@ -859,7 +862,15 @@ impl Agent { params: &serde_json::Value, job_ctx: &JobContext, ) -> Result { - execute_chat_tool_standalone(self.tools(), self.safety(), tool_name, params, job_ctx).await + execute_chat_tool_standalone( + self.tools(), + self.safety(), + &self.deps.idempotency_cache, + tool_name, + params, + job_ctx, + ) + .await } } @@ -868,9 +879,11 @@ impl Agent { /// This standalone function enables parallel invocation from spawned JoinSet /// tasks, which cannot borrow `&self`. It replicates the logic from /// `Agent::execute_chat_tool`. +#[allow(clippy::too_many_arguments)] pub(super) async fn execute_chat_tool_standalone( tools: &crate::tools::ToolRegistry, safety: &crate::safety::SafetyLayer, + idempotency_cache: &ToolIdempotencyCache, tool_name: &str, params: &serde_json::Value, job_ctx: &crate::context::JobContext, @@ -882,6 +895,15 @@ pub(super) async fn execute_chat_tool_standalone( name: tool_name.to_string(), })?; + // Check idempotency cache + let job_id_str = job_ctx.job_id.to_string(); + if tool.is_idempotent() + && let Some(cached) = idempotency_cache.get(&job_id_str, tool_name, params).await + { + tracing::debug!(tool = %tool_name, "Idempotency cache hit (chat)"); + return Ok(cached); + } + // Validate tool parameters let validation = safety.validator().validate_tool_params(params); if !validation.is_valid { @@ -953,13 +975,25 @@ pub(super) async fn execute_chat_tool_standalone( reason: e.to_string(), })?; - serde_json::to_string_pretty(&result.result).map_err(|e| { - crate::error::ToolError::ExecutionFailed { - name: tool_name.to_string(), - reason: format!("Failed to serialize result: {}", e), - } - .into() - }) + let result_str: Result = + serde_json::to_string_pretty(&result.result).map_err(|e| { + crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: format!("Failed to serialize result: {}", e), + } + .into() + }); + + // Cache successful results for idempotent tools + if let Ok(ref output_str) = result_str + && tool.is_idempotent() + { + idempotency_cache + .put(&job_id_str, tool_name, params, output_str.clone()) + .await; + } + + result_str } /// Parsed auth result fields for emitting StatusUpdate::AuthRequired. @@ -1187,6 +1221,9 @@ mod tests { http_interceptor: None, transcription: None, document_extraction: None, + idempotency_cache: crate::tools::idempotency::ToolIdempotencyCache::new( + crate::tools::idempotency::IdempotencyCacheConfig::default(), + ), }; Agent::new( @@ -1520,9 +1557,13 @@ mod tests { let job_ctx = JobContext::with_user("test", "chat", "test session"); + let cache = crate::tools::idempotency::ToolIdempotencyCache::new( + crate::tools::idempotency::IdempotencyCacheConfig::default(), + ); let result = super::execute_chat_tool_standalone( ®istry, &safety, + &cache, "echo", &serde_json::json!({"message": "hello"}), &job_ctx, @@ -1548,9 +1589,13 @@ mod tests { }); let job_ctx = JobContext::with_user("test", "chat", "test session"); + let cache = crate::tools::idempotency::ToolIdempotencyCache::new( + crate::tools::idempotency::IdempotencyCacheConfig::default(), + ); let result = super::execute_chat_tool_standalone( ®istry, &safety, + &cache, "nonexistent", &serde_json::json!({}), &job_ctx, @@ -2026,6 +2071,9 @@ mod tests { http_interceptor: None, transcription: None, document_extraction: None, + idempotency_cache: crate::tools::idempotency::ToolIdempotencyCache::new( + crate::tools::idempotency::IdempotencyCacheConfig::default(), + ), }; Agent::new( @@ -2143,6 +2191,9 @@ mod tests { http_interceptor: None, transcription: None, document_extraction: None, + idempotency_cache: crate::tools::idempotency::ToolIdempotencyCache::new( + crate::tools::idempotency::IdempotencyCacheConfig::default(), + ), }; Agent::new( diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 85f3f6eb..9a88093a 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -18,6 +18,7 @@ use crate::error::{Error, JobError}; use crate::hooks::HookRegistry; use crate::llm::LlmProvider; use crate::safety::SafetyLayer; +use crate::tools::idempotency::ToolIdempotencyCache; use crate::tools::{ApprovalContext, ToolRegistry}; /// Message to send to a worker. @@ -58,6 +59,8 @@ pub struct Scheduler { sse_tx: Option>, /// HTTP interceptor for trace recording/replay (propagated to workers). http_interceptor: Option>, + /// Idempotency cache for tool executions (shared across all workers). + idempotency_cache: ToolIdempotencyCache, /// Running jobs (main LLM-driven jobs). jobs: Arc>>, /// Running sub-tasks (tool executions, background tasks). @@ -66,6 +69,7 @@ pub struct Scheduler { impl Scheduler { /// Create a new scheduler. + #[allow(clippy::too_many_arguments)] pub fn new( config: AgentConfig, context_manager: Arc, @@ -74,6 +78,7 @@ impl Scheduler { tools: Arc, store: Option>, hooks: Arc, + idempotency_cache: ToolIdempotencyCache, ) -> Self { Self { config, @@ -85,6 +90,7 @@ impl Scheduler { hooks, sse_tx: None, http_interceptor: None, + idempotency_cache, jobs: Arc::new(RwLock::new(HashMap::new())), subtasks: Arc::new(RwLock::new(HashMap::new())), } @@ -254,6 +260,7 @@ impl Scheduler { sse_tx: self.sse_tx.clone(), approval_context, http_interceptor: self.http_interceptor.clone(), + idempotency_cache: self.idempotency_cache.clone(), }; let worker = Worker::new(job_id, deps); diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index 758e98ed..b9afb35a 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -959,6 +959,7 @@ impl Agent { for (spawn_idx, tc) in runnable.iter().enumerate() { let tools = self.tools().clone(); let safety = self.safety().clone(); + let idempotency_cache = self.deps.idempotency_cache.clone(); let channels = self.channels.clone(); let job_ctx = job_ctx.clone(); let tc = tc.clone(); @@ -979,6 +980,7 @@ impl Agent { let result = execute_chat_tool_standalone( &tools, &safety, + &idempotency_cache, &tc.name, &tc.arguments, &job_ctx, diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 19bfc8e5..481b945a 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -19,6 +19,7 @@ use crate::llm::{ ToolSelection, }; use crate::safety::SafetyLayer; +use crate::tools::idempotency::ToolIdempotencyCache; use crate::tools::rate_limiter::RateLimitResult; use crate::tools::{ApprovalContext, ToolRegistry, redact_params}; @@ -44,6 +45,8 @@ pub struct WorkerDeps { pub approval_context: Option, /// HTTP interceptor for trace recording/replay (propagated to JobContext). pub http_interceptor: Option>, + /// Idempotency cache for tool executions. + pub idempotency_cache: ToolIdempotencyCache, } /// Worker that executes a single job. @@ -285,6 +288,12 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } } + // Clean up idempotency cache entries for this job + self.deps + .idempotency_cache + .invalidate_job(&self.job_id.to_string()) + .await; + Ok(()) } @@ -737,6 +746,22 @@ Report when the job is complete or if you encounter issues you cannot resolve."# name: tool_name.to_string(), })?; + // Check idempotency cache before any expensive work + let job_id_str = job_id.to_string(); + if tool.is_idempotent() + && let Some(cached) = deps + .idempotency_cache + .get(&job_id_str, tool_name, params) + .await + { + tracing::debug!( + tool = %tool_name, + job = %job_id, + "Idempotency cache hit" + ); + return Ok(cached); + } + // Check approval: use context-aware check if available, else block all non-Never tools let requirement = tool.requires_approval(params); let blocked = @@ -967,13 +992,25 @@ Report when the job is complete or if you encounter issues you cannot resolve."# })?; // Return result as string - serde_json::to_string_pretty(&output.result).map_err(|e| { - crate::error::ToolError::ExecutionFailed { - name: tool_name.to_string(), - reason: format!("Failed to serialize result: {}", e), - } - .into() - }) + let result_str: Result = serde_json::to_string_pretty(&output.result) + .map_err(|e| { + crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: format!("Failed to serialize result: {}", e), + } + .into() + }); + + // Cache successful results for idempotent tools + if let Ok(ref output_str) = result_str + && tool.is_idempotent() + { + deps.idempotency_cache + .put(&job_id_str, tool_name, ¶ms, output_str.clone()) + .await; + } + + result_str } /// Process a tool execution result and add it to the reasoning context. @@ -1386,6 +1423,9 @@ mod tests { sse_tx: None, approval_context: None, http_interceptor: None, + idempotency_cache: crate::tools::idempotency::ToolIdempotencyCache::new( + crate::tools::idempotency::IdempotencyCacheConfig::default(), + ), }; Worker::new(job_id, deps) @@ -1648,6 +1688,9 @@ mod tests { sse_tx: None, approval_context, http_interceptor: None, + idempotency_cache: crate::tools::idempotency::ToolIdempotencyCache::new( + crate::tools::idempotency::IdempotencyCacheConfig::default(), + ), }; Worker::new(job_id, deps) diff --git a/src/main.rs b/src/main.rs index 120fa33c..c0b19506 100644 --- a/src/main.rs +++ b/src/main.rs @@ -623,6 +623,9 @@ async fn async_main() -> anyhow::Result<()> { document_extraction: Some(Arc::new( ironclaw::document_extraction::DocumentExtractionMiddleware::new(), )), + idempotency_cache: ironclaw::tools::idempotency::ToolIdempotencyCache::new( + ironclaw::tools::idempotency::IdempotencyCacheConfig::default(), + ), }; let mut agent = Agent::new( diff --git a/src/testing.rs b/src/testing.rs index 8f57cffc..e63c1b17 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -453,6 +453,9 @@ impl TestHarnessBuilder { http_interceptor: None, transcription: None, document_extraction: None, + idempotency_cache: crate::tools::idempotency::ToolIdempotencyCache::new( + crate::tools::idempotency::IdempotencyCacheConfig::default(), + ), }; TestHarness { diff --git a/src/tools/builtin/echo.rs b/src/tools/builtin/echo.rs index faf31982..ef303df0 100644 --- a/src/tools/builtin/echo.rs +++ b/src/tools/builtin/echo.rs @@ -46,4 +46,8 @@ impl Tool for EchoTool { fn requires_sanitization(&self) -> bool { false // Internal tool, no external data } + + fn is_idempotent(&self) -> bool { + true + } } diff --git a/src/tools/builtin/file.rs b/src/tools/builtin/file.rs index 72e0151c..033d6861 100644 --- a/src/tools/builtin/file.rs +++ b/src/tools/builtin/file.rs @@ -170,6 +170,10 @@ impl Tool for ReadFileTool { true // File content could contain anything } + fn is_idempotent(&self) -> bool { + true + } + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { ApprovalRequirement::UnlessAutoApproved } @@ -397,6 +401,10 @@ impl Tool for ListDirTool { false // Directory listings are safe } + fn is_idempotent(&self) -> bool { + true + } + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { ApprovalRequirement::UnlessAutoApproved } diff --git a/src/tools/builtin/job.rs b/src/tools/builtin/job.rs index f502259f..effab806 100644 --- a/src/tools/builtin/job.rs +++ b/src/tools/builtin/job.rs @@ -896,6 +896,10 @@ impl Tool for ListJobsTool { fn requires_sanitization(&self) -> bool { false } + + fn is_idempotent(&self) -> bool { + true + } } /// Tool for checking job status. @@ -975,6 +979,10 @@ impl Tool for JobStatusTool { fn requires_sanitization(&self) -> bool { false } + + fn is_idempotent(&self) -> bool { + true + } } /// Tool for canceling a job. diff --git a/src/tools/builtin/json.rs b/src/tools/builtin/json.rs index 4f29fa38..d1c54b8d 100644 --- a/src/tools/builtin/json.rs +++ b/src/tools/builtin/json.rs @@ -132,6 +132,10 @@ impl Tool for JsonTool { fn requires_sanitization(&self) -> bool { false // Internal tool, no external data } + + fn is_idempotent(&self) -> bool { + true + } } fn parse_json_input(data: &serde_json::Value) -> Result { diff --git a/src/tools/builtin/memory.rs b/src/tools/builtin/memory.rs index 71fe8a3b..5109b92c 100644 --- a/src/tools/builtin/memory.rs +++ b/src/tools/builtin/memory.rs @@ -114,6 +114,10 @@ impl Tool for MemorySearchTool { fn requires_sanitization(&self) -> bool { false // Internal memory, trusted content } + + fn is_idempotent(&self) -> bool { + true + } } /// Tool for writing to workspace memory. @@ -377,6 +381,10 @@ impl Tool for MemoryReadTool { fn requires_sanitization(&self) -> bool { false // Internal memory } + + fn is_idempotent(&self) -> bool { + true + } } /// Tool for viewing workspace structure as a tree. @@ -496,6 +504,10 @@ impl Tool for MemoryTreeTool { fn requires_sanitization(&self) -> bool { false // Internal tool } + + fn is_idempotent(&self) -> bool { + true + } } #[cfg(all(test, feature = "postgres"))] diff --git a/src/tools/builtin/time.rs b/src/tools/builtin/time.rs index bafbd4d7..b188be6d 100644 --- a/src/tools/builtin/time.rs +++ b/src/tools/builtin/time.rs @@ -95,6 +95,10 @@ impl Tool for TimeTool { fn requires_sanitization(&self) -> bool { false // Internal tool, no external data } + + fn is_idempotent(&self) -> bool { + true + } } fn execute_now( diff --git a/src/tools/idempotency.rs b/src/tools/idempotency.rs new file mode 100644 index 00000000..e010f0c1 --- /dev/null +++ b/src/tools/idempotency.rs @@ -0,0 +1,270 @@ +//! Tool execution idempotency cache. +//! +//! Caches results of idempotent tool calls (same tool + same args = same result) +//! to avoid redundant re-execution during LLM retry loops, self-repair recovery, +//! and stuck job retries. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use lru::LruCache; +use sha2::{Digest, Sha256}; +use tokio::sync::Mutex; + +/// Cached tool result with expiration. +#[derive(Clone, Debug)] +struct CachedResult { + /// The serialized tool output string. + output: String, + /// When this entry was inserted. + inserted_at: Instant, +} + +/// Configuration for the idempotency cache. +#[derive(Debug, Clone)] +pub struct IdempotencyCacheConfig { + /// Maximum total entries across all jobs. Default: 2000. + pub max_entries: usize, + /// Time-to-live for cached entries. Default: 30 minutes. + pub ttl: Duration, +} + +impl Default for IdempotencyCacheConfig { + fn default() -> Self { + Self { + max_entries: 2000, + ttl: Duration::from_secs(30 * 60), + } + } +} + +/// Global idempotency cache for tool executions. +/// +/// Uses a single LRU cache with composite keys (job_id + tool_name + args_hash). +/// Entries expire after `ttl` and the total size is bounded by `max_entries`. +#[derive(Clone)] +pub struct ToolIdempotencyCache { + inner: Arc>>, + config: IdempotencyCacheConfig, +} + +impl ToolIdempotencyCache { + /// Create a new cache with the given configuration. + pub fn new(config: IdempotencyCacheConfig) -> Self { + let cap = std::num::NonZeroUsize::new(config.max_entries) + .unwrap_or(std::num::NonZeroUsize::new(1).expect("nonzero")); + Self { + inner: Arc::new(Mutex::new(LruCache::new(cap))), + config, + } + } + + /// Look up a cached result. Returns `None` if absent or expired. + pub async fn get( + &self, + job_id: &str, + tool_name: &str, + params: &serde_json::Value, + ) -> Option { + let key = Self::cache_key(job_id, tool_name, params); + let mut cache = self.inner.lock().await; + if let Some(entry) = cache.get(&key) { + if entry.inserted_at.elapsed() < self.config.ttl { + return Some(entry.output.clone()); + } + // Expired — remove it + cache.pop(&key); + } + None + } + + /// Store a successful tool result in the cache. + pub async fn put( + &self, + job_id: &str, + tool_name: &str, + params: &serde_json::Value, + output: String, + ) { + let key = Self::cache_key(job_id, tool_name, params); + let entry = CachedResult { + output, + inserted_at: Instant::now(), + }; + let mut cache = self.inner.lock().await; + cache.put(key, entry); + } + + /// Remove all cached entries for a specific job (call on job completion). + pub async fn invalidate_job(&self, job_id: &str) { + let prefix = format!("{}:", job_id); + let mut cache = self.inner.lock().await; + // Collect keys to remove (can't mutate while iterating) + let keys_to_remove: Vec = cache + .iter() + .filter_map(|(k, _)| { + if k.starts_with(&prefix) { + Some(k.clone()) + } else { + None + } + }) + .collect(); + for key in keys_to_remove { + cache.pop(&key); + } + } + + /// Build a deterministic cache key from job_id, tool name, and params. + fn cache_key(job_id: &str, tool_name: &str, params: &serde_json::Value) -> String { + let mut hasher = Sha256::new(); + hasher.update(tool_name.as_bytes()); + hasher.update(b":"); + // Canonical JSON via serde_json::to_string (keys are sorted in serde_json maps) + let params_str = serde_json::to_string(params).unwrap_or_default(); + hasher.update(params_str.as_bytes()); + let hash = format!("{:x}", hasher.finalize()); + format!("{}:{}:{}", job_id, tool_name, hash) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config() -> IdempotencyCacheConfig { + IdempotencyCacheConfig { + max_entries: 10, + ttl: Duration::from_secs(60), + } + } + + #[tokio::test] + async fn test_cache_hit() { + let cache = ToolIdempotencyCache::new(config()); + let params = serde_json::json!({"path": "/etc/hosts"}); + cache + .put("job1", "read_file", ¶ms, "file contents".into()) + .await; + let result = cache.get("job1", "read_file", ¶ms).await; + assert_eq!(result, Some("file contents".into())); + } + + #[tokio::test] + async fn test_cache_miss() { + let cache = ToolIdempotencyCache::new(config()); + let params = serde_json::json!({"path": "/etc/hosts"}); + let result = cache.get("job1", "read_file", ¶ms).await; + assert_eq!(result, None); + } + + #[tokio::test] + async fn test_different_params_miss() { + let cache = ToolIdempotencyCache::new(config()); + let params1 = serde_json::json!({"path": "/etc/hosts"}); + let params2 = serde_json::json!({"path": "/etc/passwd"}); + cache + .put("job1", "read_file", ¶ms1, "hosts".into()) + .await; + let result = cache.get("job1", "read_file", ¶ms2).await; + assert_eq!(result, None); + } + + #[tokio::test] + async fn test_different_jobs_isolated() { + let cache = ToolIdempotencyCache::new(config()); + let params = serde_json::json!({"path": "/etc/hosts"}); + cache + .put("job1", "read_file", ¶ms, "from job1".into()) + .await; + let result = cache.get("job2", "read_file", ¶ms).await; + assert_eq!(result, None); + } + + #[tokio::test] + async fn test_invalidate_job() { + let cache = ToolIdempotencyCache::new(config()); + let params = serde_json::json!({"q": "test"}); + cache.put("job1", "echo", ¶ms, "echo1".into()).await; + cache + .put("job1", "time", &serde_json::json!({}), "now".into()) + .await; + cache.put("job2", "echo", ¶ms, "echo2".into()).await; + + cache.invalidate_job("job1").await; + + assert_eq!(cache.get("job1", "echo", ¶ms).await, None); + assert_eq!( + cache.get("job1", "time", &serde_json::json!({})).await, + None + ); + // job2 unaffected + assert_eq!( + cache.get("job2", "echo", ¶ms).await, + Some("echo2".into()) + ); + } + + #[tokio::test] + async fn test_ttl_expiry() { + let cache = ToolIdempotencyCache::new(IdempotencyCacheConfig { + max_entries: 10, + ttl: Duration::from_millis(1), + }); + let params = serde_json::json!({"x": 1}); + cache.put("job1", "echo", ¶ms, "val".into()).await; + tokio::time::sleep(Duration::from_millis(5)).await; + assert_eq!(cache.get("job1", "echo", ¶ms).await, None); + } + + #[tokio::test] + async fn test_lru_eviction() { + let cache = ToolIdempotencyCache::new(IdempotencyCacheConfig { + max_entries: 3, + ttl: Duration::from_secs(60), + }); + // Fill to capacity + for i in 0..3 { + let params = serde_json::json!({"i": i}); + cache.put("job1", "echo", ¶ms, format!("val{i}")).await; + } + // Insert one more, evicting the oldest (i=0) + let params_new = serde_json::json!({"i": 99}); + cache.put("job1", "echo", ¶ms_new, "val99".into()).await; + + assert_eq!( + cache + .get("job1", "echo", &serde_json::json!({"i": 0})) + .await, + None + ); + assert_eq!( + cache + .get("job1", "echo", &serde_json::json!({"i": 2})) + .await, + Some("val2".into()) + ); + assert_eq!( + cache.get("job1", "echo", ¶ms_new).await, + Some("val99".into()) + ); + } + + #[tokio::test] + async fn test_cache_key_determinism() { + let key1 = + ToolIdempotencyCache::cache_key("j1", "echo", &serde_json::json!({"a": 1, "b": 2})); + let key2 = + ToolIdempotencyCache::cache_key("j1", "echo", &serde_json::json!({"a": 1, "b": 2})); + assert_eq!(key1, key2); + } + + #[tokio::test] + async fn test_overwrite_existing_entry() { + let cache = ToolIdempotencyCache::new(config()); + let params = serde_json::json!({"x": 1}); + cache.put("job1", "echo", ¶ms, "old".into()).await; + cache.put("job1", "echo", ¶ms, "new".into()).await; + assert_eq!(cache.get("job1", "echo", ¶ms).await, Some("new".into())); + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index d379d474..41cf357b 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -9,6 +9,7 @@ pub mod builder; pub mod builtin; +pub mod idempotency; pub mod mcp; pub mod rate_limiter; pub mod schema_validator; diff --git a/src/tools/tool.rs b/src/tools/tool.rs index 2e1b5183..172925a3 100644 --- a/src/tools/tool.rs +++ b/src/tools/tool.rs @@ -312,6 +312,17 @@ pub trait Tool: Send + Sync { &[] } + /// Whether this tool produces the same output for the same input. + /// + /// Idempotent tools (read-only, pure functions) can have their results cached + /// to avoid re-execution when the LLM re-requests the same tool with identical + /// arguments (common during self-repair recovery or retry loops). + /// + /// Default: `false`. Override to return `true` for read-only tools. + fn is_idempotent(&self) -> bool { + false + } + /// Per-invocation rate limit for this tool. /// /// Return `Some(config)` to throttle how often this tool can be called per user. diff --git a/tests/support/test_rig.rs b/tests/support/test_rig.rs index bedc6d4a..865dfd40 100644 --- a/tests/support/test_rig.rs +++ b/tests/support/test_rig.rs @@ -624,6 +624,9 @@ impl TestRigBuilder { }, transcription: None, document_extraction: None, + idempotency_cache: ironclaw::tools::idempotency::ToolIdempotencyCache::new( + ironclaw::tools::idempotency::IdempotencyCacheConfig::default(), + ), }; // 7. Create TestChannel and ChannelManager.