diff --git a/CHANGELOG.md b/CHANGELOG.md index 85855f5..b3b9fdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -86,6 +86,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. turn sees it. The verifier impl (`cargo check`, `tsc`) is domain-specific and supplied by the consumer; loopctl ships only the trait and the middleware. +- `MemoizingMiddleware` and `PathExtractor` trait + (`middleware::memoize` module): opt-in middleware that caches + successful tool-call results keyed on `(tool_name, hash(canonical + input))` and returns cached output with a `[cached]` marker on hit. + Path-aware invalidation via the caller-supplied `PathExtractor` trait; + TTL-based expiry per `ttl_turns`. Default-off; only successful results + are cached. ### Changed @@ -107,6 +114,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. grammar-aware samplers. Default `ToolConstraint::None` reproduces prior behaviour exactly; when `response_format` is also set, it wins and `tool_constraint` is ignored. +- Removed `parking_lot` dependency entirely. All `parking_lot::Mutex` + usages migrated to `std::sync::Mutex` with the + `.unwrap_or_else(std::sync::PoisonError::into_inner)` recovery pattern. + The crate now uses a single mutex family with no external lock + dependency. +- MSRV bumped from 1.85 to 1.94. The crate uses let-chain syntax + (`if x && let Some(y) = ...`) which stabilized in Rust 1.88. - Both sequential and parallel tool dispatch now check the cancel signal between calls. Previously, a Ctrl-C during a multi-tool batch was only honored at the next turn boundary; now it aborts the remaining calls in the diff --git a/Cargo.toml b/Cargo.toml index b134fc2..04f8eb5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ documentation = "https://docs.rs/loopctl" keywords = ["agent", "framework", "llm", "loop"] categories = ["api-bindings", "development-tools"] readme = "README.md" -rust-version = "1.85" +rust-version = "1.94" [lib] name = "loopctl" @@ -27,7 +27,6 @@ tokio-util = { version = "0.7", features = ["rt"] } uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" -parking_lot = "0.12" reqwest = { version = "0.13", default-features = false, features = ["json", "stream", "rustls"], optional = true } async-stream = { version = "0.3", optional = true } httpdate = { version = "1", optional = true } diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 15ebe1c..8891ed2 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -688,7 +688,7 @@ impl BareLoop { /// let buf = Arc::clone(&buffer); /// agent.set_text_streamer(Arc::new(move |delta| { /// print!("{delta}"); - /// buf.lock().push_str(delta); + /// buf.lock().unwrap_or_else(|e| e.into_inner()).push_str(delta); /// })); /// ``` pub fn set_text_streamer(&mut self, f: Arc) { @@ -1378,20 +1378,20 @@ mod tests { use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; - use parking_lot::Mutex; + use std::sync::Mutex; #[derive(Clone)] struct MockClient { responses: Arc>>>, - model_name: Arc>, + model_name: Arc>, } impl MockClient { fn new(model: &str) -> Self { Self { responses: Arc::new(Mutex::new(Vec::new())), - model_name: Arc::new(parking_lot::Mutex::new(model.to_string())), + model_name: Arc::new(std::sync::Mutex::new(model.to_string())), } } @@ -1401,7 +1401,7 @@ mod tests { message: MessageMetadata { id: "msg_test".into(), role: "assistant".into(), - model: self.model_name.lock().clone(), + model: crate::error::recover_guard(self.model_name.lock()).clone(), }, }), StreamEvent::PartStart(PartStart { @@ -1423,11 +1423,11 @@ mod tests { }), StreamEvent::MessageStop, ]; - self.responses.lock().push(events); + crate::error::recover_guard(self.responses.lock()).push(events); } fn add_events(&self, events: Vec) { - self.responses.lock().push(events); + crate::error::recover_guard(self.responses.lock()).push(events); } fn add_tool_then_text( @@ -1443,7 +1443,7 @@ mod tests { message: MessageMetadata { id: "msg_tool".into(), role: "assistant".into(), - model: self.model_name.lock().clone(), + model: crate::error::recover_guard(self.model_name.lock()).clone(), }, }), StreamEvent::PartStart(PartStart { @@ -1459,7 +1459,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - self.responses.lock().push(tool_events); + crate::error::recover_guard(self.responses.lock()).push(tool_events); // Second response: end_turn with text let text_events = vec![ @@ -1467,7 +1467,7 @@ mod tests { message: MessageMetadata { id: "msg_final".into(), role: "assistant".into(), - model: self.model_name.lock().clone(), + model: crate::error::recover_guard(self.model_name.lock()).clone(), }, }), StreamEvent::PartStart(PartStart { @@ -1489,7 +1489,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - self.responses.lock().push(text_events); + crate::error::recover_guard(self.responses.lock()).push(text_events); } fn add_tool_only_response(&self, tool_id: &str, tool_name: &str, tool_input: Value) { @@ -1498,7 +1498,7 @@ mod tests { message: MessageMetadata { id: format!("msg_{tool_id}"), role: "assistant".into(), - model: self.model_name.lock().clone(), + model: crate::error::recover_guard(self.model_name.lock()).clone(), }, }), StreamEvent::PartStart(PartStart { @@ -1514,7 +1514,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - self.responses.lock().push(tool_events); + crate::error::recover_guard(self.responses.lock()).push(tool_events); } #[expect(dead_code)] @@ -1525,23 +1525,23 @@ mod tests { message: MessageMetadata { id: "msg_err".into(), role: "assistant".into(), - model: self.model_name.lock().clone(), + model: crate::error::recover_guard(self.model_name.lock()).clone(), }, })]; - self.responses.lock().push(events); + crate::error::recover_guard(self.responses.lock()).push(events); } } impl ApiClient for MockClient { fn model(&self) -> String { - self.model_name.lock().clone() + crate::error::recover_guard(self.model_name.lock()).clone() } fn set_model(&self, model: &str) -> bool { if model.trim().is_empty() { return false; } - *self.model_name.lock() = model.to_string(); + *crate::error::recover_guard(self.model_name.lock()) = model.to_string(); true } @@ -1552,7 +1552,7 @@ mod tests { _tools: Option>, ) -> Pin> + Send + 'static>> { - let mut guard = self.responses.lock(); + let mut guard = crate::error::recover_guard(self.responses.lock()); if let Some(events) = guard.pop_front() { let events: Vec> = events.into_iter().map(Ok).collect(); @@ -2040,7 +2040,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - client.responses.lock().push(tool_events); + crate::error::recover_guard(client.responses.lock()).push(tool_events); // Second response: end_turn client.add_text_response("Both tools executed."); @@ -2067,13 +2067,13 @@ mod tests { let received = Arc::new(Mutex::new(Vec::new())); let buf = Arc::clone(&received); agent.set_text_streamer(Arc::new(move |delta: &str| { - buf.lock().push(delta.to_string()); + crate::error::recover_guard(buf.lock()).push(delta.to_string()); })); let result = agent.run("Hi").await.unwrap(); assert!(result.success); - let received = received.lock(); + let received = crate::error::recover_guard(received.lock()); assert!(!received.is_empty(), "streamer should have fired"); assert!( received.join("").contains("Hello world"), @@ -2137,14 +2137,14 @@ mod tests { let received = Arc::new(Mutex::new(String::new())); let buf = Arc::clone(&received); agent.set_text_streamer(Arc::new(move |delta: &str| { - buf.lock().push_str(delta); + crate::error::recover_guard(buf.lock()).push_str(delta); })); agent.run("Use tool").await.unwrap(); // The InputJson delta should NOT have triggered the streamer. // Only the "Done" text response in the second turn should. - let received = received.lock(); + let received = crate::error::recover_guard(received.lock()); assert_eq!(&*received, "Done", "only text deltas should fire streamer"); } @@ -2158,7 +2158,7 @@ mod tests { "delta-recorder" } fn on_text_delta(&self, ctx: &crate::observer::TextDeltaContext) { - self.deltas.lock().push((ctx.turn, ctx.delta.clone())); + crate::error::recover_guard(self.deltas.lock()).push((ctx.turn, ctx.delta.clone())); } } @@ -2212,7 +2212,7 @@ mod tests { let result = agent.run("Hi").await.unwrap(); assert!(result.success); - let captured = captured.lock(); + let captured = crate::error::recover_guard(captured.lock()); assert_eq!(captured.len(), 3, "one on_text_delta per SSE text chunk"); let joined: String = captured.iter().map(|(_, d)| d.as_str()).collect(); assert_eq!(joined, "Hello world"); @@ -2229,10 +2229,10 @@ mod tests { "turn-recorder" } fn on_text_delta(&self, ctx: &crate::observer::TextDeltaContext) { - self.deltas.lock().push((ctx.turn, ctx.delta.clone())); + crate::error::recover_guard(self.deltas.lock()).push((ctx.turn, ctx.delta.clone())); } fn on_response(&self, ctx: &crate::observer::ResponseContext) { - self.response_turns.lock().push(ctx.turn); + crate::error::recover_guard(self.response_turns.lock()).push(ctx.turn); } } @@ -2255,8 +2255,8 @@ mod tests { assert!(result.success); assert_eq!(result.total_turns, 2); - let response_turns = response_turns.lock(); - let deltas = deltas.lock(); + let response_turns = crate::error::recover_guard(response_turns.lock()); + let deltas = crate::error::recover_guard(deltas.lock()); assert_eq!( response_turns.len(), @@ -2359,7 +2359,7 @@ mod tests { "delta-recorder" } fn on_text_delta(&self, ctx: &crate::observer::TextDeltaContext) { - self.deltas.lock().push(ctx.delta.clone()); + crate::error::recover_guard(self.deltas.lock()).push(ctx.delta.clone()); } } @@ -2376,7 +2376,7 @@ mod tests { let result = agent.run("Hi").await.unwrap(); assert!(result.success); - let captured = captured.lock(); + let captured = crate::error::recover_guard(captured.lock()); assert!( !captured.is_empty(), "observer should receive deltas with no streamer set" @@ -2395,7 +2395,7 @@ mod tests { "delta-recorder" } fn on_text_delta(&self, ctx: &crate::observer::TextDeltaContext) { - self.deltas.lock().push(ctx.delta.clone()); + crate::error::recover_guard(self.deltas.lock()).push(ctx.delta.clone()); } } @@ -2407,7 +2407,7 @@ mod tests { let streamer_buf = Arc::new(Mutex::new(Vec::new())); let buf = Arc::clone(&streamer_buf); agent.set_text_streamer(Arc::new(move |delta: &str| { - buf.lock().push(delta.to_string()); + crate::error::recover_guard(buf.lock()).push(delta.to_string()); })); let observer_buf = Arc::new(Mutex::new(Vec::new())); @@ -2419,8 +2419,8 @@ mod tests { let result = agent.run("Hi").await.unwrap(); assert!(result.success); - let streamer_buf = streamer_buf.lock(); - let observer_buf = observer_buf.lock(); + let streamer_buf = crate::error::recover_guard(streamer_buf.lock()); + let observer_buf = crate::error::recover_guard(observer_buf.lock()); assert!(!streamer_buf.is_empty(), "streamer should fire"); assert!(!observer_buf.is_empty(), "observer should fire"); assert_eq!( @@ -2494,7 +2494,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - client.responses.lock().push(tool_events); + crate::error::recover_guard(client.responses.lock()).push(tool_events); client.add_text_response("All done"); let mut registry = ToolRegistry::new(); @@ -2547,13 +2547,13 @@ mod tests { "turn-capture" } fn on_response(&self, ctx: &crate::observer::ResponseContext) { - self.response_turns.lock().push(ctx.turn); + crate::error::recover_guard(self.response_turns.lock()).push(ctx.turn); } fn on_tool_call_received(&self, ctx: &crate::observer::ToolCallReceivedContext) { - self.received_turns.lock().push(ctx.turn); + crate::error::recover_guard(self.received_turns.lock()).push(ctx.turn); } fn on_tool_pre(&self, ctx: &crate::observer::ToolPreContext) { - self.pre_turns.lock().push(ctx.turn); + crate::error::recover_guard(self.pre_turns.lock()).push(ctx.turn); } } @@ -2577,9 +2577,9 @@ mod tests { let result = agent.run("Use echo").await.unwrap(); assert!(result.success); - let received = received.lock(); - let response = response.lock(); - let pre = pre.lock(); + let received = crate::error::recover_guard(received.lock()); + let response = crate::error::recover_guard(response.lock()); + let pre = crate::error::recover_guard(pre.lock()); assert_eq!(received.len(), 1, "one tool call → one received event"); for turn in received.iter() { assert!( @@ -2756,7 +2756,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - client.responses.lock().push(tool_events); + crate::error::recover_guard(client.responses.lock()).push(tool_events); // Second response: end_turn after seeing error result client.add_text_response("Tool wasn't found, but I'll handle it."); @@ -2973,7 +2973,7 @@ mod tests { struct StreamingMockClient { model: String, - rx: parking_lot::Mutex>>>, + rx: std::sync::Mutex>>>, } impl StreamingMockClient { @@ -2987,7 +2987,7 @@ mod tests { ( Self { model: model.to_string(), - rx: parking_lot::Mutex::new(Some(rx)), + rx: std::sync::Mutex::new(Some(rx)), }, tx, ) @@ -3010,7 +3010,9 @@ mod tests { _tools: Option>, ) -> Pin> + Send + 'static>> { - let rx = self.rx.lock().take().expect("stream_messages called twice"); + let rx = crate::error::recover_guard(self.rx.lock()) + .take() + .expect("stream_messages called twice"); Box::pin(ReceiverStream { rx }) } @@ -3113,7 +3115,7 @@ mod tests { dyn std::future::Future + Send + 'a, >, > { - self.turns.lock().push(ctx.turn_number); + crate::error::recover_guard(self.turns.lock()).push(ctx.turn_number); next.dispatch(ctx) } } @@ -3139,7 +3141,7 @@ mod tests { let result = agent.run("test").await; assert!(result.is_ok()); - let turns = capture.lock().clone(); + let turns = crate::error::recover_guard(capture.lock()).clone(); // Tool was called on turn 0 (first turn) and turn 1 (second turn). assert_eq!( turns.len(), @@ -3186,8 +3188,7 @@ mod tests { } fn on_model_switched(&self, ctx: &ModelSwitchedContext) { - self.switches - .lock() + crate::error::recover_guard(self.switches.lock()) .push((ctx.from.clone(), ctx.to.clone())); } } @@ -3203,7 +3204,7 @@ mod tests { loop_.switch_model("m3").apply().unwrap(); // Observer should have received both switches. - let recorded = obs_clone.switches.lock(); + let recorded = crate::error::recover_guard(obs_clone.switches.lock()); assert_eq!(recorded.len(), 2, "should have 2 model-switch events"); assert_eq!(recorded[0], ("default".to_string(), "m2".to_string())); assert_eq!(recorded[1], ("m2".to_string(), "m3".to_string())); @@ -3212,12 +3213,12 @@ mod tests { #[tokio::test] async fn switch_model_unsupported_client() { struct StaticClient { - model_name: Arc>, + model_name: Arc>, } impl ApiClient for StaticClient { fn model(&self) -> String { - self.model_name.lock().clone() + crate::error::recover_guard(self.model_name.lock()).clone() } // Uses default set_model which returns false. @@ -3253,7 +3254,7 @@ mod tests { } let client = std::sync::Arc::new(StaticClient { - model_name: std::sync::Arc::new(parking_lot::Mutex::new("static".to_string())), + model_name: std::sync::Arc::new(std::sync::Mutex::new("static".to_string())), }); let tools = ToolRegistry::new(); let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); @@ -3405,7 +3406,7 @@ mod tests { } fn captured(&self) -> Option { - *self.reason.lock() + *crate::error::recover_guard(self.reason.lock()) } } @@ -3416,7 +3417,7 @@ mod tests { } fn on_session_end(&self, ctx: &HookSessionEndContext) { - *self.reason.lock() = Some(ctx.reason); + *crate::error::recover_guard(self.reason.lock()) = Some(ctx.reason); } } diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 6eb3a37..e34641b 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -1101,7 +1101,7 @@ mod tests { use std::sync::Arc; use std::time::Instant; - use parking_lot::Mutex; + use std::sync::Mutex; use super::*; @@ -1119,13 +1119,13 @@ mod tests { impl ApiClient for MockClient { fn model(&self) -> String { - self.model_name.lock().clone() + crate::error::recover_guard(self.model_name.lock()).clone() } fn set_model(&self, model: &str) -> bool { if model.trim().is_empty() { return false; } - *self.model_name.lock() = model.to_string(); + *crate::error::recover_guard(self.model_name.lock()) = model.to_string(); true } fn stream_messages( diff --git a/src/error.rs b/src/error.rs index 4b338ba..a6349d3 100644 --- a/src/error.rs +++ b/src/error.rs @@ -375,6 +375,43 @@ impl LoopError { } } +// =================================================== +// Mutex poison recovery helper +// =================================================== + +/// Recover a mutex guard from a `LockResult`, ignoring poison. +/// +/// Use this at lock sites protecting data where a panic mid-hold +/// cannot leave the protected value in an inconsistent state (e.g. a +/// `String`, a `Vec`, a `HashMap`, or any type whose operations are +/// individually atomic with respect to the type's invariants). +/// +/// Do **not** use this for locks protecting multi-field structs with +/// cross-field invariants (e.g. a state machine with `state` + +/// `failures` + `last_failure`). A panic between field updates +/// desynchronises the struct; recovering silently means continuing +/// with inconsistent state. In those cases, propagate the +/// `PoisonError` via `?` instead. +/// +/// # When to use this +/// +/// - `Mutex` — a `.clone()` or `=` assignment is atomic. +/// - `Mutex>` / `Mutex>` / `Mutex>` — +/// Rust's collection types maintain their invariants even if an +/// individual `push`/`insert` panics. +/// - `Mutex>` — same; `HashMap`'s internal probing table +/// is never left half-built. +/// +/// # When NOT to use this +/// +/// - `Mutex` where `MyStruct` has multiple fields that must +/// agree (e.g. a state machine with `state` + `failures` + +/// `last_failure`). A panic between field updates desynchronises the +/// struct; recovering silently continues with inconsistent state. +pub(crate) fn recover_guard(result: std::sync::LockResult) -> T { + result.unwrap_or_else(std::sync::PoisonError::into_inner) +} + #[cfg(test)] mod tests { use super::*; @@ -481,4 +518,29 @@ mod tests { let err = LoopError::Cancelled; assert!(!err.is_recoverable()); } + + #[test] + fn recover_guard_returns_value_on_ok() { + let mutex = std::sync::Mutex::new(42); + let guard = recover_guard(mutex.lock()); + assert_eq!(*guard, 42); + } + + #[test] + fn recover_guard_returns_guard_on_poison() { + let mutex = std::sync::Mutex::new(42); + // Poison the mutex by panicking while holding it. + let handle = std::sync::Arc::new(mutex); + let h2 = handle.clone(); + let result = std::thread::spawn(move || { + let _guard = h2.lock().unwrap(); + panic!("intentional"); + }) + .join(); + assert!(result.is_err(), "the panicking thread should have errored"); + // The mutex is now poisoned. recover_guard must still hand back + // the guard so the caller can use the (still-valid) inner value. + let guard = recover_guard(handle.lock()); + assert_eq!(*guard, 42); + } } diff --git a/src/memory/builtin.rs b/src/memory/builtin.rs index 82a0afd..949cc56 100644 --- a/src/memory/builtin.rs +++ b/src/memory/builtin.rs @@ -40,7 +40,7 @@ use crate::error::LoopError; use crate::memory::{ConsolidationStats, LoopMemory, MemoryEntry}; use std::future::Future; -use std::sync::{PoisonError, RwLock}; +use std::sync::RwLock; /// A simple in-memory store for loop memory entries. /// @@ -166,7 +166,7 @@ impl InMemoryStore { /// ``` #[must_use] pub fn with_entries(self, entries: Vec) -> Self { - *self.entries.write().unwrap_or_else(PoisonError::into_inner) = entries; + *crate::error::recover_guard(self.entries.write()) = entries; self } } @@ -194,10 +194,7 @@ impl LoopMemory for InMemoryStore { /// This implementation never returns an error. fn store(&self, entry: MemoryEntry) -> impl Future> + Send { async move { - self.entries - .write() - .unwrap_or_else(PoisonError::into_inner) - .push(entry); + crate::error::recover_guard(self.entries.write()).push(entry); Ok(()) } } @@ -247,7 +244,7 @@ impl LoopMemory for InMemoryStore { let query_lower = query.to_lowercase(); let query_words: Vec<&str> = query_lower.split_whitespace().collect(); - let entries = self.entries.read().unwrap_or_else(PoisonError::into_inner); + let entries = crate::error::recover_guard(self.entries.read()); let snapshot: Vec = entries.iter().cloned().collect(); drop(entries); let mut scored: Vec<(f32, MemoryEntry)> = snapshot @@ -310,7 +307,7 @@ impl LoopMemory for InMemoryStore { /// ``` fn consolidate(&self) -> impl Future> + Send { async move { - let mut entries = self.entries.write().unwrap_or_else(PoisonError::into_inner); + let mut entries = crate::error::recover_guard(self.entries.write()); let entries_before = entries.len(); entries.retain(|e| e.relevance >= 0.05); let pruned = entries_before.saturating_sub(entries.len()); @@ -329,10 +326,7 @@ impl LoopMemory for InMemoryStore { /// Used by the framework to monitor memory usage and by the /// [`is_empty`](LoopMemory::is_empty) provided method. fn len(&self) -> usize { - self.entries - .read() - .unwrap_or_else(PoisonError::into_inner) - .len() + crate::error::recover_guard(self.entries.read()).len() } } diff --git a/src/middleware.rs b/src/middleware.rs index 540adbd..ccf80d7 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -47,6 +47,7 @@ //! let result = pipeline.invoke(ctx).await; //! ``` +pub mod memoize; pub mod output_limit; pub mod permission; pub mod timeout; @@ -64,6 +65,7 @@ use std::sync::Arc; pub use crate::tool::ToolDispatchResult; +pub use memoize::{MemoizingMiddleware, PathExtractor}; pub use output_limit::OutputLimitMiddleware; pub use permission::{AskResolverFn, PermissionCheckFn, PermissionMiddleware}; pub use timeout::{TimeoutConfig, TimeoutMiddleware}; diff --git a/src/middleware/memoize.rs b/src/middleware/memoize.rs new file mode 100644 index 0000000..6916add --- /dev/null +++ b/src/middleware/memoize.rs @@ -0,0 +1,1125 @@ +//! Tool-call memoizing middleware. +//! +//! [`MemoizingMiddleware`] caches successful tool-call results keyed on +//! `(tool_name, hash(canonical input))`. On a cache hit, the middleware +//! returns the prior [`ToolDispatchResult`] with a `[cached]` marker +//! appended and skips the inner dispatch entirely — saving a token-cost +//! and context-pollution hit when a small model re-reads a file or +//! re-runs a grep it already ran. +//! +//! Cache entries are invalidated by: +//! +//! 1. **TTL** — an entry expires after `ttl_turns` turns. +//! 2. **Path-aware write invalidation** — when a write-class tool runs, +//! entries whose paths (per the [`PathExtractor`]) intersect the +//! write's paths are evicted. A write to file A does not evict a +//! cached read of file B. +//! +//! # Registration +//! +//! Register *before* `.core(registry)` so the middleware can short-circuit +//! on a hit before the inner dispatch fires: +//! +//! ```rust,ignore +//! use loopctl::middleware::{MemoizingMiddleware, PathExtractor, ToolPipeline}; +//! use std::sync::Arc; +//! +//! let extractor: Arc = /* … */; +//! let pipeline = ToolPipeline::builder() +//! .with(MemoizingMiddleware::new( +//! vec!["Read".into(), "Grep".into()], +//! vec!["Write".into(), "Edit".into()], +//! extractor, +//! 5, // ttl_turns +//! )) +//! .core(registry) +//! .build()?; +//! ``` + +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}; + +// =================================================== +// PathExtractor trait +// =================================================== + +/// Extracts the filesystem paths a tool call touches, for cache invalidation. +/// +/// Implement this per tool (or per tool family) to tell +/// [`MemoizingMiddleware`] which cached entries a write-class call should +/// invalidate. On a write-class call, the middleware intersects the +/// paths the write touches against each cached entry's paths, and evicts +/// any entry whose path set overlaps the write's. +/// +/// # Generosity +/// +/// Impls should be generous — return every path the call could affect, +/// not just the "primary" one. Over-returning is safe (just causes more +/// invalidation, never staleness); under-returning risks a stale cache +/// hit when a write should have invalidated an entry but didn't. For +/// example, a `Grep` over a directory should return every file path +/// that could be matched, so a write to any of them invalidates the +/// cached result. +/// +/// # Tools that touch no paths +/// +/// Return an empty `Vec` for tools whose output doesn't depend on the +/// filesystem (e.g. a pure-computation tool). Such entries are only +/// TTL-evicted, never path-evicted. +pub trait PathExtractor: Send + Sync { + /// Paths the call touches, as canonical strings. + /// + /// The strings are matched verbatim against other calls' path + /// strings — there is no glob, regex, or prefix logic in the + /// middleware. If a caller wants prefix matching (e.g. a directory + /// write invalidating file reads under it), the impl must emit the + /// matching prefix form itself. + fn paths(&self, tool_name: &str, input: &Value) -> Vec; +} + +// =================================================== +// Cache types (private) +// =================================================== + +/// Cache key: the tool name plus the hash of its canonical-serialized input. +/// +/// Both components are required: the same input hash could theoretically +/// collide across tools, so the tool name disambiguates. +#[derive(Debug, Eq, PartialEq, Hash, Clone)] +struct CacheKey { + /// The tool name the call was issued under. + /// + /// Compared verbatim against the configured `tools` list to decide + /// whether the call is memoized at all. Stored as part of the key so + /// that the same input hash under different tool names does not + /// produce a false cache hit. + tool_name: String, + + /// `DefaultHasher` digest of `serde_json::to_string(&input)`. + /// + /// Because `serde_json::Map` is `BTreeMap`-backed (keys sorted + /// alphabetically), two inputs that differ only in JSON key order + /// (`{"a":1,"b":2}` vs `{"b":2,"a":1}`) produce the same canonical + /// string and therefore the same hash — so semantically-equal calls + /// hit the same cache entry. + input_hash: u64, +} + +/// A single cached entry: the result, when it was inserted, and what +/// paths it touched. +struct CacheEntry { + /// The successful `ToolDispatchResult` returned on the original call. + /// + /// Cloned (with `[cached]` appended to its output and `duration` + /// zeroed) on a hit. Only successful results (`is_error == false`) + /// are cached; errors always re-run. + result: ToolDispatchResult, + + /// The `turn_number` at which the entry was inserted. + /// + /// Compared against the current turn for TTL expiry via + /// `current_turn.saturating_sub(turn_inserted) >= ttl_turns`. + /// Stale entries are evicted lazily on lookup. + turn_inserted: usize, + + /// Paths the original call touched, per the `PathExtractor`. + /// + /// Used for path-aware invalidation on write-class calls. When a + /// write tool runs, the middleware intersects the write's paths + /// against each entry's `paths` set and evicts any overlap. An + /// empty vec means the call touches no paths (e.g. a pure + /// computation tool) — such entries are only TTL-evicted, never + /// path-evicted. + paths: Vec, +} + +// =================================================== +// MemoizingMiddleware +// =================================================== + +/// Deduplicates repeat tool calls by `(tool_name, hash(canonical input))`. +/// +/// On a cache hit (memoized tool + key present + not TTL-expired), +/// returns the prior [`ToolDispatchResult`] with a `[cached]` marker +/// appended and skips the inner dispatch entirely. On a cache miss, +/// runs the tool and stores the result. On a write-class tool call, +/// invalidates cache entries whose paths (per the [`PathExtractor`]) +/// intersect the write's paths. +/// +/// Only successful results are cached (`is_error == false`); errors +/// always re-run. Entries expire after `ttl_turns` turns. +/// +/// # Cache size +/// +/// The cache grows with distinct `(tool, input)` pairs up to TTL +/// expiry. There is no LRU eviction in this middleware; if a long +/// session produces many distinct inputs, pair it with a tighter +/// `ttl_turns`. +/// +/// # Registration +/// +/// See the [module docs](self) for the registration pattern. +pub struct MemoizingMiddleware { + /// Interior-mutable cache, guarded by a `std::sync::Mutex`. + /// + /// The lock is held only briefly during lookup / insert / evict — + /// never across the inner `next.dispatch()` call (which could take + /// seconds and would serialize all tool calls). + cache: Mutex>, + + /// Exact-match tool names whose results to cache. + /// + /// Tools not in this list bypass the cache entirely. + tools: Vec, + + /// Exact-match tool names that trigger path-aware invalidation. + /// + /// After a write tool runs, the cache evicts entries whose paths + /// intersect the write's paths. Write results themselves are never + /// cached. + write_tools: Vec, + + /// The shared path extractor consulted on every cache insert and + /// every write-class invalidation. + /// + /// Required in the constructor; there is no blunt-clear fallback. + /// Impls are shared via `Arc` so the same path + /// knowledge backs every cache lookup. + path_extractor: Arc, + + /// Max turns a cache entry is valid. + /// + /// An entry expires when + /// `current_turn.saturating_sub(turn_inserted) >= ttl_turns`. So a + /// `ttl_turns` of 1 means the entry is valid on the turn after + /// insertion but expires on the second turn after. + ttl_turns: u32, +} + +impl MemoizingMiddleware { + /// Construct a memoizing middleware. + /// + /// # Arguments + /// + /// - `tools` — exact-match names of tools whose results to cache. + /// - `write_tools` — exact-match names of tools that trigger + /// path-aware invalidation. + /// - `path_extractor` — required; supplies the per-tool path + /// knowledge the framework can't infer. + /// - `ttl_turns` — max turns a cache entry is valid (entry expires + /// when `current_turn - turn_inserted >= ttl_turns`). + #[must_use] + pub fn new( + tools: Vec, + write_tools: Vec, + path_extractor: Arc, + ttl_turns: u32, + ) -> Self { + Self { + cache: Mutex::new(HashMap::new()), + 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()).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 + 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 ttl_turns = self.ttl_turns; + let current_turn = ctx.turn_number; + + 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.duration = std::time::Duration::ZERO; + 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 paths = self.path_extractor.paths(&ctx.tool_name, &ctx.input); + if let Some(key) = key { + insert(&self.cache, key, result.clone(), current_turn, paths); + } + } + + result + }) + } +} + +/// Build a `CacheKey` for `(tool_name, input)`. +/// +/// The input is canonicalized via `serde_json::to_string` (which, with +/// the default `BTreeMap`-backed `serde_json::Map`, sorts object keys +/// alphabetically — so semantically-equal inputs hash equally regardless +/// of key order) and then hashed with `DefaultHasher`. +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(), + } +} + +/// Look up a key in the cache, returning a cloned entry only if it +/// exists and is not TTL-expired. Stale entries are evicted as a side +/// effect of the lookup (lazy expiry). +fn lookup_fresh( + cache: &Mutex>, + key: &CacheKey, + current_turn: usize, + ttl_turns: u32, +) -> Option { + 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.get(key) + && expired(entry) + { + guard.remove(key); + return None; + } + guard.get(key).map(|e| e.result.clone()) +} + +/// Insert a successful tool result into the cache. +/// +/// Stores the result under the computed key alongside the turn it was +/// inserted and the paths the call touched. If an entry with the same +/// key already exists (e.g. the same call was made earlier and the +/// entry wasn't TTL-evicted or path-invalidated in between), it is +/// replaced — the newer result wins. +fn insert( + cache: &Mutex>, + key: CacheKey, + result: ToolDispatchResult, + turn_inserted: usize, + paths: Vec, +) { + let mut guard = crate::error::recover_guard(cache.lock()); + guard.insert( + key, + CacheEntry { + result, + turn_inserted, + paths, + }, + ); +} + +/// Evict entries whose paths intersect any of `write_paths`. +/// +/// Intersection is exact string equality — no prefix or glob logic. +/// An empty `write_paths` set evicts nothing (a write that touches no +/// known paths can't path-invalidate anything). +fn invalidate_paths(cache: &Mutex>, 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.retain(|_, entry| !entry.paths.iter().any(|p| write_set.contains(p.as_str()))); +} + +/// Append the `[cached]` marker to a `ToolContent`. +/// +/// For `Text`, the marker is appended to the existing string. For +/// `Multipart`, a new `Text { text: "[cached]" }` part is pushed. +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::{ToolDispatchContext, ToolPipeline}; + use crate::tool::{PermissionCheck, ToolContext, ToolRegistry}; + use std::future::Future; + use std::sync::Arc; + use std::time::Duration; + + /// A `PathExtractor` that reads `input.path` for memoized tools + /// and `input.path` for write tools. Returns `[]` otherwise. + struct PathFromInput; + + impl PathExtractor for PathFromInput { + fn paths(&self, _tool_name: &str, input: &Value) -> Vec { + input + .get("path") + .and_then(Value::as_str) + .map(std::string::ToString::to_string) + .into_iter() + .collect() + } + } + + /// A `PathExtractor` that returns `[]` for everything — no path + /// invalidation possible. Useful for TTL-only tests. + struct NoPaths; + impl PathExtractor for NoPaths { + fn paths(&self, _: &str, _: &Value) -> Vec { + Vec::new() + } + } + + /// A middleware that short-circuits with a fixed result, counting + /// how many times it was dispatched. + struct FixedOutputMiddleware { + output: ToolContent, + is_error: bool, + call_count: Arc, + } + + impl ToolMiddleware for FixedOutputMiddleware { + fn name(&self) -> &'static str { + "fixed_output" + } + fn dispatch<'a>( + &'a self, + _ctx: &'a mut ToolDispatchContext, + _next: &'a ToolPipeline, + ) -> Pin + 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: String::new(), + duration: Duration::ZERO, + } + }) + } + } + + /// Like `FixedOutputMiddleware`, but returns a different result + /// depending on whether the tool is "Read" or "Write". Needed for + /// tests that exercise the write-invalidation path in a single + /// shared pipeline (shared cache). + struct RoutingFixedOutputMiddleware { + read_output: ToolContent, + write_output: ToolContent, + write_is_error: bool, + call_count: Arc, + } + + 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 + 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: String::new(), + duration: Duration::ZERO, + } + }) + } + } + + fn make_middleware(extractor: Arc, 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) { + let call_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let registry = Arc::new(ToolRegistry::new()); + let pipeline = ToolPipeline::builder() + .with(mw) + .with(FixedOutputMiddleware { + output, + is_error, + call_count: call_count.clone(), + }) + .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; + // Inner dispatch not called on a cache hit. + 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 + ); + // First call (the cache miss) should not carry [cached]. + assert!( + !first.output.to_string().contains("[cached]"), + "first call output should not carry [cached]" + ); + } + + #[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); + + // Read(foo.rs) — cached. + 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); + + // Write(bar.rs) — should NOT invalidate the Read(foo.rs) entry. + 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); + + // Read(foo.rs) again — should still be cached (Write was to bar.rs). + 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 + ); + + // Write(foo.rs) — should invalidate the Read(foo.rs) entry. + let mut ctx = ctx_for("Write", serde_json::json!({"path": "foo.rs"}), 3); + let _ = pipeline.dispatch(&mut ctx).await; + + // Read(foo.rs) again — should re-run (cache invalidated). + 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() { + // ttl_turns = 2: entry inserted at turn 0 is fresh at turn 1, + // expired at turn 2 (>= ttl_turns). + 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); + + // Turn 1 — fresh, should hit. + 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); + + // Turn 2 — expired (1 - 0 = 1, not >= 2; but 2 - 0 = 2 >= 2). Re-runs. + 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); + + // Grep is not in the memoize list (only Read is). + 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); + + // Write is a write tool, not a memoized tool. + 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, // is_error + ); + + 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() { + // Text variant. + 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()) + } + } + + // Multipart variant. + 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); + + // Same logical input, different key order. + 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"); + } + + // ---- make_key unit tests ---- + + #[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"); + } + + // ---- lookup_fresh unit tests ---- + + #[test] + fn lookup_fresh_returns_none_on_miss() { + let cache = Mutex::new(HashMap::new()); + 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(HashMap::new()); + 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(), + }, + 0, + vec![], + ); + 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(HashMap::new()); + 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(), + }, + 0, + vec![], + ); + // ttl_turns=2: entry inserted at turn 0 is fresh at turn 1 + // (1-0=1 < 2), expired at turn 2 (2-0=2 >= 2). + 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" + ); + // Lazy eviction: the entry should have been removed by the lookup. + let guard = crate::error::recover_guard(cache.lock()); + assert!( + guard.is_empty(), + "expired entry should be evicted from cache" + ); + } + + // ---- insert unit tests ---- + + #[test] + fn insert_stores_entry_and_replace_overwrites() { + let cache = Mutex::new(HashMap::new()); + 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(), + }, + 0, + vec![], + ); + 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(), + }, + 1, + vec![], + ); + 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()) + } + } + } + + // ---- invalidate_paths unit tests ---- + + #[test] + fn invalidate_paths_removes_overlapping_entries() { + let cache = Mutex::new(HashMap::new()); + 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(), + }; + insert( + &cache, + key_a.clone(), + result.clone(), + 0, + vec!["a.rs".to_string()], + ); + insert(&cache, key_b.clone(), result, 0, vec!["b.rs".to_string()]); + + // Write to a.rs should evict key_a but not key_b. + invalidate_paths(&cache, &["a.rs".to_string()]); + + let guard = crate::error::recover_guard(cache.lock()); + assert!( + !guard.contains_key(&key_a), + "overlapping entry should be evicted" + ); + assert!( + guard.contains_key(&key_b), + "non-overlapping entry should survive" + ); + } + + #[test] + fn invalidate_paths_noop_on_empty_write_paths() { + let cache = Mutex::new(HashMap::new()); + 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(), + }, + 0, + vec!["x".to_string()], + ); + invalidate_paths(&cache, &[]); + let guard = crate::error::recover_guard(cache.lock()); + assert_eq!(guard.len(), 1, "empty write_paths should evict nothing"); + } + + #[test] + fn invalidate_paths_evicts_all_matching() { + let cache = Mutex::new(HashMap::new()); + let result = ToolDispatchResult { + tool_call_id: String::new(), + output: ToolContent::from_string("content"), + is_error: false, + duration: Duration::ZERO, + resolved_tool_name: String::new(), + }; + // Two entries both touching "shared.rs". + 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()]); + insert(&cache, k2, result, 0, vec!["shared.rs".to_string()]); + + invalidate_paths(&cache, &["shared.rs".to_string()]); + let guard = crate::error::recover_guard(cache.lock()); + assert!( + guard.is_empty(), + "both entries touch shared.rs — both evicted" + ); + } + + // ---- append_cached_marker unit tests ---- + + #[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() { + // A write that errors (permission denied, disk full) didn't + // modify the filesystem, so cached reads are still accurate and + // must NOT be invalidated. This uses a shared middleware so the + // same cache is consulted across both Read and Write calls. + let mw = make_middleware(Arc::new(PathFromInput), 10); + let registry = Arc::new(ToolRegistry::new()); + let pipeline = ToolPipeline::builder() + .with(mw) + .with(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)), + }) + .core(registry) + .build() + .expect("pipeline builds"); + + // Step 1: Read(foo.rs) — cache it. + let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0); + let r1 = pipeline.dispatch(&mut ctx).await; + assert!(!r1.is_error); + + // Step 2: Write(foo.rs) fails — should NOT invalidate. + 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"); + + // Step 3: Read(foo.rs) again — should still be cached. + 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() { + // The counterpart: a successful write to the same path MUST + // invalidate the cached read. + let mw = make_middleware(Arc::new(PathFromInput), 10); + let registry = Arc::new(ToolRegistry::new()); + let pipeline = ToolPipeline::builder() + .with(mw) + .with(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)), + }) + .core(registry) + .build() + .expect("pipeline builds"); + + // Step 1: Read(foo.rs) — cache it. + let mut ctx = ctx_for("Read", serde_json::json!({"path": "foo.rs"}), 0); + let r1 = pipeline.dispatch(&mut ctx).await; + assert!(!r1.is_error); + + // Step 2: Write(foo.rs) succeeds — should invalidate. + 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"); + + // Step 3: Read(foo.rs) again — should re-run (invalidated). + 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 + ); + } +} diff --git a/src/observer.rs b/src/observer.rs index 50b1a1e..5f36584 100644 --- a/src/observer.rs +++ b/src/observer.rs @@ -520,19 +520,20 @@ mod tests { #[test] fn host_dispatches_model_switched() { struct SwitchRecorder { - events: parking_lot::Mutex>, + events: std::sync::Mutex>, } impl LoopObserver for SwitchRecorder { fn name(&self) -> &'static str { "switch-recorder" } fn on_model_switched(&self, ctx: &ModelSwitchedContext) { - self.events.lock().push((ctx.from.clone(), ctx.to.clone())); + crate::error::recover_guard(self.events.lock()) + .push((ctx.from.clone(), ctx.to.clone())); } } let obs = Arc::new(SwitchRecorder { - events: parking_lot::Mutex::new(Vec::new()), + events: std::sync::Mutex::new(Vec::new()), }); let mut host = ObserverHost::new(); host.register(Arc::clone(&obs) as Arc); @@ -546,7 +547,7 @@ mod tests { to: "c".into(), }); - let events = obs.events.lock(); + let events = crate::error::recover_guard(obs.events.lock()); assert_eq!(events.len(), 2); assert_eq!(events[0], ("a".into(), "b".into())); assert_eq!(events[1], ("b".into(), "c".into())); @@ -591,22 +592,22 @@ mod tests { #[test] fn host_dispatches_on_text_delta_to_all_observers() { struct DeltaRecorder { - deltas: parking_lot::Mutex>, + deltas: std::sync::Mutex>, } impl LoopObserver for DeltaRecorder { fn name(&self) -> &'static str { "delta-recorder" } fn on_text_delta(&self, ctx: &TextDeltaContext) { - self.deltas.lock().push(ctx.delta.clone()); + crate::error::recover_guard(self.deltas.lock()).push(ctx.delta.clone()); } } let obs1 = Arc::new(DeltaRecorder { - deltas: parking_lot::Mutex::new(Vec::new()), + deltas: std::sync::Mutex::new(Vec::new()), }); let obs2 = Arc::new(DeltaRecorder { - deltas: parking_lot::Mutex::new(Vec::new()), + deltas: std::sync::Mutex::new(Vec::new()), }); let mut host = ObserverHost::new(); host.register(Arc::clone(&obs1) as Arc); @@ -617,8 +618,14 @@ mod tests { delta: "x".into(), }); - assert_eq!(obs1.deltas.lock().clone(), vec!["x".to_string()]); - assert_eq!(obs2.deltas.lock().clone(), vec!["x".to_string()]); + assert_eq!( + crate::error::recover_guard(obs1.deltas.lock()).clone(), + vec!["x".to_string()] + ); + assert_eq!( + crate::error::recover_guard(obs2.deltas.lock()).clone(), + vec!["x".to_string()] + ); } #[test] @@ -653,22 +660,22 @@ mod tests { #[test] fn host_dispatches_on_tool_call_received_to_all_observers() { struct ReceivedRecorder { - calls: parking_lot::Mutex>, + calls: std::sync::Mutex>, } impl LoopObserver for ReceivedRecorder { fn name(&self) -> &'static str { "received-recorder" } fn on_tool_call_received(&self, ctx: &ToolCallReceivedContext) { - self.calls.lock().push(ctx.tool.clone()); + crate::error::recover_guard(self.calls.lock()).push(ctx.tool.clone()); } } let obs1 = Arc::new(ReceivedRecorder { - calls: parking_lot::Mutex::new(Vec::new()), + calls: std::sync::Mutex::new(Vec::new()), }); let obs2 = Arc::new(ReceivedRecorder { - calls: parking_lot::Mutex::new(Vec::new()), + calls: std::sync::Mutex::new(Vec::new()), }); let mut host = ObserverHost::new(); host.register(Arc::clone(&obs1) as Arc); @@ -681,8 +688,14 @@ mod tests { input: serde_json::Value::Null, }); - assert_eq!(obs1.calls.lock().clone(), vec!["edit".to_string()]); - assert_eq!(obs2.calls.lock().clone(), vec!["edit".to_string()]); + assert_eq!( + crate::error::recover_guard(obs1.calls.lock()).clone(), + vec!["edit".to_string()] + ); + assert_eq!( + crate::error::recover_guard(obs2.calls.lock()).clone(), + vec!["edit".to_string()] + ); } #[test] diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index bcaa227..0674adc 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -91,7 +91,7 @@ pub struct AnthropicClient { /// Changed via [`ApiClient::set_model`] when the /// [`FallbackManager`](crate::fallback::FallbackManager) trips to a /// fallback model. - model: parking_lot::Mutex, + model: std::sync::Mutex, /// The maximum output tokens per response. /// @@ -208,7 +208,7 @@ impl AnthropicClient { impl ApiClient for AnthropicClient { fn model(&self) -> String { - self.model.lock().clone() + crate::error::recover_guard(self.model.lock()).clone() } fn base_url(&self) -> String { @@ -219,7 +219,7 @@ impl ApiClient for AnthropicClient { if model.trim().is_empty() { return false; } - *self.model.lock() = model.to_string(); + *crate::error::recover_guard(self.model.lock()) = model.to_string(); true } @@ -258,7 +258,7 @@ impl ApiClient for AnthropicClient { tools: Option>, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { - let model = self.model.lock().clone(); + let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = build_request_body( &RequestBodySpec { @@ -301,7 +301,7 @@ impl ApiClient for AnthropicClient { tools: Option>, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { - let model = self.model.lock().clone(); + let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = build_request_body( &RequestBodySpec { @@ -497,7 +497,7 @@ impl AnthropicClientBuilder { http, api_key, base_url: self.base_url, - model: parking_lot::Mutex::new(self.model), + model: std::sync::Mutex::new(self.model), max_tokens: self.max_tokens, }) } diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index ce1f61b..ff354db 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -86,7 +86,7 @@ pub struct GeminiClient { /// Changed via [`ApiClient::set_model`] when the /// [`FallbackManager`](crate::fallback::FallbackManager) trips to a /// fallback model. - model: parking_lot::Mutex, + model: std::sync::Mutex, } impl GeminiClient { @@ -148,7 +148,7 @@ impl GeminiClient { /// Gemini puts the model in the URL path. The API key is sent via /// the `x-goog-api-key` header, not as a query parameter. fn stream_url(&self) -> String { - let model = self.model.lock().clone(); + let model = crate::error::recover_guard(self.model.lock()).clone(); format!( "{}/models/{}:streamGenerateContent?alt=sse", self.base_url, model @@ -163,7 +163,7 @@ impl GeminiClient { /// Used by [`ApiClient::create_message`] and its `*_with_options` /// variant. fn generate_url(&self) -> String { - let model = self.model.lock().clone(); + let model = crate::error::recover_guard(self.model.lock()).clone(); format!("{}/models/{}:generateContent", self.base_url, model) } @@ -206,7 +206,7 @@ impl GeminiClient { impl ApiClient for GeminiClient { fn model(&self) -> String { - self.model.lock().clone() + crate::error::recover_guard(self.model.lock()).clone() } fn base_url(&self) -> String { @@ -217,7 +217,7 @@ impl ApiClient for GeminiClient { if model.trim().is_empty() { return false; } - *self.model.lock() = model.to_string(); + *crate::error::recover_guard(self.model.lock()) = model.to_string(); true } @@ -492,7 +492,7 @@ impl GeminiClientBuilder { http, api_key, base_url: self.base_url, - model: parking_lot::Mutex::new(self.model), + model: std::sync::Mutex::new(self.model), }) } } diff --git a/src/provider/openai.rs b/src/provider/openai.rs index ae0c157..a8f6f8b 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -92,7 +92,7 @@ pub struct OpenAiClient { /// Changed via [`ApiClient::set_model`] when the /// [`FallbackManager`](crate::fallback::FallbackManager) trips to a /// fallback model. - model: parking_lot::Mutex, + model: std::sync::Mutex, } impl OpenAiClient { @@ -205,7 +205,7 @@ impl OpenAiClient { impl ApiClient for OpenAiClient { fn model(&self) -> String { - self.model.lock().clone() + crate::error::recover_guard(self.model.lock()).clone() } fn base_url(&self) -> String { @@ -216,7 +216,7 @@ impl ApiClient for OpenAiClient { if model.trim().is_empty() { return false; } - *self.model.lock() = model.to_string(); + *crate::error::recover_guard(self.model.lock()) = model.to_string(); true } @@ -226,7 +226,7 @@ impl ApiClient for OpenAiClient { system: Option, tools: Option>, ) -> Pin> + Send + 'static>> { - let model = self.model.lock().clone(); + let model = crate::error::recover_guard(self.model.lock()).clone(); let body = RequestBody::build( &model, &messages, @@ -266,7 +266,7 @@ impl ApiClient for OpenAiClient { system: Option, tools: Option>, ) -> Pin> + Send + '_>> { - let model = self.model.lock().clone(); + let model = crate::error::recover_guard(self.model.lock()).clone(); let body = RequestBody::build( &model, &messages, @@ -303,7 +303,7 @@ impl ApiClient for OpenAiClient { tools: Option>, options: crate::structured::RequestOptions, ) -> Pin> + Send + 'static>> { - let model = self.model.lock().clone(); + let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = RequestBody::build( &model, @@ -345,7 +345,7 @@ impl ApiClient for OpenAiClient { tools: Option>, options: crate::structured::RequestOptions, ) -> Pin> + Send + '_>> { - let model = self.model.lock().clone(); + let model = crate::error::recover_guard(self.model.lock()).clone(); let rf = options.response_format.as_ref(); let body = RequestBody::build( &model, @@ -519,7 +519,7 @@ impl OpenAiClientBuilder { http, api_key, base_url: self.base_url, - model: parking_lot::Mutex::new(self.model), + model: std::sync::Mutex::new(self.model), }) } } diff --git a/src/testing.rs b/src/testing.rs index dcc4c9b..e20b898 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -104,7 +104,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use parking_lot::Mutex; +use std::sync::Mutex; use uuid::Uuid; // ================================================== @@ -176,7 +176,7 @@ pub struct MockApiClient { /// The value is returned verbatim by the [`ApiClient::model`] /// implementation. It appears in log messages and /// [`MessageMetadata`] fields. - model_name: Arc>, + model_name: Arc>, /// The queue of canned responses. /// @@ -364,7 +364,7 @@ impl MockApiClient { stop_reason: "end_turn".to_string(), }; Self { - model_name: Arc::new(parking_lot::Mutex::new(model.to_string())), + model_name: Arc::new(std::sync::Mutex::new(model.to_string())), responses: Arc::new(Mutex::new(vec![default_response])), error: None, } @@ -390,7 +390,7 @@ impl MockApiClient { /// ``` #[must_use] pub fn with_text_response(self, text: &str) -> Self { - if let Some(r) = self.responses.lock().first_mut() { + if let Some(r) = crate::error::recover_guard(self.responses.lock()).first_mut() { r.text = text.to_string(); } self @@ -418,7 +418,7 @@ impl MockApiClient { /// ``` #[must_use] pub fn with_tool_call(self, id: &str, name: &str, input: Value) -> Self { - let mut responses = self.responses.lock(); + let mut responses = crate::error::recover_guard(self.responses.lock()); if let Some(r) = responses.first_mut() { r.tool_call = Some(MockToolCall { id: id.to_string(), @@ -449,7 +449,7 @@ impl MockApiClient { /// ``` #[must_use] pub fn with_stop_reason(self, reason: &str) -> Self { - if let Some(r) = self.responses.lock().first_mut() { + if let Some(r) = crate::error::recover_guard(self.responses.lock()).first_mut() { r.stop_reason = reason.to_string(); } self @@ -487,7 +487,7 @@ impl MockApiClient { #[must_use] pub fn with_responses(self, responses: Vec) -> Self { if !responses.is_empty() { - *self.responses.lock() = responses; + *crate::error::recover_guard(self.responses.lock()) = responses; } self } @@ -538,7 +538,7 @@ impl MockApiClient { /// [R3] → pop → R3, queue stays [R3] (cloned) /// ``` fn pop_response(&self) -> MockResponse { - let mut guard = self.responses.lock(); + let mut guard = crate::error::recover_guard(self.responses.lock()); if guard.len() > 1 { guard.remove(0) } else { @@ -595,14 +595,14 @@ impl ApiClient for MockApiClient { /// assert_eq!(client.model(), "my-test-model"); /// ``` fn model(&self) -> String { - self.model_name.lock().clone() + crate::error::recover_guard(self.model_name.lock()).clone() } fn set_model(&self, model: &str) -> bool { if model.trim().is_empty() { return false; } - *self.model_name.lock() = model.to_string(); + *crate::error::recover_guard(self.model_name.lock()) = model.to_string(); true } @@ -658,7 +658,7 @@ impl ApiClient for MockApiClient { } let response = self.pop_response(); - let model = self.model_name.lock().clone(); + let model = crate::error::recover_guard(self.model_name.lock()).clone(); let mut events: Vec> = vec![Ok(StreamEvent::MessageStart(MessageStart { message: MessageMetadata { diff --git a/src/tool.rs b/src/tool.rs index 5b9d7b5..e3503a1 100644 --- a/src/tool.rs +++ b/src/tool.rs @@ -514,15 +514,51 @@ impl From<&str> for ToolOutput { /// ``` #[derive(Debug, Clone)] pub struct ToolDispatchResult { - /// Set by the engine via [`with_call_id`](Self::with_call_id). + /// The ID the model assigned to this tool call. + /// + /// Copied from the inbound `ToolCall` by the engine so the caller + /// can correlate the result back to the request. Set via + /// [`with_call_id`](Self::with_call_id) on the builder; empty when + /// constructed via [`From`](Self#impl-From). pub tool_call_id: String, - /// Preserves multipart and image content on success; wraps error in text on failure. + + /// The content returned by the tool. + /// + /// A [`ToolContent::Text`](crate::message::ToolContent::Text) on the + /// common path (plain-text result or an error message wrapped in + /// text). A + /// [`ToolContent::Multipart`](crate::message::ToolContent::Multipart) + /// when the tool produces mixed content (text + images) or multiple + /// text segments. Middlewares may rewrite this post-execution (e.g. + /// `OutputLimitMiddleware` truncates, `VerifyMiddleware` appends a + /// `[verify]` block). pub output: crate::message::ToolContent, + /// Whether the tool dispatch resulted in an error. + /// + /// `false` on success; `true` when the tool returned an error or the + /// dispatch itself failed (timeout, panic caught via `catch_unwind`, + /// permission denied). Middlewares and the recovery loop consult + /// this flag — for example, `VerifyMiddleware` skips verification + /// and `MemoizingMiddleware` skips caching when `is_error` is `true`. pub is_error: bool, - /// Wall-clock execution time. + + /// Wall-clock execution time of the tool call. + /// + /// Measured from dispatch start to completion by the engine. Zero on + /// cached results returned by `MemoizingMiddleware` (the cached call + /// didn't execute this turn). Useful for metrics, slow-tool + /// detection, and timeout diagnostics. pub duration: Duration, - /// May differ from the requested tool name if a routing middleware redirected the call. + + /// The tool name the call actually ran under. + /// + /// Usually identical to the requested name. May differ when a + /// routing middleware redirected the call (e.g. aliasing an + /// unknown tool name to the closest match). Middlewares that need + /// to reason about *which tool executed* (rather than what the + /// model requested) should read this field instead of the + /// pre-dispatch `ToolDispatchContext::tool_name`. pub resolved_tool_name: String, } diff --git a/src/tool/health.rs b/src/tool/health.rs index 9420759..2ab3f18 100644 --- a/src/tool/health.rs +++ b/src/tool/health.rs @@ -462,10 +462,7 @@ impl ToolCircuitBreaker { CircuitState::Closed => true, CircuitState::HalfOpen => false, CircuitState::Open => { - let recovered = self - .last_failure_time - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) + let recovered = crate::error::recover_guard(self.last_failure_time.lock()) .map(|t| t.elapsed() >= self.recovery_duration) .unwrap_or(false); if recovered { @@ -657,18 +654,12 @@ impl ToolHealthRegistry { /// so the caller can read counters without holding any lock. #[must_use] pub fn get_stats(&self, tool_name: &str) -> Arc { - let guard = self - .stats - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let guard = crate::error::recover_guard(self.stats.lock()); if let Some(stats) = guard.get(tool_name) { return Arc::clone(stats); } drop(guard); - let mut guard = self - .stats - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut guard = crate::error::recover_guard(self.stats.lock()); Arc::clone( guard .entry(tool_name.to_string()) @@ -682,18 +673,12 @@ impl ToolHealthRegistry { /// registry's [`CircuitBreakerConfig`]. #[must_use] pub fn get_circuit_breaker(&self, tool_name: &str) -> Arc { - let guard = self - .breakers - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let guard = crate::error::recover_guard(self.breakers.lock()); if let Some(cb) = guard.get(tool_name) { return Arc::clone(cb); } drop(guard); - let mut guard = self - .breakers - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut guard = crate::error::recover_guard(self.breakers.lock()); Arc::clone( guard .entry(tool_name.to_string()) @@ -767,10 +752,7 @@ impl ToolHealthRegistry { #[must_use] pub fn health_summary(&self) -> HashMap { let entries: Vec<(String, Arc)> = { - let guard = self - .stats - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let guard = crate::error::recover_guard(self.stats.lock()); guard .iter() .map(|(n, s)| (n.clone(), Arc::clone(s))) @@ -789,10 +771,7 @@ impl ToolHealthRegistry { /// Number of distinct tools currently tracked. #[must_use] pub fn tool_count(&self) -> usize { - self.stats - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .len() + crate::error::recover_guard(self.stats.lock()).len() } } diff --git a/src/tool/shield.rs b/src/tool/shield.rs index 077a84d..53d64b2 100644 --- a/src/tool/shield.rs +++ b/src/tool/shield.rs @@ -682,10 +682,7 @@ impl UnixShield { /// aggregate past the warn threshold — but combined with a /// single-turn hit it adds up. pub fn assess_multi_turn(&self, ctx: &ShieldContext) -> f32 { - let history = self - .turn_history - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let history = crate::error::recover_guard(self.turn_history.lock()); let same_tool_count = history .iter() .filter(|(name, _)| name == &ctx.tool_name) @@ -716,10 +713,7 @@ impl UnixShield { /// only one that can detect adversarial *sequences* (download → /// execute, write → chmod +x). pub fn assess_combination(&self, ctx: &ShieldContext) -> f32 { - let history = self - .turn_history - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let history = crate::error::recover_guard(self.turn_history.lock()); let mut max_risk = 0.0_f32; // Build a chronological candidate sequence: history entries in @@ -922,10 +916,7 @@ impl ToolSafetyShield for UnixShield { } fn record_invocation(&self, tool_name: &str, input: &Value, _success: bool) { - let mut history = self - .turn_history - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut history = crate::error::recover_guard(self.turn_history.lock()); history.push((tool_name.to_string(), input.to_string())); // Trim to last 20 entries. if history.len() > 20 {