mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
322a048b39 |
@@ -385,6 +385,9 @@ async fn run_task_isolated(params: TaskRunParams<'_>) -> TaskResult {
|
|||||||
ironclaw::agent::cost_guard::CostGuardConfig::default(),
|
ironclaw::agent::cost_guard::CostGuardConfig::default(),
|
||||||
));
|
));
|
||||||
|
|
||||||
|
let idempotency_cache = Arc::new(ironclaw::tools::ToolIdempotencyCache::new(
|
||||||
|
ironclaw::tools::IdempotencyCacheConfig::default(),
|
||||||
|
));
|
||||||
let deps = AgentDeps {
|
let deps = AgentDeps {
|
||||||
store: None,
|
store: None,
|
||||||
llm: instrumented.clone() as Arc<dyn LlmProvider>,
|
llm: instrumented.clone() as Arc<dyn LlmProvider>,
|
||||||
@@ -397,6 +400,7 @@ async fn run_task_isolated(params: TaskRunParams<'_>) -> TaskResult {
|
|||||||
skills_config: ironclaw::config::SkillsConfig::default(),
|
skills_config: ironclaw::config::SkillsConfig::default(),
|
||||||
hooks: Arc::new(ironclaw::hooks::HookRegistry::new()),
|
hooks: Arc::new(ironclaw::hooks::HookRegistry::new()),
|
||||||
cost_guard,
|
cost_guard,
|
||||||
|
idempotency_cache,
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut channels = ChannelManager::new();
|
let mut channels = ChannelManager::new();
|
||||||
|
|||||||
@@ -28,7 +28,7 @@ use crate::hooks::HookRegistry;
|
|||||||
use crate::llm::LlmProvider;
|
use crate::llm::LlmProvider;
|
||||||
use crate::safety::SafetyLayer;
|
use crate::safety::SafetyLayer;
|
||||||
use crate::skills::SkillRegistry;
|
use crate::skills::SkillRegistry;
|
||||||
use crate::tools::ToolRegistry;
|
use crate::tools::{ToolIdempotencyCache, ToolRegistry};
|
||||||
use crate::workspace::Workspace;
|
use crate::workspace::Workspace;
|
||||||
|
|
||||||
/// Collapse a tool output string into a single-line preview for display.
|
/// Collapse a tool output string into a single-line preview for display.
|
||||||
@@ -72,6 +72,8 @@ pub struct AgentDeps {
|
|||||||
pub hooks: Arc<HookRegistry>,
|
pub hooks: Arc<HookRegistry>,
|
||||||
/// Cost enforcement guardrails (daily budget, hourly rate limits).
|
/// Cost enforcement guardrails (daily budget, hourly rate limits).
|
||||||
pub cost_guard: Arc<crate::agent::cost_guard::CostGuard>,
|
pub cost_guard: Arc<crate::agent::cost_guard::CostGuard>,
|
||||||
|
/// Idempotency cache for tool executions.
|
||||||
|
pub idempotency_cache: Arc<ToolIdempotencyCache>,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The main agent that coordinates all components.
|
/// The main agent that coordinates all components.
|
||||||
@@ -115,6 +117,7 @@ impl Agent {
|
|||||||
deps.tools.clone(),
|
deps.tools.clone(),
|
||||||
deps.store.clone(),
|
deps.store.clone(),
|
||||||
deps.hooks.clone(),
|
deps.hooks.clone(),
|
||||||
|
deps.idempotency_cache.clone(),
|
||||||
));
|
));
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
|
|||||||
@@ -478,6 +478,24 @@ impl Agent {
|
|||||||
.into());
|
.into());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Check idempotency cache before executing.
|
||||||
|
// Chat tools use the job_ctx.job_id (an ephemeral UUID per chat turn).
|
||||||
|
if tool.is_idempotent()
|
||||||
|
&& let Some(cached) = self
|
||||||
|
.deps
|
||||||
|
.idempotency_cache
|
||||||
|
.get(job_ctx.job_id, tool_name, params)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
return serde_json::to_string_pretty(&cached.result).map_err(|e| {
|
||||||
|
crate::error::ToolError::ExecutionFailed {
|
||||||
|
name: tool_name.to_string(),
|
||||||
|
reason: format!("Failed to serialize cached result: {}", e),
|
||||||
|
}
|
||||||
|
.into()
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
tool = %tool_name,
|
tool = %tool_name,
|
||||||
params = %params,
|
params = %params,
|
||||||
@@ -495,6 +513,14 @@ impl Agent {
|
|||||||
|
|
||||||
match &result {
|
match &result {
|
||||||
Ok(Ok(output)) => {
|
Ok(Ok(output)) => {
|
||||||
|
// Cache successful results for idempotent tools
|
||||||
|
if tool.is_idempotent() {
|
||||||
|
self.deps
|
||||||
|
.idempotency_cache
|
||||||
|
.put(job_ctx.job_id, tool_name, params, output.clone())
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
let result_str = serde_json::to_string(&output.result)
|
let result_str = serde_json::to_string(&output.result)
|
||||||
.unwrap_or_else(|_| "<serialize error>".to_string());
|
.unwrap_or_else(|_| "<serialize error>".to_string());
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ use crate::error::{Error, JobError};
|
|||||||
use crate::hooks::HookRegistry;
|
use crate::hooks::HookRegistry;
|
||||||
use crate::llm::LlmProvider;
|
use crate::llm::LlmProvider;
|
||||||
use crate::safety::SafetyLayer;
|
use crate::safety::SafetyLayer;
|
||||||
use crate::tools::ToolRegistry;
|
use crate::tools::{ToolIdempotencyCache, ToolRegistry};
|
||||||
|
|
||||||
/// Message to send to a worker.
|
/// Message to send to a worker.
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
@@ -51,6 +51,7 @@ pub struct Scheduler {
|
|||||||
tools: Arc<ToolRegistry>,
|
tools: Arc<ToolRegistry>,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<Arc<dyn Database>>,
|
||||||
hooks: Arc<HookRegistry>,
|
hooks: Arc<HookRegistry>,
|
||||||
|
idempotency_cache: Arc<ToolIdempotencyCache>,
|
||||||
/// Running jobs (main LLM-driven jobs).
|
/// Running jobs (main LLM-driven jobs).
|
||||||
jobs: Arc<RwLock<HashMap<Uuid, ScheduledJob>>>,
|
jobs: Arc<RwLock<HashMap<Uuid, ScheduledJob>>>,
|
||||||
/// Running sub-tasks (tool executions, background tasks).
|
/// Running sub-tasks (tool executions, background tasks).
|
||||||
@@ -59,6 +60,7 @@ pub struct Scheduler {
|
|||||||
|
|
||||||
impl Scheduler {
|
impl Scheduler {
|
||||||
/// Create a new scheduler.
|
/// Create a new scheduler.
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
pub fn new(
|
pub fn new(
|
||||||
config: AgentConfig,
|
config: AgentConfig,
|
||||||
context_manager: Arc<ContextManager>,
|
context_manager: Arc<ContextManager>,
|
||||||
@@ -67,6 +69,7 @@ impl Scheduler {
|
|||||||
tools: Arc<ToolRegistry>,
|
tools: Arc<ToolRegistry>,
|
||||||
store: Option<Arc<dyn Database>>,
|
store: Option<Arc<dyn Database>>,
|
||||||
hooks: Arc<HookRegistry>,
|
hooks: Arc<HookRegistry>,
|
||||||
|
idempotency_cache: Arc<ToolIdempotencyCache>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
Self {
|
Self {
|
||||||
config,
|
config,
|
||||||
@@ -76,6 +79,7 @@ impl Scheduler {
|
|||||||
tools,
|
tools,
|
||||||
store,
|
store,
|
||||||
hooks,
|
hooks,
|
||||||
|
idempotency_cache,
|
||||||
jobs: Arc::new(RwLock::new(HashMap::new())),
|
jobs: Arc::new(RwLock::new(HashMap::new())),
|
||||||
subtasks: Arc::new(RwLock::new(HashMap::new())),
|
subtasks: Arc::new(RwLock::new(HashMap::new())),
|
||||||
}
|
}
|
||||||
@@ -123,6 +127,7 @@ impl Scheduler {
|
|||||||
tools: self.tools.clone(),
|
tools: self.tools.clone(),
|
||||||
store: self.store.clone(),
|
store: self.store.clone(),
|
||||||
hooks: self.hooks.clone(),
|
hooks: self.hooks.clone(),
|
||||||
|
idempotency_cache: self.idempotency_cache.clone(),
|
||||||
timeout: self.config.job_timeout,
|
timeout: self.config.job_timeout,
|
||||||
use_planning: self.config.use_planning,
|
use_planning: self.config.use_planning,
|
||||||
};
|
};
|
||||||
|
|||||||
+43
-1
@@ -17,7 +17,7 @@ use crate::llm::{
|
|||||||
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection,
|
ActionPlan, ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult, ToolSelection,
|
||||||
};
|
};
|
||||||
use crate::safety::SafetyLayer;
|
use crate::safety::SafetyLayer;
|
||||||
use crate::tools::ToolRegistry;
|
use crate::tools::{ToolIdempotencyCache, ToolRegistry};
|
||||||
|
|
||||||
/// Shared dependencies for worker execution.
|
/// Shared dependencies for worker execution.
|
||||||
///
|
///
|
||||||
@@ -31,6 +31,7 @@ pub struct WorkerDeps {
|
|||||||
pub tools: Arc<ToolRegistry>,
|
pub tools: Arc<ToolRegistry>,
|
||||||
pub store: Option<Arc<dyn Database>>,
|
pub store: Option<Arc<dyn Database>>,
|
||||||
pub hooks: Arc<HookRegistry>,
|
pub hooks: Arc<HookRegistry>,
|
||||||
|
pub idempotency_cache: Arc<ToolIdempotencyCache>,
|
||||||
pub timeout: Duration,
|
pub timeout: Duration,
|
||||||
pub use_planning: bool,
|
pub use_planning: bool,
|
||||||
}
|
}
|
||||||
@@ -154,6 +155,12 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Free cached tool results for this job
|
||||||
|
self.deps
|
||||||
|
.idempotency_cache
|
||||||
|
.invalidate_job(self.job_id)
|
||||||
|
.await;
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -454,6 +461,32 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
|||||||
.into());
|
.into());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Check idempotency cache before executing
|
||||||
|
if tool.is_idempotent()
|
||||||
|
&& let Some(cached) = deps.idempotency_cache.get(job_id, tool_name, ¶ms).await
|
||||||
|
{
|
||||||
|
// Record the cache hit in memory (fire-and-forget)
|
||||||
|
let _ = deps
|
||||||
|
.context_manager
|
||||||
|
.update_memory(job_id, |mem| {
|
||||||
|
let rec = mem.create_action(tool_name, params.clone()).succeed(
|
||||||
|
Some("[idempotency cache hit]".to_string()),
|
||||||
|
cached.result.clone(),
|
||||||
|
cached.duration,
|
||||||
|
);
|
||||||
|
mem.record_action(rec);
|
||||||
|
})
|
||||||
|
.await;
|
||||||
|
|
||||||
|
return serde_json::to_string_pretty(&cached.result).map_err(|e| {
|
||||||
|
crate::error::ToolError::ExecutionFailed {
|
||||||
|
name: tool_name.to_string(),
|
||||||
|
reason: format!("Failed to serialize cached result: {}", e),
|
||||||
|
}
|
||||||
|
.into()
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
tracing::debug!(
|
tracing::debug!(
|
||||||
tool = %tool_name,
|
tool = %tool_name,
|
||||||
params = %params,
|
params = %params,
|
||||||
@@ -499,6 +532,15 @@ Report when the job is complete or if you encounter issues you cannot resolve."#
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Cache successful results for idempotent tools
|
||||||
|
if let Ok(Ok(output)) = &result
|
||||||
|
&& tool.is_idempotent()
|
||||||
|
{
|
||||||
|
deps.idempotency_cache
|
||||||
|
.put(job_id, tool_name, ¶ms, output.clone())
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
|
||||||
// Record action in memory and get the ActionRecord for persistence
|
// Record action in memory and get the ActionRecord for persistence
|
||||||
let action = match &result {
|
let action = match &result {
|
||||||
Ok(Ok(output)) => {
|
Ok(Ok(output)) => {
|
||||||
|
|||||||
@@ -1372,6 +1372,9 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
max_actions_per_hour: config.agent.max_actions_per_hour,
|
max_actions_per_hour: config.agent.max_actions_per_hour,
|
||||||
},
|
},
|
||||||
));
|
));
|
||||||
|
let idempotency_cache = std::sync::Arc::new(ironclaw::tools::ToolIdempotencyCache::new(
|
||||||
|
ironclaw::tools::IdempotencyCacheConfig::default(),
|
||||||
|
));
|
||||||
let deps = AgentDeps {
|
let deps = AgentDeps {
|
||||||
store: db,
|
store: db,
|
||||||
llm,
|
llm,
|
||||||
@@ -1384,6 +1387,7 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
skills_config: config.skills.clone(),
|
skills_config: config.skills.clone(),
|
||||||
hooks,
|
hooks,
|
||||||
cost_guard,
|
cost_guard,
|
||||||
|
idempotency_cache,
|
||||||
};
|
};
|
||||||
let agent = Agent::new(
|
let agent = Agent::new(
|
||||||
config.agent.clone(),
|
config.agent.clone(),
|
||||||
|
|||||||
@@ -282,6 +282,9 @@ impl TestHarnessBuilder {
|
|||||||
max_actions_per_hour: None,
|
max_actions_per_hour: None,
|
||||||
}));
|
}));
|
||||||
|
|
||||||
|
let idempotency_cache = Arc::new(crate::tools::ToolIdempotencyCache::new(
|
||||||
|
crate::tools::IdempotencyCacheConfig::default(),
|
||||||
|
));
|
||||||
let deps = AgentDeps {
|
let deps = AgentDeps {
|
||||||
store: Some(Arc::clone(&db)),
|
store: Some(Arc::clone(&db)),
|
||||||
llm,
|
llm,
|
||||||
@@ -294,6 +297,7 @@ impl TestHarnessBuilder {
|
|||||||
skills_config: SkillsConfig::default(),
|
skills_config: SkillsConfig::default(),
|
||||||
hooks,
|
hooks,
|
||||||
cost_guard,
|
cost_guard,
|
||||||
|
idempotency_cache,
|
||||||
};
|
};
|
||||||
|
|
||||||
TestHarness {
|
TestHarness {
|
||||||
|
|||||||
@@ -46,4 +46,8 @@ impl Tool for EchoTool {
|
|||||||
fn requires_sanitization(&self) -> bool {
|
fn requires_sanitization(&self) -> bool {
|
||||||
false // Internal tool, no external data
|
false // Internal tool, no external data
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Pure function: same input always produces same output
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -269,6 +269,10 @@ impl Tool for ReadFileTool {
|
|||||||
true // Reading local files should require approval
|
true // Reading local files should require approval
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Read-only file access, safe to cache within TTL
|
||||||
|
}
|
||||||
|
|
||||||
fn domain(&self) -> ToolDomain {
|
fn domain(&self) -> ToolDomain {
|
||||||
ToolDomain::Container
|
ToolDomain::Container
|
||||||
}
|
}
|
||||||
@@ -492,6 +496,10 @@ impl Tool for ListDirTool {
|
|||||||
true // Directory listings can leak filesystem structure
|
true // Directory listings can leak filesystem structure
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Read-only directory listing, safe to cache within TTL
|
||||||
|
}
|
||||||
|
|
||||||
fn domain(&self) -> ToolDomain {
|
fn domain(&self) -> ToolDomain {
|
||||||
ToolDomain::Container
|
ToolDomain::Container
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -845,6 +845,10 @@ impl Tool for ListJobsTool {
|
|||||||
fn requires_sanitization(&self) -> bool {
|
fn requires_sanitization(&self) -> bool {
|
||||||
false
|
false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Read-only job listing, safe to cache within TTL
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Tool for checking job status.
|
/// Tool for checking job status.
|
||||||
@@ -924,6 +928,10 @@ impl Tool for JobStatusTool {
|
|||||||
fn requires_sanitization(&self) -> bool {
|
fn requires_sanitization(&self) -> bool {
|
||||||
false
|
false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Read-only status check, safe to cache within TTL
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Tool for canceling a job.
|
/// Tool for canceling a job.
|
||||||
|
|||||||
@@ -102,6 +102,10 @@ impl Tool for JsonTool {
|
|||||||
fn requires_sanitization(&self) -> bool {
|
fn requires_sanitization(&self) -> bool {
|
||||||
false // Internal tool, no external data
|
false // Internal tool, no external data
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Pure transform: same JSON in, same result out
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn parse_json_input(data: &serde_json::Value) -> Result<serde_json::Value, ToolError> {
|
fn parse_json_input(data: &serde_json::Value) -> Result<serde_json::Value, ToolError> {
|
||||||
|
|||||||
@@ -112,6 +112,10 @@ impl Tool for MemorySearchTool {
|
|||||||
fn requires_sanitization(&self) -> bool {
|
fn requires_sanitization(&self) -> bool {
|
||||||
false // Internal memory, trusted content
|
false // Internal memory, trusted content
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Read-only search, safe to cache
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Tool for writing to workspace memory.
|
/// Tool for writing to workspace memory.
|
||||||
@@ -350,6 +354,10 @@ impl Tool for MemoryReadTool {
|
|||||||
fn requires_sanitization(&self) -> bool {
|
fn requires_sanitization(&self) -> bool {
|
||||||
false // Internal memory
|
false // Internal memory
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Read-only file access, safe to cache
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Tool for viewing workspace structure as a tree.
|
/// Tool for viewing workspace structure as a tree.
|
||||||
@@ -469,6 +477,10 @@ impl Tool for MemoryTreeTool {
|
|||||||
fn requires_sanitization(&self) -> bool {
|
fn requires_sanitization(&self) -> bool {
|
||||||
false // Internal tool
|
false // Internal tool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // Read-only tree listing, safe to cache
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(all(test, feature = "postgres"))]
|
#[cfg(all(test, feature = "postgres"))]
|
||||||
|
|||||||
@@ -111,4 +111,8 @@ impl Tool for TimeTool {
|
|||||||
fn requires_sanitization(&self) -> bool {
|
fn requires_sanitization(&self) -> bool {
|
||||||
false // Internal tool, no external data
|
false // Internal tool, no external data
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
true // TTL handles staleness for time-dependent results
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,367 @@
|
|||||||
|
//! In-memory idempotency cache for tool executions.
|
||||||
|
//!
|
||||||
|
//! For tools that declare `is_idempotent() == true`, successful results are
|
||||||
|
//! cached per-job so repeated identical calls (common during self-repair
|
||||||
|
//! recovery, stuck job retries, or chat-mode retry loops) return instantly
|
||||||
|
//! without re-executing.
|
||||||
|
//!
|
||||||
|
//! ```text
|
||||||
|
//! ┌──────────────────────────────────────────────────────────┐
|
||||||
|
//! │ ToolIdempotencyCache │
|
||||||
|
//! │ │
|
||||||
|
//! │ get(job_id, tool_name, args) -> Option<ToolOutput> │
|
||||||
|
//! │ put(job_id, tool_name, args, output) │
|
||||||
|
//! │ invalidate_job(job_id) // cleanup on job completion │
|
||||||
|
//! │ │
|
||||||
|
//! │ Internal: Mutex<HashMap<(Uuid, CacheKey), CacheEntry>> │
|
||||||
|
//! │ Key: sha256(tool_name | canonical_json(args)) │
|
||||||
|
//! │ Scoped by job_id │
|
||||||
|
//! │ TTL + max entries per job with LRU eviction │
|
||||||
|
//! └──────────────────────────────────────────────────────────┘
|
||||||
|
//! ```
|
||||||
|
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
use tokio::sync::Mutex;
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::tools::ToolOutput;
|
||||||
|
|
||||||
|
/// Configuration for the idempotency cache.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct IdempotencyCacheConfig {
|
||||||
|
/// Time-to-live for cache entries.
|
||||||
|
pub ttl: Duration,
|
||||||
|
/// Maximum number of cached entries per job before LRU eviction.
|
||||||
|
pub max_entries_per_job: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for IdempotencyCacheConfig {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self {
|
||||||
|
ttl: Duration::from_secs(30 * 60), // 30 minutes
|
||||||
|
max_entries_per_job: 500,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// SHA-256 hex digest used as cache key.
|
||||||
|
type CacheKey = String;
|
||||||
|
|
||||||
|
struct CacheEntry {
|
||||||
|
output: ToolOutput,
|
||||||
|
created_at: Instant,
|
||||||
|
last_accessed: Instant,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Per-job idempotency cache for tool results.
|
||||||
|
///
|
||||||
|
/// Only caches `Ok(ToolOutput)` results. Errors are never cached so retries
|
||||||
|
/// after transient failures get a fresh execution.
|
||||||
|
pub struct ToolIdempotencyCache {
|
||||||
|
/// Map from (job_id, cache_key) -> cached result.
|
||||||
|
entries: Mutex<HashMap<(Uuid, CacheKey), CacheEntry>>,
|
||||||
|
config: IdempotencyCacheConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ToolIdempotencyCache {
|
||||||
|
/// Create a new cache with the given configuration.
|
||||||
|
pub fn new(config: IdempotencyCacheConfig) -> Self {
|
||||||
|
Self {
|
||||||
|
entries: Mutex::new(HashMap::new()),
|
||||||
|
config,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a deterministic cache key from a tool name and its arguments.
|
||||||
|
///
|
||||||
|
/// `sha256(tool_name | canonical_json(args))`. serde_json produces
|
||||||
|
/// stable output for the same input structure.
|
||||||
|
fn cache_key(tool_name: &str, args: &serde_json::Value) -> CacheKey {
|
||||||
|
let mut hasher = Sha256::new();
|
||||||
|
hasher.update(tool_name.as_bytes());
|
||||||
|
hasher.update(b"|");
|
||||||
|
if let Ok(json) = serde_json::to_string(args) {
|
||||||
|
hasher.update(json.as_bytes());
|
||||||
|
}
|
||||||
|
format!("{:x}", hasher.finalize())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Look up a cached result for a tool invocation within a job.
|
||||||
|
///
|
||||||
|
/// Returns `Some(ToolOutput)` on a cache hit (within TTL), `None` on miss.
|
||||||
|
pub async fn get(
|
||||||
|
&self,
|
||||||
|
job_id: Uuid,
|
||||||
|
tool_name: &str,
|
||||||
|
args: &serde_json::Value,
|
||||||
|
) -> Option<ToolOutput> {
|
||||||
|
let key = Self::cache_key(tool_name, args);
|
||||||
|
let now = Instant::now();
|
||||||
|
|
||||||
|
let mut guard = self.entries.lock().await;
|
||||||
|
let compound_key = (job_id, key);
|
||||||
|
|
||||||
|
if let Some(entry) = guard.get_mut(&compound_key) {
|
||||||
|
if now.duration_since(entry.created_at) < self.config.ttl {
|
||||||
|
entry.last_accessed = now;
|
||||||
|
tracing::debug!(
|
||||||
|
tool = %tool_name,
|
||||||
|
job = %job_id,
|
||||||
|
"idempotency cache hit"
|
||||||
|
);
|
||||||
|
return Some(entry.output.clone());
|
||||||
|
}
|
||||||
|
// Expired
|
||||||
|
guard.remove(&compound_key);
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::trace!(
|
||||||
|
tool = %tool_name,
|
||||||
|
job = %job_id,
|
||||||
|
"idempotency cache miss"
|
||||||
|
);
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Store a successful tool result in the cache.
|
||||||
|
///
|
||||||
|
/// Evicts expired entries and applies LRU eviction if the per-job limit
|
||||||
|
/// is exceeded.
|
||||||
|
pub async fn put(
|
||||||
|
&self,
|
||||||
|
job_id: Uuid,
|
||||||
|
tool_name: &str,
|
||||||
|
args: &serde_json::Value,
|
||||||
|
output: ToolOutput,
|
||||||
|
) {
|
||||||
|
let key = Self::cache_key(tool_name, args);
|
||||||
|
let now = Instant::now();
|
||||||
|
|
||||||
|
let mut guard = self.entries.lock().await;
|
||||||
|
|
||||||
|
// Evict expired entries for this job
|
||||||
|
guard.retain(|(jid, _), entry| {
|
||||||
|
*jid != job_id || now.duration_since(entry.created_at) < self.config.ttl
|
||||||
|
});
|
||||||
|
|
||||||
|
// Count entries for this job and evict LRU if over capacity
|
||||||
|
let job_count = guard.keys().filter(|(jid, _)| *jid == job_id).count();
|
||||||
|
if job_count >= self.config.max_entries_per_job {
|
||||||
|
// Find the LRU entry for this job
|
||||||
|
let oldest_key = guard
|
||||||
|
.iter()
|
||||||
|
.filter(|((jid, _), _)| *jid == job_id)
|
||||||
|
.min_by_key(|(_, entry)| entry.last_accessed)
|
||||||
|
.map(|(k, _)| k.clone());
|
||||||
|
|
||||||
|
if let Some(k) = oldest_key {
|
||||||
|
guard.remove(&k);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::trace!(
|
||||||
|
tool = %tool_name,
|
||||||
|
job = %job_id,
|
||||||
|
"idempotency cache store"
|
||||||
|
);
|
||||||
|
|
||||||
|
guard.insert(
|
||||||
|
(job_id, key),
|
||||||
|
CacheEntry {
|
||||||
|
output,
|
||||||
|
created_at: now,
|
||||||
|
last_accessed: now,
|
||||||
|
},
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Remove all cached entries for a job.
|
||||||
|
///
|
||||||
|
/// Call this when a job completes or fails to free memory.
|
||||||
|
pub async fn invalidate_job(&self, job_id: Uuid) {
|
||||||
|
let mut guard = self.entries.lock().await;
|
||||||
|
let before = guard.len();
|
||||||
|
guard.retain(|(jid, _), _| *jid != job_id);
|
||||||
|
let removed = before - guard.len();
|
||||||
|
if removed > 0 {
|
||||||
|
tracing::debug!(
|
||||||
|
job = %job_id,
|
||||||
|
removed,
|
||||||
|
"idempotency cache invalidated job"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Total number of entries across all jobs.
|
||||||
|
#[cfg(test)]
|
||||||
|
async fn len(&self) -> usize {
|
||||||
|
self.entries.lock().await.len()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::tools::ToolOutput;
|
||||||
|
use crate::tools::idempotency::{IdempotencyCacheConfig, ToolIdempotencyCache};
|
||||||
|
|
||||||
|
fn make_cache(ttl_ms: u64, max_per_job: usize) -> ToolIdempotencyCache {
|
||||||
|
ToolIdempotencyCache::new(IdempotencyCacheConfig {
|
||||||
|
ttl: Duration::from_millis(ttl_ms),
|
||||||
|
max_entries_per_job: max_per_job,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sample_output(text: &str) -> ToolOutput {
|
||||||
|
ToolOutput::text(text, Duration::from_millis(1))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cache_key_deterministic() {
|
||||||
|
let args = serde_json::json!({"query": "hello", "limit": 5});
|
||||||
|
let k1 = ToolIdempotencyCache::cache_key("memory_search", &args);
|
||||||
|
let k2 = ToolIdempotencyCache::cache_key("memory_search", &args);
|
||||||
|
assert_eq!(k1, k2);
|
||||||
|
assert_eq!(k1.len(), 64); // SHA-256 hex
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cache_key_varies_by_tool_name() {
|
||||||
|
let args = serde_json::json!({"query": "hello"});
|
||||||
|
let k1 = ToolIdempotencyCache::cache_key("memory_search", &args);
|
||||||
|
let k2 = ToolIdempotencyCache::cache_key("memory_read", &args);
|
||||||
|
assert_ne!(k1, k2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cache_key_varies_by_args() {
|
||||||
|
let k1 = ToolIdempotencyCache::cache_key("echo", &serde_json::json!({"message": "hello"}));
|
||||||
|
let k2 = ToolIdempotencyCache::cache_key("echo", &serde_json::json!({"message": "world"}));
|
||||||
|
assert_ne!(k1, k2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cache_hit_returns_stored_output() {
|
||||||
|
let cache = make_cache(60_000, 100);
|
||||||
|
let job = Uuid::new_v4();
|
||||||
|
let args = serde_json::json!({"message": "hi"});
|
||||||
|
|
||||||
|
// Miss
|
||||||
|
assert!(cache.get(job, "echo", &args).await.is_none());
|
||||||
|
|
||||||
|
// Store
|
||||||
|
cache.put(job, "echo", &args, sample_output("hi")).await;
|
||||||
|
|
||||||
|
// Hit
|
||||||
|
let hit = cache.get(job, "echo", &args).await;
|
||||||
|
assert!(hit.is_some());
|
||||||
|
assert_eq!(hit.unwrap().result, serde_json::json!("hi"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn cache_miss_for_different_job() {
|
||||||
|
let cache = make_cache(60_000, 100);
|
||||||
|
let job_a = Uuid::new_v4();
|
||||||
|
let job_b = Uuid::new_v4();
|
||||||
|
let args = serde_json::json!({"message": "hi"});
|
||||||
|
|
||||||
|
cache.put(job_a, "echo", &args, sample_output("hi")).await;
|
||||||
|
|
||||||
|
// Different job should miss
|
||||||
|
assert!(cache.get(job_b, "echo", &args).await.is_none());
|
||||||
|
// Same job should hit
|
||||||
|
assert!(cache.get(job_a, "echo", &args).await.is_some());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn ttl_expiry() {
|
||||||
|
let cache = make_cache(1, 100); // 1ms TTL
|
||||||
|
let job = Uuid::new_v4();
|
||||||
|
let args = serde_json::json!({"message": "hi"});
|
||||||
|
|
||||||
|
cache.put(job, "echo", &args, sample_output("hi")).await;
|
||||||
|
|
||||||
|
// Wait for TTL to expire
|
||||||
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||||
|
|
||||||
|
assert!(cache.get(job, "echo", &args).await.is_none());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn lru_eviction() {
|
||||||
|
let cache = make_cache(60_000, 2); // max 2 per job
|
||||||
|
let job = Uuid::new_v4();
|
||||||
|
|
||||||
|
// Fill with 2 entries
|
||||||
|
let args_a = serde_json::json!({"n": 1});
|
||||||
|
let args_b = serde_json::json!({"n": 2});
|
||||||
|
cache.put(job, "echo", &args_a, sample_output("a")).await;
|
||||||
|
cache.put(job, "echo", &args_b, sample_output("b")).await;
|
||||||
|
assert_eq!(cache.len().await, 2);
|
||||||
|
|
||||||
|
// Access args_a so args_b becomes the LRU
|
||||||
|
cache.get(job, "echo", &args_a).await;
|
||||||
|
|
||||||
|
// Add a third: should evict args_b (oldest accessed)
|
||||||
|
let args_c = serde_json::json!({"n": 3});
|
||||||
|
cache.put(job, "echo", &args_c, sample_output("c")).await;
|
||||||
|
assert_eq!(cache.len().await, 2);
|
||||||
|
|
||||||
|
// args_b should be gone, args_a and args_c should remain
|
||||||
|
assert!(cache.get(job, "echo", &args_b).await.is_none());
|
||||||
|
assert!(cache.get(job, "echo", &args_a).await.is_some());
|
||||||
|
assert!(cache.get(job, "echo", &args_c).await.is_some());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn invalidate_job_clears_entries() {
|
||||||
|
let cache = make_cache(60_000, 100);
|
||||||
|
let job_a = Uuid::new_v4();
|
||||||
|
let job_b = Uuid::new_v4();
|
||||||
|
let args = serde_json::json!({"message": "hi"});
|
||||||
|
|
||||||
|
cache.put(job_a, "echo", &args, sample_output("a")).await;
|
||||||
|
cache.put(job_b, "echo", &args, sample_output("b")).await;
|
||||||
|
assert_eq!(cache.len().await, 2);
|
||||||
|
|
||||||
|
cache.invalidate_job(job_a).await;
|
||||||
|
|
||||||
|
assert_eq!(cache.len().await, 1);
|
||||||
|
assert!(cache.get(job_a, "echo", &args).await.is_none());
|
||||||
|
assert!(cache.get(job_b, "echo", &args).await.is_some());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn lru_eviction_does_not_affect_other_jobs() {
|
||||||
|
let cache = make_cache(60_000, 1); // max 1 per job
|
||||||
|
let job_a = Uuid::new_v4();
|
||||||
|
let job_b = Uuid::new_v4();
|
||||||
|
|
||||||
|
let args_1 = serde_json::json!({"n": 1});
|
||||||
|
let args_2 = serde_json::json!({"n": 2});
|
||||||
|
|
||||||
|
cache.put(job_a, "echo", &args_1, sample_output("a1")).await;
|
||||||
|
cache.put(job_b, "echo", &args_1, sample_output("b1")).await;
|
||||||
|
|
||||||
|
// Adding a second entry for job_a should evict job_a's first entry,
|
||||||
|
// but not touch job_b's entry
|
||||||
|
cache.put(job_a, "echo", &args_2, sample_output("a2")).await;
|
||||||
|
|
||||||
|
assert!(cache.get(job_a, "echo", &args_1).await.is_none());
|
||||||
|
assert!(cache.get(job_a, "echo", &args_2).await.is_some());
|
||||||
|
assert!(cache.get(job_b, "echo", &args_1).await.is_some());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn default_config_is_reasonable() {
|
||||||
|
let cfg = IdempotencyCacheConfig::default();
|
||||||
|
assert_eq!(cfg.ttl, Duration::from_secs(30 * 60));
|
||||||
|
assert_eq!(cfg.max_entries_per_job, 500);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -9,6 +9,7 @@
|
|||||||
|
|
||||||
pub mod builder;
|
pub mod builder;
|
||||||
pub mod builtin;
|
pub mod builtin;
|
||||||
|
pub mod idempotency;
|
||||||
pub mod mcp;
|
pub mod mcp;
|
||||||
pub mod wasm;
|
pub mod wasm;
|
||||||
|
|
||||||
@@ -20,5 +21,6 @@ pub use builder::{
|
|||||||
LlmSoftwareBuilder, SoftwareBuilder, SoftwareType, Template, TemplateEngine, TemplateType,
|
LlmSoftwareBuilder, SoftwareBuilder, SoftwareType, Template, TemplateEngine, TemplateType,
|
||||||
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
|
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
|
||||||
};
|
};
|
||||||
|
pub use idempotency::{IdempotencyCacheConfig, ToolIdempotencyCache};
|
||||||
pub use registry::ToolRegistry;
|
pub use registry::ToolRegistry;
|
||||||
pub use tool::{Tool, ToolDomain, ToolError, ToolOutput};
|
pub use tool::{Tool, ToolDomain, ToolError, ToolOutput};
|
||||||
|
|||||||
@@ -194,6 +194,15 @@ pub trait Tool: Send + Sync {
|
|||||||
Duration::from_secs(60)
|
Duration::from_secs(60)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Whether this tool is idempotent (same args always produce the same result).
|
||||||
|
///
|
||||||
|
/// When true, successful results are cached per-job so repeated identical
|
||||||
|
/// calls return the cached result without re-executing. Tools that have
|
||||||
|
/// side effects (shell, file write, HTTP POST) should return false (the default).
|
||||||
|
fn is_idempotent(&self) -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
/// Where this tool should execute.
|
/// Where this tool should execute.
|
||||||
///
|
///
|
||||||
/// `Orchestrator` tools run in the main agent process (safe, no FS access).
|
/// `Orchestrator` tools run in the main agent process (safe, no FS access).
|
||||||
|
|||||||
Reference in New Issue
Block a user