From 2ad57aa86401bea0ea22023d1629197980e13c08 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 1 Jun 2026 21:44:55 +1200 Subject: [PATCH 01/30] refactor: observer events --- README.md | 1 - src/builtin.rs | 22 +- src/builtin/observer.rs | 642 --------------------------- src/compact.rs | 6 +- src/core.rs | 20 +- src/core/agent_observer.rs | 212 --------- src/core/observer.rs | 480 +++++++++++++++++++++ src/engine/bare.rs | 791 ++++++++-------------------------- src/engine/bare/compact.rs | 12 +- src/engine/bare/dispatch.rs | 187 +++----- src/engine/bare/emission.rs | 208 ++------- src/hooks.rs | 6 +- src/lib.rs | 6 +- src/loop_control.rs | 2 + src/loop_control/bundle.rs | 48 ++- src/loop_control/detection.rs | 8 + src/loop_control/fallback.rs | 2 - src/observability.rs | 72 ---- src/observability/console.rs | 231 ---------- src/observability/event.rs | 412 ------------------ src/observability/sink.rs | 410 ------------------ src/stream/heartbeat.rs | 13 +- 22 files changed, 843 insertions(+), 2948 deletions(-) delete mode 100644 src/builtin/observer.rs delete mode 100644 src/core/agent_observer.rs create mode 100644 src/core/observer.rs delete mode 100644 src/observability.rs delete mode 100644 src/observability/console.rs delete mode 100644 src/observability/event.rs delete mode 100644 src/observability/sink.rs diff --git a/README.md b/README.md index a85e34e..62da2d0 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,6 @@ and tool implementations; the framework handles the rest. | [`engine`](https://docs.rs/loopctl/latest/loopctl/engine/index.html) | `BareLoop` — the default agent loop engine (stream → accumulate → dispatch tools → repeat) | | [`loop_control`](https://docs.rs/loopctl/latest/loopctl/loop_control/index.html)| Loop detection, convergence detection, fallback model chains, and manager bundle | | [`message`](https://docs.rs/loopctl/latest/loopctl/message/index.html) | Conversation types: `Message`, `MessagePart`, `ToolContent`, roles | -| [`observability`](https://docs.rs/loopctl/latest/loopctl/observability/index.html) | Structured event streaming: `EventSink` trait, `ObserveEvent`, `CompositeSink`, `ConsoleSink` | | [`stream`](https://docs.rs/loopctl/latest/loopctl/stream/index.html) | Streaming event types, accumulator, stop reasons, usage tracking | | [`tool`](https://docs.rs/loopctl/latest/loopctl/tool/index.html) | `Tool` trait, `ToolRegistry`, `ToolSchema`, `ToolOutput`, `FnTool` adapter | | [`testing`](https://docs.rs/loopctl/latest/loopctl/testing/index.html) | Mock API client, mock tools, and test fixture factories (feature-gated) | diff --git a/src/builtin.rs b/src/builtin.rs index 5deb793..1211fc9 100644 --- a/src/builtin.rs +++ b/src/builtin.rs @@ -7,37 +7,23 @@ //! //! # Available Implementations //! -//! | Implementation | Trait | Purpose | -//! |---------------------|-----------------------------------------------|--------------------------------| -//! | [`InMemoryStore`] | [`AgentMemory`](crate::core::AgentMemory) | `Vec`-backed memory store | -//! | [`NoOpObserver`] | [`AgentObserver`](crate::core::AgentObserver) | No-op (default) observer | -//! | [`LoggingObserver`] | [`AgentObserver`](crate::core::AgentObserver) | Logs all events via `tracing` | -//! | [`MultiObserver`] | [`AgentObserver`](crate::core::AgentObserver) | Fans out to multiple observers | +//! | Implementation | Trait | Purpose | +//! |-------------------|-----------------------------------------------|---------------------------| +//! | [`InMemoryStore`] | [`AgentMemory`](crate::core::AgentMemory) | `Vec`-backed memory store | //! //! # Quick Start //! //! ```rust -//! use loopctl::builtin::{InMemoryStore, LoggingObserver, MultiObserver, NoOpObserver}; +//! use loopctl::builtin::InMemoryStore; //! use loopctl::core::{AgentMemory, MemoryEntry, MemoryCategory}; //! //! # tokio::runtime::Runtime::new().unwrap().block_on(async { //! // In-memory store for agent memories //! let mut store = InMemoryStore::new(); //! store.store(MemoryEntry::new(MemoryCategory::Fact, "PostgreSQL 15 is used")).await.unwrap(); -//! -//! // Logging observer -//! let observer = LoggingObserver; -//! -//! // Compose multiple observers -//! let multi = MultiObserver::new() -//! .with(LoggingObserver) -//! .with(NoOpObserver); -//! assert_eq!(multi.len(), 2); //! # }); //! ``` pub mod memory; -pub mod observer; pub use memory::InMemoryStore; -pub use observer::{LoggingObserver, MultiObserver, NoOpObserver}; diff --git a/src/builtin/observer.rs b/src/builtin/observer.rs deleted file mode 100644 index 011a372..0000000 --- a/src/builtin/observer.rs +++ /dev/null @@ -1,642 +0,0 @@ -//! Reference observer implementations — ready-to-use [`AgentObserver`] backends. -//! -//! This module provides concrete implementations of the [`AgentObserver`] trait -//! so consumers can plug in observability without implementing the trait from -//! scratch. Each implementation covers a different use case, from silent -//! no-ops to full `tracing`-based logging to composite fan-out. -//! -//! # Provided Implementations -//! -//! - **[`NoOpObserver`]** — A zero-cost no-op that silently ignores all -//! lifecycle events. Useful as a safe default and in tests. -//! - **[`LoggingObserver`]** — Emits structured log lines for every event -//! via the `tracing` crate. Session and compaction events are logged at -//! `info` level; tool events at `debug`; errors and fallbacks at `warn`. -//! - **[`MultiObserver`]** — A composite observer that fans out every event -//! to an ordered list of inner observers. Follows the composite pattern. -//! -//! # Quick Start -//! -//! ``` -//! use loopctl::builtin::observer::{LoggingObserver, MultiObserver, NoOpObserver}; -//! use loopctl::core::AgentObserver; -//! -//! // Use a single logging observer: -//! let observer = LoggingObserver; -//! -//! // Or compose multiple observers: -//! let multi = MultiObserver::new() -//! .with(LoggingObserver) -//! .with(NoOpObserver); -//! -//! multi.on_session_start(uuid::Uuid::new_v4()); -//! ``` - -use crate::core::AgentObserver; -use std::panic::{AssertUnwindSafe, catch_unwind}; -use std::sync::Arc; -use std::time::Duration; -use tracing::warn; - -/// A no-op observer that silently ignores all events. -/// -/// Every method body is empty — no allocation, no I/O, no side effects. -/// Use this as a safe default when no observability is needed, as a -/// placeholder in tests, or as a baseline when benchmarking the framework -/// overhead *without* observer noise. -/// -/// # Example -/// -/// ``` -/// use loopctl::builtin::observer::NoOpObserver; -/// use loopctl::core::AgentObserver; -/// -/// let observer = NoOpObserver; -/// observer.on_session_start(uuid::Uuid::new_v4()); // does nothing -/// observer.on_tool_call("Bash", r#"{"command": "ls"}"#); // does nothing -/// ``` -pub struct NoOpObserver; - -impl AgentObserver for NoOpObserver {} - -/// An observer that logs all lifecycle events via `tracing`. -/// -/// Emits a structured log line for each [`AgentObserver`] callback. The -/// log levels are chosen so that normal operation is visible at `info` -/// level while detailed tool execution is available at `debug` level: -/// -/// | Event | Level | -/// |---------------------------|------------------| -/// | session start / end | `info` | -/// | context compaction | `info` | -/// | turn start / end | `info` / `debug` | -/// | tool call / complete | `debug` | -/// | context warning | `warn` | -/// | fallback triggered | `warn` | -/// | errors / failure reasons | `warn` | -/// -/// Query strings in [`on_turn_start`](AgentObserver::on_turn_start) are -/// previewed to a maximum of 80 characters to avoid flooding logs with -/// long prompts. -/// -/// # Structured Fields -/// -/// Every log line attaches machine-readable fields (`session_id`, `tool`, -/// `duration`, `used_tokens`, etc.) so that log aggregation backends -/// can filter and chart without parsing free-form text. -/// -/// # Thread Safety -/// -/// [`LoggingObserver`] is a zero-sized unit struct — it carries no state -/// and can be freely shared across threads or cloned at zero cost. -/// -/// # Example -/// -/// ``` -/// use loopctl::builtin::observer::LoggingObserver; -/// use loopctl::core::AgentObserver; -/// -/// let observer = LoggingObserver; -/// observer.on_session_start(uuid::Uuid::new_v4()); -/// ``` -pub struct LoggingObserver; - -impl AgentObserver for LoggingObserver { - /// Log session start at `info` level. - /// - /// Emits the `session_id` as a structured field for correlation in - /// log aggregation systems. - fn on_session_start(&self, session_id: uuid::Uuid) { - tracing::info!(%session_id, "Session started"); - } - - /// Log session end at `info` or `warn` level depending on outcome. - /// - /// Successful sessions are logged at `info`; failures and sessions - /// with notes are logged at `warn` with the error message attached. - fn on_session_end(&self, success: bool, error: Option<&str>) { - match (success, error) { - (true, None) => tracing::info!("Session ended successfully"), - (true, Some(e)) => tracing::info!(error = e, "Session ended with note"), - (false, Some(e)) => tracing::warn!(error = e, "Session failed"), - (false, None) => tracing::warn!("Session ended unsuccessfully"), - } - } - - /// Log turn start at `info` level with a truncated query preview. - /// - /// The query is previewed to at most 80 characters to keep log lines - /// readable when dealing with long prompts. - fn on_turn_start(&self, query: &str) { - let preview = if query.len() > 80 { - let end = query - .char_indices() - .take_while(|(i, _)| *i < 80) - .last() - .map_or(0, |(i, c)| i.saturating_add(c.len_utf8())); - &query[..end] - } else { - query - }; - tracing::info!(query = preview, "Turn started"); - } - - /// Log turn completion at `debug` (success) or `warn` (failure) level. - /// - /// Successful turns emit a short `debug` line; failed turns include the - /// error reason at `warn` level for visibility in production alerts. - fn on_turn_end(&self, success: bool, error_reason: Option<&str>) { - if success { - tracing::debug!("Turn completed"); - } else { - tracing::warn!(reason = error_reason.unwrap_or("unknown"), "Turn failed"); - } - } - - /// Log a tool invocation at `debug` level. - /// - /// Emits the tool name as a structured field. The input payload is - /// accepted but not logged to avoid leaking sensitive data (e.g. file - /// contents, API keys). - fn on_tool_call(&self, tool: &str, _input: &str) { - tracing::debug!(tool, "Tool called"); - } - - /// Log tool completion at `debug` (success) or `warn` (failure) level. - /// - /// Includes the tool name, wall-clock `duration`, and — on failure — - /// the error message. The `duration` field is useful for identifying - /// slow tools in production. - fn on_tool_complete( - &self, - tool: &str, - _input: &str, - _output: &str, - duration: Duration, - success: bool, - error: Option<&str>, - ) { - if success { - tracing::debug!(tool, ?duration, "Tool completed"); - } else { - tracing::warn!( - tool, - ?duration, - error = error.unwrap_or("unknown"), - "Tool failed" - ); - } - } - - /// Log a context-window warning at `warn` level. - /// - /// Emits the `used_tokens` and `remaining_tokens` counts so operators - /// can correlate context pressure with agent behaviour. - fn on_context_warning(&self, used_tokens: u64, remaining_tokens: u64) { - tracing::warn!(used_tokens, remaining_tokens, "Context window running low"); - } - - /// Log a context compaction event at `info` level. - /// - /// Emits the message counts before and after so operators can verify - /// that compaction is reducing context size as expected. - fn on_compaction(&self, messages_before: usize, messages_after: usize) { - tracing::info!(messages_before, messages_after, "Context compacted"); - } - - /// Log a model fallback event at `warn` level. - /// - /// Emits both the `from` and `to` model identifiers so operators can - /// detect chronic primary-model failures. - fn on_fallback(&self, from: &str, to: &str) { - tracing::warn!(from, to, "Fallback triggered"); - } - - /// Log a loop-detection event at `warn` level. - /// - /// Emits the `tool` name and `repetitions` count so operators can - /// identify tools that the agent is calling in a repeated pattern. - fn on_loop_detected(&self, tool: &str, repetitions: usize) { - warn!(tool, repetitions, "loop detected in tool operations"); - } - - /// Log a convergence-detection event at `warn` level. - /// - /// Emits the configured `action` (e.g. `"stop"`, `"warn"`) so - /// operators can see when the agent's responses have become - /// semantically similar and what the framework intends to do about it. - fn on_convergence_detected(&self, action: &str) { - warn!(action, "convergence detected in agent responses"); - } -} - -/// A composite observer that fans out every event to multiple inner observers. -/// -/// Holds an ordered list of [`AgentObserver`] trait objects behind `Arc` -/// and forwards each callback to every inner observer in sequence. This -/// implements the classic **Composite** pattern, allowing you to combine -/// logging, metrics, and custom observers without writing a new struct. -/// -/// Observers are called in insertion order. If any individual observer -/// panics, the remaining observers in the list are still called — this -/// is achieved internally via [`std::panic::catch_unwind`] (or the -/// caller's panic handler). This ensures one misbehaving observer does -/// not prevent others from receiving events. -/// -/// # Architecture -/// -/// ```text -/// MultiObserver -/// ┌──────────────┐ -/// on_xxx() ──► │ observers[] │──► obs[0].on_xxx() -/// │ │──► obs[1].on_xxx() -/// │ │──► obs[2].on_xxx() -/// └──────────────┘ -/// ``` -/// -/// # Thread Safety -/// -/// Each inner observer is stored as [`Arc`]``, so the -/// same observer can be shared across multiple [`MultiObserver`] instances -/// or even other parts of the system. The fan-out loop borrows `&self`, -/// meaning all callbacks must be `&self`-safe (no `&mut self`). -/// -/// # Example -/// -/// ``` -/// use loopctl::builtin::observer::{LoggingObserver, MultiObserver, NoOpObserver}; -/// use loopctl::core::AgentObserver; -/// use std::sync::Arc; -/// -/// // Build with owned observers: -/// let multi = MultiObserver::new() -/// .with(LoggingObserver) -/// .with(NoOpObserver); -/// assert_eq!(multi.len(), 2); -/// -/// // Build with pre-Arc'd observers: -/// let shared = Arc::new(LoggingObserver); -/// let multi = MultiObserver::new() -/// .with_arc(shared.clone()) -/// .with_arc(shared); // same observer twice -/// assert_eq!(multi.len(), 2); -/// ``` -pub struct MultiObserver { - /// The ordered list of inner observers to fan out to. - /// - /// Each observer is stored as `Arc` so it can be - /// shared across threads if needed. Observers are called in insertion - /// order — first added is first notified. - /// - /// Starts empty when created via [`new`](MultiObserver::new) or - /// [`default`](MultiObserver::default). Observers are added via - /// [`with`](MultiObserver::with) or [`with_arc`](MultiObserver::with_arc). - observers: Vec>, -} - -// =================================================== -// Construction -// =================================================== - -impl MultiObserver { - /// Create an empty multi-observer with no inner observers. - /// - /// All callbacks will be no-ops until observers are added via - /// [`with`](MultiObserver::with) or [`with_arc`](MultiObserver::with_arc). - /// - /// # Example - /// - /// ``` - /// use loopctl::builtin::observer::MultiObserver; - /// - /// let multi = MultiObserver::new(); - /// assert!(multi.is_empty()); - /// ``` - #[must_use] - pub fn new() -> Self { - Self { - observers: Vec::new(), - } - } - - /// Add an owned observer to the fan-out list. - /// - /// The observer is boxed as `Arc` and appended to - /// the end of the list. Returns `self` for chaining. - /// - /// # Example - /// - /// ``` - /// use loopctl::builtin::observer::{LoggingObserver, MultiObserver, NoOpObserver}; - /// - /// let multi = MultiObserver::new() - /// .with(LoggingObserver) - /// .with(NoOpObserver); - /// assert_eq!(multi.len(), 2); - /// ``` - #[must_use] - pub fn with(mut self, observer: O) -> Self { - self.observers.push(Arc::new(observer)); - self - } - - /// Add an observer that is already behind an `Arc`. - /// - /// Useful when multiple [`MultiObserver`] instances need to share the - /// same inner observer, or when the observer is constructed externally - /// and already wrapped in an `Arc`. Accepts both `Arc` - /// and `Arc`. - /// - /// # Example - /// - /// ``` - /// use loopctl::builtin::observer::{LoggingObserver, MultiObserver}; - /// use std::sync::Arc; - /// - /// // Arc a concrete type: - /// let shared = Arc::new(LoggingObserver); - /// let multi = MultiObserver::new() - /// .with_arc(shared.clone()) - /// .with_arc(shared); - /// assert_eq!(multi.len(), 2); - /// ``` - #[must_use] - pub fn with_arc(mut self, observer: Arc) -> Self { - self.observers.push(observer); - self - } - - /// Number of observers in the fan-out list. - /// - /// Returns the count of inner observers that will receive events. - /// Mirrors the `Vec::len` of the internal observers list. - /// - /// # Example - /// - /// ``` - /// use loopctl::builtin::observer::{LoggingObserver, MultiObserver}; - /// - /// let multi = MultiObserver::new().with(LoggingObserver); - /// assert_eq!(multi.len(), 1); - /// ``` - #[must_use] - pub fn len(&self) -> usize { - self.observers.len() - } - - /// Whether there are no observers in the fan-out list. - /// - /// Returns `true` when [`len`](MultiObserver::len) is zero. When empty, - /// all callbacks are effectively no-ops (the internal loop body never - /// executes). - /// - /// # Example - /// - /// ``` - /// use loopctl::builtin::observer::MultiObserver; - /// - /// let multi = MultiObserver::new(); - /// assert!(multi.is_empty()); - /// ``` - #[must_use] - pub fn is_empty(&self) -> bool { - self.observers.is_empty() - } -} - -impl Default for MultiObserver { - /// Returns an empty multi-observer, equivalent to [`new`](MultiObserver::new). - /// - /// The default instance contains zero inner observers, so all callbacks - /// are effectively no-ops until observers are added. - /// - /// # Example - /// - /// ``` - /// use loopctl::builtin::observer::MultiObserver; - /// - /// let multi = MultiObserver::default(); - /// assert!(multi.is_empty()); - /// ``` - fn default() -> Self { - Self::new() - } -} - -// =================================================== -// AgentObserver implementation -// =================================================== - -impl AgentObserver for MultiObserver { - fn on_session_start(&self, session_id: uuid::Uuid) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_session_start(session_id); - })) { - warn!("observer panicked in on_session_start: {payload:?}"); - } - } - } - - fn on_session_end(&self, success: bool, error: Option<&str>) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_session_end(success, error); - })) { - warn!("observer panicked in on_session_end: {payload:?}"); - } - } - } - - fn on_turn_start(&self, query: &str) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_turn_start(query); - })) { - warn!("observer panicked in on_turn_start: {payload:?}"); - } - } - } - - fn on_turn_end(&self, success: bool, error_reason: Option<&str>) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_turn_end(success, error_reason); - })) { - warn!("observer panicked in on_turn_end: {payload:?}"); - } - } - } - - fn on_tool_call(&self, tool: &str, input: &str) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_tool_call(tool, input); - })) { - warn!("observer panicked in on_tool_call: {payload:?}"); - } - } - } - - fn on_tool_complete( - &self, - tool: &str, - input: &str, - output: &str, - duration: Duration, - success: bool, - error: Option<&str>, - ) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_tool_complete(tool, input, output, duration, success, error); - })) { - warn!("observer panicked in on_tool_complete: {payload:?}"); - } - } - } - - fn on_context_warning(&self, used_tokens: u64, remaining_tokens: u64) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_context_warning(used_tokens, remaining_tokens); - })) { - warn!("observer panicked in on_context_warning: {payload:?}"); - } - } - } - - fn on_compaction(&self, messages_before: usize, messages_after: usize) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_compaction(messages_before, messages_after); - })) { - warn!("observer panicked in on_compaction: {payload:?}"); - } - } - } - - fn on_fallback(&self, from: &str, to: &str) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_fallback(from, to); - })) { - warn!("observer panicked in on_fallback: {payload:?}"); - } - } - } - - fn on_loop_detected(&self, tool: &str, repetitions: usize) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_loop_detected(tool, repetitions); - })) { - warn!("observer panicked in on_loop_detected: {payload:?}"); - } - } - } - - fn on_convergence_detected(&self, action: &str) { - for obs in &self.observers { - if let Err(payload) = catch_unwind(AssertUnwindSafe(|| { - obs.on_convergence_detected(action); - })) { - warn!("observer panicked in on_convergence_detected: {payload:?}"); - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn noop_observer_does_not_panic() { - let observer = NoOpObserver; - observer.on_session_start(uuid::Uuid::new_v4()); - observer.on_session_end(true, None); - observer.on_turn_start("test query"); - observer.on_turn_end(true, None); - observer.on_tool_call("Read", r#"{"file_path": "test.rs"}"#); - observer.on_tool_complete("Read", "input", "output", Duration::ZERO, true, None); - observer.on_context_warning(1000, 500); - observer.on_compaction(10, 5); - observer.on_fallback("model-a", "model-b"); - } - - #[test] - fn logging_observer_does_not_panic() { - let observer = LoggingObserver; - observer.on_session_start(uuid::Uuid::new_v4()); - observer.on_session_end(true, None); - observer.on_session_end(false, Some("test error")); - observer.on_turn_start("test query"); - observer.on_turn_end(true, None); - observer.on_turn_end(false, Some("some reason")); - observer.on_tool_call("Bash", r#"{"command": "ls"}"#); - observer.on_tool_complete( - "Bash", - "input", - "output", - Duration::from_millis(100), - true, - None, - ); - observer.on_tool_complete( - "Bash", - "input", - "", - Duration::from_millis(50), - false, - Some("command failed"), - ); - observer.on_context_warning(5000, 2000); - observer.on_compaction(20, 10); - observer.on_fallback("primary", "fallback"); - } - - #[test] - fn multi_observer_fans_out() { - use std::sync::atomic::{AtomicUsize, Ordering}; - - let call_count = Arc::new(AtomicUsize::new(0)); - - struct CountingObserver { - count: Arc, - } - - impl AgentObserver for CountingObserver { - fn on_session_start(&self, _session_id: uuid::Uuid) { - self.count.fetch_add(1, Ordering::Relaxed); - } - } - - let multi = MultiObserver::new() - .with(CountingObserver { - count: call_count.clone(), - }) - .with(CountingObserver { - count: call_count.clone(), - }) - .with(CountingObserver { - count: call_count.clone(), - }); - - multi.on_session_start(uuid::Uuid::new_v4()); - assert_eq!(call_count.load(Ordering::Relaxed), 3); - } - - #[test] - fn multi_observer_default_is_empty() { - let multi = MultiObserver::default(); - assert!(multi.is_empty()); - assert_eq!(multi.len(), 0); - } - - #[test] - fn multi_observer_with_arc() { - let multi = MultiObserver::new().with_arc(Arc::new(NoOpObserver)); - assert_eq!(multi.len(), 1); - } -} diff --git a/src/compact.rs b/src/compact.rs index 519891c..0c3afbf 100644 --- a/src/compact.rs +++ b/src/compact.rs @@ -283,7 +283,7 @@ impl CompactionOutcome { /// /// Produced by [`ContextManager::ensure_context_fits`] when compaction /// occurs. Observers receive this via -/// [`on_compaction`](crate::core::AgentObserver::on_compaction). +/// [`on_compaction`](crate::core::observer::LoopObserver::on_compaction). #[derive(Debug, Clone)] pub struct CompactTelemetry { /// Why compaction was triggered. @@ -836,7 +836,9 @@ impl ContextManager { /// The token budget at which compaction triggers. /// - /// Equal to `context_window * threshold`. + /// Equal to `context_window * threshold`. The result is always + /// non-negative (percentage × positive count), so the f64→u64 + /// cast is safe in practice. #[must_use] #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] pub fn compact_threshold_tokens(&self) -> u64 { diff --git a/src/core.rs b/src/core.rs index a5e6587..36a0149 100644 --- a/src/core.rs +++ b/src/core.rs @@ -6,11 +6,11 @@ //! //! # Traits //! -//! | Trait | Purpose | -//! |--------------------|-----------------------------------------------------| -//! | [`AgentCore`] | Main lifecycle trait for all agent types | -//! | [`AgentMemory`] | Interface for agent memory backends | -//! | [`AgentObserver`] | Lifecycle event hooks for monitoring agents | +//! | Trait | Purpose | +//! |------------------------------------------|-----------------------------------------------------| +//! | [`AgentCore`] | Main lifecycle trait for all agent types | +//! | [`AgentMemory`] | Interface for agent memory backends | +//! | [`LoopObserver`](observer::LoopObserver) | Notification-only observer for agent lifecycle | //! //! # Supporting Types //! @@ -37,17 +37,23 @@ //! | [`ToolCall`] | A tool call requested by the agent | //! | [`ToolDispatchResult`] | Result of a single tool execution | //! | [`TurnResult`] | Result of a single agent turn | +//! +//! # Sub-modules +//! +//! - **[`observer`]** — [`LoopObserver`](observer::LoopObserver) trait, context structs, and +//! [`ObserverHost`](observer::ObserverHost) for lifecycle notification. Observers are passive — +//! they receive callbacks but cannot control flow. For flow control, see the +//! [hook system](crate::hooks). pub mod agent_core; pub mod agent_memory; -pub mod agent_observer; pub mod error; +pub mod observer; pub mod reflection; pub mod types; pub use agent_core::*; pub use agent_memory::*; -pub use agent_observer::*; pub use error::*; pub use reflection::*; pub use types::*; diff --git a/src/core/agent_observer.rs b/src/core/agent_observer.rs deleted file mode 100644 index d2ef640..0000000 --- a/src/core/agent_observer.rs +++ /dev/null @@ -1,212 +0,0 @@ -//! Agent observer trait — lifecycle event hooks for monitoring agents. -//! -//! Observers receive callbacks at key lifecycle points. -//! This enables logging, metrics, trajectory recording, -//! and other cross-cutting concerns without modifying agent logic. -//! -//! This is the **canonical** observer trait, shared between the framework -//! and production agent crates. All methods have default no-op implementations, -//! so consumers only override what they need. -//! -//! # Provided Implementations -//! -//! - **`NoOpObserver`** — Default no-op (does nothing). -//! - **`LoggingObserver`** — Logs all events via `tracing`. -//! - **`MultiObserver`** — Fans out to multiple observers. -//! -//! # Quick Start -//! -//! ``` -//! use loopctl::core::agent_observer::AgentObserver; -//! use std::time::Duration; -//! -//! struct MetricsObserver; -//! -//! impl AgentObserver for MetricsObserver { -//! fn on_turn_start(&self, _query: &str) {} -//! fn on_tool_complete(&self, tool: &str, _input: &str, _output: &str, -//! duration: Duration, success: bool, _error: Option<&str>) {} -//! fn on_session_end(&self, success: bool, _error: Option<&str>) {} -//! } -//! ``` - -use std::time::Duration; - -/// Observer trait for agent lifecycle events. -/// -/// All methods have default no-op implementations, so consumers -/// only override what they need. Implementors must be [`Send + Sync`] -/// because observers are shared across async tasks. -/// -/// # Lifecycle -/// -/// ```text -/// on_session_start(session_id) -/// → on_turn_start(query) [once per turn] -/// → on_tool_call(tool, input) [before each tool] -/// → on_tool_complete(tool, ...) [after each tool] -/// → on_turn_end(success, error) [once per turn] -/// → ... -/// on_session_end(success, error) -/// ``` -/// -/// # Implementing -/// -/// Override only the methods you care about. Every method has a no-op -/// default, so a minimal observer can be just an empty `impl`: -/// -/// ``` -/// use loopctl::core::agent_observer::AgentObserver; -/// -/// struct MyObserver; -/// impl AgentObserver for MyObserver {} // all methods are no-ops -/// ``` -/// -/// # Example -/// -/// ``` -/// use loopctl::core::agent_observer::AgentObserver; -/// use std::time::Duration; -/// -/// struct LoggingObserver; -/// -/// impl AgentObserver for LoggingObserver { -/// fn on_turn_start(&self, _query: &str) {} -/// -/// fn on_tool_complete(&self, tool: &str, _input: &str, _output: &str, -/// duration: Duration, _success: bool, _error: Option<&str>) {} -/// } -/// ``` -pub trait AgentObserver: Send + Sync { - /// Called when an agent session begins. - /// - /// Fired once at the start of `AgentCore::initialize`. The - /// `session_id` matches the one passed to the agent configuration - /// and can be used to correlate all subsequent events for this session. - fn on_session_start(&self, _session_id: uuid::Uuid) {} - - /// Called when an agent session ends. - /// - /// Fired once after `AgentCore::finalize`. `success` indicates - /// whether the session completed normally; `error` contains a - /// description when the session ended due to a failure. - fn on_session_end(&self, _success: bool, _error: Option<&str>) {} - - /// Called at the start of processing a turn. - /// - /// Fired before `AgentCore::process_turn` is invoked. The `query` - /// is the raw user input for this turn. - fn on_turn_start(&self, _query: &str) {} - - /// Called when a turn completes. - /// - /// Fired after `AgentCore::process_turn` returns. `success` - /// indicates whether the turn completed without error; when - /// `false`, `error_reason` describes what went wrong. - fn on_turn_end(&self, _success: bool, _error_reason: Option<&str>) {} - - /// Called before a tool is executed. - /// - /// Fired just before the framework dispatches a tool call. `tool` - /// is the tool name and `input` is the raw input string. Use this - /// to log tool usage or record the start time for custom timing. - fn on_tool_call(&self, _tool: &str, _input: &str) {} - - /// Called after a tool completes execution. - /// - /// Fired after the tool returns. The parameters provide the full - /// execution context: - /// - /// - `tool` — Tool name that was invoked. - /// - `input` — Raw input string passed to the tool. - /// - `output` — Result string returned by the tool. - /// - `duration` — Wall-clock execution time. - /// - `success` — Whether the tool reported success. - /// - `error` — Error message if the tool failed, `None` otherwise. - fn on_tool_complete( - &self, - _tool: &str, - _input: &str, - _output: &str, - _duration: Duration, - _success: bool, - _error: Option<&str>, - ) { - } - - /// Called when the context window is running low. - /// - /// Fired when the number of remaining tokens drops below a - /// configurable threshold defined in the agent configuration. - /// Use this to trigger - /// pre-emptive compaction or warn downstream consumers. - fn on_context_warning(&self, _used_tokens: u64, _remaining_tokens: u64) {} - - /// Called when a context compaction occurs. - /// - /// Fired after the framework compresses the conversation history. - /// `messages_before` and `messages_after` indicate how aggressive - /// the compaction was — useful for monitoring whether compaction - /// is losing too much context. - fn on_compaction(&self, _messages_before: usize, _messages_after: usize) {} - - /// Called when a fallback is triggered (e.g. switching models). - /// - /// Fired by the fallback manager when the primary model fails - /// and the framework switches to a backup. - /// `from` is the model that failed; `to` is the replacement. - fn on_fallback(&self, _from: &str, _to: &str) {} - - /// Called when a loop is detected in tool operations. - /// - /// Fired by the detection manager when the same tool call pattern - /// has been repeated beyond the configured loop threshold. - /// `tool` is the repeating tool name; `repetitions` is the count. - fn on_loop_detected(&self, _tool: &str, _repetitions: usize) {} - - /// Called when convergence is detected in agent responses. - /// - /// Fired by the detection manager when recent agent responses - /// have become semantically similar beyond the configured threshold. - /// `action` describes the configured response (e.g. `"stop"`, `"warn"`). - fn on_convergence_detected(&self, _action: &str) {} -} - -#[cfg(test)] -mod tests { - use super::*; - use std::time::Duration; - - #[test] - fn default_impl_noop_does_not_panic() { - struct Nop; - impl AgentObserver for Nop {} - - let nop = Nop; - nop.on_session_start(uuid::Uuid::nil()); - nop.on_session_end(true, None); - nop.on_turn_start("hello"); - nop.on_turn_end(true, None); - nop.on_tool_call("read_file", "/tmp/x"); - nop.on_tool_complete( - "read_file", - "in", - "out", - Duration::from_millis(5), - true, - None, - ); - nop.on_context_warning(100_000, 28_000); - nop.on_compaction(50, 20); - nop.on_fallback("llm-1", "llm-2"); - } - - #[test] - fn empty_impl_is_send_sync() { - struct EmptyObs; - impl AgentObserver for EmptyObs {} - - fn assert_send_sync() {} - assert_send_sync::(); - } -} diff --git a/src/core/observer.rs b/src/core/observer.rs new file mode 100644 index 0000000..bf8a7ca --- /dev/null +++ b/src/core/observer.rs @@ -0,0 +1,480 @@ +//! Typed lifecycle observer for agent loops. +//! +//! Defines [`LoopObserver`] — a notification-only trait with one method per +//! lifecycle point — and [`ObserverHost`], which holds registered observers +//! and dispatches notifications to each in order. +//! +//! # Context Structs +//! +//! Each callback receives a typed context struct with relevant fields: +//! +//! - [`SessionStartContext`] / [`SessionEndContext`] — session boundaries +//! - [`TurnStartContext`] / [`TurnEndContext`] — turn boundaries +//! - [`StreamContext`] / [`StreamFailureContext`] — stream success/failure +//! - [`ResponseContext`] — model response text and usage +//! - [`ToolPreContext`] / [`ToolPostContext`] — tool dispatch lifecycle +//! - [`CompactionContext`] — context window compaction +//! - [`FallbackContext`] — model fallback event +//! - [`LoopDetectedContext`] — loop detection event +//! - [`ConvergenceDetectedContext`] — convergence detection event +//! +//! # Example +//! +//! ```rust,ignore +//! use loopctl::core::observer::{LoopObserver, SessionStartContext}; +//! +//! struct MetricsObserver; +//! +//! impl LoopObserver for MetricsObserver { +//! fn name(&self) -> &str { "metrics" } +//! +//! fn on_session_start(&self, ctx: &SessionStartContext) { +//! println!("session {} started", ctx.session_id); +//! } +//! } +//! ``` + +use std::sync::Arc; + +// ================================================== +// Context structs +// ================================================== + +/// Context for [`LoopObserver::on_session_start`]. +#[derive(Debug, Clone)] +pub struct SessionStartContext { + /// Unique session identifier. + pub session_id: uuid::Uuid, +} + +/// Context for [`LoopObserver::on_session_end`]. +#[derive(Debug, Clone)] +pub struct SessionEndContext { + /// Whether the session completed successfully. + pub success: bool, + /// Error description, if the session ended due to an error. + pub error: Option, + /// Total turns completed during the session. + pub total_turns: usize, + /// Total session duration in milliseconds. + pub duration_ms: u64, +} + +/// Context for [`LoopObserver::on_turn_start`]. +#[derive(Debug, Clone)] +pub struct TurnStartContext { + /// Turn number (0-indexed). + pub turn: usize, + /// The user query that initiated this turn. + pub query: String, +} + +/// Context for [`LoopObserver::on_turn_end`]. +#[derive(Debug, Clone)] +pub struct TurnEndContext { + /// Turn number. + pub turn: usize, + /// Whether the turn completed successfully. + pub success: bool, + /// Error description, if the turn failed. + pub error: Option, + /// Wall-clock duration of the turn in milliseconds. + pub duration_ms: u64, + /// Input tokens consumed this turn. + pub input_tokens: u64, + /// Output tokens generated this turn. + pub output_tokens: u64, +} + +/// Context for [`LoopObserver::on_stream_success`]. +#[derive(Debug, Clone)] +pub struct StreamContext { + /// Turn number. + pub turn: usize, + /// Model that was streamed. + pub model: String, + /// Input tokens consumed. + pub input_tokens: u64, + /// Output tokens generated. + pub output_tokens: u64, +} + +/// Context for [`LoopObserver::on_stream_failure`]. +#[derive(Debug, Clone)] +pub struct StreamFailureContext { + /// Turn number. + pub turn: usize, + /// Model that failed. + pub model: String, + /// The error that occurred. + pub error: crate::core::AgentError, +} + +/// Context for [`LoopObserver::on_response`]. +#[derive(Debug, Clone)] +pub struct ResponseContext { + /// Turn number. + pub turn: usize, + /// The model's text response. + pub text: String, + /// Token usage for this turn, if available. + pub usage: Option, +} + +/// Context for [`LoopObserver::on_tool_pre`]. +#[derive(Debug, Clone)] +pub struct ToolPreContext { + /// Turn number. + pub turn: usize, + /// Tool name. + pub tool: String, + /// Tool call ID from the API response. + pub tool_call_id: String, +} + +/// Context for [`LoopObserver::on_tool_post`]. +#[derive(Debug, Clone)] +pub struct ToolPostContext { + /// Turn number. + pub turn: usize, + /// Tool name. + pub tool: String, + /// Deterministic hash of the tool output, if available. + pub result_hash: Option, + /// Whether the tool returned an error. + pub is_error: bool, + /// Wall-clock execution duration. + pub duration: std::time::Duration, +} + +/// Context for [`LoopObserver::on_compaction`]. +#[derive(Debug, Clone)] +pub struct CompactionContext { + /// Message count before compaction. + pub messages_before: usize, + /// Message count after compaction. + pub messages_after: usize, + /// Estimated tokens saved by compaction. + pub tokens_saved: u64, +} + +/// Context for [`LoopObserver::on_fallback`]. +#[derive(Debug, Clone)] +pub struct FallbackContext { + /// Model that failed. + pub from: String, + /// Replacement model. + pub to: String, +} + +/// Context for [`LoopObserver::on_loop_detected`]. +#[derive(Debug, Clone)] +pub struct LoopDetectedContext { + /// Description of the repeating tool pattern. + pub pattern: String, + /// Number of times the pattern was observed. + pub repetitions: usize, +} + +/// Context for [`LoopObserver::on_convergence_detected`]. +#[derive(Debug, Clone)] +pub struct ConvergenceDetectedContext { + /// Configured action to take (e.g. `"stop"`, `"warn"`, `"compact"`). + pub action: String, +} + +// ================================================== +// LoopObserver Trait +// ================================================== + +/// A notification observer that receives typed callbacks at agent loop lifecycle points. +/// +/// Observers are registered via `BareLoop` or `ManagerBundle` and called at each +/// lifecycle point in registration order. All methods are **notification-only** — they +/// return `()`. Use the [hook system](crate::hooks) if you need to control +/// flow (block/allow actions). +/// +/// All methods have default no-op implementations. Override only the callbacks you need. +pub trait LoopObserver: Send + Sync { + /// Human-readable name for diagnostics and logging. + fn name(&self) -> &str; + + /// Called when an agent session begins. + fn on_session_start(&self, _ctx: &SessionStartContext) {} + + /// Called when an agent session ends. + fn on_session_end(&self, _ctx: &SessionEndContext) {} + + /// Called at the start of processing a turn. + fn on_turn_start(&self, _ctx: &TurnStartContext) {} + + /// Called when a turn completes. + fn on_turn_end(&self, _ctx: &TurnEndContext) {} + + /// Called after the model streams a response successfully. + fn on_stream_success(&self, _ctx: &StreamContext) {} + + /// Called when the API stream fails. + fn on_stream_failure(&self, _ctx: &StreamFailureContext) {} + + /// Called after extracting the model's text response. + fn on_response(&self, _ctx: &ResponseContext) {} + + /// Called before a tool is dispatched (notification-only). + fn on_tool_pre(&self, _ctx: &ToolPreContext) {} + + /// Called after a tool completes execution. + fn on_tool_post(&self, _ctx: &ToolPostContext) {} + + /// Called after conversation compaction. + fn on_compaction(&self, _ctx: &CompactionContext) {} + + /// Called when a model fallback is triggered. + fn on_fallback(&self, _ctx: &FallbackContext) {} + + /// Called when a loop is detected in tool operations. + fn on_loop_detected(&self, _ctx: &LoopDetectedContext) {} + + /// Called when response convergence is detected. + fn on_convergence_detected(&self, _ctx: &ConvergenceDetectedContext) {} + + /// Reset observer state for a new session. + fn reset(&self) {} +} + +// ================================================== +// ObserverHost +// ================================================== + +/// Holds registered observers and dispatches notifications to each. +/// +/// Observers run in registration order. All observers are always notified — +/// there is no short-circuiting (that's the [hook system](crate::hooks)'s job). +/// +/// An empty host (no observers registered) is effectively zero-cost: +/// each notification call iterates an empty `Vec`. +#[derive(Default)] +pub struct ObserverHost { + observers: Vec>, +} + +impl ObserverHost { + /// Create an empty observer host. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Register an observer. Called in registration order at each notification point. + pub fn register(&mut self, observer: Arc) { + self.observers.push(observer); + } + + /// Number of registered observers. + #[must_use] + pub fn len(&self) -> usize { + self.observers.len() + } + + /// Whether no observers are registered. + #[must_use] + pub fn is_empty(&self) -> bool { + self.observers.is_empty() + } + + /// Reset all observers for a new session. + pub fn reset_all(&self) { + for obs in &self.observers { + obs.reset(); + } + } + + /// Dispatch [`LoopObserver::on_session_start`] to all observers. + pub fn on_session_start(&self, ctx: &SessionStartContext) { + for obs in &self.observers { + obs.on_session_start(ctx); + } + } + + /// Dispatch [`LoopObserver::on_session_end`] to all observers. + pub fn on_session_end(&self, ctx: &SessionEndContext) { + for obs in &self.observers { + obs.on_session_end(ctx); + } + } + + /// Dispatch [`LoopObserver::on_turn_start`] to all observers. + pub fn on_turn_start(&self, ctx: &TurnStartContext) { + for obs in &self.observers { + obs.on_turn_start(ctx); + } + } + + /// Dispatch [`LoopObserver::on_turn_end`] to all observers. + pub fn on_turn_end(&self, ctx: &TurnEndContext) { + for obs in &self.observers { + obs.on_turn_end(ctx); + } + } + + /// Dispatch [`LoopObserver::on_stream_success`] to all observers. + pub fn on_stream_success(&self, ctx: &StreamContext) { + for obs in &self.observers { + obs.on_stream_success(ctx); + } + } + + /// Dispatch [`LoopObserver::on_stream_failure`] to all observers. + pub fn on_stream_failure(&self, ctx: &StreamFailureContext) { + for obs in &self.observers { + obs.on_stream_failure(ctx); + } + } + + /// Dispatch [`LoopObserver::on_response`] to all observers. + pub fn on_response(&self, ctx: &ResponseContext) { + for obs in &self.observers { + obs.on_response(ctx); + } + } + + /// Dispatch [`LoopObserver::on_tool_pre`] to all observers. + pub fn on_tool_pre(&self, ctx: &ToolPreContext) { + for obs in &self.observers { + obs.on_tool_pre(ctx); + } + } + + /// Dispatch [`LoopObserver::on_tool_post`] to all observers. + pub fn on_tool_post(&self, ctx: &ToolPostContext) { + for obs in &self.observers { + obs.on_tool_post(ctx); + } + } + + /// Dispatch [`LoopObserver::on_compaction`] to all observers. + pub fn on_compaction(&self, ctx: &CompactionContext) { + for obs in &self.observers { + obs.on_compaction(ctx); + } + } + + /// Dispatch [`LoopObserver::on_fallback`] to all observers. + pub fn on_fallback(&self, ctx: &FallbackContext) { + for obs in &self.observers { + obs.on_fallback(ctx); + } + } + + /// Dispatch [`LoopObserver::on_loop_detected`] to all observers. + pub fn on_loop_detected(&self, ctx: &LoopDetectedContext) { + for obs in &self.observers { + obs.on_loop_detected(ctx); + } + } + + /// Dispatch [`LoopObserver::on_convergence_detected`] to all observers. + pub fn on_convergence_detected(&self, ctx: &ConvergenceDetectedContext) { + for obs in &self.observers { + obs.on_convergence_detected(ctx); + } + } +} + +// ================================================== +// Tests +// ================================================== + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// A test observer that counts notification invocations. + struct CountingObserver { + name: &'static str, + stream_success: AtomicUsize, + resets: AtomicUsize, + } + + impl CountingObserver { + fn new(name: &'static str) -> Self { + Self { + name, + stream_success: AtomicUsize::new(0), + resets: AtomicUsize::new(0), + } + } + } + + impl LoopObserver for CountingObserver { + fn name(&self) -> &str { + self.name + } + + fn on_stream_success(&self, _ctx: &StreamContext) { + self.stream_success.fetch_add(1, Ordering::SeqCst); + } + + fn reset(&self) { + self.resets.fetch_add(1, Ordering::SeqCst); + } + } + + #[test] + fn host_dispatches_to_single_observer() { + let obs = Arc::new(CountingObserver::new("test")); + let mut host = ObserverHost::new(); + host.register(Arc::clone(&obs) as Arc); + host.on_stream_success(&StreamContext { + turn: 0, + model: "m".into(), + input_tokens: 0, + output_tokens: 0, + }); + assert_eq!(obs.stream_success.load(Ordering::SeqCst), 1); + } + + #[test] + fn host_dispatches_to_multiple_observers() { + let obs1 = Arc::new(CountingObserver::new("a")); + let obs2 = Arc::new(CountingObserver::new("b")); + let mut host = ObserverHost::new(); + host.register(Arc::clone(&obs1) as Arc); + host.register(Arc::clone(&obs2) as Arc); + host.on_stream_success(&StreamContext { + turn: 0, + model: "m".into(), + input_tokens: 0, + output_tokens: 0, + }); + assert_eq!(obs1.stream_success.load(Ordering::SeqCst), 1); + assert_eq!(obs2.stream_success.load(Ordering::SeqCst), 1); + } + + #[test] + fn host_len_and_is_empty() { + let mut host = ObserverHost::new(); + assert!(host.is_empty()); + assert_eq!(host.len(), 0); + host.register(Arc::new(CountingObserver::new("x")) as Arc); + assert!(!host.is_empty()); + assert_eq!(host.len(), 1); + } + + #[test] + fn host_reset_all() { + let obs = Arc::new(CountingObserver::new("p")); + let mut host = ObserverHost::new(); + host.register(Arc::clone(&obs) as Arc); + host.on_stream_success(&StreamContext { + turn: 0, + model: "m".into(), + input_tokens: 0, + output_tokens: 0, + }); + assert_eq!(obs.stream_success.load(Ordering::SeqCst), 1); + host.reset_all(); + assert_eq!(obs.resets.load(Ordering::SeqCst), 1); + } +} diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 4995337..9f30acf 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -16,7 +16,7 @@ //! calls the model requests. //! - An [`AgentConfig`] governing session parameters (max turns, system //! prompt, session ID). -//! - Optional [`AgentObserver`] implementations for lifecycle instrumentation. +//! - Optional [`LoopObserver`](crate::core::observer::LoopObserver) registrations for lifecycle instrumentation. //! //! ```text //! BareLoop @@ -71,17 +71,20 @@ use crate::api_client::ApiClient; use crate::cancel::CancelSignal; use crate::compact::{ContextManager, EnsureContextResult}; +use crate::core::observer::{ + ConvergenceDetectedContext, FallbackContext, LoopDetectedContext, ResponseContext, + StreamContext, StreamFailureContext, TurnEndContext, TurnStartContext, +}; use crate::core::reflection::{ ExponentialBackoffRecovery, NoopReflector, RecoveryAction, RecoveryStrategy, ReflectionContext, Reflector, }; -use crate::core::{AgentConfig, AgentError, AgentObserver, SessionResult, ToolDispatchResult}; +use crate::core::{AgentConfig, AgentError, SessionResult, ToolDispatchResult}; use crate::engine::middleware::{ToolDispatchContext, ToolPipeline, ToolPipelineBuilder}; #[cfg(feature = "hooks")] use crate::hooks::HookAction; #[cfg(feature = "hooks")] use crate::hooks::HookExecutor; -// Hook context types are used by submodules (compact, dispatch, emission) via `use super::*`. #[cfg(feature = "hooks")] #[allow(unused_imports)] use crate::hooks::context::{ @@ -90,8 +93,7 @@ use crate::hooks::context::{ }; use crate::loop_control::bundle::ManagerBundle; use crate::loop_control::detection::{ConvergenceAction, DetectedPattern}; -use crate::message::{Message, MessagePart, Role, ToolContent, ToolContentPart}; -use crate::observability::{EventSink, NullSink, ObserveEvent}; +use crate::message::{Message, MessagePart, Role, ToolContent}; use crate::stream::handler::StreamHandler; use crate::stream::{StreamAccumulator, StreamEvent, StreamStopReason, Usage}; #[cfg(feature = "tool_health")] @@ -130,8 +132,6 @@ mod stream; /// Use one of the constructors based on what components you have: /// /// - [`new()`](BareLoop::new) — client + tools + config. -/// - [`with_observers()`](BareLoop::with_observers) — adds lifecycle -/// observers for instrumentation. /// - [`with_managers()`](BareLoop::with_managers) — full control, /// including a [`ManagerBundle`]. /// - [`from_parts()`](BareLoop::from_parts) — re-assembles from the @@ -140,7 +140,7 @@ mod stream; /// # Lifecycle /// /// ```text -/// new() / with_observers() / from_parts() +/// new() / with_managers() / from_parts() /// → run(user_input) /// → stream_turn() → dispatch_tools() → stream_turn() /// → … (repeat until end_turn or max_turns) @@ -199,13 +199,7 @@ pub struct BareLoop { /// and tool-result messages until the model signals `end_turn`. conversation: Vec, - /// Lifecycle observers for session/turn/tool events. - /// - /// All observers are notified synchronously; long-running work should - /// be offloaded to a channel or thread pool inside the observer. - observers: Vec>, - - /// Manager bundle (fallback, loop detection, convergence). + /// Manager bundle (fallback, loop detection, convergence, observers). /// /// Holds the cross-cutting managers that govern session behaviour: /// @@ -219,11 +213,6 @@ pub struct BareLoop { /// Reset at the start of every session via [`ManagerBundle::reset_all`]. managers: ManagerBundle, - /// Structured event sink for observability. - /// - /// Emits [`ObserveEvent`] variants at each lifecycle point. - event_sink: Arc, - /// Failure analyser for tool errors. /// /// When a tool call returns an error, the reflector analyses the @@ -250,9 +239,9 @@ pub struct BareLoop { /// /// When `Some`, the loop checks token usage after each turn and /// triggers compaction when usage exceeds the configured threshold. - /// Compaction replaces the conversation messages and emits - /// [`on_compaction()`](AgentObserver::on_compaction) to observers - /// and [`ObserveEvent::ContextCompacted`] to the event sink. + /// Compaction replaces the conversation messages, notifies observers + /// via [`LoopObserver::on_compaction`](crate::core::observer::LoopObserver::on_compaction), + /// and notifies observers via [`on_compaction`](crate::core::observer::LoopObserver::on_compaction). context_manager: Option>, /// Optional stream handler for resilient streaming. @@ -323,8 +312,7 @@ impl SessionBudget { /// Token counts for a single turn, captured before tool dispatch. /// /// Needed because [`SessionBudget::accumulate_usage`] mutates the running -/// totals, but the per-turn values must be reported separately in -/// [`emit_turn_complete`](BareLoop::emit_turn_complete). +/// totals, but the per-turn values must be reported separately to observers. #[derive(Clone, Copy)] struct TurnTokens { input: u64, @@ -340,7 +328,7 @@ impl TurnTokens { /// /// This is needed because [`SessionBudget::accumulate_usage`] mutates /// running totals in place, but the per-turn values must be reported - /// separately in [`emit_turn_complete`](BareLoop::emit_turn_complete). + /// separately to observers. fn from_usage(usage: Option<&Usage>) -> Self { match usage { Some(u) => Self { @@ -382,7 +370,6 @@ enum AbortReason { /// Captures completion status, an [`EndReason`] discriminant, turn/token /// counters, and wall-clock duration — everything a hook needs to log or /// react to session termination without pulling data from other sources. -#[allow(dead_code)] struct SessionEndInfo { /// Whether the session completed normally. success: bool, @@ -419,7 +406,7 @@ enum EndReason { /// /// This type is private to the module because external consumers /// interact with tool results via [`SessionResult`] or the -/// [`AgentObserver`] callbacks. +/// [`LoopObserver`](crate::core::observer::LoopObserver) callbacks. /// /// # Fields /// @@ -459,7 +446,7 @@ impl BareLoop { /// Create a new `BareLoop` with the given components. /// - /// Initializes an empty conversation history, no observers, and a + /// Initializes an empty conversation history and a /// fresh [`ManagerBundle`]. The cancellation signal starts as non-cancelled. /// /// # Parameters @@ -484,59 +471,7 @@ impl BareLoop { pipeline: None, config, conversation: Vec::new(), - observers: Vec::new(), - managers: ManagerBundle::new(), - event_sink: Arc::new(NullSink), - reflector: Arc::new(NoopReflector), - recovery: Arc::new(ExponentialBackoffRecovery::new(3)), - cancelled: Arc::new(CancelSignal::new()), - context_manager: None, - stream_handler: None, - #[cfg(feature = "hooks")] - hook_executor: None, - #[cfg(feature = "tool_health")] - health_registry: None, - } - } - - /// Create a new `BareLoop` with lifecycle observers. - /// - /// Identical to [`new()`](BareLoop::new) but accepts a `Vec` of - /// [`AgentObserver`] implementations. Observers receive callbacks for - /// session start/end, turn start/end, and tool call/complete events. - /// - /// # Parameters - /// - /// - `client` — The LLM API client, wrapped in `Arc`. - /// - `tools` — The [`ToolRegistry`] containing available tools. - /// - `config` — Session parameters. - /// - `observers` — Lifecycle observers for instrumentation. - /// - /// # Example - /// - /// ```rust,ignore - /// let agent = BareLoop::with_observers( - /// Arc::new(my_client), - /// registry, - /// config, - /// vec![Arc::new(LoggingObserver)], - /// ); - /// ``` - pub fn with_observers( - client: Arc, - tools: ToolRegistry, - config: AgentConfig, - observers: Vec>, - ) -> Self { - Self { - client, - tools: Arc::new(tools), - pipeline: None, - config, - conversation: Vec::new(), - observers, managers: ManagerBundle::new(), - event_sink: Arc::new(NullSink), reflector: Arc::new(NoopReflector), recovery: Arc::new(ExponentialBackoffRecovery::new(3)), cancelled: Arc::new(CancelSignal::new()), @@ -560,7 +495,6 @@ impl BareLoop { /// - `client` — The LLM API client, wrapped in `Arc`. /// - `tools` — The [`ToolRegistry`] containing available tools. /// - `config` — Session parameters. - /// - `observers` — Lifecycle observers. /// - `managers` — A pre-built [`ManagerBundle`]. /// /// # Example @@ -574,7 +508,6 @@ impl BareLoop { /// Arc::new(my_client), /// registry, /// config, - /// observers, /// managers, /// ); /// ``` @@ -582,7 +515,6 @@ impl BareLoop { client: Arc, tools: ToolRegistry, config: AgentConfig, - observers: Vec>, managers: ManagerBundle, ) -> Self { Self { @@ -591,9 +523,7 @@ impl BareLoop { pipeline: None, config, conversation: Vec::new(), - observers, managers, - event_sink: Arc::new(NullSink), reflector: Arc::new(NoopReflector), recovery: Arc::new(ExponentialBackoffRecovery::new(3)), cancelled: Arc::new(CancelSignal::new()), @@ -617,20 +547,18 @@ impl BareLoop { /// - `client` — The LLM API client, wrapped in `Arc`. /// - `tools` — The [`ToolRegistry`]. /// - `managers` — A [`ManagerBundle`]. - /// - `observers` — Lifecycle observers. /// - `config` — Session parameters. /// /// # Example /// /// ```rust,ignore - /// let (client, tools, managers, observers, config) = builder.into_raw_parts(); - /// let agent = BareLoop::from_parts(client, tools, managers, observers, config); + /// let (client, tools, managers, config) = builder.into_raw_parts(); + /// let agent = BareLoop::from_parts(client, tools, managers, config); /// ``` pub fn from_parts( client: Arc, tools: ToolRegistry, managers: ManagerBundle, - observers: Vec>, config: AgentConfig, ) -> Self { Self { @@ -639,9 +567,7 @@ impl BareLoop { pipeline: None, config, conversation: Vec::new(), - observers, managers, - event_sink: Arc::new(NullSink), reflector: Arc::new(NoopReflector), recovery: Arc::new(ExponentialBackoffRecovery::new(3)), cancelled: Arc::new(CancelSignal::new()), @@ -733,21 +659,6 @@ impl BareLoop { // Dependency setters // ================================================== - /// Set the [`EventSink`] for structured observability events. - /// - /// Replaces the default [`NullSink`] with a caller-supplied - /// implementation. Must be called before [`run()`](BareLoop::run). - /// - /// # Example - /// - /// ```rust,ignore - /// let mut agent = BareLoop::new(client, registry, config); - /// agent.set_event_sink(Arc::new(MySink)); - /// ``` - pub fn set_event_sink(&mut self, sink: Arc) { - self.event_sink = sink; - } - /// Set the [`Reflector`] for tool-error analysis. /// /// Replaces the default [`NoopReflector`] with a caller-supplied @@ -919,6 +830,27 @@ impl BareLoop { Ok(()) } + /// Register a [`LoopObserver`](crate::core::observer::LoopObserver) with the manager bundle's observer host. + /// + /// Plugins are called at lifecycle hook points inside the agent loop, + /// in registration order. See [`LoopObserver`](crate::core::observer::LoopObserver) + /// for the trait definition and available hooks. + /// + /// Must be called before [`run()`](BareLoop::run). + /// + /// # Example + /// + /// ```rust,ignore + /// use loopctl::core::observer::LoopObserver; + /// use std::sync::Arc; + /// + /// let mut agent = BareLoop::new(client, registry, config); + /// agent.register_observer(Arc::new(MyObserver)); + /// ``` + pub fn register_observer(&mut self, observer: Arc) { + self.managers.register_observer(observer); + } + // ================================================== // Main run loop // ================================================== @@ -986,7 +918,6 @@ impl BareLoop { self.managers.reset_all(); self.conversation.push(Message::user(user_input)); - // Main agent loop loop { if self.is_cancelled() { return self.abort_session(&budget, start.elapsed(), AbortReason::Cancelled); @@ -995,36 +926,45 @@ impl BareLoop { return self.abort_session(&budget, start.elapsed(), AbortReason::MaxTurnsExceeded); } - self.emit_turn_start(budget.turn_count, user_input); - self.notify_turn_start(user_input); + self.managers.observers().on_turn_start(&TurnStartContext { + turn: budget.turn_count, + query: user_input.to_string(), + }); let turn_start = Instant::now(); match self.stream_turn().await { Ok((assistant_msg, usage, stop_reason)) => { self.managers.fallback.record_model_success(); + self.managers.observers().on_stream_success(&StreamContext { + turn: budget.turn_count, + model: self.client.model().to_string(), + input_tokens: TurnTokens::from_usage(usage.as_ref()).input, + output_tokens: TurnTokens::from_usage(usage.as_ref()).output, + }); + budget.accumulate_usage(usage.as_ref()); let text = Self::extract_text(&assistant_msg); - - // Convergence detection let pattern = self.managers.detection.record_response(&text); + + self.managers.observers().on_response(&ResponseContext { + turn: budget.turn_count, + text: text.clone(), + usage, + }); + if let Some(result) = self.handle_detected_pattern(&pattern, budget.turn_count) { let turn_elapsed = turn_start.elapsed(); return match result { Ok(session_result) => { - self.emit_turn_complete( - budget.turn_count, - turn_elapsed, - TurnTokens::from_usage(usage.as_ref()).input, - TurnTokens::from_usage(usage.as_ref()).output, - ); - self.notify_turn_end(true, None); - self.emit_session_stop( - budget.turn_count, - start.elapsed(), - true, - "detection_early_exit", - ); + self.managers.observers().on_turn_end(&TurnEndContext { + turn: budget.turn_count, + success: true, + error: None, + duration_ms: Self::millis_u64(turn_start.elapsed()), + input_tokens: budget.input_tokens, + output_tokens: budget.output_tokens, + }); self.notify_session_end(&SessionEndInfo { success: true, reason: EndReason::Complete, @@ -1074,7 +1014,6 @@ impl BareLoop { return self.abort_session_from_error(e, start.elapsed(), &budget); } - // After tool dispatch, check if context compaction is needed. if let Err(e) = self.maybe_compact_context(budget.turn_count).await { return self.abort_turn_and_session( &budget, @@ -1091,11 +1030,21 @@ impl BareLoop { let from = self.client.model(); if let Some(to) = self.managers.fallback.fallback_model() { tracing::warn!(from, to, "fallback manager tripped"); - for obs in &self.observers { - obs.on_fallback(from, &to); - } + self.managers.observers().on_fallback(&FallbackContext { + from: from.to_string(), + to, + }); } } + + self.managers + .observers() + .on_stream_failure(&StreamFailureContext { + turn: budget.turn_count, + model: self.client.model().to_string(), + error: e.clone(), + }); + let err_str = e.to_string(); return self.abort_turn_and_session( &budget, @@ -1155,13 +1104,12 @@ impl BareLoop { "loop detected" ); - for obs in &self.observers { - obs.on_loop_detected(pattern_description, *repetitions); - } - self.event_sink.on_event(&ObserveEvent::LoopDetected { - tool: pattern_description.clone(), - repetitions: *repetitions, - }); + self.managers + .observers() + .on_loop_detected(&LoopDetectedContext { + pattern: pattern_description.clone(), + repetitions: *repetitions, + }); if *repetitions >= self.managers.detection.config().stop_threshold { tracing::error!( @@ -1192,11 +1140,9 @@ impl BareLoop { ConvergenceAction::SwitchPhase => "switch_phase", }; - for obs in &self.observers { - obs.on_convergence_detected(action_str); - } - self.event_sink - .on_event(&ObserveEvent::ConvergenceDetected { + self.managers + .observers() + .on_convergence_detected(&ConvergenceDetectedContext { action: action_str.to_string(), }); @@ -1227,35 +1173,9 @@ impl BareLoop { } } - /// Hash tool input for loop-detection deduplication. - /// - /// Produces a deterministic `u64` hash from a [`serde_json::Value`] by - /// serialising it to a canonical JSON string, then feeding the bytes - /// through [`DefaultHasher`]. The hash is passed to - /// [`DetectionManager::record_tool_call`] so that identical tool inputs - /// can be detected without storing the full JSON payload. - /// - /// # Determinism - /// - /// `serde_json::to_string` produces a stable ordering for map keys, so - /// `{"a":1,"b":2}` and `{"b":2,"a":1}` hash identically. However, - /// floating-point values are **not** normalised — `1.0` and `1` produce - /// different hashes. - /// - /// [`DetectionManager::record_tool_call`]: crate::loop_control::detection::DetectionManager::record_tool_call - fn hash_tool_input(input: &serde_json::Value) -> u64 { - use std::collections::hash_map::DefaultHasher; - use std::hash::{Hash, Hasher}; - - let serialized = serde_json::to_string(input).unwrap_or_default(); - let mut hasher = DefaultHasher::new(); - serialized.hash(&mut hasher); - hasher.finish() - } - /// Dispatch tool calls, push the result message, and record the count. /// - /// Emits turn-complete on success, turn-failed on error. + /// Notifies observers: turn-end on success, turn-end on error. /// /// # Errors /// @@ -1273,19 +1193,26 @@ impl BareLoop { budget.total_tool_calls = budget.total_tool_calls.saturating_add(results.len()); let tool_result_msg = Self::build_tool_result_message(results); self.conversation.push(tool_result_msg); - self.emit_turn_complete( - turn.idx, - turn.duration, - turn.tokens.input, - turn.tokens.output, - ); - self.notify_turn_end(true, None); + self.managers.observers().on_turn_end(&TurnEndContext { + turn: budget.turn_count, + success: true, + error: None, + duration_ms: Self::millis_u64(turn.duration), + input_tokens: turn.tokens.input, + output_tokens: turn.tokens.output, + }); Ok(()) } Err(e) => { let err_str = e.to_string(); - self.emit_turn_failed(turn.idx, turn.duration, &err_str); - self.notify_turn_end(false, Some(&err_str)); + self.managers.observers().on_turn_end(&TurnEndContext { + turn: budget.turn_count, + success: false, + error: Some(err_str), + duration_ms: Self::millis_u64(turn.duration), + input_tokens: turn.tokens.input, + output_tokens: turn.tokens.output, + }); Err(e) } } @@ -1293,8 +1220,8 @@ impl BareLoop { /// Build the final [`SessionResult`] when the model ends its turn. /// - /// Called when streaming completes with no tool calls. Emits - /// turn-complete/failed and session-stop events, notifies observers, + /// Called when streaming completes with no tool calls. Notifies + /// turn-end and session-end events, notifies observers, /// and returns the assembled result. fn finalise_session( &self, @@ -1312,28 +1239,15 @@ impl BareLoop { Some(format!("Stream stopped with reason: {stop_reason:?}")) }; - if success { - self.emit_turn_complete( - turn.idx, - turn.duration, - turn.tokens.input, - turn.tokens.output, - ); - } else { - self.emit_turn_failed( - turn.idx, - turn.duration, - error.as_deref().unwrap_or("unknown"), - ); - } - self.notify_turn_end(success, error.as_deref()); - - self.emit_session_stop( - budget.turn_count, - session_duration, + self.managers.observers().on_turn_end(&TurnEndContext { + turn: turn.idx, success, - error.as_deref().unwrap_or("completed"), - ); + error: error.as_deref().map(std::string::ToString::to_string), + duration_ms: Self::millis_u64(turn.duration), + input_tokens: turn.tokens.input, + output_tokens: turn.tokens.output, + }); + let end_reason = if success { EndReason::Complete } else { @@ -1360,7 +1274,7 @@ impl BareLoop { } } - /// Abort the session with an error — emits turn-failed + session-stop. + /// Abort the session with an error — notifies turn-end + session-end. /// /// Used when the streaming call itself fails (API error, timeout, etc.). /// @@ -1375,9 +1289,14 @@ impl BareLoop { reason: &str, error: AgentError, ) -> Result { - self.emit_turn_failed(budget.turn_count, turn_duration, reason); - self.notify_turn_end(false, Some(reason)); - self.emit_session_stop(budget.turn_count, session_duration, false, reason); + self.managers.observers().on_turn_end(&TurnEndContext { + turn: budget.turn_count, + success: false, + error: Some(reason.to_string()), + duration_ms: Self::millis_u64(turn_duration), + input_tokens: budget.input_tokens, + output_tokens: budget.output_tokens, + }); let end_reason = if matches!(error, AgentError::Cancelled) { EndReason::Cancelled } else { @@ -1396,7 +1315,7 @@ impl BareLoop { /// Abort the session after a tool-dispatch error. /// /// Handles both [`AgentError::Cancelled`] and other errors uniformly. - /// Turn-level events were already emitted inside [`dispatch_and_record`]. + /// Turn-level notifications were already sent inside [`dispatch_and_record`]. /// /// # Errors /// @@ -1407,8 +1326,6 @@ impl BareLoop { session_duration: Duration, budget: &SessionBudget, ) -> Result { - let reason = error.to_string(); - self.emit_session_stop(budget.turn_count, session_duration, false, &reason); let end_reason = if matches!(error, AgentError::Cancelled) { EndReason::Cancelled } else { @@ -1426,7 +1343,7 @@ impl BareLoop { /// Abort the session with a known reason string (cancel / max-turns). /// - /// Does not emit turn-level events since no turn was started. + /// Does not send turn-level notifications since no turn was started. /// /// # Errors /// @@ -1438,11 +1355,6 @@ impl BareLoop { session_duration: Duration, reason: AbortReason, ) -> Result { - let reason_str = match &reason { - AbortReason::Cancelled => "Cancelled", - AbortReason::MaxTurnsExceeded => "Max turns exceeded", - }; - self.emit_session_stop(budget.turn_count, session_duration, false, reason_str); let end_reason = match &reason { AbortReason::Cancelled => EndReason::Cancelled, AbortReason::MaxTurnsExceeded => EndReason::MaxTurns, @@ -1864,29 +1776,15 @@ mod tests { } // ================================================== - // Counting Observer + // Counting Plugin (test helper) // ================================================== - /// An [`AgentObserver`] that counts how many times each callback fires. + /// A [`LoopObserver`](crate::core::observer::LoopObserver) that counts + /// how many times each hook fires. /// /// Uses [`AtomicUsize`] counters with `SeqCst` ordering so that /// test assertions can read the counts from any thread after the /// agent loop completes. - /// - /// # Counters - /// - /// - [`session_starts`](CountingObserver::session_starts) — incremented - /// by [`on_session_start`](AgentObserver::on_session_start). - /// - [`session_ends`](CountingObserver::session_ends) — incremented - /// by [`on_session_end`](AgentObserver::on_session_end). - /// - [`turn_starts`](CountingObserver::turn_starts) — incremented - /// by [`on_turn_start`](AgentObserver::on_turn_start). - /// - [`turn_ends`](CountingObserver::turn_ends) — incremented - /// by [`on_turn_end`](AgentObserver::on_turn_end). - /// - [`tool_calls`](CountingObserver::tool_calls) — incremented - /// by [`on_tool_call`](AgentObserver::on_tool_call). - /// - [`tool_completes`](CountingObserver::tool_completes) — incremented - /// by [`on_tool_complete`](AgentObserver::on_tool_complete). struct CountingObserver { /// Number of times `on_session_start` was called. session_starts: AtomicUsize, @@ -1896,10 +1794,10 @@ mod tests { turn_starts: AtomicUsize, /// Number of times `on_turn_end` was called. turn_ends: AtomicUsize, - /// Number of times `on_tool_call` was called. - tool_calls: AtomicUsize, - /// Number of times `on_tool_complete` was called. - tool_completes: AtomicUsize, + /// Number of times `on_tool_pre` was called. + tool_pres: AtomicUsize, + /// Number of times `on_tool_post` was called. + tool_posts: AtomicUsize, } impl CountingObserver { @@ -1910,49 +1808,39 @@ mod tests { session_ends: AtomicUsize::new(0), turn_starts: AtomicUsize::new(0), turn_ends: AtomicUsize::new(0), - tool_calls: AtomicUsize::new(0), - tool_completes: AtomicUsize::new(0), + tool_pres: AtomicUsize::new(0), + tool_posts: AtomicUsize::new(0), } } } - impl AgentObserver for CountingObserver { - /// Increment the session-start counter. - fn on_session_start(&self, _session_id: uuid::Uuid) { + impl crate::core::observer::LoopObserver for CountingObserver { + fn name(&self) -> &str { + "counting" + } + + fn on_session_start(&self, _ctx: &crate::core::observer::SessionStartContext) { self.session_starts.fetch_add(1, Ordering::SeqCst); } - /// Increment the session-end counter. - fn on_session_end(&self, _success: bool, _error: Option<&str>) { + fn on_session_end(&self, _ctx: &crate::core::observer::SessionEndContext) { self.session_ends.fetch_add(1, Ordering::SeqCst); } - /// Increment the turn-start counter. - fn on_turn_start(&self, _query: &str) { + fn on_turn_start(&self, _ctx: &crate::core::observer::TurnStartContext) { self.turn_starts.fetch_add(1, Ordering::SeqCst); } - /// Increment the turn-end counter. - fn on_turn_end(&self, _success: bool, _error: Option<&str>) { + fn on_turn_end(&self, _ctx: &crate::core::observer::TurnEndContext) { self.turn_ends.fetch_add(1, Ordering::SeqCst); } - /// Increment the tool-call counter. - fn on_tool_call(&self, _tool: &str, _input: &str) { - self.tool_calls.fetch_add(1, Ordering::SeqCst); + fn on_tool_pre(&self, _ctx: &crate::core::observer::ToolPreContext) { + self.tool_pres.fetch_add(1, Ordering::SeqCst); } - /// Increment the tool-complete counter. - fn on_tool_complete( - &self, - _tool: &str, - _input: &str, - _output: &str, - _duration: Duration, - _success: bool, - _error: Option<&str>, - ) { - self.tool_completes.fetch_add(1, Ordering::SeqCst); + fn on_tool_post(&self, _ctx: &crate::core::observer::ToolPostContext) { + self.tool_posts.fetch_add(1, Ordering::SeqCst); } } @@ -2135,50 +2023,46 @@ mod tests { let client = MockClient::new("test-model"); client.add_text_response("Done!"); - let observer = Arc::new(CountingObserver::new()); + let plugin = Arc::new(CountingObserver::new()); let config = make_config(); - let agent = BareLoop::with_observers( - Arc::new(client), - ToolRegistry::new(), - config, - vec![observer.clone()], - ); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + agent.register_observer(plugin.clone()); let result = agent.run("Hi").await.unwrap(); assert!(result.success); - assert_eq!(observer.session_starts.load(Ordering::SeqCst), 1); - assert_eq!(observer.session_ends.load(Ordering::SeqCst), 1); - assert_eq!(observer.turn_starts.load(Ordering::SeqCst), 1); - assert_eq!(observer.turn_ends.load(Ordering::SeqCst), 1); + assert_eq!(plugin.session_starts.load(Ordering::SeqCst), 1); + assert_eq!(plugin.session_ends.load(Ordering::SeqCst), 1); + assert_eq!(plugin.turn_starts.load(Ordering::SeqCst), 1); + assert_eq!(plugin.turn_ends.load(Ordering::SeqCst), 1); } - /// Verify that a tool-using session fires `tool_call` and - /// `tool_complete` callbacks in addition to the turn callbacks. + /// Verify that a tool-using session fires `on_tool_pre` and + /// `on_tool_post` observer hooks in addition to the turn hooks. /// /// A two-turn session (tool_call + end_turn) should produce: /// - 2 turn starts, 2 turn ends - /// - 1 tool call, 1 tool complete + /// - 1 tool pre, 1 tool post #[tokio::test] async fn test_observer_tool_events() { let client = MockClient::new("test-model"); client.add_tool_then_text("tool_1", "echo", json!({"message": "test"}), "All done!"); - let observer = Arc::new(CountingObserver::new()); + let plugin = Arc::new(CountingObserver::new()); let mut registry = ToolRegistry::new(); registry.register(EchoTool); let config = make_config(); - let agent = - BareLoop::with_observers(Arc::new(client), registry, config, vec![observer.clone()]); + let mut agent = BareLoop::new(Arc::new(client), registry, config); + agent.register_observer(plugin.clone()); let result = agent.run("Echo test").await.unwrap(); assert!(result.success); - assert_eq!(observer.tool_calls.load(Ordering::SeqCst), 1); - assert_eq!(observer.tool_completes.load(Ordering::SeqCst), 1); - assert_eq!(observer.turn_starts.load(Ordering::SeqCst), 2); - assert_eq!(observer.turn_ends.load(Ordering::SeqCst), 2); + assert_eq!(plugin.tool_pres.load(Ordering::SeqCst), 1); + assert_eq!(plugin.tool_posts.load(Ordering::SeqCst), 1); + assert_eq!(plugin.turn_starts.load(Ordering::SeqCst), 2); + assert_eq!(plugin.turn_ends.load(Ordering::SeqCst), 2); } // ================================================== @@ -2367,14 +2251,7 @@ mod tests { let client = MockClient::new("test-model"); let config = make_config(); let managers = ManagerBundle::new(); - let observers: Vec> = vec![]; - let agent = BareLoop::from_parts( - Arc::new(client), - ToolRegistry::new(), - managers, - observers, - config, - ); + let agent = BareLoop::from_parts(Arc::new(client), ToolRegistry::new(), managers, config); assert!(agent.conversation().is_empty()); } @@ -2485,307 +2362,6 @@ mod tests { } // ================================================== - // EventSink + Recovery wiring tests - // ================================================== - - /// A recording [`EventSink`] that captures all emitted events. - /// - /// Uses `Mutex>` so it's `Send + Sync`. - struct RecordingSink { - events: std::sync::Mutex>, - } - - impl RecordingSink { - fn new() -> Self { - Self { - events: std::sync::Mutex::new(Vec::new()), - } - } - - /// Return a snapshot of all captured events. - fn events(&self) -> Vec { - self.events.lock().expect("lock").clone() - } - - fn count_matching(&self, pred: impl Fn(&ObserveEvent) -> bool) -> usize { - self.events - .lock() - .expect("lock") - .iter() - .filter(|e| pred(e)) - .count() - } - - /// Return true if any event matches the predicate. - fn any(&self, pred: impl Fn(&ObserveEvent) -> bool) -> bool { - self.events.lock().expect("lock").iter().any(|e| pred(e)) - } - - /// Return the first event matching the predicate, if any. - fn find(&self, pred: impl Fn(&ObserveEvent) -> bool) -> Option { - self.events - .lock() - .expect("lock") - .iter() - .find(|e| pred(e)) - .cloned() - } - } - - impl EventSink for RecordingSink { - fn on_event(&self, event: &ObserveEvent) { - self.events.lock().expect("lock").push(event.clone()); - } - } - - /// Helpers for matching [`ObserveEvent`] variants in assertions. - mod event_match { - use crate::observability::ObserveEvent; - - pub fn is_session_start(e: &ObserveEvent) -> bool { - matches!(e, ObserveEvent::SessionStart { .. }) - } - pub fn is_session_stop(e: &ObserveEvent) -> bool { - matches!(e, ObserveEvent::SessionStop { .. }) - } - pub fn is_turn_start(e: &ObserveEvent) -> bool { - matches!(e, ObserveEvent::TurnStart { .. }) - } - pub fn is_turn_complete(e: &ObserveEvent) -> bool { - matches!(e, ObserveEvent::TurnComplete { .. }) - } - pub fn is_tool_start(e: &ObserveEvent) -> bool { - matches!(e, ObserveEvent::ToolStart { .. }) - } - pub fn is_tool_complete(e: &ObserveEvent) -> bool { - matches!(e, ObserveEvent::ToolComplete { .. }) - } - } - - /// Build a [`BareLoop`] with a [`RecordingSink`] wired in. - /// - /// Returns the loop and an `Arc` for asserting events. - fn agent_with_recording_sink( - client: MockClient, - registry: ToolRegistry, - config: AgentConfig, - ) -> (BareLoop, Arc) { - let sink = Arc::new(RecordingSink::new()); - let mut agent = BareLoop::new(Arc::new(client), registry, config); - agent.event_sink = Arc::clone(&sink) as Arc; - (agent, sink) - } - - // ================================================== - // EventSink emission tests - // ================================================== - - /// Verify that a successful single-turn session emits - /// [`SessionStart`](ObserveEvent::SessionStart) and - /// [`SessionStop`](ObserveEvent::SessionStop) events. - #[tokio::test] - async fn test_sink_emits_session_start_stop_on_success() { - let client = MockClient::new("test"); - client.add_text_response("Hello!"); - - let (agent, sink) = agent_with_recording_sink(client, ToolRegistry::new(), make_config()); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); - assert!( - sink.any(event_match::is_session_start), - "missing session_start" - ); - assert!( - sink.any(event_match::is_session_stop), - "missing session_stop" - ); - assert!(sink.any(event_match::is_turn_start), "missing turn_start"); - assert!( - sink.any(event_match::is_turn_complete), - "missing turn_complete" - ); - } - - /// Verify that exceeding `max_turns` emits a - /// [`SessionStop`](ObserveEvent::SessionStop) with `success = false`. - #[tokio::test] - async fn test_sink_emits_session_stop_on_max_turns_error() { - let client = MockClient::new("test"); - // Queue a response so the mock has something to return, but - // max_turns=0 means the loop aborts before streaming. - client.add_text_response("turn 1"); - - let mut config = make_config(); - config.max_turns = 0; - - let (agent, sink) = agent_with_recording_sink(client, ToolRegistry::new(), config); - - // max_turns=0 → immediate abort, no turn executes. - let err = agent.run("Hi").await.unwrap_err(); - assert!( - matches!(err, AgentError::MaxTurnsExceeded { .. }), - "expected MaxTurnsExceeded, got {err:?}" - ); - assert!( - sink.any(event_match::is_session_start), - "missing session_start" - ); - let stop_evt = sink - .find(event_match::is_session_stop) - .expect("missing session_stop"); - assert!( - matches!(stop_evt, ObserveEvent::SessionStop { success: false, .. }), - "expected SessionStop with success=false, got {stop_evt:?}" - ); - } - - /// Verify that a tool-using session emits - /// [`ToolStart`](ObserveEvent::ToolStart) and - /// [`ToolComplete`](ObserveEvent::ToolComplete) events. - #[tokio::test] - async fn test_sink_emits_tool_events_on_tool_call() { - let mut registry = ToolRegistry::new(); - registry.register(EchoTool); - - let client = MockClient::new("test"); - client.add_tool_then_text("tool_1", "echo", json!({"message": "hi"}), "Done"); - - let (agent, sink) = agent_with_recording_sink(client, registry, make_config()); - let result = agent.run("Test").await.unwrap(); - assert!(result.success); - assert_eq!(result.tool_calls, 1); - assert!(sink.any(event_match::is_tool_start), "missing tool_start"); - assert!( - sink.any(event_match::is_tool_complete), - "missing tool_complete" - ); - } - - /// Verify that dispatching a missing tool emits - /// [`ToolComplete`](ObserveEvent::ToolComplete) with `is_error = true`. - #[tokio::test] - async fn test_sink_emits_tool_complete_with_error_on_missing_tool() { - let client = MockClient::new("test"); - client.add_tool_then_text("tool_1", "missing_tool", json!({}), "OK"); - - let (agent, sink) = agent_with_recording_sink(client, ToolRegistry::new(), make_config()); - let result = agent.run("Test").await.unwrap(); - - assert!(result.success); - assert!(sink.any(event_match::is_tool_start), "missing tool_start"); - assert!( - sink.any(event_match::is_tool_complete), - "missing tool_complete" - ); - - // The tool_complete should have is_error=true - let error_completes = sink.events().iter().any(|e| { - if let ObserveEvent::ToolComplete { is_error, .. } = e { - *is_error - } else { - false - } - }); - assert!( - error_completes, - "expected at least one tool_complete with is_error=true" - ); - } - - /// Verify that a single-turn session emits the full event sequence: - /// `SessionStart → TurnStart → TurnComplete → SessionStop`. - #[tokio::test] - async fn test_event_sequence_single_turn() { - let client = MockClient::new("test"); - client.add_text_response("Hello!"); - - let (agent, sink) = agent_with_recording_sink(client, ToolRegistry::new(), make_config()); - let result = agent.run("Hi").await.unwrap(); - assert!(result.success); - - let events = sink.events(); - // Sequence: SessionStart, TurnStart, TurnComplete, SessionStop - assert!( - event_match::is_session_start(&events[0]), - "expected SessionStart, got {:?}", - events[0] - ); - assert!( - event_match::is_turn_start(&events[1]), - "expected TurnStart, got {:?}", - events[1] - ); - assert!( - event_match::is_turn_complete(&events[2]), - "expected TurnComplete, got {:?}", - events[2] - ); - assert!( - event_match::is_session_stop(&events[3]), - "expected SessionStop, got {:?}", - events[3] - ); - assert_eq!(events.len(), 4, "expected exactly 4 events, got {events:?}"); - } - - /// Verify that a tool-using session emits the full event sequence: - /// `SessionStart → TurnStart → ToolStart → ToolComplete → TurnComplete - /// → TurnStart → TurnComplete → SessionStop`. - #[tokio::test] - async fn test_event_sequence_with_tool() { - let mut registry = ToolRegistry::new(); - registry.register(EchoTool); - - let client = MockClient::new("test"); - client.add_tool_then_text("tool_1", "echo", json!({"message": "x"}), "Done"); - - let (agent, sink) = agent_with_recording_sink(client, registry, make_config()); - let result = agent.run("Test").await.unwrap(); - assert!(result.success); - - let events = sink.events(); - // Sequence: SessionStart, TurnStart, ToolStart, ToolComplete, - // TurnComplete, TurnStart, TurnComplete, SessionStop - assert!( - event_match::is_session_start(&events[0]), - "event[0] not SessionStart" - ); - assert!( - event_match::is_turn_start(&events[1]), - "event[1] not TurnStart" - ); - assert!( - event_match::is_tool_start(&events[2]), - "event[2] not ToolStart" - ); - assert!( - event_match::is_tool_complete(&events[3]), - "event[3] not ToolComplete" - ); - assert!( - event_match::is_turn_complete(&events[4]), - "event[4] not TurnComplete" - ); - assert!( - event_match::is_turn_start(&events[5]), - "event[5] not TurnStart" - ); - assert!( - event_match::is_turn_complete(&events[6]), - "event[6] not TurnComplete" - ); - assert!( - event_match::is_session_stop(&events[7]), - "event[7] not SessionStop" - ); - assert_eq!( - events.len(), - 8, - "expected exactly 8 events, got {}", - events.len() - ); - } - // ================================================== // Recovery wiring tests // ================================================== @@ -2800,16 +2376,11 @@ mod tests { let client = MockClient::new("test"); client.add_tool_then_text("tool_1", "fail", json!({}), "Moving on"); - let (agent, sink) = agent_with_recording_sink(client, registry, make_config()); + let agent = BareLoop::new(Arc::new(client), registry, make_config()); let result = agent.run("Test").await.unwrap(); assert!(result.success); assert_eq!(result.tool_calls, 1); - assert!(sink.any(event_match::is_tool_start), "missing tool_start"); - assert!( - sink.any(event_match::is_tool_complete), - "missing tool_complete" - ); } /// Verify that when a tool is not found, the recovery wiring still @@ -2819,22 +2390,16 @@ mod tests { let client = MockClient::new("test"); client.add_tool_then_text("tool_1", "nonexistent", json!({}), "OK"); - let (agent, sink) = agent_with_recording_sink(client, ToolRegistry::new(), make_config()); + let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); let result = agent.run("Test").await.unwrap(); assert!(result.success); assert_eq!(result.tool_calls, 1); - // Should still emit tool_start and tool_complete (even for missing tools) - assert!(sink.any(event_match::is_tool_start), "missing tool_start"); - assert!( - sink.any(event_match::is_tool_complete), - "missing tool_complete" - ); } /// Verify that a failing tool with the default recovery produces - /// exactly one tool_start and one tool_complete event (NoopReflector - /// marks everything as non-recoverable, so no retries). + /// exactly one tool dispatch (NoopReflector marks everything as + /// non-recoverable, so no retries). #[tokio::test] async fn test_recovery_noop_reflector_no_retries() { let mut registry = ToolRegistry::new(); @@ -2843,21 +2408,11 @@ mod tests { let client = MockClient::new("test"); client.add_tool_then_text("tool_1", "fail", json!({}), "OK"); - let (agent, sink) = agent_with_recording_sink(client, registry, make_config()); + let agent = BareLoop::new(Arc::new(client), registry, make_config()); let result = agent.run("Test").await.unwrap(); assert!(result.success); - // NoopReflector marks everything non-recoverable → Fail → no retry - assert_eq!( - sink.count_matching(event_match::is_tool_start), - 1, - "expected exactly 1 tool_start (no retries)" - ); - assert_eq!( - sink.count_matching(event_match::is_tool_complete), - 1, - "expected exactly 1 tool_complete (no retries)" - ); + assert_eq!(result.tool_calls, 1); } /// Verify cancellation is still respected during tool recovery. @@ -2869,9 +2424,7 @@ mod tests { let client = MockClient::new("test"); client.add_tool_only_response("tc-1", "fail", json!({})); - let sink = Arc::new(RecordingSink::new()); - let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); - agent.event_sink = Arc::clone(&sink) as Arc; + let agent = BareLoop::new(Arc::new(client), registry, make_config()); // Cancel before running agent.cancel(); diff --git a/src/engine/bare/compact.rs b/src/engine/bare/compact.rs index 510eb36..bbf241f 100644 --- a/src/engine/bare/compact.rs +++ b/src/engine/bare/compact.rs @@ -4,18 +4,19 @@ //! When a [`ContextManager`] is configured, checks token usage after each //! tool dispatch and triggers compaction if usage exceeds the threshold. -use super::{AgentError, ApiClient, BareLoop, EnsureContextResult, Instant, ObserveEvent}; +use super::{AgentError, ApiClient, BareLoop, EnsureContextResult, Instant}; #[cfg(feature = "hooks")] use super::{CompactTrigger, PostCompactContext, PreCompactContext}; +use crate::core::observer::CompactionContext; + impl BareLoop { /// Check if context compaction is needed and perform it if so. /// /// When a [`ContextManager`] is configured, this method: /// 1. Calls [`ContextManager::ensure_context_fits()`] to check token usage. /// 2. If compaction occurred, replaces `self.conversation` with the compacted messages. - /// 3. Notifies observers via [`on_compaction`](AgentObserver::on_compaction). - /// 4. Emits [`ObserveEvent::ContextCompacted`] to the event sink. + /// 3. Notifies observers via [`LoopObserver::on_compaction`](crate::core::observer::LoopObserver::on_compaction). /// /// When no `ContextManager` is set, this is a no-op. /// @@ -65,10 +66,7 @@ impl BareLoop { Ok(EnsureContextResult::Compacted(outcome)) => { self.conversation = outcome.messages; let messages_after = self.conversation.len(); - for obs in &self.observers { - obs.on_compaction(messages_before, messages_after); - } - self.event_sink.on_event(&ObserveEvent::ContextCompacted { + self.managers.observers().on_compaction(&CompactionContext { messages_before, messages_after, tokens_saved: outcome.tokens_saved, diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 1db921f..1b1dc0a 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -8,12 +8,13 @@ use super::HookAction; use super::{ AgentError, ApiClient, Arc, BareLoop, Duration, Instant, PermissionCheck, RecoveryAction, - ReflectionContext, ToolCallInfo, ToolContent, ToolContentPart, ToolContext, - ToolDispatchContext, ToolDispatchResult, ToolPipeline, + ReflectionContext, ToolCallInfo, ToolContent, ToolContext, ToolDispatchContext, + ToolDispatchResult, ToolPipeline, }; #[cfg(feature = "hooks")] use super::{PostToolUseContext, PreToolUseContext}; -use crate::loop_control::loop_detector; +use crate::core::observer::{ToolPostContext, ToolPreContext}; +use crate::loop_control::loop_detector::{self, Operation}; /// Result of deciding what to do after a tool error during recovery. /// @@ -45,8 +46,8 @@ impl BareLoop { /// attempts use the delay specified by the [`RecoveryAction`]. /// /// Observers are notified before and after each tool invocation via - /// [`on_tool_call`](AgentObserver::on_tool_call) and - /// [`on_tool_complete`](AgentObserver::on_tool_complete). + /// [`LoopObserver::on_tool_pre`](crate::core::observer::LoopObserver::on_tool_pre) and + /// [`LoopObserver::on_tool_post`](crate::core::observer::LoopObserver::on_tool_post). /// /// # Errors /// @@ -100,6 +101,12 @@ impl BareLoop { return Err(AgentError::Cancelled); } + self.managers.observers().on_tool_pre(&ToolPreContext { + turn: turn_idx, + tool: tc.name.clone(), + tool_call_id: tc.id.clone(), + }); + if let Some(blocked) = self.check_pre_tool_use_hooks(tc, turn_idx) { return Ok(blocked); } @@ -108,15 +115,19 @@ impl BareLoop { return Ok(blocked); } - self.notify_tool_call(&tc.name, &tc.input.to_string()); - self.emit_tool_start(&tc.name, &tc.input.to_string()); - let start = Instant::now(); let tool_result = self .dispatch_tool(tc, &tool_context, start, turn_idx) .await?; self.post_detection(tc, &tool_result); + self.managers.observers().on_tool_post(&ToolPostContext { + turn: turn_idx, + tool: tc.name.clone(), + result_hash: loop_detector::hash_result(&tool_result.output.to_string()), + is_error: tool_result.is_error, + duration: tool_result.duration, + }); self.notify_post_tool_use_hooks(tc, &tool_result, turn_idx); self.record_tool_health(tc.name.as_str(), &tool_result); @@ -137,52 +148,62 @@ impl BareLoop { /// Check for a loop pattern before executing the tool. /// - /// Hashes the tool input, records the call with the detection manager, - /// and returns a soft-error result if the same (tool, input-hash) pair + /// Extracts the primary parameter from the tool input using the + /// configured [`ToolSignature`], records the call with the detection + /// manager, and returns a soft-error result if the same operation /// has exceeded the loop threshold. Returns `None` when dispatch should /// proceed normally. fn pre_detection(&self, tc: &ToolCallInfo, turn_idx: usize) -> Option { - let input_hash = Self::hash_tool_input(&tc.input); - let pattern = self - .managers - .detection - .record_tool_call(&tc.name, input_hash); - self.handle_detected_pattern(&pattern, turn_idx) + let operation = Operation::from_input_with_signature( + &tc.name, + &tc.input, + self.managers.detection.signature(), + ); + let pattern = self.managers.detection.record_operation(operation); + + // Check inline detection + let inline_blocked = self + .handle_detected_pattern(&pattern, turn_idx) .map(|_result| ToolDispatchResult { tool_call_id: tc.id.clone(), output: ToolContent::Text("loop detected: aborting tool dispatch".into()), is_error: true, duration: Duration::ZERO, resolved_tool_name: tc.name.clone(), - }) + }); + + if inline_blocked.is_some() { + return inline_blocked; + } + + None } /// Record the tool result with the detection manager (post-execution). /// - /// Hashes the tool's text output and feeds the (tool, input-hash, - /// result-hash) triple back to the detection manager. This lets the - /// detector distinguish "same input, same output" (stuck) from - /// "same input, different output" (progress). + /// Constructs an [`Operation`] with the result hash and records it with + /// the detection manager. This lets the detector distinguish "same input, + /// same output" (stuck) from "same input, different output" (progress). fn post_detection(&self, tc: &ToolCallInfo, tool_result: &ToolDispatchResult) { - let input_hash = Self::hash_tool_input(&tc.input); let result_hash = match &tool_result.output { ToolContent::Text(t) => loop_detector::hash_result(t), ToolContent::Multipart(_) => None, }; - if let Some(rh) = result_hash { - let _ = self.managers.detection.record_tool_call_with_result( - &tc.name, - input_hash, - Some(rh), - ); - } + let operation = Operation::from_input_with_result_and_signature( + &tc.name, + &tc.input, + result_hash, + self.managers.detection.signature(), + ); + self.managers.detection.record_operation(operation); } /// Execute a single tool call through the pipeline or registry. /// /// Tries the middleware pipeline first, then a direct registry lookup, /// then produces a not-found error result. Handles cancellation during - /// execution and emits observer/sink events. + /// execution. Observer notification is handled by the caller + /// ([`dispatch_tool_with_recovery`]). /// /// # Errors /// @@ -197,7 +218,7 @@ impl BareLoop { ) -> Result { if let Some(ref pipeline) = self.pipeline { return self - .dispatch_via_pipeline(pipeline, tc, tool_context, start, turn_idx) + .dispatch_via_pipeline(pipeline, tc, tool_context, turn_idx) .await; } @@ -206,33 +227,12 @@ impl BareLoop { let call_result = tokio::select! { r = tool.call(tc.input.clone(), tool_context) => r, () = cancel.notified() => { - let dur = start.elapsed(); - self.notify_tool_complete( - &tc.name, - &tc.input.to_string(), - "", - dur, - false, - Some("cancelled"), - ); - self.emit_tool_complete(&tc.name, "", true, dur); return Err(AgentError::Cancelled); } }; match call_result { Ok(result) => { let duration = start.elapsed(); - let output_text = result.text_content(); - let success = !result.is_error; - self.notify_tool_complete( - &tc.name, - &tc.input.to_string(), - &output_text, - duration, - success, - None, - ); - self.emit_tool_complete(&tc.name, &output_text, !success, duration); ToolDispatchResult { tool_call_id: tc.id.clone(), output: result.payload, @@ -244,15 +244,6 @@ impl BareLoop { Err(e) => { let duration = start.elapsed(); let error_msg = e.to_string(); - self.notify_tool_complete( - &tc.name, - &tc.input.to_string(), - &error_msg, - duration, - false, - Some(&error_msg), - ); - self.emit_tool_complete(&tc.name, &error_msg, true, duration); ToolDispatchResult { tool_call_id: tc.id.clone(), output: ToolContent::Text(error_msg), @@ -271,22 +262,13 @@ impl BareLoop { /// Build a soft-error result for a tool that isn't in the registry. /// - /// Notifies observers and emits a sink event with the error message + /// Notifies observers with the error message /// that lists available tool names to help the model recover. fn tool_not_found(&self, tc: &ToolCallInfo) -> ToolDispatchResult { let available: Vec = self.tools.tool_names(); let available_refs: Vec<&str> = available.iter().map(String::as_str).collect(); let error = AgentError::tool_not_found(&tc.name, &available_refs); let error_msg = error.to_string(); - self.notify_tool_complete( - &tc.name, - &tc.input.to_string(), - &error_msg, - Duration::ZERO, - false, - Some(&error_msg), - ); - self.emit_tool_complete(&tc.name, &error_msg, true, Duration::ZERO); ToolDispatchResult { tool_call_id: tc.id.clone(), output: ToolContent::Text(error_msg), @@ -356,22 +338,18 @@ impl BareLoop { }; match executor.check_pre_tool_use(&ctx) { HookAction::Allow => None, - HookAction::Block { reason } => { - self.emit_tool_complete(&tc.name, &reason, true, Duration::ZERO); - Some(ToolDispatchResult { - tool_call_id: tc.id.clone(), - output: ToolContent::Text(reason), - is_error: true, - duration: Duration::ZERO, - resolved_tool_name: tc.name.clone(), - }) - } + HookAction::Block { reason } => Some(ToolDispatchResult { + tool_call_id: tc.id.clone(), + output: ToolContent::Text(reason), + is_error: true, + duration: Duration::ZERO, + resolved_tool_name: tc.name.clone(), + }), HookAction::Ask { message } => { // In Headless mode (the default) the executor already // downgrades Ask → Block. If we reach this arm the // executor is Interactive, but BareLoop has no UI to // show a prompt, so we still treat it as Block. - self.emit_tool_complete(&tc.name, &message, true, Duration::ZERO); Some(ToolDispatchResult { tool_call_id: tc.id.clone(), output: ToolContent::Text(message), @@ -446,8 +424,9 @@ impl BareLoop { /// /// Builds a [`ToolDispatchContext`] from the tool call info, delegates /// to the pipeline's middleware chain, and converts the - /// [`ToolDispatchResult`] back to a [`ToolDispatchResult`] with proper - /// event emission. Handles cancellation via `tokio::select!`. + /// [`ToolDispatchResult`] back to a [`ToolDispatchResult`]. + /// Observer notification is handled by the caller + /// ([`dispatch_tool_with_recovery`]). /// /// # Errors /// @@ -458,7 +437,6 @@ impl BareLoop { pipeline: &ToolPipeline, tc: &ToolCallInfo, tool_context: &ToolContext, - start: Instant, turn_idx: usize, ) -> Result { let ctx = ToolDispatchContext { @@ -474,48 +452,9 @@ impl BareLoop { let dispatch_result = tokio::select! { r = pipeline.invoke(ctx) => r, () = cancel.notified() => { - let dur = start.elapsed(); - self.notify_tool_complete( - &tc.name, - &tc.input.to_string(), - "", - dur, - false, - Some("cancelled"), - ); - self.emit_tool_complete(&tc.name, "", true, dur); return Err(AgentError::Cancelled); } }; - let output_text = match &dispatch_result.output { - ToolContent::Text(t) => t.clone(), - ToolContent::Multipart(parts) => parts - .iter() - .filter_map(|p| match p { - ToolContentPart::Text { text } => Some(text.as_str()), - ToolContentPart::Image { .. } => None, - }) - .collect::>() - .join(""), - }; - self.notify_tool_complete( - &tc.name, - &tc.input.to_string(), - &output_text, - dispatch_result.duration, - !dispatch_result.is_error, - if dispatch_result.is_error { - Some(&output_text) - } else { - None - }, - ); - self.emit_tool_complete( - &tc.name, - &output_text, - dispatch_result.is_error, - dispatch_result.duration, - ); Ok(ToolDispatchResult { tool_call_id: if dispatch_result.tool_call_id.is_empty() { tc.id.clone() diff --git a/src/engine/bare/emission.rs b/src/engine/bare/emission.rs index 4a36a3f..4c7eac0 100644 --- a/src/engine/bare/emission.rs +++ b/src/engine/bare/emission.rs @@ -1,43 +1,37 @@ -//! Observer notifications and event-sink emissions. +//! Session lifecycle notifications. //! -//! Split from [`BareLoop`] for clarity — these methods are thin wrappers -//! that fan out to every registered [`AgentObserver`] and/or the -//! [`EventSink`]. Keeping them in a dedicated file makes the emission -//! surface area easy to audit and prevents the main loop file from -//! ballooning with boilerplate. +//! Split from [`BareLoop`] for clarity — these methods dispatch to +//! the [`ObserverHost`](crate::core::observer::ObserverHost) and the hook executor. //! -//! Two categories: -//! -//! - **`notify_*`** — iterate `observers` and optionally run hooks. -//! - **`emit_*`** — send a single [`ObserveEvent`] to the [`EventSink`]. +//! Only session start/end live here because they do *two* things: +//! observer notification + hook dispatch. All other observer notifications +//! are called directly at their call sites via +//! `self.managers.observers().on_*()`. -use super::{ApiClient, BareLoop, Duration, EndReason, ObserveEvent, SessionEndInfo}; +use super::{ApiClient, BareLoop, Duration, EndReason, SessionEndInfo}; +use crate::core::observer::{SessionEndContext, SessionStartContext}; #[cfg(feature = "hooks")] -use crate::hooks::context::{SessionEndContext, SessionEndReason, SessionStartContext}; +use crate::hooks::context::{ + SessionEndContext as HookSessionEndContext, SessionEndReason, + SessionStartContext as HookSessionStartContext, +}; // ================================================== -// Observer notifications +// Session lifecycle notifications // ================================================== impl BareLoop { - /// Notify all observers that the session has started. - /// - /// Called once at the beginning of [`run()`](BareLoop::run), - /// before the first turn. Iterates over every registered - /// [`AgentObserver`] and calls - /// [`on_session_start()`](AgentObserver::on_session_start) with the - /// session ID from [`AgentConfig`]. + /// Notify all observers and hooks that the session has started. pub(super) fn notify_session_start(&self) { - for obs in &self.observers { - obs.on_session_start(self.config.session_id); - } - self.event_sink.on_event(&ObserveEvent::SessionStart { - session_id: self.config.session_id, - }); + self.managers + .observers() + .on_session_start(&SessionStartContext { + session_id: self.config.session_id, + }); #[cfg(feature = "hooks")] if let Some(ref executor) = self.hook_executor { - let ctx = SessionStartContext { + let ctx = HookSessionStartContext { session_id: self.config.session_id, model: self.config.model.clone(), working_directory: std::env::current_dir() @@ -48,10 +42,7 @@ impl BareLoop { } } - /// Notify all observers that the session has ended. - /// - /// Called once when [`run()`](BareLoop::run) returns — whether - /// successfully, due to an error, or because of cancellation. + /// Notify all observers and hooks that the session has ended. pub(super) fn notify_session_end(&self, info: &SessionEndInfo) { let reason_str = match &info.reason { EndReason::Complete => None, @@ -59,9 +50,17 @@ impl BareLoop { EndReason::MaxTurns => Some("max turns exceeded"), EndReason::Error => Some("session ended with error"), }; - for obs in &self.observers { - obs.on_session_end(info.success, reason_str); - } + self.managers + .observers() + .on_session_end(&SessionEndContext { + success: info.success, + error: reason_str.map(std::string::ToString::to_string), + total_turns: info.total_turns, + duration_ms: u64::try_from( + std::time::Duration::from_secs(info.duration_secs).as_millis(), + ) + .unwrap_or(u64::MAX), + }); #[cfg(feature = "hooks")] if let Some(ref executor) = self.hook_executor { @@ -71,7 +70,7 @@ impl BareLoop { EndReason::Error => SessionEndReason::Error, EndReason::MaxTurns => SessionEndReason::MaxTurns, }; - let ctx = SessionEndContext { + let ctx = HookSessionEndContext { session_id: self.config.session_id, reason, total_turns: info.total_turns, @@ -82,147 +81,8 @@ impl BareLoop { } } - /// Notify all observers that a turn has started. - /// - /// Called at the top of every iteration of the main loop, before - /// the API streaming request. The `query` parameter is the original - /// user input (the same for every turn within a single - /// [`run()`](BareLoop::run) call). - pub(super) fn notify_turn_start(&self, query: &str) { - for obs in &self.observers { - obs.on_turn_start(query); - } - } - - /// Notify all observers that a turn has ended. - /// - /// Called after each turn completes — whether it produced a tool - /// call, ended with text, or encountered an error. The `success` - /// flag is `true` for normal turns and `false` when the API stream - /// returned an error. - pub(super) fn notify_turn_end(&self, success: bool, error: Option<&str>) { - for obs in &self.observers { - obs.on_turn_end(success, error); - } - } - - /// Notify all observers that a tool is about to be invoked. - /// - /// Called just before the tool's [`call()`](crate::tool::Tool::call) - /// method is invoked. The `tool` parameter is the tool name and - /// `input` is the JSON input serialized to a string. - pub(super) fn notify_tool_call(&self, tool: &str, input: &str) { - for obs in &self.observers { - obs.on_tool_call(tool, input); - } - } - - /// Notify all observers that a tool invocation has completed. - /// - /// Called after the tool's [`call()`](crate::tool::Tool::call) - /// method returns — whether successfully or with an error. - /// Includes the tool's output, execution `duration`, a `success` - /// flag, and an optional `error` message. - /// - /// Parameter count is dictated by the [`AgentObserver::on_tool_complete`] - /// trait method. - pub(super) fn notify_tool_complete( - &self, - tool: &str, - input: &str, - output: &str, - duration: Duration, - success: bool, - error: Option<&str>, - ) { - for obs in &self.observers { - obs.on_tool_complete(tool, input, output, duration, success, error); - } - } - - // ================================================== - // EventSink emissions - // ================================================== - /// Convert a [`Duration`] to milliseconds as `u64`. - /// - /// Clamps at `u64::MAX` if the duration exceeds ~584 million years, - /// which is safe for any practical agent session. pub(super) fn millis_u64(duration: Duration) -> u64 { u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) } - - /// Emit a [`TurnStart`](ObserveEvent::TurnStart) event. - pub(super) fn emit_turn_start(&self, turn: usize, query: &str) { - self.event_sink.on_event(&ObserveEvent::TurnStart { - turn, - query: query.to_string(), - }); - } - - /// Emit a [`TurnComplete`](ObserveEvent::TurnComplete) event. - pub(super) fn emit_turn_complete( - &self, - turn: usize, - duration: Duration, - input_tokens: u64, - output_tokens: u64, - ) { - self.event_sink.on_event(&ObserveEvent::TurnComplete { - turn, - duration_ms: Self::millis_u64(duration), - input_tokens, - output_tokens, - }); - } - - /// Emit a [`TurnFailed`](ObserveEvent::TurnFailed) event. - pub(super) fn emit_turn_failed(&self, turn: usize, duration: Duration, error: &str) { - self.event_sink.on_event(&ObserveEvent::TurnFailed { - turn, - duration_ms: Self::millis_u64(duration), - error: error.to_string(), - }); - } - - /// Emit a [`ToolStart`](ObserveEvent::ToolStart) event. - pub(super) fn emit_tool_start(&self, name: &str, input: &str) { - self.event_sink.on_event(&ObserveEvent::ToolStart { - name: name.to_string(), - input: input.to_string(), - }); - } - - /// Emit a [`ToolComplete`](ObserveEvent::ToolComplete) event. - pub(super) fn emit_tool_complete( - &self, - name: &str, - output: &str, - is_error: bool, - duration: Duration, - ) { - self.event_sink.on_event(&ObserveEvent::ToolComplete { - name: name.to_string(), - output: output.to_string(), - is_error, - duration_ms: Self::millis_u64(duration), - }); - } - - /// Emit a [`SessionStop`](ObserveEvent::SessionStop) event. - pub(super) fn emit_session_stop( - &self, - total_turns: usize, - duration: Duration, - success: bool, - reason: &str, - ) { - self.event_sink.on_event(&ObserveEvent::SessionStop { - session_id: self.config.session_id, - success, - reason: reason.to_string(), - total_turns, - duration_ms: Self::millis_u64(duration), - }); - } } diff --git a/src/hooks.rs b/src/hooks.rs index eca670c..55e1eda 100644 --- a/src/hooks.rs +++ b/src/hooks.rs @@ -1,7 +1,6 @@ //! Hook system — bidirectional lifecycle control for agent loops. //! -//! Hooks differ from observers ([`crate::core::AgentObserver`]) and event sinks -//! ([`crate::observability::EventSink`]) in two key ways: +//! Hooks differ from observers ([`crate::core::observer::LoopObserver`]) in two key ways: //! //! 1. **Return values matter.** Pre-hooks return [`HookAction`] (or [`CompactResult`]) //! to control whether an action proceeds. @@ -59,8 +58,7 @@ use context::{ /// Hook trait for bidirectional lifecycle control. /// -/// Hooks differ from observers ([`crate::core::AgentObserver`]) and event sinks -/// ([`crate::observability::EventSink`]) in two key ways: +/// Hooks differ from observers ([`crate::core::observer::LoopObserver`]) in two key ways: /// /// 1. **Return values matter.** `on_pre_*` methods return `Option` (or /// `Option`). Returning `Some(Block{...})` prevents the action diff --git a/src/lib.rs b/src/lib.rs index 1704d8d..9b0471d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,12 +9,11 @@ //! - **[`builder`]** — Fluent builder API for constructing configured agents. //! - **[`cancel`]** — Cooperative cancellation signal (`CancelSignal`). //! - **[`compact`]** — Context management and compaction (threshold detection, pluggable strategies). -//! - **[`core`]** — Foundational traits (`AgentObserver`) and error types. -//! - **[`builtin`]** — Reference implementations of core traits ([`builtin::memory::InMemoryStore`], [`builtin::observer::LoggingObserver`], etc.). +//! - **[`core`]** — Foundational traits and error types. +//! - **[`builtin`]** — Reference implementations of core traits ([`builtin::memory::InMemoryStore`], etc.). //! - **[`hooks`]** — Bidirectional lifecycle control (allow/block/ask before tool use, compaction). *Requires `hooks` feature.* //! - **[`loop_control`]** — Detection and intervention modules for agent loops. //! - **[`engine`]** — The core agentic loop that orchestrates the full agent lifecycle. -//! - **[`observability`]** — Structured event streaming (`EventSink`, `ObserveEvent`, metrics). //! - **[`stream`]** — Streaming event types for LLM API responses. //! - **[`tool`]** — Tool trait, registry, and supporting types. //! - **[`tool::health`]** — Per-tool health monitoring, circuit breakers, and self-healing routing. *Requires `tool_health` feature.* @@ -31,7 +30,6 @@ pub mod engine; pub mod hooks; pub mod loop_control; pub mod message; -pub mod observability; pub mod stream; #[cfg(feature = "testing")] pub mod testing; diff --git a/src/loop_control.rs b/src/loop_control.rs index 9ec6977..c877dcf 100644 --- a/src/loop_control.rs +++ b/src/loop_control.rs @@ -10,6 +10,8 @@ //! - **[`fallback`]** — Circuit breaker pattern for automatic API model fallback. //! - **[`detection`]** — Unified manager that orchestrates loop and convergence detection. //! - **[`bundle`]** — Aggregate struct for the agent's infrastructure managers. +//! +//! For lifecycle observation, see [`crate::core::observer`]. pub mod bundle; pub mod convergence; diff --git a/src/loop_control/bundle.rs b/src/loop_control/bundle.rs index f00bd9e..6bcb2b6 100644 --- a/src/loop_control/bundle.rs +++ b/src/loop_control/bundle.rs @@ -26,8 +26,10 @@ //! bundle.reset_all(); //! ``` +use crate::core::observer::{LoopObserver, ObserverHost}; use crate::loop_control::detection::DetectionManager; use crate::loop_control::fallback::FallbackManager; +use std::sync::Arc; /// Bundle of framework-provided manager instances. /// @@ -67,6 +69,14 @@ pub struct ManagerBundle { /// /// See [`DetectionManager`] for the full API documentation. pub detection: DetectionManager, + + /// Observer host for cross-cutting lifecycle hooks. + /// + /// Observers are registered via [`ManagerBundle::register_observer`] and + /// called at well-defined hook points inside the agent loop. + /// + /// See [`ObserverHost`] and [`LoopObserver`] for details. + observers: ObserverHost, } impl ManagerBundle { @@ -89,6 +99,7 @@ impl ManagerBundle { Self { fallback: FallbackManager::default(), detection: DetectionManager::default(), + observers: ObserverHost::new(), } } @@ -143,11 +154,39 @@ impl ManagerBundle { self } - /// Reset all managers to their initial state. + /// Register an observer with the observer host. + /// + /// Observers are called at lifecycle hook points inside the agent + /// loop, in registration order. All observers are notified at every + /// hook point (no short-circuiting). + /// + /// See [`LoopObserver`] for the trait definition and available hooks. + /// + /// # Example + /// + /// ```rust,ignore + /// use loopctl::core::observer::LoopObserver; + /// use std::sync::Arc; + /// + /// let mut bundle = ManagerBundle::new(); + /// bundle.register_observer(Arc::new(MyObserver)); + /// ``` + pub fn register_observer(&mut self, observer: Arc) { + self.observers.register(observer); + } + + /// Get a reference to the observer host. + /// + /// Used by `BareLoop` to dispatch hook calls. + pub fn observers(&self) -> &ObserverHost { + &self.observers + } + + /// Reset all managers and observers to their initial state. /// - /// Delegates to each manager's `reset()` method. Typically called at the - /// start of a new agent task or session to clear any accumulated state - /// from a previous run. + /// Delegates to each manager's `reset()` method and calls + /// [`ObserverHost::reset_all`]. Typically called at the start of + /// a new agent task or session. /// /// # Example /// @@ -161,6 +200,7 @@ impl ManagerBundle { pub fn reset_all(&self) { self.fallback.reset(); self.detection.reset(); + self.observers.reset_all(); } } diff --git a/src/loop_control/detection.rs b/src/loop_control/detection.rs index ec6390e..67d49f5 100644 --- a/src/loop_control/detection.rs +++ b/src/loop_control/detection.rs @@ -928,6 +928,14 @@ impl DetectionManager { } } + /// Returns the tool signature used for extracting primary parameters. + /// + /// Useful when callers need to construct [`Operation`]s directly using + /// the same signature the detection manager uses internally. + pub fn signature(&self) -> &dyn ToolSignature { + self.loop_detector.signature() + } + /// Record a tool call for loop detection by tool name and input hash. /// /// Creates an [`Operation`] from a `tool` name and `input_hash`, then diff --git a/src/loop_control/fallback.rs b/src/loop_control/fallback.rs index 3da944c..b2ca380 100644 --- a/src/loop_control/fallback.rs +++ b/src/loop_control/fallback.rs @@ -145,10 +145,8 @@ pub enum FallbackState { /// assert_eq!(FallbackState::from(255u8), FallbackState::Primary); // unknown → safe default /// ``` impl From for FallbackState { - #[allow(clippy::match_same_arms)] fn from(value: u8) -> Self { match value { - 0 => FallbackState::Primary, 1 => FallbackState::Fallback, 2 => FallbackState::Recovering, _ => FallbackState::Primary, diff --git a/src/observability.rs b/src/observability.rs deleted file mode 100644 index 181cc81..0000000 --- a/src/observability.rs +++ /dev/null @@ -1,72 +0,0 @@ -//! Structured event streaming for agent observability. -//! -//! This module provides the [`EventSink`] trait and [`ObserveEvent`] enum — -//! the primary observability abstractions in `loopctl`. Every consumer — -//! console logging, JSONL files, metrics, custom monitors — implements -//! [`EventSink`]. -//! -//! # Architecture -//! -//! ```text -//! ┌───────────────────────────┐ -//! │ Agent Loop │ -//! │ │ -//! │ emits ObserveEvent ├──▶ dyn EventSink::on_event() -//! │ │ │ -//! └───────────────────────────┘ ├─▶ ConsoleSink -//! ├─▶ NullSink -//! ├─▶ CompositeSink ─┬─▶ Sink 1 -//! │ ├─▶ Sink 2 -//! │ └─▶ Sink 3 -//! └─▶ MetricsSink (feature-gated) -//! ``` -//! -//! # Provided Sinks -//! -//! | Sink | Purpose | -//! |---------------------|------------------------------------------------------| -//! | [`NullSink`] | Discards all events. Useful as a default. | -//! | [`ConsoleSink`] | Prints human-readable summaries to stdout. | -//! | [`CompositeSink`] | Fans out to multiple sinks with panic isolation. | -//! -//! # Event Types -//! -//! [`ObserveEvent`] covers the full agent lifecycle: -//! -//! - Session lifecycle (`SessionStart`, `SessionStop`) -//! - Turn lifecycle (`TurnStart`, `TurnComplete`, `TurnFailed`) -//! - Tool execution (`ToolStart`, `ToolComplete`) -//! - Context management (`ContextWarning`, `ContextCompacted`) -//! - Errors (`Error`) -//! -//! # Quick Start -//! -//! ```rust -//! use loopctl::observability::{EventSink, NullSink, ObserveEvent}; -//! -//! let sink = NullSink; -//! sink.on_event(&ObserveEvent::SessionStart { -//! session_id: uuid::Uuid::nil(), -//! }); -//! ``` -//! -//! # Composing Sinks -//! -//! Use [`CompositeSink`] to fan out to multiple sinks: -//! -//! ```rust -//! use loopctl::observability::{CompositeSink, ConsoleSink, NullSink}; -//! -//! let sink = CompositeSink::new(vec![ -//! Box::new(ConsoleSink), -//! Box::new(NullSink), -//! ]); -//! ``` - -pub mod console; -pub mod event; -pub mod sink; - -pub use console::ConsoleSink; -pub use event::ObserveEvent; -pub use sink::{CompositeSink, EventSink, NullSink}; diff --git a/src/observability/console.rs b/src/observability/console.rs deleted file mode 100644 index 2408d16..0000000 --- a/src/observability/console.rs +++ /dev/null @@ -1,231 +0,0 @@ -//! A console-based [`EventSink`] that prints human-readable event summaries. -//! -//! [`ConsoleSink`] provides lightweight, readable output suitable for -//! development and debugging. Not intended for production logging — use -//! a structured sink (JSONL, metrics) for that. - -use super::event::ObserveEvent; -use super::sink::EventSink; - -/// Prints a human-readable summary of each event to stdout. -/// -/// Output is designed for developer ergonomics, not machine parsing. -/// Use this during development to see what the agent loop is doing. -/// Each event type gets a formatted one-line summary with the relevant -/// context (turn number, duration, token counts, etc.). -/// -/// # Example -/// -/// ```rust -/// use loopctl::observability::{ConsoleSink, EventSink, ObserveEvent}; -/// -/// let sink = ConsoleSink; -/// sink.on_event(&ObserveEvent::TurnComplete { -/// turn: 3, -/// duration_ms: 1200, -/// input_tokens: 450, -/// output_tokens: 200, -/// }); -/// // Prints: [turn 3] complete in 1200ms (450 in / 200 out tokens) -/// ``` -#[derive(Debug, Clone, Copy, Default)] -pub struct ConsoleSink; - -impl EventSink for ConsoleSink { - fn on_event(&self, event: &ObserveEvent) { - match event { - ObserveEvent::SessionStart { session_id } => { - println!("[session] started {session_id}"); - } - ObserveEvent::SessionStop { - session_id, - success, - total_turns, - duration_ms, - reason, - } => { - let status = if *success { "ok" } else { "failed" }; - println!( - "[session] {status} {session_id} \u{2014} {total_turns} turns in {duration_ms}ms ({reason})" - ); - } - ObserveEvent::TurnStart { turn, query } => { - let preview = truncate(query, 60); - println!("[turn {turn}] start: {preview}"); - } - ObserveEvent::TurnComplete { - turn, - duration_ms, - input_tokens, - output_tokens, - } => { - println!( - "[turn {turn}] complete in {duration_ms}ms ({input_tokens} in / {output_tokens} out tokens)" - ); - } - ObserveEvent::TurnFailed { - turn, - duration_ms, - error, - } => { - println!("[turn {turn}] failed after {duration_ms}ms: {error}"); - } - ObserveEvent::ToolStart { name, .. } => { - println!("[tool] {name} started"); - } - ObserveEvent::ToolComplete { - name, - is_error: false, - duration_ms, - .. - } => { - println!("[tool] {name} completed in {duration_ms}ms"); - } - ObserveEvent::ToolComplete { - name, - is_error: true, - duration_ms, - .. - } => { - println!("[tool] {name} failed after {duration_ms}ms"); - } - ObserveEvent::ContextWarning { - tokens_used, - tokens_remaining, - } => { - println!("[context] warning: {tokens_used} used, {tokens_remaining} remaining"); - } - ObserveEvent::ContextCompacted { - messages_before, - messages_after, - tokens_saved, - } => { - println!( - "[context] compacted {messages_before} \u{2192} {messages_after} messages ({tokens_saved} tokens saved)" - ); - } - ObserveEvent::Error { message, source } => { - println!("[error] {source}: {message}"); - } - ObserveEvent::LoopDetected { tool, repetitions } => { - println!("[detection] loop detected: {tool} repeated {repetitions} times"); - } - ObserveEvent::ConvergenceDetected { action } => { - println!("[detection] convergence detected (action: {action})"); - } - } - } -} - -/// Truncate a string to approximately `max_len` characters, appending an -/// ellipsis (`\u{2026}`) if truncated. -/// -/// Handles multi-byte characters correctly by operating on char boundaries. -fn truncate(s: &str, max_len: usize) -> String { - if s.chars().count() <= max_len { - return s.to_string(); - } - let truncated: String = s.chars().take(max_len).collect(); - format!("{truncated}\u{2026}") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn console_sink_handles_all_event_types() { - let sink = ConsoleSink; - - sink.on_event(&ObserveEvent::SessionStart { - session_id: uuid::Uuid::nil(), - }); - sink.on_event(&ObserveEvent::TurnStart { - turn: 0, - query: "hello".to_string(), - }); - sink.on_event(&ObserveEvent::ToolStart { - name: "read_file".to_string(), - input: "{}".to_string(), - }); - sink.on_event(&ObserveEvent::ToolComplete { - name: "read_file".to_string(), - output: "contents".to_string(), - is_error: false, - duration_ms: 50, - }); - sink.on_event(&ObserveEvent::ToolComplete { - name: "read_file".to_string(), - output: "not found".to_string(), - is_error: true, - duration_ms: 10, - }); - sink.on_event(&ObserveEvent::TurnComplete { - turn: 0, - duration_ms: 100, - input_tokens: 10, - output_tokens: 5, - }); - sink.on_event(&ObserveEvent::TurnFailed { - turn: 1, - duration_ms: 200, - error: "timeout".to_string(), - }); - sink.on_event(&ObserveEvent::ContextWarning { - tokens_used: 180_000, - tokens_remaining: 20_000, - }); - sink.on_event(&ObserveEvent::ContextCompacted { - messages_before: 50, - messages_after: 20, - tokens_saved: 10_000, - }); - sink.on_event(&ObserveEvent::Error { - message: "something broke".to_string(), - source: "api".to_string(), - }); - sink.on_event(&ObserveEvent::SessionStop { - session_id: uuid::Uuid::nil(), - success: true, - reason: "done".to_string(), - total_turns: 5, - duration_ms: 10_000, - }); - } - - #[test] - fn truncate_short_string_unchanged() { - assert_eq!(truncate("hello", 10), "hello"); - } - - #[test] - fn truncate_long_string() { - let long = "a".repeat(100); - let result = truncate(&long, 10); - assert_eq!(result.chars().count(), 11); // 10 chars + ellipsis - assert!(result.ends_with('\u{2026}')); - } - - #[test] - fn truncate_exact_length() { - assert_eq!(truncate("exactly!", 8), "exactly!"); - } - - #[test] - fn truncate_empty_string() { - assert_eq!(truncate("", 10), ""); - } - - #[test] - fn truncate_multibyte_chars() { - let input = "héllo😊world"; - let result = truncate(input, 5); - assert_eq!(result, "héllo…"); - } - - #[test] - fn truncate_zero_max_len() { - let result = truncate("hello", 0); - assert_eq!(result, "\u{2026}"); - } -} diff --git a/src/observability/event.rs b/src/observability/event.rs deleted file mode 100644 index 0f6ef84..0000000 --- a/src/observability/event.rs +++ /dev/null @@ -1,412 +0,0 @@ -//! Observable events emitted during the agent lifecycle. -//! -//! [`ObserveEvent`] is the central enum for all structured events in the -//! framework. The agent loop emits these at key lifecycle points, and -//! [`EventSink`](super::EventSink) implementations consume them. -//! -//! # Serialization -//! -//! Uses `#[serde(tag = "type")]` with `rename_all = "snake_case"` to produce -//! JSON objects like `{"type": "session_start", "session_id": "..."}`. -//! -//! Agents can wrap `ObserveEvent` in their own enum that adds agent-specific -//! variants. The framework's `EventSink` accepts `&ObserveEvent`; agent sinks -//! can extend with their own event types. - -use serde::{Deserialize, Serialize}; - -/// All observable events in the agent lifecycle. -/// -/// Each variant captures the relevant context for its lifecycle point. -/// Events are ordered chronologically within a session: -/// -/// ```text -/// SessionStart -/// └─ TurnStart -/// ├─ ToolStart ── ToolComplete [per tool call] -/// └─ ContextWarning? [if context is low] -/// └─ TurnComplete | TurnFailed -/// └─ ... -/// └─ ContextCompacted? [if compaction ran] -/// SessionStop -/// ``` -/// -/// # Serialization -/// -/// Each variant serializes as a JSON object with a `"type"` tag: -/// -/// ```rust -/// use loopctl::observability::ObserveEvent; -/// -/// let event = ObserveEvent::SessionStart { -/// session_id: uuid::Uuid::nil(), -/// }; -/// let json = serde_json::to_string(&event).unwrap(); -/// assert!(json.contains("\"type\":\"session_start\"")); -/// ``` -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum ObserveEvent { - /// Session has started. - /// - /// Emitted once at the beginning of an agent session, before any turns. - SessionStart { - /// Unique session identifier. - session_id: uuid::Uuid, - }, - - /// Session has ended. - /// - /// Emitted once after the agent session completes (or fails). - /// Contains aggregate statistics about the entire session. - SessionStop { - /// Unique session identifier. - session_id: uuid::Uuid, - /// Whether the session completed successfully. - success: bool, - /// Why the session stopped (`"max_turns"`, `"cancelled"`, `"error"`, etc.). - reason: String, - /// Total turns completed. - total_turns: usize, - /// Total session duration in milliseconds. - duration_ms: u64, - }, - - /// A turn has started. - /// - /// Emitted at the beginning of each turn, before the LLM is called. - TurnStart { - /// Current turn number (0-indexed). - turn: usize, - /// The user query that initiated this turn. - query: String, - }, - - /// A turn completed successfully. - /// - /// Emitted after the LLM response has been fully processed, - /// including any tool calls made during the turn. - TurnComplete { - /// Turn number. - turn: usize, - /// Wall-clock duration of the turn in milliseconds. - duration_ms: u64, - /// Input tokens consumed this turn. - input_tokens: u64, - /// Output tokens generated this turn. - output_tokens: u64, - }, - - /// A turn failed. - /// - /// Emitted when a turn encounters an unrecoverable error. - TurnFailed { - /// Turn number. - turn: usize, - /// Wall-clock duration before failure. - duration_ms: u64, - /// Error description. - error: String, - }, - - /// A tool execution has started. - /// - /// Emitted just before a tool is invoked. - ToolStart { - /// Tool name. - name: String, - /// Tool input as a JSON string. - input: String, - }, - - /// A tool execution has completed. - /// - /// Emitted after the tool returns, whether successfully or with an error. - ToolComplete { - /// Tool name. - name: String, - /// Tool output as a string. - output: String, - /// Whether the tool reported an error. - is_error: bool, - /// Wall-clock duration in milliseconds. - duration_ms: u64, - }, - - /// Context window is running low on capacity. - /// - /// Emitted when token usage exceeds a warning threshold, - /// before compaction is triggered. - ContextWarning { - /// Estimated tokens currently used. - tokens_used: u64, - /// Estimated tokens remaining. - tokens_remaining: u64, - }, - - /// A context compaction occurred. - /// - /// Emitted after the conversation history has been compacted - /// to free up context window space. - ContextCompacted { - /// Messages before compaction. - messages_before: usize, - /// Messages after compaction. - messages_after: usize, - /// Estimated tokens saved by compaction. - tokens_saved: u64, - }, - - /// A generic error event. - /// - /// Emitted for errors that don't fit a more specific category. - Error { - /// Error description. - message: String, - /// Error source or category. - source: String, - }, - - /// A loop was detected in tool operations. - /// - /// Emitted when the detection manager observes the same tool call - /// pattern repeated beyond the configured loop threshold. - LoopDetected { - /// Tool name that was repeating. - tool: String, - /// Number of repetitions observed. - repetitions: usize, - }, - - /// Convergence was detected in agent responses. - /// - /// Emitted when the detection manager observes that recent agent - /// responses have become semantically similar beyond the configured - /// threshold. - ConvergenceDetected { - /// Configured action to take (e.g. `"stop"`, `"warn"`, `"compact"`). - action: String, - }, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn session_start_round_trip() { - let event = ObserveEvent::SessionStart { - session_id: uuid::Uuid::nil(), - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"session_start\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - assert!(matches!(de, ObserveEvent::SessionStart { .. })); - } - - #[test] - fn session_stop_round_trip() { - let event = ObserveEvent::SessionStop { - session_id: uuid::Uuid::nil(), - success: true, - reason: "max_turns".to_string(), - total_turns: 5, - duration_ms: 10_000, - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"session_stop\"")); - assert!(json.contains("\"success\":true")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - assert!(matches!(de, ObserveEvent::SessionStop { .. })); - } - - #[test] - fn turn_complete_round_trip() { - let event = ObserveEvent::TurnComplete { - turn: 3, - duration_ms: 1200, - input_tokens: 450, - output_tokens: 200, - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"turn_complete\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - if let ObserveEvent::TurnComplete { - turn, - input_tokens, - output_tokens, - .. - } = de - { - assert_eq!(turn, 3); - assert_eq!(input_tokens, 450); - assert_eq!(output_tokens, 200); - } else { - panic!("expected TurnComplete"); - } - } - - #[test] - fn tool_complete_round_trip() { - let success = ObserveEvent::ToolComplete { - name: "read_file".to_string(), - output: "file contents".to_string(), - is_error: false, - duration_ms: 50, - }; - let json = serde_json::to_string(&success).unwrap(); - assert!(json.contains("\"type\":\"tool_complete\"")); - assert!(json.contains("\"is_error\":false")); - - let failure = ObserveEvent::ToolComplete { - name: "read_file".to_string(), - output: "not found".to_string(), - is_error: true, - duration_ms: 10, - }; - let json = serde_json::to_string(&failure).unwrap(); - assert!(json.contains("\"is_error\":true")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - if let ObserveEvent::ToolComplete { is_error, .. } = de { - assert!(is_error); - } else { - panic!("expected ToolComplete"); - } - } - - #[test] - fn context_compacted_round_trip() { - let event = ObserveEvent::ContextCompacted { - messages_before: 50, - messages_after: 20, - tokens_saved: 10_000, - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"context_compacted\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - if let ObserveEvent::ContextCompacted { - messages_before, - messages_after, - tokens_saved, - } = de - { - assert_eq!(messages_before, 50); - assert_eq!(messages_after, 20); - assert_eq!(tokens_saved, 10_000); - } else { - panic!("expected ContextCompacted"); - } - } - - #[test] - fn error_round_trip() { - let event = ObserveEvent::Error { - message: "something went wrong".to_string(), - source: "api".to_string(), - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"error\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - if let ObserveEvent::Error { message, source } = de { - assert_eq!(message, "something went wrong"); - assert_eq!(source, "api"); - } else { - panic!("expected Error"); - } - } - - #[test] - fn context_warning_round_trip() { - let event = ObserveEvent::ContextWarning { - tokens_used: 180_000, - tokens_remaining: 20_000, - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"context_warning\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - assert!(matches!(de, ObserveEvent::ContextWarning { .. })); - } - - #[test] - fn turn_start_round_trip() { - let event = ObserveEvent::TurnStart { - turn: 0, - query: "hello".to_string(), - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"turn_start\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - assert!(matches!(de, ObserveEvent::TurnStart { .. })); - } - - #[test] - fn turn_failed_round_trip() { - let event = ObserveEvent::TurnFailed { - turn: 2, - duration_ms: 500, - error: "timeout".to_string(), - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"turn_failed\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - assert!(matches!(de, ObserveEvent::TurnFailed { .. })); - } - - #[test] - fn tool_start_round_trip() { - let event = ObserveEvent::ToolStart { - name: "read_file".to_string(), - input: r#"{"path":"/tmp/x"}"#.to_string(), - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"tool_start\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - assert!(matches!(de, ObserveEvent::ToolStart { .. })); - } - - #[test] - fn loop_detected_round_trip() { - let event = ObserveEvent::LoopDetected { - tool: "Read(/etc/hosts)".to_string(), - repetitions: 3, - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"loop_detected\"")); - assert!(json.contains("\"repetitions\":3")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - if let ObserveEvent::LoopDetected { tool, repetitions } = de { - assert_eq!(tool, "Read(/etc/hosts)"); - assert_eq!(repetitions, 3); - } else { - panic!("expected LoopDetected"); - } - } - - #[test] - fn convergence_detected_round_trip() { - let event = ObserveEvent::ConvergenceDetected { - action: "stop".to_string(), - }; - let json = serde_json::to_string(&event).unwrap(); - assert!(json.contains("\"type\":\"convergence_detected\"")); - assert!(json.contains("\"action\":\"stop\"")); - - let de: ObserveEvent = serde_json::from_str(&json).unwrap(); - if let ObserveEvent::ConvergenceDetected { action } = de { - assert_eq!(action, "stop"); - } else { - panic!("expected ConvergenceDetected"); - } - } -} diff --git a/src/observability/sink.rs b/src/observability/sink.rs deleted file mode 100644 index f72913c..0000000 --- a/src/observability/sink.rs +++ /dev/null @@ -1,410 +0,0 @@ -//! The [`EventSink`] trait and provided implementations. -//! -//! [`EventSink`] is the primary observability abstraction in `loopctl`. -//! Every consumer — console logging, JSONL files, metrics, custom monitors — -//! implements this trait. -//! -//! # Implementations -//! -//! - [`NullSink`] — Discards all events. Useful as a default. -//! - [`CompositeSink`] — Fans out to multiple sinks with panic isolation. - -use super::event::ObserveEvent; -use std::panic::catch_unwind; -use std::sync::Arc; - -// =================================================== -// EventSink trait -// =================================================== - -/// The primary observability abstraction in `loopctl`. -/// -/// Every consumer — console logging, JSONL files, metrics, custom monitors — -/// implements `EventSink`. The agent loop calls [`on_event`](EventSink::on_event) -/// at key lifecycle points with an [`ObserveEvent`]. -/// -/// # Object Safety -/// -/// The trait is object-safe: `Box` and `Arc` work. -/// -/// # Example -/// -/// ```rust -/// use loopctl::observability::{EventSink, ObserveEvent}; -/// -/// struct PrintSink; -/// -/// impl EventSink for PrintSink { -/// fn on_event(&self, event: &ObserveEvent) { -/// match event { -/// ObserveEvent::ToolStart { name, .. } => { -/// println!("tool started: {name}"); -/// } -/// ObserveEvent::TurnComplete { turn, duration_ms, .. } => { -/// println!("turn {turn} done in {duration_ms}ms"); -/// } -/// _ => {} -/// } -/// } -/// } -/// ``` -pub trait EventSink: Send + Sync { - /// Handle an observed event. - /// - /// Called by the agent loop at each lifecycle point. Implementations - /// should be fast and non-blocking — heavy work should be offloaded - /// to a channel or background task. - fn on_event(&self, event: &ObserveEvent); -} - -// =================================================== -// NullSink -// =================================================== - -/// A sink that discards all events. -/// -/// Useful as a default when no observability is needed, or as a -/// placeholder during testing. Equivalent to `/dev/null` for events. -/// -/// # Example -/// -/// ```rust -/// use loopctl::observability::{EventSink, NullSink, ObserveEvent}; -/// -/// let sink = NullSink; -/// sink.on_event(&ObserveEvent::SessionStart { -/// session_id: uuid::Uuid::nil(), -/// }); -/// // Event is silently discarded. -/// ``` -#[derive(Debug, Clone, Copy, Default)] -pub struct NullSink; - -impl EventSink for NullSink { - fn on_event(&self, _event: &ObserveEvent) {} -} - -// =================================================== -// CompositeSink -// =================================================== - -/// A composite sink that fans out every event to multiple inner sinks. -/// -/// Holds an ordered list of [`EventSink`] trait objects behind [`Arc`] -/// and forwards each event to every inner sink in sequence. This -/// implements the classic **Composite** pattern, allowing you to combine -/// console logging, JSONL output, and custom sinks without writing a -/// new struct. -/// -/// Sinks are called in insertion order. If any individual sink panics, -/// the remaining sinks in the list are still called — this is achieved -/// internally via [`std::panic::catch_unwind`]. This ensures one -/// misbehaving sink does not prevent others from receiving events. -/// -/// # Architecture -/// -/// ```text -/// CompositeSink -/// ┌──────────────┐ -/// on_event() ─►│ sinks[] │──▶ sink[0].on_event() -/// │ │──▶ sink[1].on_event() -/// │ │──▶ sink[2].on_event() -/// └──────────────┘ -/// ``` -/// -/// # Thread Safety -/// -/// Each inner sink is stored as `Arc`, so the same sink -/// can be shared across multiple [`CompositeSink`] instances or other -/// parts of the system. The fan-out loop borrows `&self`, meaning all -/// callbacks must be `&self`-safe (no `&mut self`). -/// -/// # Example -/// -/// ```rust -/// use loopctl::observability::{CompositeSink, ConsoleSink, NullSink, EventSink}; -/// use std::sync::Arc; -/// -/// // Build with owned sinks: -/// let sink = CompositeSink::new(vec![ -/// Box::new(ConsoleSink), -/// Box::new(NullSink), -/// ]); -/// assert_eq!(sink.len(), 2); -/// -/// // Build with the builder pattern: -/// let sink = CompositeSink::new(vec![]) -/// .with(ConsoleSink) -/// .with(NullSink); -/// assert_eq!(sink.len(), 2); -/// -/// // Build with pre-Arc'd sinks: -/// let shared = Arc::new(NullSink); -/// let sink = CompositeSink::new(vec![]) -/// .with_arc(shared.clone()) -/// .with_arc(shared); // same sink twice -/// assert_eq!(sink.len(), 2); -/// ``` -pub struct CompositeSink { - /// The ordered list of inner sinks to fan out to. - /// - /// Each sink is stored as `Arc` so it can be shared - /// across threads. Sinks are called in insertion order — first added - /// is first notified. - sinks: Vec>, -} - -impl std::fmt::Debug for CompositeSink { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("CompositeSink") - .field("sink_count", &self.sinks.len()) - .finish() - } -} - -impl Default for CompositeSink { - /// Returns an empty composite sink with no inner sinks. - /// - /// The default instance contains zero inner sinks, so all events - /// are silently discarded until sinks are added. - fn default() -> Self { - Self { sinks: Vec::new() } - } -} - -impl CompositeSink { - /// Create a new composite sink that fans out to the given sinks. - /// - /// Each sink is wrapped in `Arc` and stored for - /// ordered fan-out. Returns a ready-to-use composite. - /// - /// # Example - /// - /// ```rust - /// use loopctl::observability::{CompositeSink, ConsoleSink, NullSink}; - /// - /// let sink = CompositeSink::new(vec![ - /// Box::new(ConsoleSink), - /// Box::new(NullSink), - /// ]); - /// assert_eq!(sink.len(), 2); - /// ``` - #[must_use] - pub fn new(sinks: Vec>) -> Self { - let sinks: Vec> = sinks.into_iter().map(Arc::from).collect(); - Self { sinks } - } - - /// Add an owned sink to the fan-out list. - /// - /// The sink is wrapped as `Arc` and appended to the - /// end of the list. Returns `self` for chaining. - /// - /// # Example - /// - /// ```rust - /// use loopctl::observability::{CompositeSink, ConsoleSink, NullSink}; - /// - /// let sink = CompositeSink::new(vec![]) - /// .with(ConsoleSink) - /// .with(NullSink); - /// assert_eq!(sink.len(), 2); - /// ``` - #[must_use] - pub fn with(mut self, sink: S) -> Self { - self.sinks.push(Arc::new(sink)); - self - } - - /// Add a sink that is already behind an `Arc`. - /// - /// Useful when multiple [`CompositeSink`] instances need to share the - /// same inner sink, or when the sink is constructed externally and - /// already wrapped in an `Arc`. - /// - /// # Example - /// - /// ```rust - /// use loopctl::observability::{CompositeSink, NullSink}; - /// use std::sync::Arc; - /// - /// let shared = Arc::new(NullSink); - /// let sink = CompositeSink::new(vec![]) - /// .with_arc(shared.clone()) - /// .with_arc(shared); - /// assert_eq!(sink.len(), 2); - /// ``` - #[must_use] - pub fn with_arc(mut self, sink: Arc) -> Self { - self.sinks.push(sink); - self - } - - /// Number of inner sinks in the fan-out list. - /// - /// Returns the count of sinks that will receive events. - #[must_use] - pub fn len(&self) -> usize { - self.sinks.len() - } - - /// Whether there are no inner sinks in the fan-out list. - /// - /// Returns `true` when [`len`](CompositeSink::len) is zero. When - /// empty, [`on_event`](EventSink::on_event) is effectively a no-op - /// (the internal loop body never executes). - #[must_use] - pub fn is_empty(&self) -> bool { - self.sinks.is_empty() - } - - /// Append a sink after construction. - /// - /// The sink is wrapped as `Arc` and appended to the - /// end of the fan-out list. - pub fn add(&mut self, sink: Box) { - self.sinks.push(Arc::from(sink)); - } -} - -impl EventSink for CompositeSink { - fn on_event(&self, event: &ObserveEvent) { - for sink in &self.sinks { - // Panic isolation: a failing sink must not break other sinks. - let result = catch_unwind(std::panic::AssertUnwindSafe(|| { - sink.on_event(event); - })); - if let Err(panic_payload) = result { - let msg = if let Some(s) = panic_payload.downcast_ref::<&str>() { - (*s).to_string() - } else if let Some(s) = panic_payload.downcast_ref::() { - s.clone() - } else { - "unknown panic".to_string() - }; - tracing::error!("EventSink panicked: {msg}"); - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - #[test] - fn null_sink_accepts_events() { - let sink = NullSink; - sink.on_event(&ObserveEvent::SessionStart { - session_id: uuid::Uuid::nil(), - }); - } - - #[test] - fn composite_fans_out() { - static COUNT_A: AtomicUsize = AtomicUsize::new(0); - static COUNT_B: AtomicUsize = AtomicUsize::new(0); - - struct SinkA; - struct SinkB; - - impl EventSink for SinkA { - fn on_event(&self, _event: &ObserveEvent) { - COUNT_A.fetch_add(1, Ordering::Relaxed); - } - } - - impl EventSink for SinkB { - fn on_event(&self, _event: &ObserveEvent) { - COUNT_B.fetch_add(1, Ordering::Relaxed); - } - } - - let composite = CompositeSink::new(vec![Box::new(SinkA), Box::new(SinkB)]); - composite.on_event(&ObserveEvent::SessionStart { - session_id: uuid::Uuid::nil(), - }); - - assert_eq!(COUNT_A.load(Ordering::Relaxed), 1); - assert_eq!(COUNT_B.load(Ordering::Relaxed), 1); - } - - #[test] - fn composite_isolates_panics() { - static COUNT: AtomicUsize = AtomicUsize::new(0); - - struct PanicSink; - struct CountSink; - - impl EventSink for PanicSink { - fn on_event(&self, _event: &ObserveEvent) { - panic!("boom"); - } - } - - impl EventSink for CountSink { - fn on_event(&self, _event: &ObserveEvent) { - COUNT.fetch_add(1, Ordering::Relaxed); - } - } - - // PanicSink is first; CountSink should still receive the event. - let composite = CompositeSink::new(vec![Box::new(PanicSink), Box::new(CountSink)]); - composite.on_event(&ObserveEvent::SessionStart { - session_id: uuid::Uuid::nil(), - }); - - assert_eq!(COUNT.load(Ordering::Relaxed), 1); - } - - #[test] - fn composite_with_builder() { - let composite = CompositeSink::new(vec![]) - .with(NullSink) - .with(NullSink) - .with(NullSink); - assert_eq!(composite.len(), 3); - } - - #[test] - fn composite_with_arc_builder() { - let shared = Arc::new(NullSink); - let composite = CompositeSink::new(vec![]) - .with_arc(shared.clone()) - .with_arc(shared); - assert_eq!(composite.len(), 2); - } - - #[test] - fn composite_add_and_len() { - let mut composite = CompositeSink::new(vec![Box::new(NullSink)]); - assert_eq!(composite.len(), 1); - assert!(!composite.is_empty()); - - composite.add(Box::new(NullSink)); - assert_eq!(composite.len(), 2); - } - - #[test] - fn composite_empty() { - let composite = CompositeSink::new(vec![]); - assert!(composite.is_empty()); - composite.on_event(&ObserveEvent::SessionStart { - session_id: uuid::Uuid::nil(), - }); - } - - #[test] - fn composite_default_is_empty() { - let composite = CompositeSink::default(); - assert!(composite.is_empty()); - assert_eq!(composite.len(), 0); - } - - #[test] - fn event_sink_is_object_safe() { - let _boxed: Box = Box::new(NullSink); - let _arc: Arc = Arc::new(NullSink); - } -} diff --git a/src/stream/heartbeat.rs b/src/stream/heartbeat.rs index e4a0d38..986c819 100644 --- a/src/stream/heartbeat.rs +++ b/src/stream/heartbeat.rs @@ -253,12 +253,19 @@ impl HeartbeatStream { /// ); /// // let stream = HeartbeatStream::new(inner_stream, config); /// ``` - #[allow(clippy::arithmetic_side_effects)] pub fn new(inner: S, config: HeartbeatConfig) -> Self { + /// 30 years in seconds — used as a far-future deadline fallback. + /// Computed as a const so the compiler verifies no overflow. + const THIRTY_YEARS_SECS: u64 = 86400 * 365 * 30; + let now = Instant::now(); // checked_add returns None only for extreme Duration values (hundreds of years). - // Fallback: Instant::now() + 30 years, which is effectively infinite. - let far_future = || Instant::now() + Duration::from_secs(86400 * 365 * 30); + // Fallback: 30 years from now, which is effectively infinite. + let far_future = || { + Instant::now() + .checked_add(Duration::from_secs(THIRTY_YEARS_SECS)) + .unwrap_or(Instant::now()) + }; let deadline = now.checked_add(config.timeout).unwrap_or_else(far_future); let timeout_sleep = Box::pin(tokio::time::sleep_until(tokio::time::Instant::from_std( deadline, From 2ac4687861ff275f94d7dc529964bafc673b4dff Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Sun, 21 Jun 2026 22:12:34 +1200 Subject: [PATCH 02/30] refactor: wip --- src/{api_client.rs => api.rs} | 90 +- src/{api_error.rs => api/error.rs} | 177 +-- src/builder/error.rs | 6 +- src/builtin.rs | 29 - src/cancel.rs | 23 +- src/capabilities.rs | 234 ++++ src/compact.rs | 822 ++++-------- src/compact/truncating.rs | 324 +++++ src/compact/types.rs | 262 ++++ src/config.rs | 87 ++ src/core.rs | 59 - src/core/agent_core.rs | 128 -- src/core/agent_memory.rs | 438 ------- src/core/types.rs | 1079 ---------------- src/detection.rs | 21 + .../convergence.rs | 359 +----- .../loop_detector.rs | 913 +++++++------ .../detection.rs => detection/manager.rs} | 517 ++------ src/engine.rs | 8 +- src/engine/bare.rs | 305 +++-- src/engine/bare/compact.rs | 21 +- src/engine/bare/dispatch.rs | 99 +- src/engine/bare/emission.rs | 4 +- src/engine/bare/stream.rs | 42 +- src/engine/loop_core.rs | 552 ++++++++ src/{core => }/error.rs | 160 ++- src/{loop_control => }/fallback.rs | 502 +++----- src/hooks.rs | 8 +- src/hooks/builtin/auto_commit.rs | 17 +- src/hooks/context.rs | 19 + src/hooks/executor.rs | 23 +- src/lib.rs | 57 +- src/loop_control.rs | 20 - src/loop_control/bundle.rs | 287 ----- src/memory.rs | 176 +++ src/{builtin/memory.rs => memory/builtin.rs} | 243 ++-- src/memory/entry.rs | 214 ++++ src/message.rs | 58 +- src/{engine => }/middleware.rs | 803 +----------- src/middleware/output_limit.rs | 67 + src/middleware/permission.rs | 190 +++ src/middleware/timeout.rs | 158 +++ src/middleware/tool_call.rs | 76 ++ src/middleware/unknown_tool.rs | 417 ++++++ src/{core => }/observer.rs | 246 ++-- src/observer/context.rs | 317 +++++ src/{core => }/reflection.rs | 487 ++----- src/reflection/backoff.rs | 395 ++++++ src/runtime.rs | 599 +++++++++ src/stream.rs | 18 +- src/stream/handler.rs | 45 +- src/stream/heartbeat.rs | 4 +- src/testing.rs | 56 +- src/tool.rs | 1129 +++-------------- src/tool/health.rs | 23 +- src/tool/permission.rs | 247 ++++ src/tool/registry.rs | 508 ++++++++ src/tool/shield.rs | 16 +- 58 files changed, 7036 insertions(+), 7148 deletions(-) rename src/{api_client.rs => api.rs} (77%) rename src/{api_error.rs => api/error.rs} (89%) delete mode 100644 src/builtin.rs create mode 100644 src/capabilities.rs create mode 100644 src/compact/truncating.rs create mode 100644 src/compact/types.rs create mode 100644 src/config.rs delete mode 100644 src/core.rs delete mode 100644 src/core/agent_core.rs delete mode 100644 src/core/agent_memory.rs delete mode 100644 src/core/types.rs create mode 100644 src/detection.rs rename src/{loop_control => detection}/convergence.rs (68%) rename src/{loop_control => detection}/loop_detector.rs (75%) rename src/{loop_control/detection.rs => detection/manager.rs} (71%) create mode 100644 src/engine/loop_core.rs rename src/{core => }/error.rs (73%) rename src/{loop_control => }/fallback.rs (78%) delete mode 100644 src/loop_control.rs delete mode 100644 src/loop_control/bundle.rs create mode 100644 src/memory.rs rename src/{builtin/memory.rs => memory/builtin.rs} (60%) create mode 100644 src/memory/entry.rs rename src/{engine => }/middleware.rs (60%) create mode 100644 src/middleware/output_limit.rs create mode 100644 src/middleware/permission.rs create mode 100644 src/middleware/timeout.rs create mode 100644 src/middleware/tool_call.rs create mode 100644 src/middleware/unknown_tool.rs rename src/{core => }/observer.rs (68%) create mode 100644 src/observer/context.rs rename src/{core => }/reflection.rs (61%) create mode 100644 src/reflection/backoff.rs create mode 100644 src/runtime.rs create mode 100644 src/tool/permission.rs create mode 100644 src/tool/registry.rs diff --git a/src/api_client.rs b/src/api.rs similarity index 77% rename from src/api_client.rs rename to src/api.rs index ec91d22..725b568 100644 --- a/src/api_client.rs +++ b/src/api.rs @@ -1,47 +1,12 @@ -//! API client trait — interface for LLM provider communication. +//! API client and error types — interface for LLM provider communication. //! -//! This module defines the [`ApiClient`] trait that all LLM provider -//! implementations must satisfy. It abstracts away provider-specific details -//! (authentication, request formatting, SSE parsing, error handling) behind -//! a single object-safe interface, so the rest of the framework can work -//! with any LLM backend without knowing the concrete provider. +//! - **`ApiClient`** — Trait that all LLM provider implementations must satisfy. +//! - **`error`** — `ApiError` and `ErrorCode` for all API/infrastructure errors. //! -//! # Type Aliases -//! -//! - [`BoxedApiClient`] — `Box`, useful for single ownership. -//! - [`SharedApiClient`] — `Arc`, useful for sharing across -//! tasks or threads. -//! -//! # Provided Implementations -//! -//! - Built-in providers ship with the framework. -//! - A `MockClient` is included in the `tests` module for unit testing. -//! -//! # Quick Start -//! -//! ```rust,ignore -//! use loopctl::api_client::ApiClient; -//! use loopctl::message::Message; -//! -//! // Concrete implementation is provided by downstream crates -//! let client: Box = get_client(); -//! -//! // Streaming request -//! let stream = client.stream_messages( -//! vec![Message::user("Hello!")], -//! Some("You are helpful.".into()), -//! None, -//! ); -//! -//! // Non-streaming fallback -//! let json = client.create_message( -//! vec![Message::user("Hello!")], -//! Some("You are helpful.".into()), -//! None, -//! ).await?; -//! ``` +//! See the sub-modules for detailed documentation. -use crate::api_error::ApiError; +pub mod error; +use crate::api::error::ApiError; use crate::message::Message; use crate::stream::StreamEvent; use crate::tool::ToolSchema; @@ -50,9 +15,9 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; -/// Trait for API clients that communicate with LLM providers. +/// Interface for API clients that communicate with LLM providers. /// -/// This trait defines the interface for both streaming and non-streaming +/// Defines the contract for both streaming and non-streaming /// message requests. Implementations handle provider-specific details such /// as authentication headers, request body formatting, SSE event parsing, /// and error code mapping. @@ -61,17 +26,6 @@ use std::sync::Arc; /// [`BoxedApiClient`] (`Box`) or [`SharedApiClient`] /// (`Arc`) without issues. /// -/// # Lifecycle -/// -/// ```text -/// Agent turn begins -/// → model() // identify provider -/// → stream_messages(messages, ...) // primary path -/// or create_message(messages, ...) // fallback / one-shot -/// → consume StreamEvent items -/// → turn ends -/// ``` -/// /// # Streaming Contract /// /// [`stream_messages`](ApiClient::stream_messages) returns a `'static` @@ -82,7 +36,7 @@ use std::sync::Arc; /// /// # Implementors /// -/// - Downstream crates provide Anthropic, OpenAI, Ollama, etc. implementations. +/// - Downstream crates provide implementations for concrete LLM providers. /// - Mock implementations for testing can be found in test utilities. /// /// # Example @@ -135,7 +89,7 @@ pub trait ApiClient: Send + Sync { /// Returns the provider-specific model string (e.g., /// `"llm-1"`, `"llm-2"`, `"llm-3"`). Used by the /// framework for logging, token estimation, and fallback routing via - /// [`FallbackManager`](crate::loop_control::fallback::FallbackManager). + /// [`FallbackManager`](crate::fallback::FallbackManager). /// /// Called by the framework during initialization and on each turn for /// observability purposes. @@ -205,16 +159,12 @@ pub trait ApiClient: Send + Sync { ) -> Pin> + Send + '_>>; } -// ================================================== -// Type aliases -// ================================================== - -/// Type alias for a boxed API client with single ownership. +/// Owned, single-threaded API client handle. /// -/// Represents `Box`. Use this when you need owned, -/// single-owner access to a provider client. The underlying client -/// remains `Send + Sync` (required by the trait), so the box can be -/// moved across thread boundaries — but it cannot be cloned or shared. +/// Use when you need owned, single-owner access to a provider client. +/// The underlying client remains `Send + Sync` (required by the trait), +/// so the box can be moved across thread boundaries — but it cannot be +/// cloned or shared. /// /// Ideal for test fixtures and single-agent runners. /// @@ -228,13 +178,11 @@ pub trait ApiClient: Send + Sync { /// ``` pub type BoxedApiClient = Box; -/// Type alias for a shared API client with reference-counted ownership. +/// Shared, reference-counted API client handle. /// -/// Represents `Arc`. Use this when multiple tasks or -/// threads need concurrent access to the same provider client — for -/// example, in a multi-agent system or when sharing a client between -/// an agent and a background metrics collector. `Arc` provides cheap -/// cloning (reference count increment) without duplicating the client. +/// Use when multiple tasks or threads need concurrent access to the +/// same provider client — for example, in a multi-agent system or when +/// sharing a client between an agent and a background metrics collector. /// /// For single-owner usage, see [`BoxedApiClient`]. /// diff --git a/src/api_error.rs b/src/api/error.rs similarity index 89% rename from src/api_error.rs rename to src/api/error.rs index cabccbd..f90f80d 100644 --- a/src/api_error.rs +++ b/src/api/error.rs @@ -1,10 +1,10 @@ //! API and infrastructure error types. //! -//! This module provides the error hierarchy for LLM API interactions, -//! tool execution, configuration handling, and general infrastructure -//! operations. Every failure mode an agent can encounter is captured by -//! [`ApiError`], with a corresponding [`ErrorCode`] for programmatic -//! matching, logging, and metrics. +//! Error hierarchy for LLM API interactions, tool execution, +//! configuration handling, and general infrastructure operations. Every +//! failure mode an agent can encounter is captured by [`ApiError`], with +//! a corresponding [`ErrorCode`] for programmatic matching, logging, +//! and metrics. //! //! # Error Categories //! @@ -31,27 +31,10 @@ //! via [`ApiError::code`]. //! - **[`Result`]** — A convenience alias for `std::result::Result`. //! -//! # Error Flow -//! -//! ```text -//! [API / Tool / Config / I/O] -//! │ -//! ▼ -//! ApiError (enum) -//! │ -//! ├──→ code() → ErrorCode (u32) -//! ├──→ is_retryable() → bool -//! ├──→ is_auth_error() / is_config_error() / is_context_overflow() -//! ├──→ is_io_error() / is_tool_error() -//! │ -//! ▼ -//! Result (= std::result::Result) -//! ``` -//! //! # Quick Start //! //! ```rust -//! use loopctl::api_error::{ApiError, ErrorCode}; +//! use loopctl::api::error::{ApiError, ErrorCode}; //! //! // Construct errors with the ergonomic helpers //! let err = ApiError::api("request failed"); @@ -68,10 +51,6 @@ use serde_repr::{Deserialize_repr, Serialize_repr}; use thiserror::Error; -// ================================================== -// ErrorCode -// ================================================== - /// Machine-readable error codes for programmatic error handling. /// /// Each [`ErrorCode`] variant maps to a stable numeric value (see the @@ -89,7 +68,7 @@ use thiserror::Error; /// # Example /// /// ```rust -/// use loopctl::api_error::{ApiError, ErrorCode}; +/// use loopctl::api::error::{ApiError, ErrorCode}; /// /// let err = ApiError::api_rate_limited(); /// assert_eq!(err.code(), ErrorCode::ApiRateLimited); @@ -107,15 +86,11 @@ pub enum ErrorCode { /// request could not be completed for an unspecified reason. /// Maps to numeric code **1000**. /// - /// This is the default code for [`ApiError::Api`] variants whose - /// message does not match a more specific pattern (timeout, rate - /// limit, stream, or context overflow). - /// - /// # When to use + /// Default code for [`ApiError::Api`] variants whose message does + /// not match a more specific pattern (timeout, rate limit, stream, + /// or context overflow). /// /// Use [`ApiError::api`] to construct errors that map to this code. - /// The framework's retry loop treats this code as retryable via - /// [`ApiError::is_retryable`]. ApiRequestFailed = 1000, /// The API response could not be parsed. @@ -323,13 +298,9 @@ pub enum ErrorCode { Interrupted = 1999, } -// ================================================== -// ApiError -// ================================================== - /// Main error type for API and infrastructure operations. /// -/// This enum covers all failure modes an agent might encounter when +/// Covers all failure modes an agent might encounter when /// interacting with LLM APIs, executing tools, or managing configuration /// and I/O. Each variant stores a human-readable message (or a source /// error via `#[from]`) and can be mapped to a stable [`ErrorCode`] via @@ -340,7 +311,7 @@ pub enum ErrorCode { /// Prefer the ergonomic constructor methods over enum variants directly: /// /// ```rust -/// use loopctl::api_error::{ApiError, ErrorCode}; +/// use loopctl::api::error::{ApiError, ErrorCode}; /// // Instead of ApiError::Api("...".into()) /// let err = ApiError::api("request failed"); /// @@ -462,11 +433,6 @@ impl ApiError { // ================================================== // Classifiers // ================================================== - // - // These methods inspect the error variant (and often the message - // text) to produce a classification — either an [`ErrorCode`] or - // a boolean predicate. They are called by retry loops, logging - // middleware, and metrics collectors throughout the framework. /// Derive the machine-readable [`ErrorCode`] for this error. /// @@ -512,19 +478,6 @@ impl ApiError { /// ApiError::Interrupted → Interrupted /// ApiError::Other(_) → InternalError /// ``` - /// - /// # When called - /// - /// Called by logging middleware, metrics collectors, and retry - /// logic throughout the framework. - /// - /// # Example - /// - /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; - /// let err = ApiError::api("rate limit exceeded (429)"); - /// assert_eq!(err.code(), ErrorCode::ApiRateLimited); - /// ``` #[must_use] pub fn code(&self) -> ErrorCode { match self { @@ -610,16 +563,10 @@ impl ApiError { /// length — e.g. `"context"`, `"too many tokens"`, or /// `"context length"`. /// - /// # When called - /// - /// The agent loop calls this after each failed API call to decide - /// whether to attempt context reduction (summarisation, truncation) - /// before retrying. - /// /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::api("context length exceeded"); /// assert!(err.is_context_overflow()); /// @@ -634,19 +581,7 @@ impl ApiError { } } - /// Internal helper that checks a message string for context-overflow - /// keywords. - /// - /// Not part of the public API — extracted so both [`Self::code`] - /// and [`Self::is_context_overflow`] share the same heuristic. - /// - /// The detection looks for any of the following case-insensitive - /// substrings within the message: - /// - /// - `"context"` — catches "context length exceeded" and similar - /// - `"too many tokens"` — OpenAI-style phrasing - /// - `"exceeds maximum"` — generic overflow wording - /// - `"max tokens"` — common provider phrasing + /// Check a message string for context-overflow keywords. fn is_context_overflow_internal(msg: &str) -> bool { let msg_lower = msg.to_lowercase(); msg_lower.contains("context") @@ -664,15 +599,10 @@ impl ApiError { /// [`ErrorCode::ApiTimeout`], [`ErrorCode::HttpConnectionError`], /// [`ErrorCode::HttpRequestError`], and [`ErrorCode::HttpResponseError`]. /// - /// # When called - /// - /// Retry loops in the API client layer and agent loop call this - /// before deciding whether to re-dispatch a request. - /// /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::api_rate_limited(); /// assert!(err.is_retryable()); /// @@ -700,7 +630,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::ApiError; + /// use loopctl::api::error::ApiError; /// let error = ApiError::auth("bad key"); /// if error.is_auth_error() { /// eprintln!("Please check your API key."); @@ -719,7 +649,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::ApiError; + /// use loopctl::api::error::ApiError; /// let error = ApiError::config("bad config"); /// if error.is_config_error() { /// eprintln!("Configuration problem — check loopctl.toml."); @@ -747,7 +677,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::tool_not_found("Bash"); /// assert!(err.is_tool_error()); /// ``` @@ -759,11 +689,6 @@ impl ApiError { // ================================================== // Constructors // ================================================== - // - // Convenience methods that create the appropriate [`ApiError`] - // variant with a pre-formatted message. The message formatting - // is significant because [`ApiError::code`] performs keyword - // matching on the text to select the most specific [`ErrorCode`]. /// Create a generic [`ApiError::Api`] variant. /// @@ -774,7 +699,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::api("unexpected 502 from upstream"); /// ``` pub fn api(msg: impl Into) -> Self { @@ -790,7 +715,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::auth("token expired"); /// assert_eq!(err.code(), ErrorCode::AuthFailed); /// ``` @@ -807,7 +732,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::auth_invalid_key("key expired"); /// assert_eq!(err.code(), ErrorCode::AuthInvalidKey); /// ``` @@ -824,7 +749,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::http("DNS resolution failed"); /// ``` pub fn http(msg: impl Into) -> Self { @@ -844,7 +769,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::http_with_status(503, "service unavailable"); /// assert!(err.to_string().contains("HTTP 503")); /// assert_eq!(err.code(), ErrorCode::HttpResponseError); @@ -865,7 +790,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::tool("execution failed"); /// assert_eq!(err.code(), ErrorCode::ToolExecutionFailed); /// ``` @@ -875,13 +800,13 @@ impl ApiError { /// Create a [`ApiError::Tool`] variant prefixed with the tool name. /// - /// Formats the message as `"{tool}: {msg}"` so logs clearly indicate + /// Formats the message as `"{tool}: {msg}"` so logs indicate /// which tool failed. /// /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::tool_with_name("Read", "file not found"); /// assert!(err.to_string().contains("Read")); /// ``` @@ -897,7 +822,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::tool_not_found("Bash"); /// assert_eq!(err.code(), ErrorCode::ToolNotFound); /// ``` @@ -913,7 +838,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::tool_permission("Write", "read-only filesystem"); /// assert_eq!(err.code(), ErrorCode::ToolPermissionDenied); /// ``` @@ -929,7 +854,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::tool_input_invalid("Read", "path contains null bytes"); /// assert_eq!(err.code(), ErrorCode::ToolInputInvalid); /// ``` @@ -946,7 +871,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::config("invalid TOML syntax at line 42"); /// assert_eq!(err.code(), ErrorCode::ConfigParseError); /// ``` @@ -962,7 +887,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::config_not_found("/etc/loopctl.toml"); /// assert_eq!(err.code(), ErrorCode::ConfigFileNotFound); /// ``` @@ -978,7 +903,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::config_validation("timeout must be positive"); /// assert_eq!(err.code(), ErrorCode::ConfigValidationError); /// ``` @@ -995,7 +920,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::api_timeout("no response after 30s"); /// assert!(err.is_retryable()); /// assert_eq!(err.code(), ErrorCode::ApiTimeout); @@ -1013,7 +938,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::api_rate_limited(); /// assert!(err.is_retryable()); /// assert_eq!(err.code(), ErrorCode::ApiRateLimited); @@ -1031,7 +956,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::api_stream("connection reset mid-stream"); /// assert_eq!(err.code(), ErrorCode::ApiStreamError); /// ``` @@ -1047,7 +972,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let err = ApiError::other("something unexpected happened"); /// assert_eq!(err.code(), ErrorCode::InternalError); /// ``` @@ -1057,13 +982,13 @@ impl ApiError { /// Create an I/O error for a missing file. /// - /// Convenience constructor for callers that know the operation was - /// a file lookup that failed. Maps to [`ErrorCode::IoFileNotFound`]. + /// Constructor for file-lookup failures. Maps to + /// [`ErrorCode::IoFileNotFound`]. /// /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// /// let err = ApiError::io_not_found( /// std::io::Error::new(std::io::ErrorKind::NotFound, "config.toml"), @@ -1089,7 +1014,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// /// let err = ApiError::io_read( /// std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated"), @@ -1110,7 +1035,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// /// let err = ApiError::io_write( /// std::io::Error::new(std::io::ErrorKind::WriteZero, "disk full"), @@ -1132,7 +1057,7 @@ impl ApiError { /// # Example /// /// ```rust - /// use loopctl::api_error::{ApiError, ErrorCode}; + /// use loopctl::api::error::{ApiError, ErrorCode}; /// let io_err = std::io::Error::new(std::io::ErrorKind::Other, "oops"); /// let api_err = ApiError::from_hyper(io_err); /// assert!(matches!(api_err, ApiError::Http(_))); @@ -1142,22 +1067,12 @@ impl ApiError { } } -/// Result type alias for operations that can fail with [`ApiError`]. -/// -/// Used throughout the crate so that function signatures stay concise. -/// Equivalent to `std::result::Result`. -/// -/// # When to use -/// -/// Any function in the agent framework that can fail should return -/// `Result` rather than spelling out the full path. This keeps -/// signatures uniform and allows `?` propagation across the entire -/// call stack. +/// Result type for operations that can fail with [`ApiError`]. /// /// # Example /// /// ```rust,ignore -/// use loopctl::api_error::{ApiError, ErrorCode}; +/// use loopctl::api::error::{ApiError, ErrorCode}; /// fn load_config() -> Result { /// let text = std::fs::read_to_string("loopctl.toml") /// .map_err(|e| ApiError::config_not_found("loopctl.toml"))?; @@ -1166,10 +1081,6 @@ impl ApiError { /// ``` pub type Result = std::result::Result; -// ================================================== -// Tests -// ================================================== - #[cfg(test)] /// Unit tests for the [`ApiError`] enum and [`ErrorCode`] codes. /// diff --git a/src/builder/error.rs b/src/builder/error.rs index a66cab4..2b8a384 100644 --- a/src/builder/error.rs +++ b/src/builder/error.rs @@ -32,9 +32,9 @@ /// actionable error message (e.g. which features conflict, how many /// observers exceeded the limit). /// -/// # When this is returned +/// # Validation invariants /// -/// The builder validates the following invariants at build time: +/// The builder checks the following at build time: /// /// - A core implementation must be present (enforced statically in most cases, /// but checked dynamically for edge cases). @@ -70,7 +70,7 @@ pub enum BuildError { /// Configuration is invalid. /// /// Wraps a human-readable description of what makes the current - /// `AgentConfig` invalid — for example, a + /// `LoopConfig` invalid — for example, a /// `max_turns` value of zero or a malformed model identifier. /// /// **Fix:** Adjust the config passed to `.with_config()`. diff --git a/src/builtin.rs b/src/builtin.rs deleted file mode 100644 index 1211fc9..0000000 --- a/src/builtin.rs +++ /dev/null @@ -1,29 +0,0 @@ -//! Reference implementations of framework traits. -//! -//! This module provides ready-to-use implementations of the core traits so -//! consumers can get started without writing their own. Each implementation -//! is deliberately simple — suitable for prototyping and testing — and can -//! be replaced with domain-specific versions when needed. -//! -//! # Available Implementations -//! -//! | Implementation | Trait | Purpose | -//! |-------------------|-----------------------------------------------|---------------------------| -//! | [`InMemoryStore`] | [`AgentMemory`](crate::core::AgentMemory) | `Vec`-backed memory store | -//! -//! # Quick Start -//! -//! ```rust -//! use loopctl::builtin::InMemoryStore; -//! use loopctl::core::{AgentMemory, MemoryEntry, MemoryCategory}; -//! -//! # tokio::runtime::Runtime::new().unwrap().block_on(async { -//! // In-memory store for agent memories -//! let mut store = InMemoryStore::new(); -//! store.store(MemoryEntry::new(MemoryCategory::Fact, "PostgreSQL 15 is used")).await.unwrap(); -//! # }); -//! ``` - -pub mod memory; - -pub use memory::InMemoryStore; diff --git a/src/cancel.rs b/src/cancel.rs index ae7cab2..631a992 100644 --- a/src/cancel.rs +++ b/src/cancel.rs @@ -3,7 +3,7 @@ //! [`CancelSignal`] combines an [`AtomicBool`] flag with a //! [`tokio::sync::Notify`] for sub-millisecond wake-up of waiting tasks. //! -//! # Why not just poll an `AtomicBool`? +//! # Why not poll an `AtomicBool`? //! //! Polling works but wastes CPU cycles and introduces latency proportional //! to the poll interval. By pairing the flag with a `Notify`, any call to @@ -31,10 +31,10 @@ use std::sync::atomic::{AtomicBool, Ordering}; use tokio::sync::Notify; -/// Shared cancellation signal: [`AtomicBool`] flag + [`Notify`] for instant -/// wake-up. +/// Shared cancellation signal backed by an [`AtomicBool`] flag and a +/// [`Notify`] for instant wake-up. /// -/// Wrap in `Arc` for sharing across tasks/threads. Create with +/// Wrap in `Arc` for sharing across tasks or threads. Create with /// [`CancelSignal::new`], cancel with [`CancelSignal::cancel`], and await /// instant notification with [`CancelSignal::notified`]. pub struct CancelSignal { @@ -44,6 +44,10 @@ pub struct CancelSignal { impl CancelSignal { /// Create a new, non-cancelled signal. + /// + /// Returns a [`CancelSignal`] with its internal flag set to `false`. + /// The signal is ready to be shared (via `Arc`) and awaited by + /// worker tasks until [`cancel`](Self::cancel) is called. #[must_use] pub fn new() -> Self { Self { @@ -63,11 +67,20 @@ impl CancelSignal { } /// Reset the signal so it can be reused for a new operation. + /// + /// Clears the internal cancellation flag, returning the signal to + /// its initial non-cancelled state. Any subsequent calls to + /// [`is_cancelled`](Self::is_cancelled) will return `false` until + /// [`cancel`](Self::cancel) is called again. pub fn reset(&self) { self.flag.store(false, Ordering::Release); } - /// Check whether the signal has been cancelled (non-blocking). + /// Check whether the signal has been cancelled. + /// + /// Performs a non-blocking load of the internal flag. Returns + /// `true` if [`cancel`](Self::cancel) has been called since the + /// last [`reset`](Self::reset) (or since construction). pub fn is_cancelled(&self) -> bool { self.flag.load(Ordering::Acquire) } diff --git a/src/capabilities.rs b/src/capabilities.rs new file mode 100644 index 0000000..62b7e66 --- /dev/null +++ b/src/capabilities.rs @@ -0,0 +1,234 @@ +//! Capability traits for the agent loop runtime. +//! +//! Each trait represents a single infrastructure capability that the +//! agent loop can depend on. [`LoopRuntime`](crate::runtime::LoopRuntime) +//! implements all of them, but consumers can narrow their bounds to +//! only the capabilities they need. +//! +//! # Traits +//! +//! | Trait | Purpose | +//! |-------|---------| +//! | [`Observable`] | Lifecycle event observation | +//! | [`Detectable`] | Loop and convergence detection | +//! | [`FallbackCapable`] | Model fallback / circuit breaker | +//! | [`Compactable`] | Context compaction | +//! | [`StreamCapable`] | Resilient LLM streaming | +//! | [`Hookable`] | Bidirectional lifecycle hooks | +//! | [`PipelineAware`] | Middleware pipeline dispatch | +//! | [`HealthTrackable`] | Per-tool health monitoring *(requires `tool_health` feature)* | +//! +//! # When to use +//! +//! Use these traits as bounds when you need a specific capability +//! without pulling in the full [`LoopRuntime`](crate::runtime::LoopRuntime): +//! +//! ```rust,ignore +//! fn check_patterns(runtime: &impl Detectable) { +//! let pattern = runtime.detection().record_tool_call("Read", hash); +//! } +//! ``` + +use std::sync::Arc; + +use crate::compact::ContextManager; +use crate::detection::DetectionManager; +use crate::fallback::FallbackManager; +#[cfg(feature = "hooks")] +use crate::hooks::HookExecutor; +use crate::middleware::ToolPipeline; +use crate::observer::ObserverHost; +use crate::stream::handler::StreamHandler; +#[cfg(feature = "tool_health")] +use crate::tool::health::ToolHealthRegistry; + +// ================================================== +// Capability Traits +// ================================================== + +/// Capability to emit lifecycle events to registered observers. +/// +/// Observers receive read-only notifications at well-defined hook points +/// in the agent loop. They cannot influence control flow — for that, see +/// [`Hookable`]. +/// +/// # Implementors +/// +/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// +/// # When to use +/// +/// Use this trait bound when you need to notify observers of lifecycle +/// events but don't need any other infrastructure capabilities. +/// +/// ```rust,ignore +/// fn process_turn(runtime: &impl Observable) { +/// runtime.observers().on_turn_start(&ctx); +/// // ... do work ... +/// runtime.observers().on_turn_end(&ctx); +/// } +/// ``` +pub trait Observable { + fn observers(&self) -> &ObserverHost; +} + +/// Capability to detect repetitive loops and semantic convergence. +/// +/// Loop detection catches when the agent repeats the same tool operations +/// in a cycle. Convergence detection catches when successive assistant +/// responses become semantically similar. Both are handled by +/// [`DetectionManager`]. +/// +/// # Implementors +/// +/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// +/// # When to use +/// +/// Use this trait bound when you need to record operations or responses +/// and check whether a loop or convergence pattern has been detected. +/// +/// ```rust,ignore +/// fn check_patterns(runtime: &impl Detectable) { +/// let pattern = runtime.detection().record_tool_call("Read", hash); +/// if let DetectedPattern::LoopDetected { .. } = pattern { +/// // intervention needed +/// } +/// } +/// ``` +pub trait Detectable { + fn detection(&self) -> &DetectionManager; +} + +/// Capability to fall back to an alternate model when the primary fails. +/// +/// Wraps a [`FallbackManager`] that acts as a circuit breaker: after +/// consecutive API failures exceed a threshold, requests are rerouted +/// to a fallback model until the primary stabilises. +/// +/// # Implementors +/// +/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// +/// # When to use +/// +/// Use this trait bound when you need to record API failures or check +/// whether the circuit breaker has tripped. +/// +/// ```rust,ignore +/// fn handle_stream_error(runtime: &impl FallbackCapable) { +/// let tripped = runtime.fallback().record_api_failure(); +/// if tripped { +/// if let Some(model) = runtime.fallback().fallback_model() { +/// // switch to fallback model +/// } +/// } +/// } +/// ``` +pub trait FallbackCapable { + fn fallback(&self) -> &FallbackManager; +} + +/// Capability to compact conversation context when token usage exceeds a threshold. +/// +/// When a [`ContextManager`] is configured, the loop checks token usage +/// after each turn and triggers compaction when usage exceeds the +/// configured threshold. Compaction replaces conversation messages with +/// a compressed version, preserving the most recent context. +/// +/// # Implementors +/// +/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// +/// # When to use +/// +/// Use this trait bound when you need to inspect or trigger context +/// compaction during the agent loop. Useful for custom loop +/// implementations that need to manage the context window directly. +pub trait Compactable { + fn context_manager(&self) -> Option<&Arc>; +} + +/// Capability to stream LLM responses with retry, timeout, and fallback. +/// +/// When a [`StreamHandler`] is configured, the loop delegates streaming +/// to it instead of using the basic inline logic. The handler provides +/// automatic retries, per-event timeouts, and fallback to non-streaming +/// mode. +/// +/// # Implementors +/// +/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// +/// # When to use +/// +/// Use this trait bound when you need to access the stream handler for +/// resilient streaming. Useful for custom loop implementations +/// that need to control streaming behaviour (timeouts, retries, fallback +/// to non-streaming mode). +pub trait StreamCapable { + fn stream_handler(&self) -> Option<&StreamHandler>; +} + +/// Capability to run bidirectional hooks that can block actions. +/// +/// Hooks differ from observers ([`Observable`]) in that they return +/// [`HookAction`](crate::hooks::HookAction) to control whether an action +/// proceeds. The executor stops at the first hook that returns a blocking +/// result. +/// +/// # Implementors +/// +/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// +/// # When to use +/// +/// Use this trait bound when you need to run hooks before or after +/// tool dispatch, compaction, or session start/end. +#[cfg(feature = "hooks")] +pub trait Hookable { + fn hook_executor(&self) -> Option<&HookExecutor>; +} + +/// Capability to dispatch tools through a middleware pipeline. +/// +/// When a pipeline is configured, tool calls flow through middleware +/// layers (timeouts, output limiting, etc.) before reaching the registry. +/// +/// # Implementors +/// +/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +/// +/// # When to use +/// +/// Use this trait bound when you need to dispatch a tool call through +/// optional middleware. +/// +/// ```rust,ignore +/// async fn dispatch(runtime: &impl PipelineAware, ctx: ToolDispatchContext) { +/// if let Some(pipeline) = runtime.pipeline() { +/// pipeline.invoke(ctx).await +/// } else { +/// // direct dispatch +/// } +/// } +/// ``` +pub trait PipelineAware { + fn pipeline(&self) -> Option<&ToolPipeline>; +} + +/// Capability to track per-tool health with circuit breakers. +/// +/// Records success/failure and latency for every tool dispatch. +/// Tools that exceed the failure threshold have their circuit breaker +/// opened, blocking subsequent calls until recovery. +/// +/// *Requires `tool_health` feature.* +/// +/// # Implementors +/// +/// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. +#[cfg(feature = "tool_health")] +pub trait HealthTrackable { + fn health_registry(&self) -> Option<&ToolHealthRegistry>; +} diff --git a/src/compact.rs b/src/compact.rs index 0c3afbf..432a159 100644 --- a/src/compact.rs +++ b/src/compact.rs @@ -1,8 +1,8 @@ //! Context management and compaction for agent conversations. //! //! As conversations grow, they approach the model's context window limit. -//! This module provides the infrastructure to detect when compaction is -//! needed and to carry it out through a pluggable strategy. +//! Infrastructure to detect when compaction is needed and to carry it out +//! through a pluggable strategy. //! //! # Architecture //! @@ -15,19 +15,19 @@ //! approach. //! //! ```text -//! ┌───────────────────────────────┐ -//! │ ContextManager │ -//! │ │ -//! │ estimate_tokens() │ -//! │ should_compact() │ -//! │ ensure_context_fits() │ -//! │ │ │ -//! │ ▼ │ -//! │ ┌──────────────────────┐ │ -//! │ │ dyn ContextCompactor│ │ -//! │ │ .compact() │ │ -//! │ └──────────────────────┘ │ -//! └───────────────────────────────┘ +//! ┌────────────────────────────┐ +//! │ ContextManager │ +//! │ │ +//! │ estimate_tokens() │ +//! │ should_compact() │ +//! │ ensure_context_fits() │ +//! │ │ │ +//! │ ▼ │ +//! │ ┌──────────────────────┐ │ +//! │ │ dyn ContextCompactor│ │ +//! │ │ .compact() │ │ +//! │ └──────────────────────┘ │ +//! └────────────────────────────┘ //! ``` //! //! # Provided Compactors @@ -71,13 +71,22 @@ //! [`ContextManager`]. When present, it checks token usage after each turn //! and triggers compaction automatically when usage exceeds the threshold. -use crate::message::{Message, Role}; +use crate::message::{Message, MessagePart, Role}; use std::fmt; use std::future::Future; use std::pin::Pin; use std::sync::Arc; use std::time::Instant; +pub mod truncating; +pub mod types; + +pub use truncating::{SplitResult, TokenSplitter, TruncatingCompactor}; +pub use types::{ + CompactReason, CompactTelemetry, CompactionContext, CompactionOutcome, ContextOverflow, + EnsureContextResult, PostCompactStats, PreCompactStats, +}; + // =================================================== // ContextCompactor trait // =================================================== @@ -150,7 +159,6 @@ pub trait ContextCompactor: Send + Sync { /// * `target_tokens` — The target token count for the compacted output. /// * `context` — Metadata about the compaction trigger. // The return-type boxing is required for object safety. - #[allow(clippy::type_complexity)] fn compact( &self, messages: Vec, @@ -160,565 +168,55 @@ pub trait ContextCompactor: Send + Sync { } // =================================================== -// CompactionContext +// CompactBase // =================================================== -/// Metadata passed to [`ContextCompactor::compact`] describing the -/// compaction trigger and current state. +/// Determines the base used to calculate the compaction target. /// -/// Compactors can use this information to decide how aggressively to -/// compact — e.g. an emergency compaction may use more aggressive -/// summarization than a routine threshold check. -#[derive(Debug, Clone)] -pub struct CompactionContext { - /// Estimated token count before compaction. - pub tokens_before: u64, - /// Why compaction was triggered. - pub reason: CompactReason, - /// The model's context window size. - pub context_window: u64, - /// The current turn number in the session. - pub turn: usize, -} - -// =================================================== -// CompactReason -// =================================================== - -/// Why compaction was triggered. +/// When compaction triggers, the manager asks the compactor to reduce +/// the conversation to some target token count. This enum controls +/// *what* that target is a percentage *of*: /// -/// Different triggers may warrant different compaction strategies. -/// For example, an [`Emergency`](CompactReason::Emergency) compaction -/// should be more aggressive than a routine threshold check. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum CompactReason { - /// Token usage exceeded the configured threshold percentage. - ThresholdExceeded, - /// Token usage is dangerously close to the context window limit. - Emergency, - /// Compaction was explicitly requested (e.g. by the agent or a tool). - Manual, -} - -impl fmt::Display for CompactReason { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::ThresholdExceeded => write!(f, "threshold exceeded"), - Self::Emergency => write!(f, "emergency"), - Self::Manual => write!(f, "manual"), - } - } -} - -// =================================================== -// CompactionOutcome -// =================================================== - -/// Result of a single compaction pass. +/// - [`Threshold`](CompactBase::Threshold): the target is a percentage +/// of the trigger threshold (`context_window × threshold`). +/// - [`Context`](CompactBase::Context): the target is a percentage +/// of the full context window. /// -/// Returned by [`ContextCompactor::compact`], this struct contains the -/// compacted message list along with telemetry data about what happened. -#[derive(Debug, Clone)] -pub struct CompactionOutcome { - /// The compacted message list. - pub messages: Vec, - /// How many messages were removed by compaction. - pub messages_compacted: usize, - /// Estimated token count after compaction. - pub tokens_after: u64, - /// Estimated tokens saved by compaction. - pub tokens_saved: u64, - /// Whether compaction succeeded. - pub success: bool, - /// Error message if compaction failed. - pub error: Option, -} - -impl CompactionOutcome { - /// Create an outcome representing no change (compaction was not needed). - /// - /// Use this when the compactor decides the messages don't need - /// compaction — e.g. when the message count is below the minimum. - #[must_use] - pub fn no_change(messages: Vec) -> Self { - let tokens = Self::estimate_tokens(&messages); - Self { - messages, - messages_compacted: 0, - tokens_after: tokens, - tokens_saved: 0, - success: true, - error: None, - } - } - - /// Create an outcome representing successful compaction. - #[must_use] - pub fn compacted(messages: Vec, tokens_before: u64, tokens_after: u64) -> Self { - Self { - messages_compacted: 0, // caller should set - tokens_saved: tokens_before.saturating_sub(tokens_after), - messages, - tokens_after, - success: true, - error: None, - } - } - - /// Estimate the token count for a slice of messages. - /// - /// Uses the standard 4-chars-per-token heuristic. This is the same - /// heuristic used by [`ContextManager::estimate_tokens`]. - #[must_use] - pub fn estimate_tokens(messages: &[Message]) -> u64 { - ContextManager::estimate_tokens(messages) - } -} - -// =================================================== -// CompactTelemetry -// =================================================== - -/// Telemetry data for a single compaction operation. -/// -/// Produced by [`ContextManager::ensure_context_fits`] when compaction -/// occurs. Observers receive this via -/// [`on_compaction`](crate::core::observer::LoopObserver::on_compaction). -#[derive(Debug, Clone)] -pub struct CompactTelemetry { - /// Why compaction was triggered. - pub trigger: CompactReason, - /// Conversation stats before compaction. - pub pre_compact: PreCompactStats, - /// Conversation stats after compaction. - pub post_compact: PostCompactStats, - /// Wall-clock duration of the compaction. - pub duration: std::time::Duration, -} - -/// Conversation statistics captured before compaction. -#[derive(Debug, Clone)] -pub struct PreCompactStats { - /// Total number of messages in the conversation. - pub total_messages: usize, - /// Estimated token count. - pub estimated_tokens: u64, - /// Number of user-role messages. - pub user_messages: usize, - /// Number of assistant-role messages. - pub assistant_messages: usize, - /// Number of messages containing tool calls or results. - pub tool_messages: usize, -} - -/// Conversation statistics captured after compaction. -#[derive(Debug, Clone)] -pub struct PostCompactStats { - /// Total number of messages after compaction. - pub total_messages: usize, - /// Estimated token count after compaction. - pub estimated_tokens: u64, - /// Tokens removed by compaction. - pub tokens_saved: u64, - /// Percentage of tokens saved (0–100). - pub percent_saved: u8, -} - -// =================================================== -// TruncatingCompactor -// =================================================== - -/// A simple compactor that drops the oldest messages. -/// -/// Keeps the first message (typically the system prompt) and a configurable -/// number of recent messages. No LLM calls required — useful as a fallback -/// or for contexts where summarization isn't available. -/// -/// # Strategy -/// -/// ```text -/// [System?] [Old₁, Old₂, ..., Oldₙ] [Recent₁, Recent₂, ..., Recentₘ] -/// ↑ kept ↑ discarded ↑ ↑ preserved ↑ -/// ``` -/// -/// The first message is always retained (if present) because it usually -/// contains the system prompt or conversation instructions. This prevents -/// the compactor from discarding essential context that shapes the agent's -/// behavior. If the conversation is shorter than `min_messages`, no -/// compaction occurs. +/// The percentage itself is configured via +/// [`with_compact_target_pct`](ContextManager::with_compact_target_pct). /// /// # Example /// /// ```rust -/// use loopctl::compact::TruncatingCompactor; +/// use loopctl::compact::{CompactBase, ContextManager, TruncatingCompactor}; /// use std::sync::Arc; /// -/// let compactor = TruncatingCompactor::new() -/// .with_preserve_recent(6) -/// .with_min_messages(8); -/// -/// // Pass to ContextManager: -/// // let manager = ContextManager::new(Arc::new(compactor)); -/// ``` -#[derive(Debug, Clone)] -pub struct TruncatingCompactor { - /// Number of recent messages to always preserve. - preserve_recent: usize, - /// Minimum messages before compaction is considered. - min_messages: usize, -} - -impl TruncatingCompactor { - /// Create a new truncating compactor with sensible defaults. - /// - /// Defaults: - /// - /// | Setting | Default | - /// |-------------------|---------| - /// | `preserve_recent` | 4 | - /// | `min_messages` | 6 | - #[must_use] - pub fn new() -> Self { - Self { - preserve_recent: 4, - min_messages: 6, - } - } - - /// Set how many recent messages to preserve during compaction. - /// - /// This many messages from the end of the conversation are kept - /// intact. The rest are dropped. Must be at least 1. - #[must_use] - pub fn with_preserve_recent(mut self, count: usize) -> Self { - self.preserve_recent = count.max(1); - self - } - - /// Set the minimum number of messages before compaction is attempted. - /// - /// If the conversation has fewer messages than this, compaction is - /// skipped entirely. Prevents aggressive truncation of short - /// conversations. - #[must_use] - pub fn with_min_messages(mut self, count: usize) -> Self { - self.min_messages = count.max(2); - self - } - - /// Number of recent messages that will be preserved. - #[must_use] - pub fn preserve_recent(&self) -> usize { - self.preserve_recent - } - - /// Minimum messages before compaction is attempted. - #[must_use] - pub fn min_messages(&self) -> usize { - self.min_messages - } -} - -impl Default for TruncatingCompactor { - fn default() -> Self { - Self::new() - } -} - -impl ContextCompactor for TruncatingCompactor { - fn compact( - &self, - messages: Vec, - _target_tokens: u64, - context: CompactionContext, - ) -> Pin + Send + '_>> { - Box::pin(async move { - let total = messages.len(); - if total <= self.min_messages { - return CompactionOutcome::no_change(messages); - } - - // Determine split point: keep `preserve_recent` from the end. - let split = total.saturating_sub(self.preserve_recent); - let recent: Vec = messages.get(split..).unwrap_or_default().to_vec(); - - // Always preserve the first message (typically the system prompt) - // unless it is already included in the recent slice (split == 0). - let preserved = if split > 0 { - if let Some(first) = messages.first() { - let mut v = vec![first.clone()]; - v.extend(recent); - v - } else { - recent - } - } else { - // split == 0 means recent already contains all messages. - recent - }; - - let tokens_after = CompactionOutcome::estimate_tokens(&preserved); - let messages_compacted = total.saturating_sub(preserved.len()); - - CompactionOutcome { - messages: preserved, - messages_compacted, - tokens_after, - tokens_saved: context.tokens_before.saturating_sub(tokens_after), - success: true, - error: None, - } - }) - } -} - -// =================================================== -// TokenSplitter -// =================================================== - -/// Splits a conversation into "old" and "recent" at a turn boundary. -/// -/// Used by compactors (and agent-side code) that need to know which -/// messages to compact versus preserve. Splits at role transitions for -/// coherent summarization — the split always occurs between a complete -/// request/response pair. -/// -/// # Rules -/// -/// - Never compact the last user message (it's the current request). -/// - Split at turn boundaries (role transitions) for coherent output. -/// - If the conversation is too short, `to_compact` will be empty. -/// -/// # Example -/// -/// ```rust -/// use loopctl::compact::TokenSplitter; -/// use loopctl::message::Message; -/// -/// let splitter = TokenSplitter::new() -/// .with_preserve_recent(4) -/// .with_min_messages(6); -/// -/// let messages = vec![ -/// Message::user("Hello"), -/// Message::assistant("Hi there!"), -/// Message::user("What is 2+2?"), -/// Message::assistant("4"), -/// ]; -/// -/// let result = splitter.split(&messages); -/// // With only 4 messages and min_messages=6, nothing is split off. -/// assert!(result.to_compact.is_empty()); -/// assert_eq!(result.preserved.len(), 4); +/// let compactor = TruncatingCompactor::new(); +/// let manager = ContextManager::new(Arc::new(compactor)) +/// .with_context_window(200_000) +/// .with_threshold(0.80) // triggers at 160k tokens +/// .with_compact_target(CompactBase::Context) // target = % of 200k +/// .with_compact_target_pct(0.50); // compact to 50% of 200k = 100k /// ``` -#[derive(Debug, Clone)] -pub struct TokenSplitter { - /// Number of recent messages to always preserve. - preserve_recent: usize, - /// Minimum messages before considering a split. - min_messages: usize, -} - -/// Result of splitting a conversation into old and recent portions. -#[derive(Debug, Clone)] -pub struct SplitResult { - /// Messages to compact or summarize (the "old" part). - pub to_compact: Vec, - /// Messages to preserve as-is (the "recent" part). - pub preserved: Vec, - /// Estimated tokens in `to_compact`. - pub compact_tokens: u64, - /// Estimated tokens in `preserved`. - pub preserved_tokens: u64, - /// The index in the original message list where the split occurred. - pub split_index: usize, -} - -impl TokenSplitter { - /// Create a new splitter with sensible defaults. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum CompactBase { + /// Target is a percentage of the full context window. /// - /// Defaults: + /// `target = context_window × compact_target_pct` /// - /// | Setting | Default | - /// |-------------------|---------| - /// | `preserve_recent` | 4 | - /// | `min_messages` | 6 | - #[must_use] - pub fn new() -> Self { - Self { - preserve_recent: 4, - min_messages: 6, - } - } - - /// Set how many recent messages to preserve. - #[must_use] - pub fn with_preserve_recent(mut self, count: usize) -> Self { - self.preserve_recent = count.max(1); - self - } - - /// Set the minimum messages before splitting is considered. - #[must_use] - pub fn with_min_messages(mut self, count: usize) -> Self { - self.min_messages = count.max(2); - self - } + /// Use this when you want compaction to aim for a fixed fraction + /// of the model's total capacity regardless of the trigger threshold. + #[default] + Context, - /// Split the given messages into old and recent portions. + /// Target is a percentage of the trigger threshold. /// - /// The split point is chosen at a turn boundary (a role transition - /// from assistant to user) as close as possible to leaving - /// `preserve_recent` messages in the recent portion. - /// - /// If the conversation has fewer than `min_messages`, the entire - /// conversation goes into `preserved` and `to_compact` is empty. - #[must_use] - pub fn split(&self, messages: &[Message]) -> SplitResult { - if messages.len() <= self.min_messages { - return SplitResult { - to_compact: vec![], - preserved: messages.to_vec(), - compact_tokens: 0, - preserved_tokens: ContextManager::estimate_tokens(messages), - split_index: 0, - }; - } - - // Find a split point: we want `preserve_recent` messages at the end. - // Look for a turn boundary (assistant→user transition) near the - // target split point. - let target_split = messages.len().saturating_sub(self.preserve_recent); - let split_index = Self::find_turn_boundary(messages, target_split); - let (to_compact, preserved) = messages.split_at(split_index); - SplitResult { - to_compact: to_compact.to_vec(), - preserved: preserved.to_vec(), - compact_tokens: ContextManager::estimate_tokens(to_compact), - preserved_tokens: ContextManager::estimate_tokens(preserved), - split_index, - } - } - - /// Find the nearest turn boundary at or before the target index. + /// `target = compact_threshold_tokens × compact_target_pct` /// - /// A turn boundary is a position where the previous message is - /// assistant-role and the next is user-role. This ensures we split - /// at a coherent conversation boundary. - fn find_turn_boundary(messages: &[Message], target: usize) -> usize { - if target == 0 { - return 0; - } - - // Search backwards from target for an assistant→user transition. - for i in (1..=target).rev() { - if i < messages.len() { - let Some(prev) = messages.get(i.saturating_sub(1)) else { - continue; - }; - let Some(curr) = messages.get(i) else { - continue; - }; - let prev_is_assistant = prev.role == Role::Assistant; - let curr_is_user = curr.role == Role::User; - if prev_is_assistant && curr_is_user { - return i; - } - } - } - - // Fallback: no clean boundary found, split at target. - target - } -} - -impl Default for TokenSplitter { - fn default() -> Self { - Self::new() - } -} - -// =================================================== -// ContextOverflow error -// =================================================== - -/// Error returned when the conversation cannot fit within the context -/// window, even after compaction. -/// -/// This is a terminal condition — the conversation is too large and the -/// compactor was unable to reduce it sufficiently. -#[derive(Debug, Clone)] -pub struct ContextOverflow { - /// Estimated token count of the conversation. - pub tokens_used: u64, - /// The model's context window size. - pub context_window: u64, - /// How many messages were in the conversation. - pub message_count: usize, - /// The reason compaction was attempted. - pub trigger: CompactReason, - /// Error from the compactor, if compaction was attempted. - pub compactor_error: Option, -} - -impl ContextOverflow { - /// How many tokens the conversation exceeds the window by. - #[must_use] - pub fn overflow(&self) -> u64 { - self.tokens_used.saturating_sub(self.context_window) - } - - /// The fraction of the context window used (0.0–1.0+). - #[must_use] - pub fn utilization(&self) -> f64 { - if self.context_window == 0 { - return f64::INFINITY; - } - f64::from(u32::try_from(self.tokens_used).unwrap_or(u32::MAX)) - / f64::from(u32::try_from(self.context_window).unwrap_or(u32::MAX)) - } -} - -impl fmt::Display for ContextOverflow { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "context overflow: {} tokens used of {} window ({} messages, {} overflow)", - self.tokens_used, - self.context_window, - self.message_count, - self.overflow() - ) - } -} - -impl std::error::Error for ContextOverflow {} - -// =================================================== -// EnsureContextResult -// =================================================== - -/// Result of [`ContextManager::ensure_context_fits`]. -/// -/// Tells the caller whether compaction occurred and provides the -/// (possibly compacted) message list. -#[derive(Debug, Clone)] -pub enum EnsureContextResult { - /// Compaction occurred and produced a shorter message list. - Compacted(CompactionOutcome), - /// No compaction was needed; messages returned as-is. - NoAction(Vec), -} - -impl EnsureContextResult { - /// Extract the message list from this result, regardless of variant. - #[must_use] - pub fn into_messages(self) -> Vec { - match self { - Self::Compacted(outcome) => outcome.messages, - Self::NoAction(messages) => messages, - } - } + /// This is the default. With the default `threshold = 0.80` and + /// `compact_target_pct = 0.70`, compaction targets 56% of the + /// context window (0.80 × 0.70 = 0.56). + Threshold, } // =================================================== @@ -727,14 +225,14 @@ impl EnsureContextResult { /// Manages context window usage and triggers compaction when needed. /// -/// This is the main entry point for context management. It monitors +/// Main entry point for context management. It monitors /// token usage, checks thresholds, and delegates to a pluggable /// [`ContextCompactor`] when compaction is needed. /// /// # Token Estimation /// /// Token counts are *estimates* using a 4-chars-per-token heuristic. -/// This is deliberately simple — the goal is to trigger compaction +/// Deliberately simple — the goal is to trigger compaction /// *before* hitting the actual limit, not to be perfectly accurate. /// Production systems should calibrate against their model's actual /// tokenizer. @@ -771,6 +269,10 @@ pub struct ContextManager { threshold: f64, /// Whether auto-compaction is enabled. auto_compact: bool, + /// The base used to compute the compaction target. + compact_base: CompactBase, + /// The fraction (0.0–1.0) of the target base to compact down to. + compact_target: f64, } impl ContextManager { @@ -778,11 +280,13 @@ impl ContextManager { /// /// Defaults: /// - /// | Setting | Default | - /// |------------------|---------| - /// | `context_window` | 200_000 | - /// | `threshold` | 0.80 | - /// | `auto_compact` | `true` | + /// | Setting | Default | + /// |----------------------|------------------------------| + /// | `context_window` | 200_000 | + /// | `threshold` | 0.80 | + /// | `auto_compact` | `true` | + /// | `compact_target` | [`CompactBase::Threshold`] | + /// | `compact_target_pct` | 0.70 | #[must_use] pub fn new(compactor: Arc) -> Self { Self { @@ -790,10 +294,15 @@ impl ContextManager { context_window: 200_000, threshold: 0.80, auto_compact: true, + compact_base: CompactBase::Threshold, + compact_target: 0.70, } } /// Set the model's context window size. + /// + /// Determines the upper bound on estimated tokens the manager + /// will allow before triggering compaction. #[must_use] pub fn with_context_window(mut self, tokens: u64) -> Self { self.context_window = tokens; @@ -810,48 +319,119 @@ impl ContextManager { } /// Set whether auto-compaction is enabled. + /// + /// When disabled, [`should_compact`](Self::should_compact) always + /// returns `false` and [`ensure_context_fits`](Self::ensure_context_fits) + /// will never trigger compaction. #[must_use] pub fn with_auto_compact(mut self, enabled: bool) -> Self { self.auto_compact = enabled; self } + /// Set the base used to compute the compaction target. + /// + /// See [`CompactBase`] for details. Defaults to + /// [`CompactBase::Threshold`]. + #[must_use] + pub fn with_compact_target(mut self, target: CompactBase) -> Self { + self.compact_base = target; + self + } + + /// Set the fraction (0.0–1.0) of the target base to compact down to. + /// + /// Clamped to `[0.1, 1.0]` to prevent degenerate configurations. + /// Defaults to `0.70` (70%). + #[must_use] + pub fn with_compact_target_pct(mut self, pct: f64) -> Self { + self.compact_target = pct.clamp(0.1, 1.0); + self + } + /// The model's context window size in tokens. + /// + /// This is the upper limit the manager uses to decide when compaction + /// is necessary. See [`with_context_window`](Self::with_context_window). #[must_use] pub fn context_window(&self) -> u64 { self.context_window } /// The compaction threshold (0.0–1.0). + /// + /// Compaction triggers when estimated tokens reach + /// `context_window * threshold`. See [`with_threshold`](Self::with_threshold). #[must_use] pub fn threshold(&self) -> f64 { self.threshold } /// Whether auto-compaction is enabled. + /// + /// When `false`, the manager never triggers compaction automatically. + /// See [`with_auto_compact`](Self::with_auto_compact). #[must_use] pub fn auto_compact(&self) -> bool { self.auto_compact } + /// The base used to compute the compaction target. + /// + /// See [`CompactBase`] and [`with_compact_target`](Self::with_compact_target). + #[must_use] + pub fn compact_target(&self) -> CompactBase { + self.compact_base + } + + /// The fraction (0.0–1.0) of the target base to compact down to. + /// + /// See [`with_compact_target_pct`](Self::with_compact_target_pct). + #[must_use] + pub fn compact_target_pct(&self) -> f64 { + self.compact_target + } + /// The token budget at which compaction triggers. /// /// Equal to `context_window * threshold`. The result is always /// non-negative (percentage × positive count), so the f64→u64 /// cast is safe in practice. #[must_use] - #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + #[allow( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + clippy::cast_precision_loss + )] pub fn compact_threshold_tokens(&self) -> u64 { - let threshold = - self.threshold * f64::from(u32::try_from(self.context_window).unwrap_or(u32::MAX)); - threshold as u64 + (self.threshold * self.context_window as f64) as u64 + } + + /// The token count to compact down to. + /// + /// Computed from [`compact_target`](Self::compact_target) and + /// [`compact_target_pct`](Self::compact_target_pct): + /// + /// - [`CompactBase::Threshold`]: `compact_threshold_tokens × pct` + /// - [`CompactBase::Context`]: `context_window × pct` + #[must_use] + #[allow( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + clippy::cast_precision_loss + )] + pub fn compact_target_tokens(&self) -> u64 { + let base: u64 = match self.compact_base { + CompactBase::Threshold => self.compact_threshold_tokens(), + CompactBase::Context => self.context_window, + }; + (self.compact_target * base as f64) as u64 } /// Estimate the token count for a slice of messages. /// /// Uses a 4-chars-per-token heuristic based on the text content - /// of all message parts. This is deliberately conservative — it - /// overestimates rather than underestimates. + /// of all message parts. Conservative — overestimates rather than underestimates. /// /// The estimation: /// - Counts text content from all parts (text, tool calls, tool results). @@ -867,15 +447,15 @@ impl ContextManager { let part_chars: u64 = m .parts .iter() - .map(|part| match part { - crate::message::MessagePart::Text { text } => text.len() as u64, - crate::message::MessagePart::Image { .. } => 256, // rough base64 estimate - crate::message::MessagePart::ToolCall { name, input, .. } => { + .map(|p| match p { + MessagePart::Text { text } => text.len() as u64, + MessagePart::Image { .. } => 256, // rough base64 estimate + MessagePart::ToolCall { name, input, .. } => { let name_len = name.len() as u64; let input_len = input.to_string().len() as u64; name_len.saturating_add(input_len) } - crate::message::MessagePart::ToolResult { output, .. } => match output { + MessagePart::ToolResult { output, .. } => match output { crate::message::ToolContent::Text(s) => s.len() as u64, crate::message::ToolContent::Multipart(parts) => parts .iter() @@ -911,6 +491,9 @@ impl ContextManager { } /// Check whether usage is in the emergency zone (>95% of window). + /// + /// Emergency compaction is more aggressive because the context is + /// dangerously close to overflowing the model's window. #[must_use] pub fn is_emergency(&self, used_tokens: u64) -> bool { let emergency_line = self.context_window.saturating_mul(19) / 20; // 95% @@ -918,6 +501,9 @@ impl ContextManager { } /// Determine the compaction reason for the given token count. + /// + /// Returns [`CompactReason::Emergency`] when usage exceeds 95% of + /// the window, or [`CompactReason::ThresholdExceeded`] otherwise. #[must_use] pub fn compact_reason(&self, used_tokens: u64) -> CompactReason { if self.is_emergency(used_tokens) { @@ -929,7 +515,7 @@ impl ContextManager { /// Ensure the conversation fits within the context window. /// - /// This is the main entry point called by the agent loop after each + /// Main entry point called by the agent loop after each /// turn. It: /// /// 1. Estimates the current token usage. @@ -947,14 +533,14 @@ impl ContextManager { turn: usize, ) -> Result { let tokens_before = Self::estimate_tokens(&messages); - let message_count = messages.len(); if !self.should_compact(tokens_before) { return Ok(EnsureContextResult::NoAction(messages)); } + let message_count = messages.len(); let reason = self.compact_reason(tokens_before); - let target_tokens = self.compact_threshold_tokens().saturating_mul(7) / 10; // compact to 70% of threshold + let target_tokens = self.compact_target_tokens(); let context = CompactionContext { tokens_before, reason, @@ -976,9 +562,8 @@ impl ContextManager { }); } - // Verify the compactor actually reduced the context. let tokens_after = Self::estimate_tokens(&outcome.messages); - if tokens_after > self.context_window && self.is_emergency(tokens_after) { + if tokens_after > self.context_window { return Err(ContextOverflow { tokens_used: tokens_after, context_window: self.context_window, @@ -1005,16 +590,13 @@ impl ContextManager { turn: usize, ) -> Result { let tokens_before = Self::estimate_tokens(&messages); - let message_count = messages.len(); if messages.is_empty() { return Ok(EnsureContextResult::NoAction(messages)); } - let target_tokens = self - .compact_threshold_tokens() - .saturating_mul(7) - .saturating_div(10); + let message_count = messages.len(); + let target_tokens = self.compact_target_tokens(); let context = CompactionContext { tokens_before, reason: CompactReason::Manual, @@ -1107,8 +689,10 @@ impl fmt::Debug for ContextManager { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ContextManager") .field("context_window", &self.context_window) - .field("threshold", &self.threshold) + .field("threshold", &self.threshold()) .field("auto_compact", &self.auto_compact) + .field("compact_target", &self.compact_base) + .field("compact_target_pct", &self.compact_target) .finish_non_exhaustive() } } @@ -1253,6 +837,64 @@ mod tests { assert_eq!(manager.compact_threshold_tokens(), 160_000); } + #[test] + fn test_compact_target_tokens_default() { + // Default: CompactBase::Threshold with pct 0.70 + // threshold = 200_000 * 0.80 = 160_000 + // target = 160_000 * 0.70 = 112_000 + let compactor = TruncatingCompactor::new(); + let manager = ContextManager::new(Arc::new(compactor)) + .with_context_window(200_000) + .with_threshold(0.80); + assert_eq!(manager.compact_target_tokens(), 112_000); + } + + #[test] + fn test_compact_target_tokens_context_base() { + // CompactBase::Context with pct 0.50 + // target = 200_000 * 0.50 = 100_000 + let compactor = TruncatingCompactor::new(); + let manager = ContextManager::new(Arc::new(compactor)) + .with_context_window(200_000) + .with_threshold(0.80) + .with_compact_target(CompactBase::Context) + .with_compact_target_pct(0.50); + assert_eq!(manager.compact_target_tokens(), 100_000); + } + + #[test] + fn test_compact_target_tokens_threshold_base() { + // CompactBase::Threshold with pct 0.50 + // threshold = 200_000 * 0.80 = 160_000 + // target = 160_000 * 0.50 = 80_000 + let compactor = TruncatingCompactor::new(); + let manager = ContextManager::new(Arc::new(compactor)) + .with_context_window(200_000) + .with_threshold(0.80) + .with_compact_target(CompactBase::Threshold) + .with_compact_target_pct(0.50); + assert_eq!(manager.compact_target_tokens(), 80_000); + } + + #[test] + fn test_compact_target_pct_clamped() { + let compactor = TruncatingCompactor::new(); + let manager = ContextManager::new(Arc::new(compactor)).with_compact_target_pct(0.01); + assert!((manager.compact_target_pct() - 0.1).abs() < f64::EPSILON); + + let compactor2 = TruncatingCompactor::new(); + let manager2 = ContextManager::new(Arc::new(compactor2)).with_compact_target_pct(2.0); + assert!((manager2.compact_target_pct() - 1.0).abs() < f64::EPSILON); + } + + #[test] + fn test_compact_target_default_is_threshold() { + let compactor = TruncatingCompactor::new(); + let manager = ContextManager::new(Arc::new(compactor)); + assert_eq!(manager.compact_target(), CompactBase::Threshold); + assert!((manager.compact_target_pct() - 0.70).abs() < f64::EPSILON); + } + #[tokio::test] async fn test_truncating_compactor_no_change() { let compactor = TruncatingCompactor::new() @@ -1268,7 +910,6 @@ mod tests { let outcome = compactor.compact(msgs.clone(), 500, context).await; assert!(outcome.success); assert_eq!(outcome.messages.len(), msgs.len()); - assert_eq!(outcome.messages_compacted, 0); } #[tokio::test] @@ -1289,7 +930,6 @@ mod tests { assert!(outcome.success); // 1 (first/system prompt) + 2 (preserve_recent) = 3 messages preserved. assert_eq!(outcome.messages.len(), 3); - assert_eq!(outcome.messages_compacted, 17); assert!(outcome.tokens_saved > 0); // Verify the first message was preserved. assert!(outcome.messages.first().is_some()); @@ -1428,7 +1068,6 @@ mod tests { let post = make_conversation(2); let outcome = CompactionOutcome { messages: post.clone(), - messages_compacted: 16, tokens_after: 100, tokens_saved: 800, success: true, @@ -1483,7 +1122,6 @@ mod tests { assert!(outcome.success); assert_eq!(outcome.tokens_after, 400); assert_eq!(outcome.tokens_saved, 600); // 1000 - 400 - assert_eq!(outcome.messages_compacted, 0); // caller sets assert!(outcome.error.is_none()); } diff --git a/src/compact/truncating.rs b/src/compact/truncating.rs new file mode 100644 index 0000000..74641b5 --- /dev/null +++ b/src/compact/truncating.rs @@ -0,0 +1,324 @@ +//! Truncating compactor and token splitter. +//! +//! Contents: +//! +//! - [`TruncatingCompactor`] — a simple compactor that drops the oldest messages. +//! - [`TokenSplitter`] — splits a conversation into "old" and "recent" at a turn boundary. +//! - [`SplitResult`] — result of splitting a conversation. + +use crate::compact::types::{CompactionContext, CompactionOutcome}; +use crate::compact::{ContextCompactor, ContextManager}; +use crate::message::{Message, Role}; +use std::future::Future; +use std::pin::Pin; + +// =================================================== +// TruncatingCompactor +// =================================================== + +/// A simple compactor that drops the oldest messages. +/// +/// Keeps the first message (typically the system prompt) and a configurable +/// number of recent messages. No LLM calls required — useful as a fallback +/// or for contexts where summarization isn't available. +/// +/// # Strategy +/// +/// ```text +/// [System?] [Old₁, Old₂, ..., Oldₙ] [Recent₁, Recent₂, ..., Recentₘ] +/// ↑ kept ↑ discarded ↑ ↑ preserved ↑ +/// ``` +/// +/// The first message is always retained (if present) because it usually +/// contains the system prompt or conversation instructions. This prevents +/// the compactor from discarding essential context that shapes the agent's +/// behavior. If the conversation is shorter than `min_messages`, no +/// compaction occurs. +/// +/// # Example +/// +/// ```rust +/// use loopctl::compact::TruncatingCompactor; +/// use std::sync::Arc; +/// +/// let compactor = TruncatingCompactor::new() +/// .with_preserve_recent(6) +/// .with_min_messages(8); +/// +/// // Pass to ContextManager: +/// // let manager = ContextManager::new(Arc::new(compactor)); +/// ``` +#[derive(Debug, Clone)] +pub struct TruncatingCompactor { + /// Number of recent messages to always preserve. + preserve_recent: usize, + /// Minimum messages before compaction is considered. + min_messages: usize, +} + +impl TruncatingCompactor { + /// Create a new truncating compactor with sensible defaults. + /// + /// Defaults: + /// + /// | Setting | Default | + /// |-------------------|---------| + /// | `preserve_recent` | 4 | + /// | `min_messages` | 6 | + #[must_use] + pub fn new() -> Self { + Self { + preserve_recent: 4, + min_messages: 6, + } + } + + /// Set how many recent messages to preserve during compaction. + /// + /// This many messages from the end of the conversation are kept + /// intact. The rest are dropped. Must be at least 1. + #[must_use] + pub fn with_preserve_recent(mut self, count: usize) -> Self { + self.preserve_recent = count.max(1); + self + } + + /// Set the minimum number of messages before compaction is attempted. + /// + /// If the conversation has fewer messages than this, compaction is + /// skipped entirely. Prevents aggressive truncation of short + /// conversations. + #[must_use] + pub fn with_min_messages(mut self, count: usize) -> Self { + self.min_messages = count.max(2); + self + } + + /// Number of recent messages that will be preserved. + #[must_use] + pub fn preserve_recent(&self) -> usize { + self.preserve_recent + } + + /// Minimum messages before compaction is attempted. + #[must_use] + pub fn min_messages(&self) -> usize { + self.min_messages + } +} + +impl Default for TruncatingCompactor { + fn default() -> Self { + Self::new() + } +} + +impl ContextCompactor for TruncatingCompactor { + fn compact( + &self, + messages: Vec, + _target_tokens: u64, + context: CompactionContext, + ) -> Pin + Send + '_>> { + Box::pin(async move { + let total = messages.len(); + if total <= self.min_messages { + return CompactionOutcome::no_change(messages); + } + + // Determine split point: keep `preserve_recent` from the end. + let split = total.saturating_sub(self.preserve_recent); + let recent: Vec = messages.get(split..).unwrap_or_default().to_vec(); + + // Always preserve the first message (typically the system prompt) + // unless it is already included in the recent slice (split == 0). + let preserved = if split > 0 { + if let Some(first) = messages.first() { + let mut v = vec![first.clone()]; + v.extend(recent); + v + } else { + recent + } + } else { + // split == 0 means recent already contains all messages. + recent + }; + + let tokens_after = CompactionOutcome::estimate_tokens(&preserved); + CompactionOutcome { + messages: preserved, + tokens_after, + tokens_saved: context.tokens_before.saturating_sub(tokens_after), + success: true, + error: None, + } + }) + } +} + +// =================================================== +// TokenSplitter +// =================================================== + +/// Splits a conversation into "old" and "recent" at a turn boundary. +/// +/// Used by compactors (and agent-side code) that need to know which +/// messages to compact versus preserve. Splits at role transitions for +/// coherent summarization — the split always occurs between a complete +/// request/response pair. +/// +/// # Rules +/// +/// - Never compact the last user message (it's the current request). +/// - Split at turn boundaries (role transitions) for coherent output. +/// - If the conversation is too short, `to_compact` will be empty. +/// +/// # Example +/// +/// ```rust +/// use loopctl::compact::TokenSplitter; +/// use loopctl::message::Message; +/// +/// let splitter = TokenSplitter::new() +/// .with_preserve_recent(4) +/// .with_min_messages(6); +/// +/// let messages = vec![ +/// Message::user("Hello"), +/// Message::assistant("Hi there!"), +/// Message::user("What is 2+2?"), +/// Message::assistant("4"), +/// ]; +/// +/// let result = splitter.split(&messages); +/// // With only 4 messages and min_messages=6, nothing is split off. +/// assert!(result.to_compact.is_empty()); +/// assert_eq!(result.preserved.len(), 4); +/// ``` +#[derive(Debug, Clone)] +pub struct TokenSplitter { + /// Number of recent messages to always preserve. + preserve_recent: usize, + /// Minimum messages before considering a split. + min_messages: usize, +} + +/// Result of splitting a conversation into old and recent portions. +#[derive(Debug, Clone)] +pub struct SplitResult { + /// Messages to compact or summarize (the "old" part). + pub to_compact: Vec, + /// Messages to preserve as-is (the "recent" part). + pub preserved: Vec, + /// Estimated tokens in `to_compact`. + pub compact_tokens: u64, + /// Estimated tokens in `preserved`. + pub preserved_tokens: u64, + /// The index in the original message list where the split occurred. + pub split_index: usize, +} + +impl TokenSplitter { + /// Create a new splitter with sensible defaults. + /// + /// Defaults: + /// + /// | Setting | Default | + /// |-------------------|---------| + /// | `preserve_recent` | 4 | + /// | `min_messages` | 6 | + #[must_use] + pub fn new() -> Self { + Self { + preserve_recent: 4, + min_messages: 6, + } + } + + /// Set how many recent messages to preserve. + #[must_use] + pub fn with_preserve_recent(mut self, count: usize) -> Self { + self.preserve_recent = count.max(1); + self + } + + /// Set the minimum messages before splitting is considered. + #[must_use] + pub fn with_min_messages(mut self, count: usize) -> Self { + self.min_messages = count.max(2); + self + } + + /// Split the given messages into old and recent portions. + /// + /// The split point is chosen at a turn boundary (a role transition + /// from assistant to user) as close as possible to leaving + /// `preserve_recent` messages in the recent portion. + /// + /// If the conversation has fewer than `min_messages`, the entire + /// conversation goes into `preserved` and `to_compact` is empty. + #[must_use] + pub fn split(&self, messages: &[Message]) -> SplitResult { + if messages.len() <= self.min_messages { + return SplitResult { + to_compact: vec![], + preserved: messages.to_vec(), + compact_tokens: 0, + preserved_tokens: ContextManager::estimate_tokens(messages), + split_index: 0, + }; + } + + // Find a split point: we want `preserve_recent` messages at the end. + // Look for a turn boundary (assistant→user transition) near the + // target split point. + let target_split = messages.len().saturating_sub(self.preserve_recent); + let split_index = Self::find_turn_boundary(messages, target_split); + let (to_compact, preserved) = messages.split_at(split_index); + SplitResult { + to_compact: to_compact.to_vec(), + preserved: preserved.to_vec(), + compact_tokens: ContextManager::estimate_tokens(to_compact), + preserved_tokens: ContextManager::estimate_tokens(preserved), + split_index, + } + } + + /// Find the nearest turn boundary at or before the target index. + /// + /// A turn boundary is a position where the previous message is + /// assistant-role and the next is user-role. This ensures we split + /// at a coherent conversation boundary. + fn find_turn_boundary(messages: &[Message], target: usize) -> usize { + if target == 0 { + return 0; + } + + // Search backwards from target for an assistant→user transition. + for i in (1..=target).rev() { + if i < messages.len() { + let Some(prev) = messages.get(i.saturating_sub(1)) else { + continue; + }; + let Some(curr) = messages.get(i) else { + continue; + }; + let prev_is_assistant = prev.role == Role::Assistant; + let curr_is_user = curr.role == Role::User; + if prev_is_assistant && curr_is_user { + return i; + } + } + } + + // Fallback: no clean boundary found, split at target. + target + } +} + +impl Default for TokenSplitter { + fn default() -> Self { + Self::new() + } +} diff --git a/src/compact/types.rs b/src/compact/types.rs new file mode 100644 index 0000000..daf47ff --- /dev/null +++ b/src/compact/types.rs @@ -0,0 +1,262 @@ +//! Supporting types for context compaction. +//! +//! Data types used across the compaction pipeline: +//! +//! - [`CompactReason`] — why compaction was triggered. +//! - [`CompactionContext`] — input metadata passed to compactors. +//! - [`CompactionOutcome`] — result of a single compaction pass. +//! - [`CompactTelemetry`] — telemetry data for compaction operations. +//! - [`PreCompactStats`] / [`PostCompactStats`] — stats before/after compaction. +//! - [`ContextOverflow`] — error when the conversation cannot fit. +//! - [`EnsureContextResult`] — result of [`ContextManager::ensure_context_fits`](super::ContextManager::ensure_context_fits). + +use crate::message::Message; +use std::fmt; + +// =================================================== +// CompactReason +// =================================================== + +/// Why compaction was triggered. +/// +/// Different triggers may warrant different compaction strategies. +/// For example, an [`Emergency`](CompactReason::Emergency) compaction +/// should be more aggressive than a routine threshold check. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompactReason { + /// Token usage exceeded the configured threshold percentage. + ThresholdExceeded, + /// Token usage is dangerously close to the context window limit. + Emergency, + /// Compaction was explicitly requested (e.g. by the agent or a tool). + Manual, +} + +impl fmt::Display for CompactReason { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::ThresholdExceeded => write!(f, "threshold exceeded"), + Self::Emergency => write!(f, "emergency"), + Self::Manual => write!(f, "manual"), + } + } +} + +// =================================================== +// CompactionContext +// =================================================== + +/// Metadata passed to [`ContextCompactor::compact`](super::ContextCompactor::compact) +/// describing the compaction trigger and current state. +/// +/// Compactors can use this information to decide how aggressively to +/// compact — e.g. an emergency compaction may use more aggressive +/// summarization than a routine threshold check. +#[derive(Debug, Clone)] +pub struct CompactionContext { + /// Estimated token count before compaction. + pub tokens_before: u64, + /// Why compaction was triggered. + pub reason: CompactReason, + /// The model's context window size. + pub context_window: u64, + /// The current turn number in the session. + pub turn: usize, +} + +// =================================================== +// CompactionOutcome +// =================================================== + +/// Result of a single compaction pass. +/// +/// Returned by [`ContextCompactor::compact`](super::ContextCompactor::compact), +/// this struct contains the compacted message list along with telemetry data +/// about what happened. +#[derive(Debug, Clone)] +pub struct CompactionOutcome { + /// The compacted message list. + pub messages: Vec, + /// Estimated token count after compaction. + pub tokens_after: u64, + /// Estimated tokens saved by compaction. + pub tokens_saved: u64, + /// Whether compaction succeeded. + pub success: bool, + /// Error message if compaction failed. + pub error: Option, +} + +impl CompactionOutcome { + /// Create an outcome representing no change (compaction was not needed). + /// + /// Use this when the compactor decides the messages don't need + /// compaction — e.g. when the message count is below the minimum. + #[must_use] + pub fn no_change(messages: Vec) -> Self { + let tokens = Self::estimate_tokens(&messages); + Self { + messages, + tokens_after: tokens, + tokens_saved: 0, + success: true, + error: None, + } + } + + /// Create an outcome representing successful compaction. + /// + /// Computes [`tokens_saved`](Self::tokens_saved) automatically from the + /// difference between `tokens_before` and `tokens_after`. + #[must_use] + pub fn compacted(messages: Vec, tokens_before: u64, tokens_after: u64) -> Self { + Self { + tokens_saved: tokens_before.saturating_sub(tokens_after), + messages, + tokens_after, + success: true, + error: None, + } + } + + /// Estimate the token count for a slice of messages. + /// + /// Uses the standard 4-chars-per-token heuristic, the same + /// heuristic used by [`ContextManager::estimate_tokens`](super::ContextManager::estimate_tokens). + #[must_use] + pub fn estimate_tokens(messages: &[Message]) -> u64 { + super::ContextManager::estimate_tokens(messages) + } +} + +// =================================================== +// CompactTelemetry +// =================================================== + +/// Telemetry data for a single compaction operation. +/// +/// Produced by [`ContextManager::ensure_context_fits`](super::ContextManager::ensure_context_fits) +/// when compaction occurs. Observers receive this via +/// [`on_compaction`](crate::observer::LoopObserver::on_compaction). +#[derive(Debug, Clone)] +pub struct CompactTelemetry { + /// Why compaction was triggered. + pub trigger: CompactReason, + /// Conversation stats before compaction. + pub pre_compact: PreCompactStats, + /// Conversation stats after compaction. + pub post_compact: PostCompactStats, + /// Wall-clock duration of the compaction. + pub duration: std::time::Duration, +} + +/// Conversation statistics captured before compaction. +#[derive(Debug, Clone)] +pub struct PreCompactStats { + /// Total number of messages in the conversation. + pub total_messages: usize, + /// Estimated token count. + pub estimated_tokens: u64, + /// Number of user-role messages. + pub user_messages: usize, + /// Number of assistant-role messages. + pub assistant_messages: usize, + /// Number of messages containing tool calls or results. + pub tool_messages: usize, +} + +/// Conversation statistics captured after compaction. +#[derive(Debug, Clone)] +pub struct PostCompactStats { + /// Total number of messages after compaction. + pub total_messages: usize, + /// Estimated token count after compaction. + pub estimated_tokens: u64, + /// Tokens removed by compaction. + pub tokens_saved: u64, + /// Percentage of tokens saved (0–100). + pub percent_saved: u8, +} + +// =================================================== +// ContextOverflow error +// =================================================== + +/// Error returned when the conversation cannot fit within the context +/// window, even after compaction. +/// +/// Terminal condition — the conversation is too large and the +/// compactor was unable to reduce it sufficiently. +#[derive(Debug, Clone)] +pub struct ContextOverflow { + /// Estimated token count of the conversation. + pub tokens_used: u64, + /// The model's context window size. + pub context_window: u64, + /// How many messages were in the conversation. + pub message_count: usize, + /// The reason compaction was attempted. + pub trigger: CompactReason, + /// Error from the compactor, if compaction was attempted. + pub compactor_error: Option, +} + +impl ContextOverflow { + /// How many tokens the conversation exceeds the window by. + #[must_use] + pub fn overflow(&self) -> u64 { + self.tokens_used.saturating_sub(self.context_window) + } + + /// The fraction of the context window used (0.0–1.0+). + #[must_use] + pub fn utilization(&self) -> f64 { + if self.context_window == 0 { + return f64::INFINITY; + } + f64::from(u32::try_from(self.tokens_used).unwrap_or(u32::MAX)) + / f64::from(u32::try_from(self.context_window).unwrap_or(u32::MAX)) + } +} + +impl fmt::Display for ContextOverflow { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "context overflow: {} tokens used of {} window ({} messages, {} overflow)", + self.tokens_used, + self.context_window, + self.message_count, + self.overflow() + ) + } +} + +impl std::error::Error for ContextOverflow {} + +// =================================================== +// EnsureContextResult +// =================================================== + +/// Result of [`ContextManager::ensure_context_fits`](super::ContextManager::ensure_context_fits). +/// +/// Tells the caller whether compaction occurred and provides the +/// (possibly compacted) message list. +#[derive(Debug, Clone)] +pub enum EnsureContextResult { + /// Compaction occurred and produced a shorter message list. + Compacted(CompactionOutcome), + /// No compaction was needed; messages returned as-is. + NoAction(Vec), +} + +impl EnsureContextResult { + /// Extract the message list from this result, regardless of variant. + #[must_use] + pub fn into_messages(self) -> Vec { + match self { + Self::Compacted(outcome) => outcome.messages, + Self::NoAction(messages) => messages, + } + } +} diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..5729f8c --- /dev/null +++ b/src/config.rs @@ -0,0 +1,87 @@ +//! Agent session configuration. +//! +//! Defines [`LoopConfig`] — the configuration struct that controls agent +//! session parameters such as turn limits, model selection, and context +//! window size. + +use uuid::Uuid; + +/// Configuration for an agent session. +/// +/// Holds generic agent configuration fields that apply to every session +/// regardless of the specific agent type. Domain-specific configuration +/// (e.g., ITR engine settings, `ToolShield` rules, fallback model chains) +/// should live in production-specific config types that embed or wrap +/// this struct. +/// +/// # Construction +/// +/// Use [`LoopConfig::default`] for sensible defaults or override individual +/// fields through the builder via +/// `LoopBuilder::with_config`. +/// +/// ``` +/// use loopctl::config::LoopConfig; +/// +/// let config = LoopConfig { +/// max_turns: 50, +/// model: "default".to_string(), +/// ..Default::default() +/// }; +/// ``` +#[derive(Debug, Clone)] +pub struct LoopConfig { + /// Unique session identifier (random UUID v4). + pub session_id: Uuid, + /// Model identifier (e.g. `"default"`). Passed to the API client on each request. + pub model: String, + /// Optional system prompt override. `None` means the agent core decides. + pub system_prompt: Option, + /// Maximum number of turns before forcing completion. Defaults to `200`. + pub max_turns: usize, + /// Maximum tokens for each API response. Defaults to `16_384`. + pub max_tokens: u32, + /// Context window size in tokens. Must match the actual window of [`model`](LoopConfig::model). Defaults to `200_000`. + pub context_window: u64, + /// Threshold to trigger auto-compaction (0.0–1.0). Defaults to `0.80`. + pub compact_threshold: f64, + /// Whether auto-compaction is enabled. Defaults to `true`. + pub auto_compact: bool, +} + +impl Default for LoopConfig { + /// Produce a configuration with production-ready defaults. + /// + /// | Field | Default | + /// |-------|---------| + /// | [`session_id`](LoopConfig::session_id) | Random UUID v4 | + /// | [`model`](LoopConfig::model) | `"default"` | + /// | [`system_prompt`](LoopConfig::system_prompt) | `None` | + /// | [`max_turns`](LoopConfig::max_turns) | `200` | + /// | [`max_tokens`](LoopConfig::max_tokens) | `16_384` | + /// | [`context_window`](LoopConfig::context_window) | `200_000` | + /// | [`compact_threshold`](LoopConfig::compact_threshold) | `0.80` | + /// | [`auto_compact`](LoopConfig::auto_compact) | `true` | + /// + /// # Example + /// + /// ``` + /// use loopctl::config::LoopConfig; + /// + /// let config = LoopConfig::default(); + /// assert_eq!(config.max_turns, 200); + /// assert_eq!(config.model, "default"); + /// ``` + fn default() -> Self { + Self { + session_id: Uuid::new_v4(), + model: "default".to_string(), + system_prompt: None, + max_turns: 200, + max_tokens: 16_384, + context_window: 200_000, + compact_threshold: 0.80, + auto_compact: true, + } + } +} diff --git a/src/core.rs b/src/core.rs deleted file mode 100644 index 36a0149..0000000 --- a/src/core.rs +++ /dev/null @@ -1,59 +0,0 @@ -//! Core module — foundational traits and types for building agents. -//! -//! This module defines the core abstractions that every agent built with -//! the loopctl framework depends on. Implement these traits to plug in -//! your own agent logic, memory backend, or observability layer. -//! -//! # Traits -//! -//! | Trait | Purpose | -//! |------------------------------------------|-----------------------------------------------------| -//! | [`AgentCore`] | Main lifecycle trait for all agent types | -//! | [`AgentMemory`] | Interface for agent memory backends | -//! | [`LoopObserver`](observer::LoopObserver) | Notification-only observer for agent lifecycle | -//! -//! # Supporting Types -//! -//! | Type | Purpose | -//! |--------------------------------|------------------------------------------------------| -//! | [`AgentConfig`] | Configuration for an agent session | -//! | [`AgentError`] | Unified error type for all framework operations | -//! | [`AgentState`] | Lifecycle state machine for agents | -//! | [`Correction`] | Correction produced by the reflection system | -//! | [`CompactReason`] | Why context compaction was triggered | -//! | [`CorrectionResult`] | Outcome of applying a correction | -//! | [`ConsolidationStats`] | Statistics from a memory consolidation pass | -//! | [`CorrectionType`] | Category of fix strategy for a correction | -//! | [`ExponentialBackoffRecovery`] | Retry strategy with exponential backoff | -//! | [`FailureAnalysis`] | Analysis of a failed tool call | -//! | [`FailureSeverity`] | How severe a failure is | -//! | [`MemoryCategory`] | Category of a memory entry | -//! | [`MemoryEntry`] | A single memory entry with metadata | -//! | [`NoopReflector`] | Default reflector — marks everything non-recoverable | -//! | [`RecoveryAction`] | What the framework should do after a failure | -//! | [`ReflectionContext`] | Retry state provided to the reflector | -//! | [`StopReason`] | Why the API stopped generating | -//! | [`SessionResult`] | Summary of a complete agent session | -//! | [`ToolCall`] | A tool call requested by the agent | -//! | [`ToolDispatchResult`] | Result of a single tool execution | -//! | [`TurnResult`] | Result of a single agent turn | -//! -//! # Sub-modules -//! -//! - **[`observer`]** — [`LoopObserver`](observer::LoopObserver) trait, context structs, and -//! [`ObserverHost`](observer::ObserverHost) for lifecycle notification. Observers are passive — -//! they receive callbacks but cannot control flow. For flow control, see the -//! [hook system](crate::hooks). - -pub mod agent_core; -pub mod agent_memory; -pub mod error; -pub mod observer; -pub mod reflection; -pub mod types; - -pub use agent_core::*; -pub use agent_memory::*; -pub use error::*; -pub use reflection::*; -pub use types::*; diff --git a/src/core/agent_core.rs b/src/core/agent_core.rs deleted file mode 100644 index 3bfbe74..0000000 --- a/src/core/agent_core.rs +++ /dev/null @@ -1,128 +0,0 @@ -//! Agent core trait — the main lifecycle interface for agents. -//! -//! This trait defines the fundamental operations every agent must support. -//! Different agent types (chat, coding, research) implement this trait -//! while sharing the framework's infrastructure (managers, observers, etc.). -//! -//! # Lifecycle -//! -//! ```text -//! initialize(config) -//! → process_turn(input) [repeated] -//! → process_turn(input) -//! → ... -//! → should_continue() → false -//! finalize() -//! ``` -//! -//! # Implementing -//! -//! At a minimum you must provide [`initialize`](AgentCore::initialize), -//! [`process_turn`](AgentCore::process_turn), -//! [`should_continue`](AgentCore::should_continue), -//! [`finalize`](AgentCore::finalize), -//! [`state`](AgentCore::state), and -//! [`cancel`](AgentCore::cancel). - -use crate::core::error::AgentError; -use crate::core::types::{AgentConfig, AgentState, SessionResult, TurnResult}; -use std::future::Future; -use std::pin::Pin; - -/// The core agent lifecycle trait. -/// -/// Implement this trait to create a new type of agent. The framework -/// provides shared infrastructure for context management, tool execution, -/// reflection, and observability, so implementations only need to define -/// the core processing logic. -/// -/// # Lifecycle -/// -/// ```text -/// initialize(config) -/// → process_turn(input) [repeated] -/// → process_turn(input) -/// → ... -/// → should_continue() → false -/// finalize() -/// ``` -/// -/// # Example -/// -/// ```rust,ignore -/// use loopctl::core::agent_core::AgentCore; -/// use loopctl::core::error::AgentError; -/// use loopctl::core::types::{AgentConfig, AgentState, SessionResult, TurnResult}; -/// -/// struct MyAgent { -/// state: MyState, -/// } -/// -/// impl AgentCore for MyAgent { -/// fn initialize<'a>(&'a mut self, config: &'a AgentConfig) -> Pin> + Send + 'a>> { -/// Box::pin(async { Ok(()) }) -/// } -/// fn process_turn<'a>(&'a mut self, input: &'a str) -> Pin> + Send + 'a>> { -/// Box::pin(async { Ok(TurnResult::completed("Done!")) }) -/// } -/// fn should_continue(&self) -> bool { -/// !self.state.is_complete -/// } -/// fn finalize<'a>(&'a mut self) -> Pin> + Send + 'a>> { -/// Box::pin(async { Ok(SessionResult::success(self.state.session_id)) }) -/// } -/// fn state(&self) -> AgentState { -/// AgentState::Idle -/// } -/// fn cancel(&self) {} -/// } -/// ``` -pub trait AgentCore: Send + Sync { - /// Initialize the agent with the given configuration. - /// - /// Called once before any turns are processed. Use this to set up - /// internal state, validate configuration, and prepare resources. - fn initialize<'a>( - &'a mut self, - config: &'a AgentConfig, - ) -> Pin> + Send + 'a>>; - - /// Process a single user message / turn. - /// - /// This is the main entry point for agent logic. It receives the user's - /// input and returns a [`TurnResult`] describing what happened. - fn process_turn<'a>( - &'a mut self, - input: &'a str, - ) -> Pin> + Send + 'a>>; - - /// Check whether the agent should continue processing turns. - /// - /// Called after each turn. Return `false` to end the session. - fn should_continue(&self) -> bool; - - /// Finalize the agent session and produce a summary. - /// - /// Called once after the last turn. Use this to clean up resources - /// and produce a final [`SessionResult`]. - fn finalize<'a>( - &'a mut self, - ) -> Pin> + Send + 'a>>; - - /// Get the current state of the agent. - /// - /// Used by the framework to drive the state machine and by observers - /// to report status. - fn state(&self) -> AgentState; - - /// Cancel the agent's current operation. - /// - /// Implementations must use thread-safe interior mutability (e.g. - /// [`AtomicBool`](std::sync::atomic::AtomicBool), `Mutex`) to - /// store the cancellation flag, since this method takes `&self`. The - /// flag should be set in a non-blocking fashion so that - /// [`process_turn`](AgentCore::process_turn) and - /// [`should_continue`](AgentCore::should_continue) can observe it - /// and return promptly across threads. - fn cancel(&self); -} diff --git a/src/core/agent_memory.rs b/src/core/agent_memory.rs deleted file mode 100644 index f7ffe60..0000000 --- a/src/core/agent_memory.rs +++ /dev/null @@ -1,438 +0,0 @@ -//! Agent memory trait — interface for agent memory systems. -//! -//! Memory allows agents to learn from past interactions and retrieve -//! relevant context for future tasks. This module defines the core -//! [`AgentMemory`] trait that all memory backends implement, along with -//! the [`MemoryEntry`] value type and supporting enumerations. -//! -//! # Provided Implementations -//! -//! - **`TrajectoryMemory`** — Records tool-execution trajectories and -//! retrieves relevant past experiences. -//! -//! # Quick Start -//! -//! ``` -//! use loopctl::core::agent_memory::{AgentMemory, MemoryEntry, MemoryCategory, ConsolidationStats}; -//! use loopctl::core::error::AgentError; -//! -//! struct InMemoryStore { -//! entries: Vec, -//! } -//! -//! impl AgentMemory for InMemoryStore { -//! async fn store(&mut self, entry: MemoryEntry) -> Result<(), AgentError> { -//! self.entries.push(entry); -//! Ok(()) -//! } -//! async fn retrieve(&self, query: &str, limit: usize) -> Result, AgentError> { -//! Ok(self.entries.iter() -//! .filter(|e| e.memory.contains(query)) -//! .take(limit) -//! .cloned() -//! .collect()) -//! } -//! async fn consolidate(&mut self) -> Result { -//! Ok(ConsolidationStats::default()) -//! } -//! fn len(&self) -> usize { -//! self.entries.len() -//! } -//! } -//! ``` - -use crate::core::error::AgentError; -use serde::{Deserialize, Serialize}; -use std::time::SystemTime; - -/// A memory system for agents. -/// -/// Implementations can store and retrieve entries using different -/// strategies (vector similarity, keyword matching, recency, etc.). -/// -/// # Lifecycle -/// -/// ```text -/// store(entry) [called whenever the agent learns something] -/// → retrieve(query, limit) [called before each turn to gather context] -/// → retrieve(query, limit) -/// → ... -/// → consolidate() [called periodically to prune / compress] -/// ``` -/// -/// # Implementing -/// -/// At a minimum you must provide [`store`](AgentMemory::store), -/// [`retrieve`](AgentMemory::retrieve), [`consolidate`](AgentMemory::consolidate), -/// and [`len`](AgentMemory::len). The trait supplies a default -/// [`is_empty`](AgentMemory::is_empty) implementation that delegates to `len`. -/// -/// # Example -/// -/// ``` -/// use loopctl::core::agent_memory::{AgentMemory, MemoryEntry, MemoryCategory, ConsolidationStats}; -/// use loopctl::core::error::AgentError; -/// -/// struct InMemoryStore { -/// entries: Vec, -/// } -/// -/// impl AgentMemory for InMemoryStore { -/// async fn store(&mut self, entry: MemoryEntry) -> Result<(), AgentError> { -/// self.entries.push(entry); -/// Ok(()) -/// } -/// async fn retrieve(&self, query: &str, limit: usize) -> Result, AgentError> { -/// Ok(self.entries.iter() -/// .filter(|e| e.memory.contains(query)) -/// .take(limit) -/// .cloned() -/// .collect()) -/// } -/// async fn consolidate(&mut self) -> Result { -/// let before = self.entries.len(); -/// self.entries.retain(|e| e.relevance > 0.1); -/// let after = self.entries.len(); -/// Ok(ConsolidationStats { -/// entries_before: before, -/// entries_after: after, -/// pruned: before - after, -/// ..Default::default() -/// }) -/// } -/// fn len(&self) -> usize { -/// self.entries.len() -/// } -/// } -/// ``` -#[allow(async_fn_in_trait)] -pub trait AgentMemory: Send + Sync { - /// Store a new memory entry. - /// - /// Called whenever the agent encounters information worth remembering — - /// for example after a successful tool invocation, a resolved error, or - /// an insight drawn from conversation. Implementations should persist the - /// entry in whatever backing store they use. - async fn store(&mut self, entry: MemoryEntry) -> Result<(), AgentError>; - - /// Retrieve memory entries relevant to the given query. - /// - /// Called before each turn (or on demand) to surface context the agent - /// can use. Returns up to `limit` entries ordered by relevance. The - /// definition of "relevance" is left to the implementation — common - /// strategies include vector embedding similarity, keyword overlap, - /// recency weighting, or a hybrid approach. - /// - /// Implementations that track [`MemoryEntry::access_count`] must use - /// interior mutability (e.g. `AtomicUsize`, `Mutex`) since this method - /// takes `&self`. - async fn retrieve(&self, query: &str, limit: usize) -> Result, AgentError>; - - /// Consolidate memory (e.g. prune, summarize, compress). - /// - /// Called periodically to keep the memory store healthy. Implementations - /// may remove low-relevance entries, merge duplicates, or produce - /// compressed summaries. Returns [`ConsolidationStats`] describing what - /// was done. - async fn consolidate(&mut self) -> Result; - - /// Number of entries currently stored. - /// - /// Used by the framework and by [`is_empty`](AgentMemory::is_empty). - fn len(&self) -> usize; - - /// Whether the memory is empty. - /// - /// Defaults to `self.len() == 0`. Override only if you need a cheaper - /// check than counting all entries. - fn is_empty(&self) -> bool { - self.len() == 0 - } -} - -/// A single memory entry. -/// -/// Each entry represents one discrete piece of information the agent has -/// learned. Entries carry metadata — category, tags, relevance score, -/// access count, and a validated flag — that implementations can use to -/// rank, filter, and consolidate the store. -/// -/// # Construction -/// -/// Prefer the builder-style API starting from [`MemoryEntry::new`]: -/// -/// ``` -/// use loopctl::core::agent_memory::{MemoryEntry, MemoryCategory}; -/// -/// let entry = MemoryEntry::new(MemoryCategory::Insight, "Prefer concurrent requests when possible") -/// .with_tag("performance") -/// .validated(); -/// ``` -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct MemoryEntry { - /// Unique identifier for this entry. - /// - /// Typically a UUID v4 string generated at creation time. Used for - /// deduplication and as a stable reference when merging entries during - /// consolidation. - pub id: String, - - /// The category of this memory. - /// - /// See [`MemoryCategory`] for the full set of categories. The category - /// influences how entries are ranked during retrieval and which - /// consolidation rules apply. - pub category: MemoryCategory, - - /// What the agent learned — the core payload of this entry. - /// - /// Free-form text describing the learned information. Implementations may - /// apply NLP techniques (embedding, tokenisation) to this field for - /// similarity search. - pub memory: String, - - /// Tags for categorization and retrieval. - /// - /// Arbitrary strings that act as lightweight indexes. Useful for - /// broad queries like "all performance-related memories". - pub tags: Vec, - - /// When this memory was created. - /// - /// Set to `SystemTime::now()` by [`MemoryEntry::new`]. Recency-based - /// retrieval strategies use this field directly. - pub created_at: SystemTime, - - /// A relevance score (0.0–1.0) for ranking during retrieval. - /// - /// Starts at 1.0 for new entries. Implementations may decay this - /// value over time or boost it when the entry is accessed frequently. - pub relevance: f32, - - /// How many times this memory has been accessed. - /// - /// Implementations should increment this counter each time the entry - /// is returned from [`retrieve`](AgentMemory::retrieve). Since - /// [`retrieve`](AgentMemory::retrieve) takes `&self`, implementations - /// must use interior mutability (e.g. `AtomicUsize` or `Mutex`) to - /// update this field. A high access count signals that the memory is - /// broadly useful and may be a candidate for pinning or promotion to - /// a long-term store. - pub access_count: usize, - - /// Whether this memory has been validated by successful outcomes. - /// - /// Set to `true` via the [`validated`](MemoryEntry::validated) builder - /// method or manually. Consolidation algorithms may treat validated - /// entries as higher-confidence and prefer to keep them when pruning. - pub validated: bool, -} - -impl Default for MemoryEntry { - fn default() -> Self { - Self { - id: uuid::Uuid::new_v4().to_string(), - category: MemoryCategory::Working, - memory: String::new(), - tags: Vec::new(), - created_at: SystemTime::now(), - relevance: 0.5, - access_count: 0, - validated: false, - } - } -} - -impl MemoryEntry { - /// Create a new memory entry with a fresh UUID and the current time. - /// - /// The entry starts with `relevance = 1.0`, `access_count = 0`, and - /// `validated = false`. Use the builder methods ([`with_tag`], [`validated`]) - /// to customise further. - /// - /// [`with_tag`]: MemoryEntry::with_tag - /// [`validated`]: MemoryEntry::validated - /// - /// # Example - /// - /// ``` - /// use loopctl::core::agent_memory::{MemoryEntry, MemoryCategory}; - /// - /// let entry = MemoryEntry::new( - /// MemoryCategory::ErrorPattern, - /// "Timeout on external API — retry with exponential back-off", - /// ); - /// ``` - #[must_use] - pub fn new(category: MemoryCategory, memory: impl Into) -> Self { - Self { - id: uuid::Uuid::new_v4().to_string(), - category, - memory: memory.into(), - tags: Vec::new(), - created_at: SystemTime::now(), - relevance: 1.0, - access_count: 0, - validated: false, - } - } - - /// Add a tag to this entry (builder style). - /// - /// Tags are lightweight, human-readable labels that speed up broad - /// queries. Call chain-style: - /// - /// ``` - /// use loopctl::core::agent_memory::{MemoryEntry, MemoryCategory}; - /// - /// let entry = MemoryEntry::new(MemoryCategory::Fact, "Rust 1.75 stabilised async fn in trait") - /// .with_tag("rust") - /// .with_tag("async"); - /// ``` - #[must_use] - pub fn with_tag(mut self, tag: impl Into) -> Self { - self.tags.push(tag.into()); - self - } - - /// Mark this entry as validated (builder style). - /// - /// Validated entries are treated as higher-confidence by consolidation - /// algorithms and are less likely to be pruned. - /// - /// ``` - /// use loopctl::core::agent_memory::{MemoryEntry, MemoryCategory}; - /// - /// let entry = MemoryEntry::new(MemoryCategory::Strategy, "Use parallel tool calls when independent") - /// .validated(); - /// ``` - #[must_use] - pub fn validated(mut self) -> Self { - self.validated = true; - self - } -} - -/// Category of a memory entry. -/// -/// Each category represents a distinct *kind* of knowledge. Retrieval -/// strategies may weight categories differently (e.g. preferring -/// [`ErrorPattern`](MemoryCategory::ErrorPattern) when debugging), and -/// consolidation rules may vary by category. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum MemoryCategory { - /// A recorded trajectory of tool executions. - /// - /// Captures the sequence of tool calls, their inputs, and outcomes for - /// a particular task. Useful for replaying successful strategies. - Trajectory, - - /// A pattern or insight learned from experience. - /// - /// Generalised knowledge that transcends a single interaction — e.g. - /// "users prefer concise summaries over verbose explanations". - Insight, - - /// A pattern of errors and how they were resolved. - /// - /// Pairs an observed error signature with the fix that resolved it, - /// allowing the agent to avoid repeating the same mistake. - ErrorPattern, - - /// A strategy that was proven effective. - /// - /// High-level plans or heuristics that led to good outcomes, such as - /// "when facing a large refactoring, start with tests". - Strategy, - - /// A fact or piece of knowledge. - /// - /// Static information the agent has learned — e.g. "the project uses - /// `PostgreSQL` 15". Facts are not derived from the agent's own reasoning - /// but are still valuable context. - Fact, - - /// Short-term working memory for the current session. - /// - /// Ephemeral entries that are typically discarded at the end of a - /// session. Useful for tracking intermediate state such as "the user - /// asked about file X in the previous turn". - Working, -} - -/// Statistics from a memory consolidation pass. -/// -/// Returned by [`AgentMemory::consolidate`] so callers can monitor the -/// health of the memory store over time. -/// -/// # Example -/// -/// ``` -/// use loopctl::core::agent_memory::ConsolidationStats; -/// -/// let stats = ConsolidationStats { -/// entries_before: 100, -/// entries_after: 80, -/// pruned: 15, -/// merged: 5, -/// ..Default::default() -/// }; -/// println!( -/// "Consolidated: {} → {} entries (pruned {}, merged {})", -/// stats.entries_before, stats.entries_after, stats.pruned, stats.merged, -/// ); -/// ``` -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct ConsolidationStats { - /// Number of entries before consolidation. - /// - /// Captured at the start of the [`AgentMemory::consolidate`] call. - /// Together with [`entries_after`](ConsolidationStats::entries_after) - /// it gives a quick measure of how aggressively the pass pruned or - /// merged data: `entries_before - entries_after` equals the net - /// reduction. - pub entries_before: usize, - - /// Number of entries after consolidation. - /// - /// Captured at the end of the [`AgentMemory::consolidate`] call. - /// Compare with [`entries_before`](ConsolidationStats::entries_before) - /// to determine the net reduction: `entries_before - entries_after`. - /// Note this reflects the count *after* both pruning and merging, - /// so it already accounts for entries that were merged into existing - /// ones (which reduce the count by at most one per merge). - pub entries_after: usize, - - /// Number of entries that were pruned (removed entirely). - /// - /// Pruned entries are those the implementation judged no longer - /// worth keeping — typically because their - /// [`relevance`](MemoryEntry::relevance) decayed below a threshold, - /// they are stale, or they have been superseded by newer data. - /// This count is a subset of the total reduction: - /// `entries_before - entries_after == pruned + merged` (since each - /// merge removes at least one entry). - pub pruned: usize, - - /// Number of entries that were merged into existing ones. - /// - /// Merging combines duplicate or near-duplicate entries into a single - /// representative entry, preserving the highest - /// [`relevance`](MemoryEntry::relevance) score and unioning the - /// [`tags`](MemoryEntry::tags). This is preferable to pruning when - /// the information is still valuable but is redundantly stored. - /// Each merge typically reduces the entry count by one (the source is - /// absorbed into the target). - pub merged: usize, - - /// Approximate bytes saved by consolidation. - /// - /// A best-effort estimate of the storage reclaimed, useful for - /// logging dashboards and capacity planning. Implementations should - /// sum the serialized size of pruned entries plus the source entries - /// that were absorbed during merges. The value is approximate because - /// exact byte accounting depends on the backing store's encoding and - /// overhead (e.g. index entries, padding). - pub bytes_saved: usize, -} diff --git a/src/core/types.rs b/src/core/types.rs deleted file mode 100644 index ed81b15..0000000 --- a/src/core/types.rs +++ /dev/null @@ -1,1079 +0,0 @@ -//! Common types shared across the agent framework. -//! -//! This module defines the foundational data types that every other module in -//! the framework depends on: configuration ([`AgentConfig`]), lifecycle state -//! ([`AgentState`]), turn and session results ([`TurnResult`], [`SessionResult`]), -//! tool call representations ([`ToolCall`], [`ToolDispatchResult`]), and the -//! reflection/correction system ([`Correction`], [`CorrectionType`], -//! [`CorrectionResult`]). -//! -//! These types are intentionally framework-level — they contain no -//! domain-specific or production-specific logic. Production crates should -//! extend them via composition rather than modification. -//! -//! # Quick Start -//! -//! ``` -//! use loopctl::core::types::{AgentConfig, AgentState, TurnResult, SessionResult}; -//! -//! let config = AgentConfig::default(); -//! assert_eq!(config.max_turns, 200); -//! -//! let turn = TurnResult::completed("Task done."); -//! assert!(turn.is_complete); -//! -//! let session = SessionResult::success(config.session_id); -//! assert!(session.success); -//! ``` - -use std::time::{Duration, SystemTime}; - -use crate::message::ToolContent; -use crate::tool::{ToolError, ToolOutput}; -use serde::{Deserialize, Serialize}; -use uuid::Uuid; - -// =================================================== -// Agent configuration -// =================================================== - -/// Configuration for an agent session. -/// -/// Holds generic agent configuration fields that apply to every session -/// regardless of the specific agent type. Domain-specific configuration -/// (e.g., ITR engine settings, `ToolShield` rules, fallback model chains) -/// should live in production-specific config types that embed or wrap -/// this struct. -/// -/// # Construction -/// -/// Use [`AgentConfig::default`] for sensible defaults or override individual -/// fields through the builder via -/// `AgentBuilder::with_config`. -/// -/// ``` -/// use loopctl::core::types::AgentConfig; -/// -/// let config = AgentConfig { -/// max_turns: 50, -/// model: "default".to_string(), -/// ..Default::default() -/// }; -/// ``` -#[derive(Debug, Clone)] -pub struct AgentConfig { - /// Unique session identifier. - /// - /// Automatically generated as a UUID v4 on construction. Used for - /// correlating logs, metrics, and observer events across the - /// lifetime of a single agent session. - pub session_id: Uuid, - - /// Model identifier (e.g. `"default"`). - /// - /// Production consumers should override this via the builder, but the - /// default provides a reasonable fallback. The value is passed to the - /// API client on each request. - pub model: String, - - /// Optional system prompt override. - /// - /// When `Some`, replaces the agent's default system prompt for this - /// session. When `None`, the agent core decides its own system prompt. - /// Set via `AgentBuilder::system_prompt`. - pub system_prompt: Option, - - /// Maximum number of turns before forcing completion. - /// - /// Prevents runaway sessions. After `max_turns` turns, the framework - /// forces the session to end regardless of - /// `AgentCore::should_continue`. - /// Defaults to `200`. - pub max_turns: usize, - - /// Maximum tokens for each API response. - /// - /// Controls the length of each individual LLM response. Lower values - /// reduce latency and cost; higher values allow longer responses. - /// Defaults to `16_384`. - pub max_tokens: u32, - - /// Context window size for the model (in tokens). - /// - /// Used by the auto-compaction system to estimate when the conversation - /// is approaching the model's context limit. Defaults to `200_000`. - /// - /// Must match the actual context window of the configured [`model`](AgentConfig::model). - pub context_window: u64, - - /// Threshold percentage to trigger auto-compaction (0.0–1.0). - /// - /// When the estimated token usage exceeds `compact_threshold * context_window`, - /// the auto-compaction system triggers a context compaction pass. - /// Defaults to `0.80` (80% of the context window). - pub compact_threshold: f64, - - /// Whether auto-compaction is enabled. - /// - /// When `true`, the framework automatically compacts the conversation - /// context when usage exceeds [`compact_threshold`](AgentConfig::compact_threshold). - /// Set via `AgentBuilder::auto_compact`. - /// Defaults to `true`. - pub auto_compact: bool, -} - -impl Default for AgentConfig { - /// Produce a configuration with production-ready defaults. - /// - /// | Field | Default | - /// |-------|---------| - /// | [`session_id`](AgentConfig::session_id) | Random UUID v4 | - /// | [`model`](AgentConfig::model) | `"default"` | - /// | [`system_prompt`](AgentConfig::system_prompt) | `None` | - /// | [`max_turns`](AgentConfig::max_turns) | `200` | - /// | [`max_tokens`](AgentConfig::max_tokens) | `16_384` | - /// | [`context_window`](AgentConfig::context_window) | `200_000` | - /// | [`compact_threshold`](AgentConfig::compact_threshold) | `0.80` | - /// | [`auto_compact`](AgentConfig::auto_compact) | `true` | - /// - /// # When called - /// - /// By `AgentBuilder::new` to seed the builder, or by consumers who need a - /// baseline before overriding fields. - /// - /// # Example - /// - /// ``` - /// use loopctl::core::types::AgentConfig; - /// - /// let config = AgentConfig::default(); - /// assert_eq!(config.max_turns, 200); - /// assert_eq!(config.model, "default"); - /// ``` - fn default() -> Self { - Self { - session_id: Uuid::new_v4(), - model: "default".to_string(), - system_prompt: None, - max_turns: 200, - max_tokens: 16_384, - context_window: 200_000, - compact_threshold: 0.80, - auto_compact: true, - } - } -} - -// =================================================== -// Agent lifecycle state -// =================================================== - -/// The lifecycle state of an agent. -/// -/// Models the agent as an explicit state machine, making transitions clear -/// and invalid states unrepresentable. The framework reads and writes this -/// enum to drive the agent loop and report status to observers. -/// -/// ```text -/// Idle → Processing → WaitingForTool → Processing → ... → Completed/Failed -/// ↘ Compacting ↗ -/// ↘ Reflecting ↗ -/// ``` -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum AgentState { - /// The agent is idle, waiting for a user message. - /// - /// This is the initial state after initialization and the state - /// the agent returns to between user inputs. - /// - /// No background work is performed while idle. - Idle, - - /// The agent is actively processing a turn. - /// - /// Entered when the agent core begins processing a user message - /// or a tool result. The `turn` field tracks progress for observers. - /// - /// The agent transitions here from [`Idle`](AgentState::Idle) or - /// [`WaitingForTool`](AgentState::WaitingForTool). - Processing { - /// Current turn number (0-indexed). - /// - /// Incremented at the start of each new API call cycle. Used - /// by observers to track session progress and by the framework - /// to enforce [`AgentConfig::max_turns`]. - turn: usize, - }, - - /// The agent is waiting for a tool to complete. - /// - /// Entered after the LLM requests a tool call. The framework - /// dispatches the tool and waits for its result before returning - /// to [`Processing`](AgentState::Processing). - WaitingForTool { - /// Name of the tool being executed. - /// - /// Corresponds to the `name` field of [`ToolCall`]. Used by observers for - /// latency tracking and by the detection system for loop analysis. - /// - /// Never empty — always matches a registered tool name. - tool: String, - - /// When the tool call started. - /// - /// Recorded as [`SystemTime::now`] when the tool dispatch begins. - /// Used to compute tool execution duration for metrics and - /// timeout enforcement. - started_at: SystemTime, - }, - - /// The agent is compacting its conversation context. - /// - /// Entered when token usage exceeds the - /// [`compact_threshold`](AgentConfig::compact_threshold) or when - /// compaction is explicitly requested. The agent summarizes older - /// messages to free context space, then returns to - /// [`Processing`](AgentState::Processing). - Compacting { - /// Why compaction was triggered. - /// - /// See [`CompactReason`] for the possible triggers. - /// - /// Carried through to observers for compaction analytics. - reason: CompactReason, - }, - - /// The agent is reflecting on a failure and preparing a correction. - /// - /// Entered when a tool call fails and the reflection system is - /// enabled (via `Feature::Reflection`). - /// The agent analyzes the error and produces a [`Correction`] before - /// retrying. - Reflecting { - /// Number of errors being analyzed. - /// - /// Helps the agent core gauge how many past failures to consider - /// when formulating a correction strategy. - /// - /// A higher count may indicate a systematic issue requiring - /// an [`ApproachChange`](CorrectionType::ApproachChange). - error_count: usize, - }, - - /// The agent has completed its task. - /// - /// Terminal state. The framework calls - /// `AgentCore::finalize` to produce - /// a [`SessionResult`]. - Completed { - /// Summary of what was accomplished. - /// - /// Typically the last text output from the agent core. Included - /// in [`SessionResult::final_output`] for consumer inspection. - /// - /// May be empty if the agent produced only tool calls. - summary: String, - }, - - /// The agent has failed with an unrecoverable error. - /// - /// Terminal state. The error is propagated through - /// [`SessionResult::error`]. - /// - /// No further turns will be executed after entering this state. - Failed { - /// The error that caused the failure. - /// - /// A human-readable description of the unrecoverable error. - /// Used for logging and for inclusion in - /// [`SessionResult::error`]. - error: String, - }, -} - -/// Reason for triggering context compaction. -/// -/// Indicates why the framework decided to compact the conversation context. -/// Carried inside [`AgentState::Compacting`] and used by observers for -/// compaction analytics. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum CompactReason { - /// Token usage crossed the auto-compaction threshold. - /// - /// Triggered when estimated token usage exceeds - /// [`AgentConfig::compact_threshold`] × [`AgentConfig::context_window`]. - /// - /// This is the normal, expected compaction path. - AutoThreshold, - - /// An explicit user request to compact. - /// - /// Triggered by an external API call or control signal asking the - /// agent to compact immediately, regardless of current token usage. - /// - /// Useful when the consumer knows the conversation is about to grow significantly. - UserRequested, - - /// Emergency compaction due to context window overflow. - /// - /// Triggered when the API rejects a request because the input exceeds - /// the model's context window. This is a last-resort compaction that - /// aggressively summarizes to recover. - Emergency, -} - -// =================================================== -// Turn result -// =================================================== - -/// Result of a single agent turn (one API call → response cycle). -/// -/// Produced by `AgentCore::process_turn` -/// after each LLM interaction. Contains the response text, any tool calls -/// requested, token usage, and timing information. -/// -/// # Construction -/// -/// Use [`TurnResult::completed`] for a terminal response or -/// [`TurnResult::continuing`] for a response that should keep the loop running. -/// Production code typically constructs this from the raw API response. -/// -/// ``` -/// use loopctl::core::types::TurnResult; -/// -/// let done = TurnResult::completed("All tasks finished."); -/// assert!(done.is_complete); -/// -/// let more = TurnResult::continuing("Still working..."); -/// assert!(!more.is_complete); -/// ``` -#[derive(Debug, Clone)] -pub struct TurnResult { - /// The text content of the assistant's response. - /// - /// May be empty if the response consists entirely of tool calls. - /// Combined with tool results in the session's final output. - /// - /// See [`TurnResult::tool_calls`] for the accompanying tool invocations. - pub text: String, - - /// Tool calls made during this turn. - /// - /// Each entry represents a tool the LLM requested to execute. The - /// framework dispatches them and collects results into - /// [`tool_results`](TurnResult::tool_results). - pub tool_calls: Vec, - - /// Results from tool executions during this turn. - /// - /// Populated after tool dispatch. Each result corresponds to a - /// [`ToolCall`] by matching [`tool_call_id`](ToolDispatchResult::tool_call_id) - /// to [`ToolCall::id`]. - pub tool_results: Vec, - - /// How many tokens were used in the API request. - /// - /// Includes the system prompt, conversation history, and the user - /// message. Used for cost tracking and compaction decisions. - /// - /// Reported by the provider in the response metadata. - pub input_tokens: u64, - - /// How many tokens were in the API response. - /// - /// Includes response text and tool call definitions. Used for cost - /// tracking and [`total_tokens`](TurnResult::total_tokens). - /// - /// Reported by the provider in the response metadata. - pub output_tokens: u64, - - /// Wall-clock duration of this turn. - /// - /// Measures the full time from sending the API request through receiving - /// the complete response and executing any tool calls. - pub duration: Duration, - - /// Whether the agent considers the task complete. - /// - /// When `true`, the framework will not invoke - /// `should_continue` and - /// will proceed to finalization. - pub is_complete: bool, - - /// The stop reason reported by the API. - /// - /// Indicates why the LLM stopped generating. Used by the framework - /// to decide whether to dispatch tools ([`StopReason::ToolCall`]) or - /// continue the conversation. - pub stop_reason: StopReason, -} - -impl TurnResult { - /// Create a completed turn result with a simple text response. - /// - /// Sets [`is_complete`](TurnResult::is_complete) to `true` and all - /// token counters to zero. Use this for the final turn of a session. - /// - /// # When called - /// - /// Called by agent core implementations to signal that the task is - /// done and no further turns are needed. - /// - /// # Example - /// - /// ``` - /// use loopctl::core::types::TurnResult; - /// - /// let result = TurnResult::completed("The file has been written successfully."); - /// assert!(result.is_complete); - /// assert_eq!(result.tool_calls.len(), 0); - /// ``` - #[must_use] - pub fn completed(text: impl Into) -> Self { - Self { - text: text.into(), - tool_calls: Vec::new(), - tool_results: Vec::new(), - input_tokens: 0, - output_tokens: 0, - duration: Duration::ZERO, - is_complete: true, - stop_reason: StopReason::EndTurn, - } - } - - /// Create a turn result that should continue with more turns. - /// - /// Sets [`is_complete`](TurnResult::is_complete) to `false`, indicating - /// that the agent loop should keep running. - /// - /// # When called - /// - /// Called by agent core implementations when a turn produces intermediate - /// output (e.g., tool results to process) and the session is not yet done. - /// - /// # Example - /// - /// ``` - /// use loopctl::core::types::TurnResult; - /// - /// let result = TurnResult::continuing("I need to read the file first..."); - /// assert!(!result.is_complete); - /// ``` - #[must_use] - pub fn continuing(text: impl Into) -> Self { - Self { - text: text.into(), - tool_calls: Vec::new(), - tool_results: Vec::new(), - input_tokens: 0, - output_tokens: 0, - duration: Duration::ZERO, - is_complete: false, - stop_reason: StopReason::EndTurn, - } - } - - /// Check if this turn included any tool calls. - /// - /// Returns `true` when [`tool_calls`](TurnResult::tool_calls) is - /// non-empty, indicating that the LLM requested tool execution. - /// - /// # When called - /// - /// Called by the framework to decide whether to enter the tool-dispatch - /// path or proceed to the next turn. - #[must_use] - pub fn has_tool_calls(&self) -> bool { - !self.tool_calls.is_empty() - } - - /// Total tokens (input + output) for this turn. - /// - /// Convenience method that sums [`input_tokens`](TurnResult::input_tokens) - /// and [`output_tokens`](TurnResult::output_tokens). Used for session-level - /// token accounting. - #[must_use] - pub fn total_tokens(&self) -> u64 { - self.input_tokens.saturating_add(self.output_tokens) - } -} - -// =================================================== -// Stop reason -// =================================================== - -/// Why the API stopped generating. -/// -/// Mirrors the stop reasons returned by LLM APIs. The framework uses this -/// to determine the next step: dispatch tools, continue the conversation, -/// or end the session. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum StopReason { - /// The model decided to stop (natural end of turn). - /// - /// The LLM finished its response without requesting tools or hitting - /// any limit. The framework should check - /// [`TurnResult::is_complete`] to decide whether to continue. - EndTurn, - - /// The model requested tool execution. - /// - /// The LLM response contains one or more tool calls in - /// [`TurnResult::tool_calls`]. The framework should dispatch them - /// and feed the results back. - ToolCall, - - /// The maximum token limit was reached. - /// - /// The LLM hit the [`AgentConfig::max_tokens`] limit before - /// finishing. The response may be truncated. The framework may - /// choose to continue the turn to let the model complete its output. - MaxTokens, - - /// The stop sequence was encountered. - /// - /// The model generated a configured stop sequence. Rare in - /// standard usage; typically indicates custom API configuration. - StopSequence, -} - -// =================================================== -// Tool call -// =================================================== - -/// A tool call requested by the agent. -/// -/// Represents a single tool invocation that the LLM has requested during a -/// turn. The framework matches each `ToolCall` to a registered tool, executes -/// it, and produces a [`ToolDispatchResult`] with the output. -/// -/// # Serialization -/// -/// Implements `Serialize` and `Deserialize` for persistence and inter-process -/// communication. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ToolCall { - /// Unique identifier for this tool call. - /// - /// Assigned by the LLM API. Used to correlate the call with its - /// [`ToolDispatchResult`] via [`ToolDispatchResult::tool_call_id`]. - pub id: String, - - /// Name of the tool to invoke. - /// - /// Must match a tool registered in the - /// `ToolRegistry`. Examples: `"Read"`, - /// `"Bash"`, `"Write"`. - pub tool: String, - - /// Input parameters for the tool. - /// - /// A JSON object whose schema depends on the tool. For example, - /// a `Read` tool might have `{"file_path": "/path/to/file"}`. - /// The framework passes this directly to the tool implementation. - pub input: serde_json::Value, -} - -/// The outcome of a single tool invocation. -/// -/// Produced after the framework dispatches a [`ToolCall`] and collects -/// the tool's output. Used throughout the middleware pipeline, the -/// engine dispatch layer, and returned to callers via -/// [`TurnResult::tool_results`]. -/// -/// # Fields -/// -/// | Field | Source | -/// |------------------------|-------------------------------------| -/// | `tool_call_id` | Set by the engine after dispatch | -/// | `output` | From [`ToolOutput::payload`] | -/// | `is_error` | From [`ToolOutput::is_error`] | -/// | `duration` | Measured by the dispatch layer | -/// | `resolved_tool_name` | Set by middleware or engine | -/// -/// # Construction -/// -/// Middlewares build results with [`ToolDispatchResult::ok`], -/// [`ToolDispatchResult::err`], or [`From`] combined with -/// builder methods. The engine layer attaches the `tool_call_id` via -/// [`ToolDispatchResult::with_call_id`] after the middleware pipeline -/// returns. -/// -/// ``` -/// use std::time::Duration; -/// use loopctl::core::ToolDispatchResult; -/// use loopctl::tool::ToolOutput; -/// -/// let output = ToolOutput::text("done"); -/// let result = ToolDispatchResult::from(output) -/// .with_tool_name("bash") -/// .with_duration(Duration::from_millis(42)) -/// .with_call_id("call_abc123"); -/// -/// assert_eq!(result.tool_call_id, "call_abc123"); -/// assert_eq!(result.resolved_tool_name, "bash"); -/// ``` -#[derive(Debug, Clone)] -pub struct ToolDispatchResult { - /// The tool call this result is for. - /// - /// Matches [`ToolCall::id`] to correlate results back to their - /// originating requests. Set by the engine after the middleware - /// pipeline completes via [`with_call_id`](Self::with_call_id). - pub tool_call_id: String, - - /// The tool's output content or error message. - /// - /// On success, holds the full [`ToolOutput`] - /// payload — preserving multipart and image content instead of - /// flattening to text. On failure, contains an error message wrapped - /// in [`ToolContent::Text`]. - pub output: ToolContent, - - /// Whether the tool execution resulted in an error. - /// - /// When `true`, [`output`](ToolDispatchResult::output) contains the - /// error message and the framework may trigger reflection or retry - /// logic. - pub is_error: bool, - - /// Duration of tool execution. - /// - /// Measures wall-clock time from tool dispatch start to completion. - /// Used by observers for latency tracking and by the health system - /// for timeout analysis. - pub duration: Duration, - - /// The name of the tool that actually ran. - /// - /// May differ from the requested `tool_name` if a routing - /// middleware redirected the call to an alternative tool. - pub resolved_tool_name: String, -} - -impl ToolDispatchResult { - /// Create a successful result with text output. - /// - /// Convenience constructor for the common case where a tool - /// produces a plain-text response. - #[must_use] - pub fn ok(tool_name: &str, output: String, duration: Duration) -> Self { - Self { - tool_call_id: String::new(), - output: ToolContent::Text(output), - is_error: false, - duration, - resolved_tool_name: tool_name.to_string(), - } - } - - /// Create an error result with a message. - /// - /// Used when a middleware short-circuits or the tool reports failure. - #[must_use] - pub fn err(tool_name: &str, message: String, duration: Duration) -> Self { - Self { - tool_call_id: String::new(), - output: ToolContent::Text(message), - is_error: true, - duration, - resolved_tool_name: tool_name.to_string(), - } - } - - /// Create a result from a [`ToolOutput`]. - /// - /// Converts the tool's output struct into a dispatch result, - /// preserving the error flag and content payload. - #[must_use] - pub fn from_tool_output(tool_name: &str, output: ToolOutput, duration: Duration) -> Self { - Self::from(output) - .with_tool_name(tool_name) - .with_duration(duration) - } - - /// Builder: attach the [`tool_call_id`](Self::tool_call_id). - /// - /// Called by the engine layer after the middleware pipeline returns - /// to correlate this result with the original [`ToolCall`]. - #[must_use] - pub fn with_call_id(mut self, id: impl Into) -> Self { - self.tool_call_id = id.into(); - self - } - - /// Set the [`resolved_tool_name`](Self::resolved_tool_name). - /// - /// Part of the builder chain when constructing a - /// `ToolDispatchResult` from [`From`]. - #[must_use] - pub fn with_tool_name(mut self, name: &str) -> Self { - name.clone_into(&mut self.resolved_tool_name); - self - } - - /// Set the [`duration`](Self::duration). - /// - /// Part of the builder chain when constructing a - /// `ToolDispatchResult` from [`From`]. - #[must_use] - pub fn with_duration(mut self, dur: Duration) -> Self { - self.duration = dur; - self - } - - /// Create a result from a [`ToolError`]. - /// - /// Converts the tool's error into a dispatch result with `is_error` - /// set to `true`. - #[must_use] - pub fn from_tool_error(tool_name: &str, error: &ToolError, duration: Duration) -> Self { - Self { - tool_call_id: String::new(), - output: ToolContent::Text(error.to_string()), - is_error: true, - duration, - resolved_tool_name: tool_name.to_string(), - } - } - - /// Create a result from a tool call outcome. - /// - /// Covers the common `Result` pattern produced by - /// [`Tool::call()`](crate::tool::Tool::call). Maps [`Ok`] through - /// [`from_tool_output`](Self::from_tool_output) and [`Err`] through - /// [`from_tool_error`](Self::from_tool_error). - #[must_use] - pub fn from_result( - tool_name: &str, - result: Result, - duration: Duration, - ) -> Self { - match result { - Ok(output) => Self::from_tool_output(tool_name, output, duration), - Err(e) => Self::from_tool_error(tool_name, &e, duration), - } - } -} - -/// Conversion from a bare [`ToolOutput`]. -/// -/// Produces a `ToolDispatchResult` with no call ID, [`Duration::ZERO`], -/// and an empty `resolved_tool_name`. Chain builder methods to complete -/// the fields: -/// -/// ``` -/// use std::time::Duration; -/// use loopctl::core::ToolDispatchResult; -/// use loopctl::tool::ToolOutput; -/// -/// let result = ToolDispatchResult::from(ToolOutput::text("ok")) -/// .with_call_id("call_1") -/// .with_tool_name("echo") -/// .with_duration(Duration::from_millis(5)); -/// ``` -impl From for ToolDispatchResult { - fn from(output: ToolOutput) -> Self { - Self { - tool_call_id: String::new(), - output: output.payload, - is_error: output.is_error, - duration: Duration::ZERO, - resolved_tool_name: String::new(), - } - } -} - -// =================================================== -// Session result -// =================================================== - -/// Summary of a complete agent session. -/// -/// Produced by `AgentCore::finalize` -/// after the last turn. Aggregates all session-level metrics: total turns, -/// tokens, duration, tool calls, and final output. -/// -/// # Construction -/// -/// Use [`SessionResult::success`] for a completed session or -/// [`SessionResult::failed`] for a session that ended with an error. -/// -/// ``` -/// use loopctl::core::types::SessionResult; -/// use uuid::Uuid; -/// -/// let session_id = Uuid::new_v4(); -/// -/// let ok = SessionResult::success(session_id); -/// assert!(ok.success); -/// -/// let err = SessionResult::failed(session_id, "API rate limit exceeded"); -/// assert!(!err.success); -/// assert_eq!(err.error.unwrap(), "API rate limit exceeded"); -/// ``` -#[derive(Debug, Clone)] -pub struct SessionResult { - /// The session identifier. - /// - /// Matches [`AgentConfig::session_id`]. Used for correlating this - /// result with logs, metrics, and observer events. - pub session_id: Uuid, - - /// Total number of turns executed. - /// - /// Counts every API call cycle, including tool-dispatch turns. - /// Compared against [`AgentConfig::max_turns`] to detect runaway - /// sessions. - pub total_turns: usize, - - /// Total input tokens consumed across all turns. - /// - /// Sum of [`TurnResult::input_tokens`] for every turn in the session. - /// Used for cost reporting (input tokens are typically cheaper). - pub input_tokens: u64, - - /// Total output tokens consumed across all turns. - /// - /// Sum of [`TurnResult::output_tokens`] for every turn in the session. - /// Used for cost reporting (output tokens are typically more expensive). - pub output_tokens: u64, - - /// Total wall-clock time. - /// - /// Measures from session start to finalization. Includes all API - /// calls, tool executions, and compaction passes. - pub total_duration: Duration, - - /// Number of tool calls made. - /// - /// Counts every [`ToolCall`] dispatched during the session. Useful - /// for understanding agent behavior and cost. - pub tool_calls: usize, - - /// Whether the session completed successfully. - /// - /// `true` when the agent reached [`AgentState::Completed`], `false` - /// when it reached [`AgentState::Failed`]. - pub success: bool, - - /// Final text output from the agent (if any). - /// - /// Contains the last meaningful text response from the agent core, - /// typically from [`AgentState::Completed`]. `None` if the session - /// ended without a final message. - pub final_output: Option, - - /// Error message if the session failed. - /// - /// `Some` when [`success`](SessionResult::success) is `false`, - /// containing a human-readable description of what went wrong. - /// `None` for successful sessions. - pub error: Option, -} - -impl SessionResult { - /// Create a successful session result. - /// - /// Initializes all counters to zero and sets [`success`](SessionResult::success) - /// to `true`. The framework or production code should fill in the actual - /// counters before returning. - /// - /// # When called - /// - /// Called by `AgentCore::finalize` - /// implementations when the session completed without error. - /// - /// # Example - /// - /// ``` - /// use loopctl::core::types::SessionResult; - /// use uuid::Uuid; - /// - /// let session_id = Uuid::new_v4(); - /// let result = SessionResult::success(session_id); - /// assert!(result.success); - /// assert!(result.error.is_none()); - /// ``` - #[must_use] - pub fn success(session_id: Uuid) -> Self { - Self { - session_id, - total_turns: 0, - input_tokens: 0, - output_tokens: 0, - total_duration: Duration::ZERO, - tool_calls: 0, - success: true, - final_output: None, - error: None, - } - } - - /// Create a failed session result. - /// - /// Sets [`success`](SessionResult::success) to `false` and records the - /// error message. All counters are initialized to zero. - /// - /// # When called - /// - /// Called by `AgentCore::finalize` - /// implementations when the session ended due to an unrecoverable error. - /// - /// # Example - /// - /// ``` - /// use loopctl::core::types::SessionResult; - /// use uuid::Uuid; - /// - /// let session_id = Uuid::new_v4(); - /// let result = SessionResult::failed(session_id, "API rate limit exceeded"); - /// assert!(!result.success); - /// assert_eq!(result.error.unwrap(), "API rate limit exceeded"); - /// ``` - #[must_use] - pub fn failed(session_id: Uuid, error: impl Into) -> Self { - Self { - session_id, - total_turns: 0, - input_tokens: 0, - output_tokens: 0, - total_duration: Duration::ZERO, - tool_calls: 0, - success: false, - final_output: None, - error: Some(error.into()), - } - } - - /// Total tokens (input + output) for this session. - /// - /// Convenience method that sums [`input_tokens`](SessionResult::input_tokens) - /// and [`output_tokens`](SessionResult::output_tokens). - #[must_use] - pub fn total_tokens(&self) -> u64 { - self.input_tokens.saturating_add(self.output_tokens) - } -} - -// =================================================== -// Correction system -// =================================================== - -/// A correction produced by the reflection system. -/// -/// When a tool call fails and reflection is enabled (via -/// `Feature::Reflection`), the agent -/// analyzes the error and produces a `Correction` that describes how to fix -/// the problem. The framework applies the correction and retries. -/// -/// # Serialization -/// -/// Implements `Serialize` and `Deserialize` for persistence and observability. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct Correction { - /// Type of correction to apply. - /// - /// Categorizes the fix strategy. See [`CorrectionType`] for the - /// available strategies and their semantics. - pub correction_type: CorrectionType, - - /// Human-readable description of the correction. - /// - /// Explains *what* went wrong and *how* the correction addresses it. - /// Used for logging, observability, and agent guidance. - pub description: String, - - /// Modified tool input (if applicable). - /// - /// When [`correction_type`](Correction::correction_type) is - /// [`InputFix`](CorrectionType::InputFix), contains the corrected - /// JSON input that should replace the original. `None` when the - /// correction does not modify tool input. - pub modified_input: Option, - - /// Alternative tool to use (if applicable). - /// - /// When [`correction_type`](Correction::correction_type) is - /// [`ToolChange`](CorrectionType::ToolChange), contains the name - /// of the tool that should be used instead. `None` when the - /// correction keeps the same tool. - pub alternative_tool: Option, - - /// Additional guidance for the agent. - /// - /// Free-form text that provides extra context or instructions to - /// help the agent avoid the same failure in future turns. - pub guidance: Option, -} - -/// Type of correction to apply. -/// -/// Categorizes the fix strategy that the reflection system has determined -/// is most appropriate for the observed failure. Each variant maps to a -/// different retry approach. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CorrectionType { - /// Fix the input to the tool. - /// - /// The tool was correct but its input parameters were wrong (e.g., a - /// typo in a file path). The correction provides a fixed input via - /// [`Correction::modified_input`]. - InputFix, - - /// Use a different tool. - /// - /// The chosen tool was inappropriate for the task. The correction - /// specifies an alternative via [`Correction::alternative_tool`]. - ToolChange, - - /// Fix a dependency or prerequisite. - /// - /// The tool call failed because a prerequisite was not met (e.g., - /// a directory doesn't exist). The correction describes what needs - /// to be done first. - PrerequisiteFix, - - /// Change the approach entirely. - /// - /// The current strategy is fundamentally flawed. The correction - /// provides high-level guidance for a different approach via - /// [`Correction::guidance`]. - ApproachChange, - - /// No fix possible, escalate. - /// - /// The reflection system cannot determine a correction. The - /// framework should propagate the error to the user or higher-level - /// handler. - Escalate, -} - -/// Result of applying a correction. -/// -/// Indicates whether the reflection system's correction was successfully -/// applied, failed, or was skipped. Produced after attempting to retry -/// with the corrected parameters. -#[derive(Debug, Clone)] -pub enum CorrectionResult { - /// Correction was applied successfully. - /// - /// The retry with corrected parameters succeeded and the agent can - /// continue processing normally. - Applied, - - /// Correction failed. - /// - /// The retry also failed. Contains a human-readable error message - /// describing what went wrong with the corrected attempt. - Failed(String), - - /// No correction was needed or possible. - /// - /// The reflection system decided not to apply a correction (e.g., - /// the error is transient or the correction type was - /// [`Escalate`](CorrectionType::Escalate)). - Skipped, -} diff --git a/src/detection.rs b/src/detection.rs new file mode 100644 index 0000000..747e45c --- /dev/null +++ b/src/detection.rs @@ -0,0 +1,21 @@ +//! Detection — loop and convergence detection for agent loops. +//! +//! - **[`loop_detector`]** — Detects repetitive tool-use patterns and enforces limits. +//! - **[`convergence`]** — Detects when agent responses become semantically similar. +//! - **[`manager`]** — Unified [`DetectionManager`] that orchestrates both detectors. +//! +//! For capability traits and the runtime bundle, see [`crate::runtime`]. + +pub mod convergence; +pub mod loop_detector; +pub mod manager; + +pub use convergence::{ + ConvergenceAction, ConvergenceConfig, ConvergenceConfigError, ConvergenceDetector, + ConvergenceStatus, +}; +pub use loop_detector::{ + LoopDetector, LoopDetectorConfig, LoopStatus, NoOpToolSignature, Operation, ToolSignature, + global_detector, hash_result, +}; +pub use manager::{DetectedPattern, DetectionConfig, DetectionManager, DetectionStats}; diff --git a/src/loop_control/convergence.rs b/src/detection/convergence.rs similarity index 68% rename from src/loop_control/convergence.rs rename to src/detection/convergence.rs index f368181..7ce57b4 100644 --- a/src/loop_control/convergence.rs +++ b/src/detection/convergence.rs @@ -19,19 +19,6 @@ //! [`ConvergenceConfig::window_size`], convergence is flagged and the //! configured [`ConvergenceAction`] is returned. //! -//! # Detection Flow -//! -//! ```text -//! add_response(text) -//! ├─ disabled or empty? → return no_convergence() -//! ├─ for each prev in window: -//! │ compute Jaccard similarity(text, prev) -//! │ if similarity >= threshold → increment consecutive_count -//! │ else → reset consecutive_count to 1 -//! ├─ push text into window (evict oldest if full) -//! └─ if consecutive_count >= window_size → return converged status -//! ``` -//! //! # Provided Types //! //! - **[`ConvergenceConfig`]** — Configuration: window size, similarity @@ -54,7 +41,7 @@ //! # Quick Start //! //! ```rust -//! use loopctl::loop_control::convergence::{ +//! use loopctl::detection::convergence::{ //! ConvergenceConfig, ConvergenceDetector, ConvergenceAction, //! }; //! @@ -72,7 +59,7 @@ //! let status = detector.add_response("I am working on the task."); //! let status = detector.add_response("I am working on the task."); //! assert!(status.detected); // 3 consecutive similar responses -//! # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) +//! # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) //! ``` use std::collections::HashSet; @@ -103,16 +90,11 @@ use serde::{Deserialize, Serialize}; /// | Interactive / REPL session | [`AskUser`](ConvergenceAction::AskUser) | /// | Long context / token budget | [`Compact`](ConvergenceAction::Compact) | /// -/// # Serde -/// -/// This enum derives [`Serialize`] and [`Deserialize`] so it can be -/// loaded from TOML/JSON configuration files. Variant names are serialized -/// in `snake_case` (e.g., `"switch_phase"`). /// /// # Example /// /// ```rust -/// use loopctl::loop_control::convergence::ConvergenceAction; +/// use loopctl::detection::convergence::ConvergenceAction; /// /// let action = ConvergenceAction::Warn; /// match action { @@ -129,7 +111,7 @@ pub enum ConvergenceAction { /// Stop the agent loop entirely. /// /// The agent has converged and is unlikely to make further progress. - /// This is the safest default — the caller should report the situation + /// Default action — halt the agent loop and report the situation /// to the user and await new instructions. /// /// The `DetectionManager` will @@ -153,7 +135,7 @@ pub enum ConvergenceAction { /// the converged pattern. /// /// The specific phase transition is left to the caller's discretion; - /// the detector simply signals that the current strategy has stagnated. + /// the detector signals that the current strategy has stagnated. SwitchPhase, /// Ask the user for guidance before continuing. @@ -175,7 +157,7 @@ pub enum ConvergenceAction { /// turns, keeping only the most recent N exchanges, or pruning /// low-relevance messages). /// - /// This is useful when convergence is caused by the agent repeatedly + /// Useful when convergence is caused by the agent repeatedly /// revisiting earlier context — compaction removes the redundant /// history that may be driving the repetition. Compact, @@ -195,7 +177,7 @@ pub enum ConvergenceAction { /// # Example /// /// ```rust -/// use loopctl::loop_control::convergence::{ConvergenceConfig, ConvergenceDetector, ConvergenceConfigError}; +/// use loopctl::detection::convergence::{ConvergenceConfig, ConvergenceDetector, ConvergenceConfigError}; /// /// let bad_config = ConvergenceConfig { /// window_size: 1, @@ -211,20 +193,14 @@ pub enum ConvergenceConfigError { /// Convergence requires at least one pair of consecutive responses, /// so a window of 1 (or 0) is meaningless. #[error("window_size must be at least 2, got {actual}")] - WindowTooSmall { - /// The invalid value that was provided. - actual: usize, - }, + WindowTooSmall { actual: usize }, /// `similarity_threshold` is outside the valid range `[0.0, 1.0]`. /// /// Jaccard similarity always produces a value in this range; a /// threshold outside it would never (or always) trigger. #[error("similarity_threshold must be in [0.0, 1.0], got {actual}")] - ThresholdOutOfRange { - /// The invalid value that was provided. - actual: f32, - }, + ThresholdOutOfRange { actual: f32 }, } // =================================================== @@ -236,11 +212,6 @@ pub enum ConvergenceConfigError { /// Controls the sensitivity and behavior of the [`ConvergenceDetector`]. /// Passed to [`ConvergenceDetector::new`] at construction time. /// -/// # Derives -/// -/// [`Debug`] for logging, [`Clone`] for sharing config across detectors, -/// and [`Serialize`]/[`Deserialize`] for loading from config files. -/// /// # Defaults /// /// | Field | Default | @@ -263,7 +234,7 @@ pub enum ConvergenceConfigError { /// # Example /// /// ```rust -/// use loopctl::loop_control::convergence::{ConvergenceConfig, ConvergenceAction, ConvergenceDetector}; +/// use loopctl::detection::convergence::{ConvergenceConfig, ConvergenceAction, ConvergenceDetector}; /// /// let config = ConvergenceConfig { /// enabled: true, @@ -272,66 +243,17 @@ pub enum ConvergenceConfigError { /// on_converge: ConvergenceAction::Warn, /// }; /// let detector = ConvergenceDetector::new(config)?; -/// # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) +/// # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) /// ``` #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ConvergenceConfig { - /// Whether convergence detection is active. - /// - /// When `false`, [`ConvergenceDetector::add_response`] and - /// [`ConvergenceDetector::check_convergence`] return - /// [`ConvergenceStatus::no_convergence`] immediately without performing - /// any computation. - /// - /// Defaults to `true`. - /// - /// # When to Disable - /// - /// Set to `false` in performance-critical paths where similarity - /// computation overhead is unacceptable, or when the agent's task - /// domain makes convergence impossible by design. + /// Whether convergence detection is active. Defaults to `true`. pub enabled: bool, - - /// Number of consecutive similar responses required to declare convergence. - /// - /// The detector maintains a sliding window of this size. When every - /// response in the window exceeds [`ConvergenceConfig::similarity_threshold`] - /// relative to its neighbours, convergence is detected. - /// - /// Defaults to `3`. - /// - /// # Constraints - /// - /// Must be at least `2`. Values below `2` are meaningless since - /// convergence requires at least one pair of consecutive responses. + /// Consecutive similar responses required to declare convergence. Must be ≥ 2. Defaults to `3`. pub window_size: usize, - - /// Similarity threshold (0.0–1.0) above which two responses are considered - /// "similar". - /// - /// Computed via Jaccard similarity on word-level tokens (see - /// [`ConvergenceDetector::compute_similarity`]). A value of `1.0` requires - /// exact word-set matches; `0.0` considers any two non-empty strings similar. - /// - /// Defaults to `0.95`. - /// - /// # Invariant - /// - /// Must be in the range `0.0..=1.0`. Values outside this range will - /// not panic but produce nonsensical results. + /// Jaccard similarity threshold (0.0–1.0). Defaults to `0.95`. pub similarity_threshold: f32, - - /// Action to take when convergence is detected. - /// - /// Carried through to [`ConvergenceStatus::action`] so the caller can - /// respond appropriately. See [`ConvergenceAction`] for available options. - /// - /// Defaults to [`ConvergenceAction::Stop`]. - /// - /// # Serde - /// - /// Uses `#[serde(default)]` so missing fields in TOML/JSON deserialize - /// to [`ConvergenceAction::Stop`]. + /// Action on convergence. Defaults to [`ConvergenceAction::Stop`]. #[serde(default)] pub on_converge: ConvergenceAction, } @@ -341,17 +263,17 @@ impl Default for ConvergenceConfig { /// /// The defaults are tuned for typical agent workloads: /// - /// | Field | Default | - /// |-------|---------| - /// | [`enabled`](ConvergenceConfig::enabled) | `true` | - /// | [`window_size`](ConvergenceConfig::window_size) | `3` | - /// | [`similarity_threshold`](ConvergenceConfig::similarity_threshold) | `0.95` | - /// | [`on_converge`](ConvergenceConfig::on_converge) | [`ConvergenceAction::Stop`] | + /// | Field | Default | + /// |-------------------------------------------------------------------|-----------------------------| + /// | [`enabled`](ConvergenceConfig::enabled) | `true` | + /// | [`window_size`](ConvergenceConfig::window_size) | `3` | + /// | [`similarity_threshold`](ConvergenceConfig::similarity_threshold) | `0.95` | + /// | [`on_converge`](ConvergenceConfig::on_converge) | [`ConvergenceAction::Stop`] | /// /// # Example /// /// ```rust - /// use loopctl::loop_control::convergence::ConvergenceConfig; + /// use loopctl::detection::convergence::ConvergenceConfig; /// /// let config = ConvergenceConfig::default(); /// assert!(config.enabled); @@ -380,28 +302,13 @@ impl Default for ConvergenceConfig { /// # Construction /// /// Use [`ConvergenceStatus::no_convergence`] to create a "no convergence" -/// sentinel. A "converged" status is constructed internally by the detector +/// sentinel. A "converged" status is returned by the detector /// when the [`ConvergenceConfig::window_size`] threshold is met. /// -/// # Derives -/// -/// [`Debug`] for diagnostic output, [`Clone`] for sharing results, and -/// [`Default`] which produces the same value as [`no_convergence`](ConvergenceStatus::no_convergence). -/// -/// # Fields Overview -/// -/// ```text -/// detected ──────────── true if convergence reached -/// consecutive_count ─── how many similar responses in a row -/// similarity_score ──── highest Jaccard score among compared pairs -/// similar_responses ─── the response texts that triggered detection -/// action ────────────── what the caller should do (stop, warn, ...) -/// ``` -/// /// # Example /// /// ```rust -/// use loopctl::loop_control::convergence::{ConvergenceConfig, ConvergenceDetector}; +/// use loopctl::detection::convergence::{ConvergenceConfig, ConvergenceDetector}; /// /// let config = ConvergenceConfig { /// window_size: 3, @@ -417,58 +324,19 @@ impl Default for ConvergenceConfig { /// println!("Similarity: {:.2}%", status.similarity_score * 100.0); /// println!("Action: {:?}", status.action); /// } -/// # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) +/// # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) /// ``` #[derive(Debug, Clone, Default)] pub struct ConvergenceStatus { - /// Whether convergence was detected. - /// - /// `true` when [`ConvergenceStatus::consecutive_count`] is at least - /// [`ConvergenceConfig::window_size`], meaning the agent has produced - /// enough consecutive similar responses to warrant action. - /// - /// When `false`, the other fields still contain valid data (streak count, - /// latest similarity score, etc.) but should not be treated as a - /// convergence signal. + /// `true` when `consecutive_count >= window_size`. Other fields still valid when `false`. pub detected: bool, - - /// Number of consecutive similar responses observed so far. - /// - /// Resets to `1` whenever a response falls below - /// [`ConvergenceConfig::similarity_threshold`]. Monotonically increases - /// while responses remain similar. - /// - /// Compare against [`ConvergenceConfig::window_size`] to determine - /// how close the detector is to declaring convergence. + /// Resets to `1` when similarity falls below threshold. pub consecutive_count: usize, - - /// Highest Jaccard similarity score (0.0–1.0) among the compared pairs. - /// - /// Useful for diagnostics: a score of `0.99` means the responses are - /// nearly identical; a lower score like `0.75` suggests the detector - /// is barely triggering and the threshold may need tuning. - /// - /// Computed by [`ConvergenceDetector::compute_similarity`] during - /// [`ConvergenceDetector::add_response`]. Always `0.0` when returned - /// from [`ConvergenceDetector::check_convergence`] since no new - /// comparison is performed. + /// Highest Jaccard similarity (0.0–1.0). `0.0` when no comparison was made. pub similarity_score: f32, - - /// The response strings that contributed to the convergence detection. - /// - /// Contains deduplicated copies of the recent responses that exceeded - /// the similarity threshold. Useful for logging and user-facing reports. - /// - /// Cleared whenever a dissimilar response breaks the streak, so the - /// collection only contains responses from the *current* streak. + /// Responses that exceeded the similarity threshold. Cleared on dissimilar response. pub similar_responses: Vec, - - /// The configured action to take, forwarded from - /// [`ConvergenceConfig::on_converge`]. - /// - /// The caller should inspect this field to decide how to respond (stop, - /// warn, switch phase, or ask the user). See [`ConvergenceAction`] for - /// the full list of variants and their semantics. + /// Forwarded from [`ConvergenceConfig::on_converge`]; see [`ConvergenceAction`]. pub action: ConvergenceAction, } @@ -484,7 +352,7 @@ impl ConvergenceStatus { /// # Example /// /// ```rust - /// use loopctl::loop_control::convergence::ConvergenceStatus; + /// use loopctl::detection::convergence::ConvergenceStatus; /// /// let status = ConvergenceStatus::no_convergence(); /// assert!(!status.detected); @@ -518,43 +386,10 @@ impl ConvergenceStatus { /// Prefer [`ConvergenceDetector::new`] with a custom [`ConvergenceConfig`], /// or [`ConvergenceDetector::default_detector`] for sensible defaults. /// -/// # Internal Architecture -/// -/// ```text -/// ┌─────────────────────────────────────────────────┐ -/// │ ConvergenceDetector │ -/// │ │ -/// │ config: ConvergenceConfig │ -/// │ window: VecDeque ←─ sliding window │ -/// │ consecutive_count: usize │ -/// │ similar_responses: Vec │ -/// │ │ -/// │ Methods: │ -/// │ add_response() → ConvergenceStatus │ -/// │ check_convergence() → ConvergenceStatus │ -/// │ compute_similarity() → f32 │ -/// │ clear() │ -/// └─────────────────────────────────────────────────┘ -/// ``` -/// -/// # Derives -/// -/// [`Debug`] for diagnostic output. Not [`Clone`] — the detector holds -/// mutable state and should be used as a single instance. -/// -/// # Lifecycle -/// -/// ```text -/// new(config) -/// → add_response(response) [repeated, returns ConvergenceStatus] -/// → check_convergence() [optional, peek without adding] -/// → clear() [reset for a new task] -/// ``` -/// /// # Example /// /// ```rust -/// use loopctl::loop_control::convergence::ConvergenceDetector; +/// use loopctl::detection::convergence::ConvergenceDetector; /// /// let mut detector = ConvergenceDetector::default_detector()?; /// @@ -568,52 +403,17 @@ impl ConvergenceStatus { /// // Reset for a fresh start /// detector.clear(); /// assert!(detector.window().is_empty()); -/// # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) +/// # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) /// ``` #[derive(Debug)] pub struct ConvergenceDetector { - /// Configuration controlling thresholds and actions. - /// - /// Set once at construction via [`ConvergenceDetector::new`] and read - /// on every call to [`ConvergenceDetector::add_response`]. Immutable - /// for the lifetime of the detector — create a new detector to change - /// configuration. + /// Immutable config set at construction. config: ConvergenceConfig, - - /// Sliding window of recent response strings. - /// - /// Bounded by [`ConvergenceConfig::window_size`]. When full, the oldest - /// entry is evicted before a new one is appended. Accessed by the - /// `DetectionManager` for - /// direct inspection during testing. - /// - /// The deque is ordered chronologically: index `0` is the oldest - /// response, and the last index is the most recent. - /// - /// # Capacity - /// - /// Pre-allocated to [`ConvergenceConfig::window_size`] at construction - /// to avoid frequent heap allocations during steady-state operation. + /// Bounded by `window_size`, ordered oldest→newest. pub(super) window: VecDeque, - - /// Number of consecutive responses that exceeded the similarity threshold. - /// - /// Resets to `1` when a response is sufficiently different from its - /// predecessor. When this reaches [`ConvergenceConfig::window_size`], - /// convergence is declared. - /// - /// This field is the internal counterpart of - /// [`ConvergenceStatus::consecutive_count`]. + /// Resets to `1` on dissimilar response; convergence at `window_size`. pub(super) consecutive_count: usize, - - /// Deduplicated collection of responses deemed "similar" so far. - /// - /// Grows as consecutive similar responses are observed and is cleared - /// whenever a dissimilar response breaks the streak. Reported in - /// [`ConvergenceStatus::similar_responses`] for diagnostics. - /// - /// Only unique strings are stored — duplicate responses are filtered - /// to keep the collection compact. + /// Deduplicated similar responses; cleared when streak breaks. similar_responses: Vec, } @@ -626,9 +426,6 @@ impl ConvergenceDetector { /// [`ConvergenceConfigError::ThresholdOutOfRange`] if /// [`ConvergenceConfig::similarity_threshold`] is outside `[0.0, 1.0]`. /// - /// Pre-allocates the internal window to - /// [`ConvergenceConfig::window_size`] capacity. - /// /// The detector starts with an empty window and a zero consecutive /// count — no convergence can be detected until at least /// `window_size` responses have been added. @@ -636,7 +433,7 @@ impl ConvergenceDetector { /// # Example /// /// ```rust - /// use loopctl::loop_control::convergence::{ConvergenceConfig, ConvergenceDetector}; + /// use loopctl::detection::convergence::{ConvergenceConfig, ConvergenceDetector}; /// /// let config = ConvergenceConfig { /// window_size: 5, @@ -644,13 +441,13 @@ impl ConvergenceDetector { /// ..Default::default() /// }; /// let detector = ConvergenceDetector::new(config)?; - /// # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) + /// # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) /// ``` /// /// # Errors /// - /// Returns [`ConvergenceConfigError::WindowTooSmall`] if - /// `window_size < 2`, or [`ConvergenceConfigError::ThresholdOutOfRange`] + /// Returns [`ConvergenceConfigError::WindowTooSmall`] if `window_size < 2`, + /// or [`ConvergenceConfigError::ThresholdOutOfRange`] /// if `similarity_threshold` is outside `[0.0, 1.0]`. pub fn new(config: ConvergenceConfig) -> Result { if config.window_size < 2 { @@ -692,24 +489,23 @@ impl ConvergenceDetector { /// # Example /// /// ```rust - /// use loopctl::loop_control::convergence::ConvergenceDetector; + /// use loopctl::detection::convergence::ConvergenceDetector; /// /// let detector = ConvergenceDetector::default_detector()?; - /// # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) + /// # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) /// ``` /// /// # Errors /// /// Returns [`ConvergenceConfigError`] if the default config fails - /// validation. This should never happen — the defaults are hard-coded - /// to be valid. + /// validation. pub fn default_detector() -> Result { Self::new(ConvergenceConfig::default()) } /// Add a response and check for convergence in one step. /// - /// This is the primary entry point. The response is compared against + /// Primary entry point. The response is compared against /// every prior response in the window. If any comparison exceeds /// [`ConvergenceConfig::similarity_threshold`], the consecutive count /// is incremented; otherwise it resets to `1`. @@ -728,34 +524,21 @@ impl ConvergenceDetector { /// If [`ConvergenceConfig::enabled`] is `false` or `response` is empty, /// returns [`ConvergenceStatus::no_convergence`] immediately. /// - /// # Side Effects - /// - /// Modifies the internal sliding window, consecutive counter, and - /// similar-responses collection. If the window is full, the oldest - /// entry is evicted. - /// - /// # Performance - /// - /// Each call performs `O(window_size)` similarity comparisons. - /// For a typical window of 3, this is negligible; for large windows - /// (e.g., 50), consider whether the overhead is acceptable. - /// /// # Example /// /// ```rust - /// use loopctl::loop_control::convergence::ConvergenceDetector; + /// use loopctl::detection::convergence::ConvergenceDetector; /// /// let mut detector = ConvergenceDetector::default_detector()?; /// let status = detector.add_response("Working on task..."); /// assert!(!status.detected); // First response, no comparison possible - /// # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) + /// # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) /// ``` pub fn add_response(&mut self, response: &str) -> ConvergenceStatus { if !self.config.enabled || response.is_empty() { return ConvergenceStatus::no_convergence(); } - // Check similarity with previous responses let mut max_similarity = 0.0; for prev_response in &self.window { let similarity = Self::compute_similarity(response, prev_response); @@ -769,20 +552,17 @@ impl ConvergenceDetector { self.similar_responses.push(response.to_string()); } } else { - // Response is different, reset counter self.consecutive_count = 1; self.similar_responses.clear(); self.similar_responses.push(response.to_string()); } } - // Add to window if self.window.len() >= self.config.window_size { self.window.pop_front(); } self.window.push_back(response.to_string()); - // Check if converged let detected = self.consecutive_count >= self.config.window_size; ConvergenceStatus { @@ -798,7 +578,7 @@ impl ConvergenceDetector { /// /// Inspects the current window and consecutive count to determine /// whether convergence has already been reached. Does not modify - /// internal state. + /// detector state. /// /// # Returns /// @@ -809,7 +589,7 @@ impl ConvergenceDetector { /// # Example /// /// ```rust - /// use loopctl::loop_control::convergence::{ConvergenceConfig, ConvergenceDetector}; + /// use loopctl::detection::convergence::{ConvergenceConfig, ConvergenceDetector}; /// /// let config = ConvergenceConfig { /// window_size: 3, @@ -822,7 +602,7 @@ impl ConvergenceDetector { /// detector.add_response("Response three about bananas"); /// let status = detector.check_convergence(); /// assert!(!status.detected); - /// # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) + /// # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) /// ``` #[must_use] pub fn check_convergence(&self) -> ConvergenceStatus { @@ -852,14 +632,14 @@ impl ConvergenceDetector { /// # Example /// /// ```rust - /// use loopctl::loop_control::convergence::ConvergenceDetector; + /// use loopctl::detection::convergence::ConvergenceDetector; /// /// let mut detector = ConvergenceDetector::default_detector()?; /// detector.add_response("task in progress"); /// detector.clear(); /// assert!(detector.window().is_empty()); /// assert_eq!(detector.consecutive_count(), 0); - /// # Ok::<(), loopctl::loop_control::convergence::ConvergenceConfigError>(()) + /// # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) /// ``` pub fn clear(&mut self) { self.window.clear(); @@ -874,7 +654,7 @@ impl ConvergenceDetector { /// /// # Returns /// - /// A reference to the internal [`VecDeque`] of response strings. + /// A reference to the [`VecDeque`] of response strings. /// The deque is ordered from oldest to newest. #[must_use] pub fn window(&self) -> &VecDeque { @@ -897,7 +677,7 @@ impl ConvergenceDetector { /// Get the number of consecutive similar responses. /// - /// This is the current streak counter. When it reaches + /// Current streak counter. When it reaches /// [`ConvergenceConfig::window_size`], convergence is detected. /// /// # Returns @@ -915,25 +695,7 @@ impl ConvergenceDetector { /// spaces), splits into word sets, then computes the ratio of /// intersection size to union size. /// - /// # Algorithm - /// - /// ```text - /// Jaccard(A, B) = |A ∩ B| / |A ∪ B| - /// - /// Example: A = {hello, world} - /// B = {hello, there} - /// Intersection = {hello} → |1| - /// Union = {hello, world, there} → |3| - /// Jaccard = 1/3 ≈ 0.33 - /// ``` - /// /// Returns `0.0` if either input is empty. - /// - /// # Normalization - /// - /// Before comparison, both strings are passed through - /// the internal `normalize_text` helper to produce a canonical - /// lowercased, alphanumeric-only form. #[allow(clippy::cast_precision_loss)] #[must_use] pub fn compute_similarity(a: &str, b: &str) -> f32 { @@ -962,11 +724,6 @@ impl ConvergenceDetector { intersection as f32 / union as f32 } - /// Normalize text for similarity comparison. - /// - /// Lowercases the input and replaces non-alphanumeric characters with - /// spaces, producing a canonical form suitable for word-level Jaccard - /// comparison via [`ConvergenceDetector::compute_similarity`]. fn normalize_text(text: &str) -> String { text.to_lowercase() .chars() @@ -978,16 +735,10 @@ impl ConvergenceDetector { impl Default for ConvergenceDetector { /// Produce a [`ConvergenceDetector`] with the default [`ConvergenceConfig`]. /// - /// Constructs the detector directly using struct-literal syntax, bypassing - /// [`ConvergenceDetector::new`]'s validation. This is safe because the - /// default config uses `window_size: 3` (≥ 2) and - /// `similarity_threshold: 0.95` (in `0.0..=1.0`), which satisfy all - /// validation invariants. - /// /// This impl exists so that parent types (e.g. [`DetectionManager`]) can /// derive or delegate `Default` without going through a fallible constructor. /// - /// [`DetectionManager`]: crate::loop_control::detection::DetectionManager + /// [`DetectionManager`]: crate::detection::manager::DetectionManager fn default() -> Self { let config = ConvergenceConfig::default(); let window_capacity = config.window_size; diff --git a/src/loop_control/loop_detector.rs b/src/detection/loop_detector.rs similarity index 75% rename from src/loop_control/loop_detector.rs rename to src/detection/loop_detector.rs index 162dcb1..a764dfe 100644 --- a/src/loop_control/loop_detector.rs +++ b/src/detection/loop_detector.rs @@ -11,9 +11,9 @@ //! Autonomous agents can get trapped in repetitive cycles — for example, //! repeatedly reading a file and then attempting the same edit that failed //! before. Without detection, the agent wastes tokens and time until an -//! external timeout kicks in. This module provides an early-warning system -//! that flags loops after just a few repetitions, and can force-stop the -//! agent when the repetition count becomes dangerous. +//! external timeout kicks in. An early-warning system +//! flags loops early and can force-stop the +//! agent when the repetition count exceeds a threshold. //! //! # Core Algorithm //! @@ -47,7 +47,7 @@ //! //! ```rust //! use std::sync::Arc; -//! use loopctl::loop_control::loop_detector::{ +//! use loopctl::detection::loop_detector::{ //! LoopDetector, LoopDetectorConfig, Operation, ToolSignature, //! }; //! @@ -66,32 +66,6 @@ //! } //! ``` //! -//! # Data Flow -//! -//! ```text -//! ┌──────────────┐ record() ┌─────────────────────┐ -//! │ Framework │ ────────────────► │ LoopDetector │ -//! │ (tool call) │ │ ┌───────────────┐ │ -//! └──────────────┘ │ │ sliding window│ │ -//! │ │ (VecDeque) │ │ -//! │ └───────────────┘ │ -//! │ ┌───────────────┐ │ -//! │ │ warned_ops │ │ -//! │ │ (HashSet) │ │ -//! │ └───────────────┘ │ -//! └─────────┬───────────┘ -//! │ -//! check_loop()/check_turn_limit() -//! │ -//! ▼ -//! ┌─────────────────────┐ -//! │ LoopStatus │ -//! │ is_looping │ -//! │ should_stop │ -//! │ warning │ -//! └─────────────────────┘ -//! ``` -//! //! # Edit-Recovery Workflow //! //! The detector includes special handling for the edit-recovery pattern. @@ -99,14 +73,6 @@ //! file to get updated contents, the loop warning for that file is cleared //! because the agent is making progress. This prevents false positives //! during normal edit-retry cycles. -//! -//! ```text -//! Edit(file, old_text) → FAIL (recoverable) -//! ↓ -//! Read(file) → check_and_reset_on_file_read() clears warning -//! ↓ -//! Edit(file, new_text) → SUCCESS (not counted as a loop) -//! ``` use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; @@ -129,7 +95,7 @@ use std::sync::{Arc, Mutex}; /// Rather than hard-coding tool names and JSON paths inside the detector, /// the signature pattern keeps the detector agnostic to any particular /// tool set. This allows the same [`LoopDetector`] to work with any -/// collection of tools — just swap the signature implementation. +/// collection of tools — swap the signature implementation as needed. /// /// # Implementing /// @@ -157,7 +123,7 @@ use std::sync::{Arc, Mutex}; /// # Example /// /// ```rust -/// use loopctl::loop_control::loop_detector::ToolSignature; +/// use loopctl::detection::loop_detector::ToolSignature; /// /// struct MyToolSignature; /// @@ -213,7 +179,7 @@ pub trait ToolSignature: Send + Sync { /// has changed — the agent is making progress, not looping. Returning /// `true` prevents the detector from flagging the sequence as a loop. /// - /// This is one of the most impactful methods to override correctly. + /// Overriding this method correctly is critical. /// Without recoverable-error detection, the detector will flag every /// retry cycle as a loop, even when the agent is following a healthy /// read-fail-retry pattern. @@ -230,7 +196,7 @@ pub trait ToolSignature: Send + Sync { /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::ToolSignature; + /// use loopctl::detection::loop_detector::ToolSignature; /// /// struct MySig; /// impl ToolSignature for MySig { @@ -293,7 +259,7 @@ pub trait ToolSignature: Send + Sync { /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::ToolSignature; + /// use loopctl::detection::loop_detector::ToolSignature; /// /// struct MySig; /// impl ToolSignature for MySig { @@ -319,8 +285,7 @@ pub trait ToolSignature: Send + Sync { /// /// # When Called /// - /// Called during [`LoopDetector::check_file_reads`] and - /// [`LoopDetector::check_and_reset_on_file_read`]. + /// Called during [`LoopDetector::check_file_reads`]. /// /// # Default /// @@ -338,9 +303,7 @@ pub trait ToolSignature: Send + Sync { /// /// # When Called /// - /// Called during [`LoopDetector::record`] to detect recoverable edits, - /// and during [`LoopDetector::check_and_reset_on_file_read`] to - /// identify failed edits in the operation history. + /// Called during [`LoopDetector::record`] to detect recoverable edits. /// /// # Default /// @@ -385,14 +348,12 @@ pub trait ToolSignature: Send + Sync { /// /// Called during edit-recovery logic in /// [`LoopDetector::record_from_input_with_error`] to match edit - /// operations to the same file, and during - /// [`LoopDetector::check_and_reset_on_file_read`] to find prior - /// warned edits for the file being read. + /// operations to the same file. /// /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::ToolSignature; + /// use loopctl::detection::loop_detector::ToolSignature; /// /// struct MySig; /// impl ToolSignature for MySig { @@ -424,7 +385,7 @@ pub trait ToolSignature: Send + Sync { /// Because [`extract_primary_param`](ToolSignature::extract_primary_param) /// always returns an empty string, *every* invocation of the same tool /// with the same result hash is considered identical. This means the -/// detector can still catch loops — it just can't distinguish between +/// detector can still catch loops — it cannot distinguish between /// different targets within the same tool. /// /// # When to Use @@ -447,7 +408,7 @@ pub trait ToolSignature: Send + Sync { /// /// ```rust /// use std::sync::Arc; -/// use loopctl::loop_control::loop_detector::{ +/// use loopctl::detection::loop_detector::{ /// LoopDetector, LoopDetectorConfig, NoOpToolSignature, /// }; /// @@ -466,8 +427,8 @@ pub struct NoOpToolSignature; /// returns `""`, every invocation of the same tool with the same result hash /// is considered identical for loop-detection purposes. /// -/// This implementation is intentionally empty — it relies entirely on the -/// default method bodies defined in the [`ToolSignature`] trait. See the +/// Empty implementation — relies entirely on the +/// [`ToolSignature`] trait defaults. See the /// trait-level documentation for the semantics of each default. impl ToolSignature for NoOpToolSignature {} @@ -483,7 +444,7 @@ impl ToolSignature for NoOpToolSignature {} /// then override individual fields as needed: /// /// ```rust -/// use loopctl::loop_control::loop_detector::LoopDetectorConfig; +/// use loopctl::detection::loop_detector::LoopDetectorConfig; /// use std::collections::HashMap; /// /// let config = LoopDetectorConfig { @@ -525,65 +486,32 @@ impl ToolSignature for NoOpToolSignature {} /// `"Grep" → 5`). #[derive(Debug, Clone)] pub struct LoopDetectorConfig { - /// Maximum number of operations kept in the sliding window. - /// - /// The [`LoopDetector`] maintains a [`VecDeque`] of recent operations. - /// When the deque reaches this size, the oldest entry is evicted before - /// a new one is appended. A larger window detects slower loops; a - /// smaller window is more memory-efficient and focuses on recent - /// activity. + /// Number of recent operations kept in the sliding window. /// /// **Default:** `50`. pub window_size: usize, - /// Number of identical repetitions required to flag a loop. - /// - /// An operation is considered "looping" when it appears at least this - /// many times (with the same [`Operation::result_hash`]) within the - /// current window. Can be overridden per tool via - /// [`tool_thresholds`](LoopDetectorConfig::tool_thresholds). + /// Identical repetitions required to flag a loop. Overridden per tool by `tool_thresholds`. /// /// **Default:** `3`. pub repetition_threshold: usize, - /// Maximum number of tool calls allowed in a single turn. - /// - /// Once the per-turn call count reaches this limit, - /// [`LoopDetector::check_turn_limit`] returns `true`. The turn counter - /// is reset by [`LoopDetector::reset_turn`], which the framework calls - /// at the start of each new turn. + /// Max tool calls per turn. Checked by [`check_turn_limit`](LoopDetector::check_turn_limit). /// /// **Default:** `9999` (effectively unlimited). pub max_tools_per_turn: usize, - /// Maximum number of identical file reads before a warning is raised. - /// - /// Checked by [`LoopDetector::check_file_reads`]. When a single file - /// path appears in more than this many read-type operations within the - /// window, the method returns `true`. + /// Max identical file reads before a warning. Checked by [`check_file_reads`](LoopDetector::check_file_reads). /// /// **Default:** `5`. pub max_same_file_reads: usize, - /// Number of repetitions required to force-stop the agent. - /// - /// When [`LoopDetector::check_loop`] detects repetitions ≥ this value, - /// it sets [`LoopStatus::should_stop`] to `true` and includes - /// `"STOPPING to prevent infinite loop"` in the warning message. Set - /// to `0` to disable forced stops entirely (the detector will still - /// issue warnings). + /// Repetitions required to force-stop the agent. Set to `0` to disable forced stops. /// /// **Default:** `10`. pub stop_threshold: usize, - /// Tool-specific repetition thresholds that override - /// [`repetition_threshold`](LoopDetectorConfig::repetition_threshold). - /// - /// Map keys are tool names (e.g. `"Edit"`, `"MultiEdit"`) and values - /// are the repetition count that triggers a loop for that tool. Looked - /// up by [`LoopDetectorConfig::threshold_for_tool`]. Consumers should - /// populate this with their tool names and desired thresholds. The - /// framework default is an empty map (no overrides). + /// Per-tool repetition thresholds. Looked up by [`threshold_for_tool`](LoopDetectorConfig::threshold_for_tool). /// /// **Default:** empty `HashMap`. pub tool_thresholds: HashMap, @@ -611,7 +539,7 @@ pub struct LoopDetectorConfig { /// # Example /// /// ```rust -/// use loopctl::loop_control::loop_detector::LoopDetectorConfig; +/// use loopctl::detection::loop_detector::LoopDetectorConfig; /// /// let config = LoopDetectorConfig::default(); /// assert_eq!(config.window_size, 50); @@ -627,11 +555,6 @@ pub struct LoopDetectorConfig { /// - [`LoopDetector::new`] — constructs a detector from a config. /// - [`LoopDetectorConfig::threshold_for_tool`] — per-tool threshold lookup. impl Default for LoopDetectorConfig { - /// Build a config with the default values described in the trait-level docs. - /// - /// All fields are set to their documented defaults. The operation window - /// is pre-allocated with capacity matching - /// [`window_size`](LoopDetectorConfig::window_size). fn default() -> Self { Self { window_size: 50, @@ -644,18 +567,6 @@ impl Default for LoopDetectorConfig { } } -/// Accessor methods for [`LoopDetectorConfig`]. -/// -/// This `impl` block provides helpers for looking up configuration values -/// with fallback behaviour (e.g. per-tool thresholds that delegate to the -/// generic default when no override is set). -/// -/// # Methods -/// -/// - [`threshold_for_tool`](LoopDetectorConfig::threshold_for_tool) — Returns -/// the effective repetition threshold for a given tool, checking per-tool -/// overrides first and falling back to the generic -/// [`repetition_threshold`](LoopDetectorConfig::repetition_threshold). impl LoopDetectorConfig { /// Get the effective repetition threshold for a specific tool. /// @@ -671,7 +582,7 @@ impl LoopDetectorConfig { /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::LoopDetectorConfig; + /// use loopctl::detection::loop_detector::LoopDetectorConfig; /// /// let mut config = LoopDetectorConfig::default(); /// config.tool_thresholds.insert("Edit".to_string(), 2); @@ -698,24 +609,16 @@ impl LoopDetectorConfig { /// /// # Equality Semantics /// -/// [`Operation`] derives [`PartialEq`] and [`Hash`], so two operations are +/// [`Operation`] derives [`Eq`] and [`Hash`], so two operations are /// equal only when all three fields match exactly. This means the *same* /// command producing *different* results is treated as two distinct /// operations — a key design choice that prevents false positives when /// the agent is making progress. /// -/// # Derives -/// -/// - [`Debug`] — For logging and diagnostics. -/// - [`Clone`] — Operations are stored in [`std::collections::HashSet`]s and [`VecDeque`]s -/// which may require cloning during retention scans. -/// - [`PartialEq`] + [`Eq`] — For equality comparison in loop counting. -/// - [`Hash`] — For use as keys in [`HashMap`]s during repetition counting. -/// /// # Construction /// /// ```rust -/// use loopctl::loop_control::loop_detector::{Operation, hash_result}; +/// use loopctl::detection::loop_detector::{Operation, hash_result}; /// /// // Simple construction (no result hash): /// let op = Operation::new("Read", "/src/main.rs"); @@ -734,38 +637,18 @@ impl LoopDetectorConfig { /// operations producing *different* outputs are not counted as repetitions. #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Operation { - /// Name of the tool that was invoked (e.g. `"Read"`, `"Edit"`, `"Bash"`). - /// - /// Used as the first component of the loop-detection key. Together with - /// [`primary_param`](Operation::primary_param) and - /// [`result_hash`](Operation::result_hash) it uniquely identifies a - /// repeated invocation pattern. + /// Tool name (e.g. `"Read"`, `"Edit"`, `"Bash"`). pub tool: String, - - /// Primary parameter that identifies the operation's target. - /// - /// Extracted by [`ToolSignature::extract_primary_param`]. Typically a - /// file path (for Read/Edit) or a command string (for Bash). Two - /// operations with the same `tool` but different `primary_param` are - /// *not* considered a loop (they target different resources). + /// Operation target (file path, command, etc.). Extracted by [`ToolSignature::extract_primary_param`]. pub primary_param: String, - - /// Hash of the tool result content, for result-aware loop detection. - /// - /// When `Some(hash)`, two operations are only considered identical if - /// they produced the same result. When `None`, results are not taken - /// into account — the detector relies solely on `(tool, primary_param)`. - /// Set via [`Operation::with_result_hash`] or the constructor - /// [`Operation::from_input_with_result_and_signature`]. - /// - /// Generated by the free function [`hash_result`]. + /// Hash of the tool result. `None` means results are ignored. Generated by [`hash_result`]. pub result_hash: Option, } /// Constructors and builder methods for [`Operation`]. /// /// Operations can be created in several ways depending on what information -/// is available at call sites: +/// is available: /// /// - **[`Operation::new`]** — Simplest path: tool name + primary param. /// - **[`Operation::from_input_with_signature`]** — Parses the primary @@ -774,17 +657,6 @@ pub struct Operation { /// construction with result hash, used after a tool invocation completes. /// - **[`Operation::with_result_hash`]** — Builder-style attachment of a /// result hash to an existing operation. -/// -/// # Construction Decision Tree -/// -/// ```text -/// Do you have the result yet? -/// ├── NO → Operation::from_input_with_signature(tool, input, sig) -/// │ or Operation::new(tool, param) -/// └── YES → Operation::from_input_with_result_and_signature( -/// tool, input, hash, sig) -/// or Operation::new(tool, param).with_result_hash(hash) -/// ``` impl Operation { /// Create a new operation with the given tool name and primary parameter. /// @@ -798,7 +670,7 @@ impl Operation { /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::Operation; + /// use loopctl::detection::loop_detector::Operation; /// /// let op = Operation::new("Read", "/src/main.rs"); /// assert_eq!(op.tool, "Read"); @@ -837,7 +709,7 @@ impl Operation { /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::{Operation, ToolSignature}; + /// use loopctl::detection::loop_detector::{Operation, ToolSignature}; /// /// struct MyToolSignature; /// impl ToolSignature for MyToolSignature { @@ -871,7 +743,7 @@ impl Operation { /// Create an operation from tool name, JSON input, result hash, and signature. /// /// Combines [`ToolSignature::extract_primary_param`] with an explicit - /// result hash into a single constructor call. This is the most complete + /// result hash into a single constructor call. Most complete /// construction path, used when both the input and the result are known. /// /// The `result_hash` should be computed by [`hash_result`] or set to @@ -885,7 +757,7 @@ impl Operation { /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::{Operation, ToolSignature, hash_result}; + /// use loopctl::detection::loop_detector::{Operation, ToolSignature, hash_result}; /// /// struct MyToolSignature; /// impl ToolSignature for MyToolSignature { @@ -936,7 +808,7 @@ impl Operation { /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::{Operation, hash_result}; + /// use loopctl::detection::loop_detector::{Operation, hash_result}; /// /// let hash = hash_result("file contents"); /// let op = Operation::new("Read", "/src/main.rs").with_result_hash(hash); @@ -966,7 +838,7 @@ impl Operation { /// equality comparison inside [`Operation::result_hash`]. /// /// Returns `None` if `content` is empty (no meaningful hash to compute). -/// This is intentional: an empty result usually means the tool produced +/// An empty result usually means the tool produced /// no output, and hashing it would add noise without value. /// /// # Use in Loop Detection @@ -980,11 +852,7 @@ impl Operation { /// /// # Performance /// -/// [`std::collections::hash_map::DefaultHasher`] is a fast, non-cryptographic hasher suitable for -/// runtime use. It is *not* suitable for security purposes (e.g. storing -/// passwords), but it is perfect for quick equality checks in hot paths. -/// The cost is O(n) in the length of `content`, which is acceptable because -/// tool results are typically bounded in size. +/// The cost is O(n) in the length of `content`. /// /// # Determinism /// @@ -995,7 +863,7 @@ impl Operation { /// # Example /// /// ```rust -/// use loopctl::loop_control::loop_detector::hash_result; +/// use loopctl::detection::loop_detector::hash_result; /// /// let h1 = hash_result("same output"); /// let h2 = hash_result("same output"); @@ -1031,37 +899,6 @@ pub fn hash_result(content: &str) -> Option { /// force-stopped. Also carries an optional human-readable /// [`warning`](LoopStatus::warning) message. /// -/// # Fields Overview -/// -/// - [`is_looping`](LoopStatus::is_looping) — `true` if any operation -/// exceeded its repetition threshold. -/// - [`repeated_operations`](LoopStatus::repeated_operations) — The -/// specific [`Operation`] values that triggered the loop. -/// - [`repetition_count`](LoopStatus::repetition_count) — How many times -/// the most-repeated operation appeared in the window. -/// - [`warning`](LoopStatus::warning) — Human-readable message for the -/// agent (or operator). `None` if no loop or already warned. -/// - [`should_stop`](LoopStatus::should_stop) — `true` when the -/// framework should halt the agent immediately. -/// -/// # Lifecycle -/// -/// ```text -/// check_loop() returns LoopStatus -/// │ -/// ├── is_looping = false -/// │ └── No action needed. Agent continues normally. -/// │ -/// └── is_looping = true -/// ├── should_stop = false -/// │ └── Warning emitted (if not already warned). -/// │ Agent should adjust behaviour. -/// │ -/// └── should_stop = true -/// └── STOPPING message included in warning. -/// Framework should halt the agent. -/// ``` -/// /// # Warning Deduplication /// /// The [`warning`](LoopStatus::warning) field is `None` when: @@ -1077,7 +914,7 @@ pub fn hash_result(content: &str) -> Option { /// # Example /// /// ```rust -/// use loopctl::loop_control::loop_detector::LoopDetector; +/// use loopctl::detection::loop_detector::LoopDetector; /// /// let detector = LoopDetector::default_detector(); /// let status = detector.check_loop(); @@ -1105,46 +942,15 @@ pub fn hash_result(content: &str) -> Option { /// is rarely useful for status objects. #[derive(Debug, Clone, Default)] pub struct LoopStatus { - /// Whether a loop was detected. - /// - /// `true` when at least one operation repeats beyond the applicable - /// threshold (either the generic - /// [`repetition_threshold`](LoopDetectorConfig::repetition_threshold) - /// or a per-tool override from - /// [`tool_thresholds`](LoopDetectorConfig::tool_thresholds)). + /// `true` when an operation repeats beyond the configured threshold. pub is_looping: bool, - - /// Operations that triggered the loop detection. - /// - /// Contains all operations whose repetition count equals the maximum - /// observed count and exceeds the threshold. When multiple operations - /// tie for the highest repetition count, all of them are included. - /// Empty when [`is_looping`](LoopStatus::is_looping) is `false`. + /// Operations tied for highest repetition. Empty when `is_looping` is `false`. pub repeated_operations: Vec, - - /// Number of repetitions of the most-repeated operation. - /// - /// Equals the count of the first entry in - /// [`repeated_operations`](LoopStatus::repeated_operations). Zero when - /// no loop was detected. + /// Repetitions of the most-repeated operation. Zero when no loop detected. pub repetition_count: usize, - - /// Human-readable warning message describing the detected loop. - /// - /// Contains the operation description, repetition count, a "STOPPING" - /// notice if [`should_stop`](LoopStatus::should_stop) is true, and a - /// tool-specific suggestion from - /// [`ToolSignature::get_suggestion`]. Set to `None` when no loop is - /// detected, or when the loop has already been warned about (to avoid - /// spamming the agent with duplicate warnings). + /// Loop description with count and suggestion. `None` when not looping or already warned. pub warning: Option, - - /// Whether the agent should be force-stopped due to severe looping. - /// - /// `true` when [`repetition_count`](LoopStatus::repetition_count) - /// reaches or exceeds [`LoopDetectorConfig::stop_threshold`] (and the - /// stop threshold is non-zero). The framework should halt the agent's - /// event loop when this is `true`. + /// `true` when `repetition_count >= stop_threshold` (non-zero). pub should_stop: bool, } @@ -1160,7 +966,7 @@ pub struct LoopStatus { /// /// ```rust /// use std::sync::Arc; -/// use loopctl::loop_control::loop_detector::{ +/// use loopctl::detection::loop_detector::{ /// LoopDetector, LoopDetectorConfig, NoOpToolSignature, /// }; /// @@ -1172,24 +978,7 @@ pub struct LoopStatus { /// let detector = LoopDetector::new(config, Arc::new(NoOpToolSignature)); /// /// // With custom tool signature: -/// // let detector = LoopDetector::with_signature(Arc::new(MyToolSignature)); -/// ``` -/// -/// # Lifecycle -/// -/// ```text -/// ┌─────────────────────────────────────────────────────┐ -/// │ For each tool call: │ -/// │ 1. record_from_input(tool, input, result_hash) │ -/// │ 2. check_loop() → LoopStatus │ -/// │ 3. check_turn_limit() → bool │ -/// │ │ -/// │ At turn boundary: │ -/// │ 4. reset_turn() │ -/// │ │ -/// │ At task boundary: │ -/// │ 5. reset() │ -/// └─────────────────────────────────────────────────────┘ +/// // let detector = LoopDetector::new_with_signature(Arc::new(MyToolSignature)); /// ``` /// /// # Thread Safety @@ -1198,76 +987,16 @@ pub struct LoopStatus { /// independently and handle lock poisoning gracefully (by skipping the /// operation rather than panicking). This means the detector degrades /// gracefully under contention but never blocks the agent loop. -/// -/// # Design Decisions -/// -/// - **Sliding window (not global history):** A bounded [`VecDeque`] keeps -/// memory usage predictable and focuses detection on *recent* behaviour, -/// avoiding false positives from operations far in the past. -/// - **Result-aware comparison:** Two invocations of the same tool on the -/// same target are only considered identical if they produced the same -/// result hash. This prevents the detector from flagging operations -/// where the agent is genuinely making progress. -/// - **Warning deduplication:** Once an operation has triggered a warning, -/// subsequent warnings are suppressed unless the result changes or the -/// stop threshold is reached. This prevents the agent's context from -/// being flooded with identical loop messages. -/// - **Per-turn limits:** Separate from repetition detection, the turn -/// counter catches runaway tool usage regardless of whether the tools -/// are repeating. This catches scenarios like "try 50 different bash -/// commands" that wouldn't trigger repetition-based detection. -/// -/// # Interior Fields -/// -/// The detector holds five pieces of internal state, each in its own -/// [`Mutex`] to minimise lock contention: -/// -/// | Field | Type | Purpose | -/// |----------------------|---------------------------|--------------------------------| -/// | `operations` | `Mutex>` | Sliding window of history | -/// | `config` | `LoopDetectorConfig` | Thresholds (immutable) | -/// | `turn_count` | `Mutex` | Per-turn call counter | -/// | `warned_operations` | `Mutex>` | Already-warned dedup set | -/// | `signature` | `Arc` | Tool-specific parsing logic | pub struct LoopDetector { - /// Sliding window of recent [`Operation`] records. - /// - /// Bounded by [`LoopDetectorConfig::window_size`]. When full, the - /// oldest operation is evicted before a new one is appended. The - /// window is scanned by [`LoopDetector::check_loop`] to find repeated - /// operations. + /// Sliding window of recent [`Operation`] records, bounded by [`LoopDetectorConfig::window_size`]. operations: Mutex>, - - /// Configuration controlling thresholds and limits. - /// - /// Set at construction time via [`LoopDetector::new`]. Immutable for - /// the lifetime of the detector. + /// Detector configuration (thresholds and limits). Immutable after construction. config: LoopDetectorConfig, - - /// Count of tool invocations in the current turn. - /// - /// Incremented by [`LoopDetector::record`] and reset to zero by - /// [`LoopDetector::reset_turn`]. Checked against - /// [`LoopDetectorConfig::max_tools_per_turn`] by - /// [`LoopDetector::check_turn_limit`]. + /// Per-turn tool call count. Reset by [`reset_turn`](LoopDetector::reset_turn). turn_count: Mutex, - - /// Set of operations that have already triggered a warning. - /// - /// Once an operation appears in this set, subsequent calls to - /// [`LoopDetector::check_loop`] will suppress the warning message - /// (returning `None` in [`LoopStatus::warning`]) to avoid spamming - /// the agent. Entries are cleared when the result changes (indicating - /// progress) or when [`LoopDetector::clear`] / [`LoopDetector::reset`] - /// is called. + /// Operations that already triggered a warning. Cleared on result change or [`reset`](LoopDetector::reset). warned_operations: Mutex>, - /// Tool signature for extracting tool-specific parameters. - /// - /// Wrapped in [`Arc`] because it is shared between the detector and - /// any code that needs to inspect the signature via - /// [`LoopDetector::signature`]. The trait object is `Send + Sync` so - /// it can be used from any thread. signature: Arc, } @@ -1287,15 +1016,11 @@ pub struct LoopDetector { impl LoopDetector { /// Create a new loop detector with the given configuration and tool signature. /// - /// Initialises an empty operation window, zero turn count, and an empty - /// warned-operations set. The `signature` is stored in an [`Arc`] for - /// shared access. - /// /// # Example /// /// ```rust /// use std::sync::Arc; - /// use loopctl::loop_control::loop_detector::{ + /// use loopctl::detection::loop_detector::{ /// LoopDetector, LoopDetectorConfig, NoOpToolSignature, /// }; /// @@ -1314,12 +1039,11 @@ impl LoopDetector { /// Create a detector with default configuration and a no-op tool signature. /// - /// Convenience constructor for simple use cases that don't need - /// tool-specific logic. Equivalent to: + /// Equivalent to: /// /// ```rust /// use std::sync::Arc; - /// use loopctl::loop_control::loop_detector::{LoopDetector, LoopDetectorConfig, NoOpToolSignature}; + /// use loopctl::detection::loop_detector::{LoopDetector, LoopDetectorConfig, NoOpToolSignature}; /// /// LoopDetector::new(LoopDetectorConfig::default(), Arc::new(NoOpToolSignature)); /// ``` @@ -1336,12 +1060,12 @@ impl LoopDetector { /// # Example /// /// ```rust - /// use loopctl::loop_control::loop_detector::LoopDetector; + /// use loopctl::detection::loop_detector::LoopDetector; /// /// // With a custom signature, you would write: - /// // let detector = LoopDetector::with_signature(Arc::new(MyToolSignature)); + /// // let detector = LoopDetector::new_with_signature(Arc::new(MyToolSignature)); /// ``` - pub fn with_signature(signature: Arc) -> Self { + pub fn new_with_signature(signature: Arc) -> Self { Self::new(LoopDetectorConfig::default(), signature) } @@ -1358,7 +1082,7 @@ impl LoopDetector { /// /// Uses the configured [`ToolSignature`] to extract the primary /// parameter from `input`, constructs an [`Operation`], and delegates - /// to [`record`](LoopDetector::record). This is a convenience wrapper + /// to [`record`](LoopDetector::record). Convenience wrapper /// around [`LoopDetector::record_from_input_with_error`] that passes `None` for the /// error parameter. /// @@ -1384,7 +1108,7 @@ impl LoopDetector { /// Record a tool invocation with an optional error string. /// - /// This is the primary entry point for the framework. It uses the + /// Primary entry point for the framework. Uses the /// configured [`ToolSignature`] to extract the primary parameter from /// `input`, constructs an [`Operation`], and records it. /// @@ -1410,7 +1134,7 @@ impl LoopDetector { /// /// ```rust /// use std::sync::Arc; - /// use loopctl::loop_control::loop_detector::{ + /// use loopctl::detection::loop_detector::{ /// LoopDetector, LoopDetectorConfig, ToolSignature, /// }; /// @@ -1444,8 +1168,6 @@ impl LoopDetector { result_hash, self.signature.as_ref(), ); - - // Determine recoverability accurately via the trait method. let is_recoverable = error.is_some_and(|err| self.signature.is_recoverable_error(tool, err)); @@ -1478,13 +1200,32 @@ impl LoopDetector { /// directly, via [`LoopDetector::record_from_input`], or via /// [`LoopDetector::record_from_input_with_error`]. pub fn record(&self, operation: Operation) { - // Check if this operation was previously warned with a different result - if let Ok(mut warned) = self.warned_operations.lock() { - warned.retain(|warned_op| { - !(warned_op.tool == operation.tool - && warned_op.primary_param == operation.primary_param - && warned_op.result_hash != operation.result_hash) - }); + let should_clear_history = { + let mut clear = false; + if let Ok(mut warned) = self.warned_operations.lock() { + warned.retain(|warned_op| { + let matches = warned_op.tool == operation.tool + && warned_op.primary_param == operation.primary_param + && warned_op.result_hash != operation.result_hash; + if matches { + clear = true; + } + !matches + }); + } + clear + }; + + // Remove stale operations from the sliding window so they don't + // re-trigger loop detection on the next check_loop() call. + if should_clear_history { + if let Ok(mut ops) = self.operations.lock() { + ops.retain(|op| { + !(op.tool == operation.tool + && op.primary_param == operation.primary_param + && op.result_hash != operation.result_hash) + }); + } } if let Ok(mut ops) = self.operations.lock() { @@ -1578,87 +1319,119 @@ impl LoopDetector { return LoopStatus::default(); }; - let mut repeated_operations = Vec::new(); - let mut max_repetitions = 0; - let mut op_counts: HashMap = HashMap::new(); - for op in ops.iter() { - op_counts + let (repeated_operations, max_repetitions) = + Self::find_repeated(&ops, |tool| self.config.threshold_for_tool(tool)); + let is_looping = !repeated_operations.is_empty(); + let should_stop = + self.config.stop_threshold > 0 && max_repetitions >= self.config.stop_threshold; + let warning = self.build_warning( + &repeated_operations, + max_repetitions, + is_looping, + should_stop, + ); + + LoopStatus { + is_looping, + repeated_operations, + repetition_count: max_repetitions, + warning, + should_stop, + } + } + + /// Count occurrences of each operation in the deque. + /// + /// Iterates over every operation in `ops` and builds a [`HashMap`] where + /// each key is a cloned [`Operation`] and the value is the number of + /// times it appears. Counts are capped at `usize::MAX` via saturating + /// addition to avoid overflow on extremely long deques. + fn count_operations(ops: &VecDeque) -> HashMap { + let mut counts: HashMap = HashMap::new(); + for op in ops { + counts .entry(op.clone()) .and_modify(|c| *c = c.saturating_add(1)) .or_insert(1); } - - for (op, count) in op_counts { - let threshold = self.config.threshold_for_tool(&op.tool); - if count >= threshold { - if count > max_repetitions { - max_repetitions = count; - repeated_operations.clear(); - repeated_operations.push(op); - } else if count == max_repetitions { - repeated_operations.push(op); + counts + } + + /// Find operations exceeding their per-tool threshold. + /// + /// Only the operations with the highest repetition count are + /// returned (ties included). The `threshold` closure maps a tool + /// name to its configured threshold. + fn find_repeated( + ops: &VecDeque, + threshold: impl Fn(&str) -> usize, + ) -> (Vec, usize) { + let counts = Self::count_operations(ops); + let mut repeated = Vec::new(); + let mut max = 0; + + for (op, count) in counts { + let t = threshold(&op.tool); + if count >= t { + if count > max { + max = count; + repeated.clear(); + repeated.push(op); + } else if count == max { + repeated.push(op); } } } - let is_looping = !repeated_operations.is_empty(); - let should_stop = - self.config.stop_threshold > 0 && max_repetitions >= self.config.stop_threshold; - let warning = if is_looping { - let already_warned = if let Some(first_op) = repeated_operations.first() { - if let Ok(warned) = self.warned_operations.lock() { - warned.contains(first_op) - } else { - false - } - } else { - false - }; - - if already_warned && !should_stop { - None - } else { - if let Some(first_op) = repeated_operations.first() - && let Ok(mut warned) = self.warned_operations.lock() - { - warned.insert(first_op.clone()); - } + (repeated, max) + } - let stop_msg = if should_stop { - " STOPPING to prevent infinite loop." - } else { - "" - }; - - let tool_name = repeated_operations.first().map_or("", |o| o.tool.as_str()); - - let suggestion = self - .signature - .get_suggestion(tool_name) - .unwrap_or_else(|| "Consider a different approach or tool.".to_string()); - - Some(format!( - "Loop detected: Operation '{}' repeated {} times with same result.{} {}", - repeated_operations - .first() - .map(|o| format!("{}({})", o.tool, o.primary_param)) - .unwrap_or_default(), - max_repetitions, - stop_msg, - suggestion - )) + /// Build the warning string for repeated operations. + /// + /// Returns `None` when not looping, or when already warned and + /// not stopping. When a new warning is produced, the first + /// repeated operation is recorded in `warned_operations`. + fn build_warning( + &self, + repeated_operations: &[Operation], + max_repetitions: usize, + is_looping: bool, + should_stop: bool, + ) -> Option { + if !is_looping { + return None; + } + + let first_op = repeated_operations.first()?; + let already_warned = self + .warned_operations + .lock() + .is_ok_and(|w| w.contains(first_op)); + if already_warned && !should_stop { + return None; + } + + if !already_warned { + if let Ok(mut warned) = self.warned_operations.lock() { + warned.insert(first_op.clone()); } + } + + let stop_msg = if should_stop { + " STOPPING to prevent infinite loop." } else { - None + "" }; - LoopStatus { - is_looping, - repeated_operations, - repetition_count: max_repetitions, - warning, - should_stop, - } + let suggestion = self + .signature + .get_suggestion(&first_op.tool) + .unwrap_or_else(|| "Consider a different approach or tool.".to_string()); + + Some(format!( + "Loop detected: Operation '{}({})' repeated {} times with same result.{} {}", + first_op.tool, first_op.primary_param, max_repetitions, stop_msg, suggestion + )) } /// Check whether the per-turn tool-call limit has been reached. @@ -1721,10 +1494,6 @@ impl LoopDetector { /// Called by the framework before dispatching a file-read tool call, /// to guard against excessive re-reading of the same file. /// - /// # Parameters - /// - /// - `file_path` — The file path (or a substring of it) to check. - /// /// # Returns /// /// `true` if the file has been read ≥ `max_same_file_reads` times, @@ -1737,62 +1506,15 @@ impl LoopDetector { let sig = &self.signature; let read_count = ops .iter() - .filter(|o| sig.is_file_read_tool(&o.tool) && o.primary_param.contains(file_path)) + .filter(|o| { + sig.is_file_read_tool(&o.tool) + && sig.normalize_param_for_comparison(&o.tool, &o.primary_param) == file_path + }) .count(); read_count >= self.config.max_same_file_reads } - /// Reset loop state when a file-read follows a failed edit to the same file. - /// - /// If `tool` is a read-type tool (per - /// [`ToolSignature::is_file_read_tool`]) and the operation window - /// contains a recent edit to `file_path` (per - /// [`ToolSignature::is_file_edit_tool`]), this method removes all - /// edit operations for that file from the window *and* from the - /// warned-operations set. The rationale is that the agent is - /// re-reading the file to get updated contents after a failed edit, - /// which constitutes progress rather than a loop. - /// - /// # When Called - /// - /// Called by the framework when a read tool is dispatched after a - /// failed edit, to prevent the edit-read cycle from being flagged as - /// a loop. - /// - /// # Parameters - /// - /// - `tool` — Name of the tool being dispatched. - /// - `file_path` — The file being read. - pub fn check_and_reset_on_file_read(&self, tool: &str, file_path: &str) { - if !self.signature.is_file_read_tool(tool) { - return; - } - - let sig = &self.signature; - if let Ok(mut ops) = self.operations.lock() { - let has_recent_failed_edit = ops.iter().any(|op| { - sig.is_file_edit_tool(&op.tool) - && sig.normalize_param_for_comparison(&op.tool, &op.primary_param) == file_path - }); - - if has_recent_failed_edit { - ops.retain(|op| { - let op_file = sig.normalize_param_for_comparison(&op.tool, &op.primary_param); - !(sig.is_file_edit_tool(&op.tool) && op_file == file_path) - }); - } - } - - if let Ok(mut warned) = self.warned_operations.lock() { - let sig = &self.signature; - warned.retain(|op| { - let op_file = sig.normalize_param_for_comparison(&op.tool, &op.primary_param); - !(sig.is_file_edit_tool(&op.tool) && op_file == file_path) - }); - } - } - /// Clear all recorded operations and warned-operation state. /// /// Removes every entry from the sliding window and the warned-operations @@ -1840,8 +1562,8 @@ impl LoopDetector { /// Produce a [`LoopDetector`] with default configuration and a no-op signature. /// -/// Delegates to [`LoopDetector::default_detector`]. This is the same as -/// calling `LoopDetector::default_detector()` and is provided for +/// Delegates to [`LoopDetector::default_detector`]. Same as +/// calling `LoopDetector::default_detector()` and provided for /// ergonomic compatibility with generic code that uses `Default`. /// /// The resulting detector uses [`LoopDetectorConfig::default`] thresholds @@ -1851,7 +1573,7 @@ impl LoopDetector { /// # Example /// /// ```rust -/// use loopctl::loop_control::loop_detector::LoopDetector; +/// use loopctl::detection::loop_detector::LoopDetector; /// /// let detector = LoopDetector::default(); /// let status = detector.check_loop(); @@ -1861,13 +1583,9 @@ impl LoopDetector { /// # See Also /// /// - [`LoopDetector::new`] — for custom configuration. -/// - [`LoopDetector::with_signature`] — for custom tool signatures. +/// - [`LoopDetector::new_with_signature`] — for custom tool signatures. /// - [`LoopDetector::default_detector`] — the method this delegates to. impl Default for LoopDetector { - /// Build a detector with [`LoopDetectorConfig::default`] and [`NoOpToolSignature`]. - /// - /// Equivalent to `LoopDetector::default_detector()`. The internal state - /// is empty (no operations recorded, zero turn count, no warnings). fn default() -> Self { Self::default_detector() } @@ -1878,12 +1596,12 @@ impl Default for LoopDetector { /// Provides a process-wide [`LoopDetector`] that can be accessed from /// anywhere via [`global_detector`]. The detector is created exactly once /// with default configuration ([`LoopDetectorConfig::default`]) and a -/// [`NoOpToolSignature`]. This is useful for simple agents that don't +/// [`NoOpToolSignature`]. Useful for simple agents that don't /// need tool-specific loop detection logic. /// /// For production use with custom tool signatures, prefer constructing a /// dedicated [`LoopDetector`] via [`LoopDetector::new`] or -/// [`LoopDetector::with_signature`] instead of relying on this global. +/// [`LoopDetector::new_with_signature`] instead of relying on this global. /// /// # Thread Safety /// @@ -1906,7 +1624,7 @@ static GLOBAL_DETECTOR: std::sync::OnceLock> = std::sync::Once /// # Example /// /// ```rust -/// use loopctl::loop_control::loop_detector::{global_detector, Operation}; +/// use loopctl::detection::loop_detector::{global_detector, Operation}; /// /// let detector = global_detector(); /// detector.record(Operation::new("Read", "/src/main.rs")); @@ -1917,7 +1635,7 @@ static GLOBAL_DETECTOR: std::sync::OnceLock> = std::sync::Once /// # See Also /// /// - [`LoopDetector::new`] — for custom configuration. -/// - [`LoopDetector::with_signature`] — for custom tool signatures. +/// - [`LoopDetector::new_with_signature`] — for custom tool signatures. pub fn global_detector() -> Arc { Arc::clone(GLOBAL_DETECTOR.get_or_init(|| Arc::new(LoopDetector::default_detector()))) } @@ -2373,7 +2091,7 @@ mod tests { #[test] fn test_detector_with_custom_signature() { - let detector = LoopDetector::with_signature(Arc::new(TestToolSignature)); + let detector = LoopDetector::new_with_signature(Arc::new(TestToolSignature)); detector.record(Operation::new("Read", "/test.txt")); detector.record(Operation::new("Read", "/test.txt")); @@ -2420,6 +2138,14 @@ mod tests { let hash2 = hash_result("different output - progress!"); detector.record(Operation::new("Bash", "git status").with_result_hash(hash2)); + // After recording with a different hash, both the warned set and + // the sliding window are cleared, so the warning should be gone. + let status_cleared = detector.check_loop(); + assert!( + status_cleared.warning.is_none(), + "warning should be cleared when result hash changes" + ); + for _ in 0..3 { detector.record(Operation::new("Bash", "git status").with_result_hash(hash1)); } @@ -2500,6 +2226,30 @@ mod tests { assert_eq!(status.repetition_count, 3); } + #[test] + fn test_warning_not_cleared_when_result_stays_same() { + let detector = test_detector(); + + let hash1 = hash_result("same output"); + for _ in 0..3 { + detector.record(Operation::new("Bash", "git status").with_result_hash(hash1)); + } + + let status1 = detector.check_loop(); + assert!(status1.warning.is_some()); + + // Recording again with the SAME hash should NOT clear the history. + // check_loop() suppresses already-warned ops (returns None), but the + // loop is still detected (is_looping = true) and the history is intact. + detector.record(Operation::new("Bash", "git status").with_result_hash(hash1)); + + let status2 = detector.check_loop(); + assert!( + status2.is_looping, + "loop should still be detected when result hash stays the same" + ); + } + #[test] fn test_suggestion_from_signature() { let detector = test_detector(); @@ -2517,4 +2267,203 @@ mod tests { warning ); } + + fn make_ops(pairs: &[(&str, &str)]) -> VecDeque { + pairs + .iter() + .map(|&(tool, param)| Operation::new(tool, param)) + .collect() + } + + fn make_ops_hashed(pairs: &[(&str, &str, &str)]) -> VecDeque { + pairs + .iter() + .map(|&(tool, param, result)| { + Operation::new(tool, param).with_result_hash(hash_result(result)) + }) + .collect() + } + + #[test] + fn test_count_operations_empty() { + let ops: VecDeque = VecDeque::new(); + let counts = LoopDetector::count_operations(&ops); + assert!(counts.is_empty()); + } + + #[test] + fn test_count_operations_single_op() { + let ops = make_ops(&[("Bash", "ls"), ("Bash", "ls"), ("Bash", "ls")]); + let counts = LoopDetector::count_operations(&ops); + assert_eq!(counts.len(), 1); + assert_eq!(counts.get(&Operation::new("Bash", "ls")), Some(&3)); + } + + #[test] + fn test_count_operations_distinct_ops() { + let ops = make_ops(&[("Bash", "ls"), ("Read", "file.txt"), ("Bash", "ls")]); + let counts = LoopDetector::count_operations(&ops); + assert_eq!(counts.len(), 2); + assert_eq!(counts.get(&Operation::new("Bash", "ls")), Some(&2)); + assert_eq!(counts.get(&Operation::new("Read", "file.txt")), Some(&1)); + } + + #[test] + fn test_count_operations_different_hashes_are_distinct() { + let ops = make_ops_hashed(&[("Bash", "ls", "output_a"), ("Bash", "ls", "output_b")]); + let counts = LoopDetector::count_operations(&ops); + // Different result hashes → different operations + assert_eq!(counts.len(), 2); + } + + #[test] + fn test_count_operations_same_hashes_are_grouped() { + let ops = make_ops_hashed(&[ + ("Bash", "ls", "same_output"), + ("Bash", "ls", "same_output"), + ("Bash", "ls", "same_output"), + ]); + let counts = LoopDetector::count_operations(&ops); + assert_eq!(counts.len(), 1); + let key = Operation::new("Bash", "ls").with_result_hash(hash_result("same_output")); + assert_eq!(counts.get(&key), Some(&3)); + } + + #[test] + fn test_find_repeated_none() { + let ops = make_ops(&[("Bash", "ls"), ("Read", "f.txt")]); + let (repeated, max) = LoopDetector::find_repeated(&ops, |_tool| 3); + assert!(repeated.is_empty()); + assert_eq!(max, 0); + } + + #[test] + fn test_find_repeated_single_above_threshold() { + let ops = make_ops(&[("Bash", "ls"); 5]); + let (repeated, max) = LoopDetector::find_repeated(&ops, |_tool| 3); + assert_eq!(max, 5); + assert_eq!(repeated.len(), 1); + assert_eq!(repeated[0], Operation::new("Bash", "ls")); + } + + #[test] + fn test_find_repeated_tie_keeps_both() { + let mut ops = VecDeque::new(); + for _ in 0..3 { + ops.push_back(Operation::new("Bash", "ls")); + } + for _ in 0..3 { + ops.push_back(Operation::new("Read", "f.txt")); + } + let (repeated, max) = LoopDetector::find_repeated(&ops, |_tool| 3); + assert_eq!(max, 3); + assert_eq!(repeated.len(), 2); + } + + #[test] + fn test_find_repeated_per_tool_threshold() { + let mut ops = VecDeque::new(); + for _ in 0..2 { + ops.push_back(Operation::new("Bash", "ls")); + } + for _ in 0..5 { + ops.push_back(Operation::new("Read", "f.txt")); + } + // Bash threshold = 3, Read threshold = 4 + let (repeated, max) = + LoopDetector::find_repeated(&ops, |tool| if tool == "Bash" { 3 } else { 4 }); + // Bash(2) < 3 → excluded. Read(5) >= 4 → included. + assert_eq!(max, 5); + assert_eq!(repeated.len(), 1); + assert_eq!(repeated[0].tool, "Read"); + } + + #[test] + fn test_find_repeated_higher_count_wins() { + let mut ops = VecDeque::new(); + for _ in 0..5 { + ops.push_back(Operation::new("Bash", "ls")); + } + for _ in 0..3 { + ops.push_back(Operation::new("Read", "f.txt")); + } + let (repeated, max) = LoopDetector::find_repeated(&ops, |_tool| 2); + // Only Bash(5) wins — Read(3) is below max + assert_eq!(max, 5); + assert_eq!(repeated.len(), 1); + assert_eq!(repeated[0].tool, "Bash"); + } + + #[test] + fn test_build_warning_not_looping() { + let detector = test_detector(); + let warning = detector.build_warning(&[], 0, false, false); + assert!(warning.is_none()); + } + + #[test] + fn test_build_warning_first_warning() { + let detector = test_detector(); + let ops = vec![Operation::new("Bash", "ls")]; + let warning = detector.build_warning(&ops, 3, true, false); + assert!(warning.is_some()); + let msg = warning.unwrap(); + assert!(msg.contains("Bash(ls)")); + assert!(msg.contains("3 times")); + assert!(!msg.contains("STOPPING")); + } + + #[test] + fn test_build_warning_includes_stop_message() { + let detector = test_detector(); + let ops = vec![Operation::new("Bash", "ls")]; + let warning = detector.build_warning(&ops, 5, true, true); + assert!(warning.is_some()); + let msg = warning.unwrap(); + assert!(msg.contains("STOPPING")); + } + + #[test] + fn test_build_warning_suppresses_duplicate() { + let detector = test_detector(); + let op = Operation::new("Bash", "ls"); + let ops = vec![op.clone()]; + + // First call produces a warning and records the op as warned + let w1 = detector.build_warning(&ops, 3, true, false); + assert!(w1.is_some()); + + // Second call suppresses because already warned + let w2 = detector.build_warning(&ops, 3, true, false); + assert!(w2.is_none()); + } + + #[test] + fn test_build_warning_duplicate_not_suppressed_when_stopping() { + let detector = test_detector(); + let op = Operation::new("Bash", "ls"); + let ops = vec![op.clone()]; + + // First warning + let w1 = detector.build_warning(&ops, 3, true, false); + assert!(w1.is_some()); + + // Second call with should_stop=true still produces a warning + let w2 = detector.build_warning(&ops, 3, true, true); + assert!(w2.is_some()); + assert!(w2.unwrap().contains("STOPPING")); + } + + #[test] + fn test_build_warning_includes_suggestion() { + let detector = test_detector(); + let ops = vec![Operation::new("Bash", "git status")]; + let warning = detector.build_warning(&ops, 3, true, false); + assert!(warning.is_some()); + let msg = warning.unwrap(); + assert!( + msg.contains("Check the command"), + "Should contain tool signature suggestion: {msg}" + ); + } } diff --git a/src/loop_control/detection.rs b/src/detection/manager.rs similarity index 71% rename from src/loop_control/detection.rs rename to src/detection/manager.rs index 67d49f5..b357824 100644 --- a/src/loop_control/detection.rs +++ b/src/detection/manager.rs @@ -1,6 +1,6 @@ //! Detection manager — loop and convergence detection for agent behavior. //! -//! This module provides the [`DetectionManager`] which unifies two complementary +//! [`DetectionManager`] unifies two complementary //! detection strategies into a single manager that agents consult after each turn //! to decide whether they are making progress or spinning in circles: //! @@ -19,32 +19,34 @@ //! //! # Architecture //! -//! ```text -//! DetectionManager -//! ┌───────────────────────────┐ -//! │ config: DetectionConfig │ -//! │ stats: DetectionStats │ -//! └─────┬────────────┬────────┘ -//! │ │ -//! ┌─────────────┘ └───────────┐ -//! ▼ ▼ -//! ┌──────────────┐ ┌──────────────────┐ -//! │ LoopDetector │ │ ConvergenceDet. │ -//! │ (tool call │ │ (response text │ -//! │ cycles) │ │ similarity) │ -//! └──────┬───────┘ └───────┬──────────┘ -//! │ │ -//! record_tool_call() record_response() -//! record_operation() check_convergence() -//! check_loop() │ -//! │ │ -//! └──────────────────┬────────────────────┘ -//! ▼ -//! DetectedPattern -//! ┌──────────────┼──────────────────┐ -//! ▼ ▼ ▼ -//! NoPattern LoopDetected ConvergenceDetected -//! ``` +//! `DetectionManager` is a **facade** that owns one `LoopDetector` and one +//! `ConvergenceDetector`, forwarding calls to each and merging their results +//! into a unified [`DetectedPattern`] enum. +//! +//! **Data flow** — On every agent turn the framework feeds two kinds of +//! telemetry into the manager: +//! +//! 1. *Tool calls* are sent to `record_operation` (or its convenience +//! wrappers `record_tool_call` / `record_tool_call_with_result`), which +//! hands the operation to the `LoopDetector`. The loop detector compares +//! consecutive operations by tool name, primary parameter, and optional +//! result hash; when the same signature repeats ≥ `loop_threshold` +//! times it reports a loop. +//! +//! 2. *Assistant responses* (free text) are sent to `record_response`, +//! which hands the text to the `ConvergenceDetector`. The convergence +//! detector tokenises the text into words, computes Jaccard similarity +//! against the previous response, and fires when similarity stays above +//! `convergence_threshold` for `convergence_count` consecutive turns. +//! +//! **Merging** — `check_current_pattern` queries both detectors and +//! returns the first non-`NoPattern` result. Loop detection takes priority +//! over convergence because a tool-calling loop is a stronger signal of +//! being stuck. +//! +//! **Outcome** — The three possible results are carried by [`DetectedPattern`]: +//! `NoPattern` (agent is making progress), `LoopDetected` (repeated tool +//! calls), or `ConvergenceDetected` (semantically similar responses). //! //! # Provided Types //! @@ -53,18 +55,10 @@ //! - [`DetectedPattern`] — summary enum returned by every check method. //! - [`DetectionStats`] — cumulative statistics exposed for observability. //! -//! # Re-exports -//! -//! This module re-exports key types from the underlying detector modules so -//! consumers can import everything from a single path: -//! -//! - [`ConvergenceAction`] — what to do when convergence is detected. -//! - [`LoopStatus`] — detailed loop status from [`DetectionManager::check_loop`]. -//! //! # Quick Start //! //! ```rust,ignore -//! use loopctl::loop_control::detection::{ +//! use loopctl::detection::manager::{ //! DetectionConfig, DetectionManager, DetectedPattern, //! }; //! @@ -143,12 +137,12 @@ pub use super::loop_detector::LoopStatus; /// - [`DetectionStats`] — cumulative detection counters. #[derive(Debug, Clone)] pub enum DetectedPattern { - /// The agent is repeating the same sequence of tool calls. + /// Repeated tool-call pattern detected. /// /// Emitted by [`DetectionManager::record_operation`] (and the /// convenience wrappers [`DetectionManager::record_tool_call`] / /// [`DetectionManager::record_tool_call_with_result`]) when the - /// internal [`LoopDetector`] observes the same operation at least + /// [`LoopDetector`] observes the same operation at least /// [`DetectionConfig::loop_threshold`] times in a row. /// /// A "loop" does not mean the agent is calling the *exact same* @@ -172,24 +166,9 @@ pub enum DetectedPattern { /// } /// ``` LoopDetected { - /// Number of times the pattern has repeated. - /// - /// Equal to [`LoopStatus::repetition_count`] at the time of - /// detection. When this value reaches - /// [`DetectionConfig::stop_threshold`] the agent should terminate. - /// - /// This is a `usize` ≥ 1. A value of 1 means the loop was just - /// detected (the repetition count equals the loop threshold). + /// Equal to [`LoopStatus::repetition_count`]. repetitions: usize, - /// Human-readable description of the repeating pattern. - /// - /// Typically formatted as `"ToolName(primary_param)"` extracted - /// from the first entry in [`LoopStatus::repeated_operations`]. - /// For example, if the agent keeps calling `Read("/etc/hosts")`, - /// this field would contain `"Read(/etc/hosts)"`. - /// - /// Can be empty if the loop detector did not provide any - /// repeated operations (which should not happen in practice). + /// Formatted as `"ToolName(primary_param)"`. pattern_description: String, }, /// The agent's responses have become semantically similar. @@ -201,7 +180,7 @@ pub enum DetectedPattern { /// /// Unlike [`LoopDetected`](Self::LoopDetected), which monitors tool /// call patterns, this variant tracks the *content* of the agent's - /// free-text replies. It fires when the agent keeps saying essentially + /// free-text replies. It fires when the agent keeps saying /// the same thing in different words — a strong signal that it is /// stuck even if it is calling different tools each time. /// @@ -212,38 +191,16 @@ pub enum DetectedPattern { /// - `consecutive_count` — how many consecutive response pairs exceeded /// the threshold. ConvergenceDetected { - /// Jaccard similarity score (0.0–1.0) of the most recent pair of - /// responses. - /// - /// A value of `1.0` means the responses are identical at the - /// word-token level. The default threshold is `0.95`. - /// - /// Note: this score is computed on word-level token sets after - /// lowercasing, not on raw character sequences. Two responses - /// that are paraphrases with different vocabulary may have a low - /// score even if they are semantically equivalent. + /// Jaccard similarity (0.0–1.0) of the most recent response pair. similarity: f32, - /// Number of consecutive responses that exceeded the similarity - /// threshold. - /// - /// Equal to [`ConvergenceStatus::consecutive_count`]. When this - /// value reaches [`DetectionConfig::convergence_count`], the - /// convergence detector fires. - /// - /// A value of `1` means this is the first time the threshold was - /// exceeded; higher values indicate a sustained period of - /// similarity. + /// Consecutive similar responses. consecutive_count: usize, }, - /// No pattern detected — the agent appears to be making progress. - /// - /// This is the "healthy" variant. It is returned whenever neither the - /// loop detector nor the convergence detector has fired. The agent's - /// tool calls are varying and/or its responses are diverging, which - /// indicates forward progress. + /// Neither the loop detector nor the convergence detector has fired. /// - /// Callers should simply continue the turn loop when they receive - /// this variant. + /// The agent's tool calls are varying and/or its responses are diverging, + /// which indicates forward progress. Callers should continue the + /// turn loop when they receive this variant. NoPattern, } @@ -255,7 +212,7 @@ pub enum DetectedPattern { /// /// Groups all tunables for loop detection and convergence detection in a /// single struct so consumers can construct a [`DetectionManager`] with -/// one call to [`DetectionManager::with_config`]. +/// one call to [`DetectionManager::new_with_config`]. /// /// The [`Default`] implementation provides sensible production values: /// loop threshold 3, stop threshold 10, convergence threshold 0.95, and @@ -284,7 +241,7 @@ pub enum DetectedPattern { /// convergence_threshold: 0.9, // looser similarity threshold /// ..DetectionConfig::default() /// }; -/// let dm = DetectionManager::with_config(config); +/// let dm = DetectionManager::new_with_config(config); /// ``` /// /// # See Also @@ -293,147 +250,27 @@ pub enum DetectedPattern { /// - [`DetectionConfig::to_loop_detector_config`] — converts to [`LoopDetectorConfig`]. #[derive(Debug, Clone)] pub struct DetectionConfig { - // ================================================== - // Loop detection (forwarded to LoopDetectorConfig) - // ================================================== - /// Number of consecutive similar operations before declaring a loop. - /// - /// When the [`LoopDetector`] sees the same operation this many times - /// in a row, [`DetectionManager::check_loop`] will report - /// [`LoopStatus::is_looping`] as `true`. - /// - /// This threshold applies *per tool*, so if a custom [`ToolSignature`] - /// provides per-tool overrides via [`ToolSignature::tool_thresholds`], - /// those values take precedence for the matching tool name. - /// - /// Default: **3**. + /// Consecutive similar operations before declaring a loop. Default: **3**. pub loop_threshold: usize, - - /// Number of repetitions that triggers a forced stop (0 = disabled). - /// - /// Once the loop detector has seen this many consecutive identical - /// operations the framework should terminate the session. Set to `0` - /// to disable forced stopping and rely solely on warnings. - /// - /// This maps to [`LoopDetectorConfig::stop_threshold`] during - /// construction via [`DetectionConfig::to_loop_detector_config`]. - /// When the stop threshold is reached, [`LoopStatus::should_stop`] - /// is set to `true` and the detector generates a `"STOPPING"` warning. - /// - /// Default: **10**. + /// Repetitions triggering forced stop (0 = disabled). Default: **10**. pub stop_threshold: usize, - - /// Whether loop detection is enabled. - /// - /// When `false`, [`DetectionManager::record_operation`] and - /// [`DetectionManager::record_tool_call`] return - /// [`DetectedPattern::NoPattern`] immediately without forwarding to - /// the [`LoopDetector`]. This is useful in testing scenarios or when - /// the consumer wants to rely solely on convergence detection. - /// - /// Default: **true**. + /// Whether loop detection is enabled. Default: **true**. pub enable_loop_detection: bool, - - /// Maximum number of operations to keep in the loop detector's history. - /// - /// Older operations are evicted once the ring buffer exceeds this - /// size. A larger history lets the detector recognise longer cycles, - /// at the cost of more memory. - /// - /// This maps to [`LoopDetectorConfig::window_size`] during - /// construction. For most agents the default of 100 is more than - /// sufficient — typical loops repeat within 3–10 operations. - /// - /// Default: **100**. + /// Max operations kept in loop detector history. Default: **100**. pub max_history: usize, - - // ================================================== - // Convergence detection - // ================================================== - /// Similarity threshold (0.0–1.0) for convergence detection. - /// - /// The [`ConvergenceDetector`] compares the Jaccard similarity of the - /// most recent pair of responses. If the score is at or above this - /// value for [`convergence_count`](Self::convergence_count) consecutive - /// turns, convergence is declared. - /// - /// A higher threshold (e.g., 0.99) requires near-identical responses, - /// while a lower threshold (e.g., 0.80) catches paraphrasing but may - /// produce false positives. - /// - /// Default: **0.95**. + /// Jaccard similarity threshold for convergence (0.0–1.0). Default: **0.95**. pub convergence_threshold: f32, - - /// Number of consecutive similar responses required for convergence. - /// - /// A higher value makes the detector more tolerant — the agent can - /// produce several similar responses before being flagged. Setting - /// this to 1 would flag any two similar responses immediately. - /// - /// Maps to [`ConvergenceConfig::window_size`] during construction. - /// - /// Default: **3**. + /// Consecutive similar responses for convergence. Default: **3**. pub convergence_count: usize, - - /// Whether convergence detection is enabled. - /// - /// When `false`, [`DetectionManager::record_response`] returns - /// [`DetectedPattern::NoPattern`] immediately without forwarding to - /// the [`ConvergenceDetector`]. This is useful when the consumer - /// wants to rely solely on loop detection, or in unit tests where - /// convergence noise would be distracting. - /// - /// Default: **true**. + /// Whether convergence detection is enabled. Default: **true**. pub enable_convergence_detection: bool, - - /// Action to take when convergence is detected. - /// - /// Determines whether the framework should emit a warning, force a - /// stop, or silently note the event. See [`ConvergenceAction`] for - /// available options. - /// - /// The action is forwarded to the internal [`ConvergenceDetector`] - /// during construction via [`DetectionConfig::to_convergence_config`]. - /// It is consulted by the framework when [`DetectionManager::check_convergence`] - /// returns a [`ConvergenceStatus`] with `detected == true`. - /// - /// Default: [`ConvergenceAction::default()`]. + /// Action on convergence. Default: [`ConvergenceAction::default()`]. pub on_converge: ConvergenceAction, - - /// Maximum number of responses to keep for convergence checking. - /// - /// The [`ConvergenceDetector`] maintains a sliding window of responses. - /// Older responses are evicted once this limit is reached. - /// - /// A larger window retains more history (useful for spotting slow - /// drift), while a smaller window reduces memory usage and makes the - /// detector more responsive to recent changes. - /// - /// Default: **20**. + /// Max responses kept for convergence checking. Default: **20**. pub max_response_history: usize, } impl Default for DetectionConfig { - /// Produce a configuration with production-ready defaults. - /// - /// The defaults are chosen to be conservative enough for most agent - /// workloads while still catching genuine stuck behaviours: - /// - /// | Field | Default | - /// |--------------------------------|----------------------------------| - /// | `loop_threshold` | 3 | - /// | `stop_threshold` | 10 | - /// | `enable_loop_detection` | `true` | - /// | `max_history` | 100 | - /// | `convergence_threshold` | 0.95 | - /// | `convergence_count` | 3 | - /// | `enable_convergence_detection` | `true` | - /// | `on_converge` | [`ConvergenceAction::default()`] | - /// | `max_response_history` | 20 | - /// - /// # When called - /// - /// By [`DetectionManager::new`] and anywhere a default config is needed. fn default() -> Self { Self { loop_threshold: 3, @@ -452,11 +289,6 @@ impl Default for DetectionConfig { impl DetectionConfig { /// Convert convergence settings into a [`ConvergenceConfig`]. /// - /// Maps the subset of fields that belong to the convergence subsystem - /// (`enable_convergence_detection` → `enabled`, `convergence_count` → - /// `window_size`, etc.). Called by [`DetectionManager::with_config`] - /// during construction. - /// /// # Field Mapping /// /// | `DetectionConfig` field | [`ConvergenceConfig`] field | @@ -465,11 +297,6 @@ impl DetectionConfig { /// | `convergence_count` | `window_size` | /// | `convergence_threshold` | `similarity_threshold` | /// | `on_converge` | `on_converge` | - /// - /// # When called - /// - /// Internally by [`DetectionManager`] constructors; generally not - /// needed by external callers. #[must_use] pub fn to_convergence_config(&self) -> ConvergenceConfig { ConvergenceConfig { @@ -482,10 +309,6 @@ impl DetectionConfig { /// Convert the loop-related settings into a [`LoopDetectorConfig`]. /// - /// Maps `max_history` → `window_size`, `loop_threshold` → - /// `repetition_threshold`, and `stop_threshold` directly. All other - /// [`LoopDetectorConfig`] fields inherit their defaults. - /// /// # Field Mapping /// /// | `DetectionConfig` field | [`LoopDetectorConfig`] field | @@ -493,11 +316,6 @@ impl DetectionConfig { /// | `max_history` | `window_size` | /// | `loop_threshold` | `repetition_threshold` | /// | `stop_threshold` | `stop_threshold` | - /// - /// # When called - /// - /// Internally by [`DetectionManager`] constructors; generally not - /// needed by external callers. #[must_use] pub fn to_loop_detector_config(&self) -> LoopDetectorConfig { LoopDetectorConfig { @@ -519,7 +337,7 @@ impl DetectionConfig { /// debugging. All counters are monotonically increasing within a session /// (until [`DetectionManager::reset`] is called). /// -/// This struct is [`Clone`] and [`Default`] so it can be cheaply snapshot +/// [`Clone`] and [`Default`] — can be cheaply snapshot /// and serialised for logging or UI display. /// /// # Example @@ -538,51 +356,13 @@ impl DetectionConfig { /// - [`DetectionManager::reset`] — zeroes all counters. #[derive(Debug, Clone, Default)] pub struct DetectionStats { - /// Total number of turns (operations + responses) analysed. - /// - /// Incremented by [`DetectionManager::record_operation`] each time an - /// operation is recorded, regardless of whether a loop was found. - /// This counter does **not** include responses recorded via - /// [`DetectionManager::record_response`] — only tool-call operations. - /// - /// Reset to `0` by [`DetectionManager::reset`]. + /// Tool-call operations recorded via `record_operation`. Reset by `reset`. pub turns_analyzed: usize, - - /// Number of times a loop was detected. - /// - /// Incremented only when [`DetectionManager::record_operation`] returns - /// [`DetectedPattern::LoopDetected`]. This is a subset of - /// [`turns_analyzed`](Self::turns_analyzed). - /// - /// A high ratio of `loops_detected` to `turns_analyzed` suggests the - /// agent is frequently stuck. The framework can use this metric to - /// decide whether to terminate the session early. - /// - /// Reset to `0` by [`DetectionManager::reset`]. + /// Subset of `turns_analyzed` that returned `LoopDetected`. Reset by `reset`. pub loops_detected: usize, - - /// Number of times convergence was detected. - /// - /// Incremented only when [`DetectionManager::record_response`] returns - /// [`DetectedPattern::ConvergenceDetected`]. Unlike `loops_detected`, - /// this counter tracks semantic similarity of free-text responses, - /// not tool-call repetition. - /// - /// Reset to `0` by [`DetectionManager::reset`]. + /// Times `record_response` returned `ConvergenceDetected`. Reset by `reset`. pub convergences_detected: usize, - - /// Current streak of consecutive similar operations. - /// - /// Mirrors [`LoopStatus::repetition_count`] after the most recent - /// [`DetectionManager::record_operation`] call. When this value reaches - /// [`DetectionConfig::loop_threshold`], the next call to - /// [`DetectionManager::record_operation`] will return - /// [`DetectedPattern::LoopDetected`]. - /// - /// Useful for progress bars or warning indicators that show "how close" - /// the agent is to triggering a loop detection. - /// - /// Reset to `0` by [`DetectionManager::reset`]. + /// Mirrors `LoopStatus::repetition_count`. Triggers `LoopDetected` at `loop_threshold`. Reset by `reset`. pub current_streak: usize, } @@ -601,12 +381,12 @@ pub struct DetectionStats { /// /// Choose a constructor based on your needs: /// -/// | Constructor | Use case | -/// |---|---| -/// | [`DetectionManager::new`] | Quick start with defaults | -/// | [`DetectionManager::with_config`] | Custom thresholds via [`DetectionConfig`] | -/// | [`DetectionManager::with_loop_detector`] | Inject a pre-built [`LoopDetector`] | -/// | [`DetectionManager::with_signature`] | Custom [`ToolSignature`] for JSON parsing | +/// | Constructor | Use case | +/// |------------------------------------------|-------------------------------------------| +/// | [`DetectionManager::new`] | Quick start with defaults | +/// | [`DetectionManager::new_with_config`] | Custom thresholds via [`DetectionConfig`] | +/// | [`DetectionManager::new_with_loop_detector`] | Inject a pre-built [`LoopDetector`] | +/// | [`DetectionManager::new_with_signature`] | Custom [`ToolSignature`] for JSON parsing | /// /// # Loop Detection /// @@ -618,7 +398,7 @@ pub struct DetectionStats { /// /// Record assistant responses via [`Self::record_response`] and check /// for semantic similarity using [`Self::check_convergence`]. Both -/// delegate to the internal [`ConvergenceDetector`]. +/// delegate to the [`ConvergenceDetector`]. /// /// # Direct Access /// @@ -628,15 +408,15 @@ pub struct DetectionStats { /// /// # Lifecycle /// -/// ```text -/// new() / with_config() -/// → record_operation(op) [each tool call] -/// → check_loop() [after each turn] -/// → record_response(text) [each assistant reply] -/// → check_convergence() [after each turn] -/// → stats() [observability] -/// → reset() [between tasks] -/// ``` +/// The manager progresses through a simple per-turn cycle: +/// +/// 1. **Construct** — `new()` or `with_config()`. +/// 2. **Feed** — call `record_operation(op)` for each tool invocation and +/// `record_response(text)` for each assistant reply. +/// 3. **Check** — call `check_loop()`, `check_convergence()`, or the +/// combined `check_current_pattern()` after each turn. +/// 4. **Observe** — call `stats()` at any time for cumulative counters. +/// 5. **Reset** — call `reset()` between tasks to clear all history. /// /// # Example /// @@ -672,16 +452,20 @@ pub struct DetectionStats { /// - [`DetectedPattern`] — the result type returned by check methods. /// - [`DetectionStats`] — cumulative observability counters. pub struct DetectionManager { + /// Configuration thresholds and feature flags. config: DetectionConfig, + /// Loop detector shared across compaction and analysis phases. loop_detector: Arc, + /// Convergence detector guarded by interior mutability. convergence_detector: Mutex, + /// Cumulative detection statistics. stats: Mutex, } impl DetectionManager { /// Create a new detection manager with default configuration. /// - /// Convenience wrapper around [`Self::with_config`] that passes + /// Convenience wrapper around [`Self::new_with_config`] that passes /// [`DetectionConfig::default`]. Suitable for quick prototyping or /// when the default thresholds (loop=3, stop=10, convergence=0.95) /// are acceptable. @@ -697,30 +481,13 @@ impl DetectionManager { /// Returns [`ConvergenceConfigError`] if the convergence configuration /// is invalid (e.g., threshold out of range, window too small). pub fn new() -> Result { - Self::with_config(DetectionConfig::default()) + Self::new_with_config(DetectionConfig::default()) } /// Create a new detection manager with custom configuration. /// - /// Builds a [`LoopDetector`] using a [`NoOpToolSignature`](super::loop_detector::NoOpToolSignature) - /// and a [`ConvergenceDetector`] from the relevant config fields. - /// For tool-specific JSON parsing, use [`Self::with_signature`] or - /// [`Self::with_loop_detector`] instead. - /// - /// # Construction Flow - /// - /// ```text - /// with_config(config) - /// ├─ config.to_loop_detector_config() → LoopDetectorConfig - /// │ └─ LoopDetector::new(ldc, NoOpToolSignature) - /// ├─ config.to_convergence_config() → ConvergenceConfig - /// │ └─ ConvergenceDetector::new(cc) - /// └─ DetectionStats::default() - /// ``` - /// - /// # When called - /// - /// Typically during agent initialisation, once per session. + /// For tool-specific JSON parsing, use [`Self::new_with_signature`] or + /// [`Self::new_with_loop_detector`] instead. /// /// # Example /// @@ -729,14 +496,14 @@ impl DetectionManager { /// loop_threshold: 5, /// ..DetectionConfig::default() /// }; - /// let dm = DetectionManager::with_config(config); + /// let dm = DetectionManager::new_with_config(config); /// ``` /// /// # Errors /// /// Returns [`ConvergenceConfigError`] if the convergence configuration /// is invalid (e.g., threshold out of range, window too small). - pub fn with_config(config: DetectionConfig) -> Result { + pub fn new_with_config(config: DetectionConfig) -> Result { let loop_detector = Arc::new(LoopDetector::new( config.to_loop_detector_config(), Arc::new(super::loop_detector::NoOpToolSignature), @@ -755,31 +522,25 @@ impl DetectionManager { /// /// Use this when the caller has already constructed a [`LoopDetector`] /// with a custom [`ToolSignature`] or non-default thresholds and wants - /// to inject it directly. The detector is wrapped in `Arc` internally - /// for cheap sharing. + /// to inject it directly. /// /// This constructor **does not** use the `loop_threshold`, `stop_threshold`, /// or `max_history` fields from `config` — those are already baked into /// the provided `loop_detector`. Only the convergence-related fields /// are read from `config`. /// - /// # When called - /// - /// During agent initialisation when the caller needs fine-grained - /// control over the loop detector's internals. - /// /// # Example /// /// ```rust,ignore /// let ld = LoopDetector::new(ld_config, my_signature); - /// let dm = DetectionManager::with_loop_detector(config, ld); + /// let dm = DetectionManager::new_with_loop_detector(config, ld); /// ``` /// /// # Errors /// /// Returns [`ConvergenceConfigError`] if the convergence configuration /// is invalid (e.g., threshold out of range, window too small). - pub fn with_loop_detector( + pub fn new_with_loop_detector( config: DetectionConfig, loop_detector: LoopDetector, ) -> Result { @@ -795,27 +556,15 @@ impl DetectionManager { /// Create with a specific [`ToolSignature`] for tool-specific parsing. /// - /// Constructs a [`LoopDetector`] using the provided `signature` and - /// its associated [`tool_thresholds`](ToolSignature::tool_thresholds). /// Useful when the consumer provides its own tool-aware signature /// (e.g., `DchToolSignature` for dch.sh tools that extracts /// `file_path` from JSON inputs). /// - /// The tool thresholds from the signature are merged into a - /// [`LoopDetectorConfig`] derived from [`DetectionConfig::default`], - /// overriding the generic [`DetectionConfig::loop_threshold`] for any - /// tool listed in [`ToolSignature::tool_thresholds`]. - /// - /// # When called - /// - /// During agent initialisation when the agent runtime has a domain- - /// specific [`ToolSignature`] implementation. - /// /// # Example /// /// ```rust,ignore /// let signature = Arc::new(MyToolSignature); - /// let dm = DetectionManager::with_signature(signature); + /// let dm = DetectionManager::new_with_signature(signature); /// ``` /// /// # See Also @@ -827,7 +576,7 @@ impl DetectionManager { /// /// Returns [`ConvergenceConfigError`] if the convergence configuration /// is invalid (e.g., threshold out of range, window too small). - pub fn with_signature( + pub fn new_with_signature( signature: Arc, ) -> Result { let config = DetectionConfig::default(); @@ -852,7 +601,7 @@ impl DetectionManager { /// Record an [`Operation`] for loop detection and return the current /// [`DetectedPattern`]. /// - /// This is the primary entry point for feeding data into the loop + /// Primary entry point for feeding data into the loop /// detector. It performs three steps: /// /// 1. Forwards `operation` to [`LoopDetector::record`]. @@ -900,7 +649,6 @@ impl DetectionManager { self.loop_detector.record(operation); - // Check if this triggered a loop let status = self.loop_detector.check_loop(); if status.is_looping { let mut guard = self.stats.lock().unwrap_or_else(|e| { @@ -931,7 +679,7 @@ impl DetectionManager { /// Returns the tool signature used for extracting primary parameters. /// /// Useful when callers need to construct [`Operation`]s directly using - /// the same signature the detection manager uses internally. + /// the configured [`ToolSignature`]. pub fn signature(&self) -> &dyn ToolSignature { self.loop_detector.signature() } @@ -939,7 +687,7 @@ impl DetectionManager { /// Record a tool call for loop detection by tool name and input hash. /// /// Creates an [`Operation`] from a `tool` name and `input_hash`, then - /// delegates to [`Self::record_operation`]. This is the easiest way to + /// delegates to [`Self::record_operation`]. Easiest way to /// feed data into the loop detector when you only need a numeric hash /// as the primary parameter rather than a full JSON-based signature. /// @@ -985,21 +733,19 @@ impl DetectionManager { /// when both input and result match across consecutive calls is a loop /// flagged. /// - /// This is essential for tools like `Read` where calling the same file + /// Important for tools like `Read` where calling the same file /// is perfectly fine if the file content is changing (e.g., the agent /// is editing it). /// /// # How it works /// - /// ```text - /// Turn 1: Read("/foo.txt") → result_hash=0xA ─┐ - /// Turn 2: Read("/foo.txt") → result_hash=0xA ─┤ Same hash → loop candidate - /// Turn 3: Read("/foo.txt") → result_hash=0xA ─┘ - /// → LoopDetected after loop_threshold repetitions - /// - /// Turn 1: Read("/foo.txt") → result_hash=0xA - /// Turn 2: Read("/foo.txt") → result_hash=0xB ← Different hash → not a loop - /// ``` + /// With result hashing, the detector distinguishes between a tool + /// that returns the same output every time (a genuine loop) and one + /// whose output changes between calls (the agent is making progress). + /// For example, calling `Read("/foo.txt")` three times with the same + /// result hash is flagged as a loop, but if the file is being edited + /// between reads the result hashes will differ and no loop is + /// reported. /// /// # When called /// @@ -1040,7 +786,7 @@ impl DetectionManager { self.record_operation(operation) } - /// Query the internal [`LoopDetector`] for the current loop status. + /// Query the [`LoopDetector`] for the current loop status. /// /// Returns a full [`LoopStatus`] snapshot including repetition counts, /// the `should_stop` flag, and any warning message. Unlike @@ -1076,7 +822,7 @@ impl DetectionManager { self.loop_detector.check_loop() } - /// Obtain a shared reference to the internal [`LoopDetector`]. + /// Obtain a shared reference to the [`LoopDetector`]. /// /// Useful for direct access to loop state, e.g. inspecting /// the turn count or history during diagnostics. @@ -1102,7 +848,7 @@ impl DetectionManager { &self.loop_detector } - /// Obtain a shared reference to the internal [`ConvergenceDetector`]. + /// Obtain a shared reference to the [`ConvergenceDetector`]. /// /// Returns `&Mutex` so callers can lock and /// invoke methods on the convergence detector directly — for example @@ -1137,15 +883,15 @@ impl DetectionManager { } // ================================================== - // Convergence detection (handled internally) + // Convergence detection // ================================================== /// Record an assistant response for convergence detection. /// - /// Forwards `response` to the internal [`ConvergenceDetector`] via - /// [`ConvergenceDetector::add_response`], which tokenises the text - /// into word-level tokens, computes Jaccard similarity against the - /// previous response, and updates the consecutive-similarity counter. + /// Forwards `response` to the [`ConvergenceDetector`] via + /// [`ConvergenceDetector::add_response`], which computes Jaccard + /// similarity against the previous response and updates the + /// consecutive-similarity counter. /// /// If the similarity score exceeds /// [`DetectionConfig::convergence_threshold`] for @@ -1157,14 +903,6 @@ impl DetectionManager { /// this method short-circuits and returns [`DetectedPattern::NoPattern`] /// without touching the convergence detector or statistics. /// - /// # Tokenisation - /// - /// Responses are split on whitespace into word-level tokens, - /// lowercased, and compared using Jaccard similarity (intersection - /// over union of token sets). This is fast and language-agnostic - /// but does not capture semantic equivalence — two paraphrased - /// sentences with different vocabulary will have a low score. - /// /// # When called /// /// After each assistant response during an agent turn — typically by @@ -1193,6 +931,7 @@ impl DetectionManager { if !self.config.enable_convergence_detection { return DetectedPattern::NoPattern; } + let status = self .convergence_detector .lock() @@ -1201,6 +940,7 @@ impl DetectionManager { e.into_inner() }) .add_response(response); + if status.detected { let mut guard = self.stats.lock().unwrap_or_else(|e| { tracing::warn!("stats lock poisoned, recovering"); @@ -1215,7 +955,7 @@ impl DetectionManager { DetectedPattern::NoPattern } - /// Query the internal [`ConvergenceDetector`] for the current + /// Query the [`ConvergenceDetector`] for the current /// convergence status. /// /// Returns a full [`ConvergenceStatus`] snapshot including the @@ -1264,20 +1004,17 @@ impl DetectionManager { /// /// Checks both the loop detector and the convergence detector in /// sequence and returns the first non-`NoPattern` result (loop takes - /// priority over convergence). This is a **read-only** operation — + /// priority over convergence). **Read-only** operation — /// no statistics are updated and no new data is recorded. /// /// # Priority order /// - /// ```text - /// check_current_pattern() - /// ├─ check_loop() → LoopDetected? (priority 1) - /// └─ check_convergence() → ConvergenceDetected? (priority 2) - /// └─ NoPattern (fallback) - /// ``` - /// - /// Loop detection takes priority because a looping agent is more - /// urgently stuck than a converging one. + /// Loop detection is checked first. If the loop detector reports a + /// loop, `LoopDetected` is returned immediately. Only when no loop is + /// found does the manager query the convergence detector. If neither + /// detector has fired, `NoPattern` is returned. Loop takes priority + /// because a tool-calling loop is a stronger and more urgent signal + /// that the agent is stuck. /// /// # When called /// @@ -1340,7 +1077,7 @@ impl DetectionManager { /// Take a snapshot of the cumulative detection statistics. /// - /// Locks the internal `Mutex`, clones the [`DetectionStats`] struct, + /// Clones the [`DetectionStats`] struct via the `Mutex`, /// and returns it. The snapshot reflects all operations recorded since /// the last [`Self::reset`] call (or since construction). /// @@ -1400,7 +1137,7 @@ impl DetectionManager { /// # See Also /// /// - [`DetectionConfig`] — full description of all config fields. - /// - [`Self::with_config`] — the constructor that accepts a config. + /// - [`Self::new_with_config`] — the constructor that accepts a config. #[must_use] pub fn config(&self) -> &DetectionConfig { &self.config @@ -1464,18 +1201,6 @@ impl DetectionManager { } impl Default for DetectionManager { - /// Produce a [`DetectionManager`] with the default [`DetectionConfig`]. - /// - /// Constructs the manager directly using struct-literal syntax, bypassing - /// the fallible [`DetectionManager::with_config`] constructor. This is - /// infallible because: - /// - /// - [`LoopDetector::new`] is infallible. - /// - [`ConvergenceDetector`] is initialised with a default window and - /// threshold that satisfy its validation invariants (`window_size ≥ 2`, - /// `threshold ∈ 0.0..=1.0`). - /// - /// [`ConvergenceDetector::new`]: crate::loop_control::convergence::ConvergenceDetector::new fn default() -> Self { let config = DetectionConfig::default(); let loop_config = config.to_loop_detector_config(); @@ -1586,7 +1311,7 @@ mod tests { enable_convergence_detection: false, ..Default::default() }; - let dm = DetectionManager::with_config(config).unwrap(); + let dm = DetectionManager::new_with_config(config).unwrap(); for _ in 0..10 { let result = dm.record_tool_call("read_file", 42); assert!(matches!(result, DetectedPattern::NoPattern)); @@ -1606,7 +1331,7 @@ mod tests { stop_threshold: 5, ..Default::default() }; - let dm = DetectionManager::with_config(config).unwrap(); + let dm = DetectionManager::new_with_config(config).unwrap(); // Record 5 identical operations for _ in 0..5 { diff --git a/src/engine.rs b/src/engine.rs index 492f590..632eefa 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -8,21 +8,21 @@ //! 3. **Tool dispatch** — Execute tool calls requested by the model //! 4. **Feedback** — Feed tool results back into the conversation //! 5. **Loop** — Repeat until the model stops or max turns is reached -//! 6. **Finalize** — Produce a [`SessionResult`](crate::core::SessionResult) +//! 6. **Finalize** — Produce a [`SessionResult`](crate::engine::loop_core::SessionResult) //! //! # Example //! //! ```rust,ignore //! use loopctl::engine::BareLoop; //! use loopctl::tool::ToolRegistry; -//! use loopctl::api_client::ApiClient; -//! use loopctl::core::AgentConfig; +//! use loopctl::api::ApiClient; +//! use loopctl::config::LoopConfig; //! //! let agent = BareLoop::new(client, registry, config); //! let result = agent.run("Write a hello world program").await?; //! ``` mod bare; -pub mod middleware; +pub mod loop_core; pub use bare::*; diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 9f30acf..d754aa1 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -1,6 +1,6 @@ //! `BareLoop` — the framework's default agent loop implementation. //! -//! This module provides [`BareLoop`], a generic, framework-level agent +//! [`BareLoop`] — a generic, framework-level agent //! loop that orchestrates the full lifecycle of an LLM-based agent session: //! sending messages to an LLM API, accumulating streaming responses, //! dispatching tool calls, and feeding results back into the conversation @@ -10,13 +10,13 @@ //! //! [`BareLoop`] ties together four key components: //! -//! - An [`ApiClient`](crate::api_client::ApiClient) for communicating with +//! - An [`ApiClient`](crate::api::ApiClient) for communicating with //! the LLM provider. //! - A [`ToolRegistry`](crate::tool::ToolRegistry) for dispatching tool //! calls the model requests. -//! - An [`AgentConfig`] governing session parameters (max turns, system +//! - An [`LoopConfig`] governing session parameters (max turns, system //! prompt, session ID). -//! - Optional [`LoopObserver`](crate::core::observer::LoopObserver) registrations for lifecycle instrumentation. +//! - Optional [`LoopObserver`](crate::observer::LoopObserver) registrations for lifecycle instrumentation. //! //! ```text //! BareLoop @@ -37,7 +37,7 @@ //! # Key Design Decisions //! //! - **Static dispatch** — `BareLoop` is generic over the -//! [`ApiClient`](crate::api_client::ApiClient) type parameter `C`, +//! [`ApiClient`](crate::api::ApiClient) type parameter `C`, //! avoiding `dyn` overhead for the hot path. //! - **Sequential tool dispatch** — tools within a single turn are //! executed one after another so cancellation is checked between each. @@ -52,13 +52,13 @@ //! ```rust,ignore //! use loopctl::loop_::BareLoop; //! use loopctl::tool::ToolRegistry; -//! use loopctl::core::AgentConfig; +//! use loopctl::core::LoopConfig; //! use std::sync::Arc; //! //! // 1. Build components //! let client = Arc::new(my_api_client); //! let registry = ToolRegistry::new(); -//! let config = AgentConfig::default(); +//! let config = LoopConfig::default(); //! //! // 2. Create the loop //! let agent = BareLoop::new(client, registry, config); @@ -68,19 +68,14 @@ //! println!("Agent responded in {} turns", result.total_turns); //! ``` -use crate::api_client::ApiClient; +use crate::api::ApiClient; use crate::cancel::CancelSignal; use crate::compact::{ContextManager, EnsureContextResult}; -use crate::core::observer::{ - ConvergenceDetectedContext, FallbackContext, LoopDetectedContext, ResponseContext, - StreamContext, StreamFailureContext, TurnEndContext, TurnStartContext, -}; -use crate::core::reflection::{ - ExponentialBackoffRecovery, NoopReflector, RecoveryAction, RecoveryStrategy, ReflectionContext, - Reflector, -}; -use crate::core::{AgentConfig, AgentError, SessionResult, ToolDispatchResult}; -use crate::engine::middleware::{ToolDispatchContext, ToolPipeline, ToolPipelineBuilder}; +use crate::config::LoopConfig; +use crate::detection::{ConvergenceAction, DetectedPattern}; +use crate::error::LoopError; + +use crate::engine::loop_core::SessionResult; #[cfg(feature = "hooks")] use crate::hooks::HookAction; #[cfg(feature = "hooks")] @@ -91,14 +86,22 @@ use crate::hooks::context::{ CompactTrigger, PostCompactContext, PostToolUseContext, PreCompactContext, PreToolUseContext, SessionEndContext, SessionEndReason, SessionStartContext, }; -use crate::loop_control::bundle::ManagerBundle; -use crate::loop_control::detection::{ConvergenceAction, DetectedPattern}; use crate::message::{Message, MessagePart, Role, ToolContent}; +use crate::middleware::{ToolDispatchContext, ToolPipeline, ToolPipelineBuilder}; +use crate::observer::{ + ConvergenceDetectedContext, FallbackContext, LoopDetectedContext, ResponseContext, + StreamContext, StreamFailureContext, TurnEndContext, TurnStartContext, +}; +use crate::reflection::{ + Correction, CorrectionResult, ExponentialBackoffRecovery, NoopReflector, RecoveryAction, + RecoveryStrategy, ReflectionContext, Reflector, +}; +use crate::runtime::LoopRuntime; use crate::stream::handler::StreamHandler; use crate::stream::{StreamAccumulator, StreamEvent, StreamStopReason, Usage}; #[cfg(feature = "tool_health")] use crate::tool::health::ToolHealthRegistry; -use crate::tool::{PermissionCheck, ToolContext, ToolRegistry, ToolSchema}; +use crate::tool::{PermissionCheck, ToolContext, ToolDispatchResult, ToolRegistry, ToolSchema}; use std::sync::Arc; use std::time::{Duration, Instant}; use uuid::Uuid; @@ -132,15 +135,15 @@ mod stream; /// Use one of the constructors based on what components you have: /// /// - [`new()`](BareLoop::new) — client + tools + config. -/// - [`with_managers()`](BareLoop::with_managers) — full control, -/// including a [`ManagerBundle`]. +/// - [`new_with_managers()`](BareLoop::new_with_managers) — full control, +/// including a [`LoopRuntime`]. /// - [`from_parts()`](BareLoop::from_parts) — re-assembles from the /// output of `AgentBuilder::into_raw_parts()`. /// /// # Lifecycle /// /// ```text -/// new() / with_managers() / from_parts() +/// new() / new_with_managers() / from_parts() /// → run(user_input) /// → stream_turn() → dispatch_tools() → stream_turn() /// → … (repeat until end_turn or max_turns) @@ -152,11 +155,11 @@ mod stream; /// ```rust,ignore /// use loopctl::loop_::BareLoop; /// use loopctl::tool::ToolRegistry; -/// use loopctl::core::AgentConfig; +/// use loopctl::core::LoopConfig; /// use std::sync::Arc; /// /// let registry = ToolRegistry::new(); -/// let config = AgentConfig::default(); +/// let config = LoopConfig::default(); /// /// let agent = BareLoop::new( /// Arc::new(my_client), @@ -189,8 +192,8 @@ pub struct BareLoop { /// Session parameters (max turns, model, system prompt). /// - /// See [`AgentConfig`] for the full set of options. - config: AgentConfig, + /// See [`LoopConfig`] for the full set of options. + config: LoopConfig, /// Conversation history (system + user + assistant + tool results). /// @@ -210,8 +213,8 @@ pub struct BareLoop { /// tool operations) and convergence detection (semantically /// similar responses). /// - /// Reset at the start of every session via [`ManagerBundle::reset_all`]. - managers: ManagerBundle, + /// Reset at the start of every session via [`LoopRuntime::reset_all`]. + managers: LoopRuntime, /// Failure analyser for tool errors. /// @@ -240,8 +243,8 @@ pub struct BareLoop { /// When `Some`, the loop checks token usage after each turn and /// triggers compaction when usage exceeds the configured threshold. /// Compaction replaces the conversation messages, notifies observers - /// via [`LoopObserver::on_compaction`](crate::core::observer::LoopObserver::on_compaction), - /// and notifies observers via [`on_compaction`](crate::core::observer::LoopObserver::on_compaction). + /// via [`LoopObserver::on_compaction`](crate::observer::LoopObserver::on_compaction), + /// and notifies observers via [`on_compaction`](crate::observer::LoopObserver::on_compaction). context_manager: Option>, /// Optional stream handler for resilient streaming. @@ -326,7 +329,7 @@ impl TurnTokens { /// counts for a single turn. When `usage` is `None` (provider did /// not report counts), both fields default to `0`. /// - /// This is needed because [`SessionBudget::accumulate_usage`] mutates + /// Required because [`SessionBudget::accumulate_usage`] mutates /// running totals in place, but the per-turn values must be reported /// separately to observers. fn from_usage(usage: Option<&Usage>) -> Self { @@ -356,7 +359,7 @@ struct TurnContext { /// Reason the session was aborted before normal completion. /// /// Used by [`abort_session`](BareLoop::abort_session) to select the -/// correct [`AgentError`] variant without string matching. +/// correct [`LoopError`] variant without string matching. #[derive(Clone, Copy)] enum AbortReason { /// User or external signal requested cancellation. @@ -378,6 +381,7 @@ struct SessionEndInfo { /// Total turns executed. total_turns: usize, /// Total tokens consumed (input + output). + #[cfg_attr(not(feature = "hooks"), allow(dead_code))] total_tokens: u64, /// Wall-clock session duration in seconds. duration_secs: u64, @@ -404,9 +408,8 @@ enum EndReason { /// its fields into this struct for convenient passing to /// [`dispatch_tools()`](BareLoop::dispatch_tools). /// -/// This type is private to the module because external consumers -/// interact with tool results via [`SessionResult`] or the -/// [`LoopObserver`](crate::core::observer::LoopObserver) callbacks. +/// External consumers interact with tool results via [`SessionResult`] or the +/// [`LoopObserver`](crate::observer::LoopObserver) callbacks. /// /// # Fields /// @@ -440,6 +443,76 @@ struct ToolCallInfo { input: serde_json::Value, } +impl ToolCallInfo { + /// Apply a [`Correction`] from the reflection system in-place. + /// + /// Modifies `self` according to the correction strategy: + /// + /// - [`InputFix`](crate::reflection::CorrectionType::InputFix) — replaces + /// `self.input` with `correction.modified_input` (if provided). + /// - [`ToolChange`](crate::reflection::CorrectionType::ToolChange) — + /// replaces `self.name` with `correction.alternative_tool` (if provided). + /// - Other types ([`Retry`](crate::reflection::CorrectionType::Retry), + /// [`ApproachChange`](crate::reflection::CorrectionType::ApproachChange), + /// [`Escalate`](crate::reflection::CorrectionType::Escalate)) — no + /// mutation needed; the retry proceeds with unchanged parameters. + /// + /// Returns a [`CorrectionResult`] indicating whether the correction + /// was applied, failed (e.g. missing fields), or skipped. + fn apply_correction( + &mut self, + correction: &Correction, + _prior_result: &ToolDispatchResult, + ) -> CorrectionResult { + use crate::reflection::CorrectionType; + match correction.correction_type { + CorrectionType::InputFix => { + if let Some(ref modified) = correction.modified_input { + tracing::debug!( + tool = %self.name, + "applying InputFix correction from reflector" + ); + self.input = modified.clone(); + CorrectionResult::Applied + } else { + CorrectionResult::Failed( + "InputFix correction missing modified_input".to_string(), + ) + } + } + CorrectionType::ToolChange => { + if let Some(ref alt) = correction.alternative_tool { + tracing::debug!( + old_tool = %self.name, + new_tool = %alt, + "applying ToolChange correction from reflector" + ); + self.name.clone_from(alt); + CorrectionResult::Applied + } else { + CorrectionResult::Failed( + "ToolChange correction missing alternative_tool".to_string(), + ) + } + } + CorrectionType::PrerequisiteFix | CorrectionType::ApproachChange => { + // These types do not modify the tool call parameters. + // PrerequisiteFix is advisory (the guidance may describe a + // side-effect to perform); ApproachChange is high-level + // guidance for a different strategy. The retry proceeds + // with unchanged input/tool. + CorrectionResult::Skipped + } + CorrectionType::Escalate => { + // Escalation means no correction is possible. The retry + // loop should not normally reach here because Escalate + // errors are mapped to RecoveryAction::Fail upstream. + CorrectionResult::Skipped + } + } + } +} + impl BareLoop { /// Maximum retry attempts for tool recovery before giving up. const MAX_RECOVERY_ATTEMPTS: u32 = 5; @@ -447,7 +520,7 @@ impl BareLoop { /// Create a new `BareLoop` with the given components. /// /// Initializes an empty conversation history and a - /// fresh [`ManagerBundle`]. The cancellation signal starts as non-cancelled. + /// fresh [`LoopRuntime`]. The cancellation signal starts as non-cancelled. /// /// # Parameters /// @@ -461,17 +534,17 @@ impl BareLoop { /// let agent = BareLoop::new( /// Arc::new(my_client), /// ToolRegistry::new(), - /// AgentConfig::default(), + /// LoopConfig::default(), /// ); /// ``` - pub fn new(client: Arc, tools: ToolRegistry, config: AgentConfig) -> Self { + pub fn new(client: Arc, tools: ToolRegistry, config: LoopConfig) -> Self { Self { client, tools: Arc::new(tools), pipeline: None, config, conversation: Vec::new(), - managers: ManagerBundle::new(), + managers: LoopRuntime::new(), reflector: Arc::new(NoopReflector), recovery: Arc::new(ExponentialBackoffRecovery::new(3)), cancelled: Arc::new(CancelSignal::new()), @@ -487,7 +560,7 @@ impl BareLoop { /// Create a new `BareLoop` with all components including managers. /// /// Use this constructor when you need to supply a pre-configured - /// [`ManagerBundle`] — for example, to enable loop detection or + /// [`LoopRuntime`] — for example, to enable loop detection or /// circuit-breaker policies. /// /// # Parameters @@ -495,27 +568,27 @@ impl BareLoop { /// - `client` — The LLM API client, wrapped in `Arc`. /// - `tools` — The [`ToolRegistry`] containing available tools. /// - `config` — Session parameters. - /// - `managers` — A pre-built [`ManagerBundle`]. + /// - `managers` — A pre-built [`LoopRuntime`]. /// /// # Example /// /// ```rust,ignore - /// let managers = ManagerBundle::builder() + /// let managers = LoopRuntime::builder() /// .with_loop_detection(10) /// .build(); /// - /// let agent = BareLoop::with_managers( + /// let agent = BareLoop::new_with_managers( /// Arc::new(my_client), /// registry, /// config, /// managers, /// ); /// ``` - pub fn with_managers( + pub fn new_with_managers( client: Arc, tools: ToolRegistry, - config: AgentConfig, - managers: ManagerBundle, + config: LoopConfig, + managers: LoopRuntime, ) -> Self { Self { client, @@ -538,7 +611,7 @@ impl BareLoop { /// Create from builder parts (produced by `AgentBuilder::into_raw_parts()`). /// - /// This is the most flexible constructor. It accepts all components + /// Most flexible constructor. It accepts all components /// individually, making it suitable for re-assembly after a builder /// has been consumed via `into_raw_parts()`. /// @@ -546,7 +619,7 @@ impl BareLoop { /// /// - `client` — The LLM API client, wrapped in `Arc`. /// - `tools` — The [`ToolRegistry`]. - /// - `managers` — A [`ManagerBundle`]. + /// - `managers` — A [`LoopRuntime`]. /// - `config` — Session parameters. /// /// # Example @@ -558,8 +631,8 @@ impl BareLoop { pub fn from_parts( client: Arc, tools: ToolRegistry, - managers: ManagerBundle, - config: AgentConfig, + managers: LoopRuntime, + config: LoopConfig, ) -> Self { Self { client, @@ -595,10 +668,10 @@ impl BareLoop { /// Get the agent configuration. /// - /// Returns a reference to the [`AgentConfig`] that governs session + /// Returns a reference to the [`LoopConfig`] that governs session /// parameters such as max turns, system prompt, and session ID. /// The config is immutable for the lifetime of the loop. - pub fn config(&self) -> &AgentConfig { + pub fn config(&self) -> &LoopConfig { &self.config } @@ -819,21 +892,21 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`AgentError::Config`] if the builder fails to produce a valid + /// Returns [`LoopError::Config`] if the builder fails to produce a valid /// pipeline (e.g. internal invariant violated). - pub fn set_pipeline(&mut self, builder: ToolPipelineBuilder) -> Result<(), AgentError> { + pub fn set_pipeline(&mut self, builder: ToolPipelineBuilder) -> Result<(), LoopError> { let pipeline = builder .core(Arc::clone(&self.tools)) .build() - .map_err(|e| AgentError::Config(e.to_string()))?; + .map_err(|e| LoopError::Config(e.to_string()))?; self.pipeline = Some(pipeline); Ok(()) } - /// Register a [`LoopObserver`](crate::core::observer::LoopObserver) with the manager bundle's observer host. + /// Register a [`LoopObserver`](crate::observer::LoopObserver) with the manager bundle's observer host. /// /// Plugins are called at lifecycle hook points inside the agent loop, - /// in registration order. See [`LoopObserver`](crate::core::observer::LoopObserver) + /// in registration order. See [`LoopObserver`](crate::observer::LoopObserver) /// for the trait definition and available hooks. /// /// Must be called before [`run()`](BareLoop::run). @@ -847,7 +920,7 @@ impl BareLoop { /// let mut agent = BareLoop::new(client, registry, config); /// agent.register_observer(Arc::new(MyObserver)); /// ``` - pub fn register_observer(&mut self, observer: Arc) { + pub fn register_observer(&mut self, observer: Arc) { self.managers.register_observer(observer); } @@ -857,7 +930,7 @@ impl BareLoop { /// Run the agent loop with the given user input. /// - /// This is the primary entry point. It: + /// Primary entry point. It: /// 1. Pushes the user message into the conversation /// 2. Loops: stream → accumulate → tool dispatch → feedback /// 3. Returns a [`SessionResult`] when done @@ -865,13 +938,13 @@ impl BareLoop { /// The loop terminates when one of these conditions is met: /// /// - **End turn** — the model emits `end_turn` with no tool calls. - /// - **Max turns exceeded** — [`config.max_turns`](AgentConfig::max_turns) - /// is reached, producing [`AgentError::MaxTurnsExceeded`]. + /// - **Max turns exceeded** — [`config.max_turns`](LoopConfig::max_turns) + /// is reached, producing [`LoopError::MaxTurnsExceeded`]. /// - **Cancellation** — [`cancel()`](BareLoop::cancel) was called, - /// producing [`AgentError::Cancelled`]; the caller should handle this + /// producing [`LoopError::Cancelled`]; the caller should handle this /// variant to distinguish user-initiated cancellation from other errors. /// - **API error** — the streaming request fails, producing - /// [`AgentError::Api`]. + /// [`LoopError::Api`]. /// /// # Observers /// @@ -891,7 +964,7 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`AgentError`] if: + /// Returns [`LoopError`] if: /// - The API call fails (after any retries) /// - Max turns is exceeded /// - A tool execution fails critically @@ -908,7 +981,7 @@ impl BareLoop { /// println!("Output tokens: {}", result.output_tokens); /// } /// ``` - pub async fn run(mut self, user_input: &str) -> Result { + pub async fn run(mut self, user_input: &str) -> Result { let session_id = self.config.session_id; let max_turns = self.config.max_turns; let start = Instant::now(); @@ -1083,13 +1156,13 @@ impl BareLoop { /// - `pattern` — The pattern returned by [`DetectionManager`]. /// - `turn` — Zero-based turn index, used in log messages. /// - /// [`DetectionConfig::stop_threshold`]: crate::loop_control::detection::DetectionConfig::stop_threshold - /// [`DetectionManager`]: crate::loop_control::detection::DetectionManager + /// [`DetectionConfig::stop_threshold`]: crate::detection::DetectionConfig::stop_threshold + /// [`DetectionManager`]: crate::detection::DetectionManager fn handle_detected_pattern( &self, pattern: &DetectedPattern, turn: usize, - ) -> Option> { + ) -> Option> { match pattern { DetectedPattern::NoPattern => None, @@ -1118,7 +1191,7 @@ impl BareLoop { turn, "stopping agent: loop threshold exceeded" ); - Some(Err(AgentError::LoopDetected { + Some(Err(LoopError::LoopDetected { message: format!("{pattern_description} repeated {repetitions} times"), })) } else { @@ -1147,7 +1220,7 @@ impl BareLoop { }); match action { - ConvergenceAction::Stop => Some(Err(AgentError::LoopDetected { + ConvergenceAction::Stop => Some(Err(LoopError::LoopDetected { message: "agent stopped: convergence detected".into(), })), ConvergenceAction::Warn => None, // log already happened @@ -1159,7 +1232,7 @@ impl BareLoop { } ConvergenceAction::AskUser => { // Not supported in BareLoop — treat as Stop - Some(Err(AgentError::LoopDetected { + Some(Err(LoopError::LoopDetected { message: "agent stopped: convergence detected, user input needed" .into(), })) @@ -1179,15 +1252,15 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`AgentError::Cancelled`] if the cancellation token is set. - /// Returns [`AgentError::Api`] if loop or convergence detection forces + /// Returns [`LoopError::Cancelled`] if the cancellation token is set. + /// Returns [`LoopError::Api`] if loop or convergence detection forces /// an abort, or if the underlying tool dispatch fails. async fn dispatch_and_record( &mut self, tool_calls: &[ToolCallInfo], turn: &TurnContext, budget: &mut SessionBudget, - ) -> Result<(), AgentError> { + ) -> Result<(), LoopError> { match self.dispatch_tools(tool_calls, turn.idx).await { Ok(results) => { budget.total_tool_calls = budget.total_tool_calls.saturating_add(results.len()); @@ -1280,15 +1353,15 @@ impl BareLoop { /// /// # Errors /// - /// Always returns `Err(error)`, passing through the original [`AgentError`]. + /// Always returns `Err(error)`, passing through the original [`LoopError`]. fn abort_turn_and_session( &self, budget: &SessionBudget, turn_duration: Duration, session_duration: Duration, reason: &str, - error: AgentError, - ) -> Result { + error: LoopError, + ) -> Result { self.managers.observers().on_turn_end(&TurnEndContext { turn: budget.turn_count, success: false, @@ -1297,7 +1370,7 @@ impl BareLoop { input_tokens: budget.input_tokens, output_tokens: budget.output_tokens, }); - let end_reason = if matches!(error, AgentError::Cancelled) { + let end_reason = if matches!(error, LoopError::Cancelled) { EndReason::Cancelled } else { EndReason::Error @@ -1314,19 +1387,19 @@ impl BareLoop { /// Abort the session after a tool-dispatch error. /// - /// Handles both [`AgentError::Cancelled`] and other errors uniformly. + /// Handles both [`LoopError::Cancelled`] and other errors uniformly. /// Turn-level notifications were already sent inside [`dispatch_and_record`]. /// /// # Errors /// - /// Always returns `Err(error)`, passing through the original [`AgentError`]. + /// Always returns `Err(error)`, passing through the original [`LoopError`]. fn abort_session_from_error( &self, - error: AgentError, + error: LoopError, session_duration: Duration, budget: &SessionBudget, - ) -> Result { - let end_reason = if matches!(error, AgentError::Cancelled) { + ) -> Result { + let end_reason = if matches!(error, LoopError::Cancelled) { EndReason::Cancelled } else { EndReason::Error @@ -1347,14 +1420,14 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`AgentError::Cancelled`] or [`AgentError::MaxTurnsExceeded`] + /// Returns [`LoopError::Cancelled`] or [`LoopError::MaxTurnsExceeded`] /// depending on the `reason` string. fn abort_session( &self, budget: &SessionBudget, session_duration: Duration, reason: AbortReason, - ) -> Result { + ) -> Result { let end_reason = match &reason { AbortReason::Cancelled => EndReason::Cancelled, AbortReason::MaxTurnsExceeded => EndReason::MaxTurns, @@ -1367,8 +1440,8 @@ impl BareLoop { duration_secs: session_duration.as_secs(), }); match reason { - AbortReason::Cancelled => Err(AgentError::Cancelled), - AbortReason::MaxTurnsExceeded => Err(AgentError::MaxTurnsExceeded { + AbortReason::Cancelled => Err(LoopError::Cancelled), + AbortReason::MaxTurnsExceeded => Err(LoopError::MaxTurnsExceeded { max: self.config.max_turns, }), } @@ -1382,7 +1455,7 @@ impl BareLoop { #[cfg(test)] mod tests { use super::*; - use crate::api_error::ApiError; + use crate::api::error::ApiError; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, PartStart, Usage, @@ -1661,7 +1734,7 @@ mod tests { trait PopFront { /// Remove and return the first element, shifting the rest left. /// - /// Returns `None` if the vector is empty. This is an O(n) + /// Returns `None` if the vector is empty. O(n) /// operation because it calls `Vec::remove(0)`. Acceptable /// for test-only code with small queues. fn pop_front(&mut self) -> Option; @@ -1779,7 +1852,7 @@ mod tests { // Counting Plugin (test helper) // ================================================== - /// A [`LoopObserver`](crate::core::observer::LoopObserver) that counts + /// A [`LoopObserver`](crate::observer::LoopObserver) that counts /// how many times each hook fires. /// /// Uses [`AtomicUsize`] counters with `SeqCst` ordering so that @@ -1814,32 +1887,32 @@ mod tests { } } - impl crate::core::observer::LoopObserver for CountingObserver { + impl crate::observer::LoopObserver for CountingObserver { fn name(&self) -> &str { "counting" } - fn on_session_start(&self, _ctx: &crate::core::observer::SessionStartContext) { + fn on_session_start(&self, _ctx: &crate::observer::SessionStartContext) { self.session_starts.fetch_add(1, Ordering::SeqCst); } - fn on_session_end(&self, _ctx: &crate::core::observer::SessionEndContext) { + fn on_session_end(&self, _ctx: &crate::observer::SessionEndContext) { self.session_ends.fetch_add(1, Ordering::SeqCst); } - fn on_turn_start(&self, _ctx: &crate::core::observer::TurnStartContext) { + fn on_turn_start(&self, _ctx: &crate::observer::TurnStartContext) { self.turn_starts.fetch_add(1, Ordering::SeqCst); } - fn on_turn_end(&self, _ctx: &crate::core::observer::TurnEndContext) { + fn on_turn_end(&self, _ctx: &crate::observer::TurnEndContext) { self.turn_ends.fetch_add(1, Ordering::SeqCst); } - fn on_tool_pre(&self, _ctx: &crate::core::observer::ToolPreContext) { + fn on_tool_pre(&self, _ctx: &crate::observer::ToolPreContext) { self.tool_pres.fetch_add(1, Ordering::SeqCst); } - fn on_tool_post(&self, _ctx: &crate::core::observer::ToolPostContext) { + fn on_tool_post(&self, _ctx: &crate::observer::ToolPostContext) { self.tool_posts.fetch_add(1, Ordering::SeqCst); } } @@ -1848,13 +1921,13 @@ mod tests { // Test Helpers // ================================================== - /// Create a default [`AgentConfig`] with `max_turns = 10`. + /// Create a default [`LoopConfig`] with `max_turns = 10`. /// /// Most tests use this as a baseline. Tests that need a different /// max-turns value mutate the returned config before constructing /// the loop. - fn make_config() -> AgentConfig { - AgentConfig { + fn make_config() -> LoopConfig { + LoopConfig { max_turns: 10, ..Default::default() } @@ -1905,7 +1978,7 @@ mod tests { } /// Verify that exceeding `max_turns` returns - /// [`AgentError::MaxTurnsExceeded`] and reports `success = false`. + /// [`LoopError::MaxTurnsExceeded`] and reports `success = false`. #[tokio::test] async fn test_bare_loop_max_turns_exceeded() { let client = MockClient::new("test-model"); @@ -1928,13 +2001,13 @@ mod tests { let result = agent.run("Keep going").await; assert!(result.is_err()); match result.unwrap_err() { - AgentError::MaxTurnsExceeded { max } => assert_eq!(max, 3), + LoopError::MaxTurnsExceeded { max } => assert_eq!(max, 3), other => panic!("Expected MaxTurnsExceeded, got: {other}"), } } /// Verify that calling [`cancel()`](BareLoop::cancel) mid-session - /// returns [`AgentError::Cancelled`]. + /// returns [`LoopError::Cancelled`]. #[tokio::test] async fn test_bare_loop_cancellation() { let client = MockClient::new("test-model"); @@ -1950,13 +2023,13 @@ mod tests { let result = agent.run("Hi").await; assert!(result.is_err()); match result.unwrap_err() { - AgentError::Cancelled => {} + LoopError::Cancelled => {} other => panic!("Expected Cancelled error, got: {other}"), } } /// Verify that an API error during streaming propagates as - /// [`AgentError::Api`] and marks the session as failed. + /// [`LoopError::Api`] and marks the session as failed. #[tokio::test] async fn test_bare_loop_api_error() { // The mock will return an error @@ -1966,7 +2039,7 @@ mod tests { let result = agent.run("Hi").await; assert!(result.is_err()); match result.unwrap_err() { - AgentError::Api(msg) => assert!(msg.contains("No more mock responses")), + LoopError::Api(msg) => assert!(msg.contains("No more mock responses")), other => panic!("Expected Api error, got: {other}"), } } @@ -1976,7 +2049,7 @@ mod tests { // ================================================== /// Verify that requesting a tool not present in the registry produces - /// a soft error result (not a hard [`AgentError`]), allowing the model + /// a soft error result (not a hard [`LoopError`]), allowing the model /// to see the failure and adapt. #[tokio::test] async fn test_tool_not_found_returns_error_result() { @@ -2250,7 +2323,7 @@ mod tests { fn test_from_parts() { let client = MockClient::new("test-model"); let config = make_config(); - let managers = ManagerBundle::new(); + let managers = LoopRuntime::new(); let agent = BareLoop::from_parts(Arc::new(client), ToolRegistry::new(), managers, config); assert!(agent.conversation().is_empty()); @@ -2299,7 +2372,7 @@ mod tests { } /// Verify that setting `max_turns = 0` immediately triggers - /// [`AgentError::MaxTurnsExceeded`] before any API call. + /// [`LoopError::MaxTurnsExceeded`] before any API call. #[tokio::test] async fn test_loop_terminates_with_max_turns_0() { let client = MockClient::new("test-model"); @@ -2312,7 +2385,7 @@ mod tests { let result = agent.run("Hi").await; assert!(result.is_err()); match result.unwrap_err() { - AgentError::MaxTurnsExceeded { max } => assert_eq!(max, 0), + LoopError::MaxTurnsExceeded { max } => assert_eq!(max, 0), other => panic!("Expected MaxTurnsExceeded, got: {other}"), } } @@ -2322,7 +2395,7 @@ mod tests { // ================================================== /// Verify that requesting a nonexistent tool produces a soft error - /// result (not a hard [`AgentError`]), allowing the model to see + /// result (not a hard [`LoopError`]), allowing the model to see /// the error and respond gracefully. #[tokio::test] async fn test_tool_error_is_soft_not_hard() { @@ -2472,7 +2545,7 @@ mod tests { } } - impl crate::engine::middleware::ToolMiddleware for TurnNumberCapture { + impl crate::middleware::ToolMiddleware for TurnNumberCapture { fn name(&self) -> &str { "turn_capture" } @@ -2483,9 +2556,7 @@ mod tests { next: &'a ToolPipeline, ) -> std::pin::Pin< Box< - dyn std::future::Future - + Send - + 'a, + dyn std::future::Future + Send + 'a, >, > { self.turns.lock().unwrap().push(ctx.turn_number); diff --git a/src/engine/bare/compact.rs b/src/engine/bare/compact.rs index bbf241f..8c5794d 100644 --- a/src/engine/bare/compact.rs +++ b/src/engine/bare/compact.rs @@ -4,11 +4,11 @@ //! When a [`ContextManager`] is configured, checks token usage after each //! tool dispatch and triggers compaction if usage exceeds the threshold. -use super::{AgentError, ApiClient, BareLoop, EnsureContextResult, Instant}; +use super::{ApiClient, BareLoop, EnsureContextResult, Instant, LoopError}; #[cfg(feature = "hooks")] use super::{CompactTrigger, PostCompactContext, PreCompactContext}; -use crate::core::observer::CompactionContext; +use crate::observer::CompactedContext; impl BareLoop { /// Check if context compaction is needed and perform it if so. @@ -16,20 +16,21 @@ impl BareLoop { /// When a [`ContextManager`] is configured, this method: /// 1. Calls [`ContextManager::ensure_context_fits()`] to check token usage. /// 2. If compaction occurred, replaces `self.conversation` with the compacted messages. - /// 3. Notifies observers via [`LoopObserver::on_compaction`](crate::core::observer::LoopObserver::on_compaction). + /// 3. Notifies observers via [`LoopObserver::on_compaction`](crate::observer::LoopObserver::on_compaction). /// /// When no `ContextManager` is set, this is a no-op. /// /// # Errors /// - /// Returns [`AgentError::ContextExceeded`] if compaction was needed but failed + /// Returns [`LoopError::ContextExceeded`] if compaction was needed but failed /// (i.e. the conversation exceeds the context window and the compactor /// could not reduce it sufficiently). - pub(super) async fn maybe_compact_context(&mut self, turn: usize) -> Result<(), AgentError> { + pub(super) async fn maybe_compact_context(&mut self, turn: usize) -> Result<(), LoopError> { let Some(ref ctx_manager) = self.context_manager else { return Ok(()); }; + #[cfg(feature = "hooks")] let messages_before = self.conversation.len(); // Pre-compact hook check @@ -65,10 +66,12 @@ impl BareLoop { match result { Ok(EnsureContextResult::Compacted(outcome)) => { self.conversation = outcome.messages; + #[cfg(feature = "hooks")] let messages_after = self.conversation.len(); - self.managers.observers().on_compaction(&CompactionContext { - messages_before, - messages_after, + let tokens_before = outcome.tokens_after.saturating_add(outcome.tokens_saved); + self.managers.observers().on_compaction(&CompactedContext { + tokens_before, + tokens_after: outcome.tokens_after, tokens_saved: outcome.tokens_saved, }); @@ -93,7 +96,7 @@ impl BareLoop { self.conversation = messages; Ok(()) } - Err(overflow) => Err(AgentError::ContextExceeded { + Err(overflow) => Err(LoopError::ContextExceeded { used: overflow.tokens_used, limit: overflow.context_window, }), diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 1b1dc0a..ef1cba4 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -7,14 +7,14 @@ #[cfg(feature = "hooks")] use super::HookAction; use super::{ - AgentError, ApiClient, Arc, BareLoop, Duration, Instant, PermissionCheck, RecoveryAction, - ReflectionContext, ToolCallInfo, ToolContent, ToolContext, ToolDispatchContext, - ToolDispatchResult, ToolPipeline, + ApiClient, Arc, BareLoop, Correction, CorrectionResult, Duration, Instant, LoopError, + PermissionCheck, RecoveryAction, ReflectionContext, ToolCallInfo, ToolContent, ToolContext, + ToolDispatchContext, ToolDispatchResult, ToolPipeline, }; #[cfg(feature = "hooks")] use super::{PostToolUseContext, PreToolUseContext}; -use crate::core::observer::{ToolPostContext, ToolPreContext}; -use crate::loop_control::loop_detector::{self, Operation}; +use crate::detection::loop_detector::{self, Operation}; +use crate::observer::{ToolPostContext, ToolPreContext}; /// Result of deciding what to do after a tool error during recovery. /// @@ -37,7 +37,7 @@ impl BareLoop { /// /// Tool execution is **sequential** so that cancellation can be /// checked between invocations. A tool that is not found in the - /// registry produces a soft error result (not a hard [`AgentError`]), + /// registry produces a soft error result (not a hard [`LoopError`]), /// allowing the model to recover. /// /// When a tool returns an error (execution failure or not-found), @@ -46,22 +46,22 @@ impl BareLoop { /// attempts use the delay specified by the [`RecoveryAction`]. /// /// Observers are notified before and after each tool invocation via - /// [`LoopObserver::on_tool_pre`](crate::core::observer::LoopObserver::on_tool_pre) and - /// [`LoopObserver::on_tool_post`](crate::core::observer::LoopObserver::on_tool_post). + /// [`LoopObserver::on_tool_pre`](crate::observer::LoopObserver::on_tool_pre) and + /// [`LoopObserver::on_tool_post`](crate::observer::LoopObserver::on_tool_post). /// /// # Errors /// - /// Returns [`AgentError::Cancelled`] if the cancellation flag is set + /// Returns [`LoopError::Cancelled`] if the cancellation flag is set /// between tool invocations. pub(super) async fn dispatch_tools( &self, tool_calls: &[ToolCallInfo], turn_idx: usize, - ) -> Result, AgentError> { + ) -> Result, LoopError> { let mut results = Vec::with_capacity(tool_calls.len()); for tc in tool_calls { if self.is_cancelled() { - return Err(AgentError::Cancelled); + return Err(LoopError::Cancelled); } let result = self.dispatch_tool_with_recovery(tc, turn_idx).await?; results.push(result); @@ -86,19 +86,20 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`AgentError::Cancelled`] if the cancellation signal fires + /// Returns [`LoopError::Cancelled`] if the cancellation signal fires /// during tool execution or between retry attempts. async fn dispatch_tool_with_recovery( &self, tc: &ToolCallInfo, turn_idx: usize, - ) -> Result { + ) -> Result { let tool_context = self.build_tool_context(); let mut attempt: u32 = 0; + let mut tc = tc.clone(); loop { if self.is_cancelled() { - return Err(AgentError::Cancelled); + return Err(LoopError::Cancelled); } self.managers.observers().on_tool_pre(&ToolPreContext { @@ -107,20 +108,20 @@ impl BareLoop { tool_call_id: tc.id.clone(), }); - if let Some(blocked) = self.check_pre_tool_use_hooks(tc, turn_idx) { + if let Some(blocked) = self.check_pre_tool_use_hooks(&tc, turn_idx) { return Ok(blocked); } - if let Some(blocked) = self.pre_detection(tc, turn_idx) { + if let Some(blocked) = self.pre_detection(&tc, turn_idx) { return Ok(blocked); } let start = Instant::now(); let tool_result = self - .dispatch_tool(tc, &tool_context, start, turn_idx) + .dispatch_tool(&tc, &tool_context, start, turn_idx) .await?; - self.post_detection(tc, &tool_result); + self.post_detection(&tc, &tool_result); self.managers.observers().on_tool_post(&ToolPostContext { turn: turn_idx, tool: tc.name.clone(), @@ -128,7 +129,7 @@ impl BareLoop { is_error: tool_result.is_error, duration: tool_result.duration, }); - self.notify_post_tool_use_hooks(tc, &tool_result, turn_idx); + self.notify_post_tool_use_hooks(&tc, &tool_result, turn_idx); self.record_tool_health(tc.name.as_str(), &tool_result); if !tool_result.is_error { @@ -136,12 +137,24 @@ impl BareLoop { } match self - .recovery_wait_or_return(tc, &tool_result, attempt) + .recovery_wait_or_return(&tc, &tool_result, attempt) .await { - Ok(next_attempt) => attempt = next_attempt, + Ok((next_attempt, correction)) => { + attempt = next_attempt; + if let Some(ref correction) = correction { + let correction_result = tc.apply_correction(correction, &tool_result); + if let CorrectionResult::Failed(msg) = &correction_result { + tracing::warn!( + tool = %tc.name, + error = %msg, + "correction failed to produce a usable retry" + ); + } + } + } Err(RecoveryOutcome::SoftError(returned_result)) => return Ok(returned_result), - Err(RecoveryOutcome::Cancelled) => return Err(AgentError::Cancelled), + Err(RecoveryOutcome::Cancelled) => return Err(LoopError::Cancelled), } } } @@ -207,7 +220,7 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`AgentError::Cancelled`] if the cancel signal fires + /// Returns [`LoopError::Cancelled`] if the cancel signal fires /// during tool execution. async fn dispatch_tool( &self, @@ -215,7 +228,7 @@ impl BareLoop { tool_context: &ToolContext, start: Instant, turn_idx: usize, - ) -> Result { + ) -> Result { if let Some(ref pipeline) = self.pipeline { return self .dispatch_via_pipeline(pipeline, tc, tool_context, turn_idx) @@ -227,7 +240,7 @@ impl BareLoop { let call_result = tokio::select! { r = tool.call(tc.input.clone(), tool_context) => r, () = cancel.notified() => { - return Err(AgentError::Cancelled); + return Err(LoopError::Cancelled); } }; match call_result { @@ -267,7 +280,7 @@ impl BareLoop { fn tool_not_found(&self, tc: &ToolCallInfo) -> ToolDispatchResult { let available: Vec = self.tools.tool_names(); let available_refs: Vec<&str> = available.iter().map(String::as_str).collect(); - let error = AgentError::tool_not_found(&tc.name, &available_refs); + let error = LoopError::tool_not_found(&tc.name, &available_refs); let error_msg = error.to_string(); ToolDispatchResult { tool_call_id: tc.id.clone(), @@ -282,20 +295,21 @@ impl BareLoop { /// /// Consults the reflector and recovery strategy. On `Retry`, sleeps for /// the prescribed delay (cancellation-aware) and returns the updated - /// attempt count via `Ok`. On all other recovery actions, returns the - /// original error result via `Err` (which ends the retry loop). + /// attempt count and the [`Correction`] (if any) via `Ok`. On all other + /// recovery actions, returns the original error result via `Err` (which + /// ends the retry loop). /// /// # Errors /// - /// Returns `Err(ToolDispatchResult)` when the recovery strategy decides + /// Returns `Err(RecoveryOutcome)` when the recovery strategy decides /// not to retry — the caller should return this as a soft error. async fn recovery_wait_or_return( &self, tc: &ToolCallInfo, tool_result: &ToolDispatchResult, attempt: u32, - ) -> Result { - let recovery_action = self.recover_tool_error(tc, tool_result, attempt).await; + ) -> Result<(u32, Option), RecoveryOutcome> { + let (recovery_action, correction) = self.recover_tool_error(tc, tool_result, attempt).await; match recovery_action { RecoveryAction::Retry { delay } => { let next_attempt = attempt.saturating_add(1); @@ -308,7 +322,7 @@ impl BareLoop { return Err(RecoveryOutcome::Cancelled); } } - Ok(next_attempt) + Ok((next_attempt, correction)) } RecoveryAction::Skip(_) | RecoveryAction::AskUser(_) | RecoveryAction::Fail(_) => { Err(RecoveryOutcome::SoftError(tool_result.clone())) @@ -430,7 +444,7 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`AgentError::Cancelled`] if the cancel signal fires + /// Returns [`LoopError::Cancelled`] if the cancel signal fires /// during pipeline dispatch. async fn dispatch_via_pipeline( &self, @@ -438,7 +452,7 @@ impl BareLoop { tc: &ToolCallInfo, tool_context: &ToolContext, turn_idx: usize, - ) -> Result { + ) -> Result { let ctx = ToolDispatchContext { tool_name: tc.name.clone(), input: tc.input.clone(), @@ -452,7 +466,7 @@ impl BareLoop { let dispatch_result = tokio::select! { r = pipeline.invoke(ctx) => r, () = cancel.notified() => { - return Err(AgentError::Cancelled); + return Err(LoopError::Cancelled); } }; Ok(ToolDispatchResult { @@ -473,12 +487,16 @@ impl BareLoop { /// Calls [`Reflector::analyze()`] and then [`RecoveryStrategy::decide()`]. /// If the reflector itself fails, logs the error and returns /// [`RecoveryAction::Fail`] (conservative default). + /// + /// Returns the [`RecoveryAction`] alongside the [`Correction`] (if any) + /// produced by the reflector. The correction is threaded through so the + /// retry loop can apply it before re-dispatching. async fn recover_tool_error( &self, tc: &ToolCallInfo, result: &ToolDispatchResult, attempt: u32, - ) -> RecoveryAction { + ) -> (RecoveryAction, Option) { let error_msg = match &result.output { ToolContent::Text(msg) => msg.clone(), ToolContent::Multipart(_) => result.output.to_string(), @@ -495,11 +513,14 @@ impl BareLoop { .await else { // Reflector failed — conservatively fail. - return RecoveryAction::Fail(error_msg); + return (RecoveryAction::Fail(error_msg), None); }; - self.recovery + let correction = analysis.correction.clone(); + let action = self + .recovery .decide(&analysis, attempt, Self::MAX_RECOVERY_ATTEMPTS) - .await + .await; + (action, correction) } } diff --git a/src/engine/bare/emission.rs b/src/engine/bare/emission.rs index 4c7eac0..cbb7ad5 100644 --- a/src/engine/bare/emission.rs +++ b/src/engine/bare/emission.rs @@ -1,7 +1,7 @@ //! Session lifecycle notifications. //! //! Split from [`BareLoop`] for clarity — these methods dispatch to -//! the [`ObserverHost`](crate::core::observer::ObserverHost) and the hook executor. +//! the [`ObserverHost`](crate::observer::ObserverHost) and the hook executor. //! //! Only session start/end live here because they do *two* things: //! observer notification + hook dispatch. All other observer notifications @@ -9,12 +9,12 @@ //! `self.managers.observers().on_*()`. use super::{ApiClient, BareLoop, Duration, EndReason, SessionEndInfo}; -use crate::core::observer::{SessionEndContext, SessionStartContext}; #[cfg(feature = "hooks")] use crate::hooks::context::{ SessionEndContext as HookSessionEndContext, SessionEndReason, SessionStartContext as HookSessionStartContext, }; +use crate::observer::{SessionEndContext, SessionStartContext}; // ================================================== // Session lifecycle notifications diff --git a/src/engine/bare/stream.rs b/src/engine/bare/stream.rs index 1716003..613dbb7 100644 --- a/src/engine/bare/stream.rs +++ b/src/engine/bare/stream.rs @@ -5,7 +5,7 @@ //! (retry, timeout, fallback). Otherwise, uses basic inline logic. use super::{ - AgentError, ApiClient, BareLoop, Message, StreamAccumulator, StreamEvent, StreamStopReason, + ApiClient, BareLoop, LoopError, Message, StreamAccumulator, StreamEvent, StreamStopReason, Usage, }; use crate::stream::handler::{StreamHandler, StreamHandlerError}; @@ -38,11 +38,11 @@ impl BareLoop { /// /// # Errors /// - /// Returns [`AgentError::Api`] if any stream event is an error. - /// Returns [`AgentError::Cancelled`] if the cancellation signal fires mid-stream. + /// Returns [`LoopError::Api`] if any stream event is an error. + /// Returns [`LoopError::Cancelled`] if the cancellation signal fires mid-stream. pub(super) async fn stream_turn( &self, - ) -> Result<(Message, Option, StreamStopReason), AgentError> { + ) -> Result<(Message, Option, StreamStopReason), LoopError> { // Delegate to StreamHandler if configured. if let Some(ref handler) = self.stream_handler { return self.stream_turn_via_handler(handler).await; @@ -60,7 +60,7 @@ impl BareLoop { let event_result = tokio::select! { event = stream.next() => event, () = self.cancelled.notified() => { - return Err(AgentError::Cancelled); + return Err(LoopError::Cancelled); } }; @@ -74,10 +74,10 @@ impl BareLoop { } accumulator .process(&event) - .map_err(|e| AgentError::Api(format!("stream accumulation error: {e}")))?; + .map_err(|e| LoopError::Api(format!("stream accumulation error: {e}")))?; } Some(Err(api_error)) => { - return Err(AgentError::Api(api_error.to_string())); + return Err(LoopError::Api(api_error.to_string())); } None => break, } @@ -99,15 +99,15 @@ impl BareLoop { /// # Errors /// /// Maps [`StreamHandlerError`] variants to the appropriate - /// [`AgentError`] variants: - /// - [`Cancelled`](StreamHandlerError::Cancelled) → [`AgentError::Cancelled`] - /// - [`InitFailed`](StreamHandlerError::InitFailed) → [`AgentError::Api`] - /// - [`StreamFailed`](StreamHandlerError::StreamFailed) → [`AgentError::Api`] - /// - [`FallbackFailed`](StreamHandlerError::FallbackFailed) → [`AgentError::Api`] + /// [`LoopError`] variants: + /// - [`Cancelled`](StreamHandlerError::Cancelled) → [`LoopError::Cancelled`] + /// - [`InitFailed`](StreamHandlerError::InitFailed) → [`LoopError::Api`] + /// - [`StreamFailed`](StreamHandlerError::StreamFailed) → [`LoopError::Api`] + /// - [`FallbackFailed`](StreamHandlerError::FallbackFailed) → [`LoopError::Api`] async fn stream_turn_via_handler( &self, handler: &StreamHandler, - ) -> Result<(Message, Option, StreamStopReason), AgentError> { + ) -> Result<(Message, Option, StreamStopReason), LoopError> { let system = self.config.system_prompt.clone(); let tool_schemas = self.build_tool_schemas(); let result = handler @@ -123,25 +123,25 @@ impl BareLoop { Ok((result.message, result.usage, result.stop_reason)) } - /// Map a [`StreamHandlerError`] to an [`AgentError`]. + /// Map a [`StreamHandlerError`] to an [`LoopError`]. /// /// Preserves cancellation semantics — - /// [`StreamHandlerError::Cancelled`] maps to [`AgentError::Cancelled`]. - /// All other variants map to [`AgentError::Api`] with a descriptive + /// [`StreamHandlerError::Cancelled`] maps to [`LoopError::Cancelled`]. + /// All other variants map to [`LoopError::Api`] with a descriptive /// message. - fn map_handler_error(error: StreamHandlerError) -> AgentError { + fn map_handler_error(error: StreamHandlerError) -> LoopError { match error { - StreamHandlerError::Cancelled => AgentError::Cancelled, + StreamHandlerError::Cancelled => LoopError::Cancelled, StreamHandlerError::InitFailed(outcome) => { - AgentError::Api(format!("stream init failed: {outcome}")) + LoopError::Api(format!("stream init failed: {outcome}")) } StreamHandlerError::StreamFailed(outcome) => { - AgentError::Api(format!("stream failed: {outcome}")) + LoopError::Api(format!("stream failed: {outcome}")) } StreamHandlerError::FallbackFailed { stream_outcome, fallback_error, - } => AgentError::Api(format!( + } => LoopError::Api(format!( "stream ({stream_outcome}) and fallback failed: {fallback_error}" )), } diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs new file mode 100644 index 0000000..3a4b3f2 --- /dev/null +++ b/src/engine/loop_core.rs @@ -0,0 +1,552 @@ +//! Agent core trait and foundational lifecycle types. +//! +//! Fundamental operations every agent must support, plus the data types +//! that flow through the agent lifecycle: configuration ([`LoopConfig`]), +//! lifecycle state ([`LoopState`]), turn and session results +//! ([`TurnResult`], [`SessionResult`]), and tool call representations +//! ([`ToolCall`]). +//! +//! # Lifecycle +//! +//! ```text +//! initialize(config) +//! → process_turn(input) [repeated] +//! → process_turn(input) +//! → ... +//! → should_continue() → false +//! finalize() +//! ``` +//! +//! # Implementing +//! +//! At a minimum you must provide [`initialize`](Loop::initialize), +//! [`process_turn`](Loop::process_turn), +//! [`should_continue`](Loop::should_continue), +//! [`finalize`](Loop::finalize), +//! [`state`](Loop::state), and +//! [`cancel`](Loop::cancel). +//! +//! # Quick Start +//! +//! ``` +//! use loopctl::engine::loop_core::{LoopConfig, LoopState, TurnResult, SessionResult}; +//! +//! let config = LoopConfig::default(); +//! assert_eq!(config.max_turns, 200); +//! +//! let turn = TurnResult::completed("Task done."); +//! assert!(turn.is_complete); +//! +//! let session = SessionResult::success(config.session_id); +//! assert!(session.success); +//! ``` + +use std::future::Future; +use std::pin::Pin; +use std::time::{Duration, SystemTime}; + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::error::LoopError; + +// Re-export LoopConfig for convenience — it lives in crate::config. +pub use crate::config::LoopConfig; + +// Re-export ToolDispatchResult so consumers of this module see the full +// tool-call type family in one place. +pub use crate::tool::ToolDispatchResult; + +// ================================================== +// LoopState +// ================================================== + +/// The lifecycle state of an agent. +/// +/// Models the agent as an explicit state machine, making transitions clear +/// and invalid states unrepresentable. The framework reads and writes this +/// enum to drive the agent loop and report status to observers. +/// +/// ```text +/// Idle → Processing → WaitingForTool → Processing → ... → Completed/Failed +/// ↘ Compacting ↗ +/// ↘ Reflecting ↗ +/// ``` +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LoopState { + /// The agent is idle, waiting for a user message. + /// + /// Initial state after initialization and the state + /// the agent returns to between user inputs. + /// + /// No background work is performed while idle. + Idle, + + /// The agent is actively processing a turn. + /// + /// Entered when the agent core begins processing a user message + /// or a tool result. The `turn` field tracks progress for observers. + /// + /// The agent transitions here from [`Idle`](LoopState::Idle) or + /// [`WaitingForTool`](LoopState::WaitingForTool). + Processing { + /// Current turn number (0-indexed). Used to enforce [`LoopConfig::max_turns`]. + turn: usize, + }, + + /// The agent is waiting for a tool to complete. + /// + /// Entered after the LLM requests a tool call. The framework + /// dispatches the tool and waits for its result before returning + /// to [`Processing`](LoopState::Processing). + WaitingForTool { + /// Name of the tool being executed. Never empty — always matches a registered tool name. + tool: String, + + /// When the tool call started. Used to compute execution duration. + started_at: SystemTime, + }, + + /// The agent is compacting its conversation context. + /// + /// Entered when token usage exceeds the + /// [`compact_threshold`](LoopConfig::compact_threshold) or when + /// compaction is explicitly requested. The agent summarizes older + /// messages to free context space, then returns to + /// [`Processing`](LoopState::Processing). + Compacting { + /// Why compaction was triggered. See + /// [`CompactReason`](crate::compact::types::CompactReason). + reason: crate::compact::types::CompactReason, + }, + + /// The agent is reflecting on a failure and preparing a correction. + /// + /// Entered when a tool call fails and the reflection system is + /// enabled (via `Feature::Reflection`). + /// The agent analyzes the error and produces a + /// [`Correction`](crate::reflection::Correction) before retrying. + Reflecting { + /// Number of errors being analyzed; higher counts may need + /// [`ApproachChange`](crate::reflection::CorrectionType::ApproachChange). + error_count: usize, + }, + + /// The agent has completed its task. + /// + /// Terminal state. The framework calls + /// `Loop::finalize` to produce + /// a [`SessionResult`]. + Completed { + /// May be empty if the agent produced only tool calls. + summary: String, + }, + + /// The agent has failed with an unrecoverable error. + /// + /// Terminal state. The error is propagated through + /// [`SessionResult::error`]. + /// + /// No further turns will be executed after entering this state. + Failed { error: String }, +} + +// ================================================== +// TurnResult +// ================================================== + +/// Result of a single agent turn (one API call → response cycle). +/// +/// Produced by `Loop::process_turn` +/// after each LLM interaction. Contains the response text, any tool calls +/// requested, token usage, and timing information. +/// +/// # Construction +/// +/// Use [`TurnResult::completed`] for a terminal response or +/// [`TurnResult::continuing`] for a response that should keep the loop running. +/// Production code typically constructs this from the raw API response. +/// +/// ``` +/// use loopctl::engine::loop_core::TurnResult; +/// +/// let done = TurnResult::completed("All tasks finished."); +/// assert!(done.is_complete); +/// +/// let more = TurnResult::continuing("Still working..."); +/// assert!(!more.is_complete); +/// ``` +#[derive(Debug, Clone)] +pub struct TurnResult { + /// May be empty if the response consists entirely of tool calls. + pub text: String, + /// Tool calls requested by the LLM in this turn. + pub tool_calls: Vec, + /// Results from dispatching the requested tool calls. + pub tool_results: Vec, + /// Input tokens used (system prompt + history + user message). Reported by the provider. + pub input_tokens: u64, + /// Output tokens in the API response. Reported by the provider. + pub output_tokens: u64, + /// Wall-clock duration (API request → full response + tool execution). + pub duration: Duration, + /// When `true`, the framework skips `should_continue` and proceeds to finalization. + pub is_complete: bool, + /// Used by the framework to decide whether to dispatch tools ([`StopReason::ToolCall`]) or continue. + pub stop_reason: StopReason, +} + +impl TurnResult { + /// Create a completed turn result with a simple text response. + /// + /// Sets [`is_complete`](TurnResult::is_complete) to `true` and all + /// token counters to zero. Use this for the final turn of a session. + /// + /// # Example + /// + /// ``` + /// use loopctl::engine::loop_core::TurnResult; + /// + /// let result = TurnResult::completed("The file has been written successfully."); + /// assert!(result.is_complete); + /// assert_eq!(result.tool_calls.len(), 0); + /// ``` + #[must_use] + pub fn completed(text: impl Into) -> Self { + Self { + text: text.into(), + tool_calls: Vec::new(), + tool_results: Vec::new(), + input_tokens: 0, + output_tokens: 0, + duration: Duration::ZERO, + is_complete: true, + stop_reason: StopReason::EndTurn, + } + } + + /// Create a turn result that should continue with more turns. + /// + /// Sets [`is_complete`](TurnResult::is_complete) to `false`, indicating + /// that the agent loop should keep running. + /// + /// # Example + /// + /// ``` + /// use loopctl::engine::loop_core::TurnResult; + /// + /// let result = TurnResult::continuing("I need to read the file first..."); + /// assert!(!result.is_complete); + /// ``` + #[must_use] + pub fn continuing(text: impl Into) -> Self { + Self { + text: text.into(), + tool_calls: Vec::new(), + tool_results: Vec::new(), + input_tokens: 0, + output_tokens: 0, + duration: Duration::ZERO, + is_complete: false, + stop_reason: StopReason::EndTurn, + } + } + + /// Check if this turn included any tool calls. + /// + /// Returns `true` when [`tool_calls`](TurnResult::tool_calls) is + /// non-empty, indicating that the LLM requested tool execution. + #[must_use] + pub fn has_tool_calls(&self) -> bool { + !self.tool_calls.is_empty() + } + + /// Total tokens (input + output) for this turn. + /// + /// Sums [`input_tokens`](TurnResult::input_tokens) + /// and [`output_tokens`](TurnResult::output_tokens). + #[must_use] + pub fn total_tokens(&self) -> u64 { + self.input_tokens.saturating_add(self.output_tokens) + } +} + +// ================================================== +// StopReason +// ================================================== + +/// Why the API stopped generating. +/// +/// Mirrors the stop reasons returned by LLM APIs. The framework uses this +/// to determine the next step: dispatch tools, continue the conversation, +/// or end the session. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum StopReason { + /// The model decided to stop (natural end of turn). + /// + /// The LLM finished its response without requesting tools or hitting + /// any limit. The framework should check + /// [`TurnResult::is_complete`] to decide whether to continue. + EndTurn, + + /// The model requested tool execution. + /// + /// The LLM response contains one or more tool calls in + /// [`TurnResult::tool_calls`]. The framework should dispatch them + /// and feed the results back. + ToolCall, + + /// The maximum token limit was reached. + /// + /// The LLM hit the [`LoopConfig::max_tokens`] limit before + /// finishing. The response may be truncated. The framework may + /// choose to continue the turn to let the model complete its output. + MaxTokens, + + /// The stop sequence was encountered. + /// + /// The model generated a configured stop sequence. Rare in + /// standard usage; typically indicates custom API configuration. + StopSequence, +} + +// ================================================== +// ToolCall +// ================================================== + +/// A tool call requested by the agent. +/// +/// Represents a single tool invocation that the LLM has requested during a +/// turn. The framework matches each `ToolCall` to a registered tool, executes +/// it, and produces a [`ToolDispatchResult`] with the output. +/// +/// # Serialization +/// +/// Implements `Serialize` and `Deserialize` for persistence and inter-process +/// communication. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolCall { + /// Assigned by the LLM API. Correlates with [`ToolDispatchResult::tool_call_id`]. + pub id: String, + + /// Must match a tool registered in the `ToolRegistry`. + pub tool: String, + + /// JSON object whose schema depends on the tool. + pub input: serde_json::Value, +} + +// ================================================== +// SessionResult +// ================================================== + +/// Summary of a complete agent session. +/// +/// Produced by `Loop::finalize` +/// after the last turn. Aggregates all session-level metrics: total turns, +/// tokens, duration, tool calls, and final output. +/// +/// # Construction +/// +/// Use [`SessionResult::success`] for a completed session or +/// [`SessionResult::failed`] for a session that ended with an error. +/// +/// ``` +/// use loopctl::engine::loop_core::SessionResult; +/// use uuid::Uuid; +/// +/// let session_id = Uuid::new_v4(); +/// +/// let ok = SessionResult::success(session_id); +/// assert!(ok.success); +/// +/// let err = SessionResult::failed(session_id, "API rate limit exceeded"); +/// assert!(!err.success); +/// assert_eq!(err.error.unwrap(), "API rate limit exceeded"); +/// ``` +#[derive(Debug, Clone)] +pub struct SessionResult { + /// Matches [`LoopConfig::session_id`]. + pub session_id: Uuid, + /// Total turns executed. Compared against [`LoopConfig::max_turns`]. + pub total_turns: usize, + /// Sum of [`TurnResult::input_tokens`] across all turns. + pub input_tokens: u64, + /// Sum of [`TurnResult::output_tokens`] across all turns. + pub output_tokens: u64, + /// Wall-clock time from session start to finalization. + pub total_duration: Duration, + /// Total number of tool calls executed across all turns. + pub tool_calls: usize, + /// `true` on [`LoopState::Completed`], `false` on [`LoopState::Failed`]. + pub success: bool, + /// Last meaningful text response from the agent. `None` if no final message. + pub final_output: Option, + /// `Some` on failure with a human-readable error description. + pub error: Option, +} + +impl SessionResult { + /// Create a successful session result. + /// + /// Initializes all counters to zero and sets [`success`](SessionResult::success) + /// to `true`. The framework or production code should fill in the actual + /// counters before returning. + /// + /// # Example + /// + /// ``` + /// use loopctl::engine::loop_core::SessionResult; + /// use uuid::Uuid; + /// + /// let session_id = Uuid::new_v4(); + /// let result = SessionResult::success(session_id); + /// assert!(result.success); + /// assert!(result.error.is_none()); + /// ``` + #[must_use] + pub fn success(session_id: Uuid) -> Self { + Self { + session_id, + total_turns: 0, + input_tokens: 0, + output_tokens: 0, + total_duration: Duration::ZERO, + tool_calls: 0, + success: true, + final_output: None, + error: None, + } + } + + /// Create a failed session result. + /// + /// Sets [`success`](SessionResult::success) to `false` and records the + /// error message. All counters are initialized to zero. + /// + /// # Example + /// + /// ``` + /// use loopctl::engine::loop_core::SessionResult; + /// use uuid::Uuid; + /// + /// let session_id = Uuid::new_v4(); + /// let result = SessionResult::failed(session_id, "API rate limit exceeded"); + /// assert!(!result.success); + /// assert_eq!(result.error.unwrap(), "API rate limit exceeded"); + /// ``` + #[must_use] + pub fn failed(session_id: Uuid, error: impl Into) -> Self { + Self { + session_id, + total_turns: 0, + input_tokens: 0, + output_tokens: 0, + total_duration: Duration::ZERO, + tool_calls: 0, + success: false, + final_output: None, + error: Some(error.into()), + } + } + + /// Total tokens (input + output) for this session. + /// + /// Sums [`input_tokens`](SessionResult::input_tokens) + /// and [`output_tokens`](SessionResult::output_tokens). + #[must_use] + pub fn total_tokens(&self) -> u64 { + self.input_tokens.saturating_add(self.output_tokens) + } +} + +// ================================================== +// Loop trait +// ================================================== + +/// The core agent lifecycle trait. +/// +/// Implement this trait to create a new type of agent. The framework +/// provides shared infrastructure for context management, tool execution, +/// reflection, and observability, so implementations only need to define +/// the core processing logic. +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::engine::loop_core::Loop; +/// use loopctl::error::LoopError; +/// use loopctl::engine::loop_core::{LoopConfig, LoopState, SessionResult, TurnResult}; +/// +/// struct MyAgent { +/// state: MyState, +/// } +/// +/// impl Loop for MyAgent { +/// fn initialize<'a>(&'a mut self, config: &'a LoopConfig) -> Pin> + Send + 'a>> { +/// Box::pin(async { Ok(()) }) +/// } +/// fn process_turn<'a>(&'a mut self, input: &'a str) -> Pin> + Send + 'a>> { +/// Box::pin(async { Ok(TurnResult::completed("Done!")) }) +/// } +/// fn should_continue(&self) -> bool { +/// !self.state.is_complete +/// } +/// fn finalize<'a>(&'a mut self) -> Pin> + Send + 'a>> { +/// Box::pin(async { Ok(SessionResult::success(self.state.session_id)) }) +/// } +/// fn state(&self) -> LoopState { +/// LoopState::Idle +/// } +/// fn cancel(&self) {} +/// } +/// ``` +pub trait Loop: Send + Sync { + /// Initialize the agent with the given configuration. + /// + /// Called once before any turns are processed. Use this to set up + /// internal state, validate configuration, and prepare resources. + fn initialize<'a>( + &'a mut self, + config: &'a LoopConfig, + ) -> Pin> + Send + 'a>>; + + /// Process a single user message / turn. + /// + /// Main entry point for agent logic. It receives the user's + /// input and returns a [`TurnResult`] describing what happened. + fn process_turn<'a>( + &'a mut self, + input: &'a str, + ) -> Pin> + Send + 'a>>; + + /// Check whether the agent should continue processing turns. + /// + /// Called after each turn. Return `false` to end the session. + fn should_continue(&self) -> bool; + + /// Finalize the agent session and produce a summary. + /// + /// Called once after the last turn. Use this to clean up resources + /// and produce a final [`SessionResult`]. + fn finalize<'a>( + &'a mut self, + ) -> Pin> + Send + 'a>>; + + /// Get the current state of the agent. + /// + /// Used by the framework to drive the state machine and by observers + /// to report status. + fn state(&self) -> LoopState; + + /// Cancel the agent's current operation. + /// + /// Implementations must use thread-safe interior mutability (e.g. + /// [`AtomicBool`](std::sync::atomic::AtomicBool), `Mutex`) to + /// store the cancellation flag, since this method takes `&self`. The + /// flag should be set in a non-blocking fashion so that + /// [`process_turn`](Loop::process_turn) and + /// [`should_continue`](Loop::should_continue) can observe it + /// and return promptly across threads. + fn cancel(&self); +} diff --git a/src/core/error.rs b/src/error.rs similarity index 73% rename from src/core/error.rs rename to src/error.rs index ae451bd..e8090d2 100644 --- a/src/core/error.rs +++ b/src/error.rs @@ -1,16 +1,16 @@ //! Error types for the agent framework. //! -//! Every agent operation returns `Result`, and this +//! Every agent operation returns `Result`, and this //! module defines the single unified error enum that carries structured //! context for each failure mode. The variants are fine-grained enough //! for callers to match on and recover programmatically -//! (see [`AgentError::is_recoverable`]), while still producing +//! (see [`LoopError::is_recoverable`]), while still producing //! human-readable messages via the [`thiserror`] `#[error(...)]` //! attributes. //! //! # Provided Types //! -//! - **[`AgentError`]** — The sole error enum for the framework. Each +//! - **[`LoopError`]** — The sole error enum for the framework. Each //! variant wraps the relevant context (tool names, token counts, //! phase identifiers) so that diagnostics are precise without //! requiring callers to parse free-form strings. @@ -18,11 +18,11 @@ //! # Quick Start //! //! ``` -//! use loopctl::core::error::AgentError; +//! use loopctl::error::LoopError; //! -//! fn run_tool(tool: &str, available: &[&str]) -> Result { +//! fn run_tool(tool: &str, available: &[&str]) -> Result { //! if !available.contains(&tool) { -//! return Err(AgentError::tool_not_found(tool, available)); +//! return Err(LoopError::tool_not_found(tool, available)); //! } //! // invoke the tool //! Ok("done".into()) @@ -31,7 +31,7 @@ //! fn main() { //! match run_tool("search", &["read", "write"]) { //! Ok(_) => println!("succeeded"), -//! Err(AgentError::ToolNotFound { tool, .. }) => { +//! Err(LoopError::ToolNotFound { tool, .. }) => { //! eprintln!("unknown tool: {tool}"); //! } //! Err(e) => println!("other error: {e}"), @@ -41,80 +41,76 @@ /// Unified error type for the agent framework. /// -/// All agent operations return `Result`. Each variant +/// All agent operations return `Result`. Each variant /// carries structured context to aid debugging, logging, and -/// programmatic error recovery. Use [`AgentError::is_recoverable`] +/// programmatic error recovery. Use [`LoopError::is_recoverable`] /// to decide whether a retry is feasible, or match on specific /// variants when you need targeted handling (e.g. -/// [`AgentError::Cancelled`] to distinguish user-initiated +/// [`LoopError::Cancelled`] to distinguish user-initiated /// cancellation from fatal failures). /// /// # Variants /// -/// - [`ToolNotFound`](AgentError::ToolNotFound) — The requested tool +/// - [`ToolNotFound`](LoopError::ToolNotFound) — The requested tool /// is not registered in the tool registry. -/// - [`ToolExecution`](AgentError::ToolExecution) — A registered tool +/// - [`ToolExecution`](LoopError::ToolExecution) — A registered tool /// was invoked but returned an error during execution. -/// - [`InvalidInput`](AgentError::InvalidInput) — The caller provided +/// - [`InvalidInput`](LoopError::InvalidInput) — The caller provided /// malformed or semantically invalid input. -/// - [`Api`](AgentError::Api) — An upstream LLM provider or HTTP API +/// - [`Api`](LoopError::Api) — An upstream LLM provider or HTTP API /// returned an error response. -/// - [`MaxTurnsExceeded`](AgentError::MaxTurnsExceeded) — The agent +/// - [`MaxTurnsExceeded`](LoopError::MaxTurnsExceeded) — The agent /// hit its configured turn limit without completing. -/// - [`ContextExceeded`](AgentError::ContextExceeded) — Token usage +/// - [`ContextExceeded`](LoopError::ContextExceeded) — Token usage /// overflowed the model's context window and compaction could not /// recover. -/// - [`PhaseFailed`](AgentError::PhaseFailed) — A named pipeline +/// - [`PhaseFailed`](LoopError::PhaseFailed) — A named pipeline /// phase (e.g. "reflection", "compaction") failed. -/// - [`Memory`](AgentError::Memory) — A memory +/// - [`Memory`](LoopError::Memory) — A memory /// store/retrieve/consolidate operation failed. -/// - [`Reflection`](AgentError::Reflection) — The self-correction / +/// - [`Reflection`](LoopError::Reflection) — The self-correction / /// reflection cycle encountered an error. -/// - [`Cancelled`](AgentError::Cancelled) — The user or a shutdown +/// - [`Cancelled`](LoopError::Cancelled) — The user or a shutdown /// signal cancelled the session. -/// - [`LoopDetected`](AgentError::LoopDetected) — The agent was +/// - [`LoopDetected`](LoopError::LoopDetected) — The agent was /// caught repeating the same operation without making progress. -/// - [`ToolLimitReached`](AgentError::ToolLimitReached) — The +/// - [`ToolLimitReached`](LoopError::ToolLimitReached) — The /// session or turn exceeded its tool-call budget. -/// - [`StreamError`](AgentError::StreamError) — An error occurred +/// - [`StreamError`](LoopError::StreamError) — An error occurred /// while processing a streaming response from the LLM. -/// - [`Config`](AgentError::Config) — Configuration validation +/// - [`Config`](LoopError::Config) — Configuration validation /// failed (missing fields, invalid values). -/// - [`Internal`](AgentError::Internal) — A catch-all for unexpected +/// - [`Internal`](LoopError::Internal) — A catch-all for unexpected /// or infrastructure-level errors. #[derive(Debug, Clone, thiserror::Error)] #[non_exhaustive] -pub enum AgentError { +pub enum LoopError { /// A tool was not found in the registry. /// /// Returned when the agent (or a caller) requests a tool by name /// that has not been registered. The error carries both the /// requested name and a formatted list of available tools so the - /// caller can suggest alternatives. Construct conveniently with - /// [`tool_not_found`](AgentError::tool_not_found). + /// caller can suggest alternatives. Construct with + /// [`tool_not_found`](LoopError::tool_not_found). #[error("Tool not found: {tool}. Available: {available}")] ToolNotFound { - /// The name of the requested tool. + /// The name of the requested tool (not normalised or lowercased). tool: String, - /// Comma-separated list of available tool names. - /// - /// Generated by [`AgentError::tool_not_found`] — capped at - /// ten names with a trailing "... and N more" suffix when the - /// registry is large. + /// Available tool names, capped at ten with "… and N more" suffix. available: String, }, /// Tool execution failed. /// /// The tool was found and invoked, but the tool's internal logic - /// returned an error. This is *recoverable* — the agent may retry + /// returned an error. Recoverable — the agent may retry /// with different inputs or fall back to another tool. See - /// [`is_recoverable`](AgentError::is_recoverable). + /// [`is_recoverable`](LoopError::is_recoverable). #[error("Tool execution error: {tool}: {message}")] ToolExecution { - /// Name of the tool that failed. + /// Name of the tool that failed (matches the [`ToolRegistry`](crate::tool::ToolRegistry) key). tool: String, - /// Error message produced by the tool. + /// Description of the execution failure returned by the tool. message: String, }, @@ -124,13 +120,10 @@ pub enum AgentError { /// a malformed JSON parameter, an out-of-range integer, or a /// semantically meaningless prompt. This variant is **not** /// considered recoverable by - /// [`is_recoverable`](AgentError::is_recoverable) because + /// [`is_recoverable`](LoopError::is_recoverable) because /// retrying the same input will produce the same result. #[error("Invalid input: {0}")] - InvalidInput( - /// Description of the validation failure. - String, - ), + InvalidInput(String), /// An API call failed (e.g. LLM provider error). /// @@ -140,7 +133,7 @@ pub enum AgentError { /// retry with exponential back-off. #[error("API error: {0}")] Api( - /// The upstream error message or status description. + /// Upstream error message or status description. String, ), @@ -151,7 +144,7 @@ pub enum AgentError { /// investigate why the agent is not converging. #[error("Max turns exceeded: {max}")] MaxTurnsExceeded { - /// The configured maximum turn count. + /// The configured maximum turn count. See [`LoopConfig::max_turns`](crate::config::LoopConfig::max_turns). max: usize, }, @@ -164,9 +157,9 @@ pub enum AgentError { /// `limit` fields give precise token counts for diagnostics. #[error("Context window exceeded: used {used} of {limit} tokens")] ContextExceeded { - /// Tokens currently used in the conversation. + /// Number of tokens consumed when the limit was exceeded. used: u64, - /// Maximum tokens allowed by the model. + /// Maximum tokens allowed by the model. See [`LoopConfig::context_window`](crate::config::LoopConfig::context_window). limit: u64, }, @@ -179,9 +172,9 @@ pub enum AgentError { /// without parsing error messages. #[error("Phase '{phase}' failed: {message}")] PhaseFailed { - /// Name of the phase that failed. + /// Name of the phase that failed (e.g. `"pre_process"`, `"reflection"`). phase: String, - /// Error message from the failed phase. + /// Description of the phase failure. message: String, }, @@ -191,10 +184,7 @@ pub enum AgentError { /// connection failure, a serialization error, or a capacity limit /// reached in the backing store. #[error("Memory error: {0}")] - Memory( - /// Description of the memory failure. - String, - ), + Memory(String), /// Reflection / self-correction cycle failed. /// @@ -203,15 +193,12 @@ pub enum AgentError { /// variant is *recoverable* — the framework may retry the /// reflection or fall back to a simpler strategy. #[error("Reflection error: {0}")] - Reflection( - /// Description of the reflection failure. - String, - ), + Reflection(String), /// The agent was cancelled by the user or a shutdown signal. /// - /// This is a *clean* termination, not a failure. Check for this - /// variant with [`AgentError::is_cancelled`] to avoid logging it + /// A *clean* termination, not a failure. Check for this + /// variant with [`LoopError::is_cancelled`] to avoid logging it /// as an error. The agent may have partial results available in /// its state. #[error("Agent cancelled")] @@ -226,7 +213,7 @@ pub enum AgentError { /// `message` field describes the detected pattern. #[error("Loop detected: {message}")] LoopDetected { - /// Description of the detected loop. + /// Description of the detected loop (e.g. `"Tool Read called 5 times with identical arguments"`). message: String, }, @@ -239,7 +226,7 @@ pub enum AgentError { /// to increase the limit or accept the partial result. #[error("Tool limit reached: {message}")] ToolLimitReached { - /// Description of the limit that was reached. + /// Description of the limit reached (e.g. `"Session tool limit of 100 reached"`). message: String, }, @@ -251,7 +238,7 @@ pub enum AgentError { /// whether to retry from the last complete message. #[error("Stream error: {0}")] StreamError( - /// Description of the streaming failure. + /// Description of the streaming failure (network drop, malformed SSE, timeout). String, ), @@ -264,7 +251,7 @@ pub enum AgentError { /// configuration and retry. #[error("Configuration error: {0}")] Config( - /// Description of the configuration problem. + /// Description of the configuration problem (e.g. `"max_turns must be > 0"`). String, ), @@ -275,13 +262,10 @@ pub enum AgentError { /// this only for truly unexpected conditions (e.g. a poisoned /// mutex, an allocation failure). #[error("{0}")] - Internal( - /// The internal error description. - String, - ), + Internal(String), } -impl AgentError { +impl LoopError { /// Create a tool-not-found error with a list of available tools. /// /// Called by the tool registry when a requested tool name does not @@ -292,10 +276,10 @@ impl AgentError { /// # Example /// /// ``` - /// use loopctl::core::error::AgentError; + /// use loopctl::error::LoopError; /// - /// let err = AgentError::tool_not_found("search", &["read", "write", "delete"]); - /// assert!(matches!(err, AgentError::ToolNotFound { .. })); + /// let err = LoopError::tool_not_found("search", &["read", "write", "delete"]); + /// assert!(matches!(err, LoopError::ToolNotFound { .. })); /// ``` pub fn tool_not_found(tool: impl Into, available: &[&str]) -> Self { let tool = tool.into(); @@ -326,14 +310,14 @@ impl AgentError { /// Returns `true` for variants where a retry with the same or /// modified inputs has a reasonable chance of success: /// - /// - [`ToolExecution`](AgentError::ToolExecution) — the tool may + /// - [`ToolExecution`](LoopError::ToolExecution) — the tool may /// succeed on a second attempt (transient failure, rate limit, /// etc.). - /// - [`Api`](AgentError::Api) — the upstream provider may recover + /// - [`Api`](LoopError::Api) — the upstream provider may recover /// (network blip, temporary overload). - /// - [`ContextExceeded`](AgentError::ContextExceeded) — + /// - [`ContextExceeded`](LoopError::ContextExceeded) — /// compaction may free enough tokens for a retry. - /// - [`Reflection`](AgentError::Reflection) — a second + /// - [`Reflection`](LoopError::Reflection) — a second /// reflection pass may produce a valid correction. /// /// Returns `false` for all other variants (e.g. invalid input, @@ -352,7 +336,7 @@ impl AgentError { /// Check whether the agent was explicitly cancelled. /// - /// Returns `true` only for the [`Cancelled`](AgentError::Cancelled) + /// Returns `true` only for the [`Cancelled`](LoopError::Cancelled) /// variant. Use this to distinguish user-initiated cancellation /// from genuine failures so you can log it as `info!` rather than /// `error!`. @@ -360,9 +344,9 @@ impl AgentError { /// # Example /// /// ``` - /// use loopctl::core::error::AgentError; + /// use loopctl::error::LoopError; /// - /// let err = AgentError::Cancelled; + /// let err = LoopError::Cancelled; /// assert!(err.is_cancelled()); /// ``` #[must_use] @@ -376,7 +360,7 @@ mod tests { use super::*; #[test] fn tool_not_found_empty_available_says_none_registered() { - let err = AgentError::tool_not_found("my_tool", &[]); + let err = LoopError::tool_not_found("my_tool", &[]); assert_eq!( err.to_string(), "Tool not found: my_tool. Available: none registered" @@ -394,7 +378,7 @@ mod tests { TOOLS[i - 1] }) .collect(); - let err = AgentError::tool_not_found("missing", &names); + let err = LoopError::tool_not_found("missing", &names); assert_eq!( err.to_string(), "Tool not found: missing. Available: tool_1, tool_2, tool_3, tool_4, tool_5, tool_6, tool_7, tool_8, tool_9, tool_10" @@ -407,7 +391,7 @@ mod tests { "tool_1", "tool_2", "tool_3", "tool_4", "tool_5", "tool_6", "tool_7", "tool_8", "tool_9", "tool_10", "tool_11", "tool_12", "tool_13", ]; - let err = AgentError::tool_not_found("missing", &names); + let err = LoopError::tool_not_found("missing", &names); assert_eq!( err.to_string(), "Tool not found: missing. Available: tool_1, tool_2, tool_3, tool_4, tool_5, tool_6, tool_7, tool_8, tool_9, tool_10... (and 3 more)" @@ -420,7 +404,7 @@ mod tests { "tool_1", "tool_2", "tool_3", "tool_4", "tool_5", "tool_6", "tool_7", "tool_8", "tool_9", "tool_10", "tool_11", ]; - let err = AgentError::tool_not_found("missing", &names); + let err = LoopError::tool_not_found("missing", &names); assert_eq!( err.to_string(), "Tool not found: missing. Available: tool_1, tool_2, tool_3, tool_4, tool_5, tool_6, tool_7, tool_8, tool_9, tool_10... (and 1 more)" @@ -429,7 +413,7 @@ mod tests { #[test] fn is_recoverable_true_for_tool_execution() { - let err = AgentError::ToolExecution { + let err = LoopError::ToolExecution { tool: "cat".into(), message: "something went wrong".into(), }; @@ -438,13 +422,13 @@ mod tests { #[test] fn is_recoverable_true_for_api() { - let err = AgentError::Api("rate limited".into()); + let err = LoopError::Api("rate limited".into()); assert!(err.is_recoverable()); } #[test] fn is_recoverable_true_for_context_exceeded() { - let err = AgentError::ContextExceeded { + let err = LoopError::ContextExceeded { used: 200_000, limit: 128_000, }; @@ -453,19 +437,19 @@ mod tests { #[test] fn is_recoverable_true_for_reflection() { - let err = AgentError::Reflection("need to rethink".into()); + let err = LoopError::Reflection("need to rethink".into()); assert!(err.is_recoverable()); } #[test] fn is_recoverable_false_for_tool_not_found() { - let err = AgentError::tool_not_found("nope", &["a", "b"]); + let err = LoopError::tool_not_found("nope", &["a", "b"]); assert!(!err.is_recoverable()); } #[test] fn is_recoverable_false_for_cancelled() { - let err = AgentError::Cancelled; + let err = LoopError::Cancelled; assert!(!err.is_recoverable()); } } diff --git a/src/loop_control/fallback.rs b/src/fallback.rs similarity index 78% rename from src/loop_control/fallback.rs rename to src/fallback.rs index b2ca380..ff0cc78 100644 --- a/src/loop_control/fallback.rs +++ b/src/fallback.rs @@ -15,30 +15,6 @@ //! primary with a few trial requests; if they succeed, the circuit closes and //! normal operation resumes. //! -//! # State machine -//! -//! ```text -//! ┌──────────────────────────────────────────────────────────────────────┐ -//! │ Primary (closed) │ -//! │ • Requests go to the primary model. │ -//! │ • Consecutive failures are counted. │ -//! │ │ │ -//! │ │ failures ≥ trip_threshold │ -//! │ ▼ │ -//! │ Fallback (open) │ -//! │ • Requests go to the fallback model. │ -//! │ • Waits at least `recovery_timeout` before probing primary. │ -//! │ │ │ -//! │ │ `should_try_resume_primary` → true │ -//! │ │ + `transition_to_recovering` │ -//! │ ▼ │ -//! │ Recovering (half-open) │ -//! │ • Trial requests go to primary model. │ -//! │ • successes ≥ recovery_successes_needed → Primary │ -//! │ • 1 failure → back to Fallback │ -//! └──────────────────────────────────────────────────────────────────────┘ -//! ``` -//! //! # Provided types //! //! - **[`FallbackState`]** — The three circuit-breaker states (`Primary`, `Fallback`, `Recovering`). @@ -48,7 +24,7 @@ //! # Quick Start //! //! ```rust -//! use loopctl::loop_control::fallback::{FallbackManager, FallbackConfig}; +//! use loopctl::fallback::{FallbackManager, FallbackConfig}; //! //! // Create a manager with a trip threshold of 3 failures //! let mgr = FallbackManager::new(3, 2); @@ -73,9 +49,7 @@ use tracing::{debug, info, warn}; /// Circuit breaker state for LLM model fallback. /// -/// Models the three classic circuit-breaker phases. The state is stored -/// internally as an [`AtomicU8`] for lock-free reads across threads, and -/// converted to/from `u8` via [`FallbackState::from`]. +/// Models the three classic circuit-breaker phases. /// /// # Transitions /// @@ -89,7 +63,7 @@ use tracing::{debug, info, warn}; /// # Example /// /// ```rust -/// use loopctl::loop_control::fallback::FallbackState; +/// use loopctl::fallback::FallbackState; /// /// let state = FallbackState::Primary; /// assert_eq!(state as u8, 0); @@ -97,47 +71,36 @@ use tracing::{debug, info, warn}; /// ``` #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum FallbackState { - /// Circuit is **closed** — using the primary model. - /// - /// All requests route to the primary LLM. Consecutive failures are - /// counted; once they reach the configured threshold the circuit - /// trips open and transitions to [`Fallback`](FallbackState::Fallback). + /// Operating on the primary model — no failures have tripped the + /// circuit breaker yet. Primary = 0, - /// Circuit is **open** — using the fallback model. - /// - /// Requests are routed away from the degraded primary. The manager - /// stays in this state until [`FallbackManager::should_try_resume_primary`] - /// returns `true`, at which point it transitions to - /// [`Recovering`](FallbackState::Recovering) to probe the primary. + /// A fallback model is active — the primary model failed and the + /// breaker tripped. Subsequent failures on the fallback model are + /// tracked separately. Fallback = 1, - /// Circuit is **half-open** — probing whether the primary has recovered. - /// - /// A small number of trial requests are sent to the primary model. - /// If enough consecutive successes are observed (see - /// [`FallbackConfig::recovery_successes_needed`]), the circuit closes - /// back to [`Primary`](FallbackState::Primary). A single failure - /// immediately reopens the circuit to [`Fallback`](FallbackState::Fallback). + /// Between models — the primary failed, no fallback has been + /// selected yet, or a fallback also failed and the manager is + /// searching for another candidate. Recovering = 2, } /// Converts a raw `u8` back into a [`FallbackState`]. /// -/// This is the inverse of `state as u8` and is used to decode the value -/// stored in [`FallbackManager`]'s internal `AtomicU8`. -/// Unknown values default to [`FallbackState::Primary`] for safety. +/// Inverse of `state as u8`. Unknown values default to +/// [`FallbackState::Primary`] for safety. /// /// # Safety /// /// The conversion is infallible — any out-of-range `u8` maps to -/// [`FallbackState::Primary`] so that a corrupted atomic value cannot +/// [`FallbackState::Primary`] so that a corrupted value cannot /// cause a panic. /// /// # Example /// /// ```rust -/// use loopctl::loop_control::fallback::FallbackState; +/// use loopctl::fallback::FallbackState; /// /// assert_eq!(FallbackState::from(0u8), FallbackState::Primary); /// assert_eq!(FallbackState::from(1u8), FallbackState::Fallback); @@ -165,16 +128,14 @@ impl From for FallbackState { /// # Example /// /// ```rust -/// use loopctl::loop_control::fallback::AttemptRecord; +/// use loopctl::fallback::AttemptRecord; /// /// let record = AttemptRecord::new("rate_limit"); /// assert_eq!(record.reason(), Some("rate_limit")); /// ``` #[derive(Debug, Clone)] pub struct AttemptRecord { - /// When this failure was recorded. failed_at: Instant, - /// Optional reason for the failure (e.g. `"timeout"`). reason: Option, } @@ -184,7 +145,7 @@ impl AttemptRecord { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::AttemptRecord; + /// use loopctl::fallback::AttemptRecord; /// let record = AttemptRecord::new("rate_limit"); /// assert_eq!(record.reason(), Some("rate_limit")); /// ``` @@ -201,7 +162,7 @@ impl AttemptRecord { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::AttemptRecord; + /// use loopctl::fallback::AttemptRecord; /// let record = AttemptRecord::anonymous(); /// assert!(record.reason().is_none()); /// ``` @@ -213,7 +174,7 @@ impl AttemptRecord { } } - /// Create a new attempt record with an optional reason. + /// Set the reason for this failure record. /// /// Pass `None` for an anonymous record, or `Some("reason")` for /// a labelled one. @@ -221,16 +182,14 @@ impl AttemptRecord { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::AttemptRecord; - /// let record = AttemptRecord::with_reason(Some("timeout".to_string())); + /// use loopctl::fallback::AttemptRecord; + /// let record = AttemptRecord::new("timeout").with_reason(Some("timeout".to_string())); /// assert_eq!(record.reason(), Some("timeout")); /// ``` #[must_use] - pub fn with_reason(reason: Option) -> Self { - Self { - failed_at: Instant::now(), - reason, - } + pub fn with_reason(mut self, reason: Option) -> Self { + self.reason = reason; + self } /// When this failure was recorded. @@ -238,7 +197,7 @@ impl AttemptRecord { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::AttemptRecord; + /// use loopctl::fallback::AttemptRecord; /// let record = AttemptRecord::new("timeout"); /// // failed_at is close to now /// assert!(record.failed_at().elapsed().as_secs() < 1); @@ -253,7 +212,7 @@ impl AttemptRecord { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::AttemptRecord; + /// use loopctl::fallback::AttemptRecord; /// let record = AttemptRecord::new("rate_limit"); /// assert_eq!(record.reason(), Some("rate_limit")); /// ``` @@ -274,7 +233,7 @@ impl AttemptRecord { /// # Example /// /// ```rust -/// use loopctl::loop_control::fallback::FallbackEntry; +/// use loopctl::fallback::FallbackEntry; /// /// let mut entry = FallbackEntry::new("llm-70b"); /// assert_eq!(entry.name(), "llm-70b"); @@ -287,18 +246,13 @@ impl AttemptRecord { /// ``` #[derive(Debug, Clone)] pub struct FallbackEntry { - /// Model identifier (e.g. `"llm-70b"`). + /// Model identifier (e.g. `"llm-70b"`). Must match the API client's routing identifier. name: String, - /// Whether this model is available for use. - /// - /// Set to `false` to take a model out of rotation independently of - /// failure tracking (e.g. API key revoked, model decommissioned). - /// A model that is not available is always skipped by - /// [`FallbackManager::fallback_model`]. + /// Set to `false` to take a model out of rotation independently of failure tracking. available: bool, - /// Ordered list of recorded failure attempts. + /// Recorded failure attempts for this model. attempts: Vec, - /// How many recorded failures before this entry is considered failed. + /// When `attempts.len()` reaches this threshold the model is taken out of rotation. max_fail_count: usize, } @@ -312,7 +266,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let entry = FallbackEntry::new("llm-70b"); /// assert_eq!(entry.name(), "llm-70b"); /// assert!(!entry.failed()); @@ -328,7 +282,7 @@ impl FallbackEntry { } } - /// Create a new entry with a custom `max_fail_count`. + /// Set a custom `max_fail_count` threshold. /// /// The model is only considered failed after this many attempts /// have been recorded via [`record_attempt`](Self::record_attempt). @@ -336,8 +290,8 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; - /// let mut entry = FallbackEntry::with_max_fail_count("llm-70b", 3); + /// use loopctl::fallback::FallbackEntry; + /// let mut entry = FallbackEntry::new("llm-70b").with_max_fail_count(3); /// assert_eq!(entry.max_fail_count(), 3); /// /// entry.record_attempt("timeout"); @@ -348,25 +302,20 @@ impl FallbackEntry { /// assert!(entry.failed()); // 3 of 3 /// ``` #[must_use] - pub fn with_max_fail_count(name: impl Into, max_fail_count: usize) -> Self { - Self { - name: name.into(), - available: true, - attempts: Vec::new(), - max_fail_count: max_fail_count.max(1), - } + pub fn with_max_fail_count(mut self, max_fail_count: usize) -> Self { + self.max_fail_count = max_fail_count.max(1); + self } /// Create a new entry already marked as failed. /// - /// Internally records one anonymous attempt so that - /// [`failed()`](Self::failed) returns `true` immediately. /// Useful when initializing from a known-degraded model. + /// [`failed()`](Self::failed) returns `true` immediately. /// /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let entry = FallbackEntry::new_failed("llm-70b"); /// assert!(entry.failed()); /// ``` @@ -385,7 +334,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let entry = FallbackEntry::new("llm-70b"); /// assert_eq!(entry.name(), "llm-70b"); /// ``` @@ -403,7 +352,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let mut entry = FallbackEntry::new("llm-70b"); /// assert!(!entry.failed()); /// entry.record_attempt("timeout"); @@ -420,8 +369,8 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; - /// let entry = FallbackEntry::with_max_fail_count("llm-70b", 3); + /// use loopctl::fallback::FallbackEntry; + /// let entry = FallbackEntry::new("llm-70b").with_max_fail_count(3); /// assert_eq!(entry.max_fail_count(), 3); /// ``` #[must_use] @@ -439,7 +388,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let mut entry = FallbackEntry::new("llm-70b"); /// assert!(entry.available()); /// entry.set_available(false); @@ -459,7 +408,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let mut entry = FallbackEntry::new("llm-70b"); /// entry.set_available(false); /// assert!(entry.failed()); @@ -475,7 +424,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let mut entry = FallbackEntry::new("llm-70b"); /// assert_eq!(entry.attempt_count(), 0); /// entry.record_attempt("timeout"); @@ -494,7 +443,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let mut entry = FallbackEntry::new("llm-70b"); /// entry.record_attempt("timeout"); /// entry.record_attempt("rate_limit"); @@ -509,15 +458,14 @@ impl FallbackEntry { /// Record a new failure attempt with an optional reason. /// - /// Appends to the internal attempt list. If this causes - /// [`attempt_count`](Self::attempt_count) to reach + /// If this causes [`attempt_count`](Self::attempt_count) to reach /// [`max_fail_count`](Self::max_fail_count), subsequent calls to /// [`failed()`](Self::failed) will return `true`. /// /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let mut entry = FallbackEntry::new("llm-70b"); /// entry.record_attempt("timeout"); /// entry.record_attempt("timeout"); // exceeds max_fail_count (default = 2) @@ -532,7 +480,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let mut entry = FallbackEntry::new("llm-70b"); /// entry.record_attempt_anonymous(); /// entry.record_attempt_anonymous(); // exceeds max_fail_count (default = 2) @@ -548,7 +496,7 @@ impl FallbackEntry { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackEntry; + /// use loopctl::fallback::FallbackEntry; /// let mut entry = FallbackEntry::new("llm-70b"); /// entry.record_attempt("timeout"); /// entry.record_attempt("timeout"); // exceeds max_fail_count (default = 2) @@ -579,7 +527,7 @@ impl FallbackEntry { /// # Example /// /// ```rust -/// use loopctl::loop_control::fallback::FallbackConfig; +/// use loopctl::fallback::FallbackConfig; /// use std::time::Duration; /// /// let config = FallbackConfig { @@ -591,36 +539,13 @@ impl FallbackEntry { /// ``` #[derive(Debug, Clone)] pub struct FallbackConfig { - /// Number of consecutive API failures before the circuit trips open. - /// - /// Once this many failures have been recorded in a row (without an - /// intervening success), the manager transitions from [`FallbackState::Primary`] - /// to [`FallbackState::Fallback`]. Defaults to `3`. + /// Consecutive API failures required to trip the circuit open. Defaults to `3`. pub trip_threshold: usize, - - /// Minimum time to remain in [`FallbackState::Fallback`] before probing - /// the primary model again. - /// - /// Checked by [`FallbackManager::should_try_resume_primary`]. Defaults - /// to 60 seconds. Set this higher to give a degraded API more time to - /// recover before retrying. + /// Minimum time in fallback before probing the primary model again. Defaults to 60 s. pub recovery_timeout: Duration, - - /// Number of consecutive successful requests on the primary model during - /// [`FallbackState::Recovering`] before the circuit fully closes. - /// - /// Each success is recorded via [`FallbackManager::record_model_success`]; - /// once the count reaches this threshold, the manager transitions back - /// to [`FallbackState::Primary`]. Defaults to `2`. + /// Consecutive successes during recovering before the circuit fully closes. Defaults to `2`. pub recovery_successes_needed: usize, - - /// Per-model failure threshold: how many recorded failures before a - /// fallback model is skipped. - /// - /// Each [`FallbackEntry`] in the chain tracks its own attempt history. - /// When [`FallbackEntry::attempt_count`] reaches this value, the entry - /// is considered [`failed`](FallbackEntry::failed) and is skipped by - /// [`FallbackManager::fallback_model`]. Defaults to `2`. + /// Per-model failure threshold before a fallback model is skipped. Defaults to `2`. pub max_fail_count: usize, } @@ -632,7 +557,7 @@ pub struct FallbackConfig { /// # Example /// /// ```rust -/// use loopctl::loop_control::fallback::FallbackConfig; +/// use loopctl::fallback::FallbackConfig; /// /// let config = FallbackConfig::default(); /// assert_eq!(config.trip_threshold, 3); @@ -651,18 +576,13 @@ impl Default for FallbackConfig { /// Manages circuit breaker state and model fallback transitions. /// -/// Contains all state related to detecting failures, transitioning -/// between models, and recovering back to the primary model. -/// -/// This is the **unified** fallback manager used by both the agent crate -/// and the framework. It combines the richer production API (state -/// transitions, `should_try_resume_primary`, `record_api_failure`) with -/// the framework's config-driven approach. +/// Tracks failures, transitions between models, and recovers back to +/// the primary model. Supports both config-driven construction via +/// [`FallbackManager::with_config`] and direct threshold control via +/// [`FallbackManager::new`]. /// /// # Thread safety /// -/// All mutable state is stored in atomic types ([`AtomicUsize`], -/// [`AtomicBool`], [`AtomicU8`]) or inside a [`Mutex`], so /// `&FallbackManager` is `Send + Sync` and can be freely shared across /// threads (e.g. via `Arc`). No `&mut self` is needed /// for any public method. @@ -677,11 +597,11 @@ impl Default for FallbackConfig { /// # Example /// /// ```rust -/// use loopctl::loop_control::fallback::{FallbackManager, FallbackConfig}; +/// use loopctl::fallback::{FallbackManager, FallbackConfig}; /// use std::sync::Arc; /// use std::time::Duration; /// -/// let mgr = Arc::new(FallbackManager::with_config(&FallbackConfig::default())); +/// let mgr = Arc::new(FallbackManager::new(5, 3).with_config(&FallbackConfig::default())); /// /// // Simulate failures /// assert!(!mgr.record_model_failure()); // 1 @@ -698,116 +618,27 @@ impl Default for FallbackConfig { /// } /// ``` pub struct FallbackManager { - // ================================================== - // Config - // ================================================== - /// Consecutive failure threshold before switching to fallback. - /// - /// When [`consecutive_failures`](Self::consecutive_failures) reaches - /// this value, the circuit trips from [`FallbackState::Primary`] to - /// [`FallbackState::Fallback`]. Set at construction time via - /// [`FallbackManager::new`] or [`FallbackManager::with_config`]. + /// Failures before switching to fallback. fallback_threshold: usize, - /// Successes needed on primary before resuming. - /// - /// During [`FallbackState::Recovering`], this many consecutive - /// successes must be recorded via [`record_model_success`](Self::record_model_success) - /// before the circuit closes back to [`FallbackState::Primary`]. primary_resume_threshold: usize, - - /// Per-model max failure count, passed to new [`FallbackEntry`] instances. - /// - /// When [`add_fallback_model`](Self::add_fallback_model) or - /// [`set_fallback_models`](Self::set_fallback_models) create entries, - /// they use [`FallbackEntry::with_max_fail_count`] with this value. - /// Defaults to `2`. Override via - /// [`FallbackConfig::max_fail_count`]. + /// Per-model max failure count for new [`FallbackEntry`] instances. default_max_fail_count: usize, - - // ================================================== - // Atomic state - // ================================================== /// Consecutive API failure counter. - /// - /// Incremented by [`record_model_failure`](Self::record_model_failure) - /// or [`record_api_failure`](Self::record_api_failure). Reset to `0` - /// on success ([`record_model_success`](Self::record_model_success)) - /// or by calling [`reset`](Self::reset). Starts at `0`. consecutive_failures: AtomicUsize, - /// Whether fallback has been activated (sticky flag). - /// - /// Set to `true` once the circuit has tripped at least once. Prevents - /// [`record_api_failure`](Self::record_api_failure) from re-tripping - /// the circuit on every subsequent failure. Cleared by - /// [`transition_to_primary`](Self::transition_to_primary) or - /// [`reset`](Self::reset). fallback_activated: AtomicBool, - /// Circuit breaker state (0=Primary, 1=Fallback, 2=Recovering). - /// - /// Stored as a [`FallbackState`] discriminant for lock-free reads. - /// Use [`state()`](Self::state) to access the typed value. fallback_state: AtomicU8, - /// Consecutive successes on primary during recovery. - /// - /// Incremented in [`FallbackState::Recovering`] state each time - /// [`record_model_success`](Self::record_model_success) is called. - /// Once it reaches [`primary_resume_threshold`](Self::primary_resume_threshold), - /// the circuit closes to [`FallbackState::Primary`]. Reset to `0` - /// on any state transition. primary_success_count: AtomicUsize, - - // ================================================== - // Mutex state - // ================================================== /// Original model name (before fallback). - /// - /// Set by [`for_model`](Self::for_model), - /// [`new_with_fallback`](Self::new_with_fallback), or - /// [`set_original_model`](Self::set_original_model). Retrieved by - /// [`original_model`](Self::original_model) and used by - /// [`active_model`](Self::active_model) to decide which model to use. - /// `None` until explicitly set. original_model: Mutex>, - - /// Ordered list of fallback models with their failure status. - /// - /// Set by [`set_fallback_models`](Self::set_fallback_models) or built up - /// with [`add_fallback_model`](Self::add_fallback_model). Each entry - /// tracks whether that model has been marked as failed. The - /// [`active_model`](Self::active_model) method skips failed entries and - /// returns the first non-failed model. Empty until explicitly configured. - /// - /// The models are ordered by priority: index `0` is the first fallback - /// tried, index `1` is the second, and so on. + /// Ordered fallback models with failure status. fallback_models: Mutex>, - - /// Currently active fallback model name (cached). - /// - /// Rather than scanning the entire fallback chain on every call to - /// [`fallback_model`](Self::fallback_model) or - /// [`active_model`](Self::active_model), the manager stores the name - /// of the first non-failed model here. It is recomputed whenever the - /// chain changes (add/remove/set) or when a model's failure status - /// changes ([`mark_fallback_failed`], [`clear_fallback_failed`], - /// [`set_fallback_available`]). - /// - /// `None` when no fallback model is configured or when all fallbacks - /// have failed. + /// Cached first non-failed fallback model name. active_fallback: Mutex>, - /// Time when fallback was activated. - /// - /// Set to `Some(Instant::now())` when - /// [`transition_to_fallback`](Self::transition_to_fallback) is called. - /// Cleared to `None` when returning to primary via - /// [`transition_to_primary`](Self::transition_to_primary) or - /// [`reset`](Self::reset). Checked by - /// [`should_try_resume_primary`](Self::should_try_resume_primary) - /// to enforce the cooldown period. fallback_switched_at: Mutex>, } @@ -830,7 +661,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(5, 3); /// // Trip after 5 failures, resume after 3 consecutive successes @@ -852,17 +683,15 @@ impl FallbackManager { } } - /// Create a new manager from a [`FallbackConfig`]. + /// Apply configuration from a [`FallbackConfig`] struct. /// - /// This is the preferred production constructor. It extracts - /// [`FallbackConfig::trip_threshold`] and - /// [`FallbackConfig::recovery_successes_needed`] from the config - /// and delegates to [`FallbackManager::new`]. + /// Sets the failure threshold, recovery parameters, and per-model + /// max fail count from the config. /// /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackConfig}; + /// use loopctl::fallback::{FallbackManager, FallbackConfig}; /// use std::time::Duration; /// /// let config = FallbackConfig { @@ -871,22 +700,23 @@ impl FallbackManager { /// recovery_successes_needed: 3, /// max_fail_count: 2, /// }; - /// let mgr = FallbackManager::with_config(&config); + /// let mgr = FallbackManager::new(5, 3).with_config(&config); /// ``` #[must_use] - pub fn with_config(config: &FallbackConfig) -> Self { - let mut mgr = Self::new(config.trip_threshold, config.recovery_successes_needed); - mgr.default_max_fail_count = config.max_fail_count; - mgr + pub fn with_config(mut self, config: &FallbackConfig) -> Self { + self.fallback_threshold = config.trip_threshold; + self.primary_resume_threshold = config.recovery_successes_needed; + self.default_max_fail_count = config.max_fail_count; + self } - /// Create with fallback already activated (e.g. for cloned agent loops). + /// Create with fallback already activated. /// - /// Useful when spawning a new agent loop that should start in the - /// [`FallbackState::Fallback`] state — for instance when a parent - /// agent has already determined the primary model is degraded. - /// The `original_model` is stored so the manager can attempt recovery - /// later via [`should_try_resume_primary`](Self::should_try_resume_primary). + /// Useful when a new manager should start in the + /// [`FallbackState::Fallback`] state — for instance when + /// the primary model is already known to be degraded. + /// The `original_model` is stored for later recovery via + /// [`should_try_resume_primary`](Self::should_try_resume_primary). /// /// # Parameters /// @@ -902,7 +732,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new_with_fallback("llm-70b".into(), 3); /// assert!(mgr.is_using_fallback()); @@ -926,7 +756,7 @@ impl FallbackManager { /// Create a new manager with a primary model name. /// - /// This is the framework-style constructor. The model name is stored as + /// The model name is stored as /// [`original_model`](Self::original_model) for later retrieval via /// [`active_model`](Self::active_model). Uses default thresholds /// (trip after 3 failures, resume after 2 successes). @@ -934,7 +764,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::for_model("llm-70b"); /// assert_eq!(mgr.original_model(), Some("llm-70b".to_string())); @@ -948,48 +778,16 @@ impl FallbackManager { mgr } - // ================================================== - // Private helpers - // ================================================== - - /// Recompute the cached [`active_fallback`](Self::active_fallback) from - /// the current fallback chain. - /// - /// Scans [`fallback_models`] for the first non-[`failed`](FallbackEntry::failed) - /// entry and stores its name in [`active_fallback`]. If no non-failed entry - /// exists, stores `None`. - /// - /// **Must** be called after every mutation to the fallback chain (add, - /// remove, set, mark failed, clear failed, set available) to keep the - /// cache consistent. - /// - /// [`fallback_models`]: Self::fallback_models - /// [`active_fallback`]: Self::active_fallback - fn recompute_active_fallback(&self) { - let active = self - .fallback_models - .lock() - .ok() - .and_then(|m| m.iter().find(|e| !e.failed()).map(|e| e.name.clone())); - if let Ok(mut cached) = self.active_fallback.lock() { - *cached = active; - } - } - // ================================================== // Accessors // ================================================== /// Get the current circuit breaker state. /// - /// Performs a lock-free atomic load and converts the raw `u8` into a - /// [`FallbackState`]. This is safe to call from any thread at any - /// time. - /// /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState}; /// /// let mgr = FallbackManager::new(3, 2); /// assert_eq!(mgr.state(), FallbackState::Primary); @@ -1009,7 +807,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState}; /// let mgr = FallbackManager::new(3, 2); /// assert!(!mgr.is_using_fallback()); /// ``` @@ -1028,7 +826,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// let mgr = FallbackManager::new(3, 2); /// assert!(!mgr.is_fallback_active()); /// ``` @@ -1038,14 +836,13 @@ impl FallbackManager { /// Get the number of consecutive failures. /// - /// Returns the current value of the atomic failure counter. This is - /// the count since the last success (when in + /// The count since the last success (when in /// [`FallbackState::Primary`]) or since the circuit tripped. /// /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// let mgr = FallbackManager::new(3, 2); /// assert_eq!(mgr.consecutive_failures(), 0); /// mgr.record_model_failure(); @@ -1066,7 +863,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// let mgr = FallbackManager::for_model("llm-70b"); /// assert_eq!(mgr.original_model(), Some("llm-70b".to_string())); /// ``` @@ -1076,14 +873,13 @@ impl FallbackManager { /// Set the original model name. /// - /// Overwrites the stored primary model name. Called by the framework - /// when the model is first resolved from configuration, or when the - /// user changes the model mid-session. + /// Overwrites the stored primary model name. Useful when the model + /// is resolved from configuration or changed mid-session. /// /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// let mgr = FallbackManager::new(3, 2); /// assert_eq!(mgr.original_model(), None); /// mgr.set_original_model("llm-70b".into()); @@ -1107,7 +903,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// use std::time::Duration; /// /// let mgr = FallbackManager::new(3, 2); @@ -1138,7 +934,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// let mgr = FallbackManager::for_model("llm-70b"); /// assert_eq!(mgr.active_model(), Some("llm-70b".to_string())); /// ``` @@ -1164,7 +960,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState}; /// /// let mgr = FallbackManager::for_model("llm-1"); /// mgr.add_fallback_model("llm-2"); @@ -1185,10 +981,7 @@ impl FallbackManager { pub fn set_fallback_model(&self, model: impl Into) { if let Ok(mut m) = self.fallback_models.lock() { m.clear(); - m.push(FallbackEntry::with_max_fail_count( - model, - self.default_max_fail_count, - )); + m.push(FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count)); } self.recompute_active_fallback(); } @@ -1202,7 +995,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// assert!(mgr.fallback_model().is_none()); @@ -1223,7 +1016,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1254,7 +1047,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1265,10 +1058,7 @@ impl FallbackManager { /// ``` pub fn add_fallback_model(&self, model: impl Into) { if let Ok(mut m) = self.fallback_models.lock() { - m.push(FallbackEntry::with_max_fail_count( - model, - self.default_max_fail_count, - )); + m.push(FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count)); } self.recompute_active_fallback(); } @@ -1286,7 +1076,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1298,7 +1088,7 @@ impl FallbackManager { /// ``` pub fn insert_fallback_model(&self, index: usize, model: impl Into) { if let Ok(mut m) = self.fallback_models.lock() { - let entry = FallbackEntry::with_max_fail_count(model, self.default_max_fail_count); + let entry = FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count); if index >= m.len() { m.push(entry); } else { @@ -1317,7 +1107,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1358,7 +1148,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.set_fallback_models(vec!["llm-70b".into(), "llm-120b".into(), "llm-32b".into()]); @@ -1371,7 +1161,7 @@ impl FallbackManager { if let Ok(mut m) = self.fallback_models.lock() { *m = models .into_iter() - .map(|name| FallbackEntry::with_max_fail_count(name, max_fc)) + .map(|name| FallbackEntry::new(name).with_max_fail_count(max_fc)) .collect(); } self.recompute_active_fallback(); @@ -1392,7 +1182,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1433,7 +1223,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1469,7 +1259,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1500,7 +1290,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1532,7 +1322,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1566,7 +1356,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-70b"); @@ -1600,7 +1390,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.add_fallback_model("llm-2"); @@ -1640,16 +1430,14 @@ impl FallbackManager { /// Record an API failure and check if fallback should be triggered. /// /// Called by the agent loop each time an LLM API call fails (e.g. - /// rate limit, server error, timeout). Increments the atomic failure - /// counter and returns `true` when the count reaches the configured - /// [`fallback_threshold`](FallbackManager::new) **and** the circuit - /// has not already been activated. This two-condition guard prevents - /// re-tripping the circuit on every subsequent failure after the - /// first trip. + /// rate limit, server error, timeout). Returns `true` when the + /// consecutive failure count reaches the configured + /// [`fallback_threshold`](FallbackManager::new) and the circuit + /// has not already been activated. /// /// # Returns /// - /// * `true` — the failure count has just reached the threshold for + /// * `true` — the failure count reached the threshold for /// the first time; the caller should switch to the fallback model. /// * `false` — either the threshold hasn't been reached yet, or the /// circuit has already been tripped. @@ -1657,7 +1445,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// let mgr = FallbackManager::new(3, 2); /// assert!(!mgr.record_api_failure()); // 1 /// assert!(!mgr.record_api_failure()); // 2 @@ -1699,7 +1487,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// let mgr = FallbackManager::new(3, 2); /// mgr.record_api_failure(); /// mgr.record_api_failure(); @@ -1728,7 +1516,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// let mgr = FallbackManager::new(3, 2); /// mgr.record_api_failure(); /// assert_eq!(mgr.consecutive_failures(), 1); @@ -1792,7 +1580,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState}; /// /// let mgr = FallbackManager::new(3, 2); /// assert!(!mgr.record_model_failure()); // 1 @@ -1857,7 +1645,7 @@ impl FallbackManager { /// /// ```rust /// use std::time::Duration; - /// use loopctl::loop_control::fallback::FallbackManager; + /// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::new(3, 2); /// // Not in fallback state → false @@ -1889,7 +1677,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.transition_to_fallback(); @@ -1919,7 +1707,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.transition_to_fallback(); @@ -1948,7 +1736,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState}; /// /// let mgr = FallbackManager::new(3, 2); /// mgr.transition_to_fallback(); @@ -1973,7 +1761,7 @@ impl FallbackManager { /// /// Performs a full reset: sets the state to `Primary`, zeros all /// counters, clears the `fallback_activated` flag, and clears the - /// fallback timestamp. This is a "hard reset" — it erases all + /// fallback timestamp. Hard reset — erases all /// failure history and is equivalent to creating a new manager. /// /// Use this when you want to force the circuit back to its initial @@ -1983,7 +1771,7 @@ impl FallbackManager { /// # Example /// /// ```rust - /// use loopctl::loop_control::fallback::{FallbackManager, FallbackState}; + /// use loopctl::fallback::{FallbackManager, FallbackState}; /// /// let mgr = FallbackManager::new(3, 2); /// for _ in 0..3 { mgr.record_model_failure(); } @@ -2005,6 +1793,26 @@ impl FallbackManager { } self.clear_all_fallback_failed(); } + + // ================================================== + // Private helpers + // ================================================== + + /// Recompute the cached [`active_fallback`](Self::active_fallback) from + /// the current fallback chain. + /// + /// [`fallback_models`]: Self::fallback_models + /// [`active_fallback`]: Self::active_fallback + fn recompute_active_fallback(&self) { + let active = self + .fallback_models + .lock() + .ok() + .and_then(|m| m.iter().find(|e| !e.failed()).map(|e| e.name.clone())); + if let Ok(mut cached) = self.active_fallback.lock() { + *cached = active; + } + } } /// Produces a [`FallbackManager`] with production defaults. @@ -2014,14 +1822,10 @@ impl FallbackManager { /// No model name is stored; use [`FallbackManager::for_model`] or /// [`FallbackManager::set_original_model`] to configure one. /// -/// This implementation delegates to [`FallbackManager::new`] with -/// hardcoded thresholds. The resulting manager starts in -/// [`FallbackState::Primary`] with all counters zeroed. -/// /// # Example /// /// ```rust -/// use loopctl::loop_control::fallback::FallbackManager; +/// use loopctl::fallback::FallbackManager; /// /// let mgr = FallbackManager::default(); /// assert_eq!(mgr.consecutive_failures(), 0); diff --git a/src/hooks.rs b/src/hooks.rs index 55e1eda..f1b9433 100644 --- a/src/hooks.rs +++ b/src/hooks.rs @@ -1,6 +1,6 @@ //! Hook system — bidirectional lifecycle control for agent loops. //! -//! Hooks differ from observers ([`crate::core::observer::LoopObserver`]) in two key ways: +//! Hooks differ from observers ([`crate::observer::LoopObserver`]) in two key ways: //! //! 1. **Return values matter.** Pre-hooks return [`HookAction`] (or [`CompactResult`]) //! to control whether an action proceeds. @@ -58,7 +58,7 @@ use context::{ /// Hook trait for bidirectional lifecycle control. /// -/// Hooks differ from observers ([`crate::core::observer::LoopObserver`]) in two key ways: +/// Hooks differ from observers ([`crate::observer::LoopObserver`]) in two key ways: /// /// 1. **Return values matter.** `on_pre_*` methods return `Option` (or /// `Option`). Returning `Some(Block{...})` prevents the action @@ -176,7 +176,7 @@ pub trait Hook: Send + Sync { /// /// - [`Interactivity::Headless`] — there is no human in the loop, so /// `Ask` is automatically downgraded to `Block` with a descriptive -/// reason. This is the default and the correct mode for autonomous +/// reason. Default and correct mode for autonomous /// / headless agents (e.g. `BareLoop`). /// - [`Interactivity::Interactive`] — a human is available to respond /// to prompts, so `Ask` passes through unchanged. @@ -217,7 +217,7 @@ pub enum HookAction { Block { /// Why the action was blocked. /// - /// This is returned to the model as a tool error, allowing it + /// Returned to the model as a tool error, allowing it /// to adjust its approach. reason: String, }, diff --git a/src/hooks/builtin/auto_commit.rs b/src/hooks/builtin/auto_commit.rs index c5cf1eb..f5eadf7 100644 --- a/src/hooks/builtin/auto_commit.rs +++ b/src/hooks/builtin/auto_commit.rs @@ -23,19 +23,15 @@ pub enum GitExecutorError { /// Failed to execute git command. #[error("Failed to execute git: {0}")] ExecutionFailed(String), - /// Git command returned non-zero exit code. #[error("Git error: {0}")] GitError(String), - /// No changes to commit. #[error("No changes to commit")] NoChanges, - /// Invalid repository state. #[error("Invalid repository state: {0}")] InvalidState(String), - /// Git command timed out. #[error("Git command timed out after {0:?}")] Timeout(Duration), @@ -121,7 +117,6 @@ impl GitExecutor { /// or [`GitExecutorError::GitError`] if the command fails. pub fn has_changes() -> Result { let (stdout, _stderr) = Self::run_git(&["status", "--porcelain"])?; - let status = String::from_utf8_lossy(&stdout); Ok(!status.trim().is_empty()) } @@ -428,13 +423,11 @@ impl AutoCommitHook { } } - /// Create a new auto-commit hook with custom configuration. + /// Set custom configuration. #[must_use] - pub fn with_config(config: AutoCommitConfig) -> Self { - Self { - config, - modified_files: Mutex::new(Vec::new()), - } + pub fn with_config(mut self, config: AutoCommitConfig) -> Self { + self.config = config; + self } fn track_modification(&self, file: &str) { @@ -529,7 +522,7 @@ mod tests { enabled: false, ..AutoCommitConfig::default() }; - let hook = AutoCommitHook::with_config(config); + let hook = AutoCommitHook::new().with_config(config); assert_eq!(hook.name(), "auto_commit"); } diff --git a/src/hooks/context.rs b/src/hooks/context.rs index 30c7b81..e9c7386 100644 --- a/src/hooks/context.rs +++ b/src/hooks/context.rs @@ -65,6 +65,25 @@ pub enum CompactTrigger { Manual, } +impl From for CompactTrigger { + /// Map the compaction pipeline's [`CompactReason`](crate::compact::types::CompactReason) into the hook-level + /// [`CompactTrigger`]. + /// + /// [`ThresholdExceeded`](crate::compact::types::CompactReason::ThresholdExceeded) + /// and + /// [`Emergency`](crate::compact::types::CompactReason::Emergency) + /// are both automatic triggers, while + /// [`Manual`](crate::compact::types::CompactReason::Manual) maps to + /// [`Manual`](CompactTrigger::Manual). + fn from(reason: crate::compact::types::CompactReason) -> Self { + match reason { + crate::compact::types::CompactReason::ThresholdExceeded + | crate::compact::types::CompactReason::Emergency => CompactTrigger::Auto, + crate::compact::types::CompactReason::Manual => CompactTrigger::Manual, + } + } +} + /// Context provided to `on_pre_compact` hooks. /// /// Hooks can inspect the current state and decide to abort diff --git a/src/hooks/executor.rs b/src/hooks/executor.rs index e2bac50..35cc4d5 100644 --- a/src/hooks/executor.rs +++ b/src/hooks/executor.rs @@ -78,17 +78,15 @@ impl HookExecutor { } } - /// Create an executor with the given interactivity mode and no hooks. + /// Set the interactivity mode. /// /// Use [`interactivity`](Self::interactivity) to change the mode after /// construction, or [`with_hook`](Self::with_hook) to add hooks via /// builder pattern. #[must_use] - pub fn with_interactivity(interactivity: Interactivity) -> Self { - Self { - hooks: Vec::new(), - interactivity, - } + pub fn with_interactivity(mut self, interactivity: Interactivity) -> Self { + self.interactivity = interactivity; + self } /// Set the interactivity mode (builder style). @@ -671,7 +669,8 @@ mod tests { } // Interactive executor passes Ask through unchanged. - let executor = HookExecutor::with_interactivity(Interactivity::Interactive) + let executor = HookExecutor::new() + .with_interactivity(Interactivity::Interactive) .with_hook(Arc::new(AskHook)); let ctx = dummy_pre_ctx(); let action = executor.check_pre_tool_use(&ctx); @@ -775,8 +774,9 @@ mod tests { } // with_interactivity(Headless) should downgrade. - let headless = - HookExecutor::with_interactivity(Interactivity::Headless).with_hook(Arc::new(AskHook)); + let headless = HookExecutor::new() + .with_interactivity(Interactivity::Headless) + .with_hook(Arc::new(AskHook)); assert!( headless.check_pre_tool_use(&dummy_pre_ctx()).is_block(), "Headless via with_interactivity should downgrade Ask" @@ -796,7 +796,8 @@ mod tests { } // Interactive mode does NOT alter Block actions. - let executor = HookExecutor::with_interactivity(Interactivity::Interactive) + let executor = HookExecutor::new() + .with_interactivity(Interactivity::Interactive) .with_hook(Arc::new(BlockOnlyHook)); let ctx = dummy_pre_ctx(); let action = executor.check_pre_tool_use(&ctx); @@ -813,7 +814,7 @@ mod tests { #[test] fn no_hooks_returns_allow_in_interactive() { - let executor = HookExecutor::with_interactivity(Interactivity::Interactive); + let executor = HookExecutor::new().with_interactivity(Interactivity::Interactive); let ctx = dummy_pre_ctx(); assert!(executor.check_pre_tool_use(&ctx).is_allow()); } diff --git a/src/lib.rs b/src/lib.rs index 9b0471d..0da3c39 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,33 +3,64 @@ //! //! # Module Overview //! +//! ## Foundational Types +//! //! - **[`message`]** — Core conversation types: messages, parts, tool results. -//! - **[`api_client`]** — Trait for LLM provider communication. -//! - **[`api_error`]** — API and infrastructure error types with classification. -//! - **[`builder`]** — Fluent builder API for constructing configured agents. +//! - **[`error`]** — Central error enum ([`LoopError`](error::LoopError)) for all framework operations. +//! - **[`config`]** — Session configuration ([`LoopConfig`](config::LoopConfig)). //! - **[`cancel`]** — Cooperative cancellation signal (`CancelSignal`). +//! +//! ## Subsystems +//! +//! - **[`observer`]** — Lifecycle event observation ([`LoopObserver`](observer::LoopObserver), [`ObserverHost`](observer::ObserverHost)). +//! - **[`memory`]** — Agent memory trait ([`LoopMemory`](memory::LoopMemory)) and entry types. +//! - **[`reflection`]** — Failure reflection and recovery strategies. +//! - **[`detection`]** — Loop and convergence detection ([`DetectionManager`](detection::DetectionManager)). +//! - **[`fallback`]** — Circuit breaker pattern for automatic API model fallback. //! - **[`compact`]** — Context management and compaction (threshold detection, pluggable strategies). -//! - **[`core`]** — Foundational traits and error types. -//! - **[`builtin`]** — Reference implementations of core traits ([`builtin::memory::InMemoryStore`], etc.). -//! - **[`hooks`]** — Bidirectional lifecycle control (allow/block/ask before tool use, compaction). *Requires `hooks` feature.* -//! - **[`loop_control`]** — Detection and intervention modules for agent loops. -//! - **[`engine`]** — The core agentic loop that orchestrates the full agent lifecycle. //! - **[`stream`]** — Streaming event types for LLM API responses. +//! - **[`middleware`]** — Tool dispatch middleware pipeline (timeouts, permissions, output limits). //! - **[`tool`]** — Tool trait, registry, and supporting types. //! - **[`tool::health`]** — Per-tool health monitoring, circuit breakers, and self-healing routing. *Requires `tool_health` feature.* +//! +//! ## API Layer +//! +//! - **[`api`]** — LLM API client trait ([`ApiClient`](api::ApiClient)) and error types. +//! +//! ## Runtime & Capabilities +//! +//! - **[`capabilities`]** — Capability traits ([`Observable`](capabilities::Observable), [`Detectable`](capabilities::Detectable), etc.). +//! - **[`runtime`]** — [`LoopRuntime`](runtime::LoopRuntime) — the default infrastructure bundle. +//! +//! ## Engine +//! +//! - **[`engine`]** — The core agentic loop ([`BareLoop`](engine::BareLoop)) that orchestrates the full agent lifecycle. +//! +//! ## Support +//! +//! - **[`builder`]** — Fluent builder API for constructing configured agents. +//! - **[`memory::builtin`]** — Reference [`InMemoryStore`](memory::builtin::InMemoryStore) implementation. +//! - **[`hooks`]** — Bidirectional lifecycle control (allow/block/ask before tool use, compaction). *Requires `hooks` feature.* +//! - **[`testing`]** — Test utilities and fixtures. *Requires `testing` feature.* -pub mod api_client; -pub mod api_error; +pub mod api; pub mod builder; -pub mod builtin; pub mod cancel; +pub mod capabilities; pub mod compact; -pub mod core; +pub mod config; +pub mod detection; pub mod engine; +pub mod error; +pub mod fallback; #[cfg(feature = "hooks")] pub mod hooks; -pub mod loop_control; +pub mod memory; pub mod message; +pub mod middleware; +pub mod observer; +pub mod reflection; +pub mod runtime; pub mod stream; #[cfg(feature = "testing")] pub mod testing; diff --git a/src/loop_control.rs b/src/loop_control.rs deleted file mode 100644 index c877dcf..0000000 --- a/src/loop_control.rs +++ /dev/null @@ -1,20 +0,0 @@ -//! Loop control — detection and intervention modules for agent loops. -//! -//! Provides convergence detection, loop detection, fallback management, -//! a unified detection manager, and a manager bundle for agent infrastructure. -//! -//! # Provided Modules -//! -//! - **[`convergence`]** — Detects when agent responses become semantically similar. -//! - **[`loop_detector`]** — Detects repetitive tool-use loops and enforces limits. -//! - **[`fallback`]** — Circuit breaker pattern for automatic API model fallback. -//! - **[`detection`]** — Unified manager that orchestrates loop and convergence detection. -//! - **[`bundle`]** — Aggregate struct for the agent's infrastructure managers. -//! -//! For lifecycle observation, see [`crate::core::observer`]. - -pub mod bundle; -pub mod convergence; -pub mod detection; -pub mod fallback; -pub mod loop_detector; diff --git a/src/loop_control/bundle.rs b/src/loop_control/bundle.rs deleted file mode 100644 index 6bcb2b6..0000000 --- a/src/loop_control/bundle.rs +++ /dev/null @@ -1,287 +0,0 @@ -//! Manager bundle — aggregate struct for the agent's infrastructure managers. -//! -//! The [`ManagerBundle`] owns the managers that the framework builder produces, -//! and provides a single `reset_all()` call to reinitialise every manager at the -//! start of a new task. -//! -//! # Provided Managers -//! -//! - **[`FallbackManager`]** — Circuit breaker for automatic API model fallback. -//! - **[`DetectionManager`]** — Loop and convergence detection orchestrator. -//! -//! # Quick Start -//! -//! ``` -//! use loopctl::loop_control::bundle::ManagerBundle; -//! use loopctl::loop_control::fallback::FallbackManager; -//! -//! // Create with defaults -//! let bundle = ManagerBundle::new(); -//! -//! // Or with a custom fallback manager -//! let bundle = ManagerBundle::new() -//! .with_fallback(FallbackManager::for_model("llm-70b")); -//! -//! // Reset all managers at the start of a new task -//! bundle.reset_all(); -//! ``` - -use crate::core::observer::{LoopObserver, ObserverHost}; -use crate::loop_control::detection::DetectionManager; -use crate::loop_control::fallback::FallbackManager; -use std::sync::Arc; - -/// Bundle of framework-provided manager instances. -/// -/// This struct is constructed by the `AgentBuilder` and passed to production -/// `AgentLoopBuilder` via `into_raw_parts()`. Each manager in the bundle -/// handles a specific cross-cutting concern (fallback, detection, etc.) so -/// that the agent core can remain focused on turn processing. -/// -/// # Construction -/// -/// Prefer [`ManagerBundle::new`] for defaults or the builder-style -/// [`ManagerBundle::with_fallback`] to override individual managers. -/// -/// ``` -/// # use loopctl::loop_control::bundle::ManagerBundle; -/// let bundle = ManagerBundle::new(); -/// assert!(bundle.fallback.active_model().is_none()); -/// ``` -pub struct ManagerBundle { - /// Circuit breaker for API model fallback. - /// - /// Manages automatic failover from a primary LLM model to a fallback - /// model when consecutive API failures exceed a threshold. Created - /// with [`FallbackManager::default`] which uses 3 failures to trip - /// and 2 successes to recover. - /// - /// See [`FallbackManager`] for the full state-machine documentation. - pub fallback: FallbackManager, - - /// Loop and convergence detection orchestrator. - /// - /// Wraps [`LoopDetector`](crate::loop_control::loop_detector::LoopDetector) - /// and [`ConvergenceDetector`](crate::loop_control::convergence::ConvergenceDetector) - /// behind a single interface. Records tool operations and agent - /// responses, then checks for repeated patterns or semantically - /// similar outputs. - /// - /// See [`DetectionManager`] for the full API documentation. - pub detection: DetectionManager, - - /// Observer host for cross-cutting lifecycle hooks. - /// - /// Observers are registered via [`ManagerBundle::register_observer`] and - /// called at well-defined hook points inside the agent loop. - /// - /// See [`ObserverHost`] and [`LoopObserver`] for details. - observers: ObserverHost, -} - -impl ManagerBundle { - /// Create a new bundle with default managers. - /// - /// Each manager is initialized with its default configuration. - /// Use the `with_*` builder methods to override individual managers. - /// - /// Called when building a new agent to seed the bundle with defaults. - /// - /// # Example - /// - /// ``` - /// # use loopctl::loop_control::bundle::ManagerBundle; - /// let bundle = ManagerBundle::new(); - /// assert!(bundle.fallback.active_model().is_none()); - /// ``` - #[must_use] - pub fn new() -> Self { - Self { - fallback: FallbackManager::default(), - detection: DetectionManager::default(), - observers: ObserverHost::new(), - } - } - - /// Replace the fallback manager with a custom instance. - /// - /// Consumes the current bundle and returns a new one with the given - /// [`FallbackManager`]. This is the builder-style way to override the - /// default fallback behavior. - /// - /// Typically used by callers who need custom fallback thresholds or a - /// pre-configured model name, via - /// `AgentBuilder::with_fallback`. - /// - /// # Example - /// - /// ``` - /// # use loopctl::loop_control::bundle::ManagerBundle; - /// # use loopctl::loop_control::fallback::FallbackManager; - /// let fallback = FallbackManager::for_model("llm-70b"); - /// let bundle = ManagerBundle::new().with_fallback(fallback); - /// assert_eq!(bundle.fallback.active_model().as_deref(), Some("llm-70b")); - /// ``` - #[must_use] - pub fn with_fallback(mut self, fallback: FallbackManager) -> Self { - self.fallback = fallback; - self - } - - /// Replace the detection manager with a custom instance. - /// - /// Consumes the current bundle and returns a new one with the given - /// [`DetectionManager`]. Use this when you need custom loop or - /// convergence thresholds, e.g. via - /// [`DetectionManager::with_config`]. - /// - /// # Example - /// - /// ``` - /// # use loopctl::loop_control::bundle::ManagerBundle; - /// # use loopctl::loop_control::detection::{DetectionManager, DetectionConfig}; - /// let config = DetectionConfig { - /// loop_threshold: 5, - /// ..Default::default() - /// }; - /// let detection = DetectionManager::with_config(config).unwrap(); - /// let bundle = ManagerBundle::new().with_detection(detection); - /// assert_eq!(bundle.detection.config().loop_threshold, 5); - /// ``` - #[must_use] - pub fn with_detection(mut self, detection: DetectionManager) -> Self { - self.detection = detection; - self - } - - /// Register an observer with the observer host. - /// - /// Observers are called at lifecycle hook points inside the agent - /// loop, in registration order. All observers are notified at every - /// hook point (no short-circuiting). - /// - /// See [`LoopObserver`] for the trait definition and available hooks. - /// - /// # Example - /// - /// ```rust,ignore - /// use loopctl::core::observer::LoopObserver; - /// use std::sync::Arc; - /// - /// let mut bundle = ManagerBundle::new(); - /// bundle.register_observer(Arc::new(MyObserver)); - /// ``` - pub fn register_observer(&mut self, observer: Arc) { - self.observers.register(observer); - } - - /// Get a reference to the observer host. - /// - /// Used by `BareLoop` to dispatch hook calls. - pub fn observers(&self) -> &ObserverHost { - &self.observers - } - - /// Reset all managers and observers to their initial state. - /// - /// Delegates to each manager's `reset()` method and calls - /// [`ObserverHost::reset_all`]. Typically called at the start of - /// a new agent task or session. - /// - /// # Example - /// - /// ``` - /// # use loopctl::loop_control::bundle::ManagerBundle; - /// let bundle = ManagerBundle::new(); - /// // ... after a session ... - /// bundle.reset_all(); - /// // All managers are back to their initial state - /// ``` - pub fn reset_all(&self) { - self.fallback.reset(); - self.detection.reset(); - self.observers.reset_all(); - } -} - -impl Default for ManagerBundle { - /// Produce a [`ManagerBundle`] with default managers. - /// - /// Equivalent to [`ManagerBundle::new`]. Exists so that generic code - /// can write `ManagerBundle::default()` or derive `Default` on parent - /// structs that contain a [`ManagerBundle`]. - /// - /// # Example - /// - /// ``` - /// # use loopctl::loop_control::bundle::ManagerBundle; - /// let bundle = ManagerBundle::default(); - /// // identical to ManagerBundle::new() - /// ``` - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::loop_control::detection::{DetectedPattern, DetectionManager}; - - #[test] - fn test_bundle_default() { - let bundle = ManagerBundle::default(); - assert!(bundle.fallback.active_model().is_none()); - } - - #[test] - fn test_bundle_with_custom_fallback() { - let fallback = FallbackManager::for_model("my-model"); - let bundle = ManagerBundle::new().with_fallback(fallback); - assert_eq!(bundle.fallback.active_model().as_deref(), Some("my-model")); - } - - #[test] - fn test_reset_all() { - let bundle = ManagerBundle::new(); - bundle.reset_all(); - // No panic - } - - #[test] - fn test_bundle_contains_detection_manager() { - let bundle = ManagerBundle::new(); - // Default detection config values - assert_eq!(bundle.detection.config().loop_threshold, 3); - assert_eq!(bundle.detection.config().stop_threshold, 10); - } - - #[test] - fn test_bundle_with_custom_detection() { - use crate::loop_control::detection::DetectionConfig; - - let config = DetectionConfig { - loop_threshold: 7, - stop_threshold: 20, - ..Default::default() - }; - let detection = DetectionManager::with_config(config).unwrap(); - let bundle = ManagerBundle::new().with_detection(detection); - assert_eq!(bundle.detection.config().loop_threshold, 7); - assert_eq!(bundle.detection.config().stop_threshold, 20); - // Fallback is still the default - assert!(bundle.fallback.active_model().is_none()); - } - - #[test] - fn test_reset_all_clears_detection() { - let bundle = ManagerBundle::new(); - // Record a tool call to populate detection state - let _ = bundle.detection.record_tool_call("Read", 12345); - // Reset should clear it - bundle.reset_all(); - // After reset, recording the same call should not trigger detection - let pattern = bundle.detection.record_tool_call("Read", 12345); - assert!(matches!(pattern, DetectedPattern::NoPattern)); - } -} diff --git a/src/memory.rs b/src/memory.rs new file mode 100644 index 0000000..c50f411 --- /dev/null +++ b/src/memory.rs @@ -0,0 +1,176 @@ +//! Agent memory trait — interface for agent memory systems. +//! +//! Memory allows agents to learn from past interactions and retrieve +//! relevant context for future tasks. Defines the core +//! [`LoopMemory`] trait that all memory backends implement, along with +//! the [`MemoryEntry`] value type and supporting enumerations. +//! +//! # Provided Implementations +//! +//! - **`TrajectoryMemory`** — Records tool-execution trajectories and +//! retrieves relevant past experiences. +//! +//! # Quick Start +//! +//! ``` +//! use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory, ConsolidationStats}; +//! use loopctl::error::LoopError; +//! use std::future::Future; +//! use std::pin::Pin; +//! +//! struct InMemoryStore { +//! entries: Vec, +//! } +//! +//! impl LoopMemory for InMemoryStore { +//! fn store(&mut self, entry: MemoryEntry) +//! -> Pin> + Send + '_>> +//! { +//! Box::pin(async move { self.entries.push(entry); Ok(()) }) +//! } +//! fn retrieve(&self, query: &str, limit: usize) +//! -> Pin, LoopError>> + Send + '_>> +//! { +//! let query = query.to_string(); +//! Box::pin(async move { +//! Ok(self.entries.iter() +//! .filter(|e| e.memory.contains(&query)) +//! .take(limit) +//! .cloned() +//! .collect()) +//! }) +//! } +//! fn consolidate(&mut self) +//! -> Pin> + Send + '_>> +//! { +//! Box::pin(async move { Ok(ConsolidationStats::default()) }) +//! } +//! fn len(&self) -> usize { +//! self.entries.len() +//! } +//! } +//! ``` + +use crate::error::LoopError; +use std::future::Future; +use std::pin::Pin; + +pub use builtin::InMemoryStore; +pub use entry::{ConsolidationStats, MemoryCategory, MemoryEntry}; +pub mod builtin; +pub mod entry; + +/// A memory system for loops. +/// +/// Implementations can store and retrieve entries using different +/// strategies (vector similarity, keyword matching, recency, etc.). +/// +/// # Implementing +/// +/// At a minimum you must provide [`store`](LoopMemory::store), +/// [`retrieve`](LoopMemory::retrieve), [`consolidate`](LoopMemory::consolidate), +/// and [`len`](LoopMemory::len). The trait supplies a default +/// [`is_empty`](LoopMemory::is_empty) implementation that delegates to `len`. +/// +/// # Example +/// +/// ``` +/// use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory, ConsolidationStats}; +/// use loopctl::error::LoopError; +/// use std::future::Future; +/// use std::pin::Pin; +/// +/// struct InMemoryStore { +/// entries: Vec, +/// } +/// +/// impl LoopMemory for InMemoryStore { +/// fn store(&mut self, entry: MemoryEntry) +/// -> Pin> + Send + '_>> +/// { +/// Box::pin(async move { self.entries.push(entry); Ok(()) }) +/// } +/// fn retrieve(&self, query: &str, limit: usize) +/// -> Pin, LoopError>> + Send + '_>> +/// { +/// let query = query.to_string(); +/// Box::pin(async move { +/// Ok(self.entries.iter() +/// .filter(|e| e.memory.contains(&query)) +/// .take(limit) +/// .cloned() +/// .collect()) +/// }) +/// } +/// fn consolidate(&mut self) +/// -> Pin> + Send + '_>> +/// { +/// Box::pin(async move { +/// let before = self.entries.len(); +/// self.entries.retain(|e| e.relevance > 0.1); +/// let after = self.entries.len(); +/// Ok(ConsolidationStats { +/// entries_before: before, +/// entries_after: after, +/// pruned: before - after, +/// ..Default::default() +/// }) +/// }) +/// } +/// fn len(&self) -> usize { +/// self.entries.len() +/// } +/// } +/// ``` +pub trait LoopMemory: Send + Sync { + /// Store a new memory entry. + /// + /// Called whenever the agent encounters information worth remembering — + /// for example after a successful tool invocation, a resolved error, or + /// an insight drawn from conversation. Implementations should persist the + /// entry in whatever backing store they use. + fn store( + &mut self, + entry: MemoryEntry, + ) -> Pin> + Send + '_>>; + + /// Retrieve memory entries relevant to the given query. + /// + /// Called before each turn (or on demand) to surface context the agent + /// can use. Returns up to `limit` entries ordered by relevance. The + /// definition of "relevance" is left to the implementation — common + /// strategies include vector embedding similarity, keyword overlap, + /// recency weighting, or a hybrid approach. + /// + /// Implementations that track [`MemoryEntry::access_count`] must use + /// interior mutability (e.g. `AtomicUsize`, `Mutex`) since this method + /// takes `&self`. + fn retrieve( + &self, + query: &str, + limit: usize, + ) -> Pin, LoopError>> + Send + '_>>; + + /// Consolidate memory (e.g. prune, summarize, compress). + /// + /// Called periodically to keep the memory store healthy. Implementations + /// may remove low-relevance entries, merge duplicates, or produce + /// compressed summaries. Returns [`ConsolidationStats`] describing what + /// was done. + fn consolidate( + &mut self, + ) -> Pin> + Send + '_>>; + + /// Number of entries currently stored. + /// + /// Used by the framework and by [`is_empty`](LoopMemory::is_empty). + fn len(&self) -> usize; + + /// Whether the memory is empty. + /// + /// Defaults to `self.len() == 0`. Override only if you need a cheaper + /// check than counting all entries. + fn is_empty(&self) -> bool { + self.len() == 0 + } +} diff --git a/src/builtin/memory.rs b/src/memory/builtin.rs similarity index 60% rename from src/builtin/memory.rs rename to src/memory/builtin.rs index 62d5400..0e4e32a 100644 --- a/src/builtin/memory.rs +++ b/src/memory/builtin.rs @@ -1,15 +1,14 @@ -//! Reference memory implementation — in-memory [`AgentMemory`] backend. +//! Reference memory implementation — in-memory [`LoopMemory`] backend. //! -//! This module provides [`InMemoryStore`], a simple `Vec`-backed -//! implementation of the [`AgentMemory`] trait. It is intended for -//! testing, prototyping, and as a reference for building more -//! sophisticated memory backends (e.g. vector similarity stores). +//! [`InMemoryStore`], a simple `Vec`-backed implementation of the +//! [`LoopMemory`] trait. Intended for testing, prototyping, and as a +//! reference for building more sophisticated memory backends (e.g. vector similarity stores). //! //! # Provided Implementations //! //! - **[`InMemoryStore`]** — Stores [`MemoryEntry`] values in a `Vec` and //! retrieves them via weighted keyword + tag scoring. Supports -//! [`consolidate`](AgentMemory::consolidate) by pruning entries whose +//! [`consolidate`](LoopMemory::consolidate) by pruning entries whose //! [`relevance`](MemoryEntry::relevance) drops below 0.05. //! //! # When to Use @@ -17,28 +16,14 @@ //! Use this backend when you need a zero-dependency, deterministic memory //! store — for example in unit tests, benchmarks, or single-session agents //! that don't require persistence across restarts. For production agents -//! that need durable or distributed memory, implement [`AgentMemory`] on +//! that need durable or distributed memory, implement [`LoopMemory`] on //! top of a database or vector store instead. //! -//! # Data Flow -//! -//! ```text -//! ┌──────────────────────┐ -//! agent turn ───▶ │ store(entry) │ -//! │ ▼ │ -//! │ entries: Vec<_> │ -//! │ ▼ │ -//! agent turn ◀─── │ retrieve(query, n) │ -//! │ ▼ │ -//! periodic ───▶ │ consolidate() │ -//! └──────────────────────┘ -//! ``` -//! //! # Quick Start //! //! ```rust -//! use loopctl::builtin::memory::InMemoryStore; -//! use loopctl::core::{AgentMemory, MemoryEntry, MemoryCategory}; +//! use loopctl::memory::builtin::InMemoryStore; +//! use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory}; //! //! # tokio::runtime::Runtime::new().unwrap().block_on(async { //! let mut store = InMemoryStore::new(); @@ -52,9 +37,12 @@ //! # }); //! ``` -use crate::core::{AgentError, AgentMemory, ConsolidationStats, MemoryEntry}; +use crate::error::LoopError; +use crate::memory::{ConsolidationStats, LoopMemory, MemoryEntry}; +use std::future::Future; +use std::pin::Pin; -/// A simple in-memory store for agent memory entries. +/// A simple in-memory store for loop memory entries. /// /// Stores [`MemoryEntry`] values in a flat `Vec` and retrieves them using /// a weighted scoring function that combines the entry's base @@ -69,7 +57,7 @@ use crate::core::{AgentError, AgentMemory, ConsolidationStats, MemoryEntry}; /// /// # Scoring Formula /// -/// Each candidate entry is scored during [`retrieve`](AgentMemory::retrieve) +/// Each candidate entry is scored during [`retrieve`](LoopMemory::retrieve) /// using a weighted blend of three signals: /// /// ```text @@ -86,20 +74,20 @@ use crate::core::{AgentError, AgentMemory, ConsolidationStats, MemoryEntry}; /// # Thread Safety /// /// [`InMemoryStore`] is `Send + Sync` because all mutation goes through -/// `&mut self` in the [`AgentMemory`] trait. If you need shared mutable +/// `&mut self` in the [`LoopMemory`] trait. If you need shared mutable /// access from multiple tasks, wrap it in `Arc>`. /// /// # Construction /// /// ``` -/// use loopctl::builtin::memory::InMemoryStore; -/// use loopctl::core::{MemoryEntry, MemoryCategory}; +/// use loopctl::memory::builtin::InMemoryStore; +/// use loopctl::memory::{MemoryEntry, MemoryCategory}; /// /// // Empty store: /// let store = InMemoryStore::new(); /// /// // Pre-populated: -/// let store = InMemoryStore::with_entries(vec![ +/// let store = InMemoryStore::new().with_entries(vec![ /// MemoryEntry::new(MemoryCategory::Fact, "The project uses Rust 1.95"), /// ]); /// ``` @@ -107,8 +95,8 @@ use crate::core::{AgentError, AgentMemory, ConsolidationStats, MemoryEntry}; /// # Example /// /// ```rust -/// use loopctl::builtin::memory::InMemoryStore; -/// use loopctl::core::{AgentMemory, MemoryEntry, MemoryCategory}; +/// use loopctl::memory::builtin::InMemoryStore; +/// use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory}; /// /// # tokio::runtime::Runtime::new().unwrap().block_on(async { /// let mut store = InMemoryStore::new(); @@ -120,22 +108,6 @@ use crate::core::{AgentError, AgentMemory, ConsolidationStats, MemoryEntry}; /// # }); /// ``` pub struct InMemoryStore { - /// The internal list of stored memory entries. - /// - /// Entries are appended in insertion order via [`store`](InMemoryStore::store). - /// During [`retrieve`](InMemoryStore::retrieve) the list is scanned linearly - /// and entries are scored / sorted by relevance. During - /// [`consolidate`](InMemoryStore::consolidate), entries with - /// [`relevance`](MemoryEntry::relevance) below 0.05 are removed. - /// - /// Starts empty when created via [`new`](InMemoryStore::new) or - /// [`default`](InMemoryStore::default). Pre-populated when created - /// via [`with_entries`](InMemoryStore::with_entries). - /// - /// **Constraints:** Entries are never reordered — new items are always - /// appended to the end. Removals only occur during - /// [`consolidate`](AgentMemory::consolidate) and preserve the - /// relative order of surviving entries. entries: Vec, } @@ -146,13 +118,13 @@ pub struct InMemoryStore { impl InMemoryStore { /// Create a new empty store. /// - /// Returns a fresh [`InMemoryStore`] whose [`len`](AgentMemory::len) is zero. + /// Returns a fresh [`InMemoryStore`] whose [`len`](LoopMemory::len) is zero. /// /// # Example /// /// ``` - /// use loopctl::builtin::memory::InMemoryStore; - /// use loopctl::core::AgentMemory; + /// use loopctl::memory::builtin::InMemoryStore; + /// use loopctl::memory::LoopMemory; /// /// let store = InMemoryStore::new(); /// assert!(store.is_empty()); @@ -167,66 +139,55 @@ impl InMemoryStore { /// Create a store pre-populated with the given entries. /// /// Useful for setting up test fixtures or seeding an agent with - /// prior knowledge. The entries are stored in the order provided. + /// Replace all entries with the provided list. /// /// # Example /// /// ``` - /// use loopctl::builtin::memory::InMemoryStore; - /// use loopctl::core::{AgentMemory, MemoryEntry, MemoryCategory}; + /// use loopctl::memory::builtin::InMemoryStore; + /// use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory}; /// - /// let store = InMemoryStore::with_entries(vec![ + /// let store = InMemoryStore::new().with_entries(vec![ /// MemoryEntry::new(MemoryCategory::Fact, "Rust 1.75 stabilised async fn in trait"), /// MemoryEntry::new(MemoryCategory::Strategy, "Start refactors with tests"), /// ]); /// assert_eq!(store.len(), 2); /// ``` #[must_use] - pub fn with_entries(entries: Vec) -> Self { - Self { entries } + pub fn with_entries(mut self, entries: Vec) -> Self { + self.entries = entries; + self } } impl Default for InMemoryStore { - /// Returns an empty store, equivalent to [`new`](InMemoryStore::new). - /// - /// Enables the [`Default`] trait so [`InMemoryStore`] can be used in - /// generic contexts that require `T: Default` (e.g. struct initialisation - /// with `..Default::default()`). - /// - /// # Example - /// - /// ``` - /// use loopctl::builtin::memory::InMemoryStore; - /// use loopctl::core::AgentMemory; - /// - /// let store = InMemoryStore::default(); - /// assert!(store.is_empty()); - /// ``` fn default() -> Self { Self::new() } } // =================================================== -// AgentMemory implementation +// LoopMemory implementation // =================================================== -impl AgentMemory for InMemoryStore { - /// Store a new memory entry by appending it to the internal list. +impl LoopMemory for InMemoryStore { + /// Store a new memory entry by appending it to the backing list. /// /// Called whenever the agent encounters information worth remembering — /// for example after a successful tool invocation, a resolved error, or - /// an insight drawn from conversation. The entry is simply pushed onto - /// the internal entries vector. + /// an insight drawn from conversation. /// /// # Errors /// - /// This implementation never returns an error, but the return type - /// conforms to the [`AgentMemory`] trait signature for compatibility. - async fn store(&mut self, entry: MemoryEntry) -> Result<(), AgentError> { - self.entries.push(entry); - Ok(()) + /// This implementation never returns an error. + fn store( + &mut self, + entry: MemoryEntry, + ) -> Pin> + Send + '_>> { + Box::pin(async move { + self.entries.push(entry); + Ok(()) + }) } /// Retrieve memory entries relevant to the given query. @@ -251,8 +212,8 @@ impl AgentMemory for InMemoryStore { /// # Example /// /// ```rust - /// use loopctl::builtin::memory::InMemoryStore; - /// use loopctl::core::{AgentMemory, MemoryEntry, MemoryCategory}; + /// use loopctl::memory::builtin::InMemoryStore; + /// use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory}; /// /// # tokio::runtime::Runtime::new().unwrap().block_on(async { /// let mut store = InMemoryStore::new(); @@ -264,41 +225,48 @@ impl AgentMemory for InMemoryStore { /// } /// # }); /// ``` - async fn retrieve(&self, query: &str, limit: usize) -> Result, AgentError> { - let query_lower = query.to_lowercase(); - let query_words: Vec<&str> = query_lower.split_whitespace().collect(); - - let mut scored: Vec<(f32, MemoryEntry)> = self - .entries - .iter() - .map(|entry| { - let memory_lower = entry.memory.to_lowercase(); - let tag_match = entry - .tags - .iter() - .any(|t| t.to_lowercase().contains(&query_lower)); - let word_matches = query_words - .iter() - .filter(|w| memory_lower.contains(*w)) - .count(); - let base_score = entry.relevance; - #[allow(clippy::cast_precision_loss)] - let query_bonus = if word_matches > 0 { - word_matches as f32 / query_words.len().max(1) as f32 - } else { - 0.0 - }; - let tag_bonus = if tag_match { 0.3 } else { 0.0 }; - ( - base_score * 0.5 + query_bonus * 0.4 + tag_bonus + 0.1, - entry.clone(), - ) - }) - .collect(); - - scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); - - Ok(scored.into_iter().take(limit).map(|(_, e)| e).collect()) + fn retrieve( + &self, + query: &str, + limit: usize, + ) -> Pin, LoopError>> + Send + '_>> { + let query = query.to_string(); + Box::pin(async move { + let query_lower = query.to_lowercase(); + let query_words: Vec<&str> = query_lower.split_whitespace().collect(); + + let mut scored: Vec<(f32, MemoryEntry)> = self + .entries + .iter() + .map(|entry| { + let memory_lower = entry.memory.to_lowercase(); + let tag_match = entry + .tags + .iter() + .any(|t| t.to_lowercase().contains(&query_lower)); + let word_matches = query_words + .iter() + .filter(|w| memory_lower.contains(*w)) + .count(); + let base_score = entry.relevance; + #[allow(clippy::cast_precision_loss)] + let query_bonus = if word_matches > 0 { + word_matches as f32 / query_words.len().max(1) as f32 + } else { + 0.0 + }; + let tag_bonus = if tag_match { 0.3 } else { 0.0 }; + ( + base_score * 0.5 + query_bonus * 0.4 + tag_bonus + 0.1, + entry.clone(), + ) + }) + .collect(); + + scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); + + Ok(scored.into_iter().take(limit).map(|(_, e)| e).collect()) + }) } /// Consolidate memory by pruning low-relevance entries. @@ -317,8 +285,8 @@ impl AgentMemory for InMemoryStore { /// # Example /// /// ```rust - /// use loopctl::builtin::memory::InMemoryStore; - /// use loopctl::core::AgentMemory; + /// use loopctl::memory::builtin::InMemoryStore; + /// use loopctl::memory::LoopMemory; /// /// # tokio::runtime::Runtime::new().unwrap().block_on(async { /// let mut store = InMemoryStore::new(); @@ -326,24 +294,27 @@ impl AgentMemory for InMemoryStore { /// println!("Pruned {} entries", stats.pruned); /// # }); /// ``` - async fn consolidate(&mut self) -> Result { - let entries_before = self.entries.len(); - self.entries.retain(|e| e.relevance >= 0.05); - let pruned = entries_before.saturating_sub(self.entries.len()); - Ok(ConsolidationStats { - entries_before, - entries_after: self.entries.len(), - pruned, - merged: 0, - bytes_saved: 0, + fn consolidate( + &mut self, + ) -> Pin> + Send + '_>> { + Box::pin(async move { + let entries_before = self.entries.len(); + self.entries.retain(|e| e.relevance >= 0.05); + let pruned = entries_before.saturating_sub(self.entries.len()); + Ok(ConsolidationStats { + entries_before, + entries_after: self.entries.len(), + pruned, + merged: 0, + bytes_saved: 0, + }) }) } /// Number of entries currently stored. /// - /// Returns the length of the internal entries - /// vector. Used by the framework to monitor memory usage and by the - /// [`is_empty`](AgentMemory::is_empty) provided method. + /// Used by the framework to monitor memory usage and by the + /// [`is_empty`](LoopMemory::is_empty) provided method. fn len(&self) -> usize { self.entries.len() } @@ -352,7 +323,7 @@ impl AgentMemory for InMemoryStore { #[cfg(test)] mod tests { use super::*; - use crate::core::MemoryCategory; + use crate::memory::MemoryCategory; #[tokio::test] async fn test_store_and_retrieve() { @@ -438,7 +409,7 @@ mod tests { MemoryEntry::new(MemoryCategory::Fact, "fact 1"), MemoryEntry::new(MemoryCategory::Fact, "fact 2"), ]; - let store = InMemoryStore::with_entries(entries); + let store = InMemoryStore::new().with_entries(entries); assert_eq!(store.len(), 2); } diff --git a/src/memory/entry.rs b/src/memory/entry.rs new file mode 100644 index 0000000..4962968 --- /dev/null +++ b/src/memory/entry.rs @@ -0,0 +1,214 @@ +//! Memory entry types — [`MemoryEntry`], [`MemoryCategory`], [`ConsolidationStats`]. +//! +//! These value types support the [`LoopMemory`](super::LoopMemory) trait. + +use serde::{Deserialize, Serialize}; +use std::time::SystemTime; +use uuid::Uuid; + +/// A single memory entry. +/// +/// Each entry represents one discrete piece of information the agent has +/// learned. Entries carry metadata — category, tags, relevance score, +/// access count, and a validated flag — that implementations can use to +/// rank, filter, and consolidate the store. +/// +/// # Construction +/// +/// Prefer the builder-style API starting from [`MemoryEntry::new`]: +/// +/// ``` +/// use loopctl::memory::{MemoryEntry, MemoryCategory}; +/// +/// let entry = MemoryEntry::new(MemoryCategory::Insight, "Prefer concurrent requests when possible") +/// .with_tag("performance") +/// .validated(); +/// ``` +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemoryEntry { + /// UUID v4 for deduplication and stable reference during consolidation. + pub id: Uuid, + /// Entry category, influencing ranking and consolidation rules. + pub category: MemoryCategory, + /// What the agent learned — free-form text. + pub memory: String, + /// Arbitrary labels for categorization and retrieval (e.g. `"performance"`, `"security"`). + pub tags: Vec, + /// Timestamp set by [`MemoryEntry::new`]. Recency-based strategies use this. + pub created_at: SystemTime, + /// Relevance score (0.0–1.0). Starts at 1.0; implementations may decay over time. + pub relevance: f32, + /// Number of times this entry has been retrieved. + pub access_count: usize, + /// Whether this entry has been validated. Consolidation prefers keeping validated entries. + pub validated: bool, +} + +impl Default for MemoryEntry { + fn default() -> Self { + Self { + id: Uuid::new_v4(), + category: MemoryCategory::Working, + memory: String::new(), + tags: Vec::new(), + created_at: SystemTime::now(), + relevance: 0.5, + access_count: 0, + validated: false, + } + } +} + +impl MemoryEntry { + /// Create a new memory entry with a fresh UUID and the current time. + /// + /// The entry starts with `relevance = 1.0`, `access_count = 0`, and + /// `validated = false`. Use the builder methods ([`with_tag`], [`validated`]) + /// to customise further. + /// + /// [`with_tag`]: MemoryEntry::with_tag + /// [`validated`]: MemoryEntry::validated + /// + /// # Example + /// + /// ``` + /// use loopctl::memory::{MemoryEntry, MemoryCategory}; + /// + /// let entry = MemoryEntry::new( + /// MemoryCategory::ErrorPattern, + /// "Timeout on external API — retry with exponential back-off", + /// ); + /// ``` + #[must_use] + pub fn new(category: MemoryCategory, memory: impl Into) -> Self { + Self { + id: Uuid::new_v4(), + category, + memory: memory.into(), + tags: Vec::new(), + created_at: SystemTime::now(), + relevance: 1.0, + access_count: 0, + validated: false, + } + } + + /// Add a tag to this entry (builder style). + /// + /// Tags are lightweight, human-readable labels that speed up broad + /// queries. Call chain-style: + /// + /// ``` + /// use loopctl::memory::{MemoryEntry, MemoryCategory}; + /// + /// let entry = MemoryEntry::new(MemoryCategory::Fact, "Rust 1.75 stabilised async fn in trait") + /// .with_tag("rust") + /// .with_tag("async"); + /// ``` + #[must_use] + pub fn with_tag(mut self, tag: impl Into) -> Self { + self.tags.push(tag.into()); + self + } + + /// Mark this entry as validated (builder style). + /// + /// Validated entries are treated as higher-confidence by consolidation + /// algorithms and are less likely to be pruned. + /// + /// ``` + /// use loopctl::memory::{MemoryEntry, MemoryCategory}; + /// + /// let entry = MemoryEntry::new(MemoryCategory::Strategy, "Use parallel tool calls when independent") + /// .validated(); + /// ``` + #[must_use] + pub fn validated(mut self) -> Self { + self.validated = true; + self + } +} + +/// Category of a memory entry. +/// +/// Each category represents a distinct *kind* of knowledge. Retrieval +/// strategies may weight categories differently (e.g. preferring +/// [`ErrorPattern`](MemoryCategory::ErrorPattern) when debugging), and +/// consolidation rules may vary by category. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MemoryCategory { + /// A recorded trajectory of tool executions. + /// + /// Captures the sequence of tool calls, their inputs, and outcomes for + /// a particular task. Useful for replaying successful strategies. + Trajectory, + + /// A pattern or insight learned from experience. + /// + /// Generalised knowledge that transcends a single interaction — e.g. + /// "users prefer concise summaries over verbose explanations". + Insight, + + /// A pattern of errors and how they were resolved. + /// + /// Pairs an observed error signature with the fix that resolved it, + /// allowing the agent to avoid repeating the same mistake. + ErrorPattern, + + /// A strategy that was proven effective. + /// + /// High-level plans or heuristics that led to good outcomes, such as + /// "when facing a large refactoring, start with tests". + Strategy, + + /// A fact or piece of knowledge. + /// + /// Static information the agent has learned — e.g. "the project uses + /// `PostgreSQL` 15". Facts are not derived from the agent's own reasoning + /// but are still valuable context. + Fact, + + /// Short-term working memory for the current session. + /// + /// Ephemeral entries that are typically discarded at the end of a + /// session. Useful for tracking intermediate state such as "the user + /// asked about file X in the previous turn". + Working, +} + +/// Statistics from a memory consolidation pass. +/// +/// Returned by [`LoopMemory::consolidate`](super::LoopMemory::consolidate) so callers can monitor the +/// health of the memory store over time. +/// +/// # Example +/// +/// ``` +/// use loopctl::memory::ConsolidationStats; +/// +/// let stats = ConsolidationStats { +/// entries_before: 100, +/// entries_after: 80, +/// pruned: 15, +/// merged: 5, +/// ..Default::default() +/// }; +/// println!( +/// "Consolidated: {} → {} entries (pruned {}, merged {})", +/// stats.entries_before, stats.entries_after, stats.pruned, stats.merged, +/// ); +/// ``` +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ConsolidationStats { + /// Number of entries before consolidation. + pub entries_before: usize, + /// Number of entries after consolidation. + pub entries_after: usize, + /// Entries removed (low relevance, stale, superseded). + pub pruned: usize, + /// Entries merged (duplicates combined). + pub merged: usize, + /// Estimated storage reclaimed in bytes. + pub bytes_saved: usize, +} diff --git a/src/message.rs b/src/message.rs index 08eab42..83110ef 100644 --- a/src/message.rs +++ b/src/message.rs @@ -1,7 +1,6 @@ //! Message types for agent conversations. //! -//! This module provides the core message representation used throughout -//! the loopctl framework for communication between users, agents, and +//! Core message representation used throughout the loopctl framework for communication between users, agents, and //! tools. Messages consist of a [`Role`] and a list of [`MessagePart`]s, //! supporting plain text, base64-encoded images, tool invocations, and //! tool results. @@ -200,7 +199,7 @@ impl Message { } } -/// Display a message in a human-readable format. +/// Formats a [`Message`] for display. /// /// Formats each [`MessagePart`] in the message on its own line. /// Text parts render as-is; tool-call parts render as @@ -208,8 +207,7 @@ impl Message { /// as `[Tool Result: {content}]`; image parts render as /// `[Image: {media_type}]`. /// -/// This is primarily useful for debugging and logging. For -/// structured serialization, use [`serde_json::to_string`] instead. +/// For structured serialization, use [`serde_json::to_string`] instead. /// /// # Example /// @@ -285,7 +283,7 @@ pub enum Role { /// /// This implementation is used by [`Display`](fmt::Display) to produce /// the string that LLM APIs expect: `"user"` or `"assistant"`. -/// It is also convenient for logging and building request payloads +/// It is also used for logging and building request payloads /// without pulling in the serde serializer. /// /// # Example @@ -311,8 +309,8 @@ impl fmt::Display for Role { /// /// Messages can contain multiple types of content interleaved in a /// single response: plain text, base64-encoded images, tool-call -/// invocations by the assistant, and tool results. This enum represents -/// each possible part type with tagged JSON serialization. +/// invocations by the assistant, and tool results. Each variant +/// corresponds to a possible part type with tagged JSON serialization. /// /// # Serialization /// @@ -322,7 +320,7 @@ impl fmt::Display for Role { /// /// # Construction /// -/// Use the convenience constructors [`text`](MessagePart::text), +/// Use the constructors [`text`](MessagePart::text), /// [`tool_call`](MessagePart::tool_call), and /// [`tool_result`](MessagePart::tool_result) rather than building /// variants directly. @@ -340,8 +338,8 @@ impl fmt::Display for Role { pub enum MessagePart { /// Plain text content. /// - /// The most common part type. Contains a UTF-8 text string - /// generated by the model or provided by the user. + /// Contains a UTF-8 text string generated by the model or + /// provided by the user. #[serde(rename = "text")] Text { /// The text content of this part. @@ -435,7 +433,7 @@ pub enum MessagePart { impl MessagePart { /// Create a text part. /// - /// Convenience constructor for the most common part type. + /// Constructor for the [`Text`](MessagePart::Text) variant. /// Accepts any type that implements `Into`. /// /// # Arguments @@ -533,7 +531,7 @@ impl MessagePart { /// Returns `true` if this is a [`Text`](MessagePart::Text) block. /// - /// Useful for filtering or pattern matching on parts + /// Allows filtering or pattern matching on parts /// without a full `match` expression. /// /// # Example @@ -550,7 +548,7 @@ impl MessagePart { /// Returns `true` if this is a [`ToolCall`](MessagePart::ToolCall) block. /// - /// Useful for detecting tool invocations in an assistant message + /// Detects tool invocations in an assistant message /// to decide whether to enter the tool-execution loop. /// /// # Example @@ -568,7 +566,7 @@ impl MessagePart { /// Returns `true` if this is a [`ToolResult`](MessagePart::ToolResult) block. /// - /// Useful for identifying tool results in a conversation history. + /// Identifies tool results in a conversation history. /// /// # Example /// @@ -585,7 +583,7 @@ impl MessagePart { /// Get the text content if this is a [`Text`](MessagePart::Text) block. /// /// Returns `Some(&str)` for text parts, `None` for all other - /// variants. Useful for extracting the text from a known-text part. + /// variants. Extracts the text from a known-text part. /// /// # Returns /// @@ -656,16 +654,15 @@ pub struct ImageSource { /// Base64-encoded image data. /// /// The raw bytes of the image encoded as a base64 string. Should - /// not include a data-URI prefix — just the base64 payload. + /// not include a data-URI prefix — only the base64 payload. pub data: String, } impl ImageSource { /// Create a new base64 image source. /// - /// Convenience constructor that sets the `encoding` to `"base64"` - /// automatically. This is the standard construction method for - /// image sources in loopctl API format. + /// Sets the `encoding` to `"base64"` automatically. + /// Standard construction method for image sources in loopctl API format. /// /// # Arguments /// @@ -736,7 +733,7 @@ impl ImageSource { pub enum ToolContent { /// A simple string result. /// - /// The most common case — tools that return plain text output. + /// Returned by tools that produce plain text output. Text(String), /// A multipart result with multiple parts. @@ -800,7 +797,7 @@ impl ToolContent { /// Returns `true` if this is a simple string result. /// - /// Useful for branching on how to process or display the result. + /// Branches on how to process or display the result. /// /// # Example /// @@ -818,9 +815,8 @@ impl ToolContent { /// Produces a [`Default`] value for [`ToolContent`]. /// /// Returns an empty [`String`](ToolContent::Text) variant, -/// which is the neutral starting point for tool results. This is -/// useful when building responses incrementally or initializing -/// result storage. +/// which is the neutral starting point for tool results. Used when +/// building responses incrementally or initializing result storage. /// /// # Example /// @@ -880,8 +876,8 @@ impl From<&str> for ToolContent { /// variant, concatenates all [`ToolContentPart::Text`] parts with /// newlines, silently skipping any non-text parts (e.g. images). /// -/// This is useful for quick debugging and logging. For full -/// structured serialization, use [`serde_json::to_string`]. +/// For quick debugging and logging, use [`Display`](std::fmt::Display). +/// For full structured serialization, use [`serde_json::to_string`]. /// /// # Example /// @@ -944,7 +940,7 @@ impl fmt::Display for ToolContent { pub enum ToolContentPart { /// An image part within a Multipart tool result. /// - /// Contains the base64-encoded image data. Useful for tools + /// Contains the base64-encoded image data. Used by tools /// that capture screenshots or generate images. Image { /// The image data and metadata. @@ -956,7 +952,7 @@ pub enum ToolContentPart { /// A text part within a Multipart tool result. /// - /// The most common part type. Contains a plain text string. + /// Contains a plain text string. Text { /// The text content of this part. /// @@ -969,8 +965,8 @@ pub enum ToolContentPart { impl ToolContentPart { /// Create a text part. /// - /// Convenience constructor for the most common part type. - /// Use this to build individual text segments within a + /// Constructor for the [`Text`](ToolContentPart::Text) variant. + /// Use to build individual text segments within a /// [`ToolContent::Multipart`] result. /// /// # Arguments diff --git a/src/engine/middleware.rs b/src/middleware.rs similarity index 60% rename from src/engine/middleware.rs rename to src/middleware.rs index 1e35972..d499180 100644 --- a/src/engine/middleware.rs +++ b/src/middleware.rs @@ -29,7 +29,7 @@ //! # Example //! //! ```rust,ignore -//! use loopctl::engine::middleware::{ +//! use loopctl::middleware::{ //! ToolPipeline, ToolCallMiddleware, PermissionMiddleware, TimeoutMiddleware, //! }; //! use loopctl::tool::ToolRegistry; @@ -47,95 +47,59 @@ //! let result = pipeline.invoke(ctx).await; //! ``` +pub mod output_limit; +pub mod permission; +pub mod timeout; +pub mod tool_call; +pub mod unknown_tool; + use crate::cancel::CancelSignal; -use crate::core::error::AgentError; -use crate::message::ToolContent; +use crate::error::LoopError; use crate::tool::{PermissionCheck, ToolContext, ToolRegistry}; use serde_json::Value; use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use std::time::{Duration, Instant}; -use tracing; -// =================================================== -// Dispatch context and result -// =================================================== +pub use crate::tool::ToolDispatchResult; + +pub use output_limit::OutputLimitMiddleware; +pub use permission::{AskResolverFn, PermissionCheckFn, PermissionMiddleware}; +pub use timeout::{TimeoutConfig, TimeoutMiddleware}; +pub use tool_call::ToolCallMiddleware; +pub use unknown_tool::UnknownToolMiddleware; + +// ================================================== +// Dispatch context +// ================================================== -/// Context passed through the middleware chain for a single tool invocation. +/// Context for a single tool invocation, passed through the middleware chain. /// /// Built once per tool call by the framework, then threaded through each /// middleware. Middlewares can read from and write to this struct. /// -/// # Fields -/// -/// - [`tool_name`](ToolDispatchContext::tool_name) — Which tool is being called. -/// - [`input`](ToolDispatchContext::input) — The JSON input to the tool. -/// - [`call_id`](ToolDispatchContext::call_id) — Unique ID for this tool call. -/// - [`turn_number`](ToolDispatchContext::turn_number) — Turn number within the session. -/// - [`cancel`](ToolDispatchContext::cancel) — Shared cancellation signal. -/// - [`permission`](ToolDispatchContext::permission) — Permission state for this call. -/// - [`tool_context`](ToolDispatchContext::tool_context) — Context passed to the tool. -/// /// Session identity is available via [`tool_context.session_id`](ToolContext::session_id) -/// rather than duplicated on this struct. +/// rather than duplicated here. pub struct ToolDispatchContext { - /// Which tool is being called. - /// - /// Set by the framework from the model's tool call part. - /// Middlewares can redirect to a different tool by modifying - /// this field (e.g. a routing middleware). + /// From the model's tool call. Middlewares can redirect by modifying this. pub tool_name: String, - - /// The JSON input to the tool. - /// - /// The raw JSON value from the model's tool call part. - /// Middlewares can inspect or transform this before the tool - /// is invoked (e.g. sanitisation, validation, schema upgrade). + /// Raw JSON from the model. Middlewares can inspect/transform before invocation. pub input: Value, - - /// Unique ID for this tool call. - /// - /// Generated by the model provider and used to correlate tool - /// results back to the originating tool call request. + /// Generated by the model provider, used to correlate results back to the request. pub call_id: String, - - /// Turn number within the session. - /// - /// Monotonically increasing counter for the current conversation - /// turn. Useful for logging and for middlewares that need to - /// enforce per-turn limits. + /// Monotonically increasing counter. Useful for per-turn limits. pub turn_number: usize, - - /// Shared cancellation signal — middlewares should check this - /// before doing expensive work. - /// - /// Set by the framework when the user or a deadline triggers - /// cancellation. Long-running middlewares (e.g. [`TimeoutMiddleware`]) - /// should cooperative-check this via `tokio::select!`. + /// Checked via `tokio::select!` in long-running middlewares (e.g. [`TimeoutMiddleware`]). pub cancel: Arc, - - /// Permission state for this call. - /// - /// Set by the framework based on tool registration. Middlewares - /// (e.g. [`PermissionMiddleware`]) can read and modify this. + /// Middlewares (e.g. [`PermissionMiddleware`]) can read and modify this. pub permission: PermissionCheck, - - /// Context passed to the underlying tool invocation. - /// - /// Middlewares can augment this (e.g. adding metadata) before - /// the innermost [`ToolCallMiddleware`] invokes `Tool::call()`. + /// Augmented by middlewares before [`ToolCallMiddleware`] invokes `Tool::call()`. pub tool_context: ToolContext, } -// Re-export the unified tool dispatch result from core types. -// Defined in `crate::core::types` so it's available at both -// `loopctl::core::ToolDispatchResult` and -// `loopctl::engine::middleware::ToolDispatchResult`. -pub use crate::core::types::ToolDispatchResult; -// =================================================== +// ================================================== // Middleware trait -// =================================================== +// ================================================== /// The trait for a middleware in the tool dispatch pipeline. /// @@ -201,7 +165,6 @@ pub enum PipelineError { /// No core dispatch (tool registry) was provided. #[error("pipeline requires a core dispatch (call .core() with a ToolRegistry)")] MissingCore, - /// The pipeline is empty — no middlewares and no core. #[error("pipeline has no middlewares and no core dispatch")] Empty, @@ -244,16 +207,11 @@ pub enum PipelineError { /// let result = pipeline.invoke(ctx).await; /// ``` pub struct ToolPipeline { - /// Ordered middleware layers (outermost first). + /// Outermost first. Index 0 runs first. middlewares: Arc<[Arc]>, - - /// The innermost core dispatch that calls `Tool::call()`. + /// Created from the [`ToolRegistry`] when the pipeline is built. core: Arc, - - /// Cursor position during dispatch. - /// - /// `0` at the entry point. Each middleware receives a pipeline - /// at `index + 1` as its `next`. + /// `0` at entry. Each middleware receives a pipeline at `index + 1`. index: usize, } @@ -328,27 +286,27 @@ impl ToolPipeline { /// /// Each call goes through the full middleware chain. Cancellation is /// checked between calls — if the signal is set, remaining calls - /// are skipped and an [`AgentError::Cancelled`] is returned. + /// are skipped and an [`LoopError::Cancelled`] is returned. /// /// Returns results in the same order as the input calls. /// /// # Errors /// - /// Returns [`AgentError::Cancelled`] if the cancellation signal is + /// Returns [`LoopError::Cancelled`] if the cancellation signal is /// set before all calls have been dispatched. pub async fn dispatch_all( &self, calls: Vec, - ) -> Result, AgentError> { + ) -> Result, LoopError> { let mut results = Vec::with_capacity(calls.len()); for ctx in calls { if ctx.cancel.is_cancelled() { - return Err(AgentError::Cancelled); + return Err(LoopError::Cancelled); } let cancel = Arc::clone(&ctx.cancel); let result = self.invoke(ctx).await; if cancel.is_cancelled() { - return Err(AgentError::Cancelled); + return Err(LoopError::Cancelled); } results.push(result); } @@ -420,7 +378,7 @@ impl ToolPipelineBuilder { /// Set the core tool registry for the pipeline. /// /// The registry is wrapped in a [`ToolCallMiddleware`] that performs - /// the actual tool lookup and invocation. This is the innermost + /// the actual tool lookup and invocation — the innermost /// layer of the pipeline. #[must_use] pub fn core(mut self, registry: Arc) -> Self { @@ -449,673 +407,6 @@ impl Default for ToolPipelineBuilder { } } -// =================================================== -// ToolCallMiddleware (innermost core) -// =================================================== - -/// The innermost middleware that performs the actual tool invocation. -/// -/// Looks up the tool by name in the [`ToolRegistry`], calls -/// [`crate::tool::Tool::call()`], and converts the result into a -/// [`ToolDispatchResult`]. If the tool is not found, produces a -/// soft error result (not a hard error) so the model can recover. -/// -/// This middleware is automatically created by the pipeline builder -/// and always occupies the innermost position in the chain. -pub struct ToolCallMiddleware { - registry: Arc, -} - -impl ToolCallMiddleware { - const NAME: &str = "tool_call"; - - /// Create a new core dispatch wrapping the given registry. - #[must_use] - pub fn new(registry: Arc) -> Self { - Self { registry } - } - - /// Execute the tool call — the terminal dispatch. - /// - /// Looks up the tool by name in the registry, calls `Tool::call()`, - /// and converts the result. This is not a middleware — it has no - /// `next` parameter because there is nothing to chain to. - fn dispatch( - &self, - ctx: &mut ToolDispatchContext, - ) -> Pin + Send + '_>> { - let tool_name = ctx.tool_name.clone(); - let input = ctx.input.clone(); - let tool_ctx = ctx.tool_context.clone(); - let registry = Arc::clone(&self.registry); - let cancel = Arc::clone(&ctx.cancel); - let call_id = ctx.call_id.clone(); - - Box::pin(async move { - let start = Instant::now(); - let Some(tool) = registry.get(&tool_name) else { - let available: Vec = registry.tool_names(); - let available_refs: Vec<&str> = available.iter().map(String::as_str).collect(); - let error = AgentError::tool_not_found(&tool_name, &available_refs); - return ToolDispatchResult::err(&tool_name, error.to_string(), start.elapsed()) - .with_call_id(&call_id); - }; - - let call_result = tokio::select! { - r = tool.call(input, &tool_ctx) => r, - () = cancel.notified() => { - return ToolDispatchResult::err( - &tool_name, - format!("Tool '{tool_name}' cancelled"), - start.elapsed(), - ) - .with_call_id(&call_id); - } - }; - - let duration = start.elapsed(); - ToolDispatchResult::from_result(&tool_name, call_result, duration) - .with_call_id(&call_id) - }) - } -} - -// =================================================== -// PermissionMiddleware -// =================================================== - -/// Permission check function type. -pub type PermissionCheckFn = Arc PermissionCheck + Send + Sync>; - -/// Middleware that checks tool permissions before execution. -/// -/// Inspects the [`PermissionCheck`] in the dispatch context. If the -/// permission is `Deny`, short-circuits with an error result. If -/// `Ask`, also denies (interactive prompts are outside the framework's -/// scope — agents should handle this themselves). -/// -/// # Example -/// -/// ```rust,ignore -/// // Deny all by default -/// let mw = PermissionMiddleware::deny_all(); -/// -/// // Custom logic -/// let mw = PermissionMiddleware::with_check(|ctx| { -/// if ctx.tool_name == "safe_read" { -/// PermissionCheck::Allow -/// } else { -/// PermissionCheck::Deny -/// } -/// }); -/// ``` -pub struct PermissionMiddleware { - /// Override function for permission checking. - /// - /// When `Some`, this function is called to determine the permission - /// for each dispatch. When `None`, the middleware reads - /// [`ToolDispatchContext::permission`] directly. - check_fn: Option, -} - -impl PermissionMiddleware { - /// Create a permission middleware that denies all calls. - /// - /// Every tool call will be short-circuited with a permission-denied - /// error. Useful as a safety default in restricted environments. - #[must_use] - pub fn deny_all() -> Self { - Self { - check_fn: Some(Arc::new(|_| PermissionCheck::Deny { - reason: "blocked by policy".into(), - })), - } - } - - /// Create a permission middleware that allows all calls. - /// - /// No permission checks are performed — every tool call passes - /// through to the next layer. This is the runtime-equivalent of - /// having no permission middleware at all, but it can be useful - /// for logging or metrics in permissive environments. - #[must_use] - pub fn allow_all() -> Self { - Self { - check_fn: Some(Arc::new(|_| PermissionCheck::Allow)), - } - } - - /// Create a permission middleware with a custom check function. - /// - /// The function receives a reference to the dispatch context and - /// returns the appropriate [`PermissionCheck`] for that call. - pub fn with_check( - f: impl Fn(&ToolDispatchContext) -> PermissionCheck + Send + Sync + 'static, - ) -> Self { - Self { - check_fn: Some(Arc::new(f)), - } - } - - /// Create a permission middleware that reads from the context. - /// - /// The middleware reads `ctx.permission` directly, without - /// applying any override. This is useful when the permission - /// is set by the framework or a prior middleware. - #[must_use] - pub fn from_context() -> Self { - Self { check_fn: None } - } - - fn resolve_permission(&self, ctx: &ToolDispatchContext) -> PermissionCheck { - match &self.check_fn { - Some(f) => f(ctx), - None => ctx.permission.clone(), - } - } -} - -impl ToolMiddleware for PermissionMiddleware { - fn name(&self) -> &'static str { - "permission" - } - - fn dispatch<'a>( - &'a self, - ctx: &'a mut ToolDispatchContext, - next: &'a ToolPipeline, - ) -> Pin + Send + 'a>> { - let permission = self.resolve_permission(ctx); - match permission { - PermissionCheck::Allow => next.dispatch(ctx), - PermissionCheck::Ask { .. } | PermissionCheck::Deny { .. } => { - let tool_name = ctx.tool_name.clone(); - let reason = match &permission { - PermissionCheck::Deny { reason } => reason.clone(), - PermissionCheck::Ask { prompt } => format!("permission required: {prompt}"), - _ => "blocked".to_string(), - }; - tracing::warn!( - tool = %tool_name, - permission = %reason, - "tool call blocked by permission middleware" - ); - Box::pin(std::future::ready(ToolDispatchResult::err( - &tool_name, - format!("Permission {reason} for tool '{tool_name}'"), - Duration::ZERO, - ))) - } - PermissionCheck::Modify { .. } => { - // Modify is treated as allow — the tool can proceed - // but the agent might apply modifications to the input. - next.dispatch(ctx) - } - } - } -} - -// =================================================== -// TimeoutMiddleware -// =================================================== - -/// Configuration for the [`TimeoutMiddleware`]. -#[derive(Debug, Clone)] -pub struct TimeoutConfig { - /// Timeout for the initial tool execution attempt. - /// - /// If the inner dispatch does not complete within this duration, - /// [`TimeoutMiddleware`] cancels the future and either retries - /// (if [`retry_on_timeout`](TimeoutConfig::retry_on_timeout) is - /// `true`) or returns a timeout error. - pub timeout: Duration, - - /// Whether to retry once on timeout with double the timeout. - /// - /// When `true`, the middleware makes up to - /// [`max_retries`](TimeoutConfig::max_retries) additional attempts, - /// each with double the previous timeout. - pub retry_on_timeout: bool, - - /// Maximum number of retries (0 = no retry, 1 = one retry). - /// - /// Each retry doubles the current timeout. The total number of - /// attempts is `1 + max_retries`. - pub max_retries: u32, -} - -impl Default for TimeoutConfig { - fn default() -> Self { - Self { - timeout: Duration::from_secs(120), - retry_on_timeout: false, - max_retries: 0, - } - } -} - -/// Middleware that wraps tool execution in a timeout. -/// -/// If the tool execution exceeds the configured timeout, returns an -/// error result. Optionally retries once with a longer timeout. -/// Respects the [`CancelSignal`] via `tokio::select!` so that -/// cancellation is not blocked by a slow tool. -/// -/// # Example -/// -/// ```rust,ignore -/// let mw = TimeoutMiddleware::from_secs(120); -/// let mw = TimeoutMiddleware::new(TimeoutConfig { -/// timeout: Duration::from_secs(60), -/// retry_on_timeout: false, -/// max_retries: 0, -/// }); -/// ``` -pub struct TimeoutMiddleware { - config: TimeoutConfig, -} - -impl TimeoutMiddleware { - /// Create a timeout middleware with the given configuration. - /// - /// `config.timeout` controls the per-tool execution deadline. - /// When `config.retry_on_timeout` is `true`, a timed-out call is - /// retried up to `config.max_retries` additional times with an - /// increasing back-off. - /// - /// For simpler construction see [`from_secs`](Self::from_secs) or - /// [`none`](Self::none). - #[must_use] - pub fn new(config: TimeoutConfig) -> Self { - Self { config } - } - - /// Create a timeout middleware with a fixed timeout in seconds. - /// - /// Uses default retry settings (one retry with double timeout). - #[must_use] - pub fn from_secs(secs: u64) -> Self { - Self { - config: TimeoutConfig { - timeout: Duration::from_secs(secs), - ..TimeoutConfig::default() - }, - } - } - - /// Create a timeout middleware with no timeout (pass-through). - /// - /// Useful for testing or when timeouts are handled elsewhere. - #[must_use] - pub fn none() -> Self { - Self { - config: TimeoutConfig { - timeout: Duration::MAX, - retry_on_timeout: false, - max_retries: 0, - }, - } - } -} - -impl ToolMiddleware for TimeoutMiddleware { - fn name(&self) -> &'static str { - "timeout" - } - - fn dispatch<'a>( - &'a self, - ctx: &'a mut ToolDispatchContext, - next: &'a ToolPipeline, - ) -> Pin + Send + 'a>> { - let config = self.config.clone(); - let tool_name = ctx.tool_name.clone(); - let cancel = Arc::clone(&ctx.cancel); - - Box::pin(async move { - let mut attempt = 0u32; - let mut current_timeout = config.timeout; - - loop { - let result_future = next.dispatch(ctx); - let attempt_for_log = attempt; - - tokio::select! { - result = tokio::time::timeout(current_timeout, result_future) => { - if let Ok(dispatch_result) = result { - return dispatch_result; - } - attempt = attempt.saturating_add(1); - if config.retry_on_timeout && attempt <= config.max_retries { - tracing::warn!( - tool = %tool_name, - attempt = attempt_for_log, - timeout_secs = current_timeout.as_secs(), - "tool execution timed out, retrying" - ); - current_timeout = current_timeout.saturating_mul(2); - continue; - } - tracing::error!( - tool = %tool_name, - timeout_secs = current_timeout.as_secs(), - "tool execution timed out" - ); - return ToolDispatchResult::err( - &tool_name, - format!( - "Tool '{}' timed out after {}s", - tool_name, - current_timeout.as_secs() - ), - current_timeout, - ); - } - () = cancel.notified() => { - return ToolDispatchResult::err( - &tool_name, - format!("Tool '{tool_name}' cancelled"), - Duration::ZERO, - ); - } - } - } - }) - } -} - -// =================================================== -// UnknownToolMiddleware -// =================================================== - -/// Middleware that suggests alternatives when a tool is not found. -/// -/// This middleware wraps the core dispatch. When the inner dispatch -/// produces a "tool not found" error, this middleware intercepts it, -/// computes a string-similarity score against all registered tools, -/// and appends a suggestion to the error message. -/// -/// # Similarity Metric -/// -/// Uses a simple normalized longest-common-substring ratio. This is -/// fast and effective for the common case of minor typos (e.g. -/// `"bash"` → `"basj"`). -/// -/// # Example -/// -/// ```rust,ignore -/// let mw = UnknownToolMiddleware::new(); -/// // If tool "basj" is not found, error message will say: -/// // "Tool 'basj' not found. Did you mean 'bash'?" -/// ``` -pub struct UnknownToolMiddleware { - /// The tool registry, used to enumerate available tool names - /// for suggestions. - /// - /// Owned via `Arc` so the middleware can independently list - /// tool names without needing to extract them from the chain. - registry: Arc, - - /// Minimum similarity score (0.0–1.0) to suggest an alternative. - /// - /// Only tools with a score at or above this threshold are suggested. - /// Defaults to `0.4`. - suggestion_threshold: f64, -} - -impl UnknownToolMiddleware { - /// Create a new unknown-tool middleware with default settings. - /// - /// Uses a [`suggestion_threshold`](Self::with_threshold) of `0.4`, - /// which balances catching common typos against false-positive - /// suggestions. - /// - /// For a custom threshold see [`with_threshold`](Self::with_threshold). - #[must_use] - pub fn new(registry: Arc) -> Self { - Self { - registry, - suggestion_threshold: 0.4, - } - } - - /// Create with a custom similarity threshold. - /// - /// Lower values produce more suggestions (more false positives). - /// Higher values require closer matches. - #[must_use] - pub fn with_threshold(registry: Arc, threshold: f64) -> Self { - Self { - registry, - suggestion_threshold: threshold.clamp(0.0, 1.0), - } - } - - /// Compute similarity between two strings using a normalized - /// longest-common-subsequence approach. - /// - /// Returns a value between 0.0 (completely different) and 1.0 - /// (identical). - fn similarity(a: &str, b: &str) -> f64 { - if a.is_empty() && b.is_empty() { - return 1.0; - } - if a.is_empty() || b.is_empty() { - return 0.0; - } - - let a_lower = a.to_lowercase(); - let b_lower = b.to_lowercase(); - - if a_lower == b_lower { - return 1.0; - } - - // Use longest common subsequence length as similarity metric - let a_chars: Vec = a_lower.chars().collect(); - let b_chars: Vec = b_lower.chars().collect(); - let lcs_len = Self::lcs_length(&a_chars, &b_chars); - - // Tool names are short strings, so u32 is sufficient and avoids - // usize→f64 precision loss on 64-bit targets. - let max_len = u32::try_from(a_chars.len().max(b_chars.len())).unwrap_or(u32::MAX); - let lcs_u32 = u32::try_from(lcs_len).unwrap_or(u32::MAX); - - // Check for prefix match bonus - let prefix_len = a_chars - .iter() - .zip(b_chars.iter()) - .take_while(|(a, b)| a == b) - .count(); - let prefix_bonus = if prefix_len > 0 { - let p = u32::try_from(prefix_len).unwrap_or(u32::MAX); - f64::from(p) / f64::from(max_len) * 0.1 - } else { - 0.0 - }; - - f64::from(lcs_u32) / f64::from(max_len) + prefix_bonus - } - - /// Compute the length of the longest common subsequence. - /// Compute the length of the longest common subsequence of two character slices. - /// - /// Uses an iterative dynamic-programming approach with only two rows - /// to keep memory usage O(min(a, b)). - fn lcs_length(a: &[char], b: &[char]) -> usize { - let mut prev = vec![0usize; b.len().saturating_add(1)]; - let mut curr = vec![0usize; b.len().saturating_add(1)]; - - for &a_ch in a { - for (j, &b_ch) in b.iter().enumerate() { - let j_idx = j.saturating_add(1); - *curr.get_mut(j_idx).unwrap_or(&mut 0) = if a_ch == b_ch { - prev.get(j_idx.saturating_sub(1)) - .copied() - .unwrap_or(0) - .saturating_add(1) - } else { - prev.get(j_idx) - .copied() - .unwrap_or(0) - .max(curr.get(j_idx.saturating_sub(1)).copied().unwrap_or(0)) - }; - } - std::mem::swap(&mut prev, &mut curr); - curr.fill(0); - } - - *prev.get(b.len()).unwrap_or(&0) - } - - /// Find the best matching tool name from a list, given a threshold. - /// - /// Returns the name with the highest similarity score that meets or - /// exceeds `threshold`. Returns `None` when no candidate scores high - /// enough. - fn find_best_match_inner<'a>( - requested: &str, - available: &[&'a str], - threshold: f64, - ) -> Option<(&'a str, f64)> { - let mut best: Option<(&'a str, f64)> = None; - for &name in available { - let score = Self::similarity(requested, name); - if score >= threshold { - match best { - Some((_, best_score)) if score <= best_score => {} - _ => best = Some((name, score)), - } - } - } - best - } - - /// Check if a result looks like a "tool not found" error. - /// - /// Only considers single [`Text`](ToolContent::Text) results whose - /// lowercased body contains `"not found"`. - /// [`Multipart`](ToolContent::Multipart) results and non-error results - /// always return `false`. - fn is_tool_not_found(result: &ToolDispatchResult) -> bool { - if !result.is_error { - return false; - } - let msg = match &result.output { - ToolContent::Text(t) => t.to_lowercase(), - ToolContent::Multipart(_) => return false, - }; - msg.contains("not found") - } -} - -impl ToolMiddleware for UnknownToolMiddleware { - fn name(&self) -> &'static str { - "unknown_tool" - } - - fn dispatch<'a>( - &'a self, - ctx: &'a mut ToolDispatchContext, - next: &'a ToolPipeline, - ) -> Pin + Send + 'a>> { - let tool_name = ctx.tool_name.clone(); - let registry_names = self.registry.tool_names(); - let threshold = self.suggestion_threshold; - - Box::pin(async move { - let mut result = next.dispatch(ctx).await; - - if Self::is_tool_not_found(&result) { - let available_refs: Vec<&str> = registry_names.iter().map(String::as_str).collect(); - - if let Some((suggestion, score)) = - Self::find_best_match_inner(&tool_name, &available_refs, threshold) - { - tracing::info!( - requested = %tool_name, - suggestion = %suggestion, - score = %score, - "suggesting alternative tool" - ); - // Append suggestion to the error message - if let ToolContent::Text(ref mut msg) = result.output { - *msg = format!("{msg}. Did you mean '{suggestion}'?"); - } - } - } - - result - }) - } -} - -// =================================================== -// OutputLimitMiddleware -// =================================================== - -/// Middleware that truncates tool output to a maximum character count. -/// -/// If the tool's text output exceeds the limit, it is truncated and -/// suffixed with a `[truncated]` marker. Non-text outputs ([`ToolContent::Multipart`]) -/// are passed through unchanged. -/// -/// This prevents runaway tools from flooding the conversation with -/// excessive output that would blow the context window. -/// -/// # Example -/// -/// ```rust,ignore -/// use loopctl::engine::middleware::OutputLimitMiddleware; -/// -/// let pipeline = ToolPipeline::builder() -/// .with(OutputLimitMiddleware::new(10_000)) -/// .core(registry) -/// .build()?; -/// ``` -pub struct OutputLimitMiddleware { - /// Maximum characters for text output. - max_chars: usize, -} - -impl OutputLimitMiddleware { - /// Create a new output-limiting middleware. - /// - /// `max_chars` is the maximum number of characters in the text - /// output. Outputs at or below this limit pass through unchanged. - #[must_use] - pub fn new(max_chars: usize) -> Self { - Self { max_chars } - } -} - -impl ToolMiddleware for OutputLimitMiddleware { - fn name(&self) -> &'static str { - "output_limit" - } - - fn dispatch<'a>( - &'a self, - ctx: &'a mut ToolDispatchContext, - next: &'a ToolPipeline, - ) -> Pin + Send + 'a>> { - let max_chars = self.max_chars; - Box::pin(async move { - let mut result = next.dispatch(ctx).await; - - if let ToolContent::Text(ref text) = result.output { - let char_count = text.chars().count(); - if char_count > max_chars { - let truncated: String = text.chars().take(max_chars).collect(); - result.output = ToolContent::Text(format!("{truncated}\n[truncated]")); - } - } - - result - }) - } -} - // =================================================== // Tests // =================================================== @@ -1123,12 +414,15 @@ impl ToolMiddleware for OutputLimitMiddleware { #[cfg(test)] mod tests { use super::*; + use crate::cancel::CancelSignal; + use crate::message::ToolContent; use crate::message::ToolContentPart; - use crate::tool::{Tool, ToolContext, ToolError, ToolOutput, ToolSchema}; + use crate::tool::{PermissionCheck, Tool, ToolContext, ToolError, ToolOutput, ToolSchema}; use serde_json::json; use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::{Duration, Instant}; // ================================================== // Test tools @@ -1364,7 +658,7 @@ mod tests { #[tokio::test] async fn test_permission_custom_check() { let pipeline = ToolPipeline::builder() - .with(PermissionMiddleware::with_check(|ctx| { + .with(PermissionMiddleware::from_context().with_check(|ctx| { if ctx.tool_name == "echo" { PermissionCheck::Allow } else { @@ -1559,9 +853,9 @@ mod tests { } } - // ================================================== + // =================================================== // Integration: full pipeline - // ================================================== + // =================================================== #[tokio::test] async fn test_full_pipeline_echo() { @@ -1598,10 +892,7 @@ mod tests { let pipeline = ToolPipeline::builder() .with(PermissionMiddleware::allow_all()) - .with(UnknownToolMiddleware::with_threshold( - Arc::clone(®istry), - 0.3, - )) + .with(UnknownToolMiddleware::new(Arc::clone(®istry)).with_threshold(0.3)) .core(registry) .build() .expect("valid"); diff --git a/src/middleware/output_limit.rs b/src/middleware/output_limit.rs new file mode 100644 index 0000000..9e778e4 --- /dev/null +++ b/src/middleware/output_limit.rs @@ -0,0 +1,67 @@ +//! Middleware that truncates tool output to a maximum character count. + +use super::{ToolDispatchContext, ToolDispatchResult, ToolMiddleware, ToolPipeline}; +use crate::message::ToolContent; +use std::future::Future; +use std::pin::Pin; + +/// Middleware that truncates tool output to a maximum character count. +/// +/// If the tool's text output exceeds the limit, it is truncated and +/// suffixed with a `[truncated]` marker. Non-text outputs ([`ToolContent::Multipart`]) +/// are passed through unchanged. +/// +/// This prevents runaway tools from flooding the conversation with +/// excessive output that would blow the context window. +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::middleware::OutputLimitMiddleware; +/// +/// let pipeline = ToolPipeline::builder() +/// .with(OutputLimitMiddleware::new(10_000)) +/// .core(registry) +/// .build()?; +/// ``` +pub struct OutputLimitMiddleware { + max_chars: usize, +} + +impl OutputLimitMiddleware { + /// Create a new output-limiting middleware. + /// + /// `max_chars` is the maximum number of characters in the text + /// output. Outputs at or below this limit pass through unchanged. + #[must_use] + pub fn new(max_chars: usize) -> Self { + Self { max_chars } + } +} + +impl ToolMiddleware for OutputLimitMiddleware { + fn name(&self) -> &'static str { + "output_limit" + } + + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + Send + 'a>> { + let max_chars = self.max_chars; + Box::pin(async move { + let mut result = next.dispatch(ctx).await; + + if let ToolContent::Text(ref text) = result.output { + let char_count = text.chars().count(); + if char_count > max_chars { + let truncated: String = text.chars().take(max_chars).collect(); + result.output = ToolContent::Text(format!("{truncated}\n[truncated]")); + } + } + + result + }) + } +} diff --git a/src/middleware/permission.rs b/src/middleware/permission.rs new file mode 100644 index 0000000..73845f7 --- /dev/null +++ b/src/middleware/permission.rs @@ -0,0 +1,190 @@ +//! Middleware that checks tool permissions before execution. + +use super::{ToolDispatchContext, ToolDispatchResult, ToolMiddleware, ToolPipeline}; +use crate::tool::PermissionCheck; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::Duration; + +/// Permission check function type. +pub type PermissionCheckFn = Arc PermissionCheck + Send + Sync>; + +/// Async resolver for [`PermissionCheck::Ask`]. +/// +/// Receives the prompt string and the tool name, returns `true` to allow +/// the tool call or `false` to deny it. Called by [`PermissionMiddleware`] +/// when the permission check resolves to [`PermissionCheck::Ask`]. +pub type AskResolverFn = + Arc Pin + Send>> + Send + Sync>; + +/// Middleware that checks tool permissions before execution. +/// +/// Inspects the [`PermissionCheck`] in the dispatch context: +/// +/// - `Allow` — passes through to the next layer. +/// - `Deny` — short-circuits with an error result. +/// - `Modify` — replaces `ctx.input` with the modified input, then proceeds. +/// - `Ask` — if an [`AskResolverFn`] is configured, calls it to prompt the +/// user; the tool proceeds on `true` or is denied on `false`. Without a +/// resolver, `Ask` is denied (headless mode). +/// +/// # Example +/// +/// ```rust,ignore +/// // Deny all by default +/// let mw = PermissionMiddleware::deny_all(); +/// +/// // Custom logic +/// let mw = PermissionMiddleware::with_check(|ctx| { +/// if ctx.tool_name == "safe_read" { +/// PermissionCheck::Allow +/// } else { +/// PermissionCheck::Deny +/// } +/// }); +/// ``` +pub struct PermissionMiddleware { + /// When `Some`, overrides [`ToolDispatchContext::permission`]. + check_fn: Option, + /// When `Some`, called to resolve [`PermissionCheck::Ask`] interactively. + ask_resolver: Option, +} + +impl PermissionMiddleware { + /// Create a permission middleware that denies all calls. + /// + /// Every tool call will be short-circuited with a permission-denied + /// error. Useful as a safety default in restricted environments. + #[must_use] + pub fn deny_all() -> Self { + Self { + check_fn: Some(Arc::new(|_| PermissionCheck::Deny { + reason: "blocked by policy".into(), + })), + ask_resolver: None, + } + } + + /// Create a permission middleware that allows all calls. + /// + /// No permission checks are performed — every tool call passes + /// through to the next layer. Equivalent to having no permission + /// middleware, but can be used for logging or metrics in permissive + /// environments. + #[must_use] + pub fn allow_all() -> Self { + Self { + check_fn: Some(Arc::new(|_| PermissionCheck::Allow)), + ask_resolver: None, + } + } + + /// Set a custom permission check function. + /// + /// The function receives a reference to the dispatch context and + /// returns the appropriate [`PermissionCheck`] for that call. + #[must_use] + pub fn with_check( + mut self, + f: impl Fn(&ToolDispatchContext) -> PermissionCheck + Send + Sync + 'static, + ) -> Self { + self.check_fn = Some(Arc::new(f)); + self + } + + /// Create a permission middleware that reads from the context. + /// + /// The middleware reads `ctx.permission` directly, without + /// applying any override. Use when the permission is set by the + /// framework or a prior middleware. + #[must_use] + pub fn from_context() -> Self { + Self { + check_fn: None, + ask_resolver: None, + } + } + + /// Attach an async resolver for [`PermissionCheck::Ask`]. + /// + /// When the permission check returns `Ask`, the resolver is called with + /// the prompt and tool name. The tool call proceeds if the resolver + /// returns `true`, and is denied if it returns `false`. + /// + /// Without a resolver, `Ask` is denied (headless mode). + #[must_use] + pub fn with_ask_resolver(mut self, resolver: AskResolverFn) -> Self { + self.ask_resolver = Some(resolver); + self + } + + fn resolve_permission(&self, ctx: &ToolDispatchContext) -> PermissionCheck { + match &self.check_fn { + Some(f) => f(ctx), + None => ctx.permission.clone(), + } + } +} + +impl ToolMiddleware for PermissionMiddleware { + fn name(&self) -> &'static str { + "permission" + } + + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + Send + 'a>> { + let permission = self.resolve_permission(ctx); + match permission { + PermissionCheck::Allow => next.dispatch(ctx), + PermissionCheck::Modify { modified_input } => { + ctx.input = modified_input; + next.dispatch(ctx) + } + PermissionCheck::Deny { reason } => Self::deny(ctx, &reason), + PermissionCheck::Ask { prompt } => match &self.ask_resolver { + Some(resolver) => { + let resolver = Arc::clone(resolver); + Box::pin(async move { + let tool_name = ctx.tool_name.clone(); + let approved = resolver(&prompt, &tool_name); + if approved.await { + next.dispatch(ctx).await + } else { + ToolDispatchResult::err( + &tool_name, + format!("Permission denied by user for tool '{tool_name}'"), + Duration::ZERO, + ) + } + }) + } + None => Self::deny(ctx, &format!("permission required: {prompt}")), + }, + } + } +} + +impl PermissionMiddleware { + /// Build a denied result with tracing. + fn deny<'a>( + ctx: &'a mut ToolDispatchContext, + reason: &str, + ) -> Pin + Send + 'a>> { + let tool_name = ctx.tool_name.clone(); + let reason = reason.to_string(); + tracing::warn!( + tool = %tool_name, + permission = %reason, + "tool call blocked by permission middleware" + ); + Box::pin(std::future::ready(ToolDispatchResult::err( + &tool_name, + format!("Permission {reason} for tool '{tool_name}'"), + Duration::ZERO, + ))) + } +} diff --git a/src/middleware/timeout.rs b/src/middleware/timeout.rs new file mode 100644 index 0000000..657c06d --- /dev/null +++ b/src/middleware/timeout.rs @@ -0,0 +1,158 @@ +//! Middleware that wraps tool execution in a timeout. + +use super::{ToolDispatchContext, ToolDispatchResult, ToolMiddleware, ToolPipeline}; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::Duration; + +/// Configuration for the [`TimeoutMiddleware`]. +#[derive(Debug, Clone)] +pub struct TimeoutConfig { + /// Per-tool execution deadline. + pub timeout: Duration, + /// Retry with exponential backoff up to [`max_retries`](TimeoutConfig::max_retries) times. + pub retry_on_timeout: bool, + /// Maximum number of retries (0 = no retry). Total attempts = `1 + max_retries`. + pub max_retries: u32, +} + +impl Default for TimeoutConfig { + fn default() -> Self { + Self { + timeout: Duration::from_secs(120), + retry_on_timeout: false, + max_retries: 0, + } + } +} + +/// Middleware that wraps tool execution in a timeout. +/// +/// If the tool execution exceeds the configured timeout, returns an +/// error result. Optionally retries once with a longer timeout. +/// Respects the [`CancelSignal`](crate::cancel::CancelSignal) via `tokio::select!` so that +/// cancellation is not blocked by a slow tool. +/// +/// # Example +/// +/// ```rust,ignore +/// let mw = TimeoutMiddleware::from_secs(120); +/// let mw = TimeoutMiddleware::new(TimeoutConfig { +/// timeout: Duration::from_secs(60), +/// retry_on_timeout: false, +/// max_retries: 0, +/// }); +/// ``` +pub struct TimeoutMiddleware { + config: TimeoutConfig, +} + +impl TimeoutMiddleware { + /// Create a timeout middleware with the given configuration. + /// + /// `config.timeout` controls the per-tool execution deadline. + /// When `config.retry_on_timeout` is `true`, a timed-out call is + /// retried up to `config.max_retries` additional times with an + /// increasing back-off. + /// + /// For simpler construction see [`from_secs`](Self::from_secs) or + /// [`none`](Self::none). + #[must_use] + pub fn new(config: TimeoutConfig) -> Self { + Self { config } + } + + /// Create a timeout middleware with a fixed timeout in seconds. + /// + /// Uses default retry settings (one retry with double timeout). + #[must_use] + pub fn from_secs(secs: u64) -> Self { + Self { + config: TimeoutConfig { + timeout: Duration::from_secs(secs), + ..TimeoutConfig::default() + }, + } + } + + /// Create a timeout middleware with no timeout (pass-through). + /// + /// Useful for testing or when timeouts are handled elsewhere. + #[must_use] + pub fn none() -> Self { + Self { + config: TimeoutConfig { + timeout: Duration::MAX, + retry_on_timeout: false, + max_retries: 0, + }, + } + } +} + +impl ToolMiddleware for TimeoutMiddleware { + fn name(&self) -> &'static str { + "timeout" + } + + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + Send + 'a>> { + let config = self.config.clone(); + let tool_name = ctx.tool_name.clone(); + let cancel = Arc::clone(&ctx.cancel); + + Box::pin(async move { + let mut attempt = 0u32; + let mut current_timeout = config.timeout; + + loop { + let result_future = next.dispatch(ctx); + let attempt_for_log = attempt; + + tokio::select! { + result = tokio::time::timeout(current_timeout, result_future) => { + if let Ok(dispatch_result) = result { + return dispatch_result; + } + attempt = attempt.saturating_add(1); + if config.retry_on_timeout && attempt <= config.max_retries { + tracing::warn!( + tool = %tool_name, + attempt = attempt_for_log, + timeout_secs = current_timeout.as_secs(), + "tool execution timed out, retrying" + ); + current_timeout = current_timeout.saturating_mul(2); + continue; + } + tracing::error!( + tool = %tool_name, + timeout_secs = current_timeout.as_secs(), + "tool execution timed out" + ); + return ToolDispatchResult::err( + &tool_name, + format!( + "Tool '{}' timed out after {}s", + tool_name, + current_timeout.as_secs() + ), + current_timeout, + ); + } + () = cancel.notified() => { + return ToolDispatchResult::err( + &tool_name, + format!("Tool '{tool_name}' cancelled"), + Duration::ZERO, + ); + } + } + } + }) + } +} diff --git a/src/middleware/tool_call.rs b/src/middleware/tool_call.rs new file mode 100644 index 0000000..c0fbaa6 --- /dev/null +++ b/src/middleware/tool_call.rs @@ -0,0 +1,76 @@ +//! The innermost middleware that performs the actual tool invocation. + +use super::{ToolDispatchContext, ToolDispatchResult}; +use crate::error::LoopError; +use crate::tool::ToolRegistry; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::Instant; + +/// The innermost middleware that performs the actual tool invocation. +/// +/// Looks up the tool by name in the [`ToolRegistry`], calls +/// [`crate::tool::Tool::call()`], and converts the result into a +/// [`ToolDispatchResult`]. If the tool is not found, produces a +/// soft error result (not a hard error) so the model can recover. +/// +/// This middleware is automatically created by the pipeline builder +/// and always occupies the innermost position in the chain. +pub struct ToolCallMiddleware { + pub(super) registry: Arc, +} + +impl ToolCallMiddleware { + pub(super) const NAME: &str = "tool_call"; + + /// Create a new core dispatch wrapping the given registry. + #[must_use] + pub fn new(registry: Arc) -> Self { + Self { registry } + } + + /// Execute the tool call — the terminal dispatch. + /// + /// Looks up the tool by name in the registry, calls `Tool::call()`, + /// and converts the result. There is no `next` parameter because + /// there is nothing to chain to. + pub(super) fn dispatch( + &self, + ctx: &mut ToolDispatchContext, + ) -> Pin + Send + '_>> { + let tool_name = ctx.tool_name.clone(); + let input = ctx.input.clone(); + let tool_ctx = ctx.tool_context.clone(); + let registry = Arc::clone(&self.registry); + let cancel = Arc::clone(&ctx.cancel); + let call_id = ctx.call_id.clone(); + + Box::pin(async move { + let start = Instant::now(); + let Some(tool) = registry.get(&tool_name) else { + let available: Vec = registry.tool_names(); + let available_refs: Vec<&str> = available.iter().map(String::as_str).collect(); + let error = LoopError::tool_not_found(&tool_name, &available_refs); + return ToolDispatchResult::err(&tool_name, error.to_string(), start.elapsed()) + .with_call_id(&call_id); + }; + + let call_result = tokio::select! { + r = tool.call(input, &tool_ctx) => r, + () = cancel.notified() => { + return ToolDispatchResult::err( + &tool_name, + format!("Tool '{tool_name}' cancelled"), + start.elapsed(), + ) + .with_call_id(&call_id); + } + }; + + let duration = start.elapsed(); + ToolDispatchResult::from_result(&tool_name, call_result, duration) + .with_call_id(&call_id) + }) + } +} diff --git a/src/middleware/unknown_tool.rs b/src/middleware/unknown_tool.rs new file mode 100644 index 0000000..2b216bd --- /dev/null +++ b/src/middleware/unknown_tool.rs @@ -0,0 +1,417 @@ +//! Middleware that suggests alternatives when a tool is not found. + +use super::{ToolDispatchContext, ToolDispatchResult, ToolMiddleware, ToolPipeline}; +use crate::message::ToolContent; +use crate::tool::ToolRegistry; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +/// Middleware that suggests alternatives when a tool is not found. +/// +/// This middleware wraps the core dispatch. When the inner dispatch +/// produces a "tool not found" error, this middleware intercepts it, +/// computes a string-similarity score against all registered tools, +/// and appends a suggestion to the error message. +/// +/// # Similarity Metric +/// +/// Uses a simple normalized longest-common-subsequence ratio, which is +/// fast and effective for the common case of minor typos (e.g. +/// `"bash"` → `"basj"`). +/// +/// # Example +/// +/// ```rust,ignore +/// let mw = UnknownToolMiddleware::new(); +/// // If tool "basj" is not found, error message will say: +/// // "Tool 'basj' not found. Did you mean 'bash'?" +/// ``` +pub struct UnknownToolMiddleware { + /// Tool registry used to enumerate available tool names. + registry: Arc, + /// Defaults to `0.4`. + suggestion_threshold: f64, +} + +impl UnknownToolMiddleware { + /// Create a new unknown-tool middleware with default settings. + /// + /// Uses a [`suggestion_threshold`](Self::with_threshold) of `0.4`, + /// which balances catching common typos against false-positive + /// suggestions. + /// + /// For a custom threshold see [`with_threshold`](Self::with_threshold). + #[must_use] + pub fn new(registry: Arc) -> Self { + Self { + registry, + suggestion_threshold: 0.4, + } + } + + /// Create with a custom similarity threshold. + /// + /// Lower values produce more suggestions (more false positives). + /// Higher values require closer matches. + #[must_use] + pub fn with_threshold(mut self, threshold: f64) -> Self { + self.suggestion_threshold = threshold.clamp(0.0, 1.0); + self + } + + /// Compute similarity between two strings using a normalized + /// longest-common-subsequence approach. + /// + /// Returns a value between 0.0 (completely different) and 1.0 + /// (identical). + #[must_use] + pub fn similarity(a: &str, b: &str) -> f64 { + if a.is_empty() && b.is_empty() { + return 1.0; + } + if a.is_empty() || b.is_empty() { + return 0.0; + } + + let a_lower = a.to_lowercase(); + let b_lower = b.to_lowercase(); + + if a_lower == b_lower { + return 1.0; + } + + // Use longest common subsequence length as similarity metric + let a_chars: Vec = a_lower.chars().collect(); + let b_chars: Vec = b_lower.chars().collect(); + let lcs_len = Self::lcs_length(&a_chars, &b_chars); + + // Tool names are short strings, so u32 is sufficient and avoids + // usize→f64 precision loss on 64-bit targets. + let max_len = u32::try_from(a_chars.len().max(b_chars.len())).unwrap_or(u32::MAX); + let lcs_u32 = u32::try_from(lcs_len).unwrap_or(u32::MAX); + + // Check for prefix match bonus + let prefix_len = a_chars + .iter() + .zip(b_chars.iter()) + .take_while(|(a, b)| a == b) + .count(); + let prefix_bonus = if prefix_len > 0 { + let p = u32::try_from(prefix_len).unwrap_or(u32::MAX); + f64::from(p) / f64::from(max_len) * 0.1 + } else { + 0.0 + }; + + f64::from(lcs_u32) / f64::from(max_len) + prefix_bonus + } + + /// Compute the length of the longest common subsequence of two character slices. + /// + /// Uses an iterative dynamic-programming approach with only two rows + /// to keep memory usage O(min(a, b)). + fn lcs_length(a: &[char], b: &[char]) -> usize { + let mut prev = vec![0usize; b.len().saturating_add(1)]; + let mut curr = vec![0usize; b.len().saturating_add(1)]; + + for &a_ch in a { + for (j, &b_ch) in b.iter().enumerate() { + let j_idx = j.saturating_add(1); + *curr.get_mut(j_idx).unwrap_or(&mut 0) = if a_ch == b_ch { + prev.get(j_idx.saturating_sub(1)) + .copied() + .unwrap_or(0) + .saturating_add(1) + } else { + prev.get(j_idx) + .copied() + .unwrap_or(0) + .max(curr.get(j_idx.saturating_sub(1)).copied().unwrap_or(0)) + }; + } + std::mem::swap(&mut prev, &mut curr); + curr.fill(0); + } + + *prev.get(b.len()).unwrap_or(&0) + } + + /// Find the best matching tool name from a list, given a threshold. + /// + /// Returns the name with the highest similarity score that meets or + /// exceeds `threshold`. Returns `None` when no candidate scores high + /// enough. + fn find_best_match_inner<'a>( + requested: &str, + available: &[&'a str], + threshold: f64, + ) -> Option<(&'a str, f64)> { + let mut best: Option<(&'a str, f64)> = None; + for &name in available { + let score = Self::similarity(requested, name); + if score >= threshold { + match best { + Some((_, best_score)) if score <= best_score => {} + _ => best = Some((name, score)), + } + } + } + best + } + + /// Check if a result looks like a "tool not found" error. + /// + /// Only considers single [`Text`](ToolContent::Text) results whose + /// lowercased body contains `"not found"`. + /// [`Multipart`](ToolContent::Multipart) results and non-error results + /// always return `false`. + fn is_tool_not_found(result: &ToolDispatchResult) -> bool { + if !result.is_error { + return false; + } + let msg = match &result.output { + ToolContent::Text(t) => t.to_lowercase(), + ToolContent::Multipart(_) => return false, + }; + msg.contains("not found") + } +} + +impl ToolMiddleware for UnknownToolMiddleware { + fn name(&self) -> &'static str { + "unknown_tool" + } + + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + Send + 'a>> { + let tool_name = ctx.tool_name.clone(); + let registry_names = self.registry.tool_names(); + let threshold = self.suggestion_threshold; + + Box::pin(async move { + let mut result = next.dispatch(ctx).await; + + if Self::is_tool_not_found(&result) { + let available_refs: Vec<&str> = registry_names.iter().map(String::as_str).collect(); + + if let Some((suggestion, score)) = + Self::find_best_match_inner(&tool_name, &available_refs, threshold) + { + tracing::info!( + requested = %tool_name, + suggestion = %suggestion, + score = %score, + "suggesting alternative tool" + ); + // Append suggestion to the error message + if let ToolContent::Text(ref mut msg) = result.output { + *msg = format!("{msg}. Did you mean '{suggestion}'?"); + } + } + } + + result + }) + } +} + +#[cfg(test)] +mod tests { + use super::UnknownToolMiddleware; + + #[test] + fn lcs_length_identical_strings() { + let a: Vec = "hello".chars().collect(); + let b: Vec = "hello".chars().collect(); + assert_eq!(UnknownToolMiddleware::lcs_length(&a, &b), 5); + } + + #[test] + fn lcs_length_completely_different() { + let a: Vec = "abc".chars().collect(); + let b: Vec = "xyz".chars().collect(); + assert_eq!(UnknownToolMiddleware::lcs_length(&a, &b), 0); + } + + #[test] + fn lcs_length_empty_first() { + let b: Vec = "abc".chars().collect(); + assert_eq!(UnknownToolMiddleware::lcs_length(&[], &b), 0); + } + + #[test] + fn lcs_length_empty_second() { + let a: Vec = "abc".chars().collect(); + assert_eq!(UnknownToolMiddleware::lcs_length(&a, &[]), 0); + } + + #[test] + fn lcs_length_both_empty() { + assert_eq!(UnknownToolMiddleware::lcs_length(&[], &[]), 0); + } + + #[test] + fn lcs_length_subsequence_not_substring() { + // "ace" is a subsequence of "abcde" but not a substring + let a: Vec = "ace".chars().collect(); + let b: Vec = "abcde".chars().collect(); + assert_eq!(UnknownToolMiddleware::lcs_length(&a, &b), 3); + } + + #[test] + fn lcs_length_partial_overlap() { + let a: Vec = "kitten".chars().collect(); + let b: Vec = "sitting".chars().collect(); + // LCS is "ittn" → length 4 + assert_eq!(UnknownToolMiddleware::lcs_length(&a, &b), 4); + } + + #[test] + fn lcs_length_order_independent() { + let a: Vec = "sitting".chars().collect(); + let b: Vec = "kitten".chars().collect(); + assert_eq!( + UnknownToolMiddleware::lcs_length(&a, &b), + UnknownToolMiddleware::lcs_length(&b, &a), + ); + } + + #[test] + fn lcs_length_single_character_common() { + let a: Vec = "a".chars().collect(); + let b: Vec = "a".chars().collect(); + assert_eq!(UnknownToolMiddleware::lcs_length(&a, &b), 1); + } + + #[test] + fn lcs_length_repeated_chars() { + let a: Vec = "aaa".chars().collect(); + let b: Vec = "aa".chars().collect(); + assert_eq!(UnknownToolMiddleware::lcs_length(&a, &b), 2); + } + + // ---- find_best_match_inner ---- + + #[test] + fn find_best_match_exact_hit() { + let available = ["read_file", "write_file", "list_dir"]; + let (suggestion, score) = + UnknownToolMiddleware::find_best_match_inner("read_file", &available, 0.5).unwrap(); + assert_eq!(suggestion, "read_file"); + assert_eq!(score, 1.0); + } + + #[test] + fn find_best_match_close_typo() { + let available = ["read_file", "write_file", "list_dir"]; + let (suggestion, _) = + UnknownToolMiddleware::find_best_match_inner("read_fil", &available, 0.5).unwrap(); + assert_eq!(suggestion, "read_file"); + } + + #[test] + fn find_best_match_nothing_above_threshold() { + let available = ["read_file", "write_file", "list_dir"]; + assert!(UnknownToolMiddleware::find_best_match_inner("xyz", &available, 0.5).is_none()); + } + + #[test] + fn find_best_match_empty_available() { + let available: [&str; 0] = []; + assert!( + UnknownToolMiddleware::find_best_match_inner("read_file", &available, 0.0).is_none() + ); + } + + #[test] + fn find_best_match_empty_requested() { + let available = ["read_file", "write_file"]; + // An empty request has zero overlap with everything. + assert!( + UnknownToolMiddleware::find_best_match_inner("", &available, 0.0).is_none() + || UnknownToolMiddleware::find_best_match_inner("", &available, 0.0) + .unwrap() + .1 + == 0.0 + ); + } + + #[test] + fn find_best_match_picks_highest_score() { + let available = ["read", "read_file", "rea"]; + // "read_file" requested → both "read" and "rea" are substrings, + // but "read" has a higher LCS ratio. + let (suggestion, _) = + UnknownToolMiddleware::find_best_match_inner("read_file", &available, 0.3).unwrap(); + assert_eq!(suggestion, "read_file"); + } + + #[test] + fn find_best_match_threshold_zero_returns_best() { + let available = ["abc", "xyz"]; + // With threshold 0.0 even a zero-score match could pass, but "abc" + // shares at least one char with "a" so it wins. + let result = UnknownToolMiddleware::find_best_match_inner("a", &available, 0.0); + assert!(result.is_some()); + } + + #[test] + fn find_best_match_threshold_one_requires_exact() { + let available = ["read_file", "write_file"]; + // Only an exact match can satisfy threshold 1.0. + let (suggestion, score) = + UnknownToolMiddleware::find_best_match_inner("read_file", &available, 1.0).unwrap(); + assert_eq!(suggestion, "read_file"); + assert_eq!(score, 1.0); + + assert!( + UnknownToolMiddleware::find_best_match_inner("read_fil", &available, 1.0).is_none() + ); + } + + // ---- is_tool_not_found ---- + + use crate::message::ToolContent; + use crate::tool::ToolDispatchResult; + use std::time::Duration; + + #[test] + fn is_tool_not_found_error_with_phrase() { + let result = ToolDispatchResult::err("x", "Tool 'foo' not found".into(), Duration::ZERO); + assert!(UnknownToolMiddleware::is_tool_not_found(&result)); + } + + #[test] + fn is_tool_not_found_case_insensitive() { + let result = ToolDispatchResult::err("x", "TOOL NOT FOUND".into(), Duration::ZERO); + assert!(UnknownToolMiddleware::is_tool_not_found(&result)); + } + + #[test] + fn is_tool_not_found_success_result() { + let result = ToolDispatchResult::ok("x", "done".into(), Duration::ZERO); + assert!(!UnknownToolMiddleware::is_tool_not_found(&result)); + } + + #[test] + fn is_tool_not_found_error_without_phrase() { + let result = ToolDispatchResult::err("x", "permission denied".into(), Duration::ZERO); + assert!(!UnknownToolMiddleware::is_tool_not_found(&result)); + } + + #[test] + fn is_tool_not_found_multipart_error() { + let result = ToolDispatchResult { + tool_call_id: String::new(), + output: ToolContent::Multipart(vec![]), + is_error: true, + duration: Duration::ZERO, + resolved_tool_name: "x".into(), + }; + assert!(!UnknownToolMiddleware::is_tool_not_found(&result)); + } +} diff --git a/src/core/observer.rs b/src/observer.rs similarity index 68% rename from src/core/observer.rs rename to src/observer.rs index bf8a7ca..75db134 100644 --- a/src/core/observer.rs +++ b/src/observer.rs @@ -13,7 +13,7 @@ //! - [`StreamContext`] / [`StreamFailureContext`] — stream success/failure //! - [`ResponseContext`] — model response text and usage //! - [`ToolPreContext`] / [`ToolPostContext`] — tool dispatch lifecycle -//! - [`CompactionContext`] — context window compaction +//! - [`CompactedContext`] — context window compaction //! - [`FallbackContext`] — model fallback event //! - [`LoopDetectedContext`] — loop detection event //! - [`ConvergenceDetectedContext`] — convergence detection event @@ -21,7 +21,7 @@ //! # Example //! //! ```rust,ignore -//! use loopctl::core::observer::{LoopObserver, SessionStartContext}; +//! use loopctl::observer::{LoopObserver, SessionStartContext}; //! //! struct MetricsObserver; //! @@ -36,160 +36,20 @@ use std::sync::Arc; -// ================================================== -// Context structs -// ================================================== - -/// Context for [`LoopObserver::on_session_start`]. -#[derive(Debug, Clone)] -pub struct SessionStartContext { - /// Unique session identifier. - pub session_id: uuid::Uuid, -} - -/// Context for [`LoopObserver::on_session_end`]. -#[derive(Debug, Clone)] -pub struct SessionEndContext { - /// Whether the session completed successfully. - pub success: bool, - /// Error description, if the session ended due to an error. - pub error: Option, - /// Total turns completed during the session. - pub total_turns: usize, - /// Total session duration in milliseconds. - pub duration_ms: u64, -} - -/// Context for [`LoopObserver::on_turn_start`]. -#[derive(Debug, Clone)] -pub struct TurnStartContext { - /// Turn number (0-indexed). - pub turn: usize, - /// The user query that initiated this turn. - pub query: String, -} - -/// Context for [`LoopObserver::on_turn_end`]. -#[derive(Debug, Clone)] -pub struct TurnEndContext { - /// Turn number. - pub turn: usize, - /// Whether the turn completed successfully. - pub success: bool, - /// Error description, if the turn failed. - pub error: Option, - /// Wall-clock duration of the turn in milliseconds. - pub duration_ms: u64, - /// Input tokens consumed this turn. - pub input_tokens: u64, - /// Output tokens generated this turn. - pub output_tokens: u64, -} - -/// Context for [`LoopObserver::on_stream_success`]. -#[derive(Debug, Clone)] -pub struct StreamContext { - /// Turn number. - pub turn: usize, - /// Model that was streamed. - pub model: String, - /// Input tokens consumed. - pub input_tokens: u64, - /// Output tokens generated. - pub output_tokens: u64, -} - -/// Context for [`LoopObserver::on_stream_failure`]. -#[derive(Debug, Clone)] -pub struct StreamFailureContext { - /// Turn number. - pub turn: usize, - /// Model that failed. - pub model: String, - /// The error that occurred. - pub error: crate::core::AgentError, -} - -/// Context for [`LoopObserver::on_response`]. -#[derive(Debug, Clone)] -pub struct ResponseContext { - /// Turn number. - pub turn: usize, - /// The model's text response. - pub text: String, - /// Token usage for this turn, if available. - pub usage: Option, -} - -/// Context for [`LoopObserver::on_tool_pre`]. -#[derive(Debug, Clone)] -pub struct ToolPreContext { - /// Turn number. - pub turn: usize, - /// Tool name. - pub tool: String, - /// Tool call ID from the API response. - pub tool_call_id: String, -} - -/// Context for [`LoopObserver::on_tool_post`]. -#[derive(Debug, Clone)] -pub struct ToolPostContext { - /// Turn number. - pub turn: usize, - /// Tool name. - pub tool: String, - /// Deterministic hash of the tool output, if available. - pub result_hash: Option, - /// Whether the tool returned an error. - pub is_error: bool, - /// Wall-clock execution duration. - pub duration: std::time::Duration, -} - -/// Context for [`LoopObserver::on_compaction`]. -#[derive(Debug, Clone)] -pub struct CompactionContext { - /// Message count before compaction. - pub messages_before: usize, - /// Message count after compaction. - pub messages_after: usize, - /// Estimated tokens saved by compaction. - pub tokens_saved: u64, -} - -/// Context for [`LoopObserver::on_fallback`]. -#[derive(Debug, Clone)] -pub struct FallbackContext { - /// Model that failed. - pub from: String, - /// Replacement model. - pub to: String, -} - -/// Context for [`LoopObserver::on_loop_detected`]. -#[derive(Debug, Clone)] -pub struct LoopDetectedContext { - /// Description of the repeating tool pattern. - pub pattern: String, - /// Number of times the pattern was observed. - pub repetitions: usize, -} - -/// Context for [`LoopObserver::on_convergence_detected`]. -#[derive(Debug, Clone)] -pub struct ConvergenceDetectedContext { - /// Configured action to take (e.g. `"stop"`, `"warn"`, `"compact"`). - pub action: String, -} +pub mod context; +pub use context::{ + CompactedContext, ConvergenceDetectedContext, FallbackContext, LoopDetectedContext, + ResponseContext, SessionEndContext, SessionStartContext, StreamContext, StreamFailureContext, + ToolPostContext, ToolPreContext, TurnEndContext, TurnStartContext, +}; // ================================================== // LoopObserver Trait // ================================================== /// A notification observer that receives typed callbacks at agent loop lifecycle points. /// -/// Observers are registered via `BareLoop` or `ManagerBundle` and called at each +/// Observers are registered via [`LoopRuntime`](crate::runtime::LoopRuntime) and called at each /// lifecycle point in registration order. All methods are **notification-only** — they /// return `()`. Use the [hook system](crate::hooks) if you need to control /// flow (block/allow actions). @@ -197,48 +57,92 @@ pub struct ConvergenceDetectedContext { /// All methods have default no-op implementations. Override only the callbacks you need. pub trait LoopObserver: Send + Sync { /// Human-readable name for diagnostics and logging. + /// + /// Returned in error messages and telemetry to identify which observer + /// produced a side-effect. fn name(&self) -> &str; /// Called when an agent session begins. + /// + /// Fired once per session, before the first turn starts. fn on_session_start(&self, _ctx: &SessionStartContext) {} /// Called when an agent session ends. + /// + /// Fired after the last turn completes or when a fatal error stops the loop. + /// Check [`SessionEndContext::success`] to distinguish normal exit from failure. fn on_session_end(&self, _ctx: &SessionEndContext) {} /// Called at the start of processing a turn. + /// + /// Fired before the model is called for this turn. fn on_turn_start(&self, _ctx: &TurnStartContext) {} /// Called when a turn completes. + /// + /// Fired after tool dispatch and any compaction has finished. + /// Check [`TurnEndContext::success`] to detect turn-level failures. fn on_turn_end(&self, _ctx: &TurnEndContext) {} /// Called after the model streams a response successfully. + /// + /// Provides token counts for the completed streaming request. + /// Not fired when the stream fails — see [`on_stream_failure`](Self::on_stream_failure). fn on_stream_success(&self, _ctx: &StreamContext) {} /// Called when the API stream fails. + /// + /// Fired on network errors, API errors, or stream interruptions. + /// The loop may retry or fall back to another model after this notification. fn on_stream_failure(&self, _ctx: &StreamFailureContext) {} /// Called after extracting the model's text response. + /// + /// Contains the concatenated assistant text and optional token usage. + /// Tool-call content is excluded; use [`on_tool_post`](Self::on_tool_post) + /// for tool results. fn on_response(&self, _ctx: &ResponseContext) {} - /// Called before a tool is dispatched (notification-only). + /// Called before a tool is dispatched. + /// + /// Notification-only — cannot block or modify the tool call. + /// Use the [hook system](crate::hooks) for flow control. fn on_tool_pre(&self, _ctx: &ToolPreContext) {} /// Called after a tool completes execution. + /// + /// Reports whether the tool errored and includes a hash of the result + /// for loop-detection correlation. fn on_tool_post(&self, _ctx: &ToolPostContext) {} /// Called after conversation compaction. - fn on_compaction(&self, _ctx: &CompactionContext) {} + /// + /// Reports token counts before and after compaction. Only fired when + /// compaction actually occurred — not on no-action passes. + fn on_compaction(&self, _ctx: &CompactedContext) {} /// Called when a model fallback is triggered. + /// + /// Indicates that the primary model failed and a fallback model + /// was selected for subsequent requests. fn on_fallback(&self, _ctx: &FallbackContext) {} /// Called when a loop is detected in tool operations. + /// + /// Fired when the same tool operation produces the same result + /// repeatedly, exceeding the configured threshold. fn on_loop_detected(&self, _ctx: &LoopDetectedContext) {} /// Called when response convergence is detected. + /// + /// Fired when consecutive model responses become sufficiently similar + /// as determined by the convergence detection policy. fn on_convergence_detected(&self, _ctx: &ConvergenceDetectedContext) {} /// Reset observer state for a new session. + /// + /// Called before [`on_session_start`](Self::on_session_start) to allow + /// observers to clear per-session accumulators. fn reset(&self) {} } @@ -249,7 +153,7 @@ pub trait LoopObserver: Send + Sync { /// Holds registered observers and dispatches notifications to each. /// /// Observers run in registration order. All observers are always notified — -/// there is no short-circuiting (that's the [hook system](crate::hooks)'s job). +/// there is no short-circuiting (use the [hook system](crate::hooks) for flow control). /// /// An empty host (no observers registered) is effectively zero-cost: /// each notification call iterates an empty `Vec`. @@ -260,12 +164,17 @@ pub struct ObserverHost { impl ObserverHost { /// Create an empty observer host. + /// + /// Equivalent to [`ObserverHost::default`] but more explicit. #[must_use] pub fn new() -> Self { Self::default() } - /// Register an observer. Called in registration order at each notification point. + /// Register an observer. + /// + /// Observers are called in registration order at each lifecycle point. + /// Registering the same observer twice will result in duplicate notifications. pub fn register(&mut self, observer: Arc) { self.observers.push(observer); } @@ -283,6 +192,9 @@ impl ObserverHost { } /// Reset all observers for a new session. + /// + /// Calls [`LoopObserver::reset`] on every registered observer, + /// allowing them to clear per-session accumulators. pub fn reset_all(&self) { for obs in &self.observers { obs.reset(); @@ -290,6 +202,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_session_start`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_session_start(&self, ctx: &SessionStartContext) { for obs in &self.observers { obs.on_session_start(ctx); @@ -297,6 +211,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_session_end`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_session_end(&self, ctx: &SessionEndContext) { for obs in &self.observers { obs.on_session_end(ctx); @@ -304,6 +220,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_turn_start`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_turn_start(&self, ctx: &TurnStartContext) { for obs in &self.observers { obs.on_turn_start(ctx); @@ -311,6 +229,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_turn_end`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_turn_end(&self, ctx: &TurnEndContext) { for obs in &self.observers { obs.on_turn_end(ctx); @@ -318,6 +238,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_stream_success`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_stream_success(&self, ctx: &StreamContext) { for obs in &self.observers { obs.on_stream_success(ctx); @@ -325,6 +247,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_stream_failure`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_stream_failure(&self, ctx: &StreamFailureContext) { for obs in &self.observers { obs.on_stream_failure(ctx); @@ -332,6 +256,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_response`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_response(&self, ctx: &ResponseContext) { for obs in &self.observers { obs.on_response(ctx); @@ -339,6 +265,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_tool_pre`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_tool_pre(&self, ctx: &ToolPreContext) { for obs in &self.observers { obs.on_tool_pre(ctx); @@ -346,6 +274,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_tool_post`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_tool_post(&self, ctx: &ToolPostContext) { for obs in &self.observers { obs.on_tool_post(ctx); @@ -353,13 +283,17 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_compaction`] to all observers. - pub fn on_compaction(&self, ctx: &CompactionContext) { + /// + /// Iterates registered observers in registration order. + pub fn on_compaction(&self, ctx: &CompactedContext) { for obs in &self.observers { obs.on_compaction(ctx); } } /// Dispatch [`LoopObserver::on_fallback`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_fallback(&self, ctx: &FallbackContext) { for obs in &self.observers { obs.on_fallback(ctx); @@ -367,6 +301,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_loop_detected`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_loop_detected(&self, ctx: &LoopDetectedContext) { for obs in &self.observers { obs.on_loop_detected(ctx); @@ -374,6 +310,8 @@ impl ObserverHost { } /// Dispatch [`LoopObserver::on_convergence_detected`] to all observers. + /// + /// Iterates registered observers in registration order. pub fn on_convergence_detected(&self, ctx: &ConvergenceDetectedContext) { for obs in &self.observers { obs.on_convergence_detected(ctx); diff --git a/src/observer/context.rs b/src/observer/context.rs new file mode 100644 index 0000000..39ac2c1 --- /dev/null +++ b/src/observer/context.rs @@ -0,0 +1,317 @@ +//! Typed context structs for [`LoopObserver`](crate::observer::LoopObserver) callbacks. +//! +//! Each struct carries the relevant fields for a specific lifecycle point. +//! Observers receive shared references (`&Context`) — the structs are +//! notification-only data carriers. + +// ================================================== +// Context structs +// ================================================== + +/// Context for [`LoopObserver::on_session_start`](crate::observer::LoopObserver::on_session_start). +/// +/// Carries the session identifier so observers can correlate +/// lifecycle events with a specific agent run. +#[derive(Debug, Clone)] +pub struct SessionStartContext { + /// Unique session identifier. + /// + /// Correlates all lifecycle events belonging to the same agent run. + pub session_id: uuid::Uuid, +} + +/// Context for [`LoopObserver::on_session_end`](crate::observer::LoopObserver::on_session_end). +/// +/// Captures the session's completion status, optional error description, +/// total turns executed, and wall-clock duration in milliseconds. +#[derive(Debug, Clone)] +pub struct SessionEndContext { + /// Whether the session completed successfully. + /// + /// `true` when the loop exited normally, `false` on error or cancellation. + pub success: bool, + + /// Error description, if the session ended due to an error. + /// + /// `None` when [`success`](Self::success) is `true`. + pub error: Option, + + /// Total turns completed during the session. + /// + /// Counts only turns that finished; an in-flight turn at the time of + /// a fatal error is not included. + pub total_turns: usize, + + /// Total session duration in milliseconds. + /// + /// Measured wall-clock from [`on_session_start`](crate::observer::LoopObserver::on_session_start) + /// to [`on_session_end`](crate::observer::LoopObserver::on_session_end). + pub duration_ms: u64, +} + +/// Context for [`LoopObserver::on_turn_start`](crate::observer::LoopObserver::on_turn_start). +/// +/// Provides the turn number and the user query that initiated it. +#[derive(Debug, Clone)] +pub struct TurnStartContext { + /// Turn number (0-indexed). + /// + /// Monotonically increasing within a session; resets on session restart. + pub turn: usize, + + /// The user query that initiated this turn. + /// + /// Contains the full text of the latest user message added to the + /// conversation before the turn began. + pub query: String, +} + +/// Context for [`LoopObserver::on_turn_end`](crate::observer::LoopObserver::on_turn_end). +/// +/// Reports whether the turn succeeded, any error, its duration, +/// and the token counts consumed during the turn. +#[derive(Debug, Clone)] +pub struct TurnEndContext { + /// Turn number. + /// + /// Matches the value passed to the corresponding + /// [`on_turn_start`](crate::observer::LoopObserver::on_turn_start). + pub turn: usize, + + /// Whether the turn completed successfully. + /// + /// `false` if the turn was interrupted by an error or cancellation. + pub success: bool, + + /// Error description, if the turn failed. + /// + /// `None` when [`success`](Self::success) is `true`. + pub error: Option, + + /// Wall-clock duration of the turn in milliseconds. + /// + /// Measured from [`on_turn_start`](crate::observer::LoopObserver::on_turn_start) + /// to [`on_turn_end`](crate::observer::LoopObserver::on_turn_end). + pub duration_ms: u64, + + /// Input tokens consumed this turn. + /// + /// Sum of tokens in the prompt sent to the model. + pub input_tokens: u64, + + /// Output tokens generated this turn. + /// + /// Sum of tokens across all assistant responses in the turn, + /// including intermediate tool-call rounds. + pub output_tokens: u64, +} + +/// Context for [`LoopObserver::on_stream_success`](crate::observer::LoopObserver::on_stream_success). +/// +/// Provides the model name and input/output token counts for +/// a successful streaming response. +#[derive(Debug, Clone)] +pub struct StreamContext { + /// Turn number. + /// + /// Identifies which turn this stream response belongs to. + pub turn: usize, + + /// Model that was streamed. + /// + /// The model identifier used for this request, which may differ from + /// the session default when model fallback occurred. + pub model: String, + + /// Input tokens consumed. + /// + /// Tokens in the prompt sent to the model for this request. + pub input_tokens: u64, + + /// Output tokens generated. + /// + /// Tokens in the model's streamed response. + pub output_tokens: u64, +} + +/// Context for [`LoopObserver::on_stream_failure`](crate::observer::LoopObserver::on_stream_failure). +/// +/// Carries the model name and the [`LoopError`](crate::error::LoopError) +/// that caused the streaming failure. +#[derive(Debug, Clone)] +pub struct StreamFailureContext { + /// Turn number. + /// + /// Identifies which turn this failure occurred in. + pub turn: usize, + + /// Model that failed. + /// + /// The model identifier used for the failed request. + pub model: String, + + /// The error that occurred. + /// + /// See [`LoopError`](crate::error::LoopError) for the full set of + /// failure categories. + pub error: crate::error::LoopError, +} + +/// Context for [`LoopObserver::on_response`](crate::observer::LoopObserver::on_response). +/// +/// Contains the model's text response and optional token usage +/// for the turn. +#[derive(Debug, Clone)] +pub struct ResponseContext { + /// Turn number. + /// + /// Identifies which turn produced this response. + pub turn: usize, + + /// The model's text response. + /// + /// Concatenated text content from the assistant message. + /// Tool-call content is excluded; see [`ToolPostContext`] + /// for tool result information. + pub text: String, + + /// Token usage for this turn, if available. + /// + /// Populated when the API returns usage data in the + /// streaming response. `None` when the provider does + /// not report usage. + pub usage: Option, +} + +/// Context for [`LoopObserver::on_tool_pre`](crate::observer::LoopObserver::on_tool_pre). +/// +/// Sent before a tool is executed, providing the tool name and +/// the call ID assigned by the API. +#[derive(Debug, Clone)] +pub struct ToolPreContext { + /// Turn number. + /// + /// Identifies which turn this tool call belongs to. + pub turn: usize, + + /// Tool name. + /// + /// Matches the name the tool was registered under. + pub tool: String, + + /// Tool call ID from the API response. + /// + /// Unique identifier assigned by the model for this specific + /// tool invocation, used to correlate with the tool result. + pub tool_call_id: String, +} + +/// Context for [`LoopObserver::on_tool_post`](crate::observer::LoopObserver::on_tool_post). +/// +/// Sent after a tool completes, providing the tool name, a hash +/// of the output, whether an error occurred, and the execution duration. +#[derive(Debug, Clone)] +pub struct ToolPostContext { + /// Turn number. + /// + /// Identifies which turn this tool result belongs to. + pub turn: usize, + + /// Tool name. + /// + /// Matches the name the tool was registered under. + pub tool: String, + + /// Deterministic hash of the tool output, if available. + /// + /// Used by loop detection to identify repeated tool results + /// without exposing the full output content. + pub result_hash: Option, + + /// Whether the tool returned an error. + /// + /// `true` when the tool execution resulted in an error + /// response rather than a successful output. + pub is_error: bool, + + /// Wall-clock execution duration. + /// + /// Measured from tool dispatch to completion, including any + /// permission prompts. + pub duration: std::time::Duration, +} + +/// Context for [`LoopObserver::on_compaction`](crate::observer::LoopObserver::on_compaction). +/// +/// Reports the estimated token counts before and after compaction and the +/// number of tokens saved. +#[derive(Debug, Clone)] +pub struct CompactedContext { + /// Estimated token count before compaction. + /// + /// The token count of the conversation before the compactor ran, + /// reconstructed from `tokens_after + tokens_saved`. + pub tokens_before: u64, + + /// Estimated token count after compaction. + /// + /// The token count of the compacted conversation that will be used + /// for subsequent model calls. + pub tokens_after: u64, + + /// Estimated tokens saved by compaction. + /// + /// `tokens_before - tokens_after`, the net reduction achieved by + /// the compactor. + pub tokens_saved: u64, +} + +/// Context for [`LoopObserver::on_fallback`](crate::observer::LoopObserver::on_fallback). +/// +/// Indicates which model failed (`from`) and which replacement +/// model was selected (`to`). +#[derive(Debug, Clone)] +pub struct FallbackContext { + /// Model that failed. + /// + /// The model identifier that produced the error triggering fallback. + pub from: String, + + /// Replacement model. + /// + /// The model identifier that will be used for subsequent requests. + pub to: String, +} + +/// Context for [`LoopObserver::on_loop_detected`](crate::observer::LoopObserver::on_loop_detected). +/// +/// Describes the repeating tool pattern and how many times it +/// was observed. +#[derive(Debug, Clone)] +pub struct LoopDetectedContext { + /// Description of the repeating tool pattern. + /// + /// Human-readable summary of the tool operation(s) that were + /// detected as repeating. + pub pattern: String, + + /// Number of times the pattern was observed. + /// + /// Counts consecutive repetitions of the same tool operation + /// with the same result hash. + pub repetitions: usize, +} + +/// Context for [`LoopObserver::on_convergence_detected`](crate::observer::LoopObserver::on_convergence_detected). +/// +/// Carries the configured action string (e.g. `"stop"`, `"warn"`, +/// `"compact"`) determined by the detection policy. +#[derive(Debug, Clone)] +pub struct ConvergenceDetectedContext { + /// Configured action to take (e.g. `"stop"`, `"warn"`, `"compact"`). + /// + /// Determined by the convergence detection policy configuration. + /// `"stop"` halts the loop, `"warn"` logs and continues, + /// `"compact"` triggers context compaction. + pub action: String, +} diff --git a/src/core/reflection.rs b/src/reflection.rs similarity index 61% rename from src/core/reflection.rs rename to src/reflection.rs index 8d474d5..c9aa2b7 100644 --- a/src/core/reflection.rs +++ b/src/reflection.rs @@ -1,7 +1,7 @@ //! Reflection and recovery for failed agent turns. //! //! When a tool call fails or a turn produces an unexpected result, the -//! framework needs to decide what to do. This module provides a pluggable +//! framework needs to decide what to do. A pluggable //! two-layer system: //! //! 1. **[`Reflector`]** — Analyses the failure and produces a @@ -18,33 +18,10 @@ //! - [`ExponentialBackoffRecovery`] — retries with exponential backoff up to //! a configurable limit. //! -//! # Architecture -//! -//! ```text -//! Tool call fails -//! │ -//! ▼ -//! ┌───────────────────────┐ -//! │ Reflector::analyze() │ -//! │ → FailureAnalysis │ -//! │ (recoverable?) │ -//! │ (correction?) │ -//! │ (severity) │ -//! └──────────┬────────────┘ -//! │ -//! ▼ -//! ┌───────────────────────────────┐ -//! │ RecoveryStrategy::decide() │ -//! │ → RecoveryAction │ -//! │ Retry / Skip / AskUser / │ -//! │ Fail │ -//! └───────────────────────────────┘ -//! ``` -//! //! # Quick Start //! //! ```rust -//! use loopctl::core::reflection::{ +//! use loopctl::reflection::{ //! NoopReflector, ExponentialBackoffRecovery, ReflectionContext, //! }; //! use std::sync::Arc; @@ -57,7 +34,10 @@ //! // let action = strategy.decide(&analysis, attempt, max_attempts).await; //! ``` -use crate::core::types::Correction; +pub mod backoff; +pub use backoff::ExponentialBackoffRecovery; + +use serde::{Deserialize, Serialize}; use std::fmt; use std::future::Future; use std::pin::Pin; @@ -77,7 +57,7 @@ use std::time::Duration; /// # Example /// /// ```rust -/// use loopctl::core::reflection::FailureSeverity; +/// use loopctl::reflection::FailureSeverity; /// /// assert!(FailureSeverity::Low < FailureSeverity::Critical); /// ``` @@ -86,7 +66,7 @@ use std::time::Duration; )] #[serde(rename_all = "snake_case")] pub enum FailureSeverity { - /// Minor issue — a simple retry will likely fix it. + /// Minor issue — a retry will likely fix it. Low, /// Moderate issue — may need a correction before retrying. Medium, @@ -120,7 +100,7 @@ impl fmt::Display for FailureSeverity { /// # Example /// /// ```rust -/// use loopctl::core::reflection::ReflectionContext; +/// use loopctl::reflection::ReflectionContext; /// /// let context = ReflectionContext { /// task: "Fix the bug in main.rs".to_string(), @@ -139,6 +119,103 @@ pub struct ReflectionContext { pub max_attempts: u32, } +// =================================================== +// Correction +// =================================================== + +/// Type of correction to apply. +/// +/// Categorizes the fix strategy that the reflection system has determined +/// is most appropriate for the observed failure. Each variant maps to a +/// different retry approach. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CorrectionType { + /// Fix the input to the tool. + /// + /// The tool was correct but its input parameters were wrong (e.g., a + /// typo in a file path). The correction provides a fixed input via + /// [`Correction::modified_input`]. + InputFix, + + /// Use a different tool. + /// + /// The chosen tool was inappropriate for the task. The correction + /// specifies an alternative via [`Correction::alternative_tool`]. + ToolChange, + + /// Fix a dependency or prerequisite. + /// + /// The tool call failed because a prerequisite was not met (e.g., + /// a directory doesn't exist). The correction describes what needs + /// to be done first. + PrerequisiteFix, + + /// Change the approach entirely. + /// + /// The current strategy is fundamentally flawed. The correction + /// provides high-level guidance for a different approach via + /// [`Correction::guidance`]. + ApproachChange, + + /// No fix possible, escalate. + /// + /// The reflection system cannot determine a correction. The + /// framework should propagate the error to the user or higher-level + /// handler. + Escalate, +} + +/// A correction produced by the reflection system. +/// +/// When a tool call fails and reflection is enabled (via `Feature::Reflection`), +/// the agent analyzes the error and produces a `Correction` that describes how to fix +/// the problem. The framework applies the correction and retries. +/// +/// # Serialization +/// +/// Implements `Serialize` and `Deserialize` for persistence and observability. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Correction { + /// See [`CorrectionType`] for available strategies. + pub correction_type: CorrectionType, + /// Explains *what* went wrong and *how* the correction addresses it. + pub description: String, + /// Corrected JSON input when [`CorrectionType::InputFix`]. `None` otherwise. + pub modified_input: Option, + /// Alternative tool name when [`CorrectionType::ToolChange`]. `None` otherwise. + pub alternative_tool: Option, + /// Extra context or instructions to help avoid the same failure. + pub guidance: Option, +} + +/// Result of applying a correction. +/// +/// Indicates whether the reflection system's correction was successfully +/// applied, failed, or was skipped. Produced after attempting to retry +/// with the corrected parameters. +#[derive(Debug, Clone)] +pub enum CorrectionResult { + /// Correction was applied successfully. + /// + /// The retry with corrected parameters succeeded and the agent can + /// continue processing normally. + Applied, + + /// Correction failed. + /// + /// The retry also failed. Contains a human-readable error message + /// describing what went wrong with the corrected attempt. + Failed(String), + + /// No correction was needed or possible. + /// + /// The reflection system decided not to apply a correction (e.g., + /// the error is transient or the correction type was + /// [`Escalate`](CorrectionType::Escalate)). + Skipped, +} + // =================================================== // FailureAnalysis // =================================================== @@ -152,7 +229,7 @@ pub struct ReflectionContext { /// # Example /// /// ```rust -/// use loopctl::core::reflection::{FailureAnalysis, FailureSeverity}; +/// use loopctl::reflection::{FailureAnalysis, FailureSeverity}; /// /// let analysis = FailureAnalysis { /// is_recoverable: true, @@ -165,15 +242,15 @@ pub struct ReflectionContext { /// ``` #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct FailureAnalysis { - /// Whether the framework should attempt recovery. + /// Whether the failure can be recovered from. pub is_recoverable: bool, - /// Human-readable description of the root cause. + /// Description of what went wrong. pub root_cause: String, /// How severe the failure is. pub severity: FailureSeverity, - /// Optional correction the agent can apply before retrying. + /// Suggested correction for the agent to apply before retrying. pub correction: Option, - /// Additional context about the failure (e.g., environment state). + /// Additional context (e.g., environment state at time of failure). pub context: String, } @@ -185,7 +262,7 @@ pub struct FailureAnalysis { /// /// The reflector can either produce a valid analysis, skip analysis /// (letting the framework use its default behaviour), or fail -/// internally. +/// during analysis. #[derive(Debug, thiserror::Error)] pub enum ReflectionError { /// The reflector opted out of analysing this failure. @@ -196,7 +273,7 @@ pub enum ReflectionError { /// The reflector itself encountered an error. /// - /// This is distinct from the tool failure being analysed — it means + /// Distinct from the tool failure being analysed — it means /// the reflector's own logic broke (e.g., an LLM call for /// summarisation failed). #[error("reflection internal error: {0}")] @@ -215,7 +292,7 @@ pub enum ReflectionError { /// # Example /// /// ```rust -/// use loopctl::core::reflection::RecoveryAction; +/// use loopctl::reflection::RecoveryAction; /// use std::time::Duration; /// /// let action = RecoveryAction::Retry { @@ -256,18 +333,6 @@ pub enum RecoveryAction { impl RecoveryAction { /// Returns the retry delay, if this is a [`Retry`](Self::Retry) action. - /// - /// # Example - /// - /// ```rust - /// use loopctl::core::reflection::RecoveryAction; - /// use std::time::Duration; - /// - /// let action = RecoveryAction::Retry { delay: Duration::from_secs(2) }; - /// assert_eq!(action.delay(), Some(Duration::from_secs(2))); - /// - /// assert_eq!(RecoveryAction::Fail("bad".into()).delay(), None); - /// ``` #[must_use] pub fn delay(&self) -> Option { match self { @@ -330,7 +395,7 @@ impl fmt::Display for RecoveryAction { /// # Example /// /// ```rust -/// use loopctl::core::reflection::{ +/// use loopctl::reflection::{ /// Reflector, ReflectionContext, FailureAnalysis, FailureSeverity, ReflectionError, /// }; /// use std::future::Future; @@ -410,7 +475,7 @@ pub trait Reflector: Send + Sync { /// # Example /// /// ```rust -/// use loopctl::core::reflection::{ +/// use loopctl::reflection::{ /// RecoveryStrategy, FailureAnalysis, FailureSeverity, RecoveryAction, /// }; /// use std::future::Future; @@ -509,152 +574,9 @@ impl fmt::Debug for NoopReflector { } } -// =================================================== -// ExponentialBackoffRecovery -// =================================================== - -/// Recovery strategy using exponential backoff. -/// -/// Retries up to `max_retries` times with exponential delays. If the -/// [`FailureAnalysis`] says the failure is not recoverable, returns -/// [`RecoveryAction::Fail`] immediately. -/// -/// # Backoff Formula -/// -/// `delay = min(base_delay × 2^attempt, max_delay)` -/// -/// # Example -/// -/// ```rust,ignore -/// use loopctl::core::reflection::{ExponentialBackoffRecovery, RecoveryAction, FailureAnalysis, FailureSeverity, RecoveryStrategy}; -/// use std::time::Duration; -/// -/// let strategy = ExponentialBackoffRecovery::new(3) -/// .with_base_delay(Duration::from_millis(100)) -/// .with_max_delay(Duration::from_secs(10)); -/// -/// let recoverable = FailureAnalysis { -/// is_recoverable: true, -/// root_cause: "timeout".to_string(), -/// severity: FailureSeverity::Low, -/// correction: None, -/// context: String::new(), -/// }; -/// let action = strategy.decide(&recoverable, 0, 5).await; -/// assert!(action.is_retry()); -/// assert_eq!(action.delay(), Some(Duration::from_millis(100))); -/// -/// let unrecoverable = FailureAnalysis { -/// is_recoverable: false, -/// root_cause: "invalid key".to_string(), -/// severity: FailureSeverity::Critical, -/// correction: None, -/// context: String::new(), -/// }; -/// let action = strategy.decide(&unrecoverable, 0, 5).await; -/// assert!(action.is_fail()); -/// ``` -#[derive(Debug, Clone)] -pub struct ExponentialBackoffRecovery { - /// Maximum number of retry attempts. - max_retries: u32, - /// Base delay before the first retry. - base_delay: Duration, - /// Maximum delay between retries. - max_delay: Duration, -} - -impl ExponentialBackoffRecovery { - /// Create a new strategy with the given maximum retries. - /// - /// # Example - /// - /// ```rust - /// use loopctl::core::reflection::ExponentialBackoffRecovery; - /// - /// let strategy = ExponentialBackoffRecovery::new(5); - /// ``` - #[must_use] - pub fn new(max_retries: u32) -> Self { - Self { - max_retries, - base_delay: Duration::from_millis(100), - max_delay: Duration::from_secs(30), - } - } - - /// Set the base delay (delay before the first retry). - #[must_use] - pub fn with_base_delay(mut self, delay: Duration) -> Self { - self.base_delay = delay; - self - } - - /// Set the maximum delay between retries. - #[must_use] - pub fn with_max_delay(mut self, delay: Duration) -> Self { - self.max_delay = delay; - self - } - - /// Returns the configured max retries. - #[must_use] - pub fn max_retries(&self) -> u32 { - self.max_retries - } - - /// Returns the configured base delay. - #[must_use] - pub fn base_delay(&self) -> Duration { - self.base_delay - } - - /// Returns the configured max delay. - #[must_use] - pub fn max_delay(&self) -> Duration { - self.max_delay - } - - /// Calculate the backoff delay for a given attempt. - fn delay_for_attempt(&self, attempt: u32) -> Duration { - let delay_ms = self - .base_delay - .as_millis() - .saturating_mul(1u128.checked_shl(attempt).unwrap_or(u128::MAX)); - let delay_ms = u64::try_from(delay_ms.min(self.max_delay.as_millis())).unwrap_or(u64::MAX); - Duration::from_millis(delay_ms) - } -} - -impl RecoveryStrategy for ExponentialBackoffRecovery { - fn decide( - &self, - analysis: &FailureAnalysis, - attempt: u32, - _max_attempts: u32, - ) -> Pin + Send + '_>> { - let action = if !analysis.is_recoverable { - RecoveryAction::Fail(analysis.root_cause.clone()) - } else if attempt >= self.max_retries { - RecoveryAction::Fail(format!("max retries ({}) exceeded", self.max_retries)) - } else if analysis.severity >= FailureSeverity::High && analysis.correction.is_some() { - RecoveryAction::AskUser(format!( - "high-severity failure with correction available: {}", - analysis.root_cause - )) - } else { - RecoveryAction::Retry { - delay: self.delay_for_attempt(attempt), - } - }; - Box::pin(async move { action }) - } -} - #[cfg(test)] mod tests { use super::*; - use crate::core::types::CorrectionType; // =================================================== // FailureSeverity tests @@ -840,181 +762,4 @@ mod tests { let debug = format!("{reflector:?}"); assert!(debug.contains("NoopReflector")); } - - // =================================================== - // ExponentialBackoffRecovery tests - // =================================================== - - #[test] - fn backoff_builder_defaults() { - let strategy = ExponentialBackoffRecovery::new(3); - assert_eq!(strategy.max_retries(), 3); - assert_eq!(strategy.base_delay(), Duration::from_millis(100)); - assert_eq!(strategy.max_delay(), Duration::from_secs(30)); - } - - #[test] - fn backoff_builder_custom() { - let strategy = ExponentialBackoffRecovery::new(5) - .with_base_delay(Duration::from_millis(200)) - .with_max_delay(Duration::from_secs(60)); - assert_eq!(strategy.max_retries(), 5); - assert_eq!(strategy.base_delay(), Duration::from_millis(200)); - assert_eq!(strategy.max_delay(), Duration::from_secs(60)); - } - - #[tokio::test] - async fn backoff_recoverable_first_attempt() { - let strategy = ExponentialBackoffRecovery::new(3); - let analysis = FailureAnalysis { - is_recoverable: true, - root_cause: "timeout".to_string(), - severity: FailureSeverity::Medium, - correction: None, - context: String::new(), - }; - let action = strategy.decide(&analysis, 0, 5).await; - assert!(action.is_retry()); - assert_eq!(action.delay(), Some(Duration::from_millis(100))); - } - - #[tokio::test] - async fn backoff_recoverable_second_attempt() { - let strategy = ExponentialBackoffRecovery::new(3); - let analysis = FailureAnalysis { - is_recoverable: true, - root_cause: "timeout".to_string(), - severity: FailureSeverity::Medium, - correction: None, - context: String::new(), - }; - let action = strategy.decide(&analysis, 1, 5).await; - assert!(action.is_retry()); - assert_eq!(action.delay(), Some(Duration::from_millis(200))); - } - - #[tokio::test] - async fn backoff_recoverable_third_attempt() { - let strategy = ExponentialBackoffRecovery::new(3); - let analysis = FailureAnalysis { - is_recoverable: true, - root_cause: "timeout".to_string(), - severity: FailureSeverity::Medium, - correction: None, - context: String::new(), - }; - let action = strategy.decide(&analysis, 2, 5).await; - assert!(action.is_retry()); - assert_eq!(action.delay(), Some(Duration::from_millis(400))); - } - - #[tokio::test] - async fn backoff_max_retries_exceeded() { - let strategy = ExponentialBackoffRecovery::new(3); - let analysis = FailureAnalysis { - is_recoverable: true, - root_cause: "timeout".to_string(), - severity: FailureSeverity::Low, - correction: None, - context: String::new(), - }; - let action = strategy.decide(&analysis, 3, 5).await; - assert!(action.is_fail()); - let RecoveryAction::Fail(reason) = action else { - unreachable!() - }; - assert!(reason.contains("max retries")); - } - - #[tokio::test] - async fn backoff_unrecoverable_fails_immediately() { - let strategy = ExponentialBackoffRecovery::new(3); - let analysis = FailureAnalysis { - is_recoverable: false, - root_cause: "invalid api key".to_string(), - severity: FailureSeverity::Critical, - correction: None, - context: String::new(), - }; - let action = strategy.decide(&analysis, 0, 5).await; - assert!(action.is_fail()); - let RecoveryAction::Fail(reason) = action else { - unreachable!() - }; - assert_eq!(reason, "invalid api key"); - } - - #[tokio::test] - async fn backoff_delay_capped_at_max() { - let strategy = ExponentialBackoffRecovery::new(10) - .with_base_delay(Duration::from_secs(1)) - .with_max_delay(Duration::from_secs(5)); - let analysis = FailureAnalysis { - is_recoverable: true, - root_cause: "timeout".to_string(), - severity: FailureSeverity::Medium, - correction: None, - context: String::new(), - }; - // 1 * 2^5 = 32s, capped at 5s - let action = strategy.decide(&analysis, 5, 10).await; - assert_eq!(action.delay(), Some(Duration::from_secs(5))); - } - - #[tokio::test] - async fn backoff_low_severity_retries() { - let strategy = ExponentialBackoffRecovery::new(3); - let analysis = FailureAnalysis { - is_recoverable: true, - root_cause: "transient hiccup".to_string(), - severity: FailureSeverity::Low, - correction: None, - context: String::new(), - }; - let action = strategy.decide(&analysis, 0, 5).await; - assert!(action.is_retry()); - } - - #[tokio::test] - async fn backoff_high_severity_with_correction_asks_user() { - use crate::core::types::CorrectionType; - let strategy = ExponentialBackoffRecovery::new(3); - let analysis = FailureAnalysis { - is_recoverable: true, - root_cause: "bad parameter".to_string(), - severity: FailureSeverity::High, - correction: Some(Correction { - correction_type: CorrectionType::InputFix, - description: "fix the file path".to_string(), - modified_input: None, - alternative_tool: None, - guidance: None, - }), - context: String::new(), - }; - let action = strategy.decide(&analysis, 0, 5).await; - assert!(action.is_ask_user()); - } - - #[tokio::test] - async fn backoff_high_severity_without_correction_retries() { - let strategy = ExponentialBackoffRecovery::new(3); - let analysis = FailureAnalysis { - is_recoverable: true, - root_cause: "timeout".to_string(), - severity: FailureSeverity::High, - correction: None, - context: String::new(), - }; - let action = strategy.decide(&analysis, 0, 5).await; - assert!(action.is_retry()); - } - - #[test] - fn backoff_debug_format() { - let strategy = ExponentialBackoffRecovery::new(3); - let debug = format!("{strategy:?}"); - assert!(debug.contains("ExponentialBackoffRecovery")); - assert!(debug.contains("max_retries")); - } } diff --git a/src/reflection/backoff.rs b/src/reflection/backoff.rs new file mode 100644 index 0000000..d1953fb --- /dev/null +++ b/src/reflection/backoff.rs @@ -0,0 +1,395 @@ +//! Exponential backoff recovery strategy. +//! +//! [`ExponentialBackoffRecovery`] retries recoverable failures with +//! exponentially increasing delays, up to a configurable maximum. +//! This is the default [`RecoveryStrategy`] used by `BareLoop`. +//! +//! # Quick Start +//! +//! ``` +//! use loopctl::reflection::backoff::ExponentialBackoffRecovery; +//! use loopctl::reflection::{RecoveryAction, FailureAnalysis, FailureSeverity, RecoveryStrategy}; +//! +//! let strategy = ExponentialBackoffRecovery::new(3); +//! ``` + +use super::{FailureAnalysis, FailureSeverity, RecoveryAction, RecoveryStrategy}; +use std::future::Future; +use std::pin::Pin; +use std::time::Duration; + +// =================================================== +// ExponentialBackoffRecovery +// =================================================== + +/// A recovery strategy that retries with exponential backoff. +/// +/// When a tool call fails with a recoverable error, this strategy +/// retries up to [`max_retries`](ExponentialBackoffRecovery::max_retries) times, +/// with delays that grow exponentially: `base_delay * 2^attempt`. +/// +/// # Behaviour by severity +/// +/// | Severity | Recoverable | Correction | Action | +/// |------------|-------------|------------|-------------------------------------------------------------------| +/// | Low/Medium | ✓ | — | Retry with backoff | +/// | High | ✓ | Some | [`AskUser`](super::RecoveryAction::AskUser) (correction provided) | +/// | High | ✓ | None | Retry with backoff | +/// | Critical | ✗ | — | Fail immediately | +/// | any | ✗ | — | Fail immediately | +/// +/// # Example +/// +/// ``` +/// use loopctl::reflection::backoff::ExponentialBackoffRecovery; +/// +/// let strategy = ExponentialBackoffRecovery::new(5); +/// ``` +#[derive(Debug, Clone)] +pub struct ExponentialBackoffRecovery { + /// Maximum number of retry attempts before giving up. + max_retries: u32, + /// Initial delay applied to the first retry, doubled on each subsequent attempt. + base_delay: Duration, + /// Upper bound on the computed delay, regardless of exponential growth. + max_delay: Duration, +} + +impl ExponentialBackoffRecovery { + /// Create a new strategy with the given maximum retry count. + /// + /// Defaults: `base_delay` = 100ms, `max_delay` = 30s. + #[must_use] + pub fn new(max_retries: u32) -> Self { + Self { + max_retries, + base_delay: Duration::from_millis(100), + max_delay: Duration::from_secs(30), + } + } + + /// Set a custom base delay for the exponential backoff. + /// + /// The delay for attempt `n` is `base_delay * 2^n`, capped at `max_delay`. + #[must_use] + pub fn with_base_delay(mut self, delay: Duration) -> Self { + self.base_delay = delay; + self + } + + /// Set a maximum delay cap. + /// + /// Even as the exponential grows, the delay never exceeds this value. + #[must_use] + pub fn with_max_delay(mut self, delay: Duration) -> Self { + self.max_delay = delay; + self + } + + /// Maximum retry attempts before giving up. + #[must_use] + pub fn max_retries(&self) -> u32 { + self.max_retries + } + + /// Base delay for exponential backoff calculation. + #[must_use] + pub fn base_delay(&self) -> Duration { + self.base_delay + } + + /// Maximum delay cap. + #[must_use] + pub fn max_delay(&self) -> Duration { + self.max_delay + } + + /// Calculate the backoff delay for a given attempt. + /// + /// Computes `base_delay * 2^attempt`, clamped to [`max_delay`](Self::max_delay). + /// + /// Attempt 0 yields `base_delay`, attempt 1 yields `2 * base_delay`, and + /// so on. Once the exponential exceeds `max_delay`, every subsequent + /// attempt returns `max_delay`. + /// + /// Uses saturating arithmetic to avoid overflow on very large attempt + /// numbers — once `2^attempt` would overflow, the delay is clamped to + /// `max_delay`. + fn delay_for_attempt(&self, attempt: u32) -> Duration { + let delay_ms = self + .base_delay + .as_millis() + .saturating_mul(1u128.checked_shl(attempt).unwrap_or(u128::MAX)); + let delay_ms = u64::try_from(delay_ms.min(self.max_delay.as_millis())).unwrap_or(u64::MAX); + Duration::from_millis(delay_ms) + } +} + +impl RecoveryStrategy for ExponentialBackoffRecovery { + fn decide( + &self, + analysis: &FailureAnalysis, + attempt: u32, + _max_attempts: u32, + ) -> Pin + Send + '_>> { + let action = if !analysis.is_recoverable { + RecoveryAction::Fail(analysis.root_cause.clone()) + } else if attempt >= self.max_retries { + RecoveryAction::Fail(format!("max retries ({}) exceeded", self.max_retries)) + } else if analysis.severity >= FailureSeverity::High && analysis.correction.is_some() { + RecoveryAction::AskUser(format!( + "high-severity failure with correction available: {}", + analysis.root_cause + )) + } else { + RecoveryAction::Retry { + delay: self.delay_for_attempt(attempt), + } + }; + Box::pin(async move { action }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn backoff_builder_defaults() { + let strategy = ExponentialBackoffRecovery::new(3); + assert_eq!(strategy.max_retries(), 3); + assert_eq!(strategy.base_delay(), Duration::from_millis(100)); + assert_eq!(strategy.max_delay(), Duration::from_secs(30)); + } + + #[test] + fn backoff_builder_custom() { + let strategy = ExponentialBackoffRecovery::new(5) + .with_base_delay(Duration::from_millis(200)) + .with_max_delay(Duration::from_secs(60)); + assert_eq!(strategy.max_retries(), 5); + assert_eq!(strategy.base_delay(), Duration::from_millis(200)); + assert_eq!(strategy.max_delay(), Duration::from_secs(60)); + } + + #[tokio::test] + async fn backoff_recoverable_first_attempt() { + let strategy = ExponentialBackoffRecovery::new(3); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "timeout".to_string(), + severity: FailureSeverity::Medium, + correction: None, + context: String::new(), + }; + let action = strategy.decide(&analysis, 0, 5).await; + assert!(action.is_retry()); + assert_eq!(action.delay(), Some(Duration::from_millis(100))); + } + + #[tokio::test] + async fn backoff_recoverable_second_attempt() { + let strategy = ExponentialBackoffRecovery::new(3); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "timeout".to_string(), + severity: FailureSeverity::Medium, + correction: None, + context: String::new(), + }; + let action = strategy.decide(&analysis, 1, 5).await; + assert!(action.is_retry()); + assert_eq!(action.delay(), Some(Duration::from_millis(200))); + } + + #[tokio::test] + async fn backoff_recoverable_third_attempt() { + let strategy = ExponentialBackoffRecovery::new(3); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "timeout".to_string(), + severity: FailureSeverity::Medium, + correction: None, + context: String::new(), + }; + let action = strategy.decide(&analysis, 2, 5).await; + assert!(action.is_retry()); + assert_eq!(action.delay(), Some(Duration::from_millis(400))); + } + + #[tokio::test] + async fn backoff_max_retries_exceeded() { + let strategy = ExponentialBackoffRecovery::new(3); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "timeout".to_string(), + severity: FailureSeverity::Low, + correction: None, + context: String::new(), + }; + let action = strategy.decide(&analysis, 3, 5).await; + assert!(action.is_fail()); + let RecoveryAction::Fail(reason) = action else { + unreachable!() + }; + assert!(reason.contains("max retries")); + } + + #[tokio::test] + async fn backoff_unrecoverable_fails_immediately() { + let strategy = ExponentialBackoffRecovery::new(3); + let analysis = FailureAnalysis { + is_recoverable: false, + root_cause: "invalid api key".to_string(), + severity: FailureSeverity::Critical, + correction: None, + context: String::new(), + }; + let action = strategy.decide(&analysis, 0, 5).await; + assert!(action.is_fail()); + let RecoveryAction::Fail(reason) = action else { + unreachable!() + }; + assert_eq!(reason, "invalid api key"); + } + + #[tokio::test] + async fn backoff_delay_capped_at_max() { + let strategy = ExponentialBackoffRecovery::new(10) + .with_base_delay(Duration::from_secs(1)) + .with_max_delay(Duration::from_secs(5)); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "timeout".to_string(), + severity: FailureSeverity::Medium, + correction: None, + context: String::new(), + }; + // 1 * 2^5 = 32s, capped at 5s + let action = strategy.decide(&analysis, 5, 10).await; + assert_eq!(action.delay(), Some(Duration::from_secs(5))); + } + + #[tokio::test] + async fn backoff_low_severity_retries() { + let strategy = ExponentialBackoffRecovery::new(3); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "transient hiccup".to_string(), + severity: FailureSeverity::Low, + correction: None, + context: String::new(), + }; + let action = strategy.decide(&analysis, 0, 5).await; + assert!(action.is_retry()); + } + + #[tokio::test] + async fn backoff_high_severity_with_correction_asks_user() { + use super::super::{Correction, CorrectionType}; + let strategy = ExponentialBackoffRecovery::new(3); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "bad parameter".to_string(), + severity: FailureSeverity::High, + correction: Some(Correction { + correction_type: CorrectionType::InputFix, + description: "fix the file path".to_string(), + modified_input: None, + alternative_tool: None, + guidance: None, + }), + context: String::new(), + }; + let action = strategy.decide(&analysis, 0, 5).await; + assert!(action.is_ask_user()); + } + + #[tokio::test] + async fn backoff_high_severity_without_correction_retries() { + let strategy = ExponentialBackoffRecovery::new(3); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "timeout".to_string(), + severity: FailureSeverity::High, + correction: None, + context: String::new(), + }; + let action = strategy.decide(&analysis, 0, 5).await; + assert!(action.is_retry()); + } + + #[test] + fn backoff_debug_format() { + let strategy = ExponentialBackoffRecovery::new(3); + let debug = format!("{strategy:?}"); + assert!(debug.contains("ExponentialBackoffRecovery")); + assert!(debug.contains("max_retries")); + } + + // ---- delay_for_attempt ---- + + #[test] + fn delay_for_attempt_zero_yields_base() { + let strategy = ExponentialBackoffRecovery::new(3); + assert_eq!(strategy.delay_for_attempt(0), Duration::from_millis(100)); + } + + #[test] + fn delay_for_attempt_one_doubles() { + let strategy = ExponentialBackoffRecovery::new(3); + assert_eq!(strategy.delay_for_attempt(1), Duration::from_millis(200)); + } + + #[test] + fn delay_for_attempt_two_quadruples() { + let strategy = ExponentialBackoffRecovery::new(3); + assert_eq!(strategy.delay_for_attempt(2), Duration::from_millis(400)); + } + + #[test] + fn delay_for_attempt_capped_at_max() { + let strategy = ExponentialBackoffRecovery::new(10) + .with_base_delay(Duration::from_secs(1)) + .with_max_delay(Duration::from_secs(5)); + // 1 * 2^5 = 32s, capped at 5s + assert_eq!(strategy.delay_for_attempt(5), Duration::from_secs(5)); + } + + #[test] + fn delay_for_attempt_exactly_at_max() { + let strategy = ExponentialBackoffRecovery::new(10) + .with_base_delay(Duration::from_secs(1)) + .with_max_delay(Duration::from_secs(8)); + // 1 * 2^3 = 8s, exactly at cap + assert_eq!(strategy.delay_for_attempt(3), Duration::from_secs(8)); + } + + #[test] + fn delay_for_attempt_stays_capped_for_large_attempts() { + let strategy = ExponentialBackoffRecovery::new(100) + .with_base_delay(Duration::from_millis(100)) + .with_max_delay(Duration::from_secs(30)); + let late = strategy.delay_for_attempt(50); + assert_eq!(late, Duration::from_secs(30)); + } + + #[test] + fn delay_for_attempt_custom_base() { + let strategy = + ExponentialBackoffRecovery::new(5).with_base_delay(Duration::from_millis(250)); + assert_eq!(strategy.delay_for_attempt(0), Duration::from_millis(250)); + assert_eq!(strategy.delay_for_attempt(1), Duration::from_millis(500)); + assert_eq!(strategy.delay_for_attempt(2), Duration::from_secs(1)); + } + + #[test] + fn delay_for_attempt_overflow_does_not_panic() { + let strategy = ExponentialBackoffRecovery::new(3) + .with_base_delay(Duration::from_secs(1)) + .with_max_delay(Duration::from_secs(10)); + // attempt = u32::MAX would overflow 2^n; should saturate, not panic + let delay = strategy.delay_for_attempt(u32::MAX); + assert_eq!(delay, Duration::from_secs(10)); + } +} diff --git a/src/runtime.rs b/src/runtime.rs new file mode 100644 index 0000000..7267366 --- /dev/null +++ b/src/runtime.rs @@ -0,0 +1,599 @@ +//! Loop runtime — capability traits and runtime infrastructure for agent loops. +//! +//! Two categories of types govern how agent loops +//! interact with their infrastructure: +//! +//! # Capability Traits +//! +//! Capability traits are composable interfaces that represent distinct +//! infrastructure concerns. Each trait describes a single ability — +//! observing lifecycle events, detecting loops, falling back to alternate +//! models, etc. The [`LoopRuntime`] struct implements all capability traits, +//! providing a concrete, all-in-one infrastructure bundle. +//! +//! | Trait | Purpose | +//! |--------------------------|--------------------------------------------------| +//! | [`Observable`] | Emit lifecycle events to registered observers | +//! | [`Detectable`] | Detect repetitive loops and convergence | +//! | [`FallbackCapable`] | Circuit-breaker fallback to alternate models | +//! | [`Compactable`] | Automatic context compaction when tokens exceed | +//! | [`StreamCapable`] | Resilient streaming with retries and timeouts | +//! | [`Hookable`] | Bidirectional hooks that can block actions | +//! | [`PipelineAware`] | Dispatch tools through a middleware pipeline | +//! | [`HealthTrackable`] | Per-tool health tracking with circuit breakers | +//! +//! # `LoopRuntime` +//! +//! [`LoopRuntime`] is the framework's default infrastructure bundle. +//! It bundles all managers, observers, hooks, and middleware into a single +//! struct that can be passed to any agent loop implementation: +//! +//! ```text +//! LoopRuntime +//! ├── ObserverHost → Observable +//! ├── DetectionManager → Detectable +//! ├── FallbackManager → FallbackCapable +//! ├── Option → Compactable +//! ├── Option → StreamCapable +//! ├── Option → Hookable +//! ├── Option → PipelineAware +//! └── Option → HealthTrackable +//! ``` +//! +//! # Design Philosophy +//! +//! The trait hierarchy separates **what the runtime can do** (capability +//! traits) from **how it's composed** (the `LoopRuntime` struct). This +//! allows: +//! +//! - **Generic programming** — agent loops can be written against +//! `impl Observable + Detectable` rather than a concrete type. +//! - **Testing** — swap `LoopRuntime` for a stub that implements only the +//! traits under test. +//! - **Incremental adoption** — start with `LoopRuntime::new()` and add +//! capabilities via builder methods as needed. + +use std::sync::Arc; + +use crate::compact::ContextManager; +use crate::detection::DetectionManager; +use crate::fallback::FallbackManager; +#[cfg(feature = "hooks")] +use crate::hooks::HookExecutor; +use crate::middleware::ToolPipeline; +use crate::observer::{LoopObserver, ObserverHost}; +use crate::stream::handler::StreamHandler; +#[cfg(feature = "tool_health")] +use crate::tool::health::ToolHealthRegistry; + +pub use crate::capabilities::*; + +// ================================================== +// LoopRuntime +// ================================================== + +/// The framework's default infrastructure bundle for agent loops. +/// +/// `LoopRuntime` bundles all the cross-cutting infrastructure an agent +/// loop needs: observers, detection, fallback, hooks, middleware pipeline, +/// and health tracking. It implements all capability traits so that agent +/// loops can be written against trait bounds rather than a concrete type. +/// +/// # Construction +/// +/// Use [`LoopRuntime::new`] for defaults or the builder-style `with_*` +/// methods to configure individual components: +/// +/// ```rust,ignore +/// use loopctl::runtime::LoopRuntime; +/// use loopctl::fallback::FallbackManager; +/// +/// let runtime = LoopRuntime::new() +/// .with_fallback(FallbackManager::for_model("llm-70b")); +/// ``` +/// +/// # Capability Traits +/// +/// `LoopRuntime` implements all capability traits defined in this module: +/// +/// - [`Observable`] — via the internal [`ObserverHost`] +/// - [`Detectable`] — via the internal [`DetectionManager`] +/// - [`FallbackCapable`] — via the internal [`FallbackManager`] +/// - [`Compactable`] — via an optional [`ContextManager`] +/// - [`StreamCapable`] — via an optional [`StreamHandler`] +/// - [`Hookable`] — via an optional [`HookExecutor`] ** +/// - [`PipelineAware`] — via an optional [`ToolPipeline`] +/// - [`HealthTrackable`] — via an optional [`ToolHealthRegistry`] *(requires `tool_health` feature)* +/// +/// # Reset +/// +/// Call [`reset_all`](LoopRuntime::reset_all) at the start of a new +/// session to reinitialise every manager and observer to its default state. +pub struct LoopRuntime { + /// Circuit breaker for API model fallback. + pub fallback: FallbackManager, + /// Loop and convergence detection orchestrator. + pub detection: DetectionManager, + /// Observer host for lifecycle event fan-out. + observer_host: ObserverHost, + /// Optional middleware pipeline wrapping tool dispatch. + tool_pipeline: Option, + /// Optional context manager for automatic compaction. + context_manager: Option>, + /// Optional stream handler for resilient streaming with retries. + stream_handler: Option, + /// Optional hook executor for bidirectional lifecycle interception. + #[cfg(feature = "hooks")] + hook_executor: Option>, + /// Optional per-tool health tracker with circuit breakers. *Requires `tool_health` feature.* + #[cfg(feature = "tool_health")] + health_registry: Option>, +} + +impl LoopRuntime { + /// Create a new runtime with default managers and no optional components. + /// + /// # Example + /// + /// ``` + /// # use loopctl::runtime::LoopRuntime; + /// # use loopctl::runtime::FallbackCapable; + /// let runtime = LoopRuntime::new(); + /// assert!(runtime.fallback().active_model().is_none()); + /// ``` + #[must_use] + pub fn new() -> Self { + Self { + fallback: FallbackManager::default(), + detection: DetectionManager::default(), + observer_host: ObserverHost::new(), + tool_pipeline: None, + context_manager: None, + stream_handler: None, + #[cfg(feature = "hooks")] + hook_executor: None, + #[cfg(feature = "tool_health")] + health_registry: None, + } + } + + // ================================================== + // Builder methods + // ================================================== + + /// Replace the fallback manager with a custom instance. + /// + /// # Example + /// + /// ``` + /// # use loopctl::runtime::LoopRuntime; + /// # use loopctl::runtime::FallbackCapable; + /// # use loopctl::fallback::FallbackManager; + /// let runtime = LoopRuntime::new() + /// .with_fallback(FallbackManager::for_model("llm-70b")); + /// assert_eq!(runtime.fallback().active_model().as_deref(), Some("llm-70b")); + /// ``` + #[must_use] + pub fn with_fallback(mut self, fallback: FallbackManager) -> Self { + self.fallback = fallback; + self + } + + /// Replace the detection manager with a custom instance. + /// + /// # Example + /// + /// ``` + /// # use loopctl::runtime::LoopRuntime; + /// # use loopctl::runtime::Detectable; + /// # use loopctl::detection::{DetectionManager, DetectionConfig}; + /// let config = DetectionConfig { + /// loop_threshold: 5, + /// ..Default::default() + /// }; + /// let runtime = LoopRuntime::new() + /// .with_detection(DetectionManager::new_with_config(config).unwrap()); + /// assert_eq!(runtime.detection().config().loop_threshold, 5); + /// ``` + #[must_use] + pub fn with_detection(mut self, detection: DetectionManager) -> Self { + self.detection = detection; + self + } + + // ================================================== + // Setters for optional components + // ================================================== + + /// Register an observer with the observer host. + /// + /// # Example + /// + /// ```rust,ignore + /// use loopctl::observer::LoopObserver; + /// use std::sync::Arc; + /// + /// let mut runtime = LoopRuntime::new(); + /// runtime.register_observer(Arc::new(MyObserver)); + /// ``` + pub fn register_observer(&mut self, observer: Arc) { + self.observer_host.register(observer); + } + + /// Access the observer host directly. + pub fn observers(&self) -> &ObserverHost { + &self.observer_host + } + + /// Set the middleware pipeline for tool dispatch. + /// + /// # Example + /// + /// ```rust,ignore + /// let mut runtime = LoopRuntime::new(); + /// runtime.set_pipeline(builder.build()?); + /// ``` + pub fn set_pipeline(&mut self, pipeline: ToolPipeline) { + self.tool_pipeline = Some(pipeline); + } + + /// Set the context manager for automatic compaction. + /// + /// # Example + /// + /// ```rust,ignore + /// use loopctl::compact::{ContextManager, TruncatingCompactor}; + /// use std::sync::Arc; + /// + /// let compactor = TruncatingCompactor::new() + /// .with_preserve_recent(4) + /// .with_min_messages(6); + /// let manager = ContextManager::new(Arc::new(compactor)) + /// .with_context_window(200_000) + /// .with_threshold(0.80); + /// + /// let mut runtime = LoopRuntime::new(); + /// runtime.set_context_manager(Arc::new(manager)); + /// ``` + pub fn set_context_manager(&mut self, manager: Arc) { + self.context_manager = Some(manager); + } + + /// Set the stream handler for resilient streaming with retries. + /// + /// # Example + /// + /// ```rust,ignore + /// use loopctl::stream::handler::{StreamHandler, StreamTimeoutConfig}; + /// use std::time::Duration; + /// + /// let handler = StreamHandler::with_config( + /// StreamTimeoutConfig { + /// initial_event_timeout: Duration::from_secs(60), + /// ..Default::default() + /// }, + /// Default::default(), + /// ); + /// + /// let mut runtime = LoopRuntime::new(); + /// runtime.set_stream_handler(handler); + /// ``` + pub fn set_stream_handler(&mut self, handler: StreamHandler) { + self.stream_handler = Some(handler); + } + + /// Set the hook executor for bidirectional lifecycle interception. + #[cfg(feature = "hooks")] + pub fn set_hook_executor(&mut self, executor: Arc) { + self.hook_executor = Some(executor); + } + + /// Set the tool health registry for per-tool health tracking. + /// + /// When set, records success/failure counts and latency for every + /// tool dispatch. Tools that exceed the failure threshold have their + /// circuit breaker opened, blocking subsequent calls until recovery. + /// + /// *Requires `tool_health` feature.* + #[cfg(feature = "tool_health")] + pub fn set_health_registry(&mut self, registry: Arc) { + self.health_registry = Some(registry); + } + + // ================================================== + // Lifecycle + // ================================================== + + /// Reset all managers and observers to their initial state. + /// + /// Delegates to each manager's `reset()` method and calls + /// [`ObserverHost::reset_all`]. Typically called at the start of + /// a new agent task or session. + /// + /// # Example + /// + /// ``` + /// # use loopctl::runtime::LoopRuntime; + /// let runtime = LoopRuntime::new(); + /// // ... after a session ... + /// runtime.reset_all(); + /// // All managers are back to their initial state + /// ``` + pub fn reset_all(&self) { + self.fallback.reset(); + self.detection.reset(); + self.observer_host.reset_all(); + } +} + +impl Default for LoopRuntime { + /// Produce a [`LoopRuntime`] with default managers and no optional components. + /// + /// Equivalent to [`LoopRuntime::new`]. + fn default() -> Self { + Self::new() + } +} + +// ================================================== +// Capability trait implementations +// ================================================== + +impl crate::capabilities::Observable for LoopRuntime { + fn observers(&self) -> &ObserverHost { + &self.observer_host + } +} + +impl crate::capabilities::Detectable for LoopRuntime { + fn detection(&self) -> &DetectionManager { + &self.detection + } +} + +impl crate::capabilities::FallbackCapable for LoopRuntime { + fn fallback(&self) -> &FallbackManager { + &self.fallback + } +} + +impl crate::capabilities::Compactable for LoopRuntime { + fn context_manager(&self) -> Option<&Arc> { + self.context_manager.as_ref() + } +} + +impl crate::capabilities::StreamCapable for LoopRuntime { + fn stream_handler(&self) -> Option<&StreamHandler> { + self.stream_handler.as_ref() + } +} + +#[cfg(feature = "hooks")] +impl crate::capabilities::Hookable for LoopRuntime { + fn hook_executor(&self) -> Option<&HookExecutor> { + self.hook_executor.as_deref() + } +} + +impl crate::capabilities::PipelineAware for LoopRuntime { + fn pipeline(&self) -> Option<&ToolPipeline> { + self.tool_pipeline.as_ref() + } +} + +#[cfg(feature = "tool_health")] +impl crate::capabilities::HealthTrackable for LoopRuntime { + fn health_registry(&self) -> Option<&ToolHealthRegistry> { + self.health_registry.as_deref() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::detection::{DetectedPattern, DetectionManager}; + + #[test] + fn test_runtime_default() { + let runtime = LoopRuntime::default(); + assert!(runtime.fallback().active_model().is_none()); + } + + #[test] + fn test_runtime_with_custom_fallback() { + let fallback = FallbackManager::for_model("my-model"); + let runtime = LoopRuntime::new().with_fallback(fallback); + assert_eq!( + runtime.fallback().active_model().as_deref(), + Some("my-model") + ); + } + + #[test] + fn test_reset_all() { + let runtime = LoopRuntime::new(); + runtime.reset_all(); + } + + #[test] + fn test_runtime_contains_detection_manager() { + let runtime = LoopRuntime::new(); + assert_eq!(runtime.detection().config().loop_threshold, 3); + assert_eq!(runtime.detection().config().stop_threshold, 10); + } + + #[test] + fn test_runtime_with_custom_detection() { + use crate::detection::DetectionConfig; + + let config = DetectionConfig { + loop_threshold: 7, + stop_threshold: 20, + ..Default::default() + }; + let detection = DetectionManager::new_with_config(config).unwrap(); + let runtime = LoopRuntime::new().with_detection(detection); + assert_eq!(runtime.detection().config().loop_threshold, 7); + assert_eq!(runtime.detection().config().stop_threshold, 20); + assert!(runtime.fallback().active_model().is_none()); + } + + #[test] + fn test_reset_all_clears_detection() { + let runtime = LoopRuntime::new(); + let _ = runtime.detection().record_tool_call("Read", 12345); + runtime.reset_all(); + let pattern = runtime.detection().record_tool_call("Read", 12345); + assert!(matches!(pattern, DetectedPattern::NoPattern)); + } + + #[test] + fn test_capability_traits_are_object_safe() { + // Verify that trait objects can be created + fn _assert_observable(_: &dyn Observable) {} + fn _assert_detectable(_: &dyn Detectable) {} + fn _assert_fallback(_: &dyn FallbackCapable) {} + fn _assert_pipeline(_: &dyn PipelineAware) {} + fn _assert_compactable(_: &dyn Compactable) {} + fn _assert_stream_capable(_: &dyn StreamCapable) {} + + let runtime = LoopRuntime::new(); + _assert_observable(&runtime); + _assert_detectable(&runtime); + _assert_fallback(&runtime); + _assert_pipeline(&runtime); + _assert_compactable(&runtime); + _assert_stream_capable(&runtime); + } + + #[test] + fn test_pipeline_defaults_to_none() { + let runtime = LoopRuntime::new(); + assert!(runtime.pipeline().is_none()); + } + + #[test] + fn test_observers_accessible_via_trait() { + let runtime = LoopRuntime::new(); + let _: &ObserverHost = runtime.observers(); + } + + #[test] + fn test_context_manager_defaults_to_none() { + let runtime = LoopRuntime::new(); + assert!(runtime.context_manager().is_none()); + } + + #[test] + fn test_stream_handler_defaults_to_none() { + let runtime = LoopRuntime::new(); + assert!(runtime.stream_handler().is_none()); + } + + #[test] + fn test_set_pipeline_returns_some() { + use crate::middleware::ToolPipeline; + use crate::tool::ToolRegistry; + use std::sync::Arc; + + let registry = Arc::new(ToolRegistry::new()); + let pipeline = ToolPipeline::new(registry); + let mut runtime = LoopRuntime::new(); + assert!(runtime.pipeline().is_none()); + runtime.set_pipeline(pipeline); + assert!(runtime.pipeline().is_some()); + } + + #[test] + fn test_set_context_manager_returns_some() { + use crate::compact::{ContextManager, TruncatingCompactor}; + + let compactor = TruncatingCompactor::new(); + let manager = ContextManager::new(Arc::new(compactor)); + let mut runtime = LoopRuntime::new(); + assert!(runtime.context_manager().is_none()); + runtime.set_context_manager(Arc::new(manager)); + assert!(runtime.context_manager().is_some()); + } + + #[test] + fn test_set_stream_handler_returns_some() { + use crate::stream::handler::StreamHandler; + + let handler = StreamHandler::new(); + let mut runtime = LoopRuntime::new(); + assert!(runtime.stream_handler().is_none()); + runtime.set_stream_handler(handler); + assert!(runtime.stream_handler().is_some()); + } + + #[test] + fn test_register_observer_increments_count() { + use crate::observer::LoopObserver; + use std::sync::Arc; + + struct NopObserver; + impl LoopObserver for NopObserver { + fn name(&self) -> &str { + "NopObserver" + } + } + + let mut runtime = LoopRuntime::new(); + assert!(runtime.observers().is_empty()); + runtime.register_observer(Arc::new(NopObserver)); + assert_eq!(runtime.observers().len(), 1); + runtime.register_observer(Arc::new(NopObserver)); + assert_eq!(runtime.observers().len(), 2); + } + + #[test] + fn test_reset_all_clears_fallback() { + let runtime = LoopRuntime::new(); + let _ = runtime.fallback().record_api_failure(); + runtime.reset_all(); + assert!(runtime.fallback().active_model().is_none()); + } + + #[test] + fn test_reset_all_clears_observers() { + use crate::observer::LoopObserver; + use std::sync::Arc; + + struct NopObserver; + impl LoopObserver for NopObserver { + fn name(&self) -> &str { + "NopObserver" + } + } + + let mut runtime = LoopRuntime::new(); + runtime.register_observer(Arc::new(NopObserver)); + assert_eq!(runtime.observers().len(), 1); + runtime.reset_all(); + // Observer count persists across reset — reset_all calls + // reset_all on each observer, it doesn't remove them. + assert_eq!(runtime.observers().len(), 1); + } + + #[test] + fn test_generic_bounds_accept_runtime() { + fn accepts_observable(_: &impl Observable) {} + fn accepts_detectable(_: &impl Detectable) {} + fn accepts_fallback(_: &impl FallbackCapable) {} + fn accepts_compactable(_: &impl Compactable) {} + fn accepts_stream_capable(_: &impl StreamCapable) {} + fn accepts_pipeline(_: &impl PipelineAware) {} + fn accepts_multi_bound(_: &(impl Observable + Detectable + FallbackCapable)) {} + + let runtime = LoopRuntime::new(); + accepts_observable(&runtime); + accepts_detectable(&runtime); + accepts_fallback(&runtime); + accepts_compactable(&runtime); + accepts_stream_capable(&runtime); + accepts_pipeline(&runtime); + accepts_multi_bound(&runtime); + } +} diff --git a/src/stream.rs b/src/stream.rs index 9f57a30..a52c341 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -1,6 +1,6 @@ //! Streaming event types for LLM API responses. //! -//! This module defines the types used when consuming Server-Sent Events +//! Types used when consuming Server-Sent Events //! (SSE) based streaming responses from LLM APIs. The core [`StreamEvent`] //! enum represents each discrete event in the stream lifecycle, while //! [`StreamAccumulator`] collects those events into a complete [`Message`]. @@ -8,7 +8,7 @@ //! Streaming allows the framework to process model output incrementally — //! displaying text as it arrives, detecting tool invocations as soon as //! the part starts, and reporting token usage without waiting for the -//! full response. This is essential for responsive agent behavior. +//! full response. Essential for responsive agent behavior. //! //! # Stream Lifecycle //! @@ -483,8 +483,8 @@ pub enum DeltaPart { /// Reason why the model stopped generating tokens. /// -/// This is the streaming / API-level stop reason returned by the LLM -/// provider in the [`MessageDelta`] event. It differs from the +/// Streaming / API-level stop reason returned by the LLM +/// provider in the [`MessageDelta`] event. Differs from the /// agent-level `StopReason` which is used in `TurnResult`. /// /// Use [`should_continue_tool_loop`](Self::should_continue_tool_loop) @@ -524,20 +524,20 @@ pub enum StreamStopReason { /// The response was truncated. The caller may want to request /// continuation or increase the token budget. /// - /// *Note: `AgentConfig` will be available once the builder module is complete.* + /// *Note: `LoopConfig` will be available once the builder module is complete.* MaxTokens, /// The model hit a configured stop sequence. /// /// The response ended because it matched one of the stop - /// sequences provided in the request. This is uncommon in + /// sequences provided in the request. Uncommon in /// typical agent usage. StopSequence, /// The model completed its turn naturally. /// /// The model finished generating its response without hitting - /// any limits or invoking tools. This is the normal end-of-turn + /// any limits or invoking tools. Normal end-of-turn /// signal for non-tool responses. EndTurn, } @@ -670,7 +670,7 @@ pub struct MessageDelta { /// /// `Some` when the API reports usage; `None` if usage data /// is not available or not yet received. See [`Usage`]. - /// This is typically populated in the final `MessageDelta` event + /// Typically populated in the final `MessageDelta` event /// and reflects cumulative token consumption for the entire request. pub usage: Option, } @@ -818,7 +818,7 @@ impl Usage { /// Accumulates streaming events into a complete [`Message`]. /// -/// This is a stateful builder that tracks the progress of a streaming +/// Stateful builder that tracks the progress of a streaming /// response as [`StreamEvent`]s arrive and assembles the final /// [`Message`] once all events have been processed. /// diff --git a/src/stream/handler.rs b/src/stream/handler.rs index 2ba44d8..7976246 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -1,6 +1,6 @@ //! Configuration, result types, and error types for resilient LLM stream handling. //! -//! This module defines the types that underpin [`StreamHandler`] — the framework's +//! Types that underpin [`StreamHandler`] — the framework's //! production-grade streaming resilience layer. The handler wraps //! [`ApiClient::stream_messages`] with retry, timeout, and fallback behaviour. //! @@ -38,7 +38,7 @@ //! let handler = StreamHandler::new(); //! //! // Or with custom config: -//! let handler = StreamHandler::with_config( +//! let handler = StreamHandler::new().with_config( //! StreamTimeoutConfig { //! initial_event_timeout: std::time::Duration::from_secs(60), //! ..Default::default() @@ -47,7 +47,7 @@ //! ); //! ``` -use crate::api_client::ApiClient; +use crate::api::ApiClient; use crate::cancel::CancelSignal; use crate::message::Message; use crate::stream::{StreamAccumulator, StreamEvent, StreamStopReason, Usage}; @@ -92,7 +92,7 @@ use std::time::{Duration, Instant}; pub struct StreamTimeoutConfig { /// Timeout for the first event after opening the stream. /// - /// This is the most critical timeout — if the API server never sends + /// Most critical timeout — if the API server never sends /// the first event, the stream hangs forever. Set to a generous value /// since the model may need time to begin generating. pub initial_event_timeout: Duration, @@ -353,7 +353,7 @@ impl StreamRetryConfig { pub enum StreamOutcome { /// Stream completed normally — all events received, `MessageStop` seen. /// - /// This is the happy path. The [`StreamAccumulator`] + /// Happy path. The [`StreamAccumulator`] /// contains the full response. Completed { /// Number of SSE events processed. @@ -399,7 +399,7 @@ pub enum StreamOutcome { attempts: u32, }, - /// Fell back to non-streaming [`create_message`](crate::api_client::ApiClient::create_message). + /// Fell back to non-streaming [`create_message`](crate::api::ApiClient::create_message). /// /// Streaming failed, but a non-streaming request succeeded. /// The response is complete but was not streamed incrementally. @@ -587,7 +587,7 @@ pub struct StreamProgress { /// let handler = StreamHandler::new(); /// assert_eq!(handler.timeout_config().initial_event_timeout, std::time::Duration::from_secs(120)); /// -/// let handler = StreamHandler::with_config( +/// let handler = StreamHandler::new().with_config( /// StreamTimeoutConfig { /// initial_event_timeout: std::time::Duration::from_secs(60), /// ..Default::default() @@ -650,7 +650,7 @@ impl StreamHandler { /// use loopctl::stream::handler::{StreamHandler, StreamTimeoutConfig, StreamRetryConfig}; /// use std::time::Duration; /// - /// let handler = StreamHandler::with_config( + /// let handler = StreamHandler::new().with_config( /// StreamTimeoutConfig { /// initial_event_timeout: Duration::from_secs(60), /// ..Default::default() @@ -662,11 +662,10 @@ impl StreamHandler { /// ); /// ``` #[must_use] - pub fn with_config(timeout: StreamTimeoutConfig, retry: StreamRetryConfig) -> Self { - Self { - timeout_config: timeout, - retry_config: retry, - } + pub fn with_config(mut self, timeout: StreamTimeoutConfig, retry: StreamRetryConfig) -> Self { + self.timeout_config = timeout; + self.retry_config = retry; + self } /// Returns a reference to the timeout configuration. @@ -687,7 +686,7 @@ impl StreamHandler { /// Stream one complete turn with retry, timeout, and fallback. /// - /// This is the primary entry point for resilient streaming. It + /// Primary entry point for resilient streaming. It /// orchestrates the full lifecycle: /// /// 1. Opens a stream via [`ApiClient::stream_messages`]. @@ -841,7 +840,7 @@ impl StreamHandler { total_deadline: Option, ) -> Result where - S: futures::Stream> + S: futures::Stream> + Unpin, { let mut accumulator = StreamAccumulator::new(); @@ -1212,7 +1211,7 @@ mod tests { #[test] fn handler_with_config() { - let handler = StreamHandler::with_config( + let handler = StreamHandler::new().with_config( StreamTimeoutConfig { initial_event_timeout: Duration::from_secs(60), ..Default::default() @@ -1439,7 +1438,7 @@ mod tests { // process_events async tests // =================================================== - use crate::api_error::ApiError; + use crate::api::error::ApiError; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, PartStart, StreamEvent, @@ -1537,7 +1536,7 @@ mod tests { #[tokio::test] async fn process_events_total_timeout() { // Use a very short total timeout to trigger it immediately. - let handler = StreamHandler::with_config( + let handler = StreamHandler::new().with_config( StreamTimeoutConfig { total_stream_timeout: Duration::from_millis(1), per_event_timeout: Duration::from_secs(300), @@ -1575,7 +1574,7 @@ mod tests { #[tokio::test] async fn process_events_cancelled() { - let handler = StreamHandler::with_config( + let handler = StreamHandler::new().with_config( StreamTimeoutConfig { per_event_timeout: Duration::from_secs(300), ..Default::default() @@ -1694,7 +1693,7 @@ mod tests { #[tokio::test] async fn fallback_non_streaming_success() { - let handler = StreamHandler::with_config( + let handler = StreamHandler::new().with_config( StreamTimeoutConfig { fallback_to_non_streaming: true, ..Default::default() @@ -1725,7 +1724,7 @@ mod tests { #[tokio::test] async fn fallback_non_streaming_cancelled_before_start() { - let handler = StreamHandler::with_config( + let handler = StreamHandler::new().with_config( StreamTimeoutConfig { fallback_to_non_streaming: true, ..Default::default() @@ -1749,7 +1748,7 @@ mod tests { #[tokio::test] async fn fallback_non_streaming_error() { - let handler = StreamHandler::with_config( + let handler = StreamHandler::new().with_config( StreamTimeoutConfig { fallback_to_non_streaming: true, ..Default::default() @@ -1842,7 +1841,7 @@ mod tests { // (covered above). Here we test that stream_turn returns the // error when streaming fails and the handler is configured // without fallback. - let handler = StreamHandler::with_config( + let handler = StreamHandler::new().with_config( StreamTimeoutConfig { fallback_to_non_streaming: false, ..Default::default() diff --git a/src/stream/heartbeat.rs b/src/stream/heartbeat.rs index 986c819..4b24be3 100644 --- a/src/stream/heartbeat.rs +++ b/src/stream/heartbeat.rs @@ -45,7 +45,7 @@ //! ); //! ``` -use crate::api_error::ApiError; +use crate::api::error::ApiError; use crate::stream::StreamEvent; use futures::Stream; use std::pin::Pin; @@ -197,7 +197,7 @@ impl HeartbeatConfig { /// /// `HeartbeatStream` implements `Stream` directly, so it composes with /// any other stream wrapper. Use it on any stream you've already opened -/// when you just need heartbeat/timeout without the full handler lifecycle. +/// when you need heartbeat/timeout without the full handler lifecycle. /// /// # Example /// diff --git a/src/testing.rs b/src/testing.rs index 1bcd8fb..cef3c1c 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -1,6 +1,6 @@ //! Testing utilities — mock components and fixture factories for loopctl tests. //! -//! This module provides reusable mocks and fixture factories for testing +//! Reusable mocks and fixture factories for testing //! code that depends on loopctl traits. Instead of wiring up real API //! clients and tools in tests, these stubs can be used to exercise //! agent logic in isolation, assert on streaming events, and verify tool @@ -50,8 +50,8 @@ //! - [`test_assistant_message`] — Create a test assistant [`Message`]. //! - [`test_tool_use_message`] — Create an assistant [`Message`] with //! tool-call content blocks. -//! - [`test_config`] — Create a test [`AgentConfig`] with sensible defaults. -//! - [`test_config_with_id`] — Create a test [`AgentConfig`] with a specific +//! - [`test_config`] — Create a test [`LoopConfig`] with sensible defaults. +//! - [`test_config_with_id`] — Create a test [`LoopConfig`] with a specific //! session ID. //! //! # Quick Start @@ -99,9 +99,9 @@ //! }, //! ]); -use crate::api_client::ApiClient; -use crate::api_error::ApiError; -use crate::core::AgentConfig; +use crate::api::ApiClient; +use crate::api::error::ApiError; +use crate::config::LoopConfig; use crate::message::{Message, MessagePart, Role}; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, @@ -297,7 +297,7 @@ pub struct MockToolCall { /// /// Must correspond to a tool registered in the agent's /// [`ToolRegistry`](crate::tool::ToolRegistry). The mock does not - /// validate this — it simply passes the name through. + /// validate this — passes the name through directly. /// /// During test execution, if the agent loop cannot find a tool /// with this name in the registry, it will return an error. @@ -342,14 +342,14 @@ impl MockApiClient { /// Use the builder methods to customize before passing the client /// to the code under test. /// - /// The default response is deliberately simple so that most tests + /// The default response is simple so that most tests /// only need to call [`with_text_response`](MockApiClient::with_text_response) /// to get started. /// /// # Example /// /// ```rust - /// use loopctl::api_client::ApiClient; + /// use loopctl::api::ApiClient; /// use loopctl::testing::MockApiClient; /// /// let client = MockApiClient::new("test-model"); @@ -377,7 +377,7 @@ impl MockApiClient { /// Set the text response for the first (or only) turn. /// /// Overwrites the `text` field on the initial [`MockResponse`] - /// created by [`new`](MockApiClient::new). This is the simplest way + /// created by [`new`](MockApiClient::new). Simplest way /// to configure a single-turn mock — the model will "say" the given /// text and stop. /// @@ -475,7 +475,7 @@ impl MockApiClient { /// and reused, so the mock never panics on an empty queue. /// /// If `responses` is empty the call is a no-op (the default response - /// is retained). This is the recommended way to set up complex + /// is retained). Recommended way to set up complex /// multi-turn scenarios where the model needs to reply differently /// across successive turns. /// @@ -538,10 +538,10 @@ impl MockApiClient { /// queue never empties. This ensures repeated calls to /// [`stream_messages`](ApiClient::stream_messages) always succeed. /// - /// This is a private helper used by both + /// Helper used by both /// [`stream_messages`](ApiClient::stream_messages) and /// [`create_message`](ApiClient::create_message). The method - /// acquires the `Mutex` guard internally, so callers do not need + /// acquires the `Mutex` guard, so callers do not need /// to handle locking. /// /// # Queue exhaustion strategy @@ -594,7 +594,7 @@ impl MockApiClient { /// # Ignored parameters /// /// The `_messages`, `_system`, and `_tools` parameters are accepted for -/// trait compatibility but intentionally ignored — the mock always +/// trait compatibility but ignored — the mock always /// returns its preconfigured response regardless of the input. impl ApiClient for MockApiClient { /// Return the model name this mock was created with. @@ -607,7 +607,7 @@ impl ApiClient for MockApiClient { /// # Example /// /// ```rust - /// use loopctl::api_client::ApiClient; + /// use loopctl::api::ApiClient; /// use loopctl::testing::MockApiClient; /// /// let client = MockApiClient::new("my-test-model"); @@ -633,7 +633,7 @@ impl ApiClient for MockApiClient { /// contains a single [`ApiError`] event instead. /// /// The `_messages`, `_system`, and `_tools` parameters are accepted for - /// trait compatibility but are intentionally ignored — the mock always + /// trait compatibility but ignored — the mock always /// returns the preconfigured response. /// /// # Usage tokens @@ -646,7 +646,7 @@ impl ApiClient for MockApiClient { /// /// ```rust /// # tokio::runtime::Runtime::new().unwrap().block_on(async { - /// use loopctl::api_client::ApiClient; + /// use loopctl::api::ApiClient; /// use loopctl::testing::MockApiClient; /// /// let client = MockApiClient::new("test-model").with_text_response("Hi!"); @@ -727,13 +727,13 @@ impl ApiClient for MockApiClient { /// response queue entirely. /// /// The `_messages`, `_system`, and `_tools` parameters are accepted - /// for trait compatibility but intentionally ignored. + /// for trait compatibility but ignored. /// /// # Example /// /// ```rust /// # tokio::runtime::Runtime::new().unwrap().block_on(async { - /// use loopctl::api_client::ApiClient; + /// use loopctl::api::ApiClient; /// use loopctl::testing::MockApiClient; /// /// let client = MockApiClient::new("test-model").with_text_response("Hi!"); @@ -1106,7 +1106,7 @@ impl MockTool { /// # Metadata methods /// /// The [`name`](Tool::name), [`description`](Tool::description), and -/// [`schema`](Tool::schema) methods simply return the values set at +/// [`schema`](Tool::schema) methods return the values set at /// construction time via [`MockTool::new`]. The /// [`is_concurrency_safe`](Tool::is_concurrency_safe), /// [`is_read_only`](Tool::is_read_only), and @@ -1237,7 +1237,7 @@ impl Tool for MockTool { /// directly. The returned message has [`Role::User`] and a single /// [`MessagePart::Text`] variant containing the provided string. /// -/// This is the most common fixture for constructing the "user says" +/// Most common fixture for constructing the "user says" /// part of a conversation history. /// /// # Example @@ -1284,7 +1284,7 @@ pub fn test_assistant_message(text: &str) -> Message { /// If `tool_use_id` is an empty string a unique ID of the form /// `"call_{i}"` is generated automatically (where `i` is the index). /// -/// This is the message the agent loop produces when the model requests +/// Message the agent loop produces when the model requests /// tool execution — use it to simulate the "assistant asked for a tool" /// step in multi-turn tests. Each tuple becomes a [`MessagePart::ToolCall`] /// variant in the message's `content` vector. @@ -1323,7 +1323,7 @@ pub fn test_tool_use_message(calls: &[(&str, &str, Value)]) -> Message { Message::new(Role::Assistant, blocks) } -/// Create a test [`AgentConfig`] with sensible defaults. +/// Create a test [`LoopConfig`] with sensible defaults. /// /// The returned config has: /// @@ -1353,8 +1353,8 @@ pub fn test_tool_use_message(calls: &[(&str, &str, Value)]) -> Message { /// assert_eq!(config.max_turns, 10); /// ``` #[must_use] -pub fn test_config() -> AgentConfig { - AgentConfig { +pub fn test_config() -> LoopConfig { + LoopConfig { session_id: Uuid::new_v4(), max_turns: 10, system_prompt: Some("You are a test assistant.".to_string()), @@ -1362,7 +1362,7 @@ pub fn test_config() -> AgentConfig { } } -/// Create a test [`AgentConfig`] with a specific session ID. +/// Create a test [`LoopConfig`] with a specific session ID. /// /// Same as [`test_config`] but with a caller-supplied session ID. /// Useful when tests need to assert on the ID — for example verifying @@ -1388,8 +1388,8 @@ pub fn test_config() -> AgentConfig { /// assert_eq!(config.session_id, id); /// ``` #[must_use] -pub fn test_config_with_id(id: Uuid) -> AgentConfig { - AgentConfig { +pub fn test_config_with_id(id: Uuid) -> LoopConfig { + LoopConfig { session_id: id, max_turns: 10, system_prompt: Some("You are a test assistant.".to_string()), diff --git a/src/tool.rs b/src/tool.rs index a19422a..df9cdfd 100644 --- a/src/tool.rs +++ b/src/tool.rs @@ -22,7 +22,7 @@ //! //! # Middleware Pipeline //! -//! The [`engine::middleware`](crate::engine::middleware) module provides a composable +//! The [`engine::middleware`](crate::middleware) module provides a composable //! middleware chain for tool dispatch with cross-cutting concerns //! (timeouts, output limiting, etc.). //! @@ -72,9 +72,16 @@ use std::fmt; use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::time::Duration; use crate::message::ToolContent as MessageToolContent; +pub mod permission; +pub mod registry; + +pub use permission::PermissionCheck; +pub use registry::{FnTool, ToolRegistry}; + // =================================================== // ToolSchema // =================================================== @@ -210,7 +217,7 @@ impl ToolOutput { /// the output is a structured [`MessageToolContent`] value. /// Sets [`is_error`](ToolOutput::is_error) to `false`. /// - /// This is the most general success constructor — it accepts any type + /// Most general success constructor — accepts any type /// that converts into [`MessageToolContent`]. For simple text results, /// prefer the more concise [`ToolOutput::text`] helper. /// @@ -260,7 +267,7 @@ impl ToolOutput { /// /// Convenience wrapper around [`ToolOutput::success`] that converts /// the input string into a [`MessageToolContent::Text`] variant. - /// This is the most common constructor for simple tool outputs. + /// Most common constructor for simple tool outputs. /// /// # When to use /// @@ -396,6 +403,193 @@ impl From<&str> for ToolOutput { } } +// =================================================== +// ToolDispatchResult +// =================================================== + +/// The outcome of a single tool invocation. +/// +/// Produced after the framework dispatches a tool call and collects +/// the tool's output. Used throughout the middleware pipeline, the +/// engine dispatch layer, and returned to callers. +/// +/// # Fields +/// +/// | Field | Source | +/// |------------------------|-------------------------------------| +/// | `tool_call_id` | Set by the engine after dispatch | +/// | `output` | From [`ToolOutput::payload`] | +/// | `is_error` | From [`ToolOutput::is_error`] | +/// | `duration` | Measured by the dispatch layer | +/// | `resolved_tool_name` | Set by middleware or engine | +/// +/// # Construction +/// +/// Middlewares build results with [`ToolDispatchResult::ok`], +/// [`ToolDispatchResult::err`], or [`From`] combined with +/// builder methods. The engine layer attaches the `tool_call_id` via +/// [`ToolDispatchResult::with_call_id`] after the middleware pipeline +/// returns. +/// +/// ``` +/// use std::time::Duration; +/// use loopctl::tool::{ToolDispatchResult, ToolOutput}; +/// +/// let output = ToolOutput::text("done"); +/// let result = ToolDispatchResult::from(output) +/// .with_tool_name("bash") +/// .with_duration(Duration::from_millis(42)) +/// .with_call_id("call_abc123"); +/// +/// assert_eq!(result.tool_call_id, "call_abc123"); +/// assert_eq!(result.resolved_tool_name, "bash"); +/// ``` +#[derive(Debug, Clone)] +pub struct ToolDispatchResult { + /// Set by the engine via [`with_call_id`](Self::with_call_id). + pub tool_call_id: String, + /// Preserves multipart and image content on success; wraps error in text on failure. + pub output: crate::message::ToolContent, + /// Whether the tool dispatch resulted in an error. + pub is_error: bool, + /// Wall-clock execution time. + pub duration: Duration, + /// May differ from the requested tool name if a routing middleware redirected the call. + pub resolved_tool_name: String, +} + +impl ToolDispatchResult { + /// Create a successful result with text output. + /// + /// Constructor for the common case where a tool + /// produces a plain-text response. + #[must_use] + pub fn ok(tool_name: &str, output: String, duration: Duration) -> Self { + Self { + tool_call_id: String::new(), + output: crate::message::ToolContent::Text(output), + is_error: false, + duration, + resolved_tool_name: tool_name.to_string(), + } + } + + /// Create an error result with a message. + /// + /// Used when a middleware short-circuits or the tool reports failure. + #[must_use] + pub fn err(tool_name: &str, message: String, duration: Duration) -> Self { + Self { + tool_call_id: String::new(), + output: crate::message::ToolContent::Text(message), + is_error: true, + duration, + resolved_tool_name: tool_name.to_string(), + } + } + + /// Create a result from a [`ToolOutput`]. + /// + /// Converts the tool's output struct into a dispatch result, + /// preserving the error flag and content payload. + #[must_use] + pub fn from_tool_output(tool_name: &str, output: ToolOutput, duration: Duration) -> Self { + Self::from(output) + .with_tool_name(tool_name) + .with_duration(duration) + } + + /// Builder: attach the [`tool_call_id`](Self::tool_call_id). + /// + /// Called by the engine layer after the middleware pipeline returns + /// to correlate this result with the original tool call. + #[must_use] + pub fn with_call_id(mut self, id: impl Into) -> Self { + self.tool_call_id = id.into(); + self + } + + /// Set the [`resolved_tool_name`](Self::resolved_tool_name). + /// + /// Part of the builder chain when constructing a + /// `ToolDispatchResult` from [`From`]. + #[must_use] + pub fn with_tool_name(mut self, name: &str) -> Self { + name.clone_into(&mut self.resolved_tool_name); + self + } + + /// Set the [`duration`](Self::duration). + /// + /// Part of the builder chain when constructing a + /// `ToolDispatchResult` from [`From`]. + #[must_use] + pub fn with_duration(mut self, dur: Duration) -> Self { + self.duration = dur; + self + } + + /// Create a result from a [`ToolError`]. + /// + /// Converts the tool's error into a dispatch result with `is_error` + /// set to `true`. + #[must_use] + pub fn from_tool_error(tool_name: &str, error: &ToolError, duration: Duration) -> Self { + Self { + tool_call_id: String::new(), + output: crate::message::ToolContent::Text(error.to_string()), + is_error: true, + duration, + resolved_tool_name: tool_name.to_string(), + } + } + + /// Create a result from a tool call outcome. + /// + /// Covers the common `Result` pattern produced by + /// [`Tool::call()`](Tool::call). Maps [`Ok`] through + /// [`from_tool_output`](Self::from_tool_output) and [`Err`] through + /// [`from_tool_error`](Self::from_tool_error). + #[must_use] + pub fn from_result( + tool_name: &str, + result: Result, + duration: Duration, + ) -> Self { + match result { + Ok(output) => Self::from_tool_output(tool_name, output, duration), + Err(e) => Self::from_tool_error(tool_name, &e, duration), + } + } +} + +/// Conversion from a bare [`ToolOutput`]. +/// +/// Produces a `ToolDispatchResult` with no call ID, [`Duration::ZERO`], +/// and an empty `resolved_tool_name`. Chain builder methods to complete +/// the fields: +/// +/// ``` +/// use std::time::Duration; +/// use loopctl::tool::{ToolDispatchResult, ToolOutput}; +/// +/// let result = ToolDispatchResult::from(ToolOutput::text("ok")) +/// .with_call_id("call_1") +/// .with_tool_name("echo") +/// .with_duration(Duration::from_millis(5)); +/// ``` +impl From for ToolDispatchResult { + fn from(output: ToolOutput) -> Self { + Self { + tool_call_id: String::new(), + output: output.payload, + is_error: output.is_error, + duration: Duration::ZERO, + resolved_tool_name: String::new(), + } + } +} + // =================================================== // ToolError // =================================================== @@ -780,296 +974,20 @@ impl Default for ToolContext { } } } -// =================================================== -// PermissionCheck -// =================================================== - -/// Result of a permission check before tool execution. -/// -/// Before invoking [`Tool::call`], the agent loop can run a permission -/// gate that returns one of four outcomes: allow, deny, ask the user, -/// or modify the input. This lets host applications enforce safety -/// policies without modifying individual tool implementations. -/// -/// # Lifecycle -/// -/// ```text -/// tool.call(input, ctx) -/// → permission_gate(input) -/// → Allow [proceed with original input] -/// → Deny [return ToolError::Permission] -/// → Ask { prompt } [prompt user, then Allow or Deny] -/// → Modify { .. } [proceed with modified input] -/// ``` -/// -/// # Decision tree -/// -/// ```text -/// ┌─────────────────┐ -/// │ Permission gate │ -/// └───────┬─────────┘ -/// ┌──────┼──────────┐ -/// ▼ ▼ ▼ -/// Allow Deny ┌ Ask ──┐ -/// │ │ ▼ │ -/// │ │ user │ -/// │ │ approves │ -/// │ │ | │ -/// │ │ ▼ ▼ -/// ▼ ▼ Allow Deny -/// Tool::call Err(Permission) -/// ``` -/// -/// # Example -/// -/// ```rust,ignore -/// let check = PermissionCheck::deny("dangerous operation"); -/// if check.is_deny() { -/// return Err(ToolError::Permission("blocked by policy".into())); -/// } -/// ``` -#[derive(Debug, Clone)] -pub enum PermissionCheck { - /// Allow the tool to execute unmodified. - /// - /// The agent loop proceeds with the original input and context. - Allow, - - /// Deny execution with a human-readable reason. - /// - /// The agent loop should return - /// [`ToolError::Permission`] with the given - /// `reason` so the LLM can react accordingly. - Deny { - /// Explanation of why the invocation was blocked. - /// - /// Forwarded to the LLM as part of the error message so it can - /// adjust its next action. - reason: String, - }, - - /// Prompt the user for approval before proceeding. - /// - /// In interactive sessions the agent loop should present `prompt` to - /// the user and then treat the response as either [`Allow`](PermissionCheck::Allow) - /// or [`Deny`](PermissionCheck::Deny). - Ask { - /// The question to present to the user. - /// - /// Should clearly describe the action the tool is about to take - /// and any potential side effects. - prompt: String, - }, - - /// Modify the tool's input before execution. - /// - /// The agent loop should invoke [`Tool::call`] with `modified_input` - /// instead of the original input. Useful for sanitising paths, - /// redacting secrets, or injecting default values. - Modify { - /// The sanitized or rewritten input to pass to [`Tool::call`]. - /// - /// Must conform to the tool's [`ToolSchema::input_schema`]. - modified_input: Value, - }, -} - -impl PermissionCheck { - /// Create an [`Allow`](PermissionCheck::Allow) result. - /// - /// Signals that the tool invocation may proceed without changes. - /// The `#[must_use]` attribute reminds callers to check the result - /// rather than silently discarding it. - /// - /// # When returned - /// - /// The permission gate returns this variant when the requested - /// operation is within the configured safety policy — for example, - /// a read-only tool invocation or an operation on an allowed path. - /// - /// # Example - /// - /// ```rust - /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; - /// - /// let check = PermissionCheck::allow(); - /// assert!(check.is_allow()); - /// ``` - #[must_use] - pub fn allow() -> Self { - Self::Allow - } - - /// Create a [`Deny`](PermissionCheck::Deny) result with a reason. - /// - /// The `reason` string will be forwarded to the LLM as part of the - /// error message, helping it understand why the invocation was - /// rejected and adjust its next action. - /// - /// # When returned - /// - /// The permission gate returns this variant when the requested - /// operation violates a hard safety rule — for example, executing - /// a shell command when shell access is disabled. - /// - /// # Example - /// - /// ```rust - /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; - /// - /// let check = PermissionCheck::deny("shell execution is disabled"); - /// assert!(check.is_deny()); - /// ``` - pub fn deny(reason: impl Into) -> Self { - Self::Deny { - reason: reason.into(), - } - } - - /// Create an [`Ask`](PermissionCheck::Ask) result with a prompt. - /// - /// The agent loop should present the `prompt` to the user (in - /// interactive mode) and then proceed based on the user's response. - /// - /// # When returned - /// - /// The permission gate returns this variant for operations that are - /// potentially dangerous but not outright prohibited — for example, - /// writing to a file for the first time. The user's decision is then - /// converted to [`Allow`](PermissionCheck::Allow) or - /// [`Deny`](PermissionCheck::Deny). - /// - /// # Example - /// - /// ```rust - /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; - /// - /// let check = PermissionCheck::ask("Allow write to /etc/config.yaml?"); - /// assert!(check.is_ask()); - /// ``` - pub fn ask(prompt: impl Into) -> Self { - Self::Ask { - prompt: prompt.into(), - } - } - - /// Create a [`Modify`](PermissionCheck::Modify) result with rewritten input. - /// - /// The agent loop should replace the original tool input with the - /// provided `modified_input` before invoking [`Tool::call`]. Useful - /// for sanitising paths, redacting secrets, or injecting default - /// values. - /// - /// # When returned - /// - /// The permission gate returns this variant when the requested - /// operation is acceptable but the input needs adjustment — for - /// example, resolving a relative path to an absolute one within - /// the allowed directory tree. - /// - /// # Example - /// - /// ```rust - /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; - /// use serde_json::json; - /// - /// let check = PermissionCheck::modify(json!({"path": "/safe/dir/file.txt"})); - /// assert!(check.is_modify()); - /// ``` - #[must_use] - pub fn modify(modified_input: Value) -> Self { - Self::Modify { modified_input } - } - - /// Returns `true` if this is an [`Allow`](PermissionCheck::Allow). - /// - /// Convenience predicate for the most common happy-path check. - /// Used by the agent loop to test whether to proceed with - /// [`Tool::call`] without further processing. - /// - /// # Example - /// - /// ```rust,ignore - /// if check.is_allow() { - /// let result = tool.call(input, &ctx).await; - /// } - /// ``` - #[must_use] - pub fn is_allow(&self) -> bool { - matches!(self, Self::Allow) - } - - /// Returns `true` if this is a [`Deny`](PermissionCheck::Deny). - /// - /// When `true`, the agent loop should *not* invoke the tool and - /// should instead return a permission error to the LLM. The denial - /// reason can be extracted by destructuring the variant or by - /// converting to [`ToolError::Permission`]. - /// - /// # Example - /// - /// ```rust,ignore - /// if check.is_deny() { - /// return Err(ToolError::Permission("blocked by policy".into())); - /// } - /// ``` - #[must_use] - pub fn is_deny(&self) -> bool { - matches!(self, Self::Deny { .. }) - } - - /// Returns `true` if this is an [`Ask`](PermissionCheck::Ask). - /// - /// When `true`, the agent loop should prompt the user before - /// deciding whether to allow or deny the invocation. In - /// non-interactive mode ([`ToolContext::is_non_interactive`]), the - /// loop typically treats an [`Ask`](PermissionCheck::Ask) as a - /// [`Deny`](PermissionCheck::Deny). - /// - /// # Example - /// - /// ```rust,ignore - /// if check.is_ask() { - /// println!("Tool requests approval: {}", prompt); - /// } - /// ``` - #[must_use] - pub fn is_ask(&self) -> bool { - matches!(self, Self::Ask { .. }) - } - - /// Returns `true` if this is a [`Modify`](PermissionCheck::Modify). - /// - /// When `true`, the agent loop should replace the original input - /// with the modified version before calling the tool. The modified - /// input can be extracted by matching the variant. - /// - /// # Example - /// - /// ```rust,ignore - /// if let PermissionCheck::Modify { modified_input } = check { - /// let result = tool.call(modified_input, &ctx).await; - /// } - /// ``` - #[must_use] - pub fn is_modify(&self) -> bool { - matches!(self, Self::Modify { .. }) - } -} - // =================================================== // Tool trait // =================================================== /// The trait that all agent tools must implement. /// -/// This is the framework-level tool interface. Concrete tools (defined in +/// Framework-level tool interface. Concrete tools (defined in /// downstream crates like `dch-tools`) implement this trait, and the /// [`ToolRegistry`] manages dynamic lookup by name. /// /// The trait uses a `Pin>` return type for [`call`](Tool::call) -/// to be maximally compatible with both `async fn` and manual `Future` -/// implementations, without requiring `async_fn_in_trait` stabilisation. +/// to be maximally compatible with both `async fn` bodies and manually +/// constructed futures, keeping the trait object-safe and free of +/// lifetime issues that `async fn` in traits can introduce. /// /// # Lifecycle /// @@ -1112,12 +1030,12 @@ impl PermissionCheck { /// /// # Provided methods /// -/// | Method | Default | Purpose | -/// |-------------------------------------|---------|------------------------------------| -/// | `is_concurrency_safe` | `false` | Static concurrency flag | -/// | `is_safe_for_concurrent_execution`| delegates| Per-input concurrency check | -/// | `is_read_only` | `false` | Side-effect flag for permission | -/// | `system_prompt` | `None` | Extra LLM context for this tool | +/// | Method | Default | Purpose | +/// |-------------------------------------|-----------|-------------------------------------| +/// | `is_concurrency_safe` | `false` | Static concurrency flag | +/// | `is_safe_for_concurrent_execution` | delegates | Per-input concurrency check | +/// | `is_read_only` | `false` | Side-effect flag for permission | +/// | `system_prompt` | `None` | Extra LLM context for this tool | /// /// # Example /// @@ -1194,7 +1112,7 @@ pub trait Tool: Send + Sync { /// Invoke the tool with the given input and context. /// - /// This is the main execution entry point. The `input` is a + /// Main execution entry point. The `input` is a /// [`Value`] (typically a JSON object) matching the tool's /// [`ToolSchema::input_schema`]. The [`ToolContext`] provides /// session-level data such as working directory and extensions. @@ -1206,7 +1124,7 @@ pub trait Tool: Send + Sync { /// /// The `Pin>` return type maximises compatibility /// with both `async fn` bodies and manually constructed futures, - /// without requiring `async_fn_in_trait` stabilisation. + /// keeping the trait object-safe and free of lifetime issues. /// /// # Errors /// @@ -1330,635 +1248,6 @@ pub trait Tool: Send + Sync { } } -// =================================================== -// ToolRegistry -// =================================================== - -/// Registry of available tools for dynamic lookup by name. -/// -/// The agent loop creates a [`ToolRegistry`] at session start, registers -/// all available tools via [`register`](ToolRegistry::register), and then -/// uses [`get`](ToolRegistry::get) to dispatch invocations when the LLM -/// selects a tool by name. The registry also provides bulk accessors for -/// tool schemas and concurrency-safe tool lists. -/// -/// # Data flow -/// -/// ```text -/// ┌──────────────┐ -/// │ Session init │ -/// └──────┬───────┘ -/// ▼ -/// ToolRegistry::new() -/// → register(tool_1) -/// → register(tool_2) -/// → ... -/// │ -/// ▼ -/// all_schemas() ──→ LLM API request (tool definitions) -/// get("name") ──→ Tool::call(input, ctx) (dispatch) -/// concurrent_safe_tools() ──→ parallel execution planner -/// ``` -/// -/// # Thread safety -/// -/// The registry itself is not `Sync` — it is created once during session -/// setup and then accessed immutably during tool dispatch. If you need -/// cross-thread sharing, wrap it in an `Arc>`. -/// -/// # Example -/// -/// ```rust,ignore -/// let mut registry = ToolRegistry::new(); -/// registry.register(ReadFileTool); -/// registry.register(WriteFileTool); -/// -/// // Dispatch an invocation -/// let tool = registry.get("read_file").expect("tool exists"); -/// let result = tool.call(input, &ctx).await; -/// -/// // Send schemas to the LLM -/// let schemas = registry.all_schemas(); -/// ``` -pub struct ToolRegistry { - /// Internal name → tool map. - /// - /// Each entry is a `Box` keyed by its [`Tool::name`]. - /// Populated by [`register`](ToolRegistry::register) and queried by - /// [`get`](ToolRegistry::get). - tools: HashMap>, -} - -impl ToolRegistry { - /// Create a new empty registry. - /// - /// The registry starts with no tools. Use [`register`](ToolRegistry::register) - /// to add tools before the agent loop begins processing turns. - #[must_use] - pub fn new() -> Self { - Self { - tools: HashMap::new(), - } - } - - /// Register a tool, replacing any previous tool with the same name. - /// - /// Called during session setup, before any turns are processed. If a - /// tool with the same [`Tool::name`] already exists it is silently - /// replaced. - /// - /// # Example - /// - /// ```rust,ignore - /// registry.register(ReadFileTool); - /// registry.register(WriteFileTool); - /// ``` - pub fn register(&mut self, tool: impl Tool + 'static) { - let name = tool.name().to_string(); - self.tools.insert(name, Box::new(tool)); - } - - /// Look up a tool by name. - /// - /// Returns `Some(&dyn Tool)` if a tool with the given name was - /// previously [`register`](ToolRegistry::register)ed, or `None` - /// otherwise. Called by the agent loop when dispatching an LLM tool - /// call. - /// - /// The returned reference borrows from the registry and is valid for - /// as long as the registry is alive. - /// - /// # Example - /// - /// ```rust,ignore - /// if let Some(tool) = registry.get("read_file") { - /// let result = tool.call(input, &ctx).await; - /// } - /// ``` - #[must_use] - pub fn get(&self, name: &str) -> Option<&dyn Tool> { - self.tools.get(name).map(std::convert::AsRef::as_ref) - } - - /// Check whether a tool with the given name is registered. - /// - /// Useful for pre-flight validation before attempting - /// [`get`](ToolRegistry::get). Returns `true` if the name maps to a - /// registered tool. - /// - /// # Example - /// - /// ```rust,ignore - /// if registry.contains("bash") { - /// // Safe to call registry.get("bash") - /// } - /// ``` - #[must_use] - pub fn contains(&self, name: &str) -> bool { - self.tools.contains_key(name) - } - - /// Collect [`ToolSchema`] descriptors for all registered tools. - /// - /// Called by the agent loop to build the tool list sent to the LLM - /// at the start of each session (or turn, if the tool set changes). - /// The order is unspecified. - /// - /// Each schema is freshly constructed via [`Tool::schema`], so the - /// caller does not need to worry about stale data. - /// - /// # Example - /// - /// ```rust,ignore - /// let schemas = registry.all_schemas(); - /// for schema in &schemas { - /// println!(" - {}: {}", schema.name, schema.description); - /// } - /// ``` - #[must_use] - pub fn all_schemas(&self) -> Vec { - self.tools.values().map(|t| t.schema()).collect() - } - - /// Return all registered tool names, sorted alphabetically. - /// - /// Useful for diagnostics, logging, and building error messages - /// in [`ToolError::not_found`]. - #[must_use] - pub fn tool_names(&self) -> Vec { - let mut names: Vec<_> = self.tools.keys().cloned().collect(); - names.sort(); - names - } - - /// Number of registered tools. - /// - /// Used by the framework and by - /// [`is_empty`](ToolRegistry::is_empty). Returns `0` for a freshly - /// created registry. - /// - /// # Example - /// - /// ```rust,ignore - /// assert_eq!(registry.len(), 3); // three tools registered - /// ``` - #[must_use] - pub fn len(&self) -> usize { - self.tools.len() - } - - /// Whether the registry contains no tools. - /// - /// Defaults to `self.len() == 0`. The agent loop typically checks - /// this during startup to ensure at least one tool is available. - /// - /// # Example - /// - /// ```rust,ignore - /// let registry = ToolRegistry::new(); - /// assert!(registry.is_empty()); - /// registry.register(MyTool); - /// assert!(!registry.is_empty()); - /// ``` - #[must_use] - pub fn is_empty(&self) -> bool { - self.tools.is_empty() - } - - /// Return references to all tools that are concurrency-safe. - /// - /// Filters by [`Tool::is_concurrency_safe`] returning `true`. Used - /// by the agent loop to decide which tools can be invoked in parallel - /// during a single turn. - #[must_use] - pub fn concurrent_safe_tools(&self) -> Vec<&dyn Tool> { - self.tools - .values() - .map(std::convert::AsRef::as_ref) - .filter(|t| t.is_concurrency_safe()) - .collect() - } -} - -impl Default for ToolRegistry { - /// Produce an empty registry (equivalent to [`ToolRegistry::new`]). - /// - /// Allows `ToolRegistry` to be used in contexts that require - /// [`Default`], such as struct initialization with `..Default::default()`. - /// - /// # Example - /// - /// ```rust - /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; - /// - /// let registry = ToolRegistry::default(); - /// assert!(registry.is_empty()); - /// ``` - fn default() -> Self { - Self::new() - } -} - -// =================================================== -// FnTool adapter -// =================================================== - -/// Type alias for an async tool function pointer. -/// -/// Matches the signature used by concrete tools in downstream crates: -/// `fn(Value, &ToolContext) -> Pin> + Send + 'static>>`. -/// -/// Stored in the `f` field of [`FnTool`] to adapt function-pointer-based -/// tool definitions to the [`Tool`] trait. -pub type ToolFn = - fn( - Value, - &ToolContext, - ) -> Pin> + Send + 'static>>; - -/// Type alias for a dynamic concurrency check function. -/// -/// Takes a reference to the tool input [`Value`] and returns `true` if -/// the tool is safe to run concurrently with that specific input. Used -/// by [`FnTool::with_concurrency_check`] to override the static -/// [`Tool::is_concurrency_safe`] flag on a per-call basis. -pub type ConcurrencyCheckFn = fn(&Value) -> bool; - -/// Adapter that wraps a function pointer as a [`Tool`] trait implementation. -/// -/// Use [`FnTool`] when you have a standalone async function that implements -/// tool logic and want to register it without defining a dedicated struct. -/// The adapter wraps the function pointer so it can be stored in a -/// [`ToolRegistry`] alongside any other [`Tool`] implementation. -/// -/// For complex tools with internal state, implement [`Tool`] directly on a -/// struct instead. -/// -/// # Builder API -/// -/// [`FnTool`] supports a builder pattern for optional properties: -/// -/// ```rust,ignore -/// let tool = FnTool::new("my_tool".into(), "Does a thing".into(), -/// json!({"type": "object", "properties": {"text": {"type": "string"}}}), -/// my_tool as ToolFn) -/// .concurrency_safe() // mark as safe for parallel execution -/// .read_only() // mark as side-effect-free -/// .with_system_prompt("...".into()); // inject extra LLM context -/// -/// let mut registry = ToolRegistry::new(); -/// registry.register(tool); -/// ``` -/// -/// # Builder flow -/// -/// ```text -/// FnTool::new(name, desc, schema, f) -/// │ -/// ├─ .concurrency_safe() → sets is_concurrency_safe = true -/// ├─ .with_concurrency_check(fn) → sets per-input check -/// ├─ .read_only() → sets is_read_only = true -/// └─ .with_system_prompt(s) → sets system_prompt = Some(s) -/// ``` -/// -/// # Example -/// -/// ```rust,ignore -/// fn my_tool(input: Value, _ctx: &ToolContext) -/// -> Pin> + Send + 'static>> -/// { -/// let text = input.get("text").unwrap().to_string(); -/// Box::pin(async move { Ok(ToolOutput::text(text)) }) -/// } -/// -/// let tool = FnTool::new("my_tool".into(), "Does a thing".into(), -/// json!({"type": "object", "properties": {"text": {"type": "string"}}}), -/// my_tool as ToolFn) -/// .concurrency_safe() -/// .read_only(); -/// -/// let mut registry = ToolRegistry::new(); -/// registry.register(tool); -/// ``` -pub struct FnTool { - /// The tool's unique name identifier. - /// - /// Must match the name used in the [`ToolSchema`] and serves as the - /// [`ToolRegistry`] lookup key. Set at construction time via - /// [`FnTool::new`]. - pub name: String, - - /// Human-readable description for the LLM. - /// - /// Sent to the LLM as part of the [`ToolSchema`]. A clear - /// description improves tool selection accuracy. - pub description: String, - - /// JSON Schema describing the tool's input parameters. - /// - /// Must be a valid JSON Schema object. Embedded in the - /// [`ToolSchema`] returned by [`Tool::schema`]. - pub input_schema: Value, - - /// The function pointer that implements the tool's core logic. - /// - /// Called by [`Tool::call`] with the LLM-supplied input and the - /// session's [`ToolContext`]. Must return a pinned, `Send` future - /// producing a `Result`. - pub tool_fn: ToolFn, - - /// Whether this tool is safe to run concurrently with itself. - /// - /// Set via the [`concurrency_safe`](FnTool::concurrency_safe) builder - /// method. Defaults to `false`. When `true`, the agent loop may - /// invoke this tool in parallel with other concurrent-safe tools. - pub is_concurrency_safe: bool, - - /// Optional dynamic concurrency check function. - /// - /// When set via [`with_concurrency_check`](FnTool::with_concurrency_check), - /// this function is called with the tool input to decide per-invocation - /// concurrency safety. Overrides the static - /// [`is_concurrency_safe`](FnTool::is_concurrency_safe) flag when present. - pub concurrency_check_fn: Option, - - /// Whether this tool only reads data (no side effects). - /// - /// Set via the [`read_only`](FnTool::read_only) builder method. - /// Defaults to `false`. Read-only tools can be auto-approved by - /// permission gates. - pub is_read_only: bool, - - /// Optional extra system prompt injected when this tool is available. - /// - /// Set via [`with_system_prompt`](FnTool::with_system_prompt). - /// The agent loop appends this to the system message. Defaults to - /// `None`. - pub system_prompt: Option, -} - -impl FnTool { - /// Create a new function-pointer tool with the given name, description, - /// schema, and implementation function. - /// - /// All optional properties default to their "off" values: - /// `is_concurrency_safe → false`, `concurrency_check_fn → None`, - /// `is_read_only → false`, `system_prompt → None`. Use the builder - /// methods to enable them. - /// - /// # Arguments - /// - /// | Argument | Type | Description | - /// |-----------------|-----------|------------------------------------------------| - /// | `name` | `String` | Unique tool identifier, used as registry key | - /// | `description` | `String` | Human-readable summary sent to the LLM | - /// | `input_schema` | `Value` | JSON Schema for the tool's parameters | - /// | `f` | [`ToolFn`] | The async function implementing tool logic | - /// - /// # Example - /// - /// ```rust,ignore - /// let tool = FnTool::new( - /// "grep".into(), - /// "Search files for a pattern".into(), - /// json!({"type": "object", "properties": {"pattern": {"type": "string"}}}), - /// my_grep_fn as ToolFn, - /// ); - /// ``` - pub fn new(name: String, description: String, input_schema: Value, tool_fn: ToolFn) -> Self { - Self { - name, - description, - input_schema, - tool_fn, - is_concurrency_safe: false, - concurrency_check_fn: None, - is_read_only: false, - system_prompt: None, - } - } - - /// Builder: mark this tool as concurrency-safe. - /// - /// Sets [`is_concurrency_safe`](FnTool::is_concurrency_safe) to - /// `true`, signalling that the agent loop may invoke this tool in - /// parallel with other concurrent-safe tools. - /// - /// # When to use - /// - /// Call this for tools that are pure functions or read-only — for - /// example, a file-reading tool or a math calculator. Do *not* call - /// this for tools that mutate shared state or write to the filesystem. - /// - /// # Example - /// - /// ```rust,ignore - /// let tool = FnTool::new(/* ... */) - /// .concurrency_safe(); - /// ``` - #[must_use] - pub fn concurrency_safe(mut self) -> Self { - self.is_concurrency_safe = true; - self - } - - /// Builder: set a dynamic concurrency check function. - /// - /// The provided function is called with the tool input on each - /// invocation. If it returns `true`, the tool may run concurrently - /// for that specific input. Overrides the static - /// [`is_concurrency_safe`](FnTool::is_concurrency_safe) flag. - /// - /// # Example - /// - /// ```rust,ignore - /// fn can_run_concurrently(input: &Value) -> bool { - /// // Only safe if writing to different files - /// input.get("append").is_none() - /// } - /// let tool = FnTool::new(/* ... */).with_concurrency_check(can_run_concurrently); - /// ``` - #[must_use] - pub fn with_concurrency_check(mut self, check_fn: ConcurrencyCheckFn) -> Self { - self.concurrency_check_fn = Some(check_fn); - self - } - - /// Builder: mark this tool as read-only (no side effects). - /// - /// Sets [`is_read_only`](FnTool::is_read_only) to `true`. Read-only - /// tools can be auto-approved by permission gates and are generally - /// safe to run without user confirmation. - /// - /// # When to use - /// - /// Call this for tools that only read data — file readers, search - /// tools, calculators. Do *not* call this for tools that write files, - /// execute commands, or modify external state. - /// - /// # Example - /// - /// ```rust,ignore - /// let tool = FnTool::new(/* ... */) - /// .read_only(); - /// ``` - #[must_use] - pub fn read_only(mut self) -> Self { - self.is_read_only = true; - self - } - - /// Builder: set an optional extra system prompt for this tool. - /// - /// The agent loop appends this string to the system message when the - /// tool is registered, giving the LLM additional context about how - /// to use the tool effectively. - /// - /// # When to use - /// - /// Use this when a tool benefits from usage hints or style guidance - /// — for example, a shell tool might set a prompt like "Prefer - /// single-line bash commands" to steer the LLM's behavior. - /// - /// # Example - /// - /// ```rust,ignore - /// let tool = FnTool::new(/* ... */) - /// .with_system_prompt("Always use absolute paths.".into()); - /// ``` - #[must_use] - pub fn with_system_prompt(mut self, prompt: String) -> Self { - self.system_prompt = Some(prompt); - self - } -} - -/// [`Tool`] trait implementation for [`FnTool`]. -/// -/// Delegates each trait method to the corresponding field or function -/// pointer stored in the [`FnTool`] adapter. This is the glue that lets -/// function-pointer-based tools participate in the trait system without -/// any wrapper overhead. -/// -/// # Delegation map -/// -/// | Trait method | Delegates to | -/// |--------------------------------------------|----------------------------------------------| -/// | [`Tool::name`] | [`FnTool::name`] field accessor | -/// | [`Tool::description`] | [`FnTool::description`] accessor | -/// | [`Tool::schema`] | Clones fields into [`ToolSchema`] | -/// | [`Tool::call`] | internal function pointer | -/// | [`Tool::is_concurrency_safe`] | [`FnTool::is_concurrency_safe`] | -/// | [`Tool::is_safe_for_concurrent_execution`] | [`FnTool::concurrency_check_fn`] or fallback | -/// | [`Tool::is_read_only`] | [`FnTool::is_read_only`] | -/// | [`Tool::system_prompt`] | [`FnTool::system_prompt`] clone | -impl Tool for FnTool { - /// Return the tool name from the [`FnTool::name`] field. - /// - /// This is a trivial field accessor — the name was set at construction - /// time via [`FnTool::new`] and does not change for the lifetime of - /// the adapter. The returned slice borrows from `self`. - fn name(&self) -> &str { - &self.name - } - - /// Return the tool description from the [`FnTool::description`] field. - /// - /// Like [`name`](FnTool::name), this is a field accessor for the value - /// provided at construction time. The description is sent to the LLM as - /// part of the [`ToolSchema`] to help it choose the right tool. - fn description(&self) -> &str { - &self.description - } - - /// Build a [`ToolSchema`] from the stored fields. - /// - /// Assembles the [`FnTool::name`], [`FnTool::description`], and - /// [`FnTool::input_schema`] into a [`ToolSchema`] suitable for - /// sending to the LLM. The fields are cloned so the returned schema - /// is independent of `self`. - /// - /// # When called - /// - /// Invoked by the agent loop when assembling the list of tool - /// definitions to send in an LLM API request. Typically called once - /// per session (or per turn if the tool set changes dynamically). - fn schema(&self) -> ToolSchema { - ToolSchema { - tool: self.name.clone(), - description: self.description.clone(), - input_schema: self.input_schema.clone(), - } - } - - /// Delegate execution to the stored function pointer. - /// - /// Invokes the stored function pointer with the provided `input` and `context`, - /// returning the pinned future directly. The function pointer owns - /// the full future lifecycle — the `'static` bound on [`ToolFn`] - /// ensures the future does not borrow from the tool adapter itself. - /// - /// # When called - /// - /// Called by the agent loop after the LLM selects this tool by name - /// and the permission gate (if any) returns - /// [`PermissionCheck::Allow`]. - fn call( - &self, - input: Value, - context: &ToolContext, - ) -> Pin> + Send + '_>> { - (self.tool_fn)(input, context) - } - - /// Return the static concurrency-safety flag. - /// - /// Reads [`FnTool::is_concurrency_safe`], which is set via the - /// [`concurrency_safe`](FnTool::concurrency_safe) builder method. - /// - /// This is a *static* flag — it does not consider the specific input. - /// For input-dependent checks, see - /// [`is_safe_for_concurrent_execution`](Tool::is_safe_for_concurrent_execution). - fn is_concurrency_safe(&self) -> bool { - self.is_concurrency_safe - } - - /// Dynamic concurrency check using the optional check function. - /// - /// If [`FnTool::concurrency_check_fn`] is set (via - /// [`with_concurrency_check`](FnTool::with_concurrency_check)), - /// delegates to it and returns its result. Otherwise falls back to - /// the static [`is_concurrency_safe`](Tool::is_concurrency_safe) flag. - /// - /// This allows tools to express fine-grained concurrency policies — - /// for example, allowing parallel reads to different files while - /// serializing writes to the same file. - fn is_safe_for_concurrent_execution(&self, input: &Value) -> bool { - self.concurrency_check_fn - .map_or(self.is_concurrency_safe, |f| f(input)) - } - - /// Return the read-only flag from [`FnTool::is_read_only`]. - /// - /// Set via the [`read_only`](FnTool::read_only) builder method. - /// When `true`, the agent loop's permission gate may auto-approve - /// invocations without prompting the user, since the tool has no - /// observable side effects. - fn is_read_only(&self) -> bool { - self.is_read_only - } - - /// Clone and return the optional system prompt from [`FnTool::system_prompt`]. - /// - /// The agent loop appends this string to the system message when the - /// tool is registered, giving the LLM additional context about how to - /// use the tool effectively. Returns `None` if no prompt was set via - /// [`with_system_prompt`](FnTool::with_system_prompt). - fn system_prompt(&self) -> Option { - self.system_prompt.clone() - } -} - // =================================================== // Tests // =================================================== diff --git a/src/tool/health.rs b/src/tool/health.rs index ed4ea87..3661a8b 100644 --- a/src/tool/health.rs +++ b/src/tool/health.rs @@ -1,6 +1,6 @@ //! Tool health monitoring, circuit breakers, and self-healing routing. //! -//! This module provides per-tool health tracking using lock-free atomic counters, +//! Per-tool health tracking using lock-free atomic counters, //! circuit-breaker state machines to prevent repeated calls to failing tools, //! and a registry that combines both into a unified health picture. A routing //! middleware uses the registry to redirect tool calls away from unhealthy tools @@ -587,7 +587,7 @@ impl ToolCircuitBreaker { /// health + breaker state /// - [`health_summary`](Self::health_summary) — snapshot for observability /// -/// Uses `Mutex` for the tool-name → stats/breaker maps. This is +/// Uses `Mutex` for the tool-name → stats/breaker maps. Only /// the cold path — locks are only taken when a new tool name is first seen. /// Poisoned mutex recovery follows the pattern: `unwrap_or_else(std::sync::PoisonError::into_inner)`. /// @@ -644,14 +644,11 @@ impl ToolHealthRegistry { } } - /// Create a new registry with custom circuit-breaker configuration. + /// Set custom circuit-breaker configuration. #[must_use] - pub fn with_config(config: CircuitBreakerConfig) -> Self { - Self { - stats: Mutex::new(HashMap::new()), - breakers: Mutex::new(HashMap::new()), - breaker_config: config, - } + pub fn with_config(mut self, config: CircuitBreakerConfig) -> Self { + self.breaker_config = config; + self } /// Get or create stats for a tool. @@ -1216,7 +1213,7 @@ mod tests { assert_eq!(registry.get_health_status("tool_a"), HealthStatus::Healthy); // Drive tool_b down with failures (with low threshold) - let low_threshold_registry = ToolHealthRegistry::with_config(CircuitBreakerConfig { + let low_threshold_registry = ToolHealthRegistry::new().with_config(CircuitBreakerConfig { failure_threshold: 2, recovery_duration: Duration::from_secs(30), }); @@ -1232,7 +1229,7 @@ mod tests { #[test] fn registry_is_tool_available() { - let registry = ToolHealthRegistry::with_config(CircuitBreakerConfig { + let registry = ToolHealthRegistry::new().with_config(CircuitBreakerConfig { failure_threshold: 2, recovery_duration: Duration::from_secs(30), }); @@ -1299,7 +1296,7 @@ mod tests { #[test] fn health_router_resolve_falls_back_to_alternative() { - let registry = ToolHealthRegistry::with_config(CircuitBreakerConfig { + let registry = ToolHealthRegistry::new().with_config(CircuitBreakerConfig { failure_threshold: 1, recovery_duration: Duration::from_secs(30), }); @@ -1319,7 +1316,7 @@ mod tests { #[test] fn health_router_resolve_returns_primary_when_no_healthy_alternative() { - let registry = ToolHealthRegistry::with_config(CircuitBreakerConfig { + let registry = ToolHealthRegistry::new().with_config(CircuitBreakerConfig { failure_threshold: 1, recovery_duration: Duration::from_secs(30), }); diff --git a/src/tool/permission.rs b/src/tool/permission.rs new file mode 100644 index 0000000..4d93a8b --- /dev/null +++ b/src/tool/permission.rs @@ -0,0 +1,247 @@ +//! Permission checking for tool invocations. +//! +//! [`PermissionCheck`] — the result type returned by +//! the agent loop's permission gate before a tool is allowed to execute. +//! See the [`PermissionCheck`] documentation for the full decision tree. + +use serde_json::Value; + +// =================================================== +// PermissionCheck +// =================================================== + +/// Result of a permission check before tool execution. +/// +/// Before invoking [`Tool::call`](super::Tool::call), the agent loop can run a +/// permission gate that returns one of four outcomes: allow, deny, ask the user, +/// or modify the input. This lets host applications enforce safety +/// policies without modifying individual tool implementations. +/// +/// # Example +/// +/// ```rust,ignore +/// let check = PermissionCheck::deny("dangerous operation"); +/// if check.is_deny() { +/// return Err(ToolError::Permission("blocked by policy".into())); +/// } +/// ``` +#[derive(Debug, Clone)] +pub enum PermissionCheck { + /// Allow the tool to execute unmodified. + /// + /// The agent loop proceeds with the original input and context. + Allow, + + /// Deny execution with a human-readable reason. + /// + /// The agent loop should return + /// [`ToolError::Permission`](super::ToolError::Permission) with the given + /// `reason` so the LLM can react accordingly. + Deny { + /// Explanation forwarded to the LLM as part of the error message. + reason: String, + }, + + /// Prompt the user for approval before proceeding. + /// + /// In interactive sessions the agent loop should present `prompt` to + /// the user and then treat the response as either [`Allow`](PermissionCheck::Allow) + /// or [`Deny`](PermissionCheck::Deny). + Ask { + /// Should clearly describe the action and potential side effects. + prompt: String, + }, + + /// Modify the tool's input before execution. + /// + /// The agent loop should invoke [`Tool::call`](super::Tool::call) with + /// `modified_input` instead of the original input. Useful for sanitising + /// paths, redacting secrets, or injecting default values. + Modify { + /// Must conform to the tool's [`ToolSchema::input_schema`](super::ToolSchema::input_schema). + modified_input: Value, + }, +} + +impl PermissionCheck { + /// Create an [`Allow`](PermissionCheck::Allow) result. + /// + /// Signals that the tool invocation may proceed without changes. + /// The `#[must_use]` attribute reminds callers to check the result + /// rather than silently discarding it. + /// + /// # When returned + /// + /// The permission gate returns this variant when the requested + /// operation is within the configured safety policy — for example, + /// a read-only tool invocation or an operation on an allowed path. + /// + /// # Example + /// + /// ```rust + /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; + /// + /// let check = PermissionCheck::allow(); + /// assert!(check.is_allow()); + /// ``` + #[must_use] + pub fn allow() -> Self { + Self::Allow + } + + /// Create a [`Deny`](PermissionCheck::Deny) result with a reason. + /// + /// The `reason` string will be forwarded to the LLM as part of the + /// error message, helping it understand why the invocation was + /// rejected and adjust its next action. + /// + /// # When returned + /// + /// The permission gate returns this variant when the requested + /// operation violates a hard safety rule — for example, executing + /// a shell command when shell access is disabled. + /// + /// # Example + /// + /// ```rust + /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; + /// + /// let check = PermissionCheck::deny("shell execution is disabled"); + /// assert!(check.is_deny()); + /// ``` + pub fn deny(reason: impl Into) -> Self { + Self::Deny { + reason: reason.into(), + } + } + + /// Create an [`Ask`](PermissionCheck::Ask) result with a prompt. + /// + /// The agent loop should present the `prompt` to the user (in + /// interactive mode) and then proceed based on the user's response. + /// + /// # When returned + /// + /// The permission gate returns this variant for operations that are + /// potentially dangerous but not outright prohibited — for example, + /// writing to a file for the first time. The user's decision is then + /// converted to [`Allow`](PermissionCheck::Allow) or + /// [`Deny`](PermissionCheck::Deny). + /// + /// # Example + /// + /// ```rust + /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; + /// + /// let check = PermissionCheck::ask("Allow write to /etc/config.yaml?"); + /// assert!(check.is_ask()); + /// ``` + pub fn ask(prompt: impl Into) -> Self { + Self::Ask { + prompt: prompt.into(), + } + } + + /// Create a [`Modify`](PermissionCheck::Modify) result with rewritten input. + /// + /// The agent loop should replace the original tool input with the + /// provided `modified_input` before invoking [`Tool::call`](super::Tool::call). + /// Useful for sanitising paths, redacting secrets, or injecting default + /// values. + /// + /// # When returned + /// + /// The permission gate returns this variant when the requested + /// operation is acceptable but the input needs adjustment — for + /// example, resolving a relative path to an absolute one within + /// the allowed directory tree. + /// + /// # Example + /// + /// ```rust + /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; + /// use serde_json::json; + /// + /// let check = PermissionCheck::modify(json!({"path": "/safe/dir/file.txt"})); + /// assert!(check.is_modify()); + /// ``` + #[must_use] + pub fn modify(modified_input: Value) -> Self { + Self::Modify { modified_input } + } + + /// Returns `true` if this is an [`Allow`](PermissionCheck::Allow). + /// + /// Convenience predicate for the most common happy-path check. + /// Used by the agent loop to test whether to proceed with + /// [`Tool::call`](super::Tool::call) without further processing. + /// + /// # Example + /// + /// ```rust,ignore + /// if check.is_allow() { + /// let result = tool.call(input, &ctx).await; + /// } + /// ``` + #[must_use] + pub fn is_allow(&self) -> bool { + matches!(self, Self::Allow) + } + + /// Returns `true` if this is a [`Deny`](PermissionCheck::Deny). + /// + /// When `true`, the agent loop should *not* invoke the tool and + /// should instead return a permission error to the LLM. The denial + /// reason can be extracted by destructuring the variant or by + /// converting to [`ToolError::Permission`](super::ToolError::Permission). + /// + /// # Example + /// + /// ```rust,ignore + /// if check.is_deny() { + /// return Err(ToolError::Permission("blocked by policy".into())); + /// } + /// ``` + #[must_use] + pub fn is_deny(&self) -> bool { + matches!(self, Self::Deny { .. }) + } + + /// Returns `true` if this is an [`Ask`](PermissionCheck::Ask). + /// + /// When `true`, the agent loop should prompt the user before + /// deciding whether to allow or deny the invocation. In + /// non-interactive mode ([`ToolContext::is_non_interactive`](super::ToolContext::is_non_interactive)), + /// the loop typically treats an [`Ask`](PermissionCheck::Ask) as a + /// [`Deny`](PermissionCheck::Deny). + /// + /// # Example + /// + /// ```rust,ignore + /// if check.is_ask() { + /// println!("Tool requests approval: {}", prompt); + /// } + /// ``` + #[must_use] + pub fn is_ask(&self) -> bool { + matches!(self, Self::Ask { .. }) + } + + /// Returns `true` if this is a [`Modify`](PermissionCheck::Modify). + /// + /// When `true`, the agent loop should replace the original input + /// with the modified version before calling the tool. The modified + /// input can be extracted by matching the variant. + /// + /// # Example + /// + /// ```rust,ignore + /// if let PermissionCheck::Modify { modified_input } = check { + /// let result = tool.call(modified_input, &ctx).await; + /// } + /// ``` + #[must_use] + pub fn is_modify(&self) -> bool { + matches!(self, Self::Modify { .. }) + } +} diff --git a/src/tool/registry.rs b/src/tool/registry.rs new file mode 100644 index 0000000..6034959 --- /dev/null +++ b/src/tool/registry.rs @@ -0,0 +1,508 @@ +//! Tool registry and function-pointer adapter. +//! +//! [`ToolRegistry`] for dynamic tool lookup by name, +//! and [`FnTool`] (along with [`ToolFn`] and [`ConcurrencyCheckFn`]) for +//! wrapping plain async function pointers as [`Tool`] trait +//! implementations. + +use serde_json::Value; +use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; + +use super::{Tool, ToolContext, ToolError, ToolOutput, ToolSchema}; + +// =================================================== +// ToolRegistry +// =================================================== + +/// Registry of available tools for dynamic lookup by name. +/// +/// The agent loop creates a [`ToolRegistry`] at session start, registers +/// all available tools via [`register`](ToolRegistry::register), and then +/// uses [`get`](ToolRegistry::get) to dispatch invocations when the LLM +/// selects a tool by name. The registry also provides bulk accessors for +/// tool schemas and concurrency-safe tool lists. +/// +/// # Thread safety +/// +/// The registry itself is not `Sync` — it is created once during session +/// setup and then accessed immutably during tool dispatch. If you need +/// cross-thread sharing, wrap it in an `Arc>`. +/// +/// # Example +/// +/// ```rust,ignore +/// let mut registry = ToolRegistry::new(); +/// registry.register(ReadFileTool); +/// registry.register(WriteFileTool); +/// +/// // Dispatch an invocation +/// let tool = registry.get("read_file").expect("tool exists"); +/// let result = tool.call(input, &ctx).await; +/// +/// // Send schemas to the LLM +/// let schemas = registry.all_schemas(); +/// ``` +pub struct ToolRegistry { + /// `Box` keyed by [`Tool::name`]. + tools: HashMap>, +} + +impl ToolRegistry { + /// Create a new empty registry. + /// + /// The registry starts with no tools. Use [`register`](ToolRegistry::register) + /// to add tools before the agent loop begins processing turns. + #[must_use] + pub fn new() -> Self { + Self { + tools: HashMap::new(), + } + } + + /// Register a tool, replacing any previous tool with the same name. + /// + /// Called during session setup, before any turns are processed. If a + /// tool with the same [`Tool::name`] already exists + /// it is silently replaced. + /// + /// # Example + /// + /// ```rust,ignore + /// registry.register(ReadFileTool); + /// registry.register(WriteFileTool); + /// ``` + pub fn register(&mut self, tool: impl Tool + 'static) { + let name = tool.name().to_string(); + self.tools.insert(name, Box::new(tool)); + } + + /// Look up a tool by name. + /// + /// Returns `Some(&dyn Tool)` if a tool with the given name was + /// previously [`register`](ToolRegistry::register)ed, or `None` + /// otherwise. Called by the agent loop when dispatching an LLM tool + /// call. + /// + /// The returned reference borrows from the registry and is valid for + /// as long as the registry is alive. + /// + /// # Example + /// + /// ```rust,ignore + /// if let Some(tool) = registry.get("read_file") { + /// let result = tool.call(input, &ctx).await; + /// } + /// ``` + #[must_use] + pub fn get(&self, name: &str) -> Option<&dyn Tool> { + self.tools.get(name).map(std::convert::AsRef::as_ref) + } + + /// Check whether a tool with the given name is registered. + /// + /// Useful for pre-flight validation before attempting + /// [`get`](ToolRegistry::get). Returns `true` if the name maps to a + /// registered tool. + /// + /// # Example + /// + /// ```rust,ignore + /// if registry.contains("bash") { + /// // Safe to call registry.get("bash") + /// } + /// ``` + #[must_use] + pub fn contains(&self, name: &str) -> bool { + self.tools.contains_key(name) + } + + /// Collect [`ToolSchema`] descriptors for all registered tools. + /// + /// Called by the agent loop to build the tool list sent to the LLM + /// at the start of each session (or turn, if the tool set changes). + /// The order is unspecified. + /// + /// Each schema is freshly constructed via [`Tool::schema`], + /// so the caller does not need to worry about stale data. + /// + /// # Example + /// + /// ```rust,ignore + /// let schemas = registry.all_schemas(); + /// for schema in &schemas { + /// println!(" - {}: {}", schema.name, schema.description); + /// } + /// ``` + #[must_use] + pub fn all_schemas(&self) -> Vec { + self.tools.values().map(|t| t.schema()).collect() + } + + /// Return all registered tool names, sorted alphabetically. + /// + /// Useful for diagnostics, logging, and building error messages + /// in [`ToolError::not_found`]. + #[must_use] + pub fn tool_names(&self) -> Vec { + let mut names: Vec<_> = self.tools.keys().cloned().collect(); + names.sort(); + names + } + + /// Number of registered tools. + /// + /// Used by the framework and by + /// [`is_empty`](ToolRegistry::is_empty). Returns `0` for a freshly + /// created registry. + /// + /// # Example + /// + /// ```rust,ignore + /// assert_eq!(registry.len(), 3); // three tools registered + /// ``` + #[must_use] + pub fn len(&self) -> usize { + self.tools.len() + } + + /// Whether the registry contains no tools. + /// + /// Defaults to `self.len() == 0`. The agent loop typically checks + /// this during startup to ensure at least one tool is available. + /// + /// # Example + /// + /// ```rust,ignore + /// let registry = ToolRegistry::new(); + /// assert!(registry.is_empty()); + /// registry.register(MyTool); + /// assert!(!registry.is_empty()); + /// ``` + #[must_use] + pub fn is_empty(&self) -> bool { + self.tools.is_empty() + } + + /// Return references to all tools that are concurrency-safe. + /// + /// Filters by [`Tool::is_concurrency_safe`] + /// returning `true`. Used by the agent loop to decide which tools can + /// be invoked in parallel during a single turn. + #[must_use] + pub fn concurrent_safe_tools(&self) -> Vec<&dyn Tool> { + self.tools + .values() + .map(std::convert::AsRef::as_ref) + .filter(|t| t.is_concurrency_safe()) + .collect() + } +} + +impl Default for ToolRegistry { + /// Produce an empty registry (equivalent to [`ToolRegistry::new`]). + /// + /// Allows `ToolRegistry` to be used in contexts that require + /// [`Default`], such as struct initialization with `..Default::default()`. + /// + /// # Example + /// + /// ```rust + /// use loopctl::tool::{ToolOutput, ToolError, ToolSchema, ToolContext, PermissionCheck, ToolRegistry}; + /// + /// let registry = ToolRegistry::default(); + /// assert!(registry.is_empty()); + /// ``` + fn default() -> Self { + Self::new() + } +} + +// =================================================== +// FnTool adapter +// =================================================== + +/// Type alias for an async tool function pointer. +/// +/// Matches the signature used by concrete tools in downstream crates: +/// `fn(Value, &ToolContext) -> Pin> + Send + 'static>>`. +/// +/// Stored in the `f` field of [`FnTool`] to adapt function-pointer-based +/// tool definitions to the [`Tool`] trait. +pub type ToolFn = + fn( + Value, + &ToolContext, + ) -> Pin> + Send + 'static>>; + +/// Type alias for a dynamic concurrency check function. +/// +/// Takes a reference to the tool input [`Value`] and returns `true` if +/// the tool is safe to run concurrently with that specific input. Used +/// by [`FnTool::with_concurrency_check`] to override the static +/// [`Tool::is_concurrency_safe`] flag +/// on a per-call basis. +pub type ConcurrencyCheckFn = fn(&Value) -> bool; + +/// Adapter that wraps a function pointer as a [`Tool`] trait implementation. +/// +/// Use [`FnTool`] when you have a standalone async function that implements +/// tool logic and want to register it without defining a dedicated struct. +/// The adapter wraps the function pointer so it can be stored in a +/// [`ToolRegistry`] alongside any other [`Tool`] implementation. +/// +/// For complex tools with internal state, implement [`Tool`] +/// directly on a struct instead. +/// +/// # Builder API +/// +/// [`FnTool`] supports a builder pattern for optional properties: +/// +/// ```rust,ignore +/// let tool = FnTool::new("my_tool".into(), "Does a thing".into(), +/// json!({"type": "object", "properties": {"text": {"type": "string"}}}), +/// my_tool as ToolFn) +/// .concurrency_safe() // mark as safe for parallel execution +/// .read_only() // mark as side-effect-free +/// .with_system_prompt("...".into()); // inject extra LLM context +/// +/// let mut registry = ToolRegistry::new(); +/// registry.register(tool); +/// ``` +/// +/// # Example +/// +/// ```rust,ignore +/// fn my_tool(input: Value, _ctx: &ToolContext) +/// -> Pin> + Send + 'static>> +/// { +/// let text = input.get("text").unwrap().to_string(); +/// Box::pin(async move { Ok(ToolOutput::text(text)) }) +/// } +/// +/// let tool = FnTool::new("my_tool".into(), "Does a thing".into(), +/// json!({"type": "object", "properties": {"text": {"type": "string"}}}), +/// my_tool as ToolFn) +/// .concurrency_safe() +/// .read_only(); +/// +/// let mut registry = ToolRegistry::new(); +/// registry.register(tool); +/// ``` +pub struct FnTool { + /// Must match the name used in [`ToolSchema`] and [`ToolRegistry`] lookup. + pub name: String, + /// Sent to the LLM as part of the [`ToolSchema`]. + pub description: String, + /// Must be a valid JSON Schema object. + pub input_schema: Value, + /// Called by [`Tool::call`] with the LLM-supplied + /// input and session's [`ToolContext`]. + pub tool_fn: ToolFn, + /// Set via [`concurrency_safe`](FnTool::concurrency_safe). Defaults to `false`. + pub is_concurrency_safe: bool, + /// When set, overrides the static [`is_concurrency_safe`](FnTool::is_concurrency_safe) flag. + pub concurrency_check_fn: Option, + /// Set via [`read_only`](FnTool::read_only). Defaults to `false`. + pub is_read_only: bool, + /// Set via [`with_system_prompt`](FnTool::with_system_prompt). Defaults to `None`. + pub system_prompt: Option, +} + +impl FnTool { + /// Create a new function-pointer tool with the given name, description, + /// schema, and implementation function. + /// + /// All optional properties default to their "off" values: + /// `is_concurrency_safe → false`, `concurrency_check_fn → None`, + /// `is_read_only → false`, `system_prompt → None`. Use the builder + /// methods to enable them. + /// + /// # Arguments + /// + /// - `name` — Unique tool identifier, used as the registry key. + /// - `description` — Human-readable summary sent to the LLM. + /// - `input_schema` — JSON Schema describing the tool's parameters. + /// - `tool_fn` — The async function implementing the tool logic. + /// + /// # Example + /// + /// ```rust,ignore + /// let tool = FnTool::new( + /// "grep".into(), + /// "Search files for a pattern".into(), + /// json!({"type": "object", "properties": {"pattern": {"type": "string"}}}), + /// my_grep_fn as ToolFn, + /// ); + /// ``` + pub fn new(name: String, description: String, input_schema: Value, tool_fn: ToolFn) -> Self { + Self { + name, + description, + input_schema, + tool_fn, + is_concurrency_safe: false, + concurrency_check_fn: None, + is_read_only: false, + system_prompt: None, + } + } + + /// Builder: mark this tool as concurrency-safe. + /// + /// Sets [`is_concurrency_safe`](FnTool::is_concurrency_safe) to + /// `true`, signalling that the agent loop may invoke this tool in + /// parallel with other concurrent-safe tools. + /// + /// # When to use + /// + /// Call this for tools that are pure functions or read-only — for + /// example, a file-reading tool or a math calculator. Do *not* call + /// this for tools that mutate shared state or write to the filesystem. + /// + /// # Example + /// + /// ```rust,ignore + /// let tool = FnTool::new(/* ... */) + /// .concurrency_safe(); + /// ``` + #[must_use] + pub fn concurrency_safe(mut self) -> Self { + self.is_concurrency_safe = true; + self + } + + /// Builder: set a dynamic concurrency check function. + /// + /// The provided function is called with the tool input on each + /// invocation. If it returns `true`, the tool may run concurrently + /// for that specific input. Overrides the static + /// [`is_concurrency_safe`](FnTool::is_concurrency_safe) flag. + /// + /// # Example + /// + /// ```rust,ignore + /// fn can_run_concurrently(input: &Value) -> bool { + /// // Only safe if writing to different files + /// input.get("append").is_none() + /// } + /// let tool = FnTool::new(/* ... */).with_concurrency_check(can_run_concurrently); + /// ``` + #[must_use] + pub fn with_concurrency_check(mut self, check_fn: ConcurrencyCheckFn) -> Self { + self.concurrency_check_fn = Some(check_fn); + self + } + + /// Builder: mark this tool as read-only (no side effects). + /// + /// Sets [`is_read_only`](FnTool::is_read_only) to `true`. Read-only + /// tools can be auto-approved by permission gates and are generally + /// safe to run without user confirmation. + /// + /// # When to use + /// + /// Call this for tools that only read data — file readers, search + /// tools, calculators. Do *not* call this for tools that write files, + /// execute commands, or modify external state. + /// + /// # Example + /// + /// ```rust,ignore + /// let tool = FnTool::new(/* ... */) + /// .read_only(); + /// ``` + #[must_use] + pub fn read_only(mut self) -> Self { + self.is_read_only = true; + self + } + + /// Builder: set an optional extra system prompt for this tool. + /// + /// The agent loop appends this string to the system message when the + /// tool is registered, giving the LLM additional context about how + /// to use the tool effectively. + /// + /// # When to use + /// + /// Use this when a tool benefits from usage hints or style guidance + /// — for example, a shell tool might set a prompt like "Prefer + /// single-line bash commands" to steer the LLM's behavior. + /// + /// # Example + /// + /// ```rust,ignore + /// let tool = FnTool::new(/* ... */) + /// .with_system_prompt("Always use absolute paths.".into()); + /// ``` + #[must_use] + pub fn with_system_prompt(mut self, prompt: String) -> Self { + self.system_prompt = Some(prompt); + self + } +} + +/// [`Tool`] trait implementation for [`FnTool`]. +/// +/// Delegates each trait method to the corresponding field or function +/// pointer stored in the [`FnTool`] adapter. +impl Tool for FnTool { + /// Returns the tool's unique identifier stored in [`name`](FnTool::name). + fn name(&self) -> &str { + &self.name + } + + /// Returns the human-readable description stored in + /// [`description`](FnTool::description). + fn description(&self) -> &str { + &self.description + } + + /// Builds a [`ToolSchema`] from the stored name, description, and + /// input schema. + fn schema(&self) -> ToolSchema { + ToolSchema { + tool: self.name.clone(), + description: self.description.clone(), + input_schema: self.input_schema.clone(), + } + } + + /// Delegates execution to the stored [`tool_fn`](FnTool::tool_fn) + /// function pointer, forwarding the input and context unchanged. + fn call( + &self, + input: Value, + context: &ToolContext, + ) -> Pin> + Send + '_>> { + (self.tool_fn)(input, context) + } + + /// Returns the static concurrency-safety flag set via + /// [`concurrency_safe`](FnTool::concurrency_safe). + fn is_concurrency_safe(&self) -> bool { + self.is_concurrency_safe + } + + /// Delegates to [`concurrency_check_fn`](FnTool::concurrency_check_fn) + /// when set, otherwise falls back to the static + /// [`is_concurrency_safe`](Tool::is_concurrency_safe) flag. + fn is_safe_for_concurrent_execution(&self, input: &Value) -> bool { + self.concurrency_check_fn + .map_or(self.is_concurrency_safe, |f| f(input)) + } + + /// Returns whether this tool only reads data and has no side effects, + /// as configured via [`read_only`](FnTool::read_only). + fn is_read_only(&self) -> bool { + self.is_read_only + } + + /// Returns the optional system prompt set via + /// [`with_system_prompt`](FnTool::with_system_prompt). + fn system_prompt(&self) -> Option { + self.system_prompt.clone() + } +} diff --git a/src/tool/shield.rs b/src/tool/shield.rs index 55e720c..5278a4b 100644 --- a/src/tool/shield.rs +++ b/src/tool/shield.rs @@ -1,6 +1,6 @@ //! Tool Safety Shield — multi-turn adversarial defense. //! -//! This module provides the [`ToolSafetyShield`] trait — a generic, +//! [`ToolSafetyShield`] trait — a generic, //! platform-agnostic boundary for evaluating tool call safety — and a //! reference [`UnixShield`] implementation that matches dangerous Unix //! shell patterns. @@ -356,18 +356,14 @@ impl UnixShield { } } - /// Create a shield with custom thresholds. + /// Set custom warn and block thresholds. /// /// Values are clamped to `[0.0, 1.0]`. #[must_use] - pub fn with_thresholds(warn: f32, block: f32) -> Self { - Self { - warn_threshold: warn.clamp(0.0, 1.0), - block_threshold: block.clamp(0.0, 1.0), - turn_history: Mutex::new(Vec::new()), - patterns: Self::unix_patterns(), - combination_rules: Self::unix_combination_rules(), - } + pub fn with_thresholds(mut self, warn: f32, block: f32) -> Self { + self.warn_threshold = warn.clamp(0.0, 1.0); + self.block_threshold = block.clamp(0.0, 1.0); + self } /// Create a builder for a shield with custom patterns and rules. From d4eef476e4782d5b579593e2667deeea4cd96653 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Sun, 21 Jun 2026 23:47:15 +1200 Subject: [PATCH 03/30] chore: update readme --- README.md | 96 ++++++++++++++++++++++++++++++++++--------------------- 1 file changed, 59 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index 62da2d0..b65d76e 100644 --- a/README.md +++ b/README.md @@ -17,25 +17,33 @@ and tool implementations; the framework handles the rest. | Module | Description | |--------|-------------| -| [`api_client`](https://docs.rs/loopctl/latest/loopctl/api_client/index.html) | `ApiClient` trait for LLM provider communication (streaming + non-streaming) | -| [`api_error`](https://docs.rs/loopctl/latest/loopctl/api_error/index.html) | API error types with retry classification | -| [`builder`](https://docs.rs/loopctl/latest/loopctl/builder/index.html) | Fluent builder API with type-state generics for compile-time safety | -| [`builtin`](https://docs.rs/loopctl/latest/loopctl/builtin/index.html) | Reference implementations: `InMemoryStore`, `LoggingObserver` | -| [`cancel`](https://docs.rs/loopctl/latest/loopctl/cancel/index.html) | Cooperative cancellation via `CancelSignal` (AtomicBool + tokio::Notify) | -| [`compact`](https://docs.rs/loopctl/latest/loopctl/compact/index.html) | Context compaction: `ContextCompactor` trait, `TruncatingCompactor`, `TokenSplitter` | -| [`core`](https://docs.rs/loopctl/latest/loopctl/core/index.html) | Core traits (`AgentCore`, `AgentObserver`, `AgentMemory`), config, error, and state types | -| [`engine`](https://docs.rs/loopctl/latest/loopctl/engine/index.html) | `BareLoop` — the default agent loop engine (stream → accumulate → dispatch tools → repeat) | -| [`loop_control`](https://docs.rs/loopctl/latest/loopctl/loop_control/index.html)| Loop detection, convergence detection, fallback model chains, and manager bundle | -| [`message`](https://docs.rs/loopctl/latest/loopctl/message/index.html) | Conversation types: `Message`, `MessagePart`, `ToolContent`, roles | -| [`stream`](https://docs.rs/loopctl/latest/loopctl/stream/index.html) | Streaming event types, accumulator, stop reasons, usage tracking | -| [`tool`](https://docs.rs/loopctl/latest/loopctl/tool/index.html) | `Tool` trait, `ToolRegistry`, `ToolSchema`, `ToolOutput`, `FnTool` adapter | -| [`testing`](https://docs.rs/loopctl/latest/loopctl/testing/index.html) | Mock API client, mock tools, and test fixture factories (feature-gated) | +| [`api`](https://docs.rs/loopctl/latest/loopctl/api/index.html) | `ApiClient` trait for LLM provider communication (streaming + non-streaming) | +| [`api::error`](https://docs.rs/loopctl/latest/loopctl/api/error/index.html) | API error types with retry classification | +| [`builder`](https://docs.rs/loopctl/latest/loopctl/builder/index.html) | Fluent builder API with type-state generics for compile-time safety | +| [`cancel`](https://docs.rs/loopctl/latest/loopctl/cancel/index.html) | Cooperative cancellation via `CancelSignal` (AtomicBool + tokio::Notify) | +| [`capabilities`](https://docs.rs/loopctl/latest/loopctl/capabilities/index.html) | Capability traits (`Observable`, `Detectable`, `Compactable`, etc.) | +| [`compact`](https://docs.rs/loopctl/latest/loopctl/compact/index.html) | Context compaction: `ContextCompactor` trait, `TruncatingCompactor`, `TokenSplitter` | +| [`config`](https://docs.rs/loopctl/latest/loopctl/config/index.html) | Session configuration (`LoopConfig`) | +| [`detection`](https://docs.rs/loopctl/latest/loopctl/detection/index.html) | Loop detection, convergence detection, `DetectionManager` | +| [`engine`](https://docs.rs/loopctl/latest/loopctl/engine/index.html) | `BareLoop` — the default agent loop engine (stream → accumulate → dispatch tools → repeat) | +| [`error`](https://docs.rs/loopctl/latest/loopctl/error/index.html) | Central `LoopError` enum for all framework operations | +| [`fallback`](https://docs.rs/loopctl/latest/loopctl/fallback/index.html) | Circuit-breaker pattern for automatic API model fallback (`FallbackManager`) | +| [`memory`](https://docs.rs/loopctl/latest/loopctl/memory/index.html) | `LoopMemory` trait and entry types; `memory::builtin` provides `InMemoryStore` | +| [`message`](https://docs.rs/loopctl/latest/loopctl/message/index.html) | Conversation types: `Message`, `MessagePart`, `ToolContent`, roles | +| [`middleware`](https://docs.rs/loopctl/latest/loopctl/middleware/index.html) | Tool dispatch pipeline: timeouts, permissions, output limits, unknown-tool handling | +| [`observer`](https://docs.rs/loopctl/latest/loopctl/observer/index.html) | `LoopObserver` trait and `ObserverHost` for lifecycle event observation | +| [`reflection`](https://docs.rs/loopctl/latest/loopctl/reflection/index.html) | Failure reflection and recovery strategies (`Reflector`, `RecoveryStrategy`) | +| [`runtime`](https://docs.rs/loopctl/latest/loopctl/runtime/index.html) | `LoopRuntime` — the default infrastructure bundle | +| [`stream`](https://docs.rs/loopctl/latest/loopctl/stream/index.html) | Streaming event types, accumulator, stop reasons, usage tracking | +| [`tool`](https://docs.rs/loopctl/latest/loopctl/tool/index.html) | `Tool` trait, `ToolRegistry`, `ToolSchema`, `ToolOutput`, `FnTool` adapter | +| [`hooks`](https://docs.rs/loopctl/latest/loopctl/hooks/index.html) | Bidirectional lifecycle control (allow/block/ask before tool use). *Requires `hooks` feature.* | +| [`testing`](https://docs.rs/loopctl/latest/loopctl/testing/index.html) | Mock API client, mock tools, and test fixture factories. *Requires `testing` feature.* | ## Quick Start ### Implement a Tool -```rust,ignore +```rust,no_run use loopctl::tool::{Tool, ToolContext, ToolOutput, ToolError, ToolSchema}; use serde_json::{Value, json}; use std::pin::Pin; @@ -70,30 +78,39 @@ impl Tool for EchoTool { ### Run an Agent Loop -```rust,ignore -use loopctl::engine::bare::BareLoop; +```rust,no_run +use loopctl::engine::BareLoop; use loopctl::tool::ToolRegistry; -use loopctl::core::types::AgentConfig; +use loopctl::config::LoopConfig; use std::sync::Arc; // 1. Bring your own API client (implements ApiClient trait) -let client = Arc::new(my_provider_client); +# struct MyClient; +# use loopctl::api::ApiClient; +# impl ApiClient for MyClient { +# fn model(&self) -> &str { "llm-70b" } +# fn stream_messages(&self, _req: loopctl::api::StreamRequest) +# -> std::pin::Pin> + Send>> { +# unimplemented!() +# } +# } +let client = Arc::new(MyClient); // 2. Register tools let mut registry = ToolRegistry::new(); -registry.register(EchoTool); +// registry.register(EchoTool); // 3. Configure -let config = AgentConfig { +let config = LoopConfig { max_turns: 50, - model: "gpt-4o".into(), + model: "llm-70b".into(), ..Default::default() }; // 4. Run let agent = BareLoop::new(client, registry, config); -let result = agent.run("Use the echo tool to say hello").await?; -println!("Completed in {} turns", result.total_turns); +// let result = agent.run("Use the echo tool to say hello").await?; +// println!("Completed in {} turns", result.total_turns); ``` ### Use the Testing Module @@ -103,36 +120,40 @@ println!("Completed in {} turns", result.total_turns); loopctl = { version = "0.1", features = ["testing"] } ``` -```rust,ignore +```rust,no_run use loopctl::testing::{MockApiClient, MockTool, test_config}; -use loopctl::engine::bare::BareLoop; +use loopctl::engine::BareLoop; use loopctl::tool::ToolRegistry; +use std::sync::Arc; -let mut client = MockApiClient::new(); -client.enqueue_response(/* ... */); +let mut client = MockApiClient::new("test-model"); +client = client.with_text_response("Hello from the mock"); let mut registry = ToolRegistry::new(); -registry.register(MockTool::new("demo")); +registry.register(MockTool::new("demo", "A demo tool")); let agent = BareLoop::new( - client.into_shared(), + Arc::new(client), registry, test_config(), ); -let result = agent.run("test input").await?; +// let result = agent.run("test input").await?; ``` ## Feature Flags -| Feature | Default | Description | -|-----------|---------|----------------------------------------| -| `testing` | No | Mock clients, tools, and test fixtures | +| Feature | Default | Depends on | Description | +|---------|---------|------------|-------------| +| `hooks` | No | — | Bidirectional lifecycle hooks (allow/block/ask before tool use, compaction) | +| `testing` | No | — | Mock clients, tools, and test fixtures | +| `tool_health` | No | — | Per-tool health monitoring, circuit breakers, and self-healing routing | +| `tool_shield` | No | `tool_health` | Tool permission shielding and access control | ## Architecture ```text ┌──────────────┐ - │ ApiClient │ ← you implement this + │ ApiClient │ └───────┬──────┘ │ ┌─────────────▼─────────────┐ @@ -145,9 +166,10 @@ let result = agent.run("test input").await?; └─────┬──────────┬──────────┘ │ │ ┌────────────▼──┐ ┌────▼───────────┐ - │ ToolRegistry │ │ Loop Control │ - │ (your tools) │ │ • convergence │ - └───────────────┘ │ • detection │ + │ ToolRegistry │ │ Detection & │ + │ (your tools) │ │ Fallback │ + └───────────────┘ │ • convergence │ + │ • loop detect │ │ • fallback │ └────────────────┘ ``` From a3bd73fa16ceaf3e69d1cad8813751e1731a4e85 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 22 Jun 2026 08:36:45 +1200 Subject: [PATCH 04/30] refactor: bera loop --- src/engine/bare.rs | 1246 +++++++++++------------------------ src/engine/bare/compact.rs | 12 +- src/engine/bare/dispatch.rs | 84 +-- src/engine/bare/emission.rs | 44 +- src/engine/bare/message.rs | 10 +- src/engine/bare/stream.rs | 3 +- src/engine/loop_core.rs | 146 ++++ src/runtime.rs | 407 ++++++++++-- 8 files changed, 968 insertions(+), 984 deletions(-) diff --git a/src/engine/bare.rs b/src/engine/bare.rs index d754aa1..f96deeb 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -61,7 +61,7 @@ //! let config = LoopConfig::default(); //! //! // 2. Create the loop -//! let agent = BareLoop::new(client, registry, config); +//! let mut agent = BareLoop::new(client, registry, config); //! //! // 3. Run //! let result = agent.run("Hello, agent!").await?; @@ -69,13 +69,18 @@ //! ``` use crate::api::ApiClient; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::{Duration, Instant}; + use crate::cancel::CancelSignal; -use crate::compact::{ContextManager, EnsureContextResult}; +use crate::compact::ContextManager; use crate::config::LoopConfig; -use crate::detection::{ConvergenceAction, DetectedPattern}; + use crate::error::LoopError; -use crate::engine::loop_core::SessionResult; +use crate::engine::loop_core::{LoopState, SessionResult, StopReason, ToolCall, TurnResult}; #[cfg(feature = "hooks")] use crate::hooks::HookAction; #[cfg(feature = "hooks")] @@ -89,12 +94,12 @@ use crate::hooks::context::{ use crate::message::{Message, MessagePart, Role, ToolContent}; use crate::middleware::{ToolDispatchContext, ToolPipeline, ToolPipelineBuilder}; use crate::observer::{ - ConvergenceDetectedContext, FallbackContext, LoopDetectedContext, ResponseContext, - StreamContext, StreamFailureContext, TurnEndContext, TurnStartContext, + FallbackContext, ResponseContext, StreamContext, StreamFailureContext, TurnEndContext, + TurnStartContext, }; use crate::reflection::{ - Correction, CorrectionResult, ExponentialBackoffRecovery, NoopReflector, RecoveryAction, - RecoveryStrategy, ReflectionContext, Reflector, + ExponentialBackoffRecovery, NoopReflector, RecoveryAction, RecoveryStrategy, ReflectionContext, + Reflector, }; use crate::runtime::LoopRuntime; use crate::stream::handler::StreamHandler; @@ -102,9 +107,6 @@ use crate::stream::{StreamAccumulator, StreamEvent, StreamStopReason, Usage}; #[cfg(feature = "tool_health")] use crate::tool::health::ToolHealthRegistry; use crate::tool::{PermissionCheck, ToolContext, ToolDispatchResult, ToolRegistry, ToolSchema}; -use std::sync::Arc; -use std::time::{Duration, Instant}; -use uuid::Uuid; // Phase submodules mod compact; @@ -137,13 +139,11 @@ mod stream; /// - [`new()`](BareLoop::new) — client + tools + config. /// - [`new_with_managers()`](BareLoop::new_with_managers) — full control, /// including a [`LoopRuntime`]. -/// - [`from_parts()`](BareLoop::from_parts) — re-assembles from the -/// output of `AgentBuilder::into_raw_parts()`. /// /// # Lifecycle /// /// ```text -/// new() / new_with_managers() / from_parts() +/// new() / new_with_managers() /// → run(user_input) /// → stream_turn() → dispatch_tools() → stream_turn() /// → … (repeat until end_turn or max_turns) @@ -161,7 +161,7 @@ mod stream; /// let registry = ToolRegistry::new(); /// let config = LoopConfig::default(); /// -/// let agent = BareLoop::new( +/// let mut agent = BareLoop::new( /// Arc::new(my_client), /// registry, /// config, @@ -183,13 +183,6 @@ pub struct BareLoop { /// by name in this registry and invokes it. tools: Arc, - /// Optional middleware pipeline wrapping tool dispatch. - /// - /// When `Some`, tool calls flow through the pipeline's middleware - /// chain (timeouts, output limiting, etc.) before reaching the - /// registry. When `None`, dispatches go directly to the registry. - pipeline: Option, - /// Session parameters (max turns, model, system prompt). /// /// See [`LoopConfig`] for the full set of options. @@ -197,23 +190,34 @@ pub struct BareLoop { /// Conversation history (system + user + assistant + tool results). /// - /// Grows over the session lifetime. Each call to [`run()`](BareLoop::run) + /// Grows over the session lifetime. Each call to [`run()`](crate::engine::loop_core::Loop::run) /// appends the user message, then alternates between assistant responses /// and tool-result messages until the model signals `end_turn`. conversation: Vec, - /// Manager bundle (fallback, loop detection, convergence, observers). + /// Framework runtime bundle — holds all cross-cutting infrastructure. /// - /// Holds the cross-cutting managers that govern session behaviour: + /// This is the single source of truth for: /// - /// - [`FallbackManager`] — circuit breaker that trips after repeated - /// API failures, switching to a fallback model and recovering - /// once the primary stabilises. - /// - [`DetectionManager`] — orchestrates loop detection (repeated - /// tool operations) and convergence detection (semantically - /// similar responses). + /// - [`FallbackManager`] — circuit breaker for API model fallback. + /// - [`DetectionManager`] — loop and convergence detection. + /// - [`ObserverHost`] — lifecycle event fan-out. + /// - Optional [`ToolPipeline`] — middleware pipeline for tool dispatch. + /// - Optional [`ContextManager`] — automatic context compaction. + /// - Optional [`StreamHandler`] — resilient streaming with retries. + /// - Optional [`HookExecutor`] — bidirectional lifecycle hooks. + /// - Optional [`ToolHealthRegistry`] — per-tool health tracking. /// /// Reset at the start of every session via [`LoopRuntime::reset_all`]. + /// + /// [`FallbackManager`]: crate::fallback::FallbackManager + /// [`DetectionManager`]: crate::detection::DetectionManager + /// [`ObserverHost`]: crate::observer::ObserverHost + /// [`ToolPipeline`]: crate::middleware::ToolPipeline + /// [`ContextManager`]: crate::compact::ContextManager + /// [`StreamHandler`]: crate::stream::handler::StreamHandler + /// [`HookExecutor`]: crate::hooks::HookExecutor + /// [`ToolHealthRegistry`]: crate::tool::health::ToolHealthRegistry managers: LoopRuntime, /// Failure analyser for tool errors. @@ -238,281 +242,23 @@ pub struct BareLoop { /// will wake up mid-stream when cancelled. cancelled: Arc, - /// Optional context manager for automatic compaction. - /// - /// When `Some`, the loop checks token usage after each turn and - /// triggers compaction when usage exceeds the configured threshold. - /// Compaction replaces the conversation messages, notifies observers - /// via [`LoopObserver::on_compaction`](crate::observer::LoopObserver::on_compaction), - /// and notifies observers via [`on_compaction`](crate::observer::LoopObserver::on_compaction). - context_manager: Option>, - - /// Optional stream handler for resilient streaming. - /// - /// When `Some`, replaces the inline [`stream_turn()`](BareLoop::stream_turn) - /// logic with the handler's retry, timeout, and fallback capabilities. - /// When `None`, streaming uses the basic inline logic with no retries. - stream_handler: Option, - - /// Ordered hook executor for lifecycle interception. - /// - /// When `Some`, the executor runs registered hooks before and after - /// tool dispatch, compaction, and session start/end. Hooks can - /// short-circuit with [`HookAction::Block`]. - /// [`HookAction::Ask`] is automatically downgraded to `Block` by the - /// executor in [`crate::hooks::Interactivity::Headless`] mode (the default). - /// When `None`, no lifecycle interception occurs. - /// - /// *Requires `hooks` feature.* - #[cfg(feature = "hooks")] - hook_executor: Option>, - - /// Per-tool health tracker with circuit breakers. - /// - /// When `Some`, records success/failure and latency for every tool - /// dispatch. Tools that exceed the failure threshold have their circuit - /// breaker opened, blocking subsequent calls until recovery. - /// - /// *Requires `tool_health` feature.* - #[cfg(feature = "tool_health")] - health_registry: Option>, -} - -// ================================================== -// Run-loop bookkeeping types -// ================================================== - -/// Accumulated token counts, tool-call count, and turn count for a session. -/// -/// Mutable state that flows through the [`run()`](BareLoop::run) loop. -/// Using a struct avoids scattering loose counters across the method. -#[derive(Default)] -struct SessionBudget { - input_tokens: u64, - output_tokens: u64, - total_tool_calls: usize, - turn_count: usize, -} - -impl SessionBudget { - /// Accumulate token usage from a single turn into the running totals. - /// - /// When the API reports [`Usage`] data (input and output token counts), - /// this method adds them to the session-wide accumulators using - /// saturating arithmetic to prevent overflow on very long sessions. - /// If `usage` is `None` (provider did not report counts), this is a - /// no-op. - fn accumulate_usage(&mut self, usage: Option<&Usage>) { - if let Some(u) = usage { - self.input_tokens = self.input_tokens.saturating_add(u64::from(u.input_tokens)); - self.output_tokens = self - .output_tokens - .saturating_add(u64::from(u.output_tokens)); - } - } -} - -/// Token counts for a single turn, captured before tool dispatch. -/// -/// Needed because [`SessionBudget::accumulate_usage`] mutates the running -/// totals, but the per-turn values must be reported separately to observers. -#[derive(Clone, Copy)] -struct TurnTokens { - input: u64, - output: u64, -} + /// Current lifecycle state, exposed via the [`Loop`](crate::engine::loop_core::Loop) trait. + state: LoopState, -impl TurnTokens { - /// Extract per-turn token counts from optional [`Usage`]. - /// - /// Returns a [`TurnTokens`] capturing the input and output token - /// counts for a single turn. When `usage` is `None` (provider did - /// not report counts), both fields default to `0`. + /// Session-level accumulator for turn counts, token usage, and tool calls. /// - /// Required because [`SessionBudget::accumulate_usage`] mutates - /// running totals in place, but the per-turn values must be reported - /// separately to observers. - fn from_usage(usage: Option<&Usage>) -> Self { - match usage { - Some(u) => Self { - input: u64::from(u.input_tokens), - output: u64::from(u.output_tokens), - }, - None => Self { - input: 0, - output: 0, - }, - } - } -} + /// Reused across turns in a single [`run()`](crate::engine::loop_core::Loop::run) call. Reset + /// to `SessionResult::default()` in [`initialize`](crate::engine::loop_core::Loop::initialize). + budget: SessionResult, -/// Per-turn context passed to helper methods during the [`run()`](BareLoop::run) loop. -/// -/// Bundles the zero-based turn index, wall-clock duration, and token -/// counts so that extracted methods don't need long parameter lists. -struct TurnContext { - idx: usize, - duration: Duration, - tokens: TurnTokens, -} - -/// Reason the session was aborted before normal completion. -/// -/// Used by [`abort_session`](BareLoop::abort_session) to select the -/// correct [`LoopError`] variant without string matching. -#[derive(Clone, Copy)] -enum AbortReason { - /// User or external signal requested cancellation. - Cancelled, - /// The turn budget was exhausted. - MaxTurnsExceeded, -} - -/// Aggregated session metrics passed to [`notify_session_end`](BareLoop::notify_session_end). -/// -/// Captures completion status, an [`EndReason`] discriminant, turn/token -/// counters, and wall-clock duration — everything a hook needs to log or -/// react to session termination without pulling data from other sources. -struct SessionEndInfo { - /// Whether the session completed normally. - success: bool, - /// Structured reason for the session end. - reason: EndReason, - /// Total turns executed. - total_turns: usize, - /// Total tokens consumed (input + output). - #[cfg_attr(not(feature = "hooks"), allow(dead_code))] - total_tokens: u64, - /// Wall-clock session duration in seconds. - duration_secs: u64, -} - -/// Discriminant for how a session terminated. -/// -/// Mapped to [`SessionEndReason`] inside the `#[cfg(feature = "hooks")]` -/// path so the enum itself remains feature-independent. -enum EndReason { - Complete, - Cancelled, - Error, - MaxTurns, + /// Session start time, set by [`initialize`](crate::engine::loop_core::Loop::initialize). + session_start: Option, } // ================================================== -// ToolCallInfo +// Run-loop helpers // ================================================== -/// Internal representation of a tool call extracted from a message. -/// -/// When the LLM emits a `tool_call` content part, the loop extracts -/// its fields into this struct for convenient passing to -/// [`dispatch_tools()`](BareLoop::dispatch_tools). -/// -/// External consumers interact with tool results via [`SessionResult`] or the -/// [`LoopObserver`](crate::observer::LoopObserver) callbacks. -/// -/// # Fields -/// -/// - [`id`](ToolCallInfo::id) — The unique identifier assigned by the -/// API to this tool call. Used to correlate the result back to the -/// request. -/// - [`name`](ToolCallInfo::name) — The tool name. Must match a tool -/// registered in the [`ToolRegistry`]. -/// - [`input`](ToolCallInfo::input) — The JSON input provided by the -/// model. Deserialized into a `serde_json::Value`. -#[derive(Debug, Clone)] -struct ToolCallInfo { - /// The tool call ID assigned by the API. - /// - /// Used to correlate the tool result message back to the original - /// tool call. Copied into [`ToolDispatchResult::tool_call_id`] after - /// execution. - id: String, - - /// The tool name requested by the model. - /// - /// Must exactly match a name returned by a registered - /// [`Tool::name()`](crate::tool::Tool::name). If no match is found, - /// a soft error result is produced instead of a hard error. - name: String, - - /// The tool input as a JSON value. - /// - /// Deserialized from the API's `tool_call` content part. Passed - /// directly to [`Tool::call()`](crate::tool::Tool::call). - input: serde_json::Value, -} - -impl ToolCallInfo { - /// Apply a [`Correction`] from the reflection system in-place. - /// - /// Modifies `self` according to the correction strategy: - /// - /// - [`InputFix`](crate::reflection::CorrectionType::InputFix) — replaces - /// `self.input` with `correction.modified_input` (if provided). - /// - [`ToolChange`](crate::reflection::CorrectionType::ToolChange) — - /// replaces `self.name` with `correction.alternative_tool` (if provided). - /// - Other types ([`Retry`](crate::reflection::CorrectionType::Retry), - /// [`ApproachChange`](crate::reflection::CorrectionType::ApproachChange), - /// [`Escalate`](crate::reflection::CorrectionType::Escalate)) — no - /// mutation needed; the retry proceeds with unchanged parameters. - /// - /// Returns a [`CorrectionResult`] indicating whether the correction - /// was applied, failed (e.g. missing fields), or skipped. - fn apply_correction( - &mut self, - correction: &Correction, - _prior_result: &ToolDispatchResult, - ) -> CorrectionResult { - use crate::reflection::CorrectionType; - match correction.correction_type { - CorrectionType::InputFix => { - if let Some(ref modified) = correction.modified_input { - tracing::debug!( - tool = %self.name, - "applying InputFix correction from reflector" - ); - self.input = modified.clone(); - CorrectionResult::Applied - } else { - CorrectionResult::Failed( - "InputFix correction missing modified_input".to_string(), - ) - } - } - CorrectionType::ToolChange => { - if let Some(ref alt) = correction.alternative_tool { - tracing::debug!( - old_tool = %self.name, - new_tool = %alt, - "applying ToolChange correction from reflector" - ); - self.name.clone_from(alt); - CorrectionResult::Applied - } else { - CorrectionResult::Failed( - "ToolChange correction missing alternative_tool".to_string(), - ) - } - } - CorrectionType::PrerequisiteFix | CorrectionType::ApproachChange => { - // These types do not modify the tool call parameters. - // PrerequisiteFix is advisory (the guidance may describe a - // side-effect to perform); ApproachChange is high-level - // guidance for a different strategy. The retry proceeds - // with unchanged input/tool. - CorrectionResult::Skipped - } - CorrectionType::Escalate => { - // Escalation means no correction is possible. The retry - // loop should not normally reach here because Escalate - // errors are mapped to RecoveryAction::Fail upstream. - CorrectionResult::Skipped - } - } - } -} - impl BareLoop { /// Maximum retry attempts for tool recovery before giving up. const MAX_RECOVERY_ATTEMPTS: u32 = 5; @@ -531,7 +277,7 @@ impl BareLoop { /// # Example /// /// ```rust,ignore - /// let agent = BareLoop::new( + /// let mut agent = BareLoop::new( /// Arc::new(my_client), /// ToolRegistry::new(), /// LoopConfig::default(), @@ -541,19 +287,15 @@ impl BareLoop { Self { client, tools: Arc::new(tools), - pipeline: None, config, conversation: Vec::new(), managers: LoopRuntime::new(), reflector: Arc::new(NoopReflector), recovery: Arc::new(ExponentialBackoffRecovery::new(3)), cancelled: Arc::new(CancelSignal::new()), - context_manager: None, - stream_handler: None, - #[cfg(feature = "hooks")] - hook_executor: None, - #[cfg(feature = "tool_health")] - health_registry: None, + state: LoopState::Idle, + budget: SessionResult::default(), + session_start: None, } } @@ -574,10 +316,11 @@ impl BareLoop { /// /// ```rust,ignore /// let managers = LoopRuntime::builder() - /// .with_loop_detection(10) + /// .with_detection(DetectionManager::default()) + /// .with_fallback(FallbackManager::default()) /// .build(); /// - /// let agent = BareLoop::new_with_managers( + /// let mut agent = BareLoop::new_with_managers( /// Arc::new(my_client), /// registry, /// config, @@ -593,63 +336,15 @@ impl BareLoop { Self { client, tools: Arc::new(tools), - pipeline: None, config, conversation: Vec::new(), managers, reflector: Arc::new(NoopReflector), recovery: Arc::new(ExponentialBackoffRecovery::new(3)), cancelled: Arc::new(CancelSignal::new()), - context_manager: None, - stream_handler: None, - #[cfg(feature = "hooks")] - hook_executor: None, - #[cfg(feature = "tool_health")] - health_registry: None, - } - } - - /// Create from builder parts (produced by `AgentBuilder::into_raw_parts()`). - /// - /// Most flexible constructor. It accepts all components - /// individually, making it suitable for re-assembly after a builder - /// has been consumed via `into_raw_parts()`. - /// - /// # Parameters - /// - /// - `client` — The LLM API client, wrapped in `Arc`. - /// - `tools` — The [`ToolRegistry`]. - /// - `managers` — A [`LoopRuntime`]. - /// - `config` — Session parameters. - /// - /// # Example - /// - /// ```rust,ignore - /// let (client, tools, managers, config) = builder.into_raw_parts(); - /// let agent = BareLoop::from_parts(client, tools, managers, config); - /// ``` - pub fn from_parts( - client: Arc, - tools: ToolRegistry, - managers: LoopRuntime, - config: LoopConfig, - ) -> Self { - Self { - client, - tools: Arc::new(tools), - pipeline: None, - config, - conversation: Vec::new(), - managers, - reflector: Arc::new(NoopReflector), - recovery: Arc::new(ExponentialBackoffRecovery::new(3)), - cancelled: Arc::new(CancelSignal::new()), - context_manager: None, - stream_handler: None, - #[cfg(feature = "hooks")] - hook_executor: None, - #[cfg(feature = "tool_health")] - health_registry: None, + state: LoopState::Idle, + budget: SessionResult::default(), + session_start: None, } } @@ -735,7 +430,7 @@ impl BareLoop { /// Set the [`Reflector`] for tool-error analysis. /// /// Replaces the default [`NoopReflector`] with a caller-supplied - /// implementation. Must be called before [`run()`](BareLoop::run). + /// implementation. Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// /// # Example /// @@ -751,7 +446,7 @@ impl BareLoop { /// /// Replaces the default [`ExponentialBackoffRecovery`] with a /// caller-supplied implementation. Must be called before - /// [`run()`](BareLoop::run). + /// [`run()`](crate::engine::loop_core::Loop::run). /// /// # Example /// @@ -767,7 +462,7 @@ impl BareLoop { /// /// When set, the loop checks token usage after each turn and /// triggers compaction when usage exceeds the configured threshold. - /// Must be called before [`run()`](BareLoop::run). + /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// /// # Example /// @@ -786,7 +481,7 @@ impl BareLoop { /// agent.set_context_manager(Arc::new(manager)); /// ``` pub fn set_context_manager(&mut self, manager: Arc) { - self.context_manager = Some(manager); + self.managers.set_context_manager(manager); } /// Set the [`StreamHandler`] for resilient streaming with retries, @@ -794,7 +489,7 @@ impl BareLoop { /// /// When set, the loop delegates streaming to the handler instead of /// using the inline streaming logic. Must be called before - /// [`run()`](BareLoop::run). + /// [`run()`](crate::engine::loop_core::Loop::run). /// /// # Example /// @@ -813,7 +508,7 @@ impl BareLoop { /// agent.set_stream_handler(handler); /// ``` pub fn set_stream_handler(&mut self, handler: StreamHandler) { - self.stream_handler = Some(handler); + self.managers.set_stream_handler(handler); } /// Set the [`HookExecutor`] for lifecycle interception. @@ -823,7 +518,7 @@ impl BareLoop { /// short-circuit with [`HookAction::Block`]. /// [`HookAction::Ask`] is automatically downgraded to `Block` by the /// executor in [`crate::hooks::Interactivity::Headless`] mode (the default). - /// Must be called before [`run()`](BareLoop::run). + /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// /// *Requires `hooks` feature.* /// @@ -839,7 +534,7 @@ impl BareLoop { /// ``` #[cfg(feature = "hooks")] pub fn set_hook_executor(&mut self, executor: Arc) { - self.hook_executor = Some(executor); + self.managers.set_hook_executor(executor); } /// Set the [`ToolHealthRegistry`] for per-tool health tracking. @@ -847,7 +542,7 @@ impl BareLoop { /// When set, records success/failure and latency for every tool /// dispatch. Tools that exceed the failure threshold have their /// circuit breaker opened, blocking subsequent calls until recovery. - /// Must be called before [`run()`](BareLoop::run). + /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// /// *Requires `tool_health` feature.* /// @@ -863,7 +558,7 @@ impl BareLoop { /// ``` #[cfg(feature = "tool_health")] pub fn set_health_registry(&mut self, registry: Arc) { - self.health_registry = Some(registry); + self.managers.set_health_registry(registry); } /// Set the middleware pipeline for tool dispatch. @@ -871,7 +566,7 @@ impl BareLoop { /// Replaces the default (no pipeline) with a caller-supplied /// [`ToolPipeline`]. When set, tool calls flow through the /// pipeline's middleware chain before reaching the registry. - /// Must be called before [`run()`](BareLoop::run). + /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// /// Build the pipeline using [`ToolPipeline::builder()`], adding middleware /// layers **without** calling `.core()` — the registry is injected @@ -899,7 +594,7 @@ impl BareLoop { .core(Arc::clone(&self.tools)) .build() .map_err(|e| LoopError::Config(e.to_string()))?; - self.pipeline = Some(pipeline); + self.managers.set_pipeline(pipeline); Ok(()) } @@ -909,7 +604,7 @@ impl BareLoop { /// in registration order. See [`LoopObserver`](crate::observer::LoopObserver) /// for the trait definition and available hooks. /// - /// Must be called before [`run()`](BareLoop::run). + /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// /// # Example /// @@ -925,177 +620,125 @@ impl BareLoop { } // ================================================== - // Main run loop + // Run helpers // ================================================== - /// Run the agent loop with the given user input. - /// - /// Primary entry point. It: - /// 1. Pushes the user message into the conversation - /// 2. Loops: stream → accumulate → tool dispatch → feedback - /// 3. Returns a [`SessionResult`] when done - /// - /// The loop terminates when one of these conditions is met: - /// - /// - **End turn** — the model emits `end_turn` with no tool calls. - /// - **Max turns exceeded** — [`config.max_turns`](LoopConfig::max_turns) - /// is reached, producing [`LoopError::MaxTurnsExceeded`]. - /// - **Cancellation** — [`cancel()`](BareLoop::cancel) was called, - /// producing [`LoopError::Cancelled`]; the caller should handle this - /// variant to distinguish user-initiated cancellation from other errors. - /// - **API error** — the streaming request fails, producing - /// [`LoopError::Api`]. - /// - /// # Observers - /// - /// Observers are notified at the following points: - /// - /// ```text - /// on_session_start(session_id) - /// for each turn: - /// on_turn_start(user_input) - /// [stream from API] - /// for each tool_call: - /// on_tool_call(name, input) - /// on_tool_complete(name, input, output, duration, success, error) - /// on_turn_end(success, error) - /// on_session_end(success, error) - /// ``` - /// - /// # Errors + /// Extract per-turn token counts from optional [`Usage`]. + fn usage_tokens(usage: Option<&Usage>) -> (u64, u64) { + match usage { + Some(u) => (u64::from(u.input_tokens), u64::from(u.output_tokens)), + None => (0, 0), + } + } + + /// Dispatch tool calls, push the result message, and record the count. /// - /// Returns [`LoopError`] if: - /// - The API call fails (after any retries) - /// - Max turns is exceeded - /// - A tool execution fails critically - /// - The loop is cancelled + /// Notifies observers: turn-end on success, turn-end on error. /// - /// # Example + /// # Errors /// - /// ```rust,ignore - /// let result = agent.run("Summarize this article").await?; - /// if result.success { - /// println!("Output: {:?}", result.final_output); - /// println!("Turns: {}", result.total_turns); - /// println!("Input tokens: {}", result.input_tokens); - /// println!("Output tokens: {}", result.output_tokens); - /// } - /// ``` - pub async fn run(mut self, user_input: &str) -> Result { - let session_id = self.config.session_id; - let max_turns = self.config.max_turns; - let start = Instant::now(); - let mut budget = SessionBudget::default(); + /// Returns [`LoopError::Cancelled`] if the cancellation token is set. + /// Returns [`LoopError::Api`] if loop or convergence detection forces + /// an abort, or if the underlying tool dispatch fails. + async fn dispatch_and_record( + &mut self, + tool_calls: &[ToolCall], + turn_index: usize, + turn_duration: Duration, + turn_input_tokens: u64, + turn_output_tokens: u64, + budget: &mut SessionResult, + ) -> Result<(), LoopError> { + match self.dispatch_tools(tool_calls, turn_index).await { + Ok(results) => { + budget.tool_calls = budget.tool_calls.saturating_add(results.len()); + let tool_result_msg = Self::build_tool_result_message(results); + self.conversation.push(tool_result_msg); + self.managers.observers().on_turn_end(&TurnEndContext { + turn: budget.total_turns, + success: true, + error: None, + duration_ms: Self::millis_u64(turn_duration), + input_tokens: turn_input_tokens, + output_tokens: turn_output_tokens, + }); + Ok(()) + } + Err(e) => { + let err_str = e.to_string(); + self.managers.observers().on_turn_end(&TurnEndContext { + turn: budget.total_turns, + success: false, + error: Some(err_str), + duration_ms: Self::millis_u64(turn_duration), + input_tokens: turn_input_tokens, + output_tokens: turn_output_tokens, + }); + Err(e) + } + } + } +} - self.notify_session_start(); - self.managers.reset_all(); - self.conversation.push(Message::user(user_input)); +// ================================================== +// Loop trait implementation +// ================================================== - loop { - if self.is_cancelled() { - return self.abort_session(&budget, start.elapsed(), AbortReason::Cancelled); - } - if budget.turn_count >= max_turns { - return self.abort_session(&budget, start.elapsed(), AbortReason::MaxTurnsExceeded); +impl crate::engine::loop_core::Loop for BareLoop { + fn initialize<'a>( + &'a mut self, + config: &'a crate::config::LoopConfig, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + self.state = LoopState::Processing { turn: 0 }; + self.budget = SessionResult::default(); + self.session_start = Some(Instant::now()); + self.config = config.clone(); + self.managers.reset_all(); + self.notify_session_start(); + Ok(()) + }) + } + + #[allow(clippy::too_many_lines)] + fn process_turn<'a>( + &'a mut self, + input: &'a str, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + // On the first turn, push the user's message. + if self.budget.total_turns == 0 { + self.conversation.push(Message::user(input)); } + let turn_start = Instant::now(); + self.managers.observers().on_turn_start(&TurnStartContext { - turn: budget.turn_count, - query: user_input.to_string(), + turn: self.budget.total_turns, + query: input.to_string(), }); - let turn_start = Instant::now(); - match self.stream_turn().await { - Ok((assistant_msg, usage, stop_reason)) => { + // Check cancellation before the API call. + if self.is_cancelled() { + self.state = LoopState::Failed { + error: "cancelled".into(), + }; + return Err(LoopError::Cancelled); + } + + let stream_result = self.stream_turn().await; + let (assistant_msg, usage, _stream_stop) = match stream_result { + Ok(value) => { + let (msg, usage, stop) = value; self.managers.fallback.record_model_success(); + let (in_tok, out_tok) = Self::usage_tokens(usage.as_ref()); self.managers.observers().on_stream_success(&StreamContext { - turn: budget.turn_count, + turn: self.budget.total_turns, model: self.client.model().to_string(), - input_tokens: TurnTokens::from_usage(usage.as_ref()).input, - output_tokens: TurnTokens::from_usage(usage.as_ref()).output, - }); - - budget.accumulate_usage(usage.as_ref()); - let text = Self::extract_text(&assistant_msg); - let pattern = self.managers.detection.record_response(&text); - - self.managers.observers().on_response(&ResponseContext { - turn: budget.turn_count, - text: text.clone(), - usage, + input_tokens: in_tok, + output_tokens: out_tok, }); - - if let Some(result) = self.handle_detected_pattern(&pattern, budget.turn_count) - { - let turn_elapsed = turn_start.elapsed(); - return match result { - Ok(session_result) => { - self.managers.observers().on_turn_end(&TurnEndContext { - turn: budget.turn_count, - success: true, - error: None, - duration_ms: Self::millis_u64(turn_start.elapsed()), - input_tokens: budget.input_tokens, - output_tokens: budget.output_tokens, - }); - self.notify_session_end(&SessionEndInfo { - success: true, - reason: EndReason::Complete, - total_turns: budget.turn_count, - total_tokens: budget - .input_tokens - .saturating_add(budget.output_tokens), - duration_secs: start.elapsed().as_secs(), - }); - Ok(session_result) - } - Err(e) => self.abort_turn_and_session( - &budget, - turn_elapsed, - start.elapsed(), - &e.to_string(), - e, - ), - }; - } - - let tool_calls = Self::extract_tool_calls(&assistant_msg); - self.conversation.push(assistant_msg); - budget.turn_count = budget.turn_count.saturating_add(1); - - let turn = TurnContext { - idx: budget.turn_count.saturating_sub(1), - duration: turn_start.elapsed(), - tokens: TurnTokens::from_usage(usage.as_ref()), - }; - - if tool_calls.is_empty() { - return Ok(self.finalise_session( - session_id, - text, - stop_reason, - &turn, - start.elapsed(), - &budget, - )); - } - - if let Err(e) = self - .dispatch_and_record(&tool_calls, &turn, &mut budget) - .await - { - return self.abort_session_from_error(e, start.elapsed(), &budget); - } - - if let Err(e) = self.maybe_compact_context(budget.turn_count).await { - return self.abort_turn_and_session( - &budget, - turn_start.elapsed(), - start.elapsed(), - &e.to_string(), - e, - ); - } + (msg, usage, stop) } Err(e) => { let tripped = self.managers.fallback.record_api_failure(); @@ -1113,349 +756,228 @@ impl BareLoop { self.managers .observers() .on_stream_failure(&StreamFailureContext { - turn: budget.turn_count, + turn: self.budget.total_turns, model: self.client.model().to_string(), error: e.clone(), }); - let err_str = e.to_string(); - return self.abort_turn_and_session( - &budget, - turn_start.elapsed(), - start.elapsed(), - &err_str, - e, - ); + self.state = LoopState::Failed { + error: e.to_string(), + }; + return Err(e); } + }; + + // Accumulate usage into session budget. + if let Some(u) = &usage { + self.budget.input_tokens = self + .budget + .input_tokens + .saturating_add(u64::from(u.input_tokens)); + self.budget.output_tokens = self + .budget + .output_tokens + .saturating_add(u64::from(u.output_tokens)); } - } - } - - // ================================================== - // Run helpers - // ================================================== - - /// Interpret a detected pattern and decide whether to abort the session. - /// - /// Called from both the run-loop convergence check (after extracting - /// assistant text) and the dispatch path (after recording a tool call). - /// Returns `None` to continue the loop, or `Some(result)` to abort - /// the session immediately. - /// - /// # Detection behaviour - /// - /// - **Loop detection** — only aborts when the repetition count reaches - /// the configured [`DetectionConfig::stop_threshold`]. Below that - /// threshold the pattern is logged as a warning and the loop continues. - /// - **Convergence detection** — follows the [`ConvergenceAction`] from - /// the detection config. `Stop` and `AskUser` abort the session; - /// `Warn`, `Compact`, and `SwitchPhase` allow the loop to continue. - /// - /// # Arguments - /// - /// - `pattern` — The pattern returned by [`DetectionManager`]. - /// - `turn` — Zero-based turn index, used in log messages. - /// - /// [`DetectionConfig::stop_threshold`]: crate::detection::DetectionConfig::stop_threshold - /// [`DetectionManager`]: crate::detection::DetectionManager - fn handle_detected_pattern( - &self, - pattern: &DetectedPattern, - turn: usize, - ) -> Option> { - match pattern { - DetectedPattern::NoPattern => None, - - DetectedPattern::LoopDetected { - repetitions, - pattern_description, - } => { - tracing::warn!( - repetitions, - pattern = %pattern_description, - turn, - "loop detected" - ); - - self.managers - .observers() - .on_loop_detected(&LoopDetectedContext { - pattern: pattern_description.clone(), - repetitions: *repetitions, - }); - if *repetitions >= self.managers.detection.config().stop_threshold { - tracing::error!( - repetitions, - pattern = %pattern_description, - turn, - "stopping agent: loop threshold exceeded" - ); - Some(Err(LoopError::LoopDetected { - message: format!("{pattern_description} repeated {repetitions} times"), - })) - } else { - None // warn but continue - } - } + let text = Self::extract_text(&assistant_msg); + let (turn_in, turn_out) = Self::usage_tokens(usage.as_ref()); - DetectedPattern::ConvergenceDetected { - similarity, - consecutive_count, - } => { - tracing::warn!(similarity, consecutive_count, turn, "convergence detected"); - let action = self.managers.detection.config().on_converge; - let action_str = match action { - ConvergenceAction::Stop => "stop", - ConvergenceAction::Warn => "warn", - ConvergenceAction::Compact => "compact", - ConvergenceAction::AskUser => "ask_user", - ConvergenceAction::SwitchPhase => "switch_phase", - }; + // Record the response text with the detection manager and check + // for loop/convergence patterns. + let pattern = self.managers.detection.record_response(&text); - self.managers - .observers() - .on_convergence_detected(&ConvergenceDetectedContext { - action: action_str.to_string(), - }); + self.managers.observers().on_response(&ResponseContext { + turn: self.budget.total_turns, + text: text.clone(), + usage, + }); - match action { - ConvergenceAction::Stop => Some(Err(LoopError::LoopDetected { - message: "agent stopped: convergence detected".into(), - })), - ConvergenceAction::Warn => None, // log already happened - ConvergenceAction::Compact => { - // Compact is handled by the existing compaction path, - // so just continue. The compact will happen at the - // end of the turn via maybe_compact_context(). - None - } - ConvergenceAction::AskUser => { - // Not supported in BareLoop — treat as Stop - Some(Err(LoopError::LoopDetected { - message: "agent stopped: convergence detected, user input needed" - .into(), - })) + if let Some(result) = self + .managers + .handle_detected_pattern(&pattern, self.budget.total_turns) + { + match result { + Ok(_) => { + self.state = LoopState::Completed { + summary: text.clone(), + }; + return Ok(TurnResult { + text, + tool_calls: Vec::new(), + tool_results: Vec::new(), + input_tokens: turn_in, + output_tokens: turn_out, + duration: turn_start.elapsed(), + is_complete: true, + stop_reason: StopReason::EndTurn, + }); } - ConvergenceAction::SwitchPhase => { - // Not supported in BareLoop — treat as Warn - None + Err(e) => { + self.state = LoopState::Failed { + error: e.to_string(), + }; + return Err(e); } } } - } - } - /// Dispatch tool calls, push the result message, and record the count. - /// - /// Notifies observers: turn-end on success, turn-end on error. - /// - /// # Errors - /// - /// Returns [`LoopError::Cancelled`] if the cancellation token is set. - /// Returns [`LoopError::Api`] if loop or convergence detection forces - /// an abort, or if the underlying tool dispatch fails. - async fn dispatch_and_record( - &mut self, - tool_calls: &[ToolCallInfo], - turn: &TurnContext, - budget: &mut SessionBudget, - ) -> Result<(), LoopError> { - match self.dispatch_tools(tool_calls, turn.idx).await { - Ok(results) => { - budget.total_tool_calls = budget.total_tool_calls.saturating_add(results.len()); - let tool_result_msg = Self::build_tool_result_message(results); - self.conversation.push(tool_result_msg); + let tool_calls = Self::extract_tool_calls(&assistant_msg); + self.conversation.push(assistant_msg); + self.budget.total_turns = self.budget.total_turns.saturating_add(1); + + // No tool calls → this turn is complete. + if tool_calls.is_empty() { self.managers.observers().on_turn_end(&TurnEndContext { - turn: budget.turn_count, + turn: self.budget.total_turns.saturating_sub(1), success: true, error: None, - duration_ms: Self::millis_u64(turn.duration), - input_tokens: turn.tokens.input, - output_tokens: turn.tokens.output, + duration_ms: Self::millis_u64(turn_start.elapsed()), + input_tokens: self.budget.input_tokens, + output_tokens: self.budget.output_tokens, }); - Ok(()) - } - Err(e) => { - let err_str = e.to_string(); - self.managers.observers().on_turn_end(&TurnEndContext { - turn: budget.turn_count, - success: false, - error: Some(err_str), - duration_ms: Self::millis_u64(turn.duration), - input_tokens: turn.tokens.input, - output_tokens: turn.tokens.output, + self.state = LoopState::Completed { + summary: text.clone(), + }; + return Ok(TurnResult { + text, + tool_calls: Vec::new(), + tool_results: Vec::new(), + input_tokens: turn_in, + output_tokens: turn_out, + duration: turn_start.elapsed(), + is_complete: true, + stop_reason: StopReason::EndTurn, }); - Err(e) } + + // Dispatch tool calls. + self.state = LoopState::WaitingForTool { + tool: tool_calls + .first() + .map(|tc| tc.tool.clone()) + .unwrap_or_default(), + started_at: std::time::SystemTime::now(), + }; + + // Temporarily extract budget to avoid double mutable borrow. + let mut budget = std::mem::take(&mut self.budget); + let turn_index = budget.total_turns.saturating_sub(1); + let turn_duration = turn_start.elapsed(); + if let Err(e) = self + .dispatch_and_record( + &tool_calls, + turn_index, + turn_duration, + turn_in, + turn_out, + &mut budget, + ) + .await + { + self.budget = budget; + self.state = LoopState::Failed { + error: e.to_string(), + }; + return Err(e); + } + self.budget = budget; + + // Attempt context compaction. + if let Err(e) = self.maybe_compact_context(self.budget.total_turns).await { + self.state = LoopState::Failed { + error: e.to_string(), + }; + return Err(e); + } + + self.state = LoopState::Processing { + turn: self.budget.total_turns, + }; + + Ok(TurnResult { + text, + tool_calls, + tool_results: Vec::new(), + input_tokens: turn_in, + output_tokens: turn_out, + duration: turn_start.elapsed(), + is_complete: false, + stop_reason: StopReason::ToolCall, + }) + }) + } + + fn should_continue(&self) -> bool { + if self.is_cancelled() { + return false; } + self.budget.total_turns < self.config.max_turns } - /// Build the final [`SessionResult`] when the model ends its turn. - /// - /// Called when streaming completes with no tool calls. Notifies - /// turn-end and session-end events, notifies observers, - /// and returns the assembled result. - fn finalise_session( - &self, - session_id: Uuid, - text: String, - stop_reason: StreamStopReason, - turn: &TurnContext, - session_duration: Duration, - budget: &SessionBudget, - ) -> SessionResult { - let success = stop_reason == StreamStopReason::EndTurn; - let error = if success { - None - } else { - Some(format!("Stream stopped with reason: {stop_reason:?}")) - }; - - self.managers.observers().on_turn_end(&TurnEndContext { - turn: turn.idx, - success, - error: error.as_deref().map(std::string::ToString::to_string), - duration_ms: Self::millis_u64(turn.duration), - input_tokens: turn.tokens.input, - output_tokens: turn.tokens.output, - }); - - let end_reason = if success { - EndReason::Complete - } else { - EndReason::Error - }; - self.notify_session_end(&SessionEndInfo { - success, - reason: end_reason, - total_turns: budget.turn_count, - total_tokens: budget.input_tokens.saturating_add(budget.output_tokens), - duration_secs: session_duration.as_secs(), - }); - - SessionResult { - session_id, - total_turns: budget.turn_count, - input_tokens: budget.input_tokens, - output_tokens: budget.output_tokens, - total_duration: session_duration, - tool_calls: budget.total_tool_calls, - success, - final_output: Some(text), - error, - } + fn finalize<'a>( + &'a mut self, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + let duration = self.session_start.map(|s| s.elapsed()).unwrap_or_default(); + + let success = !matches!(self.state, LoopState::Failed { .. }); + + // Fill in the final fields on the budget accumulator, then + // use it as the SessionResult. + self.budget.success = success; + self.budget.session_id = self.config.session_id; + self.budget.total_duration = duration; + + if success { + self.budget.final_output = match &self.state { + LoopState::Completed { summary } => Some(summary.clone()), + _ => Some(String::new()), + }; + } else { + self.budget.error = Some(match &self.state { + LoopState::Failed { error } => error.clone(), + _ => "session failed".to_string(), + }); + } + + // Notify observers/hooks using the final SessionResult. + self.notify_session_end(&self.budget, duration); + + Ok(self.budget.clone()) + }) } - /// Abort the session with an error — notifies turn-end + session-end. - /// - /// Used when the streaming call itself fails (API error, timeout, etc.). - /// - /// # Errors - /// - /// Always returns `Err(error)`, passing through the original [`LoopError`]. - fn abort_turn_and_session( - &self, - budget: &SessionBudget, - turn_duration: Duration, - session_duration: Duration, - reason: &str, - error: LoopError, - ) -> Result { - self.managers.observers().on_turn_end(&TurnEndContext { - turn: budget.turn_count, - success: false, - error: Some(reason.to_string()), - duration_ms: Self::millis_u64(turn_duration), - input_tokens: budget.input_tokens, - output_tokens: budget.output_tokens, - }); - let end_reason = if matches!(error, LoopError::Cancelled) { - EndReason::Cancelled - } else { - EndReason::Error - }; - self.notify_session_end(&SessionEndInfo { - success: false, - reason: end_reason, - total_turns: budget.turn_count, - total_tokens: budget.input_tokens.saturating_add(budget.output_tokens), - duration_secs: session_duration.as_secs(), - }); - Err(error) - } - - /// Abort the session after a tool-dispatch error. - /// - /// Handles both [`LoopError::Cancelled`] and other errors uniformly. - /// Turn-level notifications were already sent inside [`dispatch_and_record`]. - /// - /// # Errors - /// - /// Always returns `Err(error)`, passing through the original [`LoopError`]. - fn abort_session_from_error( - &self, - error: LoopError, - session_duration: Duration, - budget: &SessionBudget, - ) -> Result { - let end_reason = if matches!(error, LoopError::Cancelled) { - EndReason::Cancelled - } else { - EndReason::Error - }; - self.notify_session_end(&SessionEndInfo { - success: false, - reason: end_reason, - total_turns: budget.turn_count, - total_tokens: budget.input_tokens.saturating_add(budget.output_tokens), - duration_secs: session_duration.as_secs(), - }); - Err(error) - } - - /// Abort the session with a known reason string (cancel / max-turns). - /// - /// Does not send turn-level notifications since no turn was started. - /// - /// # Errors - /// - /// Returns [`LoopError::Cancelled`] or [`LoopError::MaxTurnsExceeded`] - /// depending on the `reason` string. - fn abort_session( - &self, - budget: &SessionBudget, - session_duration: Duration, - reason: AbortReason, - ) -> Result { - let end_reason = match &reason { - AbortReason::Cancelled => EndReason::Cancelled, - AbortReason::MaxTurnsExceeded => EndReason::MaxTurns, - }; - self.notify_session_end(&SessionEndInfo { - success: false, - reason: end_reason, - total_turns: budget.turn_count, - total_tokens: budget.input_tokens.saturating_add(budget.output_tokens), - duration_secs: session_duration.as_secs(), - }); - match reason { - AbortReason::Cancelled => Err(LoopError::Cancelled), - AbortReason::MaxTurnsExceeded => Err(LoopError::MaxTurnsExceeded { + fn state(&self) -> LoopState { + self.state.clone() + } + + fn cancel(&self) { + BareLoop::cancel(self); + } + + fn stop_reason(&self) -> Option { + if self.is_cancelled() { + return Some(LoopError::Cancelled); + } + if self.budget.total_turns >= self.config.max_turns { + return Some(LoopError::MaxTurnsExceeded { max: self.config.max_turns, - }), + }); } + None } -} -// ================================================== -// Tests -// ================================================== + fn config(&self) -> LoopConfig { + self.config.clone() + } +} #[cfg(test)] mod tests { use super::*; use crate::api::error::ApiError; + use crate::engine::loop_core::Loop; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, PartStart, Usage, @@ -1945,7 +1467,7 @@ mod tests { client.add_text_response("Hello! I'm done."); let config = make_config(); - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let result = agent.run("Hi").await.unwrap(); assert!(result.success); @@ -1969,7 +1491,7 @@ mod tests { registry.register(EchoTool); let config = make_config(); - let agent = BareLoop::new(Arc::new(client), registry, config); + let mut agent = BareLoop::new(Arc::new(client), registry, config); let result = agent.run("Echo hello").await.unwrap(); assert!(result.success); @@ -1997,7 +1519,7 @@ mod tests { let mut registry = ToolRegistry::new(); registry.register(EchoTool); - let agent = BareLoop::new(Arc::new(client), registry, config); + let mut agent = BareLoop::new(Arc::new(client), registry, config); let result = agent.run("Keep going").await; assert!(result.is_err()); match result.unwrap_err() { @@ -2014,7 +1536,7 @@ mod tests { client.add_text_response("Hello!"); let config = make_config(); - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); // Cancel before running agent.cancel(); @@ -2035,7 +1557,7 @@ mod tests { // The mock will return an error let client = MockClient::new("test-model"); let config = make_config(); - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let result = agent.run("Hi").await; assert!(result.is_err()); match result.unwrap_err() { @@ -2058,7 +1580,7 @@ mod tests { // Empty registry — tool won't be found let config = make_config(); - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let result = agent.run("Use nonexistent tool").await.unwrap(); // The tool-not-found should be returned as an error result in the conversation, @@ -2078,7 +1600,7 @@ mod tests { registry.register(FailingTool); let config = make_config(); - let agent = BareLoop::new(Arc::new(client), registry, config); + let mut agent = BareLoop::new(Arc::new(client), registry, config); let result = agent.run("Use failing tool").await.unwrap(); assert!(result.success); @@ -2178,7 +1700,7 @@ mod tests { ); let tool_calls = BareLoop::::extract_tool_calls(&msg_with_tools); assert_eq!(tool_calls.len(), 1); - assert_eq!(tool_calls[0].name, "echo"); + assert_eq!(tool_calls[0].tool, "echo"); } /// Verify that `build_tool_result_message` produces a user message @@ -2271,7 +1793,7 @@ mod tests { registry.register(EchoTool); let config = make_config(); - let agent = BareLoop::new(Arc::new(client), registry, config); + let mut agent = BareLoop::new(Arc::new(client), registry, config); let result = agent.run("Echo twice").await.unwrap(); @@ -2291,7 +1813,7 @@ mod tests { let client = MockClient::new("test-model"); let config = make_config(); let session_id = config.session_id; - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); assert_eq!(agent.config().session_id, session_id); assert!(agent.conversation().is_empty()); @@ -2304,7 +1826,7 @@ mod tests { fn test_cancel_signal_shared() { let client = MockClient::new("test-model"); let config = make_config(); - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let signal = agent.cancel_signal(); assert!(!signal.is_cancelled()); @@ -2313,22 +1835,6 @@ mod tests { assert!(agent.is_cancelled()); } - // ================================================== - // Tests: from_parts constructor - // ================================================== - - /// Verify that `from_parts` produces a loop with an empty - /// conversation. - #[test] - fn test_from_parts() { - let client = MockClient::new("test-model"); - let config = make_config(); - let managers = LoopRuntime::new(); - let agent = BareLoop::from_parts(Arc::new(client), ToolRegistry::new(), managers, config); - - assert!(agent.conversation().is_empty()); - } - // ================================================== // Tests: Session result fields // ================================================== @@ -2342,7 +1848,7 @@ mod tests { let config = make_config(); let session_id = config.session_id; - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let result = agent.run("Hi").await.unwrap(); assert_eq!(result.session_id, session_id); @@ -2364,7 +1870,7 @@ mod tests { let mut config = make_config(); config.max_turns = 1; - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let result = agent.run("Hi").await.unwrap(); assert!(result.success); @@ -2381,7 +1887,7 @@ mod tests { let mut config = make_config(); config.max_turns = 0; - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let result = agent.run("Hi").await; assert!(result.is_err()); match result.unwrap_err() { @@ -2429,7 +1935,7 @@ mod tests { client.add_text_response("Tool wasn't found, but I'll handle it."); let config = make_config(); - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let result = agent.run("Use missing tool").await.unwrap(); assert!(result.success); } @@ -2449,7 +1955,7 @@ mod tests { let client = MockClient::new("test"); client.add_tool_then_text("tool_1", "fail", json!({}), "Moving on"); - let agent = BareLoop::new(Arc::new(client), registry, make_config()); + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); let result = agent.run("Test").await.unwrap(); assert!(result.success); @@ -2463,7 +1969,7 @@ mod tests { let client = MockClient::new("test"); client.add_tool_then_text("tool_1", "nonexistent", json!({}), "OK"); - let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); let result = agent.run("Test").await.unwrap(); assert!(result.success); @@ -2481,7 +1987,7 @@ mod tests { let client = MockClient::new("test"); client.add_tool_then_text("tool_1", "fail", json!({}), "OK"); - let agent = BareLoop::new(Arc::new(client), registry, make_config()); + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); let result = agent.run("Test").await.unwrap(); assert!(result.success); @@ -2497,7 +2003,7 @@ mod tests { let client = MockClient::new("test"); client.add_tool_only_response("tc-1", "fail", json!({})); - let agent = BareLoop::new(Arc::new(client), registry, make_config()); + let mut agent = BareLoop::new(Arc::new(client), registry, make_config()); // Cancel before running agent.cancel(); diff --git a/src/engine/bare/compact.rs b/src/engine/bare/compact.rs index 8c5794d..5f779b8 100644 --- a/src/engine/bare/compact.rs +++ b/src/engine/bare/compact.rs @@ -4,10 +4,14 @@ //! When a [`ContextManager`] is configured, checks token usage after each //! tool dispatch and triggers compaction if usage exceeds the threshold. -use super::{ApiClient, BareLoop, EnsureContextResult, Instant, LoopError}; +use super::{ApiClient, BareLoop, Instant, LoopError}; #[cfg(feature = "hooks")] use super::{CompactTrigger, PostCompactContext, PreCompactContext}; +use crate::compact::EnsureContextResult; +use crate::capabilities::Compactable; +#[cfg(feature = "hooks")] +use crate::capabilities::Hookable; use crate::observer::CompactedContext; impl BareLoop { @@ -26,7 +30,7 @@ impl BareLoop { /// (i.e. the conversation exceeds the context window and the compactor /// could not reduce it sufficiently). pub(super) async fn maybe_compact_context(&mut self, turn: usize) -> Result<(), LoopError> { - let Some(ref ctx_manager) = self.context_manager else { + let Some(ctx_manager) = self.managers.context_manager() else { return Ok(()); }; @@ -35,7 +39,7 @@ impl BareLoop { // Pre-compact hook check #[cfg(feature = "hooks")] - if let Some(ref executor) = self.hook_executor { + if let Some(executor) = self.managers.hook_executor() { let tokens_before = crate::compact::CompactionOutcome::estimate_tokens(&self.conversation); let ctx = PreCompactContext { @@ -77,7 +81,7 @@ impl BareLoop { // Post-compact hook notification #[cfg(feature = "hooks")] - if let Some(ref executor) = self.hook_executor { + if let Some(executor) = self.managers.hook_executor() { let messages_compacted = messages_before.saturating_sub(messages_after); let ctx = PostCompactContext { trigger: CompactTrigger::Auto, diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index ef1cba4..f6b1a6f 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -7,14 +7,20 @@ #[cfg(feature = "hooks")] use super::HookAction; use super::{ - ApiClient, Arc, BareLoop, Correction, CorrectionResult, Duration, Instant, LoopError, - PermissionCheck, RecoveryAction, ReflectionContext, ToolCallInfo, ToolContent, ToolContext, - ToolDispatchContext, ToolDispatchResult, ToolPipeline, + ApiClient, Arc, BareLoop, Duration, Instant, LoopError, PermissionCheck, RecoveryAction, + ReflectionContext, ToolCall, ToolContent, ToolContext, ToolDispatchContext, ToolDispatchResult, + ToolPipeline, }; #[cfg(feature = "hooks")] use super::{PostToolUseContext, PreToolUseContext}; +#[cfg(feature = "tool_health")] +use crate::capabilities::HealthTrackable; +#[cfg(feature = "hooks")] +use crate::capabilities::Hookable; +use crate::capabilities::PipelineAware; use crate::detection::loop_detector::{self, Operation}; use crate::observer::{ToolPostContext, ToolPreContext}; +use crate::reflection::{Correction, CorrectionResult}; /// Result of deciding what to do after a tool error during recovery. /// @@ -31,7 +37,7 @@ enum RecoveryOutcome { impl BareLoop { /// Execute tool calls and return results. /// - /// Iterates over each [`ToolCallInfo`] extracted from the assistant + /// Iterates over each [`ToolCall`] extracted from the assistant /// message, looks up the corresponding tool in the [`ToolRegistry`], /// and invokes it. Each result is wrapped in a [`ToolDispatchResult`]. /// @@ -55,7 +61,7 @@ impl BareLoop { /// between tool invocations. pub(super) async fn dispatch_tools( &self, - tool_calls: &[ToolCallInfo], + tool_calls: &[ToolCall], turn_idx: usize, ) -> Result, LoopError> { let mut results = Vec::with_capacity(tool_calls.len()); @@ -90,7 +96,7 @@ impl BareLoop { /// during tool execution or between retry attempts. async fn dispatch_tool_with_recovery( &self, - tc: &ToolCallInfo, + tc: &ToolCall, turn_idx: usize, ) -> Result { let tool_context = self.build_tool_context(); @@ -104,7 +110,7 @@ impl BareLoop { self.managers.observers().on_tool_pre(&ToolPreContext { turn: turn_idx, - tool: tc.name.clone(), + tool: tc.tool.clone(), tool_call_id: tc.id.clone(), }); @@ -124,13 +130,13 @@ impl BareLoop { self.post_detection(&tc, &tool_result); self.managers.observers().on_tool_post(&ToolPostContext { turn: turn_idx, - tool: tc.name.clone(), + tool: tc.tool.clone(), result_hash: loop_detector::hash_result(&tool_result.output.to_string()), is_error: tool_result.is_error, duration: tool_result.duration, }); self.notify_post_tool_use_hooks(&tc, &tool_result, turn_idx); - self.record_tool_health(tc.name.as_str(), &tool_result); + self.record_tool_health(tc.tool.as_str(), &tool_result); if !tool_result.is_error { return Ok(tool_result); @@ -146,7 +152,7 @@ impl BareLoop { let correction_result = tc.apply_correction(correction, &tool_result); if let CorrectionResult::Failed(msg) = &correction_result { tracing::warn!( - tool = %tc.name, + tool = %tc.tool, error = %msg, "correction failed to produce a usable retry" ); @@ -166,9 +172,9 @@ impl BareLoop { /// manager, and returns a soft-error result if the same operation /// has exceeded the loop threshold. Returns `None` when dispatch should /// proceed normally. - fn pre_detection(&self, tc: &ToolCallInfo, turn_idx: usize) -> Option { + fn pre_detection(&self, tc: &ToolCall, turn_idx: usize) -> Option { let operation = Operation::from_input_with_signature( - &tc.name, + &tc.tool, &tc.input, self.managers.detection.signature(), ); @@ -176,13 +182,14 @@ impl BareLoop { // Check inline detection let inline_blocked = self + .managers .handle_detected_pattern(&pattern, turn_idx) .map(|_result| ToolDispatchResult { tool_call_id: tc.id.clone(), output: ToolContent::Text("loop detected: aborting tool dispatch".into()), is_error: true, duration: Duration::ZERO, - resolved_tool_name: tc.name.clone(), + resolved_tool_name: tc.tool.clone(), }); if inline_blocked.is_some() { @@ -197,13 +204,13 @@ impl BareLoop { /// Constructs an [`Operation`] with the result hash and records it with /// the detection manager. This lets the detector distinguish "same input, /// same output" (stuck) from "same input, different output" (progress). - fn post_detection(&self, tc: &ToolCallInfo, tool_result: &ToolDispatchResult) { + fn post_detection(&self, tc: &ToolCall, tool_result: &ToolDispatchResult) { let result_hash = match &tool_result.output { ToolContent::Text(t) => loop_detector::hash_result(t), ToolContent::Multipart(_) => None, }; let operation = Operation::from_input_with_result_and_signature( - &tc.name, + &tc.tool, &tc.input, result_hash, self.managers.detection.signature(), @@ -224,18 +231,18 @@ impl BareLoop { /// during tool execution. async fn dispatch_tool( &self, - tc: &ToolCallInfo, + tc: &ToolCall, tool_context: &ToolContext, start: Instant, turn_idx: usize, ) -> Result { - if let Some(ref pipeline) = self.pipeline { + if let Some(pipeline) = self.managers.pipeline() { return self .dispatch_via_pipeline(pipeline, tc, tool_context, turn_idx) .await; } - let tool_result = if let Some(tool) = self.tools.get(&tc.name) { + let tool_result = if let Some(tool) = self.tools.get(&tc.tool) { let cancel = Arc::clone(&self.cancelled); let call_result = tokio::select! { r = tool.call(tc.input.clone(), tool_context) => r, @@ -251,7 +258,7 @@ impl BareLoop { output: result.payload, is_error: result.is_error, duration, - resolved_tool_name: tc.name.clone(), + resolved_tool_name: tc.tool.clone(), } } Err(e) => { @@ -262,7 +269,7 @@ impl BareLoop { output: ToolContent::Text(error_msg), is_error: true, duration, - resolved_tool_name: tc.name.clone(), + resolved_tool_name: tc.tool.clone(), } } } @@ -277,17 +284,17 @@ impl BareLoop { /// /// Notifies observers with the error message /// that lists available tool names to help the model recover. - fn tool_not_found(&self, tc: &ToolCallInfo) -> ToolDispatchResult { + fn tool_not_found(&self, tc: &ToolCall) -> ToolDispatchResult { let available: Vec = self.tools.tool_names(); let available_refs: Vec<&str> = available.iter().map(String::as_str).collect(); - let error = LoopError::tool_not_found(&tc.name, &available_refs); + let error = LoopError::tool_not_found(&tc.tool, &available_refs); let error_msg = error.to_string(); ToolDispatchResult { tool_call_id: tc.id.clone(), output: ToolContent::Text(error_msg), is_error: true, duration: Duration::ZERO, - resolved_tool_name: tc.name.clone(), + resolved_tool_name: tc.tool.clone(), } } @@ -305,7 +312,7 @@ impl BareLoop { /// not to retry — the caller should return this as a soft error. async fn recovery_wait_or_return( &self, - tc: &ToolCallInfo, + tc: &ToolCall, tool_result: &ToolDispatchResult, attempt: u32, ) -> Result<(u32, Option), RecoveryOutcome> { @@ -337,15 +344,16 @@ impl BareLoop { /// blocked the call, or `None` if the call should proceed. /// /// *Requires `hooks` feature; returns `None` otherwise.* + #[allow(clippy::unused_self)] fn check_pre_tool_use_hooks( &self, - tc: &ToolCallInfo, + tc: &ToolCall, turn_idx: usize, ) -> Option { #[cfg(feature = "hooks")] - if let Some(ref executor) = self.hook_executor { + if let Some(executor) = self.managers.hook_executor() { let ctx = PreToolUseContext { - tool_name: tc.name.clone(), + tool_name: tc.tool.clone(), input: tc.input.clone(), session_id: self.config.session_id, turn_number: turn_idx, @@ -357,7 +365,7 @@ impl BareLoop { output: ToolContent::Text(reason), is_error: true, duration: Duration::ZERO, - resolved_tool_name: tc.name.clone(), + resolved_tool_name: tc.tool.clone(), }), HookAction::Ask { message } => { // In Headless mode (the default) the executor already @@ -369,7 +377,7 @@ impl BareLoop { output: ToolContent::Text(message), is_error: true, duration: Duration::ZERO, - resolved_tool_name: tc.name.clone(), + resolved_tool_name: tc.tool.clone(), }) } } @@ -386,17 +394,18 @@ impl BareLoop { /// Notify post-tool-use hooks with the execution result. /// /// *Requires `hooks` feature; no-op otherwise.* + #[allow(clippy::unused_self)] fn notify_post_tool_use_hooks( &self, - tc: &ToolCallInfo, + tc: &ToolCall, tool_result: &ToolDispatchResult, turn_idx: usize, ) { #[cfg(feature = "hooks")] - if let Some(ref executor) = self.hook_executor { + if let Some(executor) = self.managers.hook_executor() { let output_text = tool_result.output.to_string(); let ctx = PostToolUseContext { - tool_name: tc.name.clone(), + tool_name: tc.tool.clone(), input: tc.input.clone(), output: output_text, is_error: tool_result.is_error, @@ -419,9 +428,10 @@ impl BareLoop { /// Record tool health (success or failure) in the health registry. /// /// *Requires `tool_health` feature; no-op otherwise.* + #[allow(clippy::unused_self)] fn record_tool_health(&self, tool_name: &str, tool_result: &ToolDispatchResult) { #[cfg(feature = "tool_health")] - if let Some(ref health) = self.health_registry { + if let Some(health) = self.managers.health_registry() { if tool_result.is_error { health.record_failure(tool_name, tool_result.duration); } else { @@ -449,12 +459,12 @@ impl BareLoop { async fn dispatch_via_pipeline( &self, pipeline: &ToolPipeline, - tc: &ToolCallInfo, + tc: &ToolCall, tool_context: &ToolContext, turn_idx: usize, ) -> Result { let ctx = ToolDispatchContext { - tool_name: tc.name.clone(), + tool_name: tc.tool.clone(), input: tc.input.clone(), call_id: tc.id.clone(), turn_number: turn_idx, @@ -493,7 +503,7 @@ impl BareLoop { /// retry loop can apply it before re-dispatching. async fn recover_tool_error( &self, - tc: &ToolCallInfo, + tc: &ToolCall, result: &ToolDispatchResult, attempt: u32, ) -> (RecoveryAction, Option) { @@ -509,7 +519,7 @@ impl BareLoop { let Ok(analysis) = self .reflector - .analyze(&error_msg, &tc.name, &tc.input, &context) + .analyze(&error_msg, &tc.tool, &tc.input, &context) .await else { // Reflector failed — conservatively fail. diff --git a/src/engine/bare/emission.rs b/src/engine/bare/emission.rs index cbb7ad5..11cd611 100644 --- a/src/engine/bare/emission.rs +++ b/src/engine/bare/emission.rs @@ -8,7 +8,9 @@ //! are called directly at their call sites via //! `self.managers.observers().on_*()`. -use super::{ApiClient, BareLoop, Duration, EndReason, SessionEndInfo}; +use super::{ApiClient, BareLoop, Duration, SessionResult}; +#[cfg(feature = "hooks")] +use crate::capabilities::Hookable; #[cfg(feature = "hooks")] use crate::hooks::context::{ SessionEndContext as HookSessionEndContext, SessionEndReason, @@ -30,7 +32,7 @@ impl BareLoop { }); #[cfg(feature = "hooks")] - if let Some(ref executor) = self.hook_executor { + if let Some(executor) = self.managers.hook_executor() { let ctx = HookSessionStartContext { session_id: self.config.session_id, model: self.config.model.clone(), @@ -43,39 +45,29 @@ impl BareLoop { } /// Notify all observers and hooks that the session has ended. - pub(super) fn notify_session_end(&self, info: &SessionEndInfo) { - let reason_str = match &info.reason { - EndReason::Complete => None, - EndReason::Cancelled => Some("cancelled"), - EndReason::MaxTurns => Some("max turns exceeded"), - EndReason::Error => Some("session ended with error"), - }; + pub(super) fn notify_session_end(&self, result: &SessionResult, duration: Duration) { self.managers .observers() .on_session_end(&SessionEndContext { - success: info.success, - error: reason_str.map(std::string::ToString::to_string), - total_turns: info.total_turns, - duration_ms: u64::try_from( - std::time::Duration::from_secs(info.duration_secs).as_millis(), - ) - .unwrap_or(u64::MAX), + success: result.success, + error: result.error.clone(), + total_turns: result.total_turns, + duration_ms: Self::millis_u64(duration), }); #[cfg(feature = "hooks")] - if let Some(ref executor) = self.hook_executor { - let reason = match &info.reason { - EndReason::Complete => SessionEndReason::Complete, - EndReason::Cancelled => SessionEndReason::Cancelled, - EndReason::Error => SessionEndReason::Error, - EndReason::MaxTurns => SessionEndReason::MaxTurns, + if let Some(executor) = self.managers.hook_executor() { + let reason = if result.success { + SessionEndReason::Complete + } else { + SessionEndReason::Error }; let ctx = HookSessionEndContext { - session_id: self.config.session_id, + session_id: result.session_id, reason, - total_turns: info.total_turns, - total_tokens: info.total_tokens, - duration_secs: info.duration_secs, + total_turns: result.total_turns, + total_tokens: result.input_tokens.saturating_add(result.output_tokens), + duration_secs: duration.as_secs(), }; executor.notify_session_end(&ctx); } diff --git a/src/engine/bare/message.rs b/src/engine/bare/message.rs index 19a8bdb..1ce6053 100644 --- a/src/engine/bare/message.rs +++ b/src/engine/bare/message.rs @@ -5,7 +5,7 @@ //! file focuses on orchestration rather than message wrangling. use super::{ - ApiClient, BareLoop, Message, MessagePart, Role, ToolCallInfo, ToolContext, ToolDispatchResult, + ApiClient, BareLoop, Message, MessagePart, Role, ToolCall, ToolContext, ToolDispatchResult, ToolSchema, }; @@ -29,18 +29,18 @@ impl BareLoop { /// Extract tool call information from a message. /// /// Scans the message's [`MessagePart`]s for `ToolCall` variants and - /// maps each one to a [`ToolCallInfo`] containing the call ID, tool + /// maps each one to a [`ToolCall`] containing the call ID, tool /// name, and JSON input. Non-`ToolCall` parts are silently skipped. /// /// Returns an empty `Vec` when the message contains no tool calls /// (i.e. the model ended with plain text). - pub(super) fn extract_tool_calls(msg: &Message) -> Vec { + pub(super) fn extract_tool_calls(msg: &Message) -> Vec { msg.parts .iter() .filter_map(|part| match part { - MessagePart::ToolCall { id, name, input } => Some(ToolCallInfo { + MessagePart::ToolCall { id, name, input } => Some(ToolCall { id: id.clone(), - name: name.clone(), + tool: name.clone(), input: input.clone(), }), _ => None, diff --git a/src/engine/bare/stream.rs b/src/engine/bare/stream.rs index 613dbb7..ed16ebf 100644 --- a/src/engine/bare/stream.rs +++ b/src/engine/bare/stream.rs @@ -8,6 +8,7 @@ use super::{ ApiClient, BareLoop, LoopError, Message, StreamAccumulator, StreamEvent, StreamStopReason, Usage, }; +use crate::capabilities::StreamCapable; use crate::stream::handler::{StreamHandler, StreamHandlerError}; use futures::StreamExt; @@ -44,7 +45,7 @@ impl BareLoop { &self, ) -> Result<(Message, Option, StreamStopReason), LoopError> { // Delegate to StreamHandler if configured. - if let Some(ref handler) = self.stream_handler { + if let Some(handler) = self.managers.stream_handler() { return self.stream_turn_via_handler(handler).await; } diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index 3a4b3f2..d921e50 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -336,6 +336,63 @@ pub struct ToolCall { pub input: serde_json::Value, } +impl ToolCall { + /// Apply a [`Correction`](crate::reflection::Correction) from the reflection system in place. + /// + /// Modifies `self` according to the correction strategy: + /// + /// - [`InputFix`](crate::reflection::CorrectionType::InputFix) — replaces + /// `self.input` with the corrected input. + /// - [`ToolChange`](crate::reflection::CorrectionType::ToolChange) — replaces + /// `self.tool` with an alternative tool name. + /// - Other types — no mutation; the retry proceeds with unchanged parameters. + /// + /// Returns a [`CorrectionResult`](crate::reflection::CorrectionResult) indicating whether the correction + /// was applied, failed, or skipped. + pub fn apply_correction( + &mut self, + correction: &crate::reflection::Correction, + _prior_result: &crate::tool::ToolDispatchResult, + ) -> crate::reflection::CorrectionResult { + use crate::reflection::CorrectionType; + match correction.correction_type { + CorrectionType::InputFix => { + if let Some(ref modified) = correction.modified_input { + tracing::debug!( + tool = %self.tool, + "applying InputFix correction from reflector" + ); + self.input = modified.clone(); + crate::reflection::CorrectionResult::Applied + } else { + crate::reflection::CorrectionResult::Failed( + "InputFix correction missing modified_input".to_string(), + ) + } + } + CorrectionType::ToolChange => { + if let Some(ref alt) = correction.alternative_tool { + tracing::debug!( + old_tool = %self.tool, + new_tool = %alt, + "applying ToolChange correction from reflector" + ); + self.tool.clone_from(alt); + crate::reflection::CorrectionResult::Applied + } else { + crate::reflection::CorrectionResult::Failed( + "ToolChange correction missing alternative_tool".to_string(), + ) + } + } + CorrectionType::PrerequisiteFix | CorrectionType::ApproachChange => { + crate::reflection::CorrectionResult::Skipped + } + CorrectionType::Escalate => crate::reflection::CorrectionResult::Skipped, + } + } +} + // ================================================== // SessionResult // ================================================== @@ -386,6 +443,22 @@ pub struct SessionResult { pub error: Option, } +impl Default for SessionResult { + fn default() -> Self { + Self { + session_id: Uuid::nil(), + total_turns: 0, + input_tokens: 0, + output_tokens: 0, + total_duration: Duration::ZERO, + tool_calls: 0, + success: false, + final_output: None, + error: None, + } + } +} + impl SessionResult { /// Create a successful session result. /// @@ -549,4 +622,77 @@ pub trait Loop: Send + Sync { /// [`should_continue`](Loop::should_continue) can observe it /// and return promptly across threads. fn cancel(&self); + + /// Explain *why* [`should_continue`](Loop::should_continue) returned `false`. + /// + /// Called by [`run`](Loop::run) after the drive loop exits. Return: + /// + /// - `None` — the session ended normally (model finished). + /// - `Some(err)` — the session was forced to stop (`Cancelled`, + /// `MaxTurnsExceeded`, etc.). + /// + /// The default implementation returns `None` (normal completion). + fn stop_reason(&self) -> Option { + None + } + + /// Drive the full agent session: initialize → turn loop → finalize. + /// + /// This is the main entry point for running an agent. It calls + /// [`initialize`](Loop::initialize) with the agent's stored config, + /// then repeatedly calls [`process_turn`](Loop::process_turn) until either: + /// + /// - The turn result is marked `is_complete` (the model finished), or + /// - [`should_continue`](Loop::should_continue) returns `false`. + /// + /// When `should_continue` returns `false`, + /// [`stop_reason`](Loop::stop_reason) is consulted to distinguish + /// normal completion from an error (cancellation, max-turns, etc.). + /// + /// # Errors + /// + /// - [`LoopError::Cancelled`] — if the session was cancelled. + /// - [`LoopError::MaxTurnsExceeded`] — if the turn limit was reached. + /// - Any error returned by [`process_turn`](Loop::process_turn) or + /// [`finalize`](Loop::finalize). + fn run<'a>( + &'a mut self, + user_input: &'a str, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + self.initialize(&self.config()).await?; + + loop { + if !self.should_continue() { + break; + } + + match self.process_turn(user_input).await { + Ok(turn_result) if turn_result.is_complete => { + return self.finalize().await; + } + Ok(_) => { /* turn produced tool calls — continue */ } + Err(e) => { + self.finalize().await?; + return Err(e); + } + } + } + + // should_continue() returned false — ask the impl why. + if let Some(err) = self.stop_reason() { + self.finalize().await?; + return Err(err); + } + + self.finalize().await + }) + } + + /// Return the configuration that [`run`](Loop::run) passes to + /// [`initialize`](Loop::initialize). + /// + /// Implementors should return the [`LoopConfig`] they want to use + /// for the session. + fn config(&self) -> LoopConfig; } diff --git a/src/runtime.rs b/src/runtime.rs index 7267366..230ef19 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -50,8 +50,19 @@ //! `impl Observable + Detectable` rather than a concrete type. //! - **Testing** — swap `LoopRuntime` for a stub that implements only the //! traits under test. -//! - **Incremental adoption** — start with `LoopRuntime::new()` and add -//! capabilities via builder methods as needed. +//! ```rust,ignore +//! let runtime = LoopRuntime::builder() +//! .with_fallback(FallbackManager::for_model("llm-70b")) +//! .with_detection(DetectionManager::default()) +//! .with_observer(Arc::new(logging_observer)) +//! .build(); +//! +//! let agent = BareLoop::new_with_managers(client, tools, runtime, config); +//! ``` +//! +//! Every capability is optional. A runtime with no `.with_*()` calls still +//! works — it just has no observers, no hooks, no pipeline, etc. This means +//! you only pay for (and configure) the infrastructure you actually use. use std::sync::Arc; @@ -81,15 +92,20 @@ pub use crate::capabilities::*; /// /// # Construction /// -/// Use [`LoopRuntime::new`] for defaults or the builder-style `with_*` -/// methods to configure individual components: +/// Use [`LoopRuntime::builder()`] to compose only the capabilities you need, +/// or [`LoopRuntime::new()`] for a runtime with default managers and no +/// optional components: /// -/// ```rust,ignore -/// use loopctl::runtime::LoopRuntime; +/// ``` +/// # use loopctl::runtime::LoopRuntime; +/// # use loopctl::runtime::FallbackCapable; /// use loopctl::fallback::FallbackManager; /// -/// let runtime = LoopRuntime::new() -/// .with_fallback(FallbackManager::for_model("llm-70b")); +/// // Builder — explicit capability composition: +/// let runtime = LoopRuntime::builder() +/// .with_fallback(FallbackManager::for_model("llm-70b")) +/// .build(); +/// assert_eq!(runtime.fallback().active_model().as_deref(), Some("llm-70b")); /// ``` /// /// # Capability Traits @@ -130,6 +146,90 @@ pub struct LoopRuntime { health_registry: Option>, } +/// Builder for [`LoopRuntime`] — compose only the capabilities you need. +/// +/// Created via [`LoopRuntime::builder()`]. Each `.with_*()` method adds a +/// capability and returns `self` for chaining. Call `.build()` to finalize. +/// +/// # Example +/// +/// ``` +/// # use loopctl::runtime::LoopRuntime; +/// # use loopctl::runtime::{Detectable, FallbackCapable}; +/// let runtime = LoopRuntime::builder() +/// .with_fallback(Default::default()) +/// .with_detection(Default::default()) +/// .build(); +/// +/// // Only the capabilities you configured are present: +/// assert!(runtime.fallback().active_model().is_none()); +/// ``` +/// +/// All capabilities are optional — a bare `.build()` with no `.with_*()` +/// calls produces an empty runtime that still works but has no observers, +/// no hooks, no pipeline, etc. +#[must_use] +pub struct LoopRuntimeBuilder { + inner: LoopRuntime, +} + +impl LoopRuntimeBuilder { + /// Replace the fallback manager. + pub fn with_fallback(mut self, fallback: FallbackManager) -> Self { + self.inner.fallback = fallback; + self + } + + /// Replace the detection manager. + pub fn with_detection(mut self, detection: DetectionManager) -> Self { + self.inner.detection = detection; + self + } + + /// Register an observer. + pub fn with_observer(mut self, observer: Arc) -> Self { + self.inner.observer_host.register(observer); + self + } + + /// Set the middleware pipeline for tool dispatch. + pub fn with_pipeline(mut self, pipeline: ToolPipeline) -> Self { + self.inner.tool_pipeline = Some(pipeline); + self + } + + /// Set the context manager for automatic compaction. + pub fn with_context_manager(mut self, manager: Arc) -> Self { + self.inner.context_manager = Some(manager); + self + } + + /// Set the stream handler for resilient streaming. + pub fn with_stream_handler(mut self, handler: StreamHandler) -> Self { + self.inner.stream_handler = Some(handler); + self + } + + /// Set the hook executor for bidirectional lifecycle interception. + #[cfg(feature = "hooks")] + pub fn with_hook_executor(mut self, executor: Arc) -> Self { + self.inner.hook_executor = Some(executor); + self + } + + /// Set the tool health registry for per-tool health tracking. + #[cfg(feature = "tool_health")] + pub fn with_health_registry(mut self, registry: Arc) -> Self { + self.inner.health_registry = Some(registry); + self + } + + /// Finalize the builder and return the configured [`LoopRuntime`]. + pub fn build(self) -> LoopRuntime { + self.inner + } +} + impl LoopRuntime { /// Create a new runtime with default managers and no optional components. /// @@ -158,9 +258,33 @@ impl LoopRuntime { } // ================================================== - // Builder methods + // Builder methods — compose only the capabilities you need // ================================================== + /// Create a new runtime builder. + /// + /// Starts with no capabilities enabled. Use the `.with_*()` methods + /// to add only the infrastructure you need, then pass the result to + /// [`BareLoop::new_with_managers`](crate::engine::BareLoop::new_with_managers). + /// + /// # Example + /// + /// ``` + /// # use loopctl::runtime::LoopRuntime; + /// # use loopctl::runtime::{Detectable, FallbackCapable}; + /// let runtime = LoopRuntime::builder() + /// .with_fallback(Default::default()) + /// .with_detection(Default::default()) + /// .build(); + /// ``` + /// + /// This is equivalent to [`LoopRuntime::new`] — both start from an + /// empty runtime. Use `builder()` when you want the intent to be + /// explicit that you're composing capabilities. + pub fn builder() -> LoopRuntimeBuilder { + LoopRuntimeBuilder { inner: Self::new() } + } + /// Replace the fallback manager with a custom instance. /// /// # Example @@ -220,81 +344,124 @@ impl LoopRuntime { self.observer_host.register(observer); } + /// Register an observer and return `self` for chaining. + /// + /// Builder-style alias for [`register_observer`](Self::register_observer). + /// + /// # Example + /// + /// ```rust,ignore + /// let runtime = LoopRuntime::builder() + /// .with_observer(Arc::new(logging_observer)) + /// .with_observer(Arc::new(metrics_observer)) + /// .build(); + /// ``` + #[must_use] + pub fn with_observer(mut self, observer: Arc) -> Self { + self.observer_host.register(observer); + self + } + /// Access the observer host directly. pub fn observers(&self) -> &ObserverHost { &self.observer_host } - /// Set the middleware pipeline for tool dispatch. + /// Set the middleware pipeline for tool dispatch (builder-style). /// /// # Example /// /// ```rust,ignore - /// let mut runtime = LoopRuntime::new(); - /// runtime.set_pipeline(builder.build()?); + /// let runtime = LoopRuntime::builder() + /// .with_pipeline(builder.build()?) + /// .build(); /// ``` + #[must_use] + pub fn with_pipeline(mut self, pipeline: ToolPipeline) -> Self { + self.tool_pipeline = Some(pipeline); + self + } + + /// Set the middleware pipeline for tool dispatch (`&mut self` variant). pub fn set_pipeline(&mut self, pipeline: ToolPipeline) { self.tool_pipeline = Some(pipeline); } - /// Set the context manager for automatic compaction. + /// Set the context manager for automatic compaction (builder-style). /// /// # Example /// /// ```rust,ignore - /// use loopctl::compact::{ContextManager, TruncatingCompactor}; - /// use std::sync::Arc; - /// - /// let compactor = TruncatingCompactor::new() - /// .with_preserve_recent(4) - /// .with_min_messages(6); - /// let manager = ContextManager::new(Arc::new(compactor)) - /// .with_context_window(200_000) - /// .with_threshold(0.80); - /// - /// let mut runtime = LoopRuntime::new(); - /// runtime.set_context_manager(Arc::new(manager)); + /// let runtime = LoopRuntime::builder() + /// .with_context_manager(Arc::new(manager)) + /// .build(); /// ``` + #[must_use] + pub fn with_context_manager(mut self, manager: Arc) -> Self { + self.context_manager = Some(manager); + self + } + + /// Set the context manager for automatic compaction (`&mut self` variant). pub fn set_context_manager(&mut self, manager: Arc) { self.context_manager = Some(manager); } - /// Set the stream handler for resilient streaming with retries. + /// Set the stream handler for resilient streaming (builder-style). /// /// # Example /// /// ```rust,ignore - /// use loopctl::stream::handler::{StreamHandler, StreamTimeoutConfig}; - /// use std::time::Duration; - /// - /// let handler = StreamHandler::with_config( - /// StreamTimeoutConfig { - /// initial_event_timeout: Duration::from_secs(60), - /// ..Default::default() - /// }, - /// Default::default(), - /// ); - /// - /// let mut runtime = LoopRuntime::new(); - /// runtime.set_stream_handler(handler); + /// let runtime = LoopRuntime::builder() + /// .with_stream_handler(handler) + /// .build(); /// ``` + #[must_use] + pub fn with_stream_handler(mut self, handler: StreamHandler) -> Self { + self.stream_handler = Some(handler); + self + } + + /// Set the stream handler for resilient streaming (`&mut self` variant). pub fn set_stream_handler(&mut self, handler: StreamHandler) { self.stream_handler = Some(handler); } - /// Set the hook executor for bidirectional lifecycle interception. + /// Set the hook executor for bidirectional lifecycle interception (builder-style). + /// + /// *Requires `hooks` feature.* + #[must_use] + #[cfg(feature = "hooks")] + pub fn with_hook_executor(mut self, executor: Arc) -> Self { + self.hook_executor = Some(executor); + self + } + + /// Set the hook executor (`&mut self` variant). + /// + /// *Requires `hooks` feature.* #[cfg(feature = "hooks")] pub fn set_hook_executor(&mut self, executor: Arc) { self.hook_executor = Some(executor); } - /// Set the tool health registry for per-tool health tracking. + /// Set the tool health registry for per-tool health tracking (builder-style). /// /// When set, records success/failure counts and latency for every /// tool dispatch. Tools that exceed the failure threshold have their /// circuit breaker opened, blocking subsequent calls until recovery. /// /// *Requires `tool_health` feature.* + #[must_use] + #[cfg(feature = "tool_health")] + pub fn with_health_registry(mut self, registry: Arc) -> Self { + self.health_registry = Some(registry); + self + } + + /// Set the tool health registry (`&mut self` variant). + /// + /// *Requires `tool_health` feature.* #[cfg(feature = "tool_health")] pub fn set_health_registry(&mut self, registry: Arc) { self.health_registry = Some(registry); @@ -324,6 +491,164 @@ impl LoopRuntime { self.detection.reset(); self.observer_host.reset_all(); } + + // ================================================== + // Detection interpretation + // ================================================== + + /// Interpret a [`DetectedPattern`](crate::detection::DetectedPattern) and decide whether to abort. + /// + /// Generic framework logic: checks loop-detection thresholds, + /// maps convergence actions, and notifies observers. Returns + /// `None` to continue, or `Some(Err)` to abort the session. + /// + /// Called by the agent loop after each response is recorded with the + /// [`DetectionManager`]. + pub fn handle_detected_pattern( + &self, + pattern: &crate::detection::DetectedPattern, + turn: usize, + ) -> Option> { + use crate::detection::{ConvergenceAction, DetectedPattern}; + use crate::error::LoopError; + use crate::observer::{ConvergenceDetectedContext, LoopDetectedContext}; + + match pattern { + DetectedPattern::NoPattern => None, + + DetectedPattern::LoopDetected { + repetitions, + pattern_description, + } => { + tracing::warn!( + repetitions, + pattern = %pattern_description, + turn, + "loop detected" + ); + + self.observer_host.on_loop_detected(&LoopDetectedContext { + pattern: pattern_description.clone(), + repetitions: *repetitions, + }); + + if *repetitions >= self.detection.config().stop_threshold { + tracing::error!( + repetitions, + pattern = %pattern_description, + turn, + "stopping agent: loop threshold exceeded" + ); + Some(Err(LoopError::LoopDetected { + message: format!("{pattern_description} repeated {repetitions} times"), + })) + } else { + None + } + } + + DetectedPattern::ConvergenceDetected { + similarity, + consecutive_count, + } => { + tracing::warn!(similarity, consecutive_count, turn, "convergence detected"); + let action = self.detection.config().on_converge; + let action_str = match action { + ConvergenceAction::Stop => "stop", + ConvergenceAction::Warn => "warn", + ConvergenceAction::Compact => "compact", + ConvergenceAction::AskUser => "ask_user", + ConvergenceAction::SwitchPhase => "switch_phase", + }; + + self.observer_host + .on_convergence_detected(&ConvergenceDetectedContext { + action: action_str.to_string(), + }); + + match action { + ConvergenceAction::Stop => Some(Err(LoopError::LoopDetected { + message: "agent stopped: convergence detected".into(), + })), + ConvergenceAction::AskUser => Some(Err(LoopError::LoopDetected { + message: "agent stopped: convergence detected, user input needed".into(), + })), + ConvergenceAction::Warn + | ConvergenceAction::Compact + | ConvergenceAction::SwitchPhase => None, + } + } + } + } + + // ================================================== + // Session lifecycle notifications + // ================================================== + + /// Notify observers and hooks that a session has started. + /// + /// Fan-out to the [`ObserverHost`] and (if configured) the hook + /// executor. Generic for any loop implementation. + pub fn notify_session_start(&self, session_id: uuid::Uuid, #[allow(unused_variables)] model: &str) { + use crate::observer::SessionStartContext; + + self.observer_host + .on_session_start(&SessionStartContext { session_id }); + + #[cfg(feature = "hooks")] + if let Some(executor) = self.hook_executor() { + use crate::hooks::context::SessionStartContext as HookSessionStartContext; + + let ctx = HookSessionStartContext { + session_id, + model: model.to_string(), + working_directory: std::env::current_dir() + .map(|p| p.to_string_lossy().into_owned()) + .unwrap_or_default(), + }; + executor.notify_session_start(&ctx); + } + } + + /// Notify observers and hooks that a session has ended. + /// + /// Takes the final [`SessionResult`](crate::engine::loop_core::SessionResult) and duration. Fan-out to + /// the [`ObserverHost`] and (if configured) the hook executor. + pub fn notify_session_end( + &self, + result: &crate::engine::loop_core::SessionResult, + duration: std::time::Duration, + ) { + use crate::observer::SessionEndContext; + + self.observer_host.on_session_end(&SessionEndContext { + success: result.success, + error: result.error.clone(), + total_turns: result.total_turns, + duration_ms: u64::try_from(duration.as_millis()).unwrap_or(u64::MAX), + }); + + #[cfg(feature = "hooks")] + if let Some(executor) = self.hook_executor() { + use crate::hooks::context::{ + SessionEndContext as HookSessionEndContext, SessionEndReason, + }; + + let reason = if result.success { + SessionEndReason::Complete + } else { + SessionEndReason::Error + }; + let ctx = HookSessionEndContext { + session_id: result.session_id, + reason, + total_turns: result.total_turns, + total_tokens: result.input_tokens.saturating_add(result.output_tokens), + duration_secs: duration.as_secs(), + }; + executor.notify_session_end(&ctx); + } + } } impl Default for LoopRuntime { From 0ee994091c45cff1a5db23f6ac73f10a7df156d1 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 22 Jun 2026 08:37:07 +1200 Subject: [PATCH 05/30] chore: simple examples --- examples/echo-tool-cli.rs | 87 +++++++++++++++++++++++++++++++++++++++ examples/hello-cli.rs | 40 ++++++++++++++++++ examples/repl-cli.rs | 61 +++++++++++++++++++++++++++ 3 files changed, 188 insertions(+) create mode 100644 examples/echo-tool-cli.rs create mode 100644 examples/hello-cli.rs create mode 100644 examples/repl-cli.rs diff --git a/examples/echo-tool-cli.rs b/examples/echo-tool-cli.rs new file mode 100644 index 0000000..5ca3ff6 --- /dev/null +++ b/examples/echo-tool-cli.rs @@ -0,0 +1,87 @@ +//! BareLoop CLI with a custom tool — multi-turn tool dispatch. +//! +//! Registers an `echo` tool, configures the mock client to request it +//! on the first turn, then prints the final response. +//! +//! ```sh +//! cargo run --example echo-tool-cli --features testing +//! ``` + +use std::sync::Arc; + +use loopctl::config::LoopConfig; +use loopctl::engine::BareLoop; +use loopctl::engine::loop_core::Loop; +use loopctl::testing::{MockApiClient, MockResponse, MockToolCall}; +use loopctl::tool::{FnTool, ToolOutput, ToolRegistry}; +use serde_json::json; + +/// The async function backing the `echo` tool. +fn echo_fn( + input: serde_json::Value, + _ctx: &loopctl::tool::ToolContext, +) -> std::pin::Pin< + Box< + dyn std::future::Future> + + Send + + 'static, + >, +> { + let text = input + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("(empty)") + .to_string(); + Box::pin(async move { Ok(ToolOutput::text(format!("echo: {text}"))) }) +} + +#[tokio::main] +async fn main() { + // 1. Create a mock client that requests the echo tool on turn 1, + // then gives a final text response on turn 2. + let client = MockApiClient::new("echo-model").with_responses(vec![ + MockResponse { + text: String::new(), + tool_call: Some(MockToolCall { + id: "call_1".into(), + name: "echo".into(), + input: json!({"message": "Hello from the model!"}), + }), + stop_reason: "tool_use".into(), + }, + MockResponse { + text: "I echoed your message. Done!".into(), + tool_call: None, + stop_reason: "end_turn".into(), + }, + ]); + + // 2. Build the tool registry with a single `echo` tool. + let mut tools = ToolRegistry::new(); + tools.register( + FnTool::new( + "echo".into(), + "Echo back the provided message.".into(), + json!({ + "type": "object", + "properties": { + "message": {"type": "string", "description": "The text to echo"} + }, + "required": ["message"] + }), + echo_fn, + ) + .read_only(), + ); + + // 3. Construct and run the loop. + let mut agent = BareLoop::new(Arc::new(client), tools, LoopConfig::default()); + let result = agent + .run("Please echo something.") + .await + .expect("session should succeed"); + + println!("Turns: {}", result.total_turns); + println!("Tool calls: {}", result.tool_calls); + println!("Output: {}", result.final_output.unwrap_or_default()); +} diff --git a/examples/hello-cli.rs b/examples/hello-cli.rs new file mode 100644 index 0000000..d9932a0 --- /dev/null +++ b/examples/hello-cli.rs @@ -0,0 +1,40 @@ +//! Minimal BareLoop CLI — no tools, single-turn. +//! +//! Demonstrates the absolute simplest way to run a [`BareLoop`]: +//! create a mock client, build the loop, and call [`run`]. +//! +//! ```sh +//! cargo run --example hello-cli --features testing +//! ``` +//! +//! [`run`]: loopctl::engine::bare::BareLoop::run + +use std::sync::Arc; + +use loopctl::config::LoopConfig; +use loopctl::engine::BareLoop; +use loopctl::engine::loop_core::Loop; +use loopctl::testing::MockApiClient; +use loopctl::tool::ToolRegistry; + +#[tokio::main] +async fn main() { + // 1. Create a mock API client with a canned response. + let client = MockApiClient::new("hello-model").with_text_response("Hello, world!"); + + // 2. Build the components. + let tools = ToolRegistry::new(); + let config = LoopConfig::default(); + + // 3. Construct the loop. + let mut agent = BareLoop::new(Arc::new(client), tools, config); + + // 4. Run and print the result. + let result = agent + .run("Say hello!") + .await + .expect("session should succeed"); + + println!("Turns: {}", result.total_turns); + println!("Output: {}", result.final_output.unwrap_or_default()); +} diff --git a/examples/repl-cli.rs b/examples/repl-cli.rs new file mode 100644 index 0000000..ff90bb8 --- /dev/null +++ b/examples/repl-cli.rs @@ -0,0 +1,61 @@ +//! Interactive REPL CLI — read user input, run a turn, repeat. +//! +//! Demonstrates a minimal read-eval-print loop using [`BareLoop`]. +//! Each line of stdin becomes a new user message. Type `quit` or +//! `Ctrl-D` to exit. +//! +//! ```sh +//! cargo run --example repl-cli --features testing +//! ``` + +use std::io::{self, BufRead, Write}; +use std::sync::Arc; + +use loopctl::config::LoopConfig; +use loopctl::engine::BareLoop; +use loopctl::engine::loop_core::Loop; +use loopctl::testing::MockApiClient; +use loopctl::tool::ToolRegistry; + +#[tokio::main] +async fn main() { + let stdin = io::stdin(); + let mut stdout = io::stdout(); + + loop { + // Read. + write!(&mut stdout, "> ").unwrap(); + stdout.flush().unwrap(); + + let mut input = String::new(); + if stdin.lock().read_line(&mut input).unwrap_or(0) == 0 { + break; // EOF (Ctrl-D) + } + let input = input.trim(); + if input.is_empty() { + continue; + } + if input == "quit" || input == "exit" { + break; + } + + // Eval: build a fresh loop per message with a canned response. + let client = + MockApiClient::new("repl-model").with_text_response(&format!("You said: {input}")); + + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), LoopConfig::default()); + + match agent.run(input).await { + Ok(result) => { + // Print. + let output = result.final_output.unwrap_or_default(); + writeln!(&mut stdout, "{output}\n").unwrap(); + } + Err(e) => { + writeln!(&mut stdout, "Error: {e}\n").unwrap(); + } + } + } + + writeln!(&mut stdout, "Goodbye!").unwrap(); +} From 8c91ffbc3eb4280a576423eb2c4d250fa60cc5c7 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 22 Jun 2026 21:47:41 +1200 Subject: [PATCH 06/30] chore: remove builder --- src/builder.rs | 18 -- src/builder/error.rs | 277 ------------------- src/builder/features.rs | 597 ---------------------------------------- src/engine/loop_core.rs | 2 +- src/reflection.rs | 2 +- src/runtime.rs | 6 +- 6 files changed, 7 insertions(+), 895 deletions(-) delete mode 100644 src/builder.rs delete mode 100644 src/builder/error.rs delete mode 100644 src/builder/features.rs diff --git a/src/builder.rs b/src/builder.rs deleted file mode 100644 index ce77a4b..0000000 --- a/src/builder.rs +++ /dev/null @@ -1,18 +0,0 @@ -//! Builder module — fluent API for constructing configured agents. -//! -//! Provides a compile-time-safe builder that wires up all the components -//! an agent needs: `AgentCore`, memory, observers, managers, features, -//! and configuration. The builder uses **type-state generics** so that -//! missing required components are caught at compile time, not at runtime. -//! -//! # Provided Types -//! -//! - **[`BuildError`]** — Errors that can occur during builder validation. -//! - **[`Feature`]** — Named feature flag enum. -//! - **[`FeatureSet`]** — Compact set of enabled features. - -pub mod error; -pub mod features; - -pub use error::BuildError; -pub use features::{Feature, FeatureSet}; diff --git a/src/builder/error.rs b/src/builder/error.rs deleted file mode 100644 index 2b8a384..0000000 --- a/src/builder/error.rs +++ /dev/null @@ -1,277 +0,0 @@ -//! Builder error types — failures that can occur during agent construction. -//! -//! The [`BuildError`] enum enumerates every way an `AgentBuilder` -//! call to `build()` can fail. Because the builder uses -//! **type-state generics** to enforce the presence of an `AgentCore` -//! at compile time, some classes of errors (e.g. "no core set") can only be triggered -//! through `into_raw_parts()` or dynamic construction paths. -//! -//! # Provided Types -//! -//! - **[`BuildError`]** — Enum of all construction-time validation failures. -//! -//! # Quick Start -//! -//! ```rust -//! use loopctl::builder::BuildError; -//! -//! // Create a missing-dependency error with a helpful hint: -//! let err = BuildError::missing_dependency( -//! "ApiClient", -//! "Call .with_api_client() before building.", -//! ); -//! -//! // Create a feature-conflict error: -//! let err = BuildError::feature_conflict("fast_mode", "safe_mode"); -//! ``` - -/// Errors that can occur during agent construction. -/// -/// Returned by `AgentBuilder::build()` when validation of the accumulated -/// builder state fails. Each variant captures enough context to produce an -/// actionable error message (e.g. which features conflict, how many -/// observers exceeded the limit). -/// -/// # Validation invariants -/// -/// The builder checks the following at build time: -/// -/// - A core implementation must be present (enforced statically in most cases, -/// but checked dynamically for edge cases). -/// - Observer count must not exceed the framework's internal cap (32). -/// - User-provided validation closures, if any, must pass. -/// -/// # Example -/// -/// ```rust -/// use loopctl::builder::BuildError; -/// -/// // Each variant can be constructed directly or via convenience methods -/// let err = BuildError::MissingCore; -/// assert!(err.to_string().contains("AgentCore")); -/// -/// let err = BuildError::feature_conflict("fast_mode", "safe_mode"); -/// assert!(matches!(err, BuildError::FeatureConflict { .. })); -/// ``` -#[derive(Debug, Clone, thiserror::Error)] -pub enum BuildError { - /// No agent core was provided. - /// - /// The builder requires an `AgentCore` implementation - /// before it can produce a runnable agent. In the standard type-state flow this - /// is caught at compile time (the `NoCore` type parameter lacks the - /// `CoreSet` bound), but this variant covers dynamic - /// construction paths. - /// - /// **Fix:** Call `.with_core()` before `.build()`. - #[error("AgentBuilder requires an AgentCore implementation. Call .with_core() before .build()")] - MissingCore, - - /// Configuration is invalid. - /// - /// Wraps a human-readable description of what makes the current - /// `LoopConfig` invalid — for example, a - /// `max_turns` value of zero or a malformed model identifier. - /// - /// **Fix:** Adjust the config passed to `.with_config()`. - #[error("Invalid configuration: {0}")] - InvalidConfig(String), - - /// A required dependency is missing. - /// - /// Some features or managers need additional components to function. This - /// variant carries the dependency name and a hint about how to provide it. - /// - /// **Fix:** Add the missing dependency before calling `.build()`. - #[error("Missing dependency: {name}. {hint}")] - MissingDependency { - /// Name of the missing dependency (e.g. `"ApiClient"`, `"ToolRegistry"`). - name: String, - /// Human-readable hint describing how to supply the dependency. - hint: String, - }, - - /// Feature conflict — two mutually exclusive features enabled. - /// - /// Some features are logically incompatible (e.g. a "fast mode" and a - /// "safe mode" that trade off against each other). This variant names both - /// conflicting features so the caller can disable one. - /// - /// **Fix:** Disable one of the conflicting features via `.disable_feature()`. - #[error("Feature conflict: {feature_a} and {feature_b} are mutually exclusive")] - FeatureConflict { - /// Name of the first conflicting feature. - /// - /// Set via [`BuildError::feature_conflict()`]. Matches the - /// `Feature::name()` of the feature. - feature_a: String, - /// Name of the second conflicting feature. - /// - /// Set via [`BuildError::feature_conflict()`]. Matches the - /// `Feature::name()` of the feature. - feature_b: String, - }, - - /// Too many observers registered. - /// - /// The framework caps observers (currently 32) to prevent unbounded - /// memory growth and O(n) dispatch overhead on every lifecycle event. - /// - /// **Fix:** Remove observers or consolidate them into a single multiplexing - /// observer. - #[error("Too many observers: {count} registered, maximum is {max}")] - TooManyObservers { - /// Current number of observers that were registered. - /// - /// Will always be greater than [`max`](Self::TooManyObservers.max). - count: usize, - /// Maximum number of observers allowed. - /// - /// Defaults to the framework's internal cap (32). - max: usize, - }, - - /// A custom validation error from user-provided validation logic. - /// - /// Production builders can register arbitrary validation closures via - /// extension traits. When such a closure returns `Err`, this variant - /// wraps the returned message. - /// - /// **Fix:** Address the specific validation failure described in the message. - #[error("Validation failed: {0}")] - Validation(String), -} - -impl BuildError { - /// Create a missing-dependency error with a name and a hint. - /// - /// Convenience constructor for the [`MissingDependency`](BuildError::MissingDependency) - /// variant. Use this instead of constructing the tuple variant directly so the - /// caller's intent is self-documenting. - /// - /// # Arguments - /// - /// - `name` — Name of the missing dependency - /// (e.g. `"ApiClient"`). - /// - `hint` — Description of how to supply the dependency - /// (e.g. `"Call .with_api_client() before building."`). - /// - /// Accepts any `Into` so callers can pass `&str`, `String`, - /// or string literals interchangeably. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::BuildError; - /// - /// let err = BuildError::missing_dependency( - /// "ToolRegistry", - /// "Call .with_tool_registry() with a configured registry.", - /// ); - /// assert!(matches!(err, BuildError::MissingDependency { .. })); - /// ``` - #[must_use] - pub fn missing_dependency(name: impl Into, hint: impl Into) -> Self { - Self::MissingDependency { - name: name.into(), - hint: hint.into(), - } - } - - /// Create a feature-conflict error naming two incompatible features. - /// - /// Convenience constructor for the [`FeatureConflict`](BuildError::FeatureConflict) - /// variant. Accepts any `Into` so callers can pass `&str`, `String`, - /// or the result of `Feature::name()`. - /// - /// # Arguments - /// - /// - `a` — Name of the first feature (order does not matter). - /// - `b` — Name of the second feature. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::BuildError; - /// - /// let err = BuildError::feature_conflict("fast_mode", "safe_mode"); - /// if let BuildError::FeatureConflict { feature_a, feature_b } = &err { - /// assert_eq!(feature_a, "fast_mode"); - /// assert_eq!(feature_b, "safe_mode"); - /// } - /// ``` - #[must_use] - pub fn feature_conflict(a: impl Into, b: impl Into) -> Self { - Self::FeatureConflict { - feature_a: a.into(), - feature_b: b.into(), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_missing_core_display() { - let err = BuildError::MissingCore; - assert!(err.to_string().contains("AgentCore")); - assert!(err.to_string().contains(".with_core()")); - } - - #[test] - fn test_invalid_config() { - let err = BuildError::InvalidConfig("max_turns must be > 0".into()); - assert!(err.to_string().contains("max_turns")); - } - - #[test] - fn test_missing_dependency_constructor() { - let err = BuildError::missing_dependency("ApiClient", "Call .with_api_client()"); - assert!(matches!(err, BuildError::MissingDependency { .. })); - let msg = err.to_string(); - assert!(msg.contains("ApiClient")); - assert!(msg.contains(".with_api_client()")); - - // Also accepts String - let err = BuildError::missing_dependency( - String::from("ToolRegistry"), - format!("Call .with_tool_registry()"), - ); - assert!(matches!(err, BuildError::MissingDependency { .. })); - } - - #[test] - fn test_feature_conflict_constructor() { - let err = BuildError::feature_conflict("fast_mode", "safe_mode"); - if let BuildError::FeatureConflict { - feature_a, - feature_b, - } = &err - { - assert_eq!(feature_a, "fast_mode"); - assert_eq!(feature_b, "safe_mode"); - } else { - panic!("expected FeatureConflict variant"); - } - let msg = err.to_string(); - assert!(msg.contains("fast_mode")); - assert!(msg.contains("safe_mode")); - assert!(msg.contains("mutually exclusive")); - } - - #[test] - fn test_too_many_observers() { - let err = BuildError::TooManyObservers { count: 40, max: 32 }; - let msg = err.to_string(); - assert!(msg.contains('4') && msg.contains('0') && msg.contains("40")); - assert!(msg.contains("32")); - } - - #[test] - fn test_validation() { - let err = BuildError::Validation("custom check failed".into()); - assert!(err.to_string().contains("custom check failed")); - } -} diff --git a/src/builder/features.rs b/src/builder/features.rs deleted file mode 100644 index 236ebbd..0000000 --- a/src/builder/features.rs +++ /dev/null @@ -1,597 +0,0 @@ -//! Builder feature flags for agent construction. -//! -//! The [`Feature`] enum defines named feature flags that control builder -//! behaviour, and [`FeatureSet`] manages an enabled subset. Features are -//! enabled via `AgentBuilder::enable_feature()` and queried at build time -//! or during agent execution. -//! -//! # Provided Types -//! -//! - **[`Feature`]** — Named feature flag enum with 15 variants. -//! - **[`FeatureSet`]** — Compact set of enabled features with `O(1)` lookup. -//! -//! # Quick Start -//! -//! ```rust -//! use loopctl::builder::features::{Feature, FeatureSet}; -//! -//! let mut fs = FeatureSet::new(); -//! fs.enable(Feature::ToolShield); -//! assert!(fs.is_enabled(Feature::ToolShield)); -//! -//! fs.disable(Feature::ToolShield); -//! assert!(!fs.is_enabled(Feature::ToolShield)); -//! ``` - -use serde::{Deserialize, Serialize}; -use std::fmt; - -// ================================================== -// Feature -// ================================================== - -/// A named feature flag that controls agent builder behaviour. -/// -/// Each variant corresponds to a discrete piece of functionality that -/// can be enabled or disabled independently. Features are collected -/// into a [`FeatureSet`] by the builder and queried at build time or -/// during agent execution. -/// -/// # Example -/// -/// ```rust -/// use loopctl::builder::features::Feature; -/// -/// let feature = Feature::ToolShield; -/// assert_eq!(feature.name(), "ToolShield"); -/// assert!(!feature.description().is_empty()); -/// ``` -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[non_exhaustive] -pub enum Feature { - /// Sandbox tool execution behind a permission boundary. - /// - /// When enabled, tool invocations are wrapped in a sandbox policy - /// that restricts file-system access, network calls, and - /// environment-variable reads to an allow-list. - ToolShield, - - /// Detect and break out of repetitive agent loops. - /// - /// Monitors the conversation history for cycles (repeated tool - /// calls, identical assistant messages, or oscillating tool - /// arguments) and injects a corrective prompt or halts the loop. - LoopDetection, - - /// Enable short-term conversation memory management. - /// - /// Maintains a sliding window of past turns and can summarise - /// older context to keep the prompt within the model's context - /// window. - MemoryManagement, - - /// Validate tool inputs against their declared JSON Schema. - /// - /// Before dispatching a tool call, the framework validates the - /// arguments against the tool's `parameters` schema. Invalid - /// inputs are rejected before the tool is invoked. - ToolInputValidation, - - /// Record detailed execution traces for debugging. - /// - /// Captures per-turn timing, token counts, tool call arguments - /// and results, and other diagnostic data into a trace log that - /// can be inspected after the agent run completes. - ExecutionTracing, - - /// Automatic retry of transient API failures. - /// - /// When the LLM API returns a retryable error (rate limit, timeout, - /// server error), the framework retries with exponential back-off - /// up to a configurable maximum. - AutoRetry, - - /// Convergence detection — halt when the agent reaches a fixpoint. - /// - /// Compares consecutive assistant outputs for semantic equivalence - /// and stops the loop when the agent's response stops changing. - ConvergenceDetection, - - /// Fallback to a secondary model on persistent failure. - /// - /// If the primary model fails repeatedly, the framework switches - /// to a configured fallback model for subsequent turns. - ModelFallback, - - /// Prompt injection detection. - /// - /// Scans user messages for common injection patterns (jailbreak - /// prompts, role-reset attempts, etc.) and either rejects them or - /// wraps them in a safety preamble. - PromptInjectionDetection, - - /// Stream partial results to observers in real time. - /// - /// When enabled, the framework emits `on_stream_delta` events to - /// registered observers as the model generates tokens, rather than - /// waiting for the full response. - Streaming, - - /// Cache identical tool calls within a single run. - /// - /// If a tool is called with the same arguments multiple times - /// within one agent loop, the cached result is returned instead - /// of re-executing the tool. - ToolCallCaching, - - /// Rate-limit tool invocations per tool name. - /// - /// Prevents a runaway agent from overwhelming a tool (e.g. a web - /// search API) by capping the number of invocations per tool per - /// run. - ToolRateLimiting, - - /// Emit structured events to an external audit log. - /// - /// Each tool call, API request, and observer callback is logged - /// as a structured JSON event suitable for ingestion by an audit - /// pipeline. - AuditLogging, - - /// Enable cost tracking for API usage. - /// - /// Tracks token counts per model and estimates cost based on a - /// configured price table. The accumulated cost is available after - /// the agent run completes. - CostTracking, - - /// Parallel tool execution where safe. - /// - /// When multiple tool calls are requested in a single assistant - /// turn and none have data dependencies, the framework dispatches - /// them concurrently. - ParallelToolExecution, -} - -impl Feature { - /// Return all defined features as a slice. - /// - /// Useful for iterating over the full feature set or building a - /// [`FeatureSet::all()`]. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::Feature; - /// - /// let all = Feature::all(); - /// assert!(!all.is_empty()); - /// for feature in all { - /// println!("- {}: {}", feature.name(), feature.description()); - /// } - /// ``` - #[must_use] - pub const fn all() -> &'static [Feature] { - &[ - Feature::ToolShield, - Feature::LoopDetection, - Feature::MemoryManagement, - Feature::ToolInputValidation, - Feature::ExecutionTracing, - Feature::AutoRetry, - Feature::ConvergenceDetection, - Feature::ModelFallback, - Feature::PromptInjectionDetection, - Feature::Streaming, - Feature::ToolCallCaching, - Feature::ToolRateLimiting, - Feature::AuditLogging, - Feature::CostTracking, - Feature::ParallelToolExecution, - ] - } - - /// The short, `PascalCase` name of this feature. - /// - /// Matches the variant name exactly. Used as a key for - /// serialisation and display. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::Feature; - /// - /// assert_eq!(Feature::ToolShield.name(), "ToolShield"); - /// assert_eq!(Feature::AutoRetry.name(), "AutoRetry"); - /// ``` - #[must_use] - pub const fn name(self) -> &'static str { - match self { - Feature::ToolShield => "ToolShield", - Feature::LoopDetection => "LoopDetection", - Feature::MemoryManagement => "MemoryManagement", - Feature::ToolInputValidation => "ToolInputValidation", - Feature::ExecutionTracing => "ExecutionTracing", - Feature::AutoRetry => "AutoRetry", - Feature::ConvergenceDetection => "ConvergenceDetection", - Feature::ModelFallback => "ModelFallback", - Feature::PromptInjectionDetection => "PromptInjectionDetection", - Feature::Streaming => "Streaming", - Feature::ToolCallCaching => "ToolCallCaching", - Feature::ToolRateLimiting => "ToolRateLimiting", - Feature::AuditLogging => "AuditLogging", - Feature::CostTracking => "CostTracking", - Feature::ParallelToolExecution => "ParallelToolExecution", - } - } - - /// A one-line human-readable description of what this feature does. - /// - /// Suitable for display in `--help` output, UI tooltips, or debug logs. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::Feature; - /// - /// let desc = Feature::LoopDetection.description(); - /// assert!(!desc.is_empty()); - /// ``` - #[must_use] - pub const fn description(self) -> &'static str { - match self { - Feature::ToolShield => "Sandbox tool execution behind a permission boundary", - Feature::LoopDetection => "Detect and break out of repetitive agent loops", - Feature::MemoryManagement => "Enable short-term conversation memory management", - Feature::ToolInputValidation => { - "Validate tool inputs against their declared JSON Schema" - } - Feature::ExecutionTracing => "Record detailed execution traces for debugging", - Feature::AutoRetry => "Automatic retry of transient API failures", - Feature::ConvergenceDetection => "Halt when the agent reaches a fixpoint", - Feature::ModelFallback => "Fallback to a secondary model on persistent failure", - Feature::PromptInjectionDetection => "Scan user messages for injection patterns", - Feature::Streaming => "Stream partial results to observers in real time", - Feature::ToolCallCaching => "Cache identical tool calls within a single run", - Feature::ToolRateLimiting => "Rate-limit tool invocations per tool name", - Feature::AuditLogging => "Emit structured events to an external audit log", - Feature::CostTracking => "Track token counts and estimate API cost", - Feature::ParallelToolExecution => "Execute independent tools in parallel", - } - } - - /// Check whether this feature conflicts with another. - /// - /// Some features are mutually exclusive (e.g. `Streaming` and - /// `ToolCallCaching` may interfere because streaming bypasses the - /// caching layer). - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::Feature; - /// - /// // Streaming conflicts with ToolCallCaching - /// assert!(Feature::Streaming.conflicts_with(Feature::ToolCallCaching)); - /// // ToolShield does not conflict with LoopDetection - /// assert!(!Feature::ToolShield.conflicts_with(Feature::LoopDetection)); - /// ``` - #[must_use] - pub const fn conflicts_with(self, other: Feature) -> bool { - matches!( - (self, other), - (Feature::Streaming, Feature::ToolCallCaching) - | (Feature::ToolCallCaching, Feature::Streaming) - ) - } -} - -impl fmt::Display for Feature { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str(self.name()) - } -} - -// ================================================== -// FeatureSet -// ================================================== - -/// A compact set of enabled [`Feature`] flags. -/// -/// `FeatureSet` tracks which features are enabled for a given agent -/// build. It provides average `O(1)` enable/disable/query operations backed -/// by an internal `HashSet`. -/// -/// # Construction -/// -/// ```rust -/// use loopctl::builder::features::{Feature, FeatureSet}; -/// -/// // Empty set -/// let mut fs = FeatureSet::new(); -/// -/// // Enable individual features -/// fs.enable(Feature::ToolShield); -/// fs.enable(Feature::LoopDetection); -/// -/// // Or start with all features enabled -/// let all = FeatureSet::all(); -/// ``` -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct FeatureSet { - /// Enabled features stored as a set for O(1) lookup. - enabled: std::collections::HashSet, -} - -impl FeatureSet { - /// Create an empty feature set (no features enabled). - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::{Feature, FeatureSet}; - /// - /// let fs = FeatureSet::new(); - /// assert!(!fs.is_enabled(Feature::ToolShield)); - /// assert_eq!(fs.len(), 0); - /// ``` - #[must_use] - pub fn new() -> Self { - Self { - enabled: std::collections::HashSet::new(), - } - } - - /// Create a feature set with all features enabled. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::{Feature, FeatureSet}; - /// - /// let fs = FeatureSet::all(); - /// for feature in Feature::all() { - /// assert!(fs.is_enabled(*feature)); - /// } - /// ``` - #[must_use] - pub fn all() -> Self { - let mut set = Self::new(); - for feature in Feature::all() { - set.enable(*feature); - } - set - } - - /// Enable a feature. - /// - /// Idempotent — enabling an already-enabled feature is a no-op. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::{Feature, FeatureSet}; - /// - /// let mut fs = FeatureSet::new(); - /// fs.enable(Feature::ToolShield); - /// assert!(fs.is_enabled(Feature::ToolShield)); - /// ``` - pub fn enable(&mut self, feature: Feature) { - self.enabled.insert(feature); - } - - /// Disable a feature. - /// - /// Idempotent — disabling a feature that is not enabled is a no-op. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::{Feature, FeatureSet}; - /// - /// let mut fs = FeatureSet::all(); - /// fs.disable(Feature::ToolShield); - /// assert!(!fs.is_enabled(Feature::ToolShield)); - /// ``` - pub fn disable(&mut self, feature: Feature) { - self.enabled.remove(&feature); - } - - /// Check whether a feature is enabled. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::{Feature, FeatureSet}; - /// - /// let mut fs = FeatureSet::new(); - /// fs.enable(Feature::AutoRetry); - /// assert!(fs.is_enabled(Feature::AutoRetry)); - /// assert!(!fs.is_enabled(Feature::ToolShield)); - /// ``` - #[must_use] - pub fn is_enabled(&self, feature: Feature) -> bool { - self.enabled.contains(&feature) - } - - /// Return the number of enabled features. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::{Feature, FeatureSet}; - /// - /// let mut fs = FeatureSet::new(); - /// assert_eq!(fs.len(), 0); - /// fs.enable(Feature::ToolShield); - /// fs.enable(Feature::AutoRetry); - /// assert_eq!(fs.len(), 2); - /// ``` - #[must_use] - pub fn len(&self) -> usize { - self.enabled.len() - } - - /// Check whether no features are enabled. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::FeatureSet; - /// - /// let fs = FeatureSet::new(); - /// assert!(fs.is_empty()); - /// ``` - #[must_use] - pub fn is_empty(&self) -> bool { - self.enabled.is_empty() - } - - /// Return an iterator over the enabled features. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::{Feature, FeatureSet}; - /// - /// let mut fs = FeatureSet::new(); - /// fs.enable(Feature::ToolShield); - /// fs.enable(Feature::AutoRetry); - /// - /// let names: Vec<&str> = fs.iter().map(|f| f.name()).collect(); - /// assert!(names.contains(&"ToolShield")); - /// assert!(names.contains(&"AutoRetry")); - /// ``` - pub fn iter(&self) -> impl Iterator { - self.enabled.iter() - } - - /// Check whether any enabled feature conflicts with the given feature. - /// - /// Returns the first conflicting enabled feature, if any. - /// - /// # Example - /// - /// ```rust - /// use loopctl::builder::features::{Feature, FeatureSet}; - /// - /// let mut fs = FeatureSet::new(); - /// fs.enable(Feature::Streaming); - /// let conflict = fs.find_conflict(Feature::ToolCallCaching); - /// assert_eq!(conflict, Some(Feature::Streaming)); - /// ``` - #[must_use] - pub fn find_conflict(&self, feature: Feature) -> Option { - self.enabled - .iter() - .find(|enabled| enabled.conflicts_with(feature)) - .copied() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_feature_name_not_empty() { - for feature in Feature::all() { - assert!(!feature.name().is_empty()); - } - } - - #[test] - fn test_feature_description_not_empty() { - for feature in Feature::all() { - assert!(!feature.description().is_empty()); - } - } - - #[test] - fn test_feature_display() { - assert_eq!(format!("{}", Feature::ToolShield), "ToolShield"); - assert_eq!(format!("{}", Feature::AutoRetry), "AutoRetry"); - } - - #[test] - fn test_feature_conflicts() { - assert!(Feature::Streaming.conflicts_with(Feature::ToolCallCaching)); - assert!(Feature::ToolCallCaching.conflicts_with(Feature::Streaming)); - assert!(!Feature::ToolShield.conflicts_with(Feature::LoopDetection)); - assert!(!Feature::AutoRetry.conflicts_with(Feature::ModelFallback)); - } - - #[test] - fn test_feature_set_new() { - let fs = FeatureSet::new(); - assert!(fs.is_empty()); - assert_eq!(fs.len(), 0); - for feature in Feature::all() { - assert!(!fs.is_enabled(*feature)); - } - } - - #[test] - fn test_feature_set_enable_disable() { - let mut fs = FeatureSet::new(); - fs.enable(Feature::ToolShield); - assert!(fs.is_enabled(Feature::ToolShield)); - assert_eq!(fs.len(), 1); - - fs.disable(Feature::ToolShield); - assert!(!fs.is_enabled(Feature::ToolShield)); - assert!(fs.is_empty()); - } - - #[test] - fn test_feature_set_all() { - let fs = FeatureSet::all(); - assert_eq!(fs.len(), Feature::all().len()); - for feature in Feature::all() { - assert!(fs.is_enabled(*feature)); - } - } - - #[test] - fn test_feature_names() { - for feature in Feature::all() { - assert!(!feature.name().is_empty()); - } - } - - #[test] - fn test_feature_set_iter() { - let mut fs = FeatureSet::new(); - fs.enable(Feature::ToolShield); - fs.enable(Feature::AutoRetry); - - let features: Vec = fs.iter().copied().collect(); - assert_eq!(features.len(), 2); - assert!(features.contains(&Feature::ToolShield)); - assert!(features.contains(&Feature::AutoRetry)); - } - - #[test] - fn test_feature_set_find_conflict() { - let mut fs = FeatureSet::new(); - fs.enable(Feature::Streaming); - assert_eq!( - fs.find_conflict(Feature::ToolCallCaching), - Some(Feature::Streaming) - ); - assert_eq!(fs.find_conflict(Feature::ToolShield), None); - } - - #[test] - fn test_feature_set_serialization() { - let mut fs = FeatureSet::new(); - fs.enable(Feature::ToolShield); - fs.enable(Feature::AutoRetry); - - let json = serde_json::to_string(&fs).unwrap(); - let back: FeatureSet = serde_json::from_str(&json).unwrap(); - assert!(back.is_enabled(Feature::ToolShield)); - assert!(back.is_enabled(Feature::AutoRetry)); - assert!(!back.is_enabled(Feature::LoopDetection)); - } -} diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index d921e50..41e360f 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -123,7 +123,7 @@ pub enum LoopState { /// The agent is reflecting on a failure and preparing a correction. /// /// Entered when a tool call fails and the reflection system is - /// enabled (via `Feature::Reflection`). + /// enabled (via configuration). /// The agent analyzes the error and produces a /// [`Correction`](crate::reflection::Correction) before retrying. Reflecting { diff --git a/src/reflection.rs b/src/reflection.rs index c9aa2b7..03e1aed 100644 --- a/src/reflection.rs +++ b/src/reflection.rs @@ -168,7 +168,7 @@ pub enum CorrectionType { /// A correction produced by the reflection system. /// -/// When a tool call fails and reflection is enabled (via `Feature::Reflection`), +/// When a tool call fails and reflection is enabled (via configuration), /// the agent analyzes the error and produces a `Correction` that describes how to fix /// the problem. The framework applies the correction and retries. /// diff --git a/src/runtime.rs b/src/runtime.rs index 230ef19..8d62ac4 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -589,7 +589,11 @@ impl LoopRuntime { /// /// Fan-out to the [`ObserverHost`] and (if configured) the hook /// executor. Generic for any loop implementation. - pub fn notify_session_start(&self, session_id: uuid::Uuid, #[allow(unused_variables)] model: &str) { + pub fn notify_session_start( + &self, + session_id: uuid::Uuid, + #[allow(unused_variables)] model: &str, + ) { use crate::observer::SessionStartContext; self.observer_host From 14189058b0a34f1ba11233a0c4cf6cd352caeecc Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 29 Jun 2026 20:39:14 +1200 Subject: [PATCH 07/30] feat: add llm providers --- Cargo.toml | 13 + src/api/error.rs | 7 +- src/engine/bare.rs | 145 +++- src/engine/bare/stream.rs | 8 + src/lib.rs | 4 +- src/middleware/unknown_tool.rs | 4 - src/provider.rs | 469 +++++++++++++ src/provider/anthropic.rs | 1103 +++++++++++++++++++++++++++++ src/provider/gemini.rs | 909 ++++++++++++++++++++++++ src/provider/openai.rs | 1184 ++++++++++++++++++++++++++++++++ src/reflection/backoff.rs | 2 - src/stream.rs | 305 +++++--- src/testing.rs | 72 +- 13 files changed, 4088 insertions(+), 137 deletions(-) create mode 100644 src/provider.rs create mode 100644 src/provider/anthropic.rs create mode 100644 src/provider/gemini.rs create mode 100644 src/provider/openai.rs diff --git a/Cargo.toml b/Cargo.toml index 4c5b8d8..8d934ec 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,9 @@ tokio = { version = "1.52.3", features = ["sync", "macros", "time"] } uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" +reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"], optional = true } +async-stream = { version = "0.3", optional = true } + [dev-dependencies] tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time"] } proptest = "1" @@ -37,6 +40,16 @@ testing = [] tool_health = [] tool_shield = ["tool_health"] +# Providers +providers = ["dep:reqwest", "dep:async-stream"] +openai = ["providers"] +anthropic = ["providers"] +ollama = ["providers", "openai"] +deepseek = ["providers", "openai"] +grok = ["providers", "openai"] +gemini = ["providers"] +zai = ["providers", "anthropic"] + [lints.clippy] pedantic = { level = "warn", priority = -1 } unwrap_used = "deny" diff --git a/src/api/error.rs b/src/api/error.rs index f90f80d..3a682d8 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -998,11 +998,8 @@ impl ApiError { /// ``` #[must_use] pub fn io_not_found(err: std::io::Error) -> Self { - debug_assert!( - matches!(err.kind(), std::io::ErrorKind::NotFound), - "io_not_found called with non-NotFound ErrorKind: {:?}", - err.kind() - ); + // Soft-validate: if called with a non-NotFound error, still + // construct the Io variant rather than panicking. Self::Io(err) } diff --git a/src/engine/bare.rs b/src/engine/bare.rs index f96deeb..07a42a1 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -253,6 +253,14 @@ pub struct BareLoop { /// Session start time, set by [`initialize`](crate::engine::loop_core::Loop::initialize). session_start: Option, + + /// Optional callback invoked for each text delta during streaming. + /// + /// Set via [`set_text_streamer`](BareLoop::set_text_streamer). + /// When set, called from `stream_turn` on every `IndexedDelta` with + /// a `Text` payload, enabling real-time token display. + #[allow(clippy::type_complexity)] + text_streamer: Option>, } // ================================================== @@ -296,6 +304,7 @@ impl BareLoop { state: LoopState::Idle, budget: SessionResult::default(), session_start: None, + text_streamer: None, } } @@ -345,6 +354,7 @@ impl BareLoop { state: LoopState::Idle, budget: SessionResult::default(), session_start: None, + text_streamer: None, } } @@ -619,6 +629,32 @@ impl BareLoop { self.managers.register_observer(observer); } + /// Set a real-time text streaming callback. + /// + /// The callback is invoked for each text delta token as it arrives + /// from the API during [`run`](crate::engine::loop_core::Loop::run). + /// This enables real-time display of the model's output without + /// waiting for the full turn to complete. + /// + /// The callback receives a `&str` containing the delta text fragment. + /// It must be `Send + Sync` as it may be called from an async context. + /// + /// # Example + /// + /// ```rust,ignore + /// use std::sync::{Arc, Mutex}; + /// + /// let buffer = Arc::new(Mutex::new(String::new())); + /// let buf = Arc::clone(&buffer); + /// agent.set_text_streamer(Arc::new(move |delta| { + /// print!("{delta}"); + /// buf.lock().unwrap().push_str(delta); + /// })); + /// ``` + pub fn set_text_streamer(&mut self, f: Arc) { + self.text_streamer = Some(f); + } + // ================================================== // Run helpers // ================================================== @@ -1074,6 +1110,11 @@ mod tests { self.responses.lock().unwrap().push(events); } + /// Add a raw sequence of stream events as a single response turn. + fn add_events(&self, events: Vec) { + self.responses.lock().unwrap().push(events); + } + /// Add a tool_call response followed by an end_turn response. /// /// The first response contains a single `tool_call` content part @@ -1802,6 +1843,106 @@ mod tests { assert_eq!(result.tool_calls, 2); } + // ================================================== + // Tests: text streamer callback + // ================================================== + + /// Verify that `set_text_streamer` fires the callback for each text + /// delta during streaming. + #[tokio::test] + async fn test_text_streamer_fires_on_text_delta() { + let client = MockClient::new("test-model"); + client.add_text_response("Hello world"); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + + let received = Arc::new(std::sync::Mutex::new(Vec::new())); + let buf = Arc::clone(&received); + agent.set_text_streamer(Arc::new(move |delta: &str| { + buf.lock().unwrap().push(delta.to_string()); + })); + + let result = agent.run("Hi").await.unwrap(); + assert!(result.success); + + let received = received.lock().unwrap(); + assert!(!received.is_empty(), "streamer should have fired"); + assert!( + received.join("").contains("Hello world"), + "got: {:?}", + received + ); + } + + /// Verify that a run works fine without a text streamer set. + #[tokio::test] + async fn test_text_streamer_none_works() { + let client = MockClient::new("test-model"); + client.add_text_response("No streamer"); + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + + let result = agent.run("Hi").await.unwrap(); + assert!(result.success); + } + + /// Verify the streamer only fires for text deltas, not for tool-call + /// deltas or metadata events. + #[tokio::test] + async fn test_text_streamer_ignores_non_text_deltas() { + let client = MockClient::new("test-model"); + + // Build a response with tool-call events (no text). + let events = vec![ + StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: "msg-1".into(), + role: "assistant".into(), + model: "test-model".into(), + }, + }), + StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::ToolCall { + id: "call_1".into(), + name: "echo".into(), + input: Value::Null, + }), + }), + StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::InputJson { + partial_json: "{}".into(), + }, + }), + StreamEvent::PartStop, + StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some("tool_call".into()), + }, + usage: None, + }), + StreamEvent::MessageStop, + ]; + client.add_events(events); + + // Second turn: plain text response. + client.add_text_response("Done"); + + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + + let received = Arc::new(std::sync::Mutex::new(String::new())); + let buf = Arc::clone(&received); + agent.set_text_streamer(Arc::new(move |delta: &str| { + buf.lock().unwrap().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().unwrap(); + assert_eq!(&*received, "Done", "only text deltas should fire streamer"); + } + // ================================================== // Tests: Accessors // ================================================== @@ -1813,7 +1954,7 @@ mod tests { let client = MockClient::new("test-model"); let config = make_config(); let session_id = config.session_id; - let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); assert_eq!(agent.config().session_id, session_id); assert!(agent.conversation().is_empty()); @@ -1826,7 +1967,7 @@ mod tests { fn test_cancel_signal_shared() { let client = MockClient::new("test-model"); let config = make_config(); - let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); + let agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), config); let signal = agent.cancel_signal(); assert!(!signal.is_cancelled()); diff --git a/src/engine/bare/stream.rs b/src/engine/bare/stream.rs index ed16ebf..8dafc31 100644 --- a/src/engine/bare/stream.rs +++ b/src/engine/bare/stream.rs @@ -67,6 +67,14 @@ impl BareLoop { match event_result { Some(Ok(event)) => { + // Fire text streaming callback for real-time display. + if let Some(ref streamer) = self.text_streamer { + if let StreamEvent::IndexedDelta(indexed_delta) = &event { + if let crate::stream::DeltaPart::Text { text } = &indexed_delta.delta { + streamer(text); + } + } + } if let StreamEvent::MessageDelta(delta) = &event { if let Some(ref reason_str) = delta.delta.stop_reason { stop_reason = diff --git a/src/lib.rs b/src/lib.rs index 0da3c39..cc724d3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -38,13 +38,11 @@ //! //! ## Support //! -//! - **[`builder`]** — Fluent builder API for constructing configured agents. //! - **[`memory::builtin`]** — Reference [`InMemoryStore`](memory::builtin::InMemoryStore) implementation. //! - **[`hooks`]** — Bidirectional lifecycle control (allow/block/ask before tool use, compaction). *Requires `hooks` feature.* //! - **[`testing`]** — Test utilities and fixtures. *Requires `testing` feature.* pub mod api; -pub mod builder; pub mod cancel; pub mod capabilities; pub mod compact; @@ -59,6 +57,8 @@ pub mod memory; pub mod message; pub mod middleware; pub mod observer; +#[cfg(feature = "providers")] +pub mod provider; pub mod reflection; pub mod runtime; pub mod stream; diff --git a/src/middleware/unknown_tool.rs b/src/middleware/unknown_tool.rs index 2b216bd..7602e0f 100644 --- a/src/middleware/unknown_tool.rs +++ b/src/middleware/unknown_tool.rs @@ -294,8 +294,6 @@ mod tests { assert_eq!(UnknownToolMiddleware::lcs_length(&a, &b), 2); } - // ---- find_best_match_inner ---- - #[test] fn find_best_match_exact_hit() { let available = ["read_file", "write_file", "list_dir"]; @@ -373,8 +371,6 @@ mod tests { ); } - // ---- is_tool_not_found ---- - use crate::message::ToolContent; use crate::tool::ToolDispatchResult; use std::time::Duration; diff --git a/src/provider.rs b/src/provider.rs new file mode 100644 index 0000000..b6563bb --- /dev/null +++ b/src/provider.rs @@ -0,0 +1,469 @@ +//! LLM provider clients. +//! +//! This module provides ready-to-use [`ApiClient`](crate::api::ApiClient) +//! implementations for common LLM providers. Each provider is behind a +//! feature flag so you only compile what you need. +//! +//! # Feature Flags +//! +//! | Provider | Feature | API Format | +//! |--------------|------------|-------------------------| +//! | OpenAI | `openai` | OpenAI Chat Completions | +//! | Anthropic | `anthropic`| Anthropic Messages | +//! | Gemini | `gemini` | Google Gemini | +//! | Ollama | `ollama` | OpenAI-compatible | +//! | `DeepSeek` | `deepseek` | OpenAI-compatible | +//! | `Grok` (xAI) | `grok` | OpenAI-compatible | +//! | `Z.ai` | `zai` | Anthropic Messages | +//! | Self-hosted | any | OpenAI-compatible | +//! +//! Any provider that exposes an OpenAI-compatible Chat Completions API +//! can use [`OpenAiClient`] with a custom base URL. The convenience +//! constructors below pre-configure the correct endpoints. +//! +//! # Quick Start +//! +//! ```rust,ignore +//! use loopctl::provider; +//! use loopctl::engine::BareLoop; +//! use loopctl::engine::loop_core::Loop; +//! +//! // OpenAI: +//! let client = provider::OpenAiClient::from_env()?; +//! +//! // DeepSeek: +//! let client = provider::deepseek()?; +//! +//! // Anthropic: +//! let client = provider::AnthropicClient::from_env()?; +//! +//! // Gemini: +//! let client = provider::GeminiClient::from_env()?; +//! +//! // Ollama (local): +//! let client = provider::ollama("llama3")?; +//! +//! // Self-hosted (vLLM, LM Studio, etc.): +//! let client = provider::self_hosted("http://localhost:8080/v1", "my-model")?; +//! +//! let agent = BareLoop::new( +//! std::sync::Arc::new(client), +//! tool_registry, +//! config, +//! ); +//! let result = agent.run("Hello!").await?; +//! ``` + +use crate::api::error::ApiError; + +#[cfg(feature = "openai")] +pub mod openai; + +#[cfg(feature = "anthropic")] +pub mod anthropic; + +#[cfg(feature = "gemini")] +pub mod gemini; + +#[cfg(feature = "openai")] +pub use openai::OpenAiClient; + +#[cfg(feature = "anthropic")] +pub use anthropic::AnthropicClient; + +#[cfg(feature = "gemini")] +pub use gemini::GeminiClient; + +// ======================================================= +// Default endpoints / models for convenience constructors +// ======================================================= + +#[cfg(feature = "ollama")] +const OLLAMA_BASE_URL: &str = "http://localhost:11434/v1"; + +#[cfg(feature = "deepseek")] +const DEEPSEEK_BASE_URL: &str = "https://api.deepseek.com/v1"; + +#[cfg(feature = "deepseek")] +const DEEPSEEK_DEFAULT_MODEL: &str = "deepseek-chat"; + +#[cfg(feature = "grok")] +const GROK_BASE_URL: &str = "https://api.x.ai/v1"; + +#[cfg(feature = "grok")] +const GROK_DEFAULT_MODEL: &str = "grok-beta"; + +#[cfg(feature = "zai")] +const ZAI_BASE_URL: &str = "https://api.z.ai/api/anthropic"; + +#[cfg(feature = "zai")] +const ZAI_DEFAULT_MODEL: &str = "glm-4.7"; + +// ================================================== +// Internal helpers +// ================================================== + +/// Read an environment variable, falling back to a second name, then a +/// default value. +/// +/// Reduces boilerplate in the convenience constructors below where a +/// provider supports multiple env-var aliases (e.g. `XAI_API_KEY` / +/// `GROK_API_KEY`). +fn env_or_fallback(primary: &str, fallback: &str) -> Option { + std::env::var(primary) + .or_else(|_| std::env::var(fallback)) + .ok() +} + +/// Read an environment variable or return a default. +fn env_or_default(name: &str, default: &str) -> String { + std::env::var(name).unwrap_or_else(|_| default.into()) +} + +/// Look up a required API key, returning [`ApiError`] if neither env +/// var is set. +/// +/// # Errors +/// +/// Returns [`ApiError::auth_invalid_key`] if neither environment +/// variable is set. +fn require_api_key(primary: &str, fallback: Option<&str>) -> Result { + if let Some(fb) = fallback { + if let Some(val) = env_or_fallback(primary, fb) { + return Ok(val); + } + } else if let Ok(val) = std::env::var(primary) { + return Ok(val); + } + Err(ApiError::auth_invalid_key(format!("{primary} not set"))) +} + +// ============================================= +// Constructors for OpenAI-compatible providers +// ============================================= + +/// Ollama client — an [`OpenAiClient`] pointed at an Ollama server. +/// +/// Works with both local Ollama (`http://localhost:11434/v1`, no API key +/// needed) and Ollama Cloud (`https://api.ollama.com/v1`, requires +/// `OLLAMA_API_KEY`). +/// +/// Reads: +/// - `OLLAMA_API_KEY` — optional for local, required for cloud. +/// - `OLLAMA_BASE_URL` — optional, defaults to `http://localhost:11434/v1`. +/// Set to `https://api.ollama.com/v1` for Ollama Cloud. +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::provider; +/// +/// // Local: +/// let client = provider::ollama("llama3")?; +/// +/// // Cloud (set OLLAMA_API_KEY and OLLAMA_BASE_URL): +/// let client = provider::ollama("llama3")?; +/// ``` +/// +/// # Errors +/// +/// Returns [`ApiError`] if the HTTP client cannot be built. +#[cfg(feature = "ollama")] +pub fn ollama(model: &str) -> Result { + let base = env_or_default("OLLAMA_BASE_URL", OLLAMA_BASE_URL); + let api_key = env_or_default("OLLAMA_API_KEY", "ollama"); + + OpenAiClient::builder() + .api_key(api_key) + .base_url(base) + .model(model) + .build() +} + +/// `DeepSeek` client — an [`OpenAiClient`] pointed at the `DeepSeek` API. +/// +/// Reads `DEEPSEEK_API_KEY` (required) and optionally `DEEPSEEK_MODEL` +/// (defaults to `deepseek-chat`). +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::provider; +/// +/// let client = provider::deepseek()?; +/// ``` +/// +/// # Errors +/// +/// Returns [`ApiError`] if no API key is found. +#[cfg(feature = "deepseek")] +pub fn deepseek() -> Result { + let api_key = require_api_key("DEEPSEEK_API_KEY", None)?; + let model = env_or_default("DEEPSEEK_MODEL", DEEPSEEK_DEFAULT_MODEL); + + OpenAiClient::builder() + .api_key(api_key) + .base_url(DEEPSEEK_BASE_URL) + .model(model) + .build() +} + +/// `Grok` (xAI) client — an [`OpenAiClient`] pointed at the xAI API. +/// +/// Reads `XAI_API_KEY` (or `GROK_API_KEY`) (required) and optionally +/// `GROK_MODEL` (defaults to `grok-beta`). +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::provider; +/// +/// let client = provider::grok()?; +/// ``` +/// +/// # Errors +/// +/// Returns [`ApiError`] if no API key is found. +#[cfg(feature = "grok")] +pub fn grok() -> Result { + let api_key = require_api_key("XAI_API_KEY", Some("GROK_API_KEY"))?; + let model = env_or_default("GROK_MODEL", GROK_DEFAULT_MODEL); + + OpenAiClient::builder() + .api_key(api_key) + .base_url(GROK_BASE_URL) + .model(model) + .build() +} + +/// `Z.ai` (`ZhipuAI` / `BigModel`) client — an [`AnthropicClient`] pointed +/// at the `Z.ai` Anthropic-compatible API. +/// +/// `Z.ai` exposes an Anthropic Messages-compatible API at +/// `https://api.z.ai/api/anthropic`. +/// +/// Reads `ZAI_API_KEY` (or `ZHIPUAI_API_KEY`) (required) and optionally +/// `ZAI_MODEL` (defaults to `glm-4.6`). +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::provider; +/// +/// let client = provider::zai()?; +/// ``` +/// +/// # Errors +/// +/// Returns [`ApiError`] if no API key is found. +#[cfg(feature = "zai")] +pub fn zai() -> Result { + let api_key = require_api_key("ZAI_API_KEY", Some("ZHIPUAI_API_KEY"))?; + let model = env_or_default("ZAI_MODEL", ZAI_DEFAULT_MODEL); + + AnthropicClient::builder() + .api_key(api_key) + .base_url(ZAI_BASE_URL) + .model(model) + .build() +} + +/// Self-hosted client — an [`OpenAiClient`] pointed at any custom endpoint. +/// +/// Use this for `vLLM`, `LM Studio`, `text-generation-inference`, or any +/// other server that exposes an OpenAI-compatible API. +/// +/// For servers that require an API key, set it via the `OPENAI_API_KEY` +/// environment variable or use [`OpenAiClient::builder`] directly. +/// +/// # Example +/// +/// ```rust,ignore +/// use loopctl::provider; +/// +/// let client = provider::self_hosted("http://localhost:8080/v1", "my-model")?; +/// ``` +/// +/// # Errors +/// +/// Returns [`ApiError`] if the HTTP client cannot be built. +#[cfg(feature = "openai")] +pub fn self_hosted(base_url: &str, model: &str) -> Result { + let api_key = env_or_default("OPENAI_API_KEY", "self-hosted"); + + OpenAiClient::builder() + .api_key(api_key) + .base_url(base_url) + .model(model) + .build() +} + +// ================================================== +// Tests +// ================================================== + +#[cfg(test)] +mod tests { + use super::*; + + /// Helper to safely set an env var in tests (Rust 2024 requires unsafe). + macro_rules! env_set { + ($($arg:tt)*) => {{ unsafe { std::env::set_var($($arg)*) } }}; + } + + /// Helper to safely remove an env var in tests. + macro_rules! env_remove { + ($($arg:tt)*) => {{ unsafe { std::env::remove_var($($arg)*) } }}; + } + + #[test] + fn env_or_fallback_primary_set() { + env_set!("LOOPCTL_TEST_PRIMARY", "primary-val"); + env_remove!("LOOPCTL_TEST_FALLBACK"); + assert_eq!( + env_or_fallback("LOOPCTL_TEST_PRIMARY", "LOOPCTL_TEST_FALLBACK"), + Some("primary-val".into()) + ); + env_remove!("LOOPCTL_TEST_PRIMARY"); + } + + #[test] + fn env_or_fallback_fallback_used_when_primary_missing() { + env_remove!("LOOPCTL_TEST_PRIMARY2"); + env_set!("LOOPCTL_TEST_FALLBACK2", "fallback-val"); + assert_eq!( + env_or_fallback("LOOPCTL_TEST_PRIMARY2", "LOOPCTL_TEST_FALLBACK2"), + Some("fallback-val".into()) + ); + env_remove!("LOOPCTL_TEST_FALLBACK2"); + } + + #[test] + fn env_or_fallback_none_when_both_missing() { + env_remove!("LOOPCTL_TEST_NEITHER_A"); + env_remove!("LOOPCTL_TEST_NEITHER_B"); + assert_eq!( + env_or_fallback("LOOPCTL_TEST_NEITHER_A", "LOOPCTL_TEST_NEITHER_B"), + None + ); + } + + #[test] + fn env_or_default_uses_env_when_set() { + env_set!("LOOPCTL_TEST_DEFAULT", "from-env"); + assert_eq!( + env_or_default("LOOPCTL_TEST_DEFAULT", "fallback"), + "from-env" + ); + env_remove!("LOOPCTL_TEST_DEFAULT"); + } + + #[test] + fn env_or_default_uses_default_when_unset() { + env_remove!("LOOPCTL_TEST_DEFAULT2"); + assert_eq!( + env_or_default("LOOPCTL_TEST_DEFAULT2", "fallback"), + "fallback" + ); + } + + #[test] + fn require_api_key_primary_set() { + env_set!("LOOPCTL_TEST_KEY_PRIMARY", "secret"); + env_remove!("LOOPCTL_TEST_KEY_FALLBACK"); + let key = require_api_key( + "LOOPCTL_TEST_KEY_PRIMARY", + Some("LOOPCTL_TEST_KEY_FALLBACK"), + ) + .unwrap(); + assert_eq!(key, "secret"); + env_remove!("LOOPCTL_TEST_KEY_PRIMARY"); + } + + #[test] + fn require_api_key_fallback_used() { + env_remove!("LOOPCTL_TEST_KEY_PRIMARY2"); + env_set!("LOOPCTL_TEST_KEY_FALLBACK2", "fallback-secret"); + let key = require_api_key( + "LOOPCTL_TEST_KEY_PRIMARY2", + Some("LOOPCTL_TEST_KEY_FALLBACK2"), + ) + .unwrap(); + assert_eq!(key, "fallback-secret"); + env_remove!("LOOPCTL_TEST_KEY_FALLBACK2"); + } + + #[test] + fn require_api_key_no_fallback_set() { + env_set!("LOOPCTL_TEST_KEY_ONLY", "only-val"); + let key = require_api_key("LOOPCTL_TEST_KEY_ONLY", None).unwrap(); + assert_eq!(key, "only-val"); + env_remove!("LOOPCTL_TEST_KEY_ONLY"); + } + + #[test] + fn require_api_key_errors_when_missing() { + env_remove!("LOOPCTL_TEST_MISSING_KEY"); + let err = require_api_key("LOOPCTL_TEST_MISSING_KEY", None).unwrap_err(); + assert!(err.to_string().contains("LOOPCTL_TEST_MISSING_KEY")); + } + + #[test] + fn require_api_key_errors_when_both_missing() { + env_remove!("LOOPCTL_TEST_MISSING_A"); + env_remove!("LOOPCTL_TEST_MISSING_B"); + let err = + require_api_key("LOOPCTL_TEST_MISSING_A", Some("LOOPCTL_TEST_MISSING_B")).unwrap_err(); + assert!(err.to_string().contains("LOOPCTL_TEST_MISSING_A")); + } + + #[cfg(feature = "ollama")] + #[test] + fn ollama_client_builds_with_defaults() { + env_remove!("OLLAMA_BASE_URL"); + let client = ollama("llama3").unwrap(); + use crate::api::ApiClient; + assert_eq!(client.model(), "llama3"); + } + + #[cfg(feature = "ollama")] + #[test] + fn ollama_client_respects_base_url_env() { + env_set!("OLLAMA_BASE_URL", "http://my-host:1234/v1"); + let client = ollama("test-model").unwrap(); + use crate::api::ApiClient; + assert_eq!(client.model(), "test-model"); + env_remove!("OLLAMA_BASE_URL"); + } + + #[cfg(feature = "ollama")] + #[test] + fn ollama_client_uses_api_key_when_set() { + env_remove!("OLLAMA_BASE_URL"); + env_set!("OLLAMA_API_KEY", "my-cloud-key"); + // Should build successfully with the cloud key — no network call. + let client = ollama("llama3").unwrap(); + use crate::api::ApiClient; + assert_eq!(client.model(), "llama3"); + env_remove!("OLLAMA_API_KEY"); + } + + #[cfg(feature = "ollama")] + #[test] + fn ollama_client_defaults_to_local_without_key() { + env_remove!("OLLAMA_BASE_URL"); + env_remove!("OLLAMA_API_KEY"); + // Should still build — local Ollama doesn't need a real key. + let client = ollama("llama3").unwrap(); + use crate::api::ApiClient; + assert_eq!(client.model(), "llama3"); + } + + #[cfg(feature = "openai")] + #[test] + fn self_hosted_client_builds() { + let client = self_hosted("http://localhost:8080/v1", "my-model").unwrap(); + use crate::api::ApiClient; + assert_eq!(client.model(), "my-model"); + } +} diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs new file mode 100644 index 0000000..9f3ea6f --- /dev/null +++ b/src/provider/anthropic.rs @@ -0,0 +1,1103 @@ +//! Anthropic Messages API client. +//! +//! Implements [`ApiClient`] by translating between the framework's +//! [`StreamEvent`] protocol and the Anthropic Messages SSE format. +//! +//! # Construction +//! +//! ```rust,ignore +//! use loopctl::provider::AnthropicClient; +//! +//! // From environment (ANTHROPIC_API_KEY): +//! let client = AnthropicClient::from_env()?; +//! +//! // Explicit: +//! let client = AnthropicClient::builder() +//! .api_key("sk-ant-...") +//! .model("claude-sonnet-4-20250514") +//! .build()?; +//! ``` + +use std::future::Future; +use std::pin::Pin; + +use futures::stream::{Stream, StreamExt}; +use reqwest::Response; +use serde_json::Value; + +use crate::api::ApiClient; +use crate::api::error::ApiError; +use crate::message::{Message, MessagePart, Role}; +use crate::stream::{ + DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, + PartStart, StreamEvent, StreamStopReason, Usage, +}; +use crate::tool::ToolSchema; + +// ================================================== +// Constants +// ================================================== +const DEFAULT_BASE_URL: &str = "https://api.anthropic.com"; +const DEFAULT_MODEL: &str = "claude-sonnet-4-20250514"; +const ANTHROPIC_VERSION: &str = "2023-06-01"; +const SSE_EVENT_PREFIX: &str = "event: "; +const SSE_DATA_PREFIX: &str = "data: "; +const TEXT_PART_INDEX: usize = 0; +const DEFAULT_MAX_TOKENS: u32 = 8192; + +// ================================================== +// Client +// ================================================== + +/// An Anthropic Claude chat client with streaming support. +/// +/// Implements [`ApiClient`] by translating between the framework's +/// [`StreamEvent`] protocol and the Anthropic Messages SSE format. +/// +/// Also works with Anthropic-compatible endpoints such as `Z.ai` +/// — use a custom `base_url` via [`AnthropicClientBuilder::base_url`]. +pub struct AnthropicClient { + http: reqwest::Client, + api_key: String, + base_url: String, + model: String, + max_tokens: u32, +} + +impl AnthropicClient { + /// Create a builder for configuring an [`AnthropicClient`]. + #[must_use] + pub fn builder() -> AnthropicClientBuilder { + AnthropicClientBuilder::default() + } + + /// Create from environment variables. + /// + /// Reads: + /// - `ANTHROPIC_API_KEY` — required. + /// - `ANTHROPIC_BASE_URL` — optional, defaults to `https://api.anthropic.com`. + /// - `ANTHROPIC_MODEL` — optional, defaults to `claude-sonnet-4-20250514`. + /// + /// # Errors + /// + /// Returns [`ApiError`] if no API key is found. + pub fn from_env() -> Result { + let api_key = std::env::var("ANTHROPIC_API_KEY") + .map_err(|_| ApiError::auth_invalid_key("ANTHROPIC_API_KEY not set"))?; + let base_url = + std::env::var("ANTHROPIC_BASE_URL").unwrap_or_else(|_| DEFAULT_BASE_URL.into()); + let model = std::env::var("ANTHROPIC_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.into()); + + Self::builder() + .api_key(api_key) + .base_url(base_url) + .model(model) + .build() + } + + /// Send a POST request to the Messages endpoint. + /// + /// Shared by both [`ApiClient::stream_messages`] and + /// [`ApiClient::create_message`]. + /// + /// # Errors + /// + /// Returns [`ApiError`] if the request fails or the server + /// responds with a non-success status code. + async fn post_messages( + http: &reqwest::Client, + url: &str, + api_key: &str, + body: &Value, + ) -> Result { + let resp = http + .post(url) + .header("x-api-key", api_key) + .header("anthropic-version", ANTHROPIC_VERSION) + .json(body) + .send() + .await + .map_err(|e| ApiError::http(e.to_string()))?; + let status = resp.status(); + if status.is_success() { + Ok(resp) + } else { + let text = resp.text().await.unwrap_or_default(); + Err(ApiError::http_with_status(status.as_u16(), text)) + } + } + + /// Build the Messages API URL for this client. + fn messages_url(&self) -> String { + format!("{}/v1/messages", self.base_url) + } +} + +impl ApiClient for AnthropicClient { + fn model(&self) -> &str { + &self.model + } + + fn stream_messages( + &self, + messages: Vec, + system: Option, + tools: Option>, + ) -> Pin> + Send + 'static>> { + let body = build_request_body( + &self.model, + &messages, + system.as_deref(), + tools.as_deref(), + true, + self.max_tokens, + ); + let url = self.messages_url(); + let api_key = self.api_key.clone(); + let http = self.http.clone(); + + Box::pin(async_stream::try_stream! { + let resp = Self::post_messages(&http, &url, &api_key, &body).await?; + let mut sse = SseReader::from_response(resp); + let mut emitter = StreamEmitter::default(); + + while let Some((event_type, data)) = sse.next_event().await? { + emitter.process_event(&event_type, data); + for ev in emitter.drain() { + yield ev; + } + } + + for ev in emitter.finish() { + yield ev; + } + }) + } + + fn create_message( + &self, + messages: Vec, + system: Option, + tools: Option>, + ) -> Pin> + Send + '_>> { + let body = build_request_body( + &self.model, + &messages, + system.as_deref(), + tools.as_deref(), + false, + self.max_tokens, + ); + let url = self.messages_url(); + Box::pin(async move { + let resp = Self::post_messages(&self.http, &url, &self.api_key, &body).await?; + resp.json::() + .await + .map_err(|e| ApiError::http(e.to_string())) + }) + } +} + +// ================================================== +// Builder +// ================================================== + +/// Builder for [`AnthropicClient`]. +pub struct AnthropicClientBuilder { + api_key: Option, + base_url: String, + model: String, + max_tokens: u32, +} + +impl Default for AnthropicClientBuilder { + fn default() -> Self { + Self { + api_key: None, + base_url: DEFAULT_BASE_URL.into(), + model: DEFAULT_MODEL.into(), + max_tokens: DEFAULT_MAX_TOKENS, + } + } +} + +impl AnthropicClientBuilder { + /// Set the API key. + #[must_use] + pub fn api_key(mut self, key: impl Into) -> Self { + self.api_key = Some(key.into()); + self + } + + /// Set the base URL (e.g. `https://api.z.ai/api/anthropic`). + #[must_use] + pub fn base_url(mut self, url: impl Into) -> Self { + self.base_url = url.into(); + self + } + + /// Set the model name (e.g. `claude-sonnet-4-20250514`). + #[must_use] + pub fn model(mut self, model: impl Into) -> Self { + self.model = model.into(); + self + } + + /// Set the maximum output tokens per response. + /// + /// Anthropic requires this field. Defaults to 8192. + #[must_use] + pub fn max_tokens(mut self, tokens: u32) -> Self { + self.max_tokens = tokens; + self + } + + /// Build the client. + /// + /// # Errors + /// + /// Returns [`ApiError`] if no API key was set. + pub fn build(self) -> Result { + let api_key = self + .api_key + .ok_or_else(|| ApiError::auth_invalid_key("API key not provided"))?; + let http = reqwest::Client::builder() + .build() + .map_err(|e| ApiError::http(e.to_string()))?; + + Ok(AnthropicClient { + http, + api_key, + base_url: self.base_url, + model: self.model, + max_tokens: self.max_tokens, + }) + } +} + +// ================================================== +// Request body construction +// ================================================== + +/// Build the JSON request body for the Anthropic Messages API. +/// +/// Each [`Message`] is serialized via [`convert_message`], then assembled +/// with the model, `max_tokens`, system prompt, and optional tools. +fn build_request_body( + model: &str, + messages: &[Message], + system: Option<&str>, + tools: Option<&[ToolSchema]>, + stream: bool, + max_tokens: u32, +) -> Value { + let msgs: Vec = messages.iter().map(convert_message).collect(); + + let mut body = serde_json::json!({ + "model": model, + "max_tokens": max_tokens, + "messages": msgs, + "system": system.unwrap_or(""), + "stream": stream, + "tools": tools.map(convert_tools), + }); + + // Remove `tools` if None so we don't send a null field. + if tools.is_none() { + if let Some(obj) = body.as_object_mut() { + obj.remove("tools"); + } + } + + body +} + +/// Convert a single framework [`Message`] into the Anthropic JSON shape. +/// +/// - Messages with only a single text part use a plain string for `content` +/// (Anthropic's recommended optimization). +/// - Messages with tool calls or tool results use the full `content` array. +fn convert_message(m: &Message) -> Value { + let role = match m.role { + Role::User => "user", + Role::Assistant => "assistant", + }; + + // Bucket parts by Anthropic category. + let mut text_parts: Vec<&str> = Vec::new(); + let mut tool_calls: Vec = Vec::new(); + let mut tool_results: Vec = Vec::new(); + + for p in &m.parts { + match p { + MessagePart::Text { text } => text_parts.push(text.as_str()), + MessagePart::ToolCall { id, name, input } => { + tool_calls.push(serde_json::json!({ + "type": "tool_use", + "id": id, + "name": name, + "input": input, + })); + } + MessagePart::ToolResult { + call_id, output, .. + } => { + tool_results.push(serde_json::json!({ + "type": "tool_result", + "tool_use_id": call_id, + "content": output.to_string(), + })); + } + MessagePart::Image { .. } => {} + } + } + + let has_tool_content = !(tool_calls.is_empty() && tool_results.is_empty()); + + if !has_tool_content && text_parts.len() == 1 { + // Single text — Anthropic allows plain string content. + let text = text_parts.first().copied().unwrap_or_default(); + serde_json::json!({ "role": role, "content": text }) + } else if !has_tool_content { + // Multiple text parts — array of text blocks. + let blocks: Vec = text_parts + .iter() + .map(|t| serde_json::json!({"type": "text", "text": t})) + .collect(); + serde_json::json!({ "role": role, "content": blocks }) + } else { + // Mixed content — combine text + tool blocks in a single array. + let mut blocks: Vec = Vec::new(); + + if !text_parts.is_empty() { + let text = text_parts.join(""); + blocks.push(serde_json::json!({"type": "text", "text": text})); + } + blocks.extend(tool_calls); + blocks.extend(tool_results); + + serde_json::json!({ "role": role, "content": blocks }) + } +} + +/// Convert tool schemas into the Anthropic `tools` array shape. +fn convert_tools(tools: &[ToolSchema]) -> Vec { + tools + .iter() + .map(|t| { + serde_json::json!({ + "name": t.tool, + "description": &t.description, + "input_schema": t.input_schema.clone(), + }) + }) + .collect() +} + +// ================================================== +// SSE line reader +// ================================================== + +/// Minimal SSE line reader over an HTTP byte stream. +/// +/// Buffers raw bytes from the response, splits on newlines, and yields +/// `(event_type, data)` pairs. Anthropic SSE uses separate `event:` and +/// `data:` lines for each event. +struct SseReader { + bytes: Pin> + Send>>, + buf: String, +} + +impl SseReader { + /// Wrap a streaming HTTP response. + fn from_response(resp: Response) -> Self { + let bytes = resp.bytes_stream().map(|res| { + res.map(|b| String::from_utf8_lossy(&b).into_owned()) + .map_err(|e| ApiError::http(e.to_string())) + }); + Self { + bytes: Box::pin(bytes), + buf: String::new(), + } + } + + /// Extract the next SSE event as `(event_type, data_json)`. + /// + /// Returns `Ok(None)` at end-of-stream. + /// + /// # Errors + /// + /// Returns [`ApiError`] if the underlying HTTP stream fails. + async fn next_event(&mut self) -> Result)>, ApiError> { + let mut event_type = String::new(); + let mut data = String::new(); + let mut have_event = false; + + loop { + while let Some(line) = self.take_line() { + if line.is_empty() { + // Blank line = event boundary. Emit if we have one. + if have_event { + let parsed = if data.is_empty() { + None + } else { + serde_json::from_str(&data).ok() + }; + return Ok(Some((event_type, parsed))); + } + continue; + } + + if let Some(ev) = line.strip_prefix(SSE_EVENT_PREFIX) { + event_type = ev.into(); + have_event = true; + } else if let Some(d) = line.strip_prefix(SSE_DATA_PREFIX) { + data = d.into(); + have_event = true; + } + } + + match self.bytes.next().await { + Some(Ok(chunk)) => self.buf.push_str(&chunk), + Some(Err(e)) => return Err(e), + None => { + // End of stream — emit any pending event. + if have_event { + let parsed = if data.is_empty() { + None + } else { + serde_json::from_str(&data).ok() + }; + return Ok(Some((event_type, parsed))); + } + return Ok(None); + } + } + } + } + + /// Pop the first `\n`-terminated line from the buffer, if present. + fn take_line(&mut self) -> Option { + let pos = self.buf.find('\n')?; + let line = self.buf[..pos].trim().to_string(); + let rest_start = pos.saturating_add(1); + self.buf = self.buf.get(rest_start..).unwrap_or_default().to_string(); + Some(line) + } +} + +// ================================================== +// Stream event emitter +// ================================================== + +/// Stateful translator that converts Anthropic SSE events into +/// [`StreamEvent`]s. +/// +/// Encapsulates all protocol-level bookkeeping: +/// - Emitting [`MessageStart`] once. +/// - Emitting [`PartStart`] / [`IndexedDelta`] for text and tool-call content. +/// - Emitting [`PartStop`] when parts finish. +/// - Emitting the final [`MessageDelta`] with stop reason and usage. +#[derive(Default)] +struct StreamEmitter { + started: bool, + text_part_open: bool, + tool_parts_open: usize, + /// Index of the currently open tool block (from the API's `content_block_start`). + current_tool_index: Option, + finished: bool, + pending: Vec, +} + +impl StreamEmitter { + /// Process a single SSE event, appending events to the internal queue. + fn process_event(&mut self, event_type: &str, data: Option) { + match event_type { + "message_start" => self.on_message_start(data.as_ref()), + "content_block_start" => self.on_block_start(data), + "content_block_delta" => self.on_block_delta(data), + "content_block_stop" => self.on_block_stop(data), + "message_delta" => self.on_message_delta(data), + "message_stop" => self.on_message_stop(), + _ => {} + } + } + + fn on_message_start(&mut self, data: Option<&Value>) { + if self.started { + return; + } + self.started = true; + + let (id, model) = match data { + Some(v) => ( + v.pointer("/message/id") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(), + v.pointer("/message/model") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(), + ), + None => (String::new(), String::new()), + }; + + self.push(StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id, + role: "assistant".into(), + model, + }, + })); + } + + fn on_block_start(&mut self, data: Option) { + let Some(v) = data else { return }; + let block_type = v.pointer("/content_block/type").and_then(Value::as_str); + let index = v + .pointer("/index") + .and_then(Value::as_u64) + .and_then(|n| usize::try_from(n).ok()) + .unwrap_or(0); + + match block_type { + Some("tool_use") => { + let id = v + .pointer("/content_block/id") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + let name = v + .pointer("/content_block/name") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + + self.push(StreamEvent::PartStart(PartStart { + index, + part: Some(MessagePart::ToolCall { + id, + name, + input: Value::Null, + }), + })); + self.current_tool_index = Some(index); + self.tool_parts_open = self.tool_parts_open.saturating_add(1); + } + Some("text") => { + self.text_part_open = true; + self.push(StreamEvent::PartStart(PartStart { + index: TEXT_PART_INDEX, + part: Some(MessagePart::text("")), + })); + } + _ => {} + } + } + + fn on_block_delta(&mut self, data: Option) { + let Some(v) = data else { return }; + let delta_type = v.pointer("/delta/type").and_then(Value::as_str); + + match delta_type { + Some("text_delta") => { + let text = v + .pointer("/delta/text") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + if !text.is_empty() { + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: TEXT_PART_INDEX, + delta: DeltaPart::Text { text }, + })); + } + } + Some("input_json_delta") => { + let json = v + .pointer("/delta/partial_json") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + if !json.is_empty() { + // Use the index from the corresponding content_block_start. + let tool_index = self.current_tool_index.unwrap_or(TEXT_PART_INDEX); + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: tool_index, + delta: DeltaPart::InputJson { partial_json: json }, + })); + } + } + _ => {} + } + } + + fn on_block_stop(&mut self, _data: Option) { + if self.text_part_open { + self.text_part_open = false; + self.push(StreamEvent::PartStop); + } else if self.tool_parts_open > 0 { + self.tool_parts_open = self.tool_parts_open.saturating_sub(1); + self.current_tool_index = None; + self.push(StreamEvent::PartStop); + } + } + + fn on_message_delta(&mut self, data: Option) { + if self.finished { + return; + } + + let Some(v) = data else { return }; + let stop_reason = v + .pointer("/delta/stop_reason") + .and_then(Value::as_str) + .map(|s| StreamStopReason::from_api_str(s).unwrap_or(StreamStopReason::EndTurn)); + let in_tok = v + .pointer("/usage/input_tokens") + .and_then(Value::as_u64) + .and_then(|n| u32::try_from(n).ok()) + .unwrap_or(0); + let out_tok = v + .pointer("/usage/output_tokens") + .and_then(Value::as_u64) + .and_then(|n| u32::try_from(n).ok()) + .unwrap_or(0); + + let usage = if in_tok > 0 || out_tok > 0 { + Some(Usage::new(in_tok, out_tok)) + } else { + None + }; + + self.push(StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: stop_reason.map(|r| r.to_api_str().into()), + }, + usage, + })); + } + + fn on_message_stop(&mut self) { + self.finished = true; + // Close any remaining open parts. + if self.text_part_open { + self.push(StreamEvent::PartStop); + } + for _ in 0..self.tool_parts_open { + self.push(StreamEvent::PartStop); + } + self.tool_parts_open = 0; + self.text_part_open = false; + } + + /// Emit the terminal [`MessageStop`] if the stream was started. + fn finish(&mut self) -> Vec { + let mut out = self.drain(); + if self.started && !self.finished { + out.push(StreamEvent::MessageStop); + } + out + } + + /// Drain all pending events. + fn drain(&mut self) -> Vec { + std::mem::take(&mut self.pending) + } + + fn push(&mut self, ev: StreamEvent) { + self.pending.push(ev); + } +} + +// ================================================== +// Tests +// ================================================== + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::{Message, MessagePart, Role, ToolContent}; + + #[test] + fn request_body_user_text_single_string() { + let msgs = vec![Message::user("hello")]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + + let messages = body["messages"].as_array().unwrap(); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0]["role"], "user"); + assert_eq!(messages[0]["content"], "hello"); + } + + #[test] + fn request_body_includes_system() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body( + "claude-3", + &msgs, + Some("be brief"), + None, + false, + DEFAULT_MAX_TOKENS, + ); + assert_eq!(body["system"], "be brief"); + } + + #[test] + fn request_body_system_empty_when_none() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + assert_eq!(body["system"], ""); + } + + #[test] + fn request_body_model() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body( + "claude-sonnet-4", + &msgs, + None, + None, + false, + DEFAULT_MAX_TOKENS, + ); + assert_eq!(body["model"], "claude-sonnet-4"); + } + + #[test] + fn request_body_max_tokens() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + assert_eq!(body["max_tokens"], DEFAULT_MAX_TOKENS); + } + + #[test] + fn request_body_user_role() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + assert_eq!(body["messages"][0]["role"], "user"); + } + + #[test] + fn request_body_assistant_role() { + let msgs = vec![Message::new( + Role::Assistant, + vec![MessagePart::text("hello")], + )]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + assert_eq!(body["messages"][0]["role"], "assistant"); + assert_eq!(body["messages"][0]["content"], "hello"); + } + + #[test] + fn request_body_assistant_tool_calls() { + let msgs = vec![Message::new( + Role::Assistant, + vec![MessagePart::ToolCall { + id: "call_1".into(), + name: "echo".into(), + input: serde_json::json!({"msg": "hi"}), + }], + )]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + + let msg = &body["messages"][0]; + assert_eq!(msg["role"], "assistant"); + let content = msg["content"].as_array().unwrap(); + assert_eq!(content[0]["type"], "tool_use"); + assert_eq!(content[0]["id"], "call_1"); + assert_eq!(content[0]["name"], "echo"); + assert_eq!(content[0]["input"]["msg"], "hi"); + } + + #[test] + fn request_body_tool_result() { + let msgs = vec![Message::new( + Role::User, + vec![MessagePart::ToolResult { + call_id: "call_1".into(), + output: ToolContent::from_string("result text"), + is_error: None, + }], + )]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + + let msg = &body["messages"][0]; + assert_eq!(msg["role"], "user"); + let content = msg["content"].as_array().unwrap(); + assert_eq!(content[0]["type"], "tool_result"); + assert_eq!(content[0]["tool_use_id"], "call_1"); + assert_eq!(content[0]["content"], "result text"); + } + + #[test] + fn request_body_includes_tools() { + let msgs = vec![Message::user("hi")]; + let tools = vec![ToolSchema { + tool: "search".into(), + description: "Search the web".into(), + input_schema: serde_json::json!({"type": "object"}), + }]; + let body = build_request_body( + "claude-3", + &msgs, + None, + Some(&tools), + false, + DEFAULT_MAX_TOKENS, + ); + + let tools_arr = body["tools"].as_array().unwrap(); + assert_eq!(tools_arr.len(), 1); + assert_eq!(tools_arr[0]["name"], "search"); + assert_eq!(tools_arr[0]["description"], "Search the web"); + assert_eq!(tools_arr[0]["input_schema"]["type"], "object"); + } + + #[test] + fn request_body_tools_absent_when_none() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + assert!(body.get("tools").is_none()); + } + + #[test] + fn request_body_assistant_with_text_and_tool_call() { + let msgs = vec![Message::new( + Role::Assistant, + vec![ + MessagePart::text("Let me search."), + MessagePart::ToolCall { + id: "call_1".into(), + name: "search".into(), + input: serde_json::json!({"q": "rust"}), + }, + ], + )]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + + let msg = &body["messages"][0]; + assert_eq!(msg["role"], "assistant"); + let content = msg["content"].as_array().unwrap(); + assert_eq!(content.len(), 2); + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[1]["type"], "tool_use"); + } + + #[test] + fn request_body_multiple_messages() { + let msgs = vec![ + Message::user("hello"), + Message::new(Role::Assistant, vec![MessagePart::text("hi")]), + Message::user("bye"), + ]; + let body = build_request_body("claude-3", &msgs, None, None, false, DEFAULT_MAX_TOKENS); + + let messages = body["messages"].as_array().unwrap(); + assert_eq!(messages.len(), 3); + assert_eq!(messages[0]["role"], "user"); + assert_eq!(messages[1]["role"], "assistant"); + assert_eq!(messages[2]["role"], "user"); + } + + #[test] + fn convert_tools_shape() { + let tools = vec![ToolSchema { + tool: "calc".into(), + description: "Calculate".into(), + input_schema: serde_json::json!({"type": "object"}), + }]; + let out = convert_tools(&tools); + assert_eq!(out.len(), 1); + assert_eq!(out[0]["name"], "calc"); + assert_eq!(out[0]["description"], "Calculate"); + } + + #[test] + fn builder_requires_api_key() { + let result = AnthropicClient::builder().build(); + assert!(result.is_err()); + } + + #[test] + fn builder_succeeds_with_key() { + let client = AnthropicClient::builder() + .api_key("sk-test") + .build() + .unwrap(); + assert_eq!(client.model(), DEFAULT_MODEL); + } + + #[test] + fn builder_custom_base_url_and_model() { + let client = AnthropicClient::builder() + .api_key("sk-test") + .base_url("https://custom.example.com") + .model("claude-3-haiku") + .build() + .unwrap(); + assert_eq!(client.model(), "claude-3-haiku"); + } + + #[test] + fn emitter_message_start() { + let mut em = StreamEmitter::default(); + let data = serde_json::json!({ + "message": {"id": "msg_1", "model": "claude-3"} + }); + em.on_message_start(Some(&data)); + let events = em.drain(); + assert!( + events + .iter() + .any(|e| matches!(e, StreamEvent::MessageStart(_))) + ); + } + + #[test] + fn emitter_text_delta() { + let mut em = StreamEmitter::default(); + + // Start a text block. + em.on_block_start(Some(serde_json::json!({ + "index": 0, + "content_block": {"type": "text"} + }))); + em.drain(); + + // Text delta. + em.on_block_delta(Some(serde_json::json!({ + "delta": {"type": "text_delta", "text": "hi"} + }))); + let events = em.drain(); + assert_eq!(events.len(), 1); + assert!(matches!(events[0], StreamEvent::IndexedDelta(_))); + } + + #[test] + fn emitter_tool_use_block() { + let mut em = StreamEmitter::default(); + + em.on_block_start(Some(serde_json::json!({ + "index": 1, + "content_block": {"type": "tool_use", "id": "t1", "name": "echo"} + }))); + let events = em.drain(); + assert!( + events + .iter() + .any(|e| matches!(e, StreamEvent::PartStart(_))) + ); + assert_eq!(em.tool_parts_open, 1); + + // Input JSON delta. + em.on_block_delta(Some(serde_json::json!({ + "delta": {"type": "input_json_delta", "partial_json": "{\"a\":"} + }))); + let events2 = em.drain(); + assert_eq!(events2.len(), 1); + } + + #[test] + fn emitter_block_stop_closes_text() { + let mut em = StreamEmitter::default(); + em.text_part_open = true; + + em.on_block_stop(None); + let events = em.drain(); + assert!(matches!(events[0], StreamEvent::PartStop)); + assert!(!em.text_part_open); + } + + #[test] + fn emitter_message_delta_with_usage() { + let mut em = StreamEmitter::default(); + + em.on_message_delta(Some(serde_json::json!({ + "delta": {"stop_reason": "end_turn"}, + "usage": {"input_tokens": 10, "output_tokens": 20} + }))); + let events = em.drain(); + + if let StreamEvent::MessageDelta(md) = &events[0] { + assert_eq!(md.delta.stop_reason.as_deref(), Some("end_turn")); + assert_eq!(md.usage.as_ref().unwrap().input_tokens, 10); + assert_eq!(md.usage.as_ref().unwrap().output_tokens, 20); + } else { + panic!("expected MessageDelta"); + } + } + + #[test] + fn emitter_message_stop_closes_parts() { + let mut em = StreamEmitter::default(); + em.text_part_open = true; + em.tool_parts_open = 2; + + em.on_message_stop(); + let events = em.drain(); + // 1 for text + 2 for tools + assert_eq!(events.len(), 3); + assert!(events.iter().all(|e| matches!(e, StreamEvent::PartStop))); + } + + #[test] + fn emitter_finish_emits_message_stop_if_needed() { + let mut em = StreamEmitter::default(); + em.started = true; + em.finished = false; + + let events = em.finish(); + assert!(events.iter().any(|e| matches!(e, StreamEvent::MessageStop))); + } + + #[test] + fn emitter_finish_noop_if_already_stopped() { + let mut em = StreamEmitter::default(); + em.started = true; + em.finished = true; + + let events = em.finish(); + assert!(events.is_empty()); + } + + #[test] + fn sse_reader_take_line_extracts_newline() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "event: ping\n".into(), + }; + assert_eq!(reader.take_line().unwrap(), "event: ping"); + assert!(reader.buf.is_empty()); + } + + #[test] + fn sse_reader_take_line_none_without_newline() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "partial".into(), + }; + assert!(reader.take_line().is_none()); + } + + #[test] + fn sse_reader_take_line_multiple() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "line1\nline2\n".into(), + }; + assert_eq!(reader.take_line().unwrap(), "line1"); + assert_eq!(reader.take_line().unwrap(), "line2"); + } + + #[test] + fn sse_reader_take_line_trims_cr() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "data: hi\r\n".into(), + }; + assert_eq!(reader.take_line().unwrap(), "data: hi"); + } +} diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs new file mode 100644 index 0000000..b8ef4ce --- /dev/null +++ b/src/provider/gemini.rs @@ -0,0 +1,909 @@ +//! Google Gemini API client. +//! +//! Implements [`ApiClient`] by translating between the framework's +//! [`StreamEvent`] protocol and the Gemini Streaming Generate Content +//! API (SSE format). +//! +//! # Construction +//! +//! ```rust,ignore +//! use loopctl::provider::GeminiClient; +//! +//! // From environment (GEMINI_API_KEY or GOOGLE_API_KEY): +//! let client = GeminiClient::from_env()?; +//! +//! // Explicit: +//! let client = GeminiClient::builder() +//! .api_key("AIza...") +//! .model("gemini-2.0-flash") +//! .build()?; +//! ``` + +use std::future::Future; +use std::pin::Pin; + +use futures::stream::{Stream, StreamExt}; +use reqwest::Response; +use serde_json::Value; + +use crate::api::ApiClient; +use crate::api::error::ApiError; +use crate::message::{Message, MessagePart, Role}; +use crate::stream::{ + DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, + PartStart, StreamEvent, StreamStopReason, +}; +use crate::tool::ToolSchema; + +// ================================================== +// Constants +// ================================================== + +const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta"; +const DEFAULT_MODEL: &str = "gemini-2.0-flash"; +const SSE_DATA_PREFIX: &str = "data: "; +const TEXT_PART_INDEX: usize = 0; + +// ================================================== +// Client +// ================================================== + +/// A Google Gemini API client with streaming support. +/// +/// Implements [`ApiClient`] by translating between the framework's +/// [`StreamEvent`] protocol and the Gemini Streaming Generate Content API. +pub struct GeminiClient { + http: reqwest::Client, + api_key: String, + base_url: String, + model: String, +} + +impl GeminiClient { + /// Create a builder for configuring a [`GeminiClient`]. + #[must_use] + pub fn builder() -> GeminiClientBuilder { + GeminiClientBuilder::default() + } + + /// Create from environment variables. + /// + /// Reads: + /// - `GEMINI_API_KEY` (or `GOOGLE_API_KEY`) — required. + /// - `GEMINI_BASE_URL` — optional, defaults to + /// `https://generativelanguage.googleapis.com/v1beta`. + /// - `GEMINI_MODEL` — optional, defaults to `gemini-2.0-flash`. + /// + /// # Errors + /// + /// Returns [`ApiError`] if no API key is found. + pub fn from_env() -> Result { + let api_key = std::env::var("GEMINI_API_KEY") + .or_else(|_| std::env::var("GOOGLE_API_KEY")) + .map_err(|_| ApiError::auth_invalid_key("GEMINI_API_KEY not set"))?; + let base_url = std::env::var("GEMINI_BASE_URL").unwrap_or_else(|_| DEFAULT_BASE_URL.into()); + let model = std::env::var("GEMINI_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.into()); + + Self::builder() + .api_key(api_key) + .base_url(base_url) + .model(model) + .build() + } + + /// Build the streaming Generate Content URL. + /// + /// Gemini puts the model in the URL path and the API key as a query + /// parameter rather than using headers. + fn stream_url(&self) -> String { + format!( + "{}/models/{}:streamGenerateContent?alt=sse&key={}", + self.base_url, self.model, self.api_key + ) + } + + /// Build the non-streaming Generate Content URL. + fn generate_url(&self) -> String { + format!( + "{}/models/{}:generateContent?key={}", + self.base_url, self.model, self.api_key + ) + } + + /// Send a POST request and return the raw response. + /// + /// Shared by both [`ApiClient::stream_messages`] and + /// [`ApiClient::create_message`]. + /// + /// # Errors + /// + /// Returns [`ApiError`] if the request fails or the server + /// responds with a non-success status code. + async fn post_content( + http: &reqwest::Client, + url: &str, + body: &Value, + ) -> Result { + let resp = http + .post(url) + .json(body) + .send() + .await + .map_err(|e| ApiError::http(e.to_string()))?; + let status = resp.status(); + if status.is_success() { + Ok(resp) + } else { + let text = resp.text().await.unwrap_or_default(); + Err(ApiError::http_with_status(status.as_u16(), text)) + } + } +} + +impl ApiClient for GeminiClient { + fn model(&self) -> &str { + &self.model + } + + fn stream_messages( + &self, + messages: Vec, + system: Option, + tools: Option>, + ) -> Pin> + Send + 'static>> { + let body = build_request_body(&messages, system.as_deref(), tools.as_deref()); + let url = self.stream_url(); + let http = self.http.clone(); + + Box::pin(async_stream::try_stream! { + let resp = Self::post_content(&http, &url, &body).await?; + let mut sse = SseReader::from_response(resp); + let mut emitter = StreamEmitter::default(); + + while let Some(data) = sse.next_data().await? { + emitter.process_chunk(&data); + for ev in emitter.drain() { + yield ev; + } + } + + for ev in emitter.finish() { + yield ev; + } + }) + } + + fn create_message( + &self, + messages: Vec, + system: Option, + tools: Option>, + ) -> Pin> + Send + '_>> { + let body = build_request_body(&messages, system.as_deref(), tools.as_deref()); + let url = self.generate_url(); + + Box::pin(async move { + let resp = Self::post_content(&self.http, &url, &body).await?; + resp.json::() + .await + .map_err(|e| ApiError::http(e.to_string())) + }) + } +} + +// ================================================== +// Builder +// ================================================== + +/// Builder for [`GeminiClient`]. +pub struct GeminiClientBuilder { + api_key: Option, + base_url: String, + model: String, +} + +impl Default for GeminiClientBuilder { + fn default() -> Self { + Self { + api_key: None, + base_url: DEFAULT_BASE_URL.into(), + model: DEFAULT_MODEL.into(), + } + } +} + +impl GeminiClientBuilder { + /// Set the API key. + #[must_use] + pub fn api_key(mut self, key: impl Into) -> Self { + self.api_key = Some(key.into()); + self + } + + /// Set the base URL. + #[must_use] + pub fn base_url(mut self, url: impl Into) -> Self { + self.base_url = url.into(); + self + } + + /// Set the model name (e.g. `gemini-2.0-flash`). + #[must_use] + pub fn model(mut self, model: impl Into) -> Self { + self.model = model.into(); + self + } + + /// Build the client. + /// + /// # Errors + /// + /// Returns [`ApiError`] if no API key was set. + pub fn build(self) -> Result { + let api_key = self + .api_key + .ok_or_else(|| ApiError::auth_invalid_key("API key not provided"))?; + let http = reqwest::Client::builder() + .build() + .map_err(|e| ApiError::http(e.to_string()))?; + + Ok(GeminiClient { + http, + api_key, + base_url: self.base_url, + model: self.model, + }) + } +} + +// ================================================== +// Request body construction +// ================================================== + +/// Build the JSON request body for the Gemini Generate Content API. +/// +/// Unlike OpenAI/Anthropic, Gemini puts the model in the URL, not the +/// request body. Each [`Message`] is serialized via [`convert_message`]. +fn build_request_body( + messages: &[Message], + system: Option<&str>, + tools: Option<&[ToolSchema]>, +) -> Value { + let contents: Vec = messages.iter().map(convert_message).collect(); + let mut body = serde_json::json!({ "contents": contents }); + if let Some(obj) = body.as_object_mut() { + if let Some(sys) = system { + obj.insert( + "systemInstruction".into(), + serde_json::json!({"parts": [{"text": sys}]}), + ); + } + if let Some(tool_list) = tools { + obj.insert( + "tools".into(), + serde_json::json!([{"functionDeclarations": convert_tools(tool_list)}]), + ); + } + } + + body +} + +/// Convert a single framework [`Message`] into the Gemini JSON shape. +/// +/// Gemini uses `role: "user"` / `role: "model"` (not "assistant") and +/// a `parts` array for content blocks. +fn convert_message(m: &Message) -> Value { + let role = match m.role { + Role::User => "user", + Role::Assistant => "model", + }; + let parts: Vec = m.parts.iter().filter_map(convert_part).collect(); + serde_json::json!({"role": role, "parts": parts}) +} + +/// Convert a single [`MessagePart`] into a Gemini part JSON object. +/// +/// Returns `None` for image parts (not yet supported for Gemini). +fn convert_part(p: &MessagePart) -> Option { + match p { + MessagePart::Text { text } => Some(serde_json::json!({"text": text})), + MessagePart::ToolCall { name, input, .. } => Some(serde_json::json!({ + "functionCall": { + "name": name, + "args": input, + } + })), + MessagePart::ToolResult { + call_id, output, .. + } => Some(serde_json::json!({ + "functionResponse": { + "name": call_id, + "response": {"result": output.to_string()}, + } + })), + MessagePart::Image { .. } => None, + } +} + +/// Convert tool schemas into the Gemini `functionDeclarations` array. +fn convert_tools(tools: &[ToolSchema]) -> Vec { + tools + .iter() + .map(|t| { + serde_json::json!({ + "name": t.tool, + "description": &t.description, + "parameters": t.input_schema.clone(), + }) + }) + .collect() +} + +// ================================================== +// SSE line reader +// ================================================== + +/// Minimal SSE line reader over an HTTP byte stream. +/// +/// Buffers raw bytes from the response, splits on newlines, and yields +/// the JSON `data` payload of each SSE event. Gemini SSE uses only +/// `data:` lines (no `event:` type headers). +struct SseReader { + bytes: Pin> + Send>>, + buf: String, +} + +impl SseReader { + /// Wrap a streaming HTTP response. + fn from_response(resp: Response) -> Self { + let bytes = resp.bytes_stream().map(|res| { + res.map(|b| String::from_utf8_lossy(&b).into_owned()) + .map_err(|e| ApiError::http(e.to_string())) + }); + + Self { + bytes: Box::pin(bytes), + buf: String::new(), + } + } + + /// Extract the next SSE `data:` payload as parsed JSON. + /// + /// Returns `Ok(None)` at end-of-stream. + /// + /// # Errors + /// + /// Returns [`ApiError`] if the underlying HTTP stream fails. + async fn next_data(&mut self) -> Result, ApiError> { + loop { + while let Some(line) = self.take_line() { + if line.is_empty() { + continue; + } + + if let Some(data) = line.strip_prefix(SSE_DATA_PREFIX) { + if let Ok(json) = serde_json::from_str::(data) { + return Ok(Some(json)); + } + } + } + + match self.bytes.next().await { + Some(Ok(chunk)) => self.buf.push_str(&chunk), + Some(Err(e)) => return Err(e), + None => return Ok(None), + } + } + } + + /// Pop the first `\n`-terminated line from the buffer, if present. + fn take_line(&mut self) -> Option { + let pos = self.buf.find('\n')?; + let line = self.buf[..pos].trim().to_string(); + let rest_start = pos.saturating_add(1); + self.buf = self.buf.get(rest_start..).unwrap_or_default().to_string(); + + Some(line) + } +} + +// ================================================== +// Stream event emitter +// ================================================== + +/// Stateful translator that converts Gemini SSE chunks into +/// [`StreamEvent`]s. +/// +/// Gemini's SSE format is simpler than Anthropic's: each `data:` line +/// is a complete JSON object with `candidates[0].content.parts[]` for +/// text/function-call data and `candidates[0].finishReason` for the +/// stop reason. +#[derive(Default)] +struct StreamEmitter { + started: bool, + finished: bool, + pending: Vec, +} + +impl StreamEmitter { + /// Process a single Gemini SSE chunk, appending events to the queue. + fn process_chunk(&mut self, json: &Value) { + if !self.started { + self.started = true; + self.push(StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: String::new(), + role: "assistant".into(), + model: String::new(), + }, + })); + } + + self.extract_text(json); + self.extract_function_call(json); + self.extract_finish_reason(json); + } + + /// Extract text delta from `candidates[0].content.parts[0].text`. + fn extract_text(&mut self, json: &Value) { + if let Some(text) = json + .pointer("/candidates/0/content/parts/0/text") + .and_then(Value::as_str) + { + if !text.is_empty() { + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: TEXT_PART_INDEX, + delta: DeltaPart::Text { + text: text.to_string(), + }, + })); + } + } + } + + /// Extract function call from `candidates[0].content.parts[0].functionCall`. + fn extract_function_call(&mut self, json: &Value) { + if let Some(func_call) = json.pointer("/candidates/0/content/parts/0/functionCall") { + let name = func_call + .pointer("/name") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + let args = func_call.pointer("/args").cloned().unwrap_or(Value::Null); + let args_str = serde_json::to_string(&args).unwrap_or_default(); + + self.push(StreamEvent::PartStart(PartStart { + index: TEXT_PART_INDEX, + part: Some(MessagePart::ToolCall { + id: String::new(), + name, + input: args, + }), + })); + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: TEXT_PART_INDEX, + delta: DeltaPart::InputJson { + partial_json: args_str, + }, + })); + } + } + + /// Extract finish reason and emit stop events. + fn extract_finish_reason(&mut self, json: &Value) { + let Some(reason) = json + .pointer("/candidates/0/finishReason") + .and_then(Value::as_str) + else { + return; + }; + let stop = match reason { + "MAX_TOKENS" => StreamStopReason::MaxTokens, + _ => StreamStopReason::EndTurn, + }; + + self.push(StreamEvent::PartStop); + self.push(StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some(stop.to_api_str().into()), + }, + usage: None, + })); + } + + /// Drain all pending events. + fn drain(&mut self) -> Vec { + std::mem::take(&mut self.pending) + } + + /// Emit the terminal [`MessageStop`] if the stream was started. + fn finish(&mut self) -> Vec { + let mut out = self.drain(); + if self.started && !self.finished { + self.finished = true; + out.push(StreamEvent::MessageStop); + } + out + } + + fn push(&mut self, ev: StreamEvent) { + self.pending.push(ev); + } +} + +// ================================================== +// Tests +// ================================================== + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::{Message, MessagePart, Role, ToolContent}; + + #[test] + fn request_body_user_text() { + let msgs = vec![Message::user("hello")]; + let body = build_request_body(&msgs, None, None); + + let contents = body["contents"].as_array().unwrap(); + assert_eq!(contents.len(), 1); + assert_eq!(contents[0]["role"], "user"); + let parts = contents[0]["parts"].as_array().unwrap(); + assert_eq!(parts[0]["text"], "hello"); + } + + #[test] + fn request_body_includes_system_instruction() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body(&msgs, Some("be brief"), None); + + let sys = &body["systemInstruction"]; + assert!(sys.is_object()); + assert_eq!(sys["parts"][0]["text"], "be brief"); + } + + #[test] + fn request_body_no_system_instruction_when_none() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body(&msgs, None, None); + assert!(body.get("systemInstruction").is_none()); + } + + #[test] + fn request_body_assistant_maps_to_model_role() { + let msgs = vec![Message::new( + Role::Assistant, + vec![MessagePart::text("hello")], + )]; + let body = build_request_body(&msgs, None, None); + assert_eq!(body["contents"][0]["role"], "model"); + } + + #[test] + fn request_body_user_role() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body(&msgs, None, None); + assert_eq!(body["contents"][0]["role"], "user"); + } + + #[test] + fn request_body_assistant_tool_call() { + let msgs = vec![Message::new( + Role::Assistant, + vec![MessagePart::ToolCall { + id: "call_1".into(), + name: "echo".into(), + input: serde_json::json!({"msg": "hi"}), + }], + )]; + let body = build_request_body(&msgs, None, None); + + let parts = body["contents"][0]["parts"].as_array().unwrap(); + assert_eq!(parts[0]["functionCall"]["name"], "echo"); + assert_eq!(parts[0]["functionCall"]["args"]["msg"], "hi"); + } + + #[test] + fn request_body_tool_result() { + let msgs = vec![Message::new( + Role::User, + vec![MessagePart::ToolResult { + call_id: "call_1".into(), + output: ToolContent::from_string("result text"), + is_error: None, + }], + )]; + let body = build_request_body(&msgs, None, None); + + let parts = body["contents"][0]["parts"].as_array().unwrap(); + assert_eq!(parts[0]["functionResponse"]["name"], "call_1"); + assert_eq!( + parts[0]["functionResponse"]["response"]["result"], + "result text" + ); + } + + #[test] + fn request_body_includes_tools() { + let msgs = vec![Message::user("hi")]; + let tools = vec![ToolSchema { + tool: "search".into(), + description: "Search the web".into(), + input_schema: serde_json::json!({"type": "object"}), + }]; + let body = build_request_body(&msgs, None, Some(&tools)); + + let tools_arr = body["tools"].as_array().unwrap(); + assert_eq!(tools_arr.len(), 1); + let decls = tools_arr[0]["functionDeclarations"].as_array().unwrap(); + assert_eq!(decls[0]["name"], "search"); + assert_eq!(decls[0]["description"], "Search the web"); + assert_eq!(decls[0]["parameters"]["type"], "object"); + } + + #[test] + fn request_body_no_tools_when_none() { + let msgs = vec![Message::user("hi")]; + let body = build_request_body(&msgs, None, None); + assert!(body.get("tools").is_none()); + } + + #[test] + fn request_body_multiple_messages() { + let msgs = vec![ + Message::user("hello"), + Message::new(Role::Assistant, vec![MessagePart::text("hi")]), + Message::user("bye"), + ]; + let body = build_request_body(&msgs, None, None); + + let contents = body["contents"].as_array().unwrap(); + assert_eq!(contents.len(), 3); + assert_eq!(contents[0]["role"], "user"); + assert_eq!(contents[1]["role"], "model"); + assert_eq!(contents[2]["role"], "user"); + } + + #[test] + fn convert_tools_shape() { + let tools = vec![ToolSchema { + tool: "calc".into(), + description: "Calculate".into(), + input_schema: serde_json::json!({"type": "object"}), + }]; + let out = convert_tools(&tools); + assert_eq!(out.len(), 1); + assert_eq!(out[0]["name"], "calc"); + assert_eq!(out[0]["description"], "Calculate"); + } + + #[test] + fn convert_message_text_only() { + let m = Message::user("hello"); + let v = convert_message(&m); + assert_eq!(v["role"], "user"); + assert_eq!(v["parts"][0]["text"], "hello"); + } + + #[test] + fn convert_message_assistant_role() { + let m = Message::new(Role::Assistant, vec![MessagePart::text("hi")]); + let v = convert_message(&m); + assert_eq!(v["role"], "model"); + } + + #[test] + fn convert_message_skips_images() { + let m = Message::new( + Role::User, + vec![ + MessagePart::text("look"), + MessagePart::Image { + source: crate::message::ImageSource { + encoding: "base64".into(), + media_type: "image/png".into(), + data: String::new(), + }, + }, + ], + ); + let v = convert_message(&m); + let parts = v["parts"].as_array().unwrap(); + assert_eq!(parts.len(), 1); // only text, image filtered out + } + + #[test] + fn builder_requires_api_key() { + let result = GeminiClient::builder().build(); + assert!(result.is_err()); + } + + #[test] + fn builder_succeeds_with_key() { + let client = GeminiClient::builder().api_key("test-key").build().unwrap(); + assert_eq!(client.model(), DEFAULT_MODEL); + } + + #[test] + fn builder_custom_model() { + let client = GeminiClient::builder() + .api_key("test-key") + .model("gemini-1.5-pro") + .build() + .unwrap(); + assert_eq!(client.model(), "gemini-1.5-pro"); + } + + #[test] + fn builder_custom_base_url() { + let client = GeminiClient::builder() + .api_key("test-key") + .base_url("https://custom.example.com") + .model("gemini-pro") + .build() + .unwrap(); + assert_eq!(client.model(), "gemini-pro"); + } + + #[test] + fn emitter_first_chunk_emits_message_start() { + let mut em = StreamEmitter::default(); + em.process_chunk(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"text": "hi"}]}}] + })); + let events = em.drain(); + assert!( + events + .iter() + .any(|e| matches!(e, StreamEvent::MessageStart(_))) + ); + } + + #[test] + fn emitter_text_delta() { + let mut em = StreamEmitter::default(); + em.started = true; // skip MessageStart + em.process_chunk(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"text": "world"}]}}] + })); + let events = em.drain(); + assert!( + events + .iter() + .any(|e| matches!(e, StreamEvent::IndexedDelta(_))) + ); + } + + #[test] + fn emitter_empty_text_ignored() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"text": ""}]}}] + })); + let events = em.drain(); + // Only MessageStart would be here, but we set started=true so + // no events at all for empty text. + assert!( + events + .iter() + .all(|e| !matches!(e, StreamEvent::IndexedDelta(_))) + ); + } + + #[test] + fn emitter_function_call() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"functionCall": {"name": "search", "args": {"q": "rust"}}}]}}] + })); + let events = em.drain(); + assert!( + events + .iter() + .any(|e| matches!(e, StreamEvent::PartStart(_))) + ); + assert!( + events + .iter() + .any(|e| matches!(e, StreamEvent::IndexedDelta(_))) + ); + } + + #[test] + fn emitter_finish_reason_end_turn() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{"finishReason": "STOP"}] + })); + let events = em.drain(); + assert!(events.iter().any(|e| matches!(e, StreamEvent::PartStop))); + let md = events + .iter() + .find(|e| matches!(e, StreamEvent::MessageDelta(_))); + if let Some(StreamEvent::MessageDelta(d)) = md { + assert_eq!(d.delta.stop_reason.as_deref(), Some("end_turn")); + } else { + panic!("expected MessageDelta"); + } + } + + #[test] + fn emitter_finish_reason_max_tokens() { + let mut em = StreamEmitter::default(); + em.started = true; + em.process_chunk(&serde_json::json!({ + "candidates": [{"finishReason": "MAX_TOKENS"}] + })); + let events = em.drain(); + let md = events + .iter() + .find(|e| matches!(e, StreamEvent::MessageDelta(_))); + if let Some(StreamEvent::MessageDelta(d)) = md { + assert_eq!(d.delta.stop_reason.as_deref(), Some("max_tokens")); + } else { + panic!("expected MessageDelta"); + } + } + + #[test] + fn emitter_finish_emits_message_stop_if_needed() { + let mut em = StreamEmitter::default(); + em.started = true; + em.finished = false; + + let events = em.finish(); + assert!(events.iter().any(|e| matches!(e, StreamEvent::MessageStop))); + } + + #[test] + fn emitter_finish_noop_if_already_stopped() { + let mut em = StreamEmitter::default(); + em.started = true; + em.finished = true; + + let events = em.finish(); + assert!(events.is_empty()); + } + + #[test] + fn sse_reader_take_line_extracts_newline() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "data: hello\n".into(), + }; + assert_eq!(reader.take_line().unwrap(), "data: hello"); + assert!(reader.buf.is_empty()); + } + + #[test] + fn sse_reader_take_line_none_without_newline() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "partial".into(), + }; + assert!(reader.take_line().is_none()); + } + + #[test] + fn sse_reader_take_line_multiple() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "line1\nline2\n".into(), + }; + assert_eq!(reader.take_line().unwrap(), "line1"); + assert_eq!(reader.take_line().unwrap(), "line2"); + } + + #[test] + fn sse_reader_take_line_trims_cr() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "data: hi\r\n".into(), + }; + assert_eq!(reader.take_line().unwrap(), "data: hi"); + } +} diff --git a/src/provider/openai.rs b/src/provider/openai.rs new file mode 100644 index 0000000..376ba30 --- /dev/null +++ b/src/provider/openai.rs @@ -0,0 +1,1184 @@ +//! OpenAI-compatible API client. +//! +//! Works with any provider that implements the OpenAI Chat Completions +//! API with streaming: OpenAI itself, `DeepSeek`, `Grok`, Ollama (via +//! the `ollama()` constructor in the parent module), vLLM, LM Studio, etc. +//! +//! # Construction +//! +//! ```rust,ignore +//! use loopctl::provider::OpenAiClient; +//! +//! // From environment (OPENAI_API_KEY, optional OPENAI_BASE_URL): +//! let client = OpenAiClient::from_env()?; +//! +//! // Explicit: +//! let client = OpenAiClient::builder() +//! .api_key("sk-...") +//! .base_url("https://api.deepseek.com/v1") +//! .model("deepseek-chat") +//! .build()?; +//! ``` + +use std::future::Future; +use std::pin::Pin; + +use futures::stream::{Stream, StreamExt}; +use reqwest::Response; +use serde::Deserialize; +use serde_json::Value; + +use crate::api::ApiClient; +use crate::api::error::ApiError; +use crate::message::{Message, MessagePart, Role}; +use crate::stream::{ + DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, + PartStart, StreamEvent, StreamStopReason, +}; +use crate::tool::ToolSchema; + +// ================================================== +// Constants +// ================================================== + +const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1"; +const DEFAULT_MODEL: &str = "gpt-4o"; +const SSE_DONE: &str = "[DONE]"; +const SSE_DATA_PREFIX: &str = "data: "; +const TEXT_PART_INDEX: usize = 0; + +// ================================================== +// Client +// ================================================== + +/// An OpenAI-compatible chat completions client with streaming support. +/// +/// Implements [`ApiClient`] by translating between the framework's +/// [`StreamEvent`] protocol and the OpenAI Chat Completions SSE format. +/// +/// Works with any OpenAI-compatible endpoint. Use a custom `base_url` +/// to target `DeepSeek`, `Grok`, Ollama, `vLLM`, or other compatible APIs. +pub struct OpenAiClient { + http: reqwest::Client, + api_key: String, + base_url: String, + model: String, +} + +impl OpenAiClient { + /// Create a builder for configuring an [`OpenAiClient`]. + #[must_use] + pub fn builder() -> OpenAiClientBuilder { + OpenAiClientBuilder::default() + } + + /// Create from environment variables. + /// + /// Reads: + /// - `OPENAI_API_KEY` (or `API_KEY`) — required. + /// - `OPENAI_BASE_URL` (or `BASE_URL`) — optional, defaults to + /// `https://api.openai.com/v1`. + /// - `OPENAI_MODEL` (or `MODEL`) — optional, defaults to `gpt-4o`. + /// + /// # Errors + /// + /// Returns [`ApiError`] if no API key is found. + pub fn from_env() -> Result { + let api_key = std::env::var("OPENAI_API_KEY") + .or_else(|_| std::env::var("API_KEY")) + .map_err(|_| ApiError::auth_invalid_key("OPENAI_API_KEY not set"))?; + + let base_url = std::env::var("OPENAI_BASE_URL") + .or_else(|_| std::env::var("BASE_URL")) + .unwrap_or_else(|_| DEFAULT_BASE_URL.into()); + + let model = std::env::var("OPENAI_MODEL") + .or_else(|_| std::env::var("MODEL")) + .unwrap_or_else(|_| DEFAULT_MODEL.into()); + + Self::builder() + .api_key(api_key) + .base_url(base_url) + .model(model) + .build() + } + + /// Build the chat-completions URL for this client. + fn completions_url(&self) -> String { + format!("{}/chat/completions", self.base_url) + } + + /// Send a POST request to the chat-completions endpoint. + /// + /// Shared by both [`ApiClient::stream_messages`] and + /// [`ApiClient::create_message`]. Returns the raw + /// [`reqwest::Response`] after checking for HTTP errors. + /// + /// # Errors + /// + /// Returns [`ApiError`] if the request fails or the server + /// responds with a non-success status code. + async fn post_completions( + http: &reqwest::Client, + url: &str, + api_key: &str, + body: &Value, + ) -> Result { + let resp = http + .post(url) + .bearer_auth(api_key) + .json(body) + .send() + .await + .map_err(|e| ApiError::http(e.to_string()))?; + + let status = resp.status(); + if status.is_success() { + Ok(resp) + } else { + let text = resp.text().await.unwrap_or_default(); + Err(ApiError::http_with_status(status.as_u16(), text)) + } + } +} + +impl ApiClient for OpenAiClient { + fn model(&self) -> &str { + &self.model + } + + fn stream_messages( + &self, + messages: Vec, + system: Option, + tools: Option>, + ) -> Pin> + Send + 'static>> { + let body = RequestBody::build(&self.model, &messages, system.as_deref(), tools.as_deref()); + let url = self.completions_url(); + let api_key = self.api_key.clone(); + let http = self.http.clone(); + + Box::pin(async_stream::try_stream! { + let resp = Self::post_completions(&http, &url, &api_key, &body.to_json(true)).await?; + let mut sse = SseReader::from_response(resp); + let mut emitter = StreamEmitter::default(); + + while let Some(data) = sse.next_data().await? { + let Some(chunk) = OpenAiChunk::parse(&data) else { + continue; + }; + emitter.process_chunk(&chunk); + for ev in emitter.drain() { + yield ev; + } + } + + for ev in emitter.finish() { + yield ev; + } + }) + } + + fn create_message( + &self, + messages: Vec, + system: Option, + tools: Option>, + ) -> Pin> + Send + '_>> { + let body = RequestBody::build(&self.model, &messages, system.as_deref(), tools.as_deref()); + let url = self.completions_url(); + + Box::pin(async move { + let resp = + Self::post_completions(&self.http, &url, &self.api_key, &body.to_json(false)) + .await?; + resp.json::() + .await + .map_err(|e| ApiError::http(e.to_string())) + }) + } +} + +// ================================================== +// Builder +// ================================================== + +/// Builder for [`OpenAiClient`]. +pub struct OpenAiClientBuilder { + api_key: Option, + base_url: String, + model: String, +} + +impl Default for OpenAiClientBuilder { + fn default() -> Self { + Self { + api_key: None, + base_url: DEFAULT_BASE_URL.into(), + model: DEFAULT_MODEL.into(), + } + } +} + +impl OpenAiClientBuilder { + /// Set the API key. + #[must_use] + pub fn api_key(mut self, key: impl Into) -> Self { + self.api_key = Some(key.into()); + self + } + + /// Set the base URL (e.g. `https://api.deepseek.com/v1`). + #[must_use] + pub fn base_url(mut self, url: impl Into) -> Self { + self.base_url = url.into(); + self + } + + /// Set the model name (e.g. `gpt-4o`, `deepseek-chat`). + #[must_use] + pub fn model(mut self, model: impl Into) -> Self { + self.model = model.into(); + self + } + + /// Build the client. + /// + /// # Errors + /// + /// Returns [`ApiError`] if no API key was set. + pub fn build(self) -> Result { + let api_key = self + .api_key + .ok_or_else(|| ApiError::auth_invalid_key("API key not provided"))?; + + let http = reqwest::Client::builder() + .build() + .map_err(|e| ApiError::http(e.to_string()))?; + + Ok(OpenAiClient { + http, + api_key, + base_url: self.base_url, + model: self.model, + }) + } +} + +// ================================================== +// Request body construction +// ================================================== + +/// A built OpenAI Chat Completions request body. +/// +/// Separating construction from serialization lets us reuse the same +/// body for both streaming and non-streaming requests, toggling only +/// the `stream` flag via [`to_json`](Self::to_json). +struct RequestBody { + model: String, + messages: Vec, + tools: Option>, +} + +impl RequestBody { + /// Translate the framework's [`Message`] list into the OpenAI + /// Chat Completions request shape. + fn build( + model: &str, + messages: &[Message], + system: Option<&str>, + tools: Option<&[ToolSchema]>, + ) -> Self { + let mut msgs = Vec::with_capacity(messages.len().saturating_add(1)); + + if let Some(sys) = system { + msgs.push(serde_json::json!({ "role": "system", "content": sys })); + } + + for m in messages { + msgs.push(convert_message(m)); + } + + Self { + model: model.into(), + messages: msgs, + tools: tools.map(convert_tools), + } + } + + /// Serialize to a [`serde_json::Value`] with the `stream` flag + /// set as requested. + fn to_json(&self, stream: bool) -> Value { + serde_json::json!({ + "model": self.model, + "messages": self.messages, + "stream": stream, + "tools": self.tools, + }) + } +} + +/// Convert a single framework [`Message`] into the OpenAI JSON shape. +/// +/// OpenAI expects assistant messages with `tool_calls` to carry them in +/// a dedicated array, tool results to use the `tool` role, and plain +/// text to use a simple `{role, content}` pair. +fn convert_message(m: &Message) -> Value { + let role = match m.role { + Role::User => "user", + Role::Assistant => "assistant", + }; + + // Bucket parts by OpenAI category. + let mut text_parts: Vec<&str> = Vec::new(); + let mut tool_calls: Vec = Vec::new(); + let mut tool_results: Vec = Vec::new(); + + for p in &m.parts { + match p { + MessagePart::Text { text } => text_parts.push(text.as_str()), + MessagePart::ToolCall { id, name, input } => { + tool_calls.push(serde_json::json!({ + "id": id, + "type": "function", + "function": { + "name": name, + "arguments": input_to_string(input), + } + })); + } + MessagePart::ToolResult { + call_id, output, .. + } => { + tool_results.push(serde_json::json!({ + "role": "tool", + "tool_call_id": call_id, + "content": output.to_string(), + })); + } + MessagePart::Image { .. } => {} // not supported in this path + } + } + + if !tool_calls.is_empty() { + build_assistant_message(role, &tool_calls, &text_parts) + } else if !tool_results.is_empty() { + merge_tool_results(&tool_results) + } else { + serde_json::json!({ "role": role, "content": text_parts.join("") }) + } +} + +/// Build an assistant message JSON that includes `tool_calls`. +fn build_assistant_message(role: &str, tool_calls: &[Value], text_parts: &[&str]) -> Value { + let text = text_parts.join(""); + let content = if text.is_empty() { + Value::Null + } else { + Value::String(text) + }; + serde_json::json!({ + "role": role, + "content": content, + "tool_calls": tool_calls, + }) +} + +/// Merge one or more tool-result entries into a single JSON value. +/// +/// When there is only one result (the common case) we return it +/// directly; otherwise we return a JSON array so no data is lost. +fn merge_tool_results(results: &[Value]) -> Value { + if results.len() == 1 { + results.first().cloned().unwrap_or(Value::Null) + } else { + Value::Array(results.to_vec()) + } +} + +/// Convert tool schemas into the OpenAI `tools` array shape. +fn convert_tools(tools: &[ToolSchema]) -> Vec { + tools + .iter() + .map(|t| { + serde_json::json!({ + "type": "function", + "function": { + "name": t.tool, + "description": &t.description, + "parameters": t.input_schema.clone(), + } + }) + }) + .collect() +} + +/// Serialize a JSON value to a compact string for the OpenAI `arguments` field. +/// +/// OpenAI expects `arguments` to be a string containing JSON, so a raw +/// [`Value`] must be stringified. If the value is already a string we +/// pass it through unchanged. +fn input_to_string(v: &Value) -> String { + match v { + Value::String(s) => s.clone(), + other => other.to_string(), + } +} + +// ================================================== +// SSE line reader +// ================================================== + +/// Minimal SSE line reader over an HTTP byte stream. +/// +/// Buffers raw bytes from the response, splits on newlines, and yields +/// the payload of each `data:` line (without the `data: ` prefix). +/// A `[DONE]` sentinel terminates the stream. +struct SseReader { + bytes: Pin> + Send>>, + buf: String, +} + +impl SseReader { + /// Wrap a streaming HTTP response. + /// + /// The byte stream is mapped to `String` chunks up-front so the + /// rest of the reader is pure string processing. + fn from_response(resp: Response) -> Self { + let bytes = resp.bytes_stream().map(|res| { + res.map(|b| String::from_utf8_lossy(&b).into_owned()) + .map_err(|e| ApiError::http(e.to_string())) + }); + Self { + bytes: Box::pin(bytes), + buf: String::new(), + } + } + + /// Extract the next SSE `data:` payload, blocking until one is + /// available or the stream ends. + /// + /// Returns `Ok(None)` at end-of-stream (including `[DONE]`). + /// + /// # Errors + /// + /// Returns [`ApiError`] if the underlying HTTP stream fails. + async fn next_data(&mut self) -> Result, ApiError> { + loop { + // Drain any complete lines already in the buffer. + while let Some(line) = self.take_line() { + let Some(data) = line.strip_prefix(SSE_DATA_PREFIX) else { + continue; + }; + if data == SSE_DONE { + return Ok(None); + } + return Ok(Some(data.into())); + } + + // Fetch the next chunk from the network. + match self.bytes.next().await { + Some(Ok(chunk)) => self.buf.push_str(&chunk), + Some(Err(e)) => return Err(e), + None => return Ok(None), + } + } + } + + /// Pop the first `\n`-terminated line from the buffer, if present. + fn take_line(&mut self) -> Option { + let pos = self.buf.find('\n')?; + let line = self.buf[..pos].trim().to_string(); + let rest_start = pos.saturating_add(1); + self.buf = self.buf.get(rest_start..).unwrap_or_default().to_string(); + Some(line) + } +} + +// ================================================== +// OpenAI chunk types +// ================================================== + +/// A single SSE chunk from the OpenAI streaming API. +#[derive(Deserialize)] +struct OpenAiChunk { + id: String, + model: String, + choices: Vec, +} + +impl OpenAiChunk { + /// Parse a raw SSE data payload into an [`OpenAiChunk`]. + /// + /// Returns `None` for malformed payloads so the caller can skip + /// them without interrupting the stream. + fn parse(data: &str) -> Option { + serde_json::from_str(data).ok() + } +} + +#[derive(Deserialize)] +struct OpenAiChoice { + delta: Option, + finish_reason: Option, +} + +#[derive(Deserialize)] +struct OpenAiDelta { + content: Option, + tool_calls: Option>, +} + +#[derive(Deserialize)] +struct OpenAiToolCallDelta { + index: usize, + id: String, + function: Option, +} + +#[derive(Deserialize)] +struct OpenAiToolCallFunction { + name: String, + arguments: String, +} + +// ================================================== +// Stream event emitter +// ================================================== + +/// Stateful translator that converts a sequence of [`OpenAiChunk`]s +/// into [`StreamEvent`]s. +/// +/// This encapsulates all the protocol-level bookkeeping: +/// - Emitting [`MessageStart`] once. +/// - Emitting [`PartStart`] / [`IndexedDelta`] for text and tool-call content. +/// - Emitting [`PartStop`] when parts finish. +/// - Emitting the final [`MessageDelta`] with a stop reason. +/// +/// Splitting this out from the `try_stream!` macro body makes the +/// translation logic testable without a live network connection. +#[derive(Default)] +struct StreamEmitter { + started: bool, + text_part_open: bool, + open_tool_count: usize, + finished: bool, + pending: Vec, +} + +impl StreamEmitter { + /// Process a single parsed chunk, appending events to the + /// internal pending queue. + fn process_chunk(&mut self, chunk: &OpenAiChunk) { + if !self.started { + self.started = true; + self.push(StreamEvent::MessageStart(MessageStart { + message: MessageMetadata { + id: chunk.id.clone(), + role: "assistant".into(), + model: chunk.model.clone(), + }, + })); + } + + let Some(choice) = chunk.choices.first() else { + return; + }; + + if let Some(delta) = &choice.delta { + self.process_delta(delta); + } + + if let Some(reason) = &choice.finish_reason { + self.process_finish(reason); + } + } + + /// Translate a delta into text/tool-call events. + fn process_delta(&mut self, delta: &OpenAiDelta) { + if let Some(text) = &delta.content + && !text.is_empty() + { + if !self.text_part_open { + self.text_part_open = true; + self.push(StreamEvent::PartStart(PartStart { + index: TEXT_PART_INDEX, + part: Some(MessagePart::text("")), + })); + } + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: TEXT_PART_INDEX, + delta: DeltaPart::Text { text: text.clone() }, + })); + } + + if let Some(tool_calls) = &delta.tool_calls { + for tc in tool_calls { + self.process_tool_call(tc); + } + } + } + + /// Handle a single tool-call delta. + fn process_tool_call(&mut self, tc: &OpenAiToolCallDelta) { + if tc.function.is_some() { + // New tool call — emit PartStart. + self.push(StreamEvent::PartStart(PartStart { + index: tc.index, + part: Some(MessagePart::ToolCall { + id: tc.id.clone(), + name: tc + .function + .as_ref() + .map(|f| f.name.clone()) + .unwrap_or_default(), + input: Value::Null, + }), + })); + self.open_tool_count = self.open_tool_count.saturating_add(1); + } + + // Stream argument fragments. + if let Some(func) = &tc.function + && !func.arguments.is_empty() + { + self.push(StreamEvent::IndexedDelta(IndexedDelta { + index: tc.index, + delta: DeltaPart::InputJson { + partial_json: func.arguments.clone(), + }, + })); + } + } + + /// Handle a finish reason, emitting the appropriate stop events. + fn process_finish(&mut self, reason: &str) { + if self.finished { + return; + } + self.finished = true; + + // Close any open text part. + if self.text_part_open { + self.push(StreamEvent::PartStop); + } + + // Close each open tool-call part. + for _ in 0..self.open_tool_count { + self.push(StreamEvent::PartStop); + } + + let stop_reason = match reason { + "tool_calls" => StreamStopReason::ToolCall, + "length" => StreamStopReason::MaxTokens, + other => StreamStopReason::from_api_str(other).unwrap_or(StreamStopReason::EndTurn), + }; + + self.push(StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some(stop_reason.to_api_str().into()), + }, + usage: None, + })); + } + + /// Emit the terminal [`MessageStop`] if the stream was started, + /// returning all remaining events. + fn finish(&mut self) -> Vec { + let mut out = self.drain(); + if self.started { + out.push(StreamEvent::MessageStop); + } + out + } + + /// Drain all pending events. + fn drain(&mut self) -> Vec { + std::mem::take(&mut self.pending) + } + + fn push(&mut self, ev: StreamEvent) { + self.pending.push(ev); + } +} + +// ================================================== +// Tests +// ================================================== + +#[cfg(test)] +mod tests { + use super::*; + use crate::message::{Message, MessagePart, Role, ToolContent}; + use crate::tool::ToolSchema; + + #[test] + fn request_body_includes_system_message_first() { + let msgs = vec![Message::user("hello")]; + let body = RequestBody::build("gpt-4o", &msgs, Some("be brief"), None); + let json = body.to_json(true); + + let messages = json["messages"].as_array().unwrap(); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0]["role"], "system"); + assert_eq!(messages[0]["content"], "be brief"); + assert_eq!(messages[1]["role"], "user"); + } + + #[test] + fn request_body_without_system() { + let msgs = vec![Message::user("hi")]; + let body = RequestBody::build("gpt-4o", &msgs, None, None); + let json = body.to_json(false); + + let messages = json["messages"].as_array().unwrap(); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0]["role"], "user"); + } + + #[test] + fn request_body_stream_flag_toggles() { + let msgs = vec![Message::user("hi")]; + let body = RequestBody::build("gpt-4o", &msgs, None, None); + + assert_eq!(body.to_json(true)["stream"], true); + assert_eq!(body.to_json(false)["stream"], false); + } + + #[test] + fn request_body_model_and_tools() { + let msgs = vec![Message::user("hi")]; + let tools = vec![ToolSchema { + tool: "echo".into(), + description: "Echo".into(), + input_schema: serde_json::json!({"type": "object"}), + }]; + let body = RequestBody::build("my-model", &msgs, None, Some(&tools)); + let json = body.to_json(true); + + assert_eq!(json["model"], "my-model"); + let tools_arr = json["tools"].as_array().unwrap(); + assert_eq!(tools_arr.len(), 1); + assert_eq!(tools_arr[0]["type"], "function"); + assert_eq!(tools_arr[0]["function"]["name"], "echo"); + } + + #[test] + fn request_body_tools_null_when_none() { + let msgs = vec![Message::user("hi")]; + let body = RequestBody::build("gpt-4o", &msgs, None, None); + let json = body.to_json(false); + assert!(json["tools"].is_null()); + } + + #[test] + fn convert_message_user_text() { + let m = Message::user("hello world"); + let v = convert_message(&m); + assert_eq!(v["role"], "user"); + assert_eq!(v["content"], "hello world"); + } + + #[test] + fn convert_message_assistant_text() { + let m = Message::new(Role::Assistant, vec![MessagePart::text("hi there")]); + let v = convert_message(&m); + assert_eq!(v["role"], "assistant"); + assert_eq!(v["content"], "hi there"); + } + + #[test] + fn convert_message_assistant_tool_calls() { + let m = Message::new( + Role::Assistant, + vec![MessagePart::ToolCall { + id: "call_1".into(), + name: "echo".into(), + input: serde_json::json!({"message": "hi"}), + }], + ); + let v = convert_message(&m); + assert_eq!(v["role"], "assistant"); + assert!(v["content"].is_null()); + let calls = v["tool_calls"].as_array().unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0]["id"], "call_1"); + assert_eq!(calls[0]["type"], "function"); + assert_eq!(calls[0]["function"]["name"], "echo"); + // arguments should be stringified JSON + assert_eq!( + calls[0]["function"]["arguments"].as_str().unwrap(), + r#"{"message":"hi"}"# + ); + } + + #[test] + fn convert_message_tool_result() { + let m = Message::new( + Role::User, + vec![MessagePart::ToolResult { + call_id: "call_1".into(), + output: ToolContent::from_string("result text"), + is_error: None, + }], + ); + let v = convert_message(&m); + assert_eq!(v["role"], "tool"); + assert_eq!(v["tool_call_id"], "call_1"); + assert!(v["content"].is_string()); + } + + #[test] + fn convert_tools_shape() { + let tools = vec![ + ToolSchema { + tool: "search".into(), + description: "Search the web".into(), + input_schema: serde_json::json!({"type": "object"}), + }, + ToolSchema { + tool: "calc".into(), + description: "Calculate".into(), + input_schema: serde_json::json!({"type": "object"}), + }, + ]; + let out = convert_tools(&tools); + assert_eq!(out.len(), 2); + assert_eq!(out[0]["function"]["name"], "search"); + assert_eq!(out[1]["function"]["name"], "calc"); + } + + #[test] + fn input_to_string_passes_through_strings() { + assert_eq!(input_to_string(&Value::String("raw".into())), "raw"); + } + + #[test] + fn input_to_string_serializes_objects() { + let v = serde_json::json!({"a": 1}); + let s = input_to_string(&v); + assert_eq!(s, r#"{"a":1}"#); + } + + #[test] + fn input_to_string_serializes_numbers() { + let s = input_to_string(&Value::from(42)); + assert_eq!(s, "42"); + } + + #[test] + fn parse_valid_chunk() { + let data = r#"{"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"content":"hi"},"finish_reason":null}]}"#; + let chunk = OpenAiChunk::parse(data).unwrap(); + assert_eq!(chunk.id, "chatcmpl-1"); + assert_eq!(chunk.model, "gpt-4o"); + assert_eq!(chunk.choices.len(), 1); + } + + #[test] + fn parse_malformed_returns_none() { + assert!(OpenAiChunk::parse("not json").is_none()); + assert!(OpenAiChunk::parse("").is_none()); + } + + #[test] + fn emitter_emits_message_start_on_first_chunk() { + let mut em = StreamEmitter::default(); + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":""},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + let events = em.drain(); + assert!( + events + .iter() + .any(|e| matches!(e, StreamEvent::MessageStart(_))) + ); + } + + #[test] + fn emitter_text_delta_starts_part_then_deltas() { + let mut em = StreamEmitter::default(); + + // First chunk with text — should emit MessageStart + PartStart + IndexedDelta. + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":"Hel"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + + let events = em.drain(); + // MessageStart, PartStart(0), IndexedDelta(Text) + assert_eq!(events.len(), 3); + assert!(matches!( + events[1], + StreamEvent::PartStart(ref p) if p.index == TEXT_PART_INDEX + )); + assert!(matches!( + events[2], + StreamEvent::IndexedDelta(ref d) if d.index == TEXT_PART_INDEX + )); + + // Second text chunk — should only emit a delta (part already open). + let chunk2 = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":"lo"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk2); + let events2 = em.drain(); + assert_eq!(events2.len(), 1); + assert!(matches!(events2[0], StreamEvent::IndexedDelta(_))); + } + + #[test] + fn emitter_tool_call_emits_part_start_and_delta() { + let mut em = StreamEmitter::default(); + + // Start message first. + let chunk0 = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":null},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk0); + em.drain(); + + // Tool call delta. + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":1,"id":"call_1","function":{"name":"echo","arguments":"{\"msg\":"}}]},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + let events = em.drain(); + + // PartStart(tool) + IndexedDelta(InputJson) + assert_eq!(events.len(), 2); + assert!(matches!( + events[0], + StreamEvent::PartStart(ref p) if p.index == 1 + )); + assert!(matches!( + events[1], + StreamEvent::IndexedDelta(ref d) if d.index == 1 + )); + } + + #[test] + fn emitter_finish_emits_part_stops_and_message_delta() { + let mut em = StreamEmitter::default(); + + // Send some text so the text part is open. + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":"hi"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + em.drain(); + + // Now send finish_reason=stop. + let finish = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"stop"}]}"#, + ) + .unwrap(); + em.process_chunk(&finish); + let events = em.drain(); + + // PartStop + MessageDelta(stop_reason) + assert_eq!(events.len(), 2); + assert!(matches!(events[0], StreamEvent::PartStop)); + assert!(matches!(events[1], StreamEvent::MessageDelta(_))); + } + + #[test] + fn emitter_finish_with_tool_calls_stop_reason() { + let mut em = StreamEmitter::default(); + + let chunk0 = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk0); + em.drain(); + + // Open a tool call. + let tool_chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"echo","arguments":""}}]},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&tool_chunk); + em.drain(); + + // Finish with tool_calls. + let finish = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"tool_calls"}]}"#, + ) + .unwrap(); + em.process_chunk(&finish); + let events = em.drain(); + + // 1 PartStop (for tool) + MessageDelta + assert_eq!(events.len(), 2); + assert!(matches!(events[0], StreamEvent::PartStop)); + + if let StreamEvent::MessageDelta(md) = &events[1] { + assert_eq!(md.delta.stop_reason.as_deref(), Some("tool_call")); + } else { + panic!("expected MessageDelta"); + } + } + + #[test] + fn emitter_finish_appends_message_stop() { + let mut em = StreamEmitter::default(); + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":"hi"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + em.drain(); + + let finish = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"stop"}]}"#, + ) + .unwrap(); + em.process_chunk(&finish); + em.drain(); + + let final_events = em.finish(); + assert!(matches!( + final_events.last(), + Some(StreamEvent::MessageStop) + )); + } + + #[test] + fn emitter_finish_without_start_is_empty() { + let mut em = StreamEmitter::default(); + let events = em.finish(); + assert!(events.is_empty()); + } + + #[test] + fn emitter_finish_reason_length_maps_to_max_tokens() { + let mut em = StreamEmitter::default(); + let chunk0 = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":"x"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk0); + em.drain(); + + let finish = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"length"}]}"#, + ) + .unwrap(); + em.process_chunk(&finish); + let events = em.drain(); + + if let StreamEvent::MessageDelta(md) = &events[1] { + assert_eq!(md.delta.stop_reason.as_deref(), Some("max_tokens")); + } else { + panic!("expected MessageDelta"); + } + } + + #[test] + fn emitter_empty_content_does_not_open_text_part() { + let mut em = StreamEmitter::default(); + let chunk = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":""},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk); + let events = em.drain(); + + // Only MessageStart, no PartStart for empty content. + assert_eq!(events.len(), 1); + assert!(matches!(events[0], StreamEvent::MessageStart(_))); + } + + #[test] + fn emitter_double_finish_ignored() { + let mut em = StreamEmitter::default(); + let chunk0 = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":"hi"},"finish_reason":null}]}"#, + ) + .unwrap(); + em.process_chunk(&chunk0); + em.drain(); + + let finish1 = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"stop"}]}"#, + ) + .unwrap(); + em.process_chunk(&finish1); + em.drain(); + + // Second finish should not emit anything extra. + let finish2 = OpenAiChunk::parse( + r#"{"id":"c1","model":"gpt-4o","choices":[{"delta":null,"finish_reason":"stop"}]}"#, + ) + .unwrap(); + em.process_chunk(&finish2); + let events = em.drain(); + assert!(events.is_empty()); + } + + #[test] + fn merge_tool_results_single() { + let r = serde_json::json!({"role": "tool", "content": "ok"}); + let merged = merge_tool_results(&[r.clone()]); + assert_eq!(merged, r); + } + + #[test] + fn merge_tool_results_multiple() { + let results = vec![ + serde_json::json!({"role": "tool", "content": "a"}), + serde_json::json!({"role": "tool", "content": "b"}), + ]; + let merged = merge_tool_results(&results); + assert!(merged.is_array()); + assert_eq!(merged.as_array().unwrap().len(), 2); + } + + #[test] + fn sse_reader_take_line_extracts_newline_terminated() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "data: hello\n".into(), + }; + let line = reader.take_line().unwrap(); + assert_eq!(line, "data: hello"); + assert!(reader.buf.is_empty()); + } + + #[test] + fn sse_reader_take_line_returns_none_without_newline() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "partial".into(), + }; + assert!(reader.take_line().is_none()); + } + + #[test] + fn sse_reader_take_line_handles_multiple_lines() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "line1\nline2\n".into(), + }; + assert_eq!(reader.take_line().unwrap(), "line1"); + assert_eq!(reader.take_line().unwrap(), "line2"); + } + + #[test] + fn sse_reader_take_line_trims_cr() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "data: hi\r\n".into(), + }; + let line = reader.take_line().unwrap(); + assert_eq!(line, "data: hi"); + } +} diff --git a/src/reflection/backoff.rs b/src/reflection/backoff.rs index d1953fb..89dcbe7 100644 --- a/src/reflection/backoff.rs +++ b/src/reflection/backoff.rs @@ -327,8 +327,6 @@ mod tests { assert!(debug.contains("max_retries")); } - // ---- delay_for_attempt ---- - #[test] fn delay_for_attempt_zero_yields_base() { let strategy = ExponentialBackoffRecovery::new(3); diff --git a/src/stream.rs b/src/stream.rs index a52c341..5333422 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -827,27 +827,6 @@ impl Usage { /// the final message. Token usage can be retrieved at any time via /// [`usage`](Self::usage). /// -/// # State Machine -/// -/// ```text -/// MessageStart -/// ┌───────┐ ────────────▶ ┌─────────────┐ -/// │ Idle │ │ Started │ -/// └───────┘ └──────┬──────┘ -/// PartStart │ PartEnd -/// ┌─────▼──────┐ -/// │ Receiving │ -/// │ Part │ -/// └─────┬──────┘ -/// IndexedDelta │ -/// (accumulate) -/// │ -/// MessageDelta │ -/// ┌──────────▼──────────┐ -/// │ Complete │──▶ build() -/// └─────────────────────┘ -/// ``` -/// /// # Example /// /// ```rust @@ -1026,13 +1005,12 @@ impl StreamAccumulator { Ok(()) } StreamEvent::IndexedDelta(delta) => { - debug_assert_eq!( - self.current_index, - Some(delta.index), - "IndexedDelta index mismatch: expected {:?}, got {}", - self.current_index, - delta.index, - ); + // Gracefully ignore deltas whose index doesn't match the + // current part — this can happen with malformed SSE or + // provider-specific quirks. + if self.current_index != Some(delta.index) { + return Ok(()); + } match &delta.delta { DeltaPart::Text { text } => { self.current_text.push_str(text); @@ -1155,11 +1133,6 @@ impl StreamAccumulator { mod tests { use super::*; - /// Verify that [`StreamStopReason::from_api_str`] parses all known API - /// strings and returns `None` for unknown values. - /// - /// Tests `"tool_call"`, `"max_tokens"`, `"end_turn"`, `"stop_sequence"`, - /// and `"unknown"`. #[test] fn test_stream_stop_reason_from_api_str() { assert_eq!( @@ -1181,11 +1154,6 @@ mod tests { assert_eq!(StreamStopReason::from_api_str("unknown"), None); } - /// Verify that [`StreamStopReason::to_api_str`] produces the correct - /// static string for each variant. - /// - /// Ensures the round-trip `from_api_str(s) == Some(v)` implies - /// `v.to_api_str() == s`. #[test] fn test_stream_stop_reason_to_api_str() { assert_eq!(StreamStopReason::ToolCall.to_api_str(), "tool_call"); @@ -1194,11 +1162,6 @@ mod tests { assert_eq!(StreamStopReason::StopSequence.to_api_str(), "stop_sequence"); } - /// Verify that [`StreamStopReason::should_continue_tool_loop`] returns `true` - /// only for [`ToolCall`](StreamStopReason::ToolCall). - /// - /// [`EndTurn`](StreamStopReason::EndTurn) and [`MaxTokens`](StreamStopReason::MaxTokens) - /// should both return `false`. #[test] fn test_stream_stop_reason_should_continue() { assert!(StreamStopReason::ToolCall.should_continue_tool_loop()); @@ -1206,10 +1169,6 @@ mod tests { assert!(!StreamStopReason::MaxTokens.should_continue_tool_loop()); } - /// Verify [`Usage::new`] stores token counts and computes the total. - /// - /// Asserts that [`Usage::input_tokens`] and [`Usage::output_tokens`] are - /// stored as provided, and that [`Usage::total_tokens`] returns their sum. #[test] fn test_usage() { let usage = Usage::new(100, 50); @@ -1218,27 +1177,12 @@ mod tests { assert_eq!(usage.total_tokens(), 150); } - /// Verify that [`Usage::default`] produces zeroed counters. - /// - /// [`Usage::total_tokens`] should be `0` when both fields are `0`. #[test] fn test_usage_default() { let usage = Usage::default(); assert_eq!(usage.total_tokens(), 0); } - /// Verify that [`StreamAccumulator`] correctly assembles a text-only - /// streaming response into a [`Message`]. - /// - /// Feeds a full event sequence: [`MessageStart`](StreamEvent::MessageStart) → - /// [`PartStart`](StreamEvent::PartStart) → two - /// [`IndexedDelta`](StreamEvent::IndexedDelta)s → - /// [`PartStop`](StreamEvent::PartStop) → - /// [`MessageDelta`](StreamEvent::MessageDelta) → - /// [`MessageStop`](StreamEvent::MessageStop). - /// - /// Asserts the final message has [`Role::Assistant`](crate::message::Role::Assistant), - /// one part with the concatenated text "Hello world", and correct usage. #[test] fn test_accumulator_text_message() { let mut acc = StreamAccumulator::new(); @@ -1285,13 +1229,6 @@ mod tests { assert_eq!(msg.parts[0].as_text(), Some("Hello world")); } - /// Verify that [`StreamAccumulator`] correctly assembles a tool-call - /// part from streaming events. - /// - /// Feeds [`PartStart`](StreamEvent::PartStart) with a - /// [`ToolCall`](MessagePart::ToolCall) seed, then an [`InputJson`](DeltaPart::InputJson) - /// delta, then [`PartStop`](StreamEvent::PartStop). Asserts the - /// resulting message contains a tool-call part. #[test] fn test_accumulator_tool_call() { let mut acc = StreamAccumulator::new(); @@ -1318,10 +1255,6 @@ mod tests { assert!(msg.parts[0].is_tool_call()); } - /// Verify that [`StreamAccumulator::build`] on a fresh accumulator - /// produces a [`Message`] with an empty content vector. - /// - /// No events processed means no parts assembled. #[test] fn test_accumulator_empty() { let acc = StreamAccumulator::new(); @@ -1329,11 +1262,6 @@ mod tests { assert_eq!(msg.parts.len(), 0); } - /// Verify that [`StreamAccumulator::usage`] returns the [`Usage`] data - /// from a [`MessageDelta`](StreamEvent::MessageDelta) event. - /// - /// Feeds a single `MessageDelta` with known token counts and asserts - /// [`Usage::total_tokens`] returns the expected sum. #[test] fn test_accumulator_usage() { let mut acc = StreamAccumulator::new(); @@ -1345,11 +1273,6 @@ mod tests { assert_eq!(acc.usage().unwrap().total_tokens(), 150); } - /// Verify that all [`StreamEvent`] variants can be constructed. - /// - /// Constructs [`Ping`](StreamEvent::Ping), [`MessageStop`](StreamEvent::MessageStop), - /// and [`PartStop`](StreamEvent::PartStop) to confirm no compile-time - /// regressions in the enum definition. #[test] fn test_stream_event_variants() { // Just verify all variants can be constructed @@ -1358,9 +1281,6 @@ mod tests { let _ = StreamEvent::PartStop; } - /// Verify that [`StreamAccumulator::process`] returns - /// [`StreamError::InvalidToolInputJson`] when the accumulated tool-call - /// JSON is malformed at [`PartStop`](StreamEvent::PartStop). #[test] fn test_accumulator_invalid_tool_json() { let mut acc = StreamAccumulator::new(); @@ -1391,8 +1311,6 @@ mod tests { } } - /// Verify that a tool-call part with no delta input defaults to an - /// empty JSON object `{}` rather than erroring. #[test] fn test_accumulator_tool_call_empty_input() { let mut acc = StreamAccumulator::new(); @@ -1411,4 +1329,215 @@ mod tests { assert_eq!(msg.parts.len(), 1); assert!(msg.parts[0].is_tool_call()); } + + #[test] + fn test_accumulator_ignores_delta_with_mismatched_index() { + let mut acc = StreamAccumulator::new(); + acc.process(&StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::text("")), + })) + .unwrap(); + + // Delta arrives with index 1 — mismatch! Must NOT panic, must be ignored. + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 1, + delta: DeltaPart::Text { + text: "ignored".into(), + }, + })) + .unwrap(); + + // Delta with correct index 0 — should be applied. + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::Text { + text: "hello".into(), + }, + })) + .unwrap(); + + acc.process(&StreamEvent::PartStop).unwrap(); + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 1); + assert_eq!(msg.parts[0].as_text(), Some("hello")); + } + + #[test] + fn test_accumulator_ignores_input_json_with_mismatched_index() { + let mut acc = StreamAccumulator::new(); + acc.process(&StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::text("")), + })) + .unwrap(); + + // InputJson delta at wrong index — should be ignored. + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 5, + delta: DeltaPart::InputJson { + partial_json: "{\"bad\":true}".into(), + }, + })) + .unwrap(); + + acc.process(&StreamEvent::PartStop).unwrap(); + + let msg = acc.build(); + // Empty text → no parts. + assert!(msg.parts.is_empty()); + } + + #[test] + fn test_accumulator_delta_tool_call_string_value() { + let mut acc = StreamAccumulator::new(); + acc.process(&StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::tool_call("id1", "search", Value::Null)), + })) + .unwrap(); + + // DeltaPart::ToolCall carries a string JSON value. + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::ToolCall { + partial_json: Value::String("{\"q\":\"rust\"}".into()), + }, + })) + .unwrap(); + + acc.process(&StreamEvent::PartStop).unwrap(); + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 1); + if let MessagePart::ToolCall { input, .. } = &msg.parts[0] { + assert_eq!(input["q"], "rust"); + } else { + panic!("expected ToolCall"); + } + } + + #[test] + fn test_accumulator_delta_tool_call_non_string_ignored() { + let mut acc = StreamAccumulator::new(); + acc.process(&StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::tool_call("id1", "search", Value::Null)), + })) + .unwrap(); + + // Non-string JSON value — should be silently ignored. + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::ToolCall { + partial_json: Value::Number(42.into()), + }, + })) + .unwrap(); + + acc.process(&StreamEvent::PartStop).unwrap(); + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 1); + if let MessagePart::ToolCall { input, .. } = &msg.parts[0] { + assert!(input.is_object()); + } + } + + #[test] + fn test_accumulator_ping_no_op() { + let mut acc = StreamAccumulator::new(); + acc.process(&StreamEvent::Ping).unwrap(); + assert!(acc.usage().is_none()); + let msg = acc.build(); + assert!(msg.parts.is_empty()); + } + + #[test] + fn test_accumulator_message_stop_no_op() { + let mut acc = StreamAccumulator::new(); + acc.process(&StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::text("")), + })) + .unwrap(); + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::Text { text: "hi".into() }, + })) + .unwrap(); + acc.process(&StreamEvent::PartStop).unwrap(); + acc.process(&StreamEvent::MessageStop).unwrap(); + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 1); + assert_eq!(msg.parts[0].as_text(), Some("hi")); + } + + #[test] + fn test_accumulator_multiple_text_parts() { + let mut acc = StreamAccumulator::new(); + + // First text part at index 0. + acc.process(&StreamEvent::PartStart(PartStart { + index: 0, + part: Some(MessagePart::text("")), + })) + .unwrap(); + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 0, + delta: DeltaPart::Text { + text: "hello".into(), + }, + })) + .unwrap(); + acc.process(&StreamEvent::PartStop).unwrap(); + + // Second text part at index 1. + acc.process(&StreamEvent::PartStart(PartStart { + index: 1, + part: Some(MessagePart::text("")), + })) + .unwrap(); + acc.process(&StreamEvent::IndexedDelta(IndexedDelta { + index: 1, + delta: DeltaPart::Text { + text: "world".into(), + }, + })) + .unwrap(); + acc.process(&StreamEvent::PartStop).unwrap(); + + let msg = acc.build(); + assert_eq!(msg.parts.len(), 2); + assert_eq!(msg.parts[0].as_text(), Some("hello")); + assert_eq!(msg.parts[1].as_text(), Some("world")); + } + + #[test] + fn test_accumulator_message_delta_overwrites_usage() { + let mut acc = StreamAccumulator::new(); + + acc.process(&StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some("end_turn".into()), + }, + usage: Some(Usage::new(100, 50)), + })) + .unwrap(); + assert_eq!(acc.usage().unwrap().input_tokens, 100); + assert_eq!(acc.usage().unwrap().output_tokens, 50); + + // Second delta overwrites usage. + acc.process(&StreamEvent::MessageDelta(MessageDelta { + delta: MessageDeltaPayload { + stop_reason: Some("max_tokens".into()), + }, + usage: Some(Usage::new(200, 75)), + })) + .unwrap(); + assert_eq!(acc.usage().unwrap().input_tokens, 200); + assert_eq!(acc.usage().unwrap().output_tokens, 75); + } } diff --git a/src/testing.rs b/src/testing.rs index cef3c1c..726be22 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -228,7 +228,7 @@ pub struct MockApiClient { /// stop_reason: "end_turn".to_string(), /// }; /// ``` -#[derive(Clone)] +#[derive(Clone, Default)] pub struct MockResponse { /// The text content the assistant should produce. /// @@ -393,13 +393,15 @@ impl MockApiClient { /// .with_text_response("I am a test assistant."); /// ``` #[must_use] - #[allow( - clippy::unwrap_used, - clippy::indexing_slicing, - clippy::missing_panics_doc - )] pub fn with_text_response(self, text: &str) -> Self { - self.responses.lock().unwrap()[0].text = text.to_string(); + if let Some(r) = self + .responses + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .first_mut() + { + r.text = text.to_string(); + } self } @@ -424,19 +426,19 @@ impl MockApiClient { /// .with_tool_call("call_1", "bash", json!({"command": "ls"})); /// ``` #[must_use] - #[allow( - clippy::unwrap_used, - clippy::indexing_slicing, - clippy::missing_panics_doc - )] pub fn with_tool_call(self, id: &str, name: &str, input: Value) -> Self { - let mut responses = self.responses.lock().unwrap(); - responses[0].tool_call = Some(MockToolCall { - id: id.to_string(), - name: name.to_string(), - input, - }); - responses[0].stop_reason = "tool_use".to_string(); + let mut responses = self + .responses + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(r) = responses.first_mut() { + r.tool_call = Some(MockToolCall { + id: id.to_string(), + name: name.to_string(), + input, + }); + r.stop_reason = "tool_use".to_string(); + } drop(responses); self } @@ -458,13 +460,15 @@ impl MockApiClient { /// .with_stop_reason("max_tokens"); /// ``` #[must_use] - #[allow( - clippy::unwrap_used, - clippy::indexing_slicing, - clippy::missing_panics_doc - )] pub fn with_stop_reason(self, reason: &str) -> Self { - self.responses.lock().unwrap()[0].stop_reason = reason.to_string(); + if let Some(r) = self + .responses + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .first_mut() + { + r.stop_reason = reason.to_string(); + } self } @@ -498,10 +502,12 @@ impl MockApiClient { /// ]); /// ``` #[must_use] - #[allow(clippy::unwrap_used, clippy::missing_panics_doc)] pub fn with_responses(self, responses: Vec) -> Self { if !responses.is_empty() { - *self.responses.lock().unwrap() = responses; + *self + .responses + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = responses; } self } @@ -551,17 +557,15 @@ impl MockApiClient { /// [R2, R3] → pop → R2, queue becomes [R3] /// [R3] → pop → R3, queue stays [R3] (cloned) /// ``` - #[allow( - clippy::unwrap_used, - clippy::indexing_slicing, - clippy::missing_panics_doc - )] fn pop_response(&self) -> MockResponse { - let mut guard = self.responses.lock().unwrap(); + let mut guard = self + .responses + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); if guard.len() > 1 { guard.remove(0) } else { - guard[0].clone() + guard.first().cloned().unwrap_or_default() } } } From ea75f8d0b8814915a8ba1ad4e21a15c2c83c3c2a Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Mon, 29 Jun 2026 20:39:24 +1200 Subject: [PATCH 08/30] chore: add chat example --- examples/chat.rs | 518 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 518 insertions(+) create mode 100644 examples/chat.rs diff --git a/examples/chat.rs b/examples/chat.rs new file mode 100644 index 0000000..e228917 --- /dev/null +++ b/examples/chat.rs @@ -0,0 +1,518 @@ +#![allow(clippy::all, clippy::pedantic, clippy::restriction)] +//! Real provider chat CLI — talk to OpenAI, Anthropic, `DeepSeek`, `Grok`, Gemini, or Ollama. +//! +//! Set environment variables to pick the provider: +//! +//! ```sh +//! # OpenAI +//! OPENAI_API_KEY=sk-... cargo run --example chat --features openai +//! +//! # DeepSeek +//! DEEPSEEK_API_KEY=... cargo run --example chat --features deepseek +//! +//! # Grok (xAI) +//! XAI_API_KEY=... cargo run --example chat --features grok +//! +//! # Gemini (Google) +//! GEMINI_API_KEY=... cargo run --example chat --features gemini +//! +//! # Anthropic +//! ANTHROPIC_API_KEY=... cargo run --example chat --features anthropic +//! +//! # Ollama (local, no API key needed) +//! OLLAMA_MODEL=llama3 cargo run --example chat --features ollama +//! ``` + +use std::io::{self, BufRead, Write}; +use std::sync::Arc; + +use loopctl::api::ApiClient; +use loopctl::config::LoopConfig; +use loopctl::engine::BareLoop; +use loopctl::engine::loop_core::Loop; +use loopctl::observer::{LoopObserver, ToolPostContext, ToolPreContext}; +use loopctl::tool::{FnTool, ToolContext, ToolOutput, ToolRegistry}; +use serde_json::json; + +// ================================================== +// Type alias for tool function signatures +// ================================================== + +/// Shorthand for the boxed-future signature required by [`FnTool`]. +type ToolFuture = std::pin::Pin< + Box< + dyn std::future::Future> + + Send + + 'static, + >, +>; + +// ================================================== +// Observer +// ================================================== + +/// A simple observer that prints tool calls and responses to stderr. +struct PrintingObserver; + +impl LoopObserver for PrintingObserver { + fn name(&self) -> &str { + "chat-printer" + } + + fn on_tool_pre(&self, ctx: &ToolPreContext) { + eprintln!(" → calling tool: {}", ctx.tool); + } + + fn on_tool_post(&self, ctx: &ToolPostContext) { + let status = if ctx.is_error { "failed" } else { "completed" }; + eprintln!( + " ← tool {} {status} ({:.0}ms)", + ctx.tool, + ctx.duration.as_secs_f64() * 1000.0 + ); + } +} + +// ================================================== +// Usage +// ================================================== + +fn print_usage_and_exit() -> ! { + eprintln!("No provider configured.\n"); + eprintln!("Set one of:"); + eprintln!(" OPENAI_API_KEY= — OpenAI"); + eprintln!(" DEEPSEEK_API_KEY= — DeepSeek"); + eprintln!(" XAI_API_KEY= — Grok (xAI)"); + eprintln!(" GEMINI_API_KEY= — Google Gemini"); + eprintln!(" ZAI_API_KEY= — Z.ai (ZhipuAI)"); + eprintln!(" ANTHROPIC_API_KEY= — Anthropic Claude"); + eprintln!(" OLLAMA_MODEL= — Local Ollama\n"); + eprintln!(" SELF_HOSTED_BASE_URL= — Any OpenAI-compatible server"); + eprintln!(" SELF_HOSTED_MODEL= — (required with SELF_HOSTED_BASE_URL)\n"); + eprintln!( + "Build with: --features openai | deepseek | grok | zai | gemini | anthropic | ollama\n" + ); + std::process::exit(1); +} + +// ================================================== +// Tool functions +// ================================================== + +fn echo_fn(input: serde_json::Value, _ctx: &ToolContext) -> ToolFuture { + let text = input + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("(empty)") + .to_string(); + Box::pin(async move { Ok(ToolOutput::text(format!("echo: {text}"))) }) +} + +fn current_time_fn(_input: serde_json::Value, _ctx: &ToolContext) -> ToolFuture { + Box::pin(async move { + let secs = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + Ok(ToolOutput::text(format!("Unix timestamp: {secs}"))) + }) +} + +fn calculate_fn(input: serde_json::Value, _ctx: &ToolContext) -> ToolFuture { + let expr = input + .get("expression") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let result = simple_eval(&expr); + Box::pin(async move { Ok(ToolOutput::text(result)) }) +} + +/// Build the tool registry for the chat example. +fn build_tools() -> ToolRegistry { + let mut tools = ToolRegistry::new(); + + tools.register( + FnTool::new( + "echo".into(), + "Echo back the provided message.".into(), + json!({ + "type": "object", + "properties": { + "message": {"type": "string", "description": "The text to echo back"} + }, + "required": ["message"] + }), + echo_fn, + ) + .read_only(), + ); + + tools.register( + FnTool::new( + "current_time".into(), + "Get the current wall-clock time as a Unix timestamp.".into(), + json!({"type": "object", "properties": {}}), + current_time_fn, + ) + .read_only(), + ); + + tools.register( + FnTool::new( + "calculate".into(), + "Evaluate a simple arithmetic expression (e.g. \"2 + 3 * 4\"). Supports +, -, *, /, parentheses.".into(), + json!({ + "type": "object", + "properties": { + "expression": {"type": "string", "description": "The arithmetic expression to evaluate"} + }, + "required": ["expression"] + }), + calculate_fn, + ) + .read_only(), + ); + + tools +} + +// ================================================== +// Minimal arithmetic expression evaluator +// (recursive descent: + - * / and parentheses) +// ================================================== + +#[derive(Debug, Clone)] +enum Token { + Num(f64), + Plus, + Minus, + Star, + Slash, + LParen, + RParen, +} + +fn tokenize(s: &str) -> Vec { + let mut tokens = Vec::new(); + let mut chars = s.chars().peekable(); + while let Some(&c) = chars.peek() { + match c { + ' ' | '\t' => { + chars.next(); + } + '+' => { + tokens.push(Token::Plus); + chars.next(); + } + '-' => { + tokens.push(Token::Minus); + chars.next(); + } + '*' => { + tokens.push(Token::Star); + chars.next(); + } + '/' => { + tokens.push(Token::Slash); + chars.next(); + } + '(' => { + tokens.push(Token::LParen); + chars.next(); + } + ')' => { + tokens.push(Token::RParen); + chars.next(); + } + '0'..='9' | '.' => { + let mut num = String::new(); + while let Some(&d) = chars.peek() { + if d.is_ascii_digit() || d == '.' { + num.push(d); + chars.next(); + } else { + break; + } + } + if let Ok(n) = num.parse::() { + tokens.push(Token::Num(n)); + } + } + _ => { + chars.next(); + } + } + } + tokens +} + +fn simple_eval(expr: &str) -> String { + let tokens = tokenize(expr); + let mut pos = 0usize; + match parse_expr(&tokens, &mut pos) { + Ok(val) => format!("{val}"), + Err(e) => format!("Error: {e}"), + } +} + +/// Helper: peek the token at `pos`, safely. +fn peek(tokens: &[Token], pos: usize) -> Option<&Token> { + tokens.get(pos) +} + +/// Helper: advance `pos` by one. +fn advance(pos: &mut usize) { + *pos += 1; +} + +fn parse_expr(tokens: &[Token], pos: &mut usize) -> Result { + let mut left = parse_term(tokens, pos)?; + while let Some(tok) = peek(tokens, *pos) { + match tok { + Token::Plus => { + advance(pos); + left += parse_term(tokens, pos)?; + } + Token::Minus => { + advance(pos); + left -= parse_term(tokens, pos)?; + } + _ => break, + } + } + Ok(left) +} + +fn parse_term(tokens: &[Token], pos: &mut usize) -> Result { + let mut left = parse_factor(tokens, pos)?; + while let Some(tok) = peek(tokens, *pos) { + match tok { + Token::Star => { + advance(pos); + left *= parse_factor(tokens, pos)?; + } + Token::Slash => { + advance(pos); + left /= parse_factor(tokens, pos)?; + } + _ => break, + } + } + Ok(left) +} + +fn parse_factor(tokens: &[Token], pos: &mut usize) -> Result { + let Some(tok) = peek(tokens, *pos) else { + return Err("unexpected end of expression".into()); + }; + match tok { + Token::Num(n) => { + let val = *n; + advance(pos); + Ok(val) + } + Token::LParen => { + advance(pos); + let val = parse_expr(tokens, pos)?; + if matches!(peek(tokens, *pos), Some(Token::RParen)) { + advance(pos); + } + Ok(val) + } + Token::Minus => { + advance(pos); + Ok(-parse_factor(tokens, pos)?) + } + other => Err(format!("unexpected token: {other:?}")), + } +} + +// ================================================== +// REPL +// ================================================== + +#[allow(dead_code)] +async fn run_repl(client: Arc) { + eprintln!("Connected to: {}\n", client.model()); + println!("Type a message and press Enter. Type 'quit' to exit.\n"); + + // Create the agent once — conversation history persists across inputs. + let config = LoopConfig { + max_turns: 10, + ..Default::default() + }; + let mut agent = BareLoop::new(client, build_tools(), config); + agent.register_observer(Arc::new(PrintingObserver)); + + // Stream text deltas in real-time. + agent.set_text_streamer(Arc::new(|delta| { + print!("{delta}"); + let _ = std::io::stdout().flush(); + })); + + let stdin = io::stdin(); + let mut total_input: u64 = 0; + let mut total_output: u64 = 0; + + loop { + print!("> "); + #[allow(clippy::let_underscore_must_use)] + let _ = io::stdout().flush(); + + let mut input = String::new(); + if stdin.lock().read_line(&mut input).is_err() { + break; + } + let input = input.trim(); + if input.is_empty() { + continue; + } + if input == "quit" || input == "exit" { + break; + } + + match agent.run(input).await { + Ok(result) => { + total_input += result.input_tokens; + total_output += result.output_tokens; + // Text was already streamed live. Just print stats. + if result + .final_output + .as_deref() + .map_or(true, |s| s.is_empty()) + { + println!(" (empty response)"); + } + println!( + "\n\n (turns: {}, tokens: {}+{} | total: {}+{})\n", + result.total_turns, + result.input_tokens, + result.output_tokens, + total_input, + total_output + ); + } + Err(e) => { + eprintln!("\n Error: {e}\n"); + } + } + } +} + +// ================================================== +// Provider detection & main +// +// Each provider produces a different concrete type, so we can't unify +// them into a single return. Instead, the `try_provider!` macro wraps +// the repeated "check env → build-or-die → run_repl → return" pattern. +// ================================================== + +/// Build a client from the given expression, or print the error and exit. +macro_rules! build_or_die { + ($provider:expr, $label:literal) => {{ + let label: &str = $label; + $provider.unwrap_or_else(|e| { + eprintln!("{label} client error: {e}"); + std::process::exit(1); + }) + }}; +} + +#[cfg(feature = "ollama")] +fn ollama_model_from_env() -> Option { + let direct = std::env::var("OLLAMA_MODEL").or_else(|_| std::env::var("MODEL")); + match direct { + Ok(m) => Some(m), + Err(_) => { + let is_ollama_url = + std::env::var("OPENAI_BASE_URL").is_ok_and(|u| u.contains("localhost:11434")); + if is_ollama_url { + Some("llama3".into()) + } else { + None + } + } + } +} + +#[tokio::main] +async fn main() { + // Ollama (OpenAI-compatible, local) + #[cfg(feature = "ollama")] + if let Some(model) = ollama_model_from_env() { + let base = std::env::var("OLLAMA_BASE_URL") + .or_else(|_| std::env::var("OPENAI_BASE_URL")) + .unwrap_or_else(|_| "http://localhost:11434/v1".into()); + + let client = build_or_die!( + loopctl::provider::OpenAiClient::builder() + .api_key("ollama") + .base_url(base) + .model(model) + .build(), + "Ollama" + ); + run_repl(Arc::new(client)).await; + return; + } + + // DeepSeek (OpenAI-compatible) + #[cfg(feature = "deepseek")] + if std::env::var("DEEPSEEK_API_KEY").is_ok() { + let client = build_or_die!(loopctl::provider::deepseek(), "DeepSeek"); + run_repl(Arc::new(client)).await; + return; + } + + // Grok / xAI (OpenAI-compatible) + #[cfg(feature = "grok")] + if std::env::var("XAI_API_KEY").is_ok() || std::env::var("GROK_API_KEY").is_ok() { + let client = build_or_die!(loopctl::provider::grok(), "Grok"); + run_repl(Arc::new(client)).await; + return; + } + + // Gemini (Google) + #[cfg(feature = "gemini")] + if std::env::var("GEMINI_API_KEY").is_ok() || std::env::var("GOOGLE_API_KEY").is_ok() { + let client = build_or_die!(loopctl::provider::GeminiClient::from_env(), "Gemini"); + run_repl(Arc::new(client)).await; + return; + } + + // Z.ai (ZhipuAI / BigModel) + #[cfg(feature = "zai")] + if std::env::var("ZAI_API_KEY").is_ok() || std::env::var("ZHIPUAI_API_KEY").is_ok() { + let client = build_or_die!(loopctl::provider::zai(), "Z.ai"); + run_repl(Arc::new(client)).await; + return; + } + + // Self-hosted (vLLM, LM Studio, etc.) + #[cfg(feature = "openai")] + if let Ok(base) = std::env::var("SELF_HOSTED_BASE_URL") { + let model = std::env::var("SELF_HOSTED_MODEL").unwrap_or_else(|_| "default".into()); + let client = build_or_die!(loopctl::provider::self_hosted(&base, &model), "Self-hosted"); + run_repl(Arc::new(client)).await; + return; + } + + // OpenAI + #[cfg(feature = "openai")] + if std::env::var("OPENAI_API_KEY").is_ok() || std::env::var("API_KEY").is_ok() { + let client = build_or_die!(loopctl::provider::OpenAiClient::from_env(), "OpenAI"); + run_repl(Arc::new(client)).await; + return; + } + + // Anthropic + #[cfg(feature = "anthropic")] + if std::env::var("ANTHROPIC_API_KEY").is_ok() { + let client = build_or_die!(loopctl::provider::AnthropicClient::from_env(), "Anthropic"); + run_repl(Arc::new(client)).await; + return; + } + + print_usage_and_exit(); +} From 9373577b2e560b1b1f48303774b378f952d2c2eb Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 08:28:31 +1200 Subject: [PATCH 09/30] chore: replace CancelSignal with CancellationToken, fix loop detection edge cases, fix clippy pedantic warnings, update ci --- .github/workflows/ci.yml | 2 +- Cargo.toml | 1 + Makefile | 2 +- examples/chat.rs | 28 +++++++------ examples/echo-tool-cli.rs | 15 ++++--- examples/hello-cli.rs | 2 + examples/repl-cli.rs | 2 + src/api/error.rs | 9 ++--- src/cancel.rs | 74 ++++++++++------------------------ src/compact.rs | 23 ++++++++--- src/config.rs | 2 +- src/detection/convergence.rs | 23 +++++++---- src/detection/loop_detector.rs | 31 +++++++++----- src/detection/manager.rs | 2 +- src/engine/bare.rs | 58 +++++++------------------- src/engine/bare/dispatch.rs | 16 ++++++++ src/fallback.rs | 39 +++++++++++++++++- src/hooks/executor.rs | 48 +++++++++------------- src/lib.rs | 18 +++++++++ src/memory.rs | 2 +- src/memory/builtin.rs | 2 +- src/message.rs | 2 +- src/middleware.rs | 53 +++++++++++++----------- src/middleware/permission.rs | 2 +- src/middleware/timeout.rs | 3 +- src/middleware/unknown_tool.rs | 13 ++++-- src/provider.rs | 10 ++--- src/provider/openai.rs | 2 +- src/reflection/backoff.rs | 42 +++++++++++++++---- src/runtime.rs | 4 +- src/stream/handler.rs | 28 ++++++------- src/stream/heartbeat.rs | 7 ++-- src/tool.rs | 8 ++-- src/tool/health.rs | 2 +- src/tool/shield.rs | 8 +++- 35 files changed, 330 insertions(+), 253 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3a837db..248a92e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -33,7 +33,7 @@ jobs: with: components: clippy - uses: Swatinem/rust-cache@v2 - - run: cargo clippy --all-features -- -D warnings + - run: cargo clippy --all-targets --all-features -- -D warnings fmt: runs-on: ubuntu-latest diff --git a/Cargo.toml b/Cargo.toml index 8d934ec..74c13bf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ serde_json = "1" serde_repr = "0.1" thiserror = "2" tokio = { version = "1.52.3", features = ["sync", "macros", "time"] } +tokio-util = { version = "0.7", features = ["rt"] } uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" diff --git a/Makefile b/Makefile index 59ffd95..ea0ebdb 100644 --- a/Makefile +++ b/Makefile @@ -10,7 +10,7 @@ test: cargo test --doc --all-features clippy: - cargo clippy --all-features -- -D warnings + cargo clippy --all-targets --all-features -- -D warnings fmt: cargo fmt --all -- --check diff --git a/examples/chat.rs b/examples/chat.rs index e228917..1c775e2 100644 --- a/examples/chat.rs +++ b/examples/chat.rs @@ -100,12 +100,13 @@ fn print_usage_and_exit() -> ! { // ================================================== fn echo_fn(input: serde_json::Value, _ctx: &ToolContext) -> ToolFuture { - let text = input - .get("message") - .and_then(|v| v.as_str()) - .unwrap_or("(empty)") - .to_string(); - Box::pin(async move { Ok(ToolOutput::text(format!("echo: {text}"))) }) + Box::pin(async move { + let text = input + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("(empty)"); + Ok(ToolOutput::text(format!("echo: {text}"))) + }) } fn current_time_fn(_input: serde_json::Value, _ctx: &ToolContext) -> ToolFuture { @@ -119,13 +120,14 @@ fn current_time_fn(_input: serde_json::Value, _ctx: &ToolContext) -> ToolFuture } fn calculate_fn(input: serde_json::Value, _ctx: &ToolContext) -> ToolFuture { - let expr = input - .get("expression") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let result = simple_eval(&expr); - Box::pin(async move { Ok(ToolOutput::text(result)) }) + Box::pin(async move { + let expr = input + .get("expression") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let result = simple_eval(expr); + Ok(ToolOutput::text(result)) + }) } /// Build the tool registry for the chat example. diff --git a/examples/echo-tool-cli.rs b/examples/echo-tool-cli.rs index 5ca3ff6..7fe1d2a 100644 --- a/examples/echo-tool-cli.rs +++ b/examples/echo-tool-cli.rs @@ -7,6 +7,8 @@ //! cargo run --example echo-tool-cli --features testing //! ``` +#![allow(clippy::expect_used, clippy::doc_markdown)] + use std::sync::Arc; use loopctl::config::LoopConfig; @@ -27,12 +29,13 @@ fn echo_fn( + 'static, >, > { - let text = input - .get("message") - .and_then(|v| v.as_str()) - .unwrap_or("(empty)") - .to_string(); - Box::pin(async move { Ok(ToolOutput::text(format!("echo: {text}"))) }) + Box::pin(async move { + let text = input + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("(empty)"); + Ok(ToolOutput::text(format!("echo: {text}"))) + }) } #[tokio::main] diff --git a/examples/hello-cli.rs b/examples/hello-cli.rs index d9932a0..1a83e5d 100644 --- a/examples/hello-cli.rs +++ b/examples/hello-cli.rs @@ -2,6 +2,8 @@ //! //! Demonstrates the absolute simplest way to run a [`BareLoop`]: //! create a mock client, build the loop, and call [`run`]. + +#![allow(clippy::expect_used, clippy::doc_markdown)] //! //! ```sh //! cargo run --example hello-cli --features testing diff --git a/examples/repl-cli.rs b/examples/repl-cli.rs index ff90bb8..f823bc2 100644 --- a/examples/repl-cli.rs +++ b/examples/repl-cli.rs @@ -8,6 +8,8 @@ //! cargo run --example repl-cli --features testing //! ``` +#![allow(clippy::unwrap_used)] + use std::io::{self, BufRead, Write}; use std::sync::Arc; diff --git a/src/api/error.rs b/src/api/error.rs index 3a682d8..8cf1315 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -1305,12 +1305,11 @@ mod tests { #[test] fn test_result_type() { - fn returns_result() -> Result { - Ok("success".to_string()) + fn returns_result() -> String { + "success".to_string() } let result = returns_result(); - assert!(result.is_ok()); - assert_eq!(result.unwrap(), "success"); + assert_eq!(result, "success"); } #[test] @@ -1328,7 +1327,7 @@ mod tests { #[test] fn test_from_hyper() { - let error = ApiError::from_hyper(std::io::Error::new(std::io::ErrorKind::Other, "oops")); + let error = ApiError::from_hyper(std::io::Error::other("oops")); assert!(matches!(error, ApiError::Http(_))); } diff --git a/src/cancel.rs b/src/cancel.rs index 631a992..c606620 100644 --- a/src/cancel.rs +++ b/src/cancel.rs @@ -1,15 +1,8 @@ //! Cooperative cancellation signal for agent loops. //! -//! [`CancelSignal`] combines an [`AtomicBool`] flag with a -//! [`tokio::sync::Notify`] for sub-millisecond wake-up of waiting tasks. -//! -//! # Why not poll an `AtomicBool`? -//! -//! Polling works but wastes CPU cycles and introduces latency proportional -//! to the poll interval. By pairing the flag with a `Notify`, any call to -//! [`CancelSignal::cancel`] instantly wakes every task that is -//! `tokio::select!`-ing on [`CancelSignal::notified`], giving sub-µs -//! response time. +//! [`CancelSignal`] wraps a [`tokio_util::sync::CancellationToken`], which is +//! purpose-built to avoid the time-of-check-to-time-of-use (TOCTOU) race that +//! plagues hand-rolled `AtomicBool` + `Notify` combinations. //! //! # Usage //! @@ -28,61 +21,48 @@ //! } //! ``` -use std::sync::atomic::{AtomicBool, Ordering}; -use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; -/// Shared cancellation signal backed by an [`AtomicBool`] flag and a -/// [`Notify`] for instant wake-up. +/// Shared cancellation signal backed by a +/// [`tokio_util::sync::CancellationToken`]. /// /// Wrap in `Arc` for sharing across tasks or threads. Create with /// [`CancelSignal::new`], cancel with [`CancelSignal::cancel`], and await /// instant notification with [`CancelSignal::notified`]. pub struct CancelSignal { - flag: AtomicBool, - notify: Notify, + inner: CancellationToken, } impl CancelSignal { /// Create a new, non-cancelled signal. /// - /// Returns a [`CancelSignal`] with its internal flag set to `false`. - /// The signal is ready to be shared (via `Arc`) and awaited by - /// worker tasks until [`cancel`](Self::cancel) is called. + /// Returns a [`CancelSignal`] backed by a fresh + /// [`CancellationToken`]. The signal is ready to be shared (via `Arc`) + /// and awaited by worker tasks until [`cancel`](Self::cancel) is called. #[must_use] pub fn new() -> Self { Self { - flag: AtomicBool::new(false), - notify: Notify::new(), + inner: CancellationToken::new(), } } /// Fire the cancellation signal. /// - /// Sets the internal flag to `true` **and** wakes every task currently + /// Sets the internal state to cancelled **and** wakes every task currently /// awaiting [`Self::notified`]. Idempotent — calling multiple times is /// safe. pub fn cancel(&self) { - self.flag.store(true, Ordering::Release); - self.notify.notify_waiters(); - } - - /// Reset the signal so it can be reused for a new operation. - /// - /// Clears the internal cancellation flag, returning the signal to - /// its initial non-cancelled state. Any subsequent calls to - /// [`is_cancelled`](Self::is_cancelled) will return `false` until - /// [`cancel`](Self::cancel) is called again. - pub fn reset(&self) { - self.flag.store(false, Ordering::Release); + self.inner.cancel(); } /// Check whether the signal has been cancelled. /// - /// Performs a non-blocking load of the internal flag. Returns - /// `true` if [`cancel`](Self::cancel) has been called since the - /// last [`reset`](Self::reset) (or since construction). + /// Performs a non-blocking check of the internal [`CancellationToken`]. + /// Returns `true` if [`cancel`](Self::cancel) has been called since + /// construction. + #[must_use] pub fn is_cancelled(&self) -> bool { - self.flag.load(Ordering::Acquire) + self.inner.is_cancelled() } /// Return a future that completes **instantly** when [`Self::cancel`] @@ -91,6 +71,9 @@ impl CancelSignal { /// If the signal is already cancelled, the returned future completes /// immediately on the first `.await`. /// + /// This delegates to [`CancellationToken::cancelled`], which is + /// race-free by construction — no flag-then-wait loop required. + /// /// Use inside `tokio::select!` alongside the actual work future: /// /// ```rust,ignore @@ -100,11 +83,7 @@ impl CancelSignal { /// } /// ``` pub async fn notified(&self) { - let notified = self.notify.notified(); - if self.flag.load(Ordering::Acquire) { - return; - } - notified.await; + self.inner.cancelled().await; } } @@ -151,15 +130,6 @@ mod tests { assert!(signal.is_cancelled()); } - #[test] - fn test_reset() { - let signal = CancelSignal::new(); - signal.cancel(); - assert!(signal.is_cancelled()); - signal.reset(); - assert!(!signal.is_cancelled()); - } - #[tokio::test] async fn test_cancel_is_idempotent() { let signal = Arc::new(CancelSignal::new()); diff --git a/src/compact.rs b/src/compact.rs index 432a159..ab976d0 100644 --- a/src/compact.rs +++ b/src/compact.rs @@ -206,7 +206,6 @@ pub enum CompactBase { /// /// Use this when you want compaction to aim for a fixed fraction /// of the model's total capacity regardless of the trigger threshold. - #[default] Context, /// Target is a percentage of the trigger threshold. @@ -216,6 +215,7 @@ pub enum CompactBase { /// This is the default. With the default `threshold = 0.80` and /// `compact_target_pct = 0.70`, compaction targets 56% of the /// context window (0.80 × 0.70 = 0.56). + #[default] Threshold, } @@ -448,7 +448,7 @@ impl ContextManager { .parts .iter() .map(|p| match p { - MessagePart::Text { text } => text.len() as u64, + MessagePart::Text { text } => text.chars().count() as u64, MessagePart::Image { .. } => 256, // rough base64 estimate MessagePart::ToolCall { name, input, .. } => { let name_len = name.len() as u64; @@ -618,6 +618,16 @@ impl ContextManager { }); } + if outcome.tokens_after > self.context_window { + return Err(ContextOverflow { + tokens_used: outcome.tokens_after, + context_window: self.context_window, + message_count, + trigger: CompactReason::Manual, + compactor_error: None, + }); + } + Ok(EnsureContextResult::Compacted(outcome)) } @@ -668,9 +678,10 @@ impl ContextManager { tool_messages: pre_messages .iter() .filter(|m| { - m.parts - .iter() - .any(crate::message::MessagePart::is_tool_call) + m.parts.iter().any(|p| { + crate::message::MessagePart::is_tool_call(p) + || crate::message::MessagePart::is_tool_result(p) + }) }) .count(), }, @@ -932,7 +943,7 @@ mod tests { assert_eq!(outcome.messages.len(), 3); assert!(outcome.tokens_saved > 0); // Verify the first message was preserved. - assert!(outcome.messages.first().is_some()); + assert!(!outcome.messages.is_empty()); assert_eq!(outcome.messages.first().unwrap().role, first_role.unwrap()); } diff --git a/src/config.rs b/src/config.rs index 5729f8c..f5605c2 100644 --- a/src/config.rs +++ b/src/config.rs @@ -29,7 +29,7 @@ use uuid::Uuid; /// ..Default::default() /// }; /// ``` -#[derive(Debug, Clone)] +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct LoopConfig { /// Unique session identifier (random UUID v4). pub session_id: Uuid, diff --git a/src/detection/convergence.rs b/src/detection/convergence.rs index 7ce57b4..236ecad 100644 --- a/src/detection/convergence.rs +++ b/src/detection/convergence.rs @@ -540,6 +540,7 @@ impl ConvergenceDetector { } let mut max_similarity = 0.0; + let mut any_similar = false; for prev_response in &self.window { let similarity = Self::compute_similarity(response, prev_response); if similarity > max_similarity { @@ -547,17 +548,25 @@ impl ConvergenceDetector { } if similarity >= self.config.similarity_threshold { - self.consecutive_count = self.consecutive_count.saturating_add(1); + any_similar = true; if !self.similar_responses.contains(&response.to_string()) { self.similar_responses.push(response.to_string()); } - } else { - self.consecutive_count = 1; - self.similar_responses.clear(); - self.similar_responses.push(response.to_string()); } } + // Update consecutive count once per add_response call + if self.window.is_empty() { + self.consecutive_count = 1; + self.similar_responses.push(response.to_string()); + } else if any_similar { + self.consecutive_count = self.consecutive_count.saturating_add(1); + } else { + self.consecutive_count = 1; + self.similar_responses.clear(); + self.similar_responses.push(response.to_string()); + } + if self.window.len() >= self.config.window_size { self.window.pop_front(); } @@ -766,8 +775,8 @@ mod tests { let status1 = detector.add_response("alpha"); assert!(!status1.detected, "first response: no comparison possible"); assert_eq!( - status1.consecutive_count, 0, - "first response: window is empty" + status1.consecutive_count, 1, + "first response: starts a streak of 1" ); let status2 = detector.add_response("beta"); diff --git a/src/detection/loop_detector.rs b/src/detection/loop_detector.rs index a764dfe..aeec922 100644 --- a/src/detection/loop_detector.rs +++ b/src/detection/loop_detector.rs @@ -1383,6 +1383,14 @@ impl LoopDetector { } } + // Sort by tool then primary_param for deterministic ordering + // when multiple operations share the same repetition count. + repeated.sort_by(|a, b| { + a.tool + .cmp(&b.tool) + .then_with(|| a.primary_param.cmp(&b.primary_param)) + }); + (repeated, max) } @@ -1504,11 +1512,15 @@ impl LoopDetector { }; let sig = &self.signature; + // Normalize the input path the same way stored params are normalized + // so the comparison is consistent regardless of the signature used. + let normalized_input = sig.normalize_param_for_comparison("", file_path); let read_count = ops .iter() .filter(|o| { sig.is_file_read_tool(&o.tool) - && sig.normalize_param_for_comparison(&o.tool, &o.primary_param) == file_path + && sig.normalize_param_for_comparison(&o.tool, &o.primary_param) + == normalized_input }) .count(); @@ -1666,7 +1678,7 @@ mod tests { if path.is_empty() { pattern.to_string() } else { - format!("{}:{}", path, pattern) + format!("{path}:{pattern}") } } "Bash" => input @@ -1849,7 +1861,7 @@ mod tests { let detector = test_detector(); for i in 0..20 { - detector.record(Operation::new("Bash", format!("git status {}", i))); + detector.record(Operation::new("Bash", format!("git status {i}"))); } let status = detector.check_loop(); @@ -1883,7 +1895,7 @@ mod tests { let detector = test_detector(); for i in 0..5 { - let result_hash = hash_result(&format!("output {}", i)); + let result_hash = hash_result(&format!("output {i}")); detector.record(Operation::new("Bash", "git status").with_result_hash(result_hash)); } @@ -2159,7 +2171,7 @@ mod tests { let detector = test_detector(); for i in 0..5 { - let hash = hash_result(&format!("output {}", i)); + let hash = hash_result(&format!("output {i}")); detector.record(Operation::new("Bash", "git status").with_result_hash(hash)); } @@ -2263,8 +2275,7 @@ mod tests { let warning = status.warning.unwrap(); assert!( warning.contains("Check the command"), - "Bash warning should contain signature suggestion: {}", - warning + "Bash warning should contain signature suggestion: {warning}", ); } @@ -2426,8 +2437,7 @@ mod tests { #[test] fn test_build_warning_suppresses_duplicate() { let detector = test_detector(); - let op = Operation::new("Bash", "ls"); - let ops = vec![op.clone()]; + let ops = vec![Operation::new("Bash", "ls")]; // First call produces a warning and records the op as warned let w1 = detector.build_warning(&ops, 3, true, false); @@ -2441,8 +2451,7 @@ mod tests { #[test] fn test_build_warning_duplicate_not_suppressed_when_stopping() { let detector = test_detector(); - let op = Operation::new("Bash", "ls"); - let ops = vec![op.clone()]; + let ops = vec![Operation::new("Bash", "ls")]; // First warning let w1 = detector.build_warning(&ops, 3, true, false); diff --git a/src/detection/manager.rs b/src/detection/manager.rs index b357824..0ffbf6e 100644 --- a/src/detection/manager.rs +++ b/src/detection/manager.rs @@ -410,7 +410,7 @@ pub struct DetectionStats { /// /// The manager progresses through a simple per-turn cycle: /// -/// 1. **Construct** — `new()` or `with_config()`. +/// 1. **Construct** — `new()` or `new_with_config()`. /// 2. **Feed** — call `record_operation(op)` for each tool invocation and /// `record_response(text)` for each assistant reply. /// 3. **Check** — call `check_loop()`, `check_convergence()`, or the diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 07a42a1..5d8c290 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -18,22 +18,6 @@ //! prompt, session ID). //! - Optional [`LoopObserver`](crate::observer::LoopObserver) registrations for lifecycle instrumentation. //! -//! ```text -//! BareLoop -//! ┌─────────────────────────────────────────────────────────┐ -//! │ run(user_input) │ -//! │ 1. Push user message to conversation │ -//! │ 2. Loop: │ -//! │ a. stream_messages(conversation) → StreamEvents │ -//! │ b. accumulate into Message (assistant) │ -//! │ c. Extract tool calls from Message │ -//! │ d. Execute tools via ToolRegistry │ -//! │ e. Push tool_result messages to conversation │ -//! │ f. If no tool calls → break │ -//! │ 3. Return SessionResult │ -//! └─────────────────────────────────────────────────────────┘ -//! ``` -//! //! # Key Design Decisions //! //! - **Static dispatch** — `BareLoop` is generic over the @@ -50,9 +34,9 @@ //! # Quick Start //! //! ```rust,ignore -//! use loopctl::loop_::BareLoop; +//! use loopctl::engine::BareLoop; //! use loopctl::tool::ToolRegistry; -//! use loopctl::core::LoopConfig; +//! use loopctl::config::LoopConfig; //! use std::sync::Arc; //! //! // 1. Build components @@ -86,10 +70,8 @@ use crate::hooks::HookAction; #[cfg(feature = "hooks")] use crate::hooks::HookExecutor; #[cfg(feature = "hooks")] -#[allow(unused_imports)] use crate::hooks::context::{ CompactTrigger, PostCompactContext, PostToolUseContext, PreCompactContext, PreToolUseContext, - SessionEndContext, SessionEndReason, SessionStartContext, }; use crate::message::{Message, MessagePart, Role, ToolContent}; use crate::middleware::{ToolDispatchContext, ToolPipeline, ToolPipelineBuilder}; @@ -108,7 +90,6 @@ use crate::stream::{StreamAccumulator, StreamEvent, StreamStopReason, Usage}; use crate::tool::health::ToolHealthRegistry; use crate::tool::{PermissionCheck, ToolContext, ToolDispatchResult, ToolRegistry, ToolSchema}; -// Phase submodules mod compact; mod dispatch; mod emission; @@ -140,22 +121,12 @@ mod stream; /// - [`new_with_managers()`](BareLoop::new_with_managers) — full control, /// including a [`LoopRuntime`]. /// -/// # Lifecycle -/// -/// ```text -/// new() / new_with_managers() -/// → run(user_input) -/// → stream_turn() → dispatch_tools() → stream_turn() -/// → … (repeat until end_turn or max_turns) -/// → SessionResult -/// ``` -/// /// # Example /// /// ```rust,ignore -/// use loopctl::loop_::BareLoop; +/// use loopctl::engine::BareLoop; /// use loopctl::tool::ToolRegistry; -/// use loopctl::core::LoopConfig; +/// use loopctl::config::LoopConfig; /// use std::sync::Arc; /// /// let registry = ToolRegistry::new(); @@ -619,7 +590,7 @@ impl BareLoop { /// # Example /// /// ```rust,ignore - /// use loopctl::core::observer::LoopObserver; + /// use loopctl::observer::LoopObserver; /// use std::sync::Arc; /// /// let mut agent = BareLoop::new(client, registry, config); @@ -869,8 +840,8 @@ impl crate::engine::loop_core::Loop for BareLoop { success: true, error: None, duration_ms: Self::millis_u64(turn_start.elapsed()), - input_tokens: self.budget.input_tokens, - output_tokens: self.budget.output_tokens, + input_tokens: turn_in, + output_tokens: turn_out, }); self.state = LoopState::Completed { summary: text.clone(), @@ -1326,12 +1297,12 @@ mod tests { impl Tool for EchoTool { /// Return the tool name `"echo"`. - fn name(&self) -> &str { + fn name(&self) -> &'static str { "echo" } /// Return a human-readable description. - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Echoes back the input" } @@ -1377,12 +1348,12 @@ mod tests { impl Tool for FailingTool { /// Return the tool name `"fail"`. - fn name(&self) -> &str { + fn name(&self) -> &'static str { "fail" } /// Return a human-readable description. - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Always fails" } @@ -1451,7 +1422,7 @@ mod tests { } impl crate::observer::LoopObserver for CountingObserver { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "counting" } @@ -1868,8 +1839,7 @@ mod tests { assert!(!received.is_empty(), "streamer should have fired"); assert!( received.join("").contains("Hello world"), - "got: {:?}", - received + "got: {received:?}", ); } @@ -2193,7 +2163,7 @@ mod tests { } impl crate::middleware::ToolMiddleware for TurnNumberCapture { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "turn_capture" } diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index f6b1a6f..798158e 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -115,10 +115,26 @@ impl BareLoop { }); if let Some(blocked) = self.check_pre_tool_use_hooks(&tc, turn_idx) { + // Pair on_tool_pre with on_tool_post so observers see a + // complete lifecycle even when a hook blocks the call. + self.managers.observers().on_tool_post(&ToolPostContext { + turn: turn_idx, + tool: tc.tool.clone(), + result_hash: loop_detector::hash_result(&blocked.output.to_string()), + is_error: blocked.is_error, + duration: Duration::ZERO, + }); return Ok(blocked); } if let Some(blocked) = self.pre_detection(&tc, turn_idx) { + self.managers.observers().on_tool_post(&ToolPostContext { + turn: turn_idx, + tool: tc.tool.clone(), + result_hash: loop_detector::hash_result(&blocked.output.to_string()), + is_error: blocked.is_error, + duration: Duration::ZERO, + }); return Ok(blocked); } diff --git a/src/fallback.rs b/src/fallback.rs index ff0cc78..ee063b5 100644 --- a/src/fallback.rs +++ b/src/fallback.rs @@ -640,6 +640,8 @@ pub struct FallbackManager { active_fallback: Mutex>, /// Time when fallback was activated. fallback_switched_at: Mutex>, + /// How long to remain in fallback before attempting primary recovery. + recovery_timeout: Duration, } impl FallbackManager { @@ -680,6 +682,7 @@ impl FallbackManager { fallback_models: Mutex::new(Vec::new()), active_fallback: Mutex::new(None), fallback_switched_at: Mutex::new(None), + recovery_timeout: Duration::from_secs(60), } } @@ -707,9 +710,40 @@ impl FallbackManager { self.fallback_threshold = config.trip_threshold; self.primary_resume_threshold = config.recovery_successes_needed; self.default_max_fail_count = config.max_fail_count; + self.recovery_timeout = config.recovery_timeout; self } + /// Set the recovery timeout (builder style). + /// + /// This is how long the manager stays in fallback before it is willing + /// to probe the primary model again via + /// [`should_try_resume_primary`](Self::should_try_resume_primary). + /// + /// Mirrors [`FallbackConfig::recovery_timeout`] for cases where a full + /// [`with_config`](Self::with_config) is not desired. + #[must_use] + pub fn with_recovery_timeout(mut self, recovery_timeout: Duration) -> Self { + self.recovery_timeout = recovery_timeout; + self + } + + /// Configured recovery timeout. + /// + /// Returns the duration the manager will remain in fallback before it + /// is willing to probe the primary model again. Set via + /// [`with_config`](Self::with_config) (from + /// [`FallbackConfig::recovery_timeout`]) or + /// [`with_recovery_timeout`](Self::with_recovery_timeout). + /// + /// Pass this to + /// [`should_try_resume_primary`](Self::should_try_resume_primary) to + /// honour the configured timeout without hard-coding a value. + #[must_use] + pub fn recovery_timeout(&self) -> Duration { + self.recovery_timeout + } + /// Create with fallback already activated. /// /// Useful when a new manager should start in the @@ -1449,8 +1483,8 @@ impl FallbackManager { /// let mgr = FallbackManager::new(3, 2); /// assert!(!mgr.record_api_failure()); // 1 /// assert!(!mgr.record_api_failure()); // 2 - /// assert!(mgr.record_api_failure()); // 3 — threshold reached - /// assert!(mgr.record_api_failure()); // 4 — still not activated + /// assert!(mgr.record_api_failure()); // 3 — threshold reached, now activated + /// assert!(!mgr.record_api_failure()); // 4 — already activated, no re-trip /// ``` pub fn record_api_failure(&self) -> bool { let failures = self @@ -1463,6 +1497,7 @@ impl FallbackManager { threshold = self.fallback_threshold, "Fallback threshold reached" ); + self.fallback_activated.store(true, Ordering::Relaxed); true } else { false diff --git a/src/hooks/executor.rs b/src/hooks/executor.rs index 35cc4d5..3221cd2 100644 --- a/src/hooks/executor.rs +++ b/src/hooks/executor.rs @@ -80,26 +80,14 @@ impl HookExecutor { /// Set the interactivity mode. /// - /// Use [`interactivity`](Self::interactivity) to change the mode after - /// construction, or [`with_hook`](Self::with_hook) to add hooks via - /// builder pattern. + /// Use this builder method to change the mode after construction, or + /// [`with_hook`](Self::with_hook) to add hooks via the builder pattern. #[must_use] pub fn with_interactivity(mut self, interactivity: Interactivity) -> Self { self.interactivity = interactivity; self } - /// Set the interactivity mode (builder style). - /// - /// Overrides the current [`Interactivity`] mode and returns `self` - /// for chaining: - /// `HookExecutor::new().interactivity(Interactivity::Interactive).with_hook(h)`. - #[must_use] - pub fn interactivity(mut self, mode: Interactivity) -> Self { - self.interactivity = mode; - self - } - /// Register a hook (builder pattern). /// /// Hooks are called in registration order. Returns `self` @@ -316,7 +304,7 @@ mod tests { struct AllowHook; impl Hook for AllowHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "allow" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { @@ -328,7 +316,7 @@ mod tests { reason: String, } impl Hook for BlockHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "block" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { @@ -350,7 +338,7 @@ mod tests { } } impl Hook for PostRecorder { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "post_recorder" } fn on_post_tool_use(&self, _ctx: &PostToolUseContext) { @@ -449,7 +437,7 @@ mod tests { fn check_pre_compact_merges_instructions() { struct InstructionHook; impl Hook for InstructionHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "instruction" } fn on_pre_compact(&self, _ctx: &PreCompactContext) -> Option { @@ -479,7 +467,7 @@ mod tests { fn check_pre_compact_abort_takes_priority() { struct AbortHook; impl Hook for AbortHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "abort" } fn on_pre_compact(&self, _ctx: &PreCompactContext) -> Option { @@ -488,7 +476,7 @@ mod tests { } struct InstructionHook; impl Hook for InstructionHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "instruction" } fn on_pre_compact(&self, _ctx: &PreCompactContext) -> Option { @@ -520,7 +508,7 @@ mod tests { ends: AtomicUsize, } impl Hook for CounterHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "counter" } fn on_session_start(&self, _ctx: &SessionStartContext) { @@ -587,7 +575,7 @@ mod tests { count: AtomicUsize, } impl Hook for CompactRecorder { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "compact_recorder" } fn on_post_compact(&self, _ctx: &PostCompactContext) { @@ -638,7 +626,7 @@ mod tests { fn check_pre_tool_use_headless_downgrades_ask_to_block() { struct AskHook; impl Hook for AskHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "ask" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { @@ -660,7 +648,7 @@ mod tests { fn check_pre_tool_use_interactive_passes_ask_through() { struct AskHook; impl Hook for AskHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "ask" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { @@ -684,7 +672,7 @@ mod tests { fn check_pre_tool_use_interactivity_builder() { struct AskHook; impl Hook for AskHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "ask" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { @@ -694,7 +682,7 @@ mod tests { // Builder style: start Headless, switch to Interactive. let executor = HookExecutor::new() - .interactivity(Interactivity::Interactive) + .with_interactivity(Interactivity::Interactive) .with_hook(Arc::new(AskHook)); let ctx = dummy_pre_ctx(); let action = executor.check_pre_tool_use(&ctx); @@ -717,7 +705,7 @@ mod tests { fn headless_block_passes_through_unchanged() { struct BlockOnlyHook; impl Hook for BlockOnlyHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "block_only" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { @@ -737,7 +725,7 @@ mod tests { fn headless_downgrade_preserves_original_message() { struct AskHook; impl Hook for AskHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "ask" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { @@ -765,7 +753,7 @@ mod tests { fn with_interactivity_constructor_sets_mode() { struct AskHook; impl Hook for AskHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "ask" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { @@ -787,7 +775,7 @@ mod tests { fn interactive_block_passes_through_unchanged() { struct BlockOnlyHook; impl Hook for BlockOnlyHook { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "block_only" } fn on_pre_tool_use(&self, _ctx: &PreToolUseContext) -> Option { diff --git a/src/lib.rs b/src/lib.rs index cc724d3..85ec12a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -42,6 +42,24 @@ //! - **[`hooks`]** — Bidirectional lifecycle control (allow/block/ask before tool use, compaction). *Requires `hooks` feature.* //! - **[`testing`]** — Test utilities and fixtures. *Requires `testing` feature.* +// Relax strict lints in test code. The crate enforces a strict no-panic / +// no-unwrap policy in production code, but test code legitimately uses +// assertions, unwrap, indexing, etc. for readability. +#![cfg_attr( + test, + allow( + clippy::unwrap_used, + clippy::expect_used, + clippy::panic, + clippy::indexing_slicing, + clippy::missing_panics_doc, + clippy::clone_on_ref_ptr, + clippy::doc_markdown, + clippy::field_reassign_with_default, + clippy::used_underscore_items, + ) +)] + pub mod api; pub mod cancel; pub mod capabilities; diff --git a/src/memory.rs b/src/memory.rs index c50f411..6a77e88 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -7,7 +7,7 @@ //! //! # Provided Implementations //! -//! - **`TrajectoryMemory`** — Records tool-execution trajectories and +//! - **[`InMemoryStore`]** — Records tool-execution trajectories and //! retrieves relevant past experiences. //! //! # Quick Start diff --git a/src/memory/builtin.rs b/src/memory/builtin.rs index 0e4e32a..be2aefc 100644 --- a/src/memory/builtin.rs +++ b/src/memory/builtin.rs @@ -139,7 +139,7 @@ impl InMemoryStore { /// Create a store pre-populated with the given entries. /// /// Useful for setting up test fixtures or seeding an agent with - /// Replace all entries with the provided list. + /// initial context. /// /// # Example /// diff --git a/src/message.rs b/src/message.rs index 83110ef..fdb8b4a 100644 --- a/src/message.rs +++ b/src/message.rs @@ -1146,7 +1146,7 @@ mod tests { let part = ToolContentPart::text("output"); match &part { ToolContentPart::Text { text } => assert_eq!(text, "output"), - _ => panic!("expected text part"), + ToolContentPart::Image { .. } => panic!("expected text part"), } } diff --git a/src/middleware.rs b/src/middleware.rs index d499180..6e13833 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -432,10 +432,10 @@ mod tests { struct EchoTool; impl Tool for EchoTool { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "echo" } - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Echoes back the input" } fn schema(&self) -> ToolSchema { @@ -463,10 +463,10 @@ mod tests { struct ErrorTool; impl Tool for ErrorTool { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "error_tool" } - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Always returns an error" } fn schema(&self) -> ToolSchema { @@ -491,10 +491,10 @@ mod tests { } impl Tool for SlowTool { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "slow" } - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Takes a long time" } fn schema(&self) -> ToolSchema { @@ -596,7 +596,7 @@ mod tests { assert_eq!(result.resolved_tool_name, "echo"); match result.output { ToolContent::Text(ref t) => assert_eq!(t, "hello"), - other => panic!("expected Text, got {other:?}"), + other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"), } } @@ -608,7 +608,7 @@ mod tests { assert_eq!(result.resolved_tool_name, "nonexistent"); match result.output { ToolContent::Text(ref t) => assert!(t.contains("not found")), - other => panic!("expected Text, got {other:?}"), + other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"), } } @@ -639,7 +639,7 @@ mod tests { t.contains("Permission") && t.contains("blocked"), "got: {t}" ), - other => panic!("expected Text, got {other:?}"), + other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"), } } @@ -681,7 +681,7 @@ mod tests { t.contains("Permission") && t.contains("blocked"), "got: {t}" ), - other => panic!("expected Text, got {other:?}"), + other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"), } } @@ -746,7 +746,7 @@ mod tests { assert!(result.is_error); match result.output { ToolContent::Text(ref t) => assert!(t.contains("timed out")), - other => panic!("expected Text, got {other:?}"), + other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"), } } @@ -774,8 +774,8 @@ mod tests { #[test] fn test_similarity_empty() { - assert_eq!(UnknownToolMiddleware::similarity("", ""), 1.0); - assert_eq!(UnknownToolMiddleware::similarity("a", ""), 0.0); + assert!((UnknownToolMiddleware::similarity("", "") - 1.0).abs() < f64::EPSILON); + assert!((UnknownToolMiddleware::similarity("a", "") - 0.0).abs() < f64::EPSILON); } // ================================================== @@ -841,15 +841,14 @@ mod tests { // Should return immediately with permission denied, not wait for timeout assert!( elapsed < Duration::from_millis(200), - "permission should short-circuit before timeout, took {:?}", - elapsed + "permission should short-circuit before timeout, took {elapsed:?}", ); match result.output { ToolContent::Text(ref t) => assert!( t.contains("Permission") && t.contains("blocked"), "got: {t}" ), - other => panic!("expected Text, got {other:?}"), + other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"), } } @@ -877,7 +876,7 @@ mod tests { assert!(!result.is_error); match result.output { ToolContent::Text(ref t) => assert_eq!(t, "hello"), - other => panic!("expected Text, got {other:?}"), + other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"), } } @@ -908,7 +907,7 @@ mod tests { "expected suggestion, got: {msg}" ); } - other => panic!("expected Text, got {other:?}"), + other @ ToolContent::Multipart(_) => panic!("expected Text, got {other:?}"), } } @@ -922,7 +921,7 @@ mod tests { } impl ToolMiddleware for ReachTracker { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "reach_tracker" } @@ -966,11 +965,11 @@ mod tests { struct LongOutputTool; impl Tool for LongOutputTool { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "long_output" } - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Returns a long string for testing" } @@ -988,7 +987,13 @@ mod tests { _ctx: &ToolContext, ) -> Pin> + Send + '_>> { Box::pin(async move { - let count = input.get("count").and_then(|v| v.as_u64()).unwrap_or(100) as usize; + let count = usize::try_from( + input + .get("count") + .and_then(serde_json::Value::as_u64) + .unwrap_or(100), + ) + .unwrap(); let ch = input.get("char").and_then(|v| v.as_str()).unwrap_or("a"); Ok(ToolOutput::text(ch.repeat(count))) }) @@ -1000,11 +1005,11 @@ mod tests { struct MultipartTool; impl Tool for MultipartTool { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "multipart" } - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Returns multipart content" } diff --git a/src/middleware/permission.rs b/src/middleware/permission.rs index 73845f7..cd98a80 100644 --- a/src/middleware/permission.rs +++ b/src/middleware/permission.rs @@ -40,7 +40,7 @@ pub type AskResolverFn = /// if ctx.tool_name == "safe_read" { /// PermissionCheck::Allow /// } else { -/// PermissionCheck::Deny +/// PermissionCheck::Deny { reason: "not on allowlist".into() } /// } /// }); /// ``` diff --git a/src/middleware/timeout.rs b/src/middleware/timeout.rs index 657c06d..a41665a 100644 --- a/src/middleware/timeout.rs +++ b/src/middleware/timeout.rs @@ -65,7 +65,8 @@ impl TimeoutMiddleware { /// Create a timeout middleware with a fixed timeout in seconds. /// - /// Uses default retry settings (one retry with double timeout). + /// Uses default retry settings (`retry_on_timeout: false`, + /// `max_retries: 0` — i.e. no retries). #[must_use] pub fn from_secs(secs: u64) -> Self { Self { diff --git a/src/middleware/unknown_tool.rs b/src/middleware/unknown_tool.rs index 7602e0f..b362097 100644 --- a/src/middleware/unknown_tool.rs +++ b/src/middleware/unknown_tool.rs @@ -37,6 +37,9 @@ pub struct UnknownToolMiddleware { impl UnknownToolMiddleware { /// Create a new unknown-tool middleware with default settings. /// + /// `registry` is the [`ToolRegistry`] used to look up known tools and + /// generate suggestions when an unknown tool is invoked. + /// /// Uses a [`suggestion_threshold`](Self::with_threshold) of `0.4`, /// which balances catching common typos against false-positive /// suggestions. @@ -163,7 +166,9 @@ impl UnknownToolMiddleware { /// Check if a result looks like a "tool not found" error. /// /// Only considers single [`Text`](ToolContent::Text) results whose - /// lowercased body contains `"not found"`. + /// lowercased body mentions a tool that was not found. This is more + /// specific than matching a bare `"not found"` so that unrelated errors + /// like "file not found" are not mistaken for unknown tools. /// [`Multipart`](ToolContent::Multipart) results and non-error results /// always return `false`. fn is_tool_not_found(result: &ToolDispatchResult) -> bool { @@ -174,7 +179,7 @@ impl UnknownToolMiddleware { ToolContent::Text(t) => t.to_lowercase(), ToolContent::Multipart(_) => return false, }; - msg.contains("not found") + msg.contains("not found") && msg.contains("tool") } } @@ -300,7 +305,7 @@ mod tests { let (suggestion, score) = UnknownToolMiddleware::find_best_match_inner("read_file", &available, 0.5).unwrap(); assert_eq!(suggestion, "read_file"); - assert_eq!(score, 1.0); + assert!((score - 1.0).abs() < f64::EPSILON); } #[test] @@ -364,7 +369,7 @@ mod tests { let (suggestion, score) = UnknownToolMiddleware::find_best_match_inner("read_file", &available, 1.0).unwrap(); assert_eq!(suggestion, "read_file"); - assert_eq!(score, 1.0); + assert!((score - 1.0).abs() < f64::EPSILON); assert!( UnknownToolMiddleware::find_best_match_inner("read_fil", &available, 1.0).is_none() diff --git a/src/provider.rs b/src/provider.rs index b6563bb..ec967d2 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -420,18 +420,18 @@ mod tests { #[cfg(feature = "ollama")] #[test] fn ollama_client_builds_with_defaults() { + use crate::api::ApiClient; env_remove!("OLLAMA_BASE_URL"); let client = ollama("llama3").unwrap(); - use crate::api::ApiClient; assert_eq!(client.model(), "llama3"); } #[cfg(feature = "ollama")] #[test] fn ollama_client_respects_base_url_env() { + use crate::api::ApiClient; env_set!("OLLAMA_BASE_URL", "http://my-host:1234/v1"); let client = ollama("test-model").unwrap(); - use crate::api::ApiClient; assert_eq!(client.model(), "test-model"); env_remove!("OLLAMA_BASE_URL"); } @@ -439,11 +439,11 @@ mod tests { #[cfg(feature = "ollama")] #[test] fn ollama_client_uses_api_key_when_set() { + use crate::api::ApiClient; env_remove!("OLLAMA_BASE_URL"); env_set!("OLLAMA_API_KEY", "my-cloud-key"); // Should build successfully with the cloud key — no network call. let client = ollama("llama3").unwrap(); - use crate::api::ApiClient; assert_eq!(client.model(), "llama3"); env_remove!("OLLAMA_API_KEY"); } @@ -451,19 +451,19 @@ mod tests { #[cfg(feature = "ollama")] #[test] fn ollama_client_defaults_to_local_without_key() { + use crate::api::ApiClient; env_remove!("OLLAMA_BASE_URL"); env_remove!("OLLAMA_API_KEY"); // Should still build — local Ollama doesn't need a real key. let client = ollama("llama3").unwrap(); - use crate::api::ApiClient; assert_eq!(client.model(), "llama3"); } #[cfg(feature = "openai")] #[test] fn self_hosted_client_builds() { - let client = self_hosted("http://localhost:8080/v1", "my-model").unwrap(); use crate::api::ApiClient; + let client = self_hosted("http://localhost:8080/v1", "my-model").unwrap(); assert_eq!(client.model(), "my-model"); } } diff --git a/src/provider/openai.rs b/src/provider/openai.rs index 376ba30..3dc2a42 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -1127,7 +1127,7 @@ mod tests { #[test] fn merge_tool_results_single() { let r = serde_json::json!({"role": "tool", "content": "ok"}); - let merged = merge_tool_results(&[r.clone()]); + let merged = merge_tool_results(std::slice::from_ref(&r)); assert_eq!(merged, r); } diff --git a/src/reflection/backoff.rs b/src/reflection/backoff.rs index 89dcbe7..31b7e53 100644 --- a/src/reflection/backoff.rs +++ b/src/reflection/backoff.rs @@ -130,12 +130,16 @@ impl RecoveryStrategy for ExponentialBackoffRecovery { &self, analysis: &FailureAnalysis, attempt: u32, - _max_attempts: u32, + max_attempts: u32, ) -> Pin + Send + '_>> { + // Honour both ceilings: the strategy-local retry limit and the + // framework-imposed `max_attempts` budget. Whichever is reached + // first causes us to give up. + let max_retries = self.max_retries.min(max_attempts); let action = if !analysis.is_recoverable { RecoveryAction::Fail(analysis.root_cause.clone()) - } else if attempt >= self.max_retries { - RecoveryAction::Fail(format!("max retries ({}) exceeded", self.max_retries)) + } else if attempt >= max_retries { + RecoveryAction::Fail(format!("max retries ({max_retries}) exceeded")) } else if analysis.severity >= FailureSeverity::High && analysis.correction.is_some() { RecoveryAction::AskUser(format!( "high-severity failure with correction available: {}", @@ -229,12 +233,35 @@ mod tests { }; let action = strategy.decide(&analysis, 3, 5).await; assert!(action.is_fail()); - let RecoveryAction::Fail(reason) = action else { - unreachable!() + let reason = match action { + RecoveryAction::Fail(r) => r, + _ => String::new(), }; assert!(reason.contains("max retries")); } + #[tokio::test] + async fn backoff_respects_framework_max_attempts() { + // strategy allows up to 10 retries, but the framework budget is 2. + // At attempt == max_attempts we must give up, ignoring the higher + // strategy-local limit. + let strategy = ExponentialBackoffRecovery::new(10); + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "timeout".to_string(), + severity: FailureSeverity::Medium, + correction: None, + context: String::new(), + }; + let action = strategy.decide(&analysis, 2, 2).await; + assert!(action.is_fail()); + let reason = match action { + RecoveryAction::Fail(r) => r, + _ => String::new(), + }; + assert!(reason.contains("max retries (2)")); + } + #[tokio::test] async fn backoff_unrecoverable_fails_immediately() { let strategy = ExponentialBackoffRecovery::new(3); @@ -247,8 +274,9 @@ mod tests { }; let action = strategy.decide(&analysis, 0, 5).await; assert!(action.is_fail()); - let RecoveryAction::Fail(reason) = action else { - unreachable!() + let reason = match action { + RecoveryAction::Fail(r) => r, + _ => String::new(), }; assert_eq!(reason, "invalid api key"); } diff --git a/src/runtime.rs b/src/runtime.rs index 8d62ac4..571b91c 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -864,7 +864,7 @@ mod tests { struct NopObserver; impl LoopObserver for NopObserver { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "NopObserver" } } @@ -892,7 +892,7 @@ mod tests { struct NopObserver; impl LoopObserver for NopObserver { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "NopObserver" } } diff --git a/src/stream/handler.rs b/src/stream/handler.rs index 7976246..c4a341e 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -1084,7 +1084,7 @@ mod tests { ..Default::default() }; // 1000 * 2^3 = 8000, capped at 5000 - assert_eq!(config.base_delay(3), Duration::from_millis(5000)); + assert_eq!(config.base_delay(3), Duration::from_secs(5)); } #[test] @@ -1655,7 +1655,7 @@ mod tests { } impl ApiClient for HandlerMock { - fn model(&self) -> &str { + fn model(&self) -> &'static str { "test-model" } @@ -1841,21 +1841,10 @@ mod tests { // (covered above). Here we test that stream_turn returns the // error when streaming fails and the handler is configured // without fallback. - let handler = StreamHandler::new().with_config( - StreamTimeoutConfig { - fallback_to_non_streaming: false, - ..Default::default() - }, - StreamRetryConfig { - max_retries: 0, - ..Default::default() - }, - ); - /// Mock that always returns an error stream. struct ErrorMock; impl ApiClient for ErrorMock { - fn model(&self) -> &str { + fn model(&self) -> &'static str { "test-model" } fn stream_messages( @@ -1886,6 +1875,17 @@ mod tests { } } + let handler = StreamHandler::new().with_config( + StreamTimeoutConfig { + fallback_to_non_streaming: false, + ..Default::default() + }, + StreamRetryConfig { + max_retries: 0, + ..Default::default() + }, + ); + let client = ErrorMock; let cancel = Arc::new(CancelSignal::new()); diff --git a/src/stream/heartbeat.rs b/src/stream/heartbeat.rs index 4b24be3..5efe242 100644 --- a/src/stream/heartbeat.rs +++ b/src/stream/heartbeat.rs @@ -426,7 +426,7 @@ mod tests { }; let mut stream = HeartbeatStream::new(inner, config); - stream.last_heartbeat = Instant::now() - Duration::from_secs(1); + stream.last_heartbeat = Instant::now().checked_sub(Duration::from_secs(1)).unwrap(); let waker = futures::task::noop_waker(); let mut cx = Context::from_waker(&waker); @@ -446,12 +446,11 @@ mod tests { // setting start into the past, then poll. let callbacks: std::sync::Arc>> = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); - let cb = callbacks.clone(); let config = HeartbeatConfig::new( Duration::from_millis(10), Duration::from_millis(1), Box::new(move |data: HeartbeatData| { - cb.lock().unwrap().push(data); + callbacks.lock().unwrap().push(data); }), ); @@ -460,7 +459,7 @@ mod tests { let mut stream = HeartbeatStream::new(inner, config); // Manually set start into the past so timeout has elapsed. - stream.start = Instant::now() - Duration::from_secs(10); + stream.start = Instant::now().checked_sub(Duration::from_secs(10)).unwrap(); // Use a no-op waker to poll manually. let waker = futures::task::noop_waker(); diff --git a/src/tool.rs b/src/tool.rs index df9cdfd..21229ce 100644 --- a/src/tool.rs +++ b/src/tool.rs @@ -1260,10 +1260,10 @@ mod tests { struct EchoTool; impl Tool for EchoTool { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "echo" } - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Echoes back the input" } fn schema(&self) -> ToolSchema { @@ -1300,10 +1300,10 @@ mod tests { struct FailTool; impl Tool for FailTool { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "fail" } - fn description(&self) -> &str { + fn description(&self) -> &'static str { "Always fails" } fn schema(&self) -> ToolSchema { diff --git a/src/tool/health.rs b/src/tool/health.rs index 3661a8b..78fe987 100644 --- a/src/tool/health.rs +++ b/src/tool/health.rs @@ -933,7 +933,7 @@ mod tests { assert_eq!(stats.total_calls(), 0); assert_eq!(stats.success_count(), 0); assert_eq!(stats.failure_count(), 0); - assert!(stats.success_rate() == 1.0); + assert!((stats.success_rate() - 1.0).abs() < f64::EPSILON); assert!(stats.health_score() > 0.9); assert_eq!(stats.avg_duration(), Duration::ZERO); assert_eq!(stats.max_duration(), Duration::ZERO); diff --git a/src/tool/shield.rs b/src/tool/shield.rs index 5278a4b..1116121 100644 --- a/src/tool/shield.rs +++ b/src/tool/shield.rs @@ -820,7 +820,8 @@ mod tests { #[test] fn null_shield_default_trait() { - let _shield = NullShield::default(); + let shield = NullShield; + let _ = &shield; } // =================================================== @@ -898,7 +899,10 @@ mod tests { let eval_ctx = ctx("Read", json!({ "path": "/tmp/data" }), 2); let combo = shield.assess_combination(&eval_ctx); - assert_eq!(combo, 0.0, "reversed order should not match"); + assert!( + (combo - 0.0).abs() < f32::EPSILON, + "reversed order should not match" + ); // Now test correct order: Write first, then Bash(chmod +x) as the // current call. From 11e6aa3b6c2a17c81c6086cf92b91d78662664a8 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 10:36:45 +1200 Subject: [PATCH 10/30] ix: pass user_input only on first turn, soften compaction failure, add config validation, remove dead recovery limit, clean up docs --- Cargo.toml | 2 +- README.md | 1 - src/api.rs | 6 +++++ src/config.rs | 44 ++++++++++++++++++++++++++++++++++++ src/engine/bare.rs | 45 ++++++++++++++++++++++++++++--------- src/engine/bare/dispatch.rs | 3 --- src/engine/loop_core.rs | 25 ++++++++++++++++----- 7 files changed, 106 insertions(+), 20 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 74c13bf..ccfd708 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,7 +8,7 @@ repository = "https://github.com/dch-labs/loopctl" homepage = "https://github.com/dch-labs/loopctl" documentation = "https://docs.rs/loopctl" keywords = ["agent", "framework", "llm", "loop"] -categories = ["agent", "development-tools"] +categories = ["api-bindings", "development-tools"] readme = "README.md" rust-version = "1.85" diff --git a/README.md b/README.md index b65d76e..3f229f7 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,6 @@ and tool implementations; the framework handles the rest. |--------|-------------| | [`api`](https://docs.rs/loopctl/latest/loopctl/api/index.html) | `ApiClient` trait for LLM provider communication (streaming + non-streaming) | | [`api::error`](https://docs.rs/loopctl/latest/loopctl/api/error/index.html) | API error types with retry classification | -| [`builder`](https://docs.rs/loopctl/latest/loopctl/builder/index.html) | Fluent builder API with type-state generics for compile-time safety | | [`cancel`](https://docs.rs/loopctl/latest/loopctl/cancel/index.html) | Cooperative cancellation via `CancelSignal` (AtomicBool + tokio::Notify) | | [`capabilities`](https://docs.rs/loopctl/latest/loopctl/capabilities/index.html) | Capability traits (`Observable`, `Detectable`, `Compactable`, etc.) | | [`compact`](https://docs.rs/loopctl/latest/loopctl/compact/index.html) | Context compaction: `ContextCompactor` trait, `TruncatingCompactor`, `TokenSplitter` | diff --git a/src/api.rs b/src/api.rs index 725b568..1024296 100644 --- a/src/api.rs +++ b/src/api.rs @@ -118,6 +118,9 @@ pub trait ApiClient: Send + Sync { /// # Parameters /// /// - `messages` — The conversation history as a [`Vec`]. + /// Takes ownership because the returned stream must be `'static`; + /// callers (e.g. [`BareLoop`](crate::engine::BareLoop)) clone the + /// full history each turn — O(n) in the number of messages. /// - `system` — An optional system prompt to prepend. /// - `tools` — Optional tool definitions the model may invoke. /// @@ -144,6 +147,9 @@ pub trait ApiClient: Send + Sync { /// # Parameters /// /// - `messages` — The conversation history as a [`Vec`]. + /// Takes ownership because the returned stream must be `'static`; + /// callers (e.g. [`BareLoop`](crate::engine::BareLoop)) clone the + /// full history each turn — O(n) in the number of messages. /// - `system` — An optional system prompt to prepend. /// - `tools` — Optional tool definitions the model may invoke. /// diff --git a/src/config.rs b/src/config.rs index f5605c2..656fccb 100644 --- a/src/config.rs +++ b/src/config.rs @@ -85,3 +85,47 @@ impl Default for LoopConfig { } } } + +impl LoopConfig { + /// Validate the configuration fields. + /// + /// Checks that: + /// - `compact_threshold` is in the range `[0.0, 1.0]` (not `NaN`). + /// - `max_turns` is greater than zero. + /// - `context_window` is greater than zero. + /// + /// # Errors + /// + /// Returns a [`String`] describing the first invalid field, or `Ok(())` + /// if all fields are valid. + /// + /// # Example + /// + /// ``` + /// use loopctl::config::LoopConfig; + /// + /// let config = LoopConfig::default(); + /// assert!(config.validate().is_ok()); + /// + /// let bad = LoopConfig { compact_threshold: 1.5, ..config }; + /// assert!(bad.validate().is_err()); + /// ``` + pub fn validate(&self) -> Result<(), String> { + if self.max_turns == 0 { + return Err("max_turns must be greater than 0".to_string()); + } + if self.context_window == 0 { + return Err("context_window must be greater than 0".to_string()); + } + if self.compact_threshold.is_nan() + || self.compact_threshold < 0.0 + || self.compact_threshold > 1.0 + { + return Err(format!( + "compact_threshold must be in [0.0, 1.0], got {}", + self.compact_threshold + )); + } + Ok(()) + } +} diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 5d8c290..e5b259d 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -240,6 +240,12 @@ pub struct BareLoop { impl BareLoop { /// Maximum retry attempts for tool recovery before giving up. + /// + /// This is the engine-level safety ceiling passed to the + /// [`RecoveryStrategy`](crate::reflection::RecoveryStrategy) as + /// `max_attempts`. The strategy's own `max_retries` limit (typically + /// stricter) is the effective limit; this constant prevents a + /// misconfigured strategy from retrying indefinitely. const MAX_RECOVERY_ATTEMPTS: u32 = 5; /// Create a new `BareLoop` with the given components. @@ -697,6 +703,8 @@ impl crate::engine::loop_core::Loop for BareLoop { config: &'a crate::config::LoopConfig, ) -> Pin> + Send + 'a>> { Box::pin(async move { + config.validate().map_err(LoopError::Config)?; + self.state = LoopState::Processing { turn: 0 }; self.budget = SessionResult::default(); self.session_start = Some(Instant::now()); @@ -713,8 +721,10 @@ impl crate::engine::loop_core::Loop for BareLoop { input: &'a str, ) -> Pin> + Send + 'a>> { Box::pin(async move { - // On the first turn, push the user's message. - if self.budget.total_turns == 0 { + // On the first turn, `input` is the user's message (non-empty). + // On continuation turns, `input` is "" — the conversation already + // has the tool results appended from the previous turn. + if !input.is_empty() { self.conversation.push(Message::user(input)); } @@ -748,6 +758,14 @@ impl crate::engine::loop_core::Loop for BareLoop { (msg, usage, stop) } Err(e) => { + // Record the failure with the fallback circuit breaker. + // + // Note: BareLoop records API failures and trips the + // circuit breaker but does **not** automatically retry + // with the fallback model. The `FallbackManager` is + // infrastructure for downstream consumers that hold + // multiple API clients. BareLoop has a single client, + // so the error is propagated after recording. let tripped = self.managers.fallback.record_api_failure(); if tripped { let from = self.client.model(); @@ -891,11 +909,18 @@ impl crate::engine::loop_core::Loop for BareLoop { self.budget = budget; // Attempt context compaction. + // + // Compaction is best-effort: if it fails (e.g. the compactor + // cannot reduce the conversation enough), we log a warning and + // continue. The next API call may still succeed, and if it + // doesn't, the provider's context-overflow error will surface + // naturally at that point. if let Err(e) = self.maybe_compact_context(self.budget.total_turns).await { - self.state = LoopState::Failed { - error: e.to_string(), - }; - return Err(e); + tracing::warn!( + error = %e, + turn = self.budget.total_turns, + "context compaction failed; continuing with uncompactd history" + ); } self.state = LoopState::Processing { @@ -1988,8 +2013,8 @@ mod tests { assert_eq!(result.total_turns, 1); } - /// Verify that setting `max_turns = 0` immediately triggers - /// [`LoopError::MaxTurnsExceeded`] before any API call. + /// Verify that setting `max_turns = 0` immediately triggers a + /// configuration error before any API call. #[tokio::test] async fn test_loop_terminates_with_max_turns_0() { let client = MockClient::new("test-model"); @@ -2002,8 +2027,8 @@ mod tests { let result = agent.run("Hi").await; assert!(result.is_err()); match result.unwrap_err() { - LoopError::MaxTurnsExceeded { max } => assert_eq!(max, 0), - other => panic!("Expected MaxTurnsExceeded, got: {other}"), + LoopError::Config(msg) => assert!(msg.contains("max_turns")), + other => panic!("Expected Config error, got: {other}"), } } diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 798158e..1d7213d 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -336,9 +336,6 @@ impl BareLoop { match recovery_action { RecoveryAction::Retry { delay } => { let next_attempt = attempt.saturating_add(1); - if next_attempt >= Self::MAX_RECOVERY_ATTEMPTS { - return Err(RecoveryOutcome::SoftError(tool_result.clone())); - } tokio::select! { () = tokio::time::sleep(delay) => {}, () = self.cancelled.notified() => { diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index 41e360f..250705f 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -584,10 +584,15 @@ pub trait Loop: Send + Sync { config: &'a LoopConfig, ) -> Pin> + Send + 'a>>; - /// Process a single user message / turn. - /// - /// Main entry point for agent logic. It receives the user's - /// input and returns a [`TurnResult`] describing what happened. + /// Process a single turn of the loop. + /// + /// Main entry point for loop logic. On the **first** call within a + /// [`run`](Loop::run) session, `input` contains the user's message; + /// subsequent calls receive an empty string (`""`) to signal a + /// continuation turn (tool results are already in the conversation + /// history). Implementations that need the original user message should + /// capture it during [`initialize`](Loop::initialize) or the first + /// `process_turn` call. fn process_turn<'a>( &'a mut self, input: &'a str, @@ -662,12 +667,22 @@ pub trait Loop: Send + Sync { Box::pin(async move { self.initialize(&self.config()).await?; + let mut is_first_turn = true; loop { if !self.should_continue() { break; } - match self.process_turn(user_input).await { + // Pass user_input only on the first turn; subsequent turns + // receive "" to signal continuation (tool results are + // already in the conversation history). + let input = if is_first_turn { + is_first_turn = false; + user_input + } else { + "" + }; + match self.process_turn(input).await { Ok(turn_result) if turn_result.is_complete => { return self.finalize().await; } From 6e2eacfd147e0645da173a0f0199efd32505d50c Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 11:01:05 +1200 Subject: [PATCH 11/30] feat: add providers timeouts --- src/provider/anthropic.rs | 88 ++++++++++++++++++++++++++++++++++++++ src/provider/gemini.rs | 88 ++++++++++++++++++++++++++++++++++++++ src/provider/openai.rs | 89 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 265 insertions(+) diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index 9f3ea6f..68de91d 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -24,6 +24,7 @@ use std::pin::Pin; use futures::stream::{Stream, StreamExt}; use reqwest::Response; use serde_json::Value; +use std::time::Duration; use crate::api::ApiClient; use crate::api::error::ApiError; @@ -45,6 +46,11 @@ const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; const DEFAULT_MAX_TOKENS: u32 = 8192; +/// Default total request timeout (connect + response + body). +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); +/// Default TCP connection establishment timeout. +const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + // ================================================== // Client // ================================================== @@ -208,6 +214,8 @@ pub struct AnthropicClientBuilder { base_url: String, model: String, max_tokens: u32, + timeout: Duration, + connect_timeout: Duration, } impl Default for AnthropicClientBuilder { @@ -217,6 +225,8 @@ impl Default for AnthropicClientBuilder { base_url: DEFAULT_BASE_URL.into(), model: DEFAULT_MODEL.into(), max_tokens: DEFAULT_MAX_TOKENS, + timeout: DEFAULT_REQUEST_TIMEOUT, + connect_timeout: DEFAULT_CONNECT_TIMEOUT, } } } @@ -252,6 +262,27 @@ impl AnthropicClientBuilder { self } + /// Set the total request timeout (connect + response + body). + /// + /// Defaults to 120 seconds. This bounds the entire HTTP request lifecycle — + /// a hanging server will be aborted after this duration rather than + /// blocking the agent loop indefinitely. + #[must_use] + pub fn timeout(mut self, timeout: Duration) -> Self { + self.timeout = timeout; + self + } + + /// Set the TCP connection establishment timeout. + /// + /// Defaults to 10 seconds. This is the maximum time to wait for the TCP + /// connection (including TLS handshake) to be established. + #[must_use] + pub fn connect_timeout(mut self, timeout: Duration) -> Self { + self.connect_timeout = timeout; + self + } + /// Build the client. /// /// # Errors @@ -262,6 +293,8 @@ impl AnthropicClientBuilder { .api_key .ok_or_else(|| ApiError::auth_invalid_key("API key not provided"))?; let http = reqwest::Client::builder() + .timeout(self.timeout) + .connect_timeout(self.connect_timeout) .build() .map_err(|e| ApiError::http(e.to_string()))?; @@ -1100,4 +1133,59 @@ mod tests { }; assert_eq!(reader.take_line().unwrap(), "data: hi"); } + + // ================================================ + // Builder timeout tests + // ================================================ + + #[test] + fn builder_has_default_timeouts() { + // The builder should initialize with sensible non-zero defaults + // so that a hanging server cannot block the agent loop indefinitely. + let builder = AnthropicClientBuilder::default(); + assert_eq!(builder.timeout, DEFAULT_REQUEST_TIMEOUT); + assert_eq!(builder.connect_timeout, DEFAULT_CONNECT_TIMEOUT); + } + + #[test] + fn builder_custom_timeout() { + let custom = Duration::from_secs(300); + let builder = AnthropicClientBuilder::default().timeout(custom); + assert_eq!(builder.timeout, custom); + // connect_timeout should be unchanged. + assert_eq!(builder.connect_timeout, DEFAULT_CONNECT_TIMEOUT); + } + + #[test] + fn builder_custom_connect_timeout() { + let custom = Duration::from_secs(45); + let builder = AnthropicClientBuilder::default().connect_timeout(custom); + assert_eq!(builder.connect_timeout, custom); + // timeout should be unchanged. + assert_eq!(builder.timeout, DEFAULT_REQUEST_TIMEOUT); + } + + #[test] + fn builder_custom_both_timeouts() { + let req_timeout = Duration::from_secs(600); + let conn_timeout = Duration::from_secs(30); + let builder = AnthropicClientBuilder::default() + .timeout(req_timeout) + .connect_timeout(conn_timeout); + assert_eq!(builder.timeout, req_timeout); + assert_eq!(builder.connect_timeout, conn_timeout); + } + + #[test] + fn builder_timeouts_applied_on_build() { + // Verify the build succeeds — reqwest validates the configuration + // internally. If timeout/connect_timeout were somehow invalid, + // .build() would return an error. + let client = AnthropicClient::builder() + .api_key("sk-test") + .timeout(Duration::from_secs(180)) + .connect_timeout(Duration::from_secs(15)) + .build(); + assert!(client.is_ok(), "build should succeed with valid timeouts"); + } } diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index b8ef4ce..055726b 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -25,6 +25,7 @@ use std::pin::Pin; use futures::stream::{Stream, StreamExt}; use reqwest::Response; use serde_json::Value; +use std::time::Duration; use crate::api::ApiClient; use crate::api::error::ApiError; @@ -44,6 +45,11 @@ const DEFAULT_MODEL: &str = "gemini-2.0-flash"; const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; +/// Default total request timeout (connect + response + body). +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); +/// Default TCP connection establishment timeout. +const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + // ================================================== // Client // ================================================== @@ -200,6 +206,8 @@ pub struct GeminiClientBuilder { api_key: Option, base_url: String, model: String, + timeout: Duration, + connect_timeout: Duration, } impl Default for GeminiClientBuilder { @@ -208,6 +216,8 @@ impl Default for GeminiClientBuilder { api_key: None, base_url: DEFAULT_BASE_URL.into(), model: DEFAULT_MODEL.into(), + timeout: DEFAULT_REQUEST_TIMEOUT, + connect_timeout: DEFAULT_CONNECT_TIMEOUT, } } } @@ -234,6 +244,27 @@ impl GeminiClientBuilder { self } + /// Set the total request timeout (connect + response + body). + /// + /// Defaults to 120 seconds. This bounds the entire HTTP request lifecycle — + /// a hanging server will be aborted after this duration rather than + /// blocking the agent loop indefinitely. + #[must_use] + pub fn timeout(mut self, timeout: Duration) -> Self { + self.timeout = timeout; + self + } + + /// Set the TCP connection establishment timeout. + /// + /// Defaults to 10 seconds. This is the maximum time to wait for the TCP + /// connection (including TLS handshake) to be established. + #[must_use] + pub fn connect_timeout(mut self, timeout: Duration) -> Self { + self.connect_timeout = timeout; + self + } + /// Build the client. /// /// # Errors @@ -244,6 +275,8 @@ impl GeminiClientBuilder { .api_key .ok_or_else(|| ApiError::auth_invalid_key("API key not provided"))?; let http = reqwest::Client::builder() + .timeout(self.timeout) + .connect_timeout(self.connect_timeout) .build() .map_err(|e| ApiError::http(e.to_string()))?; @@ -906,4 +939,59 @@ mod tests { }; assert_eq!(reader.take_line().unwrap(), "data: hi"); } + + // ================================================ + // Builder timeout tests + // ================================================ + + #[test] + fn builder_has_default_timeouts() { + // The builder should initialize with sensible non-zero defaults + // so that a hanging server cannot block the agent loop indefinitely. + let builder = GeminiClientBuilder::default(); + assert_eq!(builder.timeout, DEFAULT_REQUEST_TIMEOUT); + assert_eq!(builder.connect_timeout, DEFAULT_CONNECT_TIMEOUT); + } + + #[test] + fn builder_custom_timeout() { + let custom = Duration::from_secs(300); + let builder = GeminiClientBuilder::default().timeout(custom); + assert_eq!(builder.timeout, custom); + // connect_timeout should be unchanged. + assert_eq!(builder.connect_timeout, DEFAULT_CONNECT_TIMEOUT); + } + + #[test] + fn builder_custom_connect_timeout() { + let custom = Duration::from_secs(45); + let builder = GeminiClientBuilder::default().connect_timeout(custom); + assert_eq!(builder.connect_timeout, custom); + // timeout should be unchanged. + assert_eq!(builder.timeout, DEFAULT_REQUEST_TIMEOUT); + } + + #[test] + fn builder_custom_both_timeouts() { + let req_timeout = Duration::from_secs(600); + let conn_timeout = Duration::from_secs(30); + let builder = GeminiClientBuilder::default() + .timeout(req_timeout) + .connect_timeout(conn_timeout); + assert_eq!(builder.timeout, req_timeout); + assert_eq!(builder.connect_timeout, conn_timeout); + } + + #[test] + fn builder_timeouts_applied_on_build() { + // Verify the build succeeds — reqwest validates the configuration + // internally. If timeout/connect_timeout were somehow invalid, + // .build() would return an error. + let client = GeminiClient::builder() + .api_key("test-key") + .timeout(Duration::from_secs(180)) + .connect_timeout(Duration::from_secs(15)) + .build(); + assert!(client.is_ok(), "build should succeed with valid timeouts"); + } } diff --git a/src/provider/openai.rs b/src/provider/openai.rs index 3dc2a42..aae31f1 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -23,6 +23,8 @@ use std::future::Future; use std::pin::Pin; +use std::time::Duration; + use futures::stream::{Stream, StreamExt}; use reqwest::Response; use serde::Deserialize; @@ -47,6 +49,11 @@ const SSE_DONE: &str = "[DONE]"; const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; +/// Default total request timeout (connect + response + body). +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); +/// Default TCP connection establishment timeout. +const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + // ================================================== // Client // ================================================== @@ -208,6 +215,8 @@ pub struct OpenAiClientBuilder { api_key: Option, base_url: String, model: String, + timeout: Duration, + connect_timeout: Duration, } impl Default for OpenAiClientBuilder { @@ -216,6 +225,8 @@ impl Default for OpenAiClientBuilder { api_key: None, base_url: DEFAULT_BASE_URL.into(), model: DEFAULT_MODEL.into(), + timeout: DEFAULT_REQUEST_TIMEOUT, + connect_timeout: DEFAULT_CONNECT_TIMEOUT, } } } @@ -242,6 +253,27 @@ impl OpenAiClientBuilder { self } + /// Set the total request timeout (connect + response + body). + /// + /// Defaults to 120 seconds. This bounds the entire HTTP request lifecycle — + /// a hanging server will be aborted after this duration rather than + /// blocking the agent loop indefinitely. + #[must_use] + pub fn timeout(mut self, timeout: Duration) -> Self { + self.timeout = timeout; + self + } + + /// Set the TCP connection establishment timeout. + /// + /// Defaults to 10 seconds. This is the maximum time to wait for the TCP + /// connection (including TLS handshake) to be established. + #[must_use] + pub fn connect_timeout(mut self, timeout: Duration) -> Self { + self.connect_timeout = timeout; + self + } + /// Build the client. /// /// # Errors @@ -253,6 +285,8 @@ impl OpenAiClientBuilder { .ok_or_else(|| ApiError::auth_invalid_key("API key not provided"))?; let http = reqwest::Client::builder() + .timeout(self.timeout) + .connect_timeout(self.connect_timeout) .build() .map_err(|e| ApiError::http(e.to_string()))?; @@ -1181,4 +1215,59 @@ mod tests { let line = reader.take_line().unwrap(); assert_eq!(line, "data: hi"); } + + // ================================================ + // Builder timeout tests + // ================================================ + + #[test] + fn builder_has_default_timeouts() { + // The builder should initialize with sensible non-zero defaults + // so that a hanging server cannot block the agent loop indefinitely. + let builder = OpenAiClientBuilder::default(); + assert_eq!(builder.timeout, DEFAULT_REQUEST_TIMEOUT); + assert_eq!(builder.connect_timeout, DEFAULT_CONNECT_TIMEOUT); + } + + #[test] + fn builder_custom_timeout() { + let custom = Duration::from_secs(300); + let builder = OpenAiClientBuilder::default().timeout(custom); + assert_eq!(builder.timeout, custom); + // connect_timeout should be unchanged. + assert_eq!(builder.connect_timeout, DEFAULT_CONNECT_TIMEOUT); + } + + #[test] + fn builder_custom_connect_timeout() { + let custom = Duration::from_secs(45); + let builder = OpenAiClientBuilder::default().connect_timeout(custom); + assert_eq!(builder.connect_timeout, custom); + // timeout should be unchanged. + assert_eq!(builder.timeout, DEFAULT_REQUEST_TIMEOUT); + } + + #[test] + fn builder_custom_both_timeouts() { + let req_timeout = Duration::from_secs(600); + let conn_timeout = Duration::from_secs(30); + let builder = OpenAiClientBuilder::default() + .timeout(req_timeout) + .connect_timeout(conn_timeout); + assert_eq!(builder.timeout, req_timeout); + assert_eq!(builder.connect_timeout, conn_timeout); + } + + #[test] + fn builder_timeouts_applied_on_build() { + // Verify the build succeeds — reqwest validates the configuration + // internally. If timeout/connect_timeout were somehow invalid, + // .build() would return an error. + let client = OpenAiClient::builder() + .api_key("sk-test") + .timeout(Duration::from_secs(180)) + .connect_timeout(Duration::from_secs(15)) + .build(); + assert!(client.is_ok(), "build should succeed with valid timeouts"); + } } From 459cbd102e71b1695acc3f6815d37b2a0d387c1f Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 11:05:18 +1200 Subject: [PATCH 12/30] fix: validate max tokens --- src/config.rs | 60 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/src/config.rs b/src/config.rs index 656fccb..0d9c40b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -93,6 +93,8 @@ impl LoopConfig { /// - `compact_threshold` is in the range `[0.0, 1.0]` (not `NaN`). /// - `max_turns` is greater than zero. /// - `context_window` is greater than zero. + /// - `max_tokens` is greater than zero. + /// - `model` is not empty. /// /// # Errors /// @@ -117,6 +119,12 @@ impl LoopConfig { if self.context_window == 0 { return Err("context_window must be greater than 0".to_string()); } + if self.max_tokens == 0 { + return Err("max_tokens must be greater than 0".to_string()); + } + if self.model.is_empty() { + return Err("model must not be empty".to_string()); + } if self.compact_threshold.is_nan() || self.compact_threshold < 0.0 || self.compact_threshold > 1.0 @@ -129,3 +137,55 @@ impl LoopConfig { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn validate_default_config_is_ok() { + let config = LoopConfig::default(); + assert!(config.validate().is_ok()); + } + + #[test] + fn validate_rejects_zero_max_tokens() { + let config = LoopConfig { + max_tokens: 0, + ..LoopConfig::default() + }; + let err = config.validate().unwrap_err(); + assert!( + err.contains("max_tokens"), + "error should mention max_tokens: {err}" + ); + } + + #[test] + fn validate_accepts_one_max_tokens() { + let config = LoopConfig { + max_tokens: 1, + ..LoopConfig::default() + }; + assert!(config.validate().is_ok()); + } + + #[test] + fn validate_rejects_empty_model() { + let config = LoopConfig { + model: String::new(), + ..LoopConfig::default() + }; + let err = config.validate().unwrap_err(); + assert!(err.contains("model"), "error should mention model: {err}"); + } + + #[test] + fn validate_accepts_nonempty_model() { + let config = LoopConfig { + model: "gpt-4".to_string(), + ..LoopConfig::default() + }; + assert!(config.validate().is_ok()); + } +} From 60cf49e4c4dcd018c34722f231f44e3198dc7cf8 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 12:28:27 +1200 Subject: [PATCH 13/30] fix: output limit middleware --- src/middleware.rs | 60 ++++++++++++++++++++++++++++++---- src/middleware/output_limit.rs | 30 ++++++++++++----- 2 files changed, 75 insertions(+), 15 deletions(-) diff --git a/src/middleware.rs b/src/middleware.rs index 6e13833..e592a34 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -1120,7 +1120,7 @@ mod tests { } #[tokio::test] - async fn test_output_limit_multipart_passes_through() { + async fn test_output_limit_multipart_text_parts_are_truncated() { let registry = long_output_registry(); let pipeline = ToolPipeline::builder() .with(OutputLimitMiddleware::new(3)) @@ -1130,12 +1130,58 @@ mod tests { let result = pipeline.invoke(test_ctx("multipart")).await; assert!(!result.is_error); - // Multipart content should pass through even though it exceeds - // the character limit — truncation only applies to Text variant. - assert!( - matches!(result.output, ToolContent::Multipart(_)), - "multipart should pass through unchanged" - ); + // Each text part in the multipart result should be individually + // truncated. "part1" (5 chars > 3) → "par\n[truncated]". + match result.output { + ToolContent::Multipart(parts) => { + assert_eq!(parts.len(), 2, "should still have 2 parts"); + for (i, part) in parts.iter().enumerate() { + match part { + ToolContentPart::Text { text } => { + assert!( + text.contains("[truncated]"), + "part {i} should be truncated: got {text:?}" + ); + assert!( + text.starts_with("par"), + "part {i} should start with first 3 chars: got {text:?}" + ); + } + ToolContentPart::Image { .. } => { + panic!("unexpected image part in MultipartTool output"); + } + } + } + } + other => panic!("expected Multipart, got {other:?}"), + } + } + + #[tokio::test] + async fn test_output_limit_multipart_short_parts_pass_through() { + // When all text parts are within the limit, multipart should + // pass through unchanged. + let registry = long_output_registry(); + let pipeline = ToolPipeline::builder() + .with(OutputLimitMiddleware::new(100)) + .core(Arc::clone(®istry)) + .build() + .expect("valid"); + + let result = pipeline.invoke(test_ctx("multipart")).await; + assert!(!result.is_error); + match result.output { + ToolContent::Multipart(parts) => { + assert_eq!(parts.len(), 2); + if let ToolContentPart::Text { text } = &parts[0] { + assert_eq!(text, "part1", "short part should not be truncated"); + } + if let ToolContentPart::Text { text } = &parts[1] { + assert_eq!(text, "part2", "short part should not be truncated"); + } + } + other => panic!("expected Multipart, got {other:?}"), + } } #[tokio::test] diff --git a/src/middleware/output_limit.rs b/src/middleware/output_limit.rs index 9e778e4..251aed0 100644 --- a/src/middleware/output_limit.rs +++ b/src/middleware/output_limit.rs @@ -1,15 +1,16 @@ //! Middleware that truncates tool output to a maximum character count. use super::{ToolDispatchContext, ToolDispatchResult, ToolMiddleware, ToolPipeline}; -use crate::message::ToolContent; +use crate::message::{ToolContent, ToolContentPart}; use std::future::Future; use std::pin::Pin; /// Middleware that truncates tool output to a maximum character count. /// /// If the tool's text output exceeds the limit, it is truncated and -/// suffixed with a `[truncated]` marker. Non-text outputs ([`ToolContent::Multipart`]) -/// are passed through unchanged. +/// suffixed with a `[truncated]` marker. For [`ToolContent::Multipart`] +/// results, each text part is individually truncated in the same way; +/// image parts are left unchanged. /// /// This prevents runaway tools from flooding the conversation with /// excessive output that would blow the context window. @@ -53,11 +54,24 @@ impl ToolMiddleware for OutputLimitMiddleware { Box::pin(async move { let mut result = next.dispatch(ctx).await; - if let ToolContent::Text(ref text) = result.output { - let char_count = text.chars().count(); - if char_count > max_chars { - let truncated: String = text.chars().take(max_chars).collect(); - result.output = ToolContent::Text(format!("{truncated}\n[truncated]")); + match result.output { + ToolContent::Text(ref text) => { + let char_count = text.chars().count(); + if char_count > max_chars { + let truncated: String = text.chars().take(max_chars).collect(); + result.output = ToolContent::Text(format!("{truncated}\n[truncated]")); + } + } + ToolContent::Multipart(ref mut parts) => { + for part in parts.iter_mut() { + if let ToolContentPart::Text { text } = part { + let char_count = text.chars().count(); + if char_count > max_chars { + let truncated: String = text.chars().take(max_chars).collect(); + *text = format!("{truncated}\n[truncated]"); + } + } + } } } From d04b2cc4c57f5adf8d4bb880af583d08eb2d135c Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 12:39:43 +1200 Subject: [PATCH 14/30] fix: orhpaned call results in truncating compactor --- src/compact/truncating.rs | 309 +++++++++++++++++++++++++++++++++++++- 1 file changed, 307 insertions(+), 2 deletions(-) diff --git a/src/compact/truncating.rs b/src/compact/truncating.rs index 74641b5..f956b14 100644 --- a/src/compact/truncating.rs +++ b/src/compact/truncating.rs @@ -8,7 +8,8 @@ use crate::compact::types::{CompactionContext, CompactionOutcome}; use crate::compact::{ContextCompactor, ContextManager}; -use crate::message::{Message, Role}; +use crate::message::{Message, MessagePart, Role}; +use std::collections::HashSet; use std::future::Future; use std::pin::Pin; @@ -127,7 +128,14 @@ impl ContextCompactor for TruncatingCompactor { } // Determine split point: keep `preserve_recent` from the end. - let split = total.saturating_sub(self.preserve_recent); + let initial_split = total.saturating_sub(self.preserve_recent); + + // Adjust split to avoid orphaning tool-call/result pairs. + // If the "recent" portion contains a ToolResult whose matching + // ToolCall would be dropped, move the split back to include the + // message containing that ToolCall. + let split = Self::adjust_for_tool_pairs(&messages, initial_split); + let recent: Vec = messages.get(split..).unwrap_or_default().to_vec(); // Always preserve the first message (typically the system prompt) @@ -157,6 +165,76 @@ impl ContextCompactor for TruncatingCompactor { } } +// =================================================== +// Tool-call/result pair protection +// ================================================== + +impl TruncatingCompactor { + /// Adjust the split index to avoid orphaning tool-call/result pairs. + /// + /// If the "recent" portion (from `split` onward) contains any + /// [`MessagePart::ToolResult`] whose matching + /// [`MessagePart::ToolCall`] (identified by `call_id` == `id`) would + /// be in the dropped portion (before `split`), the split is moved + /// backward to include the message containing the orphaned call. + fn adjust_for_tool_pairs(messages: &[Message], split: usize) -> usize { + if split == 0 { + return 0; + } + + // Collect call IDs from the recent portion — those whose calls + // are already preserved need no adjustment. + let recent = messages.get(split..).unwrap_or_default(); + let recent_call_ids: HashSet<&String> = recent + .iter() + .flat_map(|msg| msg.parts.iter()) + .filter_map(|part| match part { + MessagePart::ToolCall { id, .. } => Some(id), + _ => None, + }) + .collect(); + + // Check each result in the recent portion: if its call_id is not + // among the recent calls, the call is in the dropped portion. + let orphaned_ids: Vec<&String> = recent + .iter() + .flat_map(|msg| msg.parts.iter()) + .filter_map(|part| match part { + MessagePart::ToolResult { call_id, .. } => { + if recent_call_ids.contains(call_id) { + None + } else { + Some(call_id) + } + } + _ => None, + }) + .collect(); + + if orphaned_ids.is_empty() { + return split; + } + + // Walk backward from the split point to find the earliest message + // that contains a ToolCall matching any orphaned ID. + let mut new_split = split; + for i in (0..split).rev() { + let Some(msg) = messages.get(i) else { + continue; + }; + let has_orphaned_call = msg.parts.iter().any(|part| match part { + MessagePart::ToolCall { id, .. } => orphaned_ids.contains(&id), + _ => false, + }); + if has_orphaned_call { + new_split = i; + } + } + + new_split + } +} + // =================================================== // TokenSplitter // =================================================== @@ -322,3 +400,230 @@ impl Default for TokenSplitter { Self::new() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::compact::ContextCompactor; + use crate::compact::types::{CompactReason, CompactionContext}; + use crate::message::{Message, MessagePart, Role, ToolContent}; + use serde_json::json; + + fn tool_text(s: &str) -> ToolContent { + ToolContent::from_string(s) + } + + fn make_context(msgs: &[Message]) -> CompactionContext { + CompactionContext { + tokens_before: CompactionOutcome::estimate_tokens(msgs), + reason: CompactReason::ThresholdExceeded, + context_window: 1_000, + turn: 5, + } + } + + /// Build a conversation where a tool-call and its result straddle + /// what would be the naive split point. + /// + /// Layout (index: content): + /// + /// ```text + /// 0 user "msg0" + /// 1 assistant "reply0" + /// 2 user "msg1" + /// 3 assistant "reply1" + /// 4 user "msg2" + /// 5 assistant tool_call("call_a", "search", ...) + /// 6 user tool_result("call_a", ...) + /// 7 assistant "final reply" + /// ``` + /// + /// With `preserve_recent = 2`, the naive split would be at index 6, + /// dropping the tool-call (index 5) but keeping the result (index 6). + /// The fix should move the split back to index 5. + fn convo_with_straddling_tool_pair() -> Vec { + vec![ + Message::user("msg0"), + Message::assistant("reply0"), + Message::user("msg1"), + Message::assistant("reply1"), + Message::user("msg2"), + Message::new( + Role::Assistant, + vec![MessagePart::tool_call( + "call_a", + "search", + json!({"q": "rust"}), + )], + ), + Message::new( + Role::User, + vec![MessagePart::tool_result( + "call_a", + tool_text("result data"), + false, + )], + ), + Message::assistant("final reply"), + ] + } + + fn has_tool_call(msgs: &[Message], id: &str) -> bool { + msgs.iter() + .flat_map(|m| m.parts.iter()) + .any(|p| matches!(p, MessagePart::ToolCall { id: tool_id, .. } if tool_id == id)) + } + + fn has_tool_result(msgs: &[Message], call_id: &str) -> bool { + msgs.iter() + .flat_map(|m| m.parts.iter()) + .any(|p| matches!(p, MessagePart::ToolResult { call_id: cid, .. } if cid == call_id)) + } + + #[tokio::test] + async fn compact_preserves_tool_call_when_result_is_in_recent() { + let messages = convo_with_straddling_tool_pair(); + let compactor = TruncatingCompactor::new() + .with_preserve_recent(2) + .with_min_messages(4); + let context = make_context(&messages); + let outcome = compactor.compact(messages, 500, context).await; + + // The tool-call ("call_a") and tool-result ("call_a") must both + // be in the compacted output — neither should be orphaned. + assert!( + has_tool_call(&outcome.messages, "call_a"), + "tool-call 'call_a' must be preserved" + ); + assert!( + has_tool_result(&outcome.messages, "call_a"), + "tool-result for 'call_a' must be preserved" + ); + } + + #[tokio::test] + async fn compact_does_not_orphan_when_pairs_are_together_in_recent() { + // When both call and result are already in the recent portion, + // no adjustment is needed — the split should stay at the naive point. + let messages = vec![ + Message::user("msg0"), + Message::assistant("reply0"), + Message::user("msg1"), + Message::assistant("reply1"), + Message::user("msg2"), + Message::assistant("reply2"), + Message::new( + Role::Assistant, + vec![MessagePart::tool_call("call_b", "calc", json!({}))], + ), + Message::new( + Role::User, + vec![MessagePart::tool_result("call_b", tool_text("42"), false)], + ), + ]; + + let compactor = TruncatingCompactor::new() + .with_preserve_recent(2) + .with_min_messages(4); + let context = make_context(&messages); + let outcome = compactor.compact(messages, 500, context).await; + + assert!( + has_tool_call(&outcome.messages, "call_b"), + "tool-call 'call_b' must be preserved" + ); + assert!( + has_tool_result(&outcome.messages, "call_b"), + "tool-result for 'call_b' must be preserved" + ); + } + + #[tokio::test] + async fn compact_drops_both_call_and_result_when_in_old_portion() { + // When both call and result are entirely in the old (dropped) + // portion, the split should NOT be adjusted — both are dropped + // together, which is correct. + let messages = vec![ + Message::user("msg0"), + Message::new( + Role::Assistant, + vec![MessagePart::tool_call("call_c", "tool", json!({}))], + ), + Message::new( + Role::User, + vec![MessagePart::tool_result("call_c", tool_text("done"), false)], + ), + Message::assistant("reply1"), + Message::user("msg2"), + Message::assistant("reply2"), + Message::user("msg3"), + Message::assistant("reply3"), + ]; + + let compactor = TruncatingCompactor::new() + .with_preserve_recent(4) + .with_min_messages(4); + let context = make_context(&messages); + let outcome = compactor.compact(messages, 500, context).await; + + // Neither call_c nor its result should appear — both dropped. + assert!( + !has_tool_call(&outcome.messages, "call_c"), + "tool-call 'call_c' should be dropped" + ); + assert!( + !has_tool_result(&outcome.messages, "call_c"), + "tool-result for 'call_c' should be dropped" + ); + } + + #[test] + fn adjust_for_tool_pairs_returns_zero_when_split_is_zero() { + let messages = convo_with_straddling_tool_pair(); + assert_eq!(TruncatingCompactor::adjust_for_tool_pairs(&messages, 0), 0); + } + + #[test] + fn adjust_for_tool_pairs_no_orphans_returns_original_split() { + // No tool results in the recent portion → no adjustment. + let messages = vec![ + Message::user("a"), + Message::assistant("b"), + Message::user("c"), + Message::assistant("d"), + Message::user("e"), + Message::assistant("f"), + ]; + assert_eq!(TruncatingCompactor::adjust_for_tool_pairs(&messages, 4), 4); + } + + #[test] + fn adjust_for_tool_pairs_moves_split_back_for_orphaned_result() { + let messages = convo_with_straddling_tool_pair(); + // Naive split at index 6 would keep result (idx 6) but drop call (idx 5). + // Should adjust back to 5. + assert_eq!(TruncatingCompactor::adjust_for_tool_pairs(&messages, 6), 5); + } + + #[tokio::test] + async fn compact_short_conversation_returns_unchanged() { + // Below min_messages, the conversation should pass through unchanged. + let messages = vec![ + Message::user("hello"), + Message::new( + Role::Assistant, + vec![MessagePart::tool_call("call_d", "tool", json!({}))], + ), + Message::new( + Role::User, + vec![MessagePart::tool_result("call_d", tool_text("ok"), false)], + ), + ]; + let compactor = TruncatingCompactor::new() + .with_preserve_recent(2) + .with_min_messages(6); + let context = make_context(&messages); + let outcome = compactor.compact(messages.clone(), 500, context).await; + assert_eq!(outcome.messages.len(), messages.len()); + } +} From 36ff98e2b4b9be2bd78b954eb56c8177d7788eec Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 12:56:34 +1200 Subject: [PATCH 15/30] chore: max reponse size + timeout for llm providers --- src/middleware.rs | 4 ++-- src/provider/anthropic.rs | 45 +++++++++++++++++++++++++++++++------- src/provider/gemini.rs | 30 ++++++++++++++++++------- src/provider/openai.rs | 46 +++++++++++++++++++++++++++++++++------ 4 files changed, 100 insertions(+), 25 deletions(-) diff --git a/src/middleware.rs b/src/middleware.rs index e592a34..4dc2804 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -1153,7 +1153,7 @@ mod tests { } } } - other => panic!("expected Multipart, got {other:?}"), + other @ ToolContent::Text(_) => panic!("expected Multipart, got {other:?}"), } } @@ -1180,7 +1180,7 @@ mod tests { assert_eq!(text, "part2", "short part should not be truncated"); } } - other => panic!("expected Multipart, got {other:?}"), + other @ ToolContent::Text(_) => panic!("expected Multipart, got {other:?}"), } } diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index 68de91d..a908cd1 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -45,10 +45,8 @@ const SSE_EVENT_PREFIX: &str = "event: "; const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; const DEFAULT_MAX_TOKENS: u32 = 8192; - -/// Default total request timeout (connect + response + body). -const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); -/// Default TCP connection establishment timeout. +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); /// connect + response + body +const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -197,9 +195,18 @@ impl ApiClient for AnthropicClient { let url = self.messages_url(); Box::pin(async move { let resp = Self::post_messages(&self.http, &url, &self.api_key, &body).await?; - resp.json::() + let resp = resp + .bytes() .await - .map_err(|e| ApiError::http(e.to_string())) + .map_err(|e| ApiError::http(e.to_string()))?; + if resp.len() > MAX_RESPONSE_BODY { + return Err(ApiError::http(format!( + "response body too large: {} bytes (max {})", + resp.len(), + MAX_RESPONSE_BODY + ))); + } + serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } } @@ -474,7 +481,18 @@ impl SseReader { let parsed = if data.is_empty() { None } else { - serde_json::from_str(&data).ok() + match serde_json::from_str(&data) { + Ok(v) => Some(v), + Err(e) => { + tracing::warn!( + error = %e, + event_type = %event_type, + data_len = data.len(), + "failed to parse Anthropic SSE data, skipping" + ); + None + } + } }; return Ok(Some((event_type, parsed))); } @@ -499,7 +517,18 @@ impl SseReader { let parsed = if data.is_empty() { None } else { - serde_json::from_str(&data).ok() + match serde_json::from_str(&data) { + Ok(v) => Some(v), + Err(e) => { + tracing::warn!( + error = %e, + event_type = %event_type, + data_len = data.len(), + "failed to parse Anthropic SSE data (stream end), skipping" + ); + None + } + } }; return Ok(Some((event_type, parsed))); } diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 055726b..6270d1b 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -44,10 +44,8 @@ const DEFAULT_BASE_URL: &str = "https://generativelanguage.googleapis.com/v1beta const DEFAULT_MODEL: &str = "gemini-2.0-flash"; const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; - -/// Default total request timeout (connect + response + body). -const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); -/// Default TCP connection establishment timeout. +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body +const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -190,9 +188,18 @@ impl ApiClient for GeminiClient { Box::pin(async move { let resp = Self::post_content(&self.http, &url, &body).await?; - resp.json::() + let resp = resp + .bytes() .await - .map_err(|e| ApiError::http(e.to_string())) + .map_err(|e| ApiError::http(e.to_string()))?; + if resp.len() > MAX_RESPONSE_BODY { + return Err(ApiError::http(format!( + "response body too large: {} bytes (max {})", + resp.len(), + MAX_RESPONSE_BODY + ))); + } + serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } } @@ -416,8 +423,15 @@ impl SseReader { } if let Some(data) = line.strip_prefix(SSE_DATA_PREFIX) { - if let Ok(json) = serde_json::from_str::(data) { - return Ok(Some(json)); + match serde_json::from_str::(data) { + Ok(json) => return Ok(Some(json)), + Err(e) => { + tracing::warn!( + error = %e, + data_len = data.len(), + "failed to parse Gemini SSE data, skipping" + ); + } } } } diff --git a/src/provider/openai.rs b/src/provider/openai.rs index aae31f1..1dc557a 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -48,10 +48,8 @@ const DEFAULT_MODEL: &str = "gpt-4o"; const SSE_DONE: &str = "[DONE]"; const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; - -/// Default total request timeout (connect + response + body). -const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); -/// Default TCP connection establishment timeout. +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body +const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -199,9 +197,18 @@ impl ApiClient for OpenAiClient { let resp = Self::post_completions(&self.http, &url, &self.api_key, &body.to_json(false)) .await?; - resp.json::() + let resp = resp + .bytes() .await - .map_err(|e| ApiError::http(e.to_string())) + .map_err(|e| ApiError::http(e.to_string()))?; + if resp.len() > MAX_RESPONSE_BODY { + return Err(ApiError::http(format!( + "response body too large: {} bytes (max {})", + resp.len(), + MAX_RESPONSE_BODY + ))); + } + serde_json::from_slice::(&resp).map_err(|e| ApiError::http(e.to_string())) }) } } @@ -547,7 +554,17 @@ impl OpenAiChunk { /// Returns `None` for malformed payloads so the caller can skip /// them without interrupting the stream. fn parse(data: &str) -> Option { - serde_json::from_str(data).ok() + match serde_json::from_str(data) { + Ok(chunk) => Some(chunk), + Err(e) => { + tracing::warn!( + error = %e, + data_len = data.len(), + "failed to parse OpenAI SSE chunk, skipping" + ); + None + } + } } } @@ -915,6 +932,21 @@ mod tests { assert!(OpenAiChunk::parse("").is_none()); } + #[test] + fn parse_malformed_partial_json_returns_none() { + // Truncated JSON should also fail gracefully with a warning log. + assert!(OpenAiChunk::parse(r#"{"id":"chatcmpl-1","choices":[{"delta":{"con"#).is_none()); + } + + #[test] + fn parse_valid_chunk_with_all_fields() { + let data = r#"{"id":"abc","model":"gpt-4o","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}],"usage":null}"#; + let chunk = OpenAiChunk::parse(data).unwrap(); + assert_eq!(chunk.id, "abc"); + assert_eq!(chunk.model, "gpt-4o"); + assert_eq!(chunk.choices.len(), 1); + } + #[test] fn emitter_emits_message_start_on_first_chunk() { let mut em = StreamEmitter::default(); From 52510f0d6204b237b785baced86acabe8bc30124 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 13:05:21 +1200 Subject: [PATCH 16/30] fix: remove gemini api key from url --- src/provider/anthropic.rs | 2 +- src/provider/gemini.rs | 50 ++++++++++++++++++++++++++++++++------- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index a908cd1..e8aed1a 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -45,7 +45,7 @@ const SSE_EVENT_PREFIX: &str = "event: "; const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; const DEFAULT_MAX_TOKENS: u32 = 8192; -const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); /// connect + response + body +const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 6270d1b..fcabffd 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -101,17 +101,14 @@ impl GeminiClient { /// parameter rather than using headers. fn stream_url(&self) -> String { format!( - "{}/models/{}:streamGenerateContent?alt=sse&key={}", - self.base_url, self.model, self.api_key + "{}/models/{}:streamGenerateContent?alt=sse", + self.base_url, self.model ) } /// Build the non-streaming Generate Content URL. fn generate_url(&self) -> String { - format!( - "{}/models/{}:generateContent?key={}", - self.base_url, self.model, self.api_key - ) + format!("{}/models/{}:generateContent", self.base_url, self.model) } /// Send a POST request and return the raw response. @@ -126,10 +123,12 @@ impl GeminiClient { async fn post_content( http: &reqwest::Client, url: &str, + api_key: &str, body: &Value, ) -> Result { let resp = http .post(url) + .header("x-goog-api-key", api_key) .json(body) .send() .await @@ -158,9 +157,10 @@ impl ApiClient for GeminiClient { let body = build_request_body(&messages, system.as_deref(), tools.as_deref()); let url = self.stream_url(); let http = self.http.clone(); + let api_key = self.api_key.clone(); Box::pin(async_stream::try_stream! { - let resp = Self::post_content(&http, &url, &body).await?; + let resp = Self::post_content(&http, &url, &api_key, &body).await?; let mut sse = SseReader::from_response(resp); let mut emitter = StreamEmitter::default(); @@ -187,7 +187,7 @@ impl ApiClient for GeminiClient { let url = self.generate_url(); Box::pin(async move { - let resp = Self::post_content(&self.http, &url, &body).await?; + let resp = Self::post_content(&self.http, &url, &self.api_key, &body).await?; let resp = resp .bytes() .await @@ -793,6 +793,40 @@ mod tests { assert_eq!(client.model(), "gemini-pro"); } + #[test] + fn stream_url_does_not_expose_api_key() { + let client = GeminiClient::builder() + .api_key("secret-key-123") + .build() + .unwrap(); + let url = client.stream_url(); + assert!( + !url.contains("secret-key-123"), + "API key must not appear in stream URL: {url}" + ); + assert!( + !url.contains("key="), + "URL must not have key= query param: {url}" + ); + } + + #[test] + fn generate_url_does_not_expose_api_key() { + let client = GeminiClient::builder() + .api_key("secret-key-456") + .build() + .unwrap(); + let url = client.generate_url(); + assert!( + !url.contains("secret-key-456"), + "API key must not appear in generate URL: {url}" + ); + assert!( + !url.contains("key="), + "URL must not have key= query param: {url}" + ); + } + #[test] fn emitter_first_chunk_emits_message_start() { let mut em = StreamEmitter::default(); From 354e6a0e15aab0059085bb0c8b08b6caf5d16bc9 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 14:34:42 +1200 Subject: [PATCH 17/30] chore: non_exhaustive enums fix, missing docs --- Makefile | 7 +- src/api/error.rs | 2 + src/capabilities.rs | 8 ++ src/config.rs | 1 + src/detection/convergence.rs | 29 ++++- src/engine/loop_core.rs | 7 +- src/fallback.rs | 232 ++++++++++++++++++++++++----------- src/lib.rs | 1 + src/memory/builtin.rs | 10 ++ src/middleware.rs | 28 +++++ src/middleware/permission.rs | 14 ++- src/provider.rs | 12 +- src/provider/anthropic.rs | 83 ++++++++++++- src/provider/gemini.rs | 85 ++++++++++++- src/provider/openai.rs | 91 +++++++++++++- src/stream.rs | 2 + src/stream/handler.rs | 2 + src/tool/registry.rs | 62 ++++++++++ 18 files changed, 593 insertions(+), 83 deletions(-) diff --git a/Makefile b/Makefile index ea0ebdb..41be7e1 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ -.PHONY: check test clippy fmt docs ci lint +.PHONY: check test clippy fmt docs ci lint examples -ci: fmt check clippy test docs +ci: fmt check clippy test docs examples check: cargo check --all-features @@ -20,3 +20,6 @@ lint: docs: RUSTDOCFLAGS="-D warnings" cargo doc --no-deps --all-features + +examples: + cargo build --examples --all-features diff --git a/src/api/error.rs b/src/api/error.rs index 8cf1315..2dedf37 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -76,6 +76,7 @@ use thiserror::Error; /// ``` #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize_repr, Deserialize_repr)] #[repr(u16)] +#[non_exhaustive] pub enum ErrorCode { // ================================================== // API errors (1000-1005) @@ -345,6 +346,7 @@ pub enum ErrorCode { /// variant documentation for details on when each is produced and how /// the message text influences code selection. #[derive(Debug, Error)] +#[non_exhaustive] pub enum ApiError { /// An error returned by the LLM API provider. /// diff --git a/src/capabilities.rs b/src/capabilities.rs index 62b7e66..93cf58f 100644 --- a/src/capabilities.rs +++ b/src/capabilities.rs @@ -69,6 +69,7 @@ use crate::tool::health::ToolHealthRegistry; /// } /// ``` pub trait Observable { + /// Returns the observer host for lifecycle event fan-out. fn observers(&self) -> &ObserverHost; } @@ -97,6 +98,7 @@ pub trait Observable { /// } /// ``` pub trait Detectable { + /// Returns the detection manager for loop and convergence detection. fn detection(&self) -> &DetectionManager; } @@ -126,6 +128,7 @@ pub trait Detectable { /// } /// ``` pub trait FallbackCapable { + /// Returns the fallback manager (circuit breaker) for API model fallback. fn fallback(&self) -> &FallbackManager; } @@ -146,6 +149,7 @@ pub trait FallbackCapable { /// compaction during the agent loop. Useful for custom loop /// implementations that need to manage the context window directly. pub trait Compactable { + /// Returns the context manager, if compaction is configured. fn context_manager(&self) -> Option<&Arc>; } @@ -167,6 +171,7 @@ pub trait Compactable { /// that need to control streaming behaviour (timeouts, retries, fallback /// to non-streaming mode). pub trait StreamCapable { + /// Returns the stream handler, if resilient streaming is configured. fn stream_handler(&self) -> Option<&StreamHandler>; } @@ -187,6 +192,7 @@ pub trait StreamCapable { /// tool dispatch, compaction, or session start/end. #[cfg(feature = "hooks")] pub trait Hookable { + /// Returns the hook executor, if hooks are configured. fn hook_executor(&self) -> Option<&HookExecutor>; } @@ -214,6 +220,7 @@ pub trait Hookable { /// } /// ``` pub trait PipelineAware { + /// Returns the tool middleware pipeline, if configured. fn pipeline(&self) -> Option<&ToolPipeline>; } @@ -230,5 +237,6 @@ pub trait PipelineAware { /// - [`LoopRuntime`](crate::runtime::LoopRuntime) — the framework's default implementation. #[cfg(feature = "tool_health")] pub trait HealthTrackable { + /// Returns the tool health registry, if health tracking is configured. fn health_registry(&self) -> Option<&ToolHealthRegistry>; } diff --git a/src/config.rs b/src/config.rs index 0d9c40b..c4ce527 100644 --- a/src/config.rs +++ b/src/config.rs @@ -112,6 +112,7 @@ impl LoopConfig { /// let bad = LoopConfig { compact_threshold: 1.5, ..config }; /// assert!(bad.validate().is_err()); /// ``` + #[must_use = "validation errors should not be silently ignored"] pub fn validate(&self) -> Result<(), String> { if self.max_turns == 0 { return Err("max_turns must be greater than 0".to_string()); diff --git a/src/detection/convergence.rs b/src/detection/convergence.rs index 236ecad..1908d47 100644 --- a/src/detection/convergence.rs +++ b/src/detection/convergence.rs @@ -61,6 +61,23 @@ //! assert!(status.detected); // 3 consecutive similar responses //! # Ok::<(), loopctl::detection::convergence::ConvergenceConfigError>(()) //! ``` +//! +//! # Known Limitations +//! +//! This detector uses **Jaccard similarity** on whitespace-tokenized words: +//! +//! - **Semantic blindness**: Responses that are semantically identical but +//! use different vocabulary will not be detected as convergent. +//! - **Word-order sensitivity**: Rearranging words may reduce similarity +//! below the threshold, even when the meaning is unchanged. +//! - **Punctuation/whitespace sensitivity**: Tokenization is purely +//! whitespace-based, so minor formatting changes can affect results. +//! - **False positives**: Boilerplate-heavy responses with shared prefixes +//! (e.g., "Let me check that for you...") may trigger false positives. +//! +//! Choose a conservative `threshold` (see [`ConvergenceConfig`]) to +//! minimise false positives, and consider combining with +//! [`LoopDetector`](super::LoopDetector) for complementary pattern detection. use std::collections::HashSet; use std::collections::VecDeque; @@ -103,10 +120,12 @@ use serde::{Deserialize, Serialize}; /// ConvergenceAction::SwitchPhase => println!("Switching phase"), /// ConvergenceAction::AskUser => println!("Asking user"), /// ConvergenceAction::Compact => println!("Compacting history"), +/// _ => println!("Other action"), /// } /// ``` #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] #[serde(rename_all = "snake_case")] +#[non_exhaustive] pub enum ConvergenceAction { /// Stop the agent loop entirely. /// @@ -193,14 +212,20 @@ pub enum ConvergenceConfigError { /// Convergence requires at least one pair of consecutive responses, /// so a window of 1 (or 0) is meaningless. #[error("window_size must be at least 2, got {actual}")] - WindowTooSmall { actual: usize }, + WindowTooSmall { + /// The invalid window size that was provided. + actual: usize, + }, /// `similarity_threshold` is outside the valid range `[0.0, 1.0]`. /// /// Jaccard similarity always produces a value in this range; a /// threshold outside it would never (or always) trigger. #[error("similarity_threshold must be in [0.0, 1.0], got {actual}")] - ThresholdOutOfRange { actual: f32 }, + ThresholdOutOfRange { + /// The invalid threshold value that was provided. + actual: f32, + }, } // =================================================== diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index 250705f..580dd71 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -73,6 +73,7 @@ pub use crate::tool::ToolDispatchResult; /// ↘ Reflecting ↗ /// ``` #[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] pub enum LoopState { /// The agent is idle, waiting for a user message. /// @@ -148,7 +149,10 @@ pub enum LoopState { /// [`SessionResult::error`]. /// /// No further turns will be executed after entering this state. - Failed { error: String }, + Failed { + /// A human-readable description of the error that caused the failure. + error: String, + }, } // ================================================== @@ -281,6 +285,7 @@ impl TurnResult { /// to determine the next step: dispatch tools, continue the conversation, /// or end the session. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] pub enum StopReason { /// The model decided to stop (natural end of turn). /// diff --git a/src/fallback.rs b/src/fallback.rs index ee063b5..2753e3e 100644 --- a/src/fallback.rs +++ b/src/fallback.rs @@ -617,6 +617,28 @@ impl Default for FallbackConfig { /// mgr.record_model_success(); // → back to Primary /// } /// ``` +/// Consolidated mutex-protected fallback state. +/// +/// All mutable fallback bookkeeping lives behind a single lock so that +/// related fields are always observed together, preventing partial-state +/// reads that could occur when acquiring separate locks sequentially. +#[derive(Default)] +struct FallbackInner { + /// Original model name (before fallback). + original_model: Option, + /// Ordered fallback models with failure status. + fallback_models: Vec, + /// Cached first non-failed fallback model name. + active_fallback: Option, + /// Time when fallback was activated. + fallback_switched_at: Option, +} + +/// Circuit breaker for API model fallback. +/// +/// `&FallbackManager` is `Send + Sync` and can be freely shared across +/// threads (e.g. via `Arc`). No `&mut self` is needed +/// for any operation. pub struct FallbackManager { /// Failures before switching to fallback. fallback_threshold: usize, @@ -632,14 +654,12 @@ pub struct FallbackManager { fallback_state: AtomicU8, /// Consecutive successes on primary during recovery. primary_success_count: AtomicUsize, - /// Original model name (before fallback). - original_model: Mutex>, - /// Ordered fallback models with failure status. - fallback_models: Mutex>, - /// Cached first non-failed fallback model name. - active_fallback: Mutex>, - /// Time when fallback was activated. - fallback_switched_at: Mutex>, + /// Consolidated mutex-protected fallback state. + /// + /// Holding all related fields behind a single lock prevents partial-state + /// reads that could occur when acquiring the (formerly separate) locks one + /// at a time. + inner: Mutex, /// How long to remain in fallback before attempting primary recovery. recovery_timeout: Duration, } @@ -678,10 +698,7 @@ impl FallbackManager { fallback_activated: AtomicBool::new(false), fallback_state: AtomicU8::new(FallbackState::Primary as u8), primary_success_count: AtomicUsize::new(0), - original_model: Mutex::new(None), - fallback_models: Mutex::new(Vec::new()), - active_fallback: Mutex::new(None), - fallback_switched_at: Mutex::new(None), + inner: Mutex::new(FallbackInner::default()), recovery_timeout: Duration::from_secs(60), } } @@ -775,16 +792,14 @@ impl FallbackManager { #[must_use] pub fn new_with_fallback(original_model: String, fallback_threshold: usize) -> Self { let mgr = Self::new(fallback_threshold, 2); - if let Ok(mut m) = mgr.original_model.lock() { - *m = Some(original_model); + if let Ok(mut inner) = mgr.inner.lock() { + inner.original_model = Some(original_model); + inner.fallback_switched_at = Some(Instant::now()); } mgr.fallback_activated.store(true, Ordering::Relaxed); mgr.consecutive_failures.store(0, Ordering::Relaxed); mgr.fallback_state .store(FallbackState::Fallback as u8, Ordering::Relaxed); - if let Ok(mut t) = mgr.fallback_switched_at.lock() { - *t = Some(Instant::now()); - } mgr } @@ -806,8 +821,8 @@ impl FallbackManager { /// ``` pub fn for_model(primary_model: impl Into) -> Self { let mgr = Self::new(3, 2); - if let Ok(mut m) = mgr.original_model.lock() { - *m = Some(primary_model.into()); + if let Ok(mut inner) = mgr.inner.lock() { + inner.original_model = Some(primary_model.into()); } mgr } @@ -902,7 +917,10 @@ impl FallbackManager { /// assert_eq!(mgr.original_model(), Some("llm-70b".to_string())); /// ``` pub fn original_model(&self) -> Option { - self.original_model.lock().ok().and_then(|m| m.clone()) + self.inner + .lock() + .ok() + .and_then(|i| i.original_model.clone()) } /// Set the original model name. @@ -920,8 +938,8 @@ impl FallbackManager { /// assert_eq!(mgr.original_model(), Some("llm-70b".to_string())); /// ``` pub fn set_original_model(&self, model: String) { - if let Ok(mut m) = self.original_model.lock() { - *m = Some(model); + if let Ok(mut inner) = self.inner.lock() { + inner.original_model = Some(model); } } @@ -944,7 +962,7 @@ impl FallbackManager { /// assert!(mgr.fallback_switched_at().is_none()); /// ``` pub fn fallback_switched_at(&self) -> Option { - self.fallback_switched_at.lock().ok().and_then(|t| *t) + self.inner.lock().ok().and_then(|i| i.fallback_switched_at) } /// Get the model that should be used for the next request. @@ -1013,9 +1031,11 @@ impl FallbackManager { /// assert_eq!(mgr.active_model(), Some("llm-4".to_string())); /// ``` pub fn set_fallback_model(&self, model: impl Into) { - if let Ok(mut m) = self.fallback_models.lock() { - m.clear(); - m.push(FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count)); + if let Ok(mut inner) = self.inner.lock() { + inner.fallback_models.clear(); + inner + .fallback_models + .push(FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count)); } self.recompute_active_fallback(); } @@ -1039,7 +1059,10 @@ impl FallbackManager { /// assert_eq!(mgr.fallback_model(), Some("llm-70b".to_string())); // first in chain /// ``` pub fn fallback_model(&self) -> Option { - self.active_fallback.lock().ok().and_then(|m| m.clone()) + self.inner + .lock() + .ok() + .and_then(|i| i.active_fallback.clone()) } /// Get the full fallback model chain. @@ -1061,10 +1084,10 @@ impl FallbackManager { /// assert_eq!(chain, vec!["llm-70b", "llm-120b", "llm-32b"]); /// ``` pub fn fallback_models(&self) -> Vec { - self.fallback_models + self.inner .lock() .ok() - .map(|m| m.iter().map(|e| e.name.clone()).collect()) + .map(|i| i.fallback_models.iter().map(|e| e.name.clone()).collect()) .unwrap_or_default() } @@ -1091,8 +1114,10 @@ impl FallbackManager { /// assert_eq!(chain, vec!["llm-70b", "llm-120b"]); /// ``` pub fn add_fallback_model(&self, model: impl Into) { - if let Ok(mut m) = self.fallback_models.lock() { - m.push(FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count)); + if let Ok(mut inner) = self.inner.lock() { + inner + .fallback_models + .push(FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count)); } self.recompute_active_fallback(); } @@ -1121,12 +1146,12 @@ impl FallbackManager { /// assert_eq!(chain, vec!["llm-70b", "llm-120b", "llm-32b"]); /// ``` pub fn insert_fallback_model(&self, index: usize, model: impl Into) { - if let Ok(mut m) = self.fallback_models.lock() { + if let Ok(mut inner) = self.inner.lock() { let entry = FallbackEntry::new(model).with_max_fail_count(self.default_max_fail_count); - if index >= m.len() { - m.push(entry); + if index >= inner.fallback_models.len() { + inner.fallback_models.push(entry); } else { - m.insert(index, entry); + inner.fallback_models.insert(index, entry); } } self.recompute_active_fallback(); @@ -1154,9 +1179,9 @@ impl FallbackManager { /// assert_eq!(chain, vec!["llm-120b"]); /// ``` pub fn remove_fallback_model(&self, model: &str) -> bool { - let removed = if let Ok(mut m) = self.fallback_models.lock() { - if let Some(pos) = m.iter().position(|x| x.name == model) { - m.remove(pos); + let removed = if let Ok(mut inner) = self.inner.lock() { + if let Some(pos) = inner.fallback_models.iter().position(|x| x.name == model) { + inner.fallback_models.remove(pos); true } else { false @@ -1192,8 +1217,8 @@ impl FallbackManager { /// ``` pub fn set_fallback_models(&self, models: Vec) { let max_fc = self.default_max_fail_count; - if let Ok(mut m) = self.fallback_models.lock() { - *m = models + if let Ok(mut inner) = self.inner.lock() { + inner.fallback_models = models .into_iter() .map(|name| FallbackEntry::new(name).with_max_fail_count(max_fc)) .collect(); @@ -1232,8 +1257,8 @@ impl FallbackManager { /// assert_eq!(mgr.fallback_model(), Some("llm-3".to_string())); /// ``` pub fn mark_fallback_failed(&self, model: &str) -> bool { - let found = if let Ok(mut m) = self.fallback_models.lock() { - if let Some(entry) = m.iter_mut().find(|e| e.name == model) { + let found = if let Ok(mut inner) = self.inner.lock() { + if let Some(entry) = inner.fallback_models.iter_mut().find(|e| e.name == model) { entry.record_attempt("marked_failed"); true } else { @@ -1269,8 +1294,8 @@ impl FallbackManager { /// assert!(mgr.failed_fallbacks().is_empty()); /// ``` pub fn clear_fallback_failed(&self, model: &str) -> bool { - let found = if let Ok(mut m) = self.fallback_models.lock() { - if let Some(entry) = m.iter_mut().find(|e| e.name == model) { + let found = if let Ok(mut inner) = self.inner.lock() { + if let Some(entry) = inner.fallback_models.iter_mut().find(|e| e.name == model) { entry.clear_attempts(); true } else { @@ -1308,8 +1333,8 @@ impl FallbackManager { /// assert!(mgr.failed_fallbacks().is_empty()); /// ``` pub fn clear_all_fallback_failed(&self) { - if let Ok(mut m) = self.fallback_models.lock() { - for entry in m.iter_mut() { + if let Ok(mut inner) = self.inner.lock() { + for entry in &mut inner.fallback_models { entry.clear_attempts(); } } @@ -1336,11 +1361,12 @@ impl FallbackManager { /// assert_eq!(failed, vec!["llm-3"]); /// ``` pub fn failed_fallbacks(&self) -> Vec { - self.fallback_models + self.inner .lock() .ok() - .map(|m| { - m.iter() + .map(|i| { + i.fallback_models + .iter() .filter(|e| e.failed()) .map(|e| e.name.clone()) .collect() @@ -1368,11 +1394,12 @@ impl FallbackManager { /// assert_eq!(available, vec!["llm-2"]); /// ``` pub fn available_fallbacks(&self) -> Vec { - self.fallback_models + self.inner .lock() .ok() - .map(|m| { - m.iter() + .map(|i| { + i.fallback_models + .iter() .filter(|e| !e.failed()) .map(|e| e.name.clone()) .collect() @@ -1405,10 +1432,10 @@ impl FallbackManager { /// assert!(mgr.fallback_entry("nonexistent").is_none()); /// ``` pub fn fallback_entry(&self, name: &str) -> Option { - self.fallback_models + self.inner .lock() .ok() - .and_then(|m| m.iter().find(|e| e.name == name).cloned()) + .and_then(|i| i.fallback_models.iter().find(|e| e.name == name).cloned()) } /// Set the [`available`](FallbackEntry::available) flag on a fallback model. @@ -1441,8 +1468,8 @@ impl FallbackManager { /// assert_eq!(mgr.fallback_model(), Some("llm-2".to_string())); /// ``` pub fn set_fallback_available(&self, model: &str, available: bool) -> bool { - let found = if let Ok(mut m) = self.fallback_models.lock() { - if let Some(entry) = m.iter_mut().find(|e| e.name == model) { + let found = if let Ok(mut inner) = self.inner.lock() { + if let Some(entry) = inner.fallback_models.iter_mut().find(|e| e.name == model) { entry.set_available(available); true } else { @@ -1722,8 +1749,8 @@ impl FallbackManager { self.fallback_state .store(FallbackState::Fallback as u8, Ordering::Relaxed); self.fallback_activated.store(true, Ordering::Relaxed); - if let Ok(mut t) = self.fallback_switched_at.lock() { - *t = Some(Instant::now()); + if let Ok(mut inner) = self.inner.lock() { + inner.fallback_switched_at = Some(Instant::now()); } self.primary_success_count.store(0, Ordering::Relaxed); info!("Circuit breaker: transitioned to Fallback state"); @@ -1782,8 +1809,8 @@ impl FallbackManager { pub fn transition_to_primary(&self) { self.fallback_state .store(FallbackState::Primary as u8, Ordering::Relaxed); - if let Ok(mut t) = self.fallback_switched_at.lock() { - *t = None; + if let Ok(mut inner) = self.inner.lock() { + inner.fallback_switched_at = None; } self.primary_success_count.store(0, Ordering::Relaxed); self.consecutive_failures.store(0, Ordering::Relaxed); @@ -1823,8 +1850,8 @@ impl FallbackManager { self.consecutive_failures.store(0, Ordering::Relaxed); self.primary_success_count.store(0, Ordering::Relaxed); self.fallback_activated.store(false, Ordering::Relaxed); - if let Ok(mut t) = self.fallback_switched_at.lock() { - *t = None; + if let Ok(mut inner) = self.inner.lock() { + inner.fallback_switched_at = None; } self.clear_all_fallback_failed(); } @@ -1839,13 +1866,14 @@ impl FallbackManager { /// [`fallback_models`]: Self::fallback_models /// [`active_fallback`]: Self::active_fallback fn recompute_active_fallback(&self) { - let active = self - .fallback_models - .lock() - .ok() - .and_then(|m| m.iter().find(|e| !e.failed()).map(|e| e.name.clone())); - if let Ok(mut cached) = self.active_fallback.lock() { - *cached = active; + let active = self.inner.lock().ok().and_then(|i| { + i.fallback_models + .iter() + .find(|e| !e.failed()) + .map(|e| e.name.clone()) + }); + if let Ok(mut inner) = self.inner.lock() { + inner.active_fallback = active; } } } @@ -2084,4 +2112,70 @@ mod tests { h.join().unwrap(); } } + + /// Verify the consolidated Mutex (M5 fix) ensures multi-field updates + /// are visible atomically. When `transition_to_fallback` is called, + /// both `fallback_switched_at` and `active_fallback` should be + /// observable together. + #[test] + fn test_consolidated_mutex_fields_are_consistent() { + let mgr = FallbackManager::for_model("primary-model"); + mgr.add_fallback_model("fallback-model"); + + // Before transition: using primary model, no switch time. + assert_eq!(mgr.active_model(), Some("primary-model".to_string())); + assert!(mgr.fallback_switched_at().is_none()); + + // Transition to fallback — updates multiple fields. + mgr.transition_to_fallback(); + + // After transition: both fields should be set together. + // This verifies the consolidated Mutex prevents partial reads. + assert_eq!(mgr.active_model(), Some("fallback-model".to_string())); + assert!( + mgr.fallback_switched_at().is_some(), + "switch time should be set after transition" + ); + } + + /// Verify that `transition_to_primary` clears both fields atomically (M5). + #[test] + fn test_consolidated_mutex_clears_fields_together() { + let mgr = FallbackManager::for_model("primary-model"); + mgr.add_fallback_model("fallback-model"); + mgr.transition_to_fallback(); + + // Both fields are set. + assert_eq!(mgr.active_model(), Some("fallback-model".to_string())); + assert!(mgr.fallback_switched_at().is_some()); + + // Transition back to primary. + mgr.transition_to_primary(); + + // Both fields should be cleared together. + assert!( + mgr.fallback_switched_at().is_none(), + "switch time should be cleared after transition to primary" + ); + } + + /// Verify that `reset()` clears all fields atomically (M5). + #[test] + fn test_consolidated_mutex_reset_clears_all() { + let mgr = FallbackManager::for_model("primary-model"); + mgr.add_fallback_model("fallback-model"); + mgr.transition_to_fallback(); + mgr.record_failure(); + + // State is dirty. + assert!(mgr.consecutive_failures() > 0); + assert!(mgr.fallback_switched_at().is_some()); + + // Full reset. + mgr.reset(); + + // Everything cleared. + assert_eq!(mgr.consecutive_failures(), 0); + assert!(mgr.fallback_switched_at().is_none()); + } } diff --git a/src/lib.rs b/src/lib.rs index 85ec12a..15d7707 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -45,6 +45,7 @@ // Relax strict lints in test code. The crate enforces a strict no-panic / // no-unwrap policy in production code, but test code legitimately uses // assertions, unwrap, indexing, etc. for readability. +#![warn(missing_docs)] #![cfg_attr( test, allow( diff --git a/src/memory/builtin.rs b/src/memory/builtin.rs index be2aefc..add3dd9 100644 --- a/src/memory/builtin.rs +++ b/src/memory/builtin.rs @@ -107,6 +107,16 @@ use std::pin::Pin; /// assert_eq!(results.len(), 1); /// # }); /// ``` +/// +/// # Unbounded Growth +/// +/// `InMemoryStore` accumulates entries in a `Vec` with no automatic +/// eviction. The [`consolidate()`](InMemoryStore::consolidate) method +/// prunes entries with `relevance < 0.05`, but it must be called +/// explicitly. A long-running session that never calls `consolidate()` +/// will accumulate memory indefinitely. For production use, consider +/// calling `consolidate()` periodically or implementing a custom +/// [`LoopMemory`] with bounded capacity. pub struct InMemoryStore { entries: Vec, } diff --git a/src/middleware.rs b/src/middleware.rs index 4dc2804..a03f951 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -702,6 +702,34 @@ mod tests { assert!(result.is_error); } + /// Verify that PermissionMiddleware with Ask permission and no resolver + /// denies the tool call (M2 fix). + #[tokio::test] + async fn test_permission_ask_without_resolver_denies() { + let mut ctx = test_ctx("echo"); + ctx.permission = PermissionCheck::Ask { + prompt: "Allow echo?".to_string(), + }; + + // from_context() has no ask_resolver configured. + let pipeline = ToolPipeline::builder() + .with(PermissionMiddleware::from_context()) + .core(test_registry()) + .build() + .expect("valid"); + + let result = pipeline.invoke(ctx).await; + assert!( + result.is_error, + "Ask without resolver should deny the tool call" + ); + let msg = result.output.to_string(); + assert!( + msg.contains("permission"), + "error should mention permission: {msg}" + ); + } + // ================================================== // TimeoutMiddleware tests // ================================================== diff --git a/src/middleware/permission.rs b/src/middleware/permission.rs index cd98a80..95f1b7c 100644 --- a/src/middleware/permission.rs +++ b/src/middleware/permission.rs @@ -145,8 +145,8 @@ impl ToolMiddleware for PermissionMiddleware { next.dispatch(ctx) } PermissionCheck::Deny { reason } => Self::deny(ctx, &reason), - PermissionCheck::Ask { prompt } => match &self.ask_resolver { - Some(resolver) => { + PermissionCheck::Ask { prompt } => { + if let Some(resolver) = &self.ask_resolver { let resolver = Arc::clone(resolver); Box::pin(async move { let tool_name = ctx.tool_name.clone(); @@ -161,9 +161,15 @@ impl ToolMiddleware for PermissionMiddleware { ) } }) + } else { + tracing::warn!( + tool = %ctx.tool_name, + prompt = %prompt, + "permission Ask denied: no resolver configured" + ); + Self::deny(ctx, &format!("permission required: {prompt}")) } - None => Self::deny(ctx, &format!("permission required: {prompt}")), - }, + } } } } diff --git a/src/provider.rs b/src/provider.rs index ec967d2..19001fe 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -308,12 +308,20 @@ mod tests { /// Helper to safely set an env var in tests (Rust 2024 requires unsafe). macro_rules! env_set { - ($($arg:tt)*) => {{ unsafe { std::env::set_var($($arg)*) } }}; + ($($arg:tt)*) => {{ + // SAFETY: This is only used in single-threaded test code where + // no other task is reading or writing environment variables. + unsafe { std::env::set_var($($arg)*) } + }}; } /// Helper to safely remove an env var in tests. macro_rules! env_remove { - ($($arg:tt)*) => {{ unsafe { std::env::remove_var($($arg)*) } }}; + ($($arg:tt)*) => {{ + // SAFETY: This is only used in single-threaded test code where + // no other task is reading or writing environment variables. + unsafe { std::env::remove_var($($arg)*) } + }}; } #[test] diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index e8aed1a..859fd50 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -47,6 +47,7 @@ const TEXT_PART_INDEX: usize = 0; const DEFAULT_MAX_TOKENS: u32 = 8192; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb +const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -509,7 +510,14 @@ impl SseReader { } match self.bytes.next().await { - Some(Ok(chunk)) => self.buf.push_str(&chunk), + Some(Ok(chunk)) => { + self.buf.push_str(&chunk); + if self.buf.len() > SSE_MAX_BUFFER { + return Err(ApiError::http(format!( + "SSE buffer exceeded {SSE_MAX_BUFFER} bytes" + ))); + } + } Some(Err(e)) => return Err(e), None => { // End of stream — emit any pending event. @@ -1217,4 +1225,77 @@ mod tests { .build(); assert!(client.is_ok(), "build should succeed with valid timeouts"); } + + // ================================================== + // SSE buffer cap tests (M1) + // ================================================== + + #[tokio::test] + async fn sse_reader_take_line_splits_on_newline() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "event: message_start\ndata: {}\n\n".to_string(), + }; + assert_eq!(reader.take_line(), Some("event: message_start".to_string())); + assert_eq!(reader.take_line(), Some("data: {}".to_string())); + assert_eq!(reader.take_line(), Some(String::new())); + } + + #[tokio::test] + async fn sse_reader_next_event_extracts_payload() { + let chunk = "event: content_block_delta\ndata: {\"type\":\"text_delta\"}\n\n"; + let stream = futures::stream::iter(vec![Ok::(chunk.to_string())]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_event().await.unwrap(); + assert!(result.is_some()); + let (event_type, data) = result.unwrap(); + assert_eq!(event_type, "content_block_delta"); + assert!(data.is_some()); + } + + #[tokio::test] + async fn sse_reader_next_event_malformed_data_returns_none_value() { + // Malformed JSON data should be logged and returned as None for the + // data payload, but the event_type is still captured (H4). + let chunk = "event: ping\ndata: not valid json\n\n"; + let stream = futures::stream::iter(vec![Ok::(chunk.to_string())]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_event().await.unwrap(); + assert!(result.is_some()); + let (event_type, data) = result.unwrap(); + assert_eq!(event_type, "ping"); + assert!(data.is_none(), "malformed JSON should yield None data"); + } + + #[tokio::test] + async fn sse_reader_buffer_overflow_returns_error() { + let huge = "x".repeat(SSE_MAX_BUFFER + 1); + let stream = futures::stream::iter(vec![Ok::(huge)]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_event().await; + assert!(result.is_err(), "should error on buffer overflow"); + let err_msg = result.unwrap_err().to_string(); + assert!( + err_msg.contains("SSE buffer"), + "error should mention SSE buffer: {err_msg}" + ); + } + + // ================================================== + // Body size limit tests (H5) + // ================================================== + + #[test] + fn max_response_body_is_ten_mb() { + assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); + } } diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index fcabffd..489c26d 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -46,6 +46,7 @@ const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb +const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -437,7 +438,14 @@ impl SseReader { } match self.bytes.next().await { - Some(Ok(chunk)) => self.buf.push_str(&chunk), + Some(Ok(chunk)) => { + self.buf.push_str(&chunk); + if self.buf.len() > SSE_MAX_BUFFER { + return Err(ApiError::http(format!( + "SSE buffer exceeded {SSE_MAX_BUFFER} bytes" + ))); + } + } Some(Err(e)) => return Err(e), None => return Ok(None), } @@ -1042,4 +1050,79 @@ mod tests { .build(); assert!(client.is_ok(), "build should succeed with valid timeouts"); } + + // ================================================== + // SSE buffer cap tests (M1) + // ================================================== + + #[tokio::test] + async fn sse_reader_take_line_splits_on_newline() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "data: {\"candidates\":[]}\n\n".to_string(), + }; + assert_eq!( + reader.take_line(), + Some("data: {\"candidates\":[]}".to_string()) + ); + assert_eq!(reader.take_line(), Some(String::new())); + assert_eq!(reader.take_line(), None); + } + + #[tokio::test] + async fn sse_reader_next_data_extracts_payload() { + let chunk = "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hi\"}]}}]}\n\n"; + let stream = futures::stream::iter(vec![Ok::(chunk.to_string())]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_data().await.unwrap(); + assert!(result.is_some()); + let json = result.unwrap(); + assert!(json["candidates"].is_array()); + } + + #[tokio::test] + async fn sse_reader_next_data_malformed_returns_none() { + // Malformed JSON data should be logged (H4) and the reader + // continues looking for the next valid data line. + let chunk = "data: not valid json\n\ndata: {\"ok\":true}\n\n"; + let stream = futures::stream::iter(vec![Ok::(chunk.to_string())]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + // First call should skip malformed and return the valid one. + let result = reader.next_data().await.unwrap(); + assert!(result.is_some()); + let json = result.unwrap(); + assert_eq!(json["ok"], true); + } + + #[tokio::test] + async fn sse_reader_buffer_overflow_returns_error() { + let huge = "x".repeat(SSE_MAX_BUFFER + 1); + let stream = futures::stream::iter(vec![Ok::(huge)]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_data().await; + assert!(result.is_err(), "should error on buffer overflow"); + let err_msg = result.unwrap_err().to_string(); + assert!( + err_msg.contains("SSE buffer"), + "error should mention SSE buffer: {err_msg}" + ); + } + + // ================================================== + // Body size limit tests (H5) + // ================================================== + + #[test] + fn max_response_body_is_ten_mb() { + assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); + } } diff --git a/src/provider/openai.rs b/src/provider/openai.rs index 1dc557a..d0a78ed 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -50,6 +50,7 @@ const SSE_DATA_PREFIX: &str = "data: "; const TEXT_PART_INDEX: usize = 0; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb +const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -519,7 +520,14 @@ impl SseReader { // Fetch the next chunk from the network. match self.bytes.next().await { - Some(Ok(chunk)) => self.buf.push_str(&chunk), + Some(Ok(chunk)) => { + self.buf.push_str(&chunk); + if self.buf.len() > SSE_MAX_BUFFER { + return Err(ApiError::http(format!( + "SSE buffer exceeded {SSE_MAX_BUFFER} bytes" + ))); + } + } Some(Err(e)) => return Err(e), None => return Ok(None), } @@ -1302,4 +1310,85 @@ mod tests { .build(); assert!(client.is_ok(), "build should succeed with valid timeouts"); } + + // ================================================== + // SSE buffer cap tests (M1) + // ================================================== + + #[tokio::test] + async fn sse_reader_take_line_splits_on_newline() { + let mut reader = SseReader { + bytes: Box::pin(futures::stream::empty()), + buf: "data: hello\ndata: world\n".to_string(), + }; + assert_eq!(reader.take_line(), Some("data: hello".to_string())); + assert_eq!(reader.take_line(), Some("data: world".to_string())); + assert_eq!(reader.take_line(), None); + } + + #[tokio::test] + async fn sse_reader_next_data_extracts_payload() { + let data = "data: {\"id\":\"c1\",\"model\":\"gpt-4o\",\"choices\":[]}\n\n"; + let stream = futures::stream::iter(vec![Ok::(data.to_string())]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_data().await.unwrap(); + assert!(result.is_some()); + assert!(result.unwrap().contains("c1")); + } + + #[tokio::test] + async fn sse_reader_next_data_done_returns_none() { + let stream = + futures::stream::iter(vec![Ok::("data: [DONE]\n\n".to_string())]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_data().await.unwrap(); + assert!(result.is_none()); + } + + #[tokio::test] + async fn sse_reader_buffer_overflow_returns_error() { + // Feed a chunk larger than SSE_MAX_BUFFER without any newline so + // the buffer grows unbounded — the cap should catch it. + let huge = "x".repeat(SSE_MAX_BUFFER + 1); + let stream = futures::stream::iter(vec![Ok::(huge)]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_data().await; + assert!(result.is_err(), "should error on buffer overflow"); + let err_msg = result.unwrap_err().to_string(); + assert!( + err_msg.contains("SSE buffer"), + "error should mention SSE buffer: {err_msg}" + ); + } + + // ================================================== + // Body size limit tests (H5) + // ================================================== + + #[test] + fn max_response_body_is_ten_mb() { + assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); + } + + #[test] + fn body_size_check_rejects_oversized() { + // Verify the comparison logic used in create_message. + let oversized = MAX_RESPONSE_BODY + 1; + assert!(oversized > MAX_RESPONSE_BODY); + } + + #[test] + fn body_size_check_accepts_within_limit() { + let within = MAX_RESPONSE_BODY; + assert!(within <= MAX_RESPONSE_BODY); + } } diff --git a/src/stream.rs b/src/stream.rs index 5333422..9b82aa2 100644 --- a/src/stream.rs +++ b/src/stream.rs @@ -73,6 +73,7 @@ pub mod heartbeat; /// handled correctly — for example, when accumulated tool-call JSON /// is malformed at [`PartStop`](StreamEvent::PartStop) time. #[derive(Debug)] +#[non_exhaustive] pub enum StreamError { /// The concatenated tool-call input JSON could not be parsed. /// @@ -509,6 +510,7 @@ pub enum DeltaPart { /// assert!(!reason.should_continue_tool_loop()); /// ``` #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] pub enum StreamStopReason { /// The model decided to invoke a tool. /// diff --git a/src/stream/handler.rs b/src/stream/handler.rs index c4a341e..79589fe 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -350,6 +350,7 @@ impl StreamRetryConfig { /// Completed < TotalTimeout < EventTimeout < InitFailed < FallbackToNonStreaming < Cancelled /// ``` #[derive(Debug, Clone)] +#[non_exhaustive] pub enum StreamOutcome { /// Stream completed normally — all events received, `MessageStop` seen. /// @@ -479,6 +480,7 @@ impl fmt::Display for StreamOutcome { /// to distinguish between transient failures (retryable) and permanent /// errors (non-retryable). #[derive(Debug)] +#[non_exhaustive] pub enum StreamHandlerError { /// Streaming initialization failed after all retries. /// diff --git a/src/tool/registry.rs b/src/tool/registry.rs index 6034959..ec55ca5 100644 --- a/src/tool/registry.rs +++ b/src/tool/registry.rs @@ -75,6 +75,12 @@ impl ToolRegistry { /// ``` pub fn register(&mut self, tool: impl Tool + 'static) { let name = tool.name().to_string(); + if self.tools.contains_key(&name) { + tracing::warn!( + tool = %name, + "overwriting previously registered tool with the same name" + ); + } self.tools.insert(name, Box::new(tool)); } @@ -506,3 +512,59 @@ impl Tool for FnTool { self.system_prompt.clone() } } + +#[cfg(test)] +mod tests { + use super::*; + + /// A simple tool function for testing duplicate registration (L1 fix). + fn test_tool_fn( + _input: Value, + _ctx: &ToolContext, + ) -> Pin> + Send + 'static>> { + Box::pin(async { Ok(ToolOutput::text("ok")) }) + } + + /// A simple tool for testing duplicate registration (L1 fix). + fn make_tool(name: &str) -> FnTool { + FnTool::new( + name.into(), + "test tool".into(), + serde_json::json!({"type": "object"}), + test_tool_fn, + ) + } + + #[test] + fn register_duplicate_overwrites_previous() { + let mut registry = ToolRegistry::new(); + registry.register(make_tool("my_tool")); + assert_eq!(registry.len(), 1); + + // Registering with the same name should overwrite, not add. + registry.register(make_tool("my_tool")); + assert_eq!( + registry.len(), + 1, + "duplicate registration should not increase count" + ); + } + + #[test] + fn register_different_names_adds_both() { + let mut registry = ToolRegistry::new(); + registry.register(make_tool("tool_a")); + registry.register(make_tool("tool_b")); + assert_eq!(registry.len(), 2); + } + + #[test] + fn register_overwrite_uses_new_tool() { + let mut registry = ToolRegistry::new(); + registry.register(make_tool("my_tool")); + + // After overwriting, the tool should still be callable. + registry.register(make_tool("my_tool")); + assert!(registry.get("my_tool").is_some()); + } +} From 7bcb1e8f9608c67887792c25bdc89b23c434926e Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 17:19:42 +1200 Subject: [PATCH 18/30] fix: memory trait, config errors, doc links and clippy warnings --- Cargo.toml | 16 ++++++++++ src/config.rs | 34 +++++++++++++------- src/engine/bare.rs | 2 +- src/engine/bare/compact.rs | 4 +-- src/engine/bare/dispatch.rs | 8 ++--- src/engine/bare/stream.rs | 7 ++++ src/engine/loop_core.rs | 2 ++ src/memory.rs | 64 ++++++++++++++++++++++++------------- src/memory/builtin.rs | 59 +++++++++++++++++++--------------- src/stream/handler.rs | 1 + 10 files changed, 130 insertions(+), 67 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ccfd708..f987c30 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,22 @@ grok = ["providers", "openai"] gemini = ["providers"] zai = ["providers", "anthropic"] +[[example]] +name = "hello-cli" +required-features = ["testing"] + +[[example]] +name = "repl-cli" +required-features = ["testing"] + +[[example]] +name = "echo-tool-cli" +required-features = ["testing"] + +[[example]] +name = "chat" +required-features = ["testing", "providers"] + [lints.clippy] pedantic = { level = "warn", priority = -1 } unwrap_used = "deny" diff --git a/src/config.rs b/src/config.rs index c4ce527..31881be 100644 --- a/src/config.rs +++ b/src/config.rs @@ -98,8 +98,8 @@ impl LoopConfig { /// /// # Errors /// - /// Returns a [`String`] describing the first invalid field, or `Ok(())` - /// if all fields are valid. + /// Returns a [`Config`](crate::error::LoopError::Config) variant describing + /// or `Ok(())` if all fields are valid. /// /// # Example /// @@ -113,27 +113,35 @@ impl LoopConfig { /// assert!(bad.validate().is_err()); /// ``` #[must_use = "validation errors should not be silently ignored"] - pub fn validate(&self) -> Result<(), String> { + pub fn validate(&self) -> Result<(), crate::error::LoopError> { if self.max_turns == 0 { - return Err("max_turns must be greater than 0".to_string()); + return Err(crate::error::LoopError::Config( + "max_turns must be greater than 0".to_string(), + )); } if self.context_window == 0 { - return Err("context_window must be greater than 0".to_string()); + return Err(crate::error::LoopError::Config( + "context_window must be greater than 0".to_string(), + )); } if self.max_tokens == 0 { - return Err("max_tokens must be greater than 0".to_string()); + return Err(crate::error::LoopError::Config( + "max_tokens must be greater than 0".to_string(), + )); } if self.model.is_empty() { - return Err("model must not be empty".to_string()); + return Err(crate::error::LoopError::Config( + "model must not be empty".to_string(), + )); } if self.compact_threshold.is_nan() || self.compact_threshold < 0.0 || self.compact_threshold > 1.0 { - return Err(format!( + return Err(crate::error::LoopError::Config(format!( "compact_threshold must be in [0.0, 1.0], got {}", self.compact_threshold - )); + ))); } Ok(()) } @@ -156,9 +164,10 @@ mod tests { ..LoopConfig::default() }; let err = config.validate().unwrap_err(); + let msg = err.to_string(); assert!( - err.contains("max_tokens"), - "error should mention max_tokens: {err}" + msg.contains("max_tokens"), + "error should mention max_tokens: {msg}" ); } @@ -178,7 +187,8 @@ mod tests { ..LoopConfig::default() }; let err = config.validate().unwrap_err(); - assert!(err.contains("model"), "error should mention model: {err}"); + let msg = err.to_string(); + assert!(msg.contains("model"), "error should mention model: {msg}"); } #[test] diff --git a/src/engine/bare.rs b/src/engine/bare.rs index e5b259d..59dc527 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -703,7 +703,7 @@ impl crate::engine::loop_core::Loop for BareLoop { config: &'a crate::config::LoopConfig, ) -> Pin> + Send + 'a>> { Box::pin(async move { - config.validate().map_err(LoopError::Config)?; + config.validate()?; self.state = LoopState::Processing { turn: 0 }; self.budget = SessionResult::default(); diff --git a/src/engine/bare/compact.rs b/src/engine/bare/compact.rs index 5f779b8..5ce5007 100644 --- a/src/engine/bare/compact.rs +++ b/src/engine/bare/compact.rs @@ -17,8 +17,8 @@ use crate::observer::CompactedContext; impl BareLoop { /// Check if context compaction is needed and perform it if so. /// - /// When a [`ContextManager`] is configured, this method: - /// 1. Calls [`ContextManager::ensure_context_fits()`] to check token usage. + /// When a [`crate::compact::ContextManager`] is configured, this method: + /// 1. Calls [`ContextManager::ensure_context_fits`](crate::compact::ContextManager::ensure_context_fits) to check token usage. /// 2. If compaction occurred, replaces `self.conversation` with the compacted messages. /// 3. Notifies observers via [`LoopObserver::on_compaction`](crate::observer::LoopObserver::on_compaction). /// diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 1d7213d..9256549 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -38,7 +38,7 @@ impl BareLoop { /// Execute tool calls and return results. /// /// Iterates over each [`ToolCall`] extracted from the assistant - /// message, looks up the corresponding tool in the [`ToolRegistry`], + /// message, looks up the corresponding tool in the [`ToolRegistry`](crate::tool::ToolRegistry), /// and invokes it. Each result is wrapped in a [`ToolDispatchResult`]. /// /// Tool execution is **sequential** so that cancellation can be @@ -47,7 +47,7 @@ impl BareLoop { /// allowing the model to recover. /// /// When a tool returns an error (execution failure or not-found), - /// the framework consults the [`Reflector`] and [`RecoveryStrategy`] + /// the framework consults the [`Reflector`](crate::reflection::Reflector) and [`RecoveryStrategy`](crate::reflection::RecoveryStrategy) /// to decide whether to retry, skip, ask user, or fail. Retry /// attempts use the delay specified by the [`RecoveryAction`]. /// @@ -78,7 +78,7 @@ impl BareLoop { /// Dispatch a single tool call, using reflector + recovery on errors. /// /// If the tool call succeeds, returns the result immediately. If it - /// fails, calls [`Reflector::analyze()`] and [`RecoveryStrategy::decide()`] + /// fails, calls [`Reflector::analyze`](crate::reflection::Reflector::analyze) and [`RecoveryStrategy::decide`](crate::reflection::RecoveryStrategy::decide) /// to determine the next action: /// /// - [`Retry`](RecoveryAction::Retry) — re-dispatch the tool after the @@ -239,7 +239,7 @@ impl BareLoop { /// Tries the middleware pipeline first, then a direct registry lookup, /// then produces a not-found error result. Handles cancellation during /// execution. Observer notification is handled by the caller - /// ([`dispatch_tool_with_recovery`]). + /// (`dispatch_tool_with_recovery`). /// /// # Errors /// diff --git a/src/engine/bare/stream.rs b/src/engine/bare/stream.rs index 8dafc31..cf321d8 100644 --- a/src/engine/bare/stream.rs +++ b/src/engine/bare/stream.rs @@ -52,6 +52,13 @@ impl BareLoop { // Inline streaming (no handler). let system = self.config.system_prompt.clone(); let tool_schemas = self.build_tool_schemas(); + // Clone the conversation history for the API request. The `ApiClient` + // trait requires `'static` streams (it takes ownership of the + // messages), so a clone is unavoidable here. The in-memory clone is + // O(n) in the number of messages but is typically dwarfed by the + // cost of serialising the messages into an HTTP request body. For + // very long sessions (>200 turns with large tool outputs), consider + // enabling auto-compaction to bound the history size. let mut stream = self.client .stream_messages(self.conversation.clone(), system, tool_schemas); diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index 580dd71..d74b8be 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -181,6 +181,7 @@ pub enum LoopState { /// assert!(!more.is_complete); /// ``` #[derive(Debug, Clone)] +#[non_exhaustive] pub struct TurnResult { /// May be empty if the response consists entirely of tool calls. pub text: String, @@ -427,6 +428,7 @@ impl ToolCall { /// assert_eq!(err.error.unwrap(), "API rate limit exceeded"); /// ``` #[derive(Debug, Clone)] +#[non_exhaustive] pub struct SessionResult { /// Matches [`LoopConfig::session_id`]. pub session_id: Uuid, diff --git a/src/memory.rs b/src/memory.rs index 6a77e88..0fbeddb 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -12,41 +12,46 @@ //! //! # Quick Start //! -//! ``` +//! ```rust //! use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory, ConsolidationStats}; //! use loopctl::error::LoopError; //! use std::future::Future; //! use std::pin::Pin; +//! use std::sync::RwLock; //! -//! struct InMemoryStore { -//! entries: Vec, +//! struct MyStore { +//! entries: RwLock>, //! } //! -//! impl LoopMemory for InMemoryStore { -//! fn store(&mut self, entry: MemoryEntry) +//! impl LoopMemory for MyStore { +//! fn store(&self, entry: MemoryEntry) //! -> Pin> + Send + '_>> //! { -//! Box::pin(async move { self.entries.push(entry); Ok(()) }) +//! Box::pin(async move { +//! self.entries.write().unwrap().push(entry); +//! Ok(()) +//! }) //! } //! fn retrieve(&self, query: &str, limit: usize) //! -> Pin, LoopError>> + Send + '_>> //! { //! let query = query.to_string(); //! Box::pin(async move { -//! Ok(self.entries.iter() +//! let entries = self.entries.read().unwrap(); +//! Ok(entries.iter() //! .filter(|e| e.memory.contains(&query)) //! .take(limit) //! .cloned() //! .collect()) //! }) //! } -//! fn consolidate(&mut self) +//! fn consolidate(&self) //! -> Pin> + Send + '_>> //! { //! Box::pin(async move { Ok(ConsolidationStats::default()) }) //! } //! fn len(&self) -> usize { -//! self.entries.len() +//! self.entries.read().unwrap().len() //! } //! } //! ``` @@ -74,41 +79,47 @@ pub mod entry; /// /// # Example /// -/// ``` +/// ```rust /// use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory, ConsolidationStats}; /// use loopctl::error::LoopError; /// use std::future::Future; /// use std::pin::Pin; +/// use std::sync::RwLock; /// -/// struct InMemoryStore { -/// entries: Vec, +/// struct MyStore { +/// entries: RwLock>, /// } /// -/// impl LoopMemory for InMemoryStore { -/// fn store(&mut self, entry: MemoryEntry) +/// impl LoopMemory for MyStore { +/// fn store(&self, entry: MemoryEntry) /// -> Pin> + Send + '_>> /// { -/// Box::pin(async move { self.entries.push(entry); Ok(()) }) +/// Box::pin(async move { +/// self.entries.write().unwrap().push(entry); +/// Ok(()) +/// }) /// } /// fn retrieve(&self, query: &str, limit: usize) /// -> Pin, LoopError>> + Send + '_>> /// { /// let query = query.to_string(); /// Box::pin(async move { -/// Ok(self.entries.iter() +/// let entries = self.entries.read().unwrap(); +/// Ok(entries.iter() /// .filter(|e| e.memory.contains(&query)) /// .take(limit) /// .cloned() /// .collect()) /// }) /// } -/// fn consolidate(&mut self) +/// fn consolidate(&self) /// -> Pin> + Send + '_>> /// { /// Box::pin(async move { -/// let before = self.entries.len(); -/// self.entries.retain(|e| e.relevance > 0.1); -/// let after = self.entries.len(); +/// let mut entries = self.entries.write().unwrap(); +/// let before = entries.len(); +/// entries.retain(|e| e.relevance > 0.1); +/// let after = entries.len(); /// Ok(ConsolidationStats { /// entries_before: before, /// entries_after: after, @@ -118,7 +129,7 @@ pub mod entry; /// }) /// } /// fn len(&self) -> usize { -/// self.entries.len() +/// self.entries.read().unwrap().len() /// } /// } /// ``` @@ -129,8 +140,12 @@ pub trait LoopMemory: Send + Sync { /// for example after a successful tool invocation, a resolved error, or /// an insight drawn from conversation. Implementations should persist the /// entry in whatever backing store they use. + /// + /// Takes `&self` so that memory stores can be shared via `Arc`. + /// Implementations that need interior mutability (e.g. an in-memory `Vec`) + /// should use `Mutex`, `RwLock`, or lock-free structures internally. fn store( - &mut self, + &self, entry: MemoryEntry, ) -> Pin> + Send + '_>>; @@ -157,8 +172,11 @@ pub trait LoopMemory: Send + Sync { /// may remove low-relevance entries, merge duplicates, or produce /// compressed summaries. Returns [`ConsolidationStats`] describing what /// was done. + /// + /// Takes `&self` so that memory stores can be shared via `Arc`. + /// Implementations should use interior mutability as needed. fn consolidate( - &mut self, + &self, ) -> Pin> + Send + '_>>; /// Number of entries currently stored. diff --git a/src/memory/builtin.rs b/src/memory/builtin.rs index add3dd9..a0c99ad 100644 --- a/src/memory/builtin.rs +++ b/src/memory/builtin.rs @@ -26,7 +26,7 @@ //! use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory}; //! //! # tokio::runtime::Runtime::new().unwrap().block_on(async { -//! let mut store = InMemoryStore::new(); +//! let store = InMemoryStore::new(); //! //! store.store( //! MemoryEntry::new(MemoryCategory::Insight, "Prefer Glob over manual file search") @@ -41,6 +41,7 @@ use crate::error::LoopError; use crate::memory::{ConsolidationStats, LoopMemory, MemoryEntry}; use std::future::Future; use std::pin::Pin; +use std::sync::{PoisonError, RwLock}; /// A simple in-memory store for loop memory entries. /// @@ -73,9 +74,10 @@ use std::pin::Pin; /// /// # Thread Safety /// -/// [`InMemoryStore`] is `Send + Sync` because all mutation goes through -/// `&mut self` in the [`LoopMemory`] trait. If you need shared mutable -/// access from multiple tasks, wrap it in `Arc>`. +/// [`InMemoryStore`] is `Send + Sync`. Interior mutability is handled via +/// an internal `RwLock`, so `store` and `consolidate` only require `&self`. +/// This allows the store to be shared via `Arc` or +/// `Arc` across tasks without external locking. /// /// # Construction /// @@ -99,7 +101,7 @@ use std::pin::Pin; /// use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory}; /// /// # tokio::runtime::Runtime::new().unwrap().block_on(async { -/// let mut store = InMemoryStore::new(); +/// let store = InMemoryStore::new(); /// /// store.store(MemoryEntry::new(MemoryCategory::Insight, "Prefer Glob over manual file search")).await.unwrap(); /// @@ -118,7 +120,7 @@ use std::pin::Pin; /// calling `consolidate()` periodically or implementing a custom /// [`LoopMemory`] with bounded capacity. pub struct InMemoryStore { - entries: Vec, + entries: RwLock>, } // =================================================== @@ -142,7 +144,7 @@ impl InMemoryStore { #[must_use] pub fn new() -> Self { Self { - entries: Vec::new(), + entries: RwLock::new(Vec::new()), } } @@ -164,8 +166,8 @@ impl InMemoryStore { /// assert_eq!(store.len(), 2); /// ``` #[must_use] - pub fn with_entries(mut self, entries: Vec) -> Self { - self.entries = entries; + pub fn with_entries(self, entries: Vec) -> Self { + *self.entries.write().unwrap_or_else(PoisonError::into_inner) = entries; self } } @@ -191,11 +193,14 @@ impl LoopMemory for InMemoryStore { /// /// This implementation never returns an error. fn store( - &mut self, + &self, entry: MemoryEntry, ) -> Pin> + Send + '_>> { Box::pin(async move { - self.entries.push(entry); + self.entries + .write() + .unwrap_or_else(PoisonError::into_inner) + .push(entry); Ok(()) }) } @@ -226,7 +231,7 @@ impl LoopMemory for InMemoryStore { /// use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory}; /// /// # tokio::runtime::Runtime::new().unwrap().block_on(async { - /// let mut store = InMemoryStore::new(); + /// let store = InMemoryStore::new(); /// store.store(MemoryEntry::new(MemoryCategory::Fact, "file search uses Glob")).await.unwrap(); /// /// let results = store.retrieve("file search", 5).await.unwrap(); @@ -245,8 +250,8 @@ impl LoopMemory for InMemoryStore { let query_lower = query.to_lowercase(); let query_words: Vec<&str> = query_lower.split_whitespace().collect(); - let mut scored: Vec<(f32, MemoryEntry)> = self - .entries + let entries = self.entries.read().unwrap_or_else(PoisonError::into_inner); + let mut scored: Vec<(f32, MemoryEntry)> = entries .iter() .map(|entry| { let memory_lower = entry.memory.to_lowercase(); @@ -299,21 +304,22 @@ impl LoopMemory for InMemoryStore { /// use loopctl::memory::LoopMemory; /// /// # tokio::runtime::Runtime::new().unwrap().block_on(async { - /// let mut store = InMemoryStore::new(); + /// let store = InMemoryStore::new(); /// let stats = store.consolidate().await.unwrap(); /// println!("Pruned {} entries", stats.pruned); /// # }); /// ``` fn consolidate( - &mut self, + &self, ) -> Pin> + Send + '_>> { Box::pin(async move { - let entries_before = self.entries.len(); - self.entries.retain(|e| e.relevance >= 0.05); - let pruned = entries_before.saturating_sub(self.entries.len()); + let mut entries = self.entries.write().unwrap_or_else(PoisonError::into_inner); + let entries_before = entries.len(); + entries.retain(|e| e.relevance >= 0.05); + let pruned = entries_before.saturating_sub(entries.len()); Ok(ConsolidationStats { entries_before, - entries_after: self.entries.len(), + entries_after: entries.len(), pruned, merged: 0, bytes_saved: 0, @@ -326,7 +332,10 @@ 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.len() + self.entries + .read() + .unwrap_or_else(PoisonError::into_inner) + .len() } } @@ -337,7 +346,7 @@ mod tests { #[tokio::test] async fn test_store_and_retrieve() { - let mut store = InMemoryStore::new(); + let store = InMemoryStore::new(); store .store(MemoryEntry::new( @@ -362,7 +371,7 @@ mod tests { #[tokio::test] async fn test_retrieve_respects_limit() { - let mut store = InMemoryStore::new(); + let store = InMemoryStore::new(); for i in 0..10 { store @@ -394,7 +403,7 @@ mod tests { #[tokio::test] async fn test_consolidate_prunes_low_relevance() { - let mut store = InMemoryStore::new(); + let store = InMemoryStore::new(); let mut good_entry = MemoryEntry::new(MemoryCategory::Insight, "useful insight"); good_entry.relevance = 0.9; @@ -425,7 +434,7 @@ mod tests { #[tokio::test] async fn test_tag_matching_boosts_relevance() { - let mut store = InMemoryStore::new(); + let store = InMemoryStore::new(); let tagged = MemoryEntry::new(MemoryCategory::Strategy, "use iterators for loops").with_tag("rust"); diff --git a/src/stream/handler.rs b/src/stream/handler.rs index 79589fe..58e5265 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -1018,6 +1018,7 @@ impl StreamHandler { /// println!("Response: {:?}", result.message); /// ``` #[derive(Debug, Clone)] +#[non_exhaustive] pub struct StreamTurnResult { /// The fully accumulated assistant message. pub message: Message, From 66834428652a36bc1e3fbc3130855d7c11d31464 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Tue, 30 Jun 2026 20:14:33 +1200 Subject: [PATCH 19/30] chore: debug assert idle, config clone reference return --- src/engine/bare.rs | 68 +++++++++++++++++++++++++++++++++++++++-- src/engine/loop_core.rs | 12 +++++--- src/memory.rs | 31 +++++++++---------- src/memory/builtin.rs | 20 ++++++------ 4 files changed, 98 insertions(+), 33 deletions(-) diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 59dc527..c4575b4 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -414,11 +414,34 @@ impl BareLoop { // Dependency setters // ================================================== + /// Assert that the loop has not started running yet. + /// + /// Configuration setters must be called before [`run()`](crate::engine::loop_core::Loop::run). + /// Calling them during a running session is a logic bug — the new value + /// takes effect immediately but parts of the session may have already + /// been initialised with the old value, leading to subtle inconsistencies. + /// + /// This check is only active in debug builds (`debug_assertions`). + #[inline] + fn debug_assert_idle(&self) { + debug_assert!( + matches!(self.state, LoopState::Idle), + "BareLoop configuration setters must be called before run() — \ + current state is {:?}, expected Idle", + self.state + ); + } + /// Set the [`Reflector`] for tool-error analysis. /// /// Replaces the default [`NoopReflector`] with a caller-supplied /// implementation. Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started + /// (i.e., when [`state`](LoopState) is not [`Idle`](LoopState::Idle)). + /// /// # Example /// /// ```rust,ignore @@ -426,6 +449,7 @@ impl BareLoop { /// agent.set_reflector(Arc::new(MyReflector)); /// ``` pub fn set_reflector(&mut self, reflector: Arc) { + self.debug_assert_idle(); self.reflector = reflector; } @@ -435,6 +459,10 @@ impl BareLoop { /// caller-supplied implementation. Must be called before /// [`run()`](crate::engine::loop_core::Loop::run). /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started. + /// /// # Example /// /// ```rust,ignore @@ -442,6 +470,7 @@ impl BareLoop { /// agent.set_recovery_strategy(Arc::new(MyStrategy)); /// ``` pub fn set_recovery_strategy(&mut self, strategy: Arc) { + self.debug_assert_idle(); self.recovery = strategy; } @@ -451,6 +480,10 @@ impl BareLoop { /// triggers compaction when usage exceeds the configured threshold. /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started. + /// /// # Example /// /// ```rust,ignore @@ -468,6 +501,7 @@ impl BareLoop { /// agent.set_context_manager(Arc::new(manager)); /// ``` pub fn set_context_manager(&mut self, manager: Arc) { + self.debug_assert_idle(); self.managers.set_context_manager(manager); } @@ -478,6 +512,10 @@ impl BareLoop { /// using the inline streaming logic. Must be called before /// [`run()`](crate::engine::loop_core::Loop::run). /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started. + /// /// # Example /// /// ```rust,ignore @@ -495,6 +533,7 @@ impl BareLoop { /// agent.set_stream_handler(handler); /// ``` pub fn set_stream_handler(&mut self, handler: StreamHandler) { + self.debug_assert_idle(); self.managers.set_stream_handler(handler); } @@ -507,6 +546,10 @@ impl BareLoop { /// executor in [`crate::hooks::Interactivity::Headless`] mode (the default). /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started. + /// /// *Requires `hooks` feature.* /// /// # Example @@ -521,6 +564,7 @@ impl BareLoop { /// ``` #[cfg(feature = "hooks")] pub fn set_hook_executor(&mut self, executor: Arc) { + self.debug_assert_idle(); self.managers.set_hook_executor(executor); } @@ -531,6 +575,10 @@ impl BareLoop { /// circuit breaker opened, blocking subsequent calls until recovery. /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started. + /// /// *Requires `tool_health` feature.* /// /// # Example @@ -545,6 +593,7 @@ impl BareLoop { /// ``` #[cfg(feature = "tool_health")] pub fn set_health_registry(&mut self, registry: Arc) { + self.debug_assert_idle(); self.managers.set_health_registry(registry); } @@ -555,6 +604,10 @@ impl BareLoop { /// pipeline's middleware chain before reaching the registry. /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started. + /// /// Build the pipeline using [`ToolPipeline::builder()`], adding middleware /// layers **without** calling `.core()` — the registry is injected /// automatically from `self.tools` so that schema generation and dispatch @@ -577,6 +630,7 @@ impl BareLoop { /// Returns [`LoopError::Config`] if the builder fails to produce a valid /// pipeline (e.g. internal invariant violated). pub fn set_pipeline(&mut self, builder: ToolPipelineBuilder) -> Result<(), LoopError> { + self.debug_assert_idle(); let pipeline = builder .core(Arc::clone(&self.tools)) .build() @@ -593,6 +647,10 @@ impl BareLoop { /// /// Must be called before [`run()`](crate::engine::loop_core::Loop::run). /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started. + /// /// # Example /// /// ```rust,ignore @@ -603,6 +661,7 @@ impl BareLoop { /// agent.register_observer(Arc::new(MyObserver)); /// ``` pub fn register_observer(&mut self, observer: Arc) { + self.debug_assert_idle(); self.managers.register_observer(observer); } @@ -616,6 +675,10 @@ impl BareLoop { /// The callback receives a `&str` containing the delta text fragment. /// It must be `Send + Sync` as it may be called from an async context. /// + /// # Panics (debug only) + /// + /// In debug builds, panics if called after the session has started. + /// /// # Example /// /// ```rust,ignore @@ -629,6 +692,7 @@ impl BareLoop { /// })); /// ``` pub fn set_text_streamer(&mut self, f: Arc) { + self.debug_assert_idle(); self.text_streamer = Some(f); } @@ -1000,8 +1064,8 @@ impl crate::engine::loop_core::Loop for BareLoop { None } - fn config(&self) -> LoopConfig { - self.config.clone() + fn config(&self) -> &LoopConfig { + &self.config } } diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index d74b8be..11f4b0b 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -672,7 +672,8 @@ pub trait Loop: Send + Sync { user_input: &'a str, ) -> Pin> + Send + 'a>> { Box::pin(async move { - self.initialize(&self.config()).await?; + let config = self.config().clone(); + self.initialize(&config).await?; let mut is_first_turn = true; loop { @@ -714,7 +715,10 @@ pub trait Loop: Send + Sync { /// Return the configuration that [`run`](Loop::run) passes to /// [`initialize`](Loop::initialize). /// - /// Implementors should return the [`LoopConfig`] they want to use - /// for the session. - fn config(&self) -> LoopConfig; + /// Implementors should return a reference to the [`LoopConfig`] they + /// want to use for the session. Returning a reference (rather than an + /// owned clone) avoids a mandatory `Clone` on every call — the default + /// [`run`](Loop::run) implementation only needs the config during + /// [`initialize`], so the borrow is short-lived. + fn config(&self) -> &LoopConfig; } diff --git a/src/memory.rs b/src/memory.rs index 0fbeddb..bb4d81f 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -58,7 +58,6 @@ use crate::error::LoopError; use std::future::Future; -use std::pin::Pin; pub use builtin::InMemoryStore; pub use entry::{ConsolidationStats, MemoryCategory, MemoryEntry}; @@ -82,8 +81,6 @@ pub mod entry; /// ```rust /// use loopctl::memory::{LoopMemory, MemoryEntry, MemoryCategory, ConsolidationStats}; /// use loopctl::error::LoopError; -/// use std::future::Future; -/// use std::pin::Pin; /// use std::sync::RwLock; /// /// struct MyStore { @@ -92,30 +89,30 @@ pub mod entry; /// /// impl LoopMemory for MyStore { /// fn store(&self, entry: MemoryEntry) -/// -> Pin> + Send + '_>> +/// -> impl Future> + Send /// { -/// Box::pin(async move { +/// async move { /// self.entries.write().unwrap().push(entry); /// Ok(()) -/// }) +/// } /// } /// fn retrieve(&self, query: &str, limit: usize) -/// -> Pin, LoopError>> + Send + '_>> +/// -> impl Future, LoopError>> + Send /// { /// let query = query.to_string(); -/// Box::pin(async move { +/// async move { /// let entries = self.entries.read().unwrap(); /// Ok(entries.iter() /// .filter(|e| e.memory.contains(&query)) /// .take(limit) /// .cloned() /// .collect()) -/// }) +/// } /// } /// fn consolidate(&self) -/// -> Pin> + Send + '_>> +/// -> impl Future> + Send /// { -/// Box::pin(async move { +/// async move { /// let mut entries = self.entries.write().unwrap(); /// let before = entries.len(); /// entries.retain(|e| e.relevance > 0.1); @@ -126,7 +123,7 @@ pub mod entry; /// pruned: before - after, /// ..Default::default() /// }) -/// }) +/// } /// } /// fn len(&self) -> usize { /// self.entries.read().unwrap().len() @@ -141,13 +138,13 @@ pub trait LoopMemory: Send + Sync { /// an insight drawn from conversation. Implementations should persist the /// entry in whatever backing store they use. /// - /// Takes `&self` so that memory stores can be shared via `Arc`. + /// Takes `&self` so that memory stores can be shared via `Arc`. /// Implementations that need interior mutability (e.g. an in-memory `Vec`) /// should use `Mutex`, `RwLock`, or lock-free structures internally. fn store( &self, entry: MemoryEntry, - ) -> Pin> + Send + '_>>; + ) -> impl Future> + Send; /// Retrieve memory entries relevant to the given query. /// @@ -164,7 +161,7 @@ pub trait LoopMemory: Send + Sync { &self, query: &str, limit: usize, - ) -> Pin, LoopError>> + Send + '_>>; + ) -> impl Future, LoopError>> + Send; /// Consolidate memory (e.g. prune, summarize, compress). /// @@ -173,11 +170,11 @@ pub trait LoopMemory: Send + Sync { /// compressed summaries. Returns [`ConsolidationStats`] describing what /// was done. /// - /// Takes `&self` so that memory stores can be shared via `Arc`. + /// Takes `&self` so that memory stores can be shared via `Arc`. /// Implementations should use interior mutability as needed. fn consolidate( &self, - ) -> Pin> + Send + '_>>; + ) -> impl Future> + Send; /// Number of entries currently stored. /// diff --git a/src/memory/builtin.rs b/src/memory/builtin.rs index a0c99ad..fc7e881 100644 --- a/src/memory/builtin.rs +++ b/src/memory/builtin.rs @@ -40,7 +40,6 @@ use crate::error::LoopError; use crate::memory::{ConsolidationStats, LoopMemory, MemoryEntry}; use std::future::Future; -use std::pin::Pin; use std::sync::{PoisonError, RwLock}; /// A simple in-memory store for loop memory entries. @@ -182,6 +181,7 @@ impl Default for InMemoryStore { // LoopMemory implementation // =================================================== +#[allow(clippy::manual_async_fn)] impl LoopMemory for InMemoryStore { /// Store a new memory entry by appending it to the backing list. /// @@ -195,14 +195,14 @@ impl LoopMemory for InMemoryStore { fn store( &self, entry: MemoryEntry, - ) -> Pin> + Send + '_>> { - Box::pin(async move { + ) -> impl Future> + Send { + async move { self.entries .write() .unwrap_or_else(PoisonError::into_inner) .push(entry); Ok(()) - }) + } } /// Retrieve memory entries relevant to the given query. @@ -244,9 +244,9 @@ impl LoopMemory for InMemoryStore { &self, query: &str, limit: usize, - ) -> Pin, LoopError>> + Send + '_>> { + ) -> impl Future, LoopError>> + Send { let query = query.to_string(); - Box::pin(async move { + async move { let query_lower = query.to_lowercase(); let query_words: Vec<&str> = query_lower.split_whitespace().collect(); @@ -281,7 +281,7 @@ impl LoopMemory for InMemoryStore { scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); Ok(scored.into_iter().take(limit).map(|(_, e)| e).collect()) - }) + } } /// Consolidate memory by pruning low-relevance entries. @@ -311,8 +311,8 @@ impl LoopMemory for InMemoryStore { /// ``` fn consolidate( &self, - ) -> Pin> + Send + '_>> { - Box::pin(async move { + ) -> impl Future> + Send { + async move { let mut entries = self.entries.write().unwrap_or_else(PoisonError::into_inner); let entries_before = entries.len(); entries.retain(|e| e.relevance >= 0.05); @@ -324,7 +324,7 @@ impl LoopMemory for InMemoryStore { merged: 0, bytes_saved: 0, }) - }) + } } /// Number of entries currently stored. From ef6e4af92ccdde0dce6bee9fd3aea9d8bc6cdad4 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 08:36:54 +1200 Subject: [PATCH 20/30] feat: model switch during runtime --- Cargo.toml | 1 + src/api.rs | 30 +- src/engine/bare.rs | 745 ++++++++++++++++++++++++++------- src/engine/loop_core.rs | 154 ++++++- src/memory.rs | 9 +- src/memory/builtin.rs | 9 +- src/middleware/unknown_tool.rs | 7 +- src/observer.rs | 72 +++- src/observer/context.rs | 13 + src/provider/anthropic.rs | 22 +- src/provider/gemini.rs | 22 +- src/provider/openai.rs | 23 +- src/stream/handler.rs | 8 +- src/testing.rs | 68 +-- 14 files changed, 956 insertions(+), 227 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f987c30..7935629 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ tokio-util = { version = "0.7", features = ["rt"] } uuid = { version = "1", features = ["v4", "serde"] } tracing = "0.1" +parking_lot = "0.12" reqwest = { version = "0.12", default-features = false, features = ["json", "stream", "rustls-tls"], optional = true } async-stream = { version = "0.3", optional = true } diff --git a/src/api.rs b/src/api.rs index 1024296..50b0bf5 100644 --- a/src/api.rs +++ b/src/api.rs @@ -48,8 +48,8 @@ use std::sync::Arc; /// } /// /// impl ApiClient for MyProviderClient { -/// fn model(&self) -> &str { -/// &self.model +/// fn model(&self) -> String { +/// self.model.clone() /// } /// /// fn stream_messages( @@ -93,7 +93,18 @@ pub trait ApiClient: Send + Sync { /// /// Called by the framework during initialization and on each turn for /// observability purposes. - fn model(&self) -> &str; + fn model(&self) -> String; + + /// Attempt to switch the model at runtime. + /// + /// Returns `true` if the client supports hot-swapping and the model + /// was updated successfully. Returns `false` by default (not supported). + /// Provider implementations that store their model behind interior + /// mutability override this to enable + /// [`BareLoop::switch_model`](crate::engine::BareLoop::switch_model). + fn set_model(&self, _model: &str) -> bool { + false + } /// Stream messages from the LLM provider. /// @@ -220,8 +231,8 @@ mod tests { } impl ApiClient for MockClient { - fn model(&self) -> &str { - &self.model_name + fn model(&self) -> String { + self.model_name.clone() } fn stream_messages( @@ -322,4 +333,13 @@ mod tests { let client: SharedApiClient = Arc::new(MockClient::new("shared")); assert_eq!(client.model(), "shared"); } + + #[test] + fn default_set_model_returns_false() { + // MockClient does not override set_model, so the default impl + // should return false (unsupported). + let client = MockClient::new("test-model"); + assert!(!client.set_model("other-model")); + assert_eq!(client.model(), "test-model"); + } } diff --git a/src/engine/bare.rs b/src/engine/bare.rs index c4575b4..c4b1bdb 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -76,8 +76,8 @@ use crate::hooks::context::{ use crate::message::{Message, MessagePart, Role, ToolContent}; use crate::middleware::{ToolDispatchContext, ToolPipeline, ToolPipelineBuilder}; use crate::observer::{ - FallbackContext, ResponseContext, StreamContext, StreamFailureContext, TurnEndContext, - TurnStartContext, + FallbackContext, ModelSwitchedContext, ResponseContext, StreamContext, StreamFailureContext, + TurnEndContext, TurnStartContext, }; use crate::reflection::{ ExponentialBackoffRecovery, NoopReflector, RecoveryAction, RecoveryStrategy, ReflectionContext, @@ -688,7 +688,7 @@ impl BareLoop { /// let buf = Arc::clone(&buffer); /// agent.set_text_streamer(Arc::new(move |delta| { /// print!("{delta}"); - /// buf.lock().unwrap().push_str(delta); + /// buf.lock().push_str(delta); /// })); /// ``` pub fn set_text_streamer(&mut self, f: Arc) { @@ -696,6 +696,48 @@ impl BareLoop { self.text_streamer = Some(f); } + /// Begin a model switch operation. + /// + /// Returns a [`ModelSwitch`] builder that lets you optionally update + /// the context window and max tokens before calling `.apply()`. + /// + /// This is the preferred way to switch models when the new model has + /// a different context window or token limit: + /// + /// ```rust,ignore + /// # use loopctl::engine::BareLoop; + /// # use loopctl::config::LoopConfig; + /// # use loopctl::tool::registry::ToolRegistry; + /// # use loopctl::testing::MockApiClient; + /// # let client = std::sync::Arc::new(MockApiClient::new("model-a")); + /// # let tools = ToolRegistry::new(); + /// # let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + /// loop_.switch_model("model-b").context_window(8192).apply().unwrap(); + /// assert_eq!(loop_.config().model, "model-b"); + /// assert_eq!(loop_.config().context_window, 8192); + /// ``` + /// + /// For simple cases where you just want to swap the model name: + /// + /// ```rust,ignore + /// # use loopctl::engine::BareLoop; + /// # use loopctl::config::LoopConfig; + /// # use loopctl::tool::registry::ToolRegistry; + /// # use loopctl::testing::MockApiClient; + /// # let client = std::sync::Arc::new(MockApiClient::new("a")); + /// # let tools = ToolRegistry::new(); + /// # let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + /// loop_.switch_model("b").apply().unwrap(); + /// ``` + pub fn switch_model(&mut self, model: &str) -> ModelSwitch<'_, C> { + ModelSwitch { + loop_: self, + target_model: model.to_string(), + context_window: None, + max_tokens: None, + } + } + // ================================================== // Run helpers // ================================================== @@ -732,7 +774,7 @@ impl BareLoop { let tool_result_msg = Self::build_tool_result_message(results); self.conversation.push(tool_result_msg); self.managers.observers().on_turn_end(&TurnEndContext { - turn: budget.total_turns, + turn: turn_index, success: true, error: None, duration_ms: Self::millis_u64(turn_duration), @@ -744,7 +786,7 @@ impl BareLoop { Err(e) => { let err_str = e.to_string(); self.managers.observers().on_turn_end(&TurnEndContext { - turn: budget.total_turns, + turn: turn_index, success: false, error: Some(err_str), duration_ms: Self::millis_u64(turn_duration), @@ -755,12 +797,242 @@ impl BareLoop { } } } + // ================================================== + // Turn helpers (used by process_turn) + // ================================================== + + /// Return `Err(Cancelled)` if the cancel signal has been set. + /// + /// # Errors + /// + /// Returns [`LoopError::Cancelled`] if the cancel signal has been set. + fn check_cancellation(&mut self) -> Result<(), LoopError> { + if self.is_cancelled() { + self.state = LoopState::Failed { + error: "cancelled".into(), + }; + return Err(LoopError::Cancelled); + } + Ok(()) + } + + /// Execute the streaming API call and process the result. + /// + /// On success, records the success with the fallback manager and fires + /// [`on_stream_success`](crate::observer::LoopObserver::on_stream_success). + /// + /// On failure, records the failure with the fallback circuit breaker + /// (firing [`on_fallback`](crate::observer::LoopObserver::on_fallback) + /// if it trips), fires + /// [`on_stream_failure`](crate::observer::LoopObserver::on_stream_failure), + /// and returns the error. + /// + /// `BareLoop` records API failures and trips the circuit breaker but does + /// **not** automatically retry with the fallback model — it has a single + /// client. The `FallbackManager` is infrastructure for downstream + /// consumers that hold multiple API clients. + /// + /// # Errors + /// + /// Returns the [`LoopError`] from [`stream_turn`](Self::stream_turn) if + /// the API call fails. + async fn do_stream(&mut self) -> Result<(Message, Option, StreamStopReason), LoopError> { + match self.stream_turn().await { + Ok((msg, usage, stop)) => { + self.managers.fallback.record_model_success(); + let (in_tok, out_tok) = Self::usage_tokens(usage.as_ref()); + self.managers.observers().on_stream_success(&StreamContext { + turn: self.budget.total_turns, + model: self.client.model(), + input_tokens: in_tok, + output_tokens: out_tok, + }); + Ok((msg, usage, stop)) + } + Err(e) => { + let tripped = self.managers.fallback.record_api_failure(); + if tripped { + let from = self.client.model(); + if let Some(to) = self.managers.fallback.fallback_model() { + tracing::warn!(from = %from, to = %to, "fallback manager tripped"); + self.managers + .observers() + .on_fallback(&FallbackContext { from, to }); + } + } + + self.managers + .observers() + .on_stream_failure(&StreamFailureContext { + turn: self.budget.total_turns, + model: self.client.model(), + error: e.clone(), + }); + + self.state = LoopState::Failed { + error: e.to_string(), + }; + Err(e) + } + } + } + + /// Add this turn's token usage into the session-level budget. + fn accumulate_usage(&mut self, usage: Option<&Usage>) { + if let Some(u) = usage { + self.budget.input_tokens = self + .budget + .input_tokens + .saturating_add(u64::from(u.input_tokens)); + self.budget.output_tokens = self + .budget + .output_tokens + .saturating_add(u64::from(u.output_tokens)); + } + } + + /// Fire [`on_turn_end`](crate::observer::LoopObserver::on_turn_end). + fn finish_turn(&mut self, turn_in: u64, turn_out: u64, duration: Duration) { + self.managers.observers().on_turn_end(&TurnEndContext { + turn: self.budget.total_turns.saturating_sub(1), + success: true, + error: None, + duration_ms: Self::millis_u64(duration), + input_tokens: turn_in, + output_tokens: turn_out, + }); + } + + /// Build a [`TurnResult`] signalling turn completion. + fn turn_complete(text: String, turn_in: u64, turn_out: u64, duration: Duration) -> TurnResult { + TurnResult { + text, + tool_calls: Vec::new(), + tool_results: Vec::new(), + input_tokens: turn_in, + output_tokens: turn_out, + duration, + is_complete: true, + stop_reason: StopReason::EndTurn, + } + } + + /// Run context compaction if a [`ContextManager`] is configured. + /// + /// Best-effort: failures are logged and the turn continues with + /// un-compacted history. + async fn try_compact_context(&mut self) { + if let Err(e) = self.maybe_compact_context(self.budget.total_turns).await { + tracing::warn!( + error = %e, + turn = self.budget.total_turns, + "context compaction failed; continuing with uncompactd history" + ); + } + } } // ================================================== -// Loop trait implementation +// ModelSwitch builder // ================================================== +/// Builder for a runtime model switch on [`BareLoop`]. +/// +/// Created by [`BareLoop::switch_model`]. Allows updating +/// context-window and max-tokens alongside the model name, then applies +/// all changes atomically via [`apply`](Self::apply). +/// +/// The switch resets the fallback circuit breaker (stale failure counts +/// from the old model are meaningless for the new one) and fires +/// [`on_model_switched`](crate::observer::LoopObserver::on_model_switched) +/// to all observers. +pub struct ModelSwitch<'a, C: ApiClient> { + loop_: &'a mut BareLoop, + target_model: String, + context_window: Option, + max_tokens: Option, +} + +impl ModelSwitch<'_, C> { + /// Set the context window (in tokens) for the new model. + /// + /// If omitted, the existing `LoopConfig::context_window` is kept. + /// Updating this is important when switching to a model with a + /// significantly different context window — otherwise the + /// auto-compactor will use the wrong threshold. + #[must_use] + pub fn context_window(mut self, tokens: u64) -> Self { + self.context_window = Some(tokens); + self + } + + /// Set the max output tokens for the new model. + /// + /// If omitted, the existing `LoopConfig::max_tokens` is kept. + #[must_use] + pub fn max_tokens(mut self, tokens: u32) -> Self { + self.max_tokens = Some(tokens); + self + } + + /// Apply the model switch. + /// + /// Performs the following atomically: + /// 1. Validates the target model is non-empty. + /// 2. Delegates to [`ApiClient::set_model`] on the underlying client. + /// 3. Updates `LoopConfig::model`, `context_window`, and `max_tokens`. + /// 4. Resets the [`FallbackManager`](crate::fallback::FallbackManager) + /// circuit breaker to `Primary` and updates the original-model + /// tracker to the new model. + /// 5. Fires [`on_model_switched`](crate::observer::LoopObserver::on_model_switched). + /// + /// # Errors + /// + /// - [`LoopError::Config`] if the model name is empty/whitespace. + pub fn apply(self) -> Result<(), LoopError> { + let Self { + loop_, + target_model, + context_window, + max_tokens, + } = self; + + let trimmed = target_model.trim(); + if trimmed.is_empty() { + return Err(LoopError::Config( + "model name must not be empty or whitespace".into(), + )); + } + + let from = loop_.config.model.clone(); + loop_.client.set_model(trimmed); + loop_.config.model = trimmed.to_string(); + + if let Some(cw) = context_window { + loop_.config.context_window = cw; + } + + if let Some(mt) = max_tokens { + loop_.config.max_tokens = mt; + } + + loop_.managers.fallback.reset(); + loop_ + .managers + .fallback + .set_original_model(trimmed.to_string()); + loop_ + .managers + .observers() + .on_model_switched(&ModelSwitchedContext { + from, + to: trimmed.to_string(), + }); + + Ok(()) + } +} + impl crate::engine::loop_core::Loop for BareLoop { fn initialize<'a>( &'a mut self, @@ -779,7 +1051,6 @@ impl crate::engine::loop_core::Loop for BareLoop { }) } - #[allow(clippy::too_many_lines)] fn process_turn<'a>( &'a mut self, input: &'a str, @@ -799,81 +1070,16 @@ impl crate::engine::loop_core::Loop for BareLoop { query: input.to_string(), }); - // Check cancellation before the API call. - if self.is_cancelled() { - self.state = LoopState::Failed { - error: "cancelled".into(), - }; - return Err(LoopError::Cancelled); - } - - let stream_result = self.stream_turn().await; - let (assistant_msg, usage, _stream_stop) = match stream_result { - Ok(value) => { - let (msg, usage, stop) = value; - self.managers.fallback.record_model_success(); - let (in_tok, out_tok) = Self::usage_tokens(usage.as_ref()); - self.managers.observers().on_stream_success(&StreamContext { - turn: self.budget.total_turns, - model: self.client.model().to_string(), - input_tokens: in_tok, - output_tokens: out_tok, - }); - (msg, usage, stop) - } - Err(e) => { - // Record the failure with the fallback circuit breaker. - // - // Note: BareLoop records API failures and trips the - // circuit breaker but does **not** automatically retry - // with the fallback model. The `FallbackManager` is - // infrastructure for downstream consumers that hold - // multiple API clients. BareLoop has a single client, - // so the error is propagated after recording. - let tripped = self.managers.fallback.record_api_failure(); - if tripped { - let from = self.client.model(); - if let Some(to) = self.managers.fallback.fallback_model() { - tracing::warn!(from, to, "fallback manager tripped"); - self.managers.observers().on_fallback(&FallbackContext { - from: from.to_string(), - to, - }); - } - } + self.check_cancellation()?; - self.managers - .observers() - .on_stream_failure(&StreamFailureContext { - turn: self.budget.total_turns, - model: self.client.model().to_string(), - error: e.clone(), - }); - - self.state = LoopState::Failed { - error: e.to_string(), - }; - return Err(e); - } - }; + let (assistant_msg, usage, _stream_stop) = self.do_stream().await?; - // Accumulate usage into session budget. - if let Some(u) = &usage { - self.budget.input_tokens = self - .budget - .input_tokens - .saturating_add(u64::from(u.input_tokens)); - self.budget.output_tokens = self - .budget - .output_tokens - .saturating_add(u64::from(u.output_tokens)); - } + self.accumulate_usage(usage.as_ref()); let text = Self::extract_text(&assistant_msg); let (turn_in, turn_out) = Self::usage_tokens(usage.as_ref()); - // Record the response text with the detection manager and check - // for loop/convergence patterns. + // Record the response and check for loop/convergence patterns. let pattern = self.managers.detection.record_response(&text); self.managers.observers().on_response(&ResponseContext { @@ -886,29 +1092,25 @@ impl crate::engine::loop_core::Loop for BareLoop { .managers .handle_detected_pattern(&pattern, self.budget.total_turns) { - match result { + return match result { Ok(_) => { self.state = LoopState::Completed { summary: text.clone(), }; - return Ok(TurnResult { + Ok(Self::turn_complete( text, - tool_calls: Vec::new(), - tool_results: Vec::new(), - input_tokens: turn_in, - output_tokens: turn_out, - duration: turn_start.elapsed(), - is_complete: true, - stop_reason: StopReason::EndTurn, - }); + turn_in, + turn_out, + turn_start.elapsed(), + )) } Err(e) => { self.state = LoopState::Failed { error: e.to_string(), }; - return Err(e); + Err(e) } - } + }; } let tool_calls = Self::extract_tool_calls(&assistant_msg); @@ -917,27 +1119,16 @@ impl crate::engine::loop_core::Loop for BareLoop { // No tool calls → this turn is complete. if tool_calls.is_empty() { - self.managers.observers().on_turn_end(&TurnEndContext { - turn: self.budget.total_turns.saturating_sub(1), - success: true, - error: None, - duration_ms: Self::millis_u64(turn_start.elapsed()), - input_tokens: turn_in, - output_tokens: turn_out, - }); + self.finish_turn(turn_in, turn_out, turn_start.elapsed()); self.state = LoopState::Completed { summary: text.clone(), }; - return Ok(TurnResult { + return Ok(Self::turn_complete( text, - tool_calls: Vec::new(), - tool_results: Vec::new(), - input_tokens: turn_in, - output_tokens: turn_out, - duration: turn_start.elapsed(), - is_complete: true, - stop_reason: StopReason::EndTurn, - }); + turn_in, + turn_out, + turn_start.elapsed(), + )); } // Dispatch tool calls. @@ -972,20 +1163,8 @@ impl crate::engine::loop_core::Loop for BareLoop { } self.budget = budget; - // Attempt context compaction. - // - // Compaction is best-effort: if it fails (e.g. the compactor - // cannot reduce the conversation enough), we log a warning and - // continue. The next API call may still succeed, and if it - // doesn't, the provider's context-overflow error will surface - // naturally at that point. - if let Err(e) = self.maybe_compact_context(self.budget.total_turns).await { - tracing::warn!( - error = %e, - turn = self.budget.total_turns, - "context compaction failed; continuing with uncompactd history" - ); - } + // Attempt context compaction (best-effort). + self.try_compact_context().await; self.state = LoopState::Processing { turn: self.budget.total_turns, @@ -1085,6 +1264,8 @@ mod tests { use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; + use parking_lot::Mutex; + // ================================================== // Mock ApiClient // ================================================== @@ -1105,13 +1286,13 @@ mod tests { /// Each entry is a `Vec` representing one complete /// streaming response from the API. Popped from the front by /// [`stream_messages`](MockClient::stream_messages). - responses: Arc>>>, + responses: Arc>>>, /// Model name reported by [`ApiClient::model()`]. /// /// Copied into mock response metadata so that assertions can - /// verify the model field. - model_name: String, + /// verify which model produced a given response. + model_name: Arc>, } impl MockClient { @@ -1123,8 +1304,8 @@ mod tests { /// [`add_tool_only_response()`](MockClient::add_tool_only_response). fn new(model: &str) -> Self { Self { - responses: Arc::new(std::sync::Mutex::new(Vec::new())), - model_name: model.to_string(), + responses: Arc::new(Mutex::new(Vec::new())), + model_name: Arc::new(parking_lot::Mutex::new(model.to_string())), } } @@ -1145,7 +1326,7 @@ mod tests { message: MessageMetadata { id: "msg_test".into(), role: "assistant".into(), - model: self.model_name.clone(), + model: self.model_name.lock().clone(), }, }), StreamEvent::PartStart(PartStart { @@ -1167,12 +1348,12 @@ mod tests { }), StreamEvent::MessageStop, ]; - self.responses.lock().unwrap().push(events); + self.responses.lock().push(events); } /// Add a raw sequence of stream events as a single response turn. fn add_events(&self, events: Vec) { - self.responses.lock().unwrap().push(events); + self.responses.lock().push(events); } /// Add a tool_call response followed by an end_turn response. @@ -1202,7 +1383,7 @@ mod tests { message: MessageMetadata { id: "msg_tool".into(), role: "assistant".into(), - model: self.model_name.clone(), + model: self.model_name.lock().clone(), }, }), StreamEvent::PartStart(PartStart { @@ -1218,7 +1399,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - self.responses.lock().unwrap().push(tool_events); + self.responses.lock().push(tool_events); // Second response: end_turn with text let text_events = vec![ @@ -1226,7 +1407,7 @@ mod tests { message: MessageMetadata { id: "msg_final".into(), role: "assistant".into(), - model: self.model_name.clone(), + model: self.model_name.lock().clone(), }, }), StreamEvent::PartStart(PartStart { @@ -1248,7 +1429,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - self.responses.lock().unwrap().push(text_events); + self.responses.lock().push(text_events); } /// Add a tool_call-only response (no end_turn). @@ -1269,7 +1450,7 @@ mod tests { message: MessageMetadata { id: format!("msg_{tool_id}"), role: "assistant".into(), - model: self.model_name.clone(), + model: self.model_name.lock().clone(), }, }), StreamEvent::PartStart(PartStart { @@ -1285,7 +1466,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - self.responses.lock().unwrap().push(tool_events); + self.responses.lock().push(tool_events); } /// Add an error response. Reserved for future use. @@ -1302,17 +1483,25 @@ mod tests { message: MessageMetadata { id: "msg_err".into(), role: "assistant".into(), - model: self.model_name.clone(), + model: self.model_name.lock().clone(), }, })]; - self.responses.lock().unwrap().push(events); + self.responses.lock().push(events); } } impl ApiClient for MockClient { /// Return the model name configured at construction. - fn model(&self) -> &str { - &self.model_name + fn model(&self) -> String { + 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(); + true } /// Pop the next queued response and return it as a stream. @@ -1327,7 +1516,7 @@ mod tests { _tools: Option>, ) -> Pin> + Send + 'static>> { - let mut guard = self.responses.lock().unwrap(); + let mut guard = self.responses.lock(); if let Some(events) = guard.pop_front() { let events: Vec> = events.into_iter().map(Ok).collect(); @@ -1885,7 +2074,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - client.responses.lock().unwrap().push(tool_events); + client.responses.lock().push(tool_events); // Second response: end_turn client.add_text_response("Both tools executed."); @@ -1915,16 +2104,16 @@ mod tests { client.add_text_response("Hello world"); let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); - let received = Arc::new(std::sync::Mutex::new(Vec::new())); + let received = Arc::new(Mutex::new(Vec::new())); let buf = Arc::clone(&received); agent.set_text_streamer(Arc::new(move |delta: &str| { - buf.lock().unwrap().push(delta.to_string()); + buf.lock().push(delta.to_string()); })); let result = agent.run("Hi").await.unwrap(); assert!(result.success); - let received = received.lock().unwrap(); + let received = received.lock(); assert!(!received.is_empty(), "streamer should have fired"); assert!( received.join("").contains("Hello world"), @@ -1988,17 +2177,17 @@ mod tests { let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); - let received = Arc::new(std::sync::Mutex::new(String::new())); + let received = Arc::new(Mutex::new(String::new())); let buf = Arc::clone(&received); agent.set_text_streamer(Arc::new(move |delta: &str| { - buf.lock().unwrap().push_str(delta); + 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().unwrap(); + let received = received.lock(); assert_eq!(&*received, "Done", "only text deltas should fire streamer"); } @@ -2129,7 +2318,7 @@ mod tests { }), StreamEvent::MessageStop, ]; - client.responses.lock().unwrap().push(tool_events); + 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."); @@ -2242,11 +2431,11 @@ mod tests { /// A middleware that records the `turn_number` from each dispatch context. struct TurnNumberCapture { - turns: Arc>>, + turns: Arc>>, } impl TurnNumberCapture { - fn new(shared: Arc>>) -> Self { + fn new(shared: Arc>>) -> Self { Self { turns: shared } } } @@ -2265,7 +2454,7 @@ mod tests { dyn std::future::Future + Send + 'a, >, > { - self.turns.lock().unwrap().push(ctx.turn_number); + self.turns.lock().push(ctx.turn_number); next.dispatch(ctx) } } @@ -2285,7 +2474,7 @@ mod tests { let mut config = make_config(); config.max_turns = 10; - let capture = Arc::new(std::sync::Mutex::new(Vec::::new())); + let capture = Arc::new(Mutex::new(Vec::::new())); let mut agent = BareLoop::new(Arc::new(client), registry, config); let builder = ToolPipeline::builder().with(TurnNumberCapture::new(Arc::clone(&capture))); agent.set_pipeline(builder).unwrap(); @@ -2293,7 +2482,7 @@ mod tests { let result = agent.run("test").await; assert!(result.is_ok()); - let turns = capture.lock().unwrap().clone(); + let turns = capture.lock().clone(); // Tool was called on turn 0 (first turn) and turn 1 (second turn). assert_eq!( turns.len(), @@ -2307,4 +2496,258 @@ mod tests { "turn_number must be actual index, not max_turns (10): got {turns:?}" ); } + + // ─── Model switching tests ─── + + /// `switch_model` updates config.model, the client's model, and the + /// fallback manager's original-model tracker. + #[tokio::test] + async fn switch_model_updates_config_and_client() { + let client = MockClient::new("model-a"); + let client_arc = std::sync::Arc::new(client); + let tools = ToolRegistry::new(); + let mut config = LoopConfig::default(); + config.model = "model-a".to_string(); + + let mut loop_ = BareLoop::new(client_arc.clone(), tools, config); + + loop_.switch_model("model-b").apply().unwrap(); + + // Config was updated. + assert_eq!(loop_.config().model, "model-b"); + + // Client was also updated via set_model. + assert_eq!(client_arc.model(), "model-b"); + } + + /// `switch_model` fires `on_model_switched` to all registered observers. + #[tokio::test] + async fn switch_model_notifies_observers() { + #[derive(Default)] + struct RecordingObserver { + switches: Mutex>, + } + + impl crate::observer::LoopObserver for RecordingObserver { + fn name(&self) -> &'static str { + "recording" + } + + fn on_model_switched(&self, ctx: &ModelSwitchedContext) { + self.switches + .lock() + .push((ctx.from.clone(), ctx.to.clone())); + } + } + + let client = std::sync::Arc::new(MockClient::new("m1")); + let tools = ToolRegistry::new(); + let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + let obs = std::sync::Arc::new(RecordingObserver::default()); + let obs_clone = obs.clone(); + loop_.register_observer(obs); + + loop_.switch_model("m2").apply().unwrap(); + loop_.switch_model("m3").apply().unwrap(); + + // Observer should have received both switches. + let recorded = 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())); + } + + /// `switch_model` updates config even when the client doesn't support + /// hot-swapping. The client's internal model stays the same, but the + /// framework-level config, fallback, and observers are updated. + #[tokio::test] + async fn switch_model_unsupported_client() { + /// A client that does NOT override `set_model` (returns `false`). + struct StaticClient { + model_name: Arc>, + } + + impl ApiClient for StaticClient { + fn model(&self) -> String { + self.model_name.lock().clone() + } + // Uses default set_model which returns false. + + fn stream_messages( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn futures::stream::Stream> + + Send + + 'static, + >, + > { + Box::pin(futures::stream::empty()) + } + + fn create_message( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn std::future::Future> + + Send + + '_, + >, + > { + Box::pin(async { Ok(serde_json::Value::Null) }) + } + } + + let client = std::sync::Arc::new(StaticClient { + model_name: std::sync::Arc::new(parking_lot::Mutex::new("static".to_string())), + }); + let tools = ToolRegistry::new(); + let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + + // set_model returns false (unsupported), but apply() is best-effort + // and still updates config-level state. + loop_.switch_model("new-model").apply().unwrap(); + + // Config was updated even though client didn't support it. + assert_eq!(loop_.config().model, "new-model"); + + // Client's model remains unchanged (no interior mutability). + assert_eq!(loop_.client.model(), "static"); + } + + /// `switch_model` syncs the fallback manager's original-model tracker + /// so subsequent fallback decisions compare against the new primary. + #[tokio::test] + async fn switch_model_updates_fallback_original() { + let client = std::sync::Arc::new(MockClient::new("primary")); + let tools = ToolRegistry::new(); + let mut config = LoopConfig::default(); + config.model = "primary".to_string(); + + let mut loop_ = BareLoop::new(client, tools, config); + + // Before switch, fallback manager has no original model set. + assert_eq!(loop_.managers.fallback.original_model(), None); + + loop_.switch_model("new-primary").apply().unwrap(); + + // After switch, fallback manager tracks the new primary. + assert_eq!( + loop_.managers.fallback.original_model(), + Some("new-primary".to_string()) + ); + } + + /// `switch_model` rejects empty/whitespace-only model names. + #[tokio::test] + async fn switch_model_rejects_empty() { + let client = std::sync::Arc::new(MockClient::new("model")); + let tools = ToolRegistry::new(); + let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + + let result = loop_.switch_model("").apply(); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("empty")); + + let result = loop_.switch_model(" ").apply(); + assert!(result.is_err()); + + // Model should remain unchanged. + assert_eq!(loop_.config().model, "default"); + } + + /// `switch_model` can be called multiple times in succession. + #[tokio::test] + async fn switch_model_chained() { + let client = std::sync::Arc::new(MockClient::new("a")); + let tools = ToolRegistry::new(); + let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + + loop_.switch_model("b").apply().unwrap(); + assert_eq!(loop_.config().model, "b"); + + loop_.switch_model("c").apply().unwrap(); + assert_eq!(loop_.config().model, "c"); + + loop_.switch_model("d").apply().unwrap(); + assert_eq!(loop_.config().model, "d"); + } + + /// `switch_model` with `.context_window()` updates the config. + #[tokio::test] + async fn switch_model_updates_context_window() { + let client = std::sync::Arc::new(MockClient::new("big-model")); + let tools = ToolRegistry::new(); + let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + + let original_cw = loop_.config().context_window; + assert_ne!(original_cw, 8192); + + loop_ + .switch_model("small-model") + .context_window(8192) + .apply() + .unwrap(); + + assert_eq!(loop_.config().model, "small-model"); + assert_eq!(loop_.config().context_window, 8192); + } + + /// `switch_model` with `.max_tokens()` updates the config. + #[tokio::test] + async fn switch_model_updates_max_tokens() { + let client = std::sync::Arc::new(MockClient::new("m")); + let tools = ToolRegistry::new(); + let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + + loop_.switch_model("m2").max_tokens(4096).apply().unwrap(); + + assert_eq!(loop_.config().model, "m2"); + assert_eq!(loop_.config().max_tokens, 4096); + } + + /// `switch_model` trims whitespace from the model name. + #[tokio::test] + async fn switch_model_trims_whitespace() { + let client = std::sync::Arc::new(MockClient::new("m")); + let tools = ToolRegistry::new(); + let mut loop_ = BareLoop::new(client, tools, LoopConfig::default()); + + loop_.switch_model(" gpt-4o ").apply().unwrap(); + assert_eq!(loop_.config().model, "gpt-4o"); + } + + /// `switch_model` resets the fallback circuit breaker. + #[tokio::test] + async fn switch_model_resets_fallback_circuit() { + use crate::fallback::FallbackState; + + let client = std::sync::Arc::new(MockClient::new("primary")); + let tools = ToolRegistry::new(); + let mut config = LoopConfig::default(); + config.model = "primary".to_string(); + + let mut loop_ = BareLoop::new(client, tools, config); + + // Trip the circuit breaker. + loop_.managers.fallback.set_original_model("primary".into()); + loop_.managers.fallback.set_fallback_model("backup"); + loop_.managers.fallback.transition_to_fallback(); + assert_eq!(loop_.managers.fallback.state(), FallbackState::Fallback); + + // Switch model — circuit should reset to Primary. + loop_.switch_model("new-primary").apply().unwrap(); + + assert_eq!(loop_.managers.fallback.state(), FallbackState::Primary); + assert_eq!( + loop_.managers.fallback.original_model(), + Some("new-primary".to_string()) + ); + } } diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index 11f4b0b..f4d4c11 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -538,6 +538,156 @@ impl SessionResult { pub fn total_tokens(&self) -> u64 { self.input_tokens.saturating_add(self.output_tokens) } + + /// Construct a `SessionResult` with full control over every field. + /// + /// Intended for tests that need to assert on specific field + /// combinations that the `success()` / `failed()` constructors + /// don't cover (e.g. non-zero `tool_calls` or `total_turns`). + /// + /// Only available with the `testing` feature. + /// + /// # Example + /// + /// ``` + /// # use loopctl::engine::loop_core::SessionResult; + /// # use std::time::Duration; + /// # use uuid::Uuid; + /// let result = SessionResult::builder() + /// .session_id(Uuid::nil()) + /// .total_turns(5) + /// .tool_calls(3) + /// .success(true) + /// .build(); + /// assert_eq!(result.total_turns, 5); + /// assert_eq!(result.tool_calls, 3); + /// assert!(result.success); + /// ``` + #[cfg(feature = "testing")] + #[must_use] + pub fn builder() -> SessionResultBuilder { + SessionResultBuilder { + session_id: Uuid::nil(), + total_turns: 0, + input_tokens: 0, + output_tokens: 0, + total_duration: Duration::ZERO, + tool_calls: 0, + success: false, + final_output: None, + error: None, + } + } +} + +/// Builder for constructing [`SessionResult`] instances in tests. +/// +/// Created by [`SessionResult::builder`]. All fields default to zeroed +/// or empty values; override only the ones your test cares about, then +/// call `.build()`. +#[cfg(feature = "testing")] +#[derive(Debug, Clone)] +pub struct SessionResultBuilder { + /// Unique session identifier. + session_id: Uuid, + /// Number of turns executed. + total_turns: usize, + /// Total input tokens consumed. + input_tokens: u64, + /// Total output tokens produced. + output_tokens: u64, + /// Wall-clock duration of the session. + total_duration: Duration, + /// Number of tool calls dispatched. + tool_calls: usize, + /// Whether the session completed successfully. + success: bool, + /// Final text output, if any. + final_output: Option, + /// Error message if the session failed. + error: Option, +} + +#[cfg(feature = "testing")] +impl SessionResultBuilder { + /// Set the session ID. + #[must_use] + pub fn session_id(mut self, id: Uuid) -> Self { + self.session_id = id; + self + } + + /// Set the total turns executed. + #[must_use] + pub fn total_turns(mut self, turns: usize) -> Self { + self.total_turns = turns; + self + } + + /// Set the total input tokens. + #[must_use] + pub fn input_tokens(mut self, tokens: u64) -> Self { + self.input_tokens = tokens; + self + } + + /// Set the total output tokens. + #[must_use] + pub fn output_tokens(mut self, tokens: u64) -> Self { + self.output_tokens = tokens; + self + } + + /// Set the total session duration. + #[must_use] + pub fn total_duration(mut self, duration: Duration) -> Self { + self.total_duration = duration; + self + } + + /// Set the total tool calls made. + #[must_use] + pub fn tool_calls(mut self, calls: usize) -> Self { + self.tool_calls = calls; + self + } + + /// Set whether the session succeeded. + #[must_use] + pub fn success(mut self, success: bool) -> Self { + self.success = success; + self + } + + /// Set the final output text. + #[must_use] + pub fn final_output(mut self, output: impl Into) -> Self { + self.final_output = Some(output.into()); + self + } + + /// Set the error message. + #[must_use] + pub fn error(mut self, error: impl Into) -> Self { + self.error = Some(error.into()); + self + } + + /// Build the [`SessionResult`]. + #[must_use] + pub fn build(self) -> SessionResult { + SessionResult { + session_id: self.session_id, + total_turns: self.total_turns, + input_tokens: self.input_tokens, + output_tokens: self.output_tokens, + total_duration: self.total_duration, + tool_calls: self.tool_calls, + success: self.success, + final_output: self.final_output, + error: self.error, + } + } } // ================================================== @@ -672,6 +822,8 @@ pub trait Loop: Send + Sync { user_input: &'a str, ) -> Pin> + Send + 'a>> { Box::pin(async move { + // Clone is required: `initialize` takes `&mut self` which + // conflicts with the shared borrow from `config()`. let config = self.config().clone(); self.initialize(&config).await?; @@ -719,6 +871,6 @@ pub trait Loop: Send + Sync { /// want to use for the session. Returning a reference (rather than an /// owned clone) avoids a mandatory `Clone` on every call — the default /// [`run`](Loop::run) implementation only needs the config during - /// [`initialize`], so the borrow is short-lived. + /// initialization, so the borrow is short-lived. fn config(&self) -> &LoopConfig; } diff --git a/src/memory.rs b/src/memory.rs index bb4d81f..c1c0fd1 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -141,10 +141,7 @@ pub trait LoopMemory: Send + Sync { /// Takes `&self` so that memory stores can be shared via `Arc`. /// Implementations that need interior mutability (e.g. an in-memory `Vec`) /// should use `Mutex`, `RwLock`, or lock-free structures internally. - fn store( - &self, - entry: MemoryEntry, - ) -> impl Future> + Send; + fn store(&self, entry: MemoryEntry) -> impl Future> + Send; /// Retrieve memory entries relevant to the given query. /// @@ -172,9 +169,7 @@ pub trait LoopMemory: Send + Sync { /// /// Takes `&self` so that memory stores can be shared via `Arc`. /// Implementations should use interior mutability as needed. - fn consolidate( - &self, - ) -> impl Future> + Send; + fn consolidate(&self) -> impl Future> + Send; /// Number of entries currently stored. /// diff --git a/src/memory/builtin.rs b/src/memory/builtin.rs index fc7e881..d92ec41 100644 --- a/src/memory/builtin.rs +++ b/src/memory/builtin.rs @@ -192,10 +192,7 @@ impl LoopMemory for InMemoryStore { /// # Errors /// /// This implementation never returns an error. - fn store( - &self, - entry: MemoryEntry, - ) -> impl Future> + Send { + fn store(&self, entry: MemoryEntry) -> impl Future> + Send { async move { self.entries .write() @@ -309,9 +306,7 @@ impl LoopMemory for InMemoryStore { /// println!("Pruned {} entries", stats.pruned); /// # }); /// ``` - fn consolidate( - &self, - ) -> impl Future> + Send { + fn consolidate(&self) -> impl Future> + Send { async move { let mut entries = self.entries.write().unwrap_or_else(PoisonError::into_inner); let entries_before = entries.len(); diff --git a/src/middleware/unknown_tool.rs b/src/middleware/unknown_tool.rs index b362097..b1a7bce 100644 --- a/src/middleware/unknown_tool.rs +++ b/src/middleware/unknown_tool.rs @@ -23,7 +23,12 @@ use std::sync::Arc; /// # Example /// /// ```rust,ignore -/// let mw = UnknownToolMiddleware::new(); +/// use std::sync::Arc; +/// use loopctl::tool::registry::ToolRegistry; +/// use loopctl::middleware::unknown_tool::UnknownToolMiddleware; +/// +/// let registry = Arc::new(ToolRegistry::new()); +/// let mw = UnknownToolMiddleware::new(registry); /// // If tool "basj" is not found, error message will say: /// // "Tool 'basj' not found. Did you mean 'bash'?" /// ``` diff --git a/src/observer.rs b/src/observer.rs index 75db134..289e411 100644 --- a/src/observer.rs +++ b/src/observer.rs @@ -40,8 +40,8 @@ pub mod context; pub use context::{ CompactedContext, ConvergenceDetectedContext, FallbackContext, LoopDetectedContext, - ResponseContext, SessionEndContext, SessionStartContext, StreamContext, StreamFailureContext, - ToolPostContext, ToolPreContext, TurnEndContext, TurnStartContext, + ModelSwitchedContext, ResponseContext, SessionEndContext, SessionStartContext, StreamContext, + StreamFailureContext, ToolPostContext, ToolPreContext, TurnEndContext, TurnStartContext, }; // ================================================== // LoopObserver Trait @@ -127,6 +127,12 @@ pub trait LoopObserver: Send + Sync { /// was selected for subsequent requests. fn on_fallback(&self, _ctx: &FallbackContext) {} + /// Called when the model is hot-swapped at runtime. + /// + /// Fired by [`BareLoop::switch_model`](crate::engine::BareLoop::switch_model) + /// after the client has accepted the new model. + fn on_model_switched(&self, _ctx: &ModelSwitchedContext) {} + /// Called when a loop is detected in tool operations. /// /// Fired when the same tool operation produces the same result @@ -300,6 +306,15 @@ impl ObserverHost { } } + /// Dispatch [`LoopObserver::on_model_switched`] to all observers. + /// + /// Iterates registered observers in registration order. + pub fn on_model_switched(&self, ctx: &ModelSwitchedContext) { + for obs in &self.observers { + obs.on_model_switched(ctx); + } + } + /// Dispatch [`LoopObserver::on_loop_detected`] to all observers. /// /// Iterates registered observers in registration order. @@ -415,4 +430,57 @@ mod tests { host.reset_all(); assert_eq!(obs.resets.load(Ordering::SeqCst), 1); } + + #[test] + fn host_dispatches_model_switched() { + struct SwitchRecorder { + events: parking_lot::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())); + } + } + + let obs = Arc::new(SwitchRecorder { + events: parking_lot::Mutex::new(Vec::new()), + }); + let mut host = ObserverHost::new(); + host.register(Arc::clone(&obs) as Arc); + + host.on_model_switched(&ModelSwitchedContext { + from: "a".into(), + to: "b".into(), + }); + host.on_model_switched(&ModelSwitchedContext { + from: "b".into(), + to: "c".into(), + }); + + let events = obs.events.lock(); + assert_eq!(events.len(), 2); + assert_eq!(events[0], ("a".into(), "b".into())); + assert_eq!(events[1], ("b".into(), "c".into())); + } + + #[test] + fn model_switched_default_is_noop() { + // The default impl of on_model_switched should be a no-op + // (no panic, no crash). + struct NoopObserver; + impl LoopObserver for NoopObserver { + fn name(&self) -> &'static str { + "noop" + } + } + + let obs = NoopObserver; + obs.on_model_switched(&ModelSwitchedContext { + from: "x".into(), + to: "y".into(), + }); + } } diff --git a/src/observer/context.rs b/src/observer/context.rs index 39ac2c1..559c184 100644 --- a/src/observer/context.rs +++ b/src/observer/context.rs @@ -283,6 +283,19 @@ pub struct FallbackContext { pub to: String, } +/// Context for [`LoopObserver::on_model_switched`](crate::observer::LoopObserver::on_model_switched). +/// +/// Emitted when the model is hot-swapped via +/// [`BareLoop::switch_model`](crate::engine::BareLoop::switch_model). +/// Carries the previous and new model identifiers. +#[derive(Debug, Clone)] +pub struct ModelSwitchedContext { + /// Model identifier before the switch. + pub from: String, + /// Model identifier after the switch. + pub to: String, +} + /// Context for [`LoopObserver::on_loop_detected`](crate::observer::LoopObserver::on_loop_detected). /// /// Describes the repeating tool pattern and how many times it diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index 859fd50..2ed5f07 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -65,7 +65,7 @@ pub struct AnthropicClient { http: reqwest::Client, api_key: String, base_url: String, - model: String, + model: parking_lot::Mutex, max_tokens: u32, } @@ -139,8 +139,16 @@ impl AnthropicClient { } impl ApiClient for AnthropicClient { - fn model(&self) -> &str { - &self.model + fn model(&self) -> String { + self.model.lock().clone() + } + + fn set_model(&self, model: &str) -> bool { + if model.trim().is_empty() { + return false; + } + *self.model.lock() = model.to_string(); + true } fn stream_messages( @@ -149,8 +157,9 @@ impl ApiClient for AnthropicClient { system: Option, tools: Option>, ) -> Pin> + Send + 'static>> { + let model = self.model.lock().clone(); let body = build_request_body( - &self.model, + &model, &messages, system.as_deref(), tools.as_deref(), @@ -185,8 +194,9 @@ impl ApiClient for AnthropicClient { system: Option, tools: Option>, ) -> Pin> + Send + '_>> { + let model = self.model.lock().clone(); let body = build_request_body( - &self.model, + &model, &messages, system.as_deref(), tools.as_deref(), @@ -310,7 +320,7 @@ impl AnthropicClientBuilder { http, api_key, base_url: self.base_url, - model: self.model, + model: parking_lot::Mutex::new(self.model), max_tokens: self.max_tokens, }) } diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 489c26d..a2422ab 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -61,7 +61,7 @@ pub struct GeminiClient { http: reqwest::Client, api_key: String, base_url: String, - model: String, + model: parking_lot::Mutex, } impl GeminiClient { @@ -101,15 +101,17 @@ impl GeminiClient { /// Gemini puts the model in the URL path and the API key as a query /// parameter rather than using headers. fn stream_url(&self) -> String { + let model = self.model.lock().clone(); format!( "{}/models/{}:streamGenerateContent?alt=sse", - self.base_url, self.model + self.base_url, model ) } /// Build the non-streaming Generate Content URL. fn generate_url(&self) -> String { - format!("{}/models/{}:generateContent", self.base_url, self.model) + let model = self.model.lock().clone(); + format!("{}/models/{}:generateContent", self.base_url, model) } /// Send a POST request and return the raw response. @@ -145,8 +147,16 @@ impl GeminiClient { } impl ApiClient for GeminiClient { - fn model(&self) -> &str { - &self.model + fn model(&self) -> String { + self.model.lock().clone() + } + + fn set_model(&self, model: &str) -> bool { + if model.trim().is_empty() { + return false; + } + *self.model.lock() = model.to_string(); + true } fn stream_messages( @@ -292,7 +302,7 @@ impl GeminiClientBuilder { http, api_key, base_url: self.base_url, - model: self.model, + model: parking_lot::Mutex::new(self.model), }) } } diff --git a/src/provider/openai.rs b/src/provider/openai.rs index d0a78ed..ca33fe6 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -22,7 +22,6 @@ use std::future::Future; use std::pin::Pin; - use std::time::Duration; use futures::stream::{Stream, StreamExt}; @@ -68,7 +67,7 @@ pub struct OpenAiClient { http: reqwest::Client, api_key: String, base_url: String, - model: String, + model: parking_lot::Mutex, } impl OpenAiClient { @@ -149,8 +148,16 @@ impl OpenAiClient { } impl ApiClient for OpenAiClient { - fn model(&self) -> &str { - &self.model + fn model(&self) -> String { + self.model.lock().clone() + } + + fn set_model(&self, model: &str) -> bool { + if model.trim().is_empty() { + return false; + } + *self.model.lock() = model.to_string(); + true } fn stream_messages( @@ -159,7 +166,8 @@ impl ApiClient for OpenAiClient { system: Option, tools: Option>, ) -> Pin> + Send + 'static>> { - let body = RequestBody::build(&self.model, &messages, system.as_deref(), tools.as_deref()); + let model = self.model.lock().clone(); + let body = RequestBody::build(&model, &messages, system.as_deref(), tools.as_deref()); let url = self.completions_url(); let api_key = self.api_key.clone(); let http = self.http.clone(); @@ -191,7 +199,8 @@ impl ApiClient for OpenAiClient { system: Option, tools: Option>, ) -> Pin> + Send + '_>> { - let body = RequestBody::build(&self.model, &messages, system.as_deref(), tools.as_deref()); + let model = self.model.lock().clone(); + let body = RequestBody::build(&model, &messages, system.as_deref(), tools.as_deref()); let url = self.completions_url(); Box::pin(async move { @@ -302,7 +311,7 @@ impl OpenAiClientBuilder { http, api_key, base_url: self.base_url, - model: self.model, + model: parking_lot::Mutex::new(self.model), }) } } diff --git a/src/stream/handler.rs b/src/stream/handler.rs index 58e5265..bb4bc99 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -1658,8 +1658,8 @@ mod tests { } impl ApiClient for HandlerMock { - fn model(&self) -> &'static str { - "test-model" + fn model(&self) -> String { + "test-model".to_string() } fn stream_messages( @@ -1847,8 +1847,8 @@ mod tests { /// Mock that always returns an error stream. struct ErrorMock; impl ApiClient for ErrorMock { - fn model(&self) -> &'static str { - "test-model" + fn model(&self) -> String { + "test-model".to_string() } fn stream_messages( &self, diff --git a/src/testing.rs b/src/testing.rs index 726be22..78e36e3 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -112,7 +112,9 @@ use futures::Stream; use serde_json::{Value, json}; use std::future::Future; use std::pin::Pin; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; + +use parking_lot::Mutex; use uuid::Uuid; // ================================================== @@ -184,7 +186,7 @@ pub struct MockApiClient { /// The value is returned verbatim by the [`ApiClient::model`] /// implementation. It appears in log messages and /// [`MessageMetadata`] fields. - model_name: String, + model_name: Arc>, /// The queue of canned responses. /// @@ -368,7 +370,7 @@ impl MockApiClient { stop_reason: "end_turn".to_string(), }; Self { - model_name: model.to_string(), + model_name: Arc::new(parking_lot::Mutex::new(model.to_string())), responses: Arc::new(Mutex::new(vec![default_response])), error: None, } @@ -394,12 +396,7 @@ impl MockApiClient { /// ``` #[must_use] pub fn with_text_response(self, text: &str) -> Self { - if let Some(r) = self - .responses - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .first_mut() - { + if let Some(r) = self.responses.lock().first_mut() { r.text = text.to_string(); } self @@ -427,10 +424,7 @@ impl MockApiClient { /// ``` #[must_use] pub fn with_tool_call(self, id: &str, name: &str, input: Value) -> Self { - let mut responses = self - .responses - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut responses = self.responses.lock(); if let Some(r) = responses.first_mut() { r.tool_call = Some(MockToolCall { id: id.to_string(), @@ -461,12 +455,7 @@ impl MockApiClient { /// ``` #[must_use] pub fn with_stop_reason(self, reason: &str) -> Self { - if let Some(r) = self - .responses - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .first_mut() - { + if let Some(r) = self.responses.lock().first_mut() { r.stop_reason = reason.to_string(); } self @@ -504,10 +493,7 @@ impl MockApiClient { #[must_use] pub fn with_responses(self, responses: Vec) -> Self { if !responses.is_empty() { - *self - .responses - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = responses; + *self.responses.lock() = responses; } self } @@ -558,10 +544,7 @@ impl MockApiClient { /// [R3] → pop → R3, queue stays [R3] (cloned) /// ``` fn pop_response(&self) -> MockResponse { - let mut guard = self - .responses - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut guard = self.responses.lock(); if guard.len() > 1 { guard.remove(0) } else { @@ -617,8 +600,16 @@ impl ApiClient for MockApiClient { /// let client = MockApiClient::new("my-test-model"); /// assert_eq!(client.model(), "my-test-model"); /// ``` - fn model(&self) -> &str { - &self.model_name + fn model(&self) -> String { + 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(); + true } /// Stream a canned sequence of [`StreamEvent`]s for the next response. @@ -673,7 +664,7 @@ impl ApiClient for MockApiClient { } let response = self.pop_response(); - let model = self.model_name.clone(); + let model = self.model_name.lock().clone(); let mut events: Vec> = vec![Ok(StreamEvent::MessageStart(MessageStart { message: MessageMetadata { @@ -1606,4 +1597,21 @@ mod tests { assert_eq!(config.max_turns, 10); assert!(config.system_prompt.is_some()); } + + #[test] + fn mock_api_client_set_model() { + let client = MockApiClient::new("model-a"); + assert_eq!(client.model(), "model-a"); + + assert!(client.set_model("model-b")); + assert_eq!(client.model(), "model-b"); + } + + #[test] + fn mock_api_client_set_model_rejects_empty() { + let client = MockApiClient::new("model-a"); + assert!(!client.set_model("")); + assert!(!client.set_model(" ")); + assert_eq!(client.model(), "model-a"); + } } From 463338be0a691aff761650f590de462bedec9b2a Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 11:27:40 +1200 Subject: [PATCH 21/30] fix: loop detection hard stop propagation, cancellation guard in dispatch, fallback attempt padding, docs in tests --- src/api/error.rs | 23 --- src/compact/truncating.rs | 19 -- src/detection/loop_detector.rs | 5 - src/detection/manager.rs | 89 ++------- src/engine/bare.rs | 355 +++++++-------------------------- src/engine/bare/dispatch.rs | 44 ++-- src/fallback.rs | 127 ++++++------ src/message.rs | 53 ----- src/middleware.rs | 69 ------- src/middleware/unknown_tool.rs | 138 ++++++++++++- src/observer.rs | 1 - src/provider.rs | 2 - src/provider/anthropic.rs | 12 -- src/provider/gemini.rs | 12 -- src/provider/openai.rs | 12 -- src/reflection.rs | 24 --- src/stream/handler.rs | 38 ---- src/tool/health.rs | 28 --- src/tool/registry.rs | 2 - src/tool/shield.rs | 16 -- 20 files changed, 315 insertions(+), 754 deletions(-) diff --git a/src/api/error.rs b/src/api/error.rs index 2dedf37..79cbc7c 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -1099,9 +1099,6 @@ pub type Result = std::result::Result; mod tests { use super::*; - /// Verify that a generic [`ApiError::api`] message maps to - /// [`ErrorCode::ApiRequestFailed`] and the Display output - /// contains the expected text. #[test] fn test_api_error_code() { let error = ApiError::api("API failed"); @@ -1110,9 +1107,6 @@ mod tests { assert_eq!(error.code(), ErrorCode::ApiRequestFailed); } - /// Verify that [`ApiError::auth`] with "Invalid" in the message - /// maps to [`ErrorCode::AuthInvalidKey`] and is recognised by - /// [`ApiError::is_auth_error`]. #[test] fn test_auth_error() { let error = ApiError::auth("Invalid key"); @@ -1121,31 +1115,18 @@ mod tests { assert!(error.is_auth_error()); } - /// Verify [`ApiError::auth_invalid_key`] produces the correct error code. - /// - /// Asserts that the error carries [`ErrorCode::AuthInvalidKey`] and that - /// the prefix "invalid" in the message triggers the specialized variant - /// rather than the generic [`ErrorCode::AuthFailed`]. #[test] fn test_auth_invalid_key() { let error = ApiError::auth_invalid_key("expired"); assert_eq!(error.code(), ErrorCode::AuthInvalidKey); } - /// Verify [`ApiError::auth`] without "invalid" maps to [`ErrorCode::AuthFailed`]. - /// - /// When the message does not contain the word "invalid", the classifier - /// should fall through to the generic [`ErrorCode::AuthFailed`] variant. #[test] fn test_auth_failed_generic() { let error = ApiError::auth("token expired"); assert_eq!(error.code(), ErrorCode::AuthFailed); } - /// Verify [`ApiError::http`] maps to [`ErrorCode::HttpConnectionError`]. - /// - /// Asserts that the [`Display`](std::fmt::Display) output contains - /// "HTTP error" and that [`ApiError::code`] returns the correct variant. #[test] fn test_http_error() { let error = ApiError::http("Connection failed"); @@ -1153,10 +1134,6 @@ mod tests { assert_eq!(error.code(), ErrorCode::HttpConnectionError); } - /// Verify [`ApiError::http_with_status`] embeds the status code. - /// - /// Checks that the [`Display`](std::fmt::Display) output includes both - /// the HTTP status code (e.g., "HTTP 500") and the original error message. #[test] fn test_http_error_with_status() { // 5xx → HttpResponseError diff --git a/src/compact/truncating.rs b/src/compact/truncating.rs index f956b14..5545395 100644 --- a/src/compact/truncating.rs +++ b/src/compact/truncating.rs @@ -422,25 +422,6 @@ mod tests { } } - /// Build a conversation where a tool-call and its result straddle - /// what would be the naive split point. - /// - /// Layout (index: content): - /// - /// ```text - /// 0 user "msg0" - /// 1 assistant "reply0" - /// 2 user "msg1" - /// 3 assistant "reply1" - /// 4 user "msg2" - /// 5 assistant tool_call("call_a", "search", ...) - /// 6 user tool_result("call_a", ...) - /// 7 assistant "final reply" - /// ``` - /// - /// With `preserve_recent = 2`, the naive split would be at index 6, - /// dropping the tool-call (index 5) but keeping the result (index 6). - /// The fix should move the split back to index 5. fn convo_with_straddling_tool_pair() -> Vec { vec![ Message::user("msg0"), diff --git a/src/detection/loop_detector.rs b/src/detection/loop_detector.rs index aeec922..412ac5a 100644 --- a/src/detection/loop_detector.rs +++ b/src/detection/loop_detector.rs @@ -1657,11 +1657,6 @@ mod tests { use super::*; use serde_json::json; - /// A test tool signature that knows about "Read" and "Bash" tools. - /// - /// Implements [`ToolSignature`] for unit tests within this module. Extracts - /// the `file_path` parameter for `Read` calls and the `command` parameter - /// for `Bash` calls, falling back to an empty string for unknown tools. struct TestToolSignature; impl ToolSignature for TestToolSignature { diff --git a/src/detection/manager.rs b/src/detection/manager.rs index 0ffbf6e..d12c53e 100644 --- a/src/detection/manager.rs +++ b/src/detection/manager.rs @@ -230,8 +230,7 @@ pub enum DetectedPattern { /// [`convergence_threshold`](Self::convergence_threshold), /// [`convergence_count`](Self::convergence_count), /// [`enable_convergence_detection`](Self::enable_convergence_detection), -/// [`on_converge`](Self::on_converge), -/// [`max_response_history`](Self::max_response_history). +/// [`on_converge`](Self::on_converge). /// /// # Example /// @@ -266,8 +265,6 @@ pub struct DetectionConfig { pub enable_convergence_detection: bool, /// Action on convergence. Default: [`ConvergenceAction::default()`]. pub on_converge: ConvergenceAction, - /// Max responses kept for convergence checking. Default: **20**. - pub max_response_history: usize, } impl Default for DetectionConfig { @@ -281,7 +278,6 @@ impl Default for DetectionConfig { convergence_count: 3, enable_convergence_detection: true, on_converge: ConvergenceAction::default(), - max_response_history: 20, } } } @@ -1231,10 +1227,6 @@ impl Default for DetectionManager { mod tests { use super::*; - /// Verify that a freshly constructed [`DetectionManager`] reports [`DetectedPattern::NoPattern`]. - /// - /// Creates a new manager and immediately calls [`check_current_pattern`](DetectionManager::check_current_pattern). - /// Asserts that no pattern is detected before any data has been recorded. #[test] fn test_no_pattern_initially() { let dm = DetectionManager::new().unwrap(); @@ -1244,11 +1236,6 @@ mod tests { )); } - /// Verify that repeating the same tool call triggers [`DetectedPattern::LoopDetected`]. - /// - /// Calls [`record_tool_call`](DetectionManager::record_tool_call) with identical arguments - /// five times. Asserts that at least one call returns a loop detection result, - /// confirming the default [`DetectionConfig::loop_threshold`] of 3 is respected. #[test] fn test_loop_detection() { let dm = DetectionManager::new().unwrap(); @@ -1262,11 +1249,6 @@ mod tests { panic!("Expected loop detection after 5 identical calls"); } - /// Verify that submitting identical responses triggers [`DetectedPattern::ConvergenceDetected`]. - /// - /// Calls [`record_response`](DetectionManager::record_response) with the same string five - /// times. Asserts that convergence is detected, validating the Jaccard similarity - /// check and the [`DetectionConfig::convergence_count`] threshold. #[test] fn test_convergence_detection() { let dm = DetectionManager::new().unwrap(); @@ -1280,11 +1262,6 @@ mod tests { panic!("Expected convergence detection after 5 identical responses"); } - /// Verify that [`reset`](DetectionManager::reset) clears all detection state. - /// - /// Records a tool call and a response, then calls [`reset`](DetectionManager::reset). - /// Asserts that [`check_current_pattern`](DetectionManager::check_current_pattern) returns - /// [`DetectedPattern::NoPattern`] and that [`stats`](DetectionManager::stats) shows zero turns. #[test] fn test_reset() { let dm = DetectionManager::new().unwrap(); @@ -1299,11 +1276,6 @@ mod tests { assert_eq!(stats.turns_analyzed, 0); } - /// Verify that disabling both detectors causes all calls to return [`DetectedPattern::NoPattern`]. - /// - /// Constructs a [`DetectionManager`] with `enable_loop_detection: false` and - /// `enable_convergence_detection: false`. Records 10 identical tool calls and asserts - /// none trigger a detection, confirming the feature-flag short-circuits work. #[test] fn test_no_detection_when_disabled() { let config = DetectionConfig { @@ -1318,12 +1290,6 @@ mod tests { } } - /// Verify that the [`LoopStatus::should_stop`] flag activates after `stop_threshold` repetitions. - /// - /// Configures `loop_threshold: 3` and `stop_threshold: 5`, then records 5 identical - /// operations. Asserts that [`check_loop`](DetectionManager::check_loop) reports - /// [`is_looping`](LoopStatus::is_looping), [`should_stop`](LoopStatus::should_stop), and a - /// warning containing `"STOPPING"`. #[test] fn test_loop_status_should_stop() { let config = DetectionConfig { @@ -1345,12 +1311,6 @@ mod tests { assert!(status.warning.unwrap().contains("STOPPING")); } - /// Verify loop detection using the rich [`Operation`] API. - /// - /// Creates [`Operation`] values with [`Operation::new`] and records them via - /// [`record_operation`](DetectionManager::record_operation). Asserts that the loop - /// detector correctly identifies repeated tool + parameter pairs and reports - /// a repetition count of at least 3. #[test] fn test_record_operation_with_operation_struct() { let dm = DetectionManager::new().unwrap(); @@ -1365,11 +1325,6 @@ mod tests { assert!(status.repetition_count >= 3); } - /// Verify that [`Operation::from_input_with_signature`] integrates with loop detection. - /// - /// Parses a JSON input using [`NoOpToolSignature`](super::loop_detector::NoOpToolSignature) - /// and records 5 operations. Because `NoOpToolSignature` produces an empty `primary_param`, - /// all `"Read"` operations are considered identical and a loop is detected. #[test] fn test_record_operation_from_input() { use super::super::loop_detector::NoOpToolSignature; @@ -1390,12 +1345,6 @@ mod tests { assert!(status.is_looping); } - /// Verify that changing result hashes prevent loop detection. - /// - /// Calls [`record_tool_call_with_result`](DetectionManager::record_tool_call_with_result) - /// with the same tool and input but a different result hash each time. Asserts that - /// [`check_loop`](DetectionManager::check_loop) reports `is_looping: false`, confirming - /// that the result-aware logic treats varying outputs as progress. #[test] fn test_result_aware_detection() { let dm = DetectionManager::new().unwrap(); @@ -1410,12 +1359,6 @@ mod tests { assert!(!status.is_looping, "Different results should not be a loop"); } - /// Verify that identical result hashes trigger loop detection. - /// - /// Calls [`record_tool_call_with_result`](DetectionManager::record_tool_call_with_result) - /// with the same tool, input, *and* result hash. Asserts that the loop detector - /// identifies this as a genuine loop, confirming the result-aware path works when - /// results do not change. #[test] fn test_result_aware_same_result_is_loop() { let dm = DetectionManager::new().unwrap(); @@ -1429,12 +1372,6 @@ mod tests { assert!(status.is_looping, "Same results should be detected as loop"); } - /// Verify that [`loop_detector`](DetectionManager::loop_detector) provides access - /// to the inner [`LoopDetector`]. - /// - /// Calls [`loop_detector()`](DetectionManager::loop_detector) on a fresh manager - /// and asserts that [`turn_count`](LoopDetector::turn_count) returns 0, confirming - /// the accessor returns a usable reference. #[test] fn test_access_loop_detector() { let dm = DetectionManager::new().unwrap(); @@ -1442,11 +1379,6 @@ mod tests { assert_eq!(dm.loop_detector().turn_count(), 0); } - /// Verify that [`DetectionConfig::to_loop_detector_config`] correctly maps fields. - /// - /// Creates a [`DetectionConfig`] with custom `loop_threshold`, `stop_threshold`, and - /// `max_history`, then converts it to a [`LoopDetectorConfig`]. Asserts that each field - /// maps to the expected value (`repetition_threshold`, `stop_threshold`, `window_size`). #[test] fn test_config_to_loop_detector_config() { let config = DetectionConfig { @@ -1460,4 +1392,23 @@ mod tests { assert_eq!(ldc.stop_threshold, 15); assert_eq!(ldc.window_size, 200); } + + #[test] + fn test_detection_config_default_has_no_max_response_history() { + let config = DetectionConfig { + loop_threshold: 3, + stop_threshold: 5, + ..Default::default() + }; + assert_eq!(config.loop_threshold, 3); + assert_eq!(config.stop_threshold, 5); + let default = DetectionConfig::default(); + assert_eq!(default.loop_threshold, 3); + assert_eq!(default.stop_threshold, 10); + assert!(default.enable_loop_detection); + assert_eq!(default.max_history, 100); + assert!((default.convergence_threshold - 0.95).abs() < f32::EPSILON); + assert_eq!(default.convergence_count, 3); + assert!(default.enable_convergence_detection); + } } diff --git a/src/engine/bare.rs b/src/engine/bare.rs index c4b1bdb..a19eb5d 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -1266,42 +1266,14 @@ mod tests { use parking_lot::Mutex; - // ================================================== - // Mock ApiClient - // ================================================== - - /// A mock API client that returns configurable responses. - /// - /// Used exclusively in tests. Stores a queue of response vectors, - /// where each response is a `Vec`. Each call to - /// [`stream_messages`](MockClient::stream_messages) pops the next - /// response from the front of the queue. - /// - /// When the queue is empty, `stream_messages` returns a single - /// [`ApiError`] — this lets tests verify error-handling paths. #[derive(Clone)] struct MockClient { - /// Responses to return, in order. - /// - /// Each entry is a `Vec` representing one complete - /// streaming response from the API. Popped from the front by - /// [`stream_messages`](MockClient::stream_messages). responses: Arc>>>, - /// Model name reported by [`ApiClient::model()`]. - /// - /// Copied into mock response metadata so that assertions can - /// verify which model produced a given response. model_name: Arc>, } impl MockClient { - /// Create a new mock client with the given model name. - /// - /// The response queue starts empty. Add responses with - /// [`add_text_response()`](MockClient::add_text_response), - /// [`add_tool_then_text()`](MockClient::add_tool_then_text), or - /// [`add_tool_only_response()`](MockClient::add_tool_only_response). fn new(model: &str) -> Self { Self { responses: Arc::new(Mutex::new(Vec::new())), @@ -1309,17 +1281,6 @@ mod tests { } } - /// Add a simple text response that ends the turn. - /// - /// Generates a complete stream: `MessageStart` → - /// `PartStart(text)` → `IndexedDelta(text)` → - /// `PartStop` → `MessageDelta(end_turn, usage)` → - /// `MessageStop`. The model will emit no tool calls, so the - /// loop terminates after this response. - /// - /// # Parameters - /// - /// - `text` — The text the model will "say". fn add_text_response(&self, text: &str) { let events = vec![ StreamEvent::MessageStart(MessageStart { @@ -1351,25 +1312,10 @@ mod tests { self.responses.lock().push(events); } - /// Add a raw sequence of stream events as a single response turn. fn add_events(&self, events: Vec) { self.responses.lock().push(events); } - /// Add a tool_call response followed by an end_turn response. - /// - /// The first response contains a single `tool_call` content part - /// (causing the loop to dispatch the tool), and the second - /// response is a plain text `end_turn` (causing the loop to - /// terminate). This simulates the common two-turn pattern: - /// model requests tool → model sees result → model responds. - /// - /// # Parameters - /// - /// - `tool_id` — Unique ID for the tool call. - /// - `tool_name` — Name of the tool to invoke. - /// - `tool_input` — JSON input for the tool. - /// - `final_text` — Text the model says after seeing the result. fn add_tool_then_text( &self, tool_id: &str, @@ -1432,18 +1378,6 @@ mod tests { self.responses.lock().push(text_events); } - /// Add a tool_call-only response (no end_turn). - /// - /// The response contains a single `tool_call` content part with - /// stop reason `tool_call`. After the tool is dispatched, the - /// loop will request another turn from the API. Useful for - /// testing max-turns and multi-turn tool chains. - /// - /// # Parameters - /// - /// - `tool_id` — Unique ID for the tool call. - /// - `tool_name` — Name of the tool to invoke. - /// - `tool_input` — JSON input for the tool. fn add_tool_only_response(&self, tool_id: &str, tool_name: &str, tool_input: Value) { let tool_events = vec![ StreamEvent::MessageStart(MessageStart { @@ -1469,12 +1403,6 @@ mod tests { self.responses.lock().push(tool_events); } - /// Add an error response. Reserved for future use. - /// - /// Currently pushes an incomplete response (only `MessageStart`) - /// which does not trigger an error on its own. Reserved for - /// testing streaming-error scenarios once the accumulator - /// handles partial messages. #[expect(dead_code)] fn add_error_response(&self) { // Return an empty response that will cause the stream to error @@ -1491,7 +1419,6 @@ mod tests { } impl ApiClient for MockClient { - /// Return the model name configured at construction. fn model(&self) -> String { self.model_name.lock().clone() } @@ -1504,11 +1431,6 @@ mod tests { true } - /// Pop the next queued response and return it as a stream. - /// - /// If the response queue is empty, returns a single-element - /// stream containing an [`ApiError`]. This enables tests that - /// exhaust all responses to verify error handling. fn stream_messages( &self, _messages: Vec, @@ -1528,10 +1450,6 @@ mod tests { } } - /// Non-streaming message creation — not used in these tests. - /// - /// Returns an empty JSON object. The `BareLoop` tests - /// exercise only the streaming path. fn create_message( &self, _messages: Vec, @@ -1544,11 +1462,6 @@ mod tests { // Helper trait for Vec-like pop_front on Vec trait PopFront { - /// Remove and return the first element, shifting the rest left. - /// - /// Returns `None` if the vector is empty. O(n) - /// operation because it calls `Vec::remove(0)`. Acceptable - /// for test-only code with small queues. fn pop_front(&mut self) -> Option; } @@ -1562,31 +1475,17 @@ mod tests { } } - // ================================================== - // Mock Tool - // ================================================== - - /// A test tool that echoes back its input. - /// - /// Implements [`Tool`] with a single `message` string parameter. - /// Returns `ToolOutput::text(format!("Echo: {msg}"))` so callers - /// can verify round-trip data flow. struct EchoTool; impl Tool for EchoTool { - /// Return the tool name `"echo"`. fn name(&self) -> &'static str { "echo" } - /// Return a human-readable description. fn description(&self) -> &'static str { "Echoes back the input" } - /// Return the JSON schema for this tool. - /// - /// Requires a single string property `message`. fn schema(&self) -> ToolSchema { ToolSchema { tool: "echo".into(), @@ -1599,10 +1498,6 @@ mod tests { } } - /// Execute the tool: extract `message` from input and echo it. - /// - /// If the `message` field is missing or not a string, defaults - /// to an empty string. fn call( &self, input: Value, @@ -1617,27 +1512,17 @@ mod tests { } } - /// A test tool that always fails. - /// - /// Used to verify that tool-execution errors are handled gracefully: - /// the loop should record the error as a soft tool result and - /// continue, not abort the session. struct FailingTool; impl Tool for FailingTool { - /// Return the tool name `"fail"`. fn name(&self) -> &'static str { "fail" } - /// Return a human-readable description. fn description(&self) -> &'static str { "Always fails" } - /// Return the JSON schema for this tool. - /// - /// Accepts an empty object (no parameters). fn schema(&self) -> ToolSchema { ToolSchema { tool: "fail".into(), @@ -1646,11 +1531,6 @@ mod tests { } } - /// Execute the tool: always returns an execution error. - /// - /// Returns [`ToolError::Execution`] with a fixed message so - /// tests can assert on the error path without triggering - /// panics or unwinds. fn call( &self, _input: Value, @@ -1660,33 +1540,16 @@ mod tests { } } - // ================================================== - // Counting Plugin (test helper) - // ================================================== - - /// A [`LoopObserver`](crate::observer::LoopObserver) that counts - /// how many times each hook fires. - /// - /// Uses [`AtomicUsize`] counters with `SeqCst` ordering so that - /// test assertions can read the counts from any thread after the - /// agent loop completes. struct CountingObserver { - /// Number of times `on_session_start` was called. session_starts: AtomicUsize, - /// Number of times `on_session_end` was called. session_ends: AtomicUsize, - /// Number of times `on_turn_start` was called. turn_starts: AtomicUsize, - /// Number of times `on_turn_end` was called. turn_ends: AtomicUsize, - /// Number of times `on_tool_pre` was called. tool_pres: AtomicUsize, - /// Number of times `on_tool_post` was called. tool_posts: AtomicUsize, } impl CountingObserver { - /// Create a new observer with all counters initialized to zero. fn new() -> Self { Self { session_starts: AtomicUsize::new(0), @@ -1729,15 +1592,6 @@ mod tests { } } - // ================================================== - // Test Helpers - // ================================================== - - /// Create a default [`LoopConfig`] with `max_turns = 10`. - /// - /// Most tests use this as a baseline. Tests that need a different - /// max-turns value mutate the returned config before constructing - /// the loop. fn make_config() -> LoopConfig { LoopConfig { max_turns: 10, @@ -1745,12 +1599,6 @@ mod tests { } } - // ================================================== - // Tests: Basic lifecycle - // ================================================== - - /// Verify that a single-turn session (text response, no tool calls) - /// completes successfully and returns the model's text output. #[tokio::test] async fn test_bare_loop_single_turn() { let client = MockClient::new("test-model"); @@ -1765,8 +1613,6 @@ mod tests { assert_eq!(result.final_output.as_deref(), Some("Hello! I'm done.")); } - /// Verify that a two-turn session (tool call → tool result → end turn) - /// completes successfully and records the tool invocation. #[tokio::test] async fn test_bare_loop_with_tool_call() { let client = MockClient::new("test-model"); @@ -1789,8 +1635,6 @@ mod tests { assert_eq!(result.tool_calls, 1); } - /// Verify that exceeding `max_turns` returns - /// [`LoopError::MaxTurnsExceeded`] and reports `success = false`. #[tokio::test] async fn test_bare_loop_max_turns_exceeded() { let client = MockClient::new("test-model"); @@ -1818,8 +1662,6 @@ mod tests { } } - /// Verify that calling [`cancel()`](BareLoop::cancel) mid-session - /// returns [`LoopError::Cancelled`]. #[tokio::test] async fn test_bare_loop_cancellation() { let client = MockClient::new("test-model"); @@ -1840,8 +1682,6 @@ mod tests { } } - /// Verify that an API error during streaming propagates as - /// [`LoopError::Api`] and marks the session as failed. #[tokio::test] async fn test_bare_loop_api_error() { // The mock will return an error @@ -1856,13 +1696,6 @@ mod tests { } } - // ================================================== - // Tests: Tool dispatch - // ================================================== - - /// Verify that requesting a tool not present in the registry produces - /// a soft error result (not a hard [`LoopError`]), allowing the model - /// to see the failure and adapt. #[tokio::test] async fn test_tool_not_found_returns_error_result() { let client = MockClient::new("test-model"); @@ -1879,8 +1712,6 @@ mod tests { assert_eq!(result.total_turns, 2); } - /// Verify that a tool returning an execution error produces a soft - /// error result and the session continues to completion. #[tokio::test] async fn test_tool_execution_failure() { let client = MockClient::new("test-model"); @@ -1897,12 +1728,6 @@ mod tests { assert_eq!(result.total_turns, 2); } - // ================================================== - // Tests: Observers - // ================================================== - - /// Verify that a single-turn session fires `session_start`, - /// `turn_start`, `turn_end`, and `session_end` on the observer. #[tokio::test] async fn test_observer_lifecycle_events() { let client = MockClient::new("test-model"); @@ -1922,12 +1747,6 @@ mod tests { assert_eq!(plugin.turn_ends.load(Ordering::SeqCst), 1); } - /// Verify that a tool-using session fires `on_tool_pre` and - /// `on_tool_post` observer hooks in addition to the turn hooks. - /// - /// A two-turn session (tool_call + end_turn) should produce: - /// - 2 turn starts, 2 turn ends - /// - 1 tool pre, 1 tool post #[tokio::test] async fn test_observer_tool_events() { let client = MockClient::new("test-model"); @@ -1950,13 +1769,6 @@ mod tests { assert_eq!(plugin.turn_ends.load(Ordering::SeqCst), 2); } - // ================================================== - // Tests: Conversation management - // ================================================== - - /// Verify that `extract_tool_calls` returns an empty list for a - /// text-only message and correctly parses tool_call parts when - /// present. #[tokio::test] async fn test_conversation_built_correctly() { let client = MockClient::new("test-model"); @@ -1993,9 +1805,6 @@ mod tests { assert_eq!(tool_calls[0].tool, "echo"); } - /// Verify that `build_tool_result_message` produces a user message - /// with the correct `tool_result` content parts, including the - /// tool_call_id, output text, and is_error flag. #[tokio::test] async fn test_tool_result_message_format() { let results = vec![super::ToolDispatchResult { @@ -2025,16 +1834,6 @@ mod tests { } } - // ================================================== - // Tests: Multiple tools in one turn - // ================================================== - - /// Verify that multiple tool_call parts in a single assistant message - /// are all dispatched and counted. - /// - /// The mock emits two `tool_call` parts in one response, followed by - /// an `end_turn` response. The session should report 2 turns and 2 - /// tool calls. #[tokio::test] async fn test_multiple_tool_calls_in_one_turn() { let client = MockClient::new("test-model"); @@ -2092,12 +1891,6 @@ mod tests { assert_eq!(result.tool_calls, 2); } - // ================================================== - // Tests: text streamer callback - // ================================================== - - /// Verify that `set_text_streamer` fires the callback for each text - /// delta during streaming. #[tokio::test] async fn test_text_streamer_fires_on_text_delta() { let client = MockClient::new("test-model"); @@ -2121,7 +1914,6 @@ mod tests { ); } - /// Verify that a run works fine without a text streamer set. #[tokio::test] async fn test_text_streamer_none_works() { let client = MockClient::new("test-model"); @@ -2132,8 +1924,6 @@ mod tests { assert!(result.success); } - /// Verify the streamer only fires for text deltas, not for tool-call - /// deltas or metadata events. #[tokio::test] async fn test_text_streamer_ignores_non_text_deltas() { let client = MockClient::new("test-model"); @@ -2191,12 +1981,6 @@ mod tests { assert_eq!(&*received, "Done", "only text deltas should fire streamer"); } - // ================================================== - // Tests: Accessors - // ================================================== - - /// Verify that accessor methods return the values passed at - /// construction. #[test] fn test_accessors() { let client = MockClient::new("test-model"); @@ -2209,8 +1993,6 @@ mod tests { assert!(!agent.is_cancelled()); } - /// Verify that `cancel_signal()` returns a shared reference to the - /// same signal used by `cancel()` and `is_cancelled()`. #[test] fn test_cancel_signal_shared() { let client = MockClient::new("test-model"); @@ -2224,12 +2006,6 @@ mod tests { assert!(agent.is_cancelled()); } - // ================================================== - // Tests: Session result fields - // ================================================== - - /// Verify that the returned [`SessionResult`] has the correct - /// session ID, positive duration, and non-zero token count. #[tokio::test] async fn test_session_result_fields() { let client = MockClient::new("test-model"); @@ -2245,12 +2021,6 @@ mod tests { assert!(result.input_tokens > 0 || result.output_tokens > 0); // from mock usage } - // ================================================== - // Tests: Property — loop always terminates - // ================================================== - - /// Verify that setting `max_turns = 1` still allows a single-turn - /// session to complete normally. #[tokio::test] async fn test_loop_terminates_with_max_turns_1() { let client = MockClient::new("test-model"); @@ -2266,8 +2036,6 @@ mod tests { assert_eq!(result.total_turns, 1); } - /// Verify that setting `max_turns = 0` immediately triggers a - /// configuration error before any API call. #[tokio::test] async fn test_loop_terminates_with_max_turns_0() { let client = MockClient::new("test-model"); @@ -2285,13 +2053,6 @@ mod tests { } } - // ================================================== - // Tests: Error in tool-not-found returns error result, not hard error - // ================================================== - - /// Verify that requesting a nonexistent tool produces a soft error - /// result (not a hard [`LoopError`]), allowing the model to see - /// the error and respond gracefully. #[tokio::test] async fn test_tool_error_is_soft_not_hard() { let client = MockClient::new("test-model"); @@ -2329,13 +2090,80 @@ mod tests { assert!(result.success); } - // ================================================== - // ================================================== - // Recovery wiring tests - // ================================================== + #[tokio::test] + async fn test_loop_detection_hard_stop_propagates_loop_error() { + use crate::detection::{DetectionConfig, DetectionManager}; + use crate::runtime::LoopRuntime; + + let mut registry = ToolRegistry::new(); + registry.register(EchoTool); + let client = MockClient::new("test"); + for i in 0..10 { + client.add_tool_only_response(&format!("call_{i}"), "echo", json!({ "message": "hi" })); + } + + let runtime = LoopRuntime::new().with_detection( + DetectionManager::new_with_config(DetectionConfig { + loop_threshold: 2, + stop_threshold: 2, + ..Default::default() + }) + .expect("valid detection config"), + ); + + let mut agent = + BareLoop::new_with_managers(Arc::new(client), registry, make_config(), runtime); + let result = agent.run("test").await; + + assert!( + matches!(result, Err(LoopError::LoopDetected { .. })), + "expected Err(LoopError::LoopDetected), got {result:?}" + ); + } + + #[tokio::test] + async fn test_loop_detection_soft_block_before_stop_threshold() { + use crate::detection::{DetectionConfig, DetectionManager}; + use crate::runtime::LoopRuntime; + + let mut registry = ToolRegistry::new(); + registry.register(EchoTool); + let client = MockClient::new("test"); + client.add_tool_only_response("c1", "echo", json!({ "message": "hi" })); + client.add_tool_only_response("c2", "echo", json!({ "message": "hi" })); + client.add_text_response("Done"); + + let runtime = LoopRuntime::new().with_detection( + DetectionManager::new_with_config(DetectionConfig { + loop_threshold: 2, + stop_threshold: 10, + ..Default::default() + }) + .expect("valid detection config"), + ); + + let mut agent = + BareLoop::new_with_managers(Arc::new(client), registry, make_config(), runtime); + let result = agent.run("test").await; + + assert!(result.is_ok(), "expected Ok, got {result:?}"); + } + + #[tokio::test] + async fn test_cancelled_before_run_returns_cancelled() { + let client = MockClient::new("test"); + client.add_text_response("Hello"); + + let mut agent = BareLoop::new(Arc::new(client), ToolRegistry::new(), make_config()); + agent.cancel(); + let result = agent.run("test").await; + + assert!( + matches!(result, Err(LoopError::Cancelled)), + "expected Err(LoopError::Cancelled), got {result:?}" + ); + } - /// Verify the default `NoopReflector` + `ExponentialBackoffRecovery` - /// wiring returns soft errors (no infinite loop, no panic). #[tokio::test] async fn test_default_recovery_on_tool_error_returns_soft_result() { let mut registry = ToolRegistry::new(); @@ -2351,8 +2179,6 @@ mod tests { assert_eq!(result.tool_calls, 1); } - /// Verify that when a tool is not found, the recovery wiring still - /// produces a soft error result (no hard error propagated). #[tokio::test] async fn test_recovery_on_missing_tool_returns_soft_result() { let client = MockClient::new("test"); @@ -2365,9 +2191,6 @@ mod tests { assert_eq!(result.tool_calls, 1); } - /// Verify that a failing tool with the default recovery produces - /// exactly one tool dispatch (NoopReflector marks everything as - /// non-recoverable, so no retries). #[tokio::test] async fn test_recovery_noop_reflector_no_retries() { let mut registry = ToolRegistry::new(); @@ -2383,7 +2206,6 @@ mod tests { assert_eq!(result.tool_calls, 1); } - /// Verify cancellation is still respected during tool recovery. #[tokio::test] async fn test_recovery_respects_cancellation() { let mut registry = ToolRegistry::new(); @@ -2401,13 +2223,6 @@ mod tests { assert!(result.is_err()); } - // ================================================== - // Tests: set_pipeline uses self.tools registry - // ================================================== - - /// Verify that `set_pipeline` automatically injects `self.tools` as the - /// pipeline's core registry, so dispatch never diverges from schema - /// generation. #[tokio::test] async fn test_set_pipeline_injects_self_tools_registry() { let client = MockClient::new("test-model"); @@ -2425,11 +2240,6 @@ mod tests { assert!(result.unwrap().success); } - // ================================================== - // Tests: turn_number is actual turn index - // ================================================== - - /// A middleware that records the `turn_number` from each dispatch context. struct TurnNumberCapture { turns: Arc>>, } @@ -2459,8 +2269,6 @@ mod tests { } } - /// Verify that `turn_number` reflects the actual turn index, not - /// `config.max_turns`. #[tokio::test] async fn test_turn_number_is_actual_turn_index() { let client = MockClient::new("test-model"); @@ -2497,10 +2305,6 @@ mod tests { ); } - // ─── Model switching tests ─── - - /// `switch_model` updates config.model, the client's model, and the - /// fallback manager's original-model tracker. #[tokio::test] async fn switch_model_updates_config_and_client() { let client = MockClient::new("model-a"); @@ -2520,7 +2324,6 @@ mod tests { assert_eq!(client_arc.model(), "model-b"); } - /// `switch_model` fires `on_model_switched` to all registered observers. #[tokio::test] async fn switch_model_notifies_observers() { #[derive(Default)] @@ -2557,12 +2360,8 @@ mod tests { assert_eq!(recorded[1], ("m2".to_string(), "m3".to_string())); } - /// `switch_model` updates config even when the client doesn't support - /// hot-swapping. The client's internal model stays the same, but the - /// framework-level config, fallback, and observers are updated. #[tokio::test] async fn switch_model_unsupported_client() { - /// A client that does NOT override `set_model` (returns `false`). struct StaticClient { model_name: Arc>, } @@ -2621,8 +2420,6 @@ mod tests { assert_eq!(loop_.client.model(), "static"); } - /// `switch_model` syncs the fallback manager's original-model tracker - /// so subsequent fallback decisions compare against the new primary. #[tokio::test] async fn switch_model_updates_fallback_original() { let client = std::sync::Arc::new(MockClient::new("primary")); @@ -2644,7 +2441,6 @@ mod tests { ); } - /// `switch_model` rejects empty/whitespace-only model names. #[tokio::test] async fn switch_model_rejects_empty() { let client = std::sync::Arc::new(MockClient::new("model")); @@ -2662,7 +2458,6 @@ mod tests { assert_eq!(loop_.config().model, "default"); } - /// `switch_model` can be called multiple times in succession. #[tokio::test] async fn switch_model_chained() { let client = std::sync::Arc::new(MockClient::new("a")); @@ -2679,7 +2474,6 @@ mod tests { assert_eq!(loop_.config().model, "d"); } - /// `switch_model` with `.context_window()` updates the config. #[tokio::test] async fn switch_model_updates_context_window() { let client = std::sync::Arc::new(MockClient::new("big-model")); @@ -2699,7 +2493,6 @@ mod tests { assert_eq!(loop_.config().context_window, 8192); } - /// `switch_model` with `.max_tokens()` updates the config. #[tokio::test] async fn switch_model_updates_max_tokens() { let client = std::sync::Arc::new(MockClient::new("m")); @@ -2712,7 +2505,6 @@ mod tests { assert_eq!(loop_.config().max_tokens, 4096); } - /// `switch_model` trims whitespace from the model name. #[tokio::test] async fn switch_model_trims_whitespace() { let client = std::sync::Arc::new(MockClient::new("m")); @@ -2723,7 +2515,6 @@ mod tests { assert_eq!(loop_.config().model, "gpt-4o"); } - /// `switch_model` resets the fallback circuit breaker. #[tokio::test] async fn switch_model_resets_fallback_circuit() { use crate::fallback::FallbackState; diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 9256549..b8e1b69 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -127,7 +127,7 @@ impl BareLoop { return Ok(blocked); } - if let Some(blocked) = self.pre_detection(&tc, turn_idx) { + if let Some(blocked) = self.pre_detection(&tc, turn_idx)? { self.managers.observers().on_tool_post(&ToolPostContext { turn: turn_idx, tool: tc.tool.clone(), @@ -183,12 +183,15 @@ impl BareLoop { /// Check for a loop pattern before executing the tool. /// - /// Extracts the primary parameter from the tool input using the - /// configured [`ToolSignature`], records the call with the detection - /// manager, and returns a soft-error result if the same operation - /// has exceeded the loop threshold. Returns `None` when dispatch should - /// proceed normally. - fn pre_detection(&self, tc: &ToolCall, turn_idx: usize) -> Option { + /// # Errors + /// + /// Returns [`LoopError`] when the detection manager signals a hard + /// stop (e.g. [`LoopError::LoopDetected`]). + fn pre_detection( + &self, + tc: &ToolCall, + turn_idx: usize, + ) -> Result, LoopError> { let operation = Operation::from_input_with_signature( &tc.tool, &tc.input, @@ -196,23 +199,21 @@ impl BareLoop { ); let pattern = self.managers.detection.record_operation(operation); - // Check inline detection - let inline_blocked = self - .managers - .handle_detected_pattern(&pattern, turn_idx) - .map(|_result| ToolDispatchResult { + // Check inline detection. When the pattern triggers a hard stop + // (Err), propagate it so the agent loop terminates. When it + // produces a soft block (Ok), return the soft-error result so + // the model can see the warning and try a different approach. + match self.managers.handle_detected_pattern(&pattern, turn_idx) { + Some(Err(e)) => Err(e), + Some(Ok(_)) => Ok(Some(ToolDispatchResult { tool_call_id: tc.id.clone(), output: ToolContent::Text("loop detected: aborting tool dispatch".into()), is_error: true, duration: Duration::ZERO, resolved_tool_name: tc.tool.clone(), - }); - - if inline_blocked.is_some() { - return inline_blocked; + })), + None => Ok(None), } - - None } /// Record the tool result with the detection manager (post-execution). @@ -492,6 +493,13 @@ impl BareLoop { return Err(LoopError::Cancelled); } }; + // Guard against a late-arriving cancellation that races with the + // pipeline future resolving first. Without this check the result + // would be treated as a soft tool error by ToolCallMiddleware + // instead of a hard cancellation. + if cancel.is_cancelled() { + return Err(LoopError::Cancelled); + } Ok(ToolDispatchResult { tool_call_id: if dispatch_result.tool_call_id.is_empty() { tc.id.clone() diff --git a/src/fallback.rs b/src/fallback.rs index 2753e3e..c62482d 100644 --- a/src/fallback.rs +++ b/src/fallback.rs @@ -303,7 +303,13 @@ impl FallbackEntry { /// ``` #[must_use] pub fn with_max_fail_count(mut self, max_fail_count: usize) -> Self { - self.max_fail_count = max_fail_count.max(1); + let new_max = max_fail_count.max(1); + // If the entry was already failed, add attempts to keep it failed + // under the new threshold. + while self.attempts.len() < new_max && self.failed() { + self.attempts.push(AttemptRecord::anonymous()); + } + self.max_fail_count = new_max; self } @@ -1905,13 +1911,6 @@ mod tests { use super::*; #[test] - /// Verify that a freshly-constructed manager starts in [`FallbackState::Primary`] - /// with no failures and no fallback activation. - /// - /// Asserts initial values for [`state()`](FallbackManager::state), - /// [`is_using_fallback()`](FallbackManager::is_using_fallback), - /// [`is_fallback_active()`](FallbackManager::is_fallback_active), and - /// [`consecutive_failures()`](FallbackManager::consecutive_failures). fn test_initial_state() { let mgr = FallbackManager::new(3, 2); assert_eq!(mgr.state(), FallbackState::Primary); @@ -1921,11 +1920,6 @@ mod tests { } #[test] - /// Verify that [`record_api_failure`](FallbackManager::record_api_failure) returns - /// `true` only when the failure count first reaches the threshold. - /// - /// Calls [`record_api_failure`](FallbackManager::record_api_failure) three times - /// and asserts that only the third call returns `true`. fn test_failure_threshold() { let mgr = FallbackManager::new(3, 2); assert!(!mgr.record_api_failure()); // 1 @@ -1934,10 +1928,6 @@ mod tests { } #[test] - /// Verify that [`record_model_failure`](FallbackManager::record_model_failure) - /// transitions to [`FallbackState::Fallback`] when the threshold is reached. - /// - /// Also checks that the method returns `true` only on the threshold-crossing call. fn test_model_failure_triggers_fallback() { let mgr = FallbackManager::new(3, 2); assert!(!mgr.record_model_failure()); // 1 @@ -1947,10 +1937,6 @@ mod tests { } #[test] - /// Verify the full recovery cycle: Primary → Fallback → Recovering → Primary. - /// - /// Trips the circuit with three failures, then transitions to recovering - /// and records two successes to close the circuit back to primary. fn test_recovery() { let mgr = FallbackManager::new(3, 2); // Trigger fallback @@ -1970,11 +1956,6 @@ mod tests { } #[test] - /// Verify that a failure during [`FallbackState::Recovering`] reopens the circuit. - /// - /// After transitioning to recovering, a single call to - /// [`record_model_failure`](FallbackManager::record_model_failure) should - /// move the state back to [`FallbackState::Fallback`]. fn test_recovery_failure_goes_back_to_fallback() { let mgr = FallbackManager::new(3, 2); for _ in 0..3 { @@ -1986,12 +1967,6 @@ mod tests { } #[test] - /// Verify [`should_try_resume_primary`](FallbackManager::should_try_resume_primary) - /// enforces both the state check and the cooldown duration. - /// - /// Asserts `false` when in [`FallbackState::Primary`], `false` when in - /// fallback but the cooldown hasn't elapsed, and `true` when the cooldown - /// has passed (using a 0-second timeout). fn test_should_try_resume_primary() { let mgr = FallbackManager::new(3, 2); assert!(!mgr.should_try_resume_primary(Duration::from_secs(10))); @@ -2007,12 +1982,6 @@ mod tests { } #[test] - /// Verify that [`new_with_fallback`](FallbackManager::new_with_fallback) - /// starts in [`FallbackState::Fallback`] with the model name stored. - /// - /// Checks [`is_fallback_active()`](FallbackManager::is_fallback_active), - /// [`is_using_fallback()`](FallbackManager::is_using_fallback), and - /// [`original_model()`](FallbackManager::original_model). fn test_new_with_fallback() { let mgr = FallbackManager::new_with_fallback("llm-70b".into(), 3); assert!(mgr.is_fallback_active()); @@ -2021,11 +1990,6 @@ mod tests { } #[test] - /// Verify that [`reset`](FallbackManager::reset) clears all state back to - /// [`FallbackState::Primary`], including counters and flags. - /// - /// Trips the circuit first, then asserts that [`reset`] restores - /// every field to its initial value. fn test_reset() { let mgr = FallbackManager::new(3, 2); for _ in 0..3 { @@ -2040,11 +2004,6 @@ mod tests { } #[test] - /// Verify that [`record_api_failure`](FallbackManager::record_api_failure) - /// does not re-trip the circuit after it has already been activated. - /// - /// The [`fallback_activated`](FallbackManager::is_fallback_active) flag - /// prevents the sticky `true` return on every subsequent failure. fn test_api_failure_does_not_retrip() { let mgr = FallbackManager::new(3, 2); // Trip the circuit @@ -2060,10 +2019,6 @@ mod tests { } #[test] - /// Verify that [`record_model_success`](FallbackManager::record_model_success) - /// resets the failure counter when in [`FallbackState::Primary`]. - /// - /// After recording two failures, a single success should zero the counter. fn test_record_success_resets_on_primary() { let mgr = FallbackManager::new(3, 2); mgr.record_api_failure(); @@ -2075,10 +2030,6 @@ mod tests { } #[test] - /// Verify [`for_model`](FallbackManager::for_model) stores the model name. - /// - /// Asserts that [`FallbackManager::original_model`] returns the provided - /// model string and that the initial state is [`FallbackState::Primary`]. fn test_for_model() { let mgr = FallbackManager::for_model("llm-70b"); assert_eq!(mgr.original_model(), Some("llm-70b".to_string())); @@ -2086,11 +2037,6 @@ mod tests { } #[test] - /// Verify thread safety — concurrent reads and writes via `Arc`. - /// - /// Spawns 10 threads that each call [`record_api_failure`], [`record_model_success`], - /// [`state`], and [`consecutive_failures`] concurrently. The test passes if no - /// thread panics due to data races (guaranteed by atomic operations). fn test_concurrent_access() { use std::sync::Arc; use std::thread; @@ -2113,10 +2059,6 @@ mod tests { } } - /// Verify the consolidated Mutex (M5 fix) ensures multi-field updates - /// are visible atomically. When `transition_to_fallback` is called, - /// both `fallback_switched_at` and `active_fallback` should be - /// observable together. #[test] fn test_consolidated_mutex_fields_are_consistent() { let mgr = FallbackManager::for_model("primary-model"); @@ -2138,7 +2080,6 @@ mod tests { ); } - /// Verify that `transition_to_primary` clears both fields atomically (M5). #[test] fn test_consolidated_mutex_clears_fields_together() { let mgr = FallbackManager::for_model("primary-model"); @@ -2159,7 +2100,6 @@ mod tests { ); } - /// Verify that `reset()` clears all fields atomically (M5). #[test] fn test_consolidated_mutex_reset_clears_all() { let mgr = FallbackManager::for_model("primary-model"); @@ -2178,4 +2118,57 @@ mod tests { assert_eq!(mgr.consecutive_failures(), 0); assert!(mgr.fallback_switched_at().is_none()); } + + #[test] + fn with_max_fail_count_no_padding_when_not_failed() { + let entry = FallbackEntry::new("model-a").with_max_fail_count(5); + assert!(!entry.failed()); + assert_eq!(entry.attempt_count(), 0); + assert_eq!(entry.max_fail_count, 5); + } + + #[test] + fn with_max_fail_count_pads_already_failed_entry() { + let mut entry = FallbackEntry::new("model-b"); + entry.record_attempt("timeout"); + entry.record_attempt("timeout"); + assert!(entry.failed()); + assert_eq!(entry.attempt_count(), 2); + + let entry = entry.with_max_fail_count(5); + assert_eq!(entry.max_fail_count, 5); + assert_eq!(entry.attempt_count(), 5); + assert!(entry.failed()); + } + + #[test] + fn with_max_fail_count_pads_exactly_to_new_threshold() { + let mut entry = FallbackEntry::new("model-c"); + entry.record_attempt("err"); + entry.record_attempt("err"); + assert!(entry.failed()); + + let entry = entry.with_max_fail_count(3); + assert_eq!(entry.attempt_count(), 3); + assert!(entry.failed()); + } + + #[test] + fn with_max_fail_count_no_padding_when_lowering() { + let mut entry = FallbackEntry::new("model-d"); + entry.record_attempt("err"); + entry.record_attempt("err"); + assert!(entry.failed()); + + let entry = entry.with_max_fail_count(1); + assert_eq!(entry.max_fail_count, 1); + assert_eq!(entry.attempt_count(), 2); + assert!(entry.failed()); + } + + #[test] + fn with_max_fail_count_clamps_to_minimum_one() { + let entry = FallbackEntry::new("model-e").with_max_fail_count(0); + assert_eq!(entry.max_fail_count, 1); + } } diff --git a/src/message.rs b/src/message.rs index fdb8b4a..2b41793 100644 --- a/src/message.rs +++ b/src/message.rs @@ -1024,10 +1024,6 @@ impl ToolContentPart { mod tests { use super::*; - /// Verify that [`Message::user`] produces a message with [`Role::User`] - /// and a single [`MessagePart::Text`] containing the provided string. - /// - /// Also checks that [`MessagePart::as_text`] returns the original text. #[test] fn test_message_user_shortcut() { let msg = Message::user("Hello"); @@ -1036,10 +1032,6 @@ mod tests { assert_eq!(msg.parts[0].as_text(), Some("Hello")); } - /// Verify [`Message::assistant`] produces a message with the correct role. - /// - /// Asserts that the returned [`Message`] has [`Role::Assistant`] and contains - /// exactly one [`MessagePart::Text`] with the provided content. #[test] fn test_message_assistant_shortcut() { let msg = Message::assistant("Hi there!"); @@ -1047,20 +1039,12 @@ mod tests { assert_eq!(msg.parts.len(), 1); } - /// Verify [`Message`]'s [`Display`](fmt::Display) renders plain-text content. - /// - /// For a user message containing only [`MessagePart::Text`], the output - /// should be the raw text with no prefix or decoration. #[test] fn test_message_display() { let msg = Message::user("Hello world"); assert_eq!(msg.to_string(), "Hello world"); } - /// Verify [`Message`]'s [`Display`](fmt::Display) renders tool-call parts. - /// - /// For a message containing [`MessagePart::ToolCall`], the output should - /// include the tool name (e.g., "Tool: read_file") in the formatted string. #[test] fn test_message_display_with_tool_call() { let msg = Message { @@ -1075,20 +1059,12 @@ mod tests { assert!(display.contains("Tool: read_file")); } - /// Verify [`Role`]'s [`Display`](fmt::Display) produces the expected strings. - /// - /// Each variant should render as its lowercase name: `"user"` or `"assistant"`. #[test] fn test_role_display() { assert_eq!(Role::User.to_string(), "user"); assert_eq!(Role::Assistant.to_string(), "assistant"); } - /// Verify [`MessagePart`] helper predicates and accessors. - /// - /// Tests [`is_text`](MessagePart::is_text), [`is_tool_call`](MessagePart::is_tool_call), - /// [`is_tool_result`](MessagePart::is_tool_result), and [`as_text`](MessagePart::as_text) - /// across all three part types. #[test] fn test_part_helpers() { let text = MessagePart::text("hello"); @@ -1105,10 +1081,6 @@ mod tests { assert!(tool_result.is_tool_result()); } - /// Verify [`ImageSource::new_base64`] sets the source type and media type. - /// - /// Asserts that `encoding` is `"base64"` and that the provided MIME type - /// is stored unchanged in the [`media_type`](ImageSource::media_type) field. #[test] fn test_image_source() { let src = ImageSource::new_base64("image/png", "iVBOR..."); @@ -1116,11 +1088,6 @@ mod tests { assert_eq!(src.media_type, "image/png"); } - /// Verify `From<&str>` for [`ToolContent`] produces a string variant. - /// - /// The conversion should wrap the provided text in - /// [`ToolContent::Text`] and [`Display`](std::fmt::Display) should - /// yield the original text. #[test] fn test_tool_result_from_string() { let result: ToolContent = "hello".into(); @@ -1128,19 +1095,12 @@ mod tests { assert_eq!(result.to_string(), "hello"); } - /// Verify that [`ToolContent::default`] produces an empty-string variant. - /// - /// Equivalent to `ToolContent::from_string("")`. #[test] fn test_tool_result_default() { let result = ToolContent::default(); assert!(result.is_string()); } - /// Verify [`ToolContentPart::text`] produces the [`Text`](ToolContentPart::Text) variant. - /// - /// The constructor should create a [`ToolContentPart::Text`] containing the - /// provided string, accessible via pattern matching on the variant. #[test] fn test_tool_result_part_text() { let part = ToolContentPart::text("output"); @@ -1150,10 +1110,6 @@ mod tests { } } - /// Verify that a [`Message`] round-trips through JSON serialization. - /// - /// Ensures that `serde_json::to_string` → `serde_json::from_str` preserves - /// the message [`Role`]. #[test] fn test_message_serialization() { let msg = Message::user("test"); @@ -1162,10 +1118,6 @@ mod tests { assert_eq!(msg.role, deserialized.role); } - /// Verify that a `Vec` round-trips through JSON serialization. - /// - /// Tests the `#[serde(tag = "type")]` representation for [`Text`](MessagePart::Text) - /// and [`ToolCall`](MessagePart::ToolCall) variants. #[test] fn test_part_serialization_roundtrip() { let parts = vec![ @@ -1177,11 +1129,6 @@ mod tests { assert_eq!(parts.len(), back.len()); } - /// Verify [`ToolContent`]'s [`Display`](fmt::Display) joins text parts. - /// - /// When a Multipart [`ToolContent`] contains multiple - /// [`ToolContentPart::Text`] entries, the [`Display`](fmt::Display) impl - /// should join them with newlines. #[test] fn test_tool_result_multipart_display() { let result = ToolContent::from_multipart(vec![ diff --git a/src/middleware.rs b/src/middleware.rs index a03f951..d8fca3e 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -424,11 +424,6 @@ mod tests { use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; - // ================================================== - // Test tools - // ================================================== - - /// A simple echo tool for testing. struct EchoTool; impl Tool for EchoTool { @@ -459,7 +454,6 @@ mod tests { } } - /// A tool that always errors. struct ErrorTool; impl Tool for ErrorTool { @@ -485,7 +479,6 @@ mod tests { } } - /// A slow tool for testing timeouts. struct SlowTool { delay_ms: u64, } @@ -517,10 +510,6 @@ mod tests { } } - // ================================================== - // Helpers - // ================================================== - fn test_registry() -> Arc { let mut reg = ToolRegistry::new(); reg.register(EchoTool); @@ -541,10 +530,6 @@ mod tests { } } - // ================================================== - // Pipeline builder tests - // ================================================== - #[test] fn test_builder_requires_core() { let result = ToolPipeline::builder().build(); @@ -584,10 +569,6 @@ mod tests { ); } - // ================================================== - // Core dispatch tests - // ================================================== - #[tokio::test] async fn test_core_dispatch_echo() { let pipeline = ToolPipeline::new(test_registry()); @@ -620,10 +601,6 @@ mod tests { assert_eq!(result.resolved_tool_name, "error_tool"); } - // ================================================== - // PermissionMiddleware tests - // ================================================== - #[tokio::test] async fn test_permission_deny_all() { let pipeline = ToolPipeline::builder() @@ -702,8 +679,6 @@ mod tests { assert!(result.is_error); } - /// Verify that PermissionMiddleware with Ask permission and no resolver - /// denies the tool call (M2 fix). #[tokio::test] async fn test_permission_ask_without_resolver_denies() { let mut ctx = test_ctx("echo"); @@ -730,10 +705,6 @@ mod tests { ); } - // ================================================== - // TimeoutMiddleware tests - // ================================================== - #[tokio::test] async fn test_timeout_fast_tool_succeeds() { let registry = { @@ -778,10 +749,6 @@ mod tests { } } - // ================================================== - // UnknownToolMiddleware tests - // ================================================== - #[test] fn test_similarity_identical() { let score = UnknownToolMiddleware::similarity("bash", "bash"); @@ -806,10 +773,6 @@ mod tests { assert!((UnknownToolMiddleware::similarity("a", "") - 0.0).abs() < f64::EPSILON); } - // ================================================== - // dispatch_all tests - // ================================================== - #[tokio::test] async fn test_dispatch_all_sequential() { let pipeline = ToolPipeline::new(test_registry()); @@ -837,10 +800,6 @@ mod tests { assert!(result.is_err()); } - // ================================================== - // Ordering tests - // ================================================== - #[tokio::test] async fn test_middleware_ordering_permission_before_timeout() { // Permission denies → timeout middleware is never reached @@ -880,10 +839,6 @@ mod tests { } } - // =================================================== - // Integration: full pipeline - // =================================================== - #[tokio::test] async fn test_full_pipeline_echo() { let registry = { @@ -939,11 +894,6 @@ mod tests { } } - // ================================================== - // Short-circuit test - // ================================================== - - /// A middleware that tracks whether it was reached. struct ReachTracker { reached: Arc, } @@ -983,13 +933,6 @@ mod tests { ); } - // =================================================== - // OutputLimitMiddleware tests - // =================================================== - - /// A tool that returns a string of repeated characters for testing - /// output truncation. The tool name encodes the repeat count as - /// `long_output_N` where N is the number of characters. struct LongOutputTool; impl Tool for LongOutputTool { @@ -1028,8 +971,6 @@ mod tests { } } - /// A tool that returns [`ToolContent::Multipart`] output for testing - /// that non-text output passes through unchanged. struct MultipartTool; impl Tool for MultipartTool { @@ -1295,8 +1236,6 @@ mod tests { ); } - /// Multi-byte UTF-8 text whose byte count exceeds `max_chars` but whose - /// character count does **not** must pass through un-truncated. #[tokio::test] async fn test_output_limit_multibyte_under_char_limit_not_truncated() { // "日" is 3 bytes. 5 repetitions = 5 chars, 15 bytes. @@ -1323,10 +1262,6 @@ mod tests { assert_eq!(result.output, ToolContent::Text("日日日日日".to_string())); } - // ================================================== - // ToolDispatchResult::from_result - // ================================================== - #[test] fn from_result_maps_ok_output() { let output = ToolOutput::text("hello"); @@ -1346,10 +1281,6 @@ mod tests { assert_eq!(result.resolved_tool_name, "missing_tool"); } - // ================================================== - // From + builder methods - // ================================================== - #[test] fn from_tool_output_defaults() { let output = ToolOutput::text("ok"); diff --git a/src/middleware/unknown_tool.rs b/src/middleware/unknown_tool.rs index b1a7bce..0cf4808 100644 --- a/src/middleware/unknown_tool.rs +++ b/src/middleware/unknown_tool.rs @@ -198,7 +198,6 @@ impl ToolMiddleware for UnknownToolMiddleware { ctx: &'a mut ToolDispatchContext, next: &'a ToolPipeline, ) -> Pin + Send + 'a>> { - let tool_name = ctx.tool_name.clone(); let registry_names = self.registry.tool_names(); let threshold = self.suggestion_threshold; @@ -206,10 +205,14 @@ impl ToolMiddleware for UnknownToolMiddleware { let mut result = next.dispatch(ctx).await; if Self::is_tool_not_found(&result) { + // Read the tool name from ctx *after* dispatch so that any + // redirection applied by downstream middleware is reflected + // in the suggestion lookup. + let tool_name = ctx.tool_name.as_str(); let available_refs: Vec<&str> = registry_names.iter().map(String::as_str).collect(); if let Some((suggestion, score)) = - Self::find_best_match_inner(&tool_name, &available_refs, threshold) + Self::find_best_match_inner(tool_name, &available_refs, threshold) { tracing::info!( requested = %tool_name, @@ -420,4 +423,135 @@ mod tests { }; assert!(!UnknownToolMiddleware::is_tool_not_found(&result)); } + + use crate::cancel::CancelSignal; + use crate::middleware::{ToolDispatchContext, ToolMiddleware, ToolPipeline}; + use crate::tool::{ + PermissionCheck, Tool, ToolContext, ToolError, ToolOutput, ToolRegistry, ToolSchema, + }; + use serde_json::{Value, json}; + use std::future::Future; + use std::pin::Pin; + use std::sync::Arc; + + struct ReadFileTool; + impl Tool for ReadFileTool { + fn name(&self) -> &'static str { + "read_file" + } + fn description(&self) -> &'static str { + "Reads a file" + } + fn schema(&self) -> ToolSchema { + ToolSchema { + tool: "read_file".into(), + description: "Reads a file".into(), + input_schema: json!({"type": "object", "properties": {}}), + } + } + fn call( + &self, + _input: Value, + _ctx: &ToolContext, + ) -> Pin> + Send + '_>> { + Box::pin(async { Ok(ToolOutput::text("ok")) }) + } + } + + fn make_registry() -> Arc { + let mut reg = ToolRegistry::new(); + reg.register(ReadFileTool); + Arc::new(reg) + } + + fn make_ctx(name: &str) -> ToolDispatchContext { + ToolDispatchContext { + tool_name: name.to_string(), + input: json!({}), + call_id: "call_1".to_string(), + turn_number: 1, + cancel: Arc::new(CancelSignal::new()), + permission: PermissionCheck::Allow, + tool_context: ToolContext::default(), + } + } + + struct RenameMiddleware { + new_name: String, + } + + impl ToolMiddleware for RenameMiddleware { + fn name(&self) -> &'static str { + "rename" + } + fn dispatch<'a>( + &'a self, + ctx: &'a mut ToolDispatchContext, + next: &'a ToolPipeline, + ) -> Pin + Send + 'a>> { + ctx.tool_name = self.new_name.clone(); + Box::pin(async move { next.dispatch(ctx).await }) + } + } + + #[tokio::test] + async fn dispatch_appends_suggestion_for_typo() { + let registry = make_registry(); + let pipeline = ToolPipeline::builder() + .with(UnknownToolMiddleware::new(Arc::clone(®istry))) + .core(registry) + .build() + .expect("valid pipeline"); + + let result = pipeline.invoke(make_ctx("read_fil")).await; + assert!(result.is_error); + match result.output { + ToolContent::Text(ref msg) => { + assert!(msg.contains("not found"), "message was: {msg}"); + assert!( + msg.contains("Did you mean 'read_file'?"), + "expected suggestion in message: {msg}" + ); + } + other @ ToolContent::Multipart(_) => panic!("expected Text output, got {other:?}"), + } + } + + #[tokio::test] + async fn dispatch_suggestion_uses_post_dispatch_tool_name() { + let registry = make_registry(); + let pipeline = ToolPipeline::builder() + .with(UnknownToolMiddleware::new(Arc::clone(®istry))) + .with(RenameMiddleware { + new_name: "read_fil".into(), + }) + .core(registry) + .build() + .expect("valid pipeline"); + + let result = pipeline.invoke(make_ctx("xyz")).await; + assert!(result.is_error); + match result.output { + ToolContent::Text(ref msg) => { + assert!( + msg.contains("Did you mean 'read_file'?"), + "suggestion should reflect the renamed tool_name, message was: {msg}" + ); + } + other @ ToolContent::Multipart(_) => panic!("expected Text output, got {other:?}"), + } + } + + #[tokio::test] + async fn dispatch_known_tool_no_suggestion() { + let registry = make_registry(); + let pipeline = ToolPipeline::builder() + .with(UnknownToolMiddleware::new(Arc::clone(®istry))) + .core(registry) + .build() + .expect("valid pipeline"); + + let result = pipeline.invoke(make_ctx("read_file")).await; + assert!(!result.is_error, "known tool should not error"); + } } diff --git a/src/observer.rs b/src/observer.rs index 289e411..0b5d1d4 100644 --- a/src/observer.rs +++ b/src/observer.rs @@ -343,7 +343,6 @@ mod tests { use super::*; use std::sync::atomic::{AtomicUsize, Ordering}; - /// A test observer that counts notification invocations. struct CountingObserver { name: &'static str, stream_success: AtomicUsize, diff --git a/src/provider.rs b/src/provider.rs index 19001fe..f839622 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -306,7 +306,6 @@ pub fn self_hosted(base_url: &str, model: &str) -> Result {{ // SAFETY: This is only used in single-threaded test code where @@ -315,7 +314,6 @@ mod tests { }}; } - /// Helper to safely remove an env var in tests. macro_rules! env_remove { ($($arg:tt)*) => {{ // SAFETY: This is only used in single-threaded test code where diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index 2ed5f07..52bfd6a 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -1181,10 +1181,6 @@ mod tests { assert_eq!(reader.take_line().unwrap(), "data: hi"); } - // ================================================ - // Builder timeout tests - // ================================================ - #[test] fn builder_has_default_timeouts() { // The builder should initialize with sensible non-zero defaults @@ -1236,10 +1232,6 @@ mod tests { assert!(client.is_ok(), "build should succeed with valid timeouts"); } - // ================================================== - // SSE buffer cap tests (M1) - // ================================================== - #[tokio::test] async fn sse_reader_take_line_splits_on_newline() { let mut reader = SseReader { @@ -1300,10 +1292,6 @@ mod tests { ); } - // ================================================== - // Body size limit tests (H5) - // ================================================== - #[test] fn max_response_body_is_ten_mb() { assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index a2422ab..44832a9 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -1006,10 +1006,6 @@ mod tests { assert_eq!(reader.take_line().unwrap(), "data: hi"); } - // ================================================ - // Builder timeout tests - // ================================================ - #[test] fn builder_has_default_timeouts() { // The builder should initialize with sensible non-zero defaults @@ -1061,10 +1057,6 @@ mod tests { assert!(client.is_ok(), "build should succeed with valid timeouts"); } - // ================================================== - // SSE buffer cap tests (M1) - // ================================================== - #[tokio::test] async fn sse_reader_take_line_splits_on_newline() { let mut reader = SseReader { @@ -1127,10 +1119,6 @@ mod tests { ); } - // ================================================== - // Body size limit tests (H5) - // ================================================== - #[test] fn max_response_body_is_ten_mb() { assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); diff --git a/src/provider/openai.rs b/src/provider/openai.rs index ca33fe6..adbf979 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -1265,10 +1265,6 @@ mod tests { assert_eq!(line, "data: hi"); } - // ================================================ - // Builder timeout tests - // ================================================ - #[test] fn builder_has_default_timeouts() { // The builder should initialize with sensible non-zero defaults @@ -1320,10 +1316,6 @@ mod tests { assert!(client.is_ok(), "build should succeed with valid timeouts"); } - // ================================================== - // SSE buffer cap tests (M1) - // ================================================== - #[tokio::test] async fn sse_reader_take_line_splits_on_newline() { let mut reader = SseReader { @@ -1379,10 +1371,6 @@ mod tests { ); } - // ================================================== - // Body size limit tests (H5) - // ================================================== - #[test] fn max_response_body_is_ten_mb() { assert_eq!(MAX_RESPONSE_BODY, 10 * 1024 * 1024); diff --git a/src/reflection.rs b/src/reflection.rs index 03e1aed..fbddf4f 100644 --- a/src/reflection.rs +++ b/src/reflection.rs @@ -578,10 +578,6 @@ impl fmt::Debug for NoopReflector { mod tests { use super::*; - // =================================================== - // FailureSeverity tests - // =================================================== - #[test] fn severity_ordering() { assert!(FailureSeverity::Low < FailureSeverity::Medium); @@ -606,10 +602,6 @@ mod tests { assert_eq!(deserialized, severity); } - // =================================================== - // ReflectionContext tests - // =================================================== - #[test] fn reflection_context_fields() { let ctx = ReflectionContext { @@ -622,10 +614,6 @@ mod tests { assert_eq!(ctx.max_attempts, 5); } - // =================================================== - // FailureAnalysis tests - // =================================================== - #[test] fn failure_analysis_recoverable() { let analysis = FailureAnalysis { @@ -660,10 +648,6 @@ mod tests { assert_eq!(c.description, "fix path"); } - // =================================================== - // ReflectionError tests - // =================================================== - #[test] fn reflection_error_skipped_display() { let err = ReflectionError::Skipped("not applicable".to_string()); @@ -680,10 +664,6 @@ mod tests { assert!(s.contains("llm timeout")); } - // =================================================== - // RecoveryAction tests - // =================================================== - #[test] fn action_retry_accessors() { let action = RecoveryAction::Retry { @@ -734,10 +714,6 @@ mod tests { assert!(fail.to_string().contains("fail: bad")); } - // =================================================== - // NoopReflector tests - // =================================================== - #[tokio::test] async fn noop_reflector_marks_non_recoverable() { let reflector = NoopReflector; diff --git a/src/stream/handler.rs b/src/stream/handler.rs index bb4bc99..476c5eb 100644 --- a/src/stream/handler.rs +++ b/src/stream/handler.rs @@ -1248,10 +1248,6 @@ mod tests { assert!(debug.contains("timeout_config")); } - // =================================================== - // StreamTimeoutConfig::validate tests - // =================================================== - #[test] fn timeout_config_validate_default_ok() { assert!(StreamTimeoutConfig::default().validate().is_ok()); @@ -1309,10 +1305,6 @@ mod tests { assert!(err.contains("progress_interval")); } - // =================================================== - // StreamRetryConfig::validate tests - // =================================================== - #[test] fn retry_config_validate_default_ok() { assert!(StreamRetryConfig::default().validate().is_ok()); @@ -1406,10 +1398,6 @@ mod tests { assert!(config.validate().is_ok()); } - // =================================================== - // StreamTurnResult tests - // =================================================== - #[test] fn stream_turn_result_fields() { let result = StreamTurnResult { @@ -1437,20 +1425,12 @@ mod tests { assert!(result.usage.is_none()); } - // =================================================== - // process_events async tests - // =================================================== - use crate::api::error::ApiError; use crate::stream::{ DeltaPart, IndexedDelta, MessageDelta, MessageDeltaPayload, MessageMetadata, MessageStart, PartStart, StreamEvent, }; - /// Helper: build a minimal happy-path event stream. - /// - /// Produces: MessageStart → PartStart(text) → IndexedDelta("hi") → - /// PartStop → MessageDelta(end_turn) → MessageStop fn happy_stream_events() -> Vec> { vec![ Ok(StreamEvent::MessageStart(MessageStart { @@ -1481,7 +1461,6 @@ mod tests { ] } - /// Build a `futures::stream` from a vec of events. fn event_stream( events: Vec>, ) -> std::pin::Pin< @@ -1618,18 +1597,8 @@ mod tests { assert!(!result.from_fallback); } - // =================================================== - // fallback_non_streaming async tests - // =================================================== - - /// Minimal mock that implements [`ApiClient`] for handler tests. - /// - /// Unlike the full [`MockApiClient`](crate::testing::MockApiClient), - /// this is defined locally so it works without the `testing` feature. struct HandlerMock { - /// If set, `create_message` returns this error. create_error: Option, - /// If set, `create_message` returns this JSON. create_response: Option, } @@ -1641,7 +1610,6 @@ mod tests { } } - /// Make `create_message` succeed with the given text. fn with_text_response(mut self, text: &str) -> Self { self.create_response = Some(serde_json::json!({ "content": [{"type": "text", "text": text}], @@ -1650,7 +1618,6 @@ mod tests { self } - /// Make `create_message` fail with the given error message. fn with_create_error(mut self, msg: &str) -> Self { self.create_error = Some(msg.to_string()); self @@ -1795,10 +1762,6 @@ mod tests { } } - // =================================================== - // stream_turn async tests - // =================================================== - #[tokio::test] async fn stream_turn_happy_path() { let handler = StreamHandler::new(); @@ -1844,7 +1807,6 @@ mod tests { // (covered above). Here we test that stream_turn returns the // error when streaming fails and the handler is configured // without fallback. - /// Mock that always returns an error stream. struct ErrorMock; impl ApiClient for ErrorMock { fn model(&self) -> String { diff --git a/src/tool/health.rs b/src/tool/health.rs index 78fe987..9420759 100644 --- a/src/tool/health.rs +++ b/src/tool/health.rs @@ -923,10 +923,6 @@ impl HealthRouterBuilder { mod tests { use super::*; - // ================================================== - // ToolStats tests - // ================================================== - #[test] fn tool_stats_starts_healthy() { let stats = ToolStats::new(); @@ -1033,10 +1029,6 @@ mod tests { assert_eq!(ewma, 643_000); } - // ================================================== - // CircuitState tests - // ================================================== - #[test] fn circuit_state_from_u32() { assert_eq!(CircuitState::from(0u32), CircuitState::Closed); @@ -1052,10 +1044,6 @@ mod tests { assert_eq!(format!("{}", CircuitState::HalfOpen), "half-open"); } - // ================================================== - // ToolCircuitBreaker tests - // ================================================== - #[test] fn circuit_breaker_starts_closed() { let cb = ToolCircuitBreaker::new(Duration::from_secs(30), 3); @@ -1158,10 +1146,6 @@ mod tests { assert!(cb.is_open(), "5 failures should open breaker"); } - // ================================================== - // HealthStatus tests - // ================================================== - #[test] fn health_status_display() { assert_eq!(format!("{}", HealthStatus::Healthy), "healthy"); @@ -1169,10 +1153,6 @@ mod tests { assert_eq!(format!("{}", HealthStatus::Unhealthy), "unhealthy"); } - // ================================================== - // ToolHealthRegistry tests - // ================================================== - #[test] fn registry_starts_empty() { let registry = ToolHealthRegistry::new(); @@ -1261,10 +1241,6 @@ mod tests { assert!(*bash_score > 0.5); } - // ================================================== - // HealthRouter tests - // ================================================== - #[test] fn health_router_no_fallbacks() { let router = HealthRouter::new(); @@ -1333,10 +1309,6 @@ mod tests { assert_eq!(router.resolve_tool("bash", ®istry), "bash"); } - // ================================================== - // Concurrent stress test - // ================================================== - #[test] fn registry_concurrent_access() { use std::sync::Arc; diff --git a/src/tool/registry.rs b/src/tool/registry.rs index ec55ca5..a02968c 100644 --- a/src/tool/registry.rs +++ b/src/tool/registry.rs @@ -517,7 +517,6 @@ impl Tool for FnTool { mod tests { use super::*; - /// A simple tool function for testing duplicate registration (L1 fix). fn test_tool_fn( _input: Value, _ctx: &ToolContext, @@ -525,7 +524,6 @@ mod tests { Box::pin(async { Ok(ToolOutput::text("ok")) }) } - /// A simple tool for testing duplicate registration (L1 fix). fn make_tool(name: &str) -> FnTool { FnTool::new( name.into(), diff --git a/src/tool/shield.rs b/src/tool/shield.rs index 1116121..fe1a410 100644 --- a/src/tool/shield.rs +++ b/src/tool/shield.rs @@ -806,10 +806,6 @@ mod tests { } } - // =================================================== - // NullShield tests - // =================================================== - #[test] fn null_shield_allows_everything() { let shield = NullShield; @@ -824,10 +820,6 @@ mod tests { let _ = &shield; } - // =================================================== - // UnixShield tests - // =================================================== - #[test] fn unix_shield_allows_safe_command() { let shield = UnixShield::new(); @@ -919,10 +911,6 @@ mod tests { assert!(combo2 > 0.0, "correct order should match"); } - // =================================================== - // Builder tests - // =================================================== - #[test] fn builder_blank_has_no_patterns() { let shield = UnixShieldBuilder::blank().build(); @@ -970,10 +958,6 @@ mod tests { let _builder = UnixShieldBuilder::default(); } - // =================================================== - // Type tests - // =================================================== - #[test] fn risk_level_display() { assert_eq!(RiskLevel::Safe.to_string(), "safe"); From 664e4cb6065077471f1ffb56c50725647d7ecf21 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 14:30:21 +1200 Subject: [PATCH 22/30] fix: SSE multi-line data, tool panic isolation, error body cap --- src/engine/bare/dispatch.rs | 198 +++++++++++++++++++++++++++++++++++- src/middleware/tool_call.rs | 140 ++++++++++++++++++++++++- src/provider/anthropic.rs | 42 +++++++- src/provider/gemini.rs | 10 +- src/provider/openai.rs | 10 +- 5 files changed, 392 insertions(+), 8 deletions(-) diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index b8e1b69..fce5965 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -22,6 +22,9 @@ use crate::detection::loop_detector::{self, Operation}; use crate::observer::{ToolPostContext, ToolPreContext}; use crate::reflection::{Correction, CorrectionResult}; +use futures::FutureExt; +use std::panic::AssertUnwindSafe; + /// Result of deciding what to do after a tool error during recovery. /// /// Distinguishes between returning a soft-error result (the tool failed, @@ -261,14 +264,17 @@ impl BareLoop { let tool_result = if let Some(tool) = self.tools.get(&tc.tool) { let cancel = Arc::clone(&self.cancelled); + // Wrap the tool call in `catch_unwind` so a panicking tool + // implementation produces an error result instead of unwinding + // through and aborting the entire agent loop. let call_result = tokio::select! { - r = tool.call(tc.input.clone(), tool_context) => r, + r = AssertUnwindSafe(tool.call(tc.input.clone(), tool_context)).catch_unwind() => r, () = cancel.notified() => { return Err(LoopError::Cancelled); } }; match call_result { - Ok(result) => { + Ok(Ok(result)) => { let duration = start.elapsed(); ToolDispatchResult { tool_call_id: tc.id.clone(), @@ -278,7 +284,7 @@ impl BareLoop { resolved_tool_name: tc.tool.clone(), } } - Err(e) => { + Ok(Err(e)) => { let duration = start.elapsed(); let error_msg = e.to_string(); ToolDispatchResult { @@ -289,6 +295,28 @@ impl BareLoop { resolved_tool_name: tc.tool.clone(), } } + Err(panic_payload) => { + let duration = start.elapsed(); + let msg = panic_payload + .downcast_ref::<&'static str>() + .map(std::string::ToString::to_string) + .or_else(|| panic_payload.downcast_ref::().cloned()) + .unwrap_or_else(|| { + format!("Tool '{}' panicked (unknown payload)", tc.tool) + }); + tracing::error!( + tool = %tc.tool, + panic_message = %msg, + "tool panicked during execution" + ); + ToolDispatchResult { + tool_call_id: tc.id.clone(), + output: ToolContent::Text(format!("Tool '{}' panicked: {msg}", tc.tool)), + is_error: true, + duration, + resolved_tool_name: tc.tool.clone(), + } + } } } else { self.tool_not_found(tc) @@ -555,3 +583,167 @@ impl BareLoop { (action, correction) } } + +#[cfg(test)] +#[allow(clippy::unnecessary_literal_bound)] +mod tests { + use crate::api::error::ApiError; + use crate::config::LoopConfig; + use crate::engine::loop_core::ToolCall; + use crate::message::ToolContent; + use crate::tool::{ + Tool, ToolContext, ToolError, ToolOutput, ToolSchema, registry::ToolRegistry, + }; + use serde_json::Value; + use std::future::Future; + use std::pin::Pin; + use std::sync::Arc; + use std::time::Instant; + + use parking_lot::Mutex; + + use super::*; + + struct MockClient { + model_name: Arc>, + } + + impl MockClient { + fn new(model: &str) -> Self { + Self { + model_name: Arc::new(Mutex::new(model.to_string())), + } + } + } + + impl ApiClient for MockClient { + fn model(&self) -> String { + 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(); + true + } + fn stream_messages( + &self, + _history: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn futures::Stream> + + Send + + 'static, + >, + > { + Box::pin(futures::stream::empty()) + } + fn create_message( + &self, + _history: Vec, + _system: Option, + _tools: Option>, + ) -> Pin> + Send + '_>> + { + Box::pin(async { Err(ApiError::http("not implemented")) }) + } + } + + struct PanicTool; + + impl Tool for PanicTool { + fn name(&self) -> &str { + "panic_tool" + } + fn description(&self) -> &str { + "Panics on call" + } + fn schema(&self) -> ToolSchema { + ToolSchema { + tool: "panic_tool".into(), + description: "Panics on call".into(), + input_schema: Value::Object(serde_json::Map::new()), + } + } + fn call( + &self, + _input: Value, + _ctx: &ToolContext, + ) -> Pin> + Send + '_>> { + Box::pin(async { panic!("dispatch.rs panic tool") }) + } + } + + fn echo_fn( + _input: Value, + _ctx: &ToolContext, + ) -> Pin> + Send + 'static>> { + Box::pin(async { Ok(ToolOutput::text("ok")) }) + } + + fn make_loop(tools: ToolRegistry) -> BareLoop { + let config = LoopConfig::default(); + let client = Arc::new(MockClient::new("test")); + BareLoop::new(client, tools, config) + } + + #[tokio::test] + async fn dispatch_tool_catches_panic() { + let mut registry = ToolRegistry::new(); + registry.register(PanicTool); + let bare = make_loop(registry); + + let tc = ToolCall { + id: "tc1".into(), + tool: "panic_tool".into(), + input: Value::Null, + }; + let tool_context = ToolContext::default(); + let start = Instant::now(); + + let result = bare.dispatch_tool(&tc, &tool_context, start, 0).await; + + assert!(result.is_ok(), "panic should be caught, not propagated"); + let dispatch_result = result.unwrap(); + assert!(dispatch_result.is_error); + match &dispatch_result.output { + ToolContent::Text(text) => { + assert!(text.contains("panicked"), "expected panic message: {text}"); + } + ToolContent::Multipart(_) => panic!("expected Text"), + } + } + + #[tokio::test] + async fn dispatch_tool_normal_tool_works() { + let mut registry = ToolRegistry::new(); + registry.register(crate::tool::FnTool::new( + "echo".into(), + "echo".into(), + Value::Object(serde_json::Map::new()), + echo_fn, + )); + let bare = make_loop(registry); + + let tc = ToolCall { + id: "tc1".into(), + tool: "echo".into(), + input: Value::Null, + }; + let tool_context = ToolContext::default(); + let start = Instant::now(); + + let result = bare.dispatch_tool(&tc, &tool_context, start, 0).await; + + assert!(result.is_ok()); + let dispatch_result = result.unwrap(); + assert!(!dispatch_result.is_error); + match &dispatch_result.output { + ToolContent::Text(text) => assert_eq!(text, "ok"), + ToolContent::Multipart(_) => panic!("expected Text"), + } + } +} diff --git a/src/middleware/tool_call.rs b/src/middleware/tool_call.rs index c0fbaa6..3c22eb1 100644 --- a/src/middleware/tool_call.rs +++ b/src/middleware/tool_call.rs @@ -4,10 +4,13 @@ use super::{ToolDispatchContext, ToolDispatchResult}; use crate::error::LoopError; use crate::tool::ToolRegistry; use std::future::Future; +use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::sync::Arc; use std::time::Instant; +use futures::FutureExt; + /// The innermost middleware that performs the actual tool invocation. /// /// Looks up the tool by name in the [`ToolRegistry`], calls @@ -56,8 +59,13 @@ impl ToolCallMiddleware { .with_call_id(&call_id); }; + // Wrap the tool call in `catch_unwind` so a panicking tool + // implementation produces an error result instead of unwinding + // through and aborting the entire agent loop. Tools are + // user-supplied `dyn Tool` implementations and the framework + // cannot trust them to be panic-free. let call_result = tokio::select! { - r = tool.call(input, &tool_ctx) => r, + r = AssertUnwindSafe(tool.call(input, &tool_ctx)).catch_unwind() => r, () = cancel.notified() => { return ToolDispatchResult::err( &tool_name, @@ -69,8 +77,138 @@ impl ToolCallMiddleware { }; let duration = start.elapsed(); + + // Convert a panic payload into a tool-error result. + let call_result = match call_result { + Ok(inner) => inner, + Err(panic_payload) => { + let msg = panic_payload + .downcast_ref::<&'static str>() + .map(std::string::ToString::to_string) + .or_else(|| panic_payload.downcast_ref::().cloned()) + .unwrap_or_else(|| { + format!("Tool '{tool_name}' panicked (unknown payload)") + }); + tracing::error!(tool = %tool_name, panic_message = %msg, "tool panicked during execution"); + return ToolDispatchResult::err( + &tool_name, + format!("Tool '{tool_name}' panicked: {msg}"), + duration, + ) + .with_call_id(&call_id); + } + }; + ToolDispatchResult::from_result(&tool_name, call_result, duration) .with_call_id(&call_id) }) } } + +#[cfg(test)] +#[allow(clippy::unnecessary_literal_bound)] +mod tests { + use super::*; + use crate::cancel::CancelSignal; + use crate::message::ToolContent; + use crate::middleware::ToolDispatchContext; + use crate::tool::PermissionCheck; + use crate::tool::{ + FnTool, Tool, ToolContext, ToolError, ToolOutput, ToolSchema, registry::ToolRegistry, + }; + use std::future::Future; + use std::pin::Pin; + use std::sync::Arc; + + struct PanickingTool; + + fn echo_fn( + _input: serde_json::Value, + _ctx: &ToolContext, + ) -> Pin> + Send + 'static>> { + Box::pin(async { Ok(ToolOutput::text("ok")) }) + } + + impl Tool for PanickingTool { + fn name(&self) -> &str { + "panic_tool" + } + fn description(&self) -> &str { + "A tool that panics" + } + fn schema(&self) -> ToolSchema { + ToolSchema { + tool: "panic_tool".into(), + description: "A tool that panics".into(), + input_schema: serde_json::json!({"type": "object"}), + } + } + + fn call( + &self, + _input: serde_json::Value, + _ctx: &ToolContext, + ) -> Pin> + Send + '_>> { + Box::pin(async { + panic!("boom from panicking tool"); + }) + } + } + + fn make_ctx(tool_name: &str, cancel: Arc) -> ToolDispatchContext { + ToolDispatchContext { + tool_name: tool_name.into(), + input: serde_json::Value::Null, + call_id: "call_1".into(), + turn_number: 0, + cancel, + permission: PermissionCheck::allow(), + tool_context: ToolContext::default(), + } + } + + #[tokio::test] + async fn panicking_tool_produces_error_result_not_abort() { + let mut registry = ToolRegistry::new(); + registry.register(PanickingTool); + + let middleware = ToolCallMiddleware::new(Arc::new(registry)); + let cancel = Arc::new(CancelSignal::new()); + let mut ctx = make_ctx("panic_tool", cancel); + + let result = middleware.dispatch(&mut ctx).await; + + assert!(result.is_error, "panic should become an error result"); + match &result.output { + ToolContent::Text(text) => { + assert!( + text.contains("panicked"), + "error message should mention panic: {text}" + ); + } + ToolContent::Multipart(_) => panic!("expected Text output"), + } + } + + #[tokio::test] + async fn normal_tool_still_works() { + let mut registry = ToolRegistry::new(); + registry.register(FnTool::new( + "echo".into(), + "Echoes input".into(), + serde_json::json!({"type": "object"}), + echo_fn, + )); + + let middleware = ToolCallMiddleware::new(Arc::new(registry)); + let cancel = Arc::new(CancelSignal::new()); + let mut ctx = make_ctx("echo", cancel); + + let result = middleware.dispatch(&mut ctx).await; + assert!(!result.is_error); + match &result.output { + ToolContent::Text(text) => assert_eq!(text, "ok"), + ToolContent::Multipart(_) => panic!("expected Text output"), + } + } +} diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index 52bfd6a..8f78f9d 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -48,6 +48,9 @@ const DEFAULT_MAX_TOKENS: u32 = 8192; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb +/// Maximum bytes to read from an error response body. Prevents OOM when a +/// misconfigured or malicious server returns a multi-GB body on a 4xx/5xx. +const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -127,7 +130,12 @@ impl AnthropicClient { if status.is_success() { Ok(resp) } else { - let text = resp.text().await.unwrap_or_default(); + // Cap the error body to prevent OOM from oversized error responses. + let bytes = resp.bytes().await.unwrap_or_default(); + let text = match bytes.get(..MAX_ERROR_BODY) { + Some(truncated) => String::from_utf8_lossy(truncated).into_owned(), + None => String::from_utf8_lossy(&bytes).into_owned(), + }; Err(ApiError::http_with_status(status.as_u16(), text)) } } @@ -514,7 +522,16 @@ impl SseReader { event_type = ev.into(); have_event = true; } else if let Some(d) = line.strip_prefix(SSE_DATA_PREFIX) { - data = d.into(); + // Per the SSE specification, multiple consecutive `data:` + // lines must be concatenated with `\n` to form a single + // event payload. Using assignment here would silently + // discard earlier data lines. + if data.is_empty() { + data = d.into(); + } else { + data.push('\n'); + data.push_str(d); + } have_event = true; } } @@ -1258,6 +1275,27 @@ mod tests { assert!(data.is_some()); } + #[tokio::test] + async fn sse_reader_next_event_concatenates_multiline_data() { + let chunk = "event: content_block_delta\ndata: {\"type\":\"text_delta\",\ndata: \"text\":\"hello\"}\n\n"; + let stream = futures::stream::iter(vec![Ok::(chunk.to_string())]); + let mut reader = SseReader { + bytes: Box::pin(stream), + buf: String::new(), + }; + let result = reader.next_event().await.unwrap(); + assert!(result.is_some()); + let (event_type, data) = result.unwrap(); + assert_eq!(event_type, "content_block_delta"); + assert!( + data.is_some(), + "multi-line data should concatenate into valid JSON" + ); + let parsed = data.unwrap(); + assert_eq!(parsed["type"], "text_delta"); + assert_eq!(parsed["text"], "hello"); + } + #[tokio::test] async fn sse_reader_next_event_malformed_data_returns_none_value() { // Malformed JSON data should be logged and returned as None for the diff --git a/src/provider/gemini.rs b/src/provider/gemini.rs index 44832a9..bf85155 100644 --- a/src/provider/gemini.rs +++ b/src/provider/gemini.rs @@ -47,6 +47,9 @@ const TEXT_PART_INDEX: usize = 0; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb +/// Maximum bytes to read from an error response body. Prevents OOM when a +/// misconfigured or malicious server returns a multi-GB body on a 4xx/5xx. +const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -140,7 +143,12 @@ impl GeminiClient { if status.is_success() { Ok(resp) } else { - let text = resp.text().await.unwrap_or_default(); + // Cap the error body to prevent OOM from oversized error responses. + let bytes = resp.bytes().await.unwrap_or_default(); + let text = match bytes.get(..MAX_ERROR_BODY) { + Some(truncated) => String::from_utf8_lossy(truncated).into_owned(), + None => String::from_utf8_lossy(&bytes).into_owned(), + }; Err(ApiError::http_with_status(status.as_u16(), text)) } } diff --git a/src/provider/openai.rs b/src/provider/openai.rs index adbf979..42fd270 100644 --- a/src/provider/openai.rs +++ b/src/provider/openai.rs @@ -50,6 +50,9 @@ const TEXT_PART_INDEX: usize = 0; const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120); // connect + response + body const MAX_RESPONSE_BODY: usize = 10 * 1024 * 1024; // 10 Mb const SSE_MAX_BUFFER: usize = 1024 * 1024; // 1 Mb +/// Maximum bytes to read from an error response body. Prevents OOM when a +/// misconfigured or malicious server returns a multi-GB body on a 4xx/5xx. +const MAX_ERROR_BODY: usize = 8 * 1024; // 8 Kb const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); // ================================================== @@ -141,7 +144,12 @@ impl OpenAiClient { if status.is_success() { Ok(resp) } else { - let text = resp.text().await.unwrap_or_default(); + // Cap the error body to prevent OOM from oversized error responses. + let bytes = resp.bytes().await.unwrap_or_default(); + let text = match bytes.get(..MAX_ERROR_BODY) { + Some(truncated) => String::from_utf8_lossy(truncated).into_owned(), + None => String::from_utf8_lossy(&bytes).into_owned(), + }; Err(ApiError::http_with_status(status.as_u16(), text)) } } From 66e26d2d735020773a74bf0050ac8e6ebac8580c Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 15:20:23 +1200 Subject: [PATCH 23/30] fix: SSE parsing, tool panic isolation, error body cap, config validation, calculator robustness, and docs --- README.md | 8 ++++++++ examples/chat.rs | 15 +++++++++++++-- src/api/error.rs | 7 ++++--- src/config.rs | 13 ++++++++++++- 4 files changed, 37 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 3f229f7..cafa62b 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,14 @@ let agent = BareLoop::new( | `testing` | No | — | Mock clients, tools, and test fixtures | | `tool_health` | No | — | Per-tool health monitoring, circuit breakers, and self-healing routing | | `tool_shield` | No | `tool_health` | Tool permission shielding and access control | +| `providers` | No | — | Base provider support (`reqwest` + `async-stream`); enables `provider` module | +| `openai` | No | `providers` | OpenAI-compatible API client (`provider::openai`) | +| `anthropic` | No | `providers` | Anthropic Claude API client (`provider::anthropic`) | +| `ollama` | No | `providers`, `openai` | Ollama local model client (OpenAI-compatible) | +| `deepseek` | No | `providers`, `openai` | DeepSeek API client (OpenAI-compatible) | +| `grok` | No | `providers`, `openai` | Grok (xAI) API client (OpenAI-compatible) | +| `gemini` | No | `providers` | Google Gemini API client (`provider::gemini`) | +| `zai` | No | `providers`, `anthropic` | Z.AI API client (Anthropic-compatible) | ## Architecture diff --git a/examples/chat.rs b/examples/chat.rs index 1c775e2..890a6d6 100644 --- a/examples/chat.rs +++ b/examples/chat.rs @@ -253,7 +253,16 @@ fn simple_eval(expr: &str) -> String { let tokens = tokenize(expr); let mut pos = 0usize; match parse_expr(&tokens, &mut pos) { - Ok(val) => format!("{val}"), + Ok(val) => { + if pos < tokens.len() { + format!( + "Error: unexpected token after expression: {:?}", + tokens[pos] + ) + } else { + format!("{val}") + } + } Err(e) => format!("Error: {e}"), } } @@ -319,8 +328,10 @@ fn parse_factor(tokens: &[Token], pos: &mut usize) -> Result { let val = parse_expr(tokens, pos)?; if matches!(peek(tokens, *pos), Some(Token::RParen)) { advance(pos); + Ok(val) + } else { + Err("expected closing parenthesis".into()) } - Ok(val) } Token::Minus => { advance(pos); diff --git a/src/api/error.rs b/src/api/error.rs index 79cbc7c..cd9fd58 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -1284,11 +1284,12 @@ mod tests { #[test] fn test_result_type() { - fn returns_result() -> String { - "success".to_string() + #[allow(clippy::unnecessary_wraps)] + fn returns_result() -> super::Result { + Ok("success".to_string()) } let result = returns_result(); - assert_eq!(result, "success"); + assert_eq!(result.unwrap(), "success"); } #[test] diff --git a/src/config.rs b/src/config.rs index 31881be..408590a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -129,7 +129,7 @@ impl LoopConfig { "max_tokens must be greater than 0".to_string(), )); } - if self.model.is_empty() { + if self.model.trim().is_empty() { return Err(crate::error::LoopError::Config( "model must not be empty".to_string(), )); @@ -199,4 +199,15 @@ mod tests { }; assert!(config.validate().is_ok()); } + + #[test] + fn validate_rejects_whitespace_only_model() { + let config = LoopConfig { + model: " ".to_string(), + ..LoopConfig::default() + }; + let err = config.validate().unwrap_err(); + let msg = err.to_string(); + assert!(msg.contains("model"), "error should mention model: {msg}"); + } } From 8d7c3f085a36229abcaa381683a06cfde36a39ef Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 17:12:36 +1200 Subject: [PATCH 24/30] fix: only extend streak if previous is the same --- src/api/error.rs | 10 +-- src/detection/convergence.rs | 132 ++++++++++++++++++++++++++++++++--- src/lib.rs | 2 + 3 files changed, 132 insertions(+), 12 deletions(-) diff --git a/src/api/error.rs b/src/api/error.rs index cd9fd58..90790a6 100644 --- a/src/api/error.rs +++ b/src/api/error.rs @@ -1284,12 +1284,14 @@ mod tests { #[test] fn test_result_type() { - #[allow(clippy::unnecessary_wraps)] - fn returns_result() -> super::Result { + fn ok_path() -> super::Result { Ok("success".to_string()) } - let result = returns_result(); - assert_eq!(result.unwrap(), "success"); + fn err_path() -> super::Result { + Err(ApiError::other("failure")) + } + assert_eq!(ok_path().unwrap(), "success"); + assert!(err_path().is_err()); } #[test] diff --git a/src/detection/convergence.rs b/src/detection/convergence.rs index 1908d47..71ac4cf 100644 --- a/src/detection/convergence.rs +++ b/src/detection/convergence.rs @@ -565,26 +565,32 @@ impl ConvergenceDetector { } let mut max_similarity = 0.0; - let mut any_similar = false; for prev_response in &self.window { let similarity = Self::compute_similarity(response, prev_response); if similarity > max_similarity { max_similarity = similarity; } + } - if similarity >= self.config.similarity_threshold { - any_similar = true; - if !self.similar_responses.contains(&response.to_string()) { - self.similar_responses.push(response.to_string()); + let prev_is_similar = match self.window.back() { + Some(prev) => { + let sim = Self::compute_similarity(response, prev); + if sim >= self.config.similarity_threshold { + if !self.similar_responses.contains(&response.to_string()) { + self.similar_responses.push(response.to_string()); + } + true + } else { + false } } - } + None => false, + }; - // Update consecutive count once per add_response call if self.window.is_empty() { self.consecutive_count = 1; self.similar_responses.push(response.to_string()); - } else if any_similar { + } else if prev_is_similar { self.consecutive_count = self.consecutive_count.saturating_add(1); } else { self.consecutive_count = 1; @@ -927,6 +933,116 @@ mod tests { )); } + #[test] + fn test_streak_resets_on_dissimilar_response() { + let config = ConvergenceConfig { + window_size: 3, + similarity_threshold: 0.5, + ..Default::default() + }; + let mut detector = ConvergenceDetector::new(config).unwrap(); + + // Two similar responses build a streak of 2. + let s1 = detector.add_response("same text"); + assert_eq!(s1.consecutive_count, 1); + let s2 = detector.add_response("same text"); + assert_eq!(s2.consecutive_count, 2); + + // A dissimilar response breaks the streak and resets to 1. + let s3 = detector.add_response("completely different content here"); + assert_eq!( + s3.consecutive_count, 1, + "dissimilar response should reset streak" + ); + + // Another similar-to-immediately-previous response starts fresh. + let s4 = detector.add_response("completely different content here"); + assert_eq!(s4.consecutive_count, 2); + } + + #[test] + fn test_alternating_responses_never_converge() { + // The streak must only compare against the *immediately previous* + // response, not any similar response in the window. Alternating + // A / B / A / B should never build a streak longer than 1. + let config = ConvergenceConfig { + window_size: 3, + similarity_threshold: 0.5, + ..Default::default() + }; + let mut detector = ConvergenceDetector::new(config).unwrap(); + + detector.add_response("alpha alpha alpha"); + detector.add_response("beta beta beta"); + detector.add_response("alpha alpha alpha"); + detector.add_response("beta beta beta"); + detector.add_response("alpha alpha alpha"); + + // Even though "alpha" appeared 3 times, they were never consecutive, + // so the streak should never reach window_size. + let status = detector.check_convergence(); + assert!( + !status.detected, + "alternating responses must not trigger convergence" + ); + assert!( + status.consecutive_count <= 1, + "alternating responses should not build a streak" + ); + } + + #[test] + fn test_similar_after_gap_starts_fresh_streak() { + // A, A (streak=2), B (streak resets to 1), A (streak=1, NOT 3) + let config = ConvergenceConfig { + window_size: 3, + similarity_threshold: 0.5, + ..Default::default() + }; + let mut detector = ConvergenceDetector::new(config).unwrap(); + + detector.add_response("same text"); + let s2 = detector.add_response("same text"); + assert_eq!(s2.consecutive_count, 2); + + detector.add_response("totally different stuff"); + + // "same text" again — similar to the *previous* response? No. + // Previous was "totally different stuff", so streak should be 1. + let s4 = detector.add_response("same text"); + assert_eq!( + s4.consecutive_count, 1, + "similar response after a gap should start a fresh streak" + ); + } + + #[test] + fn test_converge_then_break_then_re_converge() { + let config = ConvergenceConfig { + window_size: 3, + similarity_threshold: 0.5, + ..Default::default() + }; + let mut detector = ConvergenceDetector::new(config).unwrap(); + + // Build to convergence. + detector.add_response("loop loop loop"); + detector.add_response("loop loop loop"); + let s3 = detector.add_response("loop loop loop"); + assert!(s3.detected); + + // Break the streak. + let s4 = detector.add_response("something entirely new and different"); + assert!(!s4.detected); + + // Build back up — need 3 consecutive again, not just 1 more. + let s5 = detector.add_response("something entirely new and different"); + assert!(!s5.detected, "only 2 consecutive so far"); + + let s6 = detector.add_response("something entirely new and different"); + assert!(s6.detected, "3 consecutive of the new response"); + } + #[test] fn test_config_threshold_boundary_valid() { // 0.0 and 1.0 are valid boundary values diff --git a/src/lib.rs b/src/lib.rs index 15d7707..cac84a3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -54,6 +54,8 @@ clippy::panic, clippy::indexing_slicing, clippy::missing_panics_doc, + clippy::missing_errors_doc, + clippy::unnecessary_wraps, clippy::clone_on_ref_ptr, clippy::doc_markdown, clippy::field_reassign_with_default, From 82e6b9e3cabbf8d813bef5c614d3885750da0008 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 17:23:34 +1200 Subject: [PATCH 25/30] fix: session end reason --- src/engine/bare.rs | 122 ++++++++++++++++++++++++++++++++++-- src/engine/bare/emission.rs | 35 +++++++++-- 2 files changed, 148 insertions(+), 9 deletions(-) diff --git a/src/engine/bare.rs b/src/engine/bare.rs index a19eb5d..760c5ae 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -65,14 +65,16 @@ use crate::config::LoopConfig; use crate::error::LoopError; use crate::engine::loop_core::{LoopState, SessionResult, StopReason, ToolCall, TurnResult}; -#[cfg(feature = "hooks")] -use crate::hooks::HookAction; -#[cfg(feature = "hooks")] -use crate::hooks::HookExecutor; +#[cfg(all(test, feature = "hooks"))] +use crate::hooks::Hook; #[cfg(feature = "hooks")] use crate::hooks::context::{ CompactTrigger, PostCompactContext, PostToolUseContext, PreCompactContext, PreToolUseContext, }; +#[cfg(all(test, feature = "hooks"))] +use crate::hooks::context::{SessionEndContext as HookSessionEndContext, SessionEndReason}; +#[cfg(feature = "hooks")] +use crate::hooks::{HookAction, HookExecutor}; use crate::message::{Message, MessagePart, Role, ToolContent}; use crate::middleware::{ToolDispatchContext, ToolPipeline, ToolPipelineBuilder}; use crate::observer::{ @@ -2541,4 +2543,116 @@ mod tests { Some("new-primary".to_string()) ); } + + #[cfg(feature = "hooks")] + struct ReasonCaptureHook { + reason: Mutex>, + } + + #[cfg(feature = "hooks")] + impl ReasonCaptureHook { + fn new() -> Arc { + Arc::new(Self { + reason: Mutex::new(None), + }) + } + + fn captured(&self) -> Option { + *self.reason.lock() + } + } + + #[cfg(feature = "hooks")] + impl Hook for ReasonCaptureHook { + fn name(&self) -> &'static str { + "ReasonCaptureHook" + } + + fn on_session_end(&self, ctx: &HookSessionEndContext) { + *self.reason.lock() = Some(ctx.reason); + } + } + + #[cfg(feature = "hooks")] + fn loop_with_reason_hook() -> (BareLoop, Arc) { + let hook = ReasonCaptureHook::new(); + let executor = Arc::new(HookExecutor::new().with_hook(hook.clone())); + let config = LoopConfig { + max_turns: 5, + ..LoopConfig::default() + }; + let mut loop_ = BareLoop::new( + Arc::new(MockClient::new("test")), + ToolRegistry::new(), + config, + ); + loop_.set_hook_executor(executor); + (loop_, hook) + } + + #[cfg(feature = "hooks")] + #[tokio::test] + async fn session_end_reason_complete() { + let (mut loop_, hook) = loop_with_reason_hook(); + // Normal completion: state is Completed, not cancelled, under max_turns. + loop_.budget.success = true; + loop_.budget.total_turns = 2; + + loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + + assert_eq!(hook.captured(), Some(SessionEndReason::Complete)); + } + + #[cfg(feature = "hooks")] + #[tokio::test] + async fn session_end_reason_cancelled() { + let (mut loop_, hook) = loop_with_reason_hook(); + // Cancel signal fired — success is true (not Failed) but cancelled. + loop_.budget.success = true; + loop_.budget.total_turns = 2; + loop_.cancelled.cancel(); + + loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + + assert_eq!(hook.captured(), Some(SessionEndReason::Cancelled)); + } + + #[cfg(feature = "hooks")] + #[tokio::test] + async fn session_end_reason_max_turns() { + let (mut loop_, hook) = loop_with_reason_hook(); + // Hit max_turns: total_turns == max_turns, not cancelled, success true. + loop_.budget.success = true; + loop_.budget.total_turns = 5; // equals config.max_turns + + loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + + assert_eq!(hook.captured(), Some(SessionEndReason::MaxTurns)); + } + + #[cfg(feature = "hooks")] + #[tokio::test] + async fn session_end_reason_error() { + let (mut loop_, hook) = loop_with_reason_hook(); + // Generic failure: success false, no context-overflow keyword. + loop_.budget.success = false; + loop_.budget.error = Some("API connection refused".to_string()); + + loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + + assert_eq!(hook.captured(), Some(SessionEndReason::Error)); + } + + #[cfg(feature = "hooks")] + #[tokio::test] + async fn session_end_reason_context_overflow() { + let (mut loop_, hook) = loop_with_reason_hook(); + // Failure with context-overflow keyword in the error message. + loop_.budget.success = false; + loop_.budget.error = Some("context length exceeded".to_string()); + + loop_.notify_session_end(&loop_.budget.clone(), Duration::from_millis(100)); + + assert_eq!(hook.captured(), Some(SessionEndReason::ContextOverflow)); + } } diff --git a/src/engine/bare/emission.rs b/src/engine/bare/emission.rs index 11cd611..50ae8cc 100644 --- a/src/engine/bare/emission.rs +++ b/src/engine/bare/emission.rs @@ -57,11 +57,7 @@ impl BareLoop { #[cfg(feature = "hooks")] if let Some(executor) = self.managers.hook_executor() { - let reason = if result.success { - SessionEndReason::Complete - } else { - SessionEndReason::Error - }; + let reason = self.session_end_reason(result.success); let ctx = HookSessionEndContext { session_id: result.session_id, reason, @@ -73,6 +69,35 @@ impl BareLoop { } } + /// Derive the structured [`SessionEndReason`] from the loop's + /// terminal state. + /// + /// Unlike a simple `success` boolean, this distinguishes + /// cancellation, max-turns exhaustion, and context overflow. + #[cfg(feature = "hooks")] + fn session_end_reason(&self, success: bool) -> SessionEndReason { + if !success { + // Context overflow is a specific failure mode distinguishable + // from a generic error by its message. + if self + .budget + .error + .as_ref() + .is_some_and(|e| e.contains("context") || e.contains("overflow")) + { + SessionEndReason::ContextOverflow + } else { + SessionEndReason::Error + } + } else if self.is_cancelled() { + SessionEndReason::Cancelled + } else if self.budget.total_turns >= self.config.max_turns { + SessionEndReason::MaxTurns + } else { + SessionEndReason::Complete + } + } + /// Convert a [`Duration`] to milliseconds as `u64`. pub(super) fn millis_u64(duration: Duration) -> u64 { u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) From 0c31280cd739f9edce7cc5746df8e0381b9e9b63 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 17:27:27 +1200 Subject: [PATCH 26/30] chore: reject tool non ubject inputfix --- src/engine/loop_core.rs | 120 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 114 insertions(+), 6 deletions(-) diff --git a/src/engine/loop_core.rs b/src/engine/loop_core.rs index f4d4c11..7c48c48 100644 --- a/src/engine/loop_core.rs +++ b/src/engine/loop_core.rs @@ -364,12 +364,18 @@ impl ToolCall { match correction.correction_type { CorrectionType::InputFix => { if let Some(ref modified) = correction.modified_input { - tracing::debug!( - tool = %self.tool, - "applying InputFix correction from reflector" - ); - self.input = modified.clone(); - crate::reflection::CorrectionResult::Applied + if modified.is_object() { + tracing::debug!( + tool = %self.tool, + "applying InputFix correction from reflector" + ); + self.input = modified.clone(); + crate::reflection::CorrectionResult::Applied + } else { + crate::reflection::CorrectionResult::Failed( + "InputFix correction modified_input must be a JSON object".to_string(), + ) + } } else { crate::reflection::CorrectionResult::Failed( "InputFix correction missing modified_input".to_string(), @@ -874,3 +880,105 @@ pub trait Loop: Send + Sync { /// initialization, so the borrow is short-lived. fn config(&self) -> &LoopConfig; } + +#[cfg(test)] +mod tests { + use super::*; + use crate::reflection::{Correction, CorrectionResult, CorrectionType}; + use crate::tool::ToolDispatchResult; + use serde_json::json; + use std::time::Duration; + + fn make_call() -> ToolCall { + ToolCall { + id: "test".to_string(), + tool: "Read".to_string(), + input: json!({"path": "/tmp"}), + } + } + + fn empty_prior_result() -> ToolDispatchResult { + ToolDispatchResult::ok("Read", String::new(), Duration::ZERO) + } + + #[test] + fn input_fix_accepts_json_object() { + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::InputFix, + description: "fix path".into(), + modified_input: Some(json!({"path": "/tmp/fixed"})), + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction, &empty_prior_result()); + assert!(matches!(result, CorrectionResult::Applied)); + assert_eq!(call.input, json!({"path": "/tmp/fixed"})); + } + + #[test] + fn input_fix_rejects_scalar() { + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::InputFix, + description: "bad fix".into(), + modified_input: Some(json!("/tmp/scalar")), + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction, &empty_prior_result()); + match result { + CorrectionResult::Failed(msg) => { + assert!(msg.contains("JSON object"), "{msg}"); + } + other => panic!("expected Failed, got {other:?}"), + } + // Input should remain unchanged. + assert_eq!(call.input, json!({"path": "/tmp"})); + } + + #[test] + fn input_fix_rejects_array() { + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::InputFix, + description: "bad fix".into(), + modified_input: Some(json!(["a", "b"])), + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction, &empty_prior_result()); + assert!(matches!(result, CorrectionResult::Failed(_))); + assert_eq!(call.input, json!({"path": "/tmp"})); + } + + #[test] + fn input_fix_rejects_null() { + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::InputFix, + description: "bad fix".into(), + modified_input: Some(serde_json::Value::Null), + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction, &empty_prior_result()); + assert!(matches!(result, CorrectionResult::Failed(_))); + assert_eq!(call.input, json!({"path": "/tmp"})); + } + + #[test] + fn input_fix_rejects_missing() { + let mut call = make_call(); + let correction = Correction { + correction_type: CorrectionType::InputFix, + description: "no input".into(), + modified_input: None, + alternative_tool: None, + guidance: None, + }; + let result = call.apply_correction(&correction, &empty_prior_result()); + assert!(matches!(result, CorrectionResult::Failed(_))); + assert_eq!(call.input, json!({"path": "/tmp"})); + } +} From 82da15bb3255cff28a2ec74a66e9f4a7d7efbcd4 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 17:30:04 +1200 Subject: [PATCH 27/30] fix: fail the multipart test on unexpected part variants --- src/middleware.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/middleware.rs b/src/middleware.rs index d8fca3e..ca31f82 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -1142,11 +1142,17 @@ mod tests { match result.output { ToolContent::Multipart(parts) => { assert_eq!(parts.len(), 2); - if let ToolContentPart::Text { text } = &parts[0] { - assert_eq!(text, "part1", "short part should not be truncated"); + match &parts[0] { + ToolContentPart::Text { text } => { + assert_eq!(text, "part1", "short part should not be truncated"); + } + ToolContentPart::Image { .. } => panic!("expected Text part[0], got Image"), } - if let ToolContentPart::Text { text } = &parts[1] { - assert_eq!(text, "part2", "short part should not be truncated"); + match &parts[1] { + ToolContentPart::Text { text } => { + assert_eq!(text, "part2", "short part should not be truncated"); + } + ToolContentPart::Image { .. } => panic!("expected Text part[1], got Image"), } } other @ ToolContent::Text(_) => panic!("expected Multipart, got {other:?}"), From 5fd7df98023e0033b65802537a0db8a888de21cc Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 17:33:58 +1200 Subject: [PATCH 28/30] fix: enforce the limit across the multipart output --- src/middleware.rs | 219 ++++++++++++++++++++++++++++++--- src/middleware/output_limit.rs | 15 ++- 2 files changed, 214 insertions(+), 20 deletions(-) diff --git a/src/middleware.rs b/src/middleware.rs index ca31f82..02299de 100644 --- a/src/middleware.rs +++ b/src/middleware.rs @@ -1099,26 +1099,38 @@ mod tests { let result = pipeline.invoke(test_ctx("multipart")).await; assert!(!result.is_error); - // Each text part in the multipart result should be individually - // truncated. "part1" (5 chars > 3) → "par\n[truncated]". + // The shared character budget (3) is consumed by the first text + // part. "part1" (5 chars > 3) → "par\n[truncated]", remaining=0. + // "part2" gets zero remaining budget → "[truncated]". match result.output { ToolContent::Multipart(parts) => { assert_eq!(parts.len(), 2, "should still have 2 parts"); - for (i, part) in parts.iter().enumerate() { - match part { - ToolContentPart::Text { text } => { - assert!( - text.contains("[truncated]"), - "part {i} should be truncated: got {text:?}" - ); - assert!( - text.starts_with("par"), - "part {i} should start with first 3 chars: got {text:?}" - ); - } - ToolContentPart::Image { .. } => { - panic!("unexpected image part in MultipartTool output"); - } + // Part 0 consumes the entire budget. + match &parts[0] { + ToolContentPart::Text { text } => { + assert!( + text.starts_with("par"), + "part 0 should start with first 3 chars: got {text:?}" + ); + assert!( + text.contains("[truncated]"), + "part 0 should be truncated: got {text:?}" + ); + } + ToolContentPart::Image { .. } => { + panic!("unexpected image part in MultipartTool output"); + } + } + // Part 1 gets zero remaining budget. + match &parts[1] { + ToolContentPart::Text { text } => { + assert_eq!( + text, "[truncated]", + "part 1 should be fully truncated with zero budget: got {text:?}" + ); + } + ToolContentPart::Image { .. } => { + panic!("unexpected image part in MultipartTool output"); } } } @@ -1315,4 +1327,177 @@ mod tests { assert_eq!(result.duration, Duration::from_millis(99)); assert!(!result.is_error); } + + /// A tool that returns 3 text parts, each individually under the limit + /// but collectively exceeding it. + struct ThreePartTextTool; + + impl Tool for ThreePartTextTool { + fn name(&self) -> &'static str { + "three_part_text" + } + + fn description(&self) -> &'static str { + "Returns 3 short text parts" + } + + fn schema(&self) -> ToolSchema { + ToolSchema { + tool: self.name().to_string(), + description: self.description().to_string(), + input_schema: json!({"type": "object", "properties": {}}), + } + } + + fn call( + &self, + _input: Value, + _ctx: &ToolContext, + ) -> Pin> + Send + '_>> { + Box::pin(async move { + Ok(ToolOutput { + payload: ToolContent::Multipart(vec![ + ToolContentPart::Text { + text: "aaaa".to_string(), + }, + ToolContentPart::Text { + text: "bbbb".to_string(), + }, + ToolContentPart::Text { + text: "cccc".to_string(), + }, + ]), + is_error: false, + }) + }) + } + } + + #[tokio::test] + async fn output_limit_multipart_shared_budget_truncates_later_parts() { + // max_chars = 6: part 0 (4) fits, leaving 2. + // Part 1 (4 > 2) → "bb\n[truncated]". Part 2 → 0 remaining. + let mut registry = ToolRegistry::new(); + registry.register(ThreePartTextTool); + let pipeline = ToolPipeline::builder() + .with(OutputLimitMiddleware::new(6)) + .core(Arc::new(registry)) + .build() + .expect("valid"); + + let result = pipeline.invoke(test_ctx("three_part_text")).await; + assert!(!result.is_error); + + match result.output { + ToolContent::Multipart(parts) => { + assert_eq!(parts.len(), 3); + match &parts[0] { + ToolContentPart::Text { text } => { + assert_eq!(text, "aaaa", "part 0 should be unmodified"); + } + ToolContentPart::Image { .. } => panic!("expected Text part 0"), + } + match &parts[1] { + ToolContentPart::Text { text } => { + assert!( + text.starts_with("bb"), + "part 1 should start with 2 chars: got {text:?}" + ); + assert!( + text.contains("[truncated]"), + "part 1 should be truncated: got {text:?}" + ); + } + ToolContentPart::Image { .. } => panic!("expected Text part 1"), + } + match &parts[2] { + ToolContentPart::Text { text } => { + assert_eq!( + text, "[truncated]", + "part 2 should be fully truncated: got {text:?}" + ); + } + ToolContentPart::Image { .. } => panic!("expected Text part 2"), + } + } + ToolContent::Text(t) => panic!("expected Multipart, got Text({t})"), + } + } + + #[tokio::test] + async fn output_limit_multipart_shared_budget_all_fit() { + // max_chars = 20: all three 4-char parts (12 total) fit. + let mut registry = ToolRegistry::new(); + registry.register(ThreePartTextTool); + let pipeline = ToolPipeline::builder() + .with(OutputLimitMiddleware::new(20)) + .core(Arc::new(registry)) + .build() + .expect("valid"); + + let result = pipeline.invoke(test_ctx("three_part_text")).await; + assert!(!result.is_error); + + match result.output { + ToolContent::Multipart(parts) => { + assert_eq!(parts.len(), 3); + for (i, part) in parts.iter().enumerate() { + match part { + ToolContentPart::Text { text } => { + assert!( + !text.contains("[truncated]"), + "part {i} should not be truncated: got {text:?}" + ); + } + ToolContentPart::Image { .. } => panic!("expected Text part {i}"), + } + } + } + ToolContent::Text(t) => panic!("expected Multipart, got Text({t})"), + } + } + + #[tokio::test] + async fn output_limit_multipart_shared_budget_exactly_consumed() { + // max_chars = 8: parts 0+1 (4+4) exactly consume the budget. + // Part 2 → 0 remaining → "[truncated]". + let mut registry = ToolRegistry::new(); + registry.register(ThreePartTextTool); + let pipeline = ToolPipeline::builder() + .with(OutputLimitMiddleware::new(8)) + .core(Arc::new(registry)) + .build() + .expect("valid"); + + let result = pipeline.invoke(test_ctx("three_part_text")).await; + assert!(!result.is_error); + + match result.output { + ToolContent::Multipart(parts) => { + assert_eq!(parts.len(), 3); + match &parts[0] { + ToolContentPart::Text { text } => { + assert_eq!(text, "aaaa", "part 0 unmodified"); + } + ToolContentPart::Image { .. } => panic!("expected Text part 0"), + } + match &parts[1] { + ToolContentPart::Text { text } => { + assert_eq!(text, "bbbb", "part 1 unmodified"); + } + ToolContentPart::Image { .. } => panic!("expected Text part 1"), + } + match &parts[2] { + ToolContentPart::Text { text } => { + assert_eq!( + text, "[truncated]", + "part 2 should be fully truncated: got {text:?}" + ); + } + ToolContentPart::Image { .. } => panic!("expected Text part 2"), + } + } + ToolContent::Text(t) => panic!("expected Multipart, got Text({t})"), + } + } } diff --git a/src/middleware/output_limit.rs b/src/middleware/output_limit.rs index 251aed0..7714011 100644 --- a/src/middleware/output_limit.rs +++ b/src/middleware/output_limit.rs @@ -63,12 +63,21 @@ impl ToolMiddleware for OutputLimitMiddleware { } } ToolContent::Multipart(ref mut parts) => { + let mut remaining = max_chars; for part in parts.iter_mut() { if let ToolContentPart::Text { text } = part { let char_count = text.chars().count(); - if char_count > max_chars { - let truncated: String = text.chars().take(max_chars).collect(); - *text = format!("{truncated}\n[truncated]"); + if char_count > remaining { + if remaining == 0 { + text.clear(); + text.push_str("[truncated]"); + } else { + let truncated: String = text.chars().take(remaining).collect(); + *text = format!("{truncated}\n[truncated]"); + } + remaining = 0; + } else { + remaining = remaining.saturating_sub(char_count); } } } From f0563a71707dd572f4efc9e976a97712eece4f00 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 17:54:41 +1200 Subject: [PATCH 29/30] chore: update readme --- README.md | 35 +++++++++++++---------------------- 1 file changed, 13 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index cafa62b..f22d7e4 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,7 @@ impl Tool for EchoTool { ```rust,no_run use loopctl::engine::BareLoop; +use loopctl::engine::loop_core::Loop; use loopctl::tool::ToolRegistry; use loopctl::config::LoopConfig; use std::sync::Arc; @@ -122,6 +123,7 @@ loopctl = { version = "0.1", features = ["testing"] } ```rust,no_run use loopctl::testing::{MockApiClient, MockTool, test_config}; use loopctl::engine::BareLoop; +use loopctl::engine::loop_core::Loop; use loopctl::tool::ToolRegistry; use std::sync::Arc; @@ -158,28 +160,17 @@ let agent = BareLoop::new( ## Architecture -```text - ┌──────────────┐ - │ ApiClient │ - └───────┬──────┘ - │ - ┌─────────────▼─────────────┐ - │ BareLoop │ - │ ┌─────────────────────┐ │ - │ │ stream → accumulate│ │ - │ │ → tool dispatch │ │ - │ │ → repeat │ │ - │ └─────────────────────┘ │ - └─────┬──────────┬──────────┘ - │ │ - ┌────────────▼──┐ ┌────▼───────────┐ - │ ToolRegistry │ │ Detection & │ - │ (your tools) │ │ Fallback │ - └───────────────┘ │ • convergence │ - │ • loop detect │ - │ • fallback │ - └────────────────┘ -``` +At the center is **BareLoop**, the default agent loop. Each turn it streams a +response from an **ApiClient** (your LLM provider), accumulates the result, and +dispatches any requested tool calls through a **ToolRegistry**. Results are fed +back into the conversation and the cycle repeats until the model ends its turn +or a configured limit is reached. + +Two cross-cutting concerns run alongside the main loop: + +- **Detection & Fallback** — convergence detection, loop detection, and + automatic model/API fallback when requests fail. +- **ToolRegistry** — holds your registered tools and routes tool calls to them. ## Development From 0a0364b6fd492cc52801323295a79049ae6dce4f Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Wed, 1 Jul 2026 18:05:01 +1200 Subject: [PATCH 30/30] fix: misleading docs, snapshot before retrieve memory --- src/api.rs | 14 +++++---- src/memory/builtin.rs | 71 ++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 75 insertions(+), 10 deletions(-) diff --git a/src/api.rs b/src/api.rs index 50b0bf5..b64ad89 100644 --- a/src/api.rs +++ b/src/api.rs @@ -129,9 +129,10 @@ pub trait ApiClient: Send + Sync { /// # Parameters /// /// - `messages` — The conversation history as a [`Vec`]. - /// Takes ownership because the returned stream must be `'static`; - /// callers (e.g. [`BareLoop`](crate::engine::BareLoop)) clone the - /// full history each turn — O(n) in the number of messages. + /// Takes ownership because the request body must be built from owned + /// data for the returned `+ '_` future; callers (e.g. + /// [`BareLoop`](crate::engine::BareLoop)) clone the full history each + /// turn — O(n) in the number of messages. /// - `system` — An optional system prompt to prepend. /// - `tools` — Optional tool definitions the model may invoke. /// @@ -158,9 +159,10 @@ pub trait ApiClient: Send + Sync { /// # Parameters /// /// - `messages` — The conversation history as a [`Vec`]. - /// Takes ownership because the returned stream must be `'static`; - /// callers (e.g. [`BareLoop`](crate::engine::BareLoop)) clone the - /// full history each turn — O(n) in the number of messages. + /// Takes ownership because the request body must be built from owned + /// data for the returned `+ '_` future; callers (e.g. + /// [`BareLoop`](crate::engine::BareLoop)) clone the full history each + /// turn — O(n) in the number of messages. /// - `system` — An optional system prompt to prepend. /// - `tools` — Optional tool definitions the model may invoke. /// diff --git a/src/memory/builtin.rs b/src/memory/builtin.rs index d92ec41..82a0afd 100644 --- a/src/memory/builtin.rs +++ b/src/memory/builtin.rs @@ -76,7 +76,7 @@ use std::sync::{PoisonError, RwLock}; /// [`InMemoryStore`] is `Send + Sync`. Interior mutability is handled via /// an internal `RwLock`, so `store` and `consolidate` only require `&self`. /// This allows the store to be shared via `Arc` or -/// `Arc` across tasks without external locking. +/// `Arc` across tasks without external locking. /// /// # Construction /// @@ -248,8 +248,10 @@ impl LoopMemory for InMemoryStore { let query_words: Vec<&str> = query_lower.split_whitespace().collect(); let entries = self.entries.read().unwrap_or_else(PoisonError::into_inner); - let mut scored: Vec<(f32, MemoryEntry)> = entries - .iter() + let snapshot: Vec = entries.iter().cloned().collect(); + drop(entries); + let mut scored: Vec<(f32, MemoryEntry)> = snapshot + .into_iter() .map(|entry| { let memory_lower = entry.memory.to_lowercase(); let tag_match = entry @@ -270,7 +272,7 @@ impl LoopMemory for InMemoryStore { let tag_bonus = if tag_match { 0.3 } else { 0.0 }; ( base_score * 0.5 + query_bonus * 0.4 + tag_bonus + 0.1, - entry.clone(), + entry, ) }) .collect(); @@ -453,4 +455,65 @@ mod tests { let store = InMemoryStore::default(); assert!(store.is_empty()); } + + #[tokio::test] + async fn test_retrieve_does_not_block_writers() { + // Populate enough entries to make scoring non-trivial. + let store = InMemoryStore::new(); + for i in 0..200 { + store + .store(MemoryEntry::new( + MemoryCategory::Fact, + format!("Fact number {i} about concurrency"), + )) + .await + .unwrap(); + } + + // Start a retrieve future (it will be polled once we await below). + let retrieve_fut = store.retrieve("concurrency", 5); + + // While retrieve is pending, a store should succeed without timing + // out — if the read lock were still held during scoring this would + // deadlock or at least block until retrieve completes. + let store_fut = store.store(MemoryEntry::new( + MemoryCategory::Insight, + "writer proceeds concurrently", + )); + + // Drive both to completion. + let (retrieved, store_res) = tokio::join!(retrieve_fut, store_fut); + let retrieved = retrieved.unwrap(); + store_res.unwrap(); + + assert!(retrieved.len() <= 5); + assert_eq!(store.len(), 201); // 200 originals + 1 concurrent store + } + + #[tokio::test] + async fn test_retrieve_ranking_preserved() { + let store = InMemoryStore::new(); + + let mut high = MemoryEntry::new(MemoryCategory::Insight, "rust rust rust rust"); + high.relevance = 0.95; + + let mut mid = MemoryEntry::new(MemoryCategory::Fact, "rust rust rust"); + mid.relevance = 0.5; + + let mut low = MemoryEntry::new(MemoryCategory::Working, "rust rust"); + low.relevance = 0.1; + + store.store(low.clone()).await.unwrap(); + store.store(high.clone()).await.unwrap(); + store.store(mid.clone()).await.unwrap(); + + let results = store.retrieve("rust", 3).await.unwrap(); + assert_eq!(results.len(), 3); + + // Entries should come back ordered by descending score. The + // highest-relevance entry must be first and the lowest last. + assert!((results[0].relevance - 0.95).abs() < 1e-6); + assert!((results[1].relevance - 0.5).abs() < 1e-6); + assert!((results[2].relevance - 0.1).abs() < 1e-6); + } }