use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use serde_json::Value;
use crate::message::{ToolContent, ToolContentPart};
use super::{ToolDispatchContext, ToolDispatchResult, ToolMiddleware, ToolPipeline};
pub trait PathExtractor: Send + Sync {
fn paths(&self, tool_name: &str, input: &Value) -> Vec<String>;
}
pub struct NoopPathExtractor;
impl PathExtractor for NoopPathExtractor {
fn paths(&self, _tool_name: &str, _input: &Value) -> Vec<String> {
Vec::new()
}
}
#[derive(Debug, Eq, PartialEq, Hash, Clone)]
struct CacheKey {
tool_name: String,
input_hash: u64,
}
struct CacheEntry {
result: ToolDispatchResult,
turn_inserted: usize,
paths: Vec<String>,
}
#[derive(Default)]
struct CacheState {
entries: HashMap<CacheKey, CacheEntry>,
write_epoch: u64,
}
pub struct MemoizingMiddleware {
cache: Mutex<CacheState>,
tools: Vec<String>,
write_tools: Vec<String>,
path_extractor: Arc<dyn PathExtractor>,
ttl_turns: u32,
}
impl MemoizingMiddleware {
#[must_use]
pub fn new(
tools: Vec<String>,
write_tools: Vec<String>,
path_extractor: Arc<dyn PathExtractor>,
ttl_turns: u32,
) -> Self {
Self {
cache: Mutex::new(CacheState::default()),
tools,
write_tools,
path_extractor,
ttl_turns,
}
}
}
impl std::fmt::Debug for MemoizingMiddleware {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let len = crate::error::recover_guard(self.cache.lock()).entries.len();
f.debug_struct("MemoizingMiddleware")
.field("cache_entries", &len)
.field("tools", &self.tools)
.field("write_tools", &self.write_tools)
.field("ttl_turns", &self.ttl_turns)
.finish_non_exhaustive()
}
}
impl ToolMiddleware for MemoizingMiddleware {
fn name(&self) -> &'static str {
"memoize"
}
fn dispatch<'a>(
&'a self,
ctx: &'a mut ToolDispatchContext,
next: &'a ToolPipeline,
) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + 'a>> {
let is_memoized = self.tools.iter().any(|t| t == &ctx.tool_name);
let is_write = self.write_tools.iter().any(|t| t == &ctx.tool_name);
let key = if is_memoized {
Some(make_key(&ctx.tool_name, &ctx.input))
} else {
None
};
let input_at_key = is_memoized.then(|| ctx.input.clone());
let ttl_turns = self.ttl_turns;
let current_turn = ctx.turn_number;
let epoch_at_dispatch = is_memoized.then(|| current_epoch(&self.cache));
Box::pin(async move {
if is_memoized
&& let Some(key) = key.as_ref()
&& let Some(cached) = lookup_fresh(&self.cache, key, current_turn, ttl_turns)
{
let mut result = cached;
append_cached_marker(&mut result.output);
result.tool_call_id.clone_from(&ctx.call_id);
return result;
}
let result = next.dispatch(ctx).await;
if is_write && !result.is_error {
let write_paths = self.path_extractor.paths(&ctx.tool_name, &ctx.input);
invalidate_paths(&self.cache, &write_paths);
} else if is_memoized && !result.is_error {
let input = input_at_key.as_ref().unwrap_or(&ctx.input);
let paths = self.path_extractor.paths(&ctx.tool_name, input);
if let (Some(key), Some(epoch)) = (key, epoch_at_dispatch) {
insert(&self.cache, key, result.clone(), current_turn, paths, epoch);
}
}
result
})
}
}
fn make_key(tool_name: &str, input: &Value) -> CacheKey {
let canonical = serde_json::to_string(input).unwrap_or_default();
let mut hasher = std::collections::hash_map::DefaultHasher::new();
canonical.hash(&mut hasher);
CacheKey {
tool_name: tool_name.to_string(),
input_hash: hasher.finish(),
}
}
fn lookup_fresh(
cache: &Mutex<CacheState>,
key: &CacheKey,
current_turn: usize,
ttl_turns: u32,
) -> Option<ToolDispatchResult> {
let mut guard = crate::error::recover_guard(cache.lock());
let expired =
|entry: &CacheEntry| current_turn.saturating_sub(entry.turn_inserted) >= ttl_turns as usize;
if let Some(entry) = guard.entries.get(key)
&& expired(entry)
{
guard.entries.remove(key);
return None;
}
guard.entries.get(key).map(|e| e.result.clone())
}
fn insert(
cache: &Mutex<CacheState>,
key: CacheKey,
result: ToolDispatchResult,
turn_inserted: usize,
paths: Vec<String>,
epoch_at_dispatch: u64,
) -> bool {
let mut guard = crate::error::recover_guard(cache.lock());
if guard.write_epoch != epoch_at_dispatch {
return false;
}
guard.entries.insert(
key,
CacheEntry {
result,
turn_inserted,
paths,
},
);
true
}
fn current_epoch(cache: &Mutex<CacheState>) -> u64 {
crate::error::recover_guard(cache.lock()).write_epoch
}
fn invalidate_paths(cache: &Mutex<CacheState>, write_paths: &[String]) {
if write_paths.is_empty() {
return;
}
let write_set: HashSet<&str> = write_paths.iter().map(String::as_str).collect();
let mut guard = crate::error::recover_guard(cache.lock());
guard
.entries
.retain(|_, entry| !entry.paths.iter().any(|p| write_set.contains(p.as_str())));
guard.write_epoch = guard.write_epoch.saturating_add(1);
}
fn append_cached_marker(output: &mut ToolContent) {
match output {
ToolContent::Text(s) => s.push_str("\n[cached]"),
ToolContent::Multipart(parts) => parts.push(ToolContentPart::Text {
text: "[cached]".to_string(),
}),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cancel::CancelSignal;
use crate::message::ToolContent;
use crate::middleware::{OutputLimitMiddleware, ToolDispatchContext, ToolPipeline};
use crate::tool::{PermissionCheck, ToolContext, ToolRegistry};
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
struct PathFromInput;
impl PathExtractor for PathFromInput {
fn paths(&self, _tool_name: &str, input: &Value) -> Vec<String> {
input
.get("path")
.and_then(Value::as_str)
.map(std::string::ToString::to_string)
.into_iter()
.collect()
}
}
struct NoPaths;
impl PathExtractor for NoPaths {
fn paths(&self, _: &str, _: &Value) -> Vec<String> {
Vec::new()
}
}
struct FixedOutputMiddleware {
output: ToolContent,
is_error: bool,
call_count: Arc<std::sync::atomic::AtomicUsize>,
}
impl ToolMiddleware for FixedOutputMiddleware {
fn name(&self) -> &'static str {
"fixed_output"
}
fn dispatch<'a>(
&'a self,
_ctx: &'a mut ToolDispatchContext,
_next: &'a ToolPipeline,
) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + 'a>> {
let output = self.output.clone();
let is_error = self.is_error;
let count = self.call_count.clone();
Box::pin(async move {
count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
ToolDispatchResult {
output,
is_error,
resolved_tool_name: String::new(),
tool_call_id: "fixture_call_id".to_string(),
duration: Duration::ZERO,
display_hint: None,
}
})
}
}
struct RoutingFixedOutputMiddleware {
read_output: ToolContent,
write_output: ToolContent,
write_is_error: bool,
call_count: Arc<std::sync::atomic::AtomicUsize>,
}
impl ToolMiddleware for RoutingFixedOutputMiddleware {
fn name(&self) -> &'static str {
"routing_fixed_output"
}
fn dispatch<'a>(
&'a self,
ctx: &'a mut ToolDispatchContext,
_next: &'a ToolPipeline,
) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + 'a>> {
let (output, is_error) = if ctx.tool_name == "Write" {
(self.write_output.clone(), self.write_is_error)
} else {
(self.read_output.clone(), false)
};
let count = self.call_count.clone();
Box::pin(async move {
count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
ToolDispatchResult {
output,
is_error,
resolved_tool_name: String::new(),
tool_call_id: "fixture_call_id".to_string(),
duration: Duration::ZERO,
display_hint: None,
}
})
}
}
fn make_middleware(extractor: Arc<dyn PathExtractor>, ttl: u32) -> MemoizingMiddleware {
MemoizingMiddleware::new(
vec!["Read".to_string()],
vec!["Write".to_string()],
extractor,
ttl,
)
}
fn pipeline(
mw: MemoizingMiddleware,
output: ToolContent,
is_error: bool,
) -> (ToolPipeline, Arc<std::sync::atomic::AtomicUsize>) {
let call_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let registry = Arc::new(ToolRegistry::new());
let pipeline = ToolPipeline::builder()
.with_middleware(mw)
.with_middleware(FixedOutputMiddleware {
output,
is_error,
call_count: call_count.clone(),
})
.with_core(registry)
.build()
.expect("pipeline builds");
(pipeline, call_count)
}
fn ctx_for(tool_name: &str, input: Value, turn: usize) -> ToolDispatchContext {
ToolDispatchContext {
tool_name: tool_name.to_string(),
input,
call_id: "c1".to_string(),
turn_number: turn,
cancel: Arc::new(CancelSignal::new()),
permission: PermissionCheck::Allow,
tool_context: ToolContext::default(),
}
}
#[tokio::test]
async fn repeat_read_returns_cached() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("file contents"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0);
let first = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 1);
let second = pipeline.dispatch(&mut ctx).await;
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
1,
"second call should hit the cache, not the inner dispatch"
);
assert!(
second.output.to_string().contains("[cached]"),
"second call output should carry [cached]: {}",
second.output
);
assert!(
!first.output.to_string().contains("[cached]"),
"first call output should not carry [cached]"
);
}
#[tokio::test]
async fn cache_hit_returns_requesting_call_id() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("file contents"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0);
ctx.call_id = "call_first".to_string();
let first = pipeline.dispatch(&mut ctx).await;
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 1);
ctx.call_id = "call_second".to_string();
let second = pipeline.dispatch(&mut ctx).await;
assert_eq!(
second.tool_call_id, "call_second",
"a cache hit must answer the requesting call, not replay the \
first call's id"
);
assert_eq!(
first.tool_call_id, "fixture_call_id",
"a miss passes the inner result through untouched"
);
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
1,
"the second call was served from the cache"
);
}
#[tokio::test]
async fn different_input_not_cached() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("content"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "a.rs"}), 0);
let _ = pipeline.dispatch(&mut ctx).await;
let mut ctx = ctx_for("Read", serde_json::json!({"path": "b.rs"}), 1);
let _ = pipeline.dispatch(&mut ctx).await;
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
2,
"different input should be a cache miss — inner dispatched twice"
);
}
#[tokio::test]
async fn write_invalidates_path_matching_cache() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("ok"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0);
let _ = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
let mut ctx = ctx_for("Write", serde_json::json!({"path": "bar.rs"}), 1);
let _ = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 2);
let r = pipeline.dispatch(&mut ctx).await;
assert!(
r.output.to_string().contains("[cached]"),
"Read(foo.rs) should still be cached after Write(bar.rs): {}",
r.output
);
let mut ctx = ctx_for("Write", serde_json::json!({"path": "foo.rs"}), 3);
let _ = pipeline.dispatch(&mut ctx).await;
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 4);
let r = pipeline.dispatch(&mut ctx).await;
assert!(
!r.output.to_string().contains("[cached]"),
"Read(foo.rs) should re-run after Write(foo.rs): {}",
r.output
);
}
#[tokio::test]
async fn ttl_expiry() {
let mw = make_middleware(Arc::new(PathFromInput), 2);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("v"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 0);
let _ = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 1);
let r = pipeline.dispatch(&mut ctx).await;
assert!(
r.output.to_string().contains("[cached]"),
"should be fresh at turn 1"
);
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 2);
let r = pipeline.dispatch(&mut ctx).await;
assert!(
!r.output.to_string().contains("[cached]"),
"should be expired at turn 2 (2-0 >= ttl_turns=2): {}",
r.output
);
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2);
}
#[tokio::test]
async fn non_memoized_tool_passes_through() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("grep result"), false);
let mut ctx = ctx_for("Grep", serde_json::json!({"pattern": "foo"}), 0);
let _ = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
let mut ctx = ctx_for("Grep", serde_json::json!({"pattern": "foo"}), 1);
let r = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2);
assert!(
!r.output.to_string().contains("[cached]"),
"non-memoized tool should never carry [cached]"
);
}
#[tokio::test]
async fn write_tool_result_not_cached() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("wrote"), false);
let mut ctx = ctx_for("Write", serde_json::json!({"path": "x"}), 0);
let r1 = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
assert!(
!r1.output.to_string().contains("[cached]"),
"write tool should never be cached"
);
let mut ctx = ctx_for("Write", serde_json::json!({"path": "x"}), 1);
let r2 = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2);
assert!(
!r2.output.to_string().contains("[cached]"),
"write tool should never be cached"
);
}
#[tokio::test]
async fn error_result_not_cached() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let (pipeline, calls) = pipeline(
mw,
ToolContent::from_string("file not found"),
true, );
let mut ctx = ctx_for("Read", serde_json::json!({"path": "missing"}), 0);
let _ = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "missing"}), 1);
let _ = pipeline.dispatch(&mut ctx).await;
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
2,
"error result should not be cached — second call re-ran"
);
}
#[tokio::test]
async fn cached_marker_format_text_and_multipart() {
let mw_text = make_middleware(Arc::new(NoPaths), 10);
let (p_text, _) = pipeline(mw_text, ToolContent::from_string("t"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "a"}), 0);
let _ = p_text.dispatch(&mut ctx).await;
let mut ctx = ctx_for("Read", serde_json::json!({"path": "a"}), 1);
let r = p_text.dispatch(&mut ctx).await;
match r.output {
ToolContent::Text(s) => assert!(s.ends_with("\n[cached]"), "text marker: {s}"),
ToolContent::Multipart(parts) => {
panic!("expected Text, got Multipart with {} parts", parts.len())
}
}
let mw2 = make_middleware(Arc::new(NoPaths), 10);
let existing = ToolContent::from_multipart(vec![ToolContentPart::Text {
text: "original".to_string(),
}]);
let (p_multi, _) = pipeline(mw2, existing, false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "b"}), 0);
let _ = p_multi.dispatch(&mut ctx).await;
let mut ctx = ctx_for("Read", serde_json::json!({"path": "b"}), 1);
let r = p_multi.dispatch(&mut ctx).await;
match r.output {
ToolContent::Multipart(parts) => {
assert_eq!(parts.len(), 2, "should have original + [cached] part");
match &parts[1] {
ToolContentPart::Text { text } => {
assert_eq!(text, "[cached]", "multipart marker part: {text}");
}
ToolContentPart::Image { .. } => panic!("expected Text part, got Image"),
}
}
ToolContent::Text(t) => panic!("expected Multipart, got Text: {t}"),
}
}
#[tokio::test]
async fn cache_key_canonicalizes_input_key_order() {
let mw = make_middleware(Arc::new(NoPaths), 10);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("ok"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x", "limit": 10}), 0);
let _ = pipeline.dispatch(&mut ctx).await;
let mut ctx = ctx_for("Read", serde_json::json!({"limit": 10, "path": "x"}), 1);
let r = pipeline.dispatch(&mut ctx).await;
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
1,
"different key order, same logical input should hit the cache"
);
assert!(
r.output.to_string().contains("[cached]"),
"should be cached despite key-order difference"
);
}
#[test]
fn memoize_middleware_name() {
let mw = MemoizingMiddleware::new(vec![], vec![], Arc::new(NoPaths), 5);
assert_eq!(mw.name(), "memoize");
}
#[test]
fn make_key_same_input_same_hash() {
let k1 = make_key("Read", &serde_json::json!({"path": "foo.rs"}));
let k2 = make_key("Read", &serde_json::json!({"path": "foo.rs"}));
assert_eq!(k1, k2, "identical inputs must produce identical keys");
}
#[test]
fn make_key_different_tool_different_hash() {
let k1 = make_key("Read", &serde_json::json!({"path": "foo.rs"}));
let k2 = make_key("Grep", &serde_json::json!({"path": "foo.rs"}));
assert_ne!(k1, k2, "different tool names must produce different keys");
}
#[test]
fn make_key_different_input_different_hash() {
let k1 = make_key("Read", &serde_json::json!({"path": "foo.rs"}));
let k2 = make_key("Read", &serde_json::json!({"path": "bar.rs"}));
assert_ne!(k1, k2, "different inputs must produce different keys");
}
#[test]
fn lookup_fresh_returns_none_on_miss() {
let cache = Mutex::new(CacheState::default());
let key = make_key("Read", &serde_json::json!({"path": "x"}));
let result = lookup_fresh(&cache, &key, 0, 10);
assert!(result.is_none(), "empty cache should miss");
}
#[test]
fn lookup_fresh_returns_result_on_hit() {
let cache = Mutex::new(CacheState::default());
let key = make_key("Read", &serde_json::json!({"path": "x"}));
insert(
&cache,
key.clone(),
ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("content"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
},
0,
vec![],
0,
);
let result = lookup_fresh(&cache, &key, 1, 10);
assert!(result.is_some(), "fresh entry should hit");
}
#[test]
fn lookup_fresh_expires_and_evicts_stale_entry() {
let cache = Mutex::new(CacheState::default());
let key = make_key("Read", &serde_json::json!({"path": "x"}));
insert(
&cache,
key.clone(),
ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("content"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
},
0,
vec![],
0,
);
assert!(
lookup_fresh(&cache, &key, 1, 2).is_some(),
"should be fresh at turn 1"
);
assert!(
lookup_fresh(&cache, &key, 2, 2).is_none(),
"should be expired at turn 2"
);
let guard = crate::error::recover_guard(cache.lock());
assert!(
guard.entries.is_empty(),
"expired entry should be evicted from cache"
);
}
#[test]
fn insert_stores_entry_and_replace_overwrites() {
let cache = Mutex::new(CacheState::default());
let key = make_key("Read", &serde_json::json!({"path": "x"}));
insert(
&cache,
key.clone(),
ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("first"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
},
0,
vec![],
0,
);
insert(
&cache,
key.clone(),
ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("second"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
},
1,
vec![],
0,
);
let result = lookup_fresh(&cache, &key, 2, 10).unwrap();
match result.output {
ToolContent::Text(s) => assert_eq!(s, "second", "second insert should overwrite"),
ToolContent::Multipart(parts) => {
panic!("expected Text, got Multipart with {} parts", parts.len())
}
}
}
#[test]
fn invalidate_paths_removes_overlapping_entries() {
let cache = Mutex::new(CacheState::default());
let key_a = make_key("Read", &serde_json::json!({"path": "a.rs"}));
let key_b = make_key("Read", &serde_json::json!({"path": "b.rs"}));
let result = ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("content"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
};
insert(
&cache,
key_a.clone(),
result.clone(),
0,
vec!["a.rs".to_string()],
0,
);
insert(
&cache,
key_b.clone(),
result,
0,
vec!["b.rs".to_string()],
0,
);
invalidate_paths(&cache, &["a.rs".to_string()]);
let guard = crate::error::recover_guard(cache.lock());
assert!(
!guard.entries.contains_key(&key_a),
"overlapping entry should be evicted"
);
assert!(
guard.entries.contains_key(&key_b),
"non-overlapping entry should survive"
);
}
#[test]
fn invalidate_paths_noop_on_empty_write_paths() {
let cache = Mutex::new(CacheState::default());
let key = make_key("Read", &serde_json::json!({"path": "x"}));
insert(
&cache,
key,
ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("content"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
},
0,
vec!["x".to_string()],
0,
);
invalidate_paths(&cache, &[]);
let guard = crate::error::recover_guard(cache.lock());
assert_eq!(
guard.entries.len(),
1,
"empty write_paths should evict nothing"
);
}
#[test]
fn invalidate_paths_evicts_all_matching() {
let cache = Mutex::new(CacheState::default());
let result = ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("content"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
};
let k1 = make_key(
"Read",
&serde_json::json!({"path": "shared.rs", "limit": 10}),
);
let k2 = make_key(
"Grep",
&serde_json::json!({"path": "shared.rs", "pattern": "foo"}),
);
insert(
&cache,
k1,
result.clone(),
0,
vec!["shared.rs".to_string()],
0,
);
insert(&cache, k2, result, 0, vec!["shared.rs".to_string()], 0);
invalidate_paths(&cache, &["shared.rs".to_string()]);
let guard = crate::error::recover_guard(cache.lock());
assert!(
guard.entries.is_empty(),
"both entries touch shared.rs — both evicted"
);
}
#[tokio::test]
async fn ttl_of_zero_caches_nothing() {
let mw = make_middleware(Arc::new(PathFromInput), 0);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("file contents"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0);
let first = pipeline.dispatch(&mut ctx).await;
assert!(
matches!(&first.output, ToolContent::Text(t) if t.contains("file contents")),
"the first call executes the tool"
);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0);
let second = pipeline.dispatch(&mut ctx).await;
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
2,
"a ttl of zero expires every entry immediately — no call is ever served from the cache"
);
assert!(
matches!(&second.output, ToolContent::Text(t) if t.contains("file contents")),
"the second call re-executes the tool"
);
}
#[test]
fn insert_after_invalidation_is_skipped() {
let cache = Mutex::new(CacheState::default());
let key = make_key("Read", &serde_json::json!({"path": "x"}));
let result = ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("content"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
};
let epoch = current_epoch(&cache);
invalidate_paths(&cache, &["x".to_string()]);
assert_eq!(
current_epoch(&cache),
epoch.saturating_add(1),
"an eviction bumps the invalidation epoch"
);
assert!(
!insert(&cache, key, result, 0, vec!["x".to_string()], epoch),
"an insert whose dispatch predated the invalidation must be skipped"
);
assert!(
crate::error::recover_guard(cache.lock()).entries.is_empty(),
"no entry may land after an unseen eviction"
);
}
#[test]
fn insert_at_the_current_epoch_is_stored() {
let cache = Mutex::new(CacheState::default());
let key = make_key("Read", &serde_json::json!({"path": "x"}));
let result = ToolDispatchResult {
tool_call_id: String::new(),
output: ToolContent::from_string("content"),
is_error: false,
duration: Duration::ZERO,
resolved_tool_name: String::new(),
display_hint: None,
};
let epoch = current_epoch(&cache);
assert!(
insert(&cache, key.clone(), result, 0, vec!["x".to_string()], epoch),
"an insert at the current epoch must be stored"
);
assert!(
crate::error::recover_guard(cache.lock())
.entries
.contains_key(&key),
"the entry must be present"
);
}
#[tokio::test]
async fn cache_hit_through_outer_cap_re_applies_the_limit() {
struct LongOutputMiddleware;
impl ToolMiddleware for LongOutputMiddleware {
fn name(&self) -> &'static str {
"long_output"
}
fn dispatch<'a>(
&'a self,
_ctx: &'a mut ToolDispatchContext,
_next: &'a ToolPipeline,
) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + 'a>> {
Box::pin(async {
ToolDispatchResult {
output: ToolContent::from_string("x".repeat(100)),
is_error: false,
resolved_tool_name: String::new(),
tool_call_id: String::new(),
duration: Duration::ZERO,
display_hint: None,
}
})
}
}
let memoize = make_middleware(Arc::new(NoopPathExtractor), 10);
let registry = Arc::new(ToolRegistry::new());
let pipeline = ToolPipeline::builder()
.with_middleware(OutputLimitMiddleware::new(20))
.with_middleware(memoize)
.with_middleware(LongOutputMiddleware)
.with_core(registry)
.build()
.expect("pipeline builds");
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 0);
let first = pipeline.dispatch(&mut ctx).await;
let first_len = match &first.output {
ToolContent::Text(text) => text.chars().count(),
ToolContent::Multipart(_) => 0,
};
assert_eq!(first_len, 20, "the miss is capped at exactly 20 chars");
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 1);
let second = pipeline.dispatch(&mut ctx).await;
let second_text = match &second.output {
ToolContent::Text(text) => text.clone(),
ToolContent::Multipart(_) => String::new(),
};
let second_len = second_text.chars().count();
assert!(
second_text.contains("[truncated]"),
"the cap truncates the cached result: {second_text:?}"
);
assert_eq!(
second_len, 20,
"the cached 100-char result flows back through the outer cap — \
re-capped to exactly 20, no bypass: got {second_text:?}"
);
}
#[tokio::test]
async fn cache_hit_preserves_the_original_call_duration() {
struct SlowOutputMiddleware;
impl ToolMiddleware for SlowOutputMiddleware {
fn name(&self) -> &'static str {
"slow_output"
}
fn dispatch<'a>(
&'a self,
_ctx: &'a mut ToolDispatchContext,
_next: &'a ToolPipeline,
) -> Pin<Box<dyn Future<Output = ToolDispatchResult> + Send + 'a>> {
Box::pin(async {
ToolDispatchResult {
output: ToolContent::from_string("content"),
is_error: false,
resolved_tool_name: String::new(),
tool_call_id: String::new(),
duration: Duration::from_millis(250),
display_hint: None,
}
})
}
}
let memoize = make_middleware(Arc::new(NoopPathExtractor), 10);
let registry = Arc::new(ToolRegistry::new());
let pipeline = ToolPipeline::builder()
.with_middleware(memoize)
.with_middleware(SlowOutputMiddleware)
.with_core(registry)
.build()
.expect("pipeline builds");
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 0);
let first = pipeline.dispatch(&mut ctx).await;
assert_eq!(
first.duration,
Duration::from_millis(250),
"the miss reports the tool's own duration"
);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 1);
let second = pipeline.dispatch(&mut ctx).await;
assert_eq!(
second.duration,
Duration::from_millis(250),
"a hit must preserve the cached duration so health statistics \
track the real tool latency, not a zero-cost lookup"
);
}
#[test]
fn append_cached_marker_text() {
let mut content = ToolContent::from_string("hello");
append_cached_marker(&mut content);
match content {
ToolContent::Text(s) => assert!(s.ends_with("\n[cached]"), "marker appended: {s}"),
ToolContent::Multipart(parts) => {
panic!("expected Text, got Multipart with {} parts", parts.len())
}
}
}
#[test]
fn append_cached_marker_multipart() {
let mut content = ToolContent::from_multipart(vec![ToolContentPart::Text {
text: "original".to_string(),
}]);
append_cached_marker(&mut content);
match content {
ToolContent::Multipart(parts) => {
assert_eq!(parts.len(), 2, "should push a new part");
match &parts[1] {
ToolContentPart::Text { text } => {
assert_eq!(text, "[cached]", "marker text: {text}");
}
ToolContentPart::Image { .. } => panic!("expected Text part, got Image"),
}
}
ToolContent::Text(t) => panic!("expected Multipart, got Text: {t}"),
}
}
#[test]
fn append_cached_marker_empty_text() {
let mut content = ToolContent::from_string("");
append_cached_marker(&mut content);
match content {
ToolContent::Text(s) => assert_eq!(s, "\n[cached]", "empty text + marker"),
ToolContent::Multipart(parts) => {
panic!("expected Text, got Multipart with {} parts", parts.len())
}
}
}
#[tokio::test]
async fn failed_write_does_not_invalidate_cache() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let registry = Arc::new(ToolRegistry::new());
let pipeline = ToolPipeline::builder()
.with_middleware(mw)
.with_middleware(RoutingFixedOutputMiddleware {
read_output: ToolContent::from_string("file contents"),
write_output: ToolContent::from_string("permission denied"),
write_is_error: true,
call_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
})
.with_core(registry)
.build()
.expect("pipeline builds");
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0);
let r1 = pipeline.dispatch(&mut ctx).await;
assert!(!r1.is_error);
let mut ctx = ctx_for("Write", serde_json::json!({"path": "foo.rs"}), 1);
let wr = pipeline.dispatch(&mut ctx).await;
assert!(wr.is_error, "write should return an error");
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 2);
let r3 = pipeline.dispatch(&mut ctx).await;
assert!(
r3.output.to_string().contains("[cached]"),
"failed write should not invalidate cached read: {}",
r3.output
);
}
#[tokio::test]
async fn successful_write_does_invalidate_cache() {
let mw = make_middleware(Arc::new(PathFromInput), 10);
let registry = Arc::new(ToolRegistry::new());
let pipeline = ToolPipeline::builder()
.with_middleware(mw)
.with_middleware(RoutingFixedOutputMiddleware {
read_output: ToolContent::from_string("file contents"),
write_output: ToolContent::from_string("wrote"),
write_is_error: false,
call_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
})
.with_core(registry)
.build()
.expect("pipeline builds");
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0);
let r1 = pipeline.dispatch(&mut ctx).await;
assert!(!r1.is_error);
let mut ctx = ctx_for("Write", serde_json::json!({"path": "foo.rs"}), 1);
let wr = pipeline.dispatch(&mut ctx).await;
assert!(!wr.is_error, "write should succeed");
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 2);
let r3 = pipeline.dispatch(&mut ctx).await;
assert!(
!r3.output.to_string().contains("[cached]"),
"successful write should invalidate cached read: {}",
r3.output
);
}
#[test]
fn noop_path_extractor_returns_empty() {
let extractor = NoopPathExtractor;
assert!(
extractor
.paths("Read", &serde_json::json!({"path": "/etc/hosts"}))
.is_empty()
);
assert!(
extractor
.paths("Grep", &serde_json::json!({"pattern": "x"}))
.is_empty()
);
assert!(extractor.paths("Write", &serde_json::json!({})).is_empty());
}
#[tokio::test]
async fn noop_path_extractor_disables_path_invalidation() {
let mw = make_middleware(Arc::new(NoopPathExtractor), 5);
let read_content = ToolContent::from_string("file body");
let (pipeline, _calls) = pipeline(mw, read_content.clone(), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 1);
let r1 = pipeline.dispatch(&mut ctx).await;
assert!(!r1.output.to_string().contains("[cached]"));
let mut ctx = ctx_for("Write", serde_json::json!({"path": "foo.rs"}), 1);
let wr = pipeline.dispatch(&mut ctx).await;
assert!(!wr.is_error, "write should succeed");
let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 2);
let r3 = pipeline.dispatch(&mut ctx).await;
assert!(
r3.output.to_string().contains("[cached]"),
"NoopPathExtractor disables path invalidation; read still cached: {}",
r3.output
);
}
#[tokio::test]
async fn ttl_of_one_expires_on_the_next_turn() {
let mw = make_middleware(Arc::new(NoPaths), 1);
let (pipeline, calls) = pipeline(mw, ToolContent::from_string("v"), false);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 0);
let _ = pipeline.dispatch(&mut ctx).await;
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 0);
let same_turn = pipeline.dispatch(&mut ctx).await;
assert!(
same_turn.output.to_string().contains("[cached]"),
"the insertion turn is still within the ttl: {}",
same_turn.output
);
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
let mut ctx = ctx_for("Read", serde_json::json!({"path": "x"}), 1);
let next_turn = pipeline.dispatch(&mut ctx).await;
assert!(
!next_turn.output.to_string().contains("[cached]"),
"a ttl_turns of 1 expires on the turn after insertion: {}",
next_turn.output
);
assert_eq!(
calls.load(std::sync::atomic::Ordering::SeqCst),
2,
"the expired entry re-runs the inner dispatch"
);
}
}